From 4b92dc6fbb73b05fcedc43b2d2a5cd1c0d753f14 Mon Sep 17 00:00:00 2001 From: Ian Harry Date: Mon, 27 Jul 2026 15:43:02 +0100 Subject: [PATCH 1/5] Running automated ruff format fix on all codes --- .../pycbc_plot_Nth_loudest_coinc_omicron.py | 235 +- pycbc/__init__.py | 106 +- pycbc/_version.py | 82 +- pycbc/_version_helper.py | 127 +- pycbc/bin_utils.py | 180 +- pycbc/boundaries.py | 109 +- pycbc/catalog/__init__.py | 122 +- pycbc/catalog/catalog.py | 45 +- pycbc/constants.py | 58 +- pycbc/conversions.py | 994 +++++--- pycbc/coordinates/__init__.py | 39 +- pycbc/coordinates/base.py | 41 +- pycbc/coordinates/space.py | 477 ++-- pycbc/cosmology.py | 166 +- pycbc/detector/ground.py | 460 ++-- pycbc/detector/space.py | 423 ++-- pycbc/distributions/__init__.py | 153 +- pycbc/distributions/angular.py | 250 +- pycbc/distributions/arbitrary.py | 120 +- pycbc/distributions/bounded.py | 142 +- pycbc/distributions/constraints.py | 68 +- pycbc/distributions/external.py | 131 +- pycbc/distributions/fixedsamples.py | 59 +- pycbc/distributions/gaussian.py | 99 +- pycbc/distributions/joint.py | 118 +- pycbc/distributions/mass.py | 133 +- pycbc/distributions/power_law.py | 76 +- pycbc/distributions/qnm.py | 120 +- pycbc/distributions/sky_location.py | 224 +- pycbc/distributions/spins.py | 232 +- pycbc/distributions/uniform.py | 42 +- pycbc/distributions/uniform_log.py | 41 +- pycbc/distributions/utils.py | 48 +- pycbc/dq.py | 342 ++- pycbc/events/__init__.py | 3 +- pycbc/events/coherent.py | 115 +- pycbc/events/coinc.py | 625 +++-- pycbc/events/coinc_rate.py | 43 +- pycbc/events/cuts.py | 246 +- pycbc/events/eventmgr.py | 978 +++---- pycbc/events/ranking.py | 284 +-- pycbc/events/significance.py | 266 +- pycbc/events/single.py | 409 +-- pycbc/events/stat.py | 401 ++- pycbc/events/threshold_cpu.py | 42 +- pycbc/events/threshold_cuda.py | 88 +- pycbc/events/threshold_cupy.py | 45 +- pycbc/events/trigger_fits.py | 64 +- pycbc/events/triggers.py | 129 +- pycbc/events/veto.py | 92 +- pycbc/fft/__init__.py | 6 +- pycbc/fft/backend_cpu.py | 10 +- pycbc/fft/backend_cuda.py | 11 +- pycbc/fft/backend_cupy.py | 7 +- pycbc/fft/backend_mkl.py | 7 +- pycbc/fft/backend_support.py | 23 +- pycbc/fft/class_api.py | 23 +- pycbc/fft/core.py | 128 +- pycbc/fft/cuda_pyfft.py | 40 +- pycbc/fft/cufft.py | 30 +- pycbc/fft/cupyfft.py | 49 +- pycbc/fft/fft_callback.py | 103 +- pycbc/fft/fftw.py | 501 ++-- pycbc/fft/fftw_pruned.py | 136 +- pycbc/fft/func_api.py | 26 +- pycbc/fft/mkl.py | 66 +- pycbc/fft/npfft.py | 62 +- pycbc/fft/parser_support.py | 45 +- pycbc/filter/autocorrelation.py | 28 +- pycbc/filter/fotonfilter.py | 52 +- pycbc/filter/matchedfilter.py | 1130 +++++---- pycbc/filter/matchedfilter_cuda.py | 30 +- pycbc/filter/matchedfilter_cupy.py | 11 +- pycbc/filter/matchedfilter_numpy.py | 1 + pycbc/filter/qtransform.py | 82 +- pycbc/filter/resample.py | 247 +- pycbc/filter/simd_correlate.py | 7 +- pycbc/filter/zpk.py | 27 +- pycbc/frame/__init__.py | 23 +- pycbc/frame/frame.py | 453 ++-- pycbc/frame/gwosc.py | 85 +- pycbc/frame/store.py | 24 +- pycbc/inference/__init__.py | 3 +- pycbc/inference/burn_in.py | 337 +-- pycbc/inference/entropy.py | 94 +- pycbc/inference/evidence.py | 69 +- pycbc/inference/gelman_rubin.py | 41 +- pycbc/inference/geweke.py | 17 +- pycbc/inference/io/__init__.py | 427 ++-- pycbc/inference/io/base_hdf.py | 357 +-- pycbc/inference/io/base_mcmc.py | 403 +-- pycbc/inference/io/base_multitemper.py | 170 +- pycbc/inference/io/base_nested_sampler.py | 18 +- pycbc/inference/io/base_sampler.py | 56 +- pycbc/inference/io/cpnest.py | 6 +- pycbc/inference/io/dynesty.py | 109 +- pycbc/inference/io/emcee.py | 40 +- pycbc/inference/io/emcee_pt.py | 55 +- pycbc/inference/io/epsie.py | 130 +- pycbc/inference/io/multinest.py | 54 +- pycbc/inference/io/nessai.py | 17 +- pycbc/inference/io/posterior.py | 34 +- pycbc/inference/io/ptemcee.py | 90 +- pycbc/inference/io/snowline.py | 6 +- pycbc/inference/io/txt.py | 28 +- pycbc/inference/io/ultranest.py | 6 +- pycbc/inference/jump/__init__.py | 28 +- pycbc/inference/jump/angular.py | 27 +- pycbc/inference/jump/bounded_normal.py | 24 +- pycbc/inference/jump/discrete.py | 35 +- pycbc/inference/jump/normal.py | 195 +- pycbc/inference/models/__init__.py | 147 +- pycbc/inference/models/analytic.py | 102 +- pycbc/inference/models/base.py | 345 ++- pycbc/inference/models/base_data.py | 49 +- pycbc/inference/models/brute_marg.py | 109 +- pycbc/inference/models/data_utils.py | 386 ++- .../inference/models/gated_gaussian_noise.py | 700 ++++-- pycbc/inference/models/gaussian_noise.py | 469 ++-- pycbc/inference/models/hierarchical.py | 451 ++-- .../models/marginalized_gaussian_noise.py | 457 ++-- pycbc/inference/models/relbin.py | 537 ++-- pycbc/inference/models/single_template.py | 125 +- pycbc/inference/models/tools.py | 491 ++-- pycbc/inference/option_utils.py | 363 ++- pycbc/inference/sampler/__init__.py | 41 +- pycbc/inference/sampler/base.py | 94 +- pycbc/inference/sampler/base_cube.py | 19 +- pycbc/inference/sampler/base_mcmc.py | 347 +-- pycbc/inference/sampler/base_multitemper.py | 178 +- pycbc/inference/sampler/cpnest.py | 109 +- pycbc/inference/sampler/dummy.py | 42 +- pycbc/inference/sampler/dynesty.py | 358 +-- pycbc/inference/sampler/emcee.py | 125 +- pycbc/inference/sampler/emcee_pt.py | 215 +- pycbc/inference/sampler/epsie.py | 215 +- pycbc/inference/sampler/games.py | 111 +- pycbc/inference/sampler/multinest.py | 233 +- pycbc/inference/sampler/nessai.py | 64 +- pycbc/inference/sampler/ptemcee.py | 281 ++- pycbc/inference/sampler/refine.py | 39 +- pycbc/inference/sampler/snowline.py | 65 +- pycbc/inference/sampler/ultranest.py | 112 +- pycbc/inject/__init__.py | 2 +- pycbc/inject/inject.py | 546 ++-- pycbc/inject/injfilterrejector.py | 265 +- pycbc/io/__init__.py | 33 +- pycbc/io/gracedb.py | 421 ++-- pycbc/io/hdf.py | 1013 ++++---- pycbc/io/ligolw.py | 177 +- pycbc/io/live.py | 55 +- pycbc/io/record.py | 593 +++-- pycbc/libutils.py | 116 +- pycbc/live/__init__.py | 2 +- pycbc/live/significance_fits.py | 132 +- pycbc/live/snr_optimizer.py | 304 ++- pycbc/live/supervision.py | 69 +- pycbc/mchirp_area.py | 355 +-- pycbc/neutron_stars/__init__.py | 10 +- pycbc/neutron_stars/eos_utils.py | 87 +- pycbc/neutron_stars/pg_isso_solver.py | 201 +- pycbc/noise/__init__.py | 2 +- pycbc/noise/gaussian.py | 58 +- pycbc/noise/reproduceable.py | 132 +- pycbc/opt.py | 50 +- pycbc/pnutils.py | 75 +- pycbc/pool.py | 105 +- pycbc/population/__init__.py | 6 +- pycbc/population/fgmc_functions.py | 443 ++-- pycbc/population/fgmc_laguerre.py | 79 +- pycbc/population/fgmc_plots.py | 224 +- pycbc/population/live_pastro.py | 165 +- pycbc/population/live_pastro_utils.py | 88 +- pycbc/population/population_models.py | 224 +- pycbc/population/rates_functions.py | 585 +++-- pycbc/population/scale_injections.py | 555 ++-- pycbc/psd/__init__.py | 866 ++++--- pycbc/psd/analytical.py | 123 +- pycbc/psd/analytical_space.py | 1088 +++++--- pycbc/psd/estimate.py | 225 +- pycbc/psd/read.py | 81 +- pycbc/psd/variation.py | 179 +- pycbc/rate.py | 147 +- pycbc/results/__init__.py | 14 +- pycbc/results/color.py | 29 +- pycbc/results/dq.py | 68 +- pycbc/results/followup.py | 75 +- pycbc/results/layout.py | 50 +- pycbc/results/metadata.py | 99 +- pycbc/results/mpld3_utils.py | 36 +- pycbc/results/plot.py | 39 +- pycbc/results/psd.py | 22 +- pycbc/results/pygrb_plotting_utils.py | 175 +- pycbc/results/pygrb_postprocessing_utils.py | 554 ++-- pycbc/results/render.py | 130 +- pycbc/results/scatter_histograms.py | 426 ++-- pycbc/results/snr.py | 24 +- pycbc/results/str_utils.py | 148 +- pycbc/results/table_utils.py | 52 +- pycbc/results/versioning.py | 129 +- pycbc/scheme.py | 176 +- pycbc/sensitivity.py | 191 +- pycbc/strain/__init__.py | 42 +- pycbc/strain/calibration.py | 89 +- pycbc/strain/gate.py | 108 +- pycbc/strain/lines.py | 98 +- pycbc/strain/recalibrate.py | 440 ++-- pycbc/strain/strain.py | 2240 ++++++++++------- pycbc/time.py | 49 +- pycbc/tmpltbank/__init__.py | 6 +- pycbc/tmpltbank/bank_conversions.py | 114 +- pycbc/tmpltbank/bank_output_utils.py | 320 +-- pycbc/tmpltbank/brute_force_methods.py | 325 ++- pycbc/tmpltbank/calc_moments.py | 378 +-- pycbc/tmpltbank/coord_utils.py | 348 +-- pycbc/tmpltbank/lambda_mapping.py | 208 +- pycbc/tmpltbank/lattice_utils.py | 68 +- pycbc/tmpltbank/option_utils.py | 1059 +++++--- pycbc/tmpltbank/partitioned_bank.py | 312 ++- pycbc/tmpltbank/sky_grid.py | 70 +- pycbc/transforms.py | 1075 +++++--- pycbc/types/__init__.py | 4 +- pycbc/types/aligned.py | 27 +- pycbc/types/array.py | 562 +++-- pycbc/types/array_cuda.py | 363 ++- pycbc/types/array_cupy.py | 76 +- pycbc/types/config.py | 165 +- pycbc/types/frequencyseries.py | 350 +-- pycbc/types/optparse.py | 306 ++- pycbc/types/timeseries.py | 729 +++--- pycbc/types/utils.py | 24 +- pycbc/vetoes/__init__.py | 4 +- pycbc/vetoes/autochisq.py | 140 +- pycbc/vetoes/bank_chisq.py | 129 +- pycbc/vetoes/chisq.py | 295 ++- pycbc/vetoes/chisq_cuda.py | 84 +- pycbc/vetoes/chisq_cupy.py | 72 +- pycbc/vetoes/sgchisq.py | 97 +- pycbc/waveform/SpinTaylorF2.py | 425 +++- pycbc/waveform/__init__.py | 18 +- pycbc/waveform/bank.py | 638 +++-- pycbc/waveform/compress.py | 510 ++-- pycbc/waveform/decompress_cpu.py | 94 +- pycbc/waveform/decompress_cuda.py | 37 +- pycbc/waveform/decompress_cupy.py | 60 +- pycbc/waveform/generator.py | 684 ++--- pycbc/waveform/multiband.py | 60 +- pycbc/waveform/nltides.py | 57 +- pycbc/waveform/parameters.py | 1217 ++++++--- pycbc/waveform/plugin.py | 124 +- pycbc/waveform/premerger.py | 34 +- pycbc/waveform/pycbc_phenomC_tmplt.py | 445 ++-- pycbc/waveform/ringdown.py | 533 ++-- pycbc/waveform/sinegauss.py | 53 +- pycbc/waveform/spa_tmplt.py | 230 +- pycbc/waveform/spa_tmplt_cuda.py | 59 +- pycbc/waveform/spa_tmplt_cupy.py | 49 +- pycbc/waveform/supernovae.py | 34 +- pycbc/waveform/utils.py | 229 +- pycbc/waveform/utils_cuda.py | 26 +- pycbc/waveform/utils_cupy.py | 19 +- pycbc/waveform/waveform.py | 1200 +++++---- pycbc/waveform/waveform_modes.py | 334 ++- pycbc/workflow/__init__.py | 25 +- pycbc/workflow/coincidence.py | 790 +++--- pycbc/workflow/configparser_test.py | 144 +- pycbc/workflow/configuration.py | 130 +- pycbc/workflow/core.py | 1240 +++++---- pycbc/workflow/datafind.py | 479 ++-- pycbc/workflow/dq.py | 86 +- pycbc/workflow/grb_utils.py | 699 ++--- pycbc/workflow/inference_followups.py | 654 +++-- pycbc/workflow/injection.py | 179 +- pycbc/workflow/jobsetup.py | 907 ++++--- pycbc/workflow/matched_filter.py | 200 +- pycbc/workflow/minifollowups.py | 1014 ++++---- pycbc/workflow/pegasus_sites.py | 274 +- pycbc/workflow/pegasus_workflow.py | 431 ++-- pycbc/workflow/plotting.py | 659 +++-- pycbc/workflow/psd.py | 121 +- pycbc/workflow/psdfiles.py | 47 +- pycbc/workflow/segment.py | 345 +-- pycbc/workflow/splittable.py | 115 +- pycbc/workflow/tmpltbank.py | 232 +- pycbc/workflow/versioning.py | 16 +- 285 files changed, 37073 insertions(+), 25597 deletions(-) diff --git a/bin/plotting/pycbc_plot_Nth_loudest_coinc_omicron.py b/bin/plotting/pycbc_plot_Nth_loudest_coinc_omicron.py index ba4d7c045ce..190814aca3c 100644 --- a/bin/plotting/pycbc_plot_Nth_loudest_coinc_omicron.py +++ b/bin/plotting/pycbc_plot_Nth_loudest_coinc_omicron.py @@ -4,95 +4,137 @@ Omicron triggers. """ -import logging -import numpy as np import argparse import glob +import logging + import matplotlib -matplotlib.use('Agg') -import matplotlib.pyplot as plt +import numpy as np +matplotlib.use("Agg") +import matplotlib.pyplot as plt from igwn_ligolw import lsctables, utils import pycbc.events +from pycbc.io.hdf import HFile +from pycbc.io.ligolw import LIGOLWContentHandler from pycbc.waveform import ( - get_td_waveform, frequency_from_polarizations, - amplitude_from_polarizations + amplitude_from_polarizations, + frequency_from_polarizations, + get_td_waveform, ) -from pycbc.io.ligolw import LIGOLWContentHandler -from pycbc.io.hdf import HFile - parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--coinc-file', type=str, required=True, - help='HDF file containing coincident CBC triggers') -parser.add_argument('--single-ifo-trigs', type=str, required=True, - help='HDF file containing single IFO CBC triggers') -parser.add_argument('--ifo', type=str, required=True, - help='IFO, L1 or H1') -parser.add_argument('--tmpltbank-file', type=str, required=True, - help='HDF file containing template information for CBC search') -parser.add_argument('--output-file', type=str, required=True, - help='Full path to output file') -parser.add_argument('--loudest-event-number', type=int, required=True, default=1, - help='Script will plot the Nth loudest coincident trigger') -parser.add_argument('--omicron-dir', type=str, required=True, - help='Directory containing Omicron triggers. Ex: /home/detchar/triggers/ER7/') -parser.add_argument('--omicron-snr-thresh', type=int, required=False, default=5, - help='SNR threshold for choosing which Omicron triggers to plot.') -parser.add_argument('--plot-window', type=float, required=False, default=32, - help='Time window to plot around CBC trigger') -parser.add_argument('--omicron-channel',type=str, required=False, default='GDS-CALIB_STRAIN', - help='Channel to plot Omicron triggers for, do not include IFO') -parser.add_argument('--analysis-level', type=str, required=False, default='foreground', - choices = ['foreground','background','background_exc'], - help='Designates which level of the analysis output to search') +parser.add_argument( + "--coinc-file", + type=str, + required=True, + help="HDF file containing coincident CBC triggers", +) +parser.add_argument( + "--single-ifo-trigs", + type=str, + required=True, + help="HDF file containing single IFO CBC triggers", +) +parser.add_argument("--ifo", type=str, required=True, help="IFO, L1 or H1") +parser.add_argument( + "--tmpltbank-file", + type=str, + required=True, + help="HDF file containing template information for CBC search", +) +parser.add_argument( + "--output-file", type=str, required=True, help="Full path to output file" +) +parser.add_argument( + "--loudest-event-number", + type=int, + required=True, + default=1, + help="Script will plot the Nth loudest coincident trigger", +) +parser.add_argument( + "--omicron-dir", + type=str, + required=True, + help="Directory containing Omicron triggers. Ex: /home/detchar/triggers/ER7/", +) +parser.add_argument( + "--omicron-snr-thresh", + type=int, + required=False, + default=5, + help="SNR threshold for choosing which Omicron triggers to plot.", +) +parser.add_argument( + "--plot-window", + type=float, + required=False, + default=32, + help="Time window to plot around CBC trigger", +) +parser.add_argument( + "--omicron-channel", + type=str, + required=False, + default="GDS-CALIB_STRAIN", + help="Channel to plot Omicron triggers for, do not include IFO", +) +parser.add_argument( + "--analysis-level", + type=str, + required=False, + default="foreground", + choices=["foreground", "background", "background_exc"], + help="Designates which level of the analysis output to search", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Reading HDF files') +logging.info("Reading HDF files") -coinc_trig_file = HFile(args.coinc_file,'r') -single_trig_file = HFile(args.single_ifo_trigs,'r') -template_file = HFile(args.tmpltbank_file,'r') +coinc_trig_file = HFile(args.coinc_file, "r") +single_trig_file = HFile(args.single_ifo_trigs, "r") +template_file = HFile(args.tmpltbank_file, "r") -logging.info('Parsing HDF files') +logging.info("Parsing HDF files") -coinc_newsnr = coinc_trig_file[args.analysis_level]['stat'][:] +coinc_newsnr = coinc_trig_file[args.analysis_level]["stat"][:] Nth_loudest_idx = np.argsort(coinc_newsnr)[-args.loudest_event_number] -if coinc_trig_file.attrs['detector_1'] == args.ifo: - idx = coinc_trig_file[args.analysis_level]['trigger_id1'][Nth_loudest_idx] +if coinc_trig_file.attrs["detector_1"] == args.ifo: + idx = coinc_trig_file[args.analysis_level]["trigger_id1"][Nth_loudest_idx] else: - idx = coinc_trig_file[args.analysis_level]['trigger_id2'][Nth_loudest_idx] + idx = coinc_trig_file[args.analysis_level]["trigger_id2"][Nth_loudest_idx] # get info about single detector triggers that comprise loudest background event # and calculate newSNR -snr = single_trig_file[args.ifo]['snr'][idx] -chisq = single_trig_file[args.ifo]['chisq'][idx] -chisq_dof = single_trig_file[args.ifo]['chisq_dof'][idx] -reduced_chisq = chisq/(2*chisq_dof - 2) -newsnr = pycbc.events.ranking.newsnr(snr,reduced_chisq) -cbc_end_time = single_trig_file[args.ifo]['end_time'][idx] -template_id = single_trig_file[args.ifo]['template_id'][idx] - -m1 = template_file['mass1'][template_id] -m2 = template_file['mass2'][template_id] -s1z = template_file['spin1z'][template_id] -s2z = template_file['spin2z'][template_id] +snr = single_trig_file[args.ifo]["snr"][idx] +chisq = single_trig_file[args.ifo]["chisq"][idx] +chisq_dof = single_trig_file[args.ifo]["chisq_dof"][idx] +reduced_chisq = chisq / (2 * chisq_dof - 2) +newsnr = pycbc.events.ranking.newsnr(snr, reduced_chisq) +cbc_end_time = single_trig_file[args.ifo]["end_time"][idx] +template_id = single_trig_file[args.ifo]["template_id"][idx] + +m1 = template_file["mass1"][template_id] +m2 = template_file["mass2"][template_id] +s1z = template_file["spin1z"][template_id] +s2z = template_file["spin2z"][template_id] omicron_start_time = cbc_end_time - args.plot_window omicron_end_time = cbc_end_time + args.plot_window -logging.info('Fetching omicron triggers') +logging.info("Fetching omicron triggers") # Generate list of directories to search over gps_era_start = str(omicron_start_time)[:5] gps_era_end = str(omicron_end_time)[:5] -eras = map(str,range(int(gps_era_start),int(gps_era_end))) +eras = map(str, range(int(gps_era_start), int(gps_era_end))) if not eras: eras = [gps_era_start] @@ -103,29 +145,49 @@ for era in eras: # Generate list of all Omicron SnglBurst xml trigger files - file_list = glob.glob(args.omicron_dir + - '/%s/%s_Omicron/%s/%s-%s_Omicron-*.xml.gz' - %(args.ifo,args.omicron_channel,era,args.ifo,args.omicron_channel.replace('-','_'))) + file_list = glob.glob( + args.omicron_dir + + "/%s/%s_Omicron/%s/%s-%s_Omicron-*.xml.gz" + % ( + args.ifo, + args.omicron_channel, + era, + args.ifo, + args.omicron_channel.replace("-", "_"), + ) + ) # Parse trigger files into SNR, time, and frequency for Omicron triggers for file_name in file_list: omicron_xml = utils.load_filename( - file_name, contenthandler=LIGOLWContentHandler) + file_name, contenthandler=LIGOLWContentHandler + ) snglburst_table = lsctables.SnglBurstTable.get_table(omicron_xml) for row in snglburst_table: - if (row.snr > args.omicron_snr_thresh and - omicron_start_time < row.peak_time < omicron_end_time): - omicron_times.append(row.peak_time + row.peak_time_ns * 10**(-9)) + if ( + row.snr > args.omicron_snr_thresh + and omicron_start_time < row.peak_time < omicron_end_time + ): + omicron_times.append(row.peak_time + row.peak_time_ns * 10 ** (-9)) omicron_snr.append(row.snr) omicron_freq.append(row.peak_frequency) # Generate inspiral waveform and calculate f(t) to plot on top of Omicron triggers -hp, hc = get_td_waveform(approximant='SEOBNRv2', mass1=m1, mass2=m2, - spin1x=0, spin1y=0, spin1z=s1z, - spin2x=0, spin2y=0, spin2z=s2z, - delta_t=(1./32768.), f_lower=30) +hp, hc = get_td_waveform( + approximant="SEOBNRv2", + mass1=m1, + mass2=m2, + spin1x=0, + spin1y=0, + spin1z=s1z, + spin2x=0, + spin2y=0, + spin2z=s2z, + delta_t=(1.0 / 32768.0), + f_lower=30, +) f = frequency_from_polarizations(hp, hc) @@ -137,24 +199,35 @@ freq = np.array(f.data) times = np.array(f.sample_times) + cbc_end_time -logging.info('Plotting') +logging.info("Plotting") plt.figure(0) -cm = plt.cm.get_cmap('Reds') -plt.scatter(omicron_times,omicron_freq,c=omicron_snr,s=30,cmap=cm,linewidth=0) -plt.grid(b=True, which='both') +cm = plt.cm.get_cmap("Reds") +plt.scatter(omicron_times, omicron_freq, c=omicron_snr, s=30, cmap=cm, linewidth=0) +plt.grid(b=True, which="both") cbar = plt.colorbar() -cbar.set_label('%s Omicron trigger SNR' % (args.ifo)) -plt.yscale('log') -plt.ylabel('Frequency (Hz)') -plt.xlabel('Time (s)') -plt.xlim(omicron_start_time,omicron_end_time) -plt.suptitle('%s CBC trigger SNR = ' % (args.ifo) + format(snr,'.2f') + - ", newSNR = " + format(newsnr,'.2f'),fontsize=12) -plt.title(format(m1,'.2f') + " - " + format(m2,'.2f') + - " solar masses at GPS time " + format(cbc_end_time,'.2f'),fontsize=12) +cbar.set_label("%s Omicron trigger SNR" % (args.ifo)) +plt.yscale("log") +plt.ylabel("Frequency (Hz)") +plt.xlabel("Time (s)") +plt.xlim(omicron_start_time, omicron_end_time) +plt.suptitle( + "%s CBC trigger SNR = " % (args.ifo) + + format(snr, ".2f") + + ", newSNR = " + + format(newsnr, ".2f"), + fontsize=12, +) +plt.title( + format(m1, ".2f") + + " - " + + format(m2, ".2f") + + " solar masses at GPS time " + + format(cbc_end_time, ".2f"), + fontsize=12, +) plt.hold(True) -plt.plot(times,freq) +plt.plot(times, freq) plt.savefig(args.output_file) -logging.info('Done! Exiting script.') +logging.info("Done! Exiting script.") diff --git a/pycbc/__init__.py b/pycbc/__init__.py index 095f6918d74..8cf304d1c9a 100644 --- a/pycbc/__init__.py +++ b/pycbc/__init__.py @@ -22,29 +22,32 @@ # # ============================================================================= # -"""PyCBC contains a toolkit for CBC gravitational wave analysis -""" -import subprocess, os, sys, signal, warnings +"""PyCBC contains a toolkit for CBC gravitational wave analysis""" + +import os +import signal +import subprocess +import sys +import warnings # Filter annoying Cython warnings that serve no good purpose. warnings.filterwarnings("ignore", message="numpy.dtype size changed") warnings.filterwarnings("ignore", message="numpy.ufunc size changed") +import importlib.machinery +import importlib.util import logging import random import string -import importlib.util -import importlib.machinery from datetime import datetime as dt try: # This will fail when pycbc is imported during the build process, # before version.py has been generated. - from .version import git_hash + from .version import PyCBCVersionAction, git_hash from .version import version as pycbc_version - from .version import PyCBCVersionAction except: - git_hash = 'none' - pycbc_version = 'none' + git_hash = "none" + pycbc_version = "none" PyCBCVersionAction = None __version__ = pycbc_version @@ -58,13 +61,14 @@ class LogFormatter(logging.Formatter): https://en.wikipedia.org/wiki/ISO_8601 e.g. 2022-11-18T09:53:01.554+00:00 """ + converter = dt.fromtimestamp def formatTime(self, record, datefmt=None): ct = self.converter(record.created).astimezone() t = ct.strftime("%Y-%m-%dT%H:%M:%S") s = f"{t}.{int(record.msecs):03d}" - timezone = ct.strftime('%z') + timezone = ct.strftime("%z") timezone_colon = f"{timezone[:-2]}:{timezone[-2:]}" s += timezone_colon return s @@ -78,33 +82,39 @@ def add_common_pycbc_options(parser): ---------- parser : argparse.ArgumentParser The argument parser to which the options will be added + """ group = parser.add_argument_group( title="PyCBC common options", description="Common options for PyCBC executables.", ) group.add_argument( - '-v', - '--verbose', - action='count', + "-v", + "--verbose", + action="count", default=0, help=( - 'Add verbosity to logging. Adding the option ' - 'multiple times makes logging progressively ' - 'more verbose, e.g. --verbose or -v provides ' - 'logging at the info level, but -vv or ' - '--verbose --verbose provides debug logging.' - ) + "Add verbosity to logging. Adding the option " + "multiple times makes logging progressively " + "more verbose, e.g. --verbose or -v provides " + "logging at the info level, but -vv or " + "--verbose --verbose provides debug logging." + ), ) group.add_argument( - '--version', + "--version", action=PyCBCVersionAction, ) -def init_logging(verbose=False, default_level=0, to_file=None, - format='%(asctime)s %(levelname)s : %(message)s'): - """Common utility for setting up logging in PyCBC. +def init_logging( + verbose=False, + default_level=0, + to_file=None, + format="%(asctime)s %(levelname)s : %(message)s", +): + """ + Common utility for setting up logging in PyCBC. Installs a signal handler such that verbosity can be activated at run-time by sending a SIGUSR1 to the process. @@ -124,16 +134,17 @@ def init_logging(verbose=False, default_level=0, to_file=None, overwritten if it already exists. format : str, optional The format to use for logging messages. + """ + def sig_handler(signum, frame): logger = logging.getLogger() log_level = logger.level if log_level == logging.DEBUG: - log_level = logging.WARN + log_level = logging.WARNING else: log_level = logging.DEBUG - logging.warning('Got signal %d, setting log level to %d', - signum, log_level) + logging.warning("Got signal %d, setting log level to %d", signum, log_level) logger.setLevel(log_level) signal.signal(signal.SIGUSR1, sig_handler) @@ -146,11 +157,10 @@ def sig_handler(signum, frame): # Otherwise, you may see duplicate messages logger.handlers.clear() - verbose_int = default_level if verbose is None \ - else int(verbose) + default_level + verbose_int = default_level if verbose is None else int(verbose) + default_level logger.setLevel(logging.WARNING - verbose_int * 10) # Initial setting if to_file is not None: - handler = logging.FileHandler(to_file, mode='w') + handler = logging.FileHandler(to_file, mode="w") else: handler = logging.StreamHandler() logger.addHandler(handler) @@ -177,61 +187,66 @@ def makedir(path): # Dynamic range factor: a large constant for rescaling # GW strains. This is 2**69 rounded to 17 sig.fig. -DYN_RANGE_FAC = 5.9029581035870565e+20 +DYN_RANGE_FAC = 5.9029581035870565e20 # String used to separate parameters in configuration file section headers. # This is used by the distributions and transforms modules -VARARGS_DELIM = '+' +VARARGS_DELIM = "+" # Check for optional CUDA support of the PyCBC Package try: - #check if pycuda is installed + # check if pycuda is installed import pycuda + # If running documentation the import doesn't fail, but it's only a mock # import, so detect that - if type(pycuda).__name__ in ('MagicMock', '_MockModule'): + if type(pycuda).__name__ in ("MagicMock", "_MockModule"): raise ImportError import pycuda.driver as _pycudadrv - #check how many CUDA device is installed + + # check how many CUDA device is installed try: _pycudadrv.init() device_count = _pycudadrv.Device.count() except Exception: device_count = 0 - #Set value to true if there is usable device - HAVE_CUDA = (device_count > 0) + # Set value to true if there is usable device + HAVE_CUDA = device_count > 0 if device_count == 0: - warnings.warn("PyCUDA imported but no CUDA device found; disabling CUDA support") + warnings.warn( + "PyCUDA imported but no CUDA device found; disabling CUDA support" + ) except ImportError: HAVE_CUDA = False # Check for MKL capability try: import pycbc.fft.mkl - HAVE_MKL=True + + HAVE_MKL = True except (ImportError, OSError): - HAVE_MKL=False + HAVE_MKL = False # Check for openmp suppport, currently we pressume it exists, unless on # platforms (mac) that are silly and don't use the standard gcc. -if sys.platform == 'darwin': +if sys.platform == "darwin": HAVE_OMP = False else: HAVE_OMP = True + # https://pynative.com/python-generate-random-string/ def random_string(stringLength=10): - """Generate a random string of fixed length """ + """Generate a random string of fixed length""" letters = string.ascii_lowercase - return ''.join(random.choice(letters) for i in range(stringLength)) + return "".join(random.choice(letters) for i in range(stringLength)) # This is needed as a backwards compatibility. The function was removed in # python 3.12. def load_source(modname, filename): loader = importlib.machinery.SourceFileLoader(modname, filename) - spec = importlib.util.spec_from_file_location(modname, filename, - loader=loader) + spec = importlib.util.spec_from_file_location(modname, filename, loader=loader) module = importlib.util.module_from_spec(spec) # The module is always executed and not cached in sys.modules. # Uncomment the following line to cache the module. @@ -239,11 +254,12 @@ def load_source(modname, filename): loader.exec_module(module) return module + # Expose some convenience functions at package level for backwards # compatibility and convenience: allow `pycbc.gps_now()` as well as # `pycbc.time.gps_now()`. try: - from .time import gps_now # noqa: F401 + from .time import gps_now except Exception: # If pycbc imported during build this may fail; silently ignore. gps_now = None diff --git a/pycbc/_version.py b/pycbc/_version.py index a68d1981160..3aac049cc77 100644 --- a/pycbc/_version.py +++ b/pycbc/_version.py @@ -19,15 +19,15 @@ extremely verbose version information for PyCBC, lal, and lalsimulation. """ -import os -import sys -import glob import argparse +import glob import inspect -import subprocess import logging +import os +import subprocess +import sys -logger = logging.getLogger('pycbc._version') +logger = logging.getLogger("pycbc._version") def print_link(library): @@ -38,17 +38,13 @@ def print_link(library): try: # Linux link = subprocess.check_output( - ['ldd', library], - stderr=subprocess.DEVNULL, - text=True + ["ldd", library], stderr=subprocess.DEVNULL, text=True ) except OSError: try: # macOS link = subprocess.check_output( - ['otool', '-L', library], - stderr=subprocess.DEVNULL, - text=True + ["otool", "-L", library], stderr=subprocess.DEVNULL, text=True ) except: link = err_msg @@ -58,47 +54,45 @@ def print_link(library): def get_lal_info(module, lib_glob): - """Return a string reporting the version and runtime library information + """ + Return a string reporting the version and runtime library information for a LAL Python import. """ module_path = inspect.getfile(module) version_str = ( - module.git_version.verbose_msg + - "\n\nImported from: " + module_path + - "\n\nRuntime libraries:\n" - ) - possible_lib_paths = glob.glob( - os.path.join(os.path.dirname(module_path), lib_glob) + module.git_version.verbose_msg + + "\n\nImported from: " + + module_path + + "\n\nRuntime libraries:\n" ) + possible_lib_paths = glob.glob(os.path.join(os.path.dirname(module_path), lib_glob)) for lib_path in possible_lib_paths: version_str += print_link(lib_path) return version_str class PyCBCVersionAction(argparse._StoreAction): - """Subclass of argparse._StoreAction that prints version information for + """ + Subclass of argparse._StoreAction that prints version information for PyCBC, and for LAL and LALSimulation depending on an integer variable. Can be supplied without the option """ + default_help = ( - 'Display PyCBC version information and exit. ' - 'Can optionally supply a modifier integer to control the ' - 'verbosity of the version information. 0 and 1 are the ' - 'same as --version; 2 provides more detailed PyCBC library ' - 'information; 3 provides information about PyCBC, ' - 'LAL and LALSimulation packages (if installed)' + "Display PyCBC version information and exit. " + "Can optionally supply a modifier integer to control the " + "verbosity of the version information. 0 and 1 are the " + "same as --version; 2 provides more detailed PyCBC library " + "information; 3 provides information about PyCBC, " + "LAL and LALSimulation packages (if installed)" ) - def __init__(self, - option_strings, - dest, - help=default_help, - **kw): + def __init__(self, option_strings, dest, help=default_help, **kw): argparse._StoreAction.__init__( self, option_strings, dest=dest, - nargs='?', + nargs="?", help=help, type=int, **kw, @@ -107,6 +101,7 @@ def __init__(self, def __call__(self, parser, namespace, values, option_string=None): version_no = 0 if values is None else values import pycbc + setattr(namespace, self.dest, version_no) if version_no <= 1: # --version called with zero or default - return the @@ -115,17 +110,20 @@ def __call__(self, parser, namespace, values, option_string=None): if version_no > 1: # --version with flag above 1 - return the verbose version string version_str = ( - "--- PyCBC Version --------------------------\n" + - pycbc.version.git_verbose_msg + "--- PyCBC Version --------------------------\n" + + pycbc.version.git_verbose_msg ) if version_no > 2: # --version called more than twice - print all version information # possible import __main__ + version_str += ( - "\n\nCurrent Executable: " + __main__.__file__ + - "\nImported from: " + inspect.getfile(pycbc) + - "\n\n--- LAL Version ----------------------------\n" + "\n\nCurrent Executable: " + + __main__.__file__ + + "\nImported from: " + + inspect.getfile(pycbc) + + "\n\n--- LAL Version ----------------------------\n" ) try: @@ -133,10 +131,7 @@ def __call__(self, parser, namespace, values, option_string=None): except ImportError: version_str += "\nLAL not installed in environment\n" else: - version_str += get_lal_info( - lal, - '_lal*.so' - ) + version_str += get_lal_info(lal, "_lal*.so") version_str += "\n\n--- LALSimulation Version-------------------\n" try: @@ -144,13 +139,10 @@ def __call__(self, parser, namespace, values, option_string=None): except ImportError: version_str += "\nLALSimulation not installed in environment\n" else: - version_str += get_lal_info( - lalsimulation, - '_lalsimulation*.so' - ) + version_str += get_lal_info(lalsimulation, "_lalsimulation*.so") print(version_str) sys.exit(0) -__all__ = ['PyCBCVersionAction'] +__all__ = ["PyCBCVersionAction"] diff --git a/pycbc/_version_helper.py b/pycbc/_version_helper.py index 27c093af341..7189ee9407c 100644 --- a/pycbc/_version_helper.py +++ b/pycbc/_version_helper.py @@ -18,19 +18,19 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -__author__ = 'Adam Mercer ' +__author__ = "Adam Mercer " -import os -import time -import subprocess -import re import distutils.version import logging +import os +import re +import subprocess +import time -logger = logging.getLogger('pycbc._version_helper') +logger = logging.getLogger("pycbc._version_helper") -class GitInfo(object): +class GitInfo: def __init__(self): self.date = None self.hash = None @@ -47,9 +47,15 @@ class GitInvocationError(LookupError): pass -def call(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - on_error='ignore', returncode=False): - """Run the given command (with shell=False) and return the output as a +def call( + command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + on_error="ignore", + returncode=False, +): + """ + Run the given command (with shell=False) and return the output as a string. Strips the output of enclosing whitespace. @@ -57,103 +63,96 @@ def call(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, If the return code is non-zero, throw GitInvocationError. """ # start external command process - p = subprocess.Popen(command, stdout=subprocess.PIPE, - stderr=subprocess.PIPE) + p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # get outputs out, _ = p.communicate() # throw exception if process failed - if p.returncode != 0 and on_error == 'raise': + if p.returncode != 0 and on_error == "raise": raise GitInvocationError('Failed to run "%s"' % " ".join(command)) - out = out.decode('utf-8').strip() + out = out.decode("utf-8").strip() if returncode: return out, p.returncode return out -def get_build_name(git_path='git'): - """Find the username of the current builder - """ - name, retcode = call(('git', 'config', 'user.name'), returncode=True) +def get_build_name(git_path="git"): + """Find the username of the current builder""" + name, retcode = call(("git", "config", "user.name"), returncode=True) if retcode: name = "Unknown User" - email, retcode = call(('git', 'config', 'user.email'), returncode=True) + email, retcode = call(("git", "config", "user.email"), returncode=True) if retcode: email = "" return f"{name} <{email}>" def get_build_date(): - """Returns the current datetime as the git build date - """ - return time.strftime(r'%Y-%m-%d %H:%M:%S +0000', time.gmtime()) + """Returns the current datetime as the git build date""" + return time.strftime(r"%Y-%m-%d %H:%M:%S +0000", time.gmtime()) -def get_last_commit(git_path='git'): - """Returns the details of the last git commit +def get_last_commit(git_path="git"): + """ + Returns the details of the last git commit Returns a tuple (hash, date, author name, author e-mail, committer name, committer e-mail). """ - hash_, udate, aname, amail, cname, cmail = call(( - git_path, - 'log', - '-1', - r'--pretty=format:%H,%ct,%an,%ae,%cn,%ce' - )).split(",") - date = time.strftime(r'%Y-%m-%d %H:%M:%S +0000', time.gmtime(float(udate))) - author = f'{aname} <{amail}>' - committer = f'{cname} <{cmail}>' + hash_, udate, aname, amail, cname, cmail = call( + (git_path, "log", "-1", r"--pretty=format:%H,%ct,%an,%ae,%cn,%ce") + ).split(",") + date = time.strftime(r"%Y-%m-%d %H:%M:%S +0000", time.gmtime(float(udate))) + author = f"{aname} <{amail}>" + committer = f"{cname} <{cmail}>" return hash_, date, author, committer -def get_git_branch(git_path='git'): - """Returns the name of the current git branch - """ - branch_match = call((git_path, 'rev-parse', '--symbolic-full-name', 'HEAD')) +def get_git_branch(git_path="git"): + """Returns the name of the current git branch""" + branch_match = call((git_path, "rev-parse", "--symbolic-full-name", "HEAD")) if branch_match == "HEAD": return None return os.path.basename(branch_match) -def get_git_tag(hash_, git_path='git'): - """Returns the name of the current git tag - """ - tag, status = call((git_path, 'describe', '--exact-match', - '--tags', hash_), returncode=True) +def get_git_tag(hash_, git_path="git"): + """Returns the name of the current git tag""" + tag, status = call( + (git_path, "describe", "--exact-match", "--tags", hash_), returncode=True + ) if status == 0: return tag return None def get_num_commits(): - return call(('git', 'rev-list', '--count', 'HEAD')) + return call(("git", "rev-list", "--count", "HEAD")) -def get_git_status(git_path='git'): - """Returns the state of the git working copy - """ - status_output = subprocess.call((git_path, 'diff-files', '--quiet')) +def get_git_status(git_path="git"): + """Returns the state of the git working copy""" + status_output = subprocess.call((git_path, "diff-files", "--quiet")) if status_output != 0: - return 'UNCLEAN: Modified working tree' + return "UNCLEAN: Modified working tree" # check index for changes - status_output = subprocess.call((git_path, 'diff-index', '--cached', - '--quiet', 'HEAD')) + status_output = subprocess.call( + (git_path, "diff-index", "--cached", "--quiet", "HEAD") + ) if status_output != 0: - return 'UNCLEAN: Modified index' - return 'CLEAN: All modifications committed' + return "UNCLEAN: Modified index" + return "CLEAN: All modifications committed" def determine_latest_release_version(): - """Query the git repository for the last released version of the code. - """ - git_path = call(('which', 'git')) + """Query the git repository for the last released version of the code.""" + git_path = call(("which", "git")) # Get all tags - tag_list = call((git_path, 'tag')).split('\n') + tag_list = call((git_path, "tag")).split("\n") # Reduce to only versions re_magic = re.compile(r"v\d+\.\d+\.\d+$") @@ -172,18 +171,16 @@ def determine_latest_release_version(): def generate_git_version_info(): - """Query the git repository information to generate a version module. - """ + """Query the git repository information to generate a version module.""" info = GitInfo() - git_path = call(('which', 'git')) + git_path = call(("which", "git")) # get build info info.builder = get_build_name() info.build_date = get_build_date() # parse git ID - info.hash, info.date, info.author, info.committer = ( - get_last_commit(git_path)) + info.hash, info.date, info.author, info.committer = get_last_commit(git_path) # determine branch info.branch = get_git_branch(git_path) @@ -193,17 +190,17 @@ def generate_git_version_info(): # determine version if info.tag: - info.version = info.tag.strip('v') - info.release = not re.search('[a-z]', info.version.lower()) + info.version = info.tag.strip("v") + info.release = not re.search("[a-z]", info.version.lower()) else: - info.version = '0.0a' + get_num_commits() + info.version = "0.0a" + get_num_commits() info.release = False # Determine *last* stable release info.last_release = determine_latest_release_version() # refresh index - call((git_path, 'update-index', '-q', '--refresh')) + call((git_path, "update-index", "-q", "--refresh")) # check working copy for changes info.status = get_git_status(git_path) diff --git a/pycbc/bin_utils.py b/pycbc/bin_utils.py index 6727d57d03b..2e715357c90 100644 --- a/pycbc/bin_utils.py +++ b/pycbc/bin_utils.py @@ -1,25 +1,27 @@ from bisect import bisect_right + try: - from fpconst import PosInf, NegInf + from fpconst import NegInf, PosInf except ImportError: # fpconst is not part of the standard library and might not be available PosInf = float("+inf") NegInf = float("-inf") -import numpy -import math import logging +import math -logger = logging.getLogger('pycbc.bin_utils') +import numpy +logger = logging.getLogger("pycbc.bin_utils") -class Bins(object): +class Bins: """ Parent class for 1-dimensional binnings. Not intended to be used directly, but to be subclassed for use in real bins classes. """ + def __init__(self, minv, maxv, n): """ Initialize a Bins instance. The three arguments are the @@ -58,9 +60,10 @@ def __getitem__(self, x): if isinstance(x, slice): if x.step is not None: raise NotImplementedError("step not supported: %s" % repr(x)) - return slice(self[x.start] if x.start is not None - else 0, self[x.stop] + 1 if x.stop is not None - else len(self)) + return slice( + self[x.start] if x.start is not None else 0, + self[x.stop] + 1 if x.stop is not None else len(self), + ) raise NotImplementedError def __iter__(self): @@ -94,7 +97,6 @@ def upper(self): class IrregularBins(Bins): - """ Bins with arbitrary, irregular spacing. We only require strict monotonicity of the bin boundaries. N boundaries define N-1 bins. @@ -122,6 +124,7 @@ class IrregularBins(Bins): >>> x == y True """ + def __init__(self, boundaries): """ Initialize a set of custom bins with the bin boundaries. @@ -143,7 +146,7 @@ def __init__(self, boundaries): def __getitem__(self, x): if isinstance(x, slice): - return super(IrregularBins, self).__getitem__(x) + return super().__getitem__(x) if self.minv <= x < self.maxv: return bisect_right(self.boundaries, x) - 1 # special measure-zero edge case @@ -162,7 +165,6 @@ def centres(self): class LinearBins(Bins): - """ Linearly-spaced bins. There are n bins of equal size, the first bin starts on the lower bound and the last bin ends on the upper @@ -200,13 +202,14 @@ class LinearBins(Bins): >>> x[10:] slice(1, 3, None) """ + def __init__(self, minv, maxv, n): - super(LinearBins, self).__init__(minv, maxv, n) + super().__init__(minv, maxv, n) self.delta = float(maxv - minv) / n def __getitem__(self, x): if isinstance(x, slice): - return super(LinearBins, self).__getitem__(x) + return super().__getitem__(x) if self.minv <= x < self.maxv: return int(math.floor((x - self.minv) / self.delta)) if x == self.maxv: @@ -218,15 +221,15 @@ def lower(self): return numpy.linspace(self.minv, self.maxv - self.delta, len(self)) def centres(self): - return numpy.linspace(self.minv + self.delta / 2., - self.maxv - self.delta / 2., len(self)) + return numpy.linspace( + self.minv + self.delta / 2.0, self.maxv - self.delta / 2.0, len(self) + ) def upper(self): return numpy.linspace(self.minv + self.delta, self.maxv, len(self)) class LinearPlusOverflowBins(Bins): - """ Linearly-spaced bins with overflow at the edges. @@ -265,15 +268,16 @@ class LinearPlusOverflowBins(Bins): >>> x[9:float("+inf")] slice(2, 5, None) """ + def __init__(self, minv, maxv, n): if n < 3: raise ValueError("n must be >= 3") - super(LinearPlusOverflowBins, self).__init__(minv, maxv, n) + super().__init__(minv, maxv, n) self.delta = float(maxv - minv) / (n - 2) def __getitem__(self, x): if isinstance(x, slice): - return super(LinearPlusOverflowBins, self).__getitem__(x) + return super().__getitem__(x) if self.minv <= x < self.maxv: return int(math.floor((x - self.minv) / self.delta)) + 1 if x >= self.maxv: @@ -286,28 +290,33 @@ def __getitem__(self, x): def lower(self): return numpy.concatenate( - (numpy.array([NegInf]), - self.minv + self.delta * numpy.arange(len(self) - 2), - numpy.array([self.maxv])) + ( + numpy.array([NegInf]), + self.minv + self.delta * numpy.arange(len(self) - 2), + numpy.array([self.maxv]), + ) ) def centres(self): return numpy.concatenate( - (numpy.array([NegInf]), - self.minv + self.delta * (numpy.arange(len(self) - 2) + 0.5), - numpy.array([PosInf])) + ( + numpy.array([NegInf]), + self.minv + self.delta * (numpy.arange(len(self) - 2) + 0.5), + numpy.array([PosInf]), + ) ) def upper(self): return numpy.concatenate( - (numpy.array([self.minv]), - self.minv + self.delta * (numpy.arange(len(self) - 2) + 1), - numpy.array([PosInf])) + ( + numpy.array([self.minv]), + self.minv + self.delta * (numpy.arange(len(self) - 2) + 1), + numpy.array([PosInf]), + ) ) class LogarithmicBins(Bins): - """ Logarithmically-spaced bins. @@ -325,16 +334,16 @@ class LogarithmicBins(Bins): >>> x[25] 2 """ + def __init__(self, minv, maxv, n): - super(LogarithmicBins, self).__init__(minv, maxv, n) + super().__init__(minv, maxv, n) self.delta = (math.log(maxv) - math.log(minv)) / n def __getitem__(self, x): if isinstance(x, slice): - return super(LogarithmicBins, self).__getitem__(x) + return super().__getitem__(x) if self.minv <= x < self.maxv: - return int(math.floor((math.log(x) - math.log(self.minv)) / - self.delta)) + return int(math.floor((math.log(x) - math.log(self.minv)) / self.delta)) if x == self.maxv: # special "measure zero" corner case return len(self) - 1 @@ -342,25 +351,28 @@ def __getitem__(self, x): def lower(self): return numpy.exp( - numpy.linspace(math.log(self.minv), math.log(self.maxv) - - self.delta, len(self)) + numpy.linspace( + math.log(self.minv), math.log(self.maxv) - self.delta, len(self) + ) ) def centres(self): return numpy.exp( - numpy.linspace(math.log(self.minv), math.log(self.maxv) - - self.delta, len(self)) + self.delta / 2. + numpy.linspace( + math.log(self.minv), math.log(self.maxv) - self.delta, len(self) + ) + + self.delta / 2.0 ) def upper(self): return numpy.exp( - numpy.linspace(math.log(self.minv) + self.delta, - math.log(self.maxv), len(self)) + numpy.linspace( + math.log(self.minv) + self.delta, math.log(self.maxv), len(self) + ) ) class LogarithmicPlusOverflowBins(Bins): - """ Logarithmically-spaced bins plus one bin at each end that goes to zero and positive infinity respectively. There are n-2 bins each @@ -392,18 +404,18 @@ class LogarithmicPlusOverflowBins(Bins): >>> x.centres() array([ 0. , 1.70997595, 5. , 14.62008869, inf]) """ + def __init__(self, minv, maxv, n): if n < 3: raise ValueError("n must be >= 3") - super(LogarithmicPlusOverflowBins, self).__init__(minv, maxv, n) + super().__init__(minv, maxv, n) self.delta = (math.log(maxv) - math.log(minv)) / (n - 2) def __getitem__(self, x): if isinstance(x, slice): - return super(LogarithmicPlusOverflowBins, self).__getitem__(x) + return super().__getitem__(x) if self.minv <= x < self.maxv: - return 1 + int(math.floor((math.log(x) - math.log(self.minv)) / - self.delta)) + return 1 + int(math.floor((math.log(x) - math.log(self.minv)) / self.delta)) if x >= self.maxv: # infinity overflow bin return len(self) - 1 @@ -414,32 +426,46 @@ def __getitem__(self, x): def lower(self): return numpy.concatenate( - (numpy.array([0.]), - numpy.exp(numpy.linspace(math.log(self.minv), math.log(self.maxv), - len(self) - 1)) - ) + ( + numpy.array([0.0]), + numpy.exp( + numpy.linspace( + math.log(self.minv), math.log(self.maxv), len(self) - 1 + ) + ), + ) ) def centres(self): return numpy.concatenate( - (numpy.array([0.]), - numpy.exp(numpy.linspace(math.log(self.minv), math.log(self.maxv) - - self.delta, len(self) - 2) + self.delta / 2.), - numpy.array([PosInf]) - ) + ( + numpy.array([0.0]), + numpy.exp( + numpy.linspace( + math.log(self.minv), + math.log(self.maxv) - self.delta, + len(self) - 2, + ) + + self.delta / 2.0 + ), + numpy.array([PosInf]), + ) ) def upper(self): return numpy.concatenate( - (numpy.exp(numpy.linspace(math.log(self.minv), math.log(self.maxv), - len(self) - 1)), - numpy.array([PosInf]) - ) + ( + numpy.exp( + numpy.linspace( + math.log(self.minv), math.log(self.maxv), len(self) - 1 + ) + ), + numpy.array([PosInf]), + ) ) class NDBins(tuple): - """ Multi-dimensional co-ordinate binning. An instance of this object is used to convert a tuple of co-ordinates into a tuple of bin @@ -471,6 +497,7 @@ class NDBins(tuple): Note that the co-ordinates to be converted must be a tuple, even if it is only a 1-dimensional co-ordinate. """ + def __new__(cls, *args): new = tuple.__new__(cls, *args) new.minv = tuple(b.minv for b in new) @@ -504,8 +531,7 @@ def __getitem__(self, coords): if len(coords) != len(self): raise ValueError("dimension mismatch") return tuple(map(lambda b, c: b[c], self, coords)) - else: - return tuple.__getitem__(self, coords) + return tuple.__getitem__(self, coords) def lower(self): """ @@ -532,16 +558,15 @@ def upper(self): return tuple(b.upper() for b in self) -class BinnedArray(object): - +class BinnedArray: """ A convenience wrapper, using the NDBins class to provide access to the elements of an array object. Technical reasons preclude providing a subclass of the array object, so the array data is made available as the "array" attribute of this class. - Examples: - + Examples + -------- Note that even for 1 dimensional arrays the index must be a tuple. >>> x = BinnedArray(NDBins((LinearBins(0, 10, 5),))) @@ -580,15 +605,16 @@ class BinnedArray(object): (0.0, 0.0) >>> x.argmax() (1.0, 1.0) + """ + def __init__(self, bins, array=None, dtype="double"): self.bins = bins if array is None: self.array = numpy.zeros(bins.shape, dtype=dtype) else: if array.shape != bins.shape: - raise ValueError("input array and input bins must have the " - "same shape") + raise ValueError("input array and input bins must have the same shape") self.array = array def __getitem__(self, coords): @@ -620,9 +646,13 @@ def argmin(self): minimum value. Same as numpy.argmin(), converting the indexes to bin co-ordinates. """ - return tuple(centres[index] for centres, index in - zip(self.centres(), numpy.unravel_index(self.array.argmin(), - self.array.shape))) + return tuple( + centres[index] + for centres, index in zip( + self.centres(), + numpy.unravel_index(self.array.argmin(), self.array.shape), + ) + ) def argmax(self): """ @@ -630,9 +660,13 @@ def argmax(self): maximum value. Same as numpy.argmax(), converting the indexes to bin co-ordinates. """ - return tuple(centres[index] for centres, index in - zip(self.centres(), numpy.unravel_index(self.array.argmax(), - self.array.shape))) + return tuple( + centres[index] + for centres, index in zip( + self.centres(), + numpy.unravel_index(self.array.argmax(), self.array.shape), + ) + ) def logregularize(self, epsilon=2**-1074): """ @@ -644,8 +678,7 @@ def logregularize(self, epsilon=2**-1074): return self -class BinnedRatios(object): - +class BinnedRatios: """ Like BinnedArray, but provides a numerator array and a denominator array. The incnumerator() method increments a bin in the numerator @@ -655,6 +688,7 @@ class BinnedRatios(object): accessible as the numerator and denominator attributes, which are both BinnedArray objects. """ + def __init__(self, bins, dtype="double"): self.numerator = BinnedArray(bins, dtype=dtype) self.denominator = BinnedArray(bins, dtype=dtype) diff --git a/pycbc/boundaries.py b/pycbc/boundaries.py index 2a22c8c9f9a..f46d1ef7bc3 100644 --- a/pycbc/boundaries.py +++ b/pycbc/boundaries.py @@ -27,10 +27,11 @@ cyclic boundaries or reflected boundaries. """ -import numpy import logging -logger = logging.getLogger('pycbc.boundaries') +import numpy + +logger = logging.getLogger("pycbc.boundaries") class _Bound(float): @@ -39,14 +40,16 @@ class _Bound(float): name = None def larger(self, other): - """A function to determine whether or not `other` is larger + """ + A function to determine whether or not `other` is larger than the bound. This raises a NotImplementedError; classes that inherit from this must define it. """ raise NotImplementedError("larger function not set") def smaller(self, other): - """A function to determine whether or not `other` is smaller + """ + A function to determine whether or not `other` is smaller than the bound. This raises a NotImplementedError; classes that inherit from this must define it. """ @@ -56,7 +59,7 @@ def smaller(self, other): class OpenBound(_Bound): """Sets larger and smaller functions to be `>` and `<`, respectively.""" - name = 'open' + name = "open" def larger(self, other): """Returns True if `other` is `>`, False otherwise""" @@ -70,7 +73,7 @@ def smaller(self, other): class ClosedBound(_Bound): """Sets larger and smaller functions to be `>=` and `<=`, respectively.""" - name = 'closed' + name = "closed" def larger(self, other): return self >= other @@ -82,10 +85,10 @@ def smaller(self, other): class ReflectedBound(ClosedBound): """Inherits from `ClosedBound`, adding reflection functions.""" - name = 'reflected' + name = "reflected" def reflect(self, value): - return 2*self - value + return 2 * self - value def reflect_left(self, value): """Only reflects the value if is > self.""" @@ -103,7 +106,7 @@ def reflect_right(self, value): boundary_types = { OpenBound.name: OpenBound, ClosedBound.name: ClosedBound, - ReflectedBound.name: ReflectedBound + ReflectedBound.name: ReflectedBound, } @@ -111,8 +114,10 @@ def reflect_right(self, value): # Helper functions for applying conditions to boundaries # + def apply_cyclic(value, bounds): - """Given a value, applies cyclic boundary conditions between the minimum + """ + Given a value, applies cyclic boundary conditions between the minimum and maximum bounds. Parameters @@ -126,11 +131,14 @@ def apply_cyclic(value, bounds): ------- float The value after the cyclic bounds are applied. + """ - return (value - bounds._min) %(bounds._max - bounds._min) + bounds._min + return (value - bounds._min) % (bounds._max - bounds._min) + bounds._min + def reflect_well(value, bounds): - """Given some boundaries, reflects the value until it falls within both + """ + Given some boundaries, reflects the value until it falls within both boundaries. This is done iteratively, reflecting left off of the `boundaries.max`, then right off of the `boundaries.min`, etc. @@ -147,6 +155,7 @@ def reflect_well(value, bounds): ------- float The value after being reflected between the two bounds. + """ while value not in bounds: value = bounds._max.reflect_left(value) @@ -163,8 +172,10 @@ def _pass(value): # Bounds class # -class Bounds(object): - """Creates and stores bounds using the given values. + +class Bounds: + """ + Creates and stores bounds using the given values. The type of boundaries used can be set using the `btype_(min|max)` parameters. These arguments set what kind of boundary is used at the @@ -280,71 +291,79 @@ class Bounds(object): >>> ax.vlines([-1., 1.], x.min(), x.max(), color='k', linestyle='--') >>> ax.set_title('reflected betewen x=-1,1') >>> fig.show() + """ - def __init__(self, min_bound=-numpy.inf, max_bound=numpy.inf, - btype_min='closed', btype_max='open', cyclic=False): + def __init__( + self, + min_bound=-numpy.inf, + max_bound=numpy.inf, + btype_min="closed", + btype_max="open", + cyclic=False, + ): # check boundary values if min_bound >= max_bound: raise ValueError("min_bound must be < max_bound") - if cyclic and not ( - numpy.isfinite(min_bound) and numpy.isfinite(max_bound)): - raise ValueError("if using cyclic, min and max bounds must both " - "be finite") + if cyclic and not (numpy.isfinite(min_bound) and numpy.isfinite(max_bound)): + raise ValueError("if using cyclic, min and max bounds must both be finite") # store bounds try: self._min = boundary_types[btype_min](min_bound) except KeyError: - raise ValueError("unrecognized btype_min {}".format(btype_min)) + raise ValueError(f"unrecognized btype_min {btype_min}") try: self._max = boundary_types[btype_max](max_bound) except KeyError: - raise ValueError("unrecognized btype_max {}".format(btype_max)) + raise ValueError(f"unrecognized btype_max {btype_max}") # store cyclic conditions self._cyclic = bool(cyclic) # store reflection conditions; we'll vectorize them here so that they # can be used with arrays - if self._min.name == 'reflected' and self._max.name == 'reflected': + if self._min.name == "reflected" and self._max.name == "reflected": self._reflect = numpy.vectorize(self._reflect_well) - self.reflected = 'well' - elif self._min.name == 'reflected': + self.reflected = "well" + elif self._min.name == "reflected": self._reflect = numpy.vectorize(self._min.reflect_right) - self.reflected = 'min' - elif self._max.name == 'reflected': + self.reflected = "min" + elif self._max.name == "reflected": self._reflect = numpy.vectorize(self._max.reflect_left) - self.reflected = 'max' + self.reflected = "max" else: self._reflect = _pass self.reflected = False def __repr__(self): - return str(self.__class__)[:-1] + " " + " ".join( - map(str, ["min", self._min, "max", self._max, - "cyclic", self._cyclic])) + ">" + return ( + str(self.__class__)[:-1] + + " " + + " ".join( + map(str, ["min", self._min, "max", self._max, "cyclic", self._cyclic]) + ) + + ">" + ) @property def min(self): - """_bounds instance: The minimum bound """ + """_bounds instance: The minimum bound""" return self._min @property def max(self): - """_bounds instance: The maximum bound """ + """_bounds instance: The maximum bound""" return self._max @property def cyclic(self): - """bool: Whether the bounds are cyclic or not. - """ + """bool: Whether the bounds are cyclic or not.""" return self._cyclic def __getitem__(self, ii): if ii == 0: return self._min - elif ii == 1: + if ii == 1: return self._max - else: - raise IndexError("index {} out of range".format(ii)) + raise IndexError(f"index {ii} out of range") def __abs__(self): return abs(self._max - self._min) @@ -353,17 +372,16 @@ def __contains__(self, value): return self._min.smaller(value) & self._max.larger(value) def _reflect_well(self, value): - """Thin wrapper around `reflect_well` that passes self as the `bounds`. - """ + """Thin wrapper around `reflect_well` that passes self as the `bounds`.""" return reflect_well(value, self) def _apply_cyclic(self, value): - """Thin wrapper around `apply_cyclic` that passes self as the `bounds`. - """ + """Thin wrapper around `apply_cyclic` that passes self as the `bounds`.""" return apply_cyclic(value, self) def apply_conditions(self, value): - """Applies any boundary conditions to the given value. + """ + Applies any boundary conditions to the given value. The value is manipulated according based on the following conditions: @@ -397,6 +415,7 @@ def apply_conditions(self, value): ------- float The value after the conditions are applied; see above for details. + """ retval = value if self._cyclic: @@ -410,7 +429,8 @@ def apply_conditions(self, value): return retval def contains_conditioned(self, value): - """Runs `apply_conditions` on the given value before testing whether it + """ + Runs `apply_conditions` on the given value before testing whether it is in bounds. Note that if `cyclic` is True, or both bounds are reflected, than this will always return True. @@ -424,5 +444,6 @@ def contains_conditioned(self, value): bool Whether or not the value is within the bounds after the boundary conditions are applied. + """ return self.apply_conditions(value) in self diff --git a/pycbc/catalog/__init__.py b/pycbc/catalog/__init__.py index 185bce02707..385190b5dfa 100644 --- a/pycbc/catalog/__init__.py +++ b/pycbc/catalog/__init__.py @@ -22,25 +22,30 @@ # # ============================================================================= # -""" This package provides information about LIGO/Virgo detections of +""" +This package provides information about LIGO/Virgo detections of compact binary mergers """ + import logging + import numpy -logger = logging.getLogger('pycbc.catalog') +logger = logging.getLogger("pycbc.catalog") _aliases = {} -_aliases['mchirp'] = 'chirp_mass_source' -_aliases['mass1'] = 'mass_1_source' -_aliases['mass2'] = 'mass_2_source' -_aliases['snr'] = 'network_matched_filter_snr' -_aliases['z'] = _aliases['redshift'] = 'redshift' -_aliases['distance'] = 'luminosity_distance' +_aliases["mchirp"] = "chirp_mass_source" +_aliases["mass1"] = "mass_1_source" +_aliases["mass2"] = "mass_2_source" +_aliases["snr"] = "network_matched_filter_snr" +_aliases["z"] = _aliases["redshift"] = "redshift" +_aliases["distance"] = "luminosity_distance" + def find_event_in_catalog(name, source=None): - """ Get event data from a catalog - + """ + Get event data from a catalog + Parameters ---------- name: str @@ -48,8 +53,9 @@ def find_event_in_catalog(name, source=None): source: str, None If not provided, look through each catalog until event is found. If provided, only check that specific catalog. + """ - from .catalog import list_catalogs, get_source + from .catalog import get_source, list_catalogs if source is None: catalogs = list_catalogs() @@ -63,39 +69,43 @@ def find_event_in_catalog(name, source=None): # Try common name data = get_source(catalog) for mname in data: - cname = data[mname]['commonName'] + cname = data[mname]["commonName"] if cname.upper() == name.upper(): name = mname data = data[name] return data - else: - raise ValueError(f'Did not find merger matching name: {name}') + raise ValueError(f"Did not find merger matching name: {name}") + -class Merger(object): +class Merger: """Informaton about a specific compact binary merger""" + def __init__(self, name, source=None): - """ Return the information of a merger + """ + Return the information of a merger Parameters ---------- name: str The name (GW prefixed date) of the merger event. + """ self.data = find_event_in_catalog(name, source=source) # Set some basic params from the dataset for key in self.data: - setattr(self, '_raw_' + key, self.data[key]) + setattr(self, "_raw_" + key, self.data[key]) for key in _aliases: setattr(self, key, self.data[_aliases[key]]) - self.common_name = self.data['commonName'] - self.time = self.data['GPS'] - self.frame = 'source' + self.common_name = self.data["commonName"] + self.time = self.data["GPS"] + self.frame = "source" def median1d(self, name, return_errors=False): - """ Return median 1d marginalized parameters + """ + Return median 1d marginalized parameters Parameters ---------- @@ -109,6 +119,7 @@ def median1d(self, name, return_errors=False): ------- param: float or tuple The requested parameter + """ if name in _aliases: name = _aliases[name] @@ -116,17 +127,17 @@ def median1d(self, name, return_errors=False): try: if return_errors: mid = self.data[name] - high = self.data[name + '_upper'] - low = self.data[name + '_lower'] + high = self.data[name + "_upper"] + low = self.data[name + "_lower"] return (mid, low, high) - else: - return self.data[name] + return self.data[name] except KeyError as e: print(e) - raise RuntimeError("Cannot get parameter {}".format(name)) + raise RuntimeError(f"Cannot get parameter {name}") def strain(self, ifo, duration=32, sample_rate=4096): - """ Return strain around the event + """ + Return strain around the event Currently this will return the strain around the event in the smallest format available. Selection of other data is not yet available. @@ -140,44 +151,47 @@ def strain(self, ifo, duration=32, sample_rate=4096): ------- strain: pycbc.types.TimeSeries Strain around the event. + """ - from pycbc.io import get_file from pycbc.frame import read_frame + from pycbc.io import get_file - for fdict in self.data['strain']: - if (fdict['detector'] == ifo and fdict['duration'] == duration and - fdict['sampling_rate'] == sample_rate and - fdict['format'] == 'gwf'): - url = fdict['url'] + for fdict in self.data["strain"]: + if ( + fdict["detector"] == ifo + and fdict["duration"] == duration + and fdict["sampling_rate"] == sample_rate + and fdict["format"] == "gwf" + ): + url = fdict["url"] break else: - raise ValueError('no strain data is available as requested ' - 'for ' + self.common_name) + raise ValueError( + "no strain data is available as requested for " + self.common_name + ) - ver = url.split('/')[-1].split('-')[1].split('_')[-1] - sampling_map = {4096: "4KHZ", - 16384: "16KHZ"} - channel = "{}:GWOSC-{}_{}_STRAIN".format( - ifo, sampling_map[sample_rate], ver) + ver = url.split("/")[-1].split("-")[1].split("_")[-1] + sampling_map = {4096: "4KHZ", 16384: "16KHZ"} + channel = f"{ifo}:GWOSC-{sampling_map[sample_rate]}_{ver}_STRAIN" filename = get_file(url, cache=True) return read_frame(filename, str(channel)) -class Catalog(object): +class Catalog: """Manage a set of binary mergers""" - def __init__(self, source='gwtc-3'): - """ Return the set of detected mergers + def __init__(self, source="gwtc-3"): + """ + Return the set of detected mergers The set of detected mergers. At some point this may have some selection abilities. """ from . import catalog - + self.data = catalog.get_source(source=source) - self.mergers = {name: Merger(name, - source=source) for name in self.data} + self.mergers = {name: Merger(name, source=source) for name in self.data} self.names = self.mergers.keys() def __len__(self): @@ -192,8 +206,7 @@ def __getitem__(self, key): if key == self.mergers[m].common_name: break else: - raise ValueError('Did not find merger matching' - ' name: {}'.format(key)) + raise ValueError(f"Did not find merger matching name: {key}") return self.mergers[m] def __setitem__(self, key, value): @@ -206,7 +219,8 @@ def __iter__(self): return iter(self.mergers) def median1d(self, param, return_errors=False): - """ Return median 1d marginalized parameters + """ + Return median 1d marginalized parameters Parameters ---------- @@ -220,11 +234,13 @@ def median1d(self, param, return_errors=False): ------- param: nump.ndarray or tuple The requested parameter + """ - v = [self.mergers[m].median1d(param, return_errors=return_errors) - for m in self.mergers] + v = [ + self.mergers[m].median1d(param, return_errors=return_errors) + for m in self.mergers + ] if return_errors: value, merror, perror = zip(*v) return numpy.array(value), numpy.array(merror), numpy.array(perror) - else: - return numpy.array(v) + return numpy.array(v) diff --git a/pycbc/catalog/catalog.py b/pycbc/catalog/catalog.py index 1e34b31ff7d..f580b2a46ec 100644 --- a/pycbc/catalog/catalog.py +++ b/pycbc/catalog/catalog.py @@ -22,15 +22,17 @@ # # ============================================================================= # -""" This modules contains information about the announced LIGO/Virgo +""" +This modules contains information about the announced LIGO/Virgo compact binary mergers """ -import logging + import json +import logging from pycbc.io import get_file -logger = logging.getLogger('pycbc.catalog.catalog') +logger = logging.getLogger("pycbc.catalog.catalog") # For the time being all quantities are the 1-d median value # FIXME with posteriors when available and we can just post-process that @@ -38,37 +40,40 @@ # LVC catalogs base_lvc_url = "https://www.gwosc.org/eventapi/jsonfull/{}/" + def lvk_catalogs(): _catalog_source = "https://gwosc.org/eventapi/json/" - catalog_list = json.load(open(get_file(_catalog_source), 'r')) + catalog_list = json.load(open(get_file(_catalog_source))) return catalog_list - + + def populate_catalogs(): - """ Refresh set of known catalogs - """ + """Refresh set of known catalogs""" global _catalogs if _catalogs is None: # update the LVK catalog information - _catalogs = {cname: 'LVK' for cname in lvk_catalogs().keys()} + _catalogs = dict.fromkeys(lvk_catalogs().keys(), "LVK") + _catalogs = None # add some aliases _aliases = {} -_aliases['gwtc-1'] = 'GWTC-1-confident' -_aliases['gwtc-2'] = 'GWTC-2' -_aliases['gwtc-2.1'] = 'GWTC-2.1-confident' -_aliases['gwtc-3'] = 'GWTC-3-confident' -_aliases['gwtc-4.0'] = 'GWTC-4.0' +_aliases["gwtc-1"] = "GWTC-1-confident" +_aliases["gwtc-2"] = "GWTC-2" +_aliases["gwtc-2.1"] = "GWTC-2.1-confident" +_aliases["gwtc-3"] = "GWTC-3-confident" +_aliases["gwtc-4.0"] = "GWTC-4.0" + def list_catalogs(): """Return a list of possible GW catalogs to query""" - populate_catalogs() + populate_catalogs() return list(_catalogs.keys()) + def get_source(source): - """Get the source data for a particular GW catalog - """ + """Get the source data for a particular GW catalog""" populate_catalogs() if source in _aliases: @@ -76,9 +81,9 @@ def get_source(source): if source in _catalogs: catalog_type = _catalogs[source] - if catalog_type == 'LVK': + if catalog_type == "LVK": fname = get_file(base_lvc_url.format(source), cache=True) - data = json.load(open(fname, 'r')) + data = json.load(open(fname)) else: - raise ValueError('Unkown catalog source {}'.format(source)) - return data['events'] + raise ValueError(f"Unkown catalog source {source}") + return data["events"] diff --git a/pycbc/constants.py b/pycbc/constants.py index c8452065576..09c834d37f5 100644 --- a/pycbc/constants.py +++ b/pycbc/constants.py @@ -12,52 +12,48 @@ import os import numpy as np -from astropy import ( - constants as aconstants, - units as aunits -) +from astropy import constants as aconstants +from astropy import units as aunits from pycbc.libutils import import_optional -lal = import_optional('lal') +lal = import_optional("lal") # We define a global logger for this module -logger = logging.getLogger('pycbc.constants') +logger = logging.getLogger("pycbc.constants") # Get the environment variable which defines which constants to use # Allowed values are 'default' or 'lal' -_CONSTANTS = os.environ.get('PYCBC_CONSTANT_SOURCE', 'default').lower() +_CONSTANTS = os.environ.get("PYCBC_CONSTANT_SOURCE", "default").lower() # first, do mappings for constants where the value is directly in astropy/numpy # All are in SI units _DEFAULT_MAPPING = { - 'C_SI': aconstants.c.value, # Speed of light - 'G_SI': aconstants.G.value, # Gravitational constant - 'MSUN_SI': aconstants.M_sun.value, # Mass of the Sun - 'PC_SI': aconstants.pc.value, # Parsec - 'REARTH_SI': aconstants.R_earth.value, # Earth equatorial radius - 'YRJUL_SI': aunits.year.to(aunits.s), # years in seconds - 'PI': np.pi, - 'TWOPI': 2 * np.pi, - 'PI_4': np.pi / 4, - 'GAMMA': np.euler_gamma, - 'LN2': np.log(2.) + "C_SI": aconstants.c.value, # Speed of light + "G_SI": aconstants.G.value, # Gravitational constant + "MSUN_SI": aconstants.M_sun.value, # Mass of the Sun + "PC_SI": aconstants.pc.value, # Parsec + "REARTH_SI": aconstants.R_earth.value, # Earth equatorial radius + "YRJUL_SI": aunits.year.to(aunits.s), # years in seconds + "PI": np.pi, + "TWOPI": 2 * np.pi, + "PI_4": np.pi / 4, + "GAMMA": np.euler_gamma, + "LN2": np.log(2.0), } # We need to define some constants from astropy values -MSUN_SI = _DEFAULT_MAPPING['MSUN_SI'] -C_SI = _DEFAULT_MAPPING['C_SI'] -G_SI = _DEFAULT_MAPPING['G_SI'] +MSUN_SI = _DEFAULT_MAPPING["MSUN_SI"] +C_SI = _DEFAULT_MAPPING["C_SI"] +G_SI = _DEFAULT_MAPPING["G_SI"] -MTSUN_SI = MSUN_SI * G_SI / (C_SI ** 3) +MTSUN_SI = MSUN_SI * G_SI / (C_SI**3) MRSUN_SI = MSUN_SI * G_SI / (C_SI * C_SI) # Add in these hybrid constants -_DEFAULT_MAPPING.update({ - 'MTSUN_SI': MTSUN_SI, - 'MRSUN_SI': MRSUN_SI -}) +_DEFAULT_MAPPING.update({"MTSUN_SI": MTSUN_SI, "MRSUN_SI": MRSUN_SI}) + # Define the Constant Lookup Function def get_constant(name): @@ -80,8 +76,9 @@ def get_constant(name): ------ NotImplementedError If the constant is not found in any of the available packages. + """ - if _CONSTANTS.lower() == 'lal': # Allow LAL + if _CONSTANTS.lower() == "lal": # Allow LAL if lal is None: raise ImportError( "PYCBC_CONSTANT_SOURCE is set to 'lal', but the 'lal' module is not installed. " @@ -89,14 +86,15 @@ def get_constant(name): ) return getattr(lal, name) - elif name in _DEFAULT_MAPPING: + if name in _DEFAULT_MAPPING: return _DEFAULT_MAPPING[name] raise NotImplementedError( - 'Contact the PyCBC team, this should never happen. You are ' - 'trying to use a constant which is not defined.' + "Contact the PyCBC team, this should never happen. You are " + "trying to use a constant which is not defined." ) + # Expose the constants as attributes of this module: for const_name in _DEFAULT_MAPPING: globals()[const_name] = get_constant(const_name) diff --git a/pycbc/conversions.py b/pycbc/conversions.py index 9bc90f5dbde..e454a19b871 100644 --- a/pycbc/conversions.py +++ b/pycbc/conversions.py @@ -29,23 +29,27 @@ """ import copy -import numpy import logging +import numpy -from pycbc.detector import Detector import pycbc.cosmology from pycbc import neutron_stars as ns -from pycbc.constants import YRJUL_SI, MSUN_SI, MTSUN_SI, C_SI, G_SI, PI +from pycbc.constants import C_SI, G_SI, MSUN_SI, MTSUN_SI, PI, YRJUL_SI +from pycbc.detector import Detector +from .coordinates import ( + cartesian_to_spherical as _cartesian_to_spherical, +) from .coordinates import ( spherical_to_cartesian as _spherical_to_cartesian, - cartesian_to_spherical as _cartesian_to_spherical) +) + +pykerr = pycbc.libutils.import_optional("pykerr") +lalsim = pycbc.libutils.import_optional("lalsimulation") -pykerr = pycbc.libutils.import_optional('pykerr') -lalsim = pycbc.libutils.import_optional('lalsimulation') +logger = logging.getLogger("pycbc.conversions") -logger = logging.getLogger('pycbc.conversions') # # ============================================================================= @@ -55,7 +59,8 @@ # ============================================================================= # def ensurearray(*args): - """Apply numpy's broadcast rules to the given arguments. + """ + Apply numpy's broadcast rules to the given arguments. This will ensure that all of the arguments are numpy arrays and that they all have the same shape. See ``numpy.broadcast_arrays`` for more details. @@ -75,6 +80,7 @@ def ensurearray(*args): arguments. The first N values are the input arguments as ``ndarrays``s. The last value is a boolean indicating whether any of the inputs was an array. + """ input_is_array = any(isinstance(arg, numpy.ndarray) for arg in args) args = list(numpy.broadcast_arrays(*args)) @@ -83,12 +89,15 @@ def ensurearray(*args): def formatreturn(arg, input_is_array=False): - """If the given argument is a numpy array with shape (1,), just returns - that value.""" + """ + If the given argument is a numpy array with shape (1,), just returns + that value. + """ if not input_is_array and arg.size == 1: arg = arg.item() return arg + # # ============================================================================= # @@ -97,8 +106,9 @@ def formatreturn(arg, input_is_array=False): # ============================================================================= # + def sec_to_year(sec): - """ Converts number of seconds to number of years """ + """Converts number of seconds to number of years""" return sec / YRJUL_SI @@ -128,21 +138,24 @@ def hypertriangle(*params, bounds=(0, 1)): ------- array The mapped parameters. Output values are in ascending order. + """ # check inputs all have the same shape ref_shape = numpy.shape(params[0]) - assert numpy.all([numpy.shape(params[i]) == ref_shape for i in range(len(params))]), \ - "All inputs must have the same number of elements" + assert numpy.all( + [numpy.shape(params[i]) == ref_shape for i in range(len(params))] + ), "All inputs must have the same number of elements" # map to numpy array params, input_is_array = ensurearray(params) # check all values lie within bounds - assert numpy.all(params >= bounds[0]) and numpy.all(params <= bounds[1]), \ + assert numpy.all(params >= bounds[0]) and numpy.all(params <= bounds[1]), ( "Input parameters lie outside of given bounds" + ) # rescale the parameters to the unit hypercube - scaled_params = (params - bounds[0])/(bounds[1] - bounds[0]) + scaled_params = (params - bounds[0]) / (bounds[1] - bounds[0]) # hypertriangulate try: @@ -153,7 +166,7 @@ def hypertriangle(*params, bounds=(0, 1)): idx = numpy.repeat(numpy.arange(K), repeats=num_pts) scaled_params.resize(K, num_pts) idx.resize(K, num_pts) - fac = numpy.power(1 - scaled_params, 1/(K - idx)) + fac = numpy.power(1 - scaled_params, 1 / (K - idx)) out_scaled_params = 1 - numpy.cumprod(fac, axis=0) # rescale to prior bounds @@ -162,6 +175,7 @@ def hypertriangle(*params, bounds=(0, 1)): out_params = [out_params[i][0] for i in range(K)] return out_params + # # ============================================================================= # @@ -208,16 +222,19 @@ def invq_from_mass1_mass2(mass1, mass2): def eta_from_mass1_mass2(mass1, mass2): """Returns the symmetric mass ratio from mass1 and mass2.""" - return mass1*mass2 / (mass1 + mass2)**2. + return mass1 * mass2 / (mass1 + mass2) ** 2.0 def mchirp_from_mass1_mass2(mass1, mass2): """Returns the chirp mass from mass1 and mass2.""" - return eta_from_mass1_mass2(mass1, mass2)**(3./5) * (mass1 + mass2) + return eta_from_mass1_mass2(mass1, mass2) ** (3.0 / 5) * (mass1 + mass2) -def eccmchirp_from_mass1_mass2_eccentricity(mass1, mass2, eccentricity, method='spa_phase'): - """Returns the effective eccentric chirp mass from mass1, mass2 and eccentricity +def eccmchirp_from_mass1_mass2_eccentricity( + mass1, mass2, eccentricity, method="spa_phase" +): + """ + Returns the effective eccentric chirp mass from mass1, mass2 and eccentricity Parameters ---------- @@ -232,71 +249,76 @@ def eccmchirp_from_mass1_mass2_eccentricity(mass1, mass2, eccentricity, method=' method : str, optiona Method to use for the calculation ("spa_phase", "fit"). See `eccmchirp_from_mchirp_eccentricity` for details. + """ allowed_methods = ("spa_phase", "fit") if method not in allowed_methods: - raise ValueError("method must be one of {}".format(allowed_methods)) + raise ValueError(f"method must be one of {allowed_methods}") mchirp = mchirp_from_mass1_mass2(mass1, mass2) - mchirp_eccentric = eccmchirp_from_mchirp_eccentricity(mchirp, eccentricity, method=method) + mchirp_eccentric = eccmchirp_from_mchirp_eccentricity( + mchirp, eccentricity, method=method + ) return mchirp_eccentric def mass1_from_mtotal_q(mtotal, q): - """Returns a component mass from the given total mass and mass ratio. + """ + Returns a component mass from the given total mass and mass ratio. If the mass ratio q is >= 1, the returned mass will be the primary (heavier) mass. If q < 1, the returned mass will be the secondary (lighter) mass. """ - return q*mtotal / (1. + q) + return q * mtotal / (1.0 + q) def mass2_from_mtotal_q(mtotal, q): - """Returns a component mass from the given total mass and mass ratio. + """ + Returns a component mass from the given total mass and mass ratio. If the mass ratio q is >= 1, the returned mass will be the secondary (lighter) mass. If q < 1, the returned mass will be the primary (heavier) mass. """ - return mtotal / (1. + q) + return mtotal / (1.0 + q) def mass1_from_mtotal_eta(mtotal, eta): - """Returns the primary mass from the total mass and symmetric mass + """ + Returns the primary mass from the total mass and symmetric mass ratio. """ - return 0.5 * mtotal * (1.0 + (1.0 - 4.0 * eta)**0.5) + return 0.5 * mtotal * (1.0 + (1.0 - 4.0 * eta) ** 0.5) def mass2_from_mtotal_eta(mtotal, eta): - """Returns the secondary mass from the total mass and symmetric mass + """ + Returns the secondary mass from the total mass and symmetric mass ratio. """ - return 0.5 * mtotal * (1.0 - (1.0 - 4.0 * eta)**0.5) + return 0.5 * mtotal * (1.0 - (1.0 - 4.0 * eta) ** 0.5) def mtotal_from_mchirp_eta(mchirp, eta): - """Returns the total mass from the chirp mass and symmetric mass ratio. - """ - return mchirp / eta**(3./5.) + """Returns the total mass from the chirp mass and symmetric mass ratio.""" + return mchirp / eta ** (3.0 / 5.0) def mass1_from_mchirp_eta(mchirp, eta): - """Returns the primary mass from the chirp mass and symmetric mass ratio. - """ + """Returns the primary mass from the chirp mass and symmetric mass ratio.""" mtotal = mtotal_from_mchirp_eta(mchirp, eta) return mass1_from_mtotal_eta(mtotal, eta) def mass2_from_mchirp_eta(mchirp, eta): - """Returns the primary mass from the chirp mass and symmetric mass ratio. - """ + """Returns the primary mass from the chirp mass and symmetric mass ratio.""" mtotal = mtotal_from_mchirp_eta(mchirp, eta) return mass2_from_mtotal_eta(mtotal, eta) def _mass2_from_mchirp_mass1(mchirp, mass1): - r"""Returns the secondary mass from the chirp mass and primary mass. + r""" + Returns the secondary mass from the chirp mass and primary mass. As this is a cubic equation this requires finding the roots and returning the one that is real. Basically it can be shown that: @@ -317,12 +339,15 @@ def _mass2_from_mchirp_mass1(mchirp, mass1): real_root = roots[(abs(roots - roots.real)).argmin()] return real_root.real + mass2_from_mchirp_mass1 = numpy.vectorize(_mass2_from_mchirp_mass1) -def _mass_from_knownmass_eta(known_mass, eta, known_is_secondary=False, - force_real=True): - r"""Returns the other component mass given one of the component masses +def _mass_from_knownmass_eta( + known_mass, eta, known_is_secondary=False, force_real=True +): + r""" + Returns the other component mass given one of the component masses and the symmetric mass ratio. This requires finding the roots of the quadratic equation: @@ -353,36 +378,42 @@ def _mass_from_knownmass_eta(known_mass, eta, known_is_secondary=False, ------- float The other component mass. + """ - roots = numpy.roots([eta, (2*eta - 1) * known_mass, eta * known_mass**2.]) + roots = numpy.roots([eta, (2 * eta - 1) * known_mass, eta * known_mass**2.0]) if force_real: roots = numpy.real(roots) if known_is_secondary: return roots[roots.argmax()] - else: - return roots[roots.argmin()] + return roots[roots.argmin()] + mass_from_knownmass_eta = numpy.vectorize(_mass_from_knownmass_eta) def mass2_from_mass1_eta(mass1, eta, force_real=True): - """Returns the secondary mass from the primary mass and symmetric mass + """ + Returns the secondary mass from the primary mass and symmetric mass ratio. """ - return mass_from_knownmass_eta(mass1, eta, known_is_secondary=False, - force_real=force_real) + return mass_from_knownmass_eta( + mass1, eta, known_is_secondary=False, force_real=force_real + ) def mass1_from_mass2_eta(mass2, eta, force_real=True): - """Returns the primary mass from the secondary mass and symmetric mass + """ + Returns the primary mass from the secondary mass and symmetric mass ratio. """ - return mass_from_knownmass_eta(mass2, eta, known_is_secondary=True, - force_real=force_real) + return mass_from_knownmass_eta( + mass2, eta, known_is_secondary=True, force_real=force_real + ) def eta_from_q(q): - r"""Returns the symmetric mass ratio from the given mass ratio. + r""" + Returns the symmetric mass ratio from the given mass ratio. This is given by: @@ -391,88 +422,86 @@ def eta_from_q(q): Note that the mass ratio may be either < 1 or > 1. """ - return q / (1. + q)**2 + return q / (1.0 + q) ** 2 def mass1_from_mchirp_q(mchirp, q): """Returns the primary mass from the given chirp mass and mass ratio.""" - mass1 = q**(2./5.) * (1.0 + q)**(1./5.) * mchirp + mass1 = q ** (2.0 / 5.0) * (1.0 + q) ** (1.0 / 5.0) * mchirp return mass1 def mass2_from_mchirp_q(mchirp, q): """Returns the secondary mass from the given chirp mass and mass ratio.""" - mass2 = q**(-3./5.) * (1.0 + q)**(1./5.) * mchirp + mass2 = q ** (-3.0 / 5.0) * (1.0 + q) ** (1.0 / 5.0) * mchirp return mass2 def _a0(f_lower): - """Used in calculating chirp times: see Cokelaer, arxiv.org:0706.4437 - appendix 1, also lalinspiral/python/sbank/tau0tau3.py. """ - return 5. / (256. * (numpy.pi * f_lower)**(8./3.)) + Used in calculating chirp times: see Cokelaer, arxiv.org:0706.4437 + appendix 1, also lalinspiral/python/sbank/tau0tau3.py. + """ + return 5.0 / (256.0 * (numpy.pi * f_lower) ** (8.0 / 3.0)) def _a3(f_lower): """Another parameter used for chirp times""" - return numpy.pi / (8. * (numpy.pi * f_lower)**(5./3.)) + return numpy.pi / (8.0 * (numpy.pi * f_lower) ** (5.0 / 3.0)) def tau0_from_mtotal_eta(mtotal, eta, f_lower): - r"""Returns :math:`\tau_0` from the total mass, symmetric mass ratio, and + r""" + Returns :math:`\tau_0` from the total mass, symmetric mass ratio, and the given frequency. """ # convert to seconds mtotal = mtotal * MTSUN_SI # formulae from arxiv.org:0706.4437 - return _a0(f_lower) / (mtotal**(5./3.) * eta) + return _a0(f_lower) / (mtotal ** (5.0 / 3.0) * eta) def tau0_from_mchirp(mchirp, f_lower): - r"""Returns :math:`\tau_0` from the chirp mass and the given frequency. - """ + r"""Returns :math:`\tau_0` from the chirp mass and the given frequency.""" # convert to seconds mchirp = mchirp * MTSUN_SI # formulae from arxiv.org:0706.4437 - return _a0(f_lower) / mchirp ** (5./3.) + return _a0(f_lower) / mchirp ** (5.0 / 3.0) def tau3_from_mtotal_eta(mtotal, eta, f_lower): - r"""Returns :math:`\tau_0` from the total mass, symmetric mass ratio, and + r""" + Returns :math:`\tau_0` from the total mass, symmetric mass ratio, and the given frequency. """ # convert to seconds mtotal = mtotal * MTSUN_SI # formulae from arxiv.org:0706.4437 - return _a3(f_lower) / (mtotal**(2./3.) * eta) + return _a3(f_lower) / (mtotal ** (2.0 / 3.0) * eta) def tau0_from_mass1_mass2(mass1, mass2, f_lower): - r"""Returns :math:`\tau_0` from the component masses and given frequency. - """ + r"""Returns :math:`\tau_0` from the component masses and given frequency.""" mtotal = mass1 + mass2 eta = eta_from_mass1_mass2(mass1, mass2) return tau0_from_mtotal_eta(mtotal, eta, f_lower) def tau3_from_mass1_mass2(mass1, mass2, f_lower): - r"""Returns :math:`\tau_3` from the component masses and given frequency. - """ + r"""Returns :math:`\tau_3` from the component masses and given frequency.""" mtotal = mass1 + mass2 eta = eta_from_mass1_mass2(mass1, mass2) return tau3_from_mtotal_eta(mtotal, eta, f_lower) def mchirp_from_tau0(tau0, f_lower): - r"""Returns chirp mass from :math:`\tau_0` and the given frequency. - """ - mchirp = (_a0(f_lower) / tau0) ** (3./5.) # in seconds + r"""Returns chirp mass from :math:`\tau_0` and the given frequency.""" + mchirp = (_a0(f_lower) / tau0) ** (3.0 / 5.0) # in seconds # convert back to solar mass units return mchirp / MTSUN_SI -def mtotal_from_tau0_tau3(tau0, tau3, f_lower, - in_seconds=False): +def mtotal_from_tau0_tau3(tau0, tau3, f_lower, in_seconds=False): r"""Returns total mass from :math:`\tau_0, \tau_3`.""" mtotal = (tau3 / _a3(f_lower)) / (tau0 / _a0(f_lower)) if not in_seconds: @@ -483,9 +512,8 @@ def mtotal_from_tau0_tau3(tau0, tau3, f_lower, def eta_from_tau0_tau3(tau0, tau3, f_lower): r"""Returns symmetric mass ratio from :math:`\tau_0, \tau_3`.""" - mtotal = mtotal_from_tau0_tau3(tau0, tau3, f_lower, - in_seconds=True) - eta = mtotal**(-2./3.) * (_a3(f_lower) / tau3) + mtotal = mtotal_from_tau0_tau3(tau0, tau3, f_lower, in_seconds=True) + eta = mtotal ** (-2.0 / 3.0) * (_a3(f_lower) / tau3) return eta @@ -504,7 +532,8 @@ def mass2_from_tau0_tau3(tau0, tau3, f_lower): def eccmchirp_from_mchirp_eccentricity(mchirp, eccentricity, method="spa_phase"): - """Return the effective eccentric chirp mass: eccmchirp. + """ + Return the effective eccentric chirp mass: eccmchirp. Parameters ---------- @@ -531,18 +560,19 @@ def eccmchirp_from_mchirp_eccentricity(mchirp, eccentricity, method="spa_phase") time-frequency track. The fitting formula is given in Eq. 8 of https://arxiv.org/abs/2107.14736. Eccentricity must be defined at the dominant (2,2) mode GW frequency of 10 Hz. + """ allowed_methods = ("spa_phase", "fit") if method not in allowed_methods: - raise ValueError("method must be one of {}".format(allowed_methods)) + raise ValueError(f"method must be one of {allowed_methods}") mchirp, eccentricity, input_is_array = ensurearray(mchirp, eccentricity) e2 = eccentricity * eccentricity if method == "spa_phase": - Emchirp = mchirp / (1.0 - 157.0 / 24.0 * e2)**(3.0 / 5.0) + Emchirp = mchirp / (1.0 - 157.0 / 24.0 * e2) ** (3.0 / 5.0) elif method == "fit": # Constants from Table 1 of https://arxiv.org/abs/2107.14736 @@ -564,19 +594,21 @@ def eccmchirp_from_mchirp_eccentricity(mchirp, eccentricity, method="spa_phase") # Calculate coefficients (Eq 9) alpha = xi * mchirp + delta - beta = mchirp2 * ( Xi_beta + - mchirp2 * ( Delta_beta + - mchirp2 * ( kappa_beta + - mchirp2 * zeta_beta ))) - gamma = mchirp2 * ( Xi_gamma + - mchirp2 * ( Delta_gamma + - mchirp2 * ( kappa_gamma + - mchirp2 * zeta_gamma ))) + beta = mchirp2 * ( + Xi_beta + + mchirp2 * (Delta_beta + mchirp2 * (kappa_beta + mchirp2 * zeta_beta)) + ) + gamma = mchirp2 * ( + Xi_gamma + + mchirp2 * (Delta_gamma + mchirp2 * (kappa_gamma + mchirp2 * zeta_gamma)) + ) Emchirp = mchirp * (1 + e2 * (alpha + e2 * (beta + e2 * gamma))) return formatreturn(Emchirp, input_is_array) + def mchirp_from_eccmchirp_eccentricity(eccmchirp, eccentricity, method="spa_phase"): - """Return the chirp mass from the effective eccentric chirp mass and eccentricity. + """ + Return the chirp mass from the effective eccentric chirp mass and eccentricity. Parameters ---------- @@ -601,18 +633,19 @@ def mchirp_from_eccmchirp_eccentricity(eccmchirp, eccentricity, method="spa_phas fitting to time-frequency track. The fitting formula is given in Eq. 8 of https://arxiv.org/abs/2107.14736. Eccentricity must be defined at the dominant (2,2) mode GW frequency of 10 Hz. + """ allowed_methods = ("spa_phase", "fit") if method not in allowed_methods: - raise ValueError("method must be one of {}".format(allowed_methods)) + raise ValueError(f"method must be one of {allowed_methods}") eccmchirp, eccentricity, input_is_array = ensurearray(eccmchirp, eccentricity) e2 = eccentricity * eccentricity if method == "spa_phase": - m = eccmchirp * ( 1 - 157/24 * e2 )**(3/5) + m = eccmchirp * (1 - 157 / 24 * e2) ** (3 / 5) elif method == "fit": # Constants from Table 1 of https://arxiv.org/abs/2107.14736 @@ -631,130 +664,142 @@ def mchirp_from_eccmchirp_eccentricity(eccmchirp, eccentricity, method="spa_phas # Initial guess using the quadratic approximation A = xi * e2 B = 1 + delta * e2 - C = - eccmchirp - m = numpy.where(A > 0, (-B + numpy.sqrt(B**2 - 4*A*C)) / (2*A), eccmchirp) + C = -eccmchirp + m = numpy.where(A > 0, (-B + numpy.sqrt(B**2 - 4 * A * C)) / (2 * A), eccmchirp) for _ in range(5): m2 = m * m alpha = xi * m + delta - beta = m2 * ( Xi_beta + m2 * ( Delta_beta + - m2 * ( kappa_beta + - m2 * zeta_beta ))) - gamma = m2 * ( Xi_gamma + - m2 * ( Delta_gamma + - m2 * ( kappa_gamma + - m2 * zeta_gamma ))) - f = m * (1 + e2 * (alpha + e2 * ( beta + e2 *gamma ))) - Emchirp + beta = m2 * ( + Xi_beta + m2 * (Delta_beta + m2 * (kappa_beta + m2 * zeta_beta)) + ) + gamma = m2 * ( + Xi_gamma + m2 * (Delta_gamma + m2 * (kappa_gamma + m2 * zeta_gamma)) + ) + f = m * (1 + e2 * (alpha + e2 * (beta + e2 * gamma))) - Emchirp d_alpha = xi - d_beta = m * ( 2 * Xi_beta + m2 * ( 4 * Delta_beta + - m2 *( 6 * kappa_beta + - m2 * 8 * zeta_beta ))) - d_gamma = m * ( 2 * Xi_gamma + m2 * ( 4 * Delta_gamma + - m2 * ( 6 * kappa_gamma + - m2 * 8 * zeta_gamma ))) - df = (1 + e2 * (alpha + e2 * ( beta + e2 * gamma )) + - m * e2 * (d_alpha + e2 * ( d_beta + e2 * d_gamma ))) + d_beta = m * ( + 2 * Xi_beta + + m2 * (4 * Delta_beta + m2 * (6 * kappa_beta + m2 * 8 * zeta_beta)) + ) + d_gamma = m * ( + 2 * Xi_gamma + + m2 * (4 * Delta_gamma + m2 * (6 * kappa_gamma + m2 * 8 * zeta_gamma)) + ) + df = ( + 1 + + e2 * (alpha + e2 * (beta + e2 * gamma)) + + m * e2 * (d_alpha + e2 * (d_beta + e2 * d_gamma)) + ) m = m - f / df return formatreturn(m, input_is_array) + def lambda_tilde(mass1, mass2, lambda1, lambda2): - """ The effective lambda parameter + """ + The effective lambda parameter The mass-weighted dominant effective lambda parameter defined in https://journals.aps.org/prd/pdf/10.1103/PhysRevD.91.043002 """ m1, m2, lambda1, lambda2, input_is_array = ensurearray( - mass1, mass2, lambda1, lambda2) + mass1, mass2, lambda1, lambda2 + ) lsum = lambda1 + lambda2 ldiff, _ = ensurearray(lambda1 - lambda2) mask = m1 < m2 ldiff[mask] = -ldiff[mask] eta = eta_from_mass1_mass2(m1, m2) - eta[eta > 0.25] = 0.25 # Account for numerical error, 0.25 is the max - p1 = (lsum) * (1 + 7. * eta - 31 * eta ** 2.0) - p2 = (1 - 4 * eta)**0.5 * (1 + 9 * eta - 11 * eta ** 2.0) * (ldiff) + eta[eta > 0.25] = 0.25 # Account for numerical error, 0.25 is the max + p1 = (lsum) * (1 + 7.0 * eta - 31 * eta**2.0) + p2 = (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2.0) * (ldiff) return formatreturn(8.0 / 13.0 * (p1 + p2), input_is_array) + def delta_lambda_tilde(mass1, mass2, lambda1, lambda2): - """ Delta lambda tilde parameter defined as + """ + Delta lambda tilde parameter defined as equation 15 in https://journals.aps.org/prd/pdf/10.1103/PhysRevD.91.043002 """ m1, m2, lambda1, lambda2, input_is_array = ensurearray( - mass1, mass2, lambda1, lambda2) + mass1, mass2, lambda1, lambda2 + ) lsum = lambda1 + lambda2 ldiff, _ = ensurearray(lambda1 - lambda2) mask = m1 < m2 ldiff[mask] = -ldiff[mask] eta = eta_from_mass1_mass2(m1, m2) - p1 = numpy.sqrt(1 - 4 * eta) * ( - 1 - (13272 / 1319) * eta + - (8944 / 1319) * eta ** 2 - ) * lsum + p1 = ( + numpy.sqrt(1 - 4 * eta) + * (1 - (13272 / 1319) * eta + (8944 / 1319) * eta**2) + * lsum + ) p2 = ( - 1 - (15910 / 1319) * eta + - (32850 / 1319) * eta ** 2 + - (3380 / 1319) * eta ** 3 + 1 - (15910 / 1319) * eta + (32850 / 1319) * eta**2 + (3380 / 1319) * eta**3 ) * ldiff return formatreturn(1 / 2 * (p1 + p2), input_is_array) -def lambda1_from_delta_lambda_tilde_lambda_tilde(delta_lambda_tilde, - lambda_tilde, - mass1, - mass2): - """ Returns lambda1 parameter by using delta lambda tilde, + +def lambda1_from_delta_lambda_tilde_lambda_tilde( + delta_lambda_tilde, lambda_tilde, mass1, mass2 +): + """ + Returns lambda1 parameter by using delta lambda tilde, lambda tilde, mass1, and mass2. """ m1, m2, delta_lambda_tilde, lambda_tilde, input_is_array = ensurearray( - mass1, mass2, delta_lambda_tilde, lambda_tilde) + mass1, mass2, delta_lambda_tilde, lambda_tilde + ) eta = eta_from_mass1_mass2(m1, m2) - p1 = 1 + 7.0*eta - 31*eta**2.0 - p2 = (1 - 4*eta)**0.5 * (1 + 9*eta - 11*eta**2.0) - p3 = (1 - 4*eta)**0.5 * (1 - 13272/1319*eta + 8944/1319*eta**2) - p4 = 1 - (15910/1319)*eta + (32850/1319)*eta**2 + (3380/1319)*eta**3 - amp = 1/((p1*p4)-(p2*p3)) - l_tilde_lambda1 = 13/16 * (p3-p4) * lambda_tilde - l_delta_tilde_lambda1 = (p1-p2) * delta_lambda_tilde + p1 = 1 + 7.0 * eta - 31 * eta**2.0 + p2 = (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2.0) + p3 = (1 - 4 * eta) ** 0.5 * (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2) + p4 = 1 - (15910 / 1319) * eta + (32850 / 1319) * eta**2 + (3380 / 1319) * eta**3 + amp = 1 / ((p1 * p4) - (p2 * p3)) + l_tilde_lambda1 = 13 / 16 * (p3 - p4) * lambda_tilde + l_delta_tilde_lambda1 = (p1 - p2) * delta_lambda_tilde lambda1 = formatreturn( - amp * (l_delta_tilde_lambda1 - l_tilde_lambda1), - input_is_array + amp * (l_delta_tilde_lambda1 - l_tilde_lambda1), input_is_array ) return lambda1 + def lambda2_from_delta_lambda_tilde_lambda_tilde( - delta_lambda_tilde, - lambda_tilde, - mass1, - mass2): - """ Returns lambda2 parameter by using delta lambda tilde, + delta_lambda_tilde, lambda_tilde, mass1, mass2 +): + """ + Returns lambda2 parameter by using delta lambda tilde, lambda tilde, mass1, and mass2. """ m1, m2, delta_lambda_tilde, lambda_tilde, input_is_array = ensurearray( - mass1, mass2, delta_lambda_tilde, lambda_tilde) + mass1, mass2, delta_lambda_tilde, lambda_tilde + ) eta = eta_from_mass1_mass2(m1, m2) - p1 = 1 + 7.0*eta - 31*eta**2.0 - p2 = (1 - 4*eta)**0.5 * (1 + 9*eta - 11*eta**2.0) - p3 = (1 - 4*eta)**0.5 * (1 - 13272/1319*eta + 8944/1319*eta**2) - p4 = 1 - (15910/1319)*eta + (32850/1319)*eta**2 + (3380/1319)*eta**3 - amp = 1/((p1*p4)-(p2*p3)) - l_tilde_lambda2 = 13/16 * (p3+p4) * lambda_tilde - l_delta_tilde_lambda2 = (p1+p2) * delta_lambda_tilde + p1 = 1 + 7.0 * eta - 31 * eta**2.0 + p2 = (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2.0) + p3 = (1 - 4 * eta) ** 0.5 * (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2) + p4 = 1 - (15910 / 1319) * eta + (32850 / 1319) * eta**2 + (3380 / 1319) * eta**3 + amp = 1 / ((p1 * p4) - (p2 * p3)) + l_tilde_lambda2 = 13 / 16 * (p3 + p4) * lambda_tilde + l_delta_tilde_lambda2 = (p1 + p2) * delta_lambda_tilde lambda2 = formatreturn( - amp * (l_tilde_lambda2 - l_delta_tilde_lambda2), - input_is_array + amp * (l_tilde_lambda2 - l_delta_tilde_lambda2), input_is_array ) return lambda2 -def lambda_from_mass_tov_file(mass, tov_file, distance=0.): - """Return the lambda parameter(s) corresponding to the input mass(es) + +def lambda_from_mass_tov_file(mass, tov_file, distance=0.0): + """ + Return the lambda parameter(s) corresponding to the input mass(es) interpolating from the mass-Lambda data for a particular EOS read in from an ASCII file. """ data = numpy.loadtxt(tov_file) mass_from_file = data[:, 0] lambda_from_file = data[:, 1] - mass_src = mass/(1.0 + pycbc.cosmology.redshift(distance)) + mass_src = mass / (1.0 + pycbc.cosmology.redshift(distance)) lambdav = numpy.interp(mass_src, mass_from_file, lambda_from_file) return lambdav @@ -781,11 +826,12 @@ def ensure_obj1_is_primary(mass1, mass2, *params): list : A list with mass1, mass2, params as arrays, with elements, each with elements re-arranged so that object 1 is the primary. + """ # Check params are 2N if len(params) % 2 != 0: raise ValueError("params must be 2N floats or arrays") - input_properties, input_is_array = ensurearray((mass1, mass2)+params) + input_properties, input_is_array = ensurearray((mass1, mass2) + params) # Check inputs are all the same length shapes = [par.shape for par in input_properties] if len(set(shapes)) != 1: @@ -798,9 +844,9 @@ def ensure_obj1_is_primary(mass1, mass2, *params): # primary (p) p = copy.copy(input_properties[i]) # secondary (s) - s = copy.copy(input_properties[i+1]) + s = copy.copy(input_properties[i + 1]) # Swap - p[mask] = input_properties[i+1][mask] + p[mask] = input_properties[i + 1][mask] s[mask] = input_properties[i][mask] # Format and include in output object output_properties.append(formatreturn(p, input_is_array)) @@ -810,9 +856,17 @@ def ensure_obj1_is_primary(mass1, mass2, *params): def remnant_mass_from_mass1_mass2_spherical_spin_eos( - mass1, mass2, spin1_a=0.0, spin1_polar=0.0, eos='2H', - spin2_a=0.0, spin2_polar=0.0, swap_companions=False, - ns_bh_mass_boundary=None, extrapolate=False): + mass1, + mass2, + spin1_a=0.0, + spin1_polar=0.0, + eos="2H", + spin2_a=0.0, + spin2_polar=0.0, + swap_companions=False, + ns_bh_mass_boundary=None, + extrapolate=False, +): """ Function that determines the remnant disk mass of an NS-BH system using the fit to numerical-relativity results discussed in @@ -825,7 +879,7 @@ def remnant_mass_from_mass1_mass2_spherical_spin_eos( Note: The NS spin does not play any role in this fit! Parameters - ----------- + ---------- mass1 : float The mass of the black hole, in solar masses. mass2 : float @@ -855,27 +909,31 @@ def remnant_mass_from_mass1_mass2_spherical_spin_eos( state prescribes the maximum possible mass2. Default is False. Returns - ---------- + ------- remnant_mass: float The remnant mass in solar masses + """ - mass1, mass2, spin1_a, spin1_polar, spin2_a, spin2_polar, \ - input_is_array = \ + mass1, mass2, spin1_a, spin1_polar, spin2_a, spin2_polar, input_is_array = ( ensurearray(mass1, mass2, spin1_a, spin1_polar, spin2_a, spin2_polar) - assert numpy.all(spin1_a >= 0) and numpy.all(spin2_a >= 0), \ + ) + assert numpy.all(spin1_a >= 0) and numpy.all(spin2_a >= 0), ( "Spin magnitude MUST be null or positive" + ) # mass1 must be greater than mass2: swap the properties of 1 and 2 or fail if swap_companions: - mass1, mass2, spin1_a, spin2_a, spin1_polar, spin2_polar = \ - ensure_obj1_is_primary(mass1, mass2, spin1_a, spin2_a, - spin1_polar, spin2_polar) + mass1, mass2, spin1_a, spin2_a, spin1_polar, spin2_polar = ( + ensure_obj1_is_primary( + mass1, mass2, spin1_a, spin2_a, spin1_polar, spin2_polar + ) + ) else: try: if any(mass2 > mass1) and input_is_array: - raise ValueError(f'Require mass1 >= mass2') + raise ValueError("Require mass1 >= mass2") except TypeError: if mass2 > mass1 and not input_is_array: - raise ValueError(f'Require mass1 >= mass2. {mass1} < {mass2}') + raise ValueError(f"Require mass1 >= mass2. {mass1} < {mass2}") eta = eta_from_mass1_mass2(mass1, mass2) # If a maximum NS mass is not provided, accept all values and # let the EOS handle this (in ns.initialize_eos) @@ -886,18 +944,29 @@ def remnant_mass_from_mass1_mass2_spherical_spin_eos( mask = mass2 <= ns_bh_mass_boundary # ...and return 0's otherwise remnant_mass = numpy.zeros(ensurearray(mass2)[0].size) - ns_compactness, ns_b_mass = ns.initialize_eos(mass2[mask], eos, - extrapolate=extrapolate) + ns_compactness, ns_b_mass = ns.initialize_eos( + mass2[mask], eos, extrapolate=extrapolate + ) remnant_mass[mask] = ns.foucart18( - eta[mask], ns_compactness, ns_b_mass, - spin1_a[mask], spin1_polar[mask]) + eta[mask], ns_compactness, ns_b_mass, spin1_a[mask], spin1_polar[mask] + ) return formatreturn(remnant_mass, input_is_array) def remnant_mass_from_mass1_mass2_cartesian_spin_eos( - mass1, mass2, spin1x=0.0, spin1y=0.0, spin1z=0.0, eos='2H', - spin2x=0.0, spin2y=0.0, spin2z=0.0, swap_companions=False, - ns_bh_mass_boundary=None, extrapolate=False): + mass1, + mass2, + spin1x=0.0, + spin1y=0.0, + spin1z=0.0, + eos="2H", + spin2x=0.0, + spin2y=0.0, + spin2z=0.0, + swap_companions=False, + ns_bh_mass_boundary=None, + extrapolate=False, +): """ Function that determines the remnant disk mass of an NS-BH system using the fit to numerical-relativity results discussed in @@ -910,7 +979,7 @@ def remnant_mass_from_mass1_mass2_cartesian_spin_eos( Note: NS spin is assumed to be 0! Parameters - ----------- + ---------- mass1 : float The mass of the black hole, in solar masses. mass2 : float @@ -944,23 +1013,30 @@ def remnant_mass_from_mass1_mass2_cartesian_spin_eos( state prescribes the maximum possible mass2. Default is False. Returns - ---------- + ------- remnant_mass: float The remnant mass in solar masses + """ spin1_a, _, spin1_polar = _cartesian_to_spherical(spin1x, spin1y, spin1z) if swap_companions: - spin2_a, _, spin2_polar = _cartesian_to_spherical(spin2x, - spin2y, spin2z) + spin2_a, _, spin2_polar = _cartesian_to_spherical(spin2x, spin2y, spin2z) else: size = ensurearray(spin1_a)[0].size spin2_a = numpy.zeros(size) spin2_polar = numpy.zeros(size) return remnant_mass_from_mass1_mass2_spherical_spin_eos( - mass1, mass2, spin1_a=spin1_a, spin1_polar=spin1_polar, eos=eos, - spin2_a=spin2_a, spin2_polar=spin2_polar, + mass1, + mass2, + spin1_a=spin1_a, + spin1_polar=spin1_polar, + eos=eos, + spin2_a=spin2_a, + spin2_polar=spin2_polar, swap_companions=swap_companions, - ns_bh_mass_boundary=ns_bh_mass_boundary, extrapolate=extrapolate) + ns_bh_mass_boundary=ns_bh_mass_boundary, + extrapolate=extrapolate, + ) # @@ -976,14 +1052,16 @@ def chi_eff(mass1, mass2, spin1z, spin2z): def chi_a(mass1, mass2, spin1z, spin2z): - """ Returns the aligned mass-weighted spin difference from mass1, mass2, + """ + Returns the aligned mass-weighted spin difference from mass1, mass2, spin1z, and spin2z. """ return (spin2z * mass2 - spin1z * mass1) / (mass2 + mass1) def chi_p(mass1, mass2, spin1x, spin1y, spin2x, spin2y): - """Returns the effective precession spin from mass1, mass2, spin1x, + """ + Returns the effective precession spin from mass1, mass2, spin1x, spin1y, spin2x, and spin2y. """ xi1 = secondary_xi(mass1, mass2, spin1x, spin1y, spin2x, spin2y) @@ -992,45 +1070,54 @@ def chi_p(mass1, mass2, spin1x, spin1y, spin2x, spin2y): def phi_a(mass1, mass2, spin1x, spin1y, spin2x, spin2y): - """ Returns the angle between the in-plane perpendicular spins.""" - phi1 = phi_from_spinx_spiny(primary_spin(mass1, mass2, spin1x, spin2x), - primary_spin(mass1, mass2, spin1y, spin2y)) - phi2 = phi_from_spinx_spiny(secondary_spin(mass1, mass2, spin1x, spin2x), - secondary_spin(mass1, mass2, spin1y, spin2y)) + """Returns the angle between the in-plane perpendicular spins.""" + phi1 = phi_from_spinx_spiny( + primary_spin(mass1, mass2, spin1x, spin2x), + primary_spin(mass1, mass2, spin1y, spin2y), + ) + phi2 = phi_from_spinx_spiny( + secondary_spin(mass1, mass2, spin1x, spin2x), + secondary_spin(mass1, mass2, spin1y, spin2y), + ) return (phi1 - phi2) % (2 * numpy.pi) def phi_s(spin1x, spin1y, spin2x, spin2y): - """ Returns the sum of the in-plane perpendicular spins.""" + """Returns the sum of the in-plane perpendicular spins.""" phi1 = phi_from_spinx_spiny(spin1x, spin1y) phi2 = phi_from_spinx_spiny(spin2x, spin2y) return (phi1 + phi2) % (2 * numpy.pi) -def chi_eff_from_spherical(mass1, mass2, spin1_a, spin1_polar, - spin2_a, spin2_polar): +def chi_eff_from_spherical(mass1, mass2, spin1_a, spin1_polar, spin2_a, spin2_polar): """Returns the effective spin using spins in spherical coordinates.""" spin1z = spin1_a * numpy.cos(spin1_polar) spin2z = spin2_a * numpy.cos(spin2_polar) return chi_eff(mass1, mass2, spin1z, spin2z) -def chi_p_from_spherical(mass1, mass2, spin1_a, spin1_azimuthal, spin1_polar, - spin2_a, spin2_azimuthal, spin2_polar): - """Returns the effective precession spin using spins in spherical +def chi_p_from_spherical( + mass1, + mass2, + spin1_a, + spin1_azimuthal, + spin1_polar, + spin2_a, + spin2_azimuthal, + spin2_polar, +): + """ + Returns the effective precession spin using spins in spherical coordinates. """ - spin1x, spin1y, _ = _spherical_to_cartesian( - spin1_a, spin1_azimuthal, spin1_polar) - spin2x, spin2y, _ = _spherical_to_cartesian( - spin2_a, spin2_azimuthal, spin2_polar) + spin1x, spin1y, _ = _spherical_to_cartesian(spin1_a, spin1_azimuthal, spin1_polar) + spin2x, spin2y, _ = _spherical_to_cartesian(spin2_a, spin2_azimuthal, spin2_polar) return chi_p(mass1, mass2, spin1x, spin1y, spin2x, spin2y) def primary_spin(mass1, mass2, spin1, spin2): """Returns the dimensionless spin of the primary mass.""" - mass1, mass2, spin1, spin2, input_is_array = ensurearray( - mass1, mass2, spin1, spin2) + mass1, mass2, spin1, spin2, input_is_array = ensurearray(mass1, mass2, spin1, spin2) sp = copy.copy(spin1) mask = mass1 < mass2 sp[mask] = spin2[mask] @@ -1039,8 +1126,7 @@ def primary_spin(mass1, mass2, spin1, spin2): def secondary_spin(mass1, mass2, spin1, spin2): """Returns the dimensionless spin of the secondary mass.""" - mass1, mass2, spin1, spin2, input_is_array = ensurearray( - mass1, mass2, spin1, spin2) + mass1, mass2, spin1, spin2, input_is_array = ensurearray(mass1, mass2, spin1, spin2) ss = copy.copy(spin2) mask = mass1 < mass2 ss[mask] = spin1[mask] @@ -1048,30 +1134,30 @@ def secondary_spin(mass1, mass2, spin1, spin2): def primary_xi(mass1, mass2, spin1x, spin1y, spin2x, spin2y): - """Returns the effective precession spin argument for the larger mass. - """ + """Returns the effective precession spin argument for the larger mass.""" spinx = primary_spin(mass1, mass2, spin1x, spin2x) spiny = primary_spin(mass1, mass2, spin1y, spin2y) return chi_perp_from_spinx_spiny(spinx, spiny) def secondary_xi(mass1, mass2, spin1x, spin1y, spin2x, spin2y): - """Returns the effective precession spin argument for the smaller mass. - """ + """Returns the effective precession spin argument for the smaller mass.""" spinx = secondary_spin(mass1, mass2, spin1x, spin2x) spiny = secondary_spin(mass1, mass2, spin1y, spin2y) return xi2_from_mass1_mass2_spin2x_spin2y(mass1, mass2, spinx, spiny) def xi1_from_spin1x_spin1y(spin1x, spin1y): - """Returns the effective precession spin argument for the larger mass. + """ + Returns the effective precession spin argument for the larger mass. This function assumes it's given spins of the primary mass. """ return chi_perp_from_spinx_spiny(spin1x, spin1y) def xi2_from_mass1_mass2_spin2x_spin2y(mass1, mass2, spin2x, spin2y): - """Returns the effective precession spin argument for the smaller mass. + """ + Returns the effective precession spin argument for the smaller mass. This function assumes it's given spins of the secondary mass. """ q = q_from_mass1_mass2(mass1, mass2) @@ -1081,13 +1167,13 @@ def xi2_from_mass1_mass2_spin2x_spin2y(mass1, mass2, spin2x, spin2y): def chi_perp_from_spinx_spiny(spinx, spiny): - """Returns the in-plane spin from the x/y components of the spin. - """ + """Returns the in-plane spin from the x/y components of the spin.""" return numpy.sqrt(spinx**2 + spiny**2) def chi_perp_from_mass1_mass2_xi2(mass1, mass2, xi2): - """Returns the in-plane spin from mass1, mass2, and xi2 for the + """ + Returns the in-plane spin from mass1, mass2, and xi2 for the secondary mass. """ q = q_from_mass1_mass2(mass1, mass2) @@ -1097,8 +1183,7 @@ def chi_perp_from_mass1_mass2_xi2(mass1, mass2, xi2): def chi_p_from_xi1_xi2(xi1, xi2): - """Returns effective precession spin from xi1 and xi2. - """ + """Returns effective precession spin from xi1 and xi2.""" xi1, xi2, input_is_array = ensurearray(xi1, xi2) chi_p = copy.copy(xi1) mask = xi1 < xi2 @@ -1107,70 +1192,66 @@ def chi_p_from_xi1_xi2(xi1, xi2): def phi1_from_phi_a_phi_s(phi_a, phi_s): - """Returns the angle between the x-component axis and the in-plane + """ + Returns the angle between the x-component axis and the in-plane spin for the primary mass from phi_s and phi_a. """ return (phi_s + phi_a) / 2.0 def phi2_from_phi_a_phi_s(phi_a, phi_s): - """Returns the angle between the x-component axis and the in-plane + """ + Returns the angle between the x-component axis and the in-plane spin for the secondary mass from phi_s and phi_a. """ return (phi_s - phi_a) / 2.0 def phi_from_spinx_spiny(spinx, spiny): - """Returns the angle between the x-component axis and the in-plane spin. - """ + """Returns the angle between the x-component axis and the in-plane spin.""" phi = numpy.arctan2(spiny, spinx) return phi % (2 * numpy.pi) def spin1z_from_mass1_mass2_chi_eff_chi_a(mass1, mass2, chi_eff, chi_a): - """Returns spin1z. - """ + """Returns spin1z.""" return (mass1 + mass2) / (2.0 * mass1) * (chi_eff - chi_a) def spin2z_from_mass1_mass2_chi_eff_chi_a(mass1, mass2, chi_eff, chi_a): - """Returns spin2z. - """ + """Returns spin2z.""" return (mass1 + mass2) / (2.0 * mass2) * (chi_eff + chi_a) def spin1x_from_xi1_phi_a_phi_s(xi1, phi_a, phi_s): - """Returns x-component spin for primary mass. - """ + """Returns x-component spin for primary mass.""" phi1 = phi1_from_phi_a_phi_s(phi_a, phi_s) return xi1 * numpy.cos(phi1) def spin1y_from_xi1_phi_a_phi_s(xi1, phi_a, phi_s): - """Returns y-component spin for primary mass. - """ + """Returns y-component spin for primary mass.""" phi1 = phi1_from_phi_a_phi_s(phi_s, phi_a) return xi1 * numpy.sin(phi1) def spin2x_from_mass1_mass2_xi2_phi_a_phi_s(mass1, mass2, xi2, phi_a, phi_s): - """Returns x-component spin for secondary mass. - """ + """Returns x-component spin for secondary mass.""" chi_perp = chi_perp_from_mass1_mass2_xi2(mass1, mass2, xi2) phi2 = phi2_from_phi_a_phi_s(phi_a, phi_s) return chi_perp * numpy.cos(phi2) def spin2y_from_mass1_mass2_xi2_phi_a_phi_s(mass1, mass2, xi2, phi_a, phi_s): - """Returns y-component spin for secondary mass. - """ + """Returns y-component spin for secondary mass.""" chi_perp = chi_perp_from_mass1_mass2_xi2(mass1, mass2, xi2) phi2 = phi2_from_phi_a_phi_s(phi_a, phi_s) return chi_perp * numpy.sin(phi2) def dquadmon_from_lambda(lambdav): - r"""Return the quadrupole moment of a neutron star given its lambda + r""" + Return the quadrupole moment of a neutron star given its lambda We use the relations defined here. https://arxiv.org/pdf/1302.4499.pdf. Note that the convention we use is that: @@ -1182,17 +1263,18 @@ def dquadmon_from_lambda(lambdav): Where :math:`\bar{Q}` (dimensionless) is the reduced quadrupole moment. """ ll = numpy.log(lambdav) - ai = .194 - bi = .0936 + ai = 0.194 + bi = 0.0936 ci = 0.0474 di = -4.21 * 10**-3.0 ei = 1.23 * 10**-4.0 - ln_quad_moment = ai + bi*ll + ci*ll**2.0 + di*ll**3.0 + ei*ll**4.0 + ln_quad_moment = ai + bi * ll + ci * ll**2.0 + di * ll**3.0 + ei * ll**4.0 return numpy.exp(ln_quad_moment) - 1 def spin_from_pulsar_freq(mass, radius, freq): - """Returns the dimensionless spin of a pulsar. + """ + Returns the dimensionless spin of a pulsar. Assumes the pulsar is a solid sphere when computing the moment of inertia. @@ -1204,10 +1286,11 @@ def spin_from_pulsar_freq(mass, radius, freq): The assumed radius of the pulsar, in kilometers. freq : float The spin frequency of the pulsar, in Hz. + """ omega = 2 * numpy.pi * freq mt = mass * MTSUN_SI - mominert = (2/5.) * mt * (radius * 1000 / C_SI)**2 + mominert = (2 / 5.0) * mt * (radius * 1000 / C_SI) ** 2 return mominert * omega / mt**2 @@ -1219,20 +1302,21 @@ def spin_from_pulsar_freq(mass, radius, freq): # ============================================================================= # def chirp_distance(dist, mchirp, ref_mass=1.4): - """Returns the chirp distance given the luminosity distance and chirp mass. - """ - return dist * (2.**(-1./5) * ref_mass / mchirp)**(5./6) + """Returns the chirp distance given the luminosity distance and chirp mass.""" + return dist * (2.0 ** (-1.0 / 5) * ref_mass / mchirp) ** (5.0 / 6) def distance_from_chirp_distance_mchirp(chirp_distance, mchirp, ref_mass=1.4): - """Returns the luminosity distance given a chirp distance and chirp mass. - """ - return chirp_distance * (2.**(-1./5) * ref_mass / mchirp)**(-5./6) + """Returns the luminosity distance given a chirp distance and chirp mass.""" + return chirp_distance * (2.0 ** (-1.0 / 5) * ref_mass / mchirp) ** (-5.0 / 6) _detector_cache = {} -def det_tc(detector_name, ra, dec, tc, ref_frame='geocentric', relative=False): - """Returns the coalescence time of a signal in the given detector. + + +def det_tc(detector_name, ra, dec, tc, ref_frame="geocentric", relative=False): + """ + Returns the coalescence time of a signal in the given detector. Parameters ---------- @@ -1252,6 +1336,7 @@ def det_tc(detector_name, ra, dec, tc, ref_frame='geocentric', relative=False): ------- float : The GPS time of the coalescence in detector `detector_name`. + """ ref_time = tc if relative: @@ -1259,28 +1344,31 @@ def det_tc(detector_name, ra, dec, tc, ref_frame='geocentric', relative=False): if ref_frame == detector_name: return tc - if detector_name == 'geocentric': + if detector_name == "geocentric": refdet = Detector(ref_frame) return tc - refdet.time_delay_from_earth_center(ra, dec, ref_time) if detector_name not in _detector_cache: _detector_cache[detector_name] = Detector(detector_name) detector = _detector_cache[detector_name] - if ref_frame == 'geocentric': + if ref_frame == "geocentric": return tc + detector.time_delay_from_earth_center(ra, dec, ref_time) - else: - other = Detector(ref_frame) - return tc + detector.time_delay_from_detector(other, ra, dec, ref_time) + other = Detector(ref_frame) + return tc + detector.time_delay_from_detector(other, ra, dec, ref_time) -def optimal_orientation_from_detector(detector_name, tc): - """ Low-level function to be called from _optimal_dec_from_detector - and _optimal_ra_from_detector""" +def optimal_orientation_from_detector(detector_name, tc): + """ + Low-level function to be called from _optimal_dec_from_detector + and _optimal_ra_from_detector + """ d = Detector(detector_name) ra, dec = d.optimal_orientation(tc) return ra, dec + def optimal_dec_from_detector(detector_name, tc): - """For a given detector and GPS time, return the optimal orientation + """ + For a given detector and GPS time, return the optimal orientation (directly overhead of the detector) in declination. @@ -1295,11 +1383,14 @@ def optimal_dec_from_detector(detector_name, tc): ------- float : The declination of the signal, in radians. + """ return optimal_orientation_from_detector(detector_name, tc)[1] + def optimal_ra_from_detector(detector_name, tc): - """For a given detector and GPS time, return the optimal orientation + """ + For a given detector and GPS time, return the optimal orientation (directly overhead of the detector) in right ascension. Parameters @@ -1313,6 +1404,7 @@ def optimal_ra_from_detector(detector_name, tc): ------- float : The declination of the signal, in radians. + """ return optimal_orientation_from_detector(detector_name, tc)[0] @@ -1325,7 +1417,8 @@ def optimal_ra_from_detector(detector_name, tc): # ============================================================================= # def snr_from_loglr(loglr): - """Returns SNR computed from the given log likelihood ratio(s). This is + """ + Returns SNR computed from the given log likelihood ratio(s). This is defined as `sqrt(2*loglr)`.If the log likelihood ratio is < 0, returns 0. Parameters @@ -1337,18 +1430,20 @@ def snr_from_loglr(loglr): ------- array or float The SNRs computed from the log likelihood ratios. + """ singleval = isinstance(loglr, float) if singleval: loglr = numpy.array([loglr]) # temporarily quiet sqrt(-1) warnings with numpy.errstate(invalid="ignore"): - snrs = numpy.sqrt(2*loglr) - snrs[numpy.isnan(snrs)] = 0. + snrs = numpy.sqrt(2 * loglr) + snrs[numpy.isnan(snrs)] = 0.0 if singleval: snrs = snrs[0] return snrs + # # ============================================================================= # @@ -1358,8 +1453,9 @@ def snr_from_loglr(loglr): # -def get_lm_f0tau(mass, spin, l, m, n=0, which='both'): - """Return the f0 and the tau for one or more overtones of an l, m mode. +def get_lm_f0tau(mass, spin, l, m, n=0, which="both"): + """ + Return the f0 and the tau for one or more overtones of an l, m mode. Parameters ---------- @@ -1386,14 +1482,14 @@ def get_lm_f0tau(mass, spin, l, m, n=0, which='both'): tau : float or array Returned if ``which`` is 'both' or 'tau'. The damping time of the QNM(s), in seconds. + """ # convert to arrays - mass, spin, l, m, n, input_is_array = ensurearray( - mass, spin, l, m, n) + mass, spin, l, m, n, input_is_array = ensurearray(mass, spin, l, m, n) # we'll ravel the arrays so we can evaluate each parameter combination # one at a a time - getf0 = which == 'both' or which == 'f0' - gettau = which == 'both' or which == 'tau' + getf0 = which == "both" or which == "f0" + gettau = which == "both" or which == "tau" out = [] if getf0: f0s = pykerr.qnmfreq(mass, spin, l, m, n) @@ -1407,7 +1503,8 @@ def get_lm_f0tau(mass, spin, l, m, n=0, which='both'): def get_lm_f0tau_allmodes(mass, spin, modes): - """Returns a dictionary of all of the frequencies and damping times for the + """ + Returns a dictionary of all of the frequencies and damping times for the requested modes. Parameters @@ -1432,10 +1529,11 @@ def get_lm_f0tau_allmodes(mass, spin, modes): tau : dict Dictionary mapping the modes to the damping times. The keys are the same as ``f0``. + """ f0, tau = {}, {} for lmn in modes: - key = '{}{}{}' + key = "{}{}{}" l, m, nmodes = int(lmn[0]), int(lmn[1]), int(lmn[2]) for n in range(nmodes): tmp_f0, tmp_tau = get_lm_f0tau(mass, spin, l, m, n) @@ -1445,7 +1543,8 @@ def get_lm_f0tau_allmodes(mass, spin, modes): def freq_from_final_mass_spin(final_mass, final_spin, l=2, m=2, n=0): - """Returns QNM frequency for the given mass and spin and mode. + """ + Returns QNM frequency for the given mass and spin and mode. Parameters ---------- @@ -1465,12 +1564,14 @@ def freq_from_final_mass_spin(final_mass, final_spin, l=2, m=2, n=0): ------- float or array The frequency of the QNM(s), in Hz. + """ - return get_lm_f0tau(final_mass, final_spin, l, m, n=n, which='f0') + return get_lm_f0tau(final_mass, final_spin, l, m, n=n, which="f0") def tau_from_final_mass_spin(final_mass, final_spin, l=2, m=2, n=0): - """Returns QNM damping time for the given mass and spin and mode. + """ + Returns QNM damping time for the given mass and spin and mode. Parameters ---------- @@ -1490,8 +1591,9 @@ def tau_from_final_mass_spin(final_mass, final_spin, l=2, m=2, n=0): ------- float or array The damping time of the QNM(s), in seconds. + """ - return get_lm_f0tau(final_mass, final_spin, l, m, n=n, which='tau') + return get_lm_f0tau(final_mass, final_spin, l, m, n=n, which="tau") # The following are from Table VIII, IX, X of Berti et al., @@ -1503,18 +1605,19 @@ def tau_from_final_mass_spin(final_mass, final_spin, l=2, m=2, n=0): (2, 1): (-0.3, 2.3561, -0.2277), (3, 3): (0.9, 2.343, -0.4810), (4, 4): (1.1929, 3.1191, -0.4825), - } +} _berti_mass_constants = { (2, 2): (1.5251, -1.1568, 0.1292), (2, 1): (0.6, -0.2339, 0.4175), (3, 3): (1.8956, -1.3043, 0.1818), (4, 4): (2.3, -1.5056, 0.2244), - } +} def final_spin_from_f0_tau(f0, tau, l=2, m=2): - """Returns the final spin based on the given frequency and damping time. + """ + Returns the final spin based on the given frequency and damping time. .. note:: Currently, only (l,m) = (2,2), (3,3), (4,4), (2,1) are supported. @@ -1537,10 +1640,11 @@ def final_spin_from_f0_tau(f0, tau, l=2, m=2): The spin of the final black hole. If the combination of frequency and damping times give an unphysical result, ``numpy.nan`` will be returned. + """ f0, tau, input_is_array = ensurearray(f0, tau) # from Berti et al. 2006 - a, b, c = _berti_spin_constants[l,m] + a, b, c = _berti_spin_constants[l, m] origshape = f0.shape # flatten inputs for storing results f0 = f0.ravel() @@ -1549,7 +1653,7 @@ def final_spin_from_f0_tau(f0, tau, l=2, m=2): for ii in range(spins.size): Q = f0[ii] * tau[ii] * numpy.pi try: - s = 1. - ((Q-a)/b)**(1./c) + s = 1.0 - ((Q - a) / b) ** (1.0 / c) except ValueError: s = numpy.nan spins[ii] = s @@ -1558,7 +1662,8 @@ def final_spin_from_f0_tau(f0, tau, l=2, m=2): def final_mass_from_f0_tau(f0, tau, l=2, m=2): - """Returns the final mass (in solar masses) based on the given frequency + """ + Returns the final mass (in solar masses) based on the given frequency and damping time. .. note:: @@ -1582,14 +1687,17 @@ def final_mass_from_f0_tau(f0, tau, l=2, m=2): The mass of the final black hole. If the combination of frequency and damping times give an unphysical result, ``numpy.nan`` will be returned. + """ # from Berti et al. 2006 spin = final_spin_from_f0_tau(f0, tau, l=l, m=m) - a, b, c = _berti_mass_constants[l,m] - return (a + b*(1-spin)**c)/(2*numpy.pi*f0*MTSUN_SI) + a, b, c = _berti_mass_constants[l, m] + return (a + b * (1 - spin) ** c) / (2 * numpy.pi * f0 * MTSUN_SI) + def freqlmn_from_other_lmn(f0, tau, current_l, current_m, new_l, new_m): - """Returns the QNM frequency (in Hz) of a chosen new (l,m) mode from the + """ + Returns the QNM frequency (in Hz) of a chosen new (l,m) mode from the given current (l,m) mode. Parameters @@ -1614,6 +1722,7 @@ def freqlmn_from_other_lmn(f0, tau, current_l, current_m, new_l, new_m): frequency and damping time provided for the current (l, m) QNM mode correspond to an unphysical Kerr black hole mass and/or spin, ``numpy.nan`` will be returned. + """ mass = final_mass_from_f0_tau(f0, tau, l=current_l, m=current_m) spin = final_spin_from_f0_tau(f0, tau, l=current_l, m=current_m) @@ -1627,7 +1736,8 @@ def freqlmn_from_other_lmn(f0, tau, current_l, current_m, new_l, new_m): def taulmn_from_other_lmn(f0, tau, current_l, current_m, new_l, new_m): - """Returns the QNM damping time (in seconds) of a chosen new (l,m) mode + """ + Returns the QNM damping time (in seconds) of a chosen new (l,m) mode from the given current (l,m) mode. Parameters @@ -1652,6 +1762,7 @@ def taulmn_from_other_lmn(f0, tau, current_l, current_m, new_l, new_m): frequency and damping time provided for the current (l, m) QNM mode correspond to an unphysical Kerr black hole mass and/or spin, ``numpy.nan`` will be returned. + """ mass = final_mass_from_f0_tau(f0, tau, l=current_l, m=current_m) spin = final_spin_from_f0_tau(f0, tau, l=current_l, m=current_m) @@ -1663,10 +1774,21 @@ def taulmn_from_other_lmn(f0, tau, current_l, current_m, new_l, new_m): new_tau = tau_from_final_mass_spin(mass, spin, l=new_l, m=new_m) return formatreturn(new_tau, input_is_array) -def get_final_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., - spin2x=0., spin2y=0., spin2z=0., - approximant='SEOBNRv4PHM', f_ref=-1): - """Estimates the final mass and spin from the given initial parameters. + +def get_final_from_initial( + mass1, + mass2, + spin1x=0.0, + spin1y=0.0, + spin1z=0.0, + spin2x=0.0, + spin2y=0.0, + spin2z=0.0, + approximant="SEOBNRv4PHM", + f_ref=-1, +): + """ + Estimates the final mass and spin from the given initial parameters. This uses the fits used by either the NRSur7dq4 or EOBNR models for converting from initial parameters to final, depending on the @@ -1709,6 +1831,7 @@ def get_final_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., The final mass, in solar masses. final_spin : float The dimensionless final spin. + """ args = (mass1, mass2, spin1x, spin1y, spin1z, spin2x, spin2y, spin2z) args = ensurearray(*args) @@ -1724,41 +1847,60 @@ def get_final_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., m2 = float(mass2[ii]) spin1 = list(map(float, [spin1x[ii], spin1y[ii], spin1z[ii]])) spin2 = list(map(float, [spin2x[ii], spin2y[ii], spin2z[ii]])) - if approximant == 'NRSur7dq4': + if approximant == "NRSur7dq4": from lalsimulation import nrfits + try: - res = nrfits.eval_nrfit(m1*MSUN_SI, - m2*MSUN_SI, - spin1, spin2, 'NRSur7dq4Remnant', - ['FinalMass', 'FinalSpin'], - f_ref=f_ref) + res = nrfits.eval_nrfit( + m1 * MSUN_SI, + m2 * MSUN_SI, + spin1, + spin2, + "NRSur7dq4Remnant", + ["FinalMass", "FinalSpin"], + f_ref=f_ref, + ) except RuntimeError: continue - final_mass[ii] = res['FinalMass'][0] / MSUN_SI - sf = res['FinalSpin'] - final_spin[ii] = (sf**2).sum()**0.5 + final_mass[ii] = res["FinalMass"][0] / MSUN_SI + sf = res["FinalSpin"] + final_spin[ii] = (sf**2).sum() ** 0.5 if sf[-1] < 0: final_spin[ii] *= -1 - elif approximant == 'SEOBNRv4': + elif approximant == "SEOBNRv4": _, fm, fs = lalsim.SimIMREOBFinalMassSpin( - m1, m2, spin1, spin2, getattr(lalsim, approximant)) + m1, m2, spin1, spin2, getattr(lalsim, approximant) + ) final_mass[ii] = fm * (m1 + m2) final_spin[ii] = fs else: _, fm, fs = lalsim.SimIMREOBFinalMassSpinPrec( - m1, m2, spin1, spin2, getattr(lalsim, approximant)) + m1, m2, spin1, spin2, getattr(lalsim, approximant) + ) final_mass[ii] = fm * (m1 + m2) final_spin[ii] = fs final_mass = final_mass.reshape(origshape) final_spin = final_spin.reshape(origshape) - return (formatreturn(final_mass, input_is_array), - formatreturn(final_spin, input_is_array)) + return ( + formatreturn(final_mass, input_is_array), + formatreturn(final_spin, input_is_array), + ) -def final_mass_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., - spin2x=0., spin2y=0., spin2z=0., - approximant='SEOBNRv4PHM', f_ref=-1): - """Estimates the final mass from the given initial parameters. +def final_mass_from_initial( + mass1, + mass2, + spin1x=0.0, + spin1y=0.0, + spin1z=0.0, + spin2x=0.0, + spin2y=0.0, + spin2z=0.0, + approximant="SEOBNRv4PHM", + f_ref=-1, +): + """ + Estimates the final mass from the given initial parameters. This uses the fits used by either the NRSur7dq4 or EOBNR models for converting from initial parameters to final, depending on the @@ -1799,16 +1941,36 @@ def final_mass_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., ------- float The final mass, in solar masses. - """ - return get_final_from_initial(mass1, mass2, spin1x, spin1y, spin1z, - spin2x, spin2y, spin2z, approximant, - f_ref=f_ref)[0] - -def final_spin_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., - spin2x=0., spin2y=0., spin2z=0., - approximant='SEOBNRv4PHM', f_ref=-1): - """Estimates the final spin from the given initial parameters. + """ + return get_final_from_initial( + mass1, + mass2, + spin1x, + spin1y, + spin1z, + spin2x, + spin2y, + spin2z, + approximant, + f_ref=f_ref, + )[0] + + +def final_spin_from_initial( + mass1, + mass2, + spin1x=0.0, + spin1y=0.0, + spin1z=0.0, + spin2x=0.0, + spin2y=0.0, + spin2z=0.0, + approximant="SEOBNRv4PHM", + f_ref=-1, +): + """ + Estimates the final spin from the given initial parameters. This uses the fits used by either the NRSur7dq4 or EOBNR models for converting from initial parameters to final, depending on the @@ -1849,10 +2011,20 @@ def final_spin_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., ------- float The dimensionless final spin. + """ - return get_final_from_initial(mass1, mass2, spin1x, spin1y, spin1z, - spin2x, spin2y, spin2z, approximant, - f_ref=f_ref)[1] + return get_final_from_initial( + mass1, + mass2, + spin1x, + spin1y, + spin1z, + spin2x, + spin2y, + spin2z, + approximant, + f_ref=f_ref, + )[1] # @@ -1863,8 +2035,10 @@ def final_spin_from_initial(mass1, mass2, spin1x=0., spin1y=0., spin1z=0., # ============================================================================= # + def velocity_to_frequency(v, M): - """ Calculate the gravitational-wave frequency from the + """ + Calculate the gravitational-wave frequency from the total mass and invariant velocity. Parameters @@ -1878,11 +2052,14 @@ def velocity_to_frequency(v, M): ------- f : float Gravitational-wave frequency + """ - return v**(3.0) / (M * MTSUN_SI * PI) + return v ** (3.0) / (M * MTSUN_SI * PI) + def frequency_to_velocity(f, M): - """ Calculate the invariant velocity from the total + """ + Calculate the invariant velocity from the total mass and gravitational-wave frequency. Parameters @@ -1896,8 +2073,9 @@ def frequency_to_velocity(f, M): ------- v : float or numpy.array Invariant velocity + """ - return (PI * M * MTSUN_SI * f)**(1.0/3.0) + return (PI * M * MTSUN_SI * f) ** (1.0 / 3.0) def f_schwarzchild_isco(M): @@ -1914,8 +2092,9 @@ def f_schwarzchild_isco(M): ------- f : float or numpy.array Frequency in Hz + """ - return velocity_to_frequency((1.0/6.0)**(0.5), M) + return velocity_to_frequency((1.0 / 6.0) ** (0.5), M) # @@ -1926,8 +2105,10 @@ def f_schwarzchild_isco(M): # ============================================================================ # + def nltides_coefs(amplitude, n, m1, m2): - """Calculate the coefficents needed to compute the + """ + Calculate the coefficents needed to compute the shift in t(f) and phi(f) due to non-linear tides. Parameters @@ -1949,8 +2130,8 @@ def nltides_coefs(amplitude, n, m1, m2): The constant factor needed to compute t(f) phi_of_f_factor: float The constant factor needed to compute phi(f) - """ + """ # Use 100.0 Hz as a reference frequency f_ref = 100.0 @@ -1959,17 +2140,17 @@ def nltides_coefs(amplitude, n, m1, m2): mc *= MSUN_SI # Calculate constants in phasing - a = (96./5.) * \ - (G_SI * PI * mc * f_ref / C_SI**3.)**(5./3.) - b = 6. * amplitude - t_of_f_factor = -1./(PI*f_ref) * b/(a*a * (n-4.)) - phi_of_f_factor = -2.*b / (a*a * (n-3.)) + a = (96.0 / 5.0) * (G_SI * PI * mc * f_ref / C_SI**3.0) ** (5.0 / 3.0) + b = 6.0 * amplitude + t_of_f_factor = -1.0 / (PI * f_ref) * b / (a * a * (n - 4.0)) + phi_of_f_factor = -2.0 * b / (a * a * (n - 3.0)) return f_ref, t_of_f_factor, phi_of_f_factor def nltides_gw_phase_difference(f, f0, amplitude, n, m1, m2): - """Calculate the gravitational-wave phase shift bwtween + """ + Calculate the gravitational-wave phase shift bwtween f and f_coalescence = infinity due to non-linear tides. To compute the phase shift between e.g. f_low and f_isco, call this function twice and compute the difference. @@ -1993,25 +2174,28 @@ def nltides_gw_phase_difference(f, f0, amplitude, n, m1, m2): ------- delta_phi: float or numpy.array Phase in radians + """ f, f0, amplitude, n, m1, m2, input_is_array = ensurearray( - f, f0, amplitude, n, m1, m2) + f, f0, amplitude, n, m1, m2 + ) delta_phi = numpy.zeros(m1.shape) f_ref, _, phi_of_f_factor = nltides_coefs(amplitude, n, m1, m2) mask = f <= f0 - delta_phi[mask] = - phi_of_f_factor[mask] * (f0[mask]/f_ref)**(n[mask]-3.) + delta_phi[mask] = -phi_of_f_factor[mask] * (f0[mask] / f_ref) ** (n[mask] - 3.0) mask = f > f0 - delta_phi[mask] = - phi_of_f_factor[mask] * (f[mask]/f_ref)**(n[mask]-3.) + delta_phi[mask] = -phi_of_f_factor[mask] * (f[mask] / f_ref) ** (n[mask] - 3.0) return formatreturn(delta_phi, input_is_array) def nltides_gw_phase_diff_isco(f_low, f0, amplitude, n, m1, m2): - """Calculate the gravitational-wave phase shift bwtween + """ + Calculate the gravitational-wave phase shift bwtween f_low and f_isco due to non-linear tides. Parameters @@ -2035,65 +2219,103 @@ def nltides_gw_phase_diff_isco(f_low, f0, amplitude, n, m1, m2): ------- delta_phi: float or numpy.array Phase in radians + """ - f0, amplitude, n, m1, m2, input_is_array = ensurearray( - f0, amplitude, n, m1, m2) + f0, amplitude, n, m1, m2, input_is_array = ensurearray(f0, amplitude, n, m1, m2) f_low = numpy.zeros(m1.shape) + f_low - phi_l = nltides_gw_phase_difference( - f_low, f0, amplitude, n, m1, m2) + phi_l = nltides_gw_phase_difference(f_low, f0, amplitude, n, m1, m2) - f_isco = f_schwarzchild_isco(m1+m2) + f_isco = f_schwarzchild_isco(m1 + m2) - phi_i = nltides_gw_phase_difference( - f_isco, f0, amplitude, n, m1, m2) + phi_i = nltides_gw_phase_difference(f_isco, f0, amplitude, n, m1, m2) return formatreturn(phi_i - phi_l, input_is_array) -__all__ = ['eccmchirp_from_mchirp_eccentricity', - 'mchirp_from_eccmchirp_eccentricity', - 'eccmchirp_from_mass1_mass2_eccentricity', - 'dquadmon_from_lambda', 'lambda_tilde', - 'lambda_from_mass_tov_file', 'primary_mass', - 'secondary_mass', 'mtotal_from_mass1_mass2', - 'q_from_mass1_mass2', 'invq_from_mass1_mass2', - 'eta_from_mass1_mass2', 'mchirp_from_mass1_mass2', - 'mass1_from_mtotal_q', 'mass2_from_mtotal_q', - 'mass1_from_mtotal_eta', 'mass2_from_mtotal_eta', - 'mtotal_from_mchirp_eta', 'mass1_from_mchirp_eta', - 'mass2_from_mchirp_eta', 'mass2_from_mchirp_mass1', - 'mass_from_knownmass_eta', 'mass2_from_mass1_eta', - 'mass1_from_mass2_eta', 'eta_from_q', 'mass1_from_mchirp_q', - 'mass2_from_mchirp_q', 'tau0_from_mtotal_eta', - 'tau3_from_mtotal_eta', 'tau0_from_mass1_mass2', - 'tau0_from_mchirp', 'mchirp_from_tau0', - 'tau3_from_mass1_mass2', 'mtotal_from_tau0_tau3', - 'eta_from_tau0_tau3', 'mass1_from_tau0_tau3', - 'mass2_from_tau0_tau3', 'primary_spin', 'secondary_spin', - 'chi_eff', 'chi_a', 'chi_p', 'phi_a', 'phi_s', - 'primary_xi', 'secondary_xi', - 'xi1_from_spin1x_spin1y', 'xi2_from_mass1_mass2_spin2x_spin2y', - 'chi_perp_from_spinx_spiny', 'chi_perp_from_mass1_mass2_xi2', - 'chi_p_from_xi1_xi2', 'phi_from_spinx_spiny', - 'phi1_from_phi_a_phi_s', 'phi2_from_phi_a_phi_s', - 'spin1z_from_mass1_mass2_chi_eff_chi_a', - 'spin2z_from_mass1_mass2_chi_eff_chi_a', - 'spin1x_from_xi1_phi_a_phi_s', 'spin1y_from_xi1_phi_a_phi_s', - 'spin2x_from_mass1_mass2_xi2_phi_a_phi_s', - 'spin2y_from_mass1_mass2_xi2_phi_a_phi_s', - 'chirp_distance', 'det_tc', 'snr_from_loglr', - 'freq_from_final_mass_spin', 'tau_from_final_mass_spin', - 'final_spin_from_f0_tau', 'final_mass_from_f0_tau', - 'final_mass_from_initial', 'final_spin_from_initial', - 'optimal_dec_from_detector', 'optimal_ra_from_detector', - 'chi_eff_from_spherical', 'chi_p_from_spherical', - 'nltides_gw_phase_diff_isco', 'spin_from_pulsar_freq', - 'freqlmn_from_other_lmn', 'taulmn_from_other_lmn', - 'remnant_mass_from_mass1_mass2_spherical_spin_eos', - 'remnant_mass_from_mass1_mass2_cartesian_spin_eos', - 'lambda1_from_delta_lambda_tilde_lambda_tilde', - 'lambda2_from_delta_lambda_tilde_lambda_tilde', - 'delta_lambda_tilde', 'hypertriangle' - ] +__all__ = [ + "chi_a", + "chi_eff", + "chi_eff_from_spherical", + "chi_p", + "chi_p_from_spherical", + "chi_p_from_xi1_xi2", + "chi_perp_from_mass1_mass2_xi2", + "chi_perp_from_spinx_spiny", + "chirp_distance", + "delta_lambda_tilde", + "det_tc", + "dquadmon_from_lambda", + "eccmchirp_from_mass1_mass2_eccentricity", + "eccmchirp_from_mchirp_eccentricity", + "eta_from_mass1_mass2", + "eta_from_q", + "eta_from_tau0_tau3", + "final_mass_from_f0_tau", + "final_mass_from_initial", + "final_spin_from_f0_tau", + "final_spin_from_initial", + "freq_from_final_mass_spin", + "freqlmn_from_other_lmn", + "hypertriangle", + "invq_from_mass1_mass2", + "lambda1_from_delta_lambda_tilde_lambda_tilde", + "lambda2_from_delta_lambda_tilde_lambda_tilde", + "lambda_from_mass_tov_file", + "lambda_tilde", + "mass1_from_mass2_eta", + "mass1_from_mchirp_eta", + "mass1_from_mchirp_q", + "mass1_from_mtotal_eta", + "mass1_from_mtotal_q", + "mass1_from_tau0_tau3", + "mass2_from_mass1_eta", + "mass2_from_mchirp_eta", + "mass2_from_mchirp_mass1", + "mass2_from_mchirp_q", + "mass2_from_mtotal_eta", + "mass2_from_mtotal_q", + "mass2_from_tau0_tau3", + "mass_from_knownmass_eta", + "mchirp_from_eccmchirp_eccentricity", + "mchirp_from_mass1_mass2", + "mchirp_from_tau0", + "mtotal_from_mass1_mass2", + "mtotal_from_mchirp_eta", + "mtotal_from_tau0_tau3", + "nltides_gw_phase_diff_isco", + "optimal_dec_from_detector", + "optimal_ra_from_detector", + "phi1_from_phi_a_phi_s", + "phi2_from_phi_a_phi_s", + "phi_a", + "phi_from_spinx_spiny", + "phi_s", + "primary_mass", + "primary_spin", + "primary_xi", + "q_from_mass1_mass2", + "remnant_mass_from_mass1_mass2_cartesian_spin_eos", + "remnant_mass_from_mass1_mass2_spherical_spin_eos", + "secondary_mass", + "secondary_spin", + "secondary_xi", + "snr_from_loglr", + "spin1x_from_xi1_phi_a_phi_s", + "spin1y_from_xi1_phi_a_phi_s", + "spin1z_from_mass1_mass2_chi_eff_chi_a", + "spin2x_from_mass1_mass2_xi2_phi_a_phi_s", + "spin2y_from_mass1_mass2_xi2_phi_a_phi_s", + "spin2z_from_mass1_mass2_chi_eff_chi_a", + "spin_from_pulsar_freq", + "tau0_from_mass1_mass2", + "tau0_from_mchirp", + "tau0_from_mtotal_eta", + "tau3_from_mass1_mass2", + "tau3_from_mtotal_eta", + "tau_from_final_mass_spin", + "taulmn_from_other_lmn", + "xi1_from_spin1x_spin1y", + "xi2_from_mass1_mass2_spin2x_spin2y", +] diff --git a/pycbc/coordinates/__init__.py b/pycbc/coordinates/__init__.py index e1d7c81d90d..78b80a98afa 100644 --- a/pycbc/coordinates/__init__.py +++ b/pycbc/coordinates/__init__.py @@ -21,17 +21,28 @@ from pycbc.coordinates.base import * from pycbc.coordinates.space import * - -__all__ = ['cartesian_to_spherical_rho', 'cartesian_to_spherical_azimuthal', - 'cartesian_to_spherical_polar', 'cartesian_to_spherical', - 'spherical_to_cartesian', - 'TIME_OFFSET_20_DEGREES', - 'localization_to_propagation_vector', - 'propagation_vector_to_localization', 'polarization_newframe', - 't_lisa_from_ssb', 't_ssb_from_t_lisa', - 'ssb_to_lisa', 'lisa_to_ssb', - 'rotation_matrix_ssb_to_lisa', 'rotation_matrix_ssb_to_geo', - 'lisa_position_ssb', 'earth_position_ssb', - 't_geo_from_ssb', 't_ssb_from_t_geo', 'ssb_to_geo', 'geo_to_ssb', - 'lisa_to_geo', 'geo_to_lisa', - ] +__all__ = [ + "TIME_OFFSET_20_DEGREES", + "cartesian_to_spherical", + "cartesian_to_spherical_azimuthal", + "cartesian_to_spherical_polar", + "cartesian_to_spherical_rho", + "earth_position_ssb", + "geo_to_lisa", + "geo_to_ssb", + "lisa_position_ssb", + "lisa_to_geo", + "lisa_to_ssb", + "localization_to_propagation_vector", + "polarization_newframe", + "propagation_vector_to_localization", + "rotation_matrix_ssb_to_geo", + "rotation_matrix_ssb_to_lisa", + "spherical_to_cartesian", + "ssb_to_geo", + "ssb_to_lisa", + "t_geo_from_ssb", + "t_lisa_from_ssb", + "t_ssb_from_t_geo", + "t_ssb_from_t_lisa", +] diff --git a/pycbc/coordinates/base.py b/pycbc/coordinates/base.py index 14e79d9197a..67585c58124 100644 --- a/pycbc/coordinates/base.py +++ b/pycbc/coordinates/base.py @@ -27,14 +27,17 @@ Base coordinate transformations, this module provides transformations between cartesian and spherical coordinates. """ + import logging + import numpy -logger = logging.getLogger('pycbc.coordinates.base') +logger = logging.getLogger("pycbc.coordinates.base") def cartesian_to_spherical_rho(x, y, z): - """ Calculates the magnitude in spherical coordinates from Cartesian + """ + Calculates the magnitude in spherical coordinates from Cartesian coordinates. Parameters @@ -50,12 +53,14 @@ def cartesian_to_spherical_rho(x, y, z): ------- rho : {numpy.array, float} The radial amplitude. + """ return numpy.sqrt(x**2 + y**2 + z**2) def cartesian_to_spherical_azimuthal(x, y): - """ Calculates the azimuthal angle in spherical coordinates from Cartesian + """ + Calculates the azimuthal angle in spherical coordinates from Cartesian coordinates. The azimuthal angle is in [0,2*pi]. Parameters @@ -69,6 +74,7 @@ def cartesian_to_spherical_azimuthal(x, y): ------- phi : {numpy.array, float} The azimuthal angle. + """ y = float(y) if isinstance(y, int) else y phi = numpy.arctan2(y, x) @@ -76,7 +82,8 @@ def cartesian_to_spherical_azimuthal(x, y): def cartesian_to_spherical_polar(x, y, z): - """ Calculates the polar angle in spherical coordinates from Cartesian + """ + Calculates the polar angle in spherical coordinates from Cartesian coordinates. The polar angle is in [0,pi]. Parameters @@ -92,17 +99,19 @@ def cartesian_to_spherical_polar(x, y, z): ------- theta : {numpy.array, float} The polar angle. + """ rho = cartesian_to_spherical_rho(x, y, z) if numpy.isscalar(rho): return numpy.arccos(z / rho) if rho else 0.0 - else: - return numpy.arccos(numpy.divide(z, rho, out=numpy.ones_like(z), - where=rho != 0)) + return numpy.arccos( + numpy.divide(z, rho, out=numpy.ones_like(z), where=rho != 0) + ) def cartesian_to_spherical(x, y, z): - """ Maps cartesian coordinates (x,y,z) to spherical coordinates + """ + Maps cartesian coordinates (x,y,z) to spherical coordinates (rho,phi,theta) where phi is in [0,2*pi] and theta is in [0,pi]. Parameters @@ -122,6 +131,7 @@ def cartesian_to_spherical(x, y, z): The azimuthal angle. theta : {numpy.array, float} The polar angle. + """ rho = cartesian_to_spherical_rho(x, y, z) phi = cartesian_to_spherical_azimuthal(x, y) @@ -130,7 +140,8 @@ def cartesian_to_spherical(x, y, z): def spherical_to_cartesian(rho, phi, theta): - """ Maps spherical coordinates (rho,phi,theta) to cartesian coordinates + """ + Maps spherical coordinates (rho,phi,theta) to cartesian coordinates (x,y,z) where phi is in [0,2*pi] and theta is in [0,pi]. Parameters @@ -150,6 +161,7 @@ def spherical_to_cartesian(rho, phi, theta): Y-coordinate. z : {numpy.array, float} Z-coordinate. + """ x = rho * numpy.cos(phi) * numpy.sin(theta) y = rho * numpy.sin(phi) * numpy.sin(theta) @@ -157,7 +169,10 @@ def spherical_to_cartesian(rho, phi, theta): return x, y, z -__all__ = ['cartesian_to_spherical_rho', 'cartesian_to_spherical_azimuthal', - 'cartesian_to_spherical_polar', 'cartesian_to_spherical', - 'spherical_to_cartesian', - ] +__all__ = [ + "cartesian_to_spherical", + "cartesian_to_spherical_azimuthal", + "cartesian_to_spherical_polar", + "cartesian_to_spherical_rho", + "spherical_to_cartesian", +] diff --git a/pycbc/coordinates/space.py b/pycbc/coordinates/space.py index f72031463bd..2a5c5b75010 100644 --- a/pycbc/coordinates/space.py +++ b/pycbc/coordinates/space.py @@ -30,19 +30,22 @@ """ import logging -import numpy as np -from scipy.spatial.transform import Rotation -from scipy.optimize import fsolve +import numpy as np from astropy import units -from astropy.constants import c, au -from astropy.time import Time -from astropy.coordinates import BarycentricMeanEcliptic, PrecessedGeocentric -from astropy.coordinates import get_body_barycentric -from astropy.coordinates import SkyCoord +from astropy.constants import au, c +from astropy.coordinates import ( + BarycentricMeanEcliptic, + PrecessedGeocentric, + SkyCoord, + get_body_barycentric, +) from astropy.coordinates.builtin_frames import ecliptic_transforms +from astropy.time import Time +from scipy.optimize import fsolve +from scipy.spatial.transform import Rotation -logger = logging.getLogger('pycbc.coordinates.space') +logger = logging.getLogger("pycbc.coordinates.space") # This constant makes sure LISA is behind the Earth by 19-23 degrees. # Making this a stand-alone constant will also make it callable by @@ -54,7 +57,8 @@ def rotation_matrix_ssb_to_lisa(alpha): - """ The rotation matrix (of frame basis) from SSB frame to LISA frame. + """ + The rotation matrix (of frame basis) from SSB frame to LISA frame. This function assumes the angle between LISA plane and the ecliptic is 60 degrees, and the period of LISA's self-rotation and orbital revolution is both one year. @@ -69,19 +73,19 @@ def rotation_matrix_ssb_to_lisa(alpha): ------- r_total : numpy.array A 3x3 rotation matrix from SSB frame to LISA frame. + """ - r = Rotation.from_rotvec([ - [0, 0, alpha], - [0, -np.pi/3, 0], - [0, 0, -alpha] - ]).as_matrix() + r = Rotation.from_rotvec( + [[0, 0, alpha], [0, -np.pi / 3, 0], [0, 0, -alpha]] + ).as_matrix() r_total = np.array(r[0]) @ np.array(r[1]) @ np.array(r[2]) return r_total def lisa_position_ssb(t_lisa, t0=TIME_OFFSET_20_DEGREES): - """ Calculating the position vector and angular displacement of LISA + """ + Calculating the position vector and angular displacement of LISA in the SSB frame, at a given time. This function assumes LISA's barycenter is orbiting around a circular orbit within the ecliptic behind the Earth. The period of it is one year. @@ -104,19 +108,22 @@ def lisa_position_ssb(t_lisa, t0=TIME_OFFSET_20_DEGREES): alpha : float The angular displacement of LISA in the SSB frame. In the unit of 'radian'. + """ OMEGA_0 = 1.99098659277e-7 R_ORBIT = au.value - alpha = np.mod(OMEGA_0 * (t_lisa + t0), 2*np.pi) - p = np.array([[R_ORBIT * np.cos(alpha)], - [R_ORBIT * np.sin(alpha)], - [0]], dtype=object) + alpha = np.mod(OMEGA_0 * (t_lisa + t0), 2 * np.pi) + p = np.array( + [[R_ORBIT * np.cos(alpha)], [R_ORBIT * np.sin(alpha)], [0]], dtype=object + ) return (p, alpha) -def localization_to_propagation_vector(longitude, latitude, - use_astropy=True, frame=None): - """ Converting the sky localization to the corresponding +def localization_to_propagation_vector( + longitude, latitude, use_astropy=True, frame=None +): + """ + Converting the sky localization to the corresponding propagation unit vector of a GW signal. Parameters @@ -136,6 +143,7 @@ def localization_to_propagation_vector(longitude, latitude, ------- [[x], [y], [z]] : numpy.array The propagation unit vector of that GW signal. + """ if use_astropy: x = -frame.cartesian.x.value @@ -151,7 +159,8 @@ def localization_to_propagation_vector(longitude, latitude, def propagation_vector_to_localization(k, use_astropy=True, frame=None): - """ Converting the propagation unit vector to the corresponding + """ + Converting the propagation unit vector to the corresponding sky localization of a GW signal. Parameters @@ -169,6 +178,7 @@ def propagation_vector_to_localization(k, use_astropy=True, frame=None): ------- (longitude, latitude) : tuple The sky localization of that GW signal. + """ if use_astropy: try: @@ -179,18 +189,21 @@ def propagation_vector_to_localization(k, use_astropy=True, frame=None): latitude = frame.dec.rad else: # latitude already within [-pi/2, pi/2] - latitude = np.float64(np.arcsin(-k[2,0])) - longitude = np.float64(np.arctan2(-k[1,0]/np.cos(latitude), - -k[0,0]/np.cos(latitude))) + latitude = np.float64(np.arcsin(-k[2, 0])) + longitude = np.float64( + np.arctan2(-k[1, 0] / np.cos(latitude), -k[0, 0] / np.cos(latitude)) + ) # longitude should within [0, 2*pi) - longitude = np.mod(longitude, 2*np.pi) + longitude = np.mod(longitude, 2 * np.pi) return (longitude, latitude) -def polarization_newframe(polarization, k, rotation_matrix, use_astropy=True, - old_frame=None, new_frame=None): - """ Converting a polarization angle from a frame to a new frame +def polarization_newframe( + polarization, k, rotation_matrix, use_astropy=True, old_frame=None, new_frame=None +): + """ + Converting a polarization angle from a frame to a new frame by using rotation matrix method. Parameters @@ -217,37 +230,42 @@ def polarization_newframe(polarization, k, rotation_matrix, use_astropy=True, ------- polarization_new_frame : float The polarization angle in the new frame of that GW signal. + """ - longitude, _ = propagation_vector_to_localization( - k, use_astropy, old_frame) + longitude, _ = propagation_vector_to_localization(k, use_astropy, old_frame) u = np.array([[np.sin(longitude)], [-np.cos(longitude)], [0]]) rotation_vector = polarization * k rotation_polarization = Rotation.from_rotvec(rotation_vector.T[0]) p = rotation_polarization.apply(u.T[0]).reshape(3, 1) p_newframe = rotation_matrix.T @ p k_newframe = rotation_matrix.T @ k - longitude_newframe, latitude_newframe = \ - propagation_vector_to_localization(k_newframe, use_astropy, new_frame) - u_newframe = np.array([[np.sin(longitude_newframe)], - [-np.cos(longitude_newframe)], [0]]) - v_newframe = np.array([ - [-np.sin(latitude_newframe) * np.cos(longitude_newframe)], - [-np.sin(latitude_newframe) * np.sin(longitude_newframe)], - [np.cos(latitude_newframe)]]) + longitude_newframe, latitude_newframe = propagation_vector_to_localization( + k_newframe, use_astropy, new_frame + ) + u_newframe = np.array( + [[np.sin(longitude_newframe)], [-np.cos(longitude_newframe)], [0]] + ) + v_newframe = np.array( + [ + [-np.sin(latitude_newframe) * np.cos(longitude_newframe)], + [-np.sin(latitude_newframe) * np.sin(longitude_newframe)], + [np.cos(latitude_newframe)], + ] + ) p_dot_u_newframe = np.vdot(p_newframe, u_newframe) p_dot_v_newframe = np.vdot(p_newframe, v_newframe) polarization_new_frame = np.arctan2(p_dot_v_newframe, p_dot_u_newframe) - polarization_new_frame = np.mod(polarization_new_frame, 2*np.pi) + polarization_new_frame = np.mod(polarization_new_frame, 2 * np.pi) # avoid the round error - if polarization_new_frame == 2*np.pi: + if polarization_new_frame == 2 * np.pi: polarization_new_frame = 0 return polarization_new_frame -def t_lisa_from_ssb(t_ssb, longitude_ssb, latitude_ssb, - t0=TIME_OFFSET_20_DEGREES): - """ Calculating the time when a GW signal arrives at the barycenter +def t_lisa_from_ssb(t_ssb, longitude_ssb, latitude_ssb, t0=TIME_OFFSET_20_DEGREES): + """ + Calculating the time when a GW signal arrives at the barycenter of LISA, by using the time and sky localization in SSB frame. Parameters @@ -270,9 +288,11 @@ def t_lisa_from_ssb(t_ssb, longitude_ssb, latitude_ssb, ------- t_lisa : float The time when a GW signal arrives at the origin of LISA frame. + """ k = localization_to_propagation_vector( - longitude_ssb, latitude_ssb, use_astropy=False) + longitude_ssb, latitude_ssb, use_astropy=False + ) def equation(t_lisa): # LISA is moving, when GW arrives at LISA center, @@ -283,9 +303,9 @@ def equation(t_lisa): return fsolve(equation, t_ssb)[0] -def t_ssb_from_t_lisa(t_lisa, longitude_ssb, latitude_ssb, - t0=TIME_OFFSET_20_DEGREES): - """ Calculating the time when a GW signal arrives at the barycenter +def t_ssb_from_t_lisa(t_lisa, longitude_ssb, latitude_ssb, t0=TIME_OFFSET_20_DEGREES): + """ + Calculating the time when a GW signal arrives at the barycenter of SSB, by using the time in LISA frame and sky localization in SSB frame. Parameters @@ -308,9 +328,11 @@ def t_ssb_from_t_lisa(t_lisa, longitude_ssb, latitude_ssb, ------- t_ssb : float The time when a GW signal arrives at the origin of SSB frame. + """ k = localization_to_propagation_vector( - longitude_ssb, latitude_ssb, use_astropy=False) + longitude_ssb, latitude_ssb, use_astropy=False + ) # LISA is moving, when GW arrives at LISA center, # time is t_lisa, not t_ssb. p = lisa_position_ssb(t_lisa, t0)[0] @@ -321,9 +343,11 @@ def equation(t_ssb): return fsolve(equation, t_lisa)[0] -def ssb_to_lisa(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, - t0=TIME_OFFSET_20_DEGREES): - """ Converting the arrive time, the sky localization, and the polarization +def ssb_to_lisa( + t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, t0=TIME_OFFSET_20_DEGREES +): + """ + Converting the arrive time, the sky localization, and the polarization from the SSB frame to the LISA frame. Parameters @@ -358,6 +382,7 @@ def ssb_to_lisa(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, polarization_lisa : float or numpy.array The polarization angle of a GW signal in LISA frame. In the unit of 'radian'. + """ if not isinstance(t_ssb, np.ndarray): t_ssb = np.array([t_ssb]) @@ -372,41 +397,47 @@ def ssb_to_lisa(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, latitude_lisa, polarization_lisa = np.zeros(num), np.zeros(num) for i in range(num): - if longitude_ssb[i] < 0 or longitude_ssb[i] >= 2*np.pi: + if longitude_ssb[i] < 0 or longitude_ssb[i] >= 2 * np.pi: raise ValueError("Longitude should within [0, 2*pi).") - if latitude_ssb[i] < -np.pi/2 or latitude_ssb[i] > np.pi/2: + if latitude_ssb[i] < -np.pi / 2 or latitude_ssb[i] > np.pi / 2: raise ValueError("Latitude should within [-pi/2, pi/2].") - if polarization_ssb[i] < 0 or polarization_ssb[i] >= 2*np.pi: + if polarization_ssb[i] < 0 or polarization_ssb[i] >= 2 * np.pi: raise ValueError("Polarization angle should within [0, 2*pi).") - t_lisa[i] = t_lisa_from_ssb(t_ssb[i], longitude_ssb[i], - latitude_ssb[i], t0) + t_lisa[i] = t_lisa_from_ssb(t_ssb[i], longitude_ssb[i], latitude_ssb[i], t0) k_ssb = localization_to_propagation_vector( - longitude_ssb[i], latitude_ssb[i], use_astropy=False) + longitude_ssb[i], latitude_ssb[i], use_astropy=False + ) # Although t_lisa calculated above using the corrected LISA position # vector by adding t0, it corresponds to the true t_ssb, not t_ssb+t0, # we need to include t0 again to correct LISA position. alpha = lisa_position_ssb(t_lisa[i], t0)[1] rotation_matrix_lisa = rotation_matrix_ssb_to_lisa(alpha) k_lisa = rotation_matrix_lisa.T @ k_ssb - longitude_lisa[i], latitude_lisa[i] = \ - propagation_vector_to_localization(k_lisa, use_astropy=False) + longitude_lisa[i], latitude_lisa[i] = propagation_vector_to_localization( + k_lisa, use_astropy=False + ) polarization_lisa[i] = polarization_newframe( - polarization_ssb[i], k_ssb, rotation_matrix_lisa, - use_astropy=False) + polarization_ssb[i], k_ssb, rotation_matrix_lisa, use_astropy=False + ) if num == 1: - params_lisa = (t_lisa[0], longitude_lisa[0], - latitude_lisa[0], polarization_lisa[0]) + params_lisa = ( + t_lisa[0], + longitude_lisa[0], + latitude_lisa[0], + polarization_lisa[0], + ) else: - params_lisa = (t_lisa, longitude_lisa, - latitude_lisa, polarization_lisa) + params_lisa = (t_lisa, longitude_lisa, latitude_lisa, polarization_lisa) return params_lisa -def lisa_to_ssb(t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, - t0=TIME_OFFSET_20_DEGREES): - """ Converting the arrive time, the sky localization, and the polarization +def lisa_to_ssb( + t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, t0=TIME_OFFSET_20_DEGREES +): + """ + Converting the arrive time, the sky localization, and the polarization from the LISA frame to the SSB frame. Parameters @@ -441,6 +472,7 @@ def lisa_to_ssb(t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, polarization_ssb : float or numpy.array The polarization angle of a GW signal in SSB frame. In the unit of 'radian'. + """ if not isinstance(t_lisa, np.ndarray): t_lisa = np.array([t_lisa]) @@ -455,37 +487,37 @@ def lisa_to_ssb(t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, latitude_ssb, polarization_ssb = np.zeros(num), np.zeros(num) for i in range(num): - if longitude_lisa[i] < 0 or longitude_lisa[i] >= 2*np.pi: + if longitude_lisa[i] < 0 or longitude_lisa[i] >= 2 * np.pi: raise ValueError("Longitude should within [0, 2*pi).") - if latitude_lisa[i] < -np.pi/2 or latitude_lisa[i] > np.pi/2: + if latitude_lisa[i] < -np.pi / 2 or latitude_lisa[i] > np.pi / 2: raise ValueError("Latitude should within [-pi/2, pi/2].") - if polarization_lisa[i] < 0 or polarization_lisa[i] >= 2*np.pi: + if polarization_lisa[i] < 0 or polarization_lisa[i] >= 2 * np.pi: raise ValueError("Polarization angle should within [0, 2*pi).") k_lisa = localization_to_propagation_vector( - longitude_lisa[i], latitude_lisa[i], use_astropy=False) + longitude_lisa[i], latitude_lisa[i], use_astropy=False + ) alpha = lisa_position_ssb(t_lisa[i], t0)[1] rotation_matrix_lisa = rotation_matrix_ssb_to_lisa(alpha) k_ssb = rotation_matrix_lisa @ k_lisa - longitude_ssb[i], latitude_ssb[i] = \ - propagation_vector_to_localization(k_ssb, use_astropy=False) - t_ssb[i] = t_ssb_from_t_lisa(t_lisa[i], longitude_ssb[i], - latitude_ssb[i], t0) + longitude_ssb[i], latitude_ssb[i] = propagation_vector_to_localization( + k_ssb, use_astropy=False + ) + t_ssb[i] = t_ssb_from_t_lisa(t_lisa[i], longitude_ssb[i], latitude_ssb[i], t0) polarization_ssb[i] = polarization_newframe( - polarization_lisa[i], k_lisa, rotation_matrix_lisa.T, - use_astropy=False) + polarization_lisa[i], k_lisa, rotation_matrix_lisa.T, use_astropy=False + ) if num == 1: - params_ssb = (t_ssb[0], longitude_ssb[0], - latitude_ssb[0], polarization_ssb[0]) + params_ssb = (t_ssb[0], longitude_ssb[0], latitude_ssb[0], polarization_ssb[0]) else: - params_ssb = (t_ssb, longitude_ssb, - latitude_ssb, polarization_ssb) + params_ssb = (t_ssb, longitude_ssb, latitude_ssb, polarization_ssb) return params_ssb def rotation_matrix_ssb_to_geo(epsilon=np.deg2rad(23.439281)): - """ The rotation matrix (of frame basis) from SSB frame to + """ + The rotation matrix (of frame basis) from SSB frame to geocentric frame. Parameters @@ -497,16 +529,16 @@ def rotation_matrix_ssb_to_geo(epsilon=np.deg2rad(23.439281)): ------- r : numpy.array A 3x3 rotation matrix from SSB frame to geocentric frame. + """ - r = Rotation.from_rotvec([ - [-epsilon, 0, 0] - ]).as_matrix() + r = Rotation.from_rotvec([[-epsilon, 0, 0]]).as_matrix() return np.array(r[0]) def earth_position_ssb(t_geo): - """ Calculating the position vector and angular displacement of the Earth + """ + Calculating the position vector and angular displacement of the Earth in the SSB frame, at a given time. By using Astropy. Parameters @@ -523,14 +555,14 @@ def earth_position_ssb(t_geo): alpha : float The angular displacement of the Earth in the SSB frame. In the unit of 'radian'. + """ - t = Time(t_geo, format='gps') - pos = get_body_barycentric('earth', t) + t = Time(t_geo, format="gps") + pos = get_body_barycentric("earth", t) # BarycentricMeanEcliptic doesn't have obstime attribute, # it's a good inertial frame, but ICRS is not. - icrs_coord = SkyCoord(pos, frame='icrs', obstime=t) - bme_coord = icrs_coord.transform_to( - BarycentricMeanEcliptic(equinox='J2000')) + icrs_coord = SkyCoord(pos, frame="icrs", obstime=t) + bme_coord = icrs_coord.transform_to(BarycentricMeanEcliptic(equinox="J2000")) x = bme_coord.cartesian.x.to(units.m).value y = bme_coord.cartesian.y.to(units.m).value z = bme_coord.cartesian.z.to(units.m).value @@ -540,9 +572,9 @@ def earth_position_ssb(t_geo): return (p, alpha) -def t_geo_from_ssb(t_ssb, longitude_ssb, latitude_ssb, - use_astropy=True, frame=None): - """ Calculating the time when a GW signal arrives at the barycenter +def t_geo_from_ssb(t_ssb, longitude_ssb, latitude_ssb, use_astropy=True, frame=None): + """ + Calculating the time when a GW signal arrives at the barycenter of the Earth, by using the time and sky localization in SSB frame. Parameters @@ -561,9 +593,11 @@ def t_geo_from_ssb(t_ssb, longitude_ssb, latitude_ssb, ------- t_geo : float The time when a GW signal arrives at the origin of geocentric frame. + """ k = localization_to_propagation_vector( - longitude_ssb, latitude_ssb, use_astropy, frame) + longitude_ssb, latitude_ssb, use_astropy, frame + ) def equation(t_geo): # Earth is moving, when GW arrives at Earth center, @@ -574,9 +608,9 @@ def equation(t_geo): return fsolve(equation, t_ssb)[0] -def t_ssb_from_t_geo(t_geo, longitude_ssb, latitude_ssb, - use_astropy=True, frame=None): - """ Calculating the time when a GW signal arrives at the barycenter +def t_ssb_from_t_geo(t_geo, longitude_ssb, latitude_ssb, use_astropy=True, frame=None): + """ + Calculating the time when a GW signal arrives at the barycenter of SSB, by using the time in geocentric frame and sky localization in SSB frame. @@ -596,9 +630,11 @@ def t_ssb_from_t_geo(t_geo, longitude_ssb, latitude_ssb, ------- t_ssb : float The time when a GW signal arrives at the origin of SSB frame. + """ k = localization_to_propagation_vector( - longitude_ssb, latitude_ssb, use_astropy, frame) + longitude_ssb, latitude_ssb, use_astropy, frame + ) # Earth is moving, when GW arrives at Earth center, # time is t_geo, not t_ssb. p = earth_position_ssb(t_geo)[0] @@ -609,9 +645,9 @@ def equation(t_ssb): return fsolve(equation, t_geo)[0] -def ssb_to_geo(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, - use_astropy=True): - """ Converting the arrive time, the sky localization, and the polarization +def ssb_to_geo(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, use_astropy=True): + """ + Converting the arrive time, the sky localization, and the polarization from the SSB frame to the geocentric frame. Parameters @@ -647,6 +683,7 @@ def ssb_to_geo(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, polarization_geo : float or numpy.array The polarization angle of a GW signal in geocentric frame. In the unit of 'radian'. + """ if not isinstance(t_ssb, np.ndarray): t_ssb = np.array([t_ssb]) @@ -663,71 +700,79 @@ def ssb_to_geo(t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, polarization_geo = np.full(num, np.nan) for i in range(num): - if longitude_ssb[i] < 0 or longitude_ssb[i] >= 2*np.pi: + if longitude_ssb[i] < 0 or longitude_ssb[i] >= 2 * np.pi: raise ValueError("Longitude should within [0, 2*pi).") - if latitude_ssb[i] < -np.pi/2 or latitude_ssb[i] > np.pi/2: + if latitude_ssb[i] < -np.pi / 2 or latitude_ssb[i] > np.pi / 2: raise ValueError("Latitude should within [-pi/2, pi/2].") - if polarization_ssb[i] < 0 or polarization_ssb[i] >= 2*np.pi: + if polarization_ssb[i] < 0 or polarization_ssb[i] >= 2 * np.pi: raise ValueError("Polarization angle should within [0, 2*pi).") if use_astropy: # BarycentricMeanEcliptic doesn't have obstime attribute, # it's a good inertial frame, but PrecessedGeocentric is not. bme_coord = BarycentricMeanEcliptic( - lon=longitude_ssb[i]*units.radian, - lat=latitude_ssb[i]*units.radian, - equinox='J2000') - t_geo[i] = t_geo_from_ssb(t_ssb[i], longitude_ssb[i], - latitude_ssb[i], use_astropy, bme_coord) - geo_sky = bme_coord.transform_to(PrecessedGeocentric( - equinox='J2000', obstime=Time(t_geo[i], format='gps'))) + lon=longitude_ssb[i] * units.radian, + lat=latitude_ssb[i] * units.radian, + equinox="J2000", + ) + t_geo[i] = t_geo_from_ssb( + t_ssb[i], longitude_ssb[i], latitude_ssb[i], use_astropy, bme_coord + ) + geo_sky = bme_coord.transform_to( + PrecessedGeocentric( + equinox="J2000", obstime=Time(t_geo[i], format="gps") + ) + ) longitude_geo[i] = geo_sky.ra.rad latitude_geo[i] = geo_sky.dec.rad k_geo = localization_to_propagation_vector( - longitude_geo[i], latitude_geo[i], - use_astropy, geo_sky) + longitude_geo[i], latitude_geo[i], use_astropy, geo_sky + ) k_ssb = localization_to_propagation_vector( - None, None, use_astropy, bme_coord) - rotation_matrix_geo = \ - ecliptic_transforms.icrs_to_baryecliptic( - from_coo=None, - to_frame=BarycentricMeanEcliptic(equinox='J2000')) + None, None, use_astropy, bme_coord + ) + rotation_matrix_geo = ecliptic_transforms.icrs_to_baryecliptic( + from_coo=None, to_frame=BarycentricMeanEcliptic(equinox="J2000") + ) polarization_geo[i] = polarization_newframe( - polarization_ssb[i], k_ssb, - rotation_matrix_geo, use_astropy, - old_frame=bme_coord, - new_frame=geo_sky) + polarization_ssb[i], + k_ssb, + rotation_matrix_geo, + use_astropy, + old_frame=bme_coord, + new_frame=geo_sky, + ) else: - t_geo[i] = t_geo_from_ssb(t_ssb[i], longitude_ssb[i], - latitude_ssb[i], use_astropy) + t_geo[i] = t_geo_from_ssb( + t_ssb[i], longitude_ssb[i], latitude_ssb[i], use_astropy + ) rotation_matrix_geo = rotation_matrix_ssb_to_geo() k_ssb = localization_to_propagation_vector( - longitude_ssb[i], latitude_ssb[i], - use_astropy) + longitude_ssb[i], latitude_ssb[i], use_astropy + ) k_geo = rotation_matrix_geo.T @ k_ssb - longitude_geo[i], latitude_geo[i] = \ - propagation_vector_to_localization(k_geo, use_astropy) + longitude_geo[i], latitude_geo[i] = propagation_vector_to_localization( + k_geo, use_astropy + ) polarization_geo[i] = polarization_newframe( - polarization_ssb[i], k_ssb, - rotation_matrix_geo, use_astropy) + polarization_ssb[i], k_ssb, rotation_matrix_geo, use_astropy + ) # As mentioned in LDC manual, the p,q vectors are opposite between # LDC and LAL conventions, see Sec 4.1.5 in . - polarization_geo[i] = np.mod(polarization_geo[i]+np.pi, 2*np.pi) + polarization_geo[i] = np.mod(polarization_geo[i] + np.pi, 2 * np.pi) if num == 1: - params_geo = (t_geo[0], longitude_geo[0], - latitude_geo[0], polarization_geo[0]) + params_geo = (t_geo[0], longitude_geo[0], latitude_geo[0], polarization_geo[0]) else: - params_geo = (t_geo, longitude_geo, - latitude_geo, polarization_geo) + params_geo = (t_geo, longitude_geo, latitude_geo, polarization_geo) return params_geo -def geo_to_ssb(t_geo, longitude_geo, latitude_geo, polarization_geo, - use_astropy=True): - """ Converting the arrive time, the sky localization, and the polarization +def geo_to_ssb(t_geo, longitude_geo, latitude_geo, polarization_geo, use_astropy=True): + """ + Converting the arrive time, the sky localization, and the polarization from the geocentric frame to the SSB frame. Parameters @@ -763,6 +808,7 @@ def geo_to_ssb(t_geo, longitude_geo, latitude_geo, polarization_geo, polarization_ssb : float or numpy.array The polarization angle of a GW signal in SSB frame. In the unit of 'radian'. + """ if not isinstance(t_geo, np.ndarray): t_geo = np.array([t_geo]) @@ -779,73 +825,83 @@ def geo_to_ssb(t_geo, longitude_geo, latitude_geo, polarization_geo, polarization_ssb = np.full(num, np.nan) for i in range(num): - if longitude_geo[i] < 0 or longitude_geo[i] >= 2*np.pi: + if longitude_geo[i] < 0 or longitude_geo[i] >= 2 * np.pi: raise ValueError("Longitude should within [0, 2*pi).") - if latitude_geo[i] < -np.pi/2 or latitude_geo[i] > np.pi/2: + if latitude_geo[i] < -np.pi / 2 or latitude_geo[i] > np.pi / 2: raise ValueError("Latitude should within [-pi/2, pi/2].") - if polarization_geo[i] < 0 or polarization_geo[i] >= 2*np.pi: + if polarization_geo[i] < 0 or polarization_geo[i] >= 2 * np.pi: raise ValueError("Polarization angle should within [0, 2*pi).") if use_astropy: # BarycentricMeanEcliptic doesn't have obstime attribute, # it's a good inertial frame, but PrecessedGeocentric is not. geo_coord = PrecessedGeocentric( - ra=longitude_geo[i]*units.radian, - dec=latitude_geo[i]*units.radian, - equinox='J2000', - obstime=Time(t_geo[i], format='gps')) - ssb_sky = geo_coord.transform_to( - BarycentricMeanEcliptic(equinox='J2000')) + ra=longitude_geo[i] * units.radian, + dec=latitude_geo[i] * units.radian, + equinox="J2000", + obstime=Time(t_geo[i], format="gps"), + ) + ssb_sky = geo_coord.transform_to(BarycentricMeanEcliptic(equinox="J2000")) longitude_ssb[i] = ssb_sky.lon.rad latitude_ssb[i] = ssb_sky.lat.rad k_ssb = localization_to_propagation_vector( - longitude_ssb[i], latitude_ssb[i], - use_astropy, ssb_sky) + longitude_ssb[i], latitude_ssb[i], use_astropy, ssb_sky + ) k_geo = localization_to_propagation_vector( - None, None, use_astropy, geo_coord) - rotation_matrix_geo = \ - ecliptic_transforms.icrs_to_baryecliptic( - from_coo=None, - to_frame=BarycentricMeanEcliptic(equinox='J2000')) - t_ssb[i] = t_ssb_from_t_geo(t_geo[i], longitude_ssb[i], - latitude_ssb[i], use_astropy, - ssb_sky) + None, None, use_astropy, geo_coord + ) + rotation_matrix_geo = ecliptic_transforms.icrs_to_baryecliptic( + from_coo=None, to_frame=BarycentricMeanEcliptic(equinox="J2000") + ) + t_ssb[i] = t_ssb_from_t_geo( + t_geo[i], longitude_ssb[i], latitude_ssb[i], use_astropy, ssb_sky + ) polarization_ssb[i] = polarization_newframe( - polarization_geo[i], k_geo, - rotation_matrix_geo.T, - use_astropy, - old_frame=geo_coord, - new_frame=ssb_sky) + polarization_geo[i], + k_geo, + rotation_matrix_geo.T, + use_astropy, + old_frame=geo_coord, + new_frame=ssb_sky, + ) else: rotation_matrix_geo = rotation_matrix_ssb_to_geo() k_geo = localization_to_propagation_vector( - longitude_geo[i], latitude_geo[i], use_astropy) + longitude_geo[i], latitude_geo[i], use_astropy + ) k_ssb = rotation_matrix_geo @ k_geo - longitude_ssb[i], latitude_ssb[i] = \ - propagation_vector_to_localization(k_ssb, use_astropy) - t_ssb[i] = t_ssb_from_t_geo(t_geo[i], longitude_ssb[i], - latitude_ssb[i], use_astropy) + longitude_ssb[i], latitude_ssb[i] = propagation_vector_to_localization( + k_ssb, use_astropy + ) + t_ssb[i] = t_ssb_from_t_geo( + t_geo[i], longitude_ssb[i], latitude_ssb[i], use_astropy + ) polarization_ssb[i] = polarization_newframe( - polarization_geo[i], k_geo, - rotation_matrix_geo.T, use_astropy) + polarization_geo[i], k_geo, rotation_matrix_geo.T, use_astropy + ) # As mentioned in LDC manual, the p,q vectors are opposite between # LDC and LAL conventions, see Sec 4.1.5 in . - polarization_ssb[i] = np.mod(polarization_ssb[i]-np.pi, 2*np.pi) + polarization_ssb[i] = np.mod(polarization_ssb[i] - np.pi, 2 * np.pi) if num == 1: - params_ssb = (t_ssb[0], longitude_ssb[0], - latitude_ssb[0], polarization_ssb[0]) + params_ssb = (t_ssb[0], longitude_ssb[0], latitude_ssb[0], polarization_ssb[0]) else: - params_ssb = (t_ssb, longitude_ssb, - latitude_ssb, polarization_ssb) + params_ssb = (t_ssb, longitude_ssb, latitude_ssb, polarization_ssb) return params_ssb -def lisa_to_geo(t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, - t0=TIME_OFFSET_20_DEGREES, use_astropy=True): - """ Converting the arrive time, the sky localization, and the polarization +def lisa_to_geo( + t_lisa, + longitude_lisa, + latitude_lisa, + polarization_lisa, + t0=TIME_OFFSET_20_DEGREES, + use_astropy=True, +): + """ + Converting the arrive time, the sky localization, and the polarization from the LISA frame to the geocentric frame. Parameters @@ -883,18 +939,28 @@ def lisa_to_geo(t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, polarization_geo : float or numpy.array The polarization angle of a GW signal in geocentric frame. In the unit of 'radian'. + """ t_ssb, longitude_ssb, latitude_ssb, polarization_ssb = lisa_to_ssb( - t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, t0) + t_lisa, longitude_lisa, latitude_lisa, polarization_lisa, t0 + ) t_geo, longitude_geo, latitude_geo, polarization_geo = ssb_to_geo( - t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, use_astropy) + t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, use_astropy + ) return (t_geo, longitude_geo, latitude_geo, polarization_geo) -def geo_to_lisa(t_geo, longitude_geo, latitude_geo, polarization_geo, - t0=TIME_OFFSET_20_DEGREES, use_astropy=True): - """ Converting the arrive time, the sky localization, and the polarization +def geo_to_lisa( + t_geo, + longitude_geo, + latitude_geo, + polarization_geo, + t0=TIME_OFFSET_20_DEGREES, + use_astropy=True, +): + """ + Converting the arrive time, the sky localization, and the polarization from the geocentric frame to the LISA frame. Parameters @@ -932,22 +998,35 @@ def geo_to_lisa(t_geo, longitude_geo, latitude_geo, polarization_geo, polarization_geo : float or numpy.array The polarization angle of a GW signal in LISA frame. In the unit of 'radian'. + """ t_ssb, longitude_ssb, latitude_ssb, polarization_ssb = geo_to_ssb( - t_geo, longitude_geo, latitude_geo, polarization_geo, use_astropy) + t_geo, longitude_geo, latitude_geo, polarization_geo, use_astropy + ) t_lisa, longitude_lisa, latitude_lisa, polarization_lisa = ssb_to_lisa( - t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, t0) + t_ssb, longitude_ssb, latitude_ssb, polarization_ssb, t0 + ) return (t_lisa, longitude_lisa, latitude_lisa, polarization_lisa) -__all__ = ['TIME_OFFSET_20_DEGREES', - 'localization_to_propagation_vector', - 'propagation_vector_to_localization', 'polarization_newframe', - 't_lisa_from_ssb', 't_ssb_from_t_lisa', - 'ssb_to_lisa', 'lisa_to_ssb', - 'rotation_matrix_ssb_to_lisa', 'rotation_matrix_ssb_to_geo', - 'lisa_position_ssb', 'earth_position_ssb', - 't_geo_from_ssb', 't_ssb_from_t_geo', 'ssb_to_geo', 'geo_to_ssb', - 'lisa_to_geo', 'geo_to_lisa', - ] +__all__ = [ + "TIME_OFFSET_20_DEGREES", + "earth_position_ssb", + "geo_to_lisa", + "geo_to_ssb", + "lisa_position_ssb", + "lisa_to_geo", + "lisa_to_ssb", + "localization_to_propagation_vector", + "polarization_newframe", + "propagation_vector_to_localization", + "rotation_matrix_ssb_to_geo", + "rotation_matrix_ssb_to_lisa", + "ssb_to_geo", + "ssb_to_lisa", + "t_geo_from_ssb", + "t_lisa_from_ssb", + "t_ssb_from_t_geo", + "t_ssb_from_t_lisa", +] diff --git a/pycbc/cosmology.py b/pycbc/cosmology.py index e0cbfb58ae1..bf19beecb10 100644 --- a/pycbc/cosmology.py +++ b/pycbc/cosmology.py @@ -30,20 +30,23 @@ """ import logging -import numpy -from scipy import interpolate + import astropy.cosmology +import numpy from astropy import units from astropy.cosmology import CosmologyError, parameters +from scipy import interpolate + import pycbc.conversions -logger = logging.getLogger('pycbc.cosmology') +logger = logging.getLogger("pycbc.cosmology") -DEFAULT_COSMOLOGY = 'Planck15' +DEFAULT_COSMOLOGY = "Planck15" def get_cosmology(cosmology=None, **kwargs): - r"""Gets an astropy cosmology class. + r""" + Gets an astropy cosmology class. Parameters ---------- @@ -88,8 +91,10 @@ def get_cosmology(cosmology=None, **kwargs): """ if kwargs and cosmology is not None: - raise ValueError("if providing custom cosmological parameters, do " - "not provide a `cosmology` argument") + raise ValueError( + "if providing custom cosmological parameters, do " + "not provide a `cosmology` argument" + ) if isinstance(cosmology, astropy.cosmology.FlatLambdaCDM): # just return return cosmology @@ -99,13 +104,14 @@ def get_cosmology(cosmology=None, **kwargs): if cosmology is None: cosmology = DEFAULT_COSMOLOGY if cosmology not in parameters.available: - raise ValueError("unrecognized cosmology {}".format(cosmology)) + raise ValueError(f"unrecognized cosmology {cosmology}") cosmology = getattr(astropy.cosmology, cosmology) return cosmology -def z_at_value(func, fval, unit, zmax=1000., **kwargs): - r"""Wrapper around astropy.cosmology.z_at_value to handle numpy arrays. +def z_at_value(func, fval, unit, zmax=1000.0, **kwargs): + r""" + Wrapper around astropy.cosmology.z_at_value to handle numpy arrays. Getting a z for a cosmological quantity involves numerically inverting ``func``. The ``zmax`` argument sets how large of a z to guess (see @@ -131,22 +137,22 @@ def z_at_value(func, fval, unit, zmax=1000., **kwargs): ------- float The redshift at the requested values. + """ fval, input_is_array = pycbc.conversions.ensurearray(fval) # make sure fval is atleast 1D if fval.size == 1 and fval.ndim == 0: fval = fval.reshape(1) zs = numpy.zeros(fval.shape, dtype=float) # the output array - if 'method' not in kwargs: + if "method" not in kwargs: # workaround for https://github.com/astropy/astropy/issues/14249 # FIXME remove when fixed in astropy/scipy - kwargs['method'] = 'bounded' - for (ii, val) in enumerate(fval): + kwargs["method"] = "bounded" + for ii, val in enumerate(fval): try: - zs[ii] = astropy.cosmology.z_at_value(func, val*unit, zmax=zmax, - **kwargs) + zs[ii] = astropy.cosmology.z_at_value(func, val * unit, zmax=zmax, **kwargs) except CosmologyError: - if ii == len(zs)-1: + if ii == len(zs) - 1: # if zs[ii] is less than but very close to zmax, let's say # zs[ii] is the last element in the [zmin, zmax], # `z_at_value` will also returns "CosmologyError", please @@ -154,7 +160,7 @@ def z_at_value(func, fval, unit, zmax=1000., **kwargs): # cosmology.z_at_value.html), in order to avoid bumping up # zmax, just set zs equals to previous value, we assume # the `func` is smooth - zs[ii] = zs[ii-1] + zs[ii] = zs[ii - 1] else: # we'll get this if the z was larger than zmax; in that # case we'll try bumping up zmax later to get a value @@ -166,14 +172,15 @@ def z_at_value(func, fval, unit, zmax=1000., **kwargs): # we'll keep bumping up the maxz until we can get a result counter = 0 # to prevent running forever while replacemask.any(): - kwargs['zmin'] = zmax + kwargs["zmin"] = zmax zmax = 10 * zmax idx = numpy.where(replacemask) for ii in idx: val = fval[ii] try: zs[ii] = astropy.cosmology.z_at_value( - func, val*unit, zmax=zmax, **kwargs) + func, val * unit, zmax=zmax, **kwargs + ) replacemask[ii] = False except CosmologyError: # didn't work, try on next loop @@ -181,16 +188,19 @@ def z_at_value(func, fval, unit, zmax=1000., **kwargs): counter += 1 if counter == 5: # give up and warn the user - logger.warning("One or more values correspond to a " - "redshift > {0:.1e}. The redshift for these " - "have been set to inf. If you would like " - "better precision, call God.".format(zmax)) + logger.warning( + "One or more values correspond to a " + f"redshift > {zmax:.1e}. The redshift for these " + "have been set to inf. If you would like " + "better precision, call God." + ) break return pycbc.conversions.formatreturn(zs, input_is_array) def _redshift(distance, **kwargs): - r"""Uses astropy to get redshift from the given luminosity distance. + r""" + Uses astropy to get redshift from the given luminosity distance. Parameters ---------- @@ -205,13 +215,15 @@ def _redshift(distance, **kwargs): ------- float : The redshift corresponding to the given luminosity distance. + """ cosmology = get_cosmology(**kwargs) return z_at_value(cosmology.luminosity_distance, distance, units.Mpc) -class DistToZ(object): - r"""Interpolates luminosity distance as a function of redshift to allow for +class DistToZ: + r""" + Interpolates luminosity distance as a function of redshift to allow for fast conversion. The :mod:`astropy.cosmology` module provides methods for converting any @@ -240,8 +252,10 @@ class speeds that up by pre-interpolating :math:`D(z)`. It works by setting All other keyword args are passed to :py:func:`get_cosmology` to select a cosmology. If none provided, will use :py:attr:`DEFAULT_COSMOLOGY`. + """ - def __init__(self, default_maxz=1000., numpoints=10000, **kwargs): + + def __init__(self, default_maxz=1000.0, numpoints=10000, **kwargs): self.numpoints = int(numpoints) self.default_maxz = default_maxz self.cosmology = get_cosmology(**kwargs) @@ -254,22 +268,22 @@ def __init__(self, default_maxz=1000., numpoints=10000, **kwargs): def setup_interpolant(self): """Initializes the z(d) interpolation.""" # for computing nearby (z < 1) redshifts - zs = numpy.linspace(0., 1., num=self.numpoints) + zs = numpy.linspace(0.0, 1.0, num=self.numpoints) ds = self.cosmology.luminosity_distance(zs).value - self.nearby_d2z = interpolate.interp1d(ds, zs, kind='linear', - bounds_error=False) + self.nearby_d2z = interpolate.interp1d( + ds, zs, kind="linear", bounds_error=False + ) # for computing far away (z > 1) redshifts - zs = numpy.logspace(0, numpy.log10(self.default_maxz), - num=self.numpoints) + zs = numpy.logspace(0, numpy.log10(self.default_maxz), num=self.numpoints) ds = self.cosmology.luminosity_distance(zs).value - self.faraway_d2z = interpolate.interp1d(ds, zs, kind='linear', - bounds_error=False) + self.faraway_d2z = interpolate.interp1d( + ds, zs, kind="linear", bounds_error=False + ) # store the default maximum distance self.default_maxdist = ds.max() def get_redshift(self, dist): - """Returns the redshift for the given distance. - """ + """Returns the redshift for the given distance.""" dist, input_is_array = pycbc.conversions.ensurearray(dist) try: zs = self.nearby_d2z(dist) @@ -287,10 +301,9 @@ def get_redshift(self, dist): # furthest default; fall back to using astropy if replacemask.any(): # well... check that the distance is positive and finite first - if not (dist > 0.).all() and numpy.isfinite(dist).all(): + if not (dist > 0.0).all() and numpy.isfinite(dist).all(): raise ValueError("distance must be finite and > 0") - zs[replacemask] = _redshift(dist[replacemask], - cosmology=self.cosmology) + zs[replacemask] = _redshift(dist[replacemask], cosmology=self.cosmology) return pycbc.conversions.formatreturn(zs, input_is_array) def __call__(self, dist): @@ -298,12 +311,12 @@ def __call__(self, dist): # set up D(z) interpolating classes for the standard cosmologies -_d2zs = {_c: DistToZ(cosmology=_c) - for _c in parameters.available} +_d2zs = {_c: DistToZ(cosmology=_c) for _c in parameters.available} def redshift(distance, **kwargs): - r"""Returns the redshift associated with the given luminosity distance. + r""" + Returns the redshift associated with the given luminosity distance. If the requested cosmology is one of the pre-defined ones in :py:attr:`astropy.cosmology.parameters.available`, :py:class:`DistToZ` is @@ -323,6 +336,7 @@ def redshift(distance, **kwargs): ------- float : The redshift corresponding to the given distance. + """ cosmology = get_cosmology(**kwargs) try: @@ -333,8 +347,9 @@ def redshift(distance, **kwargs): return z -class ComovingVolInterpolator(object): - r"""Interpolates comoving volume to distance or redshift. +class ComovingVolInterpolator: + r""" + Interpolates comoving volume to distance or redshift. The :mod:`astropy.cosmology` module provides methods for converting any cosmological parameter (like luminosity distance) to redshift. This can be @@ -366,9 +381,12 @@ class speeds that up by pre-interpolating :math:`D(z)`. It works by setting All other keyword args are passed to :py:func:`get_cosmology` to select a cosmology. If none provided, will use :py:attr:`DEFAULT_COSMOLOGY`. + """ - def __init__(self, parameter, default_maxz=10., numpoints=1000, - vol_func=None, **kwargs): + + def __init__( + self, parameter, default_maxz=10.0, numpoints=1000, vol_func=None, **kwargs + ): self.parameter = parameter self.numpoints = int(numpoints) self.default_maxz = default_maxz @@ -391,31 +409,29 @@ def _create_interpolant(self, minz, maxz): zs = z_at_value(self.vol_func, numpy.exp(logvs), self.vol_units, maxz) - if self.parameter != 'redshift': + if self.parameter != "redshift": ys = cosmological_quantity_from_redshift(zs, self.parameter) else: ys = zs - return interpolate.interp1d(logvs, ys, kind='linear', - bounds_error=False) + return interpolate.interp1d(logvs, ys, kind="linear", bounds_error=False) def setup_interpolant(self): """Initializes the z(d) interpolation.""" # get VC bounds # for computing nearby (z < 1) redshifts minz = 0.001 - maxz = 1. + maxz = 1.0 self.nearby_interp = self._create_interpolant(minz, maxz) # for computing far away (z > 1) redshifts - minz = 1. + minz = 1.0 maxz = self.default_maxz self.faraway_interp = self._create_interpolant(minz, maxz) # store the default maximum volume self.default_maxvol = numpy.log(self.vol_func(maxz).value) def get_value_from_logv(self, logv): - """Returns the redshift for the given distance. - """ + """Returns the redshift for the given distance.""" logv, input_is_array = pycbc.conversions.ensurearray(logv) try: vals = self.nearby_interp(logv) @@ -435,13 +451,11 @@ def get_value_from_logv(self, logv): # well... check that the logv is finite first if not numpy.isfinite(logv).all(): raise ValueError("comoving volume must be finite and > 0") - zs = z_at_value(self.vol_func, - numpy.exp(logv[replacemask]), self.vol_units) - if self.parameter == 'redshift': + zs = z_at_value(self.vol_func, numpy.exp(logv[replacemask]), self.vol_units) + if self.parameter == "redshift": vals[replacemask] = zs else: - vals[replacemask] = \ - getattr(self.cosmology, self.parameter)(zs).value + vals[replacemask] = getattr(self.cosmology, self.parameter)(zs).value return pycbc.conversions.formatreturn(vals, input_is_array) def get_value(self, volume): @@ -452,15 +466,19 @@ def __call__(self, volume): # set up D(z) interpolating classes for the standard cosmologies -_v2ds = {_c: ComovingVolInterpolator('luminosity_distance', cosmology=_c) - for _c in parameters.available} +_v2ds = { + _c: ComovingVolInterpolator("luminosity_distance", cosmology=_c) + for _c in parameters.available +} -_v2zs = {_c: ComovingVolInterpolator('redshift', cosmology=_c) - for _c in parameters.available} +_v2zs = { + _c: ComovingVolInterpolator("redshift", cosmology=_c) for _c in parameters.available +} def redshift_from_comoving_volume(vc, interp=True, **kwargs): - r"""Returns the redshift from the given comoving volume. + r""" + Returns the redshift from the given comoving volume. Parameters ---------- @@ -485,6 +503,7 @@ def redshift_from_comoving_volume(vc, interp=True, **kwargs): ------- float : The redshift at the given comoving volume. + """ cosmology = get_cosmology(**kwargs) lookup = _v2zs if interp else {} @@ -498,7 +517,8 @@ def redshift_from_comoving_volume(vc, interp=True, **kwargs): def distance_from_comoving_volume(vc, interp=True, **kwargs): - r"""Returns the luminosity distance from the given comoving volume. + r""" + Returns the luminosity distance from the given comoving volume. Parameters ---------- @@ -522,6 +542,7 @@ def distance_from_comoving_volume(vc, interp=True, **kwargs): ------- float : The luminosity distance at the given comoving volume. + """ cosmology = get_cosmology(**kwargs) lookup = _v2ds if interp else {} @@ -535,9 +556,9 @@ def distance_from_comoving_volume(vc, interp=True, **kwargs): return dist -def cosmological_quantity_from_redshift(z, quantity, strip_unit=True, - **kwargs): - r"""Returns the value of a cosmological quantity (e.g., age) at a redshift. +def cosmological_quantity_from_redshift(z, quantity, strip_unit=True, **kwargs): + r""" + Returns the value of a cosmological quantity (e.g., age) at a redshift. Parameters ---------- @@ -559,6 +580,7 @@ def cosmological_quantity_from_redshift(z, quantity, strip_unit=True, The value of the quantity at the requested value. If ``strip_unit`` is ``True``, will return the value. Otherwise, will return the value with units. + """ cosmology = get_cosmology(**kwargs) val = getattr(cosmology, quantity)(z) @@ -567,7 +589,9 @@ def cosmological_quantity_from_redshift(z, quantity, strip_unit=True, return val -__all__ = ['redshift', 'redshift_from_comoving_volume', - 'distance_from_comoving_volume', - 'cosmological_quantity_from_redshift', - ] +__all__ = [ + "cosmological_quantity_from_redshift", + "distance_from_comoving_volume", + "redshift", + "redshift_from_comoving_volume", +] diff --git a/pycbc/detector/ground.py b/pycbc/detector/ground.py index 26b84fb253c..d7fd704fd72 100644 --- a/pycbc/detector/ground.py +++ b/pycbc/detector/ground.py @@ -25,36 +25,41 @@ # # ============================================================================= # -"""This module provides utilities for calculating detector responses and timing +""" +This module provides utilities for calculating detector responses and timing between ground-based observatories. """ -import os + import logging -import numpy as np -from numpy import cos, sin +import os import lal +import numpy as np from astropy import constants, coordinates, units from astropy.coordinates.matrix_utilities import rotation_matrix -from astropy.units.si import sday, meter +from astropy.units.si import meter, sday +from numpy import cos, sin import pycbc.libutils +from pycbc.time import gmst_accurate from pycbc.types import TimeSeries from pycbc.types.config import InterpolatingConfigParser -from pycbc.time import gmst_accurate -logger = logging.getLogger('pycbc.detector') +logger = logging.getLogger("pycbc.detector") # Response functions are modelled after those in lalsuite and as also # presented in https://arxiv.org/pdf/gr-qc/0008066.pdf + def get_available_detectors(): - """ List the available detectors """ + """List the available detectors""" dets = list(_ground_detectors.keys()) return dets + def get_available_lal_detectors(): - """Return list of detectors known in the currently sourced lalsuite. + """ + Return list of detectors known in the currently sourced lalsuite. This function will query lalsuite about which detectors are known to lalsuite. Detectors are identified by a two character string e.g. 'K1', but also by a longer, and clearer name, e.g. KAGRA. This function returns @@ -67,20 +72,30 @@ def get_available_lal_detectors(): ld = lal.__dict__ known_lal_names = [j for j in ld.keys() if "DETECTOR_PREFIX" in j] known_prefixes = [ld[k] for k in known_lal_names] - known_names = [ld[k.replace('PREFIX', 'NAME')] for k in known_lal_names] + known_names = [ld[k.replace("PREFIX", "NAME")] for k in known_lal_names] return list(zip(known_prefixes, known_names)) + _ground_detectors = {} -def add_detector_on_earth(name, longitude, latitude, - yangle=0, xangle=None, height=0, - xlength=4000, ylength=4000, - xaltitude=0, yaltitude=0): - """ Add a new detector on the earth + +def add_detector_on_earth( + name, + longitude, + latitude, + yangle=0, + xangle=None, + height=0, + xlength=4000, + ylength=4000, + xaltitude=0, + yaltitude=0, +): + """ + Add a new detector on the earth Parameters ---------- - name: str two-letter name to identify the detector longitude: float @@ -99,6 +114,7 @@ def add_detector_on_earth(name, longitude, latitude, height: float The height in meters of the detector above the standard reference ellipsoidal earth + """ if xangle is None: # assume right angle detector if no separate xarm direction given @@ -106,46 +122,49 @@ def add_detector_on_earth(name, longitude, latitude, # baseline response of a single arm pointed in the -X direction resp = np.array([[-1, 0, 0], [0, 0, 0], [0, 0, 0]]) - rm2 = rotation_matrix(-longitude * units.rad, 'z') - rm1 = rotation_matrix(-1.0 * (np.pi / 2.0 - latitude) * units.rad, 'y') - + rm2 = rotation_matrix(-longitude * units.rad, "z") + rm1 = rotation_matrix(-1.0 * (np.pi / 2.0 - latitude) * units.rad, "y") + # Calculate response in earth centered coordinates # by rotation of response in coordinates aligned # with the detector arms resps = [] vecs = [] for angle, azi in [(yangle, yaltitude), (xangle, xaltitude)]: - rm0 = rotation_matrix(angle * units.rad, 'z') - rmN = rotation_matrix(-azi * units.rad, 'y') + rm0 = rotation_matrix(angle * units.rad, "z") + rmN = rotation_matrix(-azi * units.rad, "y") rm = rm2 @ rm1 @ rm0 @ rmN # apply rotation resps.append(rm @ resp @ rm.T / 2.0) vecs.append(rm @ np.array([-1, 0, 0])) - full_resp = (resps[0] - resps[1]) - loc = coordinates.EarthLocation.from_geodetic(longitude * units.rad, - latitude * units.rad, - height=height*units.meter) + full_resp = resps[0] - resps[1] + loc = coordinates.EarthLocation.from_geodetic( + longitude * units.rad, latitude * units.rad, height=height * units.meter + ) loc = np.array([loc.x.value, loc.y.value, loc.z.value]) - _ground_detectors[name] = {'location': loc, - 'response': full_resp, - 'xresp': resps[1], - 'yresp': resps[0], - 'xvec': vecs[1], - 'yvec': vecs[0], - 'yangle': yangle, - 'xangle': xangle, - 'height': height, - 'xaltitude': xaltitude, - 'yaltitude': yaltitude, - 'ylength': ylength, - 'xlength': xlength, - } + _ground_detectors[name] = { + "location": loc, + "response": full_resp, + "xresp": resps[1], + "yresp": resps[0], + "xvec": vecs[1], + "yvec": vecs[0], + "yangle": yangle, + "xangle": xangle, + "height": height, + "xaltitude": xaltitude, + "yaltitude": yaltitude, + "ylength": ylength, + "xlength": xlength, + } + # Notation matches # Eq 4 of https://link.aps.org/accepted/10.1103/PhysRevD.96.084004 def single_arm_frequency_response(f, n, arm_length): - """ The relative amplitude factor of the arm response due to + """ + The relative amplitude factor of the arm response due to signal delay. This is relevant where the long-wavelength approximation no longer applies) """ @@ -156,61 +175,68 @@ def single_arm_frequency_response(f, n, arm_length): c = np.exp(-2.0 * phase) * (1 - np.exp(phase * (1 + n))) / (1 + n) return a * (b - c) * 2.0 # We'll make this relative to the static resp + def load_detector_config(config_files): - """ Add custom detectors from a configuration file + """ + Add custom detectors from a configuration file Parameters ---------- config_files: str or list of strs The config file(s) which specify new detectors + """ - methods = {'earth_normal': (add_detector_on_earth, - ['longitude', 'latitude'])} + methods = {"earth_normal": (add_detector_on_earth, ["longitude", "latitude"])} conf = InterpolatingConfigParser(config_files) - dets = conf.get_subsections('detector') + dets = conf.get_subsections("detector") for det in dets: - kwds = dict(conf.items('detector-{}'.format(det))) + kwds = dict(conf.items(f"detector-{det}")) try: - method, arg_names = methods[kwds.pop('method')] + method, arg_names = methods[kwds.pop("method")] except KeyError: - raise ValueError("Missing or unkown method, " - "options are {}".format(methods.keys())) + raise ValueError( + f"Missing or unkown method, options are {methods.keys()}" + ) for k in kwds: kwds[k] = float(kwds[k]) try: args = [kwds.pop(arg) for arg in arg_names] - except KeyError as e: - raise ValueError("missing required detector argument" - " {} are required".format(arg_names)) + except KeyError: + raise ValueError( + f"missing required detector argument {arg_names} are required" + ) method(det.upper(), *args, **kwds) # prepopulate using detectors hardcoded into lalsuite for pref, name in get_available_lal_detectors(): - lalsim = pycbc.libutils.import_optional('lalsimulation') + lalsim = pycbc.libutils.import_optional("lalsimulation") lal_det = lalsim.DetectorPrefixToLALDetector(pref).frDetector - add_detector_on_earth(pref, - lal_det.vertexLongitudeRadians, - lal_det.vertexLatitudeRadians, - height=lal_det.vertexElevation, - xangle=lal_det.xArmAzimuthRadians, - yangle=lal_det.yArmAzimuthRadians, - xlength=lal_det.xArmMidpoint * 2, - ylength=lal_det.yArmMidpoint * 2, - xaltitude=lal_det.xArmAltitudeRadians, - yaltitude=lal_det.yArmAltitudeRadians, - ) + add_detector_on_earth( + pref, + lal_det.vertexLongitudeRadians, + lal_det.vertexLatitudeRadians, + height=lal_det.vertexElevation, + xangle=lal_det.xArmAzimuthRadians, + yangle=lal_det.yArmAzimuthRadians, + xlength=lal_det.xArmMidpoint * 2, + ylength=lal_det.yArmMidpoint * 2, + xaltitude=lal_det.xArmAltitudeRadians, + yaltitude=lal_det.yArmAltitudeRadians, + ) # autoload detector config files -if 'PYCBC_DETECTOR_CONFIG' in os.environ: - load_detector_config(os.environ['PYCBC_DETECTOR_CONFIG'].split(':')) +if "PYCBC_DETECTOR_CONFIG" in os.environ: + load_detector_config(os.environ["PYCBC_DETECTOR_CONFIG"].split(":")) -class Detector(object): - """A gravitational wave detector - """ +class Detector: + """A gravitational wave detector""" + def __init__(self, detector_name, reference_time=1126259462.0): - """ Create class representing a gravitational-wave detector + """ + Create class representing a gravitational-wave detector + Parameters ---------- detector_name: str @@ -220,21 +246,21 @@ def __init__(self, detector_name, reference_time=1126259462.0): will be estimated from a reference time. If 'None', we will calculate the time for each gps time requested explicitly using a slower but higher precision method. + """ self.name = str(detector_name) - + lal_detectors = [pfx for pfx, name in get_available_lal_detectors()] if detector_name in _ground_detectors: self.info = _ground_detectors[detector_name] - self.response = self.info['response'] - self.location = self.info['location'] + self.response = self.info["response"] + self.location = self.info["location"] else: - raise ValueError("Unkown detector {}".format(detector_name)) + raise ValueError(f"Unkown detector {detector_name}") - loc = coordinates.EarthLocation(self.location[0], - self.location[1], - self.location[2], - unit=meter) + loc = coordinates.EarthLocation( + self.location[0], self.location[1], self.location[2], unit=meter + ) self.latitude = loc.lat.rad self.longitude = loc.lon.rad @@ -247,27 +273,29 @@ def set_gmst_reference(self): self.sday = float(sday.si.scale) self.gmst_reference = gmst_accurate(self.reference_time) else: - raise RuntimeError("Can't get accurate sidereal time without GPS " - "reference time!") + raise RuntimeError( + "Can't get accurate sidereal time without GPS reference time!" + ) def lal(self): - """ Return lal data type detector instance """ + """Return lal data type detector instance""" import lal + d = lal.FrDetector() d.vertexLongitudeRadians = self.longitude d.vertexLatitudeRadians = self.latitude - d.vertexElevation = self.info['height'] - d.xArmAzimuthRadians = self.info['xangle'] - d.yArmAzimuthRadians = self.info['yangle'] - d.xArmAltitudeRadians = self.info['xaltitude'] - d.yArmAltitudeRadians = self.info['yaltitude'] + d.vertexElevation = self.info["height"] + d.xArmAzimuthRadians = self.info["xangle"] + d.yArmAzimuthRadians = self.info["yangle"] + d.xArmAltitudeRadians = self.info["xaltitude"] + d.yArmAltitudeRadians = self.info["yaltitude"] # This is somewhat abused by lalsimulation at the moment # to determine a filter kernel size. We set this only so that # value gets a similar number of samples as other detectors # it is used for nothing else - d.yArmMidpoint = self.info['ylength'] / 2.0 - d.xArmMidpoint = self.info['xlength'] / 2.0 + d.yArmMidpoint = self.info["ylength"] / 2.0 + d.xArmMidpoint = self.info["xlength"] / 2.0 x = lal.Detector() r = lal.CreateDetector(x, d, lal.LALDETECTORTYPE_IFODIFF) @@ -285,23 +313,34 @@ def gmst_estimate(self, gps_time): return gmst def light_travel_time_to_detector(self, det): - """ Return the light travel time from this detector + """ + Return the light travel time from this detector + Parameters ---------- det: Detector The other detector to determine the light travel time to. + Returns ------- time: float The light travel time in seconds + """ d = self.location - det.location - return float(d.dot(d)**0.5 / constants.c.value) - - def antenna_pattern(self, right_ascension, declination, polarization, t_gps, - frequency=0, - polarization_type='tensor'): - """Return the detector response. + return float(d.dot(d) ** 0.5 / constants.c.value) + + def antenna_pattern( + self, + right_ascension, + declination, + polarization, + t_gps, + frequency=0, + polarization_type="tensor", + ): + """ + Return the detector response. Parameters ---------- @@ -320,6 +359,7 @@ def antenna_pattern(self, right_ascension, declination, polarization, t_gps, The plus or vector-x or breathing polarization factor for this sky location / orientation fcross(default) or fy or fl : float or numpy.ndarray The cross or vector-y or longitudnal polarization factor for this sky location / orientation + """ if isinstance(t_gps, lal.LIGOTimeGPS): t_gps = float(t_gps) @@ -338,14 +378,12 @@ def antenna_pattern(self, right_ascension, declination, polarization, t_gps, e2 = sin(declination) nhat = np.array([e0, e1, e2], dtype=object) - nx = nhat.dot(self.info['xvec']) - ny = nhat.dot(self.info['yvec']) + nx = nhat.dot(self.info["xvec"]) + ny = nhat.dot(self.info["yvec"]) - rx = single_arm_frequency_response(frequency, nx, - self.info['xlength']) - ry = single_arm_frequency_response(frequency, ny, - self.info['ylength']) - resp = ry * self.info['yresp'] - rx * self.info['xresp'] + rx = single_arm_frequency_response(frequency, nx, self.info["xlength"]) + ry = single_arm_frequency_response(frequency, ny, self.info["ylength"]) + resp = ry * self.info["yresp"] - rx * self.info["xresp"] ttype = np.complex128 else: resp = self.response @@ -353,27 +391,27 @@ def antenna_pattern(self, right_ascension, declination, polarization, t_gps, x0 = -cospsi * singha - sinpsi * cosgha * sindec x1 = -cospsi * cosgha + sinpsi * singha * sindec - x2 = sinpsi * cosdec + x2 = sinpsi * cosdec x = np.array([x0, x1, x2], dtype=object) dx = resp.dot(x) - y0 = sinpsi * singha - cospsi * cosgha * sindec - y1 = sinpsi * cosgha + cospsi * singha * sindec - y2 = cospsi * cosdec + y0 = sinpsi * singha - cospsi * cosgha * sindec + y1 = sinpsi * cosgha + cospsi * singha * sindec + y2 = cospsi * cosdec y = np.array([y0, y1, y2], dtype=object) dy = resp.dot(y) - if polarization_type != 'tensor': + if polarization_type != "tensor": z0 = -cosdec * cosgha z1 = cosdec * singha z2 = -sindec z = np.array([z0, z1, z2], dtype=object) dz = resp.dot(z) - if polarization_type == 'tensor': - if hasattr(dx, 'shape'): + if polarization_type == "tensor": + if hasattr(dx, "shape"): fplus = (x * dx - y * dy).sum(axis=0).astype(ttype) fcross = (x * dy + y * dx).sum(axis=0).astype(ttype) else: @@ -381,8 +419,8 @@ def antenna_pattern(self, right_ascension, declination, polarization, t_gps, fcross = (x * dy + y * dx).sum() return fplus, fcross - elif polarization_type == 'vector': - if hasattr(dx, 'shape'): + if polarization_type == "vector": + if hasattr(dx, "shape"): fx = (z * dx + x * dz).sum(axis=0).astype(ttype) fy = (z * dy + y * dz).sum(axis=0).astype(ttype) else: @@ -391,8 +429,8 @@ def antenna_pattern(self, right_ascension, declination, polarization, t_gps, return fx, fy - elif polarization_type == 'scalar': - if hasattr(dx, 'shape'): + if polarization_type == "scalar": + if hasattr(dx, "shape"): fb = (x * dx + y * dy).sum(axis=0).astype(ttype) fl = (z * dz).sum(axis=0) else: @@ -401,16 +439,16 @@ def antenna_pattern(self, right_ascension, declination, polarization, t_gps, return fb, fl def time_delay_from_earth_center(self, right_ascension, declination, t_gps): - """Return the time delay from the earth center + """Return the time delay from the earth center""" + return self.time_delay_from_location( + np.array([0, 0, 0]), right_ascension, declination, t_gps + ) + + def time_delay_from_location( + self, other_location, right_ascension, declination, t_gps + ): """ - return self.time_delay_from_location(np.array([0, 0, 0]), - right_ascension, - declination, - t_gps) - - def time_delay_from_location(self, other_location, right_ascension, - declination, t_gps): - """Return the time delay from the given location to detector for + Return the time delay from the given location to detector for a signal with the given sky location In other words return `t1 - t2` where `t1` is the arrival time in this detector and `t2` is the arrival time in the @@ -431,6 +469,7 @@ def time_delay_from_location(self, other_location, right_ascension, ------- float The arrival time difference between the detectors. + """ ra_angle = self.gmst_estimate(t_gps) - right_ascension cosd = cos(declination) @@ -443,13 +482,16 @@ def time_delay_from_location(self, other_location, right_ascension, dx = other_location - self.location return dx.dot(ehat).astype(np.float64) / constants.c.value - def time_delay_from_detector(self, other_detector, right_ascension, - declination, t_gps): - """Return the time delay from the given to detector for a signal with + def time_delay_from_detector( + self, other_detector, right_ascension, declination, t_gps + ): + """ + Return the time delay from the given to detector for a signal with the given sky location; i.e. return `t1 - t2` where `t1` is the arrival time in this detector and `t2` is the arrival time in the other detector. Note that this would return the same value as `time_delay_from_earth_center` if `other_detector` was geocentric. + Parameters ---------- other_detector : detector.Detector @@ -460,19 +502,21 @@ def time_delay_from_detector(self, other_detector, right_ascension, The declination (in rad) of the signal. t_gps : float The GPS time (in s) of the signal. + Returns ------- float The arrival time difference between the detectors. + + """ + return self.time_delay_from_location( + other_detector.location, right_ascension, declination, t_gps + ) + + def arrival_time(self, ref_tc, ra, dec, ref_frame="geocentric"): """ - return self.time_delay_from_location(other_detector.location, - right_ascension, - declination, - t_gps) - - def arrival_time(self, ref_tc, ra, dec, ref_frame='geocentric'): - """Compute the arrival time in this detector. - + Compute the arrival time in this detector. + Parameters ---------- ref_tc : {float, lal.LIGOTimeGPS} @@ -484,34 +528,36 @@ def arrival_time(self, ref_tc, ra, dec, ref_frame='geocentric'): ref_frame : str (optional) The detector to convert from, in which ref_tc is sampled. Default 'geocentric'. - + Returns ------- - float : + float : The coalescence time converted to the current detector frame. + """ - if ref_frame == 'geocentric': + if ref_frame == "geocentric": # from geocenter - tc = ref_tc + \ - self.time_delay_from_earth_center(ra, dec, ref_tc) + tc = ref_tc + self.time_delay_from_earth_center(ra, dec, ref_tc) elif ref_frame == self.name: # no time shift; sampling in current det tc = ref_tc elif ref_frame in get_available_detectors(): # from sampling det refdet = Detector(ref_frame) - tc = ref_tc + \ - self.time_delay_from_detector(refdet, ra, dec, ref_tc) + tc = ref_tc + self.time_delay_from_detector(refdet, ra, dec, ref_tc) else: - raise ValueError(f'Unrecognized ref_frame argument {ref_frame}. ' - 'Accepted arguments are: "geocentric", ' - f'{get_available_detectors()}') + raise ValueError( + f"Unrecognized ref_frame argument {ref_frame}. " + 'Accepted arguments are: "geocentric", ' + f"{get_available_detectors()}" + ) return tc - def project_wave(self, hp, hc, ra, dec, polarization, - method='lal', - reference_time=None): - """Return the strain of a waveform as measured by the detector. + def project_wave( + self, hp, hc, ra, dec, polarization, method="lal", reference_time=None + ): + """ + Return the strain of a waveform as measured by the detector. Apply the time shift for the given detector relative to the assumed geocentric frame and apply the antenna patterns to the plus and cross polarizations. @@ -535,24 +581,35 @@ def project_wave(self, hp, hc, ra, dec, polarization, The time to use as, a reference for some methods of projection. Used by 'constant' and 'vary_polarization' methods. Uses average time if not provided. + """ # The robust and most fefature rich method which includes # time changing antenna patterns and doppler shifts due to the # earth rotation and orbit - if method == 'lal': + if method == "lal": import lalsimulation + h_lal = lalsimulation.SimDetectorStrainREAL8TimeSeries( - hp.astype(np.float64).lal(), hc.astype(np.float64).lal(), - ra, dec, polarization, self.lal()) + hp.astype(np.float64).lal(), + hc.astype(np.float64).lal(), + ra, + dec, + polarization, + self.lal(), + ) ts = TimeSeries( - h_lal.data.data, delta_t=h_lal.deltaT, epoch=h_lal.epoch, - dtype=np.float64, copy=False) + h_lal.data.data, + delta_t=h_lal.deltaT, + epoch=h_lal.epoch, + dtype=np.float64, + copy=False, + ) # 'constant' assume fixed orientation relative to source over the # duration of the signal, accurate for short duration signals # 'fixed_polarization' applies only time changing orientation # but no doppler corrections - elif method in ['constant', 'vary_polarization']: + elif method in ["constant", "vary_polarization"]: if reference_time is not None: rtime = reference_time else: @@ -562,13 +619,14 @@ def project_wave(self, hp, hc, ra, dec, polarization, # the midpoint time. rtime = (float(hp.end_time) + float(hp.start_time)) / 2.0 - if method == 'constant': + if method == "constant": time = rtime - elif method == 'vary_polarization': - if (not isinstance(hp, TimeSeries) or - not isinstance(hc, TimeSeries)): - raise TypeError('Waveform polarizations must be given' - ' as time series for this method') + elif method == "vary_polarization": + if not isinstance(hp, TimeSeries) or not isinstance(hc, TimeSeries): + raise TypeError( + "Waveform polarizations must be given" + " as time series for this method" + ) # this is more granular than needed, may be optimized later # assume earth rotation in ~30 ms needed for earth ceneter @@ -583,11 +641,12 @@ def project_wave(self, hp, hc, ra, dec, polarization, # add in only the correction for the time variance in the polarization # due to the earth's rotation, no doppler correction applied else: - raise ValueError("Unkown projection method {}".format(method)) + raise ValueError(f"Unkown projection method {method}") return ts def optimal_orientation(self, t_gps): - """Return the optimal orientation in right ascension and declination + """ + Return the optimal orientation in right ascension and declination for a given GPS time. Parameters @@ -601,30 +660,39 @@ def optimal_orientation(self, t_gps): Right ascension that is optimally oriented for the detector dec: float Declination that is optimally oriented for the detector + """ - ra = self.longitude + (self.gmst_estimate(t_gps) % (2.0*np.pi)) + ra = self.longitude + (self.gmst_estimate(t_gps) % (2.0 * np.pi)) dec = self.latitude return ra, dec def get_icrs_pos(self): - """ Transforms GCRS frame to ICRS frame + """ + Transforms GCRS frame to ICRS frame Returns - ---------- + ------- loc: numpy.ndarray shape (3,1) units: AU ICRS coordinates in cartesian system + """ loc = self.location - loc = coordinates.SkyCoord(x=loc[0], y=loc[1], z=loc[2], unit=units.m, - frame='gcrs', representation_type='cartesian').transform_to('icrs') - loc.representation_type = 'cartesian' - conv = np.float32(((loc.x.unit/units.AU).decompose()).to_string()) - loc = np.array([np.float32(loc.x), np.float32(loc.y), - np.float32(loc.z)])*conv + loc = coordinates.SkyCoord( + x=loc[0], + y=loc[1], + z=loc[2], + unit=units.m, + frame="gcrs", + representation_type="cartesian", + ).transform_to("icrs") + loc.representation_type = "cartesian" + conv = np.float32(((loc.x.unit / units.AU).decompose()).to_string()) + loc = np.array([np.float32(loc.x), np.float32(loc.y), np.float32(loc.z)]) * conv return loc def effective_distance(self, distance, ra, dec, pol, time, inclination): - """ Distance scaled to account for amplitude factors + """ + Distance scaled to account for amplitude factors The effective distance of the source. This scales the distance so that the amplitude is equal to a source which is optimally oriented with @@ -650,61 +718,73 @@ def effective_distance(self, distance, ra, dec, pol, time, inclination): ------- eff_dist: float The effective distance of the source + """ fp, fc = self.antenna_pattern(ra, dec, pol, time) ic = np.cos(inclination) - ip = 0.5 * (1. + ic * ic) + ip = 0.5 * (1.0 + ic * ic) scale = ((fp * ip) ** 2.0 + (fc * ic) ** 2.0) ** 0.5 return distance / scale + def overhead_antenna_pattern(right_ascension, declination, polarization): - """Return the antenna pattern factors F+ and Fx as a function of sky + """ + Return the antenna pattern factors F+ and Fx as a function of sky location and polarization angle for a hypothetical interferometer located at the north pole. Angles are in radians. Declinations of ±π/2 correspond to the normal to the detector plane (i.e. overhead and underneath) while the point with zero right ascension and declination is the direction of one of the interferometer arms. + Parameters ---------- right_ascension: float declination: float polarization: float + Returns ------- f_plus: float f_cros: float + """ # convert from declination coordinate to polar (angle dropped from north axis) theta = np.pi / 2.0 - declination - f_plus = - (1.0/2.0) * (1.0 + cos(theta)*cos(theta)) * \ - cos (2.0 * right_ascension) * cos (2.0 * polarization) - \ - cos(theta) * sin(2.0*right_ascension) * sin (2.0 * polarization) + f_plus = -(1.0 / 2.0) * (1.0 + cos(theta) * cos(theta)) * cos( + 2.0 * right_ascension + ) * cos(2.0 * polarization) - cos(theta) * sin(2.0 * right_ascension) * sin( + 2.0 * polarization + ) - f_cross = (1.0/2.0) * (1.0 + cos(theta)*cos(theta)) * \ - cos (2.0 * right_ascension) * sin (2.0* polarization) - \ - cos(theta) * sin(2.0*right_ascension) * cos (2.0 * polarization) + f_cross = (1.0 / 2.0) * (1.0 + cos(theta) * cos(theta)) * cos( + 2.0 * right_ascension + ) * sin(2.0 * polarization) - cos(theta) * sin(2.0 * right_ascension) * cos( + 2.0 * polarization + ) return f_plus, f_cross -def ppdets(ifos, separator=', '): - """Pretty-print a list (or set) of detectors: return a string listing +def ppdets(ifos, separator=", "): + """ + Pretty-print a list (or set) of detectors: return a string listing the given detectors alphabetically and separated by the given string (comma by default). """ if ifos: return separator.join(sorted(ifos)) - return 'no detectors' + return "no detectors" + __all__ = [ - 'Detector', - 'get_available_detectors', - 'get_available_lal_detectors', - 'add_detector_on_earth', - 'single_arm_frequency_response', - 'ppdets', - 'overhead_antenna_pattern', - 'load_detector_config', - '_ground_detectors', + "Detector", + "_ground_detectors", + "add_detector_on_earth", + "get_available_detectors", + "get_available_lal_detectors", + "load_detector_config", + "overhead_antenna_pattern", + "ppdets", + "single_arm_frequency_response", ] diff --git a/pycbc/detector/space.py b/pycbc/detector/space.py index ac1233cb8e7..675ddc3e418 100644 --- a/pycbc/detector/space.py +++ b/pycbc/detector/space.py @@ -10,26 +10,33 @@ This module provides utilities for simulating the GW response of space-based observatories. """ + +import logging from abc import ABC, abstractmethod -from pycbc.coordinates.space import TIME_OFFSET_20_DEGREES -from pycbc.types import TimeSeries + import numpy -from numpy import cos, sin from astropy import constants -import logging +from numpy import cos, sin + +from pycbc.coordinates.space import TIME_OFFSET_20_DEGREES +from pycbc.types import TimeSeries + def get_available_space_detectors(): """List the available space detectors""" dets = list(_space_detectors.keys()) aliases = [] for i in dets: - aliases.extend(_space_detectors[i]['aliases']) + aliases.extend(_space_detectors[i]["aliases"]) return dets + aliases + def parse_det_name(detector_name): - """Parse a string into a detector name and TDI channel. - The input is assumed to look like '{detector name}_{channel name}.'""" - out = detector_name.split('_', 1) + """ + Parse a string into a detector name and TDI channel. + The input is assumed to look like '{detector name}_{channel name}.' + """ + out = detector_name.split("_", 1) det = out[0] try: chan = out[1] @@ -38,6 +45,7 @@ def parse_det_name(detector_name): chan = None return det, chan + def apply_polarization(hp, hc, polarization): """ Apply polarization rotation matrix. @@ -58,17 +66,26 @@ def apply_polarization(hp, hc, polarization): (array, array) The plus and cross polarizations of the GW rotated by the polarization angle. + """ - cphi = cos(2*polarization) - sphi = sin(2*polarization) + cphi = cos(2 * polarization) + sphi = sin(2 * polarization) - hp_ssb = hp*cphi - hc*sphi - hc_ssb = hp*sphi + hc*cphi + hp_ssb = hp * cphi - hc * sphi + hc_ssb = hp * sphi + hc * cphi return hp_ssb, hc_ssb -def check_signal_times(hp, hc, orbit_start_time, orbit_end_time, - offset=TIME_OFFSET_20_DEGREES, pad_data=False, t0=1e4): + +def check_signal_times( + hp, + hc, + orbit_start_time, + orbit_end_time, + offset=TIME_OFFSET_20_DEGREES, + pad_data=False, + t0=1e4, +): """ Ensure that input signal lies within the provided orbital window. This assumes that the start times of hp and hc are relative to the detector @@ -110,6 +127,7 @@ def check_signal_times(hp, hc, orbit_start_time, orbit_end_time, (pycbc.types.TimeSeries, pycbc.types.TimeSeries) The plus and cross polarizations of the GW in the SSB frame, padded as requested and/or truncated to fit in the orbital window. + """ dt = hp.delta_t @@ -119,7 +137,7 @@ def check_signal_times(hp, hc, orbit_start_time, orbit_end_time, # pad the data with zeros if pad_data: - pad_idx = int(t0/dt) + pad_idx = int(t0 / dt) hp.prepend_zeros(pad_idx) hp.append_zeros(pad_idx) hc.prepend_zeros(pad_idx) @@ -127,16 +145,20 @@ def check_signal_times(hp, hc, orbit_start_time, orbit_end_time, # make sure signal lies within orbit length if hp.duration + hp.start_time > orbit_end_time: - logging.warning('Time of signal end is greater than end of orbital ' + - f'data. Cutting signal at {orbit_end_time}.') + logging.warning( + "Time of signal end is greater than end of orbital " + f"data. Cutting signal at {orbit_end_time}." + ) # cut off data succeeding orbit end time end_idx = numpy.argwhere(hp.sample_times.numpy() <= orbit_end_time)[-1][0] hp = hp[:end_idx] hc = hc[:end_idx] if hp.start_time < orbit_start_time: - logging.warning('Time of signal start is less than start of orbital ' + - f'data. Cutting signal at {orbit_start_time}.') + logging.warning( + "Time of signal start is less than start of orbital " + f"data. Cutting signal at {orbit_start_time}." + ) # cut off data preceding orbit start time start_idx = numpy.argwhere(hp.sample_times.numpy() >= orbit_start_time)[0][0] hp = hp[start_idx:] @@ -144,6 +166,7 @@ def check_signal_times(hp, hc, orbit_start_time, orbit_end_time, return hp, hc + def cut_channels(tdi_dict, remove_garbage=False, t0=1e4): """ Cut TDI channels if needed. @@ -163,12 +186,13 @@ def cut_channels(tdi_dict, remove_garbage=False, t0=1e4): t0 : float (optional) Time in seconds to cut/zero from data if remove_garbage is True/'zero'. Default 1e4. + """ for chan in tdi_dict.keys(): if remove_garbage: dt = tdi_dict[chan].delta_t - pad_idx = int(t0/dt) - if remove_garbage == 'zero': + pad_idx = int(t0 / dt) + if remove_garbage == "zero": # zero the edge data tdi_dict[chan][:pad_idx] = 0 tdi_dict[chan][-pad_idx:] = 0 @@ -181,11 +205,14 @@ def cut_channels(tdi_dict, remove_garbage=False, t0=1e4): return tdi_dict -_space_detectors = {'LISA': {'armlength': 2.5e9, - 'aliases': ['LISA_A', 'LISA_E', 'LISA_T', - 'LISA_X', 'LISA_Y', 'LISA_Z'], - }, - } + +_space_detectors = { + "LISA": { + "armlength": 2.5e9, + "aliases": ["LISA_A", "LISA_E", "LISA_T", "LISA_X", "LISA_Y", "LISA_Z"], + }, +} + class AbsSpaceDet(ABC): """ @@ -196,18 +223,22 @@ class AbsSpaceDet(ABC): detector_name : str The name of the detector. Accepts any output from `get_available_space_detectors`. - + reference_time : float (optional) The reference time in seconds of the signal in the SSB frame. This is defined such that the detector mission start time corresponds to 0. Default None. + """ + def __init__(self, detector_name, reference_time=None, **kwargs): self.det, self.chan = parse_det_name(detector_name) if detector_name not in get_available_space_detectors(): - raise NotImplementedError('Unrecognized detector. ', - 'Currently accepts: ', - f'{get_available_space_detectors()}') + raise NotImplementedError( + "Unrecognized detector. ", + "Currently accepts: ", + f"{get_available_space_detectors()}", + ) self.reference_time = reference_time @property @@ -239,7 +270,7 @@ class _LDC_detector(AbsSpaceDet): detector_name : str The name of the detector. Accepts any output from `get_available_space_detectors`. - + reference_time : float (optional) The reference time in seconds of the signal in the SSB frame. This is defined such that the detector mission start time corresponds to 0. @@ -253,23 +284,31 @@ class _LDC_detector(AbsSpaceDet): offset : float (optional) The time in seconds by which to offset the input waveform if apply_offset is True. Default 7365189.431698299. - + orbits : str (optional) The constellation orbital data used for generating projections and TDI. See self.orbits_init for accepted inputs. Default 'EqualArmlength'. + """ - def __init__(self, detector_name, reference_time=None, apply_offset=False, - offset=TIME_OFFSET_20_DEGREES, - orbits='EqualArmlength', **kwargs): + + def __init__( + self, + detector_name, + reference_time=None, + apply_offset=False, + offset=TIME_OFFSET_20_DEGREES, + orbits="EqualArmlength", + **kwargs, + ): super().__init__(detector_name, reference_time, **kwargs) - assert self.det == 'LISA', 'LDC backend only works with LISA detector' + assert self.det == "LISA", "LDC backend only works with LISA detector" # specify whether to apply offsets to GPS times if apply_offset: self.offset = offset else: - self.offset = 0. + self.offset = 0.0 # orbits properties self.orbits = orbits @@ -289,13 +328,13 @@ def __init__(self, detector_name, reference_time=None, apply_offset=False, # class initialization self.proj_init = None self.tdi_init = None - self.tdi_chan = 'AET' - if self.chan is not None and self.chan in 'XYZ': - self.tdi_chan = 'XYZ' + self.tdi_chan = "AET" + if self.chan is not None and self.chan in "XYZ": + self.tdi_chan = "XYZ" @property def sky_coords(self): - return 'eclipticlongitude', 'eclipticlatitude' + return "eclipticlongitude", "eclipticlatitude" def orbits_init(self, orbits, size=316, dt=100000.0, t_init=0.0): """ @@ -322,38 +361,43 @@ def orbits_init(self, orbits, size=316, dt=100000.0, t_init=0.0): t_init : float (optional) The start time in seconds to use if generating a new orbit file. Default 0. + """ - defaults = ['EqualArmlength', 'Keplerian'] - assert type(orbits) == str, ('Must input either a file path as ', - 'str, "EqualArmlength", or "Keplerian"') + defaults = ["EqualArmlength", "Keplerian"] + assert type(orbits) == str, ( + "Must input either a file path as ", + 'str, "EqualArmlength", or "Keplerian"', + ) # generate a new file if orbits in defaults: try: import lisaorbits except ImportError: - raise ImportError('lisaorbits not found') - if orbits == 'EqualArmlength': + raise ImportError("lisaorbits not found") + if orbits == "EqualArmlength": o = lisaorbits.EqualArmlengthOrbits() - if orbits == 'Keplerian': + if orbits == "Keplerian": o = lisaorbits.KeplerianOrbits() - o.write('orbits.h5', dt=dt, size=size, t0=t_init, mode='w') - ofile = 'orbits.h5' + o.write("orbits.h5", dt=dt, size=size, t0=t_init, mode="w") + ofile = "orbits.h5" self.orbits_start_time = t_init - self.orbits_end_time = t_init + size*dt + self.orbits_end_time = t_init + size * dt self.orbits = ofile # read in from an existing file path else: import h5py + ofile = orbits - with h5py.File(ofile, 'r') as f: - self.orbits_start_time = f.attrs['t0'] - self.orbits_end_time = self.orbit_start_time + \ - f.attrs['dt']*f.attrs['size'] + with h5py.File(ofile, "r") as f: + self.orbits_start_time = f.attrs["t0"] + self.orbits_end_time = ( + self.orbit_start_time + f.attrs["dt"] * f.attrs["size"] + ) # add light travel buffer times - lisa_arm = _space_detectors['LISA']['armlength'] + lisa_arm = _space_detectors["LISA"]["armlength"] ltt_au = constants.au.value / constants.c.value ltt_arm = lisa_arm / constants.c.value self.orbits_start_time += ltt_arm + ltt_au @@ -378,30 +422,31 @@ def strain_container(self, response, orbits=None): dict, array The arguments and measurements associated with the link and orbital data. + """ try: from pytdi import Data except ImportError: - raise ImportError('pyTDI required for TDI combinations') + raise ImportError("pyTDI required for TDI combinations") - links = ['12', '23', '31', '13', '32', '21'] + links = ["12", "23", "31", "13", "32", "21"] # format the measurements from link data measurements = {} for i, link in enumerate(links): - measurements[f'isi_{link}'] = response[:, i] - measurements[f'isi_sb_{link}'] = response[:, i] - measurements[f'tmi_{link}'] = 0. - measurements[f'rfi_{link}'] = 0. - measurements[f'rfi_sb_{link}'] = 0. + measurements[f"isi_{link}"] = response[:, i] + measurements[f"isi_sb_{link}"] = response[:, i] + measurements[f"tmi_{link}"] = 0.0 + measurements[f"rfi_{link}"] = 0.0 + measurements[f"rfi_sb_{link}"] = 0.0 - df = 1/self.dt + df = 1 / self.dt t_init = self.orbits_start_time # call in the orbital data using pyTDI if orbits is None: orbits = self.orbits - return Data.from_orbits(orbits, df, t_init, 'tcb/ltt', **measurements) + return Data.from_orbits(orbits, df, t_init, "tcb/ltt", **measurements) def get_links(self, hp, hc, lamb, beta, polarization): """ @@ -429,20 +474,27 @@ def get_links(self, hp, hc, lamb, beta, polarization): ndarray The waveform projected to the LISA laser links. Shape is (6, N) for input waveforms with N total samples. + """ try: from lisagwresponse import ReadStrain except ImportError: - raise ImportError('LISA GW Response not found') + raise ImportError("LISA GW Response not found") if self.dt is None: self.dt = hp.delta_t # configure orbits and signal self.orbits_init(orbits=self.orbits) - hp, hc = check_signal_times(hp, hc, self.orbits_start_time, - self.orbits_end_time, offset=self.offset, - pad_data=self.pad_data, t0=self.t0) + hp, hc = check_signal_times( + hp, + hc, + self.orbits_start_time, + self.orbits_end_time, + offset=self.offset, + pad_data=self.pad_data, + t0=self.t0, + ) self.start_time = hp.start_time - self.offset self.sample_times = hp.sample_times.numpy() @@ -451,9 +503,14 @@ def get_links(self, hp, hc, lamb, beta, polarization): if self.proj_init is None: # initialize the class - self.proj_init = ReadStrain(self.sample_times, hp, hc, - gw_beta=beta, gw_lambda=lamb, - orbits=self.orbits) + self.proj_init = ReadStrain( + self.sample_times, + hp, + hc, + gw_beta=beta, + gw_lambda=lamb, + orbits=self.orbits, + ) else: # update params in the initialized class self.proj_init.gw_beta = beta @@ -461,14 +518,26 @@ def get_links(self, hp, hc, lamb, beta, polarization): self.proj_init.set_strain(self.sample_times, hp, hc) # project the signal - wf_proj = self.proj_init.compute_gw_response(self.sample_times, - self.proj_init.LINKS) + wf_proj = self.proj_init.compute_gw_response( + self.sample_times, self.proj_init.LINKS + ) return wf_proj - def project_wave(self, hp, hc, lamb, beta, polarization=0, - tdi=1.5, tdi_chan=None, pad_data=False, - remove_garbage=False, t0=1e4, **kwargs): + def project_wave( + self, + hp, + hc, + lamb, + beta, + polarization=0, + tdi=1.5, + tdi_chan=None, + pad_data=False, + remove_garbage=False, + t0=1e4, + **kwargs, + ): """ Evaluate the TDI observables. @@ -526,11 +595,12 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, dict ({str: pycbc.types.TimeSeries}) The TDI observables as TimeSeries objects keyed by their corresponding TDI channel name. + """ try: from pytdi import michelson except ImportError: - raise ImportError('pyTDI not found') + raise ImportError("pyTDI not found") # set TDI generation if tdi == 1.5: @@ -538,8 +608,9 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, elif tdi == 2: X, Y, Z = michelson.X2, michelson.Y2, michelson.Z2 else: - raise ValueError('Unrecognized TDI generation input. ' + - 'Please input either 1 or 2.') + raise ValueError( + "Unrecognized TDI generation input. " + "Please input either 1 or 2." + ) # set TDI channels if tdi_chan is None: @@ -549,8 +620,7 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, self.pad_data = pad_data self.remove_garbage = remove_garbage self.t0 = t0 - response = self.get_links(hp, hc, lamb, beta, - polarization=polarization) + response = self.get_links(hp, hc, lamb, beta, polarization=polarization) # load in data using response measurements self.tdi_init = self.strain_container(response, self.orbits) @@ -561,30 +631,31 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, chanz = Z.build(**self.tdi_init.args)(self.tdi_init.measurements) # convert to AET if specified - if tdi_chan == 'XYZ': - tdi_dict = {'LISA_X': TimeSeries(chanx, delta_t=self.dt, - epoch=self.start_time), - 'LISA_Y': TimeSeries(chany, delta_t=self.dt, - epoch=self.start_time), - 'LISA_Z': TimeSeries(chanz, delta_t=self.dt, - epoch=self.start_time)} - elif tdi_chan == 'AET': - chana = (chanz - chanx)/numpy.sqrt(2) - chane = (chanx - 2*chany + chanz)/numpy.sqrt(6) - chant = (chanx + chany + chanz)/numpy.sqrt(3) - tdi_dict = {'LISA_A': TimeSeries(chana, delta_t=self.dt, - epoch=self.start_time), - 'LISA_E': TimeSeries(chane, delta_t=self.dt, - epoch=self.start_time), - 'LISA_T': TimeSeries(chant, delta_t=self.dt, - epoch=self.start_time)} + if tdi_chan == "XYZ": + tdi_dict = { + "LISA_X": TimeSeries(chanx, delta_t=self.dt, epoch=self.start_time), + "LISA_Y": TimeSeries(chany, delta_t=self.dt, epoch=self.start_time), + "LISA_Z": TimeSeries(chanz, delta_t=self.dt, epoch=self.start_time), + } + elif tdi_chan == "AET": + chana = (chanz - chanx) / numpy.sqrt(2) + chane = (chanx - 2 * chany + chanz) / numpy.sqrt(6) + chant = (chanx + chany + chanz) / numpy.sqrt(3) + tdi_dict = { + "LISA_A": TimeSeries(chana, delta_t=self.dt, epoch=self.start_time), + "LISA_E": TimeSeries(chane, delta_t=self.dt, epoch=self.start_time), + "LISA_T": TimeSeries(chant, delta_t=self.dt, epoch=self.start_time), + } else: - raise ValueError('Unrecognized TDI channel input. ' + - 'Please input either "XYZ" or "AET".') + raise ValueError( + "Unrecognized TDI channel input. " + 'Please input either "XYZ" or "AET".' + ) # processing - tdi_dict = cut_channels(tdi_dict, remove_garbage=self.remove_garbage, - t0=self.t0) + tdi_dict = cut_channels( + tdi_dict, remove_garbage=self.remove_garbage, t0=self.t0 + ) return tdi_dict @@ -600,7 +671,7 @@ class _FLR_detector(AbsSpaceDet): detector_name : str The name of the detector. Accepts any output from `get_available_space_detectors`. - + reference_time : float (optional) The reference time in seconds of the signal in the SSB frame. This is defined such that the detector mission start time corresponds to 0. @@ -622,22 +693,33 @@ class _FLR_detector(AbsSpaceDet): The constellation orbital data used for generating projections and TDI. See self.orbits_init for accepted inputs. Default 'EqualArmlength'. + """ - def __init__(self, detector_name, reference_time=None, apply_offset=False, - offset=TIME_OFFSET_20_DEGREES, - orbits='EqualArmlength', use_gpu=False, **kwargs): - logging.warning('WARNING: FastLISAResponse TDI implementation is a ', - 'work in progress. Currently unable to reproduce LDC ', - 'or BBHx waveforms.') + + def __init__( + self, + detector_name, + reference_time=None, + apply_offset=False, + offset=TIME_OFFSET_20_DEGREES, + orbits="EqualArmlength", + use_gpu=False, + **kwargs, + ): + logging.warning( + "WARNING: FastLISAResponse TDI implementation is a ", + "work in progress. Currently unable to reproduce LDC ", + "or BBHx waveforms.", + ) self.use_gpu = use_gpu super().__init__(detector_name, reference_time, **kwargs) - assert self.det == 'LISA', 'FLR backend only works with LISA detector' + assert self.det == "LISA", "FLR backend only works with LISA detector" # specify whether to apply offsets to GPS times if apply_offset: self.offset = offset else: - self.offset = 0. + self.offset = 0.0 # orbits properties self.orbits = orbits @@ -656,14 +738,14 @@ def __init__(self, detector_name, reference_time=None, apply_offset=False, # class initialization self.tdi_init = None - self.tdi_chan = 'AET' - if 'tdi_chan' in kwargs.keys(): - if kwargs['tdi_chan'] is not None and kwargs['tdi_chan'] in 'XYZ': - self.tdi_chan = 'XYZ' + self.tdi_chan = "AET" + if "tdi_chan" in kwargs: + if kwargs["tdi_chan"] is not None and kwargs["tdi_chan"] in "XYZ": + self.tdi_chan = "XYZ" @property def sky_coords(self): - return 'eclipticlongitude', 'eclipticlatitude' + return "eclipticlongitude", "eclipticlatitude" def orbits_init(self, orbits): """ @@ -676,6 +758,7 @@ def orbits_init(self, orbits): corresponding Orbits class from LISA Analysis Tools is called. Else, the input is treated as a file path following LISA Orbits format. + """ # if orbits are already a class instance, skip this if type(self.orbits) is not (str or None): @@ -687,18 +770,20 @@ def orbits_init(self, orbits): raise ImportError("LISA Analysis Tools required for FLR orbits") # load an orbit from lisatools - defaults = ['EqualArmlength', 'ESA'] + defaults = ["EqualArmlength", "ESA"] if orbits in defaults: - if orbits == 'EqualArmlength': + if orbits == "EqualArmlength": o = detector.EqualArmlengthOrbits() - if orbits == 'ESA': + if orbits == "ESA": o = detector.ESAOrbits() # create a new orbits instance for file input else: + class CustomOrbits(detector.Orbits): def __init__(self): super().__init__(orbits) + o = CustomOrbits() self.orbits = o @@ -735,23 +820,30 @@ def get_links(self, hp, hc, lamb, beta, polarization=0, use_gpu=None): ndarray The waveform projected to the LISA laser links. Shape is (6, N) for input waveforms with N total samples. + """ try: from fastlisaresponse import pyResponseTDI except ImportError: - raise ImportError('FastLISAResponse not found') + raise ImportError("FastLISAResponse not found") if self.dt is None: self.dt = hp.delta_t # configure the orbit and signal self.orbits_init(orbits=self.orbits) - hp, hc = check_signal_times(hp, hc, self.orbits_start_time, - self.orbits_end_time, offset=self.offset, - pad_data=self.pad_data, t0=self.t0) + hp, hc = check_signal_times( + hp, + hc, + self.orbits_start_time, + self.orbits_end_time, + offset=self.offset, + pad_data=self.pad_data, + t0=self.t0, + ) self.start_time = hp.start_time - self.offset self.sample_times = hp.sample_times.numpy() - + # apply polarization hp, hc = apply_polarization(hp, hc, polarization) @@ -761,7 +853,7 @@ def get_links(self, hp, hc, lamb, beta, polarization=0, use_gpu=None): # format wf to hp + i*hc hp = hp.numpy() hc = hc.numpy() - wf = hp + 1j*hc + wf = hp + 1j * hc if use_gpu is None: use_gpu = self.use_gpu @@ -769,16 +861,17 @@ def get_links(self, hp, hc, lamb, beta, polarization=0, use_gpu=None): # convert to cupy if needed if use_gpu: import cupy + wf = cupy.asarray(wf) if self.tdi_init is None: # initialize the class - self.tdi_init = pyResponseTDI(1/self.dt, len(wf), - orbits=self.orbits, - use_gpu=use_gpu) + self.tdi_init = pyResponseTDI( + 1 / self.dt, len(wf), orbits=self.orbits, use_gpu=use_gpu + ) else: # update params in the initialized class - self.tdi_init.sampling_frequency = 1/self.dt + self.tdi_init.sampling_frequency = 1 / self.dt self.tdi_init.num_pts = len(wf) self.tdi_init.orbits = self.orbits self.tdi_init.use_gpu = use_gpu @@ -789,9 +882,21 @@ def get_links(self, hp, hc, lamb, beta, polarization=0, use_gpu=None): return wf_proj - def project_wave(self, hp, hc, lamb, beta, polarization=0, - tdi=1.5, tdi_chan=None, use_gpu=None, pad_data=False, - remove_garbage=False, t0=1e4, **kwargs): + def project_wave( + self, + hp, + hc, + lamb, + beta, + polarization=0, + tdi=1.5, + tdi_chan=None, + use_gpu=None, + pad_data=False, + remove_garbage=False, + t0=1e4, + **kwargs, + ): """ Evaluate the TDI observables. @@ -852,6 +957,7 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, dict ({str: pycbc.types.TimeSeries}) The TDI observables as TimeSeries objects keyed by their corresponding TDI channel name. + """ # set use_gpu if use_gpu is None: @@ -861,14 +967,13 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, self.pad_data = pad_data self.remove_garbage = remove_garbage self.t0 = t0 - self.get_links(hp, hc, lamb, beta, polarization=polarization, - use_gpu=use_gpu) + self.get_links(hp, hc, lamb, beta, polarization=polarization, use_gpu=use_gpu) # set TDI configuration (let FLR handle if not 1 or 2) if tdi == 1.5: - tdi_opt = '1st generation' + tdi_opt = "1st generation" elif tdi == 2: - tdi_opt = '2nd generation' + tdi_opt = "2nd generation" else: tdi_opt = tdi @@ -881,10 +986,10 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, if tdi_chan is None: tdi_chan = self.tdi_chan - if tdi_chan in ['XYZ', 'AET', 'AE']: + if tdi_chan in ["XYZ", "AET", "AE"]: self.tdi_init.tdi_chan = tdi_chan else: - raise ValueError('TDI channels must be one of: XYZ, AET, AE') + raise ValueError("TDI channels must be one of: XYZ, AET, AE") # generate the TDI channels tdi_obs = self.tdi_init.get_tdi_delays() @@ -893,18 +998,23 @@ def project_wave(self, hp, hc, lamb, beta, polarization=0, tdi_dict = {} for i, chan in enumerate(tdi_chan): # save as TimeSeries - tdi_dict[f'LISA_{chan}'] = TimeSeries(tdi_obs[i], delta_t=self.dt, - epoch=self.start_time) + tdi_dict[f"LISA_{chan}"] = TimeSeries( + tdi_obs[i], delta_t=self.dt, epoch=self.start_time + ) - tdi_dict = cut_channels(tdi_dict, remove_garbage=self.remove_garbage, - t0=self.t0) + tdi_dict = cut_channels( + tdi_dict, remove_garbage=self.remove_garbage, t0=self.t0 + ) return tdi_dict -_backends = {'LISA': {'LDC': _LDC_detector, - 'FLR': _FLR_detector, - }, - } +_backends = { + "LISA": { + "LDC": _LDC_detector, + "FLR": _FLR_detector, + }, +} + class SpaceDetector(AbsSpaceDet): """ @@ -915,7 +1025,7 @@ class SpaceDetector(AbsSpaceDet): detector_name : str The name of the detector. Accepts any output from `get_available_space_detectors`. - + reference_time : float (optional) The reference time in seconds of the signal in the SSB frame. This is defined such that the detector mission start time corresponds to 0. @@ -924,17 +1034,19 @@ class SpaceDetector(AbsSpaceDet): backend : str (optional) The backend architecture to use for generating TDI. Accepts 'LDC' or 'FLR'. Default 'LDC'. + """ - def __init__(self, detector_name, reference_time=None, backend='LDC', - **kwargs): + + def __init__(self, detector_name, reference_time=None, backend="LDC", **kwargs): super().__init__(detector_name, reference_time, **kwargs) if backend in _backends[self.det].keys(): c = _backends[self.det][backend] self.backend = c(detector_name, reference_time, **kwargs) else: - raise ValueError(f'Detector {self.det} does not support backend ', - f'{backend}.This detector accepts: ' - f'{_backends[self.det].keys()}') + raise ValueError( + f"Detector {self.det} does not support backend ", + f"{backend}.This detector accepts: {_backends[self.det].keys()}", + ) @property def sky_coords(self): @@ -947,5 +1059,8 @@ def project_wave(self, hp, hc, lamb, beta, *args, **kwargs): return self.backend.project_wave(hp, hc, lamb, beta, *args, **kwargs) -__all__ = ['get_available_space_detectors', 'SpaceDetector', - '_space_detectors',] \ No newline at end of file +__all__ = [ + "SpaceDetector", + "_space_detectors", + "get_available_space_detectors", +] diff --git a/pycbc/distributions/__init__.py b/pycbc/distributions/__init__.py index 3ee7f2416c1..e10196deb5e 100644 --- a/pycbc/distributions/__init__.py +++ b/pycbc/distributions/__init__.py @@ -17,57 +17,69 @@ This modules provides classes and functions for drawing and calculating the probability density function of distributions. """ + # imports needed for functions below import configparser as _ConfigParser -from pycbc.distributions import constraints -from pycbc import VARARGS_DELIM as _VARARGS_DELIM -# Promote some classes/functions to the distributions name space -from pycbc.distributions.utils import draw_samples_from_config -from pycbc.distributions.angular import UniformAngle, SinAngle, CosAngle, \ - UniformSolidAngle +from pycbc import VARARGS_DELIM as _VARARGS_DELIM +from pycbc.distributions import constraints +from pycbc.distributions.angular import ( + CosAngle, + SinAngle, + UniformAngle, + UniformSolidAngle, +) from pycbc.distributions.arbitrary import Arbitrary, FromFile +from pycbc.distributions.external import DistributionFunctionFromFile, External +from pycbc.distributions.fixedsamples import FixedSamples from pycbc.distributions.gaussian import Gaussian +from pycbc.distributions.joint import JointDistribution +from pycbc.distributions.mass import MchirpfromUniformMass1Mass2, QfromUniformMass1Mass2 from pycbc.distributions.power_law import UniformPowerLaw, UniformRadius -from pycbc.distributions.sky_location import UniformSky, UniformDiskSky, FisherSky, HealpixSky +from pycbc.distributions.qnm import UniformF0Tau +from pycbc.distributions.sky_location import ( + FisherSky, + HealpixSky, + UniformDiskSky, + UniformSky, +) +from pycbc.distributions.spins import IndependentChiPChiEff from pycbc.distributions.uniform import Uniform from pycbc.distributions.uniform_log import UniformLog10 -from pycbc.distributions.spins import IndependentChiPChiEff -from pycbc.distributions.qnm import UniformF0Tau -from pycbc.distributions.joint import JointDistribution -from pycbc.distributions.external import External, DistributionFunctionFromFile -from pycbc.distributions.fixedsamples import FixedSamples -from pycbc.distributions.mass import MchirpfromUniformMass1Mass2, \ - QfromUniformMass1Mass2 + +# Promote some classes/functions to the distributions name space +from pycbc.distributions.utils import draw_samples_from_config # a dict of all available distributions distribs = { - IndependentChiPChiEff.name : IndependentChiPChiEff, - Arbitrary.name : Arbitrary, - FromFile.name : FromFile, - Gaussian.name : Gaussian, - UniformPowerLaw.name : UniformPowerLaw, - UniformRadius.name : UniformRadius, - Uniform.name : Uniform, - UniformAngle.name : UniformAngle, - CosAngle.name : CosAngle, - SinAngle.name : SinAngle, - UniformSolidAngle.name : UniformSolidAngle, - UniformSky.name : UniformSky, - UniformDiskSky.name : UniformDiskSky, - UniformLog10.name : UniformLog10, - UniformF0Tau.name : UniformF0Tau, + IndependentChiPChiEff.name: IndependentChiPChiEff, + Arbitrary.name: Arbitrary, + FromFile.name: FromFile, + Gaussian.name: Gaussian, + UniformPowerLaw.name: UniformPowerLaw, + UniformRadius.name: UniformRadius, + Uniform.name: Uniform, + UniformAngle.name: UniformAngle, + CosAngle.name: CosAngle, + SinAngle.name: SinAngle, + UniformSolidAngle.name: UniformSolidAngle, + UniformSky.name: UniformSky, + UniformDiskSky.name: UniformDiskSky, + UniformLog10.name: UniformLog10, + UniformF0Tau.name: UniformF0Tau, External.name: External, DistributionFunctionFromFile.name: DistributionFunctionFromFile, FixedSamples.name: FixedSamples, MchirpfromUniformMass1Mass2.name: MchirpfromUniformMass1Mass2, QfromUniformMass1Mass2.name: QfromUniformMass1Mass2, FisherSky.name: FisherSky, - HealpixSky.name: HealpixSky + HealpixSky.name: HealpixSky, } + def read_distributions_from_config(cp, section="prior"): - """Returns a list of PyCBC distribution instances for a section in the + """ + Returns a list of PyCBC distribution instances for a section in the given configuration file. Parameters @@ -81,6 +93,7 @@ def read_distributions_from_config(cp, section="prior"): ------- list A list of the parsed distributions. + """ dists = [] variable_args = [] @@ -96,7 +109,8 @@ def read_distributions_from_config(cp, section="prior"): def _convert_liststring_to_list(lstring): - """Checks if an argument of the configuration file is a string of a list + """ + Checks if an argument of the configuration file is a string of a list and returns the corresponding list (of strings). The argument is considered to be a list if it starts with '[' and ends @@ -105,16 +119,22 @@ def _convert_liststring_to_list(lstring): the argument does not start and end with '[' and ']', the argument will just be returned as is. """ - if lstring[0]=='[' and lstring[-1]==']': - lstring = [str(lstring[1:-1].split(',')[n].strip().strip("'")) - for n in range(len(lstring[1:-1].split(',')))] + if lstring[0] == "[" and lstring[-1] == "]": + lstring = [ + str(lstring[1:-1].split(",")[n].strip().strip("'")) + for n in range(len(lstring[1:-1].split(","))) + ] return lstring -def read_params_from_config(cp, prior_section='prior', - vargs_section='variable_params', - sargs_section='static_params'): - """Loads static and variable parameters from a configuration file. +def read_params_from_config( + cp, + prior_section="prior", + vargs_section="variable_params", + sargs_section="static_params", +): + """ + Loads static and variable parameters from a configuration file. Parameters ---------- @@ -135,34 +155,44 @@ def read_params_from_config(cp, prior_section='prior', The names of the parameters to vary in the PE run. static_args : dict Dictionary of names -> values giving the parameters to keep fixed. + """ # sanity check that each parameter in [variable_params] has a prior section variable_args = cp.options(vargs_section) subsections = cp.get_subsections(prior_section) - tags = set([p for tag in subsections for p in tag.split('+')]) + tags = set([p for tag in subsections for p in tag.split("+")]) missing_prior = set(variable_args) - tags if any(missing_prior): - raise KeyError("You are missing a priors section in the config file " - "for parameter(s): {}".format(', '.join(missing_prior))) + raise KeyError( + "You are missing a priors section in the config file " + "for parameter(s): {}".format(", ".join(missing_prior)) + ) # sanity check that each parameter with a priors section is in # [variable_args] missing_variable = tags - set(variable_args) if any(missing_variable): - raise KeyError("Prior section found for parameter(s) {} but not " - "listed as variable parameter(s)." - .format(', '.join(missing_variable))) + raise KeyError( + "Prior section found for parameter(s) {} but not " + "listed as variable parameter(s).".format(", ".join(missing_variable)) + ) # get static args try: - static_args = dict([(key, cp.get_opt_tags(sargs_section, key, [])) - for key in cp.options(sargs_section)]) + static_args = dict( + [ + (key, cp.get_opt_tags(sargs_section, key, [])) + for key in cp.options(sargs_section) + ] + ) except _ConfigParser.NoSectionError: static_args = {} # sanity check that each parameter in [variable_args] # is not repeated in [static_args] for arg in variable_args: if arg in static_args: - raise KeyError("Parameter {} found both in static_args and in " - "variable_args sections.".format(arg)) + raise KeyError( + f"Parameter {arg} found both in static_args and in " + "variable_args sections." + ) # try converting values to float for key in static_args: val = static_args[key] @@ -177,9 +207,11 @@ def read_params_from_config(cp, prior_section='prior', return variable_args, static_args -def read_constraints_from_config(cp, transforms=None, static_args=None, - constraint_section='constraint'): - """Loads parameter constraints from a configuration file. +def read_constraints_from_config( + cp, transforms=None, static_args=None, constraint_section="constraint" +): + """ + Loads parameter constraints from a configuration file. Parameters ---------- @@ -197,17 +229,20 @@ def read_constraints_from_config(cp, transforms=None, static_args=None, ------- list List of ``Constraint`` objects. Empty if no constraints were provided. + """ cons = [] for subsection in cp.get_subsections(constraint_section): name = cp.get_opt_tag(constraint_section, "name", subsection) constraint_arg = cp.get_opt_tag( - constraint_section, "constraint_arg", subsection) + constraint_section, "constraint_arg", subsection + ) # get any other keyword arguments kwargs = {} section = constraint_section + "-" + subsection - extra_opts = [key for key in cp.options(section) - if key not in ["name", "constraint_arg"]] + extra_opts = [ + key for key in cp.options(section) if key not in ["name", "constraint_arg"] + ] for key in extra_opts: val = cp.get(section, key) if key == "required_parameters": @@ -218,8 +253,10 @@ def read_constraints_from_config(cp, transforms=None, static_args=None, except ValueError: pass kwargs[key] = val - cons.append(constraints.constraints[name]( - constraint_arg, static_args=static_args, transforms=transforms, - **kwargs)) + cons.append( + constraints.constraints[name]( + constraint_arg, static_args=static_args, transforms=transforms, **kwargs + ) + ) return cons diff --git a/pycbc/distributions/angular.py b/pycbc/distributions/angular.py index 83ed77b621b..2ee9f45959d 100644 --- a/pycbc/distributions/angular.py +++ b/pycbc/distributions/angular.py @@ -15,20 +15,21 @@ """ This modules provides classes for evaluating angular distributions. """ + import logging from configparser import Error + import numpy -from pycbc import VARARGS_DELIM -from pycbc import boundaries -from pycbc.distributions import bounded -from pycbc.distributions import uniform +from pycbc import VARARGS_DELIM, boundaries +from pycbc.distributions import bounded, uniform -logger = logging.getLogger('pycbc.distributions.angular') +logger = logging.getLogger("pycbc.distributions.angular") class UniformAngle(uniform.Uniform): - r"""A uniform distribution in which the dependent variable is between + r""" + A uniform distribution in which the dependent variable is between `[0,2pi)`. The domain of the distribution may optionally be made cyclic using the @@ -51,21 +52,24 @@ class UniformAngle(uniform.Uniform): be passed; in that case, the domain bounds will be used. Notes - ------ + ----- For more information, see Uniform. + """ - name = 'uniform_angle' - _domainbounds = (0, 2*numpy.pi) + name = "uniform_angle" + + _domainbounds = (0, 2 * numpy.pi) def __init__(self, cyclic_domain=False, **params): # _domain is a bounds instance used to apply cyclic conditions; this is # applied first, before any bounds specified in the initialization # are used - self._domain = boundaries.Bounds(self._domainbounds[0], - self._domainbounds[1], cyclic=cyclic_domain) + self._domain = boundaries.Bounds( + self._domainbounds[0], self._domainbounds[1], cyclic=cyclic_domain + ) - for p,bnds in params.items(): + for p, bnds in params.items(): if bnds is None: bnds = self._domain elif isinstance(bnds, boundaries.Bounds): @@ -77,14 +81,13 @@ def __init__(self, cyclic_domain=False, **params): bnds = boundaries.Bounds(bnds[0], bnds[1]) # check that the bounds are in the domain if bnds.min < self._domain.min or bnds.max > self._domain.max: - raise ValueError("bounds must be in [{x},{y}); " - "got [{a},{b})".format(x=self._domain.min, - y=self._domain.max, a=bnds.min, - b=bnds.max)) + raise ValueError( + f"bounds must be in [{self._domain.min},{self._domain.max}); got [{bnds.min},{bnds.max})" + ) # update params[p] = bnds - super(UniformAngle, self).__init__(**params) + super().__init__(**params) @property def domain(self): @@ -92,7 +95,8 @@ def domain(self): return self._domain def apply_boundary_conditions(self, **kwargs): - r"""Maps values to be in [0, 2pi) (the domain) first, before applying + r""" + Maps values to be in [0, 2pi) (the domain) first, before applying any additional boundary conditions. Parameters @@ -106,16 +110,23 @@ def apply_boundary_conditions(self, **kwargs): ------- dict A dictionary of the parameter names and the conditioned values. + """ # map values to be within the domain - kwargs = dict([[p, self._domain.apply_conditions(val)] - for p,val in kwargs.items() if p in self._bounds]) + kwargs = dict( + [ + [p, self._domain.apply_conditions(val)] + for p, val in kwargs.items() + if p in self._bounds + ] + ) # now apply additional conditions - return super(UniformAngle, self).apply_boundary_conditions(**kwargs) + return super().apply_boundary_conditions(**kwargs) @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. By default, @@ -157,17 +168,25 @@ def from_config(cls, cp, section, variable_args): ------- UniformAngle A distribution instance from the pycbc.inference.prior module. + """ # we'll retrieve the setting for cyclic_domain directly - additional_opts = {'cyclic_domain': cp.has_option_tag(section, - 'cyclic_domain', variable_args)} - return bounded.bounded_from_config(cls, cp, section, variable_args, - bounds_required=False, - additional_opts=additional_opts) + additional_opts = { + "cyclic_domain": cp.has_option_tag(section, "cyclic_domain", variable_args) + } + return bounded.bounded_from_config( + cls, + cp, + section, + variable_args, + bounds_required=False, + additional_opts=additional_opts, + ) class SinAngle(UniformAngle): - r"""A sine distribution; the pdf of each parameter `\theta` is given by: + r""" + A sine distribution; the pdf of each parameter `\theta` is given by: ..math:: p(\theta) = \frac{\sin \theta}{\cos\theta_0 - \cos\theta_1}, \theta_0 \leq \theta < \theta_1, @@ -188,58 +207,72 @@ class SinAngle(UniformAngle): `boundaries.Bounds` instances or tuples. The bounds must be in [0,PI]. These are converted to radians for storage. None may also be passed; in that case, the domain bounds will be used. + """ - name = 'sin_angle' + + name = "sin_angle" _func = numpy.cos _dfunc = numpy.sin _arcfunc = numpy.arccos _domainbounds = (0, numpy.pi) def __init__(self, **params): - super(SinAngle, self).__init__(**params) + super().__init__(**params) # replace the domain - self._domain = boundaries.Bounds(self._domainbounds[0], - self._domainbounds[1], btype_min='closed', btype_max='closed', - cyclic=False) - self._lognorm = -sum([numpy.log( - abs(self._func(bnd[1]) - self._func(bnd[0]))) \ - for bnd in self._bounds.values()]) + self._domain = boundaries.Bounds( + self._domainbounds[0], + self._domainbounds[1], + btype_min="closed", + btype_max="closed", + cyclic=False, + ) + self._lognorm = -sum( + [ + numpy.log(abs(self._func(bnd[1]) - self._func(bnd[0]))) + for bnd in self._bounds.values() + ] + ) self._norm = numpy.exp(self._lognorm) def _cdfinv_param(self, arg, value): - """Return inverse of cdf for mapping unit interval to parameter bounds. - """ - scale = (numpy.cos(self._bounds[arg][0]) - - numpy.cos(self._bounds[arg][1])) - offset = 1. + numpy.cos(self._bounds[arg][1]) / scale + """Return inverse of cdf for mapping unit interval to parameter bounds.""" + scale = numpy.cos(self._bounds[arg][0]) - numpy.cos(self._bounds[arg][1]) + offset = 1.0 + numpy.cos(self._bounds[arg][1]) / scale new_value = numpy.arccos(-scale * (value - offset)) return new_value def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs not in self: - return 0. - return self._norm * \ - self._dfunc(numpy.array([kwargs[p] for p in self._params])).prod() - + return 0.0 + return ( + self._norm + * self._dfunc(numpy.array([kwargs[p] for p in self._params])).prod() + ) def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs not in self: return -numpy.inf - return self._lognorm + \ - numpy.log(self._dfunc( - numpy.array([kwargs[p] for p in self._params]))).sum() + return ( + self._lognorm + + numpy.log( + self._dfunc(numpy.array([kwargs[p] for p in self._params])) + ).sum() + ) class CosAngle(SinAngle): - r"""A cosine distribution. This is the same thing as a sine distribution, + r""" + A cosine distribution. This is the same thing as a sine distribution, but with the domain shifted to `[-pi/2, pi/2]`. See SinAngle for more details. @@ -250,24 +283,27 @@ class CosAngle(SinAngle): (optionally) their corresponding bounds, as either `boundaries.Bounds` instances or tuples. The bounds must be in [-PI/2, PI/2]. + """ - name = 'cos_angle' + + name = "cos_angle" _func = numpy.sin _dfunc = numpy.cos _arcfunc = numpy.arcsin - _domainbounds = (-numpy.pi/2, numpy.pi/2) + _domainbounds = (-numpy.pi / 2, numpy.pi / 2) def _cdfinv_param(self, param, value): a = self._bounds[param][0] b = self._bounds[param][1] scale = numpy.sin(b) - numpy.sin(a) - offset = 1. - numpy.sin(b)/(numpy.sin(b) - numpy.sin(a)) + offset = 1.0 - numpy.sin(b) / (numpy.sin(b) - numpy.sin(a)) new_value = numpy.arcsin((value - offset) * scale) return new_value class UniformSolidAngle(bounded.BoundedDist): - r"""A distribution that is uniform in the solid angle of a sphere. The names + r""" + A distribution that is uniform in the solid angle of a sphere. The names of the two angluar parameters can be specified on initalization. Parameters @@ -292,25 +328,34 @@ class UniformSolidAngle(bounded.BoundedDist): values are constrained to be in [0, 2pi) using cyclic boundaries prior to applying any other boundary conditions and prior to evaluating the pdf. Default is False. + """ - name = 'uniform_solidangle' + + name = "uniform_solidangle" _polardistcls = SinAngle _azimuthaldistcls = UniformAngle - _default_polar_angle = 'theta' - _default_azimuthal_angle = 'phi' - - def __init__(self, polar_angle=None, azimuthal_angle=None, - polar_bounds=None, azimuthal_bounds=None, - azimuthal_cyclic_domain=False): + _default_polar_angle = "theta" + _default_azimuthal_angle = "phi" + + def __init__( + self, + polar_angle=None, + azimuthal_angle=None, + polar_bounds=None, + azimuthal_bounds=None, + azimuthal_cyclic_domain=False, + ): if polar_angle is None: polar_angle = self._default_polar_angle if azimuthal_angle is None: azimuthal_angle = self._default_azimuthal_angle - self._polardist = self._polardistcls(**{ - polar_angle: polar_bounds}) - self._azimuthaldist = self._azimuthaldistcls(**{ - azimuthal_angle: azimuthal_bounds, - 'cyclic_domain': azimuthal_cyclic_domain}) + self._polardist = self._polardistcls(**{polar_angle: polar_bounds}) + self._azimuthaldist = self._azimuthaldistcls( + **{ + azimuthal_angle: azimuthal_bounds, + "cyclic_domain": azimuthal_cyclic_domain, + } + ) self._polar_angle = polar_angle self._azimuthal_angle = azimuthal_angle self._bounds = self._polardist.bounds.copy() @@ -319,11 +364,13 @@ def __init__(self, polar_angle=None, azimuthal_angle=None, @property def bounds(self): - """dict: The bounds on each angle. The keys are the names of the polar + """ + dict: The bounds on each angle. The keys are the names of the polar and azimuthal angles, the values are the minimum and maximum of each, in radians. For example, if the distribution was initialized with `polar_angle='theta', polar_bounds=(0,0.5)` then the bounds will have - `'theta': 0, 1.5707963267948966` as an entry.""" + `'theta': 0, 1.5707963267948966` as an entry. + """ return self._bounds @property @@ -337,14 +384,15 @@ def azimuthal_angle(self): return self._azimuthal_angle def _cdfinv_param(self, param, value): - """ Return the cdfinv for a single given parameter """ + """Return the cdfinv for a single given parameter""" if param == self.polar_angle: return self._polardist._cdfinv_param(param, value) - elif param == self.azimuthal_angle: + if param == self.azimuthal_angle: return self._azimuthaldist._cdfinv_param(param, value) def apply_boundary_conditions(self, **kwargs): - r"""Maps the given values to be within the domain of the azimuthal and + r""" + Maps the given values to be within the domain of the azimuthal and polar angles, before applying any other boundary conditions. Parameters @@ -359,6 +407,7 @@ def apply_boundary_conditions(self, **kwargs): ------- dict A dictionary of the parameter names and the conditioned values. + """ polarval = kwargs[self._polar_angle] azval = kwargs[self._azimuthal_angle] @@ -370,7 +419,6 @@ def apply_boundary_conditions(self, **kwargs): azval = self._bounds[self._azimuthal_angle].apply_conditions(azval) return {self._polar_angle: polarval, self._azimuthal_angle: azval} - def _pdf(self, **kwargs): r""" Returns the pdf at the given angles. @@ -386,10 +434,9 @@ def _pdf(self, **kwargs): ------- float The value of the pdf at the given values. - """ - return self._polardist._pdf(**kwargs) * \ - self._azimuthaldist._pdf(**kwargs) + """ + return self._polardist._pdf(**kwargs) * self._azimuthaldist._pdf(**kwargs) def _logpdf(self, **kwargs): r""" @@ -406,13 +453,14 @@ def _logpdf(self, **kwargs): ------- float The value of the pdf at the given values. + """ - return self._polardist._logpdf(**kwargs) +\ - self._azimuthaldist._logpdf(**kwargs) + return self._polardist._logpdf(**kwargs) + self._azimuthaldist._logpdf(**kwargs) @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. + """ + Returns a distribution based on a configuration file. The section must have the names of the polar and azimuthal angles in the tag part of the section header. For example: @@ -469,43 +517,53 @@ def from_config(cls, cp, section, variable_args): ------- UniformSolidAngle A distribution instance from the pycbc.inference.prior module. + """ tag = variable_args variable_args = variable_args.split(VARARGS_DELIM) # get the variables that correspond to the polar/azimuthal angles try: - polar_angle = cp.get_opt_tag(section, 'polar-angle', tag) + polar_angle = cp.get_opt_tag(section, "polar-angle", tag) except Error: polar_angle = cls._default_polar_angle try: - azimuthal_angle = cp.get_opt_tag(section, 'azimuthal-angle', tag) + azimuthal_angle = cp.get_opt_tag(section, "azimuthal-angle", tag) except Error: azimuthal_angle = cls._default_azimuthal_angle if polar_angle not in variable_args: - raise Error("polar-angle %s is not one of the variable args (%s)"%( - polar_angle, ', '.join(variable_args))) + raise Error( + "polar-angle %s is not one of the variable args (%s)" + % (polar_angle, ", ".join(variable_args)) + ) if azimuthal_angle not in variable_args: - raise Error("azimuthal-angle %s is not one of the variable args "%( - azimuthal_angle) + "(%s)"%(', '.join(variable_args))) + raise Error( + "azimuthal-angle %s is not one of the variable args " + % (azimuthal_angle) + + "(%s)" % (", ".join(variable_args)) + ) # get the bounds, if provided polar_bounds = bounded.get_param_bounds_from_config( - cp, section, tag, - polar_angle) + cp, section, tag, polar_angle + ) azimuthal_bounds = bounded.get_param_bounds_from_config( - cp, section, tag, - azimuthal_angle) + cp, section, tag, azimuthal_angle + ) # see if the a cyclic domain is desired for the azimuthal angle - azimuthal_cyclic_domain = cp.has_option_tag(section, - 'azimuthal_cyclic_domain', tag) + azimuthal_cyclic_domain = cp.has_option_tag( + section, "azimuthal_cyclic_domain", tag + ) - return cls(polar_angle=polar_angle, azimuthal_angle=azimuthal_angle, - polar_bounds=polar_bounds, - azimuthal_bounds=azimuthal_bounds, - azimuthal_cyclic_domain=azimuthal_cyclic_domain) + return cls( + polar_angle=polar_angle, + azimuthal_angle=azimuthal_angle, + polar_bounds=polar_bounds, + azimuthal_bounds=azimuthal_bounds, + azimuthal_cyclic_domain=azimuthal_cyclic_domain, + ) -__all__ = ['UniformAngle', 'SinAngle', 'CosAngle', 'UniformSolidAngle'] +__all__ = ["CosAngle", "SinAngle", "UniformAngle", "UniformSolidAngle"] diff --git a/pycbc/distributions/arbitrary.py b/pycbc/distributions/arbitrary.py index f9597ecbe6f..e644542e2b8 100644 --- a/pycbc/distributions/arbitrary.py +++ b/pycbc/distributions/arbitrary.py @@ -16,18 +16,22 @@ This modules provides classes for evaluating arbitrary distributions from a file. """ + import logging + import numpy import scipy.stats -from pycbc.distributions import bounded import pycbc.transforms +from pycbc.distributions import bounded from pycbc.io.hdf import HFile -logger = logging.getLogger('pycbc.distributions.arbitrary') +logger = logging.getLogger("pycbc.distributions.arbitrary") + class Arbitrary(bounded.BoundedDist): - r"""A distribution constructed from a set of parameter values using a kde. + r""" + A distribution constructed from a set of parameter values using a kde. Bounds may be optionally provided to limit the range. Parameters @@ -43,26 +47,29 @@ class Arbitrary(bounded.BoundedDist): a list of their parameter values. If multiple parameters are provided, a single kde will be produced with dimension equal to the number of parameters. + """ - name = 'arbitrary' + + name = "arbitrary" def __init__(self, bounds=None, bandwidth="scott", **kwargs): # initialize the bounds if bounds is None: bounds = {} bounds.update({p: None for p in kwargs if p not in bounds}) - super(Arbitrary, self).__init__(**bounds) + super().__init__(**bounds) # check that all parameters specified in bounds have samples if set(self.params) != set(kwargs.keys()): - raise ValueError("Must provide samples for all parameters given " - "in the bounds dictionary") + raise ValueError( + "Must provide samples for all parameters given in the bounds dictionary" + ) # if bounds are provided use logit transform to move the points # to +/- inifinity self._transforms = {} self._tparams = {} - for param,bnds in self.bounds.items(): + for param, bnds in self.bounds.items(): if numpy.isfinite(bnds[1] - bnds[0]): - tparam = 'logit'+param + tparam = "logit" + param samples = kwargs[param] t = pycbc.transforms.Logit(param, tparam, domain=bnds) self._transforms[tparam] = t @@ -74,8 +81,7 @@ def __init__(self, bounds=None, bandwidth="scott", **kwargs): # transform the sample points kwargs[param] = t.transform({param: samples})[tparam] elif not (~numpy.isfinite(bnds[0]) and ~numpy.isfinite(bnds[1])): - raise ValueError("if specifying bounds, both bounds must " - "be finite") + raise ValueError("if specifying bounds, both bounds must be finite") # build the kde self._kde = self.get_kde_from_arrays(*[kwargs[p] for p in self.params]) self.set_bandwidth(bandwidth) @@ -89,17 +95,17 @@ def kde(self): return self._kde def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ for p in self._params: - if p not in kwargs.keys(): - raise ValueError('Missing parameter {} to construct pdf.' - .format(p)) + if p not in kwargs: + raise ValueError(f"Missing parameter {p} to construct pdf.") if kwargs in self: # transform into the kde space - jacobian = 1. + jacobian = 1.0 for param, tparam in self._tparams.items(): t = self._transforms[tparam] try: @@ -108,9 +114,8 @@ def _pdf(self, **kwargs): # can get a value error if the value is exactly == to # the bounds, in which case, just return 0. if kwargs[param] in self.bounds[param]: - return 0. - else: - raise ValueError(e) + return 0.0 + raise ValueError(e) kwargs[param] = samples[tparam] # update the jacobian for the transform; if p is the pdf # in the params frame (the one we want) and p' is the pdf @@ -118,30 +123,28 @@ def _pdf(self, **kwargs): # p = J * p', where J is the Jacobian of going from p to p' jacobian *= t.jacobian(samples) # for scipy < 0.15.0, gaussian_kde.pdf = gaussian_kde.evaluate - this_pdf = jacobian * self._kde.evaluate([kwargs[p] - for p in self._params]) + this_pdf = jacobian * self._kde.evaluate([kwargs[p] for p in self._params]) if len(this_pdf) == 1: return float(this_pdf) - else: - return this_pdf - else: - return 0. + return this_pdf + return 0.0 def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs not in self: return -numpy.inf - else: - return numpy.log(self._pdf(**kwargs)) + return numpy.log(self._pdf(**kwargs)) def set_bandwidth(self, set_bw="scott"): self._kde.set_bandwidth(set_bw) def rvs(self, size=1, param=None): - """Gives a set of random values drawn from the kde. + """ + Gives a set of random values drawn from the kde. Parameters ---------- @@ -158,6 +161,7 @@ def rvs(self, size=1, param=None): specified, the array will only have an element corresponding to the given parameter. Otherwise, the array will have an element for each parameter in self's params. + """ if param is not None: dtype = [(param, float)] @@ -166,14 +170,13 @@ def rvs(self, size=1, param=None): size = int(size) arr = numpy.zeros(size, dtype=dtype) draws = self._kde.resample(size) - draws = {param: draws[ii,:] for ii,param in enumerate(self.params)} - for (param,_) in dtype: + draws = {param: draws[ii, :] for ii, param in enumerate(self.params)} + for param, _ in dtype: try: # transform back to param space tparam = self._tparams[param] tdraws = {tparam: draws[param]} - draws[param] = self._transforms[tparam].inverse_transform( - tdraws)[param] + draws[param] = self._transforms[tparam].inverse_transform(tdraws)[param] except KeyError: pass arr[param] = draws[param] @@ -181,7 +184,8 @@ def rvs(self, size=1, param=None): @staticmethod def get_kde_from_arrays(*arrays): - r"""Constructs a KDE from the given arrays. + r""" + Constructs a KDE from the given arrays. \*arrays : Each argument should be a 1D numpy array to construct the kde from. @@ -192,15 +196,19 @@ def get_kde_from_arrays(*arrays): @classmethod def from_config(cls, cp, section, variable_args): - """Raises a NotImplementedError; to load from a config file, use + """ + Raises a NotImplementedError; to load from a config file, use `FromFile`. """ - raise NotImplementedError("This class does not support loading from a " - "config file. Use `FromFile` instead.") + raise NotImplementedError( + "This class does not support loading from a " + "config file. Use `FromFile` instead." + ) class FromFile(Arbitrary): - r"""A distribution that reads the values of the parameter(s) from an hdf + r""" + A distribution that reads the values of the parameter(s) from an hdf file, computes the kde to construct the pdf, and draws random variables from it. @@ -232,11 +240,14 @@ class FromFile(Arbitrary): The log of the normalization. kde : The kde obtained from the values in the file. + """ - name = 'fromfile' + + name = "fromfile" + def __init__(self, filename=None, datagroup=None, **params): if filename is None: - raise ValueError('A file must be specified for this distribution.') + raise ValueError("A file must be specified for this distribution.") self._filename = filename self.datagroup = datagroup # Get the parameter names to pass to get_kde_from_file @@ -245,17 +256,16 @@ def __init__(self, filename=None, datagroup=None, **params): else: ps = list(params.keys()) param_vals, bw = self.get_arrays_from_file(filename, params=ps) - super(FromFile, self).__init__(bounds=params, bandwidth=bw, - **param_vals) + super().__init__(bounds=params, bandwidth=bw, **param_vals) @property def filename(self): - """str: The path to the file containing values for the parameter(s). - """ + """str: The path to the file containing values for the parameter(s).""" return self._filename def get_arrays_from_file(self, params_file, params=None): - """Reads the values of one or more parameters from an hdf file and + """ + Reads the values of one or more parameters from an hdf file and returns as a dictionary. Parameters @@ -269,11 +279,12 @@ def get_arrays_from_file(self, params_file, params=None): ------- dict A dictionary of the parameters mapping `param_name -> array`. + """ try: - f = HFile(params_file, 'r') + f = HFile(params_file, "r") except: - raise ValueError('File not found.') + raise ValueError("File not found.") if self.datagroup is not None: get = f[self.datagroup] else: @@ -283,8 +294,7 @@ def get_arrays_from_file(self, params_file, params=None): params = [params] for p in params: if p not in get.keys(): - raise ValueError('Parameter {} is not in {}' - .format(p, params_file)) + raise ValueError(f"Parameter {p} is not in {params_file}") else: params = [str(k) for k in get.keys()] params_values = {p: get[p][()] for p in params} @@ -298,7 +308,8 @@ def get_arrays_from_file(self, params_file, params=None): @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled @@ -332,8 +343,11 @@ def from_config(cls, cp, section, variable_args): ------- BoundedDist A distribution instance from the pycbc.inference.prior module. + """ - return bounded.bounded_from_config(cls, cp, section, variable_args, - bounds_required=False) + return bounded.bounded_from_config( + cls, cp, section, variable_args, bounds_required=False + ) + -__all__ = ['Arbitrary', 'FromFile'] +__all__ = ["Arbitrary", "FromFile"] diff --git a/pycbc/distributions/bounded.py b/pycbc/distributions/bounded.py index 0f0efa0d9c2..0c7390ea741 100644 --- a/pycbc/distributions/bounded.py +++ b/pycbc/distributions/bounded.py @@ -15,21 +15,24 @@ """ This modules provides classes for evaluating distributions with bounds. """ + import logging import warnings from configparser import Error + import numpy -from pycbc import boundaries -from pycbc import VARARGS_DELIM +from pycbc import VARARGS_DELIM, boundaries + +logger = logging.getLogger("pycbc.distributions.bounded") -logger = logging.getLogger('pycbc.distributions.bounded') # # Distributions for priors # def get_param_bounds_from_config(cp, section, tag, param): - """Gets bounds for the given parameter from a section in a config file. + """ + Gets bounds for the given parameter from a section in a config file. Minimum and maximum values for bounds are specified by adding `min-{param}` and `max-{param}` options, where `{param}` is the name of @@ -82,43 +85,46 @@ def get_param_bounds_from_config(cp, section, tag, param): bounds : {Bounds instance | None} If bounds were provided, a `boundaries.Bounds` instance representing the bounds. Otherwise, `None`. + """ try: - minbnd = float(cp.get_opt_tag(section, 'min-'+param, tag)) + minbnd = float(cp.get_opt_tag(section, "min-" + param, tag)) except Error: minbnd = None try: - maxbnd = float(cp.get_opt_tag(section, 'max-'+param, tag)) + maxbnd = float(cp.get_opt_tag(section, "max-" + param, tag)) except Error: maxbnd = None if minbnd is None and maxbnd is None: bnds = None elif minbnd is None or maxbnd is None: - raise ValueError("if specifying bounds for %s, " %(param) + - "you must provide both a minimum and a maximum") + raise ValueError( + "if specifying bounds for %s, " % (param) + + "you must provide both a minimum and a maximum" + ) else: - bndargs = {'min_bound': minbnd, 'max_bound': maxbnd} + bndargs = {"min_bound": minbnd, "max_bound": maxbnd} # try to get any other conditions, if provided try: - minbtype = cp.get_opt_tag(section, 'btype-min-{}'.format(param), - tag) + minbtype = cp.get_opt_tag(section, f"btype-min-{param}", tag) except Error: - minbtype = 'closed' + minbtype = "closed" try: - maxbtype = cp.get_opt_tag(section, 'btype-max-{}'.format(param), - tag) + maxbtype = cp.get_opt_tag(section, f"btype-max-{param}", tag) except Error: - maxbtype = 'open' - bndargs.update({'btype_min': minbtype, 'btype_max': maxbtype}) - cyclic = cp.has_option_tag(section, 'cyclic-{}'.format(param), tag) - bndargs.update({'cyclic': cyclic}) + maxbtype = "open" + bndargs.update({"btype_min": minbtype, "btype_max": maxbtype}) + cyclic = cp.has_option_tag(section, f"cyclic-{param}", tag) + bndargs.update({"cyclic": cyclic}) bnds = boundaries.Bounds(**bndargs) return bnds -def bounded_from_config(cls, cp, section, variable_args, - bounds_required=False, additional_opts=None): - """Returns a bounded distribution based on a configuration file. The +def bounded_from_config( + cls, cp, section, variable_args, bounds_required=False, additional_opts=None +): + """ + Returns a bounded distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -151,6 +157,7 @@ def bounded_from_config(cls, cp, section, variable_args, ------- cls An instance of the given class. + """ tag = variable_args variable_args = variable_args.split(VARARGS_DELIM) @@ -159,26 +166,26 @@ def bounded_from_config(cls, cp, section, variable_args, additional_opts = {} # list of args that are used to construct distribution - special_args = ["name"] + \ - ['min-{}'.format(arg) for arg in variable_args] + \ - ['max-{}'.format(arg) for arg in variable_args] + \ - ['btype-min-{}'.format(arg) for arg in variable_args] + \ - ['btype-max-{}'.format(arg) for arg in variable_args] + \ - ['cyclic-{}'.format(arg) for arg in variable_args] + \ - list(additional_opts.keys()) + special_args = ( + ["name"] + + [f"min-{arg}" for arg in variable_args] + + [f"max-{arg}" for arg in variable_args] + + [f"btype-min-{arg}" for arg in variable_args] + + [f"btype-max-{arg}" for arg in variable_args] + + [f"cyclic-{arg}" for arg in variable_args] + + list(additional_opts.keys()) + ) # get a dict with bounds as value dist_args = {} for param in variable_args: bounds = get_param_bounds_from_config(cp, section, tag, param) if bounds_required and bounds is None: - raise ValueError("min and/or max missing for parameter %s"%( - param)) + raise ValueError("min and/or max missing for parameter %s" % (param)) dist_args[param] = bounds # add any additional options that user put in that section for key in cp.options("-".join([section, tag])): - # ignore options that are already included if key in special_args: continue @@ -191,7 +198,7 @@ def bounded_from_config(cls, cp, section, variable_args, pass # add option - dist_args.update({key:val}) + dist_args.update({key: val}) dist_args.update(additional_opts) @@ -199,7 +206,7 @@ def bounded_from_config(cls, cp, section, variable_args, return cls(**dist_args) -class BoundedDist(object): +class BoundedDist: r""" A generic class for storing common properties of distributions in which each parameter has a minimum and maximum value. @@ -210,27 +217,31 @@ class BoundedDist(object): The keyword arguments should provide the names of parameters and their corresponding bounds, as either tuples or a `boundaries.Bounds` instance. + """ + def __init__(self, **params): # convert input bounds to Bounds class, if necessary - for param,bnds in params.items(): + for param, bnds in params.items(): if bnds is None: params[param] = boundaries.Bounds() elif not isinstance(bnds, boundaries.Bounds): params[param] = boundaries.Bounds(bnds[0], bnds[1]) # warn the user about reflected boundaries if isinstance(bnds, boundaries.Bounds) and ( - bnds.min.name == 'reflected' or - bnds.max.name == 'reflected'): - warnings.warn("Param {} has one or more ".format(param) + - "reflected boundaries. Reflected boundaries " - "can cause issues when used in an MCMC.") + bnds.min.name == "reflected" or bnds.max.name == "reflected" + ): + warnings.warn( + f"Param {param} has one or more " + "reflected boundaries. Reflected boundaries " + "can cause issues when used in an MCMC." + ) self._bounds = params self._params = sorted(list(params.keys())) @property def params(self): - """list of strings: The list of parameter names.""" + """List of strings: The list of parameter names.""" return self._params @property @@ -240,14 +251,17 @@ def bounds(self): def __contains__(self, params): try: - return all(self._bounds[p].contains_conditioned(params[p]) - for p in self._params) + return all( + self._bounds[p].contains_conditioned(params[p]) for p in self._params + ) except KeyError: - raise ValueError("must provide all parameters [%s]" %( - ', '.join(self._params))) + raise ValueError( + "must provide all parameters [%s]" % (", ".join(self._params)) + ) def apply_boundary_conditions(self, **kwargs): - r"""Applies any boundary conditions to the given values (e.g., applying + r""" + Applies any boundary conditions to the given values (e.g., applying cyclic conditions, and/or reflecting values off of boundaries). This is done by running `apply_conditions` of each bounds in self on the corresponding value. See `boundaries.Bounds.apply_conditions` for @@ -265,12 +279,19 @@ def apply_boundary_conditions(self, **kwargs): ------- dict A dictionary of the parameter names and the conditioned values. + """ - return dict([[p, self._bounds[p].apply_conditions(val)] - for p,val in kwargs.items() if p in self._bounds]) + return dict( + [ + [p, self._bounds[p].apply_conditions(val)] + for p, val in kwargs.items() + if p in self._bounds + ] + ) def pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. Any boundary conditions are applied to the values before the pdf is evaluated. @@ -278,14 +299,16 @@ def pdf(self, **kwargs): return self._pdf(**self.apply_boundary_conditions(**kwargs)) def _pdf(self, **kwargs): - """The underlying pdf function called by `self.pdf`. This must be set + """ + The underlying pdf function called by `self.pdf`. This must be set by any class that inherits from this class. Otherwise, a `NotImplementedError` is raised. """ raise NotImplementedError("pdf function not set") def logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. Any boundary conditions are applied to the values before the pdf is evaluated. @@ -293,7 +316,8 @@ def logpdf(self, **kwargs): return self._logpdf(**self.apply_boundary_conditions(**kwargs)) def _logpdf(self, **kwargs): - """The underlying log pdf function called by `self.logpdf`. This must + """ + The underlying log pdf function called by `self.logpdf`. This must be set by any class that inherits from this class. Otherwise, a `NotImplementedError` is raised. """ @@ -302,11 +326,12 @@ def _logpdf(self, **kwargs): __call__ = logpdf def _cdfinv_param(self, param, value): - """Return the cdfinv for a single given parameter """ + """Return the cdfinv for a single given parameter""" raise NotImplementedError("inverse cdf not set") def cdfinv(self, **kwds): - """Return the inverse cdf to map the unit interval to parameter bounds. + """ + Return the inverse cdf to map the unit interval to parameter bounds. You must provide a keyword for every parameter. """ updated = {} @@ -315,7 +340,7 @@ def cdfinv(self, **kwds): return updated def rvs(self, size=1, **kwds): - "Draw random value" + """Draw random value""" dtype = [(p, float) for p in self.params] arr = numpy.zeros(size, dtype=dtype) draw = {} @@ -328,7 +353,8 @@ def rvs(self, size=1, **kwds): @classmethod def from_config(cls, cp, section, variable_args, bounds_required=False): - """Returns a distribution based on a configuration file. The parameters + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -354,6 +380,8 @@ def from_config(cls, cp, section, variable_args, bounds_required=False): ------- BoundedDist A distribution instance from the pycbc.distribution subpackage. + """ - return bounded_from_config(cls, cp, section, variable_args, - bounds_required=bounds_required) + return bounded_from_config( + cls, cp, section, variable_args, bounds_required=bounds_required + ) diff --git a/pycbc/distributions/constraints.py b/pycbc/distributions/constraints.py index b35ee18af07..bea4aa9a57b 100644 --- a/pycbc/distributions/constraints.py +++ b/pycbc/distributions/constraints.py @@ -15,30 +15,35 @@ """ This modules provides classes for evaluating multi-dimensional constraints. """ + import logging import re -import scipy.spatial + import numpy +import scipy.spatial from pycbc import transforms -from pycbc.io import record, HFile +from pycbc.io import HFile, record -logger = logging.getLogger('pycbc.distributions.constraints') +logger = logging.getLogger("pycbc.distributions.constraints") -class Constraint(object): - """Creates a constraint that evaluates to True if parameters obey +class Constraint: + """ + Creates a constraint that evaluates to True if parameters obey the constraint and False if they do not. """ + name = "custom" - def __init__(self, constraint_arg, static_args=None, transforms=None, - **kwargs): + def __init__(self, constraint_arg, static_args=None, transforms=None, **kwargs): static_args = ( - {} if static_args is None - else dict(sorted( - static_args.items(), key=lambda x: len(x[0]), reverse=True)) + {} + if static_args is None + else dict( + sorted(static_args.items(), key=lambda x: len(x[0]), reverse=True) ) + ) for arg, val in static_args.items(): swp = f"'{val}'" if isinstance(val, str) else str(val) # Substitute static arg name for value if it appears in the @@ -47,17 +52,14 @@ def __init__(self, constraint_arg, static_args=None, transforms=None, # This ensures that static_args that are also kwargs in function calls are # handled correctly, i.e., the kwarg is not touched while its value is replaced # with the static_arg value. - constraint_arg = re.sub( - r'\b{}(?!\_|\=)'.format(arg), swp, constraint_arg) + constraint_arg = re.sub(rf"\b{arg}(?!\_|\=)", swp, constraint_arg) self.constraint_arg = constraint_arg self.transforms = transforms - for kwarg in kwargs.keys(): + for kwarg in kwargs: setattr(self, kwarg, kwargs[kwarg]) def __call__(self, params): - """Evaluates constraint. - """ - + """Evaluates constraint.""" if isinstance(params, dict): params = record.FieldArray.from_kwargs(**params) elif not isinstance(params, record.FieldArray): @@ -66,41 +68,41 @@ def __call__(self, params): try: out = self._constraint(params) except (NameError, AttributeError, TypeError): - if self.transforms: params = transforms.apply_transforms(params, self.transforms) out = self._constraint(params) - if isinstance(out, record.FieldArray): out = out.item() if params.size == 1 else out return out def _constraint(self, params): - """ Evaluates constraint function. - """ + """Evaluates constraint function.""" return params[self.constraint_arg] class SupernovaeConvexHull(Constraint): - """Pre defined constraint for core-collapse waveforms that checks + """ + Pre defined constraint for core-collapse waveforms that checks whether a given set of coefficients lie within the convex hull of the coefficients of the principal component basis vectors. """ + name = "supernovae_convex_hull" required_parameters = ["coeff_0", "coeff_1"] def __init__(self, constraint_arg, transforms=None, **kwargs): - super(SupernovaeConvexHull, - self).__init__(constraint_arg, transforms=transforms, **kwargs) + super().__init__( + constraint_arg, transforms=transforms, **kwargs + ) - if 'principal_components_file' in kwargs: - pc_filename = kwargs['principal_components_file'] - hull_dimention = numpy.array(kwargs['hull_dimention']) + if "principal_components_file" in kwargs: + pc_filename = kwargs["principal_components_file"] + hull_dimention = numpy.array(kwargs["hull_dimention"]) self.hull_dimention = int(hull_dimention) - pc_file = HFile(pc_filename, 'r') - pc_coefficients = numpy.array(pc_file.get('coefficients')) + pc_file = HFile(pc_filename, "r") + pc_coefficients = numpy.array(pc_file.get("coefficients")) pc_file.close() hull_points = [] for dim in range(self.hull_dimention): @@ -112,17 +114,15 @@ def __init__(self, constraint_arg, transforms=None, **kwargs): def _constraint(self, params): output_array = [] - points = numpy.array([params["coeff_0"], - params["coeff_1"], - params["coeff_2"]]) + points = numpy.array([params["coeff_0"], params["coeff_1"], params["coeff_2"]]) for coeff_index in range(len(params["coeff_0"])): - point = points[:, coeff_index][:self.hull_dimention] + point = points[:, coeff_index][: self.hull_dimention] output_array.append(self._hull.find_simplex(point) >= 0) return numpy.array(output_array) # list of all constraints constraints = { - Constraint.name : Constraint, - SupernovaeConvexHull.name : SupernovaeConvexHull, + Constraint.name: Constraint, + SupernovaeConvexHull.name: SupernovaeConvexHull, } diff --git a/pycbc/distributions/external.py b/pycbc/distributions/external.py index 6829a17de84..da6042e4365 100644 --- a/pycbc/distributions/external.py +++ b/pycbc/distributions/external.py @@ -16,20 +16,22 @@ This modules provides classes for evaluating PDF, logPDF, CDF and inverse CDF from external arbitrary distributions, and drawing samples from them. """ -import logging + import importlib -import numpy as np +import logging +import numpy as np import scipy.integrate as scipy_integrate import scipy.interpolate as scipy_interpolate from pycbc import VARARGS_DELIM -logger = logging.getLogger('pycbc.distributions.external') +logger = logging.getLogger("pycbc.distributions.external") -class External(object): - """ Distribution defined by external cdfinv and logpdf functions +class External: + """ + Distribution defined by external cdfinv and logpdf functions To add to an inference configuration file: @@ -69,11 +71,12 @@ class External(object): ... return kwds >>> e = External(['x', 'y'], logpdf, cdfinv=cdfinv) >>> e.rvs(size=10) + """ + name = "external" - def __init__(self, params=None, logpdf=None, - rvs=None, cdfinv=None, **kwds): + def __init__(self, params=None, logpdf=None, rvs=None, cdfinv=None, **kwds): self.params = params self.logpdf = logpdf self.cdfinv = cdfinv @@ -83,11 +86,10 @@ def __init__(self, params=None, logpdf=None, raise ValueError("Must provide either rvs or cdfinv") def rvs(self, size=1, **kwds): - "Draw random value" + """Draw random value""" if self._rvs: return self._rvs(size=size) - samples = {param: np.random.uniform(0, 1, size=size) - for param in self.params} + samples = {param: np.random.uniform(0, 1, size=size) for param in self.params} return self.cdfinv(**samples) def apply_boundary_conditions(self, **params): @@ -100,26 +102,27 @@ def __call__(self, **kwds): def from_config(cls, cp, section, variable_args): tag = variable_args params = variable_args.split(VARARGS_DELIM) - modulestr = cp.get_opt_tag(section, 'module', tag) + modulestr = cp.get_opt_tag(section, "module", tag) mod = importlib.import_module(modulestr) - logpdfstr = cp.get_opt_tag(section, 'logpdf', tag) + logpdfstr = cp.get_opt_tag(section, "logpdf", tag) logpdf = getattr(mod, logpdfstr) cdfinv = rvs = None - if cp.has_option_tag(section, 'cdfinv', tag): - cdfinvstr = cp.get_opt_tag(section, 'cdfinv', tag) + if cp.has_option_tag(section, "cdfinv", tag): + cdfinvstr = cp.get_opt_tag(section, "cdfinv", tag) cdfinv = getattr(mod, cdfinvstr) - if cp.has_option_tag(section, 'rvs', tag): - rvsstr = cp.get_opt_tag(section, 'rvs', tag) + if cp.has_option_tag(section, "rvs", tag): + rvsstr = cp.get_opt_tag(section, "rvs", tag) rvs = getattr(mod, rvsstr) return cls(params=params, logpdf=logpdf, rvs=rvs, cdfinv=cdfinv) class DistributionFunctionFromFile(External): - r"""Evaluating PDF, logPDF, CDF and inverse CDF from the external + r""" + Evaluating PDF, logPDF, CDF and inverse CDF from the external density function. To add to an inference configuration file: @@ -153,19 +156,20 @@ class DistributionFunctionFromFile(External): This class is different from `pycbc.distributions.arbitrary.FromFile`, which needs samples from the hdf file to construct the PDF by using KDE. This class reads in any continuous functions of the parameter. + """ + name = "external_func_fromfile" - def __init__(self, params=None, file_path=None, - column_index=None, **kwargs): + def __init__(self, params=None, file_path=None, column_index=None, **kwargs): super().__init__(cdfinv=self._cdfinv, logpdf=self.logpdf) self.params = params - self.data = np.loadtxt(fname=file_path, unpack=True, comments='#') + self.data = np.loadtxt(fname=file_path, unpack=True, comments="#") self.column_index = int(column_index) - self.epsabs = kwargs.get('epsabs', 1.49e-05) - self.epsrel = kwargs.get('epsrel', 1.49e-05) + self.epsabs = kwargs.get("epsabs", 1.49e-05) + self.epsrel = kwargs.get("epsrel", 1.49e-05) self.x_list = np.linspace(self.data[0][0], self.data[0][-1], 1000) - self.interp = {'pdf': callable, 'cdf': callable, 'cdfinv': callable} + self.interp = {"pdf": callable, "cdf": callable, "cdfinv": callable} if not file_path: raise ValueError("Must provide the path to density function file.") @@ -174,19 +178,30 @@ def logpdf(self, **kwargs): return self._logpdf(x, **kwargs) def _pdf(self, x010, **kwargs): - """Calculate and interpolate the PDF by using the given density - function, then return the corresponding value at the given x.""" - if self.interp['pdf'] == callable: + """ + Calculate and interpolate the PDF by using the given density + function, then return the corresponding value at the given x. + """ + if self.interp["pdf"] == callable: func_unnorm = scipy_interpolate.interp1d( - self.data[0], self.data[self.column_index]) + self.data[0], self.data[self.column_index] + ) norm_const = scipy_integrate.quad( - func_unnorm, self.data[0][0], self.data[0][-1], - epsabs=self.epsabs, epsrel=self.epsrel, limit=500, - **kwargs)[0] - self.interp['pdf'] = scipy_interpolate.interp1d( - self.data[0], self.data[self.column_index]/norm_const, - bounds_error=False, fill_value=0) - pdf_val = np.float64(self.interp['pdf'](x010)) + func_unnorm, + self.data[0][0], + self.data[0][-1], + epsabs=self.epsabs, + epsrel=self.epsrel, + limit=500, + **kwargs, + )[0] + self.interp["pdf"] = scipy_interpolate.interp1d( + self.data[0], + self.data[self.column_index] / norm_const, + bounds_error=False, + fill_value=0, + ) + pdf_val = np.float64(self.interp["pdf"](x010)) return pdf_val def _logpdf(self, x010, **kwargs): @@ -195,41 +210,49 @@ def _logpdf(self, x010, **kwargs): return z def _cdf(self, x, **kwargs): - """Calculate and interpolate the CDF, then return the corresponding - value at the given x.""" - if self.interp['cdf'] == callable: + """ + Calculate and interpolate the CDF, then return the corresponding + value at the given x. + """ + if self.interp["cdf"] == callable: cdf_list = [] for x_val in self.x_list: cdf_x = scipy_integrate.quad( - self._pdf, self.data[0][0], x_val, epsabs=self.epsabs, - epsrel=self.epsrel, limit=500, **kwargs)[0] + self._pdf, + self.data[0][0], + x_val, + epsabs=self.epsabs, + epsrel=self.epsrel, + limit=500, + **kwargs, + )[0] cdf_list.append(cdf_x) - self.interp['cdf'] = \ - scipy_interpolate.interp1d(self.x_list, cdf_list) - cdf_val = np.float64(self.interp['cdf'](x)) + self.interp["cdf"] = scipy_interpolate.interp1d(self.x_list, cdf_list) + cdf_val = np.float64(self.interp["cdf"](x)) return cdf_val def _cdfinv(self, **kwargs): - """Calculate and interpolate the inverse CDF, then return the - corresponding parameter value at the given CDF value.""" - if self.interp['cdfinv'] == callable: + """ + Calculate and interpolate the inverse CDF, then return the + corresponding parameter value at the given CDF value. + """ + if self.interp["cdfinv"] == callable: cdf_list = [] for x_value in self.x_list: cdf_list.append(self._cdf(x_value)) - self.interp['cdfinv'] = \ - scipy_interpolate.interp1d(cdf_list, self.x_list) - cdfinv_val = {self.params[0]: np.float64( - self.interp['cdfinv'](kwargs[self.params[0]]))} + self.interp["cdfinv"] = scipy_interpolate.interp1d(cdf_list, self.x_list) + cdfinv_val = { + self.params[0]: np.float64(self.interp["cdfinv"](kwargs[self.params[0]])) + } return cdfinv_val @classmethod def from_config(cls, cp, section, variable_args): tag = variable_args params = variable_args.split(VARARGS_DELIM) - file_path = cp.get_opt_tag(section, 'file_path', tag) - column_index = cp.get_opt_tag(section, 'column_index', tag) - return cls(params=params, file_path=file_path, - column_index=column_index) + file_path = cp.get_opt_tag(section, "file_path", tag) + column_index = cp.get_opt_tag(section, "column_index", tag) + return cls(params=params, file_path=file_path, column_index=column_index) -__all__ = ['External', 'DistributionFunctionFromFile'] +__all__ = ["DistributionFunctionFromFile", "External"] diff --git a/pycbc/distributions/fixedsamples.py b/pycbc/distributions/fixedsamples.py index dbf7cd878f5..8b6ed8141a4 100644 --- a/pycbc/distributions/fixedsamples.py +++ b/pycbc/distributions/fixedsamples.py @@ -16,16 +16,18 @@ This modules provides classes for evaluating distributions based on a fixed set of points """ + import logging + import numpy import numpy.random from pycbc import VARARGS_DELIM -logger = logging.getLogger('pycbc.distributions.fixedsamples') +logger = logging.getLogger("pycbc.distributions.fixedsamples") -class FixedSamples(object): +class FixedSamples: """ A distribution consisting of a collection of a large number of fixed points. Only these values can be drawn from, so the number of points may need to be @@ -44,6 +46,7 @@ class FixedSamples(object): Sampled points of the distribution. May contain transformed parameters which are different from the original distribution. If so, an inverse mapping is provided to associate points with other parameters provided. + """ name = "fixed_samples" @@ -53,18 +56,19 @@ def __init__(self, params, samples): self.samples = samples self.p1 = self.samples[params[0]] - self.frac = len(self.p1)**0.5 / len(self.p1) + self.frac = len(self.p1) ** 0.5 / len(self.p1) self.sort = self.p1.argsort() self.p1sorted = self.p1[self.sort] assert len(numpy.unique(self.p1)) == len(self.p1) if len(params) > 2: - raise ValueError("Only one or two parameters supported " - "for fixed sample distribution") + raise ValueError( + "Only one or two parameters supported for fixed sample distribution" + ) def rvs(self, size=1, **kwds): - "Draw random value" + """Draw random value""" i = numpy.random.randint(0, high=len(self.p1), size=size) return {p: self.samples[p][i] for p in self.params} @@ -72,13 +76,12 @@ def cdfinv(self, **original): """Map unit cube to parameters in the space""" new = {} - #First dimension + # First dimension u1 = original[self.params[0]] i1 = int(round(u1 * len(self.p1))) if i1 >= len(self.p1): i1 = len(self.p1) - 1 - if i1 < 0: - i1 = 0 + i1 = max(i1, 0) new[self.params[0]] = p1v = self.p1sorted[i1] if len(self.params) == 1: return new @@ -100,8 +103,7 @@ def cdfinv(self, **original): i2 = int(round(u2 * len(p2part))) if i2 >= len(p2part): i2 = len(p2part) - 1 - if i2 < 0: - i2 = 0 + i2 = max(i2, 0) new[self.params[1]] = p2part[i2] p1part = numpy.array(self.p1[region[l]], ndmin=1) @@ -109,16 +111,17 @@ def cdfinv(self, **original): return new def apply_boundary_conditions(self, **params): - """ Apply boundary conditions (none here) """ + """Apply boundary conditions (none here)""" return params def __call__(self, **kwds): - """ Dummy function, not the actual pdf """ + """Dummy function, not the actual pdf""" return 0 @classmethod def from_config(cls, cp, section, tag): - """ Return instance based on config file + """ + Return instance based on config file Return a new instance based on the config file. This will draw from a single distribution section provided in the config file and @@ -127,25 +130,30 @@ def from_config(cls, cp, section, tag): file. """ from pycbc.distributions import read_distributions_from_config - from pycbc.transforms import (read_transforms_from_config, - apply_transforms, BaseTransform) + from pycbc.transforms import ( + BaseTransform, + apply_transforms, + read_transforms_from_config, + ) from pycbc.transforms import transforms as global_transforms params = tag.split(VARARGS_DELIM) - subname = cp.get_opt_tag(section, 'subname', tag) - size = cp.get_opt_tag(section, 'sample-size', tag) + subname = cp.get_opt_tag(section, "subname", tag) + size = cp.get_opt_tag(section, "sample-size", tag) - distsec = '{}_sample'.format(subname) + distsec = f"{subname}_sample" dist = read_distributions_from_config(cp, section=distsec) if len(dist) > 1: - raise ValueError("Fixed sample distrubtion only supports a single" - " distribution to sample from.") + raise ValueError( + "Fixed sample distrubtion only supports a single" + " distribution to sample from." + ) - logger.info('Drawing samples for fixed sample distribution:%s', params) + logger.info("Drawing samples for fixed sample distribution:%s", params) samples = dist[0].rvs(size=int(float(size))) samples = {p: samples[p] for p in samples.dtype.names} - transec = '{}_transform'.format(subname) + transec = f"{subname}_transform" trans = read_transforms_from_config(cp, section=transec) if len(trans) > 0: trans = trans[0] @@ -161,11 +169,14 @@ class Thook(BaseTransform): p1name = params[0] sort = p1.argsort() p1sorted = p1[sort] + def transform(self, maps): idx = numpy.searchsorted(self.p1sorted, maps[self.p1name]) out = {p: samples[p][self.sort[idx]] for p in self.outputs} return self.format_output(maps, out) + global_transforms[Thook.name] = Thook return cls(params, samples) -__all__ = ['FixedSamples'] + +__all__ = ["FixedSamples"] diff --git a/pycbc/distributions/gaussian.py b/pycbc/distributions/gaussian.py index 1f235bc4d43..f2233145068 100644 --- a/pycbc/distributions/gaussian.py +++ b/pycbc/distributions/gaussian.py @@ -15,17 +15,21 @@ """ This modules provides classes for evaluating Gaussian distributions. """ + import logging + import numpy -from scipy.special import erf, erfinv import scipy.stats +from scipy.special import erf, erfinv from pycbc.distributions import bounded -logger = logging.getLogger('pycbc.distributions.gaussian') +logger = logging.getLogger("pycbc.distributions.gaussian") + class Gaussian(bounded.BoundedDist): - r"""A Gaussian distribution on the given parameters; the parameters are + r""" + A Gaussian distribution on the given parameters; the parameters are independent of each other. Bounds can be provided on each parameter, in which case the distribution @@ -77,7 +81,9 @@ class Gaussian(bounded.BoundedDist): Create a bounded Gaussian distribution with the same parameters, but with cyclic boundary conditions: >>> dist = distributions.Gaussian(mass1=Bounds(1,10, cyclic=True), mass1_mean=3, mass1_var=2) + """ + name = "gaussian" def __init__(self, **params): @@ -92,46 +98,45 @@ def __init__(self, **params): self._lognorm = {} self._expnorm = {} # pull out specified means, variance - mean_args = [p for p in params if p.endswith('_mean')] - var_args = [p for p in params if p.endswith('_var')] + mean_args = [p for p in params if p.endswith("_mean")] + var_args = [p for p in params if p.endswith("_var")] self._mean = dict([[p[:-5], params.pop(p)] for p in mean_args]) self._var = dict([[p[:-4], params.pop(p)] for p in var_args]) # initialize the bounds - super(Gaussian, self).__init__(**params) + super().__init__(**params) # check that there are no params in mean/var that are not in params missing = set(self._mean.keys()) - set(params.keys()) if any(missing): - raise ValueError("means provided for unknow params {}".format( - ', '.join(missing))) + raise ValueError( + "means provided for unknow params {}".format(", ".join(missing)) + ) missing = set(self._var.keys()) - set(params.keys()) if any(missing): - raise ValueError("vars provided for unknow params {}".format( - ', '.join(missing))) + raise ValueError( + "vars provided for unknow params {}".format(", ".join(missing)) + ) # set default mean/var for params not specified - self._mean.update(dict([[p, 0.] - for p in params if p not in self._mean])) - self._var.update(dict([[p, 1.] - for p in params if p not in self._var])) + self._mean.update(dict([[p, 0.0] for p in params if p not in self._mean])) + self._var.update(dict([[p, 1.0] for p in params if p not in self._var])) # compute norms - for p,bnds in self._bounds.items(): + for p, bnds in self._bounds.items(): sigmasq = self._var[p] mu = self._mean[p] - a,b = bnds - invnorm = scipy.stats.norm.cdf(b, loc=mu, scale=sigmasq**0.5) \ - - scipy.stats.norm.cdf(a, loc=mu, scale=sigmasq**0.5) - invnorm *= numpy.sqrt(2*numpy.pi*sigmasq) - self._norm[p] = 1./invnorm + a, b = bnds + invnorm = scipy.stats.norm.cdf( + b, loc=mu, scale=sigmasq**0.5 + ) - scipy.stats.norm.cdf(a, loc=mu, scale=sigmasq**0.5) + invnorm *= numpy.sqrt(2 * numpy.pi * sigmasq) + self._norm[p] = 1.0 / invnorm self._lognorm[p] = numpy.log(self._norm[p]) - self._expnorm[p] = -1./(2*sigmasq) - + self._expnorm[p] = -1.0 / (2 * sigmasq) @property def mean(self): return self._mean - @property def var(self): return self._var @@ -140,7 +145,7 @@ def _normalcdf(self, param, value): """The CDF of the normal distribution, without bounds.""" mu = self._mean[param] var = self._var[param] - return 0.5*(1. + erf((value - mu)/(2*var)**0.5)) + return 0.5 * (1.0 + erf((value - mu) / (2 * var) ** 0.5)) def cdf(self, param, value): """Returns the CDF of the given parameter value.""" @@ -148,58 +153,62 @@ def cdf(self, param, value): if a != -numpy.inf: phi_a = self._normalcdf(param, a) else: - phi_a = 0. + phi_a = 0.0 if b != numpy.inf: phi_b = self._normalcdf(param, b) else: - phi_b = 1. + phi_b = 1.0 phi_x = self._normalcdf(param, value) - return (phi_x - phi_a)/(phi_b - phi_a) + return (phi_x - phi_a) / (phi_b - phi_a) def _normalcdfinv(self, param, p): """The inverse CDF of the normal distribution, without bounds.""" mu = self._mean[param] var = self._var[param] - return mu + (2*var)**0.5 * erfinv(2*p - 1.) + return mu + (2 * var) ** 0.5 * erfinv(2 * p - 1.0) def _cdfinv_param(self, param, p): - """Return inverse of the CDF. - """ + """Return inverse of the CDF.""" a, b = self._bounds[param] if a != -numpy.inf: phi_a = self._normalcdf(param, a) else: - phi_a = 0. + phi_a = 0.0 if b != numpy.inf: phi_b = self._normalcdf(param, b) else: - phi_b = 1. + phi_b = 1.0 adjusted_p = phi_a + p * (phi_b - phi_a) return self._normalcdfinv(param, adjusted_p) def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ return numpy.exp(self._logpdf(**kwargs)) - def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs in self: - return sum([self._lognorm[p] + - self._expnorm[p]*(kwargs[p]-self._mean[p])**2. - for p in self._params]) - else: - return -numpy.inf + return sum( + [ + self._lognorm[p] + + self._expnorm[p] * (kwargs[p] - self._mean[p]) ** 2.0 + for p in self._params + ] + ) + return -numpy.inf @classmethod def from_config(cls, cp, section, variable_args): - """Returns a Gaussian distribution based on a configuration file. The + """ + Returns a Gaussian distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -235,9 +244,11 @@ def from_config(cls, cp, section, variable_args): ------- Gaussian A distribution instance from the pycbc.inference.prior module. + """ - return bounded.bounded_from_config(cls, cp, section, variable_args, - bounds_required=False) + return bounded.bounded_from_config( + cls, cp, section, variable_args, bounds_required=False + ) -__all__ = ['Gaussian'] +__all__ = ["Gaussian"] diff --git a/pycbc/distributions/joint.py b/pycbc/distributions/joint.py index 920d91cdca9..9de3a895a49 100644 --- a/pycbc/distributions/joint.py +++ b/pycbc/distributions/joint.py @@ -12,17 +12,18 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This module provides classes to describe joint distributions -""" +"""This module provides classes to describe joint distributions""" + import logging + import numpy from pycbc.io.record import FieldArray -logger = logging.getLogger('pycbc.distributions.joint') +logger = logging.getLogger("pycbc.distributions.joint") -class JointDistribution(object): +class JointDistribution: r""" Callable class that calculates the joint distribution built from a set of distributions. @@ -73,7 +74,8 @@ class JointDistribution(object): >>> print(prior_eval(mass1=20, mass2=1)) """ - name = 'joint' + + name = "joint" def __init__(self, variable_args, *distributions, **kwargs): @@ -85,8 +87,9 @@ def __init__(self, variable_args, *distributions, **kwargs): # store the constraints on the parameters defined inside the # distributions list - self._constraints = kwargs["constraints"] \ - if "constraints" in kwargs.keys() else [] + self._constraints = ( + kwargs["constraints"] if "constraints" in kwargs else [] + ) # store kwargs self.kwargs = kwargs @@ -100,20 +103,25 @@ def __init__(self, variable_args, *distributions, **kwargs): varset = set(self.variable_args) missing_params = distparams - varset if missing_params: - raise ValueError("provided variable_args do not include " - "parameters %s" %(','.join(missing_params)) + " which are " - "required by the provided distributions") + raise ValueError( + "provided variable_args do not include " + "parameters %s" % (",".join(missing_params)) + " which are " + "required by the provided distributions" + ) extra_params = varset - distparams if extra_params: - raise ValueError("variable_args %s " %(','.join(extra_params)) + - "are not in any of the provided distributions") + raise ValueError( + "variable_args %s " % (",".join(extra_params)) + + "are not in any of the provided distributions" + ) # if there are constraints then find the renormalization factor # since a constraint will cut out part of the space # do this by random sampling the full space and find the percent # of samples rejected - n_test_samples = kwargs["n_test_samples"] \ - if "n_test_samples" in kwargs else int(1e6) + n_test_samples = ( + kwargs["n_test_samples"] if "n_test_samples" in kwargs else int(1e6) + ) if self._constraints: logger.info("Renormalizing distribution for constraints") @@ -132,9 +140,11 @@ def __init__(self, variable_args, *distributions, **kwargs): # the fraction of acceptances in random sampling of entire space self._pdf_scale = result.sum() / float(n_test_samples) if self._pdf_scale == 0.0: - raise ValueError("None of the random draws for pdf " + raise ValueError( + "None of the random draws for pdf " "renormalization satisfied the constraints. " - " You can try increasing the 'n_test_samples' keyword.") + " You can try increasing the 'n_test_samples' keyword." + ) else: self._pdf_scale = 1.0 @@ -144,7 +154,8 @@ def __init__(self, variable_args, *distributions, **kwargs): self._logpdf_scale = numpy.log(self._pdf_scale) def apply_boundary_conditions(self, **params): - """Applies each distributions' boundary conditions to the given list + """ + Applies each distributions' boundary conditions to the given list of parameters, returning a new list with the conditions applied. Parameters @@ -158,6 +169,7 @@ def apply_boundary_conditions(self, **params): dict A dictionary of the parameters after each distribution's `apply_boundary_conditions` function has been applied. + """ for dist in self.distributions: params.update(dist.apply_boundary_conditions(**params)) @@ -165,7 +177,8 @@ def apply_boundary_conditions(self, **params): @staticmethod def _return_atomic(params): - """Determines if an array or atomic value should be returned given a + """ + Determines if an array or atomic value should be returned given a set of input params. Parameters @@ -178,24 +191,25 @@ def _return_atomic(params): bool : Whether or not functions run on the parameters should be returned as atomic types or not. + """ if isinstance(params, dict): - return not any(isinstance(val, numpy.ndarray) - for val in params.values()) - elif isinstance(params, numpy.record): + return not any(isinstance(val, numpy.ndarray) for val in params.values()) + if isinstance(params, numpy.record): return True - elif isinstance(params, numpy.ndarray): + if isinstance(params, numpy.ndarray): return False params = params.view(type=FieldArray) - elif isinstance(params, FieldArray): + if isinstance(params, FieldArray): return False - else: - raise ValueError("params must be either dict, FieldArray, " - "record, or structured array") + raise ValueError( + "params must be either dict, FieldArray, record, or structured array" + ) @staticmethod def _ensure_fieldarray(params): - """Ensures the given params are a ``FieldArray``. + """ + Ensures the given params are a ``FieldArray``. Parameters ---------- @@ -207,22 +221,23 @@ def _ensure_fieldarray(params): ------- FieldArray The given values as a FieldArray. + """ if isinstance(params, dict): return FieldArray.from_kwargs(**params) - elif isinstance(params, numpy.record): - return FieldArray.from_records(tuple(params), - names=params.dtype.names) - elif isinstance(params, numpy.ndarray): + if isinstance(params, numpy.record): + return FieldArray.from_records(tuple(params), names=params.dtype.names) + if isinstance(params, numpy.ndarray): return params.view(type=FieldArray) - elif isinstance(params, FieldArray): + if isinstance(params, FieldArray): return params - else: - raise ValueError("params must be either dict, FieldArray, " - "record, or structured array") + raise ValueError( + "params must be either dict, FieldArray, record, or structured array" + ) def within_constraints(self, params): - """Evaluates whether the given parameters satisfy the constraints. + """ + Evaluates whether the given parameters satisfy the constraints. Parameters ---------- @@ -235,6 +250,7 @@ def within_constraints(self, params): If params was an array, or if params a dictionary and one or more of the parameters are arrays, will return an array of booleans. Otherwise, a boolean. + """ params = self._ensure_fieldarray(params) return_atomic = self._return_atomic(params) @@ -247,7 +263,8 @@ def within_constraints(self, params): return result def contains(self, params): - """Evaluates whether the given parameters satisfy the boundary + """ + Evaluates whether the given parameters satisfy the boundary conditions, boundaries, and constraints. This method is different from `within_constraints`, that method only check the constraints. @@ -262,6 +279,7 @@ def contains(self, params): If params was an array, or if params a dictionary and one or more of the parameters are arrays, will return an array of booleans. Otherwise, a boolean. + """ params = self.apply_boundary_conditions(**params) result = True @@ -279,8 +297,7 @@ def contains(self, params): return result def __call__(self, **params): - """Evaluate joint distribution for parameters. - """ + """Evaluate joint distribution for parameters.""" return_atomic = self._return_atomic(params) # check if statisfies constraints if len(self._constraints) != 0: @@ -308,8 +325,7 @@ def __call__(self, **params): return logp - self._logpdf_scale def rvs(self, size=1): - """ Rejection samples the parameter space. - """ + """Rejection samples the parameter space.""" # create output FieldArray dtype = [(arg, float) for arg in self.variable_args] out = FieldArray(size, dtype=dtype) @@ -331,28 +347,26 @@ def rvs(self, size=1): nkeep = keep.sum() kmin = size - remaining kmax = min(nkeep, remaining) - out[kmin:kmin+kmax] = scratch[keep][:kmax] + out[kmin : kmin + kmax] = scratch[keep][:kmax] remaining = max(0, remaining - nkeep) # to try to speed up next go around, we'll increase the draw # size by the fraction of values that were kept, but cap at 1e6 - ndraw = int(min(1e6, ndraw * numpy.ceil(ndraw / (nkeep + 1.)))) + ndraw = int(min(1e6, ndraw * numpy.ceil(ndraw / (nkeep + 1.0)))) return out @property def well_reflected(self): - """ Get list of which parameters are well reflected - """ + """Get list of which parameters are well reflected""" reflect = [] bounds = self.bounds for param in bounds: - if bounds[param].reflected == 'well': + if bounds[param].reflected == "well": reflect.append(param) return reflect @property def cyclic(self): - """ Get list of which parameters are cyclic - """ + """Get list of which parameters are cyclic""" cyclic = [] bounds = self.bounds for param in bounds: @@ -362,16 +376,16 @@ def cyclic(self): @property def bounds(self): - """ Get the dict of boundaries - """ + """Get the dict of boundaries""" bnds = {} for dist in self.distributions: - if hasattr(dist, 'bounds'): + if hasattr(dist, "bounds"): bnds.update(dist.bounds) return bnds def cdfinv(self, **original): - """ Apply the inverse cdf to the array of values [0, 1]. Every + """ + Apply the inverse cdf to the array of values [0, 1]. Every variable parameter must be given as a keyword argument. """ updated = {} diff --git a/pycbc/distributions/mass.py b/pycbc/distributions/mass.py index f32ac4c846c..4faada7b6be 100644 --- a/pycbc/distributions/mass.py +++ b/pycbc/distributions/mass.py @@ -13,23 +13,25 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This modules provides classes for evaluating distributions for mchirp and +""" +This modules provides classes for evaluating distributions for mchirp and q (i.e., mass ratio) from uniform component mass. """ + import logging -import numpy +import numpy from scipy.interpolate import interp1d from scipy.special import hyp2f1 -from pycbc.distributions import power_law -from pycbc.distributions import bounded +from pycbc.distributions import bounded, power_law -logger = logging.getLogger('pycbc.distributions.mass') +logger = logging.getLogger("pycbc.distributions.mass") class MchirpfromUniformMass1Mass2(power_law.UniformPowerLaw): - r"""A distribution for chirp mass from uniform component mass + + r""" + A distribution for chirp mass from uniform component mass + constraints given by chirp mass. This is a special case for UniformPowerLaw with index 1. For more details see UniformPowerLaw. @@ -74,7 +76,6 @@ class MchirpfromUniformMass1Mass2(power_law.UniformPowerLaw): Examples -------- - Generate 10000 random numbers from this distribution in [5,100] >>> from pycbc import distributions as dist @@ -98,16 +99,18 @@ class MchirpfromUniformMass1Mass2(power_law.UniformPowerLaw): The keyword arguments should provide the names of parameters and their corresponding bounds, as either tuples or a `boundaries.Bounds` instance. + """ name = "mchirp_from_uniform_mass1_mass2" def __init__(self, dim=2, **params): - super(MchirpfromUniformMass1Mass2, self).__init__(dim=2, **params) + super().__init__(dim=2, **params) class QfromUniformMass1Mass2(bounded.BoundedDist): - r"""A distribution for mass ratio (i.e., q) from uniform component mass + r""" + A distribution for mass ratio (i.e., q) from uniform component mass + constraints given by q. The parameters (i.e. `**params`) are independent of each other. Instances @@ -127,7 +130,6 @@ class QfromUniformMass1Mass2(bounded.BoundedDist): Examples -------- - Generate 10000 random numbers from this distribution in [1,8] >>> from pycbc import distributions as dist @@ -136,15 +138,16 @@ class QfromUniformMass1Mass2(bounded.BoundedDist): """ - name = 'q_from_uniform_mass1_mass2' + name = "q_from_uniform_mass1_mass2" def __init__(self, **params): - super(QfromUniformMass1Mass2, self).__init__(**params) + super().__init__(**params) self._norm = 1.0 self._lognorm = 0.0 for p in self._params: - self._norm /= self._cdf_param(p, self._bounds[p][1]) - \ - self._cdf_param(p, self._bounds[p][0]) + self._norm /= self._cdf_param(p, self._bounds[p][1]) - self._cdf_param( + p, self._bounds[p][0] + ) self._lognorm = numpy.log(self._norm) @property @@ -158,75 +161,85 @@ def lognorm(self): return self._lognorm def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ for p in self._params: - if p not in kwargs.keys(): - raise ValueError( - 'Missing parameter {} to construct pdf.'.format(p)) + if p not in kwargs: + raise ValueError(f"Missing parameter {p} to construct pdf.") if kwargs in self: - pdf = self._norm * \ - numpy.prod([(1.+kwargs[p])**(2./5)/kwargs[p]**(6./5) - for p in self._params]) + pdf = self._norm * numpy.prod( + [ + (1.0 + kwargs[p]) ** (2.0 / 5) / kwargs[p] ** (6.0 / 5) + for p in self._params + ] + ) return float(pdf) - else: - return 0.0 + return 0.0 def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ for p in self._params: - if p not in kwargs.keys(): - raise ValueError( - 'Missing parameter {} to construct logpdf.'.format(p)) + if p not in kwargs: + raise ValueError(f"Missing parameter {p} to construct logpdf.") if kwargs in self: return numpy.log(self._pdf(**kwargs)) - else: - return -numpy.inf + return -numpy.inf def _cdf_param(self, param, value): - r""">>> from sympy import * - >>> x = Symbol('x') - >>> integrate((1+x)**(2/5)/x**(6/5)) - Output: - _ - -0.2 |_ /-0.4, -0.2 | I*pi\ + r""" + >>> from sympy import * + >>> x = Symbol('x') + >>> integrate((1+x)**(2/5)/x**(6/5)) + Output: + _ + -0.2 |_ /-0.4, -0.2 | I*pi\ -5.0*x * | | | x*e | - 2 1 \ 0.8 | / + 2 1 \ 0.8 | / """ if param in self._params: - return -5. * value**(-1./5) * hyp2f1(-2./5, -1./5, 4./5, -value) - else: - raise ValueError('{} is not contructed yet.'.format(param)) + return ( + -5.0 * value ** (-1.0 / 5) * hyp2f1(-2.0 / 5, -1.0 / 5, 4.0 / 5, -value) + ) + raise ValueError(f"{param} is not contructed yet.") def _cdfinv_param(self, param, value): - """Return the inverse cdf to map the unit interval to parameter bounds. - Note that value should be uniform in [0,1].""" + """ + Return the inverse cdf to map the unit interval to parameter bounds. + Note that value should be uniform in [0,1]. + """ if (numpy.array(value) < 0).any() or (numpy.array(value) > 1).any(): - raise ValueError( - 'q_from_uniform_m1_m2 cdfinv requires input in [0,1].') + raise ValueError("q_from_uniform_m1_m2 cdfinv requires input in [0,1].") if param in self._params: lower_bound = self._bounds[param][0] upper_bound = self._bounds[param][1] - q_array = numpy.linspace( - lower_bound, upper_bound, num=1000, endpoint=True) - q_invcdf_interp = interp1d(self._cdf_param(param, q_array), - q_array, kind='cubic', - bounds_error=True) + q_array = numpy.linspace(lower_bound, upper_bound, num=1000, endpoint=True) + q_invcdf_interp = interp1d( + self._cdf_param(param, q_array), + q_array, + kind="cubic", + bounds_error=True, + ) return q_invcdf_interp( - (self._cdf_param(param, upper_bound) - - self._cdf_param(param, lower_bound)) * value + - self._cdf_param(param, lower_bound)) - else: - raise ValueError('{} is not contructed yet.'.format(param)) + ( + self._cdf_param(param, upper_bound) + - self._cdf_param(param, lower_bound) + ) + * value + + self._cdf_param(param, lower_bound) + ) + raise ValueError(f"{param} is not contructed yet.") def rvs(self, size=1, param=None): - """Gives a set of random values drawn from this distribution. + """ + Gives a set of random values drawn from this distribution. Parameters ---------- @@ -243,20 +256,22 @@ def rvs(self, size=1, param=None): specified, the array will only have an element corresponding to the given parameter. Otherwise, the array will have an element for each parameter in self's params. + """ if param is not None: dtype = [(param, float)] else: dtype = [(p, float) for p in self.params] arr = numpy.zeros(size, dtype=dtype) - for (p, _) in dtype: + for p, _ in dtype: uniformcdfvalue = numpy.random.uniform(0, 1, size=size) arr[p] = self._cdfinv_param(p, uniformcdfvalue) return arr @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. The parameters + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -288,9 +303,11 @@ def from_config(cls, cp, section, variable_args): QfromUniformMass1Mass2 A distribution instance from the pycbc.distributions.bounded module. + """ - return super(QfromUniformMass1Mass2, cls).from_config( - cp, section, variable_args, bounds_required=True) + return super().from_config( + cp, section, variable_args, bounds_required=True + ) __all__ = ["MchirpfromUniformMass1Mass2", "QfromUniformMass1Mass2"] diff --git a/pycbc/distributions/power_law.py b/pycbc/distributions/power_law.py index f87f2124f32..a86796fbdd5 100644 --- a/pycbc/distributions/power_law.py +++ b/pycbc/distributions/power_law.py @@ -16,12 +16,15 @@ This modules provides classes for evaluating distributions where the probability density function is a power law. """ + import logging + import numpy from pycbc.distributions import bounded -logger = logging.getLogger('pycbc.distributions.power_law') +logger = logging.getLogger("pycbc.distributions.power_law") + class UniformPowerLaw(bounded.BoundedDist): r""" @@ -113,17 +116,20 @@ class UniformPowerLaw(bounded.BoundedDist): dim : int The dimension of volume space. In the notation above `dim` is :math:`n+1`. For a 3-dimensional sphere this is 3. + """ + name = "uniform_power_law" + def __init__(self, dim=None, **params): - super(UniformPowerLaw, self).__init__(**params) + super().__init__(**params) self.dim = dim self._norm = 1.0 self._lognorm = 0.0 for p in self._params: - self._norm *= self.dim / \ - (self._bounds[p][1]**(self.dim) - - self._bounds[p][0]**(self.dim)) + self._norm *= self.dim / ( + self._bounds[p][1] ** (self.dim) - self._bounds[p][0] ** (self.dim) + ) self._lognorm = numpy.log(self._norm) @property @@ -137,51 +143,52 @@ def lognorm(self): return self._lognorm def _cdfinv_param(self, param, value): - """Return inverse of cdf to map unit interval to parameter bounds. - """ + """Return inverse of cdf to map unit interval to parameter bounds.""" n = self.dim - 1 r_l = self._bounds[param][0] r_h = self._bounds[param][1] - new_value = ((r_h**(n+1) - r_l**(n+1))*value + r_l**(n+1))**(1./(n+1)) + new_value = ((r_h ** (n + 1) - r_l ** (n + 1)) * value + r_l ** (n + 1)) ** ( + 1.0 / (n + 1) + ) return new_value def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ for p in self._params: - if p not in kwargs.keys(): - raise ValueError( - 'Missing parameter {} to construct pdf.'.format(p)) + if p not in kwargs: + raise ValueError(f"Missing parameter {p} to construct pdf.") if kwargs in self: - pdf = self._norm * \ - numpy.prod([(kwargs[p])**(self.dim - 1) - for p in self._params]) + pdf = self._norm * numpy.prod( + [(kwargs[p]) ** (self.dim - 1) for p in self._params] + ) return float(pdf) - else: - return 0.0 + return 0.0 def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ for p in self._params: - if p not in kwargs.keys(): - raise ValueError( - 'Missing parameter {} to construct pdf.'.format(p)) + if p not in kwargs: + raise ValueError(f"Missing parameter {p} to construct pdf.") if kwargs in self: - log_pdf = self._lognorm + \ - (self.dim - 1) * \ - numpy.log([kwargs[p] for p in self._params]).sum() + log_pdf = ( + self._lognorm + + (self.dim - 1) * numpy.log([kwargs[p] for p in self._params]).sum() + ) return log_pdf - else: - return -numpy.inf + return -numpy.inf @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. The parameters + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -201,20 +208,25 @@ def from_config(cls, cp, section, variable_args): ------- Uniform A distribution instance from the pycbc.inference.prior module. + """ - return super(UniformPowerLaw, cls).from_config(cp, section, - variable_args, - bounds_required=True) + return super().from_config( + cp, section, variable_args, bounds_required=True + ) class UniformRadius(UniformPowerLaw): - """ For a uniform distribution in volume using spherical coordinates, this + """ + For a uniform distribution in volume using spherical coordinates, this is the distriubtion to use for the radius. For more details see UniformPowerLaw. """ + name = "uniform_radius" + def __init__(self, dim=3, **params): - super(UniformRadius, self).__init__(dim=3, **params) + super().__init__(dim=3, **params) + __all__ = ["UniformPowerLaw", "UniformRadius"] diff --git a/pycbc/distributions/qnm.py b/pycbc/distributions/qnm.py index f9413c89037..bb6338fa8a0 100644 --- a/pycbc/distributions/qnm.py +++ b/pycbc/distributions/qnm.py @@ -15,18 +15,20 @@ import logging import re + import numpy import pycbc -from pycbc import conversions, boundaries +from pycbc import boundaries, conversions -from . import uniform, bounded +from . import bounded, uniform -logger = logging.getLogger('pycbc.distributions.qnm') +logger = logging.getLogger("pycbc.distributions.qnm") class UniformF0Tau(uniform.Uniform): - """A distribution uniform in QNM frequency and damping time. + """ + A distribution uniform in QNM frequency and damping time. Constraints may be placed to exclude frequencies and damping times corresponding to specific masses and spins. @@ -65,7 +67,6 @@ class UniformF0Tau(uniform.Uniform): Examples -------- - Create a distribution: >>> dist = UniformF0Tau(f0=(10., 2048.), tau=(1e-4,1e-2)) @@ -103,11 +104,19 @@ class UniformF0Tau(uniform.Uniform): """ - name = 'uniform_f0_tau' - - def __init__(self, f0=None, tau=None, final_mass=None, final_spin=None, - rdfreq='f0', damping_time='tau', norm_tolerance=1e-3, - norm_seed=0): + name = "uniform_f0_tau" + + def __init__( + self, + f0=None, + tau=None, + final_mass=None, + final_spin=None, + rdfreq="f0", + damping_time="tau", + norm_tolerance=1e-3, + norm_seed=0, + ): if f0 is None: raise ValueError("must provide a range for f0") if tau is None: @@ -115,35 +124,39 @@ def __init__(self, f0=None, tau=None, final_mass=None, final_spin=None, self.rdfreq = rdfreq self.damping_time = damping_time parent_args = {rdfreq: f0, damping_time: tau} - super(UniformF0Tau, self).__init__(**parent_args) + super().__init__(**parent_args) if final_mass is None: - final_mass = (0., numpy.inf) + final_mass = (0.0, numpy.inf) if final_spin is None: final_spin = (-0.996, 0.996) self.final_mass_bounds = boundaries.Bounds( - min_bound=final_mass[0], max_bound=final_mass[1]) + min_bound=final_mass[0], max_bound=final_mass[1] + ) self.final_spin_bounds = boundaries.Bounds( - min_bound=final_spin[0], max_bound=final_spin[1]) + min_bound=final_spin[0], max_bound=final_spin[1] + ) # Re-normalize to account for cuts: we'll do this by just sampling # a large number of spaces f0 taus, and seeing how many are in the # desired range. # perseve the current random state s = numpy.random.get_state() numpy.random.seed(norm_seed) - nsamples = int(1./norm_tolerance**2) - draws = super(UniformF0Tau, self).rvs(size=nsamples) + nsamples = int(1.0 / norm_tolerance**2) + draws = super().rvs(size=nsamples) # reset the random state numpy.random.set_state(s) num_in = self._constraints(draws).sum() # if num_in is 0, than the requested tolerance is too large if num_in == 0: - raise ValueError("the normalization is < then the norm_tolerance; " - "try again with a smaller nrom_tolerance") + raise ValueError( + "the normalization is < then the norm_tolerance; " + "try again with a smaller nrom_tolerance" + ) self._lognorm += numpy.log(num_in) - numpy.log(nsamples) self._norm = numpy.exp(self._lognorm) def __contains__(self, params): - isin = super(UniformF0Tau, self).__contains__(params) + isin = super().__contains__(params) if isin: isin &= self._constraints(params) return isin @@ -152,8 +165,8 @@ def _constraints(self, params): f0 = params[self.rdfreq] tau = params[self.damping_time] # check if we need to specify a particular mode (l,m) != (2,2) - if re.match(r'f_\d{3}', self.rdfreq): - mode = self.rdfreq.strip('f_') + if re.match(r"f_\d{3}", self.rdfreq): + mode = self.rdfreq.strip("f_") l, m = int(mode[0]), int(mode[1]) else: l, m = 2, 2 @@ -163,11 +176,13 @@ def _constraints(self, params): mf = conversions.final_mass_from_f0_tau(f0, tau, l=l, m=m) sf = conversions.final_spin_from_f0_tau(f0, tau, l=l, m=m) isin = (self.final_mass_bounds.__contains__(mf)) & ( - self.final_spin_bounds.__contains__(sf)) + self.final_spin_bounds.__contains__(sf) + ) return isin def rvs(self, size=1): - """Draw random samples from this distribution. + """ + Draw random samples from this distribution. Parameters ---------- @@ -178,6 +193,7 @@ def rvs(self, size=1): ------- array A structured array of the random draws. + """ size = int(size) dtype = [(p, float) for p in self.params] @@ -185,17 +201,18 @@ def rvs(self, size=1): remaining = size keepidx = 0 while remaining: - draws = super(UniformF0Tau, self).rvs(size=remaining) + draws = super().rvs(size=remaining) mask = self._constraints(draws) addpts = mask.sum() - arr[keepidx:keepidx+addpts] = draws[mask] + arr[keepidx : keepidx + addpts] = draws[mask] keepidx += addpts remaining = size - keepidx return arr @classmethod def from_config(cls, cp, section, variable_args): - """Initialize this class from a config file. + """ + Initialize this class from a config file. Bounds on ``f0``, ``tau``, ``final_mass`` and ``final_spin`` should be specified by providing ``min-{param}`` and ``max-{param}``. If @@ -233,38 +250,45 @@ def from_config(cls, cp, section, variable_args): UniformF0Tau : This class initialized with the parameters provided in the config file. + """ tag = variable_args variable_args = set(variable_args.split(pycbc.VARARGS_DELIM)) # get f0 and tau - f0 = bounded.get_param_bounds_from_config(cp, section, tag, 'f0') - tau = bounded.get_param_bounds_from_config(cp, section, tag, 'tau') + f0 = bounded.get_param_bounds_from_config(cp, section, tag, "f0") + tau = bounded.get_param_bounds_from_config(cp, section, tag, "tau") # see if f0 and tau should be renamed - if cp.has_option_tag(section, 'rdfreq', tag): - rdfreq = cp.get_opt_tag(section, 'rdfreq', tag) + if cp.has_option_tag(section, "rdfreq", tag): + rdfreq = cp.get_opt_tag(section, "rdfreq", tag) else: - rdfreq = 'f0' - if cp.has_option_tag(section, 'damping_time', tag): - damping_time = cp.get_opt_tag(section, 'damping_time', tag) + rdfreq = "f0" + if cp.has_option_tag(section, "damping_time", tag): + damping_time = cp.get_opt_tag(section, "damping_time", tag) else: - damping_time = 'tau' + damping_time = "tau" # check that they match whats in the variable args if not variable_args == set([rdfreq, damping_time]): - raise ValueError("variable args do not match rdfreq and " - "damping_time names") + raise ValueError("variable args do not match rdfreq and damping_time names") # get the final mass and spin values, if provided final_mass = bounded.get_param_bounds_from_config( - cp, section, tag, 'final_mass') + cp, section, tag, "final_mass" + ) final_spin = bounded.get_param_bounds_from_config( - cp, section, tag, 'final_spin') + cp, section, tag, "final_spin" + ) extra_opts = {} - if cp.has_option_tag(section, 'norm_tolerance', tag): - extra_opts['norm_tolerance'] = float( - cp.get_opt_tag(section, 'norm_tolerance', tag)) - if cp.has_option_tag(section, 'norm_seed', tag): - extra_opts['norm_seed'] = int( - cp.get_opt_tag(section, 'norm_seed', tag)) - return cls(f0=f0, tau=tau, - final_mass=final_mass, final_spin=final_spin, - rdfreq=rdfreq, damping_time=damping_time, - **extra_opts) + if cp.has_option_tag(section, "norm_tolerance", tag): + extra_opts["norm_tolerance"] = float( + cp.get_opt_tag(section, "norm_tolerance", tag) + ) + if cp.has_option_tag(section, "norm_seed", tag): + extra_opts["norm_seed"] = int(cp.get_opt_tag(section, "norm_seed", tag)) + return cls( + f0=f0, + tau=tau, + final_mass=final_mass, + final_spin=final_spin, + rdfreq=rdfreq, + damping_time=damping_time, + **extra_opts, + ) diff --git a/pycbc/distributions/sky_location.py b/pycbc/distributions/sky_location.py index 6ecf6522091..9456476702c 100644 --- a/pycbc/distributions/sky_location.py +++ b/pycbc/distributions/sky_location.py @@ -13,50 +13,52 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This modules provides classes for evaluating sky distributions in +""" +This modules provides classes for evaluating sky distributions in right ascension and declination. """ import copy import logging import warnings + import numpy from scipy.spatial.transform import Rotation -from pycbc.distributions import angular from pycbc import VARARGS_DELIM +from pycbc.distributions import angular from pycbc.io import FieldArray -from pycbc.types import angle_as_radians from pycbc.libutils import import_optional +from pycbc.types import angle_as_radians - -logger = logging.getLogger('pycbc.distributions.sky_location') +logger = logging.getLogger("pycbc.distributions.sky_location") class UniformSky(angular.UniformSolidAngle): - """A distribution that is uniform on the sky. This is the same as + """ + A distribution that is uniform on the sky. This is the same as UniformSolidAngle, except that the polar angle varies from pi/2 (the north pole) to -pi/2 (the south pole) instead of 0 to pi. Also, the default names are "dec" (declination) for the polar angle and "ra" (right ascension) for the azimuthal angle, instead of "theta" and "phi". """ - name = 'uniform_sky' + name = "uniform_sky" _polardistcls = angular.CosAngle - _default_polar_angle = 'dec' - _default_azimuthal_angle = 'ra' + _default_polar_angle = "dec" + _default_azimuthal_angle = "ra" def to_uniform_patch(self, coverage): if coverage < 1: logging.warning( - 'Attempt to convert UniformSky to a ' - 'uniform patch assumes 100% coverage' + "Attempt to convert UniformSky to a uniform patch assumes 100% coverage" ) return self class UniformDiskSky: - """A distribution that represents a uniform disk on the sky. The declination + """ + A distribution that represents a uniform disk on the sky. The declination varies from π/2 to -π/2 and the right ascension varies from 0 to 2π. Parameters @@ -70,31 +72,29 @@ class UniformDiskSky: radius: float or str Radius of the disk. Use the rad or deg suffix to specify units, otherwise radians are assumed. + """ - name = 'uniform_disk_sky' - _params = ['ra', 'dec'] + + name = "uniform_disk_sky" + _params = ["ra", "dec"] def __init__(self, **params): - mean_ra = angle_as_radians(params['mean_ra']) - mean_dec = angle_as_radians(params['mean_dec']) - radius = angle_as_radians(params['radius']) + mean_ra = angle_as_radians(params["mean_ra"]) + mean_dec = angle_as_radians(params["mean_dec"]) + radius = angle_as_radians(params["radius"]) if mean_ra < 0 or mean_ra > 2 * numpy.pi: raise ValueError( - f'The mean RA must be between 0 and 2π, {mean_ra} rad given' + f"The mean RA must be between 0 and 2π, {mean_ra} rad given" ) if mean_dec < -numpy.pi / 2 or mean_dec > numpy.pi / 2: raise ValueError( - 'The mean declination must be between ' - f'-π/2 and π/2, {mean_dec} rad given' + "The mean declination must be between " + f"-π/2 and π/2, {mean_dec} rad given" ) if radius < 0 or radius > 2 * numpy.pi: - raise ValueError( - 'Radius must be non-negative and smaller than 2π' - ) + raise ValueError("Radius must be non-negative and smaller than 2π") # Prepare a rotation that puts the North Pole at the mean position - self.rotation = Rotation.from_euler( - 'yz', [numpy.pi / 2 - mean_dec, mean_ra] - ) + self.rotation = Rotation.from_euler("yz", [numpy.pi / 2 - mean_dec, mean_ra]) self.mean_ra, self.mean_dec, self.radius = mean_ra, mean_dec, radius @property @@ -110,9 +110,9 @@ def from_config(cls, cp, section, variable_args): "Not all parameters used by this distribution " "included in tag portion of section name" ) - mean_ra = cp.get_opt_tag(section, 'mean_ra', tag) - mean_dec = cp.get_opt_tag(section, 'mean_dec', tag) - radius = cp.get_opt_tag(section, 'radius', tag) + mean_ra = cp.get_opt_tag(section, "mean_ra", tag) + mean_dec = cp.get_opt_tag(section, "mean_dec", tag) + radius = cp.get_opt_tag(section, "radius", tag) return cls( mean_ra=mean_ra, mean_dec=mean_dec, @@ -125,9 +125,11 @@ def get_max_prob_point(self): def rvs(self, size): # Draw samples from a distribution centered on the North pole np_ra = numpy.random.uniform(low=0, high=(2 * numpy.pi), size=size) - np_dec = angular.SinAngle( - polar_bounds=(0,self.radius)).rvs( - size=size).astype(numpy.float64) + np_dec = ( + angular.SinAngle(polar_bounds=(0, self.radius)) + .rvs(size=size) + .astype(numpy.float64) + ) # Convert the samples to intermediate cartesian representation np_cart = numpy.empty(shape=(size, 3)) @@ -141,24 +143,25 @@ def rvs(self, size): # Convert the samples back to spherical coordinates. # Some unpleasant conditional operations are needed # to get the correct angle convention. - rot_radec = FieldArray(size, dtype=[('ra', ' 2 * numpy.pi: raise ValueError( - f'The mean RA must be between 0 and 2π, {mean_ra} rad given' + f"The mean RA must be between 0 and 2π, {mean_ra} rad given" ) if mean_dec < -numpy.pi / 2 or mean_dec > numpy.pi / 2: raise ValueError( - 'The mean declination must be between ' - f'-π/2 and π/2, {mean_dec} rad given' + "The mean declination must be between " + f"-π/2 and π/2, {mean_dec} rad given" ) if sigma < 0 or sigma > 2 * numpy.pi: raise ValueError( - 'Sigma must be non-negative and smaller than 2π ' - '(preferably much smaller)' + "Sigma must be non-negative and smaller than 2π " + "(preferably much smaller)" ) if sigma > 0.35: logger.warning( - 'Warning: sigma = %s rad is probably too large for the ' - 'Fisher approximation to be valid', + "Warning: sigma = %s rad is probably too large for the " + "Fisher approximation to be valid", sigma, ) self.rayleigh_scale = 0.66 * sigma # Prepare a rotation that puts the North Pole at the mean position - self.rotation = Rotation.from_euler( - 'yz', [numpy.pi / 2 - mean_dec, mean_ra] - ) + self.rotation = Rotation.from_euler("yz", [numpy.pi / 2 - mean_dec, mean_ra]) # storing center position for `to_uniform_patch()` self.mean_ra, self.mean_dec = mean_ra, mean_dec @@ -241,9 +243,9 @@ def from_config(cls, cp, section, variable_args): "Not all parameters used by this distribution " "included in tag portion of section name" ) - mean_ra = cp.get_opt_tag(section, 'mean_ra', tag) - mean_dec = cp.get_opt_tag(section, 'mean_dec', tag) - sigma = cp.get_opt_tag(section, 'sigma', tag) + mean_ra = cp.get_opt_tag(section, "mean_ra", tag) + mean_dec = cp.get_opt_tag(section, "mean_dec", tag) + sigma = cp.get_opt_tag(section, "sigma", tag) return cls( mean_ra=mean_ra, mean_dec=mean_dec, @@ -270,20 +272,23 @@ def rvs(self, size): # Convert the samples back to spherical coordinates. # Some unpleasant conditional operations are needed # to get the correct angle convention. - rot_radec = FieldArray(size, dtype=[('ra', ' numpy.pi / 2 - boundaries_phi[straddler_mask,:] -= numpy.pi / 2 - boundaries_phi[straddler_mask,:] = self.normalize_azimuth( - boundaries_phi[straddler_mask,:] + boundaries_phi[straddler_mask, :] -= numpy.pi / 2 + boundaries_phi[straddler_mask, :] = self.normalize_azimuth( + boundaries_phi[straddler_mask, :] + ) + boundaries_phi_min[straddler_mask] = boundaries_phi[straddler_mask, :].min( + axis=1 + ) + boundaries_phi_max[straddler_mask] = boundaries_phi[straddler_mask, :].max( + axis=1 ) - boundaries_phi_min[straddler_mask] = boundaries_phi[straddler_mask,:].min(axis=1) - boundaries_phi_max[straddler_mask] = boundaries_phi[straddler_mask,:].max(axis=1) del boundaries_phi final_thetas = numpy.array([]) @@ -443,9 +449,7 @@ def rvs(self, size): random_thetas = numpy.arccos( numpy.random.uniform(boundaries_z_min, boundaries_z_max) ) - random_phis = numpy.random.uniform( - boundaries_phi_min, boundaries_phi_max - ) + random_phis = numpy.random.uniform(boundaries_phi_min, boundaries_phi_max) # Now we can undo the shift for the straddlers. Ugly! random_phis[straddler_mask] += numpy.pi / 2 @@ -454,16 +458,12 @@ def rvs(self, size): ) # Find which points fall inside their pixels. - sampled_indices = self.healpix_map.ang2pix( - random_thetas, random_phis - ) + sampled_indices = self.healpix_map.ang2pix(random_thetas, random_phis) acceptance_mask = sampled_indices == pix_indices final_thetas = numpy.concatenate( (final_thetas, random_thetas[acceptance_mask]) ) - final_phis = numpy.concatenate( - (final_phis, random_phis[acceptance_mask]) - ) + final_phis = numpy.concatenate((final_phis, random_phis[acceptance_mask])) # Iterate if we are still missing some pixels. rej_mask = numpy.logical_not(acceptance_mask) @@ -475,10 +475,10 @@ def rvs(self, size): straddler_mask = straddler_mask[rej_mask] # Convert back to the radec convention - radec = FieldArray(size, dtype=[('ra', ' 1. or xi_bounds[0] < 0.) or ( - xi_bounds[1] > 1. or xi_bounds[1] < 0.): + xi_bounds = (0, 1.0) + if (xi_bounds[0] > 1.0 or xi_bounds[0] < 0.0) or ( + xi_bounds[1] > 1.0 or xi_bounds[1] < 0.0 + ): raise ValueError("xi bounds must be in [0, 1)") self.xi1_distr = UniformPowerLaw(dim=0.5, xi1=xi_bounds) self.xi2_distr = UniformPowerLaw(dim=0.5, xi2=xi_bounds) # the angles - self.phia_distr = UniformAngle(phi_a=(0,2)) - self.phis_distr = UniformAngle(phi_s=(0,2)) - self.distributions = {'mass1': self.mass1_distr, - 'mass2': self.mass2_distr, - 'xi1': self.xi1_distr, - 'xi2': self.xi2_distr, - 'chi_eff': self.chieff_distr, - 'chi_a': self.chia_distr, - 'phi_a': self.phia_distr, - 'phi_s': self.phis_distr} + self.phia_distr = UniformAngle(phi_a=(0, 2)) + self.phis_distr = UniformAngle(phi_s=(0, 2)) + self.distributions = { + "mass1": self.mass1_distr, + "mass2": self.mass2_distr, + "xi1": self.xi1_distr, + "xi2": self.xi2_distr, + "chi_eff": self.chieff_distr, + "chi_a": self.chia_distr, + "phi_a": self.phia_distr, + "phi_s": self.phis_distr, + } # create random variables for the kde if nsamples is None: nsamples = 1e4 @@ -127,16 +145,24 @@ def __init__(self, mass1=None, mass2=None, chi_eff=None, chi_a=None, rvals = self.rvs(size=int(nsamples)) # reset the random state back to what it was numpy.random.set_state(rstate) - bounds = dict(b for distr in self.distributions.values() - for b in distr.bounds.items()) - super(IndependentChiPChiEff, self).__init__(mass1=rvals['mass1'], - mass2=rvals['mass2'], xi1=rvals['xi1'], xi2=rvals['xi2'], - chi_eff=rvals['chi_eff'], chi_a=rvals['chi_a'], - phi_a=rvals['phi_a'], phi_s=rvals['phi_s'], - bounds=bounds) + bounds = dict( + b for distr in self.distributions.values() for b in distr.bounds.items() + ) + super().__init__( + mass1=rvals["mass1"], + mass2=rvals["mass2"], + xi1=rvals["xi1"], + xi2=rvals["xi2"], + chi_eff=rvals["chi_eff"], + chi_a=rvals["chi_a"], + phi_a=rvals["phi_a"], + phi_s=rvals["phi_s"], + bounds=bounds, + ) def _constraints(self, values): - """Applies physical constraints to the given parameter values. + """ + Applies physical constraints to the given parameter values. Parameters ---------- @@ -147,28 +173,42 @@ def _constraints(self, values): ------- bool Whether or not the values satisfy physical + """ - mass1, mass2, phi_a, phi_s, chi_eff, chi_a, xi1, xi2, _ = \ - conversions.ensurearray(values['mass1'], values['mass2'], - values['phi_a'], values['phi_s'], - values['chi_eff'], values['chi_a'], - values['xi1'], values['xi2']) + mass1, mass2, phi_a, phi_s, chi_eff, chi_a, xi1, xi2, _ = ( + conversions.ensurearray( + values["mass1"], + values["mass2"], + values["phi_a"], + values["phi_s"], + values["chi_eff"], + values["chi_a"], + values["xi1"], + values["xi2"], + ) + ) s1x = conversions.spin1x_from_xi1_phi_a_phi_s(xi1, phi_a, phi_s) - s2x = conversions.spin2x_from_mass1_mass2_xi2_phi_a_phi_s(mass1, mass2, - xi2, phi_a, phi_s) + s2x = conversions.spin2x_from_mass1_mass2_xi2_phi_a_phi_s( + mass1, mass2, xi2, phi_a, phi_s + ) s1y = conversions.spin1y_from_xi1_phi_a_phi_s(xi1, phi_a, phi_s) - s2y = conversions.spin2y_from_mass1_mass2_xi2_phi_a_phi_s(mass1, mass2, - xi2, phi_a, phi_s) - s1z = conversions.spin1z_from_mass1_mass2_chi_eff_chi_a(mass1, mass2, - chi_eff, chi_a) - s2z = conversions.spin2z_from_mass1_mass2_chi_eff_chi_a(mass1, mass2, - chi_eff, chi_a) - test = ((s1x**2. + s1y**2. + s1z**2.) < 1.) & \ - ((s2x**2. + s2y**2. + s2z**2.) < 1.) + s2y = conversions.spin2y_from_mass1_mass2_xi2_phi_a_phi_s( + mass1, mass2, xi2, phi_a, phi_s + ) + s1z = conversions.spin1z_from_mass1_mass2_chi_eff_chi_a( + mass1, mass2, chi_eff, chi_a + ) + s2z = conversions.spin2z_from_mass1_mass2_chi_eff_chi_a( + mass1, mass2, chi_eff, chi_a + ) + test = ((s1x**2.0 + s1y**2.0 + s1z**2.0) < 1.0) & ( + (s2x**2.0 + s2y**2.0 + s2z**2.0) < 1.0 + ) return test def __contains__(self, params): - """Determines whether the given values are in each parameter's bounds + """ + Determines whether the given values are in each parameter's bounds and satisfy the constraints. """ isin = all([params in dist for dist in self.distributions.values()]) @@ -178,63 +218,60 @@ def __contains__(self, params): return self._constraints(params) def _draw(self, size=1, **kwargs): - """Draws random samples without applying physical constrains. - """ + """Draws random samples without applying physical constrains.""" # draw masses try: - mass1 = kwargs['mass1'] + mass1 = kwargs["mass1"] except KeyError: - mass1 = self.mass1_distr.rvs(size=size)['mass1'] + mass1 = self.mass1_distr.rvs(size=size)["mass1"] try: - mass2 = kwargs['mass2'] + mass2 = kwargs["mass2"] except KeyError: - mass2 = self.mass2_distr.rvs(size=size)['mass2'] + mass2 = self.mass2_distr.rvs(size=size)["mass2"] # draw angles try: - phi_a = kwargs['phi_a'] + phi_a = kwargs["phi_a"] except KeyError: - phi_a = self.phia_distr.rvs(size=size)['phi_a'] + phi_a = self.phia_distr.rvs(size=size)["phi_a"] try: - phi_s = kwargs['phi_s'] + phi_s = kwargs["phi_s"] except KeyError: - phi_s = self.phis_distr.rvs(size=size)['phi_s'] + phi_s = self.phis_distr.rvs(size=size)["phi_s"] # draw chi_eff, chi_a try: - chi_eff = kwargs['chi_eff'] + chi_eff = kwargs["chi_eff"] except KeyError: - chi_eff = self.chieff_distr.rvs(size=size)['chi_eff'] + chi_eff = self.chieff_distr.rvs(size=size)["chi_eff"] try: - chi_a = kwargs['chi_a'] + chi_a = kwargs["chi_a"] except KeyError: - chi_a = self.chia_distr.rvs(size=size)['chi_a'] + chi_a = self.chia_distr.rvs(size=size)["chi_a"] # draw xis try: - xi1 = kwargs['xi1'] + xi1 = kwargs["xi1"] except KeyError: - xi1 = self.xi1_distr.rvs(size=size)['xi1'] + xi1 = self.xi1_distr.rvs(size=size)["xi1"] try: - xi2 = kwargs['xi2'] + xi2 = kwargs["xi2"] except KeyError: - xi2 = self.xi2_distr.rvs(size=size)['xi2'] + xi2 = self.xi2_distr.rvs(size=size)["xi2"] dtype = [(p, float) for p in self.params] arr = numpy.zeros(size, dtype=dtype) - arr['mass1'] = mass1 - arr['mass2'] = mass2 - arr['phi_a'] = phi_a - arr['phi_s'] = phi_s - arr['chi_eff'] = chi_eff - arr['chi_a'] = chi_a - arr['xi1'] = xi1 - arr['xi2'] = xi2 + arr["mass1"] = mass1 + arr["mass2"] = mass2 + arr["phi_a"] = phi_a + arr["phi_s"] = phi_s + arr["chi_eff"] = chi_eff + arr["chi_a"] = chi_a + arr["xi1"] = xi1 + arr["xi2"] = xi2 return arr def apply_boundary_conditions(self, **kwargs): return kwargs - def rvs(self, size=1, **kwargs): - """Returns random values for all of the parameters. - """ + """Returns random values for all of the parameters.""" size = int(size) dtype = [(p, float) for p in self.params] arr = numpy.zeros(size, dtype=dtype) @@ -244,15 +281,15 @@ def rvs(self, size=1, **kwargs): draws = self._draw(size=remaining, **kwargs) mask = self._constraints(draws) addpts = mask.sum() - arr[keepidx:keepidx+addpts] = draws[mask] + arr[keepidx : keepidx + addpts] = draws[mask] keepidx += addpts remaining = size - keepidx return arr - @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. The parameters + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -272,23 +309,30 @@ def from_config(cls, cp, section, variable_args): ------- IndependentChiPChiEff A distribution instance. + """ tag = variable_args variable_args = variable_args.split(VARARGS_DELIM) if not set(variable_args) == set(cls._params): - raise ValueError("Not all parameters used by this distribution " - "included in tag portion of section name") + raise ValueError( + "Not all parameters used by this distribution " + "included in tag portion of section name" + ) # get the bounds for the setable parameters - mass1 = get_param_bounds_from_config(cp, section, tag, 'mass1') - mass2 = get_param_bounds_from_config(cp, section, tag, 'mass2') - chi_eff = get_param_bounds_from_config(cp, section, tag, 'chi_eff') - chi_a = get_param_bounds_from_config(cp, section, tag, 'chi_a') - xi_bounds = get_param_bounds_from_config(cp, section, tag, 'xi_bounds') - if cp.has_option('-'.join([section, tag]), 'nsamples'): - nsamples = int(cp.get('-'.join([section, tag]), 'nsamples')) + mass1 = get_param_bounds_from_config(cp, section, tag, "mass1") + mass2 = get_param_bounds_from_config(cp, section, tag, "mass2") + chi_eff = get_param_bounds_from_config(cp, section, tag, "chi_eff") + chi_a = get_param_bounds_from_config(cp, section, tag, "chi_a") + xi_bounds = get_param_bounds_from_config(cp, section, tag, "xi_bounds") + if cp.has_option("-".join([section, tag]), "nsamples"): + nsamples = int(cp.get("-".join([section, tag]), "nsamples")) else: nsamples = None - return cls(mass1=mass1, mass2=mass2, chi_eff=chi_eff, chi_a=chi_a, - xi_bounds=xi_bounds, nsamples=nsamples) - - + return cls( + mass1=mass1, + mass2=mass2, + chi_eff=chi_eff, + chi_a=chi_a, + xi_bounds=xi_bounds, + nsamples=nsamples, + ) diff --git a/pycbc/distributions/uniform.py b/pycbc/distributions/uniform.py index 1e3bdde1c91..5f5efcd5bd4 100644 --- a/pycbc/distributions/uniform.py +++ b/pycbc/distributions/uniform.py @@ -15,12 +15,14 @@ """ This modules provides classes for evaluating uniform distributions. """ + import logging + import numpy from pycbc.distributions import bounded -logger = logging.getLogger('pycbc.distributions.uniform') +logger = logging.getLogger("pycbc.distributions.uniform") class Uniform(bounded.BoundedDist): @@ -78,15 +80,19 @@ class Uniform(bounded.BoundedDist): >>> dist.pdf(phi=60.) 0.025 + """ - name = 'uniform' + + name = "uniform" + def __init__(self, **params): - super(Uniform, self).__init__(**params) + super().__init__(**params) # compute the norm and save # temporarily suppress numpy divide by 0 warning with numpy.errstate(divide="ignore"): - self._lognorm = -sum([numpy.log(abs(bnd[1]-bnd[0])) - for bnd in self._bounds.values()]) + self._lognorm = -sum( + [numpy.log(abs(bnd[1] - bnd[0])) for bnd in self._bounds.values()] + ) self._norm = numpy.exp(self._lognorm) @property @@ -100,35 +106,35 @@ def lognorm(self): return self._lognorm def _cdfinv_param(self, param, value): - """Return the inverse cdf to map the unit interval to parameter bounds. - """ + """Return the inverse cdf to map the unit interval to parameter bounds.""" lower_bound = self._bounds[param][0] upper_bound = self._bounds[param][1] return (upper_bound - lower_bound) * value + lower_bound def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs in self: return self._norm - else: - return 0. + return 0.0 def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs in self: return self._lognorm - else: - return -numpy.inf + return -numpy.inf @classmethod def from_config(cls, cp, section, variable_args): - """Returns a distribution based on a configuration file. The parameters + """ + Returns a distribution based on a configuration file. The parameters for the distribution are retrieved from the section titled "[`section`-`variable_args`]" in the config file. @@ -148,9 +154,11 @@ def from_config(cls, cp, section, variable_args): ------- Uniform A distribution instance from the pycbc.inference.prior module. + """ - return super(Uniform, cls).from_config(cp, section, variable_args, - bounds_required=True) + return super().from_config( + cp, section, variable_args, bounds_required=True + ) -__all__ = ['Uniform'] +__all__ = ["Uniform"] diff --git a/pycbc/distributions/uniform_log.py b/pycbc/distributions/uniform_log.py index 94bfaebc93b..7652ffc824f 100644 --- a/pycbc/distributions/uniform_log.py +++ b/pycbc/distributions/uniform_log.py @@ -12,19 +12,23 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This modules provides classes for evaluating distributions whose logarithm +""" +This modules provides classes for evaluating distributions whose logarithm are uniform. """ + import logging + import numpy from pycbc.distributions import uniform -logger = logging.getLogger('pycbc.distributions.uniform_log') +logger = logging.getLogger("pycbc.distributions.uniform_log") class UniformLog10(uniform.Uniform): - r""" A uniform distribution on the log base 10 of the given parameters. + r""" + A uniform distribution on the log base 10 of the given parameters. The parameters are independent of each other. Instances of this class can be called like a function. By default, logpdf will be called. @@ -34,41 +38,46 @@ class UniformLog10(uniform.Uniform): The keyword arguments should provide the names of parameters and their corresponding bounds, as either tuples or a `boundaries.Bounds` instance. + """ + name = "uniform_log10" def __init__(self, **params): - super(UniformLog10, self).__init__(**params) - self._norm = numpy.prod([numpy.log10(bnd[1]) - numpy.log10(bnd[0]) - for bnd in self._bounds.values()]) + super().__init__(**params) + self._norm = numpy.prod( + [numpy.log10(bnd[1]) - numpy.log10(bnd[0]) for bnd in self._bounds.values()] + ) self._lognorm = numpy.log(self._norm) def _cdfinv_param(self, param, value): - """Return the cdfinv for a single given parameter """ + """Return the cdfinv for a single given parameter""" lower_bound = numpy.log10(self._bounds[param][0]) upper_bound = numpy.log10(self._bounds[param][1]) - return 10. ** ((upper_bound - lower_bound) * value + lower_bound) + return 10.0 ** ((upper_bound - lower_bound) * value + lower_bound) def _pdf(self, **kwargs): - """Returns the pdf at the given values. The keyword arguments must + """ + Returns the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs in self: - vals = numpy.array([numpy.log(10) * self._norm * kwargs[param] - for param in kwargs.keys()]) + vals = numpy.array( + [numpy.log(10) * self._norm * kwargs[param] for param in kwargs] + ) return 1.0 / numpy.prod(vals) - else: - return 0. + return 0.0 def _logpdf(self, **kwargs): - """Returns the log of the pdf at the given values. The keyword + """ + Returns the log of the pdf at the given values. The keyword arguments must contain all of parameters in self's params. Unrecognized arguments are ignored. """ if kwargs in self: return numpy.log(self._pdf(**kwargs)) - else: - return -numpy.inf + return -numpy.inf + __all__ = ["UniformLog10"] diff --git a/pycbc/distributions/utils.py b/pycbc/distributions/utils.py index bf63a0422e5..c8ed007f411 100644 --- a/pycbc/distributions/utils.py +++ b/pycbc/distributions/utils.py @@ -25,18 +25,20 @@ This module provides functions for drawing samples from a standalone .ini file in a Python script, rather than in the command line. """ + import logging + import numpy as np +from pycbc import distributions, transforms from pycbc.types.config import InterpolatingConfigParser -from pycbc import transforms -from pycbc import distributions -logger = logging.getLogger('pycbc.distributions.utils') +logger = logging.getLogger("pycbc.distributions.utils") -def prior_from_config(cp, prior_section='prior'): - """Loads a prior distribution from the given config file. +def prior_from_config(cp, prior_section="prior"): + """ + Loads a prior distribution from the given config file. Parameters ---------- @@ -50,35 +52,42 @@ def prior_from_config(cp, prior_section='prior'): ------- distributions.JointDistribution The prior distribution. - """ + """ # Read variable and static parameters from the config file variable_params, static_params = distributions.read_params_from_config( - cp, prior_section=prior_section, vargs_section='variable_params', - sargs_section='static_params') + cp, + prior_section=prior_section, + vargs_section="variable_params", + sargs_section="static_params", + ) # Read waveform_transforms to apply to priors from the config file - if any(cp.get_subsections('waveform_transforms')): + if any(cp.get_subsections("waveform_transforms")): waveform_transforms = transforms.read_transforms_from_config( - cp, 'waveform_transforms') + cp, "waveform_transforms" + ) else: waveform_transforms = None # Read constraints to apply to priors from the config file constraints = distributions.read_constraints_from_config( - cp, transforms=waveform_transforms, static_args=static_params) + cp, transforms=waveform_transforms, static_args=static_params + ) # Get PyCBC distribution instances for each variable parameter in the # config file dists = distributions.read_distributions_from_config(cp, prior_section) # construct class that will return draws from the prior - return distributions.JointDistribution(variable_params, *dists, - **{"constraints": constraints}) + return distributions.JointDistribution( + variable_params, *dists, constraints=constraints + ) def draw_samples_from_config(path, num=1, seed=150914): - r""" Generate sampling points from a standalone .ini file. + r""" + Generate sampling points from a standalone .ini file. Parameters ---------- @@ -90,7 +99,7 @@ def draw_samples_from_config(path, num=1, seed=150914): The random seed for sampling. Returns - -------- + ------- samples : pycbc.io.record.FieldArray The parameter values and names of sample(s). @@ -112,14 +121,14 @@ def draw_samples_from_config(path, num=1, seed=150914): >>> print(sample) >>> # Print a certain parameter, for example 'mass1'. >>> print(sample[0]['mass1']) - """ + """ np.random.seed(seed) # Initialise InterpolatingConfigParser class. config_parser = InterpolatingConfigParser() # Read the file - file = open(path, 'r') + file = open(path) config_parser.read_file(file) file.close() @@ -129,9 +138,10 @@ def draw_samples_from_config(path, num=1, seed=150914): samples = prior_dists.rvs(size=int(num)) # Apply parameter transformation. - if any(config_parser.get_subsections('waveform_transforms')): + if any(config_parser.get_subsections("waveform_transforms")): waveform_transforms = transforms.read_transforms_from_config( - config_parser, 'waveform_transforms') + config_parser, "waveform_transforms" + ) samples = transforms.apply_transforms(samples, waveform_transforms) return samples diff --git a/pycbc/dq.py b/pycbc/dq.py index cf0b530ff29..c43e31393c1 100644 --- a/pycbc/dq.py +++ b/pycbc/dq.py @@ -21,21 +21,26 @@ # # ============================================================================= # -""" Utilities to query archival instrument status information of +""" +Utilities to query archival instrument status information of gravitational-wave detectors from public sources and/or dqsegdb. """ -import logging import json +import logging + import numpy -from igwn_segments import segmentlist, segment +from igwn_segments import segment, segmentlist + from pycbc.frame.gwosc import get_run from pycbc.io import get_file -logger = logging.getLogger('pycbc.dq') +logger = logging.getLogger("pycbc.dq") + def parse_veto_definer(veto_def_filename, ifos): - """ Parse a veto definer file from the filename and return a dictionary + """ + Parse a veto definer file from the filename and return a dictionary indexed by ifo and veto definer category level. Parameters @@ -47,33 +52,35 @@ def parse_veto_definer(veto_def_filename, ifos): definer file Returns - -------- + ------- parsed_definition: dict Returns a dictionary first indexed by ifo, then category level, and finally a list of veto definitions. + """ - from igwn_ligolw import ligolw, utils as ligolw_utils + from igwn_ligolw import ligolw + from igwn_ligolw import utils as ligolw_utils + from pycbc.io.ligolw import LIGOLWContentHandler as h data = {} for ifo_name in ifos: data[ifo_name] = {} - data[ifo_name]['CAT_H'] = [] + data[ifo_name]["CAT_H"] = [] for cat_num in range(1, 5): - data[ifo_name]['CAT_{}'.format(cat_num)] = [] + data[ifo_name][f"CAT_{cat_num}"] = [] - indoc = ligolw_utils.load_filename(veto_def_filename, False, - contenthandler=h) - veto_table = ligolw.Table.get_table(indoc, 'veto_definer') + indoc = ligolw_utils.load_filename(veto_def_filename, False, contenthandler=h) + veto_table = ligolw.Table.get_table(indoc, "veto_definer") - ifo = veto_table.getColumnByName('ifo') - name = veto_table.getColumnByName('name') - version = numpy.array(veto_table.getColumnByName('version')) - category = numpy.array(veto_table.getColumnByName('category')) - start = numpy.array(veto_table.getColumnByName('start_time')) - end = numpy.array(veto_table.getColumnByName('end_time')) - start_pad = numpy.array(veto_table.getColumnByName('start_pad')) - end_pad = numpy.array(veto_table.getColumnByName('end_pad')) + ifo = veto_table.getColumnByName("ifo") + name = veto_table.getColumnByName("name") + version = numpy.array(veto_table.getColumnByName("version")) + category = numpy.array(veto_table.getColumnByName("category")) + start = numpy.array(veto_table.getColumnByName("start_time")) + end = numpy.array(veto_table.getColumnByName("end_time")) + start_pad = numpy.array(veto_table.getColumnByName("start_pad")) + end_pad = numpy.array(veto_table.getColumnByName("end_pad")) for i in range(len(veto_table)): if ifo[i] not in data: @@ -84,50 +91,62 @@ def parse_veto_definer(veto_def_filename, ifos): # often used any more). So we remap 3 to H and anything above 3 to # N-1. 2 and 1 correspond to 2 and 1 (YAY!) if category[i] > 3: - curr_cat = "CAT_{}".format(category[i]-1) + curr_cat = f"CAT_{category[i] - 1}" elif category[i] == 3: curr_cat = "CAT_H" else: - curr_cat = "CAT_{}".format(category[i]) - - veto_info = {'name': name[i], - 'version': version[i], - 'full_name': name[i]+':'+str(version[i]), - 'start': start[i], - 'end': end[i], - 'start_pad': start_pad[i], - 'end_pad': end_pad[i], - } + curr_cat = f"CAT_{category[i]}" + + veto_info = { + "name": name[i], + "version": version[i], + "full_name": name[i] + ":" + str(version[i]), + "start": start[i], + "end": end[i], + "start_pad": start_pad[i], + "end_pad": end_pad[i], + } data[ifo[i]][curr_cat].append(veto_info) return data -GWOSC_URL = 'https://www.gwosc.org/timeline/segments/json/{}/{}_{}/{}/{}/' +GWOSC_URL = "https://www.gwosc.org/timeline/segments/json/{}/{}_{}/{}/{}/" def query_dqsegdb2(detector, flag_name, start_time, end_time, server): - """Utility function for better error reporting when calling dqsegdb2. - """ + """Utility function for better error reporting when calling dqsegdb2.""" from dqsegdb2.query import query_segments - complete_flag = detector + ':' + flag_name + complete_flag = detector + ":" + flag_name try: - query_res = query_segments(complete_flag, - int(start_time), - int(end_time), - host=server) - return query_res['active'] + query_res = query_segments( + complete_flag, int(start_time), int(end_time), host=server + ) + return query_res["active"] except Exception as e: - logger.error('Could not query segment database, check name ' - '(%s), times (%d-%d) and server (%s)', - complete_flag, int(start_time), int(end_time), - server) + logger.error( + "Could not query segment database, check name " + "(%s), times (%d-%d) and server (%s)", + complete_flag, + int(start_time), + int(end_time), + server, + ) raise e -def query_flag(ifo, segment_name, start_time, end_time, - source='any', server="https://segments.ligo.org", - veto_definer=None, cache=False): - """Return the times where the flag is active + +def query_flag( + ifo, + segment_name, + start_time, + end_time, + source="any", + server="https://segments.ligo.org", + veto_definer=None, + cache=False, +): + """ + Return the times where the flag is active Parameters ---------- @@ -151,47 +170,61 @@ def query_flag(ifo, segment_name, start_time, end_time, If true cache the query. Default is not to cache Returns - --------- + ------- segments: igwn_segments.segmentlist List of segments + """ flag_segments = segmentlist([]) - if source in ['GWOSC', 'any']: + if source in ["GWOSC", "any"]: # Special cases as the GWOSC convention is backwards from normal # LIGO / Virgo operation!!!! - if (('_HW_INJ' in segment_name and 'NO' not in segment_name) or - 'VETO' in segment_name): - data = query_flag(ifo, 'DATA', start_time, end_time, cache=cache) + if ( + "_HW_INJ" in segment_name and "NO" not in segment_name + ) or "VETO" in segment_name: + data = query_flag(ifo, "DATA", start_time, end_time, cache=cache) - if '_HW_INJ' in segment_name: - name = 'NO_' + segment_name + if "_HW_INJ" in segment_name: + name = "NO_" + segment_name else: - name = segment_name.replace('_VETO', '') + name = segment_name.replace("_VETO", "") negate = query_flag(ifo, name, start_time, end_time, cache=cache) return (data - negate).coalesce() duration = end_time - start_time try: - url = GWOSC_URL.format(get_run(start_time + duration/2, ifo), - ifo, segment_name, - int(start_time), int(duration)) + url = GWOSC_URL.format( + get_run(start_time + duration / 2, ifo), + ifo, + segment_name, + int(start_time), + int(duration), + ) fname = get_file(url, cache=cache, timeout=10) - data = json.load(open(fname, 'r')) - if 'segments' in data: - flag_segments = data['segments'] - - except Exception as e: - if source != 'any': - raise ValueError("Unable to find {} segments in GWOSC, check " - "flag name or times".format(segment_name)) - - return query_flag(ifo, segment_name, start_time, end_time, - source='dqsegdb', server=server, - veto_definer=veto_definer) - - elif source == 'dqsegdb': + data = json.load(open(fname)) + if "segments" in data: + flag_segments = data["segments"] + + except Exception: + if source != "any": + raise ValueError( + f"Unable to find {segment_name} segments in GWOSC, check " + "flag name or times" + ) + + return query_flag( + ifo, + segment_name, + start_time, + end_time, + source="dqsegdb", + server=server, + veto_definer=veto_definer, + ) + + elif source == "dqsegdb": # The veto definer will allow the use of MACRO names # These directly correspond to the name in the veto definer file if veto_definer is not None: @@ -202,50 +235,59 @@ def query_flag(ifo, segment_name, start_time, end_time, if veto_definer is not None and segment_name in veto_def[ifo]: for flag in veto_def[ifo][segment_name]: partial = segmentlist([]) - segs = query_dqsegdb2(ifo, flag['full_name'], - start_time, end_time, server) + segs = query_dqsegdb2( + ifo, flag["full_name"], start_time, end_time, server + ) # Apply padding to each segment for rseg in segs: - seg_start = rseg[0] + flag['start_pad'] - seg_end = rseg[1] + flag['end_pad'] + seg_start = rseg[0] + flag["start_pad"] + seg_end = rseg[1] + flag["end_pad"] partial.append(segment(seg_start, seg_end)) # Limit to the veto definer stated valid region of this flag - flag_start = flag['start'] - flag_end = flag['end'] + flag_start = flag["start"] + flag_end = flag["end"] # Corner case: if the flag end time is 0 it means 'no limit' # so use the query end time if flag_end == 0: flag_end = int(end_time) send = segmentlist([segment(flag_start, flag_end)]) - flag_segments += (partial.coalesce() & send) + flag_segments += partial.coalesce() & send else: # Standard case just query directly - segs = query_dqsegdb2(ifo, segment_name, start_time, end_time, - server) + segs = query_dqsegdb2(ifo, segment_name, start_time, end_time, server) for rseg in segs: flag_segments.append(segment(rseg[0], rseg[1])) # dqsegdb output is not guaranteed to lie entirely within start # and end times, hence restrict to this range - flag_segments = flag_segments.coalesce() & \ - segmentlist([segment(int(start_time), int(end_time))]) + flag_segments = flag_segments.coalesce() & segmentlist( + [segment(int(start_time), int(end_time))] + ) else: - raise ValueError("Source must be `dqsegdb`, `GWOSC` or `any`." - " Got {}".format(source)) + raise ValueError( + f"Source must be `dqsegdb`, `GWOSC` or `any`. Got {source}" + ) return segmentlist(flag_segments).coalesce() -def query_cumulative_flags(ifo, segment_names, start_time, end_time, - source='any', server="https://segments.ligo.org", - veto_definer=None, - bounds=None, - padding=None, - override_ifos=None, - cache=False): - """Return the times where any flag is active +def query_cumulative_flags( + ifo, + segment_names, + start_time, + end_time, + source="any", + server="https://segments.ligo.org", + veto_definer=None, + bounds=None, + padding=None, + override_ifos=None, + cache=False, +): + """ + Return the times where any flag is active Parameters ---------- @@ -277,9 +319,10 @@ def query_cumulative_flags(ifo, segment_names, start_time, end_time, basis. Returns - --------- + ------- segments: igwn_segments.segmentlist List of segments + """ total_segs = segmentlist([]) for flag_name in segment_names: @@ -287,10 +330,16 @@ def query_cumulative_flags(ifo, segment_names, start_time, end_time, if override_ifos is not None and flag_name in override_ifos: ifo_name = override_ifos[flag_name] - segs = query_flag(ifo_name, flag_name, start_time, end_time, - source=source, server=server, - veto_definer=veto_definer, - cache=cache) + segs = query_flag( + ifo_name, + flag_name, + start_time, + end_time, + source=source, + server=server, + veto_definer=veto_definer, + cache=cache, + ) if padding and flag_name in padding: s, e = padding[flag_name] @@ -309,7 +358,8 @@ def query_cumulative_flags(ifo, segment_names, start_time, end_time, def parse_flag_str(flag_str): - """ Parse a dq flag query string + """ + Parse a dq flag query string Parameters ---------- @@ -330,8 +380,9 @@ def parse_flag_str(flag_str): The boundary of a given flag padding: dict Any padding that should be applied to the segments for a given flag + """ - flags = flag_str.replace(' ', '').strip().split(',') + flags = flag_str.replace(" ", "").strip().split(",") signs = {} ifos = {} @@ -341,35 +392,35 @@ def parse_flag_str(flag_str): for flag in flags: # Check if the flag should add or subtract time - if not (flag[0] == '+' or flag[0] == '-'): + if not (flag[0] == "+" or flag[0] == "-"): err_msg = "DQ flags must begin with a '+' or a '-' character. " - err_msg += "You provided {}. ".format(flag) + err_msg += f"You provided {flag}. " err_msg += "See http://pycbc.org/pycbc/latest/html/workflow/segments.html" err_msg += " for more information." raise ValueError(err_msg) - sign = flag[0] == '+' + sign = flag[0] == "+" flag = flag[1:] ifo = pad = bound = None # Check for non-default IFO - if len(flag.split(':')[0]) == 2: - ifo = flag.split(':')[0] + if len(flag.split(":")[0]) == 2: + ifo = flag.split(":")[0] flag = flag[3:] # Check for padding options - if '<' in flag: - popt = flag.split('<')[1].split('>')[0] - spad, epad = popt.split(':') + if "<" in flag: + popt = flag.split("<")[1].split(">")[0] + spad, epad = popt.split(":") pad = (float(spad), float(epad)) - flag = flag.replace(popt, '').replace('<>', '') + flag = flag.replace(popt, "").replace("<>", "") # Check if there are bounds on the flag - if '[' in flag: - bopt = flag.split('[')[1].split(']')[0] - start, end = bopt.split(':') + if "[" in flag: + bopt = flag.split("[")[1].split("]")[0] + start, end = bopt.split(":") bound = (int(start), int(end)) - flag = flag.replace(bopt, '').replace('[]', '') + flag = flag.replace(bopt, "").replace("[]", "") if ifo: ifos[flag] = ifo @@ -383,10 +434,18 @@ def parse_flag_str(flag_str): return bflags, signs, ifos, bounds, padding -def query_str(ifo, flag_str, start_time, end_time, source='any', - server="https://segments.ligo.org", veto_definer=None, - cache=False): - """ Query for flags based on a special str syntax +def query_str( + ifo, + flag_str, + start_time, + end_time, + source="any", + server="https://segments.ligo.org", + veto_definer=None, + cache=False, +): + """ + Query for flags based on a special str syntax Parameters ---------- @@ -418,30 +477,41 @@ def query_str(ifo, flag_str, start_time, end_time, source='any', ------- segs: segmentlist A list of segments corresponding to the flag query string + """ flags, sign, ifos, bounds, padding = parse_flag_str(flag_str) up = [f for f in flags if sign[f]] down = [f for f in flags if not sign[f]] if len(up) + len(down) != len(flags): - raise ValueError('Not all flags could be parsed, check +/- prefix') - segs = query_cumulative_flags(ifo, up, start_time, end_time, - source=source, - server=server, - veto_definer=veto_definer, - bounds=bounds, - padding=padding, - override_ifos=ifos, - cache=cache) - - mseg = query_cumulative_flags(ifo, down, start_time, end_time, - source=source, - server=server, - veto_definer=veto_definer, - bounds=bounds, - padding=padding, - override_ifos=ifos, - cache=cache) + raise ValueError("Not all flags could be parsed, check +/- prefix") + segs = query_cumulative_flags( + ifo, + up, + start_time, + end_time, + source=source, + server=server, + veto_definer=veto_definer, + bounds=bounds, + padding=padding, + override_ifos=ifos, + cache=cache, + ) + + mseg = query_cumulative_flags( + ifo, + down, + start_time, + end_time, + source=source, + server=server, + veto_definer=veto_definer, + bounds=bounds, + padding=padding, + override_ifos=ifos, + cache=cache, + ) segs = (segs - mseg).coalesce() return segs diff --git a/pycbc/events/__init__.py b/pycbc/events/__init__.py index ba4f18b6461..4dfde47750e 100644 --- a/pycbc/events/__init__.py +++ b/pycbc/events/__init__.py @@ -2,7 +2,6 @@ This packages contains modules for clustering events """ +from .coinc import * from .eventmgr import * from .veto import * -from .coinc import * - diff --git a/pycbc/events/coherent.py b/pycbc/events/coherent.py index fd549915d53..01a91f7029e 100644 --- a/pycbc/events/coherent.py +++ b/pycbc/events/coherent.py @@ -21,19 +21,23 @@ # # ============================================================================= # -""" This module contains functions for calculating and manipulating coherent +""" +This module contains functions for calculating and manipulating coherent triggers. """ + import logging + import numpy as np from .eventmgr_cython import get_coinc_indexes_cython_twodet_twocoinc -logger = logging.getLogger('pycbc.events.coherent') +logger = logging.getLogger("pycbc.events.coherent") def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos, wraparound_dict): - """Return the indexes corresponding to coincident triggers. If only one + """ + Return the indexes corresponding to coincident triggers. If only one detector is available in the network, the list of its unique indexes is simply returned. @@ -56,6 +60,7 @@ def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos, wraparound_dict): coinc_idx: list List of indexes for triggers in geocent time that appear in multiple detectors + """ if min_nifos == 2 and len(idx_dict) == 2: ifos = list(idx_dict.keys()) @@ -75,7 +80,7 @@ def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos, wraparound_dict): time_delay_idx[ifos[1]], wraparound_dict[ifos[0]], wraparound_dict[ifos[1]], - outarr + outarr, ) return outarr[:num_idxs] coinc_list = np.array([], dtype=int) @@ -88,7 +93,10 @@ def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos, wraparound_dict): # these represent triggers appearing in multiple detectors. if len(idx_dict[ifo]) != 0: coinc_list = np.hstack( - [coinc_list, (idx_dict[ifo] - time_delay_idx[ifo]) % wraparound_dict[ifo]] + [ + coinc_list, + (idx_dict[ifo] - time_delay_idx[ifo]) % wraparound_dict[ifo], + ] ) # Search through coinc_idx for repeated indexes. These must have been loud # in at least min_nifos detectors if the analysis uses more than 1 @@ -101,7 +109,8 @@ def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos, wraparound_dict): def get_coinc_triggers(snrs, idx, t_delay_idx): - """Returns a dictionary, indexed by IFO, that collects the individual + """ + Returns a dictionary, indexed by IFO, that collects the individual IFO SNRs of coincident triggers by using the indices of such triggers within the complete SNR timeseries of each IFO. @@ -119,17 +128,17 @@ def get_coinc_triggers(snrs, idx, t_delay_idx): ------- coincs: dict Dictionary of coincident trigger SNRs in each detector + """ # loops through snrs # %len(snrs[ifo]) was included as part of a wrap-around solution - coincs = { - ifo: snrs[ifo][(idx + t_delay_idx[ifo]) % len(snrs[ifo])] - for ifo in snrs} + coincs = {ifo: snrs[ifo][(idx + t_delay_idx[ifo]) % len(snrs[ifo])] for ifo in snrs} return coincs def coincident_snr(snr_dict, index, threshold, time_delay_idx): - """Calculate the coincident SNR for all coincident triggers above + """ + Calculate the coincident SNR for all coincident triggers above threshold Parameters @@ -155,13 +164,12 @@ def coincident_snr(snr_dict, index, threshold, time_delay_idx): coinc_triggers: dict Dictionary of individual detector SNRs for triggers that survive cuts + """ # Restrict the snr timeseries to just the interesting points coinc_triggers = get_coinc_triggers(snr_dict, index, time_delay_idx) # Calculate the coincident snr - snr_array = np.array( - [coinc_triggers[ifo] for ifo in coinc_triggers.keys()] - ) + snr_array = np.array([coinc_triggers[ifo] for ifo in coinc_triggers.keys()]) rho_coinc = abs(np.sqrt(np.sum(snr_array * snr_array.conj(), axis=0))) # Apply threshold thresh_indexes = rho_coinc > threshold @@ -172,7 +180,8 @@ def coincident_snr(snr_dict, index, threshold, time_delay_idx): def get_projection_matrix(f_plus, f_cross, sigma, projection="standard"): - """Calculate the matrix that projects the signal onto the network. + """ + Calculate the matrix that projects the signal onto the network. Definitions can be found in Fairhurst (2018) [arXiv:1712.04724]. For the standard projection see Eq. 8, and for left/right circular projections see Eq. 21, with further discussion in @@ -199,6 +208,7 @@ def get_projection_matrix(f_plus, f_cross, sigma, projection="standard"): ------- projection_matrix: np.ndarray The matrix that projects the signal onto the detector network + """ # Calculate the weighted antenna responses keys = sorted(sigma.keys()) @@ -227,16 +237,16 @@ def get_projection_matrix(f_plus, f_cross, sigma, projection="standard"): ) / (np.dot(w_p, w_p) + np.dot(w_c, w_c)) else: raise ValueError( - f'Unknown projection: {projection}. Allowed values are: ' - '"standard", "left", and "right"') + f"Unknown projection: {projection}. Allowed values are: " + '"standard", "left", and "right"' + ) return projection_matrix -def coherent_snr( - snr_triggers, index, threshold, projection_matrix, coinc_snr=None -): - """Calculate the coherent SNR for a given set of triggers. See +def coherent_snr(snr_triggers, index, threshold, projection_matrix, coinc_snr=None): + """ + Calculate the coherent SNR for a given set of triggers. See Eq. 2.26 of Harry & Fairhurst (2011) [arXiv:1012.4939]. @@ -264,11 +274,10 @@ def coherent_snr( coinc_snr: list or None (default: None) The coincident SNR values for triggers surviving the coherent cut + """ # Calculate rho_coh - snr_array = np.array( - [snr_triggers[ifo] for ifo in sorted(snr_triggers.keys())] - ) + snr_array = np.array([snr_triggers[ifo] for ifo in sorted(snr_triggers.keys())]) snr_proj = np.inner(snr_array.conj().transpose(), projection_matrix) rho_coh2 = sum(snr_proj.transpose() * snr_array) rho_coh = abs(np.sqrt(rho_coh2)) @@ -278,16 +287,14 @@ def coherent_snr( coinc_snr = [] if coinc_snr is None else coinc_snr if len(coinc_snr) != 0: coinc_snr = coinc_snr[above] - snrv = { - ifo: snr_triggers[ifo][above] - for ifo in snr_triggers.keys() - } + snrv = {ifo: snr_triggers[ifo][above] for ifo in snr_triggers.keys()} rho_coh = rho_coh[above] return rho_coh, index, snrv, coinc_snr def network_chisq(chisq, chisq_dof, snr_dict): - """Calculate the network chi-squared statistic. This is the sum of + """ + Calculate the network chi-squared statistic. This is the sum of SNR-weighted individual detector chi-squared values. See Eq. 5.4 of Dorrington (2019) [http://orca.cardiff.ac.uk/id/eprint/128124]. @@ -305,6 +312,7 @@ def network_chisq(chisq, chisq_dof, snr_dict): ------- net_chisq: list Network chi-squared values + """ ifos = sorted(snr_dict.keys()) chisq_per_dof = dict.fromkeys(ifos) @@ -321,10 +329,17 @@ def network_chisq(chisq, chisq_dof, snr_dict): def null_snr( - rho_coh, rho_coinc, apply_cut=True, null_min=5.25, null_grad=0.2, - null_step=20.0, index=None, snrv=None + rho_coh, + rho_coinc, + apply_cut=True, + null_min=5.25, + null_grad=0.2, + null_step=20.0, + index=None, + snrv=None, ): - """Calculate the null SNR and optionally apply threshold cut where + """ + Calculate the null SNR and optionally apply threshold cut where null SNR > null_min where coherent SNR < null_step and null SNR > (null_grad * rho_coh + null_min) elsewhere. See Eq. 3.1 of Harry & Fairhurst (2011) [arXiv:1012.4939] or @@ -367,24 +382,22 @@ def null_snr( Indexes for surviving triggers snrv: dict Single detector SNRs for surviving triggers + """ index = {} if index is None else index snrv = {} if snrv is None else snrv # Calculate null SNRs - null2 = rho_coinc ** 2 - rho_coh ** 2 + null2 = rho_coinc**2 - rho_coh**2 # Numerical errors may make this negative and break the sqrt, so set # negative values to 0. null2[null2 < 0] = 0 - null = null2 ** 0.5 + null = null2**0.5 if apply_cut: # Make cut on null. - keep = ( - ((null < null_min) & (rho_coh <= null_step)) - | ( - (null < ((rho_coh - null_step) * null_grad + null_min)) - & (rho_coh > null_step) - ) - ) + keep = ((null < null_min) & (rho_coh <= null_step)) | ( + (null < ((rho_coh - null_step) * null_grad + null_min)) + & (rho_coh > null_step) + ) index = index[keep] rho_coh = rho_coh[keep] snrv = {ifo: snrv[ifo][keep] for ifo in snrv} @@ -394,9 +407,10 @@ def null_snr( def reweight_snr_by_null( - network_snr, null, coherent, null_min=5.25, null_grad=0.2, - null_step=20.0): - """Re-weight the detection statistic as a function of the null SNR. + network_snr, null, coherent, null_min=5.25, null_grad=0.2, null_step=20.0 +): + """ + Re-weight the detection statistic as a function of the null SNR. See Eq. 16 of Williamson et al. (2014) [arXiv:1410.6042] and note that the 4.25 appearing there is actually linked to the 5.25 of Eq. 12, hence the -1 carried out in this function. @@ -414,19 +428,16 @@ def reweight_snr_by_null( ------- rw_snr: numpy.ndarray Re-weighted SNR for each trigger + """ - downweight = ( - ((null > null_min - 1) & (coherent <= null_step)) - | ( - (null > (coherent * null_grad + null_min - 1)) - & (coherent > null_step) - ) - ) + downweight = ((null > null_min - 1) & (coherent <= null_step)) | ( + (null > (coherent * null_grad + null_min - 1)) & (coherent > null_step) + ) rw_fac = np.where( coherent > null_step, 1 + null - (null_min - 1) - (coherent - null_step) * null_grad, - 1 + null - (null_min - 1) - ) + 1 + null - (null_min - 1), + ) rw_snr = np.where(downweight, network_snr / rw_fac, network_snr) return rw_snr diff --git a/pycbc/events/coinc.py b/pycbc/events/coinc.py index a36807e1366..37561aa8a82 100644 --- a/pycbc/events/coinc.py +++ b/pycbc/events/coinc.py @@ -21,34 +21,39 @@ # # ============================================================================= # -""" This module contains functions for calculating and manipulating +""" +This module contains functions for calculating and manipulating coincident triggers. """ -import numpy -import logging import copy -import time as timemod +import logging import threading +import time as timemod + +import numpy import pycbc.pnutils -from pycbc.detector import Detector, ppdets from pycbc import conversions as conv +from pycbc.detector import Detector, ppdets from . import stat as pycbcstat -from .eventmgr_cython import coincbuffer_expireelements -from .eventmgr_cython import coincbuffer_numgreater -from .eventmgr_cython import timecoincidence_constructidxs -from .eventmgr_cython import timecoincidence_constructfold -from .eventmgr_cython import timecoincidence_getslideint -from .eventmgr_cython import timecoincidence_findidxlen -from .eventmgr_cython import timecluster_cython +from .eventmgr_cython import ( + coincbuffer_expireelements, + coincbuffer_numgreater, + timecluster_cython, + timecoincidence_constructfold, + timecoincidence_constructidxs, + timecoincidence_findidxlen, + timecoincidence_getslideint, +) -logger = logging.getLogger('pycbc.events.coinc') +logger = logging.getLogger("pycbc.events.coinc") def background_bin_from_string(background_bins, data): - """ Return template ids for each bin as defined by the format string + """ + Return template ids for each bin as defined by the format string Parameters ---------- @@ -63,6 +68,7 @@ def background_bin_from_string(background_bins, data): ------- bins: dict Dictionary of location indices indexed by a bin name + """ used = numpy.array([], dtype=numpy.uint32) bins = {} @@ -71,68 +77,66 @@ def background_bin_from_string(background_bins, data): cached_values = {} for mbin in background_bins: locs = None - name, bin_type_list, boundary_list = tuple(mbin.split(':')) + name, bin_type_list, boundary_list = tuple(mbin.split(":")) - bin_type_list = bin_type_list.split(',') - boundary_list = boundary_list.split(',') + bin_type_list = bin_type_list.split(",") + boundary_list = boundary_list.split(",") for bin_type, boundary in zip(bin_type_list, boundary_list): - if boundary[0:2] == 'lt': - member_func = lambda vals, bd=boundary : vals < float(bd[2:]) - elif boundary[0:2] == 'gt': - member_func = lambda vals, bd=boundary : vals > float(bd[2:]) + if boundary[0:2] == "lt": + member_func = lambda vals, bd=boundary: vals < float(bd[2:]) + elif boundary[0:2] == "gt": + member_func = lambda vals, bd=boundary: vals > float(bd[2:]) else: - raise RuntimeError("Can't parse boundary condition! Must begin " - "with 'lt' or 'gt'") + raise RuntimeError( + "Can't parse boundary condition! Must begin with 'lt' or 'gt'" + ) if bin_type in cached_values: vals = cached_values[bin_type] - elif bin_type == 'component' and boundary[0:2] == 'lt': + elif bin_type == "component" and boundary[0:2] == "lt": # maximum component mass is less than boundary value - vals = numpy.maximum(data['mass1'], data['mass2']) - elif bin_type == 'component' and boundary[0:2] == 'gt': + vals = numpy.maximum(data["mass1"], data["mass2"]) + elif bin_type == "component" and boundary[0:2] == "gt": # minimum component mass is greater than bdary - vals = numpy.minimum(data['mass1'], data['mass2']) - elif bin_type == 'total': - vals = data['mass1'] + data['mass2'] - elif bin_type == 'chirp': + vals = numpy.minimum(data["mass1"], data["mass2"]) + elif bin_type == "total": + vals = data["mass1"] + data["mass2"] + elif bin_type == "chirp": vals = pycbc.pnutils.mass1_mass2_to_mchirp_eta( - data['mass1'], data['mass2'])[0] - elif bin_type == 'ratio': - vals = conv.q_from_mass1_mass2(data['mass1'], data['mass2']) - elif bin_type == 'eta': + data["mass1"], data["mass2"] + )[0] + elif bin_type == "ratio": + vals = conv.q_from_mass1_mass2(data["mass1"], data["mass2"]) + elif bin_type == "eta": vals = pycbc.pnutils.mass1_mass2_to_mchirp_eta( - data['mass1'], - data['mass2'] + data["mass1"], data["mass2"] )[1] - elif bin_type == 'chi_eff': + elif bin_type == "chi_eff": vals = conv.chi_eff( - data['mass1'], - data['mass2'], - data['spin1z'], - data['spin2z'] + data["mass1"], data["mass2"], data["spin1z"], data["spin2z"] ) - elif bin_type.endswith('Peak'): + elif bin_type.endswith("Peak"): vals = pycbc.pnutils.get_freq( - 'f' + bin_type, - data['mass1'], - data['mass2'], - data['spin1z'], - data['spin2z'] + "f" + bin_type, + data["mass1"], + data["mass2"], + data["spin1z"], + data["spin2z"], ) cached_values[bin_type] = vals - elif bin_type.endswith('duration'): + elif bin_type.endswith("duration"): vals = pycbc.pnutils.get_imr_duration( - data['mass1'], - data['mass2'], - data['spin1z'], - data['spin2z'], - data['f_lower'], - approximant=bin_type.replace('duration', '') + data["mass1"], + data["mass2"], + data["spin1z"], + data["spin2z"], + data["f_lower"], + approximant=bin_type.replace("duration", ""), ) cached_values[bin_type] = vals else: - raise ValueError('Invalid bin type %s' % bin_type) + raise ValueError("Invalid bin type %s" % bin_type) sub_locs = member_func(vals) sub_locs = numpy.where(sub_locs)[0] @@ -151,7 +155,8 @@ def background_bin_from_string(background_bins, data): def timeslide_durations(start1, start2, end1, end2, timeslide_offsets): - """ Find the coincident time for each timeslide. + """ + Find the coincident time for each timeslide. Find the coincident time for each timeslide, where the first time vector is slid to the right by the offset in the given timeslide_offsets vector. @@ -170,11 +175,13 @@ def timeslide_durations(start1, start2, end1, end2, timeslide_offsets): Array of offsets (in seconds) for each timeslide Returns - -------- + ------- durations: numpy.ndarray Array of coincident time for each timeslide in the offset array + """ from . import veto + durations = [] seg2 = veto.start_end_to_segments(start2, end2) for offset in timeslide_offsets: @@ -184,7 +191,8 @@ def timeslide_durations(start1, start2, end1, end2, timeslide_offsets): def time_coincidence(t1, t2, window, slide_step=0): - """ Find coincidences by time window + """ + Find coincidences by time window Parameters ---------- @@ -206,14 +214,16 @@ def time_coincidence(t1, t2, window, slide_step=0): Array of indices into the t2 array slide : numpy.ndarray Array of slide ids + """ if slide_step: length1 = len(t1) length2 = len(t2) fold1 = numpy.zeros(length1, dtype=numpy.float64) fold2 = numpy.zeros(length2, dtype=numpy.float64) - timecoincidence_constructfold(fold1, fold2, t1, t2, slide_step, - length1, length2) + timecoincidence_constructfold( + fold1, fold2, t1, t2, slide_step, length1, length2 + ) else: fold1 = t1 fold2 = t2 @@ -225,8 +235,7 @@ def time_coincidence(t1, t2, window, slide_step=0): if slide_step: # FIXME explain this - fold2 = numpy.concatenate([fold2 - slide_step, fold2, - fold2 + slide_step]) + fold2 = numpy.concatenate([fold2 - slide_step, fold2, fold2 + slide_step]) left = fold2.searchsorted(fold1 - window) right = fold2.searchsorted(fold1 + window) @@ -234,8 +243,9 @@ def time_coincidence(t1, t2, window, slide_step=0): lenidx = timecoincidence_findidxlen(left, right, len(left)) idx1 = numpy.zeros(lenidx, dtype=numpy.uint32) idx2 = numpy.zeros(lenidx, dtype=numpy.uint32) - timecoincidence_constructidxs(idx1, idx2, sort1, sort2, left, right, - len(left), len(sort2)) + timecoincidence_constructidxs( + idx1, idx2, sort1, sort2, left, right, len(left), len(sort2) + ) slide = numpy.zeros(lenidx, dtype=numpy.int32) if slide_step: @@ -246,9 +256,9 @@ def time_coincidence(t1, t2, window, slide_step=0): return idx1, idx2, slide -def time_multi_coincidence(times, slide_step=0, slop=.003, - pivot='H1', fixed='L1'): - """ Find multi detector coincidences. +def time_multi_coincidence(times, slide_step=0, slop=0.003, pivot="H1", fixed="L1"): + """ + Find multi detector coincidences. Parameters ---------- @@ -274,16 +284,18 @@ def time_multi_coincidence(times, slide_step=0, slop=.003, recorded slide: array of int Slide ids of coincident triggers in pivot ifo + """ + def win(ifo1, ifo2): d1 = Detector(ifo1) d2 = Detector(ifo2) return d1.light_travel_time_to_detector(d2) + slop # Find coincs between the 'pivot' and 'fixed' detectors as in 2-ifo case - pivot_id, fix_id, slide = time_coincidence(times[pivot], times[fixed], - win(pivot, fixed), - slide_step=slide_step) + pivot_id, fix_id, slide = time_coincidence( + times[pivot], times[fixed], win(pivot, fixed), slide_step=slide_step + ) # Additional detectors do not slide independently of the 'fixed' one # Each trigger in an additional detector must be concident with both @@ -316,7 +328,7 @@ def win(ifo1, ifo2): # tested against fixed and pivot are now present for testing with new # dependent ifos for ifo2 in ids: - logger.info('added ifo %s, testing against %s', ifo1, ifo2) + logger.info("added ifo %s, testing against %s", ifo1, ifo2) w = win(ifo1, ifo2) left = time1.searchsorted(ctimes[ifo2] - w) right = time1.searchsorted(ctimes[ifo2] + w) @@ -333,9 +345,11 @@ def win(ifo1, ifo2): # However there are rare corner cases at starts/ends of inspiral # jobs. For these, arbitrarily keep the first trigger and # discard the second (and any subsequent ones). - logger.warning('Triggers in %s are closer than coincidence ' - 'window, 1 or more coincs will be discarded. ' - 'This is a warning, not an error.' % ifo1) + logger.warning( + "Triggers in %s are closer than coincidence " + "window, 1 or more coincs will be discarded. " + "This is a warning, not an error." % ifo1 + ) # identify indices of times in ifo1 that form coincs with ifo2 dep_ids = left[nz] # slide is array of slide ids attached to pivot ifo @@ -356,7 +370,8 @@ def win(ifo1, ifo2): def cluster_coincs(stat, time1, time2, timeslide_id, slide, window, **kwargs): - """Cluster coincident events for each timeslide separately, across + """ + Cluster coincident events for each timeslide separately, across templates, based on the ranking statistic Parameters @@ -378,9 +393,10 @@ def cluster_coincs(stat, time1, time2, timeslide_id, slide, window, **kwargs): ------- cindex: numpy.ndarray The set of indices corresponding to the surviving coincidences. + """ if len(time1) == 0 or len(time2) == 0: - logger.info('No coinc triggers in one, or both, ifos.') + logger.info("No coinc triggers in one, or both, ifos.") return numpy.array([]) if numpy.isfinite(slide): @@ -395,15 +411,15 @@ def cluster_coincs(stat, time1, time2, timeslide_id, slide, window, **kwargs): span = (time.max() - time.min()) + window * 10 time = time + span * tslide - logger.info('Clustering events over %s s window', window) + logger.info("Clustering events over %s s window", window) cidx = cluster_over_time(stat, time, window, **kwargs) - logger.info('%d triggers remaining', len(cidx)) + logger.info("%d triggers remaining", len(cidx)) return cidx -def cluster_coincs_multiifo(stat, time_coincs, timeslide_id, slide, window, - **kwargs): - """Cluster coincident events for each timeslide separately, across +def cluster_coincs_multiifo(stat, time_coincs, timeslide_id, slide, window, **kwargs): + """ + Cluster coincident events for each timeslide separately, across templates, based on the ranking statistic Parameters @@ -423,14 +439,15 @@ def cluster_coincs_multiifo(stat, time_coincs, timeslide_id, slide, window, ------- cindex: numpy.ndarray The set of indices corresponding to the surviving coincidences + """ time_coinc_zip = list(zip(*time_coincs)) if len(time_coinc_zip) == 0: - logger.info('No coincident triggers.') + logger.info("No coincident triggers.") return numpy.array([]) time_avg_num = [] - #find number of ifos and mean time over participating ifos for each coinc + # find number of ifos and mean time over participating ifos for each coinc for tc in time_coinc_zip: time_avg_num.append(mean_if_greater_than_zero(tc)) @@ -442,23 +459,24 @@ def cluster_coincs_multiifo(stat, time_coincs, timeslide_id, slide, window, # shift all but the pivot ifo by (num_ifos-1) * timeslide_id * slide # this leads to a mean coinc time located around pivot time if numpy.isfinite(slide): - nifos_minusone = (num_ifos - numpy.ones_like(num_ifos)) - time_avg = time_avg + (nifos_minusone * timeslide_id * slide)/num_ifos + nifos_minusone = num_ifos - numpy.ones_like(num_ifos) + time_avg = time_avg + (nifos_minusone * timeslide_id * slide) / num_ifos tslide = timeslide_id.astype(numpy.longdouble) time_avg = time_avg.astype(numpy.longdouble) span = (time_avg.max() - time_avg.min()) + window * 10 time_avg = time_avg + span * tslide - logger.info('Clustering events over %s s window', window) + logger.info("Clustering events over %s s window", window) cidx = cluster_over_time(stat, time_avg, window, **kwargs) - logger.info('%d triggers remaining', len(cidx)) + logger.info("%d triggers remaining", len(cidx)) return cidx def mean_if_greater_than_zero(vals): - """ Calculate mean over numerical values, ignoring values less than zero. + """ + Calculate mean over numerical values, ignoring values less than zero. E.g. used for mean time over coincident triggers when timestamps are set to -1 for ifos not included in the coincidence. @@ -474,15 +492,16 @@ def mean_if_greater_than_zero(vals): greater than zero num_above_zero: int The number of entries in the vector which are above zero + """ vals = numpy.array(vals) above_zero = vals > 0 return vals[above_zero].mean(), above_zero.sum() -def cluster_over_time(stat, time, window, method='python', - argmax=numpy.argmax): - """Cluster generalized transient events over time via maximum stat over a +def cluster_over_time(stat, time, window, method="python", argmax=numpy.argmax): + """ + Cluster generalized transient events over time via maximum stat over a symmetric sliding window Parameters @@ -503,8 +522,8 @@ def cluster_over_time(stat, time, window, method='python', ------- cindex: numpy.ndarray The set of indices corresponding to the surviving coincidences. - """ + """ indices = [] time_sorting = time.argsort() stat = stat[time_sorting] @@ -514,11 +533,11 @@ def cluster_over_time(stat, time, window, method='python', right = time.searchsorted(time + window) indices = numpy.zeros(len(left), dtype=numpy.uint32) - logger.debug('%d triggers before clustering', len(time)) + logger.debug("%d triggers before clustering", len(time)) - if method == 'cython': + if method == "cython": j = timecluster_cython(indices, left, right, stat, len(left)) - elif method == 'python': + elif method == "python": # i is the index we are inspecting, j is the next one to save i = 0 j = 0 @@ -550,19 +569,26 @@ def cluster_over_time(stat, time, window, method='python', elif max_loc < i: i += 1 else: - raise ValueError(f'Do not recognize method {method}') + raise ValueError(f"Do not recognize method {method}") indices = indices[:j] - logger.debug('%d triggers remaining', len(indices)) + logger.debug("%d triggers remaining", len(indices)) return time_sorting[indices] -class MultiRingBuffer(object): +class MultiRingBuffer: """Dynamic size n-dimensional ring buffer that can expire elements.""" - def __init__(self, num_rings, max_time, dtype, min_buffer_size=16, - buffer_increment=8, resize_invalid_fraction=0.4): + def __init__( + self, + num_rings, + max_time, + dtype, + min_buffer_size=16, + buffer_increment=8, + resize_invalid_fraction=0.4, + ): """ Parameters ---------- @@ -588,6 +614,7 @@ def __init__(self, num_rings, max_time, dtype, min_buffer_size=16, options, be careful changing default values, it is possible to get stuck in a mode where the buffers are always being resized. + """ self.max_time = max_time self.buffer = [] @@ -599,8 +626,7 @@ def __init__(self, num_rings, max_time, dtype, min_buffer_size=16, self.resize_invalid_fraction = resize_invalid_fraction for _ in range(num_rings): self.buffer.append(numpy.zeros(self.min_buffer_size, dtype=dtype)) - self.buffer_expire.append(numpy.zeros(self.min_buffer_size, - dtype=int)) + self.buffer_expire.append(numpy.zeros(self.min_buffer_size, dtype=int)) self.valid_ends.append(0) self.valid_starts.append(0) self.time = 0 @@ -627,14 +653,14 @@ def discard_last(self, indices): self.valid_ends[i] -= 1 def advance_time(self): - """Advance the internal time increment by 1, expiring any triggers + """ + Advance the internal time increment by 1, expiring any triggers that are now too old. """ self.time += 1 def add(self, indices, values): - """Add triggers in 'values' to the buffers indicated by the indices - """ + """Add triggers in 'values' to the buffers indicated by the indices""" for i, v in zip(indices, values): # Expand ring buffer size if needed if self.valid_ends[i] == len(self.buffer[i]): @@ -647,15 +673,15 @@ def add(self, indices, values): self.buffer[i], max( len(self.buffer[i]) + self.buffer_increment, - self.min_buffer_size - ) + self.min_buffer_size, + ), ) self.buffer_expire[i] = numpy.resize( self.buffer_expire[i], max( len(self.buffer[i]) + self.buffer_increment, - self.min_buffer_size - ) + self.min_buffer_size, + ), ) curr_pos = self.valid_ends[i] self.buffer[i][curr_pos] = v @@ -666,13 +692,12 @@ def add(self, indices, values): def valid_slice(self, buffer_index): """Return the valid slice for this buffer index""" ret_slice = slice( - self.valid_starts[buffer_index], - self.valid_ends[buffer_index] + self.valid_starts[buffer_index], self.valid_ends[buffer_index] ) return ret_slice def expire_vector(self, buffer_index): - """Return the expiration vector of a given ring buffer """ + """Return the expiration vector of a given ring buffer""" return self.buffer_expire[buffer_index][self.valid_slice(buffer_index)] def update_valid_start(self, buffer_index): @@ -688,7 +713,8 @@ def update_valid_start(self, buffer_index): self.valid_starts[buffer_index] = j def check_expired_triggers(self, buffer_index): - """Check if we should free memory for this buffer index. + """ + Check if we should free memory for this buffer index. Check what fraction of triggers are expired in the specified buffer and if it is more than the allowed fraction (set by @@ -701,8 +727,12 @@ def check_expired_triggers(self, buffer_index): if (buf_len - val_end) + val_start > invalid_limit: # If self.resize_invalid_fraction of stored triggers are expired # or are not set, free up memory - self.buffer_expire[buffer_index] = self.buffer_expire[buffer_index][val_start:val_end].copy() - self.buffer[buffer_index] = self.buffer[buffer_index][val_start:val_end].copy() + self.buffer_expire[buffer_index] = self.buffer_expire[buffer_index][ + val_start:val_end + ].copy() + self.buffer[buffer_index] = self.buffer[buffer_index][ + val_start:val_end + ].copy() self.valid_ends[buffer_index] -= val_start self.valid_starts[buffer_index] = 0 @@ -714,13 +744,13 @@ def data(self, buffer_index): return self.buffer[buffer_index][self.valid_slice(buffer_index)] -class CoincExpireBuffer(object): - """Unordered dynamic sized buffer that handles +class CoincExpireBuffer: + """ + Unordered dynamic sized buffer that handles multiple expiration vectors. """ - def __init__(self, expiration, ifos, - initial_size=2**20, dtype=numpy.float32): + def __init__(self, expiration, ifos, initial_size=2**20, dtype=numpy.float32): """ Parameters ---------- @@ -733,8 +763,8 @@ def __init__(self, expiration, ifos, The initial size of the buffer. dtype: numpy.dtype The dtype of each element of the buffer. - """ + """ self.expiration = expiration self.buffer = numpy.zeros(initial_size, dtype=dtype) self.index = 0 @@ -751,8 +781,7 @@ def __len__(self): @property def nbytes(self): - """Returns the approximate memory usage of self. - """ + """Returns the approximate memory usage of self.""" nbs = [self.timer[ifo].nbytes for ifo in self.ifos] nbs.append(self.buffer.nbytes) return sum(nbs) @@ -766,7 +795,8 @@ def remove(self, num): self.index -= num def add(self, values, times, ifos): - """Add values to the internal buffer + """ + Add values to the internal buffer Parameters ---------- @@ -776,8 +806,8 @@ def add(self, values, times, ifos): The current time to use for each element being added. ifos: list of strs The set of timers to be incremented. - """ + """ for ifo in ifos: self.time[ifo] += 1 @@ -788,10 +818,10 @@ def add(self, values, times, ifos): self.timer[ifo].resize(newlen) self.buffer.resize(newlen, refcheck=False) - self.buffer[self.index:self.index+len(values)] = values + self.buffer[self.index : self.index + len(values)] = values if len(values) > 0: for ifo in self.ifos: - self.timer[ifo][self.index:self.index+len(values)] = times[ifo] + self.timer[ifo][self.index : self.index + len(values)] = times[ifo] self.index += len(values) @@ -805,18 +835,18 @@ def add(self, values, times, ifos): self.time[ifos[0]], self.time[ifos[1]], self.expiration, - self.index + self.index, ) else: # Numpy version for >2 ifo case keep = None for ifo in ifos: - kt = self.timer[ifo][:self.index] >= self.time[ifo] - self.expiration + kt = self.timer[ifo][: self.index] >= self.time[ifo] - self.expiration keep = numpy.logical_and(keep, kt) if keep is not None else kt - self.buffer[:keep.sum()] = self.buffer[:self.index][keep] + self.buffer[: keep.sum()] = self.buffer[: self.index][keep] for ifo in self.ifos: - self.timer[ifo][:keep.sum()] = self.timer[ifo][:self.index][keep] + self.timer[ifo][: keep.sum()] = self.timer[ifo][: self.index][keep] self.index = keep.sum() def num_greater(self, value): @@ -826,20 +856,27 @@ def num_greater(self, value): @property def data(self): """Return the array of elements""" - return self.buffer[:self.index] + return self.buffer[: self.index] -class LiveCoincTimeslideBackgroundEstimator(object): +class LiveCoincTimeslideBackgroundEstimator: """Rolling buffer background estimation.""" - def __init__(self, num_templates, analysis_block, background_statistic, - sngl_ranking, stat_files, ifos, - ifar_limit=100, - timeslide_interval=.035, - coinc_window_pad=.002, - statistic_refresh_rate=None, - return_background=False, - **kwargs): + def __init__( + self, + num_templates, + analysis_block, + background_statistic, + sngl_ranking, + stat_files, + ifos, + ifar_limit=100, + timeslide_interval=0.035, + coinc_window_pad=0.002, + statistic_refresh_rate=None, + return_background=False, + **kwargs, + ): """ Parameters ---------- @@ -874,17 +911,13 @@ class (in seconds), default not do do this kwargs: dict Additional options for the statistic to use. See stat.py for more details on statistic options. + """ self.num_templates = num_templates self.analysis_block = analysis_block stat_class = pycbcstat.get_statistic(background_statistic) - self.stat_calculator = stat_class( - sngl_ranking, - stat_files, - ifos=ifos, - **kwargs - ) + self.stat_calculator = stat_class(sngl_ranking, stat_files, ifos=ifos, **kwargs) self.time_stat_refreshed = timemod.time() self.stat_calculator_lock = threading.Lock() @@ -898,13 +931,17 @@ class (in seconds), default not do do this if len(self.ifos) != 2: raise ValueError("Only a two ifo analysis is supported at this time") - self.lookback_time = (ifar_limit / conv.sec_to_year(1.) * timeslide_interval) ** 0.5 + self.lookback_time = ( + ifar_limit / conv.sec_to_year(1.0) * timeslide_interval + ) ** 0.5 self.buffer_size = int(numpy.ceil(self.lookback_time / analysis_block)) self.dets = {ifo: Detector(ifo) for ifo in ifos} - self.time_window = self.dets[ifos[0]].light_travel_time_to_detector( - self.dets[ifos[1]]) + coinc_window_pad + self.time_window = ( + self.dets[ifos[0]].light_travel_time_to_detector(self.dets[ifos[1]]) + + coinc_window_pad + ) self.coincs = CoincExpireBuffer(self.buffer_size, self.ifos) self.singles = {} @@ -915,7 +952,8 @@ class (in seconds), default not do do this @classmethod def pick_best_coinc(cls, coinc_results): - """Choose the best two-ifo coinc by ifar first, then statistic if needed. + """ + Choose the best two-ifo coinc by ifar first, then statistic if needed. This function picks which of the available double-ifo coincs to use. It chooses the best (highest) ifar. The ranking statistic is used as @@ -933,6 +971,7 @@ def pick_best_coinc(cls, coinc_results): best: coinc results dict If there is a coinc, this will contain the 'best' one. Otherwise it will return the provided dict. + """ mstat = 0 mifar = 0 @@ -944,13 +983,13 @@ def pick_best_coinc(cls, coinc_results): for result in coinc_results: # Check that a coinc was possible. See the 'add_singles' method # to see where this flag was added into the results dict - if 'coinc_possible' in result: + if "coinc_possible" in result: trials += 1 # Check that a coinc exists - if 'foreground/ifar' in result: - ifar = result['foreground/ifar'] - stat = result['foreground/stat'] + if "foreground/ifar" in result: + ifar = result["foreground/ifar"] + stat = result["foreground/stat"] if ifar > mifar or (ifar == mifar and stat > mstat): mifar = ifar mstat = stat @@ -958,15 +997,16 @@ def pick_best_coinc(cls, coinc_results): # apply trials factor for the best coinc if mresult: - mresult['foreground/ifar'] = mifar / float(trials) - logger.info('Found %s coinc with ifar %s', - mresult['foreground/type'], - mresult['foreground/ifar']) + mresult["foreground/ifar"] = mifar / float(trials) + logger.info( + "Found %s coinc with ifar %s", + mresult["foreground/type"], + mresult["foreground/ifar"], + ) return mresult # If no coinc, just return one of the results dictionaries. They will # all contain the same results (i.e. single triggers) in this case. - else: - return coinc_results[0] + return coinc_results[0] @classmethod def from_cli(cls, args, num_templates, analysis_chunk, ifos): @@ -984,41 +1024,62 @@ def from_cli(cls, args, num_templates, analysis_chunk, ifos): stat_keywords, ) - return cls(num_templates, analysis_chunk, - args.ranking_statistic, - args.sngl_ranking, - stat_files, - return_background=args.store_background, - ifar_limit=args.background_ifar_limit, - timeslide_interval=args.timeslide_interval, - ifos=ifos, - coinc_window_pad=args.coinc_window_pad, - statistic_refresh_rate=args.statistic_refresh_rate, - **kwargs) + return cls( + num_templates, + analysis_chunk, + args.ranking_statistic, + args.sngl_ranking, + stat_files, + return_background=args.store_background, + ifar_limit=args.background_ifar_limit, + timeslide_interval=args.timeslide_interval, + ifos=ifos, + coinc_window_pad=args.coinc_window_pad, + statistic_refresh_rate=args.statistic_refresh_rate, + **kwargs, + ) @staticmethod def insert_args(parser): pycbcstat.insert_statistic_option_group(parser) - group = parser.add_argument_group('Coincident Background Estimation') - group.add_argument('--store-background', action='store_true', - help="Return background triggers with zerolag coincidencs") - group.add_argument('--background-ifar-limit', type=float, + group = parser.add_argument_group("Coincident Background Estimation") + group.add_argument( + "--store-background", + action="store_true", + help="Return background triggers with zerolag coincidencs", + ) + group.add_argument( + "--background-ifar-limit", + type=float, help="The limit on inverse false alarm rate to calculate " - "background in years", default=100.0) - group.add_argument('--timeslide-interval', type=float, - help="The interval between timeslides in seconds", default=0.1) - group.add_argument('--ifar-remove-threshold', type=float, - help="NOT YET IMPLEMENTED", default=100.0) + "background in years", + default=100.0, + ) + group.add_argument( + "--timeslide-interval", + type=float, + help="The interval between timeslides in seconds", + default=0.1, + ) + group.add_argument( + "--ifar-remove-threshold", + type=float, + help="NOT YET IMPLEMENTED", + default=100.0, + ) @staticmethod def verify_args(args, parser): """Verify that psd-var-related options are consistent""" - if ((hasattr(args, 'psd_variation') and not args.psd_variation) - and 'psdvar' in args.sngl_ranking): - parser.error(f"The single ifo ranking stat {args.sngl_ranking} " - "requires --psd-variation.") + if ( + hasattr(args, "psd_variation") and not args.psd_variation + ) and "psdvar" in args.sngl_ranking: + parser.error( + f"The single ifo ranking stat {args.sngl_ranking} " + "requires --psd-variation." + ) @property def background_time(self): @@ -1031,16 +1092,19 @@ def background_time(self): def save_state(self, filename): """Save the current state of the background buffers""" import pickle + pickle.dump(self, filename) @staticmethod def restore_state(filename): """Restore state of the background buffers from a file""" import pickle + return pickle.load(filename) def ifar(self, coinc_stat): - """Map a given value of the coincident ranking statistic to an inverse + """ + Map a given value of the coincident ranking statistic to an inverse false-alarm rate (IFAR) using the interally stored background sample. Parameters @@ -1055,13 +1119,15 @@ def ifar(self, coinc_stat): ifar_saturated: bool True if `coinc_stat` is larger than all the available background, in which case `ifar` is to be considered an upper limit. + """ n = self.coincs.num_greater(coinc_stat) ifar = conv.sec_to_year(self.background_time) / (n + 1) return ifar, n == 0 def set_singles_buffer(self, results): - """Create the singles buffer + """ + Create the singles buffer This creates the singles buffer for each ifo. The dtype is determined by a representative sample of the single triggers in the results. @@ -1070,13 +1136,17 @@ def set_singles_buffer(self, results): ---------- results: dict of dict Dict indexed by ifo and then trigger column. + """ # Determine the dtype from a sample of the data. self.singles_dtype = [] data = False for ifo in self.ifos: - if ifo in results and results[ifo] is not False \ - and len(results[ifo]['snr']): + if ( + ifo in results + and results[ifo] is not False + and len(results[ifo]["snr"]) + ): data = results[ifo] break @@ -1086,17 +1156,18 @@ def set_singles_buffer(self, results): for key in data: self.singles_dtype.append((key, data[key].dtype)) - if 'stat' not in data: - self.singles_dtype.append(('stat', self.stat_calculator.single_dtype)) + if "stat" not in data: + self.singles_dtype.append(("stat", self.stat_calculator.single_dtype)) # Create a ring buffer for each template ifo combination for ifo in self.ifos: - self.singles[ifo] = MultiRingBuffer(self.num_templates, - self.buffer_size, - self.singles_dtype) + self.singles[ifo] = MultiRingBuffer( + self.num_templates, self.buffer_size, self.singles_dtype + ) def _add_singles_to_buffer(self, results, ifos): - """Add single detector triggers to the internal buffer + """ + Add single detector triggers to the internal buffer Parameters ---------- @@ -1110,6 +1181,7 @@ def _add_singles_to_buffer(self, results, ifos): updated_singles: dict of numpy.ndarrays Array of indices that have been just updated in the internal buffers of single detector triggers. + """ if len(self.singles.keys()) == 0: self.set_singles_buffer(results) @@ -1125,29 +1197,31 @@ def _add_singles_to_buffer(self, results, ifos): for ifo in ifos: trigs = results[ifo] - if len(trigs['snr'] > 0): + if len(trigs["snr"] > 0): trigsc = copy.copy(trigs) - trigsc['ifo'] = ifo - trigsc['chisq'] = trigs['chisq'] * trigs['chisq_dof'] - trigsc['chisq_dof'] = (trigs['chisq_dof'] + 2) / 2 + trigsc["ifo"] = ifo + trigsc["chisq"] = trigs["chisq"] * trigs["chisq_dof"] + trigsc["chisq_dof"] = (trigs["chisq_dof"] + 2) / 2 single_stat = self.stat_calculator.single(trigsc) - del trigsc['ifo'] + del trigsc["ifo"] else: - single_stat = numpy.array([], ndmin=1, - dtype=self.stat_calculator.single_dtype) - trigs['stat'] = single_stat + single_stat = numpy.array( + [], ndmin=1, dtype=self.stat_calculator.single_dtype + ) + trigs["stat"] = single_stat # add each single detector trigger to the and advance the buffer data = numpy.zeros(len(single_stat), dtype=self.singles_dtype) for key, value in trigs.items(): data[key] = value - self.singles[ifo].add(trigs['template_id'], data) - updated_indices[ifo] = trigs['template_id'] + self.singles[ifo].add(trigs["template_id"], data) + updated_indices[ifo] = trigs["template_id"] return updated_indices def _find_coincs(self, results, valid_ifos): - """Look for coincs within the set of single triggers + """ + Look for coincs within the set of single triggers Parameters ---------- @@ -1166,6 +1240,7 @@ def _find_coincs(self, results, valid_ifos): Number of time shifted coincidences found. coinc_results: dict of arrays A dictionary of arrays containing the coincident results. + """ # For each new single detector trigger find the allowed coincidences # Record the template and the index of the single trigger that forms @@ -1174,10 +1249,10 @@ def _find_coincs(self, results, valid_ifos): # Initialize cstat = [[]] offsets = [] - ctimes = {self.ifos[0]:[], self.ifos[1]:[]} - single_expire = {self.ifos[0]:[], self.ifos[1]:[]} + ctimes = {self.ifos[0]: [], self.ifos[1]: []} + single_expire = {self.ifos[0]: [], self.ifos[1]: []} template_ids = [[]] - trigger_ids = {self.ifos[0]:[[]], self.ifos[1]:[[]]} + trigger_ids = {self.ifos[0]: [[]], self.ifos[1]: [[]]} # Calculate all the permutations of coincident triggers for each # new single detector trigger collected @@ -1190,7 +1265,7 @@ def _find_coincs(self, results, valid_ifos): for fixed_ifo, shift_ifo, shift_vec in zip( [self.ifos[0], self.ifos[1]], [self.ifos[1], self.ifos[0]], - [[0, -1], [-1, 0]] + [[0, -1], [-1, 0]], ): if fixed_ifo not in valid_ifos: # This ifo is not online now, so no new triggers or coincs @@ -1198,51 +1273,48 @@ def _find_coincs(self, results, valid_ifos): # Find newly added triggers in fixed_ifo trigs = results[fixed_ifo] # Calculate mchirp as a vectorized operation - mchirps = conv.mchirp_from_mass1_mass2( - trigs['mass1'], trigs['mass2'] - ) + mchirps = conv.mchirp_from_mass1_mass2(trigs["mass1"], trigs["mass2"]) # Loop over them one trigger at a time - for i in range(len(trigs['end_time'])): - trig_stat = trigs['stat'][i] - trig_time = trigs['end_time'][i] - template = trigs['template_id'][i] + for i in range(len(trigs["end_time"])): + trig_stat = trigs["stat"][i] + trig_time = trigs["end_time"][i] + template = trigs["template_id"][i] mchirp = mchirps[i] # Get current shift_ifo triggers in the same template - times = self.singles[shift_ifo].data(template)['end_time'] - stats = self.singles[shift_ifo].data(template)['stat'] + times = self.singles[shift_ifo].data(template)["end_time"] + stats = self.singles[shift_ifo].data(template)["stat"] # Perform coincidence. i1 is the list of trigger indices in the # shift_ifo which make coincs, slide is the corresponding slide # index. # (The second output would just be a list of zeroes as we only # have one trigger in the fixed_ifo.) - i1, _, slide = time_coincidence(times, - numpy.array(trig_time, ndmin=1, - dtype=numpy.float64), - self.time_window, - self.timeslide_interval) + i1, _, slide = time_coincidence( + times, + numpy.array(trig_time, ndmin=1, dtype=numpy.float64), + self.time_window, + self.timeslide_interval, + ) # Make a copy of the fixed ifo trig_stat for each coinc. # NB for some statistics the "stat" entry holds more than just # a ranking number. E.g. for the phase time consistency test, # it must also contain the phase, time and sensitivity. if self.trig_stat_memory is None: - self.trig_stat_memory = numpy.zeros( - 1, - dtype=trig_stat.dtype - ) + self.trig_stat_memory = numpy.zeros(1, dtype=trig_stat.dtype) while len(self.trig_stat_memory) < len(i1): self.trig_stat_memory = numpy.resize( - self.trig_stat_memory, - len(self.trig_stat_memory)*2 + self.trig_stat_memory, len(self.trig_stat_memory) * 2 ) - self.trig_stat_memory[:len(i1)] = trig_stat + self.trig_stat_memory[: len(i1)] = trig_stat # Force data into form needed by stat.py and then compute the # ranking statistic values. - sngls_list = [[fixed_ifo, self.trig_stat_memory[:len(i1)]], - [shift_ifo, stats[i1]]] + sngls_list = [ + [fixed_ifo, self.trig_stat_memory[: len(i1)]], + [shift_ifo, stats[i1]], + ] c = self.stat_calculator.rank_stat_coinc( sngls_list, @@ -1251,7 +1323,7 @@ def _find_coincs(self, results, valid_ifos): shift_vec, time_addition=self.coinc_window_pad, mchirp=mchirp, - dets=self.dets + dets=self.dets, ) # Store data about new triggers: slide index, stat value and @@ -1259,8 +1331,7 @@ def _find_coincs(self, results, valid_ifos): offsets.append(slide) cstat.append(c) ctimes[shift_ifo].append(times[i1]) - ctimes[fixed_ifo].append(numpy.zeros(len(c), - dtype=numpy.float64)) + ctimes[fixed_ifo].append(numpy.zeros(len(c), dtype=numpy.float64)) ctimes[fixed_ifo][-1].fill(trig_time) # As background triggers are removed after a certain time, we @@ -1268,11 +1339,8 @@ def _find_coincs(self, results, valid_ifos): single_expire[shift_ifo].append( self.singles[shift_ifo].expire_vector(template)[i1] ) - single_expire[fixed_ifo].append(numpy.zeros(len(c), - dtype=numpy.int32)) - single_expire[fixed_ifo][-1].fill( - self.singles[fixed_ifo].time - 1 - ) + single_expire[fixed_ifo].append(numpy.zeros(len(c), dtype=numpy.int32)) + single_expire[fixed_ifo][-1].fill(self.singles[fixed_ifo].time - 1) # Save the template and trigger ids to keep association # to singles. The trigger was just added so it must be in @@ -1288,8 +1356,7 @@ def _find_coincs(self, results, valid_ifos): trigger_ids[ifo] = numpy.concatenate(trigger_ids[ifo]).astype(numpy.int32) logger.info( - "%s: %s background and zerolag coincs", - ppdets(self.ifos, "-"), len(cstat) + "%s: %s background and zerolag coincs", ppdets(self.ifos, "-"), len(cstat) ) # Cluster the triggers we've found @@ -1301,13 +1368,18 @@ def _find_coincs(self, results, valid_ifos): ctime0 = numpy.concatenate(ctimes[self.ifos[0]]).astype(numpy.float64) ctime1 = numpy.concatenate(ctimes[self.ifos[1]]).astype(numpy.float64) logger.info("Clustering %s coincs", ppdets(self.ifos, "-")) - cidx = cluster_coincs(cstat, ctime0, ctime1, offsets, - self.timeslide_interval, - self.analysis_block + 2*self.time_window, - method='cython') + cidx = cluster_coincs( + cstat, + ctime0, + ctime1, + offsets, + self.timeslide_interval, + self.analysis_block + 2 * self.time_window, + method="cython", + ) offsets = offsets[cidx] - zerolag_idx = (offsets == 0) - bkg_idx = (offsets != 0) + zerolag_idx = offsets == 0 + bkg_idx = offsets != 0 for ifo in self.ifos: single_expire[ifo] = numpy.concatenate(single_expire[ifo]) @@ -1327,32 +1399,33 @@ def _find_coincs(self, results, valid_ifos): zerolag_cstat = cstat[cidx][zerolag_idx] ifar, ifar_sat = self.ifar(zerolag_cstat[0]) zerolag_results = { - 'foreground/ifar': ifar, - 'foreground/ifar_saturated': ifar_sat, - 'foreground/stat': zerolag_cstat, - 'foreground/type': '-'.join(self.ifos) + "foreground/ifar": ifar, + "foreground/ifar_saturated": ifar_sat, + "foreground/stat": zerolag_cstat, + "foreground/type": "-".join(self.ifos), } template = template_ids[idx] for ifo in self.ifos: trig_id = trigger_ids[ifo][idx] single_data = self.singles[ifo].data(template)[trig_id] for key in single_data.dtype.names: - path = f'foreground/{ifo}/{key}' + path = f"foreground/{ifo}/{key}" zerolag_results[path] = single_data[key] coinc_results.update(zerolag_results) # Save some summary statistics about the background - coinc_results['background/time'] = numpy.array([self.background_time]) - coinc_results['background/count'] = len(self.coincs.data) + coinc_results["background/time"] = numpy.array([self.background_time]) + coinc_results["background/count"] = len(self.coincs.data) # Save all the background triggers if self.return_background: - coinc_results['background/stat'] = self.coincs.data + coinc_results["background/stat"] = self.coincs.data return num_background, coinc_results def backout_last(self, updated_singles, num_coincs): - """Remove the recently added singles and coincs + """ + Remove the recently added singles and coincs Parameters ---------- @@ -1362,13 +1435,15 @@ def backout_last(self, updated_singles, num_coincs): num_coincs: int The number of coincs that were just added to the internal buffer of coincident triggers + """ for ifo in updated_singles: self.singles[ifo].discard_last(updated_singles[ifo]) self.coincs.remove(num_coincs) def add_singles(self, results): - """Add singles to the background estimate and find candidates + """ + Add singles to the background estimate and find candidates Parameters ---------- @@ -1381,16 +1456,20 @@ def add_singles(self, results): ------- coinc_results: dict of arrays A dictionary of arrays containing the coincident results. + """ # Let's see how large everything is logger.info( "%s: %s coincs, %s bytes", - ppdets(self.ifos, "-"), len(self.coincs), self.coincs.nbytes + ppdets(self.ifos, "-"), + len(self.coincs), + self.coincs.nbytes, ) # If there are no results just return valid_ifos = [k for k in results.keys() if results[k] and k in self.ifos] - if len(valid_ifos) == 0: return {} + if len(valid_ifos) == 0: + return {} with self.stat_calculator_lock: # Add single triggers to the internal buffer @@ -1401,7 +1480,7 @@ def add_singles(self, results): # record if a coinc is possible in this chunk if len(valid_ifos) == 2: - coinc_results['coinc_possible'] = True + coinc_results["coinc_possible"] = True return coinc_results @@ -1410,18 +1489,14 @@ def start_refresh_thread(self): Start a thread managing whether the stat_calculator will be updated """ if self.statistic_refresh_rate is None: - logger.info( - "Statistic refresh disabled for %s", ppdets(self.ifos, "-") - ) + logger.info("Statistic refresh disabled for %s", ppdets(self.ifos, "-")) return thread = threading.Thread( target=self.refresh_statistic, daemon=True, - name="Stat refresh " + ppdets(self.ifos, "-") - ) - logger.info( - "Starting %s statistic refresh thread", ppdets(self.ifos, "-") + name="Stat refresh " + ppdets(self.ifos, "-"), ) + logger.info("Starting %s statistic refresh thread", ppdets(self.ifos, "-")) thread.start() def refresh_statistic(self): @@ -1452,15 +1527,15 @@ def refresh_statistic(self): __all__ = [ + "CoincExpireBuffer", + "LiveCoincTimeslideBackgroundEstimator", + "MultiRingBuffer", "background_bin_from_string", - "timeslide_durations", - "time_coincidence", - "time_multi_coincidence", "cluster_coincs", "cluster_coincs_multiifo", - "mean_if_greater_than_zero", "cluster_over_time", - "MultiRingBuffer", - "CoincExpireBuffer", - "LiveCoincTimeslideBackgroundEstimator" + "mean_if_greater_than_zero", + "time_coincidence", + "time_multi_coincidence", + "timeslide_durations", ] diff --git a/pycbc/events/coinc_rate.py b/pycbc/events/coinc_rate.py index a05646f050a..12ee35dd184 100644 --- a/pycbc/events/coinc_rate.py +++ b/pycbc/events/coinc_rate.py @@ -5,16 +5,19 @@ # # ============================================================================= # -""" This module contains functions for calculating expected rates of noise - and signal coincidences. +""" +This module contains functions for calculating expected rates of noise +and signal coincidences. """ import itertools import logging + import numpy + import pycbc.detector -logger = logging.getLogger('pycbc.events.coinc_rate') +logger = logging.getLogger("pycbc.events.coinc_rate") def multiifo_noise_lograte(log_rates, slop): @@ -36,16 +39,16 @@ def multiifo_noise_lograte(log_rates, slop): expected_log_rates: dict Key: ifo combination string Value: expected log coincidence rate in the combination, units log Hz + """ expected_log_rates = {} # Order of ifos must be stable in output dict keys, so sort them ifos = sorted(list(log_rates.keys())) - ifostring = ' '.join(ifos) + ifostring = " ".join(ifos) # Calculate coincidence for all-ifo combination - expected_log_rates[ifostring] = \ - combination_noise_lograte(log_rates, slop) + expected_log_rates[ifostring] = combination_noise_lograte(log_rates, slop) # If more than one possible coincidence type exists, # calculate coincidence for subsets through recursion @@ -83,10 +86,13 @@ def combination_noise_rate(rates, slop): ------- numpy array Expected coincidence rate in the combination, units Hz + """ - logger.warning('combination_noise_rate() is liable to numerical ' - 'underflows, use combination_noise_lograte ' - 'instead') + logger.warning( + "combination_noise_rate() is liable to numerical " + "underflows, use combination_noise_lograte " + "instead" + ) log_rates = {k: numpy.log(r) for (k, r) in rates.items()} # exp may underflow return numpy.exp(combination_noise_lograte(log_rates, slop)) @@ -113,11 +119,10 @@ def combination_noise_lograte(log_rates, slop, dets=None): ------- numpy array Expected log coincidence rate in the combination, units Hz + """ # multiply product of trigger rates by the overlap time - allowed_area = multiifo_noise_coincident_area(list(log_rates), - slop, - dets=dets) + allowed_area = multiifo_noise_coincident_area(list(log_rates), slop, dets=dets) # list(dict.values()) is python-3-proof rateprod = numpy.sum(list(log_rates.values()), axis=0) return numpy.log(allowed_area) + rateprod @@ -144,6 +149,7 @@ def multiifo_noise_coincident_area(ifos, slop, dets=None): ------- allowed_area: float area in units of seconds^(n_ifos-1) that coincident values can fall in + """ # set up detector objects if dets is None: @@ -151,8 +157,9 @@ def multiifo_noise_coincident_area(ifos, slop, dets=None): n_ifos = len(ifos) if n_ifos == 2: - allowed_area = 2. * \ - (dets[ifos[0]].light_travel_time_to_detector(dets[ifos[1]]) + slop) + allowed_area = 2.0 * ( + dets[ifos[0]].light_travel_time_to_detector(dets[ifos[1]]) + slop + ) elif n_ifos == 3: tofs = numpy.zeros(n_ifos) ifo2_num = [] @@ -168,7 +175,7 @@ def multiifo_noise_coincident_area(ifos, slop, dets=None): # combine these to calculate allowed area allowed_area = 0 for i, _ in enumerate(ifos): - allowed_area += 2 * tofs[i] * tofs[ifo2_num[i]] - tofs[i]**2 + allowed_area += 2 * tofs[i] * tofs[ifo2_num[i]] - tofs[i] ** 2 else: raise NotImplementedError("Not able to deal with more than 3 ifos") @@ -188,6 +195,7 @@ def multiifo_signal_coincident_area(ifos): ------- allowed_area: float area in units of seconds^(n_ifos-1) that coincident signals will occupy + """ n_ifos = len(ifos) @@ -211,8 +219,9 @@ def multiifo_signal_coincident_area(ifos): tofs[i] = det0.light_travel_time_to_detector(det1) # calculate allowed area - phi_12 = numpy.arccos((tofs[0]**2 + tofs[1]**2 - tofs[2]**2) - / (2 * tofs[0] * tofs[1])) + phi_12 = numpy.arccos( + (tofs[0] ** 2 + tofs[1] ** 2 - tofs[2] ** 2) / (2 * tofs[0] * tofs[1]) + ) allowed_area = numpy.pi * tofs[0] * tofs[1] * numpy.sin(phi_12) else: raise NotImplementedError("Not able to deal with more than 3 ifos") diff --git a/pycbc/events/cuts.py b/pycbc/events/cuts.py index 587f502f580..cb6d079cef8 100644 --- a/pycbc/events/cuts.py +++ b/pycbc/events/cuts.py @@ -26,43 +26,46 @@ This module contains functions for reading in command line options and applying cuts to triggers or templates in the offline search """ -import logging + import copy +import logging + import numpy as np + from pycbc.events import ranking -from pycbc.io import hdf -from pycbc.tmpltbank import bank_conversions as bank_conv -from pycbc.io import get_chisq_from_file_choice +from pycbc.io import get_chisq_from_file_choice, hdf # Only used to check isinstance: from pycbc.io.hdf import ReadByTemplate +from pycbc.tmpltbank import bank_conversions as bank_conv -logger = logging.getLogger('pycbc.events.cuts') +logger = logging.getLogger("pycbc.events.cuts") # sngl_rank_keys are the allowed names of reweighted SNR functions sngl_rank_keys = ranking.sngls_ranking_function_dict.keys() trigger_param_choices = list(sngl_rank_keys) -trigger_param_choices += [cc + '_chisq' for cc in hdf.chisq_choices] -trigger_param_choices += ['end_time', 'psd_var_val', 'sigmasq', - 'sigma_multiple'] - -template_fit_param_choices = ['fit_by_fit_coeff', 'smoothed_fit_coeff', - 'fit_by_count_above_thresh', - 'smoothed_fit_count_above_thresh', - 'fit_by_count_in_template', - 'smoothed_fit_count_in_template'] -template_param_choices = bank_conv.conversion_options + \ - template_fit_param_choices +trigger_param_choices += [cc + "_chisq" for cc in hdf.chisq_choices] +trigger_param_choices += ["end_time", "psd_var_val", "sigmasq", "sigma_multiple"] + +template_fit_param_choices = [ + "fit_by_fit_coeff", + "smoothed_fit_coeff", + "fit_by_count_above_thresh", + "smoothed_fit_count_above_thresh", + "fit_by_count_in_template", + "smoothed_fit_count_in_template", +] +template_param_choices = bank_conv.conversion_options + template_fit_param_choices # What are the inequalities associated with the cuts? # 'upper' means upper limit, and so requires value < threshold # to keep a trigger ineq_functions = { - 'upper': np.less, - 'lower': np.greater, - 'upper_inc': np.less_equal, - 'lower_inc': np.greater_equal + "upper": np.less, + "lower": np.greater, + "upper_inc": np.less_equal, + "lower_inc": np.greater_equal, } ineq_choices = list(ineq_functions.keys()) @@ -71,22 +74,29 @@ def insert_cuts_option_group(parser): """ Add options to the parser for cuts to the templates/triggers """ - parser.add_argument('--trigger-cuts', nargs='+', - help="Cuts to apply to the triggers, supplied as " - "PARAMETER:VALUE:LIMIT, where, PARAMETER is the " - "parameter to be cut, VALUE is the value at " - "which it is cut, and LIMIT is one of '" - + "', '".join(ineq_choices) + - "' to indicate the inequality needed. " - "PARAMETER is one of:'" - + "', '".join(trigger_param_choices) + - "'. For example snr:6:LOWER removes triggers " - "with matched filter SNR < 6") - parser.add_argument('--template-cuts', nargs='+', - help="Cuts to apply to the triggers, supplied as " - "PARAMETER:VALUE:LIMIT. Format is the same as in " - "--trigger-cuts. PARAMETER can be one of '" - + "', '".join(template_param_choices) + "'.") + parser.add_argument( + "--trigger-cuts", + nargs="+", + help="Cuts to apply to the triggers, supplied as " + "PARAMETER:VALUE:LIMIT, where, PARAMETER is the " + "parameter to be cut, VALUE is the value at " + "which it is cut, and LIMIT is one of '" + + "', '".join(ineq_choices) + + "' to indicate the inequality needed. " + "PARAMETER is one of:'" + + "', '".join(trigger_param_choices) + + "'. For example snr:6:LOWER removes triggers " + "with matched filter SNR < 6", + ) + parser.add_argument( + "--template-cuts", + nargs="+", + help="Cuts to apply to the triggers, supplied as " + "PARAMETER:VALUE:LIMIT. Format is the same as in " + "--trigger-cuts. PARAMETER can be one of '" + + "', '".join(template_param_choices) + + "'.", + ) def convert_inputstr(inputstr, choices): @@ -97,26 +107,32 @@ def convert_inputstr(inputstr, choices): Do input checks """ try: - cut_param, cut_value_str, cut_limit = inputstr.split(':') + cut_param, cut_value_str, cut_limit = inputstr.split(":") except ValueError as value_e: - logger.warning("ERROR: Cut string format not correct, please " - "supply as PARAMETER:VALUE:LIMIT") + logger.warning( + "ERROR: Cut string format not correct, please " + "supply as PARAMETER:VALUE:LIMIT" + ) raise value_e if cut_param.lower() not in choices: - raise NotImplementedError("Cut parameter " + cut_param.lower() + " " - "not recognised, choose from " - + ", ".join(choices)) + raise NotImplementedError( + "Cut parameter " + cut_param.lower() + " " + "not recognised, choose from " + ", ".join(choices) + ) if cut_limit.lower() not in ineq_choices: - raise NotImplementedError("Cut inequality " + cut_limit.lower() + " " - "not recognised, choose from " - + ", ".join(ineq_choices)) + raise NotImplementedError( + "Cut inequality " + cut_limit.lower() + " " + "not recognised, choose from " + ", ".join(ineq_choices) + ) try: cut_value = float(cut_value_str) except ValueError as value_e: - logger.warning("ERROR: Cut value must be convertible into a float, " - "got '%s'.", cut_value_str) + logger.warning( + "ERROR: Cut value must be convertible into a float, got '%s'.", + cut_value_str, + ) raise value_e return {(cut_param, ineq_functions[cut_limit]): cut_value} @@ -135,13 +151,17 @@ def check_update_cuts(cut_dict, new_cut): new_cut: single-entry dictionary dictionary to define the new cut which is being considered to add + """ new_cut_key = list(new_cut.keys())[0] if new_cut_key in cut_dict: # The cut has already been called - logger.warning("WARNING: Cut parameter %s and function %s have " - "already been used. Utilising the strictest cut.", - new_cut_key[0], new_cut_key[1].__name__) + logger.warning( + "WARNING: Cut parameter %s and function %s have " + "already been used. Utilising the strictest cut.", + new_cut_key[0], + new_cut_key[1].__name__, + ) # Extract the function and work out which is strictest cut_function = new_cut_key[1] value_new = list(new_cut.values())[0] @@ -150,17 +170,25 @@ def check_update_cuts(cut_dict, new_cut): # The new threshold would survive the cut of the # old threshold, therefore the new threshold is stricter # - update it - logger.warning("WARNING: New threshold of %.3f is " - "stricter than old threshold %.3f, " - "using cut at %.3f.", - value_new, value_old, value_new) + logger.warning( + "WARNING: New threshold of %.3f is " + "stricter than old threshold %.3f, " + "using cut at %.3f.", + value_new, + value_old, + value_new, + ) cut_dict.update(new_cut) else: # New cut would not make a difference, ignore it - logger.warning("WARNING: New threshold of %.3f is less " - "strict than old threshold %.3f, using " - "cut at %.3f.", - value_new, value_old, value_old) + logger.warning( + "WARNING: New threshold of %.3f is less " + "strict than old threshold %.3f, using " + "cut at %.3f.", + value_new, + value_old, + value_old, + ) else: # This is a new cut - add it cut_dict.update(new_cut) @@ -194,14 +222,12 @@ def ingest_cuts_option_group(args): return trigger_cut_dict, template_cut_dict -def sigma_multiple_cut_thresh(template_ids, statistic, - cut_thresh, ifo): +def sigma_multiple_cut_thresh(template_ids, statistic, cut_thresh, ifo): """ Apply cuts based on a multiple of the median sigma value for the template Parameters ---------- - template_ids: template_id values for each of the triggers to be considered, this will be used to associate a sigma threshold for each trigger @@ -219,14 +245,18 @@ def sigma_multiple_cut_thresh(template_ids, statistic, idx_out: numpy array An array of the indices of triggers which meet the criteria set by the dictionary + """ statistic_classname = statistic.__class__.__name__ - if not hasattr(statistic, 'fits_by_tid'): - raise ValueError("Cut parameter 'sigma_muliple' cannot " - "be used when the ranking statistic " + - statistic_classname + " does not use " - "template fitting.") - tid_med_sigma = statistic.fits_by_tid[ifo]['median_sigma'] + if not hasattr(statistic, "fits_by_tid"): + raise ValueError( + "Cut parameter 'sigma_muliple' cannot " + "be used when the ranking statistic " + + statistic_classname + + " does not use " + "template fitting." + ) + tid_med_sigma = statistic.fits_by_tid[ifo]["median_sigma"] return cut_thresh * tid_med_sigma[template_ids] @@ -253,8 +283,9 @@ def apply_trigger_cuts(triggers, trigger_cut_dict, statistic=None): idx_out: numpy array An array of the indices which meet the criteria set by the dictionary + """ - idx_out = np.arange(len(triggers['snr'])) + idx_out = np.arange(len(triggers["snr"])) # Loop through the different cuts, and apply them for parameter_cut_function, cut_thresh in trigger_cut_dict.items(): @@ -262,32 +293,31 @@ def apply_trigger_cuts(triggers, trigger_cut_dict, statistic=None): parameter, cut_function = parameter_cut_function # What kind of parameter is it? - if parameter.endswith('_chisq'): + if parameter.endswith("_chisq"): # parameter is a chisq-type thing - chisq_choice = parameter.split('_')[0] + chisq_choice = parameter.split("_")[0] # Currently calculated for all triggers - this seems inefficient value = get_chisq_from_file_choice(triggers, chisq_choice) # Apply any previous cuts to the value for comparison value = value[idx_out] elif parameter == "sigma_multiple": if isinstance(triggers, ReadByTemplate): - value = np.sqrt(triggers['sigmasq'][idx_out]) + value = np.sqrt(triggers["sigmasq"][idx_out]) # Get a cut threshold value, this will be different # depending on the template ID, so we rewrite cut_thresh # as a value for each trigger, numpy comparison functions # allow this - cut_thresh = sigma_multiple_cut_thresh(triggers.template_num, - statistic, - cut_thresh, - triggers.ifo) + cut_thresh = sigma_multiple_cut_thresh( + triggers.template_num, statistic, cut_thresh, triggers.ifo + ) else: err_msg = "Cuts on 'sigma_multiple' are only implemented for " err_msg += "triggers in a ReadByTemplate format. This code " err_msg += f"uses a {type(triggers).__name__} format." raise NotImplementedError(err_msg) - elif ((not hasattr(triggers, "file") and parameter in triggers) - or (hasattr(triggers, "file") - and parameter in triggers.file[triggers.ifo])): + elif (not hasattr(triggers, "file") and parameter in triggers) or ( + hasattr(triggers, "file") and parameter in triggers.file[triggers.ifo] + ): # parameter can be read direct from the trigger dictionary / file value = triggers[parameter] # Apply any previous cuts to the value for comparison @@ -299,17 +329,20 @@ def apply_trigger_cuts(triggers, trigger_cut_dict, statistic=None): # Apply any previous cuts to the value for comparison value = value[idx_out] else: - raise NotImplementedError("Parameter '" + parameter + "' not " - "recognised. Input sanitisation means " - "this shouldn't have happened?!") + raise NotImplementedError( + "Parameter '" + parameter + "' not " + "recognised. Input sanitisation means " + "this shouldn't have happened?!" + ) idx_out = idx_out[cut_function(value, cut_thresh)] return idx_out -def apply_template_fit_cut(statistic, ifos, parameter_cut_function, cut_thresh, - template_ids): +def apply_template_fit_cut( + statistic, ifos, parameter_cut_function, cut_thresh, template_ids +): """ Apply cuts to template fit parameters, these have a few more checks needed, so we separate out from apply_template_cuts defined later @@ -341,22 +374,25 @@ def apply_template_fit_cut(statistic, ifos, parameter_cut_function, cut_thresh, ------- tids_out: numpy array Array of template_ids which have passed this cut + """ parameter, cut_function = parameter_cut_function statistic_classname = statistic.__class__.__name__ # We can only apply template fit cuts if template fits have been done - if not hasattr(statistic, 'fits_by_tid'): - raise ValueError("Cut parameter " + parameter + " cannot " - "be used when the ranking statistic " + - statistic_classname + " does not use " - "template fitting.") + if not hasattr(statistic, "fits_by_tid"): + raise ValueError( + "Cut parameter " + parameter + " cannot " + "be used when the ranking statistic " + + statistic_classname + + " does not use " + "template fitting." + ) # Is the parameter actually in the fits dictionary? if parameter not in statistic.fits_by_tid[ifos[0]]: # Shouldn't get here due to input sanitisation - raise ValueError("Cut parameter " + parameter + " not " - "available in fits file.") + raise ValueError("Cut parameter " + parameter + " not available in fits file.") # Template IDs array to cut down in each IFO tids_out = copy.copy(template_ids) @@ -370,8 +406,9 @@ def apply_template_fit_cut(statistic, ifos, parameter_cut_function, cut_thresh, return tids_out -def apply_template_cuts(bank, template_cut_dict, template_ids=None, - statistic=None, ifos=None): +def apply_template_cuts( + bank, template_cut_dict, template_ids=None, statistic=None, ifos=None +): """ Fetch/calculate the parameter for the templates, possibly already preselected by template_ids, and then apply the cuts defined @@ -412,14 +449,17 @@ def apply_template_cuts(bank, template_cut_dict, template_ids=None, ------- tids_out: numpy array Array of template_ids which have passed all cuts + """ # Get the initial list of templates: - tids_out = np.arange(bank['mass1'].size) \ - if template_ids is None else template_ids[:] + tids_out = ( + np.arange(bank["mass1"].size) if template_ids is None else template_ids[:] + ) if (statistic is None) ^ (ifos is None): - raise NotImplementedError("Either both or neither of statistic and " - "ifos must be supplied.") + raise NotImplementedError( + "Either both or neither of statistic and ifos must be supplied." + ) if not template_cut_dict: # No cuts are defined in the dictionary: just return the @@ -438,13 +478,13 @@ def apply_template_cuts(bank, template_cut_dict, template_ids=None, tids_out = tids_out[cut_function(values, cut_thresh)] elif parameter in template_fit_param_choices: if statistic and ifos: - tids_out = apply_template_fit_cut(statistic, - ifos, - parameter_cut_function, - cut_thresh, - tids_out) + tids_out = apply_template_fit_cut( + statistic, ifos, parameter_cut_function, cut_thresh, tids_out + ) else: - raise ValueError("Cut parameter " + parameter + " not recognised." - " This shouldn't happen with input sanitisation") + raise ValueError( + "Cut parameter " + parameter + " not recognised." + " This shouldn't happen with input sanitisation" + ) return tids_out diff --git a/pycbc/events/eventmgr.py b/pycbc/events/eventmgr.py index 1eb3006bec7..4ffb4183209 100644 --- a/pycbc/events/eventmgr.py +++ b/pycbc/events/eventmgr.py @@ -21,44 +21,49 @@ # # ============================================================================= # -"""This modules defines functions for clustering and thresholding timeseries to +""" +This modules defines functions for clustering and thresholding timeseries to produces event triggers """ -import os.path + import copy import itertools import logging +import os.path import pickle -import numpy + import h5py +import numpy -from pycbc.types import Array -from pycbc.scheme import schemed from pycbc.detector import Detector +from pycbc.scheme import schemed +from pycbc.types import Array from . import coinc, ranking - from .eventmgr_cython import findchirp_cluster_over_window_cython -logger = logging.getLogger('pycbc.events.eventmgr') +logger = logging.getLogger("pycbc.events.eventmgr") + @schemed("pycbc.events.threshold_") def threshold(series, value): - """Return list of values and indices values over threshold in series. - """ + """Return list of values and indices values over threshold in series.""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + @schemed("pycbc.events.threshold_") def threshold_only(series, value): - """Return list of values and indices whose values in series are - larger (in absolute value) than value + """ + Return list of values and indices whose values in series are + larger (in absolute value) than value """ err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + # FIXME: This should be under schemed, but I don't understand that yet! def threshold_real_numpy(series, value): arr = series.data @@ -66,14 +71,15 @@ def threshold_real_numpy(series, value): vals = arr[locs] return locs, vals + @schemed("pycbc.events.threshold_") def threshold_and_cluster(series, threshold, window): - """Return list of values and indices values over threshold in series. - """ + """Return list of values and indices values over threshold in series.""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + @schemed("pycbc.events.threshold_") def _threshold_cluster_factory(series): err_msg = "This class is a stub that should be overridden using the " @@ -81,18 +87,21 @@ def _threshold_cluster_factory(series): raise ValueError(err_msg) -class ThresholdCluster(object): - """Create a threshold and cluster engine +class ThresholdCluster: + """ + Create a threshold and cluster engine Parameters - ----------- + ---------- series : complex64 Input pycbc.types.Array (or subclass); it will be searched for points above threshold that are then clustered + """ + def __new__(cls, *args, **kwargs): real_cls = _threshold_cluster_factory(*args, **kwargs) - return real_cls(*args, **kwargs) # pylint:disable=not-callable + return real_cls(*args, **kwargs) # pylint:disable=not-callable # The class below should serve as the parent for all schemed classes. @@ -103,36 +112,37 @@ def __new__(cls, *args, **kwargs): # http://stackoverflow.com/questions/2025562/inherit-docstrings-in-python-class-inheritance # # will work? Is there a better way? -class _BaseThresholdCluster(object): +class _BaseThresholdCluster: def threshold_and_cluster(self, threshold, window): """ Threshold and cluster the memory specified at instantiation with the threshold and window size specified at creation. Parameters - ----------- + ---------- threshold : float32 The minimum absolute value of the series given at object initialization to return when thresholding and clustering. window : uint32 The size (in number of samples) of the window over which to cluster - Returns: - -------- + Returns + ------- event_vals : complex64 Numpy array, complex values of the clustered events event_locs : uint32 Numpy array, indices into series of location of events + """ - pass def findchirp_cluster_over_window(times, values, window_length): - """ Reduce the events by clustering over a window using + """ + Reduce the events by clustering over a window using the FindChirp clustering algorithm Parameters - ----------- + ---------- indices: Array The list of indices of the SNR values snr: Array @@ -144,24 +154,27 @@ def findchirp_cluster_over_window(times, values, window_length): ------- indices: Array The reduced list of indices of the SNR values + """ - assert window_length > 0, 'Clustering window length is not positive' + assert window_length > 0, "Clustering window length is not positive" indices = numpy.zeros(len(times), dtype=numpy.int32) tlen = len(times) absvalues = numpy.asarray(abs(values)) times = numpy.asarray(times, dtype=numpy.int32) - k = findchirp_cluster_over_window_cython(times, absvalues, window_length, - indices, tlen) + k = findchirp_cluster_over_window_cython( + times, absvalues, window_length, indices, tlen + ) - return indices[0:k+1] + return indices[0 : k + 1] def cluster_reduce(idx, snr, window_size): - """ Reduce the events by clustering over a window + """ + Reduce the events by clustering over a window Parameters - ----------- + ---------- indices: Array The list of indices of the SNR values snr: Array @@ -175,56 +188,44 @@ def cluster_reduce(idx, snr, window_size): The list of indices of the SNR values snr: Array The list of SNR values + """ ind = findchirp_cluster_over_window(idx, snr, window_size) return idx.take(ind), snr.take(ind) -class H5FileSyntSugar(object): - """Convenience class that adds some syntactic sugar to h5py.File. - """ - def __init__(self, name, prefix=''): - self.f = h5py.File(name, 'w') +class H5FileSyntSugar: + """Convenience class that adds some syntactic sugar to h5py.File.""" + + def __init__(self, name, prefix=""): + self.f = h5py.File(name, "w") self.prefix = prefix def __setitem__(self, name, data): self.f.create_dataset( - self.prefix + '/' + name, + self.prefix + "/" + name, data=data, - compression='gzip', + compression="gzip", compression_opts=9, - shuffle=True + shuffle=True, ) -class EventManager(object): - def __init__( - self, - opt, - column, - column_types, - array_minsize=10000, - **kwds - ): +class EventManager: + def __init__(self, opt, column, column_types, array_minsize=10000, **kwds): self.opt = opt self.global_params = kwds - self.array_minsize=array_minsize + self.array_minsize = array_minsize - self.event_dtype = [('template_id', int)] + self.event_dtype = [("template_id", int)] for col, coltype in zip(column, column_types): self.event_dtype.append((col, coltype)) - self._events = numpy.zeros( - [self.array_minsize], - dtype=self.event_dtype - ) + self._events = numpy.zeros([self.array_minsize], dtype=self.event_dtype) self._events_size = 0 self.template_params = [] self.template_index = -1 - self.template_events = numpy.zeros( - [self.array_minsize], - dtype=self.event_dtype - ) + self.template_events = numpy.zeros([self.array_minsize], dtype=self.event_dtype) self.template_event_size = 0 self.write_performance = False @@ -233,21 +234,21 @@ def save_state(self, tnum_finished, filename): from pycbc.io.hdf import dump_state self.tnum_finished = tnum_finished - logger.info('Writing checkpoint file at template %s', tnum_finished) - fp = h5py.File(filename, 'w') + logger.info("Writing checkpoint file at template %s", tnum_finished) + fp = h5py.File(filename, "w") dump_state(self, fp, protocol=pickle.HIGHEST_PROTOCOL) fp.close() @property def events(self): - return self._events[:self._events_size] + return self._events[: self._events_size] @staticmethod def restore_state(filename): """Restore state of the background buffers from a file""" from pycbc.io.hdf import load_state - fp = h5py.File(filename, 'r') + fp = h5py.File(filename, "r") try: mgr = load_state(fp) except Exception as e: @@ -255,7 +256,7 @@ def restore_state(filename): raise e fp.close() next_template = mgr.tnum_finished + 1 - logger.info('Restoring with checkpoint at template %s', next_template) + logger.info("Restoring with checkpoint at template %s", next_template) return mgr.tnum_finished + 1, mgr @classmethod @@ -275,70 +276,75 @@ def from_multi_ifo_interface(cls, opt, ifo, column, column_types, **kwds): def cut_events_via_mask(self, keep): # keep should be a boolean array of len self._events_size num_keep = keep.sum() - self._events[:num_keep] = self._events[:self._events_size][keep] + self._events[:num_keep] = self._events[: self._events_size][keep] self._events_size = num_keep def cut_events_via_indices(self, indices): # indices should be a list of indices to keep num_keep = len(indices) - self._events[:num_keep] = self._events[:self._events_size][indices] + self._events[:num_keep] = self._events[: self._events_size][indices] self._events_size = num_keep def chisq_threshold(self, value, num_bins, delta=0): keep = numpy.ones(self._events_size, dtype=bool) for i, event in enumerate(self.events): - xi = event['chisq'] / (event['chisq_dof'] + - delta * event['snr'].conj() * event['snr']) + xi = event["chisq"] / ( + event["chisq_dof"] + delta * event["snr"].conj() * event["snr"] + ) if xi > value: keep[i] = 0 self.cut_events_via_mask(keep) def newsnr_threshold(self, threshold): - """ Remove events with newsnr smaller than given threshold - """ + """Remove events with newsnr smaller than given threshold""" if not self.opt.chisq_bins: - raise RuntimeError('Chi-square test must be enabled in order to ' - 'use newsnr threshold') + raise RuntimeError( + "Chi-square test must be enabled in order to use newsnr threshold" + ) - nsnrs = ranking.newsnr(abs(self.events['snr']), - self.events['chisq'] / self.events['chisq_dof']) + nsnrs = ranking.newsnr( + abs(self.events["snr"]), self.events["chisq"] / self.events["chisq_dof"] + ) self.cut_events_via_mask(nsnrs >= threshold) def keep_near_injection(self, window, injections): from pycbc.events.veto import indices_within_times + if len(self.events) == 0: return inj_time = numpy.array(injections.end_times()) - gpstime = self.events['time_index'].astype(numpy.float64) + gpstime = self.events["time_index"].astype(numpy.float64) gpstime = gpstime / self.opt.sample_rate + self.opt.gps_start_time i = indices_within_times(gpstime, inj_time - window, inj_time + window) self.cut_events_via_indices(i) - def keep_loudest_in_interval(self, window, num_keep, statname="newsnr", - log_chirp_width=None): + def keep_loudest_in_interval( + self, window, num_keep, statname="newsnr", log_chirp_width=None + ): if len(self.events) == 0: return e_copy = self.events.copy() # Here self.events['snr'] is the complex SNR - e_copy['snr'] = abs(e_copy['snr']) + e_copy["snr"] = abs(e_copy["snr"]) # Messy step because pycbc inspiral's internal 'chisq_dof' is 2p-2 # but stat.py / ranking.py functions use 'chisq_dof' = p - e_copy['chisq_dof'] = e_copy['chisq_dof'] / 2 + 1 + e_copy["chisq_dof"] = e_copy["chisq_dof"] / 2 + 1 statv = ranking.get_sngls_ranking_from_trigs(e_copy, statname) # Convert trigger time to integer bin number # NB time_index and window are in units of samples - wtime = (e_copy['time_index'] / window).astype(numpy.int32) + wtime = (e_copy["time_index"] / window).astype(numpy.int32) bins = numpy.unique(wtime) if log_chirp_width: from pycbc.conversions import mchirp_from_mass1_mass2 - m1 = numpy.array([p['tmplt'].mass1 for p in self.template_params]) - m2 = numpy.array([p['tmplt'].mass2 for p in self.template_params]) - mc = mchirp_from_mass1_mass2(m1, m2)[e_copy['template_id']] + + m1 = numpy.array([p["tmplt"].mass1 for p in self.template_params]) + m2 = numpy.array([p["tmplt"].mass2 for p in self.template_params]) + mc = mchirp_from_mass1_mass2(m1, m2)[e_copy["template_id"]] # convert chirp mass to integer bin number imc = (numpy.log(mc) / log_chirp_width).astype(numpy.int32) @@ -352,7 +358,7 @@ def keep_loudest_in_interval(self, window, num_keep, statname="newsnr", bloudest = statv[bloc].argsort()[-num_keep:] keep.append(bloc[bloudest]) else: - bloc = numpy.where((wtime == b))[0] + bloc = numpy.where(wtime == b)[0] bloudest = statv[bloc].argsort()[-num_keep:] keep.append(bloc[bloudest]) @@ -360,7 +366,7 @@ def keep_loudest_in_interval(self, window, num_keep, statname="newsnr", self.cut_events_via_indices(keep) def add_template_events(self, columns, vectors): - """ Add a vector indexed """ + """Add a vector indexed""" # initialize with zeros - since vectors can be None, look for the # longest one that isn't new_events = None @@ -370,30 +376,26 @@ def add_template_events(self, columns, vectors): break # they shouldn't all be None assert new_events is not None - new_events['template_id'] = self.template_index + new_events["template_id"] = self.template_index for c, v in zip(columns, vectors): if v is not None: if isinstance(v, Array): new_events[c] = v.numpy() else: new_events[c] = v - new_size = self.template_event_size+len(new_events) + new_size = self.template_event_size + len(new_events) if new_size > len(self.template_events): - self.template_events.resize( - new_size + self.array_minsize, - refcheck=False - ) - self.template_events[self.template_event_size:new_size] = new_events + self.template_events.resize(new_size + self.array_minsize, refcheck=False) + self.template_events[self.template_event_size : new_size] = new_events self.template_event_size += len(new_events) def cluster_template_events(self, tcolumn, column, window_size): - """ Cluster the internal events over the named column - """ + """Cluster the internal events over the named column""" if window_size > 0: - cvec = self.template_events[column][:self.template_event_size] - tvec = self.template_events[tcolumn][:self.template_event_size] + cvec = self.template_events[column][: self.template_event_size] + tvec = self.template_events[tcolumn][: self.template_event_size] indices = findchirp_cluster_over_window(tvec, cvec, window_size) - self.template_events[:len(indices)] = self.template_events[indices] + self.template_events[: len(indices)] = self.template_events[indices] self.template_event_size = len(indices) def new_template(self, **kwds): @@ -407,13 +409,13 @@ def finalize_template_events(self): new_event_size = self.template_event_size + self._events_size if new_event_size > len(self._events): self._events.resize( - (new_event_size//self.array_minsize + 1) * self.array_minsize, - refcheck=False + (new_event_size // self.array_minsize + 1) * self.array_minsize, + refcheck=False, ) - self._events[self._events_size:new_event_size] = ( - self.template_events[:self.template_event_size] - ) + self._events[self._events_size : new_event_size] = self.template_events[ + : self.template_event_size + ] self._events_size = new_event_size self.template_event_size = 0 @@ -422,8 +424,7 @@ def consolidate_events(self, opt, gwstrain=None): logger.info("We currently have %d triggers", len(self.events)) if opt.chisq_threshold and opt.chisq_bins: logger.info("Removing triggers with poor chisq") - self.chisq_threshold(opt.chisq_threshold, opt.chisq_bins, - opt.chisq_delta) + self.chisq_threshold(opt.chisq_threshold, opt.chisq_bins, opt.chisq_delta) logger.info("%d remaining triggers", len(self.events)) if opt.newsnr_threshold and opt.chisq_bins: @@ -432,24 +433,28 @@ def consolidate_events(self, opt, gwstrain=None): logger.info("%d remaining triggers", len(self.events)) if opt.keep_loudest_interval: - logger.info("Removing triggers not within the top %s " - "loudest of a %s second interval by %s", - opt.keep_loudest_num, opt.keep_loudest_interval, - opt.keep_loudest_stat) - self.keep_loudest_in_interval\ - (opt.keep_loudest_interval * opt.sample_rate, - opt.keep_loudest_num, statname=opt.keep_loudest_stat, - log_chirp_width=opt.keep_loudest_log_chirp_window) + logger.info( + "Removing triggers not within the top %s " + "loudest of a %s second interval by %s", + opt.keep_loudest_num, + opt.keep_loudest_interval, + opt.keep_loudest_stat, + ) + self.keep_loudest_in_interval( + opt.keep_loudest_interval * opt.sample_rate, + opt.keep_loudest_num, + statname=opt.keep_loudest_stat, + log_chirp_width=opt.keep_loudest_log_chirp_window, + ) logger.info("%d remaining triggers", len(self.events)) - if opt.injection_window and hasattr(gwstrain, 'injections'): - logger.info("Keeping triggers within %s seconds of injection", - opt.injection_window) - self.keep_near_injection(opt.injection_window, - gwstrain.injections) + if opt.injection_window and hasattr(gwstrain, "injections"): + logger.info( + "Keeping triggers within %s seconds of injection", opt.injection_window + ) + self.keep_near_injection(opt.injection_window, gwstrain.injections) logger.info("%d remaining triggers", len(self.events)) - def finalize_events(self): # At the moment this method does nothing, but it is called by all # front-end codes using this, and may have stuff added to it in the @@ -458,12 +463,11 @@ def finalize_events(self): def make_output_dir(self, outname): path = os.path.dirname(outname) - if path != '': + if path != "": if not os.path.exists(path) and path is not None: os.makedirs(path) - def save_performance(self, ncores, nfilters, ntemplates, run_time, - setup_time): + def save_performance(self, ncores, nfilters, ntemplates, run_time, setup_time): """ Calls variables from pycbc_inspiral to be used in a timing calculation """ @@ -475,66 +479,71 @@ def save_performance(self, ncores, nfilters, ntemplates, run_time, self.write_performance = True def write_events(self, outname): - """Write the found events to a file. The only currently supported + """ + Write the found events to a file. The only currently supported format is HDF5, indicated by an .hdf or .h5 extension. """ self.make_output_dir(outname) - if outname.endswith(('.hdf', '.h5')): + if outname.endswith((".hdf", ".h5")): self.write_to_hdf(outname) return - raise ValueError('Unsupported event output file format') + raise ValueError("Unsupported event output file format") def write_to_hdf(self, outname): - self.events.sort(order='template_id') - th = numpy.array([p['tmplt'].template_hash for p in - self.template_params]) - tid = self.events['template_id'] + self.events.sort(order="template_id") + th = numpy.array([p["tmplt"].template_hash for p in self.template_params]) + tid = self.events["template_id"] f = H5FileSyntSugar(outname, self.opt.channel_name[0:2]) if len(self.events): - f['snr'] = abs(self.events['snr']) + f["snr"] = abs(self.events["snr"]) try: # Precessing - f['u_vals'] = self.events['u_vals'] - f['coa_phase'] = self.events['coa_phase'] - f['hplus_cross_corr'] = self.events['hplus_cross_corr'] + f["u_vals"] = self.events["u_vals"] + f["coa_phase"] = self.events["coa_phase"] + f["hplus_cross_corr"] = self.events["hplus_cross_corr"] except Exception: # Not precessing - f['coa_phase'] = numpy.angle(self.events['snr']) - f['chisq'] = self.events['chisq'] - f['bank_chisq'] = self.events['bank_chisq'] - f['bank_chisq_dof'] = self.events['bank_chisq_dof'] - f['cont_chisq'] = self.events['cont_chisq'] - f['end_time'] = self.events['time_index'] / \ - float(self.opt.sample_rate) \ - + self.opt.gps_start_time + f["coa_phase"] = numpy.angle(self.events["snr"]) + f["chisq"] = self.events["chisq"] + f["bank_chisq"] = self.events["bank_chisq"] + f["bank_chisq_dof"] = self.events["bank_chisq_dof"] + f["cont_chisq"] = self.events["cont_chisq"] + f["end_time"] = ( + self.events["time_index"] / float(self.opt.sample_rate) + + self.opt.gps_start_time + ) try: # Precessing template_sigmasq_plus = numpy.array( - [t['sigmasq_plus'] for t in self.template_params], - dtype=numpy.float32) - f['sigmasq_plus'] = template_sigmasq_plus[tid] + [t["sigmasq_plus"] for t in self.template_params], + dtype=numpy.float32, + ) + f["sigmasq_plus"] = template_sigmasq_plus[tid] template_sigmasq_cross = numpy.array( - [t['sigmasq_cross'] for t in self.template_params], - dtype=numpy.float32) - f['sigmasq_cross'] = template_sigmasq_cross[tid] + [t["sigmasq_cross"] for t in self.template_params], + dtype=numpy.float32, + ) + f["sigmasq_cross"] = template_sigmasq_cross[tid] # FIXME: I want to put something here, but I haven't yet # figured out what it should be. I think we would also # need information from the plus and cross correlation # (both real and imaginary(?)) to get this. - f['sigmasq'] = template_sigmasq_plus[tid] + f["sigmasq"] = template_sigmasq_plus[tid] except Exception: # Not precessing - f['sigmasq'] = self.events['sigmasq'] + f["sigmasq"] = self.events["sigmasq"] # Template durations should ideally be stored in the bank file. # At present, however, a few plotting/visualization codes # downstream in the offline search workflow rely on durations being # stored in the trigger files instead. - template_durations = [p['tmplt'].template_duration for p in - self.template_params] - f['template_duration'] = numpy.array(template_durations, - dtype=numpy.float32)[tid] + template_durations = [ + p["tmplt"].template_duration for p in self.template_params + ] + f["template_duration"] = numpy.array( + template_durations, dtype=numpy.float32 + )[tid] # FIXME: Can we get this value from the autochisq instance? cont_dof = self.opt.autochi_number_points @@ -544,94 +553,94 @@ def write_to_hdf(self, outname): cont_dof = cont_dof * 2 if self.opt.autochi_max_valued_dof: cont_dof = self.opt.autochi_max_valued_dof - f['cont_chisq_dof'] = numpy.repeat(cont_dof, len(self.events)) + f["cont_chisq_dof"] = numpy.repeat(cont_dof, len(self.events)) - if 'chisq_dof' in self.events.dtype.names: - f['chisq_dof'] = self.events['chisq_dof'] / 2 + 1 + if "chisq_dof" in self.events.dtype.names: + f["chisq_dof"] = self.events["chisq_dof"] / 2 + 1 else: - f['chisq_dof'] = numpy.zeros(len(self.events)) + f["chisq_dof"] = numpy.zeros(len(self.events)) - f['template_hash'] = th[tid] + f["template_hash"] = th[tid] - if 'sg_chisq' in self.events.dtype.names: - f['sg_chisq'] = self.events['sg_chisq'] + if "sg_chisq" in self.events.dtype.names: + f["sg_chisq"] = self.events["sg_chisq"] if self.opt.psdvar_segment is not None: - f['psd_var_val'] = self.events['psd_var_val'] + f["psd_var_val"] = self.events["psd_var_val"] if self.opt.trig_start_time: - f['search/start_time'] = numpy.array([self.opt.trig_start_time]) + f["search/start_time"] = numpy.array([self.opt.trig_start_time]) search_start_time = float(self.opt.trig_start_time) else: - f['search/start_time'] = numpy.array([self.opt.gps_start_time + - self.opt.segment_start_pad]) - search_start_time = float(self.opt.gps_start_time + - self.opt.segment_start_pad) + f["search/start_time"] = numpy.array( + [self.opt.gps_start_time + self.opt.segment_start_pad] + ) + search_start_time = float( + self.opt.gps_start_time + self.opt.segment_start_pad + ) if self.opt.trig_end_time: - f['search/end_time'] = numpy.array([self.opt.trig_end_time]) + f["search/end_time"] = numpy.array([self.opt.trig_end_time]) search_end_time = float(self.opt.trig_end_time) else: - f['search/end_time'] = numpy.array([self.opt.gps_end_time - - self.opt.segment_end_pad]) - search_end_time = float(self.opt.gps_end_time - - self.opt.segment_end_pad) + f["search/end_time"] = numpy.array( + [self.opt.gps_end_time - self.opt.segment_end_pad] + ) + search_end_time = float(self.opt.gps_end_time - self.opt.segment_end_pad) if self.write_performance: self.analysis_time = search_end_time - search_start_time time_ratio = float(self.analysis_time) / float(self.run_time) temps_per_core = float(self.ntemplates) / float(self.ncores) filters_per_core = float(self.nfilters) / float(self.ncores) - f['search/templates_per_core'] = \ - numpy.array([temps_per_core * time_ratio]) - f['search/filter_rate_per_core'] = \ - numpy.array([filters_per_core / float(self.run_time)]) - f['search/setup_time_fraction'] = \ - numpy.array([float(self.setup_time) / float(self.run_time)]) - f['search/run_time'] = numpy.array([float(self.run_time)]) - - if 'q_trans' in self.global_params: - qtrans = self.global_params['q_trans'] + f["search/templates_per_core"] = numpy.array([temps_per_core * time_ratio]) + f["search/filter_rate_per_core"] = numpy.array( + [filters_per_core / float(self.run_time)] + ) + f["search/setup_time_fraction"] = numpy.array( + [float(self.setup_time) / float(self.run_time)] + ) + f["search/run_time"] = numpy.array([float(self.run_time)]) + + if "q_trans" in self.global_params: + qtrans = self.global_params["q_trans"] for key in qtrans: - if key == 'qtiles': + if key == "qtiles": for seg in qtrans[key]: for q in qtrans[key][seg]: - f['qtransform/%s/%s/%s' % (key, seg, q)] = \ - qtrans[key][seg][q] - elif key == 'qplanes': + f["qtransform/%s/%s/%s" % (key, seg, q)] = qtrans[key][seg][ + q + ] + elif key == "qplanes": for seg in qtrans[key]: - f['qtransform/%s/%s' % (key, seg)] = qtrans[key][seg] + f["qtransform/%s/%s" % (key, seg)] = qtrans[key][seg] - if 'gating_info' in self.global_params: - gating_info = self.global_params['gating_info'] - for gate_type in ['file', 'auto']: + if "gating_info" in self.global_params: + gating_info = self.global_params["gating_info"] + for gate_type in ["file", "auto"]: if gate_type in gating_info: - f['gating/' + gate_type + '/time'] = \ - numpy.array([float(g[0]) for g in gating_info[gate_type]]) - f['gating/' + gate_type + '/width'] = \ - numpy.array([g[1] for g in gating_info[gate_type]]) - f['gating/' + gate_type + '/pad'] = \ - numpy.array([g[2] for g in gating_info[gate_type]]) + f["gating/" + gate_type + "/time"] = numpy.array( + [float(g[0]) for g in gating_info[gate_type]] + ) + f["gating/" + gate_type + "/width"] = numpy.array( + [g[1] for g in gating_info[gate_type]] + ) + f["gating/" + gate_type + "/pad"] = numpy.array( + [g[2] for g in gating_info[gate_type]] + ) f.f.close() class EventManagerMultiDetBase(EventManager): def __init__( - self, - opt, - ifos, - column, - column_types, - psd=None, - array_minsize=10000, - **kwargs + self, opt, ifos, column, column_types, psd=None, array_minsize=10000, **kwargs ): self.opt = opt self.ifos = ifos self.global_params = kwargs self.array_minsize = array_minsize if psd is not None: - self.global_params['psd'] = psd[ifos[0]] + self.global_params["psd"] = psd[ifos[0]] # The events array does not like holding the ifo as string, # so create a mapping dict and hold as an int @@ -641,14 +650,11 @@ def __init__( self.ifo_dict[ifo] = i self.ifo_reverse[i] = ifo - self.event_dtype = [('template_id', int), ('event_id', int)] + self.event_dtype = [("template_id", int), ("event_id", int)] for col, coltype in zip(column, column_types): self.event_dtype.append((col, coltype)) - self._events = numpy.zeros( - [self.array_minsize], - dtype=self.event_dtype - ) + self._events = numpy.zeros([self.array_minsize], dtype=self.event_dtype) self._events_size = 0 self.event_id_map = {} @@ -660,13 +666,12 @@ def __init__( self.write_performance = False for ifo in ifos: self.template_event_dict[ifo] = numpy.zeros( - [self.array_minsize], - dtype=self.event_dtype + [self.array_minsize], dtype=self.event_dtype ) self.template_event_size_dict[ifo] = 0 def add_template_events_to_ifo(self, ifo, columns, vectors): - """ Add a vector indexed """ + """Add a vector indexed""" # Just call through to the standard function self.template_events = self.template_event_dict[ifo] self.template_event_size = self.template_event_size_dict[ifo] @@ -677,7 +682,8 @@ def add_template_events_to_ifo(self, ifo, columns, vectors): self.template_event_size = None def write_gating_info_to_hdf(self, hf): - """Write per-detector gating information to an h5py file object. + """ + Write per-detector gating information to an h5py file object. The information is laid out according to the following groups and datasets: `//gating/{file, auto}/{time, width, pad}` where "file" and "auto" indicate respectively externally-provided gates and @@ -685,57 +691,66 @@ def write_gating_info_to_hdf(self, hf): indicate the gate center times, total durations and padding durations in seconds respectively. """ - if 'gating_info' not in self.global_params: + if "gating_info" not in self.global_params: return - gates = self.global_params['gating_info'] - for ifo, gate_type in itertools.product(self.ifos, ['file', 'auto']): + gates = self.global_params["gating_info"] + for ifo, gate_type in itertools.product(self.ifos, ["file", "auto"]): if gate_type not in gates[ifo]: continue - hf[f'{ifo}/gating/{gate_type}/time'] = numpy.array( + hf[f"{ifo}/gating/{gate_type}/time"] = numpy.array( [float(g[0]) for g in gates[ifo][gate_type]] ) - hf[f'{ifo}/gating/{gate_type}/width'] = numpy.array( + hf[f"{ifo}/gating/{gate_type}/width"] = numpy.array( [g[1] for g in gates[ifo][gate_type]] ) - hf[f'{ifo}/gating/{gate_type}/pad'] = numpy.array( + hf[f"{ifo}/gating/{gate_type}/pad"] = numpy.array( [g[2] for g in gates[ifo][gate_type]] ) class EventManagerCoherent(EventManagerMultiDetBase): - def __init__(self, opt, ifos, column, column_types, network_column, - network_column_types, segments, time_slides, - psd=None, **kwargs): - super(EventManagerCoherent, self).__init__( - opt, ifos, column, column_types, psd=None, **kwargs) - self.network_event_dtype = \ - [(ifo + '_event_id', int) for ifo in self.ifos] - self.network_event_dtype.append(('template_id', int)) - self.network_event_dtype.append(('event_id', int)) + def __init__( + self, + opt, + ifos, + column, + column_types, + network_column, + network_column_types, + segments, + time_slides, + psd=None, + **kwargs, + ): + super().__init__( + opt, ifos, column, column_types, psd=None, **kwargs + ) + self.network_event_dtype = [(ifo + "_event_id", int) for ifo in self.ifos] + self.network_event_dtype.append(("template_id", int)) + self.network_event_dtype.append(("event_id", int)) for col, coltype in zip(network_column, network_column_types): self.network_event_dtype.append((col, coltype)) self.network_events = numpy.zeros( - [self.array_minsize], - dtype=self.network_event_dtype + [self.array_minsize], dtype=self.network_event_dtype ) self.network_events_size = 0 self.event_index = {} for ifo in self.ifos: self.event_index[ifo] = 0 - self.event_index['network'] = 0 - self.template_event_dict['network'] = numpy.zeros( - [self.array_minsize], - dtype=self.network_event_dtype + self.event_index["network"] = 0 + self.template_event_dict["network"] = numpy.zeros( + [self.array_minsize], dtype=self.network_event_dtype ) - self.template_event_size_dict['network'] = 0 + self.template_event_size_dict["network"] = 0 self.segments = segments self.time_slides = time_slides - def cluster_template_network_events(self, tcolumn, column, window_size, - slide=0): - """ Cluster the internal events over the named column + def cluster_template_network_events(self, tcolumn, column, window_size, slide=0): + """ + Cluster the internal events over the named column + Parameters - ---------------- + ---------- tcolumn Indicates which column contains the time. column @@ -744,12 +759,11 @@ def cluster_template_network_events(self, tcolumn, column, window_size, The size of the window. slide Default is 0. + """ - net_event_size = self.template_event_size_dict['network'] - net_event_dict = self.template_event_dict['network'] - slide_indices = ( - net_event_dict[:net_event_size]['slide_id'] == slide - ) + net_event_size = self.template_event_size_dict["network"] + net_event_dict = self.template_event_dict["network"] + slide_indices = net_event_dict[:net_event_size]["slide_id"] == slide cvec = net_event_dict[:net_event_size][column][slide_indices] tvec = net_event_dict[:net_event_size][tcolumn][slide_indices] if not window_size == 0: @@ -766,70 +780,69 @@ def cluster_template_network_events(self, tcolumn, column, window_size, # if a slide_indices = 0 if any(~slide_indices): - indices = numpy.concatenate(( + indices = numpy.concatenate( + ( numpy.flatnonzero(~slide_indices), - numpy.flatnonzero(slide_indices)[indices])) + numpy.flatnonzero(slide_indices)[indices], + ) + ) indices.sort() # get value of key for where you have indicies for key in self.template_event_dict: curr_ted = self.template_event_dict[key] - curr_ted[:len(indices)] = curr_ted[indices] + curr_ted[: len(indices)] = curr_ted[indices] self.template_event_size_dict[key] = len(indices) def add_template_network_events(self, columns, vectors): - """ Add a vector indexed """ + """Add a vector indexed""" # initialize with zeros - since vectors can be None, look for the # longest one that isn't new_events = numpy.zeros( max([len(v) for v in vectors if v is not None]), - dtype=self.network_event_dtype + dtype=self.network_event_dtype, ) # they shouldn't all be None assert new_events is not None - new_events['template_id'] = self.template_index + new_events["template_id"] = self.template_index for c, v in zip(columns, vectors): if v is not None: if isinstance(v, Array): new_events[c] = v.numpy() else: new_events[c] = v - new_size = self.template_event_size+len(new_events) + new_size = self.template_event_size + len(new_events) if new_size > len(self.template_events): - self.template_events.resize( - new_size + self.array_minsize, - refcheck=False - ) - self.template_events[self.template_event_size:new_size] = new_events + self.template_events.resize(new_size + self.array_minsize, refcheck=False) + self.template_events[self.template_event_size : new_size] = new_events self.template_event_size += len(new_events) def add_template_events_to_network(self, columns, vectors): - """ Add a vector indexed """ + """Add a vector indexed""" # Just call through to the standard function - self.template_events = self.template_event_dict['network'] - self.template_event_size = self.template_event_size_dict['network'] + self.template_events = self.template_event_dict["network"] + self.template_event_size = self.template_event_size_dict["network"] self.add_template_network_events(columns, vectors) - self.template_event_dict['network'] = self.template_events - self.template_event_size_dict['network'] = self.template_event_size + self.template_event_dict["network"] = self.template_events + self.template_event_size_dict["network"] = self.template_event_size self.template_events = None self.template_event_size = None def write_to_hdf(self, outname): - sort_order = self.events.argsort(order='template_id') - self._events[:self._events_size] = self.events[sort_order] - th = numpy.array( - [p['tmplt'].template_hash for p in self.template_params]) + sort_order = self.events.argsort(order="template_id") + self._events[: self._events_size] = self.events[sort_order] + th = numpy.array([p["tmplt"].template_hash for p in self.template_params]) f = H5FileSyntSugar(outname) self.write_gating_info_to_hdf(f) # Output network stuff - f.prefix = 'network' - network_events = self.network_events[:self.network_events_size] + f.prefix = "network" + network_events = self.network_events[: self.network_events_size] for col in network_events.dtype.names: - if col == 'time_index': - f['end_time_gc'] = ( + if col == "time_index": + f["end_time_gc"] = ( network_events[col] / float(self.opt.sample_rate[self.ifos[0].lower()]) + self.opt.gps_start_time[self.ifos[0].lower()] - ) + ) else: f[col] = network_events[col] starts = [] @@ -837,57 +850,61 @@ def write_to_hdf(self, outname): for seg in self.segments[self.ifos[0]]: starts.append(int(seg.start_time)) ends.append(int(seg.end_time)) - f['search/segments/start_times'] = starts - f['search/segments/end_times'] = ends + f["search/segments/start_times"] = starts + f["search/segments/end_times"] = ends # Individual ifo stuff for i, ifo in enumerate(self.ifos): - tid = self.events['template_id'][self.events['ifo'] == i] + tid = self.events["template_id"][self.events["ifo"] == i] f.prefix = ifo - ifo_events = numpy.array([e for e in self.events - if e['ifo'] == self.ifo_dict[ifo]], dtype=self.event_dtype) + ifo_events = numpy.array( + [e for e in self.events if e["ifo"] == self.ifo_dict[ifo]], + dtype=self.event_dtype, + ) if len(ifo_events): - f['snr'] = abs(ifo_events['snr']) - f['event_id'] = ifo_events['event_id'] + f["snr"] = abs(ifo_events["snr"]) + f["event_id"] = ifo_events["event_id"] try: # Precessing - f['u_vals'] = ifo_events['u_vals'] - f['coa_phase'] = ifo_events['coa_phase'] - f['hplus_cross_corr'] = ifo_events['hplus_cross_corr'] + f["u_vals"] = ifo_events["u_vals"] + f["coa_phase"] = ifo_events["coa_phase"] + f["hplus_cross_corr"] = ifo_events["hplus_cross_corr"] except Exception: - f['coa_phase'] = numpy.angle(ifo_events['snr']) - f['chisq'] = ifo_events['chisq'] - f['bank_chisq'] = ifo_events['bank_chisq'] - f['bank_chisq_dof'] = ifo_events['bank_chisq_dof'] - f['auto_chisq'] = ifo_events['auto_chisq'] - f['auto_chisq_dof'] = ifo_events['auto_chisq_dof'] - f['end_time'] = ifo_events['time_index'] / \ - float(self.opt.sample_rate[ifo]) + \ - self.opt.gps_start_time[ifo] - f['time_index'] = ifo_events['time_index'] - f['slide_id'] = ifo_events['slide_id'] + f["coa_phase"] = numpy.angle(ifo_events["snr"]) + f["chisq"] = ifo_events["chisq"] + f["bank_chisq"] = ifo_events["bank_chisq"] + f["bank_chisq_dof"] = ifo_events["bank_chisq_dof"] + f["auto_chisq"] = ifo_events["auto_chisq"] + f["auto_chisq_dof"] = ifo_events["auto_chisq_dof"] + f["end_time"] = ( + ifo_events["time_index"] / float(self.opt.sample_rate[ifo]) + + self.opt.gps_start_time[ifo] + ) + f["time_index"] = ifo_events["time_index"] + f["slide_id"] = ifo_events["slide_id"] try: # Precessing template_sigmasq_plus = numpy.array( - [t['sigmasq_plus'] for t in self.template_params], - dtype=numpy.float32 + [t["sigmasq_plus"] for t in self.template_params], + dtype=numpy.float32, ) - f['sigmasq_plus'] = template_sigmasq_plus[tid] + f["sigmasq_plus"] = template_sigmasq_plus[tid] template_sigmasq_cross = numpy.array( - [t['sigmasq_cross'] for t in self.template_params], - dtype=numpy.float32 + [t["sigmasq_cross"] for t in self.template_params], + dtype=numpy.float32, ) - f['sigmasq_cross'] = template_sigmasq_cross[tid] + f["sigmasq_cross"] = template_sigmasq_cross[tid] # FIXME: I want to put something here, but I haven't yet # figured out what it should be. I think we would also # need information from the plus and cross correlation # (both real and imaginary(?)) to get this. - f['sigmasq'] = template_sigmasq_plus[tid] + f["sigmasq"] = template_sigmasq_plus[tid] except Exception: # Not precessing template_sigmasq = numpy.array( - [t['sigmasq'][ifo] for t in self.template_params], - dtype=numpy.float32) - f['sigmasq'] = template_sigmasq[tid] + [t["sigmasq"][ifo] for t in self.template_params], + dtype=numpy.float32, + ) + f["sigmasq"] = template_sigmasq[tid] # FIXME: Can we get this value from the autochisq instance? # cont_dof = self.opt.autochi_number_points @@ -899,45 +916,54 @@ def write_to_hdf(self, outname): # cont_dof = self.opt.autochi_max_valued_dof # f['cont_chisq_dof'] = numpy.repeat(cont_dof, len(ifo_events)) - if 'chisq_dof' in ifo_events.dtype.names: - f['chisq_dof'] = ifo_events['chisq_dof'] / 2 + 1 + if "chisq_dof" in ifo_events.dtype.names: + f["chisq_dof"] = ifo_events["chisq_dof"] / 2 + 1 else: - f['chisq_dof'] = numpy.zeros(len(ifo_events)) + f["chisq_dof"] = numpy.zeros(len(ifo_events)) - f['template_hash'] = th[tid] - f['search/time_slides'] = numpy.array(self.time_slides[ifo]) + f["template_hash"] = th[tid] + f["search/time_slides"] = numpy.array(self.time_slides[ifo]) if self.opt.trig_start_time: - f['search/start_time'] = numpy.array([ - self.opt.trig_start_time[ifo]], dtype=numpy.int32) + f["search/start_time"] = numpy.array( + [self.opt.trig_start_time[ifo]], dtype=numpy.int32 + ) search_start_time = float(self.opt.trig_start_time[ifo]) else: - f['search/start_time'] = numpy.array([ - self.opt.gps_start_time[ifo] + - self.opt.segment_start_pad[ifo]], dtype=numpy.int32) - search_start_time = float(self.opt.gps_start_time[ifo] + - self.opt.segment_start_pad[ifo]) + f["search/start_time"] = numpy.array( + [self.opt.gps_start_time[ifo] + self.opt.segment_start_pad[ifo]], + dtype=numpy.int32, + ) + search_start_time = float( + self.opt.gps_start_time[ifo] + self.opt.segment_start_pad[ifo] + ) if self.opt.trig_end_time: - f['search/end_time'] = numpy.array([ - self.opt.trig_end_time[ifo]], dtype=numpy.int32) + f["search/end_time"] = numpy.array( + [self.opt.trig_end_time[ifo]], dtype=numpy.int32 + ) search_end_time = float(self.opt.trig_end_time[ifo]) else: - f['search/end_time'] = numpy.array( - [self.opt.gps_end_time[ifo] - - self.opt.segment_end_pad[ifo]], dtype=numpy.int32) - search_end_time = float(self.opt.gps_end_time[ifo] - - self.opt.segment_end_pad[ifo]) + f["search/end_time"] = numpy.array( + [self.opt.gps_end_time[ifo] - self.opt.segment_end_pad[ifo]], + dtype=numpy.int32, + ) + search_end_time = float( + self.opt.gps_end_time[ifo] - self.opt.segment_end_pad[ifo] + ) if self.write_performance: self.analysis_time = search_end_time - search_start_time time_ratio = float(self.analysis_time) / float(self.run_time) temps_per_core = float(self.ntemplates) / float(self.ncores) filters_per_core = float(self.nfilters) / float(self.ncores) - f['search/templates_per_core'] = \ - numpy.array([temps_per_core * time_ratio]) - f['search/filter_rate_per_core'] = \ - numpy.array([filters_per_core / float(self.run_time)]) - f['search/setup_time_fraction'] = \ - numpy.array([float(self.setup_time) / float(self.run_time)]) + f["search/templates_per_core"] = numpy.array( + [temps_per_core * time_ratio] + ) + f["search/filter_rate_per_core"] = numpy.array( + [filters_per_core / float(self.run_time)] + ) + f["search/setup_time_fraction"] = numpy.array( + [float(self.setup_time) / float(self.run_time)] + ) def finalize_template_events(self): # Check that none of the template events have the same time index as an @@ -951,97 +977,112 @@ def finalize_template_events(self): new_template_event_mask = {} existing_template_event_mask = {} for i, ifo in enumerate(self.ifos): - ifo_events = numpy.where(self.events['ifo'] == i) - existing_times[ifo] = self.events['time_index'][ifo_events] + ifo_events = numpy.where(self.events["ifo"] == i) + existing_times[ifo] = self.events["time_index"][ifo_events] curr_ev_dict = self.template_event_dict[ifo] curr_ev_size = self.template_event_size_dict[ifo] - new_times[ifo] = curr_ev_dict[:curr_ev_size]['time_index'] - existing_template_id[ifo] = self.events['template_id'][ifo_events] - new_template_id[ifo] = curr_ev_dict[:curr_ev_size]['template_id'] + new_times[ifo] = curr_ev_dict[:curr_ev_size]["time_index"] + existing_template_id[ifo] = self.events["template_id"][ifo_events] + new_template_id[ifo] = curr_ev_dict[:curr_ev_size]["template_id"] # This is true for each existing event that has the same time index # and template id as a template trigger. existing_events_mask[ifo] = numpy.argwhere( numpy.logical_and( numpy.isin(existing_times[ifo], new_times[ifo]), - numpy.isin(existing_template_id[ifo], new_template_id[ifo]) - )).reshape(-1,) + numpy.isin(existing_template_id[ifo], new_template_id[ifo]), + ) + ).reshape( + -1, + ) # This is true for each template event that has either a new # trigger time or a new template id. new_template_event_mask[ifo] = numpy.argwhere( numpy.logical_or( - ~numpy.isin(new_times[ifo], existing_times[ifo]), - ~numpy.isin(new_template_id[ifo], existing_template_id[ifo]) - )).reshape(-1,) + ~numpy.isin(new_times[ifo], existing_times[ifo]), + ~numpy.isin(new_template_id[ifo], existing_template_id[ifo]), + ) + ).reshape( + -1, + ) # This is true for each template event that has the same time index # and template id as an exisitng event trigger. existing_template_event_mask[ifo] = numpy.argwhere( numpy.logical_and( numpy.isin(new_times[ifo], existing_times[ifo]), - numpy.isin(new_template_id[ifo], existing_template_id[ifo]) - )).reshape(-1,) + numpy.isin(new_template_id[ifo], existing_template_id[ifo]), + ) + ).reshape( + -1, + ) # Set ids (These show how each trigger in the single ifo trigger # list correspond to the network triggers) num_events = len(new_template_event_mask[ifo]) - new_event_ids = numpy.arange(self.event_index[ifo], - self.event_index[ifo] + num_events) + new_event_ids = numpy.arange( + self.event_index[ifo], self.event_index[ifo] + num_events + ) # Every template event that corresponds to a new trigger gets a new # id. Triggers that have been found before are not saved. - curr_ev_dict[:curr_ev_size]['event_id'][new_template_event_mask[ifo]] = new_event_ids - net_ev_dict = self.template_event_dict['network'] - net_ev_size = self.template_event_size_dict['network'] - net_ev_dict[:net_ev_size][ifo + '_event_id'][new_template_event_mask[ifo]] = new_event_ids + curr_ev_dict[:curr_ev_size]["event_id"][new_template_event_mask[ifo]] = ( + new_event_ids + ) + net_ev_dict = self.template_event_dict["network"] + net_ev_size = self.template_event_size_dict["network"] + net_ev_dict[:net_ev_size][ifo + "_event_id"][ + new_template_event_mask[ifo] + ] = new_event_ids # Template events that have been found before get the event id of # the first time they were found. - net_ev_dict[:net_ev_size][ifo + '_event_id'][ - existing_template_event_mask[ifo]] = \ - self.events[self.events['ifo'] == i][ - existing_events_mask[ifo]]['event_id'] + net_ev_dict[:net_ev_size][ifo + "_event_id"][ + existing_template_event_mask[ifo] + ] = self.events[self.events["ifo"] == i][existing_events_mask[ifo]][ + "event_id" + ] self.event_index[ifo] = self.event_index[ifo] + num_events # Add the network event ids for the events with this template. - num_events = self.template_event_size_dict['network'] - new_event_ids = numpy.arange(self.event_index['network'], - self.event_index['network'] + num_events) - self.event_index['network'] = self.event_index['network'] + num_events - self.template_event_dict['network'][:num_events]['event_id'] = new_event_ids + num_events = self.template_event_size_dict["network"] + new_event_ids = numpy.arange( + self.event_index["network"], self.event_index["network"] + num_events + ) + self.event_index["network"] = self.event_index["network"] + num_events + self.template_event_dict["network"][:num_events]["event_id"] = new_event_ids # Move template events for each ifo to the events list for ifo in self.ifos: new_event_size = len(new_template_event_mask[ifo]) + self._events_size if new_event_size > len(self._events): - new_chunk_size = (new_event_size//self.array_minsize + 1) - self._events.resize( - new_chunk_size * self.array_minsize, - refcheck=False - ) - self._events[self._events_size:new_event_size] = ( - self.template_event_dict[ifo][:self.template_event_size_dict[ifo]][new_template_event_mask[ifo]] - ) + new_chunk_size = new_event_size // self.array_minsize + 1 + self._events.resize(new_chunk_size * self.array_minsize, refcheck=False) + self._events[self._events_size : new_event_size] = self.template_event_dict[ + ifo + ][: self.template_event_size_dict[ifo]][new_template_event_mask[ifo]] self._events_size = new_event_size self.template_event_size_dict[ifo] = 0 # Move the template events for the network to the network events list - new_network_size = self.template_event_size_dict['network'] + self.network_events_size + new_network_size = ( + self.template_event_size_dict["network"] + self.network_events_size + ) while new_network_size > len(self.network_events): self.network_events.resize( - len(self.network_events) + self.array_minsize, - refcheck=False + len(self.network_events) + self.array_minsize, refcheck=False ) - self.network_events[self.network_events_size:new_network_size] = ( - self.template_event_dict['network'][:self.template_event_size_dict['network']] + self.network_events[self.network_events_size : new_network_size] = ( + self.template_event_dict["network"][ + : self.template_event_size_dict["network"] + ] ) self.network_events_size = new_network_size - self.template_event_size_dict['network'] = 0 + self.template_event_size_dict["network"] = 0 class EventManagerMultiDet(EventManagerMultiDetBase): def __init__(self, opt, ifos, column, column_types, psd=None, **kwargs): - super(EventManagerMultiDet, self).__init__( - opt, ifos, column, column_types, psd=None, **kwargs) + super().__init__( + opt, ifos, column, column_types, psd=None, **kwargs + ) self.event_index = 0 - def cluster_template_events_single_ifo( - self, tcolumn, column, window_size, ifo): - """ Cluster the internal events over the named column - """ + def cluster_template_events_single_ifo(self, tcolumn, column, window_size, ifo): + """Cluster the internal events over the named column""" # Just call through to the standard function self.template_events = self.template_event_dict[ifo] self.template_event_size = self.template_event_size_dict[ifo] @@ -1051,17 +1092,17 @@ def cluster_template_events_single_ifo( self.template_events = None self.template_event_size = None - def finalize_template_events(self, perform_coincidence=True, - coinc_window=0.0): + def finalize_template_events(self, perform_coincidence=True, coinc_window=0.0): # Set ids for ifo in self.ifos: temp_ev_dict = self.template_event_dict[ifo] temp_ev_size = self.template_event_size_dict[ifo] num_events = len(temp_ev_dict[:temp_ev_size]) - new_event_ids = numpy.arange(self.event_index, - self.event_index+num_events) - temp_ev_dict[:temp_ev_size]['event_id'] = new_event_ids - self.event_index = self.event_index+num_events + new_event_ids = numpy.arange( + self.event_index, self.event_index + num_events + ) + temp_ev_dict[:temp_ev_size]["event_id"] = new_event_ids + self.event_index = self.event_index + num_events if perform_coincidence: if not len(self.ifos) == 2: @@ -1073,19 +1114,26 @@ def finalize_template_events(self, perform_coincidence=True, curr_tev_dict2 = self.template_event_dict[ifo2] curr_tev_size1 = self.template_event_size_dict[ifo1] curr_tev_size2 = self.template_event_size_dict[ifo2] - end_times1 = curr_tev_dict1[:curr_tev_size1]['time_index'] /\ - float(self.opt.sample_rate[ifo1]) + self.opt.gps_start_time[ifo1] - end_times2 = curr_tev_dict2[:curr_tev_size2]['time_index'] /\ - float(self.opt.sample_rate[ifo2]) + self.opt.gps_start_time[ifo2] + end_times1 = ( + curr_tev_dict1[:curr_tev_size1]["time_index"] + / float(self.opt.sample_rate[ifo1]) + + self.opt.gps_start_time[ifo1] + ) + end_times2 = ( + curr_tev_dict2[:curr_tev_size2]["time_index"] + / float(self.opt.sample_rate[ifo2]) + + self.opt.gps_start_time[ifo2] + ) light_travel_time = Detector(ifo1).light_travel_time_to_detector( - Detector(ifo2)) + Detector(ifo2) + ) coinc_window = coinc_window + light_travel_time # FIXME: Remove!!! coinc_window = 2.0 if len(end_times1) and len(end_times2): - idx_list1, idx_list2, _ = \ - coinc.time_coincidence(end_times1, end_times2, - coinc_window) + idx_list1, idx_list2, _ = coinc.time_coincidence( + end_times1, end_times2, coinc_window + ) if len(idx_list1): for idx1, idx2 in zip(idx_list1, idx_list2): event1 = curr_tev_dict1[:curr_tev_size1][idx1] @@ -1094,65 +1142,68 @@ def finalize_template_events(self, perform_coincidence=True, for ifo in self.ifos: new_event_size = self.template_event_size_dict[ifo] + self._events_size if new_event_size > len(self._events): - new_chunk_size = (new_event_size//self.array_minsize + 1) - self._events.resize( - new_chunk_size * self.array_minsize, - refcheck=False - ) - self._events[self._events_size:new_event_size] = ( - self.template_event_dict[ifo][:self.template_event_size_dict[ifo]] - ) + new_chunk_size = new_event_size // self.array_minsize + 1 + self._events.resize(new_chunk_size * self.array_minsize, refcheck=False) + self._events[self._events_size : new_event_size] = self.template_event_dict[ + ifo + ][: self.template_event_size_dict[ifo]] self.template_event_size_dict[ifo] = 0 self._events_size = new_event_size def write_to_hdf(self, outname): - sort_order = self.events.argsort(order='template_id') - self._events[:self._events_size] = self.events[sort_order] - th = numpy.array([p['tmplt'].template_hash for p in - self.template_params]) - tid = self.events['template_id'] + sort_order = self.events.argsort(order="template_id") + self._events[: self._events_size] = self.events[sort_order] + th = numpy.array([p["tmplt"].template_hash for p in self.template_params]) + tid = self.events["template_id"] f = H5FileSyntSugar(outname) self.write_gating_info_to_hdf(f) for ifo in self.ifos: f.prefix = ifo - ifo_events = numpy.array([e for e in self.events if - e['ifo'] == self.ifo_dict[ifo]], - dtype=self.event_dtype) + ifo_events = numpy.array( + [e for e in self.events if e["ifo"] == self.ifo_dict[ifo]], + dtype=self.event_dtype, + ) if len(ifo_events): - f['snr'] = abs(ifo_events['snr']) + f["snr"] = abs(ifo_events["snr"]) try: # Precessing - f['u_vals'] = ifo_events['u_vals'] - f['coa_phase'] = ifo_events['coa_phase'] - f['hplus_cross_corr'] = ifo_events['hplus_cross_corr'] + f["u_vals"] = ifo_events["u_vals"] + f["coa_phase"] = ifo_events["coa_phase"] + f["hplus_cross_corr"] = ifo_events["hplus_cross_corr"] except Exception: - f['coa_phase'] = numpy.angle(ifo_events['snr']) - f['chisq'] = ifo_events['chisq'] - f['bank_chisq'] = ifo_events['bank_chisq'] - f['bank_chisq_dof'] = ifo_events['bank_chisq_dof'] - f['cont_chisq'] = ifo_events['cont_chisq'] - f['end_time'] = ifo_events['time_index'] / \ - float(self.opt.sample_rate[ifo]) + \ - self.opt.gps_start_time[ifo] + f["coa_phase"] = numpy.angle(ifo_events["snr"]) + f["chisq"] = ifo_events["chisq"] + f["bank_chisq"] = ifo_events["bank_chisq"] + f["bank_chisq_dof"] = ifo_events["bank_chisq_dof"] + f["cont_chisq"] = ifo_events["cont_chisq"] + f["end_time"] = ( + ifo_events["time_index"] / float(self.opt.sample_rate[ifo]) + + self.opt.gps_start_time[ifo] + ) try: # Precessing - template_sigmasq_plus = numpy.array([t['sigmasq_plus'] for - t in self.template_params], dtype=numpy.float32) - f['sigmasq_plus'] = template_sigmasq_plus[tid] - template_sigmasq_cross = numpy.array([t['sigmasq_cross'] - for t in self.template_params], dtype=numpy.float32) - f['sigmasq_cross'] = template_sigmasq_cross[tid] + template_sigmasq_plus = numpy.array( + [t["sigmasq_plus"] for t in self.template_params], + dtype=numpy.float32, + ) + f["sigmasq_plus"] = template_sigmasq_plus[tid] + template_sigmasq_cross = numpy.array( + [t["sigmasq_cross"] for t in self.template_params], + dtype=numpy.float32, + ) + f["sigmasq_cross"] = template_sigmasq_cross[tid] # FIXME: I want to put something here, but I haven't yet # figured out what it should be. I think we would also # need information from the plus and cross correlation # (both real and imaginary(?)) to get this. - f['sigmasq'] = template_sigmasq_plus[tid] + f["sigmasq"] = template_sigmasq_plus[tid] except Exception: # Not precessing - template_sigmasq = numpy.array([t['sigmasq'][ifo] for t in - self.template_params], - dtype=numpy.float32) - f['sigmasq'] = template_sigmasq[tid] + template_sigmasq = numpy.array( + [t["sigmasq"][ifo] for t in self.template_params], + dtype=numpy.float32, + ) + f["sigmasq"] = template_sigmasq[tid] # FIXME: Can we get this value from the autochisq instance? cont_dof = self.opt.autochi_number_points @@ -1162,53 +1213,70 @@ def write_to_hdf(self, outname): # cont_dof = cont_dof * 2 # if self.opt.autochi_max_valued_dof: # cont_dof = self.opt.autochi_max_valued_dof - f['cont_chisq_dof'] = numpy.repeat(cont_dof, len(ifo_events)) + f["cont_chisq_dof"] = numpy.repeat(cont_dof, len(ifo_events)) - if 'chisq_dof' in ifo_events.dtype.names: - f['chisq_dof'] = ifo_events['chisq_dof'] / 2 + 1 + if "chisq_dof" in ifo_events.dtype.names: + f["chisq_dof"] = ifo_events["chisq_dof"] / 2 + 1 else: - f['chisq_dof'] = numpy.zeros(len(ifo_events)) + f["chisq_dof"] = numpy.zeros(len(ifo_events)) - f['template_hash'] = th[tid] + f["template_hash"] = th[tid] if self.opt.psdvar_segment is not None: - f['psd_var_val'] = ifo_events['psd_var_val'] + f["psd_var_val"] = ifo_events["psd_var_val"] if self.opt.trig_start_time: - f['search/start_time'] = numpy.array( - [self.opt.trig_start_time[ifo]], dtype=numpy.int32) + f["search/start_time"] = numpy.array( + [self.opt.trig_start_time[ifo]], dtype=numpy.int32 + ) search_start_time = float(self.opt.trig_start_time[ifo]) else: - f['search/start_time'] = numpy.array( - [self.opt.gps_start_time[ifo] + - self.opt.segment_start_pad[ifo]], dtype=numpy.int32) - search_start_time = float(self.opt.gps_start_time[ifo] + - self.opt.segment_start_pad[ifo]) + f["search/start_time"] = numpy.array( + [self.opt.gps_start_time[ifo] + self.opt.segment_start_pad[ifo]], + dtype=numpy.int32, + ) + search_start_time = float( + self.opt.gps_start_time[ifo] + self.opt.segment_start_pad[ifo] + ) if self.opt.trig_end_time: - f['search/end_time'] = numpy.array( - [self.opt.trig_end_time[ifo]], dtype=numpy.int32) + f["search/end_time"] = numpy.array( + [self.opt.trig_end_time[ifo]], dtype=numpy.int32 + ) search_end_time = float(self.opt.trig_end_time[ifo]) else: - f['search/end_time'] = numpy.array( - [self.opt.gps_end_time[ifo] - - self.opt.segment_end_pad[ifo]], dtype=numpy.int32) - search_end_time = float(self.opt.gps_end_time[ifo] - - self.opt.segment_end_pad[ifo]) + f["search/end_time"] = numpy.array( + [self.opt.gps_end_time[ifo] - self.opt.segment_end_pad[ifo]], + dtype=numpy.int32, + ) + search_end_time = float( + self.opt.gps_end_time[ifo] - self.opt.segment_end_pad[ifo] + ) if self.write_performance: self.analysis_time = search_end_time - search_start_time time_ratio = float(self.analysis_time) / float(self.run_time) temps_per_core = float(self.ntemplates) / float(self.ncores) filters_per_core = float(self.nfilters) / float(self.ncores) - f['search/templates_per_core'] = \ - numpy.array([temps_per_core * time_ratio]) - f['search/filter_rate_per_core'] = \ - numpy.array([filters_per_core / float(self.run_time)]) - f['search/setup_time_fraction'] = \ - numpy.array([float(self.setup_time) / float(self.run_time)]) - - -__all__ = ['threshold_and_cluster', 'findchirp_cluster_over_window', - 'threshold', 'cluster_reduce', 'ThresholdCluster', - 'threshold_real_numpy', 'threshold_only', - 'EventManager', 'EventManagerMultiDet', 'EventManagerCoherent'] + f["search/templates_per_core"] = numpy.array( + [temps_per_core * time_ratio] + ) + f["search/filter_rate_per_core"] = numpy.array( + [filters_per_core / float(self.run_time)] + ) + f["search/setup_time_fraction"] = numpy.array( + [float(self.setup_time) / float(self.run_time)] + ) + + +__all__ = [ + "EventManager", + "EventManagerCoherent", + "EventManagerMultiDet", + "ThresholdCluster", + "cluster_reduce", + "findchirp_cluster_over_window", + "threshold", + "threshold_and_cluster", + "threshold_only", + "threshold_real_numpy", +] diff --git a/pycbc/events/ranking.py b/pycbc/events/ranking.py index 2d408c340e8..a1f1ccdaaad 100644 --- a/pycbc/events/ranking.py +++ b/pycbc/events/ranking.py @@ -1,30 +1,30 @@ -""" This module contains functions for calculating single-ifo ranking +""" +This module contains functions for calculating single-ifo ranking statistic values """ + import logging + import numpy -logger = logging.getLogger('pycbc.events.ranking') +logger = logging.getLogger("pycbc.events.ranking") -def effsnr(snr, reduced_x2, fac=250., - **kwargs): # pylint:disable=unused-argument - """Calculate the effective SNR statistic. See (S5y1 paper) for definition. - """ +def effsnr(snr, reduced_x2, fac=250.0, **kwargs): # pylint:disable=unused-argument + """Calculate the effective SNR statistic. See (S5y1 paper) for definition.""" snr = numpy.array(snr, ndmin=1, dtype=numpy.float64) rchisq = numpy.array(reduced_x2, ndmin=1, dtype=numpy.float64) - esnr = snr / (1 + snr ** 2 / fac) ** 0.25 / rchisq ** 0.25 + esnr = snr / (1 + snr**2 / fac) ** 0.25 / rchisq**0.25 # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return esnr - else: - return esnr[0] + return esnr[0] -def newsnr(snr, reduced_x2, q=6., n=2., - **kwargs): # pylint:disable=unused-argument - """Calculate the re-weighted SNR statistic ('newSNR') from given SNR and +def newsnr(snr, reduced_x2, q=6.0, n=2.0, **kwargs): # pylint:disable=unused-argument + """ + Calculate the re-weighted SNR statistic ('newSNR') from given SNR and reduced chi-squared values. See http://arxiv.org/abs/1208.3491 for definition. Previous implementation in glue/ligolw/lsctables.py """ @@ -32,68 +32,63 @@ def newsnr(snr, reduced_x2, q=6., n=2., reduced_x2 = numpy.array(reduced_x2, ndmin=1, dtype=numpy.float64) # newsnr is only different from snr if reduced chisq > 1 - ind = numpy.where(reduced_x2 > 1.)[0] - nsnr[ind] *= (0.5 * (1. + reduced_x2[ind] ** (q/n))) ** (-1./q) + ind = numpy.where(reduced_x2 > 1.0)[0] + nsnr[ind] *= (0.5 * (1.0 + reduced_x2[ind] ** (q / n))) ** (-1.0 / q) # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return nsnr - else: - return nsnr[0] + return nsnr[0] def newsnr_sgveto(snr, brchisq, sgchisq, **kwargs): - """ Combined SNR derived from NewSNR and Sine-Gaussian Chisq""" - nsnr = numpy.array( - newsnr( - snr, - brchisq, - **kwargs), - ndmin=1) + """Combined SNR derived from NewSNR and Sine-Gaussian Chisq""" + nsnr = numpy.array(newsnr(snr, brchisq, **kwargs), ndmin=1) sgchisq = numpy.array(sgchisq, ndmin=1) t = numpy.array(sgchisq > 4, ndmin=1) if len(t): nsnr[t] = nsnr[t] / (sgchisq[t] / 4.0) ** 0.5 # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return nsnr - else: - return nsnr[0] + return nsnr[0] -def newsnr_sgveto_psdvar(snr, brchisq, sgchisq, psd_var_val, - min_expected_psdvar=0.65, - **kwargs): - """ Combined SNR derived from SNR, reduced Allen chisq, sine-Gaussian chisq and - PSD variation statistic""" +def newsnr_sgveto_psdvar( + snr, brchisq, sgchisq, psd_var_val, min_expected_psdvar=0.65, **kwargs +): + """ + Combined SNR derived from SNR, reduced Allen chisq, sine-Gaussian chisq and + PSD variation statistic + """ # If PSD var is lower than the 'minimum usually expected value' stop this # being used in the statistic. This low value might arise because a # significant fraction of the "short" PSD period was gated (for instance). psd_var_val = numpy.array(psd_var_val, copy=True) - psd_var_val[psd_var_val < min_expected_psdvar] = 1. - scaled_snr = snr * (psd_var_val ** -0.5) - scaled_brchisq = brchisq * (psd_var_val ** -1.) - nsnr = newsnr_sgveto( - scaled_snr, - scaled_brchisq, - sgchisq, - **kwargs - ) + psd_var_val[psd_var_val < min_expected_psdvar] = 1.0 + scaled_snr = snr * (psd_var_val**-0.5) + scaled_brchisq = brchisq * (psd_var_val**-1.0) + nsnr = newsnr_sgveto(scaled_snr, scaled_brchisq, sgchisq, **kwargs) # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return nsnr - else: - return nsnr[0] - - -def newsnr_sgveto_psdvar_threshold(snr, brchisq, sgchisq, psd_var_val, - min_expected_psdvar=0.65, - brchisq_threshold=10.0, - psd_var_val_threshold=10.0, - **kwargs): - """ newsnr_sgveto_psdvar with thresholds applied. + return nsnr[0] + + +def newsnr_sgveto_psdvar_threshold( + snr, + brchisq, + sgchisq, + psd_var_val, + min_expected_psdvar=0.65, + brchisq_threshold=10.0, + psd_var_val_threshold=10.0, + **kwargs, +): + """ + newsnr_sgveto_psdvar with thresholds applied. This is the newsnr_sgveto_psdvar statistic with additional options to threshold on chi-squared or PSD variation. @@ -104,65 +99,53 @@ def newsnr_sgveto_psdvar_threshold(snr, brchisq, sgchisq, psd_var_val, sgchisq, psd_var_val, min_expected_psdvar=min_expected_psdvar, - **kwargs + **kwargs, ) nsnr = numpy.array(nsnr, ndmin=1) - nsnr[brchisq > brchisq_threshold] = 1. - nsnr[psd_var_val > psd_var_val_threshold] = 1. + nsnr[brchisq > brchisq_threshold] = 1.0 + nsnr[psd_var_val > psd_var_val_threshold] = 1.0 # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return nsnr - else: - return nsnr[0] - - -def newsnr_sgveto_psdvar_scaled(snr, brchisq, sgchisq, psd_var_val, - scaling=0.33, min_expected_psdvar=0.65, - **kwargs): - """ Combined SNR derived from NewSNR, Sine-Gaussian Chisq and scaled PSD - variation statistic. """ - nsnr = numpy.array( - newsnr_sgveto( - snr, - brchisq, - sgchisq, - **kwargs), - ndmin=1) + return nsnr[0] + + +def newsnr_sgveto_psdvar_scaled( + snr, brchisq, sgchisq, psd_var_val, scaling=0.33, min_expected_psdvar=0.65, **kwargs +): + """ + Combined SNR derived from NewSNR, Sine-Gaussian Chisq and scaled PSD + variation statistic. + """ + nsnr = numpy.array(newsnr_sgveto(snr, brchisq, sgchisq, **kwargs), ndmin=1) psd_var_val = numpy.array(psd_var_val, ndmin=1, copy=True) - psd_var_val[psd_var_val < min_expected_psdvar] = 1. + psd_var_val[psd_var_val < min_expected_psdvar] = 1.0 # Default scale is 0.33 as tuned from analysis of data from O2 chunks - nsnr = nsnr / psd_var_val ** scaling + nsnr = nsnr / psd_var_val**scaling # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return nsnr - else: - return nsnr[0] + return nsnr[0] -def newsnr_sgveto_psdvar_scaled_threshold(snr, bchisq, sgchisq, psd_var_val, - threshold=2.0, - **kwargs): - """ Combined SNR derived from NewSNR and Sine-Gaussian Chisq, and +def newsnr_sgveto_psdvar_scaled_threshold( + snr, bchisq, sgchisq, psd_var_val, threshold=2.0, **kwargs +): + """ + Combined SNR derived from NewSNR and Sine-Gaussian Chisq, and scaled psd variation. """ - nsnr = newsnr_sgveto_psdvar_scaled( - snr, - bchisq, - sgchisq, - psd_var_val, - **kwargs - ) + nsnr = newsnr_sgveto_psdvar_scaled(snr, bchisq, sgchisq, psd_var_val, **kwargs) nsnr = numpy.array(nsnr, ndmin=1) - nsnr[bchisq > threshold] = 1. + nsnr[bchisq > threshold] = 1.0 # If snr input is float, return a float. Otherwise return numpy array. - if hasattr(snr, '__len__'): + if hasattr(snr, "__len__"): return nsnr - else: - return nsnr[0] + return nsnr[0] def get_snr(trigs, **kwargs): # pylint:disable=unused-argument @@ -179,8 +162,9 @@ def get_snr(trigs, **kwargs): # pylint:disable=unused-argument ------- numpy.ndarray Array of snr values + """ - return numpy.array(trigs['snr'][:], ndmin=1, dtype=numpy.float32) + return numpy.array(trigs["snr"][:], ndmin=1, dtype=numpy.float32) def get_newsnr(trigs, **kwargs): @@ -197,13 +181,10 @@ def get_newsnr(trigs, **kwargs): ------- numpy.ndarray Array of newsnr values + """ - dof = 2. * trigs['chisq_dof'][:] - 2. - nsnr = newsnr( - trigs['snr'][:], - trigs['chisq'][:] / dof, - **kwargs - ) + dof = 2.0 * trigs["chisq_dof"][:] - 2.0 + nsnr = newsnr(trigs["snr"][:], trigs["chisq"][:] / dof, **kwargs) return numpy.array(nsnr, ndmin=1, dtype=numpy.float32) @@ -221,13 +202,11 @@ def get_newsnr_sgveto(trigs, **kwargs): ------- numpy.ndarray Array of newsnr values + """ - dof = 2. * trigs['chisq_dof'][:] - 2. + dof = 2.0 * trigs["chisq_dof"][:] - 2.0 nsnr_sg = newsnr_sgveto( - trigs['snr'][:], - trigs['chisq'][:] / dof, - trigs['sg_chisq'][:], - **kwargs + trigs["snr"][:], trigs["chisq"][:] / dof, trigs["sg_chisq"][:], **kwargs ) return numpy.array(nsnr_sg, ndmin=1, dtype=numpy.float32) @@ -247,14 +226,15 @@ def get_newsnr_sgveto_psdvar(trigs, **kwargs): ------- numpy.ndarray Array of newsnr values + """ - dof = 2. * trigs['chisq_dof'][:] - 2. + dof = 2.0 * trigs["chisq_dof"][:] - 2.0 nsnr_sg_psd = newsnr_sgveto_psdvar( - trigs['snr'][:], - trigs['chisq'][:] / dof, - trigs['sg_chisq'][:], - trigs['psd_var_val'][:], - **kwargs + trigs["snr"][:], + trigs["chisq"][:] / dof, + trigs["sg_chisq"][:], + trigs["psd_var_val"][:], + **kwargs, ) return numpy.array(nsnr_sg_psd, ndmin=1, dtype=numpy.float32) @@ -274,13 +254,15 @@ def get_newsnr_sgveto_psdvar_threshold(trigs, **kwargs): ------- numpy.ndarray Array of newsnr values + """ - dof = 2. * trigs['chisq_dof'][:] - 2. + dof = 2.0 * trigs["chisq_dof"][:] - 2.0 nsnr_sg_psdt = newsnr_sgveto_psdvar_threshold( - trigs['snr'][:], trigs['chisq'][:] / dof, - trigs['sg_chisq'][:], - trigs['psd_var_val'][:], - **kwargs + trigs["snr"][:], + trigs["chisq"][:] / dof, + trigs["sg_chisq"][:], + trigs["psd_var_val"][:], + **kwargs, ) return numpy.array(nsnr_sg_psdt, ndmin=1, dtype=numpy.float32) @@ -300,14 +282,15 @@ def get_newsnr_sgveto_psdvar_scaled(trigs, **kwargs): ------- numpy.ndarray Array of newsnr values + """ - dof = 2. * trigs['chisq_dof'][:] - 2. + dof = 2.0 * trigs["chisq_dof"][:] - 2.0 nsnr_sg_psdscale = newsnr_sgveto_psdvar_scaled( - trigs['snr'][:], - trigs['chisq'][:] / dof, - trigs['sg_chisq'][:], - trigs['psd_var_val'][:], - **kwargs + trigs["snr"][:], + trigs["chisq"][:] / dof, + trigs["sg_chisq"][:], + trigs["psd_var_val"][:], + **kwargs, ) return numpy.array(nsnr_sg_psdscale, ndmin=1, dtype=numpy.float32) @@ -328,44 +311,42 @@ def get_newsnr_sgveto_psdvar_scaled_threshold(trigs, **kwargs): ------- numpy.ndarray Array of newsnr values + """ - dof = 2. * trigs['chisq_dof'][:] - 2. + dof = 2.0 * trigs["chisq_dof"][:] - 2.0 nsnr_sg_psdt = newsnr_sgveto_psdvar_scaled_threshold( - trigs['snr'][:], - trigs['chisq'][:] / dof, - trigs['sg_chisq'][:], - trigs['psd_var_val'][:], - **kwargs + trigs["snr"][:], + trigs["chisq"][:] / dof, + trigs["sg_chisq"][:], + trigs["psd_var_val"][:], + **kwargs, ) return numpy.array(nsnr_sg_psdt, ndmin=1, dtype=numpy.float32) sngls_ranking_function_dict = { - 'snr': get_snr, - 'newsnr': get_newsnr, - 'new_snr': get_newsnr, - 'newsnr_sgveto': get_newsnr_sgveto, - 'newsnr_sgveto_psdvar': get_newsnr_sgveto_psdvar, - 'newsnr_sgveto_psdvar_threshold': get_newsnr_sgveto_psdvar_threshold, - 'newsnr_sgveto_psdvar_scaled': get_newsnr_sgveto_psdvar_scaled, - 'newsnr_sgveto_psdvar_scaled_threshold': - get_newsnr_sgveto_psdvar_scaled_threshold, + "snr": get_snr, + "newsnr": get_newsnr, + "new_snr": get_newsnr, + "newsnr_sgveto": get_newsnr_sgveto, + "newsnr_sgveto_psdvar": get_newsnr_sgveto_psdvar, + "newsnr_sgveto_psdvar_threshold": get_newsnr_sgveto_psdvar_threshold, + "newsnr_sgveto_psdvar_scaled": get_newsnr_sgveto_psdvar_scaled, + "newsnr_sgveto_psdvar_scaled_threshold": get_newsnr_sgveto_psdvar_scaled_threshold, } # Lists of datasets required in the trigs object for each function reqd_datasets = {} -reqd_datasets['snr'] = ['snr'] -reqd_datasets['newsnr'] = reqd_datasets['snr'] + ['chisq', 'chisq_dof'] -reqd_datasets['new_snr'] = reqd_datasets['newsnr'] -reqd_datasets['newsnr_sgveto'] = reqd_datasets['newsnr'] + ['sg_chisq'] -reqd_datasets['newsnr_sgveto_psdvar'] = \ - reqd_datasets['newsnr_sgveto'] + ['psd_var_val'] -reqd_datasets['newsnr_sgveto_psdvar_threshold'] = \ - reqd_datasets['newsnr_sgveto_psdvar'] -reqd_datasets['newsnr_sgveto_psdvar_scaled'] = \ - reqd_datasets['newsnr_sgveto_psdvar'] -reqd_datasets['newsnr_sgveto_psdvar_scaled_threshold'] = \ - reqd_datasets['newsnr_sgveto_psdvar'] +reqd_datasets["snr"] = ["snr"] +reqd_datasets["newsnr"] = reqd_datasets["snr"] + ["chisq", "chisq_dof"] +reqd_datasets["new_snr"] = reqd_datasets["newsnr"] +reqd_datasets["newsnr_sgveto"] = reqd_datasets["newsnr"] + ["sg_chisq"] +reqd_datasets["newsnr_sgveto_psdvar"] = reqd_datasets["newsnr_sgveto"] + ["psd_var_val"] +reqd_datasets["newsnr_sgveto_psdvar_threshold"] = reqd_datasets["newsnr_sgveto_psdvar"] +reqd_datasets["newsnr_sgveto_psdvar_scaled"] = reqd_datasets["newsnr_sgveto_psdvar"] +reqd_datasets["newsnr_sgveto_psdvar_scaled_threshold"] = reqd_datasets[ + "newsnr_sgveto_psdvar" +] def get_sngls_ranking_from_trigs(trigs, statname, **kwargs): @@ -376,17 +357,18 @@ def get_sngls_ranking_from_trigs(trigs, statname, **kwargs): specific statname. Parameters - ----------- + ---------- trigs: dict of numpy.ndarrays, SingleDetTriggers or ReadByTemplate Dictionary holding single detector trigger information. statname: The statistic to use. + """ # Identify correct function try: sngl_func = sngls_ranking_function_dict[statname] except KeyError as exc: - err_msg = 'Single-detector ranking {} not recognized'.format(statname) + err_msg = f"Single-detector ranking {statname} not recognized" raise ValueError(err_msg) from exc # NOTE: In the sngl_funcs all the kwargs are explicitly stated, so any diff --git a/pycbc/events/significance.py b/pycbc/events/significance.py index 0341aefcac5..14d65c526e1 100644 --- a/pycbc/events/significance.py +++ b/pycbc/events/significance.py @@ -26,18 +26,21 @@ through different estimation methods of the background, and functions that read in the associated options to do so. """ -import logging + import copy +import logging + import numpy as np -from pycbc.events import trigger_fits as trstats + from pycbc import conversions as conv +from pycbc.events import trigger_fits as trstats -logger = logging.getLogger('pycbc.events.significance') +logger = logging.getLogger("pycbc.events.significance") -def count_n_louder(bstat, fstat, dec, - **kwargs): # pylint:disable=unused-argument - """ Calculate for each foreground event the number of background events +def count_n_louder(bstat, fstat, dec, **kwargs): # pylint:disable=unused-argument + """ + Calculate for each foreground event the number of background events that are louder than it. Parameters @@ -57,6 +60,7 @@ def count_n_louder(bstat, fstat, dec, The number of background triggers above each foreground trigger {} : (empty) dictionary Ensure we return the same tuple of objects as n_louder_from_fit() + """ sort = bstat.argsort() bstat = copy.deepcopy(bstat)[sort] @@ -71,16 +75,15 @@ def count_n_louder(bstat, fstat, dec, # We need to subtract one from the index, to be consistent with definition # of n_louder, as here we do want to include the background value at the # found index - idx = np.searchsorted(bstat, fstat, side='left') - 1 + idx = np.searchsorted(bstat, fstat, side="left") - 1 # If the foreground are *quieter* than the background or at the same value # then the search sorted algorithm will choose position -1, which does not # exist. We force it back to zero. if isinstance(idx, np.ndarray): # Case where our input is an array idx[idx < 0] = 0 - else: # Case where our input is just a scalar value - if idx < 0: - idx = 0 + elif idx < 0: + idx = 0 fore_n_louder = n_louder[idx] @@ -91,9 +94,14 @@ def count_n_louder(bstat, fstat, dec, return back_cum_num, fore_n_louder, {} -def n_louder_from_fit(back_stat, fore_stat, dec_facs, - fit_function='exponential', fit_threshold=0, - **kwargs): # pylint:disable=unused-argument +def n_louder_from_fit( + back_stat, + fore_stat, + dec_facs, + fit_function="exponential", + fit_threshold=0, + **kwargs, +): # pylint:disable=unused-argument """ Use a fit to events in back_stat in order to estimate the distribution for use in recovering the estimate count of louder @@ -123,14 +131,11 @@ def n_louder_from_fit(back_stat, fore_stat, dec_facs, foreground event sig_info : a dictionary Information regarding the significance fit - """ + """ # Calculate the fitting factor of the ranking statistic distribution alpha, sig_alpha = trstats.fit_above_thresh( - fit_function, - back_stat, - thresh=fit_threshold, - weights=dec_facs + fit_function, back_stat, thresh=fit_threshold, weights=dec_facs ) # Count background events above threshold as the cum_fit is @@ -145,14 +150,12 @@ def n_louder_from_fit(back_stat, fore_stat, dec_facs, fg_n_louder = np.zeros_like(fore_stat) # Ue the fit above the threshold - bg_n_louder[bg_above] = n_above * trstats.cum_fit(fit_function, - back_stat[bg_above], - alpha, - fit_threshold) - fg_n_louder[fg_above] = n_above * trstats.cum_fit(fit_function, - fore_stat[fg_above], - alpha, - fit_threshold) + bg_n_louder[bg_above] = n_above * trstats.cum_fit( + fit_function, back_stat[bg_above], alpha, fit_threshold + ) + fg_n_louder[fg_above] = n_above * trstats.cum_fit( + fit_function, fore_stat[fg_above], alpha, fit_threshold + ) # Below the fit threshold, we expect there to be sufficient events # to use the count_n_louder method, and the distribution may deviate @@ -163,9 +166,7 @@ def n_louder_from_fit(back_stat, fore_stat, dec_facs, # Count the number of below-threshold background events louder than the # bg and foreground bg_n_louder[bg_below], fg_n_louder[fg_below], _ = count_n_louder( - back_stat[bg_below], - fore_stat[fg_below], - dec_facs[bg_below] + back_stat[bg_below], fore_stat[fg_below], dec_facs[bg_below] ) # As we have only counted the louder below-threshold events, need to @@ -174,25 +175,23 @@ def n_louder_from_fit(back_stat, fore_stat, dec_facs, bg_n_louder[bg_below] += n_above fg_n_louder[fg_below] += n_above - sig_info = {'alpha': alpha, 'sig_alpha': sig_alpha, 'n_above': n_above} + sig_info = {"alpha": alpha, "sig_alpha": sig_alpha, "n_above": n_above} return bg_n_louder, fg_n_louder, sig_info -_significance_meth_dict = { - 'trigger_fit': n_louder_from_fit, - 'n_louder': count_n_louder -} +_significance_meth_dict = {"trigger_fit": n_louder_from_fit, "n_louder": count_n_louder} _default_opt_dict = { - 'method': 'n_louder', - 'fit_threshold': None, - 'fit_function': None, - 'far_limit': 0.} + "method": "n_louder", + "fit_threshold": None, + "fit_function": None, + "far_limit": 0.0, +} -def get_n_louder(back_stat, fore_stat, dec_facs, - method=_default_opt_dict['method'], - **kwargs): # pylint:disable=unused-argument +def get_n_louder( + back_stat, fore_stat, dec_facs, method=_default_opt_dict["method"], **kwargs +): # pylint:disable=unused-argument """ Wrapper to find the correct n_louder calculation method using standard inputs @@ -208,16 +207,18 @@ def get_n_louder(back_stat, fore_stat, dec_facs, ) back_stat_nonan = np.where(nanmask, -np.inf, back_stat) return _significance_meth_dict[method]( - back_stat_nonan, - fore_stat, - dec_facs, - **kwargs) + back_stat_nonan, fore_stat, dec_facs, **kwargs + ) -def get_far(back_stat, fore_stat, dec_facs, - background_time, - method=_default_opt_dict['method'], - **kwargs): # pylint:disable=unused-argument +def get_far( + back_stat, + fore_stat, + dec_facs, + background_time, + method=_default_opt_dict["method"], + **kwargs, +): # pylint:disable=unused-argument """ Return the appropriate FAR given the significance calculation method @@ -242,17 +243,13 @@ def get_far(back_stat, fore_stat, dec_facs, """ bg_n_louder, fg_n_louder, significance_info = get_n_louder( - back_stat, - fore_stat, - dec_facs, - method=method, - **kwargs + back_stat, fore_stat, dec_facs, method=method, **kwargs ) # If we are counting the number of louder events in the background, # we add one. This is part of the p-value calculation in Usman 2015. # If we are doing trigger fit extrapolation, this is not needed - if method == 'n_louder': + if method == "n_louder": bg_n_louder += 1 fg_n_louder += 1 @@ -270,35 +267,49 @@ def insert_significance_option_group(parser): Add some options for use when a significance is being estimated from events or event distributions. """ - parser.add_argument('--far-calculation-method', nargs='+', - default=[], - help="Method used for FAR calculation in each " - "detector combination, given as " - "combination:method pairs, i.e. " - "H1:trigger_fit H1L1:n_louder H1L1V1:n_louder " - "etc. Method options are [" - + ",".join(_significance_meth_dict.keys()) + - "]. Default = n_louder for all not given") - parser.add_argument('--fit-threshold', nargs='+', default=[], - help="Trigger statistic fit thresholds for FAN " - "estimation, given as combination-value pairs " - "ex. H1:0 L1:0 V1:-4 for all combinations with " - "--far-calculation-method = trigger_fit") - parser.add_argument("--fit-function", nargs='+', default=[], - help="Functional form for the statistic slope fit if " - "--far-calculation-method is 'trigger_fit'. " - "Given as combination:function pairs, i.e. " - "H1:exponential H1L1:n_louder H1L1V1:n_louder. " - "Options: [" - + ",".join(trstats.fitalpha_dict.keys()) + "]. " - "Default = exponential for all") - parser.add_argument('--limit-ifar', nargs='+', default=[], - help="Impose upper limits on IFAR values (years)" - ". Given as combination:value pairs, eg " - "H1L1:10000 L1:1000. Used to avoid under/" - "overflows for loud signals and injections " - "using the fit extrapolation method. A value" - " 0 or no value means unlimited IFAR") + parser.add_argument( + "--far-calculation-method", + nargs="+", + default=[], + help="Method used for FAR calculation in each " + "detector combination, given as " + "combination:method pairs, i.e. " + "H1:trigger_fit H1L1:n_louder H1L1V1:n_louder " + "etc. Method options are [" + + ",".join(_significance_meth_dict.keys()) + + "]. Default = n_louder for all not given", + ) + parser.add_argument( + "--fit-threshold", + nargs="+", + default=[], + help="Trigger statistic fit thresholds for FAN " + "estimation, given as combination-value pairs " + "ex. H1:0 L1:0 V1:-4 for all combinations with " + "--far-calculation-method = trigger_fit", + ) + parser.add_argument( + "--fit-function", + nargs="+", + default=[], + help="Functional form for the statistic slope fit if " + "--far-calculation-method is 'trigger_fit'. " + "Given as combination:function pairs, i.e. " + "H1:exponential H1L1:n_louder H1L1V1:n_louder. " + "Options: [" + ",".join(trstats.fitalpha_dict.keys()) + "]. " + "Default = exponential for all", + ) + parser.add_argument( + "--limit-ifar", + nargs="+", + default=[], + help="Impose upper limits on IFAR values (years)" + ". Given as combination:value pairs, eg " + "H1L1:10000 L1:1000. Used to avoid under/" + "overflows for loud signals and injections " + "using the fit extrapolation method. A value" + " 0 or no value means unlimited IFAR", + ) def positive_float(inp): @@ -308,8 +319,9 @@ def positive_float(inp): """ fl_in = float(inp) if fl_in < 0: - logger.warning("Value provided to positive_float is less than zero, " - "this is not allowed") + logger.warning( + "Value provided to positive_float is less than zero, this is not allowed" + ) raise ValueError return fl_in @@ -321,24 +333,23 @@ def check_significance_options(args, parser): """ # Check that the combo:method/function/threshold are in the # right format, and are in allowed combinations - lists_to_check = [(args.far_calculation_method, str, - _significance_meth_dict.keys()), - (args.fit_function, str, - trstats.fitalpha_dict.keys()), - (args.fit_threshold, float, None), - (args.limit_ifar, positive_float, None)] + lists_to_check = [ + (args.far_calculation_method, str, _significance_meth_dict.keys()), + (args.fit_function, str, trstats.fitalpha_dict.keys()), + (args.fit_threshold, float, None), + (args.limit_ifar, positive_float, None), + ] for list_to_check, type_to_convert, allowed_values in lists_to_check: combo_list = [] for combo_value in list_to_check: try: - combo, value = tuple(combo_value.split(':')) + combo, value = tuple(combo_value.split(":")) except ValueError: parser.error("Need combo:value format, got %s" % combo_value) if combo in combo_list: - parser.error("Duplicate combo %s in a significance " - "option" % combo) + parser.error("Duplicate combo %s in a significance option" % combo) combo_list.append(combo) try: @@ -347,8 +358,10 @@ def check_significance_options(args, parser): err_fmat = "Value {} of combo {} can't be converted" parser.error(err_fmat.format(value, combo)) - if allowed_values is not None and \ - type_to_convert(value) not in allowed_values: + if ( + allowed_values is not None + and type_to_convert(value) not in allowed_values + ): err_fmat = "Value {} of combo {} is not in allowed values: {}" parser.error(err_fmat.format(value, combo, allowed_values)) @@ -356,24 +369,28 @@ def check_significance_options(args, parser): methods = {} # A method has been specified for combo_value in args.far_calculation_method: - combo, value = tuple(combo_value.split(':')) + combo, value = tuple(combo_value.split(":")) methods[combo] = value # A function or threshold has been specified function_or_thresh_given = [] for combo_value in args.fit_function + args.fit_threshold: - combo, _ = tuple(combo_value.split(':')) + combo, _ = tuple(combo_value.split(":")) if combo not in methods: # Assign the default method for use in further tests - methods[combo] = _default_opt_dict['method'] + methods[combo] = _default_opt_dict["method"] function_or_thresh_given.append(combo) for combo, value in methods.items(): - if value != 'trigger_fit' and combo in function_or_thresh_given: + if value != "trigger_fit" and combo in function_or_thresh_given: # Function/Threshold given for combo not using trigger_fit method - parser.error("--fit-function and/or --fit-threshold given for " - + combo + " which has method " + value) - elif value == 'trigger_fit' and combo not in function_or_thresh_given: + parser.error( + "--fit-function and/or --fit-threshold given for " + + combo + + " which has method " + + value + ) + elif value == "trigger_fit" and combo not in function_or_thresh_given: # Threshold not given for trigger_fit combo parser.error("Threshold required for combo " + combo) @@ -390,7 +407,7 @@ def ifar_opt_to_far_limit(ifar_str): """ ifar_float = positive_float(ifar_str) - far_hz = 0. if (ifar_float == 0.) else conv.sec_to_year(1. / ifar_float) + far_hz = 0.0 if (ifar_float == 0.0) else conv.sec_to_year(1.0 / ifar_float) return far_hz @@ -402,7 +419,6 @@ def digest_significance_options(combo_keys, args): Parameters ---------- - combo_keys: list of strings list of detector combinations for which options are needed @@ -414,12 +430,14 @@ def digest_significance_options(combo_keys, args): significance_dict: dictionary Dictionary containing method, threshold and function for trigger fits as appropriate, and any limit on FAR (Hz) - """ - lists_to_unpack = [('method', args.far_calculation_method, str), - ('fit_function', args.fit_function, str), - ('fit_threshold', args.fit_threshold, float), - ('far_limit', args.limit_ifar, ifar_opt_to_far_limit)] + """ + lists_to_unpack = [ + ("method", args.far_calculation_method, str), + ("fit_function", args.fit_function, str), + ("fit_threshold", args.fit_threshold, float), + ("far_limit", args.limit_ifar, ifar_opt_to_far_limit), + ] significance_dict = {} # Set everything as a default to start with: @@ -429,13 +447,14 @@ def digest_significance_options(combo_keys, args): # Unpack everything from the arguments into the dictionary for argument_key, arg_to_unpack, conv_func in lists_to_unpack: for combo_value in arg_to_unpack: - combo, value = tuple(combo_value.split(':')) + combo, value = tuple(combo_value.split(":")) if combo not in significance_dict: # Allow options for detector combos that are not actually # used/required for a given job. Such options have # no effect, but emit a warning for (e.g.) diagnostic checks - logger.warning("Key %s not used by this code, uses %s", - combo, combo_keys) + logger.warning( + "Key %s not used by this code, uses %s", combo, combo_keys + ) significance_dict[combo] = copy.deepcopy(_default_opt_dict) significance_dict[combo][argument_key] = conv_func(value) @@ -467,25 +486,24 @@ def apply_far_limit(far, significance_dict, combo=None): far_out = copy.deepcopy(far) if isinstance(combo, str): # Single IFO combo used - if significance_dict[combo]['far_limit'] == 0: + if significance_dict[combo]["far_limit"] == 0: return far_out far_limit_str = f"{significance_dict[combo]['far_limit']:.3e}" - logger.info("Applying FAR limit of %s to %s events", - far_limit_str, combo) - far_out = np.maximum(far, significance_dict[combo]['far_limit']) + logger.info("Applying FAR limit of %s to %s events", far_limit_str, combo) + far_out = np.maximum(far, significance_dict[combo]["far_limit"]) else: # IFO combo supplied as an array, by e.g. pycbc_add_statmap # Need to check which events are in which IFO combo in order to # apply the right limit to each for ifo_combo in significance_dict: - if significance_dict[ifo_combo]['far_limit'] == 0: + if significance_dict[ifo_combo]["far_limit"] == 0: continue far_limit_str = f"{significance_dict[ifo_combo]['far_limit']:.3e}" - logger.info("Applying FAR limit of %s to %s events", - far_limit_str, ifo_combo) - this_combo_idx = combo == ifo_combo.encode('utf-8') + logger.info( + "Applying FAR limit of %s to %s events", far_limit_str, ifo_combo + ) + this_combo_idx = combo == ifo_combo.encode("utf-8") far_out[this_combo_idx] = np.maximum( - far[this_combo_idx], - significance_dict[ifo_combo]['far_limit'] + far[this_combo_idx], significance_dict[ifo_combo]["far_limit"] ) return far_out diff --git a/pycbc/events/single.py b/pycbc/events/single.py index 94134a7b6e5..d80ab494a3a 100644 --- a/pycbc/events/single.py +++ b/pycbc/events/single.py @@ -1,36 +1,41 @@ -""" utilities for assigning FAR to single detector triggers -""" -import logging +"""utilities for assigning FAR to single detector triggers""" + import copy +import logging import threading import time + import numpy as np -from pycbc.events import trigger_fits as fits, stat -from pycbc.types import MultiDetOptionAction +from pycbc import bin_utils from pycbc import conversions as conv +from pycbc.events import stat +from pycbc.events import trigger_fits as fits from pycbc.io.hdf import HFile -from pycbc import bin_utils +from pycbc.types import MultiDetOptionAction -logger = logging.getLogger('pycbc.events.single') - - -class LiveSingle(object): - def __init__(self, ifo, - ranking_threshold=10.0, - reduced_chisq_threshold=5, - duration_threshold=0, - fit_file=None, - sngl_ifar_est_dist=None, - fixed_ifar=None, - maximum_ifar=None, - statistic=None, - stat_features=None, - stat_keywords=None, - sngl_ranking=None, - stat_files=None, - statistic_refresh_rate=None, - **kwargs): +logger = logging.getLogger("pycbc.events.single") + + +class LiveSingle: + def __init__( + self, + ifo, + ranking_threshold=10.0, + reduced_chisq_threshold=5, + duration_threshold=0, + fit_file=None, + sngl_ifar_est_dist=None, + fixed_ifar=None, + maximum_ifar=None, + statistic=None, + stat_features=None, + stat_keywords=None, + sngl_ranking=None, + stat_files=None, + statistic_refresh_rate=None, + **kwargs, + ): """ Parameters ---------- @@ -74,6 +79,7 @@ class (in seconds), default not do do this kwargs: dict Additional options for the statistic to use. See stat.py for more details on statistic options. + """ self.ifo = ifo self.fit_file = fit_file @@ -91,93 +97,127 @@ class (in seconds), default not do do this stat_keywords, ) self.stat_calculator = stat_class( - sngl_ranking, - stat_files, - ifos=[ifo], - **stat_extras + sngl_ranking, stat_files, ifos=[ifo], **stat_extras ) self.thresholds = { "ranking": ranking_threshold, "reduced_chisq": reduced_chisq_threshold, - "duration": duration_threshold} + "duration": duration_threshold, + } @staticmethod def insert_args(parser): - parser.add_argument('--single-ranking-threshold', nargs='+', - type=float, action=MultiDetOptionAction, - help='Single ranking threshold for ' - 'single-detector events. Can be given ' - 'as a single value or as detector-value ' - 'pairs, e.g. H1:6 L1:7 V1:6.5') - parser.add_argument('--single-reduced-chisq-threshold', nargs='+', - type=float, action=MultiDetOptionAction, - help='Maximum reduced chi-squared threshold for ' - 'single triggers. Calcuated after any PSD ' - 'variation reweighting is applied. Can be ' - 'given as a single value or as ' - 'detector-value pairs, e.g. H1:2 L1:2 V1:3') - parser.add_argument('--single-duration-threshold', nargs='+', - type=float, action=MultiDetOptionAction, - help='Minimum duration threshold for single ' - 'triggers. Can be given as a single value ' - 'or as detector-value pairs, e.g. H1:6 L1:6 ' - 'V1:8') - parser.add_argument('--single-fixed-ifar', nargs='+', - type=float, action=MultiDetOptionAction, - help='A fixed value for IFAR, still uses cuts ' - 'defined by command line. Can be given as ' - 'a single value or as detector-value pairs, ' - 'e.g. H1:0.001 L1:0.001 V1:0.0005') - parser.add_argument('--single-maximum-ifar', nargs='+', - type=float, action=MultiDetOptionAction, - help='A maximum possible value for IFAR for ' - 'single-detector events. Can be given as ' - 'a single value or as detector-value pairs, ' - 'e.g. H1:100 L1:1000 V1:50') - parser.add_argument('--single-fit-file', - help='File which contains definitons of fit ' - 'coefficients and counts for specific ' - 'single trigger IFAR fitting.') - parser.add_argument('--sngl-ifar-est-dist', nargs='+', - action=MultiDetOptionAction, - help='Which trigger distribution to use when ' - 'calculating IFAR of single triggers. ' - 'Can be given as a single value or as ' - 'detector-value pairs, e.g. H1:mean ' - 'L1:mean V1:conservative') + parser.add_argument( + "--single-ranking-threshold", + nargs="+", + type=float, + action=MultiDetOptionAction, + help="Single ranking threshold for " + "single-detector events. Can be given " + "as a single value or as detector-value " + "pairs, e.g. H1:6 L1:7 V1:6.5", + ) + parser.add_argument( + "--single-reduced-chisq-threshold", + nargs="+", + type=float, + action=MultiDetOptionAction, + help="Maximum reduced chi-squared threshold for " + "single triggers. Calcuated after any PSD " + "variation reweighting is applied. Can be " + "given as a single value or as " + "detector-value pairs, e.g. H1:2 L1:2 V1:3", + ) + parser.add_argument( + "--single-duration-threshold", + nargs="+", + type=float, + action=MultiDetOptionAction, + help="Minimum duration threshold for single " + "triggers. Can be given as a single value " + "or as detector-value pairs, e.g. H1:6 L1:6 " + "V1:8", + ) + parser.add_argument( + "--single-fixed-ifar", + nargs="+", + type=float, + action=MultiDetOptionAction, + help="A fixed value for IFAR, still uses cuts " + "defined by command line. Can be given as " + "a single value or as detector-value pairs, " + "e.g. H1:0.001 L1:0.001 V1:0.0005", + ) + parser.add_argument( + "--single-maximum-ifar", + nargs="+", + type=float, + action=MultiDetOptionAction, + help="A maximum possible value for IFAR for " + "single-detector events. Can be given as " + "a single value or as detector-value pairs, " + "e.g. H1:100 L1:1000 V1:50", + ) + parser.add_argument( + "--single-fit-file", + help="File which contains definitons of fit " + "coefficients and counts for specific " + "single trigger IFAR fitting.", + ) + parser.add_argument( + "--sngl-ifar-est-dist", + nargs="+", + action=MultiDetOptionAction, + help="Which trigger distribution to use when " + "calculating IFAR of single triggers. " + "Can be given as a single value or as " + "detector-value pairs, e.g. H1:mean " + "L1:mean V1:conservative", + ) @staticmethod def verify_args(args, parser, ifos): - sngl_opts = [args.single_reduced_chisq_threshold, - args.single_duration_threshold, - args.single_ranking_threshold, - args.sngl_ifar_est_dist] - - sngl_opts_str = ("--single-reduced-chisq-threshold, " - "--single-duration-threshold, " - "--single-ranking-threshold, " - "--sngl-ifar-est-dist") + sngl_opts = [ + args.single_reduced_chisq_threshold, + args.single_duration_threshold, + args.single_ranking_threshold, + args.sngl_ifar_est_dist, + ] + + sngl_opts_str = ( + "--single-reduced-chisq-threshold, " + "--single-duration-threshold, " + "--single-ranking-threshold, " + "--sngl-ifar-est-dist" + ) if any(sngl_opts) and not all(sngl_opts): - parser.error(f"Single detector trigger options ({sngl_opts_str}) " - "must either all be given or none.") - - if args.enable_single_detector_upload \ - and not args.enable_gracedb_upload: - parser.error("--enable-single-detector-upload requires " - "--enable-gracedb-upload to be set.") - - sngl_optional_opts = [args.single_fixed_ifar, - args.single_fit_file, - args.single_maximum_ifar] - sngl_optional_opts_str = ("--single-fixed-ifar, " - "--single-fit-file," - "--single-maximum-ifar") + parser.error( + f"Single detector trigger options ({sngl_opts_str}) " + "must either all be given or none." + ) + + if args.enable_single_detector_upload and not args.enable_gracedb_upload: + parser.error( + "--enable-single-detector-upload requires " + "--enable-gracedb-upload to be set." + ) + + sngl_optional_opts = [ + args.single_fixed_ifar, + args.single_fit_file, + args.single_maximum_ifar, + ] + sngl_optional_opts_str = ( + "--single-fixed-ifar, --single-fit-file,--single-maximum-ifar" + ) if any(sngl_optional_opts) and not all(sngl_opts): - parser.error("Optional singles options " - f"({sngl_optional_opts_str}) given but not all " - f"required options ({sngl_opts_str}) are.") + parser.error( + "Optional singles options " + f"({sngl_optional_opts_str}) given but not all " + f"required options ({sngl_opts_str}) are." + ) for ifo in ifos: # Check which option(s) are needed for each IFO and if they exist: @@ -186,38 +226,47 @@ def verify_args(args, parser, ifos): # args.sngl_ifar_est_dist.default_set is True if single value has # been set to be the same for all values # bool(args.sngl_ifar_est_dist) is True if option is given - if args.sngl_ifar_est_dist and \ - not args.sngl_ifar_est_dist.default_set \ - and not args.sngl_ifar_est_dist[ifo]: + if ( + args.sngl_ifar_est_dist + and not args.sngl_ifar_est_dist.default_set + and not args.sngl_ifar_est_dist[ifo] + ): # Option has been given, different for each IFO, # and this one is not present - parser.error("All IFOs required in --single-ifar-est-dist " - "if IFO-specific options are given.") + parser.error( + "All IFOs required in --single-ifar-est-dist " + "if IFO-specific options are given." + ) if args.sngl_ifar_est_dist[ifo] is None: # Default - no singles being used continue - if not args.sngl_ifar_est_dist[ifo] == 'fixed': + if not args.sngl_ifar_est_dist[ifo] == "fixed": if not args.single_fit_file: # Fixed IFAR option doesnt need the fits file - parser.error(f"Single detector trigger fits file must be " - "given if --single-ifar-est-dist is not " - f"fixed for all ifos (at least {ifo} has " - f"option {args.sngl_ifar_est_dist[ifo]}).") + parser.error( + f"Single detector trigger fits file must be " + "given if --single-ifar-est-dist is not " + f"fixed for all ifos (at least {ifo} has " + f"option {args.sngl_ifar_est_dist[ifo]})." + ) if ifo in args.single_fixed_ifar: - parser.error(f"Value {args.single_fixed_ifar[ifo]} given " - f"for {ifo} in --single-fixed-ifar, but " - f"--single-ifar-est-dist for {ifo} " - f"is {args.sngl_ifar_est_dist[ifo]}, not " - "fixed.") - else: - # Check that the fixed IFAR value has actually been - # given if using this instead of a distribution - if not args.single_fixed_ifar[ifo]: - parser.error(f"--single-fixed-ifar must be " - "given if --single-ifar-est-dist is fixed. " - f"This is true for at least {ifo}.") + parser.error( + f"Value {args.single_fixed_ifar[ifo]} given " + f"for {ifo} in --single-fixed-ifar, but " + f"--single-ifar-est-dist for {ifo} " + f"is {args.sngl_ifar_est_dist[ifo]}, not " + "fixed." + ) + # Check that the fixed IFAR value has actually been + # given if using this instead of a distribution + elif not args.single_fixed_ifar[ifo]: + parser.error( + f"--single-fixed-ifar must be " + "given if --single-ifar-est-dist is fixed. " + f"This is true for at least {ifo}." + ) # Return value is a boolean whether we are analysing singles or not # The checks already performed mean that all(sngl_opts) is okay @@ -234,44 +283,44 @@ def from_cli(cls, args, ifo): stat_files = sum(stat_files, []) return cls( - ifo, ranking_threshold=args.single_ranking_threshold[ifo], - reduced_chisq_threshold=args.single_reduced_chisq_threshold[ifo], - duration_threshold=args.single_duration_threshold[ifo], - fixed_ifar=args.single_fixed_ifar, - maximum_ifar=args.single_maximum_ifar[ifo], - fit_file=args.single_fit_file, - sngl_ifar_est_dist=args.sngl_ifar_est_dist[ifo], - statistic=args.ranking_statistic, - sngl_ranking=args.sngl_ranking, - stat_features=args.statistic_features, - stat_keywords=args.statistic_keywords, - stat_files=stat_files, - statistic_refresh_rate=args.statistic_refresh_rate, - ) + ifo, + ranking_threshold=args.single_ranking_threshold[ifo], + reduced_chisq_threshold=args.single_reduced_chisq_threshold[ifo], + duration_threshold=args.single_duration_threshold[ifo], + fixed_ifar=args.single_fixed_ifar, + maximum_ifar=args.single_maximum_ifar[ifo], + fit_file=args.single_fit_file, + sngl_ifar_est_dist=args.sngl_ifar_est_dist[ifo], + statistic=args.ranking_statistic, + sngl_ranking=args.sngl_ranking, + stat_features=args.statistic_features, + stat_keywords=args.statistic_keywords, + stat_files=stat_files, + statistic_refresh_rate=args.statistic_refresh_rate, + ) def check(self, trigs, data_reader): - """ Look for a single detector trigger that passes the thresholds in + """ + Look for a single detector trigger that passes the thresholds in the current data. """ - # Apply cuts to trigs before clustering # Cut on snr so that triggers which could not reach the ranking # threshold do not have ranking calculated - if 'psd_var_val' in trigs: + if "psd_var_val" in trigs: # We should apply the PSD variation rescaling, as this can # re-weight the SNR to be above SNR - trig_chisq = trigs['chisq'] / trigs['psd_var_val'] - trig_snr = trigs['snr'] / (trigs['psd_var_val'] ** 0.5) + trig_chisq = trigs["chisq"] / trigs["psd_var_val"] + trig_snr = trigs["snr"] / (trigs["psd_var_val"] ** 0.5) else: - trig_chisq = trigs['chisq'] - trig_snr = trigs['snr'] - - valid_idx = (trigs['template_duration'] > - self.thresholds['duration']) & \ - (trig_chisq < - self.thresholds['reduced_chisq']) & \ - (trig_snr > - self.thresholds['ranking']) + trig_chisq = trigs["chisq"] + trig_snr = trigs["snr"] + + valid_idx = ( + (trigs["template_duration"] > self.thresholds["duration"]) + & (trig_chisq < self.thresholds["reduced_chisq"]) + & (trig_snr > self.thresholds["ranking"]) + ) if not np.any(valid_idx): return None @@ -280,20 +329,19 @@ def check(self, trigs, data_reader): # Convert back from the pycbc live convention of chisq always # meaning the reduced chisq. trigsc = copy.copy(cut_trigs) - trigsc['ifo'] = self.ifo - trigsc['chisq'] = cut_trigs['chisq'] * cut_trigs['chisq_dof'] - trigsc['chisq_dof'] = (cut_trigs['chisq_dof'] + 2) / 2 + trigsc["ifo"] = self.ifo + trigsc["chisq"] = cut_trigs["chisq"] * cut_trigs["chisq_dof"] + trigsc["chisq_dof"] = (cut_trigs["chisq_dof"] + 2) / 2 # Calculate the ranking reweighted SNR for cutting with self.stat_calculator_lock: single_rank = self.stat_calculator.get_sngl_ranking(trigsc) - sngl_idx = single_rank > self.thresholds['ranking'] + sngl_idx = single_rank > self.thresholds["ranking"] if not np.any(sngl_idx): return None - cutall_trigs = {k: trigsc[k][sngl_idx] - for k in trigs} + cutall_trigs = {k: trigsc[k][sngl_idx] for k in trigs} # Calculate the ranking statistic with self.stat_calculator_lock: @@ -304,21 +352,17 @@ def check(self, trigs, data_reader): i = rank.argmax() # calculate the (inverse) false-alarm rate - ifar = self.calculate_ifar( - rank[i], - trigsc['template_duration'][i] - ) + ifar = self.calculate_ifar(rank[i], trigsc["template_duration"][i]) if ifar is None: return None # fill in a new candidate event candidate = { - f'foreground/{self.ifo}/{k}': cut_trigs[k][sngl_idx][i] - for k in trigs + f"foreground/{self.ifo}/{k}": cut_trigs[k][sngl_idx][i] for k in trigs } - candidate['foreground/stat'] = rank[i] - candidate['foreground/ifar'] = ifar - candidate['HWINJ'] = data_reader.near_hwinj() + candidate["foreground/stat"] = rank[i] + candidate["foreground/ifar"] = ifar + candidate["HWINJ"] = data_reader.near_hwinj() return candidate def calculate_ifar(self, sngl_ranking, duration): @@ -327,18 +371,18 @@ def calculate_ifar(self, sngl_ranking, duration): return self.fixed_ifar[self.ifo] try: - with HFile(self.fit_file, 'r') as fit_file: - bin_edges = fit_file['bins_edges'][:] - live_time = fit_file[self.ifo].attrs['live_time'] - thresh = fit_file[self.ifo].attrs['fit_threshold'] + with HFile(self.fit_file, "r") as fit_file: + bin_edges = fit_file["bins_edges"][:] + live_time = fit_file[self.ifo].attrs["live_time"] + thresh = fit_file[self.ifo].attrs["fit_threshold"] dist_grp = fit_file[self.ifo][self.sngl_ifar_est_dist] - rates = dist_grp['counts'][:] / live_time - coeffs = dist_grp['fit_coeff'][:] + rates = dist_grp["counts"][:] / live_time + coeffs = dist_grp["fit_coeff"][:] except FileNotFoundError: logger.error( - 'Single fit file %s not found; ' - 'dropping a potential single-detector candidate!', - self.fit_file + "Single fit file %s not found; " + "dropping a potential single-detector candidate!", + self.fit_file, ) return None @@ -354,17 +398,14 @@ def calculate_ifar(self, sngl_ranking, duration): ) return None - rate_louder = rate * fits.cum_fit( - 'exponential', - [sngl_ranking], - coeff, - thresh - )[0] + rate_louder = ( + rate * fits.cum_fit("exponential", [sngl_ranking], coeff, thresh)[0] + ) # apply a trials factor of the number of duration bins rate_louder *= len(rates) - return min(conv.sec_to_year(1. / rate_louder), self.maximum_ifar) + return min(conv.sec_to_year(1.0 / rate_louder), self.maximum_ifar) def start_refresh_thread(self): """ @@ -374,9 +415,7 @@ def start_refresh_thread(self): logger.info("Statistic refresh disabled for %s", self.ifo) return thread = threading.Thread( - target=self.refresh_statistic, - daemon=True, - name="Stat refresh " + self.ifo + target=self.refresh_statistic, daemon=True, name="Stat refresh " + self.ifo ) logger.info("Starting %s statistic refresh thread", self.ifo) thread.start() @@ -390,9 +429,7 @@ def refresh_statistic(self): since_stat_refresh = time.time() - self.time_stat_refreshed if since_stat_refresh > self.statistic_refresh_rate: self.time_stat_refreshed = time.time() - logger.info( - "Checking %s statistic for updated files", self.ifo - ) + logger.info("Checking %s statistic for updated files", self.ifo) with self.stat_calculator_lock: self.stat_calculator.check_update_files() # Sleep one second for safety @@ -402,6 +439,6 @@ def refresh_statistic(self): logger.debug( "%s statistic: Waiting %.3fs for next refresh", self.ifo, - self.statistic_refresh_rate - since_stat_refresh + self.statistic_refresh_rate - since_stat_refresh, ) time.sleep(self.statistic_refresh_rate - since_stat_refresh) diff --git a/pycbc/events/stat.py b/pycbc/events/stat.py index faa020ee9a0..8829417c87f 100644 --- a/pycbc/events/stat.py +++ b/pycbc/events/stat.py @@ -25,16 +25,19 @@ This module contains functions for calculating coincident ranking statistic values. """ + import logging -from hashlib import sha1 from datetime import datetime as dt -import numpy +from hashlib import sha1 + import h5py +import numpy -from . import ranking -from . import coinc_rate -from .eventmgr_cython import logsignalrateinternals_computepsignalbins -from .eventmgr_cython import logsignalrateinternals_compute2detrate +from . import coinc_rate, ranking +from .eventmgr_cython import ( + logsignalrateinternals_compute2detrate, + logsignalrateinternals_computepsignalbins, +) logger = logging.getLogger("pycbc.events.stat") @@ -48,7 +51,7 @@ ] -class Stat(object): +class Stat: """Base class which should be extended to provide a statistic""" def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): @@ -66,8 +69,8 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): statistic class. ifos: list of strs, needed for some statistics The list of detector names - """ + """ self.files = {} files = files or [] for filename in files: @@ -112,18 +115,16 @@ def get_file_hashes(self): """ Get sha1 hashes for all the files """ - logger.debug( - "Getting file hashes" - ) + logger.debug("Getting file hashes") start = dt.now() file_hashes = {} for stat, filename in self.files.items(): - with open(filename, 'rb') as file_binary: + with open(filename, "rb") as file_binary: file_hashes[stat] = sha1(file_binary.read()).hexdigest() logger.debug( "Got file hashes for %d files, took %.3es", len(self.files), - (dt.now() - start).total_seconds() + (dt.now() - start).total_seconds(), ) return file_hashes @@ -136,7 +137,7 @@ def files_changed(self): if changed_file_hashes[stat] != old_hash: logger.info( "%s statistic file %s has changed", - ''.join(self.ifos), + "".join(self.ifos), stat, ) else: @@ -144,10 +145,7 @@ def files_changed(self): del changed_file_hashes[stat] if changed_file_hashes == {}: - logger.debug( - "No %s statistic files have changed", - ''.join(self.ifos) - ) + logger.debug("No %s statistic files have changed", "".join(self.ifos)) return list(changed_file_hashes.keys()) @@ -182,11 +180,10 @@ def get_sngl_ranking(self, trigs): ------- numpy.ndarray The array of single detector values + """ return ranking.get_sngls_ranking_from_trigs( - trigs, - self.sngl_ranking, - **self.sngl_ranking_kwargs + trigs, self.sngl_ranking, **self.sngl_ranking_kwargs ) def single(self, trigs): # pylint:disable=unused-argument @@ -202,6 +199,7 @@ def single(self, trigs): # pylint:disable=unused-argument ------- numpy.ndarray The array of single detector values + """ err_msg = "This function is a stub that should be overridden by the " err_msg += "sub-classes. You shouldn't be seeing this error!" @@ -221,14 +219,13 @@ def rank_stat_single(self, single_info): # pylint:disable=unused-argument ------- numpy.ndarray The array of single detector statistics + """ err_msg = "This function is a stub that should be overridden by the " err_msg += "sub-classes. You shouldn't be seeing this error!" raise NotImplementedError(err_msg) - def rank_stat_coinc( - self, s, slide, step, to_shift, **kwargs - ): # pylint:disable=unused-argument + def rank_stat_coinc(self, s, slide, step, to_shift, **kwargs): # pylint:disable=unused-argument """ Calculate the coincident detection statistic. """ @@ -246,9 +243,10 @@ def _check_coinc_lim_subclass(self, allowed_names): check you will see the failure message here. Parameters - ----------- + ---------- allowed_names : list list of allowed classes for the specific sub-classed method. + """ if type(self).__name__ not in allowed_names: err_msg = "This is being called from a subclass which has not " @@ -257,16 +255,13 @@ def _check_coinc_lim_subclass(self, allowed_names): err_msg += "list of allowed_names above." raise NotImplementedError(err_msg) - def coinc_lim_for_thresh( - self, s, thresh, limifo, **kwargs - ): # pylint:disable=unused-argument + def coinc_lim_for_thresh(self, s, thresh, limifo, **kwargs): # pylint:disable=unused-argument """ Optimization function to identify coincs too quiet to be of interest Calculate the required single detector statistic to exceed the threshold for each of the input triggers. """ - err_msg = "This function is a stub that should be overridden by the " err_msg += "sub-classes. You shouldn't be seeing this error!" raise NotImplementedError(err_msg) @@ -290,6 +285,7 @@ def single(self, trigs): ------- numpy.ndarray The array of single detector values + """ return self.get_sngl_ranking(trigs) @@ -310,12 +306,11 @@ def rank_stat_single(self, single_info): ------- numpy.ndarray The array of single detector statistics + """ return single_info[1] - def rank_stat_coinc( - self, sngls_list, slide, step, to_shift, **kwargs - ): # pylint:disable=unused-argument + def rank_stat_coinc(self, sngls_list, slide, step, to_shift, **kwargs): # pylint:disable=unused-argument """ Calculate the coincident detection statistic. @@ -333,16 +328,15 @@ def rank_stat_coinc( ------- numpy.ndarray Array of coincident ranking statistic values + """ - cstat = sum(sngl[1] ** 2. for sngl in sngls_list) ** 0.5 + cstat = sum(sngl[1] ** 2.0 for sngl in sngls_list) ** 0.5 # For single-detector "cuts" the single ranking is set to -1 for sngls in sngls_list: cstat[sngls == -1] = 0 return cstat - def coinc_lim_for_thresh( - self, s, thresh, limifo, **kwargs - ): # pylint:disable=unused-argument + def coinc_lim_for_thresh(self, s, thresh, limifo, **kwargs): # pylint:disable=unused-argument """ Optimization function to identify coincs too quiet to be of interest @@ -364,14 +358,15 @@ def coinc_lim_for_thresh( numpy.ndarray Array of limits on the limifo single statistic to exceed thresh. + """ # Safety against subclassing and not rethinking this allowed_names = ["QuadratureSumStatistic"] self._check_coinc_lim_subclass(allowed_names) - s0 = thresh ** 2. - sum(sngl[1] ** 2. for sngl in s) + s0 = thresh**2.0 - sum(sngl[1] ** 2.0 for sngl in s) s0[s0 < 0] = 0 - return s0 ** 0.5 + return s0**0.5 class PhaseTDStatistic(QuadratureSumStatistic): @@ -408,6 +403,7 @@ def __init__( pregenerate_hist: bool, optional If False, do not pregenerate histogram on class instantiation. Default is True. + """ QuadratureSumStatistic.__init__( self, sngl_ranking, files=files, ifos=ifos, **kwargs @@ -424,7 +420,7 @@ def __init__( # Assign attribute so that it can be replaced with other functions self.has_hist = False self.hist_ifos = None - self.ref_snr = 5. + self.ref_snr = 5.0 self.relsense = {} self.swidth = self.pwidth = self.twidth = None self.srbmin = self.srbmax = None @@ -445,7 +441,7 @@ def __init__( # Is the histogram needed to be pre-generated? hist_needed = pregenerate_hist hist_needed &= not len(ifos) == 1 - hist_needed &= (type(self).__name__ == "PhaseTD" or self.kwargs["phasetd"]) + hist_needed &= type(self).__name__ == "PhaseTD" or self.kwargs["phasetd"] if hist_needed: self.get_hist() @@ -453,7 +449,7 @@ def __init__( # remove all phasetd files from self.files and self.file_hashes, # as they are not needed for k in list(self.files.keys()): - if 'phasetd_newsnr' in k: + if "phasetd_newsnr" in k: del self.files[k] del self.file_hashes[k] @@ -466,6 +462,7 @@ def get_hist(self, ifos=None): ifos: list The list of ifos. Needed if not given when initializing the class instance. + """ ifos = ifos or self.ifos @@ -487,8 +484,11 @@ def get_hist(self, ifos=None): # If there are other phasetd_newsnr files, they aren't needed. # So tidy them out of the self.files dictionary - rejected = [key for key in self.files.keys() - if 'phasetd_newsnr' in key and not key == selected] + rejected = [ + key + for key in self.files.keys() + if "phasetd_newsnr" in key and not key == selected + ] for k in rejected: del self.files[k] del self.file_hashes[k] @@ -529,23 +529,19 @@ def get_hist(self, ifos=None): n_ifos = len(self.hist_ifos) bin_volume = (self.twidth * self.pwidth * self.swidth) ** (n_ifos - 1) - self.hist_max = -1. * numpy.inf + self.hist_max = -1.0 * numpy.inf # Read histogram for each ifo, to use if that ifo has smallest SNR in # the coinc for ifo in self.hist_ifos: - # renormalise to PDF - self.weights[ifo] = \ - (weights[ifo] / (weights[ifo].sum() * bin_volume)) + self.weights[ifo] = weights[ifo] / (weights[ifo].sum() * bin_volume) self.weights[ifo] = self.weights[ifo].astype(numpy.float32) if param[ifo].dtype == numpy.int8: # Older style, incorrectly sorted histogram file ncol = param[ifo].shape[1] - self.pdtype = [ - ("c%s" % i, param[ifo].dtype) for i in range(ncol) - ] + self.pdtype = [("c%s" % i, param[ifo].dtype) for i in range(ncol)] self.param_bin[ifo] = numpy.zeros( len(self.weights[ifo]), dtype=self.pdtype ) @@ -632,19 +628,13 @@ def update_file(self, key): If others are used (i.e. this statistic is inherited), they will need updated separately """ - if 'phasetd_newsnr' in key and not len(self.ifos) == 1: - if ''.join(sorted(self.ifos)) not in key: + if "phasetd_newsnr" in key and not len(self.ifos) == 1: + if "".join(sorted(self.ifos)) not in key: logger.debug( - "%s file is not used for %s statistic", - key, - ''.join(self.ifos) + "%s file is not used for %s statistic", key, "".join(self.ifos) ) return False - logger.info( - "Updating %s statistic %s file", - ''.join(self.ifos), - key - ) + logger.info("Updating %s statistic %s file", "".join(self.ifos), key) # This is a PhaseTDStatistic file which needs updating self.get_hist() return True @@ -667,6 +657,7 @@ def logsignalrate(self, stats, shift, to_shift): ------- value: log of coinc signal rate density for the given single-ifo triggers and time shifts + """ # Convert time shift vector to dict, as hist ifos and self.ifos may # not be in same order @@ -682,24 +673,26 @@ def logsignalrate(self, stats, shift, to_shift): ) smin = snrs.argmin(axis=0) # Store a list of the triggers using each ifo as reference - rtypes = { - ifo: numpy.where(smin == j)[0] for j, ifo in enumerate(self.ifos) - } + rtypes = {ifo: numpy.where(smin == j)[0] for j, ifo in enumerate(self.ifos)} # Get reference ifo information rate = numpy.zeros(len(shift), dtype=numpy.float32) - ps = {ifo: numpy.array(stats[ifo]['coa_phase'], - dtype=numpy.float32, ndmin=1) - for ifo in self.ifos} - ts = {ifo: numpy.array(stats[ifo]['end_time'], - dtype=numpy.float64, ndmin=1) - for ifo in self.ifos} - ss = {ifo: numpy.array(stats[ifo]['snr'], - dtype=numpy.float32, ndmin=1) - for ifo in self.ifos} - sigs = {ifo: numpy.array(stats[ifo]['sigmasq'], - dtype=numpy.float32, ndmin=1) - for ifo in self.ifos} + ps = { + ifo: numpy.array(stats[ifo]["coa_phase"], dtype=numpy.float32, ndmin=1) + for ifo in self.ifos + } + ts = { + ifo: numpy.array(stats[ifo]["end_time"], dtype=numpy.float64, ndmin=1) + for ifo in self.ifos + } + ss = { + ifo: numpy.array(stats[ifo]["snr"], dtype=numpy.float32, ndmin=1) + for ifo in self.ifos + } + sigs = { + ifo: numpy.array(stats[ifo]["sigmasq"], dtype=numpy.float32, ndmin=1) + for ifo in self.ifos + } for ref_ifo in self.ifos: rtype = rtypes[ref_ifo] pref = ps[ref_ifo] @@ -788,12 +781,10 @@ def logsignalrate(self, stats, shift, to_shift): # These weren't in our histogram so give them max penalty # instead of random value - missed = numpy.where( - self.param_bin[ref_ifo][loc] != nbinned - )[0] + missed = numpy.where(self.param_bin[ref_ifo][loc] != nbinned)[0] rate[rtype[missed]] = self.max_penalty # Scale by signal population SNR - rate[rtype] *= (sref[rtype] / self.ref_snr) ** -4. + rate[rtype] *= (sref[rtype] / self.ref_snr) ** -4.0 return numpy.log(rate) @@ -814,6 +805,7 @@ def single(self, trigs): ------- numpy.ndarray Array of single detector parameter values + """ sngl_stat = self.get_sngl_ranking(trigs) singles = numpy.zeros(len(sngl_stat), dtype=self.single_dtype) @@ -841,26 +833,23 @@ def rank_stat_single(self, single_info): ------- numpy.ndarray The array of single detector statistics + """ return single_info[1]["snglstat"] - def rank_stat_coinc( - self, sngls_list, slide, step, to_shift, **kwargs - ): # pylint:disable=unused-argument + def rank_stat_coinc(self, sngls_list, slide, step, to_shift, **kwargs): # pylint:disable=unused-argument """ Calculate the coincident detection statistic, defined in Eq 2 of [Nitz et al, 2017](https://doi.org/10.3847/1538-4357/aa8f50). """ rstat = sum(s[1]["snglstat"] ** 2 for s in sngls_list) - cstat = rstat + 2. * self.logsignalrate( + cstat = rstat + 2.0 * self.logsignalrate( dict(sngls_list), slide * step, to_shift ) cstat[cstat < 0] = 0 return cstat**0.5 - def coinc_lim_for_thresh( - self, sngls_list, thresh, limifo, **kwargs - ): # pylint:disable=unused-argument + def coinc_lim_for_thresh(self, sngls_list, thresh, limifo, **kwargs): # pylint:disable=unused-argument """ Optimization function to identify coincs too quiet to be of interest. Calculate the required single detector statistic to exceed the @@ -873,12 +862,10 @@ def coinc_lim_for_thresh( if not self.has_hist: self.get_hist() - fixed_stat_sq = sum( - [b["snglstat"] ** 2 for a, b in sngls_list if a != limifo] - ) - s1 = thresh ** 2. - fixed_stat_sq + fixed_stat_sq = sum([b["snglstat"] ** 2 for a, b in sngls_list if a != limifo]) + s1 = thresh**2.0 - fixed_stat_sq # Assume best case scenario and use maximum signal rate - s1 -= 2. * self.hist_max + s1 -= 2.0 * self.hist_max s1[s1 < 0] = 0 return s1**0.5 @@ -910,21 +897,18 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): ifos: list of strs, not used here The list of detector names kwargs: values and features needed for the statistic + """ if not files: raise RuntimeError("Files not specified") - PhaseTDStatistic.__init__( - self, sngl_ranking, files=files, ifos=ifos, **kwargs - ) + PhaseTDStatistic.__init__(self, sngl_ranking, files=files, ifos=ifos, **kwargs) # Get the single-detector rates fit files # the stat file attributes are hard-coded as '%{ifo}-fit_coeffs' parsed_attrs = [f.split("-") for f in self.files.keys()] self.bg_ifos = [ - at[0] - for at in parsed_attrs - if (len(at) == 2 and at[1] == "fit_coeffs") + at[0] for at in parsed_attrs if (len(at) == 2 and at[1] == "fit_coeffs") ] if not len(self.bg_ifos): raise RuntimeError( @@ -948,17 +932,11 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): self.min_snr = numpy.inf # Some modifiers for the statistic to get it into a nice range - self.benchmark_lograte = float( - self.kwargs.get("benchmark_lograte", -14.6) - ) - self.min_stat = float( - self.kwargs.get("minimum_statistic_cutoff", -30.) - ) + self.benchmark_lograte = float(self.kwargs.get("benchmark_lograte", -14.6)) + self.min_stat = float(self.kwargs.get("minimum_statistic_cutoff", -30.0)) # Modifier to get a sensible value of the fit slope below threshold - self.alphabelow = float( - self.kwargs.get("alpha_below_thresh", numpy.inf) - ) + self.alphabelow = float(self.kwargs.get("alpha_below_thresh", numpy.inf)) # This will be used to keep track of the template number being used self.curr_tnum = None @@ -966,9 +944,7 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): # Applies a constant offset to all statistic values in a given instance. # This can be used to e.g. change relative rankings between different # event types. Default is zero offset. - self.stat_correction = float( - self.kwargs.get("statistic_correction", 0) - ) + self.stat_correction = float(self.kwargs.get("statistic_correction", 0)) # Go through the keywords and add class information as needed: if self.kwargs["sensitive_volume"]: @@ -981,7 +957,7 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): [self.fits_by_tid[ifo]["median_sigma"] for ifo in ref_ifos], axis=0, ) - self.benchmark_logvol = 3. * numpy.log(hl_net_med_sigma) + self.benchmark_logvol = 3.0 * numpy.log(hl_net_med_sigma) if self.kwargs["dq"]: # Reweight the noise rate by the dq reweighting factor @@ -989,7 +965,7 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): self.dq_bin_by_tid = {} self.dq_state_segments = None self.low_latency = False - self.single_dtype.append(('dq_state', int)) + self.single_dtype.append(("dq_state", int)) for ifo in self.ifos: key = f"{ifo}-dq_stat_info" @@ -1030,9 +1006,10 @@ def assign_template_bins(self, key): statistic file key string Returns - --------- + ------- bin_dict: dict of strs Dictionary containing the bin name for each template id + """ ifo = key.split("-")[0] with h5py.File(self.files[key], "r") as dq_file: @@ -1058,7 +1035,7 @@ def assign_dq_rates(self, key): statistic file key string Returns - --------- + ------- dq_dict: dict of {time: dq_value} dicts for each bin Dictionary containing the mapping between the time and the dq value for each individual bin. @@ -1083,8 +1060,7 @@ def setup_segments(self, key): dq_state_segs_dict = {} for k in ifo_grp["dq_segments"].keys(): seg_dict = {} - seg_dict["start"] = \ - ifo_grp[f"dq_segments/{k}/segment_starts"][:] + seg_dict["start"] = ifo_grp[f"dq_segments/{k}/segment_starts"][:] seg_dict["end"] = ifo_grp[f"dq_segments/{k}/segment_ends"][:] dq_state_segs_dict[k] = seg_dict @@ -1092,7 +1068,6 @@ def setup_segments(self, key): def find_dq_noise_rate(self, trigs, dq_state): """Get dq values for a specific ifo and dq states""" - dq_val = numpy.ones(len(dq_state)) if self.curr_ifo in self.dq_rates_by_state: @@ -1121,28 +1096,29 @@ def find_dq_state_by_time(self, ifo, times): def check_low_latency(self, key): """ Check if the statistic file indicates low latency mode. + Parameters ---------- key: str Statistic file key string. + Returns ------- None + """ - ifo = key.split('-')[0] - with h5py.File(self.files[key], 'r') as dq_file: + ifo = key.split("-")[0] + with h5py.File(self.files[key], "r") as dq_file: ifo_grp = dq_file[ifo] - if 'dq_segments' not in ifo_grp.keys(): + if "dq_segments" not in ifo_grp.keys(): # if segs are not in file, we must be in LL if self.dq_state_segments is not None: raise ValueError( - 'Either all dq stat files must have segments or none' + "Either all dq stat files must have segments or none" ) self.low_latency = True elif self.low_latency: - raise ValueError( - 'Either all dq stat files must have segments or none' - ) + raise ValueError("Either all dq stat files must have segments or none") def reassign_rate(self, ifo): """ @@ -1150,27 +1126,28 @@ def reassign_rate(self, ifo): normalised number. Parameters - ----------- + ---------- ifo: str The ifo to consider. + """ - with h5py.File(self.files[f'{ifo}-fit_coeffs'], 'r') as coeff_file: - analysis_time = float(coeff_file.attrs['analysis_time']) - fbt = 'fit_by_template' in coeff_file + with h5py.File(self.files[f"{ifo}-fit_coeffs"], "r") as coeff_file: + analysis_time = float(coeff_file.attrs["analysis_time"]) + fbt = "fit_by_template" in coeff_file - self.fits_by_tid[ifo]['smoothed_rate_above_thresh'] /= analysis_time - self.fits_by_tid[ifo]['smoothed_rate_in_template'] /= analysis_time + self.fits_by_tid[ifo]["smoothed_rate_above_thresh"] /= analysis_time + self.fits_by_tid[ifo]["smoothed_rate_in_template"] /= analysis_time # The by-template fits may have been stored in the smoothed fits file if fbt: - self.fits_by_tid[ifo]['fit_by_rate_above_thresh'] /= analysis_time - self.fits_by_tid[ifo]['fit_by_rate_in_template'] /= analysis_time + self.fits_by_tid[ifo]["fit_by_rate_above_thresh"] /= analysis_time + self.fits_by_tid[ifo]["fit_by_rate_in_template"] /= analysis_time def assign_fits(self, ifo): """ Extract fits from single-detector rate fit files Parameters - ----------- + ---------- ifo: str The detector to get fits for. @@ -1179,6 +1156,7 @@ def assign_fits(self, ifo): rate_dict: dict A dictionary containing the fit information in the `alpha`, `rate` and `thresh` keys. + """ coeff_file = h5py.File(self.files[f"{ifo}-fit_coeffs"], "r") template_id = coeff_file["template_id"][:] @@ -1187,15 +1165,13 @@ def assign_fits(self, ifo): tid_sort = numpy.argsort(template_id) fits_by_tid_dict = {} - fits_by_tid_dict["smoothed_fit_coeff"] = coeff_file["fit_coeff"][:][ - tid_sort - ] + fits_by_tid_dict["smoothed_fit_coeff"] = coeff_file["fit_coeff"][:][tid_sort] fits_by_tid_dict["smoothed_rate_above_thresh"] = coeff_file[ "count_above_thresh" ][:][tid_sort].astype(float) - fits_by_tid_dict["smoothed_rate_in_template"] = coeff_file[ - "count_in_template" - ][:][tid_sort].astype(float) + fits_by_tid_dict["smoothed_rate_in_template"] = coeff_file["count_in_template"][ + : + ][tid_sort].astype(float) if self.kwargs["sensitive_volume"]: fits_by_tid_dict["median_sigma"] = coeff_file["median_sigma"][:][ tid_sort @@ -1204,9 +1180,7 @@ def assign_fits(self, ifo): # The by-template fits may have been stored in the smoothed fits file if "fit_by_template" in coeff_file: coeff_fbt = coeff_file["fit_by_template"] - fits_by_tid_dict["fit_by_fit_coeff"] = coeff_fbt["fit_coeff"][:][ - tid_sort - ] + fits_by_tid_dict["fit_by_fit_coeff"] = coeff_fbt["fit_coeff"][:][tid_sort] fits_by_tid_dict["fit_by_rate_above_thresh"] = coeff_fbt[ "count_above_thresh" ][:][tid_sort].astype(float) @@ -1214,7 +1188,6 @@ def assign_fits(self, ifo): "count_in_template" ][:][tid_sort].astype(float) - # Keep the fit threshold in fits_by_tid fits_by_tid_dict["thresh"] = coeff_file.attrs["stat_threshold"] @@ -1233,7 +1206,7 @@ def update_file(self, key): if PhaseTDStatistic.update_file(self, key): return True - if key.endswith('-fit_coeffs'): + if key.endswith("-fit_coeffs"): # This is a ExpFitStatistic file which needs updating # Which ifo is it? ifo = key[:2] @@ -1241,32 +1214,20 @@ def update_file(self, key): if self.kwargs["normalize_fit_rate"]: self.reassign_rate(ifo) self.get_ref_vals(ifo) - logger.info( - "Updating %s statistic %s file", - ''.join(self.ifos), - key - ) + logger.info("Updating %s statistic %s file", "".join(self.ifos), key) return True # Is the key a KDE statistic file that we update here? - if key.endswith('kde_file'): - logger.info( - "Updating %s statistic %s file", - ''.join(self.ifos), - key - ) - kde_style = key.split('-')[0] + if key.endswith("kde_file"): + logger.info("Updating %s statistic %s file", "".join(self.ifos), key) + kde_style = key.split("-")[0] self.assign_kdes(kde_style) return True # We also need to check if the DQ files have updated - if key.endswith('dq_stat_info'): - ifo = key.split('-')[0] - logger.info( - "Updating %s statistic %s file", - ifo, - key - ) + if key.endswith("dq_stat_info"): + ifo = key.split("-")[0] + logger.info("Updating %s statistic %s file", ifo, key) self.dq_rates_by_state[ifo] = self.assign_dq_rates(key) self.dq_bin_by_tid[ifo] = self.assign_template_bins(key) return True @@ -1280,11 +1241,12 @@ def get_ref_vals(self, ifo): This is stored in `self.alphamax[ifo]` in the class instance. Parameters - ----------- + ---------- ifo: str The detector to get fits for. + """ - self.alphamax[ifo] = self.fits_by_tid[ifo]['smoothed_fit_coeff'].max() + self.alphamax[ifo] = self.fits_by_tid[ifo]["smoothed_fit_coeff"].max() def find_fits(self, trigs): """ @@ -1300,18 +1262,19 @@ def find_fits(self, trigs): If multiple templates are in play we must return arrays. Returns - -------- + ------- alphai: float or numpy array The alpha fit value(s) ratei: float or numpy array The rate fit value(s) thresh: float or numpy array The thresh fit value(s) + """ try: ifo = trigs.ifo except AttributeError: - ifo = trigs.get('ifo', None) + ifo = trigs.get("ifo", None) if ifo is None: ifo = self.ifos[0] assert ifo in self.ifos @@ -1332,9 +1295,7 @@ def find_kdes(self): # and 'template-kde_file' parsed_attrs = [f.split("-") for f in self.files.keys()] self.kde_names = [ - at[0] - for at in parsed_attrs - if (len(at) == 2 and at[1] == "kde_file") + at[0] for at in parsed_attrs if (len(at) == 2 and at[1] == "kde_file") ] assert sorted(self.kde_names) == ["signal", "template"], ( "Two KDE stat files are required, they should have stat attr " @@ -1346,9 +1307,10 @@ def assign_kdes(self, kname): Extract values from KDE files Parameters - ----------- + ---------- kname: str Used to label the kde files. + """ with h5py.File(self.files[kname + "-kde_file"], "r") as kde_file: self.kde_by_tid[kname + "_kdevals"] = kde_file["data_kde"][:] @@ -1371,14 +1333,15 @@ def lognoiserate(self, trigs): and rescale by the fitted coefficients alpha and rate Parameters - ----------- + ---------- trigs: dict of numpy.ndarrays, h5py group or similar dict-like object Object holding single detector trigger information. Returns - --------- + ------- lognoisel: numpy.array Array of log noise rate density for each input trigger. + """ # What is the template number currently being used? try: @@ -1392,9 +1355,7 @@ def lognoiserate(self, trigs): alphai, ratei, thresh = self.find_fits(trigs) sngl_stat = self.get_sngl_ranking(trigs) lognoisel = ( - -alphai * (sngl_stat - thresh) - + numpy.log(alphai) - + numpy.log(ratei) + -alphai * (sngl_stat - thresh) + numpy.log(alphai) + numpy.log(ratei) ) if not numpy.isinf(self.alphabelow): @@ -1426,12 +1387,12 @@ def single(self, trigs): ------- numpy.ndarray The array of single detector values - """ + """ try: self.curr_ifo = trigs.ifo except AttributeError: - self.curr_ifo = trigs.get('ifo', None) + self.curr_ifo = trigs.get("ifo", None) if self.curr_ifo is None: self.curr_ifo = self.ifos[0] assert self.curr_ifo in self.ifos @@ -1465,22 +1426,22 @@ def single(self, trigs): from pycbc.conversions import mchirp_from_mass1_mass2 try: - mass1 = trigs.param['mass1'] - mass2 = trigs.param['mass2'] + mass1 = trigs.param["mass1"] + mass2 = trigs.param["mass2"] except AttributeError: - mass1 = trigs['mass1'] - mass2 = trigs['mass2'] + mass1 = trigs["mass1"] + mass2 = trigs["mass2"] self.curr_mchirp = mchirp_from_mass1_mass2(mass1, mass2) if self.kwargs["dq"]: if self.low_latency: # trigs should already have a dq state assigned - singles['dq_state'] = trigs['dq_state'][:] + singles["dq_state"] = trigs["dq_state"][:] else: - singles['dq_state'] = self.find_dq_state_by_time( - self.curr_ifo, trigs['end_time'][:] + singles["dq_state"] = self.find_dq_state_by_time( + self.curr_ifo, trigs["end_time"][:] ) - dq_rate = self.find_dq_noise_rate(trigs, singles['dq_state']) + dq_rate = self.find_dq_noise_rate(trigs, singles["dq_state"]) dq_rate = numpy.maximum(dq_rate, 1) sngl_stat += numpy.log(dq_rate) @@ -1507,6 +1468,7 @@ def sensitive_volume_factor(self, sngls): Array of values for the network log-volume factor. This is the log of the cube of the sensitive distance (sigma), divided by a benchmark volume. + """ # Get benchmark log volume as single-ifo information : # benchmark_logvol for a given template is not ifo-dependent, so @@ -1517,16 +1479,14 @@ def sensitive_volume_factor(self, sngls): # any are nan, they are all nan if any(numpy.isnan(benchmark_logvol)): # This can be the case in pycbc live if there are no triggers - # from this template in the trigger fits file. If so, assume + # from this template in the trigger fits file. If so, assume # that sigma for the triggers being ranked is # representative of the benchmark network. return 0 # Network sensitivity for a given coinc type is approximately # determined by the least sensitive ifo - network_sigmasq = numpy.amin( - [sngl[1]["sigmasq"] for sngl in sngls], axis=0 - ) + network_sigmasq = numpy.amin([sngl[1]["sigmasq"] for sngl in sngls], axis=0) # Volume \propto sigma^3 or sigmasq^1.5 network_logvol = 1.5 * numpy.log(network_sigmasq) - benchmark_logvol @@ -1548,6 +1508,7 @@ def logsignalrate_shared(self, sngls_info): sr_factor: numpy.ndarray Array of values to be added to the ranking statistic when taking various signal rate factors into account + """ # Other features affecting the signal rate sr_factor = 0 @@ -1563,7 +1524,7 @@ def logsignalrate_shared(self, sngls_info): else: # curr_mchirp will be a number mchirp = min(self.curr_mchirp, self.mcm) - sr_factor += numpy.log((mchirp / 20.) ** (11. / 3.)) + sr_factor += numpy.log((mchirp / 20.0) ** (11.0 / 3.0)) if self.kwargs["kde"]: # KDE reweighting @@ -1585,6 +1546,7 @@ def rank_stat_single(self, single_info): ------- numpy.ndarray The array of single detector statistics + """ sngls = single_info[1] @@ -1607,9 +1569,7 @@ def rank_stat_single(self, single_info): loglr[loglr < self.min_stat] = self.min_stat return loglr - def rank_stat_coinc( - self, s, slide, step, to_shift, **kwargs - ): # pylint:disable=unused-argument + def rank_stat_coinc(self, s, slide, step, to_shift, **kwargs): # pylint:disable=unused-argument """ Calculate the coincident detection statistic. @@ -1644,7 +1604,7 @@ def rank_stat_coinc( ln_noise_rate = coinc_rate.combination_noise_lograte( sngl_dict, kwargs["time_addition"], - dets=kwargs.get('dets', None), + dets=kwargs.get("dets"), ) ln_noise_rate -= self.benchmark_lograte @@ -1663,7 +1623,7 @@ def rank_stat_coinc( noise_twindow = coinc_rate.multiifo_noise_coincident_area( self.hist_ifos, kwargs["time_addition"], - dets=kwargs.get('dets', None), + dets=kwargs.get("dets"), ) # Volume is the allowed time difference window, multiplied by 2pi # for each phase difference dimension and by allowed range of SNR @@ -1671,9 +1631,7 @@ def rank_stat_coinc( # dimensions for both phase and SNR n_ifos = len(self.hist_ifos) snr_range = (self.srbmax - self.srbmin) * self.swidth - hist_vol = noise_twindow * (2. * numpy.pi * snr_range) ** ( - n_ifos - 1 - ) + hist_vol = noise_twindow * (2.0 * numpy.pi * snr_range) ** (n_ifos - 1) # Noise PDF is 1/volume, assuming a uniform distribution of noise # coincs ln_noise_rate -= numpy.log(hist_vol) @@ -1696,9 +1654,7 @@ def rank_stat_coinc( return loglr - def coinc_lim_for_thresh( - self, s, thresh, limifo, **kwargs - ): # pylint:disable=unused-argument + def coinc_lim_for_thresh(self, s, thresh, limifo, **kwargs): # pylint:disable=unused-argument """ Optimization function to identify coincs too quiet to be of interest @@ -1732,6 +1688,7 @@ def coinc_lim_for_thresh( numpy.array The minimum snglstat values required in the 'pivot' detector in order to reach the threshold specified + """ # Safety against subclassing and not rethinking this allowed_names = ["ExpFitStatistic"] @@ -1760,9 +1717,7 @@ def coinc_lim_for_thresh( if not self.has_hist: self.get_hist() # Assume best-case scenario and use maximum signal rate - ln_s = numpy.log( - self.hist_max * (self.min_snr / self.ref_snr) ** -4. - ) + ln_s = numpy.log(self.hist_max * (self.min_snr / self.ref_snr) ** -4.0) # Shared info is the same as in the coinc calculation ln_s += self.logsignalrate_shared(s) @@ -1801,19 +1756,16 @@ def __init__(self, sngl_ranking, files=None, ifos=None, **kwargs): ifos: list of strs, not used here The list of detector names + """ - ExpFitStatistic.__init__( - self, sngl_ranking, files=files, ifos=ifos, **kwargs - ) + ExpFitStatistic.__init__(self, sngl_ranking, files=files, ifos=ifos, **kwargs) # for low-mass templates the exponential slope alpha \approx 6 - self.alpharef = 6. + self.alpharef = 6.0 self.single_increasing = True self.single_dtype = numpy.float32 # Modifier to get a sensible value of the fit slope below threshold - self.alphabelow = float( - self.kwargs.get("alpha_below_thresh", numpy.inf) - ) + self.alphabelow = float(self.kwargs.get("alpha_below_thresh", numpy.inf)) def single(self, trigs): """ @@ -1828,11 +1780,12 @@ def single(self, trigs): ------- numpy.ndarray The array of single detector values + """ logr_n = self.lognoiserate(trigs) _, _, thresh = self.find_fits(trigs) # shift by log of reference slope alpha - logr_n += -1. * numpy.log(self.alpharef) + logr_n += -1.0 * numpy.log(self.alpharef) # add threshold and rescale by reference slope stat = thresh - (logr_n / self.alpharef) return numpy.array(stat, ndmin=1, dtype=numpy.float32) @@ -1851,16 +1804,15 @@ def rank_stat_single(self, single_info): ------- numpy.ndarray The array of single detector statistics + """ if self.single_increasing: sngl_multiifo = single_info[1] else: - sngl_multiifo = -1. * single_info[1] + sngl_multiifo = -1.0 * single_info[1] return sngl_multiifo - def rank_stat_coinc( - self, s, slide, step, to_shift, **kwargs - ): # pylint:disable=unused-argument + def rank_stat_coinc(self, s, slide, step, to_shift, **kwargs): # pylint:disable=unused-argument """ Calculate the coincident detection statistic. @@ -1878,13 +1830,12 @@ def rank_stat_coinc( ------- numpy.ndarray Array of coincident ranking statistic values + """ # scale by 1/sqrt(number of ifos) to resemble network SNR return sum(sngl[1] for sngl in s) / (len(s) ** 0.5) - def coinc_lim_for_thresh( - self, s, thresh, limifo, **kwargs - ): # pylint:disable=unused-argument + def coinc_lim_for_thresh(self, s, thresh, limifo, **kwargs): # pylint:disable=unused-argument """ Optimization function to identify coincs too quiet to be of interest @@ -1906,6 +1857,7 @@ def coinc_lim_for_thresh( numpy.ndarray Array of limits on the limifo single statistic to exceed thresh. + """ # Safety against subclassing and not rethinking this allowed_names = ["ExpFitCombinedSNR"] @@ -1941,6 +1893,7 @@ def get_statistic(stat): ------ RuntimeError If the string is not recognized as corresponding to a Stat subclass + """ try: return statistic_dict[stat] @@ -1955,7 +1908,7 @@ def insert_statistic_option_group(parser, default_ranking_statistic=None): Adds the options used to initialize a PyCBC Stat class. Parameters - ----------- + ---------- parser : object OptionParser instance. default_ranking_statisic : str @@ -1963,9 +1916,10 @@ def insert_statistic_option_group(parser, default_ranking_statistic=None): option. The option is no longer required if a default is provided. Returns - -------- + ------- strain_opt_group : optparser.argument_group The argument group that is added to the parser. + """ statistic_opt_group = parser.add_argument_group( "Options needed to initialize a PyCBC Stat class for computing the " @@ -2033,6 +1987,7 @@ def parse_statistic_feature_options(stat_features, stat_kwarg_list): ------- stat_kwarg_dict : dict Statistic keywords in dict format + """ stat_kwarg_dict = {} @@ -2080,6 +2035,7 @@ def get_statistic_from_opts(opts, ifos): ------- class Subclass of Stat base class + """ # Allow None inputs if opts.statistic_files is None: @@ -2089,8 +2045,7 @@ def get_statistic_from_opts(opts, ifos): # flatten the list of lists of filenames to a single list (may be empty) # if needed (e.g. not calling get_statistic_from_opts in a loop) - if len(opts.statistic_files) > 0 and \ - isinstance(opts.statistic_files[0], list): + if len(opts.statistic_files) > 0 and isinstance(opts.statistic_files[0], list): opts.statistic_files = sum(opts.statistic_files, []) extra_kwargs = parse_statistic_feature_options( diff --git a/pycbc/events/threshold_cpu.py b/pycbc/events/threshold_cpu.py index 9a812d791eb..65b0caec7d2 100644 --- a/pycbc/events/threshold_cpu.py +++ b/pycbc/events/threshold_cpu.py @@ -22,20 +22,23 @@ # ============================================================================= # import logging + import numpy -from .simd_threshold_cython import parallel_thresh_cluster, parallel_threshold -from .eventmgr import _BaseThresholdCluster + from .. import opt +from .eventmgr import _BaseThresholdCluster +from .simd_threshold_cython import parallel_thresh_cluster, parallel_threshold -logger = logging.getLogger('pycbc.events.threshold_cpu') +logger = logging.getLogger("pycbc.events.threshold_cpu") l2_cache_size = opt.get_l2_cache_size() if l2_cache_size is not None: - default_segsize = l2_cache_size // numpy.dtype('complex64').itemsize + default_segsize = l2_cache_size // numpy.dtype("complex64").itemsize else: # Seems to work for Sandy Bridge/Ivy Bridge/Haswell, for now? default_segsize = 32768 + def threshold_numpy(series, value): arr = series.data locs = numpy.where(arr.real**2 + arr.imag**2 > value**2)[0] @@ -46,6 +49,8 @@ def threshold_numpy(series, value): outl = None outv = None count = None + + def threshold_inline(series, value): arr = numpy.array(series.data, copy=False, dtype=numpy.complex64) global outl, outv, count @@ -60,8 +65,8 @@ def threshold_inline(series, value): num = count[0] if num > 0: return outl[0:num], outv[0:num] - else: - return numpy.array([], numpy.uint32), numpy.array([], numpy.float32) + return numpy.array([], numpy.uint32), numpy.array([], numpy.float32) + # threshold_numpy can also be used here, but for now we use the inline code # in all instances. Not sure why we're defining threshold *and* threshold_only @@ -69,25 +74,30 @@ def threshold_inline(series, value): threshold = threshold_inline threshold_only = threshold_inline + class CPUThresholdCluster(_BaseThresholdCluster): def __init__(self, series): - self.series = numpy.array(series.data, copy=False, - dtype=numpy.complex64) + self.series = numpy.array(series.data, copy=False, dtype=numpy.complex64) self.slen = numpy.uint32(len(series)) self.outv = numpy.zeros(self.slen, numpy.complex64) self.outl = numpy.zeros(self.slen, numpy.uint32) self.segsize = numpy.uint32(default_segsize) def threshold_and_cluster(self, threshold, window): - self.count = parallel_thresh_cluster(self.series, self.slen, - self.outv, self.outl, - numpy.float32(threshold), - numpy.uint32(window), - self.segsize) + self.count = parallel_thresh_cluster( + self.series, + self.slen, + self.outv, + self.outl, + numpy.float32(threshold), + numpy.uint32(window), + self.segsize, + ) if self.count > 0: - return self.outv[0:self.count], self.outl[0:self.count] - else: - return numpy.array([], dtype = numpy.complex64), numpy.array([], dtype = numpy.uint32) + return self.outv[0 : self.count], self.outl[0 : self.count] + return numpy.array([], dtype=numpy.complex64), numpy.array( + [], dtype=numpy.uint32 + ) def _threshold_cluster_factory(series): diff --git a/pycbc/events/threshold_cuda.py b/pycbc/events/threshold_cuda.py index 9911a9f85f6..11f3f9ab4df 100644 --- a/pycbc/events/threshold_cuda.py +++ b/pycbc/events/threshold_cuda.py @@ -22,14 +22,18 @@ # ============================================================================= # import logging -import numpy, mako.template -from pycuda.tools import dtype_to_ctype -from pycuda.elementwise import ElementwiseKernel + +import mako.template +import numpy from pycuda.compiler import SourceModule -from .eventmgr import _BaseThresholdCluster +from pycuda.elementwise import ElementwiseKernel +from pycuda.tools import dtype_to_ctype + import pycbc.scheme -logger = logging.getLogger('pycbc.events.threshold_cuda') +from .eventmgr import _BaseThresholdCluster + +logger = logging.getLogger("pycbc.events.threshold_cuda") threshold_op = """ if (i == 0) @@ -45,35 +49,45 @@ """ threshold_kernel = ElementwiseKernel( - " %(tp_in)s *in, %(tp_out1)s *outv, %(tp_out2)s *outl, %(tp_th)s threshold, %(tp_n)s *bn" % { - "tp_in": dtype_to_ctype(numpy.complex64), - "tp_out1": dtype_to_ctype(numpy.complex64), - "tp_out2": dtype_to_ctype(numpy.uint32), - "tp_th": dtype_to_ctype(numpy.float32), - "tp_n": dtype_to_ctype(numpy.uint32), - }, - threshold_op, - "getstuff") + " %(tp_in)s *in, %(tp_out1)s *outv, %(tp_out2)s *outl, %(tp_th)s threshold, %(tp_n)s *bn" + % { + "tp_in": dtype_to_ctype(numpy.complex64), + "tp_out1": dtype_to_ctype(numpy.complex64), + "tp_out2": dtype_to_ctype(numpy.uint32), + "tp_th": dtype_to_ctype(numpy.float32), + "tp_n": dtype_to_ctype(numpy.uint32), + }, + threshold_op, + "getstuff", +) import pycuda.driver as drv -class T(): + +class T: pass + tn = T() tv = T() tl = T() # This avoids this code running if in the documentation build process, # and we don't have pycuda installed -if type(drv).__name__ not in ('MagicMock', '_MockModule'): - n = drv.pagelocked_empty((1), numpy.uint32, mem_flags=drv.host_alloc_flags.DEVICEMAP) +if type(drv).__name__ not in ("MagicMock", "_MockModule"): + n = drv.pagelocked_empty( + (1), numpy.uint32, mem_flags=drv.host_alloc_flags.DEVICEMAP + ) nptr = numpy.intp(n.base.get_device_pointer()) - val = drv.pagelocked_empty((4096*256), numpy.complex64, mem_flags=drv.host_alloc_flags.DEVICEMAP) + val = drv.pagelocked_empty( + (4096 * 256), numpy.complex64, mem_flags=drv.host_alloc_flags.DEVICEMAP + ) vptr = numpy.intp(val.base.get_device_pointer()) - loc = drv.pagelocked_empty((4096*256), numpy.int32, mem_flags=drv.host_alloc_flags.DEVICEMAP) + loc = drv.pagelocked_empty( + (4096 * 256), numpy.int32, mem_flags=drv.host_alloc_flags.DEVICEMAP + ) lptr = numpy.intp(loc.base.get_device_pointer()) tn.gpudata = nptr @@ -215,11 +229,15 @@ class T(): """) tfn_cache = {} + + def get_tkernel(slen, window): if window < 32: - raise ValueError("GPU threshold kernel does not support a window smaller than 32 samples") + raise ValueError( + "GPU threshold kernel does not support a window smaller than 32 samples" + ) - elif window <= 4096: + if window <= 4096: nt = 128 elif window <= 16384: nt = 256 @@ -245,6 +263,7 @@ def get_tkernel(slen, window): tfn_cache[(nt, nb)] = (fn, fn2) return tfn_cache[(nt, nb)], nt, nb + def threshold_and_cluster(series, threshold, window): outl = tl.gpudata outv = tv.gpudata @@ -257,12 +276,21 @@ def threshold_and_cluster(series, threshold, window): cl = loc[0:nb] cv = val[0:nb] - fn.prepared_call((nb, 1), (nt, 1, 1), series, outv, outl, window, threshold,) + fn.prepared_call( + (nb, 1), + (nt, 1, 1), + series, + outv, + outl, + window, + threshold, + ) fn2.prepared_call((1, 1), (nb, 1, 1), outv, outl, threshold, window) pycbc.scheme.mgr.state.context.synchronize() - w = (cl != -1) + w = cl != -1 return cv[w], cl[w] + class CUDAThresholdCluster(_BaseThresholdCluster): def __init__(self, series): self.series = series.data.gpudata @@ -281,12 +309,20 @@ def threshold_and_cluster(self, threshold, window): cl = loc[0:nb] cv = val[0:nb] - fn((nb, 1), (nt, 1, 1), self.series, self.outv, self.outl, window, threshold,) + fn( + (nb, 1), + (nt, 1, 1), + self.series, + self.outv, + self.outl, + window, + threshold, + ) fn2((1, 1), (nb, 1, 1), self.outv, self.outl, threshold, window) pycbc.scheme.mgr.state.context.synchronize() - w = (cl != -1) + w = cl != -1 return cv[w], cl[w] + def _threshold_cluster_factory(series): return CUDAThresholdCluster - diff --git a/pycbc/events/threshold_cupy.py b/pycbc/events/threshold_cupy.py index 547ea9326a4..8def999dfc2 100644 --- a/pycbc/events/threshold_cupy.py +++ b/pycbc/events/threshold_cupy.py @@ -22,9 +22,11 @@ # ============================================================================= # -import cupy as cp import functools + +import cupy as cp import mako.template + from .eventmgr import _BaseThresholdCluster val = None @@ -162,12 +164,15 @@ } """) -@functools.lru_cache(maxsize=None) + +@functools.cache def get_tkernel(slen, window): if window < 32: - raise ValueError("GPU threshold kernel does not support a window smaller than 32 samples") + raise ValueError( + "GPU threshold kernel does not support a window smaller than 32 samples" + ) - elif window <= 4096: + if window <= 4096: nt = 128 elif window <= 16384: nt = 256 @@ -182,24 +187,21 @@ def get_tkernel(slen, window): raise ValueError("More than 1024 blocks not supported yet") fn = cp.RawKernel( - tkernel1.render(chunk=nt), - 'threshold_and_cluster', - backend='nvcc' + tkernel1.render(chunk=nt), "threshold_and_cluster", backend="nvcc" ) fn2 = cp.RawKernel( - tkernel2.render(blocks=nb), - 'threshold_and_cluster2', - backend='nvcc' + tkernel2.render(blocks=nb), "threshold_and_cluster2", backend="nvcc" ) return (fn, fn2), nt, nb + def threshold_and_cluster(series, threshold, window): global val global loc if val is None: - val = cp.zeros(4096*256, dtype=cp.complex64) + val = cp.zeros(4096 * 256, dtype=cp.complex64) if loc is None: - loc = cp.zeros(4096*256, cp.int32) + loc = cp.zeros(4096 * 256, cp.int32) outl = loc outv = val @@ -214,9 +216,10 @@ def threshold_and_cluster(series, threshold, window): fn((nb,), (nt,), (series.data, outv, outl, window, threshold)) fn2((1,), (nb,), (outv, outl, threshold, window)) - w = (cl != -1) + w = cl != -1 return cv[w], cl[w] + class CUDAThresholdCluster(_BaseThresholdCluster): def __init__(self, series): self.series = series @@ -224,9 +227,9 @@ def __init__(self, series): global val global loc if val is None: - val = cp.zeros(4096*256, dtype=cp.complex64) + val = cp.zeros(4096 * 256, dtype=cp.complex64) if loc is None: - loc = cp.zeros(4096*256, cp.int32) + loc = cp.zeros(4096 * 256, cp.int32) self.outl = loc self.outv = val @@ -243,16 +246,12 @@ def threshold_and_cluster(self, threshold, window): fn( (nt, 1, 1), (nb, 1), - (self.series.data, self.outv, self.outl, window, threshold) - ) - fn2( - (nb, 1, 1), - (1, 1), - (self.outv, self.outl, threshold, window) + (self.series.data, self.outv, self.outl, window, threshold), ) - w = (cl != -1) + fn2((nb, 1, 1), (1, 1), (self.outv, self.outl, threshold, window)) + w = cl != -1 return cp.asnumpy(cv[w]), cp.asnumpy(cl[w]) + def _threshold_cluster_factory(series): return CUDAThresholdCluster - diff --git a/pycbc/events/trigger_fits.py b/pycbc/events/trigger_fits.py index e2a6002e05b..3e510a33536 100644 --- a/pycbc/events/trigger_fits.py +++ b/pycbc/events/trigger_fits.py @@ -49,17 +49,19 @@ # Public License for more details. import logging + import numpy from scipy.stats import kstest -logger = logging.getLogger('pycbc.events.trigger_fits') +logger = logging.getLogger("pycbc.events.trigger_fits") + def exponential_fitalpha(vals, thresh, w): """ Maximum likelihood estimator for the fit factor for an exponential decrease model """ - return 1. / (numpy.average(vals, weights=w) - thresh) + return 1.0 / (numpy.average(vals, weights=w) - thresh) def rayleigh_fitalpha(vals, thresh, w): @@ -67,7 +69,7 @@ def rayleigh_fitalpha(vals, thresh, w): Maximum likelihood estimator for the fit factor for a Rayleigh distribution of events """ - return 2. / (numpy.average(vals ** 2., weights=w) - thresh ** 2.) + return 2.0 / (numpy.average(vals**2.0, weights=w) - thresh**2.0) def power_fitalpha(vals, thresh, w): @@ -75,22 +77,23 @@ def power_fitalpha(vals, thresh, w): Maximum likelihood estimator for the fit factor for a power law model """ - return numpy.average(numpy.log(vals/thresh), weights=w) ** -1. + 1. + return numpy.average(numpy.log(vals / thresh), weights=w) ** -1.0 + 1.0 fitalpha_dict = { - 'exponential' : exponential_fitalpha, - 'rayleigh' : rayleigh_fitalpha, - 'power' : power_fitalpha + "exponential": exponential_fitalpha, + "rayleigh": rayleigh_fitalpha, + "power": power_fitalpha, } # measurement standard deviation = (-d^2 log L/d alpha^2)^(-1/2) fitstd_dict = { - 'exponential' : lambda weights, alpha : alpha / sum(weights) ** 0.5, - 'rayleigh' : lambda weights, alpha : alpha / sum(weights) ** 0.5, - 'power' : lambda weights, alpha : (alpha - 1.) / sum(weights) ** 0.5 + "exponential": lambda weights, alpha: alpha / sum(weights) ** 0.5, + "rayleigh": lambda weights, alpha: alpha / sum(weights) ** 0.5, + "power": lambda weights, alpha: (alpha - 1.0) / sum(weights) ** 0.5, } + def fit_above_thresh(distr, vals, thresh=None, weights=None): """ Maximum likelihood fit for the coefficient alpha @@ -120,6 +123,7 @@ def fit_above_thresh(distr, vals, thresh=None, weights=None): Fitted value sigma_alpha : float Standard error in fitted value + """ vals = numpy.array(vals) if thresh is None: @@ -129,9 +133,12 @@ def fit_above_thresh(distr, vals, thresh=None, weights=None): above_thresh = vals >= thresh if numpy.count_nonzero(above_thresh) == 0: # Nothing is above threshold - warn and return -1 - logger.warning("No values are above the threshold, %.2f, " - "maximum is %.2f.", thresh, vals.max()) - return -1., -1. + logger.warning( + "No values are above the threshold, %.2f, maximum is %.2f.", + thresh, + vals.max(), + ) + return -1.0, -1.0 vals = vals[above_thresh] @@ -151,12 +158,12 @@ def fit_above_thresh(distr, vals, thresh=None, weights=None): # a: slope parameter of the fit # t: lower threshold stat value fitfn_dict = { - 'exponential' : lambda x, a, t : a * numpy.exp(-a * (x - t)), - 'rayleigh' : lambda x, a, t : (a * x * \ - numpy.exp(-a * (x ** 2 - t ** 2) / 2.)), - 'power' : lambda x, a, t : (a - 1.) * x ** (-a) * t ** (a - 1.) + "exponential": lambda x, a, t: a * numpy.exp(-a * (x - t)), + "rayleigh": lambda x, a, t: a * x * numpy.exp(-a * (x**2 - t**2) / 2.0), + "power": lambda x, a, t: (a - 1.0) * x ** (-a) * t ** (a - 1.0), } + def fit_fn(distr, xvals, alpha, thresh): """ The fitted function normalized to 1 above threshold @@ -176,20 +183,22 @@ def fit_fn(distr, xvals, alpha, thresh): ------- fit : array of floats Fitted function at the requested xvals + """ xvals = numpy.array(xvals) fit = fitfn_dict[distr](xvals, alpha, thresh) # set fitted values below threshold to 0 - numpy.putmask(fit, xvals < thresh, 0.) + numpy.putmask(fit, xvals < thresh, 0.0) return fit cum_fndict = { - 'exponential' : lambda x, alpha, t : numpy.exp(-alpha * (x - t)), - 'rayleigh' : lambda x, alpha, t : numpy.exp(-alpha * (x ** 2. - t ** 2.) / 2.), - 'power' : lambda x, alpha, t : x ** (1. - alpha) * t ** (alpha - 1.) + "exponential": lambda x, alpha, t: numpy.exp(-alpha * (x - t)), + "rayleigh": lambda x, alpha, t: numpy.exp(-alpha * (x**2.0 - t**2.0) / 2.0), + "power": lambda x, alpha, t: x ** (1.0 - alpha) * t ** (alpha - 1.0), } + def cum_fit(distr, xvals, alpha, thresh): """ Integral of the fitted function above a given value (reverse CDF) @@ -209,18 +218,20 @@ def cum_fit(distr, xvals, alpha, thresh): ------- cum_fit : array of floats Reverse CDF of fitted function at the requested xvals + """ xvals = numpy.array(xvals) cum_fit = cum_fndict[distr](xvals, alpha, thresh) # set fitted values below threshold to 0 - numpy.putmask(cum_fit, xvals < thresh, 0.) + numpy.putmask(cum_fit, xvals < thresh, 0.0) return cum_fit + def tail_threshold(vals, N=1000): """Determine a threshold above which there are N louder values""" vals = numpy.array(vals) if len(vals) < N: - raise RuntimeError('Not enough input values to determine threshold') + raise RuntimeError("Not enough input values to determine threshold") vals.sort() return min(vals[-N:]) @@ -252,14 +263,17 @@ def KS_test(distr, vals, alpha, thresh=None): KS test statistic p-value : float p-value, assumed to be two-tailed + """ vals = numpy.array(vals) if thresh is None: thresh = min(vals) else: vals = vals[vals >= thresh] + def cdf_fn(x): return 1 - cum_fndict[distr](x, alpha, thresh) + return kstest(vals, cdf_fn) @@ -287,8 +301,9 @@ def which_bin(par, minpar, maxpar, nbins, log=False): ------- binind : int Bin index + """ - assert (par >= minpar and par <= maxpar) + assert par >= minpar and par <= maxpar if log: par, minpar, maxpar = numpy.log(par), numpy.log(minpar), numpy.log(maxpar) # par lies some fraction of the way between min and max @@ -304,4 +319,3 @@ def which_bin(par, minpar, maxpar, nbins, log=False): if par == maxpar: binind = nbins - 1 return binind - diff --git a/pycbc/events/triggers.py b/pycbc/events/triggers.py index c21cda49c5a..08f0b726866 100644 --- a/pycbc/events/triggers.py +++ b/pycbc/events/triggers.py @@ -13,48 +13,59 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This modules contains functions for reading single and coincident triggers +""" +This modules contains functions for reading single and coincident triggers from the command line. """ + import logging + import h5py import numpy +import pycbc.detector from pycbc import conversions, pnutils from pycbc.events import coinc -import pycbc.detector -logger = logging.getLogger('pycbc.events.triggers') +logger = logging.getLogger("pycbc.events.triggers") def insert_bank_bins_option_group(parser): - """ Add options to the optparser object for selecting templates in bins. + """ + Add options to the optparser object for selecting templates in bins. Parameters - ----------- + ---------- parser : object OptionParser instance. + """ - bins_group = parser.add_argument_group( - "Options for selecting templates in bins.") - bins_group.add_argument("--bank-bins", nargs="+", default=None, - help="Ordered list of mass bin upper boundaries. " - "An ordered list of type-boundary pairs, " - "applied sequentially. Must provide a name " - "(can be any unique string for tagging " - "purposes), the parameter to bin " - "on, and the membership condition via " - "'lt' / 'gt' operators. " - "Ex. name1:component:lt2 name2:total:lt15") - bins_group.add_argument("--bank-file", default=None, - help="HDF format template bank file.") - bins_group.add_argument("--f-lower", default=None, - help="Low frequency cutoff in Hz.") + bins_group = parser.add_argument_group("Options for selecting templates in bins.") + bins_group.add_argument( + "--bank-bins", + nargs="+", + default=None, + help="Ordered list of mass bin upper boundaries. " + "An ordered list of type-boundary pairs, " + "applied sequentially. Must provide a name " + "(can be any unique string for tagging " + "purposes), the parameter to bin " + "on, and the membership condition via " + "'lt' / 'gt' operators. " + "Ex. name1:component:lt2 name2:total:lt15", + ) + bins_group.add_argument( + "--bank-file", default=None, help="HDF format template bank file." + ) + bins_group.add_argument( + "--f-lower", default=None, help="Low frequency cutoff in Hz." + ) return bins_group def bank_bins_from_cli(opts): - """ Parses the CLI options related to binning templates in the bank. + """ + Parses the CLI options related to binning templates in the bank. Parameters ---------- @@ -67,6 +78,7 @@ def bank_bins_from_cli(opts): A dict with bin names as key and an array of their indices as value. bank : dict A dict of the datasets from the bank file. + """ bank = {} fp = h5py.File(opts.bank_file) @@ -76,7 +88,7 @@ def bank_bins_from_cli(opts): if opts.bank_bins: bins_idx = coinc.background_bin_from_string(opts.bank_bins, bank) else: - bins_idx = {"all" : numpy.arange(0, len(bank[tuple(fp.keys())[0]]))} + bins_idx = {"all": numpy.arange(0, len(bank[tuple(fp.keys())[0]]))} fp.close() return bins_idx, bank @@ -96,11 +108,12 @@ def get_mass_spin(bank, tid): ------- m1, m2, s1z, s2z : tuple of floats or arrays of floats Parameter values of the bank entries + """ - m1 = bank['mass1'][:][tid] - m2 = bank['mass2'][:][tid] - s1z = bank['spin1z'][:][tid] - s2z = bank['spin2z'][:][tid] + m1 = bank["mass1"][:][tid] + m2 = bank["mass2"][:][tid] + s1z = bank["spin1z"][:][tid] + s2z = bank["spin2z"][:][tid] return m1, m2, s1z, s2z @@ -121,26 +134,28 @@ def get_param(par, args, m1, m2, s1z, s2z): ------- parvals : float or array of floats Calculated parameter values + """ - if par == 'mchirp': + if par == "mchirp": parvals = conversions.mchirp_from_mass1_mass2(m1, m2) - elif par == 'mtotal': + elif par == "mtotal": parvals = m1 + m2 - elif par == 'eta': + elif par == "eta": parvals = conversions.eta_from_mass1_mass2(m1, m2) - elif par in ['chi_eff', 'effective_spin']: + elif par in ["chi_eff", "effective_spin"]: parvals = conversions.chi_eff(m1, m2, s1z, s2z) - elif par == 'template_duration': + elif par == "template_duration": # default to SEOBNRv4 duration function - if not hasattr(args, 'approximant') or args.approximant is None: + if not hasattr(args, "approximant") or args.approximant is None: args.approximant = "SEOBNRv4" - parvals = pnutils.get_imr_duration(m1, m2, s1z, s2z, args.f_lower, - args.approximant) + parvals = pnutils.get_imr_duration( + m1, m2, s1z, s2z, args.f_lower, args.approximant + ) if args.min_duration: parvals += args.min_duration - elif par == 'tau0': + elif par == "tau0": parvals = conversions.tau0_from_mass1_mass2(m1, m2, args.f_lower) - elif par == 'tau3': + elif par == "tau3": parvals = conversions.tau3_from_mass1_mass2(m1, m2, args.f_lower) elif par in pnutils.named_frequency_cutoffs.keys(): parvals = pnutils.frequency_cutoff_from_name(par, m1, m2, s1z, s2z) @@ -175,6 +190,7 @@ def get_found_param(injfile, bankfile, trigfile, param, ifo, args=None): [return value]: NumPy array of floats, array of boolean The calculated parameter values and a Boolean mask indicating which injections were found in the given ifo (if supplied) + """ foundtmp = injfile["found_after_vetoes/template_id"][:] # will record whether inj was found in the given ifo @@ -183,8 +199,11 @@ def get_found_param(injfile, bankfile, trigfile, param, ifo, args=None): try: # old 2-ifo behaviour # get the name of the ifo in the injection file, eg "detector_1" # and the integer from that name - ifolabel = [name for name, val in injfile.attrs.items() if \ - "detector" in name and val == ifo][0] + ifolabel = [ + name + for name, val in injfile.attrs.items() + if "detector" in name and val == ifo + ][0] foundtrg = injfile["found_after_vetoes/trigger_id" + ifolabel[-1]] except IndexError: # multi-ifo foundtrg = injfile["found_after_vetoes/%s/trigger_id" % ifo] @@ -192,14 +211,13 @@ def get_found_param(injfile, bankfile, trigfile, param, ifo, args=None): found_in_ifo = foundtrg[:] != -1 if bankfile is not None and param in bankfile.keys(): return bankfile[param][:][foundtmp], found_in_ifo - elif trigfile is not None and param in trigfile[ifo].keys(): + if trigfile is not None and param in trigfile[ifo].keys(): return trigfile[ifo][param][:][foundtrg], found_in_ifo - else: - assert bankfile - b = bankfile - return get_param(param, args, b['mass1'][:], b['mass2'][:], - b['spin1z'][:], b['spin2z'][:])[foundtmp],\ - found_in_ifo + assert bankfile + b = bankfile + return get_param( + param, args, b["mass1"][:], b["mass2"][:], b["spin1z"][:], b["spin2z"][:] + )[foundtmp], found_in_ifo def get_inj_param(injfile, param, ifo, args=None): @@ -222,6 +240,7 @@ def get_inj_param(injfile, param, ifo, args=None): ------- [return value]: NumPy array of floats The calculated parameter values + """ det = pycbc.detector.Detector(ifo) @@ -229,11 +248,15 @@ def get_inj_param(injfile, param, ifo, args=None): if param in inj.keys(): return inj[param][:] - if param == "end_time_"+ifo[0].lower(): - return inj['end_time'][:] + det.time_delay_from_earth_center( - inj['longitude'][:], - inj['latitude'][:], - inj['end_time'][:]) - else: - return get_param(param, args, inj['mass1'][:], inj['mass2'][:], - inj['spin1z'][:], inj['spin2z'][:]) + if param == "end_time_" + ifo[0].lower(): + return inj["end_time"][:] + det.time_delay_from_earth_center( + inj["longitude"][:], inj["latitude"][:], inj["end_time"][:] + ) + return get_param( + param, + args, + inj["mass1"][:], + inj["mass2"][:], + inj["spin1z"][:], + inj["spin2z"][:], + ) diff --git a/pycbc/events/veto.py b/pycbc/events/veto.py index e8ded61c185..d4a94a8c4c4 100644 --- a/pycbc/events/veto.py +++ b/pycbc/events/veto.py @@ -1,20 +1,26 @@ -""" This module contains utilities to manipulate trigger lists based on +""" +This module contains utilities to manipulate trigger lists based on segment. """ + import logging + import numpy +from igwn_ligolw import ligolw, lsctables +from igwn_ligolw import utils as ligolw_utils from igwn_segments import segment, segmentlist -from igwn_ligolw import ligolw, lsctables, utils as ligolw_utils -logger = logging.getLogger('pycbc.events.veto') +logger = logging.getLogger("pycbc.events.veto") + def start_end_to_segments(start, end): return segmentlist([segment(s, e) for s, e in zip(start, end)]) + def segments_to_start_end(segs): segs.coalesce() - return (numpy.array([s[0] for s in segs]), - numpy.array([s[1] for s in segs])) + return (numpy.array([s[0] for s in segs]), numpy.array([s[1] for s in segs])) + def start_end_from_segments(segment_file): """ @@ -28,15 +34,16 @@ def start_end_from_segments(segment_file): ------- start: numpy.ndarray end: numpy.ndarray + """ from pycbc.io.ligolw import LIGOLWContentHandler as h indoc = ligolw_utils.load_filename(segment_file, False, contenthandler=h) - segment_table = lsctables.SegmentTable.get_table(indoc) - start = numpy.array(segment_table.getColumnByName('start_time')) - start_ns = numpy.array(segment_table.getColumnByName('start_time_ns')) - end = numpy.array(segment_table.getColumnByName('end_time')) - end_ns = numpy.array(segment_table.getColumnByName('end_time_ns')) + segment_table = lsctables.SegmentTable.get_table(indoc) + start = numpy.array(segment_table.getColumnByName("start_time")) + start_ns = numpy.array(segment_table.getColumnByName("start_time_ns")) + end = numpy.array(segment_table.getColumnByName("end_time")) + end_ns = numpy.array(segment_table.getColumnByName("end_time_ns")) return start + start_ns * 1e-9, end + end_ns * 1e-9 @@ -57,6 +64,7 @@ def indices_within_times(times, start, end): ------- indices: numpy.ndarray Array of indices into times + """ # coalesce the start/end segments start, end = segments_to_start_end(start_end_to_segments(start, end).coalesce()) @@ -71,6 +79,7 @@ def indices_within_times(times, start, end): return tsort[numpy.hstack([numpy.r_[s:e] for s, e in zip(left, right)])] + def indices_outside_times(times, start, end): """ Return an index array into times that like outside the durations defined by start end arrays @@ -88,13 +97,16 @@ def indices_outside_times(times, start, end): ------- indices: numpy.ndarray Array of indices into times + """ exclude = indices_within_times(times, start, end) indices = numpy.arange(0, len(times)) return numpy.delete(indices, exclude) + def select_segments_by_definer(segment_file, segment_name=None, ifo=None): - """ Return the list of segments that match the segment name + """ + Return the list of segments that match the segment name Parameters ---------- @@ -108,16 +120,17 @@ def select_segments_by_definer(segment_file, segment_name=None, ifo=None): Returns ------- seg: list of segments + """ from pycbc.io.ligolw import LIGOLWContentHandler as h indoc = ligolw_utils.load_filename(segment_file, False, contenthandler=h) - segment_table = ligolw.Table.get_table(indoc, 'segment') + segment_table = ligolw.Table.get_table(indoc, "segment") - seg_def_table = ligolw.Table.get_table(indoc, 'segment_definer') - def_ifos = seg_def_table.getColumnByName('ifos') - def_names = seg_def_table.getColumnByName('name') - def_ids = seg_def_table.getColumnByName('segment_def_id') + seg_def_table = ligolw.Table.get_table(indoc, "segment_definer") + def_ifos = seg_def_table.getColumnByName("ifos") + def_names = seg_def_table.getColumnByName("name") + def_ids = seg_def_table.getColumnByName("segment_def_id") valid_id = [] for def_ifo, def_name, def_id in zip(def_ifos, def_names, def_ids): @@ -127,21 +140,22 @@ def select_segments_by_definer(segment_file, segment_name=None, ifo=None): continue valid_id += [def_id] - start = numpy.array(segment_table.getColumnByName('start_time')) - start_ns = numpy.array(segment_table.getColumnByName('start_time_ns')) - end = numpy.array(segment_table.getColumnByName('end_time')) - end_ns = numpy.array(segment_table.getColumnByName('end_time_ns')) + start = numpy.array(segment_table.getColumnByName("start_time")) + start_ns = numpy.array(segment_table.getColumnByName("start_time_ns")) + end = numpy.array(segment_table.getColumnByName("end_time")) + end_ns = numpy.array(segment_table.getColumnByName("end_time_ns")) start, end = start + 1e-9 * start_ns, end + 1e-9 * end_ns - did = segment_table.getColumnByName('segment_def_id') + did = segment_table.getColumnByName("segment_def_id") keep = numpy.array([d in valid_id for d in did]) if sum(keep) > 0: return start_end_to_segments(start[keep], end[keep]) - else: - return segmentlist([]) + return segmentlist([]) + def indices_within_segments(times, segment_files, ifo=None, segment_name=None): - """ Return the list of indices that should be vetoed by the segments in the + """ + Return the list of indices that should be vetoed by the segments in the list of veto_files. Parameters @@ -155,12 +169,14 @@ def indices_within_segments(times, segment_files, ifo=None, segment_name=None): The ifo to retrieve segments for from the segment files segment_name: str, optional name of segment + Returns ------- indices: numpy.ndarray The array of index values within the segments segmentlist: The segment list corresponding to the selected time. + """ veto_segs = segmentlist([]) indices = numpy.array([], dtype=numpy.uint32) @@ -175,8 +191,10 @@ def indices_within_segments(times, segment_files, ifo=None, segment_name=None): return indices, veto_segs.coalesce() + def indices_outside_segments(times, segment_files, ifo=None, segment_name=None): - """ Return the list of indices that are outside the segments in the + """ + Return the list of indices that are outside the segments in the list of segment files. Parameters @@ -190,39 +208,39 @@ def indices_outside_segments(times, segment_files, ifo=None, segment_name=None): The ifo to retrieve segments for from the segment files segment_name: str, optional name of segment + Returns - -------- + ------- indices: numpy.ndarray The array of index values outside the segments segmentlist: The segment list corresponding to the selected time. + """ - exclude, segs = indices_within_segments(times, segment_files, - ifo=ifo, segment_name=segment_name) + exclude, segs = indices_within_segments( + times, segment_files, ifo=ifo, segment_name=segment_name + ) indices = numpy.arange(0, len(times)) return numpy.delete(indices, exclude), segs + def get_segment_definer_comments(xml_file, include_version=True): """Returns a dict with the comment column as the value for each segment""" - from pycbc.io.ligolw import LIGOLWContentHandler as h # read segment definer table - xmldoc = ligolw_utils.load_fileobj(xml_file, - compress='auto', - contenthandler=h) + xmldoc = ligolw_utils.load_fileobj(xml_file, compress="auto", contenthandler=h) seg_def_table = lsctables.SegmentDefTable.get_table(xmldoc) # put comment column into a dict comment_dict = {} for seg_def in seg_def_table: if include_version: - full_channel_name = ':'.join([str(seg_def.ifos), - str(seg_def.name), - str(seg_def.version)]) + full_channel_name = ":".join( + [str(seg_def.ifos), str(seg_def.name), str(seg_def.version)] + ) else: - full_channel_name = ':'.join([str(seg_def.ifos), - str(seg_def.name)]) + full_channel_name = ":".join([str(seg_def.ifos), str(seg_def.name)]) comment_dict[full_channel_name] = seg_def.comment diff --git a/pycbc/fft/__init__.py b/pycbc/fft/__init__.py index 10862e6b4ba..6c4fa9f27a0 100644 --- a/pycbc/fft/__init__.py +++ b/pycbc/fft/__init__.py @@ -14,7 +14,7 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -from .parser_support import insert_fft_option_group, verify_fft_options, from_cli -from .func_api import fft, ifft -from .class_api import FFT, IFFT from .backend_support import get_backend_names +from .class_api import FFT, IFFT +from .func_api import fft, ifft +from .parser_support import from_cli, insert_fft_option_group, verify_fft_options diff --git a/pycbc/fft/backend_cpu.py b/pycbc/fft/backend_cpu.py index 3a5c96a1ea8..cee5fc62a84 100644 --- a/pycbc/fft/backend_cpu.py +++ b/pycbc/fft/backend_cpu.py @@ -17,15 +17,14 @@ from .core import _list_available -_backend_dict = {'fftw' : 'fftw', - 'mkl' : 'mkl', - 'numpy' : 'npfft'} -_backend_list = ['mkl', 'fftw', 'numpy'] +_backend_dict = {"fftw": "fftw", "mkl": "mkl", "numpy": "npfft"} +_backend_list = ["mkl", "fftw", "numpy"] _alist, _adict = _list_available(_backend_list, _backend_dict) cpu_backend = None + def set_backend(backend_list): global cpu_backend for backend in backend_list: @@ -33,8 +32,9 @@ def set_backend(backend_list): cpu_backend = backend break + def get_backend(): return _adict[cpu_backend] -set_backend(_backend_list) +set_backend(_backend_list) diff --git a/pycbc/fft/backend_cuda.py b/pycbc/fft/backend_cuda.py index 10adee52ad8..61fe3ea5eb9 100644 --- a/pycbc/fft/backend_cuda.py +++ b/pycbc/fft/backend_cuda.py @@ -16,20 +16,21 @@ # MA 02111-1307 USA import pycbc + from .core import _list_available -_backend_dict = {'cuda' : 'cufft', - 'pyfft' : 'cuda_pyfft'} -_backend_list = ['cuda','pyfft'] +_backend_dict = {"cuda": "cufft", "pyfft": "cuda_pyfft"} +_backend_list = ["cuda", "pyfft"] _alist = [] _adict = {} if pycbc.HAVE_CUDA: - _alist, _adict = _list_available(_backend_list,_backend_dict) + _alist, _adict = _list_available(_backend_list, _backend_dict) cuda_backend = None + def set_backend(backend_list): global cuda_backend for backend in backend_list: @@ -37,7 +38,9 @@ def set_backend(backend_list): cuda_backend = backend break + def get_backend(): return _adict[cuda_backend] + set_backend(_backend_list) diff --git a/pycbc/fft/backend_cupy.py b/pycbc/fft/backend_cupy.py index fec13ad746b..555a0483090 100644 --- a/pycbc/fft/backend_cupy.py +++ b/pycbc/fft/backend_cupy.py @@ -17,13 +17,14 @@ from .core import _list_available -_backend_dict = {'cupy' : 'cupyfft'} -_backend_list = ['cupy'] +_backend_dict = {"cupy": "cupyfft"} +_backend_list = ["cupy"] _alist, _adict = _list_available(_backend_list, _backend_dict) cupy_backend = None + def set_backend(backend_list): global cupy_backend for backend in backend_list: @@ -31,7 +32,9 @@ def set_backend(backend_list): cupy_backend = backend break + def get_backend(): return _adict[cupy_backend] + set_backend(_backend_list) diff --git a/pycbc/fft/backend_mkl.py b/pycbc/fft/backend_mkl.py index bf4d3408668..76b082b5d6a 100644 --- a/pycbc/fft/backend_mkl.py +++ b/pycbc/fft/backend_mkl.py @@ -17,13 +17,14 @@ from .core import _list_available -_backend_dict = {'mkl' : 'mkl'} -_backend_list = ['mkl'] +_backend_dict = {"mkl": "mkl"} +_backend_list = ["mkl"] _alist, _adict = _list_available(_backend_list, _backend_dict) mkl_backend = None + def set_backend(backend_list): global mkl_backend for backend in backend_list: @@ -31,7 +32,9 @@ def set_backend(backend_list): mkl_backend = backend break + def get_backend(): return _adict[mkl_backend] + set_backend(_backend_list) diff --git a/pycbc/fft/backend_support.py b/pycbc/fft/backend_support.py index ed1c76c7fe8..8b776467390 100644 --- a/pycbc/fft/backend_support.py +++ b/pycbc/fft/backend_support.py @@ -29,7 +29,6 @@ import pycbc import pycbc.scheme - # These are global variables, that are modified by the various scheme- # dependent submodules, to maintain a list of all possible backends # for all possible schemes that are available at runtime. This list @@ -44,19 +43,24 @@ # in the global list, and we assume that the keys to the dict are in one-to-one # correspondence with the items in the list. + def _update_global_available(new_list, new_dict, global_list, global_dict): for item in new_list: if item not in global_list: global_list.append(item) - global_dict.update({item:new_dict[item]}) + global_dict.update({item: new_dict[item]}) + def get_backend_modules(): return _all_backends_dict.values() + def get_backend_names(): return list(_all_backends_dict.keys()) -BACKEND_PREFIX="pycbc.fft.backend_" + +BACKEND_PREFIX = "pycbc.fft.backend_" + @pycbc.scheme.schemed(BACKEND_PREFIX) def set_backend(backend_list): @@ -64,20 +68,23 @@ def set_backend(backend_list): err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + @pycbc.scheme.schemed(BACKEND_PREFIX) def get_backend(): err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + # Import all scheme-dependent backends, to get _all_backends accurate: for scheme_name in ["cpu", "mkl", "cuda", "cupy"]: try: - mod = __import__('pycbc.fft.backend_' + scheme_name, fromlist = ['_alist', '_adict']) - _alist = getattr(mod, "_alist") - _adict = getattr(mod, "_adict") - _update_global_available(_alist, _adict, _all_backends_list, - _all_backends_dict) + mod = __import__( + "pycbc.fft.backend_" + scheme_name, fromlist=["_alist", "_adict"] + ) + _alist = mod._alist + _adict = mod._adict + _update_global_available(_alist, _adict, _all_backends_list, _all_backends_dict) except ImportError: pass diff --git a/pycbc/fft/class_api.py b/pycbc/fft/class_api.py index 598a432ef9d..04645f9cd2c 100644 --- a/pycbc/fft/class_api.py +++ b/pycbc/fft/class_api.py @@ -28,18 +28,22 @@ from .backend_support import get_backend + def _fft_factory(invec, outvec, nbatch=1, size=None): backend = get_backend() - cls = getattr(backend, 'FFT') + cls = backend.FFT return cls + def _ifft_factory(invec, outvec, nbatch=1, size=None): backend = get_backend() - cls = getattr(backend, 'IFFT') + cls = backend.IFFT return cls -class FFT(object): - """ Create a forward FFT engine + +class FFT: + """ + Create a forward FFT engine Parameters ---------- @@ -58,13 +62,17 @@ class FFT(object): The addresses in memory of both vectors should be divisible by pycbc.PYCBC_ALIGNMENT. + """ + def __new__(cls, *args, **kwargs): real_cls = _fft_factory(*args, **kwargs) return real_cls(*args, **kwargs) -class IFFT(object): - """ Create a reverse FFT engine + +class IFFT: + """ + Create a reverse FFT engine Parameters ---------- @@ -83,8 +91,9 @@ class IFFT(object): The addresses in memory of both vectors should be divisible by pycbc.PYCBC_ALIGNMENT. + """ + def __new__(cls, *args, **kwargs): real_cls = _ifft_factory(*args, **kwargs) return real_cls(*args, **kwargs) - diff --git a/pycbc/fft/core.py b/pycbc/fft/core.py index 92e3cdcb229..6e6bc2b0f80 100644 --- a/pycbc/fft/core.py +++ b/pycbc/fft/core.py @@ -25,15 +25,17 @@ This package provides a front-end to various fast Fourier transform implementations within PyCBC. """ + from pycbc.types import Array as _Array -from pycbc.types import TimeSeries as _TimeSeries from pycbc.types import FrequencySeries as _FrequencySeries +from pycbc.types import TimeSeries as _TimeSeries # The following helper function is in this top-level module because it # is used by the scheme-dependent files to write their version of the # _available_backends() function. It cannot go in backend_support as # that woulc cause circular imports + def _list_available(possible_list, possible_dict): # It possibly is strange that we have both a list and a dict. # The reason for this is that the name the user specfies for a @@ -47,13 +49,16 @@ def _list_available(possible_list, possible_dict): available_dict = {} for backend in possible_list: try: - mod = __import__('pycbc.fft.' + possible_dict[backend], fromlist = ['pycbc.fft']) - available_dict.update({backend:mod}) + mod = __import__( + "pycbc.fft." + possible_dict[backend], fromlist=["pycbc.fft"] + ) + available_dict.update({backend: mod}) available_list.append(backend) except (ImportError, OSError): pass return available_list, available_dict + # The main purpose of the top-level module is to present a # uniform interface for a forward and reverse FFT, independent of # the underlying backend. We perform sanity checking here, at the @@ -61,28 +66,21 @@ def _list_available(possible_list, possible_dict): # facilitate this checking, we define dicts mapping the numpy dtype # to the corresponding precisions and types. + def _check_fft_args(invec, outvec): - if not isinstance(invec,_Array): + if not isinstance(invec, _Array): raise TypeError("Input is not a PyCBC Array") - if not isinstance(outvec,_Array): + if not isinstance(outvec, _Array): raise TypeError("Output is not a PyCBC Array") - if isinstance(invec,_TimeSeries) and not isinstance( - outvec,_FrequencySeries): - raise TypeError( - "When input is TimeSeries output must be FrequencySeries") - if isinstance(outvec,_TimeSeries) and not isinstance( - invec,_FrequencySeries): - raise TypeError( - "When output is TimeSeries input must be FrequencySeries") - if isinstance(invec,_FrequencySeries) and not isinstance( - outvec,_TimeSeries): - raise TypeError( - "When input is FrequencySeries output must be TimeSeries") - if isinstance(outvec,_FrequencySeries) and not isinstance( - invec,_TimeSeries): - raise TypeError( - "When output is FrequencySeries input must be TimeSeries") + if isinstance(invec, _TimeSeries) and not isinstance(outvec, _FrequencySeries): + raise TypeError("When input is TimeSeries output must be FrequencySeries") + if isinstance(outvec, _TimeSeries) and not isinstance(invec, _FrequencySeries): + raise TypeError("When output is TimeSeries input must be FrequencySeries") + if isinstance(invec, _FrequencySeries) and not isinstance(outvec, _TimeSeries): + raise TypeError("When input is FrequencySeries output must be TimeSeries") + if isinstance(outvec, _FrequencySeries) and not isinstance(invec, _TimeSeries): + raise TypeError("When output is FrequencySeries input must be TimeSeries") iprec = invec.precision oprec = outvec.precision @@ -91,7 +89,8 @@ def _check_fft_args(invec, outvec): itype = invec.kind otype = outvec.kind - return [iprec,itype,otype] + return [iprec, itype, otype] + def _check_fwd_args(invec, itype, outvec, otype, nbatch, size): ilen = len(invec) @@ -102,28 +101,32 @@ def _check_fwd_args(invec, itype, outvec, otype, nbatch, size): raise ValueError("When nbatch > 1, size cannot be 'None'") if size is None: size = ilen - inplace = (invec.ptr == outvec.ptr) + inplace = invec.ptr == outvec.ptr if (ilen % nbatch) != 0: raise ValueError("Input length must be divisible by nbatch") if (olen % nbatch) != 0: raise ValueError("Output length must be divisible by nbatch") - if itype == 'complex' and otype == 'complex': - if (ilen/nbatch) != size: + if itype == "complex" and otype == "complex": + if (ilen / nbatch) != size: raise ValueError("For C2C FFT, len(invec) must be nbatch*size") - if (olen/nbatch) != size: + if (olen / nbatch) != size: raise ValueError("For C2C FFT, len(outvec) must be nbatch*size") - elif itype == 'real' and otype == 'complex': - if (olen/nbatch) != int(size/2 + 1): + elif itype == "real" and otype == "complex": + if (olen / nbatch) != int(size / 2 + 1): raise ValueError("For R2C FFT, len(outvec) must be nbatch*(size/2 + 1)") if inplace: - if (ilen/nbatch) != int(2*(size/2 + 1)): - raise ValueError("For R2C in-place FFT, len(invec) must be nbatch*2*(size/2+1)") - else: - if (ilen/nbatch) != size: - raise ValueError("For R2C out-of-place FFT, len(invec) must be nbatch*size") + if (ilen / nbatch) != int(2 * (size / 2 + 1)): + raise ValueError( + "For R2C in-place FFT, len(invec) must be nbatch*2*(size/2+1)" + ) + elif (ilen / nbatch) != size: + raise ValueError( + "For R2C out-of-place FFT, len(invec) must be nbatch*size" + ) else: raise ValueError("Inconsistent dtypes for forward FFT") + def _check_inv_args(invec, itype, outvec, otype, nbatch, size): ilen = len(invec) olen = len(outvec) @@ -133,25 +136,29 @@ def _check_inv_args(invec, itype, outvec, otype, nbatch, size): raise ValueError("When nbatch > 1, size cannot be 'None'") if size is None: size = olen - inplace = (invec.ptr == outvec.ptr) + inplace = invec.ptr == outvec.ptr if (ilen % nbatch) != 0: raise ValueError("Input length must be divisible by nbatch") if (olen % nbatch) != 0: raise ValueError("Output length must be divisible by nbatch") - if itype == 'complex' and otype == 'complex': - if (ilen/nbatch) != size: + if itype == "complex" and otype == "complex": + if (ilen / nbatch) != size: raise ValueError("For C2C IFFT, len(invec) must be nbatch*size") - if (olen/nbatch) != size: + if (olen / nbatch) != size: raise ValueError("For C2C IFFT, len(outvec) must be nbatch*size") - elif itype == 'complex' and otype == 'real': - if (ilen/nbatch) != int(size/2 + 1): + elif itype == "complex" and otype == "real": + if (ilen / nbatch) != int(size / 2 + 1): raise ValueError("For C2R IFFT, len(invec) must be nbatch*(size/2 + 1)") if inplace: - if (olen/nbatch) != 2*int(size/2 + 1): - raise ValueError("For C2R in-place IFFT, len(outvec) must be nbatch*2*(size/2+1)") - else: - if (olen/nbatch) != size: - raise ValueError("For C2R out-of-place IFFT, len(outvec) must be nbatch*size") + if (olen / nbatch) != 2 * int(size / 2 + 1): + raise ValueError( + "For C2R in-place IFFT, len(outvec) must be nbatch*2*(size/2+1)" + ) + elif (olen / nbatch) != size: + raise ValueError( + "For C2R out-of-place IFFT, len(outvec) must be nbatch*size" + ) + # The class-based approach requires the following: @@ -169,14 +176,15 @@ def _check_inv_args(invec, itype, outvec, otype, nbatch, size): # nontrivial work and hence should be called inside the __init__ method of all child classes, # before anything else. -class _BaseFFT(object): + +class _BaseFFT: def __init__(self, invec, outvec, nbatch, size): _, itype, otype = _check_fft_args(invec, outvec) _check_fwd_args(invec, itype, outvec, otype, nbatch, size) self.forward = True self.invec = invec self.outvec = outvec - self.inplace = (self.invec.ptr == self.outvec.ptr) + self.inplace = self.invec.ptr == self.outvec.ptr self.nbatch = nbatch if nbatch > 1: self.size = size @@ -184,15 +192,15 @@ def __init__(self, invec, outvec, nbatch, size): self.size = len(invec) # Whether we are complex-to-complex or real-to-complex is determined # by itype: - if itype == 'complex': + if itype == "complex": # Complex-to-complex case: self.idist = self.size self.odist = self.size else: # Real-to-complex case: - self.odist = int(self.size/2 + 1) + self.odist = int(self.size / 2 + 1) if self.inplace: - self.idist = 2*int(self.size/2 + 1) + self.idist = 2 * int(self.size / 2 + 1) else: self.idist = self.size @@ -200,11 +208,11 @@ def __init__(self, invec, outvec, nbatch, size): # we should divide by, whether C2C or R2HC transform if isinstance(self.invec, _TimeSeries): self.outvec._epoch = self.invec._epoch - self.outvec._delta_f = 1.0/(self.invec._delta_t * len(self.invec)) + self.outvec._delta_f = 1.0 / (self.invec._delta_t * len(self.invec)) self.scale = self.invec._delta_t elif isinstance(self.invec, _FrequencySeries): self.outvec._epoch = self.invec._epoch - self.outvec._delta_t = 1.0/(self.invec._delta_f * len(self.invec)) + self.outvec._delta_t = 1.0 / (self.invec._delta_f * len(self.invec)) self.scale = self.invec._delta_f def execute(self): @@ -220,16 +228,16 @@ def execute(self): its output by the input vector's delta_t (when input is a TimeSeries) or delta_f (when input is a FrequencySeries). """ - pass -class _BaseIFFT(object): + +class _BaseIFFT: def __init__(self, invec, outvec, nbatch, size): _, itype, otype = _check_fft_args(invec, outvec) _check_inv_args(invec, itype, outvec, otype, nbatch, size) self.forward = False self.invec = invec self.outvec = outvec - self.inplace = (self.invec.ptr == self.outvec.ptr) + self.inplace = self.invec.ptr == self.outvec.ptr self.nbatch = nbatch if nbatch > 1: self.size = size @@ -237,15 +245,15 @@ def __init__(self, invec, outvec, nbatch, size): self.size = len(outvec) # Whether we are complex-to-complex or complex-to-real is determined # by otype: - if otype == 'complex': + if otype == "complex": # Complex-to-complex case: self.idist = self.size self.odist = self.size else: # Complex-to-real case: - self.idist = int(self.size/2 + 1) + self.idist = int(self.size / 2 + 1) if self.inplace: - self.odist = 2*int(self.size/2 + 1) + self.odist = 2 * int(self.size / 2 + 1) else: self.odist = self.size @@ -253,11 +261,11 @@ def __init__(self, invec, outvec, nbatch, size): # we should divide by, whether C2C or HC2R transform if isinstance(self.invec, _TimeSeries): self.outvec._epoch = self.invec._epoch - self.outvec._delta_f = 1.0/(self.invec._delta_t * len(self.outvec)) + self.outvec._delta_f = 1.0 / (self.invec._delta_t * len(self.outvec)) self.scale = self.invec._delta_t elif isinstance(self.invec, _FrequencySeries): self.outvec._epoch = self.invec._epoch - self.outvec._delta_t = 1.0/(self.invec._delta_f * len(self.outvec)) + self.outvec._delta_t = 1.0 / (self.invec._delta_f * len(self.outvec)) self.scale = self.invec._delta_f def execute(self): @@ -273,5 +281,3 @@ def execute(self): its output by the input vector's delta_t (when input is a TimeSeries) or delta_f (when input is a FrequencySeries). """ - pass - diff --git a/pycbc/fft/cuda_pyfft.py b/pycbc/fft/cuda_pyfft.py index 94092ccd430..960afdcd9b6 100644 --- a/pycbc/fft/cuda_pyfft.py +++ b/pycbc/fft/cuda_pyfft.py @@ -26,41 +26,45 @@ for the PyCBC package. """ -import pycbc.scheme from pyfft.cuda import Plan +import pycbc.scheme + _plans = {} -#These dicts need to be cleared before the cuda context is destroyed + +# These dicts need to be cleared before the cuda context is destroyed def _clear_plan_dict(): _plans.clear() + pycbc.scheme.register_clean_cuda(_clear_plan_dict) -#itype and otype are actual dtypes here, not strings -def _get_plan(itype,otype,inlen): +# itype and otype are actual dtypes here, not strings +def _get_plan(itype, otype, inlen): try: - theplan = _plans[(itype,otype,inlen)] + theplan = _plans[(itype, otype, inlen)] except KeyError: - theplan = Plan(inlen,dtype = itype,normalize=False,fast_math=True) - _plans.update({(itype,otype,inlen) : theplan }) + theplan = Plan(inlen, dtype=itype, normalize=False, fast_math=True) + _plans.update({(itype, otype, inlen): theplan}) return theplan -def fft(invec,outvec,prec,itype,otype): - if itype =='complex' and otype == 'complex': - pyplan=_get_plan(invec.dtype, outvec.dtype, len(invec)) - pyplan.execute(invec.data,outvec.data) - elif itype=='real' and otype=='complex': +def fft(invec, outvec, prec, itype, otype): + if itype == "complex" and otype == "complex": + pyplan = _get_plan(invec.dtype, outvec.dtype, len(invec)) + pyplan.execute(invec.data, outvec.data) + + elif itype == "real" and otype == "complex": raise NotImplementedError("Only Complex to Complex FFTs for pyfft currently.") -def ifft(invec,outvec,prec,itype,otype): - if itype =='complex' and otype == 'complex': - pyplan=_get_plan(invec.dtype,outvec.dtype,len(invec)) - pyplan.execute(invec.data,outvec.data,inverse=True) - elif itype=='complex' and otype=='real': - raise NotImplementedError("Only Complex to Complex IFFTs for pyfft currently.") +def ifft(invec, outvec, prec, itype, otype): + if itype == "complex" and otype == "complex": + pyplan = _get_plan(invec.dtype, outvec.dtype, len(invec)) + pyplan.execute(invec.data, outvec.data, inverse=True) + elif itype == "complex" and otype == "real": + raise NotImplementedError("Only Complex to Complex IFFTs for pyfft currently.") diff --git a/pycbc/fft/cufft.py b/pycbc/fft/cufft.py index 87a935da8c0..6d47606b9df 100644 --- a/pycbc/fft/cufft.py +++ b/pycbc/fft/cufft.py @@ -27,6 +27,7 @@ """ import pycbc.scheme + # The following is a hack, to ensure that any error in importing # cufft is treated as the module being unavailable at runtime. # Ideally, the real error and its traceback would be appended to @@ -35,38 +36,43 @@ try: import skcuda.fft as cu_fft except: - raise ImportError("Unable to import skcuda.fft; try direct import" - " to get full traceback") + raise ImportError( + "Unable to import skcuda.fft; try direct import to get full traceback" + ) from .core import _BaseFFT, _BaseIFFT _forward_plans = {} _reverse_plans = {} -#These dicts need to be cleared before the cuda context is destroyed + +# These dicts need to be cleared before the cuda context is destroyed def _clear_plan_dicts(): _forward_plans.clear() _reverse_plans.clear() + pycbc.scheme.register_clean_cuda(_clear_plan_dicts) -#itype and otype are actual dtypes here, not strings + +# itype and otype are actual dtypes here, not strings def _get_fwd_plan(itype, otype, inlen, batch=1): try: theplan = _forward_plans[(itype, otype, inlen, batch)] except KeyError: theplan = cu_fft.Plan((inlen,), itype, otype, batch=batch) - _forward_plans.update({(itype, otype, inlen) : theplan }) + _forward_plans.update({(itype, otype, inlen): theplan}) return theplan -#The complex to real plan wants the actual size, not the N/2+1 -#That's why the inverse plans use the outvec length, instead of the invec + +# The complex to real plan wants the actual size, not the N/2+1 +# That's why the inverse plans use the outvec length, instead of the invec def _get_inv_plan(itype, otype, outlen, batch=1): try: theplan = _reverse_plans[(itype, otype, outlen, batch)] except KeyError: theplan = cu_fft.Plan((outlen,), itype, otype, batch=batch) - _reverse_plans.update({(itype, otype, outlen) : theplan }) + _reverse_plans.update({(itype, otype, outlen): theplan}) return theplan @@ -75,13 +81,15 @@ def fft(invec, outvec, prec, itype, otype): cuplan = _get_fwd_plan(invec.dtype, outvec.dtype, len(invec)) cu_fft.fft(invec.data, outvec.data, cuplan) + def ifft(invec, outvec, prec, itype, otype): cuplan = _get_inv_plan(invec.dtype, outvec.dtype, len(outvec)) cu_fft.ifft(invec.data, outvec.data, cuplan) + class FFT(_BaseFFT): def __init__(self, invec, outvec, nbatch=1, size=None): - super(FFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.plan = _get_fwd_plan(invec.dtype, outvec.dtype, len(invec), batch=nbatch) self.invec = invec.data self.outvec = outvec.data @@ -89,9 +97,10 @@ def __init__(self, invec, outvec, nbatch=1, size=None): def execute(self): cu_fft.fft(self.invec, self.outvec, self.plan) + class IFFT(_BaseIFFT): def __init__(self, invec, outvec, nbatch=1, size=None): - super(IFFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.plan = _get_inv_plan(invec.dtype, outvec.dtype, len(outvec), batch=nbatch) self.invec = invec.data @@ -99,4 +108,3 @@ def __init__(self, invec, outvec, nbatch=1, size=None): def execute(self): cu_fft.ifft(self.invec, self.outvec, self.plan) - diff --git a/pycbc/fft/cupyfft.py b/pycbc/fft/cupyfft.py index 9fa0e468133..ed81353ed17 100644 --- a/pycbc/fft/cupyfft.py +++ b/pycbc/fft/cupyfft.py @@ -27,37 +27,40 @@ """ import cupy.fft -from .core import _check_fft_args -from .core import _BaseFFT, _BaseIFFT -_INV_FFT_MSG = ("I cannot perform an {} between data with an input type of " - "{} and an output type of {}") +from .core import _BaseFFT, _BaseIFFT, _check_fft_args + +_INV_FFT_MSG = ( + "I cannot perform an {} between data with an input type of " + "{} and an output type of {}" +) + def fft(invec, outvec, _, itype, otype): if invec.ptr == outvec.ptr: - raise NotImplementedError("cupy backend of pycbc.fft does not " - "support in-place transforms") - if itype == 'complex' and otype == 'complex': - outvec.data[:] = cupy.asarray(cupy.fft.fft(invec.data), - dtype=outvec.dtype) - elif itype == 'real' and otype == 'complex': - outvec.data[:] = cupy.asarray(cupy.fft.rfft(invec.data), - dtype=outvec.dtype) + raise NotImplementedError( + "cupy backend of pycbc.fft does not support in-place transforms" + ) + if itype == "complex" and otype == "complex": + outvec.data[:] = cupy.asarray(cupy.fft.fft(invec.data), dtype=outvec.dtype) + elif itype == "real" and otype == "complex": + outvec.data[:] = cupy.asarray(cupy.fft.rfft(invec.data), dtype=outvec.dtype) else: raise ValueError(_INV_FFT_MSG.format("FFT", itype, otype)) def ifft(invec, outvec, _, itype, otype): if invec.ptr == outvec.ptr: - raise NotImplementedError("cupy backend of pycbc.fft does not " - "support in-place transforms") - if itype == 'complex' and otype == 'complex': - outvec.data[:] = cupy.asarray(cupy.fft.ifft(invec.data), - dtype=outvec.dtype) + raise NotImplementedError( + "cupy backend of pycbc.fft does not support in-place transforms" + ) + if itype == "complex" and otype == "complex": + outvec.data[:] = cupy.asarray(cupy.fft.ifft(invec.data), dtype=outvec.dtype) outvec *= len(outvec) - elif itype == 'complex' and otype == 'real': - outvec.data[:] = cupy.asarray(cupy.fft.irfft(invec.data,len(outvec)), - dtype=outvec.dtype) + elif itype == "complex" and otype == "real": + outvec.data[:] = cupy.asarray( + cupy.fft.irfft(invec.data, len(outvec)), dtype=outvec.dtype + ) outvec *= len(outvec) else: raise ValueError(_INV_FFT_MSG.format("IFFT", itype, otype)) @@ -67,8 +70,9 @@ class FFT(_BaseFFT): """ Class for performing FFTs via the cupy interface. """ + def __init__(self, invec, outvec, nbatch=1, size=None): - super(FFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.prec, self.itype, self.otype = _check_fft_args(invec, outvec) def execute(self): @@ -79,8 +83,9 @@ class IFFT(_BaseIFFT): """ Class for performing IFFTs via the cupy interface. """ + def __init__(self, invec, outvec, nbatch=1, size=None): - super(IFFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.prec, self.itype, self.otype = _check_fft_args(invec, outvec) def execute(self): diff --git a/pycbc/fft/fft_callback.py b/pycbc/fft/fft_callback.py index 594713321f7..3dc1d3da32a 100644 --- a/pycbc/fft/fft_callback.py +++ b/pycbc/fft/fft_callback.py @@ -1,5 +1,8 @@ #!/usr/bin/python -import os, subprocess, ctypes +import ctypes +import os +import subprocess + from mako.template import Template full_corr = """ @@ -160,18 +163,20 @@ } """) + def compile(source, name): - """ Compile the string source code into a shared object linked against + """ + Compile the string source code into a shared object linked against the static version of cufft for callback support. """ # If we start using this again, we should find a better place for the cache - cache = os.path.join('/tmp', name) + cache = os.path.join("/tmp", name) hash_file = cache + ".hash" lib_file = cache + ".so" obj_file = cache + ".o" try: - if int(open(hash_file, "r").read()) == hash(source): + if int(open(hash_file).read()) == hash(source): return lib_file raise ValueError except: @@ -182,15 +187,34 @@ def compile(source, name): fsrc.write(source) fsrc.close() - cmd = ["nvcc", "-ccbin", "g++", "-dc", "-m64", - "--compiler-options", "'-fPIC'", - "-o", obj_file, - "-c", src_file] + cmd = [ + "nvcc", + "-ccbin", + "g++", + "-dc", + "-m64", + "--compiler-options", + "'-fPIC'", + "-o", + obj_file, + "-c", + src_file, + ] print(" ".join(cmd)) subprocess.check_call(cmd) - cmd = ["nvcc", "-shared", "-ccbin", "g++", "-m64", - "-o", lib_file, obj_file, "-lcufft_static", "-lculibos"] + cmd = [ + "nvcc", + "-shared", + "-ccbin", + "g++", + "-m64", + "-o", + lib_file, + obj_file, + "-lcufft_static", + "-lculibos", + ] print(" ".join(cmd)) subprocess.check_call(cmd) @@ -200,12 +224,14 @@ def compile(source, name): fhash.write(str(hash(source))) return lib_file -def get_fn_plan(callback=None, out_callback=None, name='pycbc_cufft', parameters=None): - """ Get the IFFT execute and plan functions - """ + +def get_fn_plan(callback=None, out_callback=None, name="pycbc_cufft", parameters=None): + """Get the IFFT execute and plan functions""" if parameters is None: parameters = [] - source = fftsrc.render(input_callback=callback, output_callback=out_callback, parameters=parameters) + source = fftsrc.render( + input_callback=callback, output_callback=out_callback, parameters=parameters + ) path = compile(source, name) lib = ctypes.cdll.LoadLibrary(path) fn = lib.execute @@ -215,38 +241,53 @@ def get_fn_plan(callback=None, out_callback=None, name='pycbc_cufft', parameters plan.argyptes = [ctypes.c_uint] return fn, plan + _plans = {} + class param(ctypes.Structure): _fields_ = [("htilde", ctypes.c_void_p)] + + hparam = param() + def c2c_correlate_ifft(htilde, stilde, outvec): - key = 'cnf' + key = "cnf" if key not in _plans: - fn, pfn = get_fn_plan(callback=full_corr, parameters = [("void*", "htilde")]) + fn, pfn = get_fn_plan(callback=full_corr, parameters=[("void*", "htilde")]) plan = pfn(len(outvec), int(htilde.data.gpudata)) _plans[key] = (fn, plan, int(htilde.data.gpudata)) fn, plan, _ = _plans[key] hparam.htilde = htilde.data.gpudata fn(plan, int(stilde.data.gpudata), int(outvec.data.gpudata), ctypes.pointer(hparam)) + class param2(ctypes.Structure): - _fields_ = [("htilde", ctypes.c_void_p), - ("in_kmax", ctypes.c_uint), - ("out_kmin", ctypes.c_uint), - ("out_kmax", ctypes.c_uint)] + _fields_ = [ + ("htilde", ctypes.c_void_p), + ("in_kmax", ctypes.c_uint), + ("out_kmin", ctypes.c_uint), + ("out_kmax", ctypes.c_uint), + ] + + hparam_zeros = param2() + def c2c_half_correlate_ifft(htilde, stilde, outvec): - key = 'cn' + key = "cn" if key not in _plans: - fn, pfn = get_fn_plan(callback=zero_corr, - parameters = [("void*", "htilde"), - ("unsigned int", "in_kmax"), - ("unsigned int", "out_kmin"), - ("unsigned int", "out_kmax")], - out_callback=zero_out) + fn, pfn = get_fn_plan( + callback=zero_corr, + parameters=[ + ("void*", "htilde"), + ("unsigned int", "in_kmax"), + ("unsigned int", "out_kmin"), + ("unsigned int", "out_kmax"), + ], + out_callback=zero_out, + ) plan = pfn(len(outvec), int(htilde.data.gpudata)) _plans[key] = (fn, plan, int(htilde.data.gpudata)) fn, plan, _ = _plans[key] @@ -254,5 +295,9 @@ def c2c_half_correlate_ifft(htilde, stilde, outvec): hparam_zeros.in_kmax = htilde.end_idx hparam_zeros.out_kmin = stilde.analyze.start hparam_zeros.out_kmax = stilde.analyze.stop - fn(plan, int(stilde.data.gpudata), int(outvec.data.gpudata), ctypes.pointer(hparam_zeros)) - + fn( + plan, + int(stilde.data.gpudata), + int(outvec.data.gpudata), + ctypes.pointer(hparam_zeros), + ) diff --git a/pycbc/fft/fftw.py b/pycbc/fft/fftw.py index 4c9a5c8ed4d..3df49a26e88 100644 --- a/pycbc/fft/fftw.py +++ b/pycbc/fft/fftw.py @@ -1,11 +1,14 @@ +import ctypes import os -from pycbc.types import zeros + import numpy as _np -import ctypes + import pycbc.scheme as _scheme from pycbc.libutils import get_ctypes_library -from .core import _BaseFFT, _BaseIFFT +from pycbc.types import zeros + from ..types import check_aligned +from .core import _BaseFFT, _BaseIFFT # IMPORTANT NOTE TO PYCBC DEVELOPERS: # Because this module is loaded automatically when present, and because @@ -16,12 +19,12 @@ # NOTE: # When loading FFTW we use os.RTLD_DEEPBIND to avoid potential segfaults due # to conflicts with MKL if both are present. -if hasattr(os, 'RTLD_DEEPBIND'): +if hasattr(os, "RTLD_DEEPBIND"): FFTW_RTLD_MODE = os.RTLD_DEEPBIND else: FFTW_RTLD_MODE = ctypes.DEFAULT_MODE -#FFTW constants, these are pulled from fftw3.h +# FFTW constants, these are pulled from fftw3.h FFTW_FORWARD = -1 FFTW_BACKWARD = 1 @@ -39,8 +42,8 @@ # We need to construct them directly with CDLL so # we can give the RTLD_GLOBAL mode, which we must do # in order to use the threaded libraries as well. -double_lib = get_ctypes_library('fftw3', ['fftw3'], mode=FFTW_RTLD_MODE) -float_lib = get_ctypes_library('fftw3f', ['fftw3f'], mode=FFTW_RTLD_MODE) +double_lib = get_ctypes_library("fftw3", ["fftw3"], mode=FFTW_RTLD_MODE) +float_lib = get_ctypes_library("fftw3f", ["fftw3f"], mode=FFTW_RTLD_MODE) if (double_lib is None) or (float_lib is None): raise ImportError("Unable to find FFTW libraries") @@ -63,13 +66,15 @@ # directly, but only by functions that get the value they use from # scheme.mgr.num_threads + def _fftw_plan_with_nthreads(nthreads): global _fftw_current_nthreads if not HAVE_FFTW_THREADED: - if (nthreads > 1): - raise ValueError("Threading is NOT enabled, but {0} > 1 threads specified".format(nthreads)) - else: - _pycbc_current_threads = nthreads + if nthreads > 1: + raise ValueError( + f"Threading is NOT enabled, but {nthreads} > 1 threads specified" + ) + _pycbc_current_threads = nthreads else: dplanwthr = _double_threaded_lib.fftw_plan_with_nthreads fplanwthr = _float_threaded_lib.fftwf_plan_with_nthreads @@ -79,12 +84,16 @@ def _fftw_plan_with_nthreads(nthreads): fplanwthr(nthreads) _fftw_current_nthreads = nthreads + # This is a global dict-of-dicts used when initializing threads and # setting the threading library -_fftw_threading_libnames = { 'unthreaded' : {'double' : None, 'float' : None}, - 'openmp' : {'double' : 'fftw3_omp', 'float' : 'fftw3f_omp'}, - 'pthreads' : {'double' : 'fftw3_threads', 'float' : 'fftw3f_threads'}} +_fftw_threading_libnames = { + "unthreaded": {"double": None, "float": None}, + "openmp": {"double": "fftw3_omp", "float": "fftw3f_omp"}, + "pthreads": {"double": "fftw3_threads", "float": "fftw3f_threads"}, +} + def _init_threads(backend): # This function actually sets the backend and initializes. It returns zero on @@ -98,35 +107,34 @@ def _init_threads(backend): global _float_threaded_lib if _fftw_threaded_set: raise RuntimeError( - "Threading backend for FFTW already set to {0}; cannot be changed".format(_fftw_threaded_lib)) + f"Threading backend for FFTW already set to {_fftw_threaded_lib}; cannot be changed" + ) try: - double_threaded_libname = _fftw_threading_libnames[backend]['double'] - float_threaded_libname = _fftw_threading_libnames[backend]['float'] + double_threaded_libname = _fftw_threading_libnames[backend]["double"] + float_threaded_libname = _fftw_threading_libnames[backend]["float"] except KeyError: - raise ValueError("Backend {0} for FFTW threading does not exist!".format(backend)) + raise ValueError( + f"Backend {backend} for FFTW threading does not exist!" + ) if double_threaded_libname is not None: try: # For reasons Ian doesn't understand we should not load libgomp # first using RTLD_DEEPBIND, so force loading it here if needed - if backend == 'openmp': - get_ctypes_library('gomp', [], mode=ctypes.DEFAULT_MODE) + if backend == "openmp": + get_ctypes_library("gomp", [], mode=ctypes.DEFAULT_MODE) # Note that the threaded libraries don't have their own pkg-config # files we must look for them wherever we look for double or single # FFTW itself. _double_threaded_lib = get_ctypes_library( - double_threaded_libname, - ['fftw3'], - mode=FFTW_RTLD_MODE + double_threaded_libname, ["fftw3"], mode=FFTW_RTLD_MODE ) - _float_threaded_lib = get_ctypes_library( - float_threaded_libname, - ['fftw3f'], - mode=FFTW_RTLD_MODE + _float_threaded_lib = get_ctypes_library( + float_threaded_libname, ["fftw3f"], mode=FFTW_RTLD_MODE ) if (_double_threaded_lib is None) or (_float_threaded_lib is None): - err_str = 'Unable to load threaded libraries' - err_str += f'{double_threaded_libname} or ' - err_str += f'{float_threaded_libname}' + err_str = "Unable to load threaded libraries" + err_str += f"{double_threaded_libname} or " + err_str += f"{float_threaded_libname}" raise RuntimeError(err_str) dret = _double_threaded_lib.fftw_init_threads() fret = _float_threaded_lib.fftwf_init_threads() @@ -146,6 +154,7 @@ def _init_threads(backend): _fftw_threaded_lib = backend return 0 + def set_threads_backend(backend=None): # This is the user facing function. If given a backend it just # calls _init_threads and lets it do the work. If not (the default) @@ -154,28 +163,35 @@ def set_threads_backend(backend=None): retval = _init_threads(backend) # Since the user specified this backend raise an exception if the above failed if retval != 0: - raise RuntimeError("Could not initialize FFTW threading backend {0}".format(backend)) + raise RuntimeError( + f"Could not initialize FFTW threading backend {backend}" + ) else: # Note that we pop() from the end, so 'pthreads' # is the first thing tried - _backend_list = ['unthreaded','openmp', 'pthreads'] + _backend_list = ["unthreaded", "openmp", "pthreads"] while not _fftw_threaded_set: _next_backend = _backend_list.pop() retval = _init_threads(_next_backend) + # Function to import system-wide wisdom files. + def import_sys_wisdom(): if not _fftw_threaded_set: set_threads_backend() double_lib.fftw_import_system_wisdom() float_lib.fftwf_import_system_wisdom() + # We provide an interface for changing the "measure level" # By default this is 0, which does no planning, # but we provide functions to read and set it _default_measurelvl = 0 + + def get_measure_level(): """ Get the current 'measure level' used in deciding how much effort to put into @@ -184,6 +200,7 @@ def get_measure_level(): """ return _default_measurelvl + def set_measure_level(mlvl): """ Set the current 'measure level' used in deciding how much effort to expend @@ -191,49 +208,62 @@ def set_measure_level(mlvl): to 3 (most effort and time). """ global _default_measurelvl - if mlvl not in (0,1,2,3): + if mlvl not in (0, 1, 2, 3): raise ValueError("Measure level can only be one of 0, 1, 2, or 3") _default_measurelvl = mlvl -_flag_dict = {0: FFTW_ESTIMATE, - 1: FFTW_MEASURE, - 2: FFTW_MEASURE|FFTW_PATIENT, - 3: FFTW_MEASURE|FFTW_PATIENT|FFTW_EXHAUSTIVE} -def get_flag(mlvl,aligned): + +_flag_dict = { + 0: FFTW_ESTIMATE, + 1: FFTW_MEASURE, + 2: FFTW_MEASURE | FFTW_PATIENT, + 3: FFTW_MEASURE | FFTW_PATIENT | FFTW_EXHAUSTIVE, +} + + +def get_flag(mlvl, aligned): if aligned: return _flag_dict[mlvl] - else: - return (_flag_dict[mlvl]|FFTW_UNALIGNED) + return _flag_dict[mlvl] | FFTW_UNALIGNED + # Add the ability to read/store wisdom to filenames + def wisdom_io(filename, precision, action): - """Import or export an FFTW plan for single or double precision. - """ + """Import or export an FFTW plan for single or double precision.""" if not _fftw_threaded_set: set_threads_backend() - fmap = {('float', 'import'): float_lib.fftwf_import_wisdom_from_filename, - ('float', 'export'): float_lib.fftwf_export_wisdom_to_filename, - ('double', 'import'): double_lib.fftw_import_wisdom_from_filename, - ('double', 'export'): double_lib.fftw_export_wisdom_to_filename} + fmap = { + ("float", "import"): float_lib.fftwf_import_wisdom_from_filename, + ("float", "export"): float_lib.fftwf_export_wisdom_to_filename, + ("double", "import"): double_lib.fftw_import_wisdom_from_filename, + ("double", "export"): double_lib.fftw_export_wisdom_to_filename, + } f = fmap[(precision, action)] f.argtypes = [ctypes.c_char_p] retval = f(filename.encode()) if retval == 0: - raise RuntimeError(('Could not {0} wisdom ' - 'from file {1}').format(action, filename)) + raise RuntimeError( + f"Could not {action} wisdom from file {filename}" + ) + def import_single_wisdom_from_filename(filename): - wisdom_io(filename, 'float', 'import') + wisdom_io(filename, "float", "import") + def import_double_wisdom_from_filename(filename): - wisdom_io(filename, 'double', 'import') + wisdom_io(filename, "double", "import") + def export_single_wisdom_to_filename(filename): - wisdom_io(filename, 'float', 'export') + wisdom_io(filename, "float", "export") + def export_double_wisdom_to_filename(filename): - wisdom_io(filename, 'double', 'export') + wisdom_io(filename, "double", "export") + def set_planning_limit(time): if not _fftw_threaded_set: @@ -247,22 +277,34 @@ def set_planning_limit(time): f.argtypes = [ctypes.c_double] f(time) + # Create function maps for the dtypes -plan_function = {'float32': {'complex64': float_lib.fftwf_plan_dft_r2c_1d}, - 'float64': {'complex128': double_lib.fftw_plan_dft_r2c_1d}, - 'complex64': {'float32': float_lib.fftwf_plan_dft_c2r_1d, - 'complex64': float_lib.fftwf_plan_dft_1d}, - 'complex128': {'float64': double_lib.fftw_plan_dft_c2r_1d, - 'complex128': double_lib.fftw_plan_dft_1d} - } - -execute_function = {'float32': {'complex64': float_lib.fftwf_execute_dft_r2c}, - 'float64': {'complex128': double_lib.fftw_execute_dft_r2c}, - 'complex64': {'float32': float_lib.fftwf_execute_dft_c2r, - 'complex64': float_lib.fftwf_execute_dft}, - 'complex128': {'float64': double_lib.fftw_execute_dft_c2r, - 'complex128': double_lib.fftw_execute_dft} - } +plan_function = { + "float32": {"complex64": float_lib.fftwf_plan_dft_r2c_1d}, + "float64": {"complex128": double_lib.fftw_plan_dft_r2c_1d}, + "complex64": { + "float32": float_lib.fftwf_plan_dft_c2r_1d, + "complex64": float_lib.fftwf_plan_dft_1d, + }, + "complex128": { + "float64": double_lib.fftw_plan_dft_c2r_1d, + "complex128": double_lib.fftw_plan_dft_1d, + }, +} + +execute_function = { + "float32": {"complex64": float_lib.fftwf_execute_dft_r2c}, + "float64": {"complex128": double_lib.fftw_execute_dft_r2c}, + "complex64": { + "float32": float_lib.fftwf_execute_dft_c2r, + "complex64": float_lib.fftwf_execute_dft, + }, + "complex128": { + "float64": double_lib.fftw_execute_dft_c2r, + "complex128": double_lib.fftw_execute_dft, + }, +} + def plan(size, idtype, odtype, direction, mlvl, aligned, nthreads, inplace): if not _fftw_threaded_set: @@ -270,24 +312,24 @@ def plan(size, idtype, odtype, direction, mlvl, aligned, nthreads, inplace): if nthreads != _fftw_current_nthreads: _fftw_plan_with_nthreads(nthreads) # Convert a measure-level to flags - flags = get_flag(mlvl,aligned) + flags = get_flag(mlvl, aligned) # We make our arrays of the necessary type and size. Things can be # tricky, especially for in-place transforms with one of input or # output real. - if (idtype == odtype): + if idtype == odtype: # We're in the complex-to-complex case, so lengths are the same ip = zeros(size, dtype=idtype) if inplace: op = ip else: op = zeros(size, dtype=odtype) - elif (idtype.kind == 'c') and (odtype.kind == 'f'): + elif (idtype.kind == "c") and (odtype.kind == "f"): # Complex-to-real (reverse), so size is length of real array. # However the complex array may be larger (in bytes) and # should therefore be allocated first and reused for an in-place # transform - ip = zeros(size/2+1, dtype=idtype) + ip = zeros(size / 2 + 1, dtype=idtype) if inplace: op = ip.view(dtype=odtype)[0:size] else: @@ -297,7 +339,7 @@ def plan(size, idtype, odtype, direction, mlvl, aligned, nthreads, inplace): # However it is still true that the complex array may be larger # (in bytes) and should therefore be allocated first and reused # for an in-place transform - op = zeros(size/2+1, dtype=odtype) + op = zeros(size / 2 + 1, dtype=odtype) if inplace: ip = op.view(dtype=idtype)[0:size] else: @@ -311,20 +353,24 @@ def plan(size, idtype, odtype, direction, mlvl, aligned, nthreads, inplace): # handle the C2C cases (forward and reverse) if idtype.kind == odtype.kind: - f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int, ctypes.c_int] + f.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ] theplan = f(size, ip.ptr, op.ptr, direction, flags) # handle the R2C and C2R case else: - f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int] + f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int] theplan = f(size, ip.ptr, op.ptr, flags) # We don't need ip or op anymore del ip, op # Make the destructors - if idtype.char in ['f', 'F']: + if idtype.char in ["f", "F"]: destroy = float_lib.fftwf_destroy_plan else: destroy = double_lib.fftw_destroy_plan @@ -341,17 +387,33 @@ def execute(plan, invec, outvec): f.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] f(plan, invec.ptr, outvec.ptr) + def fft(invec, outvec, prec, itype, otype): - theplan, destroy = plan(len(invec), invec.dtype, outvec.dtype, FFTW_FORWARD, - get_measure_level(),(check_aligned(invec.data) and check_aligned(outvec.data)), - _scheme.mgr.state.num_threads, (invec.ptr == outvec.ptr)) + theplan, destroy = plan( + len(invec), + invec.dtype, + outvec.dtype, + FFTW_FORWARD, + get_measure_level(), + (check_aligned(invec.data) and check_aligned(outvec.data)), + _scheme.mgr.state.num_threads, + (invec.ptr == outvec.ptr), + ) execute(theplan, invec, outvec) destroy(theplan) + def ifft(invec, outvec, prec, itype, otype): - theplan, destroy = plan(len(outvec), invec.dtype, outvec.dtype, FFTW_BACKWARD, - get_measure_level(),(check_aligned(invec.data) and check_aligned(outvec.data)), - _scheme.mgr.state.num_threads, (invec.ptr == outvec.ptr)) + theplan, destroy = plan( + len(outvec), + invec.dtype, + outvec.dtype, + FFTW_BACKWARD, + get_measure_level(), + (check_aligned(invec.data) and check_aligned(outvec.data)), + _scheme.mgr.state.num_threads, + (invec.ptr == outvec.ptr), + ) execute(theplan, invec, outvec) destroy(theplan) @@ -360,61 +422,126 @@ def ifft(invec, outvec, prec, itype, otype): # First, set up a lot of different ctypes functions: plan_many_c2c_f = float_lib.fftwf_plan_many_dft -plan_many_c2c_f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_int, ctypes.c_uint] +plan_many_c2c_f.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_uint, +] plan_many_c2c_f.restype = ctypes.c_void_p plan_many_c2c_d = double_lib.fftw_plan_many_dft -plan_many_c2c_d.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_int, ctypes.c_uint] +plan_many_c2c_d.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_uint, +] plan_many_c2c_d.restype = ctypes.c_void_p plan_many_c2r_f = float_lib.fftwf_plan_many_dft_c2r -plan_many_c2r_f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_uint] +plan_many_c2r_f.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_uint, +] plan_many_c2r_f.restype = ctypes.c_void_p plan_many_c2r_d = double_lib.fftw_plan_many_dft_c2r -plan_many_c2r_d.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_uint] +plan_many_c2r_d.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_uint, +] plan_many_c2r_d.restype = ctypes.c_void_p plan_many_r2c_f = float_lib.fftwf_plan_many_dft_r2c -plan_many_r2c_f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_uint] +plan_many_r2c_f.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_uint, +] plan_many_r2c_f.restype = ctypes.c_void_p plan_many_r2c_d = double_lib.fftw_plan_many_dft_r2c -plan_many_r2c_d.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, - ctypes.c_uint] +plan_many_r2c_d.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_uint, +] plan_many_r2c_d.restype = ctypes.c_void_p # Now set up a dictionary indexed by (str(input_dtype), str(output_dtype)) to # translate input and output dtypes into the correct planning function. -_plan_funcs_dict = { ('complex64', 'complex64') : plan_many_c2c_f, - ('float32', 'complex64') : plan_many_r2c_f, - ('complex64', 'float32') : plan_many_c2r_f, - ('complex128', 'complex128') : plan_many_c2c_d, - ('float64', 'complex128') : plan_many_r2c_d, - ('complex128', 'float64') : plan_many_c2r_d } +_plan_funcs_dict = { + ("complex64", "complex64"): plan_many_c2c_f, + ("float32", "complex64"): plan_many_r2c_f, + ("complex64", "float32"): plan_many_c2r_f, + ("complex128", "complex128"): plan_many_c2c_d, + ("float64", "complex128"): plan_many_r2c_d, + ("complex128", "float64"): plan_many_c2r_d, +} # To avoid multiple-inheritance, we set up a function that returns much # of the initialization that will need to be handled in __init__ of both # classes. + def _fftw_setup(fftobj): n = _np.asarray([fftobj.size], dtype=_np.int32) inembed = _np.asarray([len(fftobj.invec)], dtype=_np.int32) @@ -427,32 +554,54 @@ def _fftw_setup(fftobj): mlvl = get_measure_level() aligned = check_aligned(fftobj.invec.data) and check_aligned(fftobj.outvec.data) flags = get_flag(mlvl, aligned) - plan_func = _plan_funcs_dict[ (str(fftobj.invec.dtype), str(fftobj.outvec.dtype)) ] - tmpin = zeros(len(fftobj.invec), dtype = fftobj.invec.dtype) - tmpout = zeros(len(fftobj.outvec), dtype = fftobj.outvec.dtype) + plan_func = _plan_funcs_dict[(str(fftobj.invec.dtype), str(fftobj.outvec.dtype))] + tmpin = zeros(len(fftobj.invec), dtype=fftobj.invec.dtype) + tmpout = zeros(len(fftobj.outvec), dtype=fftobj.outvec.dtype) # C2C - if fftobj.outvec.kind == 'complex' and fftobj.invec.kind == 'complex': + if fftobj.outvec.kind == "complex" and fftobj.invec.kind == "complex": if fftobj.forward: ffd = FFTW_FORWARD else: ffd = FFTW_BACKWARD - plan = plan_func(1, n.ctypes.data, fftobj.nbatch, - tmpin.ptr, inembed.ctypes.data, 1, fftobj.idist, - tmpout.ptr, onembed.ctypes.data, 1, fftobj.odist, - ffd, flags) + plan = plan_func( + 1, + n.ctypes.data, + fftobj.nbatch, + tmpin.ptr, + inembed.ctypes.data, + 1, + fftobj.idist, + tmpout.ptr, + onembed.ctypes.data, + 1, + fftobj.odist, + ffd, + flags, + ) # R2C or C2R (hence no direction argument for plan creation) else: - plan = plan_func(1, n.ctypes.data, fftobj.nbatch, - tmpin.ptr, inembed.ctypes.data, 1, fftobj.idist, - tmpout.ptr, onembed.ctypes.data, 1, fftobj.odist, - flags) + plan = plan_func( + 1, + n.ctypes.data, + fftobj.nbatch, + tmpin.ptr, + inembed.ctypes.data, + 1, + fftobj.idist, + tmpout.ptr, + onembed.ctypes.data, + 1, + fftobj.odist, + flags, + ) del tmpin del tmpout return plan + class FFT(_BaseFFT): def __init__(self, invec, outvec, nbatch=1, size=None): - super(FFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.iptr = self.invec.ptr self.optr = self.outvec.ptr self._efunc = execute_function[str(self.invec.dtype)][str(self.outvec.dtype)] @@ -462,9 +611,10 @@ def __init__(self, invec, outvec, nbatch=1, size=None): def execute(self): self._efunc(self.plan, self.iptr, self.optr) + class IFFT(_BaseIFFT): def __init__(self, invec, outvec, nbatch=1, size=None): - super(IFFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.iptr = self.invec.ptr self.optr = self.outvec.ptr self._efunc = execute_function[str(self.invec.dtype)][str(self.outvec.dtype)] @@ -474,6 +624,7 @@ def __init__(self, invec, outvec, nbatch=1, size=None): def execute(self): self._efunc(self.plan, self.iptr, self.optr) + def insert_fft_options(optgroup): """ Inserts the options that affect the behavior of this backend @@ -482,32 +633,50 @@ def insert_fft_options(optgroup): ---------- optgroup: fft_option OptionParser argument group whose options are extended + + """ + optgroup.add_argument( + "--fftw-measure-level", + help="Determines the measure level used in planning " + "FFTW FFTs; allowed values are: " + str([0, 1, 2, 3]), + type=int, + default=_default_measurelvl, + ) + optgroup.add_argument( + "--fftw-threads-backend", + help="Give 'openmp', 'pthreads' or 'unthreaded' to specify which threaded FFTW to use", + default=None, + ) + optgroup.add_argument( + "--fftw-input-float-wisdom-file", + help="Filename from which to read single-precision wisdom", + default=None, + ) + optgroup.add_argument( + "--fftw-input-double-wisdom-file", + help="Filename from which to read double-precision wisdom", + default=None, + ) + optgroup.add_argument( + "--fftw-output-float-wisdom-file", + help="Filename to which to write single-precision wisdom", + default=None, + ) + optgroup.add_argument( + "--fftw-output-double-wisdom-file", + help="Filename to which to write double-precision wisdom", + default=None, + ) + optgroup.add_argument( + "--fftw-import-system-wisdom", + help="If given, call fftw[f]_import_system_wisdom()", + action="store_true", + ) + + +def verify_fft_options(opt, parser): """ - optgroup.add_argument("--fftw-measure-level", - help="Determines the measure level used in planning " - "FFTW FFTs; allowed values are: " + str([0,1,2,3]), - type=int, default=_default_measurelvl) - optgroup.add_argument("--fftw-threads-backend", - help="Give 'openmp', 'pthreads' or 'unthreaded' to specify which threaded FFTW to use", - default=None) - optgroup.add_argument("--fftw-input-float-wisdom-file", - help="Filename from which to read single-precision wisdom", - default=None) - optgroup.add_argument("--fftw-input-double-wisdom-file", - help="Filename from which to read double-precision wisdom", - default=None) - optgroup.add_argument("--fftw-output-float-wisdom-file", - help="Filename to which to write single-precision wisdom", - default=None) - optgroup.add_argument("--fftw-output-double-wisdom-file", - help="Filename to which to write double-precision wisdom", - default=None) - optgroup.add_argument("--fftw-import-system-wisdom", - help = "If given, call fftw[f]_import_system_wisdom()", - action = "store_true") - -def verify_fft_options(opt,parser): - """Parses the FFT options and verifies that they are + Parses the FFT options and verifies that they are reasonable. Parameters @@ -517,18 +686,28 @@ def verify_fft_options(opt,parser): required attributes. parser : object OptionParser instance. - """ - if opt.fftw_measure_level not in [0,1,2,3]: - parser.error("{0} is not a valid FFTW measure level.".format(opt.fftw_measure_level)) - if opt.fftw_import_system_wisdom and ((opt.fftw_input_float_wisdom_file is not None) - or (opt.fftw_input_double_wisdom_file is not None)): - parser.error("If --fftw-import-system-wisdom is given, then you cannot give" - " either of --fftw-input-float-wisdom-file or --fftw-input-double-wisdom-file") + """ + if opt.fftw_measure_level not in [0, 1, 2, 3]: + parser.error( + f"{opt.fftw_measure_level} is not a valid FFTW measure level." + ) + + if opt.fftw_import_system_wisdom and ( + (opt.fftw_input_float_wisdom_file is not None) + or (opt.fftw_input_double_wisdom_file is not None) + ): + parser.error( + "If --fftw-import-system-wisdom is given, then you cannot give" + " either of --fftw-input-float-wisdom-file or --fftw-input-double-wisdom-file" + ) if opt.fftw_threads_backend is not None: - if opt.fftw_threads_backend not in ['openmp','pthreads','unthreaded']: - parser.error("Invalid threads backend; must be 'openmp', 'pthreads' or 'unthreaded'") + if opt.fftw_threads_backend not in ["openmp", "pthreads", "unthreaded"]: + parser.error( + "Invalid threads backend; must be 'openmp', 'pthreads' or 'unthreaded'" + ) + def from_cli(opt): # Since opt.fftw_threads_backend defaults to None, the following is always diff --git a/pycbc/fft/fftw_pruned.py b/pycbc/fft/fftw_pruned.py index d4fd0fcdab9..8d630100193 100644 --- a/pycbc/fft/fftw_pruned.py +++ b/pycbc/fft/fftw_pruned.py @@ -1,4 +1,5 @@ -"""This module provides a functions to perform a pruned FFT based on FFTW +""" +This module provides a functions to perform a pruned FFT based on FFTW This should be considered a test and example module, as the functionality can and should be generalized to other FFT backends, and precisions. @@ -11,20 +12,28 @@ I use a similar naming convention here, with minor simplifications to the twiddle factors. """ -import numpy, ctypes, pycbc.types -from pycbc.libutils import get_ctypes_library + +import ctypes import logging + +import numpy + +import pycbc.types +from pycbc.libutils import get_ctypes_library + from .fftw_pruned_cython import second_phase_cython -logger = logging.getLogger('pycbc.events.fftw_pruned') +logger = logging.getLogger("pycbc.events.fftw_pruned") -warn_msg = ("The FFTW_pruned module can be used to speed up computing SNR " - "timeseries by computing first at a low sample rate and then " - "computing at full sample rate only at certain samples. This code " - "has not yet been used in production, and has no test case. " - "This was also ported to Cython in this state. " - "This code would need verification before trusting results. " - "Please do contribute test cases.") +warn_msg = ( + "The FFTW_pruned module can be used to speed up computing SNR " + "timeseries by computing first at a low sample rate and then " + "computing at full sample rate only at certain samples. This code " + "has not yet been used in production, and has no test case. " + "This was also ported to Cython in this state. " + "This code would need verification before trusting results. " + "Please do contribute test cases." +) logger.warning(warn_msg) @@ -34,31 +43,32 @@ FFTW_MEASURE = 0 FFTW_PATIENT = 1 << 5 FFTW_ESTIMATE = 1 << 6 -float_lib = get_ctypes_library('fftw3f', ['fftw3f'],mode=ctypes.RTLD_GLOBAL) +float_lib = get_ctypes_library("fftw3f", ["fftw3f"], mode=ctypes.RTLD_GLOBAL) fexecute = float_lib.fftwf_execute_dft fexecute.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] ftexecute = float_lib.fftwf_execute_dft ftexecute.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + def plan_transpose(N1, N2): """ Create a plan for transposing internally to the pruned_FFT calculation. (Alex to provide a write up with more details.) Parameters - ----------- + ---------- N1 : int Number of rows. N2 : int Number of columns. Returns - -------- + ------- plan : FFTWF plan The plan for performing the FFTW transpose. - """ + """ rows = N1 cols = N2 @@ -70,58 +80,90 @@ def plan_transpose(N1, N2): iodim[4] = rows iodim[5] = 1 - N = N1*N2 + N = N1 * N2 vin = pycbc.types.zeros(N, dtype=numpy.complex64) vout = pycbc.types.zeros(N, dtype=numpy.complex64) f = float_lib.fftwf_plan_guru_dft - f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int] + f.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ] f.restype = ctypes.c_void_p return f(0, None, 2, iodim.ctypes.data, vin.ptr, vout.ptr, None, FFTW_MEASURE) + def plan_first_phase(N1, N2): """ Create a plan for the first stage of the pruned FFT operation. (Alex to provide a write up with more details.) Parameters - ----------- + ---------- N1 : int Number of rows. N2 : int Number of columns. Returns - -------- + ------- plan : FFTWF plan The plan for performing the first phase FFT. + """ - N = N1*N2 + N = N1 * N2 vin = pycbc.types.zeros(N, dtype=numpy.complex64) vout = pycbc.types.zeros(N, dtype=numpy.complex64) f = float_lib.fftwf_plan_many_dft - f.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int, ctypes.c_int, - ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_int, ctypes.c_int, - ctypes.c_int, ctypes.c_int] + f.argtypes = [ + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ] f.restype = ctypes.c_void_p - return f(1, ctypes.byref(ctypes.c_int(N2)), N1, - vin.ptr, None, 1, N2, - vout.ptr, None, 1, N2, FFTW_BACKWARD, FFTW_MEASURE) + return f( + 1, + ctypes.byref(ctypes.c_int(N2)), + N1, + vin.ptr, + None, + 1, + N2, + vout.ptr, + None, + 1, + N2, + FFTW_BACKWARD, + FFTW_MEASURE, + ) + _theplan = None + + def first_phase(invec, outvec, N1, N2): """ This implements the first phase of the FFT decomposition, using the standard FFT many plans. Parameters - ----------- + ---------- invec : array The input array. outvec : array @@ -130,12 +172,14 @@ def first_phase(invec, outvec, N1, N2): Number of rows. N2 : int Number of columns. + """ global _theplan if _theplan is None: _theplan = plan_first_phase(N1, N2) fexecute(_theplan, invec.ptr, outvec.ptr) + def second_phase(invec, indices, N1, N2): """ This is the second phase of the FFT decomposition that actually performs @@ -157,9 +201,10 @@ def second_phase(invec, indices, N1, N2): Returns ------- out : array of floats + """ invec = numpy.array(invec.data, copy=False) - NI = len(indices) # pylint:disable=unused-variable + NI = len(indices) # pylint:disable=unused-variable N1 = int(N1) N2 = int(N2) out = numpy.zeros(len(indices), dtype=numpy.complex64) @@ -170,21 +215,25 @@ def second_phase(invec, indices, N1, N2): return out + _thetransposeplan = None + + def fft_transpose_fftw(vec): """ Perform an FFT transpose from vec into outvec. (Alex to provide more details in a write-up.) Parameters - ----------- + ---------- vec : array Input array. Returns - -------- + ------- outvec : array Transposed output array. + """ global _thetransposeplan outvec = pycbc.types.zeros(len(vec), dtype=vec.dtype) @@ -192,17 +241,19 @@ def fft_transpose_fftw(vec): N1, N2 = splay(vec) _thetransposeplan = plan_transpose(N1, N2) ftexecute(_thetransposeplan, vec.ptr, outvec.ptr) - return outvec + return outvec + fft_transpose = fft_transpose_fftw + def splay(vec): - """ Determine two lengths to split stride the input vector by - """ - N2 = 2 ** int(numpy.log2( len(vec) ) / 2) + """Determine two lengths to split stride the input vector by""" + N2 = 2 ** int(numpy.log2(len(vec)) / 2) N1 = len(vec) / N2 return N1, N2 + def pruned_c2cifft(invec, outvec, indices, pretransposed=False): """ Perform a pruned iFFT, only valid for power of 2 iffts as the @@ -211,7 +262,7 @@ def pruned_c2cifft(invec, outvec, indices, pretransposed=False): of 2. (Alex to provide more details in write up. Parameters - ----------- + ---------- invec : array The input vector. This should be the correlation between the data and the template at full sample rate. Ideally this is pre-transposed, but @@ -224,9 +275,10 @@ def pruned_c2cifft(invec, outvec, indices, pretransposed=False): Used to indicate whether or not invec is pretransposed. Returns - -------- + ------- SNRs : array The complex SNRs at the indexes given by indices. + """ N1, N2 = splay(invec) diff --git a/pycbc/fft/func_api.py b/pycbc/fft/func_api.py index 87355d24c09..a26c8a9c21d 100644 --- a/pycbc/fft/func_api.py +++ b/pycbc/fft/func_api.py @@ -26,13 +26,16 @@ implementations within PyCBC. """ -from pycbc.types import TimeSeries as _TimeSeries from pycbc.types import FrequencySeries as _FrequencySeries -from .core import _check_fft_args, _check_fwd_args, _check_inv_args +from pycbc.types import TimeSeries as _TimeSeries + from .backend_support import get_backend +from .core import _check_fft_args, _check_fwd_args, _check_inv_args + def fft(invec, outvec): - """ Fourier transform from invec to outvec. + """ + Fourier transform from invec to outvec. Perform a fourier transform. The type of transform is determined by the dtype of invec and outvec. @@ -43,6 +46,7 @@ def fft(invec, outvec): The input vector. outvec : TimeSeries or FrequencySeries The output. + """ prec, itype, otype = _check_fft_args(invec, outvec) _check_fwd_args(invec, itype, outvec, otype, 1, None) @@ -54,15 +58,17 @@ def fft(invec, outvec): # we should divide by, whether C2C or R2HC transform if isinstance(invec, _TimeSeries): outvec._epoch = invec._epoch - outvec._delta_f = 1.0/(invec._delta_t * len(invec)) + outvec._delta_f = 1.0 / (invec._delta_t * len(invec)) outvec *= invec._delta_t elif isinstance(invec, _FrequencySeries): outvec._epoch = invec._epoch - outvec._delta_t = 1.0/(invec._delta_f * len(invec)) + outvec._delta_t = 1.0 / (invec._delta_f * len(invec)) outvec *= invec._delta_f + def ifft(invec, outvec): - """ Inverse fourier transform from invec to outvec. + """ + Inverse fourier transform from invec to outvec. Perform an inverse fourier transform. The type of transform is determined by the dtype of invec and outvec. @@ -73,6 +79,7 @@ def ifft(invec, outvec): The input vector. outvec : TimeSeries or FrequencySeries The output. + """ prec, itype, otype = _check_fft_args(invec, outvec) _check_inv_args(invec, itype, outvec, otype, 1, None) @@ -84,10 +91,9 @@ def ifft(invec, outvec): # we should divide by, whether C2C or HC2R transform if isinstance(invec, _TimeSeries): outvec._epoch = invec._epoch - outvec._delta_f = 1.0/(invec._delta_t * len(outvec)) + outvec._delta_f = 1.0 / (invec._delta_t * len(outvec)) outvec *= invec._delta_t - elif isinstance(invec,_FrequencySeries): + elif isinstance(invec, _FrequencySeries): outvec._epoch = invec._epoch - outvec._delta_t = 1.0/(invec._delta_f * len(outvec)) + outvec._delta_t = 1.0 / (invec._delta_f * len(outvec)) outvec *= invec._delta_f - diff --git a/pycbc/fft/mkl.py b/pycbc/fft/mkl.py index 08bcf72699b..df0d2f15862 100644 --- a/pycbc/fft/mkl.py +++ b/pycbc/fft/mkl.py @@ -1,18 +1,21 @@ -import ctypes, pycbc.libutils +import ctypes + +import pycbc.libutils +import pycbc.scheme as _scheme from pycbc.types import zeros + from .core import _BaseFFT, _BaseIFFT -import pycbc.scheme as _scheme -lib = pycbc.libutils.get_ctypes_library('mkl_rt', []) +lib = pycbc.libutils.get_ctypes_library("mkl_rt", []) if lib is None: raise ImportError -#MKL constants taken from mkl_df_defines.h +# MKL constants taken from mkl_df_defines.h DFTI_FORWARD_DOMAIN = 0 DFTI_DIMENSION = 1 DFTI_LENGTHS = 2 DFTI_PRECISION = 3 -DFTI_FORWARD_SCALE = 4 +DFTI_FORWARD_SCALE = 4 DFTI_BACKWARD_SCALE = 5 DFTI_NUMBER_OF_TRANSFORMS = 7 DFTI_COMPLEX_STORAGE = 8 @@ -54,18 +57,23 @@ DFTI_PERM_FORMAT = 56 DFTI_CCE_FORMAT = 57 -mkl_domain = {'real': {'complex': DFTI_REAL}, - 'complex': {'real': DFTI_REAL, - 'complex':DFTI_COMPLEX, - } - } +mkl_domain = { + "real": {"complex": DFTI_REAL}, + "complex": { + "real": DFTI_REAL, + "complex": DFTI_COMPLEX, + }, +} + +mkl_descriptor = { + "single": lib.DftiCreateDescriptor_s_1d, + "double": lib.DftiCreateDescriptor_d_1d, +} -mkl_descriptor = {'single': lib.DftiCreateDescriptor_s_1d, - 'double': lib.DftiCreateDescriptor_d_1d, - } def check_status(status): - """ Check the status of a mkl functions and raise a python exeption if + """ + Check the status of a mkl functions and raise a python exeption if there is an error. """ if status: @@ -73,6 +81,7 @@ def check_status(status): msg = lib.DftiErrorMessage(status) raise RuntimeError(msg) + def create_descriptor(size, idtype, odtype, inplace): invec = zeros(1, dtype=idtype) outvec = zeros(1, dtype=odtype) @@ -99,26 +108,38 @@ def create_descriptor(size, idtype, odtype, inplace): return desc + def fft(invec, outvec, prec, itype, otype): - descr = create_descriptor(max(len(invec), len(outvec)), invec.dtype, - outvec.dtype, (invec.ptr == outvec.ptr)) + descr = create_descriptor( + max(len(invec), len(outvec)), + invec.dtype, + outvec.dtype, + (invec.ptr == outvec.ptr), + ) f = lib.DftiComputeForward f.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] status = f(descr, invec.ptr, outvec.ptr) lib.DftiFreeDescriptor(ctypes.byref(descr)) check_status(status) + def ifft(invec, outvec, prec, itype, otype): - descr = create_descriptor(max(len(invec), len(outvec)), invec.dtype, - outvec.dtype, (invec.ptr == outvec.ptr)) + descr = create_descriptor( + max(len(invec), len(outvec)), + invec.dtype, + outvec.dtype, + (invec.ptr == outvec.ptr), + ) f = lib.DftiComputeBackward f.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] status = f(descr, invec.ptr, outvec.ptr) lib.DftiFreeDescriptor(ctypes.byref(descr)) check_status(status) + # Class based API + def _get_desc(fftobj): desc = ctypes.c_void_p(1) domain = mkl_domain[str(fftobj.invec.kind)][str(fftobj.outvec.kind)] @@ -134,8 +155,7 @@ def _get_desc(fftobj): lib.DftiSetValue.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int] # The following only matters if the transform is C2R or R2C - status = lib.DftiSetValue(desc, DFTI_CONJUGATE_EVEN_STORAGE, - DFTI_COMPLEX_COMPLEX) + status = lib.DftiSetValue(desc, DFTI_CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX) check_status(status) # In-place or out-of-place: @@ -165,9 +185,10 @@ def _get_desc(fftobj): return desc + class FFT(_BaseFFT): def __init__(self, invec, outvec, nbatch=1, size=None): - super(FFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.iptr = self.invec.ptr self.optr = self.outvec.ptr self._efunc = lib.DftiComputeForward @@ -177,9 +198,10 @@ def __init__(self, invec, outvec, nbatch=1, size=None): def execute(self): self._efunc(self.desc, self.iptr, self.optr) + class IFFT(_BaseIFFT): def __init__(self, invec, outvec, nbatch=1, size=None): - super(IFFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) self.iptr = self.invec.ptr self.optr = self.outvec.ptr self._efunc = lib.DftiComputeBackward diff --git a/pycbc/fft/npfft.py b/pycbc/fft/npfft.py index 77439dca58c..ec1dcb658c9 100644 --- a/pycbc/fft/npfft.py +++ b/pycbc/fft/npfft.py @@ -27,57 +27,64 @@ """ import logging + import numpy.fft -from .core import _check_fft_args -from .core import _BaseFFT, _BaseIFFT -logger = logging.getLogger('pycbc.events.npfft') +from .core import _BaseFFT, _BaseIFFT, _check_fft_args + +logger = logging.getLogger("pycbc.events.npfft") + +_INV_FFT_MSG = ( + "I cannot perform an {} between data with an input type of " + "{} and an output type of {}" +) -_INV_FFT_MSG = ("I cannot perform an {} between data with an input type of " - "{} and an output type of {}") def fft(invec, outvec, _, itype, otype): if invec.ptr == outvec.ptr: - raise NotImplementedError("numpy backend of pycbc.fft does not " - "support in-place transforms") - if itype == 'complex' and otype == 'complex': - outvec.data[:] = numpy.asarray(numpy.fft.fft(invec.data), - dtype=outvec.dtype) - elif itype == 'real' and otype == 'complex': - outvec.data[:] = numpy.asarray(numpy.fft.rfft(invec.data), - dtype=outvec.dtype) + raise NotImplementedError( + "numpy backend of pycbc.fft does not support in-place transforms" + ) + if itype == "complex" and otype == "complex": + outvec.data[:] = numpy.asarray(numpy.fft.fft(invec.data), dtype=outvec.dtype) + elif itype == "real" and otype == "complex": + outvec.data[:] = numpy.asarray(numpy.fft.rfft(invec.data), dtype=outvec.dtype) else: raise ValueError(_INV_FFT_MSG.format("FFT", itype, otype)) def ifft(invec, outvec, _, itype, otype): if invec.ptr == outvec.ptr: - raise NotImplementedError("numpy backend of pycbc.fft does not " - "support in-place transforms") - if itype == 'complex' and otype == 'complex': - outvec.data[:] = numpy.asarray(numpy.fft.ifft(invec.data), - dtype=outvec.dtype) + raise NotImplementedError( + "numpy backend of pycbc.fft does not support in-place transforms" + ) + if itype == "complex" and otype == "complex": + outvec.data[:] = numpy.asarray(numpy.fft.ifft(invec.data), dtype=outvec.dtype) outvec *= len(outvec) - elif itype == 'complex' and otype == 'real': - outvec.data[:] = numpy.asarray(numpy.fft.irfft(invec.data,len(outvec)), - dtype=outvec.dtype) + elif itype == "complex" and otype == "real": + outvec.data[:] = numpy.asarray( + numpy.fft.irfft(invec.data, len(outvec)), dtype=outvec.dtype + ) outvec *= len(outvec) else: raise ValueError(_INV_FFT_MSG.format("IFFT", itype, otype)) -WARN_MSG = ("You are using the class-based PyCBC FFT API, with the numpy " - "backed. This is provided for convenience only. If performance is " - "important use the class-based API with one of the other backends " - "(for e.g. MKL or FFTW)") +WARN_MSG = ( + "You are using the class-based PyCBC FFT API, with the numpy " + "backed. This is provided for convenience only. If performance is " + "important use the class-based API with one of the other backends " + "(for e.g. MKL or FFTW)" +) class FFT(_BaseFFT): """ Class for performing FFTs via the numpy interface. """ + def __init__(self, invec, outvec, nbatch=1, size=None): - super(FFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) logger.warning(WARN_MSG) self.prec, self.itype, self.otype = _check_fft_args(invec, outvec) @@ -89,8 +96,9 @@ class IFFT(_BaseIFFT): """ Class for performing IFFTs via the numpy interface. """ + def __init__(self, invec, outvec, nbatch=1, size=None): - super(IFFT, self).__init__(invec, outvec, nbatch, size) + super().__init__(invec, outvec, nbatch, size) logger.warning(WARN_MSG) self.prec, self.itype, self.otype = _check_fft_args(invec, outvec) diff --git a/pycbc/fft/parser_support.py b/pycbc/fft/parser_support.py index 7879e5cc984..0ba842e389a 100644 --- a/pycbc/fft/parser_support.py +++ b/pycbc/fft/parser_support.py @@ -26,12 +26,17 @@ implementations within PyCBC. """ -from .backend_support import get_backend_modules, get_backend_names -from .backend_support import set_backend, get_backend +from .backend_support import ( + get_backend, + get_backend_modules, + get_backend_names, + set_backend, +) # Next we add all of the machinery to set backends and their options # from the command line. + def insert_fft_option_group(parser): """ Adds the options used to choose an FFT backend. This should be used @@ -45,17 +50,23 @@ def insert_fft_option_group(parser): ---------- parser : object OptionParser instance + """ - fft_group = parser.add_argument_group("Options for selecting the" - " FFT backend and controlling its performance" - " in this program.") + fft_group = parser.add_argument_group( + "Options for selecting the" + " FFT backend and controlling its performance" + " in this program." + ) # We have one argument to specify the backends. This becomes the default list used # if none is specified for a particular call of fft() of ifft(). Note that this # argument expects a *list* of inputs, as indicated by the nargs='*'. - fft_group.add_argument("--fft-backends", - help="Preference list of the FFT backends. " - "Choices are: \n" + str(get_backend_names()), - nargs='*', default=[]) + fft_group.add_argument( + "--fft-backends", + help="Preference list of the FFT backends. " + "Choices are: \n" + str(get_backend_names()), + nargs="*", + default=[], + ) for backend in get_backend_modules(): try: @@ -63,8 +74,10 @@ def insert_fft_option_group(parser): except AttributeError: pass + def verify_fft_options(opt, parser): - """Parses the FFT options and verifies that they are + """ + Parses the FFT options and verifies that they are reasonable. Parameters @@ -74,13 +87,13 @@ def verify_fft_options(opt, parser): required attributes. parser : object OptionParser instance. - """ + """ if len(opt.fft_backends) > 0: _all_backends = get_backend_names() for backend in opt.fft_backends: if backend not in _all_backends: - parser.error("Backend {0} is not available".format(backend)) + parser.error(f"Backend {backend} is not available") for backend in get_backend_modules(): try: @@ -88,13 +101,16 @@ def verify_fft_options(opt, parser): except AttributeError: pass + # The following function is the only one that is designed # only to work with the active scheme. We'd like to fix that, # eventually, but it's non-trivial because of how poorly MKL # and FFTW cooperate. + def from_cli(opt): - """Parses the command line options and sets the FFT backend + """ + Parses the command line options and sets the FFT backend for each (available) scheme. Aside from setting the default backed for this context, this function will also call (if it exists) the from_cli function of the specified backends in @@ -110,8 +126,9 @@ def from_cli(opt): the required attributes. Returns - """ + ------- + """ set_backend(opt.fft_backends) # Eventually, we need to be able to parse command lines diff --git a/pycbc/filter/autocorrelation.py b/pycbc/filter/autocorrelation.py index 6e0e54efade..4442921b509 100644 --- a/pycbc/filter/autocorrelation.py +++ b/pycbc/filter/autocorrelation.py @@ -27,11 +27,14 @@ """ import numpy + from pycbc.filter.matchedfilter import correlate from pycbc.types import FrequencySeries, TimeSeries, zeros + def calculate_acf(data, delta_t=1.0, unbiased=False): - r"""Calculates the one-sided autocorrelation function. + r""" + Calculates the one-sided autocorrelation function. Calculates the autocorrelation function (ACF) and returns the one-sided ACF. The ACF is defined as the autocovariance divided by the variance. The @@ -46,7 +49,7 @@ def calculate_acf(data, delta_t=1.0, unbiased=False): the variance of :math:`X_{t}`. Parameters - ----------- + ---------- data : TimeSeries or numpy.array A TimeSeries or numpy.array of data. delta_t : float @@ -61,8 +64,8 @@ def calculate_acf(data, delta_t=1.0, unbiased=False): acf : numpy.array If data is a TimeSeries then acf will be a TimeSeries of the one-sided ACF. Else acf is a numpy.array. - """ + """ # if given a TimeSeries instance then get numpy.array if isinstance(data, TimeSeries): y = data.numpy() @@ -75,7 +78,7 @@ def calculate_acf(data, delta_t=1.0, unbiased=False): ny_orig = len(y) npad = 1 - while npad < 2*ny_orig: + while npad < 2 * ny_orig: npad = npad << 1 ypad = numpy.zeros(npad) ypad[:ny_orig] = y @@ -85,8 +88,9 @@ def calculate_acf(data, delta_t=1.0, unbiased=False): # correlate # do not need to give the congjugate since correlate function does it - cdata = FrequencySeries(zeros(len(fdata), dtype=fdata.dtype), - delta_f=fdata.delta_f, copy=False) + cdata = FrequencySeries( + zeros(len(fdata), dtype=fdata.dtype), delta_f=fdata.delta_f, copy=False + ) correlate(fdata, fdata, cdata) # IFFT correlated data to get unnormalized autocovariance time series @@ -96,19 +100,19 @@ def calculate_acf(data, delta_t=1.0, unbiased=False): # normalize the autocovariance # note that dividing by acf[0] is the same as ( y.var() * len(acf) ) if unbiased: - acf /= ( y.var() * numpy.arange(len(acf), 0, -1) ) + acf /= y.var() * numpy.arange(len(acf), 0, -1) else: acf /= acf[0] # return input datatype if isinstance(data, TimeSeries): return TimeSeries(acf, delta_t=delta_t) - else: - return acf + return acf def calculate_acl(data, m=5, dtype=int): - r"""Calculates the autocorrelation length (ACL). + r""" + Calculates the autocorrelation length (ACL). Given a normalized autocorrelation function :math:`\rho[i]` (by normalized, we mean that :math:`\rho[0] = 1`), the ACL :math:`\tau` is: @@ -133,7 +137,7 @@ def calculate_acl(data, m=5, dtype=int): N. Madras and A.D. Sokal, J. Stat. Phys. 50, 109 (1988). Parameters - ----------- + ---------- data : TimeSeries or array A TimeSeries of data. m : int @@ -148,8 +152,8 @@ def calculate_acl(data, m=5, dtype=int): acl : int or float The autocorrelation length. If the ACL cannot be estimated, returns ``numpy.inf``. - """ + """ # sanity check output data type if dtype not in [int, float]: raise ValueError("The dtype must be either int or float.") diff --git a/pycbc/filter/fotonfilter.py b/pycbc/filter/fotonfilter.py index 3cb4a6c4930..0513e606ead 100644 --- a/pycbc/filter/fotonfilter.py +++ b/pycbc/filter/fotonfilter.py @@ -16,17 +16,21 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import logging -import numpy import sys -from pycbc import frame + +import numpy # import dependencies that are not standard to pycbc from foton import Filter, iir2z -logger = logging.getLogger('pycbc.filter.fotonfilter') +from pycbc import frame + +logger = logging.getLogger("pycbc.filter.fotonfilter") + def get_swstat_bits(frame_filenames, swstat_channel_name, start_time, end_time): - ''' This function just checks the first time in the SWSTAT channel + """ + This function just checks the first time in the SWSTAT channel to see if the filter was on, it doesn't check times beyond that. This is just for a first test on a small chunck of data. @@ -39,11 +43,11 @@ def get_swstat_bits(frame_filenames, swstat_channel_name, start_time, end_time): Bit 12 = Filter module output switch on/off Bit 13 = Filter module limit switch on/off Bit 14 = Filter module history reset momentary switch - ''' - + """ # read frames - swstat = frame.read_frame(frame_filenames, swstat_channel_name, - start_time=start_time, end_time=end_time) + swstat = frame.read_frame( + frame_filenames, swstat_channel_name, start_time=start_time, end_time=end_time + ) # convert number in channel to binary bits = bin(int(swstat[0])) @@ -56,27 +60,26 @@ def get_swstat_bits(frame_filenames, swstat_channel_name, start_time, end_time): return bits[-10:], filterbank_off -def filter_data(data, filter_name, filter_file, bits, filterbank_off=False, - swstat_channel_name=None): - ''' +def filter_data( + data, filter_name, filter_file, bits, filterbank_off=False, swstat_channel_name=None +): + """ A naive function to determine if the filter was on at the time and then filter the data. - ''' - + """ # if filterbank is off then return a time series of zeroes if filterbank_off: return numpy.zeros(len(data)) # loop over the 10 filters in the filterbank for i in range(10): - # read the filter filter = Filter(filter_file[filter_name][i]) # if bit is on then filter the data - bit = int(bits[-(i+1)]) + bit = int(bits[-(i + 1)]) if bit: - logger.info('filtering with filter module %d', i) + logger.info("filtering with filter module %d", i) # if there are second-order sections then filter with them if len(filter.sections): @@ -86,22 +89,21 @@ def filter_data(data, filter_name, filter_file, bits, filterbank_off=False, else: coeffs = iir2z(filter_file[filter_name][i]) if len(coeffs) > 1: - logger.info( - 'Gain-only filter module return more than one number' - ) + logger.info("Gain-only filter module return more than one number") sys.exit() gain = coeffs[0] data = gain * data - return data + return data + def read_gain_from_frames(frame_filenames, gain_channel_name, start_time, end_time): - ''' + """ Returns the gain from the file. - ''' - + """ # get timeseries from frame - gain = frame.read_frame(frame_filenames, gain_channel_name, - start_time=start_time, end_time=end_time) + gain = frame.read_frame( + frame_filenames, gain_channel_name, start_time=start_time, end_time=end_time + ) return gain[0] diff --git a/pycbc/filter/matchedfilter.py b/pycbc/filter/matchedfilter.py index e974ab2eb39..7dab2af2f89 100644 --- a/pycbc/filter/matchedfilter.py +++ b/pycbc/filter/matchedfilter.py @@ -28,19 +28,27 @@ import logging from math import sqrt + import numpy -from pycbc.types import TimeSeries, FrequencySeries, zeros, Array -from pycbc.types import complex_same_precision_as, real_same_precision_as -from pycbc.fft import fft, ifft, IFFT +import pycbc import pycbc.scheme from pycbc import events from pycbc.events import ranking -import pycbc +from pycbc.fft import IFFT, fft, ifft +from pycbc.types import ( + Array, + FrequencySeries, + TimeSeries, + complex_same_precision_as, + real_same_precision_as, + zeros, +) + +logger = logging.getLogger("pycbc.filter.matchedfilter") -logger = logging.getLogger('pycbc.filter.matchedfilter') +BACKEND_PREFIX = "pycbc.filter.matchedfilter_" -BACKEND_PREFIX="pycbc.filter.matchedfilter_" @pycbc.scheme.schemed(BACKEND_PREFIX) def correlate(x, y, z): @@ -49,11 +57,12 @@ def correlate(x, y, z): raise ValueError(err_msg) -class BatchCorrelator(object): - """ Create a batch correlation engine - """ +class BatchCorrelator: + """Create a batch correlation engine""" + def __init__(self, xs, zs, size): - """ Correlate x and y, store in z. Arrays need not be equal length, but + """ + Correlate x and y, store in z. Arrays need not be equal length, but must be at least size long and of the same dtype. No error checking will be performed, so be careful. All dtypes must be complex64. Note, must be created within the processing context that it will be used in. @@ -84,11 +93,12 @@ def _correlate_factory(x, y, z): raise ValueError(err_msg) -class Correlator(object): - """ Create a correlator engine +class Correlator: + """ + Create a correlator engine Parameters - --------- + ---------- x : complex64 Input pycbc.types.Array (or subclass); it will be conjugated y : complex64 @@ -99,10 +109,12 @@ class Correlator(object): The addresses in memory of the data of all three parameter vectors must be the same modulo pycbc.PYCBC_ALIGNMENT + """ + def __new__(cls, *args, **kwargs): real_cls = _correlate_factory(*args, **kwargs) - return real_cls(*args, **kwargs) # pylint:disable=not-callable + return real_cls(*args, **kwargs) # pylint:disable=not-callable # The class below should serve as the parent for all schemed classes. @@ -113,7 +125,7 @@ def __new__(cls, *args, **kwargs): # http://stackoverflow.com/questions/2025562/inherit-docstrings-in-python-class-inheritance # # will work? Is there a better way? -class _BaseCorrelator(object): +class _BaseCorrelator: def correlate(self): """ Compute the correlation of the vectors specified at object @@ -123,15 +135,28 @@ def correlate(self): changing between invocations, but not their locations in memory or length. """ - pass -class MatchedFilterControl(object): - def __init__(self, low_frequency_cutoff, high_frequency_cutoff, snr_threshold, tlen, - delta_f, dtype, segment_list, template_output, use_cluster, - downsample_factor=1, upsample_threshold=1, upsample_method='pruned_fft', - gpu_callback_method='none', cluster_function='symmetric'): - """ Create a matched filter engine. +class MatchedFilterControl: + def __init__( + self, + low_frequency_cutoff, + high_frequency_cutoff, + snr_threshold, + tlen, + delta_f, + dtype, + segment_list, + template_output, + use_cluster, + downsample_factor=1, + upsample_threshold=1, + upsample_method="pruned_fft", + gpu_callback_method="none", + cluster_function="symmetric", + ): + """ + Create a matched filter engine. Parameters ---------- @@ -162,20 +187,23 @@ def __init__(self, low_frequency_cutoff, high_frequency_cutoff, snr_threshold, t sliding forward window; if 'symmetric', each window's peak is compared to the windows before and after it, and only kept as a trigger if larger than both. + """ # Assuming analysis time is constant across templates and segments, also # delta_f is constant across segments. self.tlen = tlen self.flen = self.tlen / 2 + 1 self.delta_f = delta_f - self.delta_t = 1.0/(self.delta_f * self.tlen) + self.delta_t = 1.0 / (self.delta_f * self.tlen) self.dtype = dtype self.snr_threshold = snr_threshold self.flow = low_frequency_cutoff self.fhigh = high_frequency_cutoff self.gpu_callback_method = gpu_callback_method - if cluster_function not in ['symmetric', 'findchirp']: - raise ValueError("MatchedFilter: 'cluster_function' must be either 'symmetric' or 'findchirp'") + if cluster_function not in ["symmetric", "findchirp"]: + raise ValueError( + "MatchedFilter: 'cluster_function' must be either 'symmetric' or 'findchirp'" + ) self.cluster_function = cluster_function self.segments = segment_list self.htilde = template_output @@ -184,48 +212,57 @@ def __init__(self, low_frequency_cutoff, high_frequency_cutoff, snr_threshold, t self.snr_mem = zeros(self.tlen, dtype=self.dtype) self.corr_mem = zeros(self.tlen, dtype=self.dtype) - if use_cluster and (cluster_function == 'symmetric'): - self.matched_filter_and_cluster = self.full_matched_filter_and_cluster_symm + if use_cluster and (cluster_function == "symmetric"): + self.matched_filter_and_cluster = ( + self.full_matched_filter_and_cluster_symm + ) # setup the threasholding/clustering operations for each segment self.threshold_and_clusterers = [] for seg in self.segments: thresh = events.ThresholdCluster(self.snr_mem[seg.analyze]) self.threshold_and_clusterers.append(thresh) - elif use_cluster and (cluster_function == 'findchirp'): - self.matched_filter_and_cluster = self.full_matched_filter_and_cluster_fc + elif use_cluster and (cluster_function == "findchirp"): + self.matched_filter_and_cluster = ( + self.full_matched_filter_and_cluster_fc + ) else: self.matched_filter_and_cluster = self.full_matched_filter_thresh_only # Assuming analysis time is constant across templates and segments, also # delta_f is constant across segments. - self.kmin, self.kmax = get_cutoff_indices(self.flow, self.fhigh, - self.delta_f, self.tlen) + self.kmin, self.kmax = get_cutoff_indices( + self.flow, self.fhigh, self.delta_f, self.tlen + ) # Set up the correlation operations for each analysis segment corr_slice = slice(self.kmin, self.kmax) self.correlators = [] for seg in self.segments: - corr = Correlator(self.htilde[corr_slice], - seg[corr_slice], - self.corr_mem[corr_slice]) + corr = Correlator( + self.htilde[corr_slice], seg[corr_slice], self.corr_mem[corr_slice] + ) self.correlators.append(corr) # setup up the ifft we will do self.ifft = IFFT(self.corr_mem, self.snr_mem) elif downsample_factor >= 1: - self.matched_filter_and_cluster = self.hierarchical_matched_filter_and_cluster + self.matched_filter_and_cluster = ( + self.hierarchical_matched_filter_and_cluster + ) self.downsample_factor = downsample_factor self.upsample_method = upsample_method self.upsample_threshold = upsample_threshold N_full = self.tlen N_red = N_full / downsample_factor - self.kmin_full, self.kmax_full = get_cutoff_indices(self.flow, - self.fhigh, self.delta_f, N_full) + self.kmin_full, self.kmax_full = get_cutoff_indices( + self.flow, self.fhigh, self.delta_f, N_full + ) - self.kmin_red, _ = get_cutoff_indices(self.flow, - self.fhigh, self.delta_f, N_red) + self.kmin_red, _ = get_cutoff_indices( + self.flow, self.fhigh, self.delta_f, N_red + ) if self.kmax_full < N_red: self.kmax_red = self.kmax_full @@ -233,15 +270,20 @@ def __init__(self, low_frequency_cutoff, high_frequency_cutoff, snr_threshold, t self.kmax_red = N_red - 1 self.snr_mem = zeros(N_red, dtype=self.dtype) - self.corr_mem_full = FrequencySeries(zeros(N_full, dtype=self.dtype), delta_f=self.delta_f) + self.corr_mem_full = FrequencySeries( + zeros(N_full, dtype=self.dtype), delta_f=self.delta_f + ) self.corr_mem = Array(self.corr_mem_full[0:N_red], copy=False) self.inter_vec = zeros(N_full, dtype=self.dtype) else: raise ValueError("Invalid downsample factor") - def full_matched_filter_and_cluster_symm(self, segnum, template_norm, window, epoch=None): - """ Returns the complex snr timeseries, normalization of the complex snr, + def full_matched_filter_and_cluster_symm( + self, segnum, template_norm, window, epoch=None + ): + """ + Returns the complex snr timeseries, normalization of the complex snr, the correlation vector frequency series, the list of indices of the triggers, and the snr values at the trigger locations. Returns empty lists for these for points that are not above the threshold. @@ -270,11 +312,14 @@ def full_matched_filter_and_cluster_symm(self, segnum, template_norm, window, ep List of indices of the triggers. snrv : Array The snr values at the trigger locations. + """ norm = (4.0 * self.delta_f) / sqrt(template_norm) self.correlators[segnum].correlate() self.ifft.execute() - snrv, idx = self.threshold_and_clusterers[segnum].threshold_and_cluster(self.snr_threshold / norm, window) + snrv, idx = self.threshold_and_clusterers[segnum].threshold_and_cluster( + self.snr_threshold / norm, window + ) if len(idx) == 0: return [], [], [], [], [] @@ -285,8 +330,11 @@ def full_matched_filter_and_cluster_symm(self, segnum, template_norm, window, ep corr = FrequencySeries(self.corr_mem, delta_f=self.delta_f, copy=False) return snr, norm, corr, idx, snrv - def full_matched_filter_and_cluster_fc(self, segnum, template_norm, window, epoch=None): - """ Returns the complex snr timeseries, normalization of the complex snr, + def full_matched_filter_and_cluster_fc( + self, segnum, template_norm, window, epoch=None + ): + """ + Returns the complex snr timeseries, normalization of the complex snr, the correlation vector frequency series, the list of indices of the triggers, and the snr values at the trigger locations. Returns empty lists for these for points that are not above the threshold. @@ -315,12 +363,14 @@ def full_matched_filter_and_cluster_fc(self, segnum, template_norm, window, epoc List of indices of the triggers. snrv : Array The snr values at the trigger locations. + """ norm = (4.0 * self.delta_f) / sqrt(template_norm) self.correlators[segnum].correlate() self.ifft.execute() - idx, snrv = events.threshold(self.snr_mem[self.segments[segnum].analyze], - self.snr_threshold / norm) + idx, snrv = events.threshold( + self.snr_mem[self.segments[segnum].analyze], self.snr_threshold / norm + ) idx, snrv = events.cluster_reduce(idx, snrv, window) if len(idx) == 0: @@ -332,8 +382,11 @@ def full_matched_filter_and_cluster_fc(self, segnum, template_norm, window, epoc corr = FrequencySeries(self.corr_mem, delta_f=self.delta_f, copy=False) return snr, norm, corr, idx, snrv - def full_matched_filter_thresh_only(self, segnum, template_norm, window=None, epoch=None): - """ Returns the complex snr timeseries, normalization of the complex snr, + def full_matched_filter_thresh_only( + self, segnum, template_norm, window=None, epoch=None + ): + """ + Returns the complex snr timeseries, normalization of the complex snr, the correlation vector frequency series, the list of indices of the triggers, and the snr values at the trigger locations. Returns empty lists for these for points that are not above the threshold. @@ -363,12 +416,14 @@ def full_matched_filter_thresh_only(self, segnum, template_norm, window=None, ep List of indices of the triggers. snrv : Array The snr values at the trigger locations. + """ norm = (4.0 * self.delta_f) / sqrt(template_norm) self.correlators[segnum].correlate() self.ifft.execute() - idx, snrv = events.threshold_only(self.snr_mem[self.segments[segnum].analyze], - self.snr_threshold / norm) + idx, snrv = events.threshold_only( + self.snr_mem[self.segments[segnum].analyze], self.snr_threshold / norm + ) logger.info("%d points above threshold", len(idx)) snr = TimeSeries(self.snr_mem, epoch=epoch, delta_t=self.delta_t, copy=False) @@ -376,7 +431,8 @@ def full_matched_filter_thresh_only(self, segnum, template_norm, window=None, ep return snr, norm, corr, idx, snrv def hierarchical_matched_filter_and_cluster(self, segnum, template_norm, window): - """ Returns the complex snr timeseries, normalization of the complex snr, + """ + Returns the complex snr timeseries, normalization of the complex snr, the correlation vector frequency series, the list of indices of the triggers, and the snr values at the trigger locations. Returns empty lists for these for points that are not above the threshold. @@ -404,78 +460,105 @@ def hierarchical_matched_filter_and_cluster(self, segnum, template_norm, window) List of indices of the triggers. snrv : Array The snr values at the trigger locations. + """ - from pycbc.fft.fftw_pruned import pruned_c2cifft, fft_transpose + from pycbc.fft.fftw_pruned import fft_transpose, pruned_c2cifft + htilde = self.htilde stilde = self.segments[segnum] norm = (4.0 * stilde.delta_f) / sqrt(template_norm) - correlate(htilde[self.kmin_red:self.kmax_red], - stilde[self.kmin_red:self.kmax_red], - self.corr_mem[self.kmin_red:self.kmax_red]) + correlate( + htilde[self.kmin_red : self.kmax_red], + stilde[self.kmin_red : self.kmax_red], + self.corr_mem[self.kmin_red : self.kmax_red], + ) ifft(self.corr_mem, self.snr_mem) - if not hasattr(stilde, 'red_analyze'): - stilde.red_analyze = \ - slice(stilde.analyze.start/self.downsample_factor, - stilde.analyze.stop/self.downsample_factor) - + if not hasattr(stilde, "red_analyze"): + stilde.red_analyze = slice( + stilde.analyze.start / self.downsample_factor, + stilde.analyze.stop / self.downsample_factor, + ) - idx_red, snrv_red = events.threshold(self.snr_mem[stilde.red_analyze], - self.snr_threshold / norm * self.upsample_threshold) + idx_red, snrv_red = events.threshold( + self.snr_mem[stilde.red_analyze], + self.snr_threshold / norm * self.upsample_threshold, + ) if len(idx_red) == 0: return [], None, [], [], [] - idx_red, _ = events.cluster_reduce(idx_red, snrv_red, window / self.downsample_factor) - logger.info("%d points above threshold at reduced resolution", - len(idx_red)) + idx_red, _ = events.cluster_reduce( + idx_red, snrv_red, window / self.downsample_factor + ) + logger.info("%d points above threshold at reduced resolution", len(idx_red)) # The fancy upsampling is here - if self.upsample_method=='pruned_fft': - idx = (idx_red + stilde.analyze.start/self.downsample_factor)\ - * self.downsample_factor + if self.upsample_method == "pruned_fft": + idx = ( + idx_red + stilde.analyze.start / self.downsample_factor + ) * self.downsample_factor idx = smear(idx, self.downsample_factor) # cache transposed versions of htilde and stilde - if not hasattr(self.corr_mem_full, 'transposed'): - self.corr_mem_full.transposed = zeros(len(self.corr_mem_full), dtype=self.dtype) + if not hasattr(self.corr_mem_full, "transposed"): + self.corr_mem_full.transposed = zeros( + len(self.corr_mem_full), dtype=self.dtype + ) - if not hasattr(htilde, 'transposed'): + if not hasattr(htilde, "transposed"): htilde.transposed = zeros(len(self.corr_mem_full), dtype=self.dtype) - htilde.transposed[self.kmin_full:self.kmax_full] = htilde[self.kmin_full:self.kmax_full] + htilde.transposed[self.kmin_full : self.kmax_full] = htilde[ + self.kmin_full : self.kmax_full + ] htilde.transposed = fft_transpose(htilde.transposed) - if not hasattr(stilde, 'transposed'): + if not hasattr(stilde, "transposed"): stilde.transposed = zeros(len(self.corr_mem_full), dtype=self.dtype) - stilde.transposed[self.kmin_full:self.kmax_full] = stilde[self.kmin_full:self.kmax_full] + stilde.transposed[self.kmin_full : self.kmax_full] = stilde[ + self.kmin_full : self.kmax_full + ] stilde.transposed = fft_transpose(stilde.transposed) - correlate(htilde.transposed, stilde.transposed, self.corr_mem_full.transposed) - snrv = pruned_c2cifft(self.corr_mem_full.transposed, self.inter_vec, idx, pretransposed=True) + correlate( + htilde.transposed, stilde.transposed, self.corr_mem_full.transposed + ) + snrv = pruned_c2cifft( + self.corr_mem_full.transposed, self.inter_vec, idx, pretransposed=True + ) idx = idx - stilde.analyze.start - idx2, snrv = events.threshold(Array(snrv, copy=False), self.snr_threshold / norm) + idx2, snrv = events.threshold( + Array(snrv, copy=False), self.snr_threshold / norm + ) if len(idx2) > 0: - correlate(htilde[self.kmax_red:self.kmax_full], - stilde[self.kmax_red:self.kmax_full], - self.corr_mem_full[self.kmax_red:self.kmax_full]) + correlate( + htilde[self.kmax_red : self.kmax_full], + stilde[self.kmax_red : self.kmax_full], + self.corr_mem_full[self.kmax_red : self.kmax_full], + ) idx, snrv = events.cluster_reduce(idx[idx2], snrv, window) else: idx, snrv = [], [] logger.info("%d points at full rate and clustering", len(idx)) return self.snr_mem, norm, self.corr_mem_full, idx, snrv - else: - raise ValueError("Invalid upsample method") - - -def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, - hpnorm=None, hcnorm=None, - out=None, thresh=0, - analyse_slice=None): + raise ValueError("Invalid upsample method") + + +def compute_max_snr_over_sky_loc_stat( + hplus, + hcross, + hphccorr, + hpnorm=None, + hcnorm=None, + out=None, + thresh=0, + analyse_slice=None, +): """ Matched filter maximised over polarization and orbital phase. @@ -484,7 +567,7 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, in this statistic before using it. Parameters - ----------- + ---------- hplus : TimeSeries This is the IFFTed complex SNR time series of (h+, data). If not normalized, supply the normalization factor so this can be done! @@ -512,9 +595,10 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, If given, use this array to store the output. Returns - -------- + ------- det_stat : TimeSeries The SNR maximized over sky location + """ # NOTE: Not much optimization has been done here! This may need to be # Cythonized. @@ -522,37 +606,35 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, if out is None: out = zeros(len(hplus)) out.non_zero_locs = numpy.array([], dtype=out.dtype) + elif not hasattr(out, "non_zero_locs"): + # Doing this every time is not a zero-cost operation + out.data[:] = 0 + out.non_zero_locs = numpy.array([], dtype=out.dtype) else: - if not hasattr(out, 'non_zero_locs'): - # Doing this every time is not a zero-cost operation - out.data[:] = 0 - out.non_zero_locs = numpy.array([], dtype=out.dtype) - else: - # Only set non zero locations to zero - out.data[out.non_zero_locs] = 0 - + # Only set non zero locations to zero + out.data[out.non_zero_locs] = 0 # If threshold is given we can limit the points at which to compute the # full statistic if thresh: # This is the statistic that always overestimates the SNR... # It allows some unphysical freedom that the full statistic does not - idx_p, _ = events.threshold_only(hplus[analyse_slice], - thresh / (2**0.5 * hpnorm)) - idx_c, _ = events.threshold_only(hcross[analyse_slice], - thresh / (2**0.5 * hcnorm)) + idx_p, _ = events.threshold_only( + hplus[analyse_slice], thresh / (2**0.5 * hpnorm) + ) + idx_c, _ = events.threshold_only( + hcross[analyse_slice], thresh / (2**0.5 * hcnorm) + ) idx_p = idx_p + analyse_slice.start idx_c = idx_c + analyse_slice.start hp_red = hplus[idx_p] * hpnorm hc_red = hcross[idx_p] * hcnorm - stat_p = hp_red.real**2 + hp_red.imag**2 + \ - hc_red.real**2 + hc_red.imag**2 - locs_p = idx_p[stat_p > (thresh*thresh)] + stat_p = hp_red.real**2 + hp_red.imag**2 + hc_red.real**2 + hc_red.imag**2 + locs_p = idx_p[stat_p > (thresh * thresh)] hp_red = hplus[idx_c] * hpnorm hc_red = hcross[idx_c] * hcnorm - stat_c = hp_red.real**2 + hp_red.imag**2 + \ - hc_red.real**2 + hc_red.imag**2 - locs_c = idx_c[stat_c > (thresh*thresh)] + stat_c = hp_red.real**2 + hp_red.imag**2 + hc_red.real**2 + hc_red.imag**2 + locs_c = idx_c[stat_c > (thresh * thresh)] locs = numpy.unique(numpy.concatenate((locs_p, locs_c))) hplus = hplus[locs] @@ -561,18 +643,16 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, hplus = hplus * hpnorm hcross = hcross * hcnorm - # Calculate and sanity check the denominator - denom = 1 - hphccorr*hphccorr + denom = 1 - hphccorr * hphccorr if denom < 0: if hphccorr > 1: - err_msg = "Overlap between hp and hc is given as %f. " %(hphccorr) + err_msg = "Overlap between hp and hc is given as %f. " % (hphccorr) err_msg += "How can an overlap be bigger than 1?" raise ValueError(err_msg) - else: - err_msg = "There really is no way to raise this error!?! " - err_msg += "If you're seeing this, it is bad." - raise ValueError(err_msg) + err_msg = "There really is no way to raise this error!?! " + err_msg += "If you're seeing this, it is bad." + raise ValueError(err_msg) if denom == 0: # This case, of hphccorr==1, makes the statistic degenerate # This case should not physically be possible luckily. @@ -581,18 +661,25 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, err_msg += "so why are you seeing this?" raise ValueError(err_msg) - assert(len(hplus) == len(hcross)) + assert len(hplus) == len(hcross) # Now the stuff where comp. cost may be a problem - hplus_magsq = numpy.real(hplus) * numpy.real(hplus) + \ - numpy.imag(hplus) * numpy.imag(hplus) - hcross_magsq = numpy.real(hcross) * numpy.real(hcross) + \ - numpy.imag(hcross) * numpy.imag(hcross) - rho_pluscross = numpy.real(hplus) * numpy.real(hcross) + numpy.imag(hplus)*numpy.imag(hcross) - - sqroot = (hplus_magsq - hcross_magsq)**2 - sqroot += 4 * (hphccorr * hplus_magsq - rho_pluscross) * \ - (hphccorr * hcross_magsq - rho_pluscross) + hplus_magsq = numpy.real(hplus) * numpy.real(hplus) + numpy.imag( + hplus + ) * numpy.imag(hplus) + hcross_magsq = numpy.real(hcross) * numpy.real(hcross) + numpy.imag( + hcross + ) * numpy.imag(hcross) + rho_pluscross = numpy.real(hplus) * numpy.real(hcross) + numpy.imag( + hplus + ) * numpy.imag(hcross) + + sqroot = (hplus_magsq - hcross_magsq) ** 2 + sqroot += ( + 4 + * (hphccorr * hplus_magsq - rho_pluscross) + * (hphccorr * hcross_magsq - rho_pluscross) + ) # Sometimes this can be less than 0 due to numeric imprecision, catch this. if (sqroot < 0).any(): indices = numpy.arange(len(sqroot))[sqroot < 0] @@ -602,8 +689,11 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, raise ValueError(err_msg) sqroot[indices] = 0 sqroot = numpy.sqrt(sqroot) - det_stat_sq = 0.5 * (hplus_magsq + hcross_magsq - \ - 2 * rho_pluscross*hphccorr + sqroot) / denom + det_stat_sq = ( + 0.5 + * (hplus_magsq + hcross_magsq - 2 * rho_pluscross * hphccorr + sqroot) + / denom + ) det_stat = numpy.sqrt(det_stat_sq) @@ -611,12 +701,14 @@ def compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, out.data[locs] = det_stat out.non_zero_locs = locs return out - else: - return Array(det_stat, copy=False) + return Array(det_stat, copy=False) -def compute_u_val_for_sky_loc_stat(hplus, hcross, hphccorr, - hpnorm=None, hcnorm=None, indices=None): - """The max-over-sky location detection statistic maximizes over a phase, + +def compute_u_val_for_sky_loc_stat( + hplus, hcross, hphccorr, hpnorm=None, hcnorm=None, indices=None +): + """ + The max-over-sky location detection statistic maximizes over a phase, an amplitude and the ratio of F+ and Fx, encoded in a variable called u. Here we return the value of u for the given indices. """ @@ -631,42 +723,52 @@ def compute_u_val_for_sky_loc_stat(hplus, hcross, hphccorr, # Sanity checking in func. above should already have identified any points # which are bad, and should be used to construct indices for input here - hplus_magsq = numpy.real(hplus) * numpy.real(hplus) + \ - numpy.imag(hplus) * numpy.imag(hplus) - hcross_magsq = numpy.real(hcross) * numpy.real(hcross) + \ - numpy.imag(hcross) * numpy.imag(hcross) - rho_pluscross = numpy.real(hplus) * numpy.real(hcross) + \ - numpy.imag(hplus)*numpy.imag(hcross) + hplus_magsq = numpy.real(hplus) * numpy.real(hplus) + numpy.imag( + hplus + ) * numpy.imag(hplus) + hcross_magsq = numpy.real(hcross) * numpy.real(hcross) + numpy.imag( + hcross + ) * numpy.imag(hcross) + rho_pluscross = numpy.real(hplus) * numpy.real(hcross) + numpy.imag( + hplus + ) * numpy.imag(hcross) a = hphccorr * hplus_magsq - rho_pluscross b = hplus_magsq - hcross_magsq c = rho_pluscross - hphccorr * hcross_magsq - sq_root = b*b - 4*a*c + sq_root = b * b - 4 * a * c sq_root = sq_root**0.5 sq_root = -sq_root # Catch the a->0 case - bad_lgc = (a == 0) + bad_lgc = a == 0 dbl_bad_lgc = numpy.logical_and(c == 0, b == 0) dbl_bad_lgc = numpy.logical_and(bad_lgc, dbl_bad_lgc) # Initialize u - u = sq_root * 0. + u = sq_root * 0.0 # In this case u is completely degenerate, so set it to 1 - u[dbl_bad_lgc] = 1. + u[dbl_bad_lgc] = 1.0 # If a->0 avoid overflow by just setting to a large value - u[bad_lgc & ~dbl_bad_lgc] = 1E17 + u[bad_lgc & ~dbl_bad_lgc] = 1e17 # Otherwise normal statistic - u[~bad_lgc] = (-b[~bad_lgc] + sq_root[~bad_lgc]) / (2*a[~bad_lgc]) + u[~bad_lgc] = (-b[~bad_lgc] + sq_root[~bad_lgc]) / (2 * a[~bad_lgc]) snr_cplx = hplus * u + hcross coa_phase = numpy.angle(snr_cplx) return u, coa_phase -def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, - hpnorm=None, hcnorm=None, - out=None, thresh=0, - analyse_slice=None): + +def compute_max_snr_over_sky_loc_stat_no_phase( + hplus, + hcross, + hphccorr, + hpnorm=None, + hcnorm=None, + out=None, + thresh=0, + analyse_slice=None, +): """ Matched filter maximised over polarization phase. @@ -680,7 +782,7 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, collapses to the normal statistic (at twice the computational cost!) Parameters - ----------- + ---------- hplus : TimeSeries This is the IFFTed complex SNR time series of (h+, data). If not normalized, supply the normalization factor so this can be done! @@ -708,9 +810,10 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, If given, use this array to store the output. Returns - -------- + ------- det_stat : TimeSeries The SNR maximized over sky location + """ # NOTE: Not much optimization has been done here! This may need to be # Cythonized. @@ -718,14 +821,13 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, if out is None: out = zeros(len(hplus)) out.non_zero_locs = numpy.array([], dtype=out.dtype) + elif not hasattr(out, "non_zero_locs"): + # Doing this every time is not a zero-cost operation + out.data[:] = 0 + out.non_zero_locs = numpy.array([], dtype=out.dtype) else: - if not hasattr(out, 'non_zero_locs'): - # Doing this every time is not a zero-cost operation - out.data[:] = 0 - out.non_zero_locs = numpy.array([], dtype=out.dtype) - else: - # Only set non zero locations to zero - out.data[out.non_zero_locs] = 0 + # Only set non zero locations to zero + out.data[out.non_zero_locs] = 0 # If threshold is given we can limit the points at which to compute the # full statistic @@ -736,22 +838,22 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, # For now this is copied from the max-over-phase statistic. One could # probably make this faster by removing the imaginary components of # the matched filter, as these are not used here. - idx_p, _ = events.threshold_only(hplus[analyse_slice], - thresh / (2**0.5 * hpnorm)) - idx_c, _ = events.threshold_only(hcross[analyse_slice], - thresh / (2**0.5 * hcnorm)) + idx_p, _ = events.threshold_only( + hplus[analyse_slice], thresh / (2**0.5 * hpnorm) + ) + idx_c, _ = events.threshold_only( + hcross[analyse_slice], thresh / (2**0.5 * hcnorm) + ) idx_p = idx_p + analyse_slice.start idx_c = idx_c + analyse_slice.start hp_red = hplus[idx_p] * hpnorm hc_red = hcross[idx_p] * hcnorm - stat_p = hp_red.real**2 + hp_red.imag**2 + \ - hc_red.real**2 + hc_red.imag**2 - locs_p = idx_p[stat_p > (thresh*thresh)] + stat_p = hp_red.real**2 + hp_red.imag**2 + hc_red.real**2 + hc_red.imag**2 + locs_p = idx_p[stat_p > (thresh * thresh)] hp_red = hplus[idx_c] * hpnorm hc_red = hcross[idx_c] * hcnorm - stat_c = hp_red.real**2 + hp_red.imag**2 + \ - hc_red.real**2 + hc_red.imag**2 - locs_c = idx_c[stat_c > (thresh*thresh)] + stat_c = hp_red.real**2 + hp_red.imag**2 + hc_red.real**2 + hc_red.imag**2 + locs_c = idx_c[stat_c > (thresh * thresh)] locs = numpy.unique(numpy.concatenate((locs_p, locs_c))) hplus = hplus[locs] @@ -760,18 +862,16 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, hplus = hplus * hpnorm hcross = hcross * hcnorm - # Calculate and sanity check the denominator - denom = 1 - hphccorr*hphccorr + denom = 1 - hphccorr * hphccorr if denom < 0: if hphccorr > 1: - err_msg = "Overlap between hp and hc is given as %f. " %(hphccorr) + err_msg = "Overlap between hp and hc is given as %f. " % (hphccorr) err_msg += "How can an overlap be bigger than 1?" raise ValueError(err_msg) - else: - err_msg = "There really is no way to raise this error!?! " - err_msg += "If you're seeing this, it is bad." - raise ValueError(err_msg) + err_msg = "There really is no way to raise this error!?! " + err_msg += "If you're seeing this, it is bad." + raise ValueError(err_msg) if denom == 0: # This case, of hphccorr==1, makes the statistic degenerate # This case should not physically be possible luckily. @@ -780,14 +880,14 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, err_msg += "so why are you seeing this?" raise ValueError(err_msg) - assert(len(hplus) == len(hcross)) + assert len(hplus) == len(hcross) # Now the stuff where comp. cost may be a problem hplus_magsq = numpy.real(hplus) * numpy.real(hplus) hcross_magsq = numpy.real(hcross) * numpy.real(hcross) rho_pluscross = numpy.real(hplus) * numpy.real(hcross) - det_stat_sq = (hplus_magsq + hcross_magsq - 2 * rho_pluscross*hphccorr) + det_stat_sq = hplus_magsq + hcross_magsq - 2 * rho_pluscross * hphccorr det_stat = numpy.sqrt(det_stat_sq / denom) @@ -795,12 +895,14 @@ def compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, hphccorr, out.data[locs] = det_stat out.non_zero_locs = locs return out - else: - return Array(det_stat, copy=False) + return Array(det_stat, copy=False) -def compute_u_val_for_sky_loc_stat_no_phase(hplus, hcross, hphccorr, - hpnorm=None , hcnorm=None, indices=None): - """The max-over-sky location (no phase) detection statistic maximizes over + +def compute_u_val_for_sky_loc_stat_no_phase( + hplus, hcross, hphccorr, hpnorm=None, hcnorm=None, indices=None +): + """ + The max-over-sky location (no phase) detection statistic maximizes over an amplitude and the ratio of F+ and Fx, encoded in a variable called u. Here we return the value of u for the given indices. @@ -815,27 +917,38 @@ def compute_u_val_for_sky_loc_stat_no_phase(hplus, hcross, hphccorr, if hcnorm is not None: hcross = hcross * hcnorm - rhoplusre=numpy.real(hplus) - rhocrossre=numpy.real(hcross) - overlap=numpy.real(hphccorr) + rhoplusre = numpy.real(hplus) + rhocrossre = numpy.real(hcross) + overlap = numpy.real(hphccorr) - denom = (-rhocrossre+overlap*rhoplusre) + denom = -rhocrossre + overlap * rhoplusre # Initialize tan_kappa array - u_val = denom * 0. + u_val = denom * 0.0 # Catch the denominator -> 0 case - numpy.putmask(u_val, denom == 0, 1E17) + numpy.putmask(u_val, denom == 0, 1e17) # Otherwise do normal statistic - numpy.putmask(u_val, denom != 0, (-rhoplusre+overlap*rhocrossre)/(-rhocrossre+overlap*rhoplusre)) + numpy.putmask( + u_val, + denom != 0, + (-rhoplusre + overlap * rhocrossre) / (-rhocrossre + overlap * rhoplusre), + ) coa_phase = numpy.zeros(len(indices), dtype=numpy.float32) return u_val, coa_phase -class MatchedFilterSkyMaxControl(object): +class MatchedFilterSkyMaxControl: # FIXME: This seems much more simplistic than the aligned-spin class. # E.g. no correlators. Is this worth updating? - def __init__(self, low_frequency_cutoff, high_frequency_cutoff, - snr_threshold, tlen, delta_f, dtype): + def __init__( + self, + low_frequency_cutoff, + high_frequency_cutoff, + snr_threshold, + tlen, + delta_f, + dtype, + ): """ Create a matched filter engine. @@ -849,6 +962,7 @@ def __init__(self, low_frequency_cutoff, high_frequency_cutoff, to the nyquist frequency. snr_threshold : float The minimum snr to return when filtering + """ self.tlen = tlen self.delta_f = delta_f @@ -857,8 +971,7 @@ def __init__(self, low_frequency_cutoff, high_frequency_cutoff, self.flow = low_frequency_cutoff self.fhigh = high_frequency_cutoff - self.matched_filter_and_cluster = \ - self.full_matched_filter_and_cluster + self.matched_filter_and_cluster = self.full_matched_filter_and_cluster self.snr_plus_mem = zeros(self.tlen, dtype=self.dtype) self.corr_plus_mem = zeros(self.tlen, dtype=self.dtype) self.snr_cross_mem = zeros(self.tlen, dtype=self.dtype) @@ -869,9 +982,9 @@ def __init__(self, low_frequency_cutoff, high_frequency_cutoff, self.cached_hplus_hcross_hcross = None self.cached_hplus_hcross_psd = None - - def full_matched_filter_and_cluster(self, hplus, hcross, hplus_norm, - hcross_norm, psd, stilde, window): + def full_matched_filter_and_cluster( + self, hplus, hcross, hplus_norm, hcross_norm, psd, stilde, window + ): """ Return the complex snr and normalization. @@ -898,22 +1011,27 @@ def full_matched_filter_and_cluster(self, hplus, hcross, hplus_norm, List of indices of the triggers. snrv : Array The snr values at the trigger locations. - """ - - I_plus, Iplus_corr, Iplus_norm = matched_filter_core(hplus, stilde, - h_norm=hplus_norm, - low_frequency_cutoff=self.flow, - high_frequency_cutoff=self.fhigh, - out=self.snr_plus_mem, - corr_out=self.corr_plus_mem) + """ + I_plus, Iplus_corr, Iplus_norm = matched_filter_core( + hplus, + stilde, + h_norm=hplus_norm, + low_frequency_cutoff=self.flow, + high_frequency_cutoff=self.fhigh, + out=self.snr_plus_mem, + corr_out=self.corr_plus_mem, + ) - I_cross, Icross_corr, Icross_norm = matched_filter_core(hcross, - stilde, h_norm=hcross_norm, - low_frequency_cutoff=self.flow, - high_frequency_cutoff=self.fhigh, - out=self.snr_cross_mem, - corr_out=self.corr_cross_mem) + I_cross, Icross_corr, Icross_norm = matched_filter_core( + hcross, + stilde, + h_norm=hcross_norm, + low_frequency_cutoff=self.flow, + high_frequency_cutoff=self.fhigh, + out=self.snr_cross_mem, + corr_out=self.corr_cross_mem, + ) # The information on the complex side of this overlap is important # we may want to use this in the future. @@ -924,12 +1042,16 @@ def full_matched_filter_and_cluster(self, hplus, hcross, hplus_norm, if not id(psd) == self.cached_hplus_hcross_psd: self.cached_hplus_hcross_correlation = None if self.cached_hplus_hcross_correlation is None: - hplus_cross_corr = overlap_cplx(hplus, hcross, psd=psd, - low_frequency_cutoff=self.flow, - high_frequency_cutoff=self.fhigh, - normalized=False) + hplus_cross_corr = overlap_cplx( + hplus, + hcross, + psd=psd, + low_frequency_cutoff=self.flow, + high_frequency_cutoff=self.fhigh, + normalized=False, + ) hplus_cross_corr = numpy.real(hplus_cross_corr) - hplus_cross_corr = hplus_cross_corr / (hcross_norm*hplus_norm)**0.5 + hplus_cross_corr = hplus_cross_corr / (hcross_norm * hplus_norm) ** 0.5 self.cached_hplus_hcross_correlation = hplus_cross_corr self.cached_hplus_hcross_hplus = id(hplus) self.cached_hplus_hcross_hcross = id(hcross) @@ -937,63 +1059,77 @@ def full_matched_filter_and_cluster(self, hplus, hcross, hplus_norm, else: hplus_cross_corr = self.cached_hplus_hcross_correlation - snr = self._maximized_snr(I_plus,I_cross, - hplus_cross_corr, - hpnorm=Iplus_norm, - hcnorm=Icross_norm, - out=self.snr_mem, - thresh=self.snr_threshold, - analyse_slice=stilde.analyze) + snr = self._maximized_snr( + I_plus, + I_cross, + hplus_cross_corr, + hpnorm=Iplus_norm, + hcnorm=Icross_norm, + out=self.snr_mem, + thresh=self.snr_threshold, + analyse_slice=stilde.analyze, + ) # FIXME: This should live further down # Convert output to pycbc TimeSeries delta_t = 1.0 / (self.tlen * stilde.delta_f) - snr = TimeSeries(snr, epoch=stilde.start_time, delta_t=delta_t, - copy=False) + snr = TimeSeries(snr, epoch=stilde.start_time, delta_t=delta_t, copy=False) - idx, snrv = events.threshold_real_numpy(snr[stilde.analyze], - self.snr_threshold) + idx, snrv = events.threshold_real_numpy(snr[stilde.analyze], self.snr_threshold) if len(idx) == 0: return [], 0, 0, [], [], [], [], 0, 0, 0 logger.info("%d points above threshold", len(idx)) - idx, snrv = events.cluster_reduce(idx, snrv, window) logger.info("%d clustered points", len(idx)) # erased self. - u_vals, coa_phase = self._maximized_extrinsic_params\ - (I_plus.data, I_cross.data, hplus_cross_corr, - indices=idx+stilde.analyze.start, hpnorm=Iplus_norm, - hcnorm=Icross_norm) - - + u_vals, coa_phase = self._maximized_extrinsic_params( + I_plus.data, + I_cross.data, + hplus_cross_corr, + indices=idx + stilde.analyze.start, + hpnorm=Iplus_norm, + hcnorm=Icross_norm, + ) - return snr, Iplus_corr, Icross_corr, idx, snrv, u_vals, coa_phase,\ - hplus_cross_corr, Iplus_norm, Icross_norm + return ( + snr, + Iplus_corr, + Icross_corr, + idx, + snrv, + u_vals, + coa_phase, + hplus_cross_corr, + Iplus_norm, + Icross_norm, + ) def _maximized_snr(self, hplus, hcross, hphccorr, **kwargs): - return compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, - **kwargs) + return compute_max_snr_over_sky_loc_stat(hplus, hcross, hphccorr, **kwargs) def _maximized_extrinsic_params(self, hplus, hcross, hphccorr, **kwargs): - return compute_u_val_for_sky_loc_stat(hplus, hcross, hphccorr, - **kwargs) + return compute_u_val_for_sky_loc_stat(hplus, hcross, hphccorr, **kwargs) class MatchedFilterSkyMaxControlNoPhase(MatchedFilterSkyMaxControl): # Basically the same as normal SkyMaxControl, except we use a slight # variation in the internal SNR functions. def _maximized_snr(self, hplus, hcross, hphccorr, **kwargs): - return compute_max_snr_over_sky_loc_stat_no_phase(hplus, hcross, - hphccorr, **kwargs) + return compute_max_snr_over_sky_loc_stat_no_phase( + hplus, hcross, hphccorr, **kwargs + ) def _maximized_extrinsic_params(self, hplus, hcross, hphccorr, **kwargs): - return compute_u_val_for_sky_loc_stat_no_phase(hplus, hcross, hphccorr, - **kwargs) + return compute_u_val_for_sky_loc_stat_no_phase( + hplus, hcross, hphccorr, **kwargs + ) + def make_frequency_series(vec): - """Return a frequency series of the input vector. + """ + Return a frequency series of the input vector. If the input is a frequency series it is returned, else if the input vector is a real time series it is fourier transformed and returned as a @@ -1007,6 +1143,7 @@ def make_frequency_series(vec): ------- Frequency Series: FrequencySeries A frequency domain version of the input vector. + """ if isinstance(vec, FrequencySeries): return vec @@ -1014,16 +1151,19 @@ def make_frequency_series(vec): N = len(vec) n = N // 2 + 1 delta_f = 1.0 / N / vec.delta_t - vectilde = FrequencySeries(zeros(n, dtype=complex_same_precision_as(vec)), - delta_f=delta_f, copy=False) + vectilde = FrequencySeries( + zeros(n, dtype=complex_same_precision_as(vec)), delta_f=delta_f, copy=False + ) fft(vec, vectilde) return vectilde - else: - raise TypeError("Can only convert a TimeSeries to a FrequencySeries") + raise TypeError("Can only convert a TimeSeries to a FrequencySeries") -def sigmasq_series(htilde, psd=None, low_frequency_cutoff=None, - high_frequency_cutoff=None): - """Return a cumulative sigmasq frequency series. + +def sigmasq_series( + htilde, psd=None, low_frequency_cutoff=None, high_frequency_cutoff=None +): + """ + Return a cumulative sigmasq frequency series. Return a frequency series containing the accumulated power in the input up to that frequency. @@ -1045,15 +1185,20 @@ def sigmasq_series(htilde, psd=None, low_frequency_cutoff=None, ------- Frequency Series: FrequencySeries A frequency series containing the cumulative sigmasq. + """ htilde = make_frequency_series(htilde) - N = (len(htilde)-1) * 2 + N = (len(htilde) - 1) * 2 norm = 4.0 * htilde.delta_f - kmin, kmax = get_cutoff_indices(low_frequency_cutoff, - high_frequency_cutoff, htilde.delta_f, N) + kmin, kmax = get_cutoff_indices( + low_frequency_cutoff, high_frequency_cutoff, htilde.delta_f, N + ) - sigma_vec = FrequencySeries(zeros(len(htilde), dtype=real_same_precision_as(htilde)), - delta_f = htilde.delta_f, copy=False) + sigma_vec = FrequencySeries( + zeros(len(htilde), dtype=real_same_precision_as(htilde)), + delta_f=htilde.delta_f, + copy=False, + ) mag = htilde.squared_norm() @@ -1062,12 +1207,12 @@ def sigmasq_series(htilde, psd=None, low_frequency_cutoff=None, sigma_vec[kmin:kmax] = mag[kmin:kmax].cumsum() - return sigma_vec*norm + return sigma_vec * norm -def sigmasq(htilde, psd = None, low_frequency_cutoff=None, - high_frequency_cutoff=None): - """Return the loudness of the waveform. This is defined (see Duncan +def sigmasq(htilde, psd=None, low_frequency_cutoff=None, high_frequency_cutoff=None): + """ + Return the loudness of the waveform. This is defined (see Duncan Brown's thesis) as the unnormalized matched-filter of the input waveform, htilde, with itself. This quantity is usually referred to as (sigma)^2 and is then used to normalize matched-filters with the data. @@ -1086,19 +1231,21 @@ def sigmasq(htilde, psd = None, low_frequency_cutoff=None, Returns ------- sigmasq: float + """ htilde = make_frequency_series(htilde) - N = (len(htilde)-1) * 2 + N = (len(htilde) - 1) * 2 norm = 4.0 * htilde.delta_f - kmin, kmax = get_cutoff_indices(low_frequency_cutoff, - high_frequency_cutoff, htilde.delta_f, N) + kmin, kmax = get_cutoff_indices( + low_frequency_cutoff, high_frequency_cutoff, htilde.delta_f, N + ) ht = htilde[kmin:kmax] if psd: try: numpy.testing.assert_almost_equal(ht.delta_f, psd.delta_f) except AssertionError: - raise ValueError('Waveform does not have same delta_f as psd') + raise ValueError("Waveform does not have same delta_f as psd") if psd is None: sq = ht.inner(ht) @@ -1107,9 +1254,10 @@ def sigmasq(htilde, psd = None, low_frequency_cutoff=None, return sq.real * norm -def sigma(htilde, psd = None, low_frequency_cutoff=None, - high_frequency_cutoff=None): - """ Return the sigma of the waveform. See sigmasq for more details. + +def sigma(htilde, psd=None, low_frequency_cutoff=None, high_frequency_cutoff=None): + """ + Return the sigma of the waveform. See sigmasq for more details. Parameters ---------- @@ -1125,9 +1273,11 @@ def sigma(htilde, psd = None, low_frequency_cutoff=None, Returns ------- sigmasq: float + """ return sqrt(sigmasq(htilde, psd, low_frequency_cutoff, high_frequency_cutoff)) + def get_cutoff_indices(flow, fhigh, df, N): """ Gets the indices of a frequency series at which to stop an overlap @@ -1149,37 +1299,47 @@ def get_cutoff_indices(flow, fhigh, df, N): ------- kmin: int kmax: int + """ if flow: kmin = int(flow / df) if kmin < 0: err_msg = "Start frequency cannot be negative. " - err_msg += "Supplied value and kmin {} and {}".format(flow, kmin) + err_msg += f"Supplied value and kmin {flow} and {kmin}" raise ValueError(err_msg) else: kmin = 1 if fhigh: kmax = int(fhigh / df) - if kmax > int((N + 1)/2.): - kmax = int((N + 1)/2.) + kmax = min(kmax, int((N + 1) / 2.0)) else: # int() truncates towards 0, so this is # equivalent to the floor of the float - kmax = int((N + 1)/2.) + kmax = int((N + 1) / 2.0) if kmax <= kmin: err_msg = "Kmax cannot be less than or equal to kmin. " err_msg += "Provided values of freqencies (min,max) were " - err_msg += "{} and {} ".format(flow, fhigh) + err_msg += f"{flow} and {fhigh} " err_msg += "corresponding to (kmin, kmax) of " - err_msg += "{} and {}.".format(kmin, kmax) + err_msg += f"{kmin} and {kmax}." raise ValueError(err_msg) - return kmin,kmax + return kmin, kmax + -def matched_filter_core(template, data, psd=None, low_frequency_cutoff=None, - high_frequency_cutoff=None, h_norm=None, out=None, corr_out=None): - """ Return the complex snr and normalization. +def matched_filter_core( + template, + data, + psd=None, + low_frequency_cutoff=None, + high_frequency_cutoff=None, + h_norm=None, + out=None, + corr_out=None, +): + """ + Return the complex snr and normalization. Return the complex snr, along with its associated normalization of the template, matched filtered against the data. @@ -1216,6 +1376,7 @@ def matched_filter_core(template, data, psd=None, low_frequency_cutoff=None, A frequency series containing the correlation vector. norm : float The normalization of the complex snr. + """ htilde = make_frequency_series(template) stilde = make_frequency_series(data) @@ -1223,9 +1384,10 @@ def matched_filter_core(template, data, psd=None, low_frequency_cutoff=None, if len(htilde) != len(stilde): raise ValueError("Length of template and data must match") - N = (len(stilde)-1) * 2 - kmin, kmax = get_cutoff_indices(low_frequency_cutoff, - high_frequency_cutoff, stilde.delta_f, N) + N = (len(stilde) - 1) * 2 + kmin, kmax = get_cutoff_indices( + low_frequency_cutoff, high_frequency_cutoff, stilde.delta_f, N + ) if corr_out is not None: qtilde = corr_out @@ -1234,10 +1396,10 @@ def matched_filter_core(template, data, psd=None, low_frequency_cutoff=None, if out is None: _q = zeros(N, dtype=complex_same_precision_as(data)) - elif (len(out) == N) and type(out) is Array and out.kind =='complex': + elif (len(out) == N) and type(out) is Array and out.kind == "complex": _q = out else: - raise TypeError('Invalid Output Vector: wrong length or dtype') + raise TypeError("Invalid Output Vector: wrong length or dtype") correlate(htilde[kmin:kmax], stilde[kmin:kmax], qtilde[kmin:kmax]) @@ -1256,11 +1418,16 @@ def matched_filter_core(template, data, psd=None, low_frequency_cutoff=None, if h_norm is None: h_norm = sigmasq(htilde, psd, low_frequency_cutoff, high_frequency_cutoff) - norm = (4.0 * stilde.delta_f) / sqrt( h_norm) + norm = (4.0 * stilde.delta_f) / sqrt(h_norm) + + return ( + TimeSeries(_q, epoch=stilde._epoch, delta_t=stilde.delta_t, copy=False), + FrequencySeries( + qtilde, epoch=stilde._epoch, delta_f=stilde.delta_f, copy=False + ), + norm, + ) - return (TimeSeries(_q, epoch=stilde._epoch, delta_t=stilde.delta_t, copy=False), - FrequencySeries(qtilde, epoch=stilde._epoch, delta_f=stilde.delta_f, copy=False), - norm) def smear(idx, factor): """ @@ -1270,28 +1437,35 @@ def smear(idx, factor): E.g.: smear([5,7,100],2) = [3,4,5,6,7,8,9,98,99,100,101,102] Parameters - ----------- + ---------- idx : numpy.array of ints The indexes to be smeared. factor : idx The factor by which to smear out the input array. Returns - -------- + ------- new_idx : numpy.array of ints The smeared array of indexes. - """ - + """ s = [idx] - for i in range(factor+1): - a = i - factor/2 + for i in range(factor + 1): + a = i - factor / 2 s += [idx + a] return numpy.unique(numpy.concatenate(s)) -def matched_filter(template, data, psd=None, low_frequency_cutoff=None, - high_frequency_cutoff=None, sigmasq=None): - """ Return the complex snr. + +def matched_filter( + template, + data, + psd=None, + low_frequency_cutoff=None, + high_frequency_cutoff=None, + sigmasq=None, +): + """ + Return the complex snr. Return the complex snr, along with its associated normalization of the template, matched filtered against the data. @@ -1318,13 +1492,22 @@ def matched_filter(template, data, psd=None, low_frequency_cutoff=None, ------- snr : TimeSeries A time series containing the complex snr. + """ - snr, _, norm = matched_filter_core(template, data, psd=psd, - low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff, h_norm=sigmasq) + snr, _, norm = matched_filter_core( + template, + data, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + h_norm=sigmasq, + ) return snr * norm + _snr = None + + def match( vec1, vec2, @@ -1336,7 +1519,8 @@ def match( subsample_interpolation=False, return_phase=False, ): - """Return the match between the two TimeSeries or FrequencySeries. + """ + Return the match between the two TimeSeries or FrequencySeries. Return the match between two waveforms. This is equivalent to the overlap maximized over time and phase. @@ -1380,8 +1564,8 @@ def match( The number of samples to shift to get the match. phi: float Phase to rotate complex waveform to get the match, if desired. - """ + """ htilde = make_frequency_series(vec1) stilde = make_frequency_series(vec2) @@ -1419,12 +1603,19 @@ def match( rounded_max_id = int(round(max_id)) phi = numpy.angle(snr[rounded_max_id]) return maxsnr * snr_norm / sqrt(v2_norm), max_id, phi - else: - return maxsnr * snr_norm / sqrt(v2_norm), max_id + return maxsnr * snr_norm / sqrt(v2_norm), max_id -def overlap(vec1, vec2, psd=None, low_frequency_cutoff=None, - high_frequency_cutoff=None, normalized=True): - """ Return the overlap between the two TimeSeries or FrequencySeries. + +def overlap( + vec1, + vec2, + psd=None, + low_frequency_cutoff=None, + high_frequency_cutoff=None, + normalized=True, +): + """ + Return the overlap between the two TimeSeries or FrequencySeries. Parameters ---------- @@ -1444,16 +1635,28 @@ def overlap(vec1, vec2, psd=None, low_frequency_cutoff=None, Returns ------- overlap: float + """ + return overlap_cplx( + vec1, + vec2, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + normalized=normalized, + ).real - return overlap_cplx(vec1, vec2, psd=psd, \ - low_frequency_cutoff=low_frequency_cutoff,\ - high_frequency_cutoff=high_frequency_cutoff,\ - normalized=normalized).real -def overlap_cplx(vec1, vec2, psd=None, low_frequency_cutoff=None, - high_frequency_cutoff=None, normalized=True): - """Return the complex overlap between the two TimeSeries or FrequencySeries. +def overlap_cplx( + vec1, + vec2, + psd=None, + low_frequency_cutoff=None, + high_frequency_cutoff=None, + normalized=True, +): + """ + Return the complex overlap between the two TimeSeries or FrequencySeries. Parameters ---------- @@ -1473,12 +1676,17 @@ def overlap_cplx(vec1, vec2, psd=None, low_frequency_cutoff=None, Returns ------- overlap: complex + """ htilde = make_frequency_series(vec1) stilde = make_frequency_series(vec2) - kmin, kmax = get_cutoff_indices(low_frequency_cutoff, - high_frequency_cutoff, stilde.delta_f, (len(stilde)-1) * 2) + kmin, kmax = get_cutoff_indices( + low_frequency_cutoff, + high_frequency_cutoff, + stilde.delta_f, + (len(stilde) - 1) * 2, + ) if psd: inner = (htilde[kmin:kmax]).weighted_inner(stilde[kmin:kmax], psd[kmin:kmax]) @@ -1486,18 +1694,28 @@ def overlap_cplx(vec1, vec2, psd=None, low_frequency_cutoff=None, inner = (htilde[kmin:kmax]).inner(stilde[kmin:kmax]) if normalized: - sig1 = sigma(vec1, psd=psd, low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff) - sig2 = sigma(vec2, psd=psd, low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff) + sig1 = sigma( + vec1, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + ) + sig2 = sigma( + vec2, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + ) norm = 1 / sig1 / sig2 else: norm = 1 return 4 * htilde.delta_f * inner * norm + def quadratic_interpolate_peak(left, middle, right): - """ Interpolate the peak and offset using a quadratic approximation + """ + Interpolate the peak and offset using a quadratic approximation Parameters ---------- @@ -1514,22 +1732,29 @@ def quadratic_interpolate_peak(left, middle, right): Array of bins offsets, each in the range [-1/2, 1/2] peak_values : numpy array Array of the estimated peak values at the interpolated offset + """ - bin_offset = 1.0/2.0 * (left - right) / (left - 2 * middle + right) + bin_offset = 1.0 / 2.0 * (left - right) / (left - 2 * middle + right) peak_value = middle - 0.25 * (left - right) * bin_offset return bin_offset, peak_value -class LiveBatchMatchedFilter(object): - +class LiveBatchMatchedFilter: """Calculate SNR and signal consistency tests in a batched progression""" - def __init__(self, templates, snr_threshold, chisq_bins, sg_chisq, - maxelements=2**27, - snr_abort_threshold=None, - newsnr_threshold=None, - max_triggers_in_batch=None): - """Create a batched matchedfilter instance + def __init__( + self, + templates, + snr_threshold, + chisq_bins, + sg_chisq, + maxelements=2**27, + snr_abort_threshold=None, + newsnr_threshold=None, + max_triggers_in_batch=None, + ): + """ + Create a batched matchedfilter instance Parameters ---------- @@ -1553,6 +1778,7 @@ def __init__(self, templates, snr_threshold, chisq_bins, sg_chisq, Record X number of the loudest triggers by SNR in each MPI process. Signal consistency values will also only be calculated for these triggers. + """ self.snr_threshold = snr_threshold self.snr_abort_threshold = snr_abort_threshold @@ -1560,6 +1786,7 @@ def __init__(self, templates, snr_threshold, chisq_bins, sg_chisq, self.max_triggers_in_batch = max_triggers_in_batch from pycbc import vetoes + self.power_chisq = vetoes.SingleDetPowerChisq(chisq_bins, None) self.sg_chisq = sg_chisq @@ -1601,9 +1828,12 @@ def __init__(self, templates, snr_threshold, chisq_bins, sg_chisq, dur, count = i self.out_mem[i] = zeros(size, dtype=numpy.complex64) self.cout_mem[i] = zeros(size, dtype=numpy.complex64) - self.ifts[i] = IFFT(self.cout_mem[i], self.out_mem[i], - nbatch=count, - size=len(self.cout_mem[i]) // count) + self.ifts[i] = IFFT( + self.cout_mem[i], + self.out_mem[i], + nbatch=count, + size=len(self.cout_mem[i]) // count, + ) # Split the templates into their processing groups for dur, count in mem_ids: @@ -1624,7 +1854,9 @@ def __init__(self, templates, snr_threshold, chisq_bins, sg_chisq, htilde.cout = self.cout_mem[mid][s:e] s += psize e += psize - self.corr.append(BatchCorrelator(tgroup, [t.cout for t in tgroup], len(tgroup[0]))) + self.corr.append( + BatchCorrelator(tgroup, [t.cout for t in tgroup], len(tgroup[0])) + ) def set_data(self, data): """Set the data reader object to use""" @@ -1649,15 +1881,17 @@ def process_all(self): veto_info = [] while 1: result, veto = self._process_batch() - if result is False: return False - if result is None: break + if result is False: + return False + if result is None: + break results.append(result) veto_info += veto result = self.combine_results(results) if self.max_triggers_in_batch: - sort = result['snr'].argsort()[::-1][:self.max_triggers_in_batch] + sort = result["snr"].argsort()[::-1][: self.max_triggers_in_batch] for key in result: result[key] = result[key][sort] @@ -1671,27 +1905,28 @@ def _process_vetoes(self, results, veto_info): """Calculate signal based vetoes""" chisq = numpy.array(numpy.zeros(len(veto_info)), numpy.float32, ndmin=1) dof = numpy.array(numpy.zeros(len(veto_info)), numpy.uint32, ndmin=1) - sg_chisq = numpy.array(numpy.zeros(len(veto_info)), numpy.float32, - ndmin=1) - results['chisq'] = chisq - results['chisq_dof'] = dof - results['sg_chisq'] = sg_chisq + sg_chisq = numpy.array(numpy.zeros(len(veto_info)), numpy.float32, ndmin=1) + results["chisq"] = chisq + results["chisq_dof"] = dof + results["sg_chisq"] = sg_chisq keep = [] for i, (snrv, norm, l, htilde, stilde) in enumerate(veto_info): correlate(htilde, stilde, htilde.cout) - c, d = self.power_chisq.values(htilde.cout, snrv, - norm, stilde.psd, [l], htilde) + c, d = self.power_chisq.values( + htilde.cout, snrv, norm, stilde.psd, [l], htilde + ) chisq[i] = c[0] / d[0] dof[i] = d[0] - sgv = self.sg_chisq.values(stilde, htilde, stilde.psd, - snrv, norm, c, d, [l]) + sgv = self.sg_chisq.values( + stilde, htilde, stilde.psd, snrv, norm, c, d, [l] + ) if sgv is not None: sg_chisq[i] = sgv[0] if self.newsnr_threshold: - newsnr = ranking.newsnr(results['snr'][i], chisq[i]) + newsnr = ranking.newsnr(results["snr"][i], chisq[i]) if newsnr >= self.newsnr_threshold: keep.append(i) @@ -1740,14 +1975,14 @@ def _process_batch(self): # Find the peaks in our SNR times series from the various templates i = 0 for htilde in tgroup: - if hasattr(htilde, 'time_offset'): - if 'time_offset' not in result: - result['time_offset'] = [] + if hasattr(htilde, "time_offset"): + if "time_offset" not in result: + result["time_offset"] = [] l = htilde.out[seg].abs_arg_max() sgm = htilde.sigmasq(psd) - norm = 4.0 * htilde.delta_f / (sgm ** 0.5) + norm = 4.0 * htilde.delta_f / (sgm**0.5) l += valid_start snrv = numpy.array([htilde.out[l]]) @@ -1762,8 +1997,10 @@ def _process_batch(self): # We have an SNR so high that we will drop the entire analysis # of this chunk of time! if self.snr_abort_threshold is not None and s > self.snr_abort_threshold: - logger.info("We are seeing some *really* high SNRs, let's " - "assume they aren't signals and just give up") + logger.info( + "We are seeing some *really* high SNRs, let's " + "assume they aren't signals and just give up" + ) return False, [] veto_info.append((snrv, norm, l, htilde, stilde)) @@ -1771,7 +2008,7 @@ def _process_batch(self): snr[i] = snrv[0] * norm sigmasq[i] = sgm templates[i] = htilde.id - if not hasattr(htilde, 'dict_params'): + if not hasattr(htilde, "dict_params"): htilde.dict_params = {} for key in tkeys: htilde.dict_params[key] = htilde.params[key] @@ -1779,30 +2016,38 @@ def _process_batch(self): for key in tkeys: result[key].append(htilde.dict_params[key]) - if hasattr(htilde, 'time_offset'): - result['time_offset'].append(htilde.time_offset) + if hasattr(htilde, "time_offset"): + result["time_offset"].append(htilde.time_offset) i += 1 - result['snr'] = abs(snr[0:i]) - result['coa_phase'] = numpy.angle(snr[0:i]) - result['end_time'] = time[0:i] - result['template_id'] = templates[0:i] - result['sigmasq'] = sigmasq[0:i] + result["snr"] = abs(snr[0:i]) + result["coa_phase"] = numpy.angle(snr[0:i]) + result["end_time"] = time[0:i] + result["template_id"] = templates[0:i] + result["sigmasq"] = sigmasq[0:i] for key in tkeys: result[key] = numpy.array(result[key]) - if 'time_offset' in result: - result['time_offset'] = numpy.array(result['time_offset']) + if "time_offset" in result: + result["time_offset"] = numpy.array(result["time_offset"]) return result, veto_info -def followup_event_significance(ifo, data_reader, bank, - template_id, coinc_times, - coinc_threshold=0.005, - lookback=150, duration=0.095): - """Given a detector, a template waveform and a set of candidate event + +def followup_event_significance( + ifo, + data_reader, + bank, + template_id, + coinc_times, + coinc_threshold=0.005, + lookback=150, + duration=0.095, +): + """ + Given a detector, a template waveform and a set of candidate event times in different detectors, perform an on-source/off-source analysis to determine if the SNR in the first detector has a significant peak in the on-source window. The significance is given in terms of a @@ -1851,14 +2096,16 @@ def followup_event_significance(ifo, data_reader, bank, followup_info: dict or None Results of the followup calculation (see above) or None if `ifo` did not have usable data. + """ from pycbc.waveform import get_waveform_filter_length_in_time + tmplt = bank.table[template_id] - length_in_time = get_waveform_filter_length_in_time(tmplt['approximant'], - tmplt) + length_in_time = get_waveform_filter_length_in_time(tmplt["approximant"], tmplt) # calculate onsource time range from pycbc.detector import Detector + onsource_start = -numpy.inf onsource_end = numpy.inf fdet = Detector(ifo) @@ -1866,10 +2113,8 @@ def followup_event_significance(ifo, data_reader, bank, for cifo in coinc_times: time = coinc_times[cifo] dtravel = Detector(cifo).light_travel_time_to_detector(fdet) - if time - dtravel > onsource_start: - onsource_start = time - dtravel - if time + dtravel < onsource_end: - onsource_end = time + dtravel + onsource_start = max(onsource_start, time - dtravel) + onsource_end = min(onsource_end, time + dtravel) # Source must be within this time window to be considered a possible # coincidence @@ -1887,13 +2132,11 @@ def followup_event_significance(ifo, data_reader, bank, trim_pad = data_reader.trim_padding * data_reader.strain.delta_t buffer_duration = lookback + 2 * trim_pad + length_in_time buffer_samples = bank.round_up(int(buffer_duration * bank.sample_rate)) - max_safe_buffer_samples = int( - 0.9 * data_reader.strain.duration * bank.sample_rate - ) + max_safe_buffer_samples = int(0.9 * data_reader.strain.duration * bank.sample_rate) if buffer_samples > max_safe_buffer_samples: buffer_samples = max_safe_buffer_samples - new_lookback = ( - buffer_samples / bank.sample_rate - (2 * trim_pad + length_in_time) + new_lookback = buffer_samples / bank.sample_rate - ( + 2 * trim_pad + length_in_time ) # Require a minimum lookback time of twice the onsource window or SNR # time series (whichever is longer) so we have enough data for the @@ -1902,18 +2145,18 @@ def followup_event_significance(ifo, data_reader, bank, min_required_lookback = 2 * max(onsource_end - onsource_start, duration) if new_lookback > min_required_lookback: logging.warning( - 'Strain buffer too short for a lookback time of %f s, ' - 'reducing lookback to %f s', + "Strain buffer too short for a lookback time of %f s, " + "reducing lookback to %f s", lookback, - new_lookback + new_lookback, ) else: logging.error( - 'Strain buffer too short to compute the followup SNR time ' - 'series for template %d, will not use %s for followup. ' - 'Either use shorter templates, or raise --max-length.', + "Strain buffer too short to compute the followup SNR time " + "series for template %d, will not use %s for followup. " + "Either use shorter templates, or raise --max-length.", template_id, - ifo + ifo, ) return None buffer_duration = buffer_samples / bank.sample_rate @@ -1925,13 +2168,11 @@ def followup_event_significance(ifo, data_reader, bank, - data_reader.reduced_pad * data_reader.strain.delta_t - buffer_duration ) - if not data_reader.state.is_extent_valid( - state_start_time, buffer_duration - ): + if not data_reader.state.is_extent_valid(state_start_time, buffer_duration): logging.info( - '%s strain buffer contains invalid data during lookback, ' - 'will not use for followup', - ifo + "%s strain buffer contains invalid data during lookback, " + "will not use for followup", + ifo, ) return None @@ -1942,9 +2183,9 @@ def followup_event_significance(ifo, data_reader, bank, dq_duration = onsource_end - onsource_start + duration if not data_reader.dq.is_extent_valid(dq_start_time, dq_duration): logging.info( - '%s DQ buffer indicates invalid data during onsource window, ' - 'will not use for followup', - ifo + "%s DQ buffer indicates invalid data during onsource window, " + "will not use for followup", + ifo, ) return None @@ -1969,7 +2210,7 @@ def followup_event_significance(ifo, data_reader, bank, window = int((onsource_end - onsource_start) * snr.sample_rate) nsamples = int(len(bkg) / window) - peaks = bkg[:nsamples*window].reshape(nsamples, window).max(axis=1) + peaks = bkg[: nsamples * window].reshape(nsamples, window).max(axis=1) num_louder_bg = (peaks >= peak_value).sum() pvalue = (1 + num_louder_bg) / float(1 + nsamples) pvalue_saturated = num_louder_bg == 0 @@ -1977,25 +2218,25 @@ def followup_event_significance(ifo, data_reader, bank, # Return recentered source SNR for bayestar, along with p-value, and trig peak_full = int((peak_time - snr.start_time) / snr.delta_t) half_dur_samples = int(snr.sample_rate * duration / 2) - snr_slice = slice(peak_full - half_dur_samples, - peak_full + half_dur_samples + 1) + snr_slice = slice(peak_full - half_dur_samples, peak_full + half_dur_samples + 1) baysnr = snr[snr_slice] - logger.info('Adding %s to candidate, pvalue %s, %s samples', ifo, - pvalue, nsamples) + logger.info("Adding %s to candidate, pvalue %s, %s samples", ifo, pvalue, nsamples) return { - 'snr_series': baysnr * norm, - 'peak_time': peak_time, - 'pvalue': pvalue, - 'pvalue_saturated': pvalue_saturated, - 'sigma2': sigma2 + "snr_series": baysnr * norm, + "peak_time": peak_time, + "pvalue": pvalue, + "pvalue_saturated": pvalue_saturated, + "sigma2": sigma2, } -def compute_followup_snr_series(data_reader, htilde, trig_time, - duration=0.095, check_state=True, - coinc_window=0.05): - """Given a StrainBuffer, a template frequency series and a trigger time, + +def compute_followup_snr_series( + data_reader, htilde, trig_time, duration=0.095, check_state=True, coinc_window=0.05 +): + """ + Given a StrainBuffer, a template frequency series and a trigger time, compute a portion of the SNR time series centered on the trigger for its rapid sky localization and followup. @@ -2032,6 +2273,7 @@ def compute_followup_snr_series(data_reader, htilde, trig_time, snr : TimeSeries The portion of SNR around the trigger. None if the detector is offline or has bad data quality, and check_state is True. + """ if check_state: # was the detector observing for the full amount of involved data? @@ -2039,8 +2281,7 @@ def compute_followup_snr_series(data_reader, htilde, trig_time, state_end_time = trig_time + duration / 2 state_duration = state_end_time - state_start_time if data_reader.state is not None: - if not data_reader.state.is_extent_valid(state_start_time, - state_duration): + if not data_reader.state.is_extent_valid(state_start_time, state_duration): return None # was the data quality ok for the full amount of involved data? @@ -2051,8 +2292,9 @@ def compute_followup_snr_series(data_reader, htilde, trig_time, return None stilde = data_reader.overwhitened_data(htilde.delta_f) - snr, _, norm = matched_filter_core(htilde, stilde, - h_norm=htilde.sigmasq(stilde.psd)) + snr, _, norm = matched_filter_core( + htilde, stilde, h_norm=htilde.sigmasq(stilde.psd) + ) valid_end = int(len(snr) - data_reader.trim_padding) valid_start = int(valid_end - data_reader.blocksize * snr.sample_rate) @@ -2061,17 +2303,18 @@ def compute_followup_snr_series(data_reader, htilde, trig_time, coinc_samples = int(snr.sample_rate * coinc_window) valid_start -= half_dur_samples + coinc_samples valid_end += half_dur_samples - if valid_start < 0 or valid_end > len(snr)-1: - raise ValueError(('Requested SNR duration ({0} s)' - ' too long').format(duration)) + if valid_start < 0 or valid_end > len(snr) - 1: + raise ValueError(f"Requested SNR duration ({duration} s) too long") # Onsource slice for Bayestar followup onsource_idx = float(trig_time - snr.start_time) * snr.sample_rate onsource_idx = int(round(onsource_idx)) - onsource_slice = slice(onsource_idx - half_dur_samples, - onsource_idx + half_dur_samples + 1) + onsource_slice = slice( + onsource_idx - half_dur_samples, onsource_idx + half_dur_samples + 1 + ) return snr[onsource_slice] * norm + def optimized_match( vec1, vec2, @@ -2082,7 +2325,8 @@ def optimized_match( v2_norm=None, return_phase=False, ): - """Given two waveforms (as numpy arrays), + """ + Given two waveforms (as numpy arrays), compute the optimized match between them, making use of scipy.minimize_scalar. @@ -2117,8 +2361,8 @@ def optimized_match( The number of samples to shift to get the match. phi: float Phase to rotate complex waveform to get the match, if desired. - """ + """ from scipy.optimize import minimize_scalar htilde = make_frequency_series(vec1) @@ -2188,27 +2432,35 @@ def to_minimize(dt): norm = numpy.sqrt(norm_1 * norm_2) - res = minimize_scalar( - to_minimize, - method="brent", - bracket=(-delta_t, delta_t) - ) + res = minimize_scalar(to_minimize, method="brent", bracket=(-delta_t, delta_t)) m, angle = product_offset(res.x) if return_phase: return m / norm, res.x / delta_t + max_id, -angle - else: - return m / norm, res.x / delta_t + max_id - - -__all__ = ['match', 'optimized_match', 'matched_filter', 'sigmasq', 'sigma', 'get_cutoff_indices', - 'sigmasq_series', 'make_frequency_series', 'overlap', - 'overlap_cplx', 'matched_filter_core', 'correlate', - 'MatchedFilterControl', 'LiveBatchMatchedFilter', - 'MatchedFilterSkyMaxControl', 'MatchedFilterSkyMaxControlNoPhase', - 'compute_max_snr_over_sky_loc_stat_no_phase', - 'compute_max_snr_over_sky_loc_stat', - 'compute_followup_snr_series', - 'compute_u_val_for_sky_loc_stat_no_phase', - 'compute_u_val_for_sky_loc_stat', - 'followup_event_significance'] + return m / norm, res.x / delta_t + max_id + + +__all__ = [ + "LiveBatchMatchedFilter", + "MatchedFilterControl", + "MatchedFilterSkyMaxControl", + "MatchedFilterSkyMaxControlNoPhase", + "compute_followup_snr_series", + "compute_max_snr_over_sky_loc_stat", + "compute_max_snr_over_sky_loc_stat_no_phase", + "compute_u_val_for_sky_loc_stat", + "compute_u_val_for_sky_loc_stat_no_phase", + "correlate", + "followup_event_significance", + "get_cutoff_indices", + "make_frequency_series", + "match", + "matched_filter", + "matched_filter_core", + "optimized_match", + "overlap", + "overlap_cplx", + "sigma", + "sigmasq", + "sigmasq_series", +] diff --git a/pycbc/filter/matchedfilter_cuda.py b/pycbc/filter/matchedfilter_cuda.py index 7e0b2ff7ab6..ba5becf73f4 100644 --- a/pycbc/filter/matchedfilter_cuda.py +++ b/pycbc/filter/matchedfilter_cuda.py @@ -22,27 +22,32 @@ # ============================================================================= # from pycuda.elementwise import ElementwiseKernel -from pycuda.tools import context_dependent_memoize -from pycuda.tools import dtype_to_ctype from pycuda.gpuarray import _get_common_dtype +from pycuda.tools import context_dependent_memoize, dtype_to_ctype + from .matchedfilter import _BaseCorrelator + @context_dependent_memoize -def get_correlate_kernel(dtype_x, dtype_y,dtype_out): +def get_correlate_kernel(dtype_x, dtype_y, dtype_out): return ElementwiseKernel( - "%(tp_x)s *x, %(tp_y)s *y, %(tp_z)s *z" % { - "tp_x": dtype_to_ctype(dtype_x), - "tp_y": dtype_to_ctype(dtype_y), - "tp_z": dtype_to_ctype(dtype_out), - }, - "z[i] = conj(x[i]) * y[i]", - "correlate") + "%(tp_x)s *x, %(tp_y)s *y, %(tp_z)s *z" + % { + "tp_x": dtype_to_ctype(dtype_x), + "tp_y": dtype_to_ctype(dtype_y), + "tp_z": dtype_to_ctype(dtype_out), + }, + "z[i] = conj(x[i]) * y[i]", + "correlate", + ) + def correlate(a, b, out, stream=None): - dtype_out = _get_common_dtype(a,b) + dtype_out = _get_common_dtype(a, b) krnl = get_correlate_kernel(a.dtype, b.dtype, dtype_out) krnl(a.data, b.data, out.data) + class CUDACorrelator(_BaseCorrelator): def __init__(self, x, y, z): self.x = x.data @@ -54,7 +59,6 @@ def __init__(self, x, y, z): def correlate(self): self.krnl(self.x, self.y, self.z) + def _correlate_factory(x, y, z): return CUDACorrelator - - diff --git a/pycbc/filter/matchedfilter_cupy.py b/pycbc/filter/matchedfilter_cupy.py index 1a5a01a65fa..eba90733dfe 100644 --- a/pycbc/filter/matchedfilter_cupy.py +++ b/pycbc/filter/matchedfilter_cupy.py @@ -23,6 +23,7 @@ # import cupy as cp + from .matchedfilter import _BaseCorrelator # Here X,Y,Z are "type placeholder"s, so this covers 32 and 64 bit inputs. @@ -30,15 +31,14 @@ # I've also made this work for mixed types (32 bit -> 64 bit), but this means # we need to always supply the output, which we do. correlate_kernel = cp.ElementwiseKernel( - "X x, Y y", - "Z z", - "z = conj(x) * y", - "correlate_kernel" + "X x, Y y", "Z z", "z = conj(x) * y", "correlate_kernel" ) + def correlate(a, b, out): correlate_kernel(a.data, b.data, out.data) + class CUPYCorrelator(_BaseCorrelator): def __init__(self, x, y, z): self.x = x.data @@ -48,7 +48,6 @@ def __init__(self, x, y, z): def correlate(self): correlate_kernel(self.x, self.y, self.z) + def _correlate_factory(x, y, z): return CUPYCorrelator - - diff --git a/pycbc/filter/matchedfilter_numpy.py b/pycbc/filter/matchedfilter_numpy.py index 721d9dc6a47..1dc7ff48394 100644 --- a/pycbc/filter/matchedfilter_numpy.py +++ b/pycbc/filter/matchedfilter_numpy.py @@ -15,6 +15,7 @@ import numpy + def correlate(x, y, z): z.data[:] = numpy.conjugate(x.data)[:] z *= y diff --git a/pycbc/filter/qtransform.py b/pycbc/filter/qtransform.py index fb06ce17b89..8a668f6261d 100644 --- a/pycbc/filter/qtransform.py +++ b/pycbc/filter/qtransform.py @@ -30,13 +30,16 @@ """ import numpy -from numpy import ceil, log, exp -from pycbc.types.timeseries import FrequencySeries, TimeSeries +from numpy import ceil, exp, log + from pycbc.fft import ifft from pycbc.types import zeros +from pycbc.types.timeseries import FrequencySeries, TimeSeries + def qplane(qplane_tile_dict, fseries, return_complex=False): - """Performs q-transform on each tile for each q-plane and selects + """ + Performs q-transform on each tile for each q-plane and selects tile with the maximum energy. Q-transform can then be interpolated to a desired frequency and time resolution. @@ -59,6 +62,7 @@ def qplane(qplane_tile_dict, fseries, return_complex=False): The frequencies that the qtransform is samled. qplane : numpy.ndarray (2d) The two dimensional interpolated qtransform of this time series. + """ # store q-transforms for each q in a dict qplanes = {} @@ -83,8 +87,10 @@ def qplane(qplane_tile_dict, fseries, return_complex=False): plane = numpy.array([v.numpy() for v in plane]) return max_key, times, frequencies, numpy.array(plane) + def qtiling(fseries, qrange, frange, mismatch=0.2): - """Iterable constructor of QTile tuples + """ + Iterable constructor of QTile tuples Parameters ---------- @@ -101,6 +107,7 @@ def qtiling(fseries, qrange, frange, mismatch=0.2): ------- qplane_tile_dict: 'dict' dictionary containing Q-tile tuples for a set of Q-planes + """ qplane_tile_dict = {} qs = list(_iter_qs(qrange, deltam_f(mismatch))) @@ -110,8 +117,10 @@ def qtiling(fseries, qrange, frange, mismatch=0.2): return qplane_tile_dict + def deltam_f(mismatch): - """Fractional mismatch between neighbouring tiles + """ + Fractional mismatch between neighbouring tiles Parameters ---------- @@ -121,11 +130,14 @@ def deltam_f(mismatch): Returns ------- :type: 'float' + """ - return 2 * (mismatch / 3.) ** (1/2.) + return 2 * (mismatch / 3.0) ** (1 / 2.0) + def _iter_qs(qrange, deltam): - """Iterate over the Q values + """ + Iterate over the Q values Parameters ---------- @@ -138,18 +150,19 @@ def _iter_qs(qrange, deltam): ------- Q-value: Q value for Q-tile - """ + """ # work out how many Qs we need - cumum = log(float(qrange[1]) / qrange[0]) / 2**(1/2.) + cumum = log(float(qrange[1]) / qrange[0]) / 2 ** (1 / 2.0) nplanes = int(max(ceil(cumum / deltam), 1)) dq = cumum / nplanes for i in range(nplanes): - yield qrange[0] * exp(2**(1/2.) * dq * (i + .5)) - return + yield qrange[0] * exp(2 ** (1 / 2.0) * dq * (i + 0.5)) + def _iter_frequencies(q, frange, mismatch, dur): - """Iterate over the frequencies of this 'QPlane' + """ + Iterate over the frequencies of this 'QPlane' Parameters ---------- @@ -166,22 +179,27 @@ def _iter_frequencies(q, frange, mismatch, dur): ------- frequencies: Q-Tile frequency + """ # work out how many frequencies we need minf, maxf = frange - fcum_mismatch = log(float(maxf) / minf) * (2 + q**2)**(1/2.) / 2. + fcum_mismatch = log(float(maxf) / minf) * (2 + q**2) ** (1 / 2.0) / 2.0 nfreq = int(max(1, ceil(fcum_mismatch / deltam_f(mismatch)))) fstep = fcum_mismatch / nfreq - fstepmin = 1. / dur + fstepmin = 1.0 / dur # for each frequency, yield a QTile for i in range(nfreq): - yield (float(minf) * - exp(2 / (2 + q**2)**(1/2.) * (i + .5) * fstep) // - fstepmin * fstepmin) - return + yield ( + float(minf) + * exp(2 / (2 + q**2) ** (1 / 2.0) * (i + 0.5) * fstep) + // fstepmin + * fstepmin + ) + def qseries(fseries, Q, f0, return_complex=False): - """Calculate the energy 'TimeSeries' for the given fseries + """ + Calculate the energy 'TimeSeries' for the given fseries Parameters ---------- @@ -199,33 +217,33 @@ def qseries(fseries, Q, f0, return_complex=False): energy: '~pycbc.types.TimeSeries' A 'TimeSeries' of the normalized energy from the Q-transform of this tile against the data. + """ # normalize and generate bi-square window - qprime = Q / 11**(1/2.) - norm = numpy.sqrt(315. * qprime / (128. * f0)) + qprime = Q / 11 ** (1 / 2.0) + norm = numpy.sqrt(315.0 * qprime / (128.0 * f0)) window_size = 2 * int(f0 / qprime * fseries.duration) + 1 - xfrequencies = numpy.linspace(-1., 1., window_size) + xfrequencies = numpy.linspace(-1.0, 1.0, window_size) start = int((f0 - (f0 / qprime)) * fseries.duration) end = int(start + window_size) center = (start + end) // 2 - windowed = fseries[start:end] * (1 - xfrequencies ** 2) ** 2 * norm + windowed = fseries[start:end] * (1 - xfrequencies**2) ** 2 * norm - tlen = (len(fseries)-1) * 2 + tlen = (len(fseries) - 1) * 2 windowed.resize(tlen) windowed.roll(-center) # calculate the time series for this q -value - windowed = FrequencySeries(windowed, delta_f=fseries.delta_f, - epoch=fseries.start_time) - ctseries = TimeSeries(zeros(tlen, dtype=numpy.complex128), - delta_t=fseries.delta_t) + windowed = FrequencySeries( + windowed, delta_f=fseries.delta_f, epoch=fseries.start_time + ) + ctseries = TimeSeries(zeros(tlen, dtype=numpy.complex128), delta_t=fseries.delta_t) ifft(windowed, ctseries) if return_complex: return ctseries - else: - energy = ctseries.squared_norm() - medianenergy = numpy.median(energy.numpy()) - return energy / float(medianenergy) + energy = ctseries.squared_norm() + medianenergy = numpy.median(energy.numpy()) + return energy / float(medianenergy) diff --git a/pycbc/filter/resample.py b/pycbc/filter/resample.py index eaa0a17ff51..48bce87b8d7 100644 --- a/pycbc/filter/resample.py +++ b/pycbc/filter/resample.py @@ -22,19 +22,31 @@ # ============================================================================= # import functools + import lal import numpy import scipy.signal -from pycbc.types import TimeSeries, Array, zeros, FrequencySeries, real_same_precision_as -from pycbc.types import complex_same_precision_as -from pycbc.fft import ifft, fft -_resample_func = {numpy.dtype('float32'): lal.ResampleREAL4TimeSeries, - numpy.dtype('float64'): lal.ResampleREAL8TimeSeries} +from pycbc.fft import fft, ifft +from pycbc.types import ( + Array, + FrequencySeries, + TimeSeries, + complex_same_precision_as, + real_same_precision_as, + zeros, +) + +_resample_func = { + numpy.dtype("float32"): lal.ResampleREAL4TimeSeries, + numpy.dtype("float64"): lal.ResampleREAL8TimeSeries, +} + @functools.lru_cache(maxsize=20) def cached_firwin(*args, **kwargs): - """Cache the FIR filter coefficients. + """ + Cache the FIR filter coefficients. This is mostly done for PyCBC Live, which rapidly and repeatedly resamples data. """ return scipy.signal.firwin(*args, **kwargs) @@ -51,8 +63,10 @@ def cached_firwin(*args, **kwargs): LFILTER_UNIQUE_ID_2 = 154687641 LFILTER_UNIQUE_ID_3 = 548946442 + def lfilter(coefficients, timeseries): - """ Apply filter coefficients to a time series + """ + Apply filter coefficients to a time series Parameters ---------- @@ -65,19 +79,23 @@ def lfilter(coefficients, timeseries): ------- tseries: numpy.ndarray filtered array + """ from pycbc.filter import correlate + fillen = len(coefficients) # If there aren't many points just use the default scipy method if len(timeseries) < 2**7: series = scipy.signal.lfilter(coefficients, 1.0, timeseries) - return TimeSeries(series, - epoch=timeseries.start_time, - delta_t=timeseries.delta_t) - elif (len(timeseries) < fillen * 10) or (len(timeseries) < 2**18): - from pycbc.strain.strain import create_memory_and_engine_for_class_based_fft - from pycbc.strain.strain import execute_cached_fft + return TimeSeries( + series, epoch=timeseries.start_time, delta_t=timeseries.delta_t + ) + if (len(timeseries) < fillen * 10) or (len(timeseries) < 2**18): + from pycbc.strain.strain import ( + create_memory_and_engine_for_class_based_fft, + execute_cached_fft, + ) cseries = (Array(coefficients[::-1] * 1)).astype(timeseries.dtype) cseries.resize(len(timeseries)) @@ -100,21 +118,24 @@ def lfilter(coefficients, timeseries): npoints = len(cseries) # NOTE: This function is cached! ifftouts = create_memory_and_engine_for_class_based_fft( - npoints, - timeseries.dtype, - ifft=True, - uid=LFILTER_UNIQUE_ID_1 + npoints, timeseries.dtype, ifft=True, uid=LFILTER_UNIQUE_ID_1 ) # FFT contents of cseries into cfreq - cfreq = execute_cached_fft(cseries, uid=LFILTER_UNIQUE_ID_2, - copy_output=False, - normalize_by_rate=False) + cfreq = execute_cached_fft( + cseries, + uid=LFILTER_UNIQUE_ID_2, + copy_output=False, + normalize_by_rate=False, + ) # FFT contents of timeseries into tfreq - tfreq = execute_cached_fft(timeseries, uid=LFILTER_UNIQUE_ID_3, - copy_output=False, - normalize_by_rate=False) + tfreq = execute_cached_fft( + timeseries, + uid=LFILTER_UNIQUE_ID_3, + copy_output=False, + normalize_by_rate=False, + ) cout, out, fft_class = ifftouts @@ -123,21 +144,25 @@ def lfilter(coefficients, timeseries): # IFFT correlation output into out fft_class.execute() - return TimeSeries(out.numpy() / len(out), epoch=timeseries.start_time, - delta_t=timeseries.delta_t) - else: - # recursively perform which saves a bit on memory usage - # but must keep within recursion limit - chunksize = max(fillen * 5, len(timeseries) // 2) - part1 = lfilter(coefficients, timeseries[0:chunksize]) - part2 = lfilter(coefficients, timeseries[chunksize - fillen:]) - out = timeseries.copy() - out[:len(part1)] = part1 - out[len(part1):] = part2[fillen:] - return out + return TimeSeries( + out.numpy() / len(out), + epoch=timeseries.start_time, + delta_t=timeseries.delta_t, + ) + # recursively perform which saves a bit on memory usage + # but must keep within recursion limit + chunksize = max(fillen * 5, len(timeseries) // 2) + part1 = lfilter(coefficients, timeseries[0:chunksize]) + part2 = lfilter(coefficients, timeseries[chunksize - fillen :]) + out = timeseries.copy() + out[: len(part1)] = part1 + out[len(part1) :] = part2[fillen:] + return out + def fir_zero_filter(coeff, timeseries): - """Filter the timeseries with a set of FIR coefficients + """ + Filter the timeseries with a set of FIR coefficients Parameters ---------- @@ -151,6 +176,7 @@ def fir_zero_filter(coeff, timeseries): filtered_series: pycbc.types.TimeSeries Return the filtered timeseries, which has been properly shifted to account for the FIR filter delay and the corrupted regions zeroed out. + """ # apply the filter series = lfilter(coeff, timeseries) @@ -159,12 +185,14 @@ def fir_zero_filter(coeff, timeseries): # corruption regions contain zeros # If the number of filter coefficients is odd, the central point *should* # be included in the output so we only zero out a region of len(coeff) - 1 - series[:(len(coeff) // 2) * 2] = 0 - series.roll(-len(coeff)//2) + series[: (len(coeff) // 2) * 2] = 0 + series.roll(-len(coeff) // 2) return series -def resample_to_delta_t(timeseries, delta_t, method='butterworth'): - """Resmple the time_series to delta_t + +def resample_to_delta_t(timeseries, delta_t, method="butterworth"): + """ + Resmple the time_series to delta_t Resamples the TimeSeries instance time_series to the given time step, delta_t. Only powers of two and real valued time series are supported @@ -192,41 +220,40 @@ def resample_to_delta_t(timeseries, delta_t, method='butterworth'): Examples -------- - >>> h_plus_sampled = resample_to_delta_t(h_plus, 1.0/2048) + """ - if not isinstance(timeseries,TimeSeries): + if not isinstance(timeseries, TimeSeries): raise TypeError("Can only resample time series") - if timeseries.kind != 'real': + if timeseries.kind != "real": raise TypeError("Time series must be real") if timeseries.sample_rate_close(1.0 / delta_t): return timeseries * 1 - if method == 'butterworth': + if method == "butterworth": lal_data = timeseries.lal() _resample_func[timeseries.dtype](lal_data, delta_t) data = lal_data.data.data - elif method == 'ldas': + elif method == "ldas": factor = int(round(delta_t / timeseries.delta_t)) numtaps = factor * 20 + 1 # The kaiser window has been testing using the LDAS implementation # and is in the same configuration as used in the original lalinspiral - filter_coefficients = cached_firwin(numtaps, 1.0 / factor, - window=('kaiser', 5)) + filter_coefficients = cached_firwin(numtaps, 1.0 / factor, window=("kaiser", 5)) # apply the filter and decimate data = fir_zero_filter(filter_coefficients, timeseries)[::factor] else: - raise ValueError('Invalid resampling method: %s' % method) + raise ValueError("Invalid resampling method: %s" % method) - ts = TimeSeries(data, delta_t = delta_t, - dtype=timeseries.dtype, - epoch=timeseries._epoch) + ts = TimeSeries( + data, delta_t=delta_t, dtype=timeseries.dtype, epoch=timeseries._epoch + ) # From the construction of the LDAS FIR filter there will be 10 corrupted samples # explanation here https://lscsoft.docs.ligo.org/lalsuite/lal/group___resample_time_series__c.html @@ -234,14 +261,19 @@ def resample_to_delta_t(timeseries, delta_t, method='butterworth'): return ts -_highpass_func = {numpy.dtype('float32'): lal.HighPassREAL4TimeSeries, - numpy.dtype('float64'): lal.HighPassREAL8TimeSeries} -_lowpass_func = {numpy.dtype('float32'): lal.LowPassREAL4TimeSeries, - numpy.dtype('float64'): lal.LowPassREAL8TimeSeries} +_highpass_func = { + numpy.dtype("float32"): lal.HighPassREAL4TimeSeries, + numpy.dtype("float64"): lal.HighPassREAL8TimeSeries, +} +_lowpass_func = { + numpy.dtype("float32"): lal.LowPassREAL4TimeSeries, + numpy.dtype("float64"): lal.LowPassREAL8TimeSeries, +} def notch_fir(timeseries, f1, f2, order, beta=5.0): - """ notch filter the time series using an FIR filtered generated from + """ + Notch filter the time series using an FIR filtered generated from the ideal response passed through a time-domain kaiser window (beta = 5.0) The suppression of the notch filter is related to the bandwidth and @@ -264,14 +296,17 @@ def notch_fir(timeseries, f1, f2, order, beta=5.0): (Extent of the filter on either side of zero) beta: float Beta parameter of the kaiser window that sets the side lobe attenuation. + """ - k1 = f1 / float((int(1.0 / timeseries.delta_t) / 2)) - k2 = f2 / float((int(1.0 / timeseries.delta_t) / 2)) - coeff = cached_firwin(order * 2 + 1, [k1, k2], window=('kaiser', beta)) + k1 = f1 / float(int(1.0 / timeseries.delta_t) / 2) + k2 = f2 / float(int(1.0 / timeseries.delta_t) / 2) + coeff = cached_firwin(order * 2 + 1, [k1, k2], window=("kaiser", beta)) return fir_zero_filter(coeff, timeseries) + def lowpass_fir(timeseries, frequency, order, beta=5.0): - """ Lowpass filter the time series using an FIR filtered generated from + """ + Lowpass filter the time series using an FIR filtered generated from the ideal response passed through a kaiser window (beta = 5.0) Parameters @@ -284,13 +319,16 @@ def lowpass_fir(timeseries, frequency, order, beta=5.0): Number of corrupted samples on each side of the time series beta: float Beta parameter of the kaiser window that sets the side lobe attenuation. + """ - k = frequency / float((int(1.0 / timeseries.delta_t) / 2)) - coeff = cached_firwin(order * 2 + 1, k, window=('kaiser', beta)) + k = frequency / float(int(1.0 / timeseries.delta_t) / 2) + coeff = cached_firwin(order * 2 + 1, k, window=("kaiser", beta)) return fir_zero_filter(coeff, timeseries) + def highpass_fir(timeseries, frequency, order, beta=5.0): - """ Highpass filter the time series using an FIR filtered generated from + """ + Highpass filter the time series using an FIR filtered generated from the ideal response passed through a kaiser window (beta = 5.0) Parameters @@ -303,13 +341,16 @@ def highpass_fir(timeseries, frequency, order, beta=5.0): Number of corrupted samples on each side of the time series beta: float Beta parameter of the kaiser window that sets the side lobe attenuation. + """ - k = frequency / float((int(1.0 / timeseries.delta_t) / 2)) - coeff = cached_firwin(order * 2 + 1, k, window=('kaiser', beta), pass_zero=False) + k = frequency / float(int(1.0 / timeseries.delta_t) / 2) + coeff = cached_firwin(order * 2 + 1, k, window=("kaiser", beta), pass_zero=False) return fir_zero_filter(coeff, timeseries) + def highpass(timeseries, frequency, filter_order=8, attenuation=0.1): - """Return a new timeseries that is highpassed. + """ + Return a new timeseries that is highpassed. Return a new time series that is highpassed above the `frequency`. @@ -337,22 +378,26 @@ def highpass(timeseries, frequency, filter_order=8, attenuation=0.1): time_series is not real valued """ - if not isinstance(timeseries, TimeSeries): raise TypeError("Can only resample time series") - if timeseries.kind != 'real': + if timeseries.kind != "real": raise TypeError("Time series must be real") lal_data = timeseries.lal() - _highpass_func[timeseries.dtype](lal_data, frequency, - 1-attenuation, filter_order) + _highpass_func[timeseries.dtype](lal_data, frequency, 1 - attenuation, filter_order) + + return TimeSeries( + lal_data.data.data, + delta_t=lal_data.deltaT, + dtype=timeseries.dtype, + epoch=timeseries._epoch, + ) - return TimeSeries(lal_data.data.data, delta_t = lal_data.deltaT, - dtype=timeseries.dtype, epoch=timeseries._epoch) def lowpass(timeseries, frequency, filter_order=8, attenuation=0.1): - """Return a new timeseries that is lowpassed. + """ + Return a new timeseries that is lowpassed. Return a new time series that is lowpassed below the `frequency`. @@ -378,24 +423,28 @@ def lowpass(timeseries, frequency, filter_order=8, attenuation=0.1): time_series is not an instance of TimeSeries. TypeError: time_series is not real valued - """ + """ if not isinstance(timeseries, TimeSeries): raise TypeError("Can only resample time series") - if timeseries.kind != 'real': + if timeseries.kind != "real": raise TypeError("Time series must be real") lal_data = timeseries.lal() - _lowpass_func[timeseries.dtype](lal_data, frequency, - 1-attenuation, filter_order) + _lowpass_func[timeseries.dtype](lal_data, frequency, 1 - attenuation, filter_order) - return TimeSeries(lal_data.data.data, delta_t = lal_data.deltaT, - dtype=timeseries.dtype, epoch=timeseries._epoch) + return TimeSeries( + lal_data.data.data, + delta_t=lal_data.deltaT, + dtype=timeseries.dtype, + epoch=timeseries._epoch, + ) -def interpolate_complex_frequency(series, delta_f, zeros_offset=0, side='right'): - """Interpolate complex frequency series to desired delta_f. +def interpolate_complex_frequency(series, delta_f, zeros_offset=0, side="right"): + """ + Interpolate complex frequency series to desired delta_f. Return a new complex frequency series that has been interpolated to the desired delta_f. @@ -415,30 +464,42 @@ def interpolate_complex_frequency(series, delta_f, zeros_offset=0, side='right') ------- interpolated series : FrequencySeries A new FrequencySeries that has been interpolated. + """ - new_n = int( (len(series)-1) * series.delta_f / delta_f + 1) - old_N = int( (len(series)-1) * 2 ) - new_N = int( (new_n - 1) * 2 ) - time_series = TimeSeries(zeros(old_N), delta_t =1.0/(series.delta_f*old_N), - dtype=real_same_precision_as(series)) + new_n = int((len(series) - 1) * series.delta_f / delta_f + 1) + old_N = int((len(series) - 1) * 2) + new_N = int((new_n - 1) * 2) + time_series = TimeSeries( + zeros(old_N), + delta_t=1.0 / (series.delta_f * old_N), + dtype=real_same_precision_as(series), + ) ifft(series, time_series) time_series.roll(-zeros_offset) time_series.resize(new_N) - if side == 'left': + if side == "left": time_series.roll(zeros_offset + new_N - old_N) - elif side == 'right': + elif side == "right": time_series.roll(zeros_offset) - out_series = FrequencySeries(zeros(new_n), epoch=series.epoch, - delta_f=delta_f, dtype=series.dtype) + out_series = FrequencySeries( + zeros(new_n), epoch=series.epoch, delta_f=delta_f, dtype=series.dtype + ) fft(time_series, out_series) return out_series -__all__ = ['resample_to_delta_t', 'highpass', 'lowpass', - 'interpolate_complex_frequency', 'highpass_fir', - 'lowpass_fir', 'notch_fir', 'fir_zero_filter'] +__all__ = [ + "fir_zero_filter", + "highpass", + "highpass_fir", + "interpolate_complex_frequency", + "lowpass", + "lowpass_fir", + "notch_fir", + "resample_to_delta_t", +] diff --git a/pycbc/filter/simd_correlate.py b/pycbc/filter/simd_correlate.py index faeacccae0d..70d2012e53b 100644 --- a/pycbc/filter/simd_correlate.py +++ b/pycbc/filter/simd_correlate.py @@ -13,10 +13,12 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -from pycbc.types import float32, complex64 import numpy as _np + +from pycbc.types import complex64, float32 + from .. import opt -from .simd_correlate_cython import ccorrf_simd, ccorrf_parallel +from .simd_correlate_cython import ccorrf_parallel, ccorrf_simd """ This module interfaces to C functions for multiplying @@ -63,6 +65,7 @@ def correlate_simd(ht, st, qt): # Seems to work for Sandy Bridge/Ivy Bridge/Haswell, for now? default_segsize = 8192 + def correlate_parallel(ht, st, qt): htilde = _np.array(ht.data, copy=False, dtype=complex64) stilde = _np.array(st.data, copy=False, dtype=complex64) diff --git a/pycbc/filter/zpk.py b/pycbc/filter/zpk.py index 635088634a7..ca74392dc5d 100644 --- a/pycbc/filter/zpk.py +++ b/pycbc/filter/zpk.py @@ -23,12 +23,14 @@ # import numpy as np +from scipy.signal import sosfilt, zpk2sos -from scipy.signal import zpk2sos, sosfilt from pycbc.types import TimeSeries + def filter_zpk(timeseries, z, p, k): - """Return a new timeseries that was filtered with a zero-pole-gain filter. + """ + Return a new timeseries that was filtered with a zero-pole-gain filter. The transfer function in the s-domain looks like: .. math:: \\frac{H(s) = (s - s_1) * (s - s_3) * ... * (s - s_n)}{(s - s_2) * (s - s_4) * ... * (s - s_m)}, m >= n @@ -68,8 +70,8 @@ def filter_zpk(timeseries, z, p, k): To apply a 5 zeroes at 100Hz, 5 poles at 1Hz, and a gain of 1e-10 filter to a TimeSeries instance, do: >>> filtered_data = zpk_filter(timeseries, [100]*5, [1]*5, 1e-10) - """ + """ # sanity check type if not isinstance(timeseries, TimeSeries): raise TypeError("Can only filter TimeSeries instances.") @@ -77,8 +79,10 @@ def filter_zpk(timeseries, z, p, k): # sanity check casual filter degree = len(p) - len(z) if degree < 0: - raise TypeError("May not have more zeroes than poles. \ - Filter is not casual.") + raise TypeError( + "May not have more zeroes than poles. \ + Filter is not casual." + ) # cast zeroes and poles as arrays and gain as a float z = np.array(z) @@ -94,14 +98,14 @@ def filter_zpk(timeseries, z, p, k): fs = 2.0 * timeseries.sample_rate # zeroes in the z-domain - z_zd = (1 + z/fs) / (1 - z/fs) + z_zd = (1 + z / fs) / (1 - z / fs) # any zeros that were at infinity are moved to the Nyquist frequency z_zd = z_zd[np.isfinite(z_zd)] z_zd = np.append(z_zd, -np.ones(degree)) # poles in the z-domain - p_zd = (1 + p/fs) / (1 - p/fs) + p_zd = (1 + p / fs) / (1 - p / fs) # gain change in z-domain k_zd = k * np.prod(fs - z) / np.prod(fs - p) @@ -112,6 +116,9 @@ def filter_zpk(timeseries, z, p, k): # filter filtered_data = sosfilt(sos, timeseries.numpy()) - return TimeSeries(filtered_data, delta_t = timeseries.delta_t, - dtype=timeseries.dtype, - epoch=timeseries._epoch) + return TimeSeries( + filtered_data, + delta_t=timeseries.delta_t, + dtype=timeseries.dtype, + epoch=timeseries._epoch, + ) diff --git a/pycbc/frame/__init__.py b/pycbc/frame/__init__.py index 0ef88b9c568..ee265a35738 100644 --- a/pycbc/frame/__init__.py +++ b/pycbc/frame/__init__.py @@ -1,9 +1,14 @@ -from . frame import (locations_to_cache, read_frame, - query_and_read_frame, frame_paths, write_frame, - DataBuffer, StatusBuffer, iDQBuffer) - -from . store import (read_store) - +from .frame import ( + DataBuffer, + StatusBuffer, + frame_paths, + iDQBuffer, + locations_to_cache, + query_and_read_frame, + read_frame, + write_frame, +) +from .store import read_store # Status flags for the calibration state vector # See e.g. https://dcc.ligo.org/LIGO-G1700234 @@ -23,8 +28,7 @@ KAPPA_C_OK = 4096 FCC_OK = 8192 NO_GAP = 16384 -NO_HWINJ = NO_STOCH_HW_INJ | NO_CBC_HW_INJ | \ - NO_BURST_HW_INJ | NO_DETCHAR_HW_INJ +NO_HWINJ = NO_STOCH_HW_INJ | NO_CBC_HW_INJ | NO_BURST_HW_INJ | NO_DETCHAR_HW_INJ # relevant bits in the LIGO O2/O3 low-latency DQ vector # If the bit is 0 then we should veto @@ -40,7 +44,8 @@ def flag_names_to_bitmask(flags): - """Takes a list of flag names corresponding to bits in a status channel + """ + Takes a list of flag names corresponding to bits in a status channel and returns the corresponding bit mask. """ mask = 0 diff --git a/pycbc/frame/frame.py b/pycbc/frame/frame.py index bea7386418a..080c0420bbe 100644 --- a/pycbc/frame/frame.py +++ b/pycbc/frame/frame.py @@ -17,95 +17,104 @@ This modules contains functions for reading in data from frame files or caches """ -import logging -import warnings -import os.path import glob -import time +import logging import math +import os.path import re +import time +import warnings from urllib.parse import urlparse -import numpy -import lalframe import lal +import lalframe +import numpy from gwdatafind import find_urls as find_frame_urls import pycbc from pycbc.types import TimeSeries, zeros -logger = logging.getLogger('pycbc.frame.frame') +logger = logging.getLogger("pycbc.frame.frame") # map LAL series types to corresponding functions and Numpy types _fr_type_map = { lal.S_TYPE_CODE: [ - lalframe.FrStreamReadREAL4TimeSeries, numpy.float32, + lalframe.FrStreamReadREAL4TimeSeries, + numpy.float32, lal.CreateREAL4TimeSeries, lalframe.FrStreamGetREAL4TimeSeriesMetadata, lal.CreateREAL4Sequence, - lalframe.FrameAddREAL4TimeSeriesProcData + lalframe.FrameAddREAL4TimeSeriesProcData, ], lal.D_TYPE_CODE: [ - lalframe.FrStreamReadREAL8TimeSeries, numpy.float64, + lalframe.FrStreamReadREAL8TimeSeries, + numpy.float64, lal.CreateREAL8TimeSeries, lalframe.FrStreamGetREAL8TimeSeriesMetadata, lal.CreateREAL8Sequence, - lalframe.FrameAddREAL8TimeSeriesProcData + lalframe.FrameAddREAL8TimeSeriesProcData, ], lal.C_TYPE_CODE: [ - lalframe.FrStreamReadCOMPLEX8TimeSeries, numpy.complex64, + lalframe.FrStreamReadCOMPLEX8TimeSeries, + numpy.complex64, lal.CreateCOMPLEX8TimeSeries, lalframe.FrStreamGetCOMPLEX8TimeSeriesMetadata, lal.CreateCOMPLEX8Sequence, - lalframe.FrameAddCOMPLEX8TimeSeriesProcData + lalframe.FrameAddCOMPLEX8TimeSeriesProcData, ], lal.Z_TYPE_CODE: [ - lalframe.FrStreamReadCOMPLEX16TimeSeries, numpy.complex128, + lalframe.FrStreamReadCOMPLEX16TimeSeries, + numpy.complex128, lal.CreateCOMPLEX16TimeSeries, lalframe.FrStreamGetCOMPLEX16TimeSeriesMetadata, lal.CreateCOMPLEX16Sequence, - lalframe.FrameAddCOMPLEX16TimeSeriesProcData + lalframe.FrameAddCOMPLEX16TimeSeriesProcData, ], lal.U4_TYPE_CODE: [ - lalframe.FrStreamReadUINT4TimeSeries, numpy.uint32, + lalframe.FrStreamReadUINT4TimeSeries, + numpy.uint32, lal.CreateUINT4TimeSeries, lalframe.FrStreamGetUINT4TimeSeriesMetadata, lal.CreateUINT4Sequence, - lalframe.FrameAddUINT4TimeSeriesProcData + lalframe.FrameAddUINT4TimeSeriesProcData, ], lal.I4_TYPE_CODE: [ - lalframe.FrStreamReadINT4TimeSeries, numpy.int32, + lalframe.FrStreamReadINT4TimeSeries, + numpy.int32, lal.CreateINT4TimeSeries, lalframe.FrStreamGetINT4TimeSeriesMetadata, lal.CreateINT4Sequence, - lalframe.FrameAddINT4TimeSeriesProcData + lalframe.FrameAddINT4TimeSeriesProcData, ], } + def _read_channel(channel, stream, start, duration): - """ Get channel using lalframe """ + """Get channel using lalframe""" channel_type = lalframe.FrStreamGetTimeSeriesType(channel, stream) read_func = _fr_type_map[channel_type][0] d_type = _fr_type_map[channel_type][1] data = read_func(stream, channel, start, duration, 0) - return TimeSeries(data.data.data, delta_t=data.deltaT, epoch=start, - dtype=d_type) + return TimeSeries(data.data.data, delta_t=data.deltaT, epoch=start, dtype=d_type) def _is_gwf(file_path): - """Test if a file is a frame file by checking if its contents begins with - the magic string 'IGWD'.""" + """ + Test if a file is a frame file by checking if its contents begins with + the magic string 'IGWD'. + """ try: - with open(file_path, 'rb') as f: - if f.read(4) == b'IGWD': + with open(file_path, "rb") as f: + if f.read(4) == b"IGWD": return True - except IOError: + except OSError: pass return False def locations_to_cache(locations, latest=False): - """ Return a cumulative cache file build from the list of locations + """ + Return a cumulative cache file build from the list of locations Parameters ---------- @@ -121,11 +130,13 @@ def locations_to_cache(locations, latest=False): cache : lal.Cache A cumulative lal cache object containing the files derived from the list of locations. + """ cum_cache = lal.Cache() for source in locations: flist = glob.glob(source) if latest: + def relaxed_getctime(fn): # when building a cache from a directory of temporary # low-latency frames, files might disappear between @@ -134,8 +145,9 @@ def relaxed_getctime(fn): return os.path.getctime(fn) except OSError: return 0 + if not flist: - raise ValueError('no frame or cache files found in ' + source) + raise ValueError("no frame or cache files found in " + source) flist = [max(flist, key=relaxed_getctime)] for file_path in flist: @@ -152,10 +164,18 @@ def relaxed_getctime(fn): cum_cache = lal.CacheMerge(cum_cache, cache) return cum_cache -def read_frame(location, channels, start_time=None, - end_time=None, duration=None, check_integrity=False, - sieve=None): - """Read time series from frame data. + +def read_frame( + location, + channels, + start_time=None, + end_time=None, + duration=None, + check_integrity=False, + sieve=None, +): + """ + Read time series from frame data. Using the `location`, which can either be a frame file ".gwf" or a frame cache ".gwf", read in the data for the given channel(s) and output @@ -189,8 +209,8 @@ def read_frame(location, channels, start_time=None, Frame Data: TimeSeries or list of TimeSeries A TimeSeries or a list of TimeSeries, corresponding to the data from the frame file/cache for a given channel or channels. - """ + """ if end_time and duration: raise ValueError("end time and duration are mutually exclusive") @@ -207,14 +227,15 @@ def read_frame(location, channels, start_time=None, # Before sieving, check if this is sane. Otherwise it will fail later. if (int(math.ceil(end_time)) - int(start_time)) <= 0: raise ValueError("Negative or null duration") - lal.CacheSieve(cum_cache, int(start_time), int(math.ceil(end_time)), - None, None, None) + lal.CacheSieve( + cum_cache, int(start_time), int(math.ceil(end_time)), None, None, None + ) stream = lalframe.FrStreamCacheOpen(cum_cache) stream.mode = lalframe.FR_STREAM_VERBOSE_MODE if check_integrity: - stream.mode = (stream.mode | lalframe.FR_STREAM_CHECKSUM_MODE) + stream.mode = stream.mode | lalframe.FR_STREAM_CHECKSUM_MODE lalframe.FrStreamSetMode(stream, stream.mode) @@ -228,13 +249,12 @@ def read_frame(location, channels, start_time=None, channel_type = lalframe.FrStreamGetTimeSeriesType(first_channel, stream) create_series_func = _fr_type_map[channel_type][2] get_series_metadata_func = _fr_type_map[channel_type][3] - series = create_series_func(first_channel, stream.epoch, 0, 0, - lal.ADCCountUnit, 0) + series = create_series_func(first_channel, stream.epoch, 0, 0, lal.ADCCountUnit, 0) get_series_metadata_func(series, stream) data_duration = (data_length + 0.5) * series.deltaT if start_time is None: - start_time = stream.epoch*1 + start_time = stream.epoch * 1 if end_time is None: end_time = start_time + data_duration @@ -251,7 +271,7 @@ def read_frame(location, channels, start_time=None, # lalframe behaves dangerously with invalid duration so catch it here if duration <= 0: raise ValueError("Negative or null duration") - #if duration > data_duration: + # if duration > data_duration: # raise ValueError("Requested duration longer than available data") if type(channels) is list: @@ -261,13 +281,14 @@ def read_frame(location, channels, start_time=None, lalframe.FrStreamSeek(stream, start_time) all_data.append(channel_data) return all_data - else: - return _read_channel(channels, stream, start_time, duration) + return _read_channel(channels, stream, start_time, duration) + def frame_paths( - frame_type, start_time, end_time, server=None, url_type='file', site=None + frame_type, start_time, end_time, server=None, url_type="file", site=None ): - """Return the paths to a span of frame files. + """ + Return the paths to a span of frame files. Parameters ---------- @@ -297,21 +318,23 @@ def frame_paths( Examples -------- >>> paths = frame_paths('H1_LDAS_C02_L2', 968995968, 968995968+2048) + """ if site is None: # this case is tolerated for backward compatibility site = frame_type[0] warnings.warn( - f'Guessing site {site} from frame type {frame_type}', - DeprecationWarning + f"Guessing site {site} from frame type {frame_type}", DeprecationWarning ) - cache = find_frame_urls(site, frame_type, start_time, end_time, - urltype=url_type, host=server) + cache = find_frame_urls( + site, frame_type, start_time, end_time, urltype=url_type, host=server + ) return [urlparse(entry).path for entry in cache] def get_site_from_type_or_channel(frame_type, channels): - """Determine the site for querying gwdatafind (H, L, V, etc) based on + """ + Determine the site for querying gwdatafind (H, L, V, etc) based on substrings of the frame type and channel(s). The type should begin with S: or SN:, in which case S is taken as the @@ -333,11 +356,12 @@ def get_site_from_type_or_channel(frame_type, channels): The site letter. frame_type : string The frame type with the site prefix (if any) removed. + """ - site_re = '^([^:])[^:]?:' + site_re = "^([^:])[^:]?:" m = re.match(site_re, frame_type) if m: - return m.groups(1)[0], frame_type[m.end():] + return m.groups(1)[0], frame_type[m.end() :] chan = channels if isinstance(chan, list): chan = channels[0] @@ -345,15 +369,17 @@ def get_site_from_type_or_channel(frame_type, channels): if m: return m.groups(1)[0], frame_type warnings.warn( - f'Guessing site {frame_type[0]} from frame type {frame_type}', - DeprecationWarning + f"Guessing site {frame_type[0]} from frame type {frame_type}", + DeprecationWarning, ) return frame_type[0], frame_type -def query_and_read_frame(frame_type, channels, start_time, end_time, - sieve=None, check_integrity=False): - """Read time series from frame data. +def query_and_read_frame( + frame_type, channels, start_time, end_time, sieve=None, check_integrity=False +): + """ + Read time series from frame data. Query for the location of physical frames matching the frame type. Return a time series containing the channel between the given start and end times. @@ -389,42 +415,42 @@ def query_and_read_frame(frame_type, channels, start_time, end_time, -------- >>> ts = query_and_read_frame('H1_LDAS_C02_L2', 'H1:LDAS-STRAIN', >>> 968995968, 968995968+2048) + """ site, frame_type = get_site_from_type_or_channel(frame_type, channels) # Allows compatibility with our standard tools # We may want to place this into a higher level frame getting tool - if frame_type in ['LOSC_STRAIN', 'GWOSC_STRAIN']: + if frame_type in ["LOSC_STRAIN", "GWOSC_STRAIN"]: from pycbc.frame.gwosc import read_strain_gwosc + if not isinstance(channels, list): channels = [channels] - data = [read_strain_gwosc(c[:2], start_time, end_time) - for c in channels] + data = [read_strain_gwosc(c[:2], start_time, end_time) for c in channels] return data if len(data) > 1 else data[0] - if frame_type in ['LOSC', 'GWOSC']: + if frame_type in ["LOSC", "GWOSC"]: from pycbc.frame.gwosc import read_frame_gwosc + return read_frame_gwosc(channels, start_time, end_time) - logger.info('Querying datafind server') + logger.info("Querying datafind server") paths = frame_paths( - frame_type, - int(start_time), - int(numpy.ceil(end_time)), - site=site + frame_type, int(start_time), int(numpy.ceil(end_time)), site=site ) - logger.info('Found frame file paths: %s', ' '.join(paths)) + logger.info("Found frame file paths: %s", " ".join(paths)) return read_frame( paths, channels, start_time=start_time, end_time=end_time, sieve=sieve, - check_integrity=check_integrity + check_integrity=check_integrity, ) def write_frame(location, channels, timeseries): - """Write a list of time series to a single frame file. + """ + Write a list of time series to a single frame file. Parameters ---------- @@ -436,6 +462,7 @@ def write_frame(location, channels, timeseries): timeseries: TimeSeries A TimeSeries or list of TimeSeries, corresponding to the data to be written to the frame file for a given channel. + """ # check if a single channel or a list of channels if type(channels) is list and type(timeseries) is list: @@ -459,13 +486,18 @@ def write_frame(location, channels, timeseries): raise ValueError("Start and end times of TimeSeries must be integer seconds.") # create frame - frame = lalframe.FrameNew(epoch=gps_start_time, duration=duration, - project='', run=1, frnum=1, - detectorFlags=lal.LALDETECTORTYPE_ABSENT) + frame = lalframe.FrameNew( + epoch=gps_start_time, + duration=duration, + project="", + run=1, + frnum=1, + detectorFlags=lal.LALDETECTORTYPE_ABSENT, + ) - for i,tseries in enumerate(timeseries): + for i, tseries in enumerate(timeseries): # get data type - for seriestype in _fr_type_map.keys(): + for seriestype in _fr_type_map: if _fr_type_map[seriestype][1] == tseries.dtype: create_series_func = _fr_type_map[seriestype][2] create_sequence_func = _fr_type_map[seriestype][4] @@ -473,9 +505,14 @@ def write_frame(location, channels, timeseries): break # add time series to frame - series = create_series_func(channels[i], tseries.start_time, - 0, tseries.delta_t, lal.ADCCountUnit, - len(tseries.numpy())) + series = create_series_func( + channels[i], + tseries.start_time, + 0, + tseries.delta_t, + lal.ADCCountUnit, + len(tseries.numpy()), + ) series.data = create_sequence_func(len(tseries.numpy())) series.data.data = tseries.numpy() add_series_func(frame, series) @@ -483,22 +520,25 @@ def write_frame(location, channels, timeseries): # write frame lalframe.FrameWrite(frame, location) -class DataBuffer(object): - """A linear buffer that acts as a FILO for reading in frame data - """ +class DataBuffer: + """A linear buffer that acts as a FILO for reading in frame data""" - def __init__(self, frame_src, - channel_name, - start_time, - max_buffer=2048, - force_update_cache=True, - increment_update_cache=None, - dtype=numpy.float64): - """ Create a rolling buffer of frame data + def __init__( + self, + frame_src, + channel_name, + start_time, + max_buffer=2048, + force_update_cache=True, + increment_update_cache=None, + dtype=numpy.float64, + ): + """ + Create a rolling buffer of frame data Parameters - --------- + ---------- frame_src: str of list of strings Strings that indicate where to read from files from. This can be a list of frame files, a glob, etc. @@ -510,25 +550,31 @@ def __init__(self, frame_src, Length of the buffer in seconds dtype: {dtype, numpy.float32}, Optional Data type to use for the interal buffer + """ self.frame_src = frame_src self.channel_name = channel_name self.read_pos = start_time self.force_update_cache = force_update_cache self.increment_update_cache = increment_update_cache - self.detector = channel_name.split(':')[0] + self.detector = channel_name.split(":")[0] self.update_cache() - self.channel_type, self.raw_sample_rate = self._retrieve_metadata(self.stream, self.channel_name) + self.channel_type, self.raw_sample_rate = self._retrieve_metadata( + self.stream, self.channel_name + ) raw_size = self.raw_sample_rate * max_buffer - self.raw_buffer = TimeSeries(zeros(raw_size, dtype=dtype), - copy=False, - epoch=start_time - max_buffer, - delta_t=1.0/self.raw_sample_rate) + self.raw_buffer = TimeSeries( + zeros(raw_size, dtype=dtype), + copy=False, + epoch=start_time - max_buffer, + delta_t=1.0 / self.raw_sample_rate, + ) def update_cache(self): - """Reset the lal cache. This can be used to update the cache if the + """ + Reset the lal cache. This can be used to update the cache if the result may change due to more files being added to the filesystem, for example. """ @@ -538,7 +584,8 @@ def update_cache(self): @staticmethod def _retrieve_metadata(stream, channel_name): - """Retrieve basic metadata by reading the first file in the cache + """ + Retrieve basic metadata by reading the first file in the cache Parameters ---------- @@ -553,18 +600,21 @@ def _retrieve_metadata(stream, channel_name): Enum value which indicates the dtype of the channel sample_rate: int The sample rate of the data within this channel + """ lalframe.FrStreamGetVectorLength(channel_name, stream) channel_type = lalframe.FrStreamGetTimeSeriesType(channel_name, stream) create_series_func = _fr_type_map[channel_type][2] get_series_metadata_func = _fr_type_map[channel_type][3] - series = create_series_func(channel_name, stream.epoch, 0, 0, - lal.ADCCountUnit, 0) + series = create_series_func( + channel_name, stream.epoch, 0, 0, lal.ADCCountUnit, 0 + ) get_series_metadata_func(series, stream) - return channel_type, int(1.0/series.deltaT) + return channel_type, int(1.0 / series.deltaT) def _read_frame(self, blocksize): - """Try to read the block of data blocksize seconds long + """ + Try to read the block of data blocksize seconds long Parameters ---------- @@ -580,49 +630,56 @@ def _read_frame(self, blocksize): ------ RuntimeError: If data cannot be read for any reason + """ try: read_func = _fr_type_map[self.channel_type][0] dtype = _fr_type_map[self.channel_type][1] - data = read_func(self.stream, self.channel_name, - self.read_pos, int(blocksize), 0) - return TimeSeries(data.data.data, delta_t=data.deltaT, - epoch=self.read_pos, - dtype=dtype) + data = read_func( + self.stream, self.channel_name, self.read_pos, int(blocksize), 0 + ) + return TimeSeries( + data.data.data, delta_t=data.deltaT, epoch=self.read_pos, dtype=dtype + ) except Exception: - raise RuntimeError('Cannot read {0} frame data'.format(self.channel_name)) + raise RuntimeError(f"Cannot read {self.channel_name} frame data") def null_advance(self, blocksize): - """Advance and insert zeros + """ + Advance and insert zeros Parameters ---------- blocksize: int The number of seconds to attempt to read from the channel + """ self.raw_buffer.roll(-int(blocksize * self.raw_sample_rate)) self.read_pos += blocksize self.raw_buffer.start_time += blocksize def advance(self, blocksize): - """Add blocksize seconds more to the buffer, push blocksize seconds + """ + Add blocksize seconds more to the buffer, push blocksize seconds from the beginning. Parameters ---------- blocksize: int The number of seconds to attempt to read from the channel + """ ts = self._read_frame(blocksize) self.raw_buffer.roll(-len(ts)) - self.raw_buffer[-len(ts):] = ts[:] + self.raw_buffer[-len(ts) :] = ts[:] self.read_pos += blocksize self.raw_buffer.start_time += blocksize return ts def update_cache_by_increment(self, blocksize): - """Update the internal cache by starting from the first frame + """ + Update the internal cache by starting from the first frame and incrementing. Guess the next frame file name by incrementing from the first found @@ -633,29 +690,32 @@ def update_cache_by_increment(self, blocksize): ---------- blocksize: int Number of seconds to increment the next frame file. + """ start = float(self.raw_buffer.end_time) end = float(start + blocksize) - if not hasattr(self, 'dur'): + if not hasattr(self, "dur"): fname = glob.glob(self.frame_src[0])[0] - fname = os.path.splitext(os.path.basename(fname))[0].split('-') + fname = os.path.splitext(os.path.basename(fname))[0].split("-") - self.beg = '-'.join([fname[0], fname[1]]) + self.beg = "-".join([fname[0], fname[1]]) self.ref = int(fname[2]) self.dur = int(fname[3]) - fstart = int(self.ref + numpy.floor((start - self.ref) / float(self.dur)) * self.dur) + fstart = int( + self.ref + numpy.floor((start - self.ref) / float(self.dur)) * self.dur + ) starts = numpy.arange(fstart, end, self.dur).astype(int) keys = [] for s in starts: pattern = self.increment_update_cache - if 'GPS' in pattern: - n = int(pattern[int(pattern.index('GPS') + 3)]) - pattern = pattern.replace('GPS%s' % n, str(s)[0:n]) + if "GPS" in pattern: + n = int(pattern[int(pattern.index("GPS") + 3)]) + pattern = pattern.replace("GPS%s" % n, str(s)[0:n]) - name = f'{pattern}/{self.beg}-{s}-{self.dur}.gwf' + name = f"{pattern}/{self.beg}-{s}-{self.dur}.gwf" # check that file actually exists, else abort now if not os.path.exists(name): raise RuntimeError @@ -664,11 +724,13 @@ def update_cache_by_increment(self, blocksize): cache = locations_to_cache(keys) stream = lalframe.FrStreamCacheOpen(cache) self.stream = stream - self.channel_type, self.raw_sample_rate = \ - self._retrieve_metadata(self.stream, self.channel_name) + self.channel_type, self.raw_sample_rate = self._retrieve_metadata( + self.stream, self.channel_name + ) def attempt_advance(self, blocksize, timeout=10): - """ Attempt to advance the frame buffer. Retry upon failure, except + """ + Attempt to advance the frame buffer. Retry upon failure, except if the frame file is beyond the timeout limit. Parameters @@ -682,6 +744,7 @@ def attempt_advance(self, blocksize, timeout=10): ------- data: TimeSeries TimeSeries containg 'blocksize' seconds of frame data + """ if self.force_update_cache: self.update_cache() @@ -700,22 +763,26 @@ def attempt_advance(self, blocksize, timeout=10): # so we should try again time.sleep(0.1) -class StatusBuffer(DataBuffer): - - """ Read state vector or DQ information from a frame file """ - def __init__(self, frame_src, - channel_name, - start_time, - max_buffer=2048, - valid_mask=3, - force_update_cache=False, - increment_update_cache=None, - valid_on_zero=False): - """ Create a rolling buffer of status data from a frame +class StatusBuffer(DataBuffer): + """Read state vector or DQ information from a frame file""" + + def __init__( + self, + frame_src, + channel_name, + start_time, + max_buffer=2048, + valid_mask=3, + force_update_cache=False, + increment_update_cache=None, + valid_on_zero=False, + ): + """ + Create a rolling buffer of status data from a frame Parameters - --------- + ---------- frame_src: str of list of strings Strings that indicate where to read from files from. This can be a list of frame files, a glob, etc. @@ -730,18 +797,25 @@ def __init__(self, frame_src, valid_on_zero: bool If True, `valid_mask` is ignored and the status is considered "good" simply when the channel is zero. + """ - DataBuffer.__init__(self, frame_src, channel_name, start_time, - max_buffer=max_buffer, - force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache, - dtype=numpy.int32) + DataBuffer.__init__( + self, + frame_src, + channel_name, + start_time, + max_buffer=max_buffer, + force_update_cache=force_update_cache, + increment_update_cache=increment_update_cache, + dtype=numpy.int32, + ) self.valid_mask = valid_mask self.valid_on_zero = valid_on_zero def check_valid(self, values, flag=None): - """Check if the data contains any non-valid status information + """ + Check if the data contains any non-valid status information Parameters ---------- @@ -755,6 +829,7 @@ def check_valid(self, values, flag=None): status: boolean Returns True if all of the status information if valid, False if any is not. + """ if self.valid_on_zero: valid = values.numpy() == 0 @@ -765,7 +840,8 @@ def check_valid(self, values, flag=None): return bool(numpy.all(valid)) def is_extent_valid(self, start_time, duration, flag=None): - """Check if the duration contains any non-valid frames + """ + Check if the duration contains any non-valid frames Parameters ---------- @@ -781,6 +857,7 @@ def is_extent_valid(self, start_time, duration, flag=None): status: boolean Returns True if all of the status information if valid, False if any is not. + """ sr = self.raw_buffer.sample_rate s = int((start_time - self.raw_buffer.start_time) * sr) @@ -789,7 +866,8 @@ def is_extent_valid(self, start_time, duration, flag=None): return self.check_valid(data, flag=flag) def indices_of_flag(self, start_time, duration, times, padding=0): - """ Return the indices of the times lying in the flagged region + """ + Return the indices of the times lying in the flagged region Parameters ---------- @@ -806,8 +884,10 @@ def indices_of_flag(self, start_time, duration, times, padding=0): indices: numpy.ndarray Array of indices marking the location of triggers within valid time. + """ from pycbc.events.veto import indices_outside_times + sr = self.raw_buffer.sample_rate s = int((start_time - self.raw_buffer.start_time - padding) * sr) - 1 e = s + int((duration + padding) * sr) + 1 @@ -817,8 +897,9 @@ def indices_of_flag(self, start_time, duration, times, padding=0): if self.valid_on_zero: invalid = data.numpy() != 0 else: - invalid = numpy.bitwise_and(data.numpy(), self.valid_mask) \ - != self.valid_mask + invalid = ( + numpy.bitwise_and(data.numpy(), self.valid_mask) != self.valid_mask + ) starts = stamps[invalid] - padding ends = starts + 1.0 / sr + padding * 2.0 @@ -826,7 +907,8 @@ def indices_of_flag(self, start_time, duration, times, padding=0): return idx def advance(self, blocksize): - """ Add blocksize seconds more to the buffer, push blocksize seconds + """ + Add blocksize seconds more to the buffer, push blocksize seconds from the beginning. Parameters @@ -839,6 +921,7 @@ def advance(self, blocksize): status: boolean Returns True if all of the status information if valid, False if any is not. + """ try: if self.increment_update_cache: @@ -849,18 +932,21 @@ def advance(self, blocksize): self.null_advance(blocksize) return False -class iDQBuffer(object): - - """ Read iDQ timeseries from a frame file """ - def __init__(self, frame_src, - idq_channel_name, - idq_status_channel_name, - idq_threshold, - start_time, - max_buffer=512, - force_update_cache=False, - increment_update_cache=None): +class iDQBuffer: + """Read iDQ timeseries from a frame file""" + + def __init__( + self, + frame_src, + idq_channel_name, + idq_status_channel_name, + idq_threshold, + start_time, + max_buffer=512, + force_update_cache=False, + increment_update_cache=None, + ): """ Parameters ---------- @@ -885,19 +971,29 @@ def __init__(self, frame_src, is an alternate to the forced updated of the frame cache, and apptempts to predict the next frame file name without probing the filesystem. + """ self.threshold = idq_threshold - self.idq = DataBuffer(frame_src, idq_channel_name, start_time, - max_buffer=max_buffer, - force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache) - self.idq_state = DataBuffer(frame_src, idq_status_channel_name, start_time, - max_buffer=max_buffer, - force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache) + self.idq = DataBuffer( + frame_src, + idq_channel_name, + start_time, + max_buffer=max_buffer, + force_update_cache=force_update_cache, + increment_update_cache=increment_update_cache, + ) + self.idq_state = DataBuffer( + frame_src, + idq_status_channel_name, + start_time, + max_buffer=max_buffer, + force_update_cache=force_update_cache, + increment_update_cache=increment_update_cache, + ) def flag_at_times(self, start_time, duration, times, padding=0): - """ Check whether the idq flag was on at given times + """ + Check whether the idq flag was on at given times Parameters ---------- @@ -915,6 +1011,7 @@ def flag_at_times(self, start_time, duration, times, padding=0): ------- flag_state: numpy.ndarray Boolean array of whether flag was on at given times + """ from pycbc.events.veto import indices_within_times @@ -946,7 +1043,8 @@ def flag_at_times(self, start_time, duration, times, padding=0): return flagged_bool def advance(self, blocksize): - """ Add blocksize seconds more to the buffer, push blocksize seconds + """ + Add blocksize seconds more to the buffer, push blocksize seconds from the beginning. Parameters @@ -959,30 +1057,33 @@ def advance(self, blocksize): status: boolean Returns True if advance is succesful, False if not. + """ idq_ts = self.idq.attempt_advance(blocksize) idq_state_ts = self.idq_state.attempt_advance(blocksize) return (idq_ts is not None) and (idq_state_ts is not None) def null_advance(self, blocksize): - """Advance and insert zeros + """ + Advance and insert zeros Parameters ---------- blocksize: int The number of seconds to advance the buffers + """ self.idq.null_advance(blocksize) self.idq_state.null_advance(blocksize) __all__ = [ - 'locations_to_cache', - 'read_frame', - 'query_and_read_frame', - 'frame_paths', - 'write_frame', - 'DataBuffer', - 'StatusBuffer', - 'iDQBuffer' + "DataBuffer", + "StatusBuffer", + "frame_paths", + "iDQBuffer", + "locations_to_cache", + "query_and_read_frame", + "read_frame", + "write_frame", ] diff --git a/pycbc/frame/gwosc.py b/pycbc/frame/gwosc.py index 499a495eac7..f003f171dad 100644 --- a/pycbc/frame/gwosc.py +++ b/pycbc/frame/gwosc.py @@ -17,19 +17,21 @@ This modules contains functions for getting data from the Gravitational Wave Open Science Center (GWOSC). """ -import logging + import json +import logging -from pycbc.io import get_file from pycbc.frame import read_frame +from pycbc.io import get_file -logger = logging.getLogger('pycbc.frame.gwosc') +logger = logging.getLogger("pycbc.frame.gwosc") _GWOSC_URL = "https://www.gwosc.org/archive/links/%s/%s/%s/%s/json/" def get_run(time, ifo=None): - """Return the run name for a given time. + """ + Return the run name for a given time. Parameters ---------- @@ -40,37 +42,39 @@ def get_run(time, ifo=None): except for some special times where data releases were made for a single detector under unusual circumstances. For example, to get the data around GW170608 in the Hanford detector. + """ cases = [ ( # ifo is only needed in this special case, otherwise, # the run name is the same for all ifos - 1180911618 <= time <= 1180982427 and ifo == 'H1', - 'BKGW170608_16KHZ_R1' + 1180911618 <= time <= 1180982427 and ifo == "H1", + "BKGW170608_16KHZ_R1", ), - (1396417050 <= time <= 1422118818, 'O4b_16KHZ_R1'), - (1368195220 <= time <= 1389456018, 'O4a_16KHZ_R1'), - (1253977219 <= time <= 1320363336, 'O3b_16KHZ_R1'), - (1238166018 <= time <= 1253977218, 'O3a_16KHZ_R1'), - (1164556817 <= time <= 1187733618, 'O2_16KHZ_R1'), - (1126051217 <= time <= 1137254417, 'O1'), - (815011213 <= time <= 875318414, 'S5'), - (930787215 <= time <= 971568015, 'S6') + (1396417050 <= time <= 1422118818, "O4b_16KHZ_R1"), + (1368195220 <= time <= 1389456018, "O4a_16KHZ_R1"), + (1253977219 <= time <= 1320363336, "O3b_16KHZ_R1"), + (1238166018 <= time <= 1253977218, "O3a_16KHZ_R1"), + (1164556817 <= time <= 1187733618, "O2_16KHZ_R1"), + (1126051217 <= time <= 1137254417, "O1"), + (815011213 <= time <= 875318414, "S5"), + (930787215 <= time <= 971568015, "S6"), ] for condition, name in cases: if condition: return name - raise ValueError(f'Time {time} not available in a public dataset') + raise ValueError(f"Time {time} not available in a public dataset") def _get_channel(time): if time < 1164556817: - return 'LOSC-STRAIN' - return 'GWOSC-16KHZ_R1_STRAIN' + return "LOSC-STRAIN" + return "GWOSC-16KHZ_R1_STRAIN" def gwosc_frame_json(ifo, start_time, end_time): - """Get the information about the public data files in a duration of time. + """ + Get the information about the public data files in a duration of time. Parameters ---------- @@ -86,27 +90,31 @@ def gwosc_frame_json(ifo, start_time, end_time): info: dict A dictionary containing information about the files that span the requested times. + """ run = get_run(start_time) run2 = get_run(end_time) if run != run2: raise ValueError( - 'Spanning multiple runs is not currently supported. ' - f'You have requested data that uses both {run} and {run2}' + "Spanning multiple runs is not currently supported. " + f"You have requested data that uses both {run} and {run2}" ) url = _GWOSC_URL % (run, ifo, int(start_time), int(end_time)) try: - return json.load(open(get_file(url, cache=False), 'r')) + return json.load(open(get_file(url, cache=False))) except Exception as exc: - msg = ('Failed to find gwf files for ' - f'ifo={ifo}, run={run}, between {start_time}-{end_time}') + msg = ( + "Failed to find gwf files for " + f"ifo={ifo}, run={run}, between {start_time}-{end_time}" + ) raise ValueError(msg) from exc def gwosc_frame_urls(ifo, start_time, end_time): - """Get a list of URLs to GWOSC frame files. + """ + Get a list of URLs to GWOSC frame files. Parameters ---------- @@ -122,13 +130,15 @@ def gwosc_frame_urls(ifo, start_time, end_time): frame_files: list A dictionary containing information about the files that span the requested times. + """ - data = gwosc_frame_json(ifo, start_time, end_time)['strain'] - return [d['url'] for d in data if d['format'] == 'gwf'] + data = gwosc_frame_json(ifo, start_time, end_time)["strain"] + return [d["url"] for d in data if d["format"] == "gwf"] def read_frame_gwosc(channels, start_time, end_time): - """Read channels from GWOSC data. + """ + Read channels from GWOSC data. Parameters ---------- @@ -143,6 +153,7 @@ def read_frame_gwosc(channels, start_time, end_time): ------- ts: TimeSeries Returns a timeseries or list of timeseries with the requested data. + """ if not isinstance(channels, list): channels = [channels] @@ -151,8 +162,9 @@ def read_frame_gwosc(channels, start_time, end_time): for ifo in ifos: urls[ifo] = gwosc_frame_urls(ifo, start_time, end_time) if len(urls[ifo]) == 0: - raise ValueError("No data found for %s so we " - "can't produce a time series" % ifo) + raise ValueError( + "No data found for %s so we can't produce a time series" % ifo + ) fnames = {ifo: [] for ifo in ifos} for ifo in ifos: @@ -160,16 +172,20 @@ def read_frame_gwosc(channels, start_time, end_time): fname = get_file(url, cache=True) fnames[ifo].append(fname) - ts_list = [read_frame(fnames[channel[0:2]], channel, - start_time=start_time, end_time=end_time) - for channel in channels] + ts_list = [ + read_frame( + fnames[channel[0:2]], channel, start_time=start_time, end_time=end_time + ) + for channel in channels + ] if len(ts_list) == 1: return ts_list[0] return ts_list def read_strain_gwosc(ifo, start_time, end_time): - """Get the strain data from the GWOSC data. + """ + Get the strain data from the GWOSC data. Parameters ---------- @@ -184,6 +200,7 @@ def read_strain_gwosc(ifo, start_time, end_time): ------- ts: TimeSeries Returns a timeseries with the strain data. + """ channel = _get_channel(start_time) - return read_frame_gwosc(f'{ifo}:{channel}', start_time, end_time) + return read_frame_gwosc(f"{ifo}:{channel}", start_time, end_time) diff --git a/pycbc/frame/store.py b/pycbc/frame/store.py index 1965cfeb68c..847b9b7a0ea 100644 --- a/pycbc/frame/store.py +++ b/pycbc/frame/store.py @@ -16,17 +16,20 @@ """ This modules contains functions for reading in data from hdf stores """ + import logging + import numpy -from pycbc.types import TimeSeries from pycbc.io.hdf import HFile +from pycbc.types import TimeSeries -logger = logging.getLogger('pycbc.frame.store') +logger = logging.getLogger("pycbc.frame.store") def read_store(fname, channel, start_time, end_time): - """ Read time series data from hdf store + """ + Read time series data from hdf store Parameters ---------- @@ -45,13 +48,13 @@ def read_store(fname, channel, start_time, end_time): Time series containing the requested data """ - fhandle = HFile(fname, 'r') + fhandle = HFile(fname, "r") if channel not in fhandle: - raise ValueError('Could not find channel name {}'.format(channel)) + raise ValueError(f"Could not find channel name {channel}") # Determine which segment data lies in (can only read contiguous data now) - starts = fhandle[channel]['segments']['start'][:] - ends = fhandle[channel]['segments']['end'][:] + starts = fhandle[channel]["segments"]["start"][:] + ends = fhandle[channel]["segments"]["end"][:] diff = start_time - starts loc = numpy.where(diff >= 0)[0] @@ -61,15 +64,14 @@ def read_store(fname, channel, start_time, end_time): etime = ends[sidx] if stime > start_time: - raise ValueError("Cannot read data segment before {}".format(stime)) + raise ValueError(f"Cannot read data segment before {stime}") if etime < end_time: - raise ValueError("Cannot read data segment past {}".format(etime)) + raise ValueError(f"Cannot read data segment past {etime}") data = fhandle[channel][str(sidx)] sample_rate = len(data) / (etime - stime) start = int((start_time - stime) * sample_rate) end = int((end_time - stime) * sample_rate) - return TimeSeries(data[start:end], delta_t=1.0/sample_rate, - epoch=start_time) + return TimeSeries(data[start:end], delta_t=1.0 / sample_rate, epoch=start_time) diff --git a/pycbc/inference/__init__.py b/pycbc/inference/__init__.py index 55c3d58ec84..065b872b76f 100644 --- a/pycbc/inference/__init__.py +++ b/pycbc/inference/__init__.py @@ -1,3 +1,2 @@ # pylint: disable=unused-import -from . import (models, sampler, io) -from . import (burn_in, entropy, gelman_rubin, geweke, option_utils) +from . import burn_in, entropy, gelman_rubin, geweke, io, models, option_utils, sampler diff --git a/pycbc/inference/burn_in.py b/pycbc/inference/burn_in.py index 83bea886048..34d3da341d4 100644 --- a/pycbc/inference/burn_in.py +++ b/pycbc/inference/burn_in.py @@ -26,15 +26,15 @@ have burned in. """ - import logging from abc import ABCMeta, abstractmethod + import numpy from scipy.stats import ks_2samp from pycbc.io.record import get_vars_from_arg -logger = logging.getLogger('pycbc.inference.burn_in') +logger = logging.getLogger("pycbc.inference.burn_in") # The value to use for a burn-in iteration if a chain is not burned in NOT_BURNED_IN_ITER = -1 @@ -50,7 +50,8 @@ def ks_test(samples1, samples2, threshold=0.9): - """Applies a KS test to determine if two sets of samples are the same. + """ + Applies a KS test to determine if two sets of samples are the same. The ks test is applied parameter-by-parameter. If the two-tailed p-value returned by the test is greater than ``threshold``, the samples are @@ -70,10 +71,12 @@ def ks_test(samples1, samples2, threshold=0.9): dict : Dictionary mapping parameter names to booleans indicating whether the given parameter passes the KS test. + """ is_the_same = {} assert set(samples1.keys()) == set(samples2.keys()), ( - "samples1 and 2 must have the same parameters") + "samples1 and 2 must have the same parameters" + ) # iterate over the parameters for param in samples1: s1 = samples1[param] @@ -84,7 +87,8 @@ def ks_test(samples1, samples2, threshold=0.9): def max_posterior(lnps_per_walker, dim): - """Burn in based on samples being within dim/2 of maximum posterior. + """ + Burn in based on samples being within dim/2 of maximum posterior. Parameters ---------- @@ -101,13 +105,13 @@ def max_posterior(lnps_per_walker, dim): index will be be equal to the length of the chain. is_burned_in : array of bool Whether or not a walker is burned in. + """ if len(lnps_per_walker.shape) != 2: - raise ValueError("lnps_per_walker must have shape " - "nwalkers x niterations") + raise ValueError("lnps_per_walker must have shape nwalkers x niterations") # find the value to compare against max_p = lnps_per_walker.max() - criteria = max_p - dim/2. + criteria = max_p - dim / 2.0 nwalkers, _ = lnps_per_walker.shape burn_in_idx = numpy.empty(nwalkers, dtype=int) is_burned_in = numpy.empty(nwalkers, dtype=bool) @@ -125,7 +129,8 @@ def max_posterior(lnps_per_walker, dim): def posterior_step(logposts, dim): - """Finds the last time a chain made a jump > dim/2. + """ + Finds the last time a chain made a jump > dim/2. Parameters ---------- @@ -139,10 +144,11 @@ def posterior_step(logposts, dim): int The index of the last time the logpost made a jump > dim/2. If that never happened, returns 0. + """ if logposts.ndim > 1: raise ValueError("logposts must be a 1D array") - criteria = dim/2. + criteria = dim / 2.0 dp = numpy.diff(logposts) indices = numpy.where(dp >= criteria)[0] if indices.size > 0: @@ -153,7 +159,8 @@ def posterior_step(logposts, dim): def nacl(nsamples, acls, nacls=5): - """Burn in based on ACL. + """ + Burn in based on ACL. This applies the following test to determine burn in: @@ -181,13 +188,15 @@ def nacl(nsamples, acls, nacls=5): Dictionary of parameter -> boolean(s) indicating if the chain(s) pass the test. If an array of values was provided for the acls, the values will be arrays of booleans. + """ - kstart = int(nsamples / 2.) + kstart = int(nsamples / 2.0) return {param: (nacls * acl) < kstart for (param, acl) in acls.items()} def evaluate_tests(burn_in_test, test_is_burned_in, test_burn_in_iter): - """Evaluates burn in data from multiple tests. + """ + Evaluates burn in data from multiple tests. The iteration to use for burn-in depends on the logic in the burn-in test string. For example, if the test was 'max_posterior | nacl' and @@ -221,15 +230,16 @@ def evaluate_tests(burn_in_test, test_is_burned_in, test_burn_in_iter): The iteration at which all the tests pass. If the tests did not all pass (``is_burned_in`` is false), then returns :py:data:`NOT_BURNED_IN_ITER`. + """ burn_in_iters = numpy.unique(list(test_burn_in_iter.values())) burn_in_iters.sort() for ii in burn_in_iters: - test_results = {t: (test_is_burned_in[t] & - 0 <= test_burn_in_iter[t] <= ii) - for t in test_is_burned_in} - is_burned_in = eval(burn_in_test, {"__builtins__": None}, - test_results) + test_results = { + t: (test_is_burned_in[t] & 0 <= test_burn_in_iter[t] <= ii) + for t in test_is_burned_in + } + is_burned_in = eval(burn_in_test, {"__builtins__": None}, test_results) if is_burned_in: break if not is_burned_in: @@ -249,9 +259,13 @@ def evaluate_tests(burn_in_test, test_is_burned_in, test_burn_in_iter): class BaseBurnInTests(metaclass=ABCMeta): """Base class for burn in tests.""" - available_tests = ('halfchain', 'min_iterations', 'max_posterior', - 'posterior_step', 'nacl', - ) + available_tests = ( + "halfchain", + "min_iterations", + "max_posterior", + "posterior_step", + "nacl", + ) # pylint: disable=unnecessary-pass @@ -267,27 +281,28 @@ def __init__(self, sampler, burn_in_test, **kwargs): self.test_aux_info = {} # any additional information the test stores # Arguments specific to each test... # for nacl: - self._nacls = int(kwargs.pop('nacls', 5)) + self._nacls = int(kwargs.pop("nacls", 5)) # for max_posterior and posterior_step - self._ndim = int(kwargs.pop('ndim', len(sampler.variable_params))) + self._ndim = int(kwargs.pop("ndim", len(sampler.variable_params))) # for min iterations - self._min_iterations = int(kwargs.pop('min_iterations', 0)) + self._min_iterations = int(kwargs.pop("min_iterations", 0)) @abstractmethod def burn_in_index(self, filename): - """The burn in index (retrieved from the iteration). + """ + The burn in index (retrieved from the iteration). This is an abstract method because how this is evaluated depends on if this is an ensemble MCMC or not. """ - pass def _getniters(self, filename): - """Convenience function to get the number of iterations in the file. + """ + Convenience function to get the number of iterations in the file. If `niterations` hasn't been written to the file yet, just returns 0. """ - with self.sampler.io(filename, 'r') as fp: + with self.sampler.io(filename, "r") as fp: try: niters = fp.niterations except KeyError: @@ -295,11 +310,12 @@ def _getniters(self, filename): return niters def _getnsamples(self, filename): - """Convenience function to get the number of samples saved in the file. + """ + Convenience function to get the number of samples saved in the file. If no samples have been written to the file yet, just returns 0. """ - with self.sampler.io(filename, 'r') as fp: + with self.sampler.io(filename, "r") as fp: try: group = fp[fp.samples_group] # we'll just use the first parameter @@ -310,22 +326,23 @@ def _getnsamples(self, filename): return nsamples def _index2iter(self, filename, index): - """Converts the index in some samples at which burn in occurs to the + """ + Converts the index in some samples at which burn in occurs to the iteration of the sampler that corresponds to. """ - with self.sampler.io(filename, 'r') as fp: + with self.sampler.io(filename, "r") as fp: thin_interval = fp.thinned_by return index * thin_interval def _iter2index(self, filename, iteration): - """Converts an iteration to the index it corresponds to. - """ - with self.sampler.io(filename, 'r') as fp: + """Converts an iteration to the index it corresponds to.""" + with self.sampler.io(filename, "r") as fp: thin_interval = fp.thinned_by return iteration // thin_interval def _getlogposts(self, filename): - """Convenience function for retrieving log posteriors. + """ + Convenience function for retrieving log posteriors. Parameters ---------- @@ -337,21 +354,25 @@ def _getlogposts(self, filename): array The log posterior values. They are not flattened, so have dimension nwalkers x niterations. + """ - with self.sampler.io(filename, 'r') as fp: + with self.sampler.io(filename, "r") as fp: samples = fp.read_raw_samples( - ['loglikelihood', 'logprior'], thin_start=0, thin_interval=1, - flatten=False) - logposts = samples['loglikelihood'] + samples['logprior'] + ["loglikelihood", "logprior"], + thin_start=0, + thin_interval=1, + flatten=False, + ) + logposts = samples["loglikelihood"] + samples["logprior"] return logposts def _getacls(self, filename, start_index): - """Convenience function for calculating acls for the given filename. - """ + """Convenience function for calculating acls for the given filename.""" return self.sampler.compute_acl(filename, start_index=start_index) def _getaux(self, test): - """Convenience function for getting auxilary information. + """ + Convenience function for getting auxilary information. Parameters ---------- @@ -364,6 +385,7 @@ def _getaux(self, test): The ``test_aux_info[test]`` dictionary. If a dictionary does not exist yet for the given test, an empty dictionary will be created and saved to ``test_aux_info[test]``. + """ try: aux = self.test_aux_info[test] @@ -372,16 +394,16 @@ def _getaux(self, test): return aux def halfchain(self, filename): - """Just uses half the chain as the burn-in iteration. - """ + """Just uses half the chain as the burn-in iteration.""" niters = self._getniters(filename) # this test cannot determine when something will burn in # only when it was not burned in in the past - self.test_is_burned_in['halfchain'] = True - self.test_burn_in_iteration['halfchain'] = niters//2 + self.test_is_burned_in["halfchain"] = True + self.test_burn_in_iteration["halfchain"] = niters // 2 def min_iterations(self, filename): - """Just checks that the sampler has been run for the minimum number + """ + Just checks that the sampler has been run for the minimum number of iterations. """ niters = self._getniters(filename) @@ -390,33 +412,31 @@ def min_iterations(self, filename): burn_in_iter = self._min_iterations else: burn_in_iter = NOT_BURNED_IN_ITER - self.test_is_burned_in['min_iterations'] = is_burned_in - self.test_burn_in_iteration['min_iterations'] = burn_in_iter + self.test_is_burned_in["min_iterations"] = is_burned_in + self.test_burn_in_iteration["min_iterations"] = burn_in_iter @abstractmethod def max_posterior(self, filename): """Carries out the max posterior test and stores the results.""" - pass @abstractmethod def posterior_step(self, filename): """Carries out the posterior step test and stores the results.""" - pass @abstractmethod def nacl(self, filename): """Carries out the nacl test and stores the results.""" - pass @abstractmethod def evaluate(self, filename): - """Performs all tests and evaluates the results to determine if and + """ + Performs all tests and evaluates the results to determine if and when all tests pass. """ - pass def write(self, fp, path=None): - """Writes burn-in info to an open HDF file. + """ + Writes burn-in info to an open HDF file. Parameters ---------- @@ -427,20 +447,21 @@ def write(self, fp, path=None): Path in the HDF file to write the data to. Default is (None) is to write to the path given by the file's ``sampler_group`` attribute. + """ if path is None: path = fp.sampler_group - fp.write_data('burn_in_test', self.burn_in_test, path) - fp.write_data('is_burned_in', self.is_burned_in, path) - fp.write_data('burn_in_iteration', self.burn_in_iteration, path) - testgroup = 'burn_in_tests' + fp.write_data("burn_in_test", self.burn_in_test, path) + fp.write_data("is_burned_in", self.is_burned_in, path) + fp.write_data("burn_in_iteration", self.burn_in_iteration, path) + testgroup = "burn_in_tests" # write individual test data for tst in self.do_tests: - subpath = '/'.join([path, testgroup, tst]) - fp.write_data('is_burned_in', self.test_is_burned_in[tst], subpath) - fp.write_data('burn_in_iteration', - self.test_burn_in_iteration[tst], - subpath) + subpath = "/".join([path, testgroup, tst]) + fp.write_data("is_burned_in", self.test_is_burned_in[tst], subpath) + fp.write_data( + "burn_in_iteration", self.test_burn_in_iteration[tst], subpath + ) # write auxiliary info if tst in self.test_aux_info: for name, data in self.test_aux_info[tst].items(): @@ -455,25 +476,26 @@ def _extra_tests_from_config(cp, section, tag): @classmethod def from_config(cls, cp, sampler): """Loads burn in from section [sampler-burn_in].""" - section = 'sampler' - tag = 'burn_in' - burn_in_test = cp.get_opt_tag(section, 'burn-in-test', tag) + section = "sampler" + tag = "burn_in" + burn_in_test = cp.get_opt_tag(section, "burn-in-test", tag) kwargs = {} - if cp.has_option_tag(section, 'nacl', tag): - kwargs['nacl'] = int(cp.get_opt_tag(section, 'nacl', tag)) - if cp.has_option_tag(section, 'ndim', tag): - kwargs['ndim'] = int( - cp.get_opt_tag(section, 'ndim', tag)) - if cp.has_option_tag(section, 'min-iterations', tag): - kwargs['min_iterations'] = int( - cp.get_opt_tag(section, 'min-iterations', tag)) + if cp.has_option_tag(section, "nacl", tag): + kwargs["nacl"] = int(cp.get_opt_tag(section, "nacl", tag)) + if cp.has_option_tag(section, "ndim", tag): + kwargs["ndim"] = int(cp.get_opt_tag(section, "ndim", tag)) + if cp.has_option_tag(section, "min-iterations", tag): + kwargs["min_iterations"] = int( + cp.get_opt_tag(section, "min-iterations", tag) + ) # load any class specific tests kwargs.update(cls._extra_tests_from_config(cp, section, tag)) return cls(sampler, burn_in_test, **kwargs) class MCMCBurnInTests(BaseBurnInTests): - """Burn-in tests for collections of independent MCMC chains. + """ + Burn-in tests for collections of independent MCMC chains. This differs from EnsembleMCMCBurnInTests in that chains are treated as being independent of each other. The ``is_burned_in`` attribute will be @@ -481,8 +503,9 @@ class MCMCBurnInTests(BaseBurnInTests): all chains must pass the burn in tests). In other words, independent samples can be collected even if all of the chains are not burned in. """ + def __init__(self, sampler, burn_in_test, **kwargs): - super(MCMCBurnInTests, self).__init__(sampler, burn_in_test, **kwargs) + super().__init__(sampler, burn_in_test, **kwargs) try: nchains = sampler.nchains except AttributeError: @@ -506,34 +529,33 @@ def max_posterior(self, filename): burn_in_iter = self._index2iter(filename, burn_in_idx) burn_in_iter[~is_burned_in] = NOT_BURNED_IN_ITER # save - test = 'max_posterior' + test = "max_posterior" self.test_is_burned_in[test] = is_burned_in self.test_burn_in_iteration[test] = burn_in_iter def posterior_step(self, filename): """Applies the posterior-step test.""" logposts = self._getlogposts(filename) - burn_in_idx = numpy.array([posterior_step(logps, self._ndim) - for logps in logposts]) + burn_in_idx = numpy.array( + [posterior_step(logps, self._ndim) for logps in logposts] + ) # this test cannot determine when something will burn in # only when it was not burned in in the past - test = 'posterior_step' + test = "posterior_step" if test not in self.test_is_burned_in: self.test_is_burned_in[test] = numpy.ones(self.nchains, dtype=bool) # convert index to iterations - self.test_burn_in_iteration[test] = self._index2iter(filename, - burn_in_idx) + self.test_burn_in_iteration[test] = self._index2iter(filename, burn_in_idx) def nacl(self, filename): """Applies the :py:func:`nacl` test.""" nsamples = self._getnsamples(filename) - acls = self._getacls(filename, start_index=nsamples//2) + acls = self._getacls(filename, start_index=nsamples // 2) is_burned_in = nacl(nsamples, acls, self._nacls) # stack the burn in results into an nparams x nchains array - burn_in_per_chain = numpy.stack(list(is_burned_in.values())).all( - axis=0) + burn_in_per_chain = numpy.stack(list(is_burned_in.values())).all(axis=0) # store - test = 'nacl' + test = "nacl" self.test_is_burned_in[test] = burn_in_per_chain try: burn_in_iter = self.test_burn_in_iteration[test] @@ -541,8 +563,7 @@ def nacl(self, filename): # hasn't been stored yet burn_in_iter = numpy.repeat(NOT_BURNED_IN_ITER, self.nchains) self.test_burn_in_iteration[test] = burn_in_iter - burn_in_iter[burn_in_per_chain] = self._index2iter(filename, - nsamples//2) + burn_in_iter[burn_in_per_chain] = self._index2iter(filename, nsamples // 2) # add the status for each parameter as additional information self.test_aux_info[test] = is_burned_in @@ -556,19 +577,26 @@ def evaluate(self, filename): for ci in range(self.nchains): # some tests (like halfchain) just store a single bool for all # chains - tibi = {t: r[ci] if isinstance(r, numpy.ndarray) else r - for t, r in self.test_is_burned_in.items()} - tbi = {t: r[ci] if isinstance(r, numpy.ndarray) else r - for t, r in self.test_burn_in_iteration.items()} - is_burned_in, burn_in_iter = evaluate_tests(self.burn_in_test, - tibi, tbi) + tibi = { + t: r[ci] if isinstance(r, numpy.ndarray) else r + for t, r in self.test_is_burned_in.items() + } + tbi = { + t: r[ci] if isinstance(r, numpy.ndarray) else r + for t, r in self.test_burn_in_iteration.items() + } + is_burned_in, burn_in_iter = evaluate_tests(self.burn_in_test, tibi, tbi) self.is_burned_in[ci] = is_burned_in self.burn_in_iteration[ci] = burn_in_iter - logger.info("Number of chains burned in: %i of %i", - self.is_burned_in.sum(), self.nchains) + logger.info( + "Number of chains burned in: %i of %i", + self.is_burned_in.sum(), + self.nchains, + ) def write(self, fp, path=None): - """Writes burn-in info to an open HDF file. + """ + Writes burn-in info to an open HDF file. Parameters ---------- @@ -579,21 +607,24 @@ def write(self, fp, path=None): Path in the HDF file to write the data to. Default is (None) is to write to the path given by the file's ``sampler_group`` attribute. + """ if path is None: path = fp.sampler_group - super(MCMCBurnInTests, self).write(fp, path) + super().write(fp, path) # add number of chains burned in as additional metadata - fp.write_data('nchains_burned_in', self.is_burned_in.sum(), path) + fp.write_data("nchains_burned_in", self.is_burned_in.sum(), path) class MultiTemperedMCMCBurnInTests(MCMCBurnInTests): - """Adds support for multiple temperatures to + """ + Adds support for multiple temperatures to :py:class:`MCMCBurnInTests`. """ def _getacls(self, filename, start_index): - """Convenience function for calculating acls for the given filename. + """ + Convenience function for calculating acls for the given filename. This function is used by the ``n_acl`` burn-in test. That function expects the returned ``acls`` dict to just report a single ACL for @@ -614,14 +645,15 @@ def _getacls(self, filename, start_index): ------- dict : Dictionary of parameter names -> array giving ACL for each chain. + """ - acls = super(MultiTemperedMCMCBurnInTests, self)._getacls( - filename, start_index) + acls = super()._getacls(filename, start_index) # acls will have shape ntemps x nchains, flatten to nchains return {param: vals.max(axis=0) for (param, vals) in acls.items()} def _getlogposts(self, filename): - """Convenience function for retrieving log posteriors. + """ + Convenience function for retrieving log posteriors. This just gets the coldest temperature chain, and returns arrays with shape nwalkers x niterations, so the parent class can run the same @@ -633,15 +665,19 @@ def _getlogposts(self, filename): class EnsembleMCMCBurnInTests(BaseBurnInTests): """Provides methods for estimating burn-in of an ensemble MCMC.""" - available_tests = ('halfchain', 'min_iterations', 'max_posterior', - 'posterior_step', 'nacl', 'ks_test', - ) + available_tests = ( + "halfchain", + "min_iterations", + "max_posterior", + "posterior_step", + "nacl", + "ks_test", + ) def __init__(self, sampler, burn_in_test, **kwargs): - super(EnsembleMCMCBurnInTests, self).__init__( - sampler, burn_in_test, **kwargs) + super().__init__(sampler, burn_in_test, **kwargs) # for kstest - self._ksthreshold = float(kwargs.pop('ks_threshold', 0.9)) + self._ksthreshold = float(kwargs.pop("ks_threshold", 0.9)) def burn_in_index(self, filename): """The burn in index (retrieved from the iteration).""" @@ -661,73 +697,74 @@ def max_posterior(self, filename): else: burn_in_iter = NOT_BURNED_IN_ITER # store - test = 'max_posterior' + test = "max_posterior" self.test_is_burned_in[test] = all_burned_in self.test_burn_in_iteration[test] = burn_in_iter aux = self._getaux(test) # additional info - aux['iteration_per_walker'] = self._index2iter(filename, burn_in_idx) - aux['status_per_walker'] = is_burned_in + aux["iteration_per_walker"] = self._index2iter(filename, burn_in_idx) + aux["status_per_walker"] = is_burned_in def posterior_step(self, filename): """Applies the posterior-step test.""" logposts = self._getlogposts(filename) - burn_in_idx = numpy.array([posterior_step(logps, self._ndim) - for logps in logposts]) + burn_in_idx = numpy.array( + [posterior_step(logps, self._ndim) for logps in logposts] + ) burn_in_iters = self._index2iter(filename, burn_in_idx) # this test cannot determine when something will burn in # only when it was not burned in in the past - test = 'posterior_step' + test = "posterior_step" self.test_is_burned_in[test] = True self.test_burn_in_iteration[test] = burn_in_iters.max() # store the iteration per walker as additional info aux = self._getaux(test) - aux['iteration_per_walker'] = burn_in_iters + aux["iteration_per_walker"] = burn_in_iters def nacl(self, filename): """Applies the :py:func:`nacl` test.""" nsamples = self._getnsamples(filename) - acls = self._getacls(filename, start_index=nsamples//2) + acls = self._getacls(filename, start_index=nsamples // 2) is_burned_in = nacl(nsamples, acls, self._nacls) all_burned_in = all(is_burned_in.values()) if all_burned_in: - burn_in_iter = self._index2iter(filename, nsamples//2) + burn_in_iter = self._index2iter(filename, nsamples // 2) else: burn_in_iter = NOT_BURNED_IN_ITER # store - test = 'nacl' + test = "nacl" self.test_is_burned_in[test] = all_burned_in self.test_burn_in_iteration[test] = burn_in_iter # store the status per parameter as additional info aux = self._getaux(test) - aux['status_per_parameter'] = is_burned_in + aux["status_per_parameter"] = is_burned_in def ks_test(self, filename): """Applies ks burn-in test.""" nsamples = self._getnsamples(filename) - with self.sampler.io(filename, 'r') as fp: + with self.sampler.io(filename, "r") as fp: # get the samples from the mid point samples1 = fp.read_raw_samples( - ['loglikelihood', 'logprior'], iteration=int(nsamples/2.)) + ["loglikelihood", "logprior"], iteration=int(nsamples / 2.0) + ) # get the last samples - samples2 = fp.read_raw_samples( - ['loglikelihood', 'logprior'], iteration=-1) + samples2 = fp.read_raw_samples(["loglikelihood", "logprior"], iteration=-1) # do the test # is_the_same is a dictionary of params --> bool indicating whether or # not the 1D marginal is the same at the half way point is_the_same = ks_test(samples1, samples2, threshold=self._ksthreshold) is_burned_in = all(is_the_same.values()) if is_burned_in: - burn_in_iter = self._index2iter(filename, int(nsamples//2)) + burn_in_iter = self._index2iter(filename, int(nsamples // 2)) else: burn_in_iter = NOT_BURNED_IN_ITER # store - test = 'ks_test' + test = "ks_test" self.test_is_burned_in[test] = is_burned_in self.test_burn_in_iteration[test] = burn_in_iter # store the test per parameter as additional info aux = self._getaux(test) - aux['status_per_parameter'] = is_the_same + aux["status_per_parameter"] = is_the_same def evaluate(self, filename): """Runs all of the burn-in tests.""" @@ -736,32 +773,32 @@ def evaluate(self, filename): logger.info("Evaluating %s burn-in test", tst) getattr(self, tst)(filename) is_burned_in, burn_in_iter = evaluate_tests( - self.burn_in_test, self.test_is_burned_in, - self.test_burn_in_iteration) + self.burn_in_test, self.test_is_burned_in, self.test_burn_in_iteration + ) self.is_burned_in = is_burned_in self.burn_in_iteration = burn_in_iter logger.info("Is burned in: %r", self.is_burned_in) if self.is_burned_in: - logger.info("Burn-in iteration: %i", - int(self.burn_in_iteration)) + logger.info("Burn-in iteration: %i", int(self.burn_in_iteration)) @staticmethod def _extra_tests_from_config(cp, section, tag): """Loads the ks test settings from the config file.""" kwargs = {} - if cp.has_option_tag(section, 'ks-threshold', tag): - kwargs['ks_threshold'] = float( - cp.get_opt_tag(section, 'ks-threshold', tag)) + if cp.has_option_tag(section, "ks-threshold", tag): + kwargs["ks_threshold"] = float(cp.get_opt_tag(section, "ks-threshold", tag)) return kwargs class EnsembleMultiTemperedMCMCBurnInTests(EnsembleMCMCBurnInTests): - """Adds support for multiple temperatures to + """ + Adds support for multiple temperatures to :py:class:`EnsembleMCMCBurnInTests`. """ def _getacls(self, filename, start_index): - """Convenience function for calculating acls for the given filename. + """ + Convenience function for calculating acls for the given filename. This function is used by the ``n_acl`` burn-in test. That function expects the returned ``acls`` dict to just report a single ACL for @@ -771,13 +808,15 @@ def _getacls(self, filename, start_index): Since we calculate the acls, this will also store it to the sampler. """ - acls = super(EnsembleMultiTemperedMCMCBurnInTests, self)._getacls( - filename, start_index) + acls = super()._getacls( + filename, start_index + ) # return the max for each parameter return {param: vals.max() for (param, vals) in acls.items()} def _getlogposts(self, filename): - """Convenience function for retrieving log posteriors. + """ + Convenience function for retrieving log posteriors. This just gets the coldest temperature chain, and returns arrays with shape nwalkers x niterations, so the parent class can run the same @@ -788,13 +827,17 @@ def _getlogposts(self, filename): def _multitemper_getlogposts(sampler, filename): """Retrieve log posteriors for multi tempered samplers.""" - with sampler.io(filename, 'r') as fp: + with sampler.io(filename, "r") as fp: samples = fp.read_raw_samples( - ['loglikelihood', 'logprior'], thin_start=0, thin_interval=1, - temps=0, flatten=False) + ["loglikelihood", "logprior"], + thin_start=0, + thin_interval=1, + temps=0, + flatten=False, + ) # reshape to drop the first dimension - for (stat, arr) in samples.items(): + for stat, arr in samples.items(): _, nwalkers, niterations = arr.shape samples[stat] = arr.reshape((nwalkers, niterations)) - logposts = samples['loglikelihood'] + samples['logprior'] + logposts = samples["loglikelihood"] + samples["logprior"] return logposts diff --git a/pycbc/inference/entropy.py b/pycbc/inference/entropy.py index 9f71183fd55..bb1024bcc1c 100644 --- a/pycbc/inference/entropy.py +++ b/pycbc/inference/entropy.py @@ -1,4 +1,5 @@ -""" The module contains functions for calculating the +""" +The module contains functions for calculating the Kullback-Leibler divergence. """ @@ -7,7 +8,8 @@ def check_hist_params(samples, hist_min, hist_max, hist_bins): - """ Checks that the bound values given for the histogram are consistent, + """ + Checks that the bound values given for the histogram are consistent, returning the range if they are or raising an error if they are not. Also checks that if hist_bins is a str, it corresponds to a method available in numpy.histogram @@ -33,15 +35,15 @@ def check_hist_params(samples, hist_min, hist_max, hist_bins): The bounds (hist_min, hist_max) or None. hist_bins : int or str Number of bins or method for optimal width bin calculation. - """ - hist_methods = ['auto', 'fd', 'doane', 'scott', 'stone', 'rice', - 'sturges', 'sqrt'] + """ + hist_methods = ["auto", "fd", "doane", "scott", "stone", "rice", "sturges", "sqrt"] if not hist_bins: - hist_bins = 'fd' + hist_bins = "fd" elif isinstance(hist_bins, str) and hist_bins not in hist_methods: - raise ValueError('Method for calculating bins width must be one of' - ' {}'.format(hist_methods)) + raise ValueError( + f"Method for calculating bins width must be one of {hist_methods}" + ) # No bounds given, return None if not hist_min and not hist_max: @@ -54,7 +56,7 @@ def check_hist_params(samples, hist_min, hist_max, hist_bins): hist_min = samples.min() # Both bounds given elif hist_min and hist_max and hist_min >= hist_max: - raise ValueError('hist_min must be lower than hist_max.') + raise ValueError("hist_min must be lower than hist_max.") hist_range = (hist_min, hist_max) @@ -62,7 +64,8 @@ def check_hist_params(samples, hist_min, hist_max, hist_bins): def compute_pdf(samples, method, bins, hist_min, hist_max): - """ Computes the probability density function for a set of samples. + """ + Computes the probability density function for a set of samples. Parameters ---------- @@ -90,26 +93,27 @@ def compute_pdf(samples, method, bins, hist_min, hist_max): ------- pdf : numpy.array Discrete probability distribution calculated from samples. - """ - if method == 'kde': + """ + if method == "kde": samples_kde = stats.gaussian_kde(samples) - npts = 10000 if len(samples) <= 10000 else len(samples) + npts = max(10000, len(samples)) draw = samples_kde.resample(npts) pdf = samples_kde.evaluate(draw) - elif method == 'hist': - hist_range, hist_bins = check_hist_params(samples, hist_min, - hist_max, bins) - pdf, _ = numpy.histogram(samples, bins=hist_bins, - range=hist_range, density=True) + elif method == "hist": + hist_range, hist_bins = check_hist_params(samples, hist_min, hist_max, bins) + pdf, _ = numpy.histogram( + samples, bins=hist_bins, range=hist_range, density=True + ) else: - raise ValueError('Method not recognized.') + raise ValueError("Method not recognized.") return pdf def entropy(pdf1, base=numpy.e): - """ Computes the information entropy for a single parameter + """ + Computes the information entropy for a single parameter from one probability density function. Parameters @@ -124,14 +128,24 @@ def entropy(pdf1, base=numpy.e): ------- numpy.float64 The information entropy value. - """ + """ return stats.entropy(pdf1, base=base) -def kl(samples1, samples2, pdf1=False, pdf2=False, kde=False, - bins=None, hist_min=None, hist_max=None, base=numpy.e): - """ Computes the Kullback-Leibler divergence for a single parameter +def kl( + samples1, + samples2, + pdf1=False, + pdf2=False, + kde=False, + bins=None, + hist_min=None, + hist_max=None, + base=numpy.e, +): + """ + Computes the Kullback-Leibler divergence for a single parameter from two distributions. Parameters @@ -171,27 +185,31 @@ def kl(samples1, samples2, pdf1=False, pdf2=False, kde=False, ------- numpy.float64 The Kullback-Leibler divergence value. + """ if pdf1 and pdf2 and kde: - raise ValueError('KDE can only be used when at least one of pdf1 or ' - 'pdf2 is False.') + raise ValueError( + "KDE can only be used when at least one of pdf1 or pdf2 is False." + ) - sample_groups = {'P': (samples1, pdf1), 'Q': (samples2, pdf2)} + sample_groups = {"P": (samples1, pdf1), "Q": (samples2, pdf2)} pdfs = {} for n in sample_groups: samples, pdf = sample_groups[n] if pdf: pdfs[n] = samples else: - method = 'kde' if kde else 'hist' + method = "kde" if kde else "hist" pdfs[n] = compute_pdf(samples, method, bins, hist_min, hist_max) - return stats.entropy(pdfs['P'], qk=pdfs['Q'], base=base) + return stats.entropy(pdfs["P"], qk=pdfs["Q"], base=base) -def js(samples1, samples2, kde=False, bins=None, hist_min=None, hist_max=None, - base=numpy.e): - """ Computes the Jensen-Shannon divergence for a single parameter +def js( + samples1, samples2, kde=False, bins=None, hist_min=None, hist_max=None, base=numpy.e +): + """ + Computes the Jensen-Shannon divergence for a single parameter from two distributions. Parameters @@ -224,19 +242,19 @@ def js(samples1, samples2, kde=False, bins=None, hist_min=None, hist_max=None, ------- numpy.float64 The Jensen-Shannon divergence value. - """ - sample_groups = {'P': samples1, 'Q': samples2} + """ + sample_groups = {"P": samples1, "Q": samples2} pdfs = {} for n in sample_groups: samples = sample_groups[n] - method = 'kde' if kde else 'hist' + method = "kde" if kde else "hist" pdfs[n] = compute_pdf(samples, method, bins, hist_min, hist_max) - pdfs['M'] = (1./2) * (pdfs['P'] + pdfs['Q']) + pdfs["M"] = (1.0 / 2) * (pdfs["P"] + pdfs["Q"]) js_div = 0 - for pdf in (pdfs['P'], pdfs['Q']): - js_div += (1./2) * kl(pdf, pdfs['M'], pdf1=True, pdf2=True, base=base) + for pdf in (pdfs["P"], pdfs["Q"]): + js_div += (1.0 / 2) * kl(pdf, pdfs["M"], pdf1=True, pdf2=True, base=base) return js_div diff --git a/pycbc/inference/evidence.py b/pycbc/inference/evidence.py index a89b56689fb..18e91e2b58a 100644 --- a/pycbc/inference/evidence.py +++ b/pycbc/inference/evidence.py @@ -16,6 +16,7 @@ This modules provides functions for estimating the marginal likelihood or evidence of a model. """ + import numpy from scipy import integrate @@ -27,7 +28,8 @@ def arithmetic_mean_estimator(log_likelihood): - """Returns the log evidence via the prior arithmetic mean estimator (AME). + """ + Returns the log evidence via the prior arithmetic mean estimator (AME). The logarithm form of AME is used. This is the most basic evidence estimator, and often requires O(billions) of samples @@ -43,11 +45,12 @@ def arithmetic_mean_estimator(log_likelihood): ------- float : Estimation of the log of the evidence. + """ num_samples = len(log_likelihood) logl_max = numpy.max(log_likelihood) - log_evidence = 0. + log_evidence = 0.0 for i, _ in enumerate(log_likelihood): log_evidence += numpy.exp(log_likelihood[i] - logl_max) @@ -58,7 +61,8 @@ def arithmetic_mean_estimator(log_likelihood): def harmonic_mean_estimator(log_likelihood): - """Returns the log evidence via posterior harmonic mean estimator (HME). + """ + Returns the log evidence via posterior harmonic mean estimator (HME). The logarithm form of HME is used. This method is not recommended for general use. It is very slow to converge, @@ -76,24 +80,25 @@ def harmonic_mean_estimator(log_likelihood): ------- float : Estimation of the log of the evidence. + """ num_samples = len(log_likelihood) - logl_max = numpy.max(-1.0*log_likelihood) + logl_max = numpy.max(-1.0 * log_likelihood) - log_evidence = 0. + log_evidence = 0.0 for i, _ in enumerate(log_likelihood): - log_evidence += numpy.exp(-1.0*log_likelihood[i] + logl_max) + log_evidence += numpy.exp(-1.0 * log_likelihood[i] + logl_max) - log_evidence = -1.0*numpy.log(log_evidence) + log_evidence = -1.0 * numpy.log(log_evidence) log_evidence += logl_max log_evidence += numpy.log(num_samples) return log_evidence -def thermodynamic_integration(log_likelihood, betas, - method="simpsons"): - """Returns the log evidence of the model via thermodynamic integration. +def thermodynamic_integration(log_likelihood, betas, method="simpsons"): + """ + Returns the log evidence of the model via thermodynamic integration. Also returns an estimated standard deviation for the log evidence. Current options are integration through the trapezoid rule, a @@ -125,13 +130,13 @@ def thermodynamic_integration(log_likelihood, betas, mcmc_std : float The standard deviation of the log evidence estimate from Monte-Carlo spread. + """ # Check if the method of integration is in the list of choices method_list = ["trapezoid", "trapezoid_corrected", "simpsons"] if method not in method_list: - raise ValueError("Method %s not supported. Expected %s" - % (method, method_list)) + raise ValueError("Method %s not supported. Expected %s" % (method, method_list)) # Read in the data and ensure ordering of data. # Ascending order sort order = numpy.argsort(betas) @@ -140,9 +145,9 @@ def thermodynamic_integration(log_likelihood, betas, # Assume log likelihood is given in shape of beta, walker, # and iteration. - log_likelihood = numpy.reshape(log_likelihood, - (len(betas), - len(log_likelihood[0].flatten()))) + log_likelihood = numpy.reshape( + log_likelihood, (len(betas), len(log_likelihood[0].flatten())) + ) average_logl = numpy.average(log_likelihood, axis=1) @@ -155,9 +160,9 @@ def thermodynamic_integration(log_likelihood, betas, # https://link.springer.com/article/10.1007/s11222-013-9397-1 var_correction = 0 for i in range(len(betas) - 1): - delta_beta = betas[i+1] - betas[i] - pre_fac_var = (1. / 12.) * (delta_beta ** 2.0) - var_diff = numpy.var(log_likelihood[i+1]) + delta_beta = betas[i + 1] - betas[i] + pre_fac_var = (1.0 / 12.0) * (delta_beta**2.0) + var_diff = numpy.var(log_likelihood[i + 1]) var_diff -= numpy.var(log_likelihood[i]) var_correction -= pre_fac_var * var_diff @@ -170,8 +175,7 @@ def thermodynamic_integration(log_likelihood, betas, # so we can sacrifice precision there, rather than near # beta -> 1. Option even="last" puts trapezoid rule at # first few points. - log_evidence = integrate.simps(average_logl, betas, - even="last") + log_evidence = integrate.simps(average_logl, betas, even="last") # Estimate the Monte Carlo variance of the evidence calculation # See (Evans, Annis, 2019.) @@ -189,8 +193,7 @@ def thermodynamic_integration(log_likelihood, betas, elif method == "simpsons": for i, _ in enumerate(log_likelihood[0]): - ti_vec[i] = integrate.simps(logl_per_samp[i], betas, - even="last") + ti_vec[i] = integrate.simps(logl_per_samp[i], betas, even="last") # Standard error is sample std / sqrt(number of samples) mcmc_std = numpy.std(ti_vec) / numpy.sqrt(float(len(log_likelihood[0]))) @@ -199,7 +202,8 @@ def thermodynamic_integration(log_likelihood, betas, def stepping_stone_algorithm(log_likelihood, betas): - """Returns the log evidence of the model via stepping stone algorithm. + """ + Returns the log evidence of the model via stepping stone algorithm. Also returns an estimated standard deviation for the log evidence. Parameters @@ -219,6 +223,7 @@ def stepping_stone_algorithm(log_likelihood, betas): mcmc_std : float The standard deviation of the log evidence estimate from Monte-Carlo spread. + """ # Reverse order sort order = numpy.argsort(betas)[::-1] @@ -227,17 +232,17 @@ def stepping_stone_algorithm(log_likelihood, betas): # Assume log likelihood is given in shape of beta, # walker, iteration. - log_likelihood = numpy.reshape(log_likelihood, - (len(betas), - len(log_likelihood[0].flatten()))) + log_likelihood = numpy.reshape( + log_likelihood, (len(betas), len(log_likelihood[0].flatten())) + ) log_rk_pb = numpy.zeros(len(betas) - 1) for i in range(len(betas) - 1): - delta_beta = betas[i] - betas[i+1] + delta_beta = betas[i] - betas[i + 1] # Max log likelihood for beta [i+1] - max_logl_pb = numpy.max(log_likelihood[i+1]) + max_logl_pb = numpy.max(log_likelihood[i + 1]) val_1 = delta_beta * max_logl_pb - val_2 = delta_beta * (log_likelihood[i+1] - max_logl_pb) + val_2 = delta_beta * (log_likelihood[i + 1] - max_logl_pb) val_2 = numpy.log(numpy.average(numpy.exp(val_2))) log_rk_pb[i] = val_1 + val_2 @@ -247,10 +252,10 @@ def stepping_stone_algorithm(log_likelihood, betas): # Calculate the Monte Carlo variation mcmc_std = 0 for i in range(len(betas) - 1): - delta_beta = betas[i] - betas[i+1] - pre_fact = (delta_beta * log_likelihood[i+1]) - log_rk_pb[i] + delta_beta = betas[i] - betas[i + 1] + pre_fact = (delta_beta * log_likelihood[i + 1]) - log_rk_pb[i] pre_fact = numpy.exp(pre_fact) - 1.0 - val = numpy.sum(pre_fact ** 2) + val = numpy.sum(pre_fact**2) mcmc_std += val diff --git a/pycbc/inference/gelman_rubin.py b/pycbc/inference/gelman_rubin.py index 9c354f5df3f..0833fb01d02 100644 --- a/pycbc/inference/gelman_rubin.py +++ b/pycbc/inference/gelman_rubin.py @@ -12,7 +12,8 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This modules provides functions for evaluating the Gelman-Rubin convergence +""" +This modules provides functions for evaluating the Gelman-Rubin convergence diagnostic statistic. """ @@ -20,7 +21,8 @@ def walk(chains, start, end, step): - """ Calculates Gelman-Rubin conervergence statistic along chains of data. + """ + Calculates Gelman-Rubin conervergence statistic along chains of data. This function will advance along the chains and calculate the statistic for each step. @@ -45,8 +47,8 @@ def walk(chains, start, end, step): stats : numpy.array Array with convergence statistic. It has shape (nparameters, ncalculations). - """ + """ # get number of chains, parameters, and iterations chains = numpy.array(chains) _, nparameters, _ = chains.shape @@ -67,7 +69,8 @@ def walk(chains, start, end, step): def gelman_rubin(chains, auto_burn_in=True): - """ Calculates the univariate Gelman-Rubin convergence statistic + """ + Calculates the univariate Gelman-Rubin convergence statistic which compares the evolution of multiple chains in a Markov-Chain Monte Carlo process and computes their difference to determine their convergence. The between-chain and within-chain variances are computed for each sampling @@ -88,14 +91,13 @@ def gelman_rubin(chains, auto_burn_in=True): psrf : numpy.array A numpy.array of shape (nparameters) that has the point estimates of the potential scale reduction factor. - """ + """ # remove first half of samples # this will have shape (nchains, nparameters, niterations) if auto_burn_in: _, _, niterations = numpy.array(chains).shape - chains = numpy.array([chain[:, niterations // 2 + 1:] - for chain in chains]) + chains = numpy.array([chain[:, niterations // 2 + 1 :] for chain in chains]) # get number of chains, parameters, and iterations chains = numpy.array(chains) @@ -146,38 +148,37 @@ def gelman_rubin(chains, auto_burn_in=True): # get V the combined variance of all chains # this will have shape (nparameters) - v = ((niterations - 1.) * w_diag / niterations + - (1. + 1. / nchains) * b_diag / niterations) + v = (niterations - 1.0) * w_diag / niterations + ( + 1.0 + 1.0 / nchains + ) * b_diag / niterations # get factors in variance of V calculation # this will have shape (nparameters) k = 2 * b_diag**2 / (nchains - 1) - mid_term = numpy.cov( - var, means**2)[nparameters:2*nparameters, 0:nparameters].T - end_term = numpy.cov( - var, means)[nparameters:2*nparameters, 0:nparameters].T + mid_term = numpy.cov(var, means**2)[nparameters : 2 * nparameters, 0:nparameters].T + end_term = numpy.cov(var, means)[nparameters : 2 * nparameters, 0:nparameters].T wb = niterations / nchains * numpy.diag(mid_term - 2 * mu_hat * end_term) # get variance of V # this will have shape (nparameters) var_v = ( - (niterations - 1.) ** 2 * s + - (1. + 1. / nchains) ** 2 * k + - 2. * (niterations - 1.) * (1. + 1. / nchains) * wb + (niterations - 1.0) ** 2 * s + + (1.0 + 1.0 / nchains) ** 2 * k + + 2.0 * (niterations - 1.0) * (1.0 + 1.0 / nchains) * wb ) / niterations**2 # get degrees of freedom # this will have shape (nparameters) - dof = (2. * v**2) / var_v + dof = (2.0 * v**2) / var_v # more degrees of freedom factors # this will have shape (nparameters) - df_adj = (dof + 3.) / (dof + 1.) + df_adj = (dof + 3.0) / (dof + 1.0) # estimate R # this will have shape (nparameters) - r2_fixed = (niterations - 1.) / niterations - r2_random = (1. + 1. / nchains) * (1. / niterations) * (b_diag / w_diag) + r2_fixed = (niterations - 1.0) / niterations + r2_random = (1.0 + 1.0 / nchains) * (1.0 / niterations) * (b_diag / w_diag) r2_estimate = r2_fixed + r2_random # calculate PSRF the potential scale reduction factor diff --git a/pycbc/inference/geweke.py b/pycbc/inference/geweke.py index 727e29b80b1..b33ed5aa91e 100644 --- a/pycbc/inference/geweke.py +++ b/pycbc/inference/geweke.py @@ -13,15 +13,14 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Functions for computing the Geweke convergence statistic. -""" +"""Functions for computing the Geweke convergence statistic.""" import numpy -def geweke(x, seg_length, seg_stride, end_idx, ref_start, - ref_end=None, seg_start=0): - """ Calculates Geweke conervergence statistic for a chain of data. +def geweke(x, seg_length, seg_stride, end_idx, ref_start, ref_end=None, seg_start=0): + """ + Calculates Geweke conervergence statistic for a chain of data. This function will advance along the chain and calculate the statistic for each step. @@ -52,8 +51,8 @@ def geweke(x, seg_length, seg_stride, end_idx, ref_start, The end index of the first segment in the chain. stats : numpy.array The Geweke convergence diagnostic statistic for the segment. - """ + """ # lists to hold statistic and end index stats = [] ends = [] @@ -66,7 +65,6 @@ def geweke(x, seg_length, seg_stride, end_idx, ref_start, # loop over all segments for start in starts: - # find the end of the first segment x_start_end = int(start + seg_length) @@ -74,8 +72,9 @@ def geweke(x, seg_length, seg_stride, end_idx, ref_start, x_start = x[start:x_start_end] # compute statistic - stats.append((x_start.mean() - x_end.mean()) / numpy.sqrt( - x_start.var() + x_end.var())) + stats.append( + (x_start.mean() - x_end.mean()) / numpy.sqrt(x_start.var() + x_end.var()) + ) # store end of first segment ends.append(x_start_end) diff --git a/pycbc/inference/io/__init__.py b/pycbc/inference/io/__init__.py index 23df1697631..56c4d88ec15 100644 --- a/pycbc/inference/io/__init__.py +++ b/pycbc/inference/io/__init__.py @@ -14,33 +14,33 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""I/O utilities for pycbc inference -""" +"""I/O utilities for pycbc inference""" - -import os import argparse +import logging +import os import shutil import textwrap -import numpy -import logging + import h5py as _h5py -from pycbc.io.record import (FieldArray, _numpy_function_lib) +import numpy + from pycbc import waveform as _waveform -from pycbc.io.hdf import (dump_state, load_state) +from pycbc.inference.option_utils import ParseLabelArg, ParseParametersArg +from pycbc.io.hdf import dump_state, load_state +from pycbc.io.record import FieldArray, _numpy_function_lib -from pycbc.inference.option_utils import (ParseLabelArg, ParseParametersArg) +from .cpnest import CPNestFile +from .dynesty import DynestyFile from .emcee import EmceeFile from .emcee_pt import EmceePTFile -from .ptemcee import PTEmceeFile -from .cpnest import CPNestFile from .multinest import MultinestFile -from .dynesty import DynestyFile -from .ultranest import UltranestFile -from .snowline import SnowlineFile from .nessai import NessaiFile from .posterior import PosteriorFile +from .ptemcee import PTEmceeFile +from .snowline import SnowlineFile from .txt import InferenceTXTFile +from .ultranest import UltranestFile filetypes = { EmceeFile.name: EmceeFile, @@ -57,13 +57,15 @@ try: from .epsie import EpsieFile + filetypes[EpsieFile.name] = EpsieFile except ImportError: pass def get_file_type(filename): - """ Returns I/O object to use for file. + """ + Returns I/O object to use for file. Parameters ---------- @@ -74,13 +76,14 @@ def get_file_type(filename): ------- file_type : {InferenceFile, InferenceTXTFile} The type of inference file object to use. + """ txt_extensions = [".txt", ".dat", ".csv"] hdf_extensions = [".hdf", ".h5", ".bkup", ".checkpoint"] for ext in hdf_extensions: if filename.endswith(ext): - with _h5py.File(filename, 'r') as fp: - filetype = fp.attrs['filetype'] + with _h5py.File(filename, "r") as fp: + filetype = fp.attrs["filetype"] try: filetype = str(filetype.decode()) except AttributeError: @@ -93,7 +96,8 @@ def get_file_type(filename): def loadfile(path, mode=None, filetype=None, **kwargs): - """Loads the given file using the appropriate InferenceFile class. + """ + Loads the given file using the appropriate InferenceFile class. If ``filetype`` is not provided, this will try to retreive the ``filetype`` from the file's ``attrs``. If the file does not exist yet, an IOError will @@ -116,15 +120,18 @@ def loadfile(path, mode=None, filetype=None, **kwargs): An open file handler to the file. The class used for IO with the file is determined by the ``filetype`` keyword (if provided) or the ``filetype`` stored in the file (if not provided). + """ if filetype is None: # try to read the file to get its filetype try: fileclass = get_file_type(path) - except IOError: + except OSError: # file doesn't exist, filetype must be provided - raise IOError("The file appears not to exist. In this case, " - "filetype must be provided.") + raise OSError( + "The file appears not to exist. In this case, " + "filetype must be provided." + ) else: fileclass = filetypes[filetype] return fileclass(path, mode=mode, **kwargs) @@ -140,7 +147,8 @@ def loadfile(path, mode=None, filetype=None, **kwargs): def check_integrity(filename): - """Checks the integrity of an InferenceFile. + """ + Checks the integrity of an InferenceFile. Checks done are: @@ -164,35 +172,38 @@ def check_integrity(filename): If the samples group does not exist. IOError If any of the checks fail. + """ # check that the file exists if not os.path.exists(filename): - raise ValueError("file {} does not exist".format(filename)) + raise ValueError(f"file {filename} does not exist") # if the file is corrupted such that it cannot be opened, the next line # will raise an IOError - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: # check that all datasets in samples have the same shape parameters = list(fp[fp.samples_group].keys()) # but only do the check if parameters have been written if len(parameters) > 0: - group = fp.samples_group + '/{}' + group = fp.samples_group + "/{}" # use the first parameter as a reference shape ref_shape = fp[group.format(parameters[0])].shape - if not all(fp[group.format(param)].shape == ref_shape - for param in parameters): - raise IOError("not all datasets in the samples group have " - "the same shape") + if not all( + fp[group.format(param)].shape == ref_shape for param in parameters + ): + raise OSError( + "not all datasets in the samples group have the same shape" + ) # check that we can read the first/last sample - firstidx = tuple([0]*len(ref_shape)) - lastidx = tuple([-1]*len(ref_shape)) + firstidx = tuple([0] * len(ref_shape)) + lastidx = tuple([-1] * len(ref_shape)) for param in parameters: _ = fp[group.format(param)][firstidx] _ = fp[group.format(param)][lastidx] -def validate_checkpoint_files(checkpoint_file, backup_file, - check_nsamples=True): - """Checks if the given checkpoint and/or backup files are valid. +def validate_checkpoint_files(checkpoint_file, backup_file, check_nsamples=True): + """ + Checks if the given checkpoint and/or backup files are valid. The checkpoint file is considered valid if: @@ -221,37 +232,38 @@ def validate_checkpoint_files(checkpoint_file, backup_file, checkpoint_valid : bool Whether or not the checkpoint (and backup) file may be used for loading samples. + """ # check if checkpoint file exists and is valid try: check_integrity(checkpoint_file) checkpoint_valid = True - except (ValueError, KeyError, IOError): + except (OSError, ValueError, KeyError): checkpoint_valid = False # backup file try: check_integrity(backup_file) backup_valid = True - except (ValueError, KeyError, IOError): + except (OSError, ValueError, KeyError): backup_valid = False # since we can open the file, run self diagnostics if checkpoint_valid: - with loadfile(checkpoint_file, 'r') as fp: + with loadfile(checkpoint_file, "r") as fp: checkpoint_valid = fp.validate() if backup_valid: - with loadfile(backup_file, 'r') as fp: + with loadfile(backup_file, "r") as fp: backup_valid = fp.validate() if check_nsamples: # This check is not required by nested samplers # check that the checkpoint and backup have the same number of samples; # if not, assume the checkpoint has the correct number if checkpoint_valid and backup_valid: - with loadfile(checkpoint_file, 'r') as fp: + with loadfile(checkpoint_file, "r") as fp: group = list(fp[fp.samples_group].keys())[0] nsamples = fp[fp.samples_group][group].size - with loadfile(backup_file, 'r') as fp: + with loadfile(backup_file, "r") as fp: group = list(fp[fp.samples_group].keys())[0] backup_nsamples = fp[fp.samples_group][group].size backup_valid = nsamples == backup_nsamples @@ -277,7 +289,8 @@ def validate_checkpoint_files(checkpoint_file, backup_file, # ============================================================================= # def get_common_parameters(input_files, collection=None): - """Gets a list of variable params that are common across all input files. + """ + Gets a list of variable params that are common across all input files. If no common parameters are found, a ``ValueError`` is raised. @@ -295,13 +308,14 @@ def get_common_parameters(input_files, collection=None): ------- list : List of the parameter names. + """ if collection is None: collection = "all" parameters = [] for fn in input_files: - fp = loadfile(fn, 'r') - if collection == 'all': + fp = loadfile(fn, "r") + if collection == "all": ps = fp[fp.samples_group].keys() else: ps = fp.attrs[collection] @@ -309,8 +323,11 @@ def get_common_parameters(input_files, collection=None): fp.close() parameters = list(set.intersection(*parameters)) if parameters == []: - raise ValueError("no common parameters found for collection {} in " - "files {}".format(collection, ', '.join(input_files))) + raise ValueError( + "no common parameters found for collection {} in files {}".format( + collection, ", ".join(input_files) + ) + ) # if using python 3 to read a file created in python 2, need to convert # parameters to strs try: @@ -321,13 +338,16 @@ def get_common_parameters(input_files, collection=None): class NoInputFileError(Exception): - """Raised in custom argparse Actions by arguments needing input-files when - no file(s) were provided.""" - pass + """ + Raised in custom argparse Actions by arguments needing input-files when + no file(s) were provided. + """ + class PrintFileParams(argparse.Action): - """Argparse action that will load input files and print possible parameters + """ + Argparse action that will load input files and print possible parameters to screen. Once this is done, the program is forced to exit immediately. The behvior is similar to --help, except that the input-file is read. @@ -336,10 +356,11 @@ class PrintFileParams(argparse.Action): The ``input_file`` attribute must be set in the parser namespace before this action is called. Otherwise, a ``NoInputFileError`` is raised. """ + def __init__(self, skip_args=None, nargs=0, **kwargs): if nargs != 0: raise ValueError("nargs for this action must be 0") - super(PrintFileParams, self).__init__(nargs=nargs, **kwargs) + super().__init__(nargs=nargs, **kwargs) self.skip_args = skip_args def __call__(self, parser, namespace, values, option_string=None): @@ -353,57 +374,75 @@ def __call__(self, parser, namespace, values, option_string=None): raise_err = True if raise_err: raise NoInputFileError("must provide at least one input file") - else: - # just return to stop further processing - return + # just return to stop further processing + return filesbytype = {} fileparsers = {} for fn in input_files: - fp = loadfile(fn, 'r') + fp = loadfile(fn, "r") try: filesbytype[fp.name].append(fn) except KeyError: filesbytype[fp.name] = [fn] # get any extra options fileparsers[fp.name], _ = fp.extra_args_parser( - skip_args=self.skip_args, add_help=False) + skip_args=self.skip_args, add_help=False + ) fp.close() # now print information about the intersection of all parameters - parameters = get_common_parameters(input_files, collection='all') - print("\n"+textwrap.fill("Parameters available with this (these) " - "input file(s):"), end="\n\n") - print(textwrap.fill(' '.join(sorted(parameters))), - end="\n\n") + parameters = get_common_parameters(input_files, collection="all") + print( + "\n" + + textwrap.fill("Parameters available with this (these) input file(s):"), + end="\n\n", + ) + print(textwrap.fill(" ".join(sorted(parameters))), end="\n\n") # information about the pycbc functions pfuncs = sorted(FieldArray.functionlib.fget(FieldArray).keys()) - print(textwrap.fill("Available pycbc functions (see " - "http://pycbc.org/pycbc/latest/html for " - "more details):"), end="\n\n") - print(textwrap.fill(', '.join(pfuncs)), end="\n\n") + print( + textwrap.fill( + "Available pycbc functions (see " + "http://pycbc.org/pycbc/latest/html for " + "more details):" + ), + end="\n\n", + ) + print(textwrap.fill(", ".join(pfuncs)), end="\n\n") # numpy funcs - npfuncs = sorted([name for (name, obj) in _numpy_function_lib.items() - if isinstance(obj, numpy.ufunc)]) - print(textwrap.fill("Available numpy functions:"), - end="\n\n") - print(textwrap.fill(', '.join(npfuncs)), end="\n\n") + npfuncs = sorted( + [ + name + for (name, obj) in _numpy_function_lib.items() + if isinstance(obj, numpy.ufunc) + ] + ) + print(textwrap.fill("Available numpy functions:"), end="\n\n") + print(textwrap.fill(", ".join(npfuncs)), end="\n\n") # misc consts = "e euler_gamma inf nan pi" - print(textwrap.fill("Recognized constants:"), - end="\n\n") + print(textwrap.fill("Recognized constants:"), end="\n\n") print(consts, end="\n\n") - print(textwrap.fill("Python arthimetic (+ - * / // ** %), " - "binary (&, |, etc.), and comparison (>, <, >=, " - "etc.) operators may also be used."), end="\n\n") + print( + textwrap.fill( + "Python arthimetic (+ - * / // ** %), " + "binary (&, |, etc.), and comparison (>, <, >=, " + "etc.) operators may also be used." + ), + end="\n\n", + ) # print out the extra arguments that may be used - outstr = textwrap.fill("The following are additional command-line " - "options that may be provided, along with the " - "input files that understand them:") - print("\n"+outstr, end="\n\n") + outstr = textwrap.fill( + "The following are additional command-line " + "options that may be provided, along with the " + "input files that understand them:" + ) + print("\n" + outstr, end="\n\n") for ftype, fparser in fileparsers.items(): - fnames = ', '.join(filesbytype[ftype]) + fnames = ", ".join(filesbytype[ftype]) if fparser is None: outstr = textwrap.fill( - "File(s) {} use no additional options.".format(fnames)) + f"File(s) {fnames} use no additional options." + ) print(outstr, end="\n\n") else: fparser.usage = fnames @@ -412,7 +451,8 @@ def __call__(self, parser, namespace, values, option_string=None): class ResultsArgumentParser(argparse.ArgumentParser): - r"""Wraps argument parser, and preloads arguments needed for loading samples + r""" + Wraps argument parser, and preloads arguments needed for loading samples from a file. This parser class should be used by any program that wishes to use the @@ -455,10 +495,13 @@ class ResultsArgumentParser(argparse.ArgumentParser): Passed to ``add_results_option_group``; see that function for details. \**kwargs : All other keyword arguments are passed to ``argparse.ArgumentParser``. + """ - def __init__(self, skip_args=None, defaultparams=None, - autoparamlabels=True, **kwargs): - super(ResultsArgumentParser, self).__init__(**kwargs) + + def __init__( + self, skip_args=None, defaultparams=None, autoparamlabels=True, **kwargs + ): + super().__init__(**kwargs) # add attribute to communicate to arguments what to do when there is # no input files self.no_input_file_err = False @@ -466,29 +509,28 @@ def __init__(self, skip_args=None, defaultparams=None, skip_args = [] self.skip_args = skip_args if defaultparams is None: - defaultparams = 'variable_params' + defaultparams = "variable_params" self.defaultparams = defaultparams # add the results option grup self.add_results_option_group(autoparamlabels=autoparamlabels) @property def actions(self): - """Exposes the actions this parser can do as a dictionary. + """ + Exposes the actions this parser can do as a dictionary. The dictionary maps the ``dest`` to actions. """ return {act.dest: act for act in self._actions} def _unset_required(self): - """Convenience function to turn off required arguments for first parse. - """ + """Convenience function to turn off required arguments for first parse.""" self._required_args = [act for act in self._actions if act.required] for act in self._required_args: act.required = False def _reset_required(self): - """Convenience function to turn required arguments back on. - """ + """Convenience function to turn required arguments back on.""" for act in self._required_args: act.required = True @@ -499,30 +541,33 @@ def parse_known_args(self, args=None, namespace=None): # pass self.no_input_file_err = True self._unset_required() - opts, extra_opts = super(ResultsArgumentParser, self).parse_known_args( - args, namespace) + opts, extra_opts = super().parse_known_args( + args, namespace + ) # now do it again self.no_input_file_err = False self._reset_required() - opts, extra_opts = super(ResultsArgumentParser, self).parse_known_args( - args, opts) + opts, extra_opts = super().parse_known_args( + args, opts + ) # populate the parameters option if it wasn't specified - if opts.parameters is None or opts.parameters == ['*']: - parameters = get_common_parameters(opts.input_file, - collection=self.defaultparams) + if opts.parameters is None or opts.parameters == ["*"]: + parameters = get_common_parameters( + opts.input_file, collection=self.defaultparams + ) # now call parse parameters action to re-populate the namespace - self.actions['parameters'](self, opts, parameters) + self.actions["parameters"](self, opts, parameters) # check if we're being greedy or not - elif '*' in opts.parameters: + elif "*" in opts.parameters: # remove the * from the parameters and the labels - opts.parameters = [p for p in opts.parameters if p != '*'] - opts.parameters_labels.pop('*', None) + opts.parameters = [p for p in opts.parameters if p != "*"] + opts.parameters_labels.pop("*", None) # add the rest of the parameters not used - all_params = get_common_parameters(opts.input_file, - collection=self.defaultparams) + all_params = get_common_parameters( + opts.input_file, collection=self.defaultparams + ) # extract the used parameters from the parameters option - used_params = FieldArray.parse_parameters(opts.parameters, - all_params) + used_params = FieldArray.parse_parameters(opts.parameters, all_params) add_params = set(all_params) - set(used_params) # repopulate the name space with the additional parameters if add_params: @@ -532,11 +577,12 @@ def parse_known_args(self, args=None, namespace=None): # parse the sampler-specific options and check for any unknowns unknown = [] for fn in opts.input_file: - fp = loadfile(fn, 'r') + fp = loadfile(fn, "r") sampler_parser, _ = fp.extra_args_parser(skip_args=self.skip_args) if sampler_parser is not None: opts, still_unknown = sampler_parser.parse_known_args( - extra_opts, namespace=opts) + extra_opts, namespace=opts + ) unknown.append(set(still_unknown)) # the intersection of the unknowns are options not understood by # any of the files @@ -545,7 +591,8 @@ def parse_known_args(self, args=None, namespace=None): return opts, list(unknown) def add_results_option_group(self, autoparamlabels=True): - """Adds the options used to call pycbc.inference.io.results_from_cli + """ + Adds the options used to call pycbc.inference.io.results_from_cli function to the parser. These are options releated to loading the results from a run of @@ -561,86 +608,104 @@ def add_results_option_group(self, autoparamlabels=True): ``waveform.parameters`` if a parameter name is the same as a parameter there. Otherwise, will just use whatever label is provided. Default is True. + """ results_reading_group = self.add_argument_group( title="Arguments for loading results", description="Additional, file-specific arguments may also be " "provided, depending on what input-files are given. See " - "--file-help for details.") + "--file-help for details.", + ) results_reading_group.add_argument( - "--input-file", type=str, required=True, nargs="+", - action=ParseLabelArg, metavar='FILE[:LABEL]', + "--input-file", + type=str, + required=True, + nargs="+", + action=ParseLabelArg, + metavar="FILE[:LABEL]", help="Path to input HDF file(s). A label may be specified for " - "each input file to use for plots when multiple files are " - "specified.") + "each input file to use for plots when multiple files are " + "specified.", + ) # advanced help results_reading_group.add_argument( - "-H", "--file-help", - action=PrintFileParams, skip_args=self.skip_args, + "-H", + "--file-help", + action=PrintFileParams, + skip_args=self.skip_args, help="Based on the provided input-file(s), print all available " - "parameters that may be retrieved and all possible functions " - "on those parameters. Also print available additional " - "arguments that may be passed. This option is like an " - "advanced --help: if run, the program will just print the " - "information to screen, then exit.") + "parameters that may be retrieved and all possible functions " + "on those parameters. Also print available additional " + "arguments that may be passed. This option is like an " + "advanced --help: if run, the program will just print the " + "information to screen, then exit.", + ) if autoparamlabels: paramparser = ParseParametersArg lblhelp = ( "If LABEL is the same as a parameter in " "pycbc.waveform.parameters, the label " "property of that parameter will be used (e.g., if LABEL " - "were 'mchirp' then {} would be used). " - .format(_waveform.parameters.mchirp.label)) + f"were 'mchirp' then {_waveform.parameters.mchirp.label} would be used). " + ) else: paramparser = ParseLabelArg - lblhelp = '' + lblhelp = "" results_reading_group.add_argument( - "--parameters", type=str, nargs="+", metavar="PARAM[:LABEL]", + "--parameters", + type=str, + nargs="+", + metavar="PARAM[:LABEL]", action=paramparser, help="Name of parameters to load; default is to load all. The " - "parameters can be any of the model params or posterior " - "stats (loglikelihood, logprior, etc.) in the input file(s), " - "derived parameters from them, or any function of them. If " - "multiple files are provided, any parameter common to all " - "files may be used. Syntax for functions is python; any math " - "functions in the numpy libary may be used. Can optionally " - "also specify a LABEL for each parameter. If no LABEL is " - "provided, PARAM will used as the LABEL. {}" - "To see all possible parameters that may be used with the " - "given input file(s), as well as all avaiable functions, " - "run --file-help, along with one or more input files. " - "If '*' is provided in addition to other parameters names, " - "then parameters will be loaded in a greedy fashion; i.e., " - "all other parameters that exist in the file(s) that are not " - "explicitly mentioned will also be loaded. For example, " - "if the input-file(s) contains 'srcmass1', " - "'srcmass2', and 'distance', and " - "\"'primary_mass(srcmass1, srcmass2):mass1' '*'\", is given " - "then 'mass1' and 'distance' will be loaded. Otherwise, " - "without the '*', only 'mass1' would be loaded. " - "Note that any parameter that is used in a function " - "will not automatically be added. Tip: enclose " - "arguments in single quotes, or else special characters will " - "be interpreted as shell commands. For example, the " - "wildcard should be given as either '*' or \\*, otherwise " - "bash will expand the * into the names of all the files in " - "the current directory." - .format(lblhelp)) + "parameters can be any of the model params or posterior " + "stats (loglikelihood, logprior, etc.) in the input file(s), " + "derived parameters from them, or any function of them. If " + "multiple files are provided, any parameter common to all " + "files may be used. Syntax for functions is python; any math " + "functions in the numpy libary may be used. Can optionally " + "also specify a LABEL for each parameter. If no LABEL is " + f"provided, PARAM will used as the LABEL. {lblhelp}" + "To see all possible parameters that may be used with the " + "given input file(s), as well as all avaiable functions, " + "run --file-help, along with one or more input files. " + "If '*' is provided in addition to other parameters names, " + "then parameters will be loaded in a greedy fashion; i.e., " + "all other parameters that exist in the file(s) that are not " + "explicitly mentioned will also be loaded. For example, " + "if the input-file(s) contains 'srcmass1', " + "'srcmass2', and 'distance', and " + "\"'primary_mass(srcmass1, srcmass2):mass1' '*'\", is given " + "then 'mass1' and 'distance' will be loaded. Otherwise, " + "without the '*', only 'mass1' would be loaded. " + "Note that any parameter that is used in a function " + "will not automatically be added. Tip: enclose " + "arguments in single quotes, or else special characters will " + "be interpreted as shell commands. For example, the " + "wildcard should be given as either '*' or \\*, otherwise " + "bash will expand the * into the names of all the files in " + "the current directory.", + ) results_reading_group.add_argument( - "--constraint", type=str, nargs="+", metavar="CONSTRAINT[:FILE]", + "--constraint", + type=str, + nargs="+", + metavar="CONSTRAINT[:FILE]", help="Apply a constraint to the samples. If a file is provided " - "after the constraint, it will only be applied to the given " - "file. Otherwise, the constraint will be applied to all " - "files. Only one constraint may be applied to a file. " - "Samples that violate the constraint will be removed. Syntax " - "is python; any parameter or function of parameter can be " - "used, similar to the parameters argument. Multiple " - "constraints may be combined by using '&' and '|'.") + "after the constraint, it will only be applied to the given " + "file. Otherwise, the constraint will be applied to all " + "files. Only one constraint may be applied to a file. " + "Samples that violate the constraint will be removed. Syntax " + "is python; any parameter or function of parameter can be " + "used, similar to the parameters argument. Multiple " + "constraints may be combined by using '&' and '|'.", + ) return results_reading_group def results_from_cli(opts, load_samples=True, **kwargs): - r"""Loads an inference result file along with any labels associated with it + r""" + Loads an inference result file along with any labels associated with it from the command line options. Parameters @@ -665,8 +730,8 @@ def results_from_cli(opts, load_samples=True, **kwargs): \**kwargs : Any other keyword arguments that are passed to read samples using samples_from_cli - """ + """ # lists for files and samples from all input files fp_all = [] samples_all = [] @@ -679,17 +744,19 @@ def results_from_cli(opts, load_samples=True, **kwargs): constraints = {} if opts.constraint is not None: for constraint in opts.constraint: - if len(constraint.split(':')) == 2: - constraint, fn = constraint.split(':') + if len(constraint.split(":")) == 2: + constraint, fn = constraint.split(":") constraints[fn] = constraint # no file provided, make sure there's only one constraint elif len(opts.constraint) > 1: - raise ValueError("must provide a file to apply constraints " - "to if providing more than one constraint") + raise ValueError( + "must provide a file to apply constraints " + "to if providing more than one constraint" + ) else: # this means no file, only one constraint, apply to all # files - constraints = {fn: constraint for fn in input_files} + constraints = dict.fromkeys(input_files, constraint) # loop over all input files for input_file in input_files: @@ -703,19 +770,18 @@ def results_from_cli(opts, load_samples=True, **kwargs): logging.info("Loading samples") # read samples from file - samples = fp.samples_from_cli(opts, parameters=opts.parameters, - **kwargs) - logging.info("Loaded {} samples".format(samples.size)) + samples = fp.samples_from_cli(opts, parameters=opts.parameters, **kwargs) + logging.info(f"Loaded {samples.size} samples") if input_file in constraints: logging.info("Applying constraints") mask = samples[constraints[input_file]] samples = samples[mask] if samples.size == 0: - raise ValueError("No samples remain after constraint {} " - "applied".format(constraints[input_file])) - logging.info("{} samples remain".format(samples.size)) - + raise ValueError( + f"No samples remain after constraint {constraints[input_file]} applied" + ) + logging.info(f"{samples.size} samples remain") # else do not read samples else: @@ -735,7 +801,8 @@ def results_from_cli(opts, load_samples=True, **kwargs): def injections_from_cli(opts): - """Gets injection parameters from the inference file(s). + """ + Gets injection parameters from the inference file(s). If the opts have a ``injection_samples_map`` option, the injection parameters will be remapped accordingly. See @@ -752,10 +819,11 @@ def injections_from_cli(opts): FieldArray Array of the injection parameters from all of the input files given by ``opts.input_file``. + """ # see if a mapping was provided - if hasattr(opts, 'injection_samples_map') and opts.injection_samples_map: - param_map = [opt.split(':') for opt in opts.injection_samples_map] + if hasattr(opts, "injection_samples_map") and opts.injection_samples_map: + param_map = [opt.split(":") for opt in opts.injection_samples_map] else: param_map = [] input_files = opts.input_file @@ -764,7 +832,7 @@ def injections_from_cli(opts): injections = None # loop over all input files getting the injection files for input_file in input_files: - fp = loadfile(input_file, 'r') + fp = loadfile(input_file, "r") these_injs = fp.read_injections() # apply mapping if it was provided if param_map: @@ -776,8 +844,7 @@ def injections_from_cli(opts): these_injs[param] = mapvals.pop(param) # add the rest as new fields ps = list(mapvals.keys()) - these_injs = these_injs.add_fields([mapvals[p] for p in ps], - names=ps) + these_injs = these_injs.add_fields([mapvals[p] for p in ps], names=ps) if injections is None: injections = these_injs else: diff --git a/pycbc/inference/io/base_hdf.py b/pycbc/inference/io/base_hdf.py index a62da414b2f..803ddc3a98b 100644 --- a/pycbc/inference/io/base_hdf.py +++ b/pycbc/inference/io/base_hdf.py @@ -21,29 +21,28 @@ # # ============================================================================= # -"""This modules defines functions for reading and writing samples that the +""" +This modules defines functions for reading and writing samples that the inference samplers generate. """ - -import sys import logging +import sys +from abc import ABCMeta, abstractmethod from io import StringIO -from abc import (ABCMeta, abstractmethod) - -import numpy import h5py +import numpy -from pycbc.io import FieldArray from pycbc.inject import InjectionSet -from pycbc.io import (dump_state, load_state) -from pycbc.workflow import WorkflowConfigParser +from pycbc.io import FieldArray, dump_state, load_state from pycbc.types import FrequencySeries +from pycbc.workflow import WorkflowConfigParser def format_attr(val): - """Formats an attr so that it can be read in either python 2 or 3. + """ + Formats an attr so that it can be read in either python 2 or 3. In python 2, strings that are saved as an attribute in an hdf file default to unicode. Since unicode was removed in python 3, if you load that file @@ -64,6 +63,7 @@ def format_attr(val): If ``val`` was a byte string, the value as a ``str``. If the value was a numpy array of ``bytes_``, the value as a list of ``str``. Otherwise, just returns the value. + """ try: val = str(val.decode()) @@ -75,36 +75,38 @@ def format_attr(val): class BaseInferenceFile(h5py.File, metaclass=ABCMeta): - """Base class for all inference hdf files. + """ + Base class for all inference hdf files. This is a subclass of the h5py.File object. It adds functions for handling reading and writing the samples from the samplers. Parameters - ----------- + ---------- path : str The path to the HDF file. mode : {None, str} The mode to open the file, eg. "w" for write and "r" for read. + """ name = None - samples_group = 'samples' - sampler_group = 'sampler_info' - data_group = 'data' - injections_group = 'injections' - config_group = 'config_file' + samples_group = "samples" + sampler_group = "sampler_info" + data_group = "data" + injections_group = "injections" + config_group = "config_file" def __init__(self, path, mode=None, **kwargs): - super(BaseInferenceFile, self).__init__(path, mode, **kwargs) + super().__init__(path, mode, **kwargs) # check that file type matches self try: - filetype = self.attrs['filetype'] + filetype = self.attrs["filetype"] except KeyError: - if mode == 'w': + if mode == "w": # first time creating the file, add this class's name filetype = self.name - self.attrs['filetype'] = filetype + self.attrs["filetype"] = filetype else: filetype = None try: @@ -112,13 +114,16 @@ def __init__(self, path, mode=None, **kwargs): except AttributeError: pass if filetype != self.name: - raise ValueError("This file has filetype {}, whereas this class " - "is named {}. This indicates that the file was " - "not written by this class, and so cannot be " - "read by this class.".format(filetype, self.name)) + raise ValueError( + f"This file has filetype {filetype}, whereas this class " + f"is named {self.name}. This indicates that the file was " + "not written by this class, and so cannot be " + "read by this class." + ) def __getattr__(self, attr): - """Things stored in ``.attrs`` are promoted to instance attributes. + """ + Things stored in ``.attrs`` are promoted to instance attributes. Note that properties will be called before this, so if there are any properties that share the same name as something in ``.attrs``, that @@ -127,7 +132,8 @@ def __getattr__(self, attr): return self.attrs[attr] def getattrs(self, group=None, create_missing=True): - """Convenience function for getting the `attrs` from the file or group. + """ + Convenience function for getting the `attrs` from the file or group. Parameters ---------- @@ -142,6 +148,7 @@ def getattrs(self, group=None, create_missing=True): ------- h5py.File.attrs An attrs instance of the file or requested group. + """ if group is None or group == "/": attrs = self.attrs @@ -158,7 +165,8 @@ def getattrs(self, group=None, create_missing=True): @abstractmethod def write_samples(self, samples, **kwargs): - r"""This should write all of the provided samples. + r""" + This should write all of the provided samples. This function should be used to write both samples and model stats. @@ -168,11 +176,12 @@ def write_samples(self, samples, **kwargs): Samples should be provided as a dictionary of numpy arrays. \**kwargs : Any other keyword args the sampler needs to write data. + """ - pass def parse_parameters(self, parameters, array_class=None): - """Parses a parameters arg to figure out what fields need to be loaded. + """ + Parses a parameters arg to figure out what fields need to be loaded. Parameters ---------- @@ -191,6 +200,7 @@ def parse_parameters(self, parameters, array_class=None): ------- list : A list of strings giving the fields to load from the file. + """ # get the type of array class to use if array_class is None: @@ -200,7 +210,8 @@ def parse_parameters(self, parameters, array_class=None): return array_class.parse_parameters(parameters, possible_fields) def read_samples(self, parameters, array_class=None, **kwargs): - r"""Reads samples for the given parameter(s). + r""" + Reads samples for the given parameter(s). The ``parameters`` can be the name of any dataset in ``samples_group``, a virtual field or method of ``FieldArray`` (as long as the file @@ -215,7 +226,7 @@ def read_samples(self, parameters, array_class=None, **kwargs): ``FieldArray``. Parameters - ----------- + ---------- parameters : (list of) strings The parameter(s) to retrieve. array_class : FieldArray-like class, optional @@ -229,6 +240,7 @@ def read_samples(self, parameters, array_class=None, **kwargs): ------- FieldArray : The samples as a ``FieldArray``. + """ # get the type of array class to use if array_class is None: @@ -240,9 +252,10 @@ def read_samples(self, parameters, array_class=None, **kwargs): # convert to FieldArray samples = array_class.from_kwargs(**samples) # add the static params and attributes - addatrs = (list(self.static_params.items()) + - list(self[self.samples_group].attrs.items())) - for (p, val) in addatrs: + addatrs = list(self.static_params.items()) + list( + self[self.samples_group].attrs.items() + ) + for p, val in addatrs: if p in loadfields: continue setattr(samples, format_attr(p), format_attr(val)) @@ -250,15 +263,16 @@ def read_samples(self, parameters, array_class=None, **kwargs): @abstractmethod def read_raw_samples(self, fields, **kwargs): - """Low level function for reading datasets in the samples group. + """ + Low level function for reading datasets in the samples group. This should return a dictionary of numpy arrays. """ - pass @staticmethod def extra_args_parser(parser=None, skip_args=None, **kwargs): - r"""Provides a parser that can be used to parse sampler-specific command + r""" + Provides a parser that can be used to parse sampler-specific command line options for loading samples. This is optional. Inheriting classes may override this if they want to @@ -286,12 +300,14 @@ def extra_args_parser(parser=None, skip_args=None, **kwargs): for the ``parser`` argument (default is None). actions : list of argparse.Action List of the actions that were added. + """ return parser, [] @staticmethod def _get_optional_args(args, opts, err_on_missing=False, **kwargs): - r"""Convenience function to retrieve arguments from an argparse + r""" + Convenience function to retrieve arguments from an argparse namespace. Parameters @@ -314,6 +330,7 @@ def _get_optional_args(args, opts, err_on_missing=False, **kwargs): Dictionary mapping arguments to values retrieved from ``opts``. If keyword arguments were provided, these will also be included in the dictionary. + """ parsed = {} for arg in args: @@ -322,13 +339,13 @@ def _get_optional_args(args, opts, err_on_missing=False, **kwargs): except AttributeError as e: if err_on_missing: raise AttributeError(e) - else: - continue + continue parsed.update(kwargs) return parsed def samples_from_cli(self, opts, parameters=None, **kwargs): - r"""Reads samples from the given command-line options. + r""" + Reads samples from the given command-line options. Parameters ---------- @@ -346,6 +363,7 @@ def samples_from_cli(self, opts, parameters=None, **kwargs): ------- FieldArray : Array of the loaded samples. + """ if parameters is None and opts.parameters is None: parameters = self.variable_params @@ -359,27 +377,28 @@ def samples_from_cli(self, opts, parameters=None, **kwargs): @property def static_params(self): - """Returns a dictionary of the static_params. The keys are the argument + """ + Returns a dictionary of the static_params. The keys are the argument names, values are the value they were set to. """ return {arg: self.attrs[arg] for arg in self.attrs["static_params"]} @property def effective_nsamples(self): - """Returns the effective number of samples stored in the file. - """ + """Returns the effective number of samples stored in the file.""" try: - return self.attrs['effective_nsamples'] + return self.attrs["effective_nsamples"] except KeyError: return 0 def write_effective_nsamples(self, effective_nsamples): """Writes the effective number of samples stored in the file.""" - self.attrs['effective_nsamples'] = effective_nsamples + self.attrs["effective_nsamples"] = effective_nsamples @property def thin_start(self): - """The default start index to use when reading samples. + """ + The default start index to use when reading samples. Unless overridden by sub-class attribute, just returns 0. """ @@ -387,7 +406,8 @@ def thin_start(self): @property def thin_interval(self): - """The default interval to use when reading samples. + """ + The default interval to use when reading samples. Unless overridden by sub-class attribute, just returns 1. """ @@ -395,7 +415,8 @@ def thin_interval(self): @property def thin_end(self): - """The defaut end index to use when reading samples. + """ + The defaut end index to use when reading samples. Unless overriden by sub-class attribute, just return ``None``. """ @@ -403,7 +424,8 @@ def thin_end(self): @property def cmd(self): - """Returns the (last) saved command line. + """ + Returns the (last) saved command line. If the file was created from a run that resumed from a checkpoint, only the last command line used is returned. @@ -412,6 +434,7 @@ def cmd(self): ------- cmd : string The command line that created this InferenceFile. + """ cmd = self.attrs["cmd"] if isinstance(cmd, numpy.ndarray): @@ -419,7 +442,8 @@ def cmd(self): return cmd def write_logevidence(self, lnz, dlnz): - """Writes the given log evidence and its error. + """ + Writes the given log evidence and its error. Results are saved to file's 'log_evidence' and 'dlog_evidence' attributes. @@ -430,19 +454,22 @@ def write_logevidence(self, lnz, dlnz): The log of the evidence. dlnz : float The error in the estimate of the log evidence. + """ - self.attrs['log_evidence'] = lnz - self.attrs['dlog_evidence'] = dlnz + self.attrs["log_evidence"] = lnz + self.attrs["dlog_evidence"] = dlnz @property def log_evidence(self): - """Returns the log of the evidence and its error, if they exist in the + """ + Returns the log of the evidence and its error, if they exist in the file. Raises a KeyError otherwise. """ return self.attrs["log_evidence"], self.attrs["dlog_evidence"] def write_random_state(self, group=None, state=None): - """Writes the state of the random number generator from the file. + """ + Writes the state of the random number generator from the file. The random state is written to ``sampler_group``/random_state. @@ -453,6 +480,7 @@ def write_random_state(self, group=None, state=None): state : tuple, optional Specify the random state to write. If None, will use ``numpy.random.get_state()``. + """ # Write out the default numpy random state group = self.sampler_group if group is None else group @@ -463,8 +491,9 @@ def write_random_state(self, group=None, state=None): if dataset_name in self: self[dataset_name][:] = arr else: - self.create_dataset(dataset_name, arr.shape, fletcher32=True, - dtype=arr.dtype) + self.create_dataset( + dataset_name, arr.shape, fletcher32=True, dtype=arr.dtype + ) self[dataset_name][:] = arr self[dataset_name].attrs["s"] = s self[dataset_name].attrs["pos"] = pos @@ -472,7 +501,8 @@ def write_random_state(self, group=None, state=None): self[dataset_name].attrs["cached_gauss"] = cached_gauss def read_random_state(self, group=None): - """Reads the state of the random number generator from the file. + """ + Reads the state of the random number generator from the file. Parameters ---------- @@ -483,6 +513,7 @@ def read_random_state(self, group=None): ------- tuple A tuple with 5 elements that can be passed to numpy.set_state. + """ # Read numpy randomstate group = self.sampler_group if group is None else group @@ -496,73 +527,79 @@ def read_random_state(self, group=None): return state def write_strain(self, strain_dict, group=None): - """Writes strain for each IFO to file. + """ + Writes strain for each IFO to file. Parameters - ----------- + ---------- strain : {dict, FrequencySeries} A dict of FrequencySeries where the key is the IFO. group : {None, str} The group to write the strain to. If None, will write to the top level. + """ subgroup = self.data_group + "/{ifo}/strain" if group is None: group = subgroup else: - group = '/'.join([group, subgroup]) + group = "/".join([group, subgroup]) for ifo, strain in strain_dict.items(): self[group.format(ifo=ifo)] = strain - self[group.format(ifo=ifo)].attrs['delta_t'] = strain.delta_t - self[group.format(ifo=ifo)].attrs['start_time'] = \ - float(strain.start_time) + self[group.format(ifo=ifo)].attrs["delta_t"] = strain.delta_t + self[group.format(ifo=ifo)].attrs["start_time"] = float(strain.start_time) def write_stilde(self, stilde_dict, group=None): - """Writes stilde for each IFO to file. + """ + Writes stilde for each IFO to file. Parameters - ----------- + ---------- stilde : {dict, FrequencySeries} A dict of FrequencySeries where the key is the IFO. group : {None, str} The group to write the strain to. If None, will write to the top level. + """ subgroup = self.data_group + "/{ifo}/stilde" if group is None: group = subgroup else: - group = '/'.join([group, subgroup]) + group = "/".join([group, subgroup]) for ifo, stilde in stilde_dict.items(): self[group.format(ifo=ifo)] = stilde - self[group.format(ifo=ifo)].attrs['delta_f'] = stilde.delta_f - self[group.format(ifo=ifo)].attrs['epoch'] = float(stilde.epoch) + self[group.format(ifo=ifo)].attrs["delta_f"] = stilde.delta_f + self[group.format(ifo=ifo)].attrs["epoch"] = float(stilde.epoch) def write_psd(self, psds, group=None): - """Writes PSD for each IFO to file. + """ + Writes PSD for each IFO to file. PSDs are written to ``[{group}/]data/{detector}/psds/0``, where {group} is the optional keyword argument. Parameters - ----------- + ---------- psds : dict A dict of detector name -> FrequencySeries. group : str, optional Specify a top-level group to write the data to. If ``None`` (the default), data will be written to the file's top level. + """ subgroup = self.data_group + "/{ifo}/psds/0" if group is None: group = subgroup else: - group = '/'.join([group, subgroup]) + group = "/".join([group, subgroup]) for ifo in psds: self[group.format(ifo=ifo)] = psds[ifo] - self[group.format(ifo=ifo)].attrs['delta_f'] = psds[ifo].delta_f + self[group.format(ifo=ifo)].attrs["delta_f"] = psds[ifo].delta_f def write_injections(self, injection_file, group=None): - """Writes injection parameters from the given injection file. + """ + Writes injection parameters from the given injection file. Everything in the injection file is copied to ``[{group}/]injections_group``, where ``{group}`` is the optional @@ -576,23 +613,22 @@ def write_injections(self, injection_file, group=None): Specify a top-level group to write the injections group to. If ``None`` (the default), injections group will be written to the file's top level. + """ logging.info("Writing injection file to output") - if group is None or group == '/': + if group is None or group == "/": group = self.injections_group else: - group = '/'.join([group, self.injections_group]) + group = "/".join([group, self.injections_group]) try: with h5py.File(injection_file, "r") as fp: - super(BaseInferenceFile, self).copy(fp, group) - except IOError: - logging.warning( - "Could not read %s as an HDF file", - injection_file - ) + super().copy(fp, group) + except OSError: + logging.warning("Could not read %s as an HDF file", injection_file) def read_injections(self, group=None): - """Gets injection parameters. + """ + Gets injection parameters. Injections are retrieved from ``[{group}/]injections``. @@ -606,17 +642,19 @@ def read_injections(self, group=None): ------- FieldArray Array of the injection parameters. + """ - if group is None or group == '/': + if group is None or group == "/": group = self.injections_group else: - group = '/'.join([group, self.injections_group]) + group = "/".join([group, self.injections_group]) injset = InjectionSet(self.filename, hdf_group=group) injections = injset.table.view(FieldArray) return injections def write_command_line(self): - """Writes command line to attributes. + """ + Writes command line to attributes. The command line is written to the file's ``attrs['cmd']``. If this attribute already exists in the file (this can happen when resuming @@ -637,7 +675,8 @@ def write_command_line(self): @staticmethod def get_slice(thin_start=None, thin_interval=None, thin_end=None): - """Formats a slice to retrieve a thinned array from an HDF file. + """ + Formats a slice to retrieve a thinned array from an HDF file. Parameters ---------- @@ -652,6 +691,7 @@ def get_slice(thin_start=None, thin_interval=None, thin_end=None): ------- slice : The slice needed. + """ if thin_start is not None: thin_start = int(thin_start) @@ -662,7 +702,8 @@ def get_slice(thin_start=None, thin_interval=None, thin_end=None): return slice(thin_start, thin_end, thin_interval) def copy_metadata(self, other): - """Copies all metadata from this file to the other file. + """ + Copies all metadata from this file to the other file. Metadata is defined as everything in the top-level ``.attrs``. @@ -670,6 +711,7 @@ def copy_metadata(self, other): ---------- other : InferenceFile An open inference file to write the data to. + """ logging.info("Copying metadata") # copy attributes @@ -677,7 +719,8 @@ def copy_metadata(self, other): other.attrs[key] = self.attrs[key] def copy_info(self, other, ignore=None): - """Copies "info" from this file to the other. + """ + Copies "info" from this file to the other. "Info" is defined all groups that are not the samples group. @@ -687,6 +730,7 @@ def copy_info(self, other, ignore=None): The output file. Must be an hdf file. ignore : (list of) str Don't copy the given groups. + """ logging.info("Copying info") # copy non-samples/stats data @@ -697,11 +741,18 @@ def copy_info(self, other, ignore=None): ignore = set(ignore + [self.samples_group]) copy_groups = set(self.keys()) - ignore for key in copy_groups: - super(BaseInferenceFile, self).copy(key, other) - - def copy_samples(self, other, parameters=None, parameter_names=None, - read_args=None, write_args=None): - """Should copy samples to the other files. + super().copy(key, other) + + def copy_samples( + self, + other, + parameters=None, + parameter_names=None, + read_args=None, + write_args=None, + ): + """ + Should copy samples to the other files. Parameters ---------- @@ -717,6 +768,7 @@ def copy_samples(self, other, parameters=None, parameter_names=None, Arguments to pass to ``read_samples``. write_args : dict, optional Arguments to pass to ``write_samples``. + """ # select the samples to copy logging.info("Reading samples to copy") @@ -724,27 +776,33 @@ def copy_samples(self, other, parameters=None, parameter_names=None, parameters = self.variable_params # if list of desired parameters is different, rename if set(parameters) != set(self.variable_params): - other.attrs['variable_params'] = parameters + other.attrs["variable_params"] = parameters if read_args is None: read_args = {} samples = self.read_samples(parameters, **read_args) - logging.info("Copying {} samples".format(samples.size)) + logging.info(f"Copying {samples.size} samples") # if different parameter names are desired, get them from the samples if parameter_names: arrs = {pname: samples[p] for p, pname in parameter_names.items()} - arrs.update({p: samples[p] for p in parameters if - p not in parameter_names}) + arrs.update({p: samples[p] for p in parameters if p not in parameter_names}) samples = FieldArray.from_kwargs(**arrs) - other.attrs['variable_params'] = samples.fieldnames + other.attrs["variable_params"] = samples.fieldnames logging.info("Writing samples") if write_args is None: write_args = {} - other.write_samples({p: samples[p] for p in samples.fieldnames}, - **write_args) - - def copy(self, other, ignore=None, parameters=None, parameter_names=None, - read_args=None, write_args=None): - """Copies metadata, info, and samples in this file to another file. + other.write_samples({p: samples[p] for p in samples.fieldnames}, **write_args) + + def copy( + self, + other, + ignore=None, + parameters=None, + parameter_names=None, + read_args=None, + write_args=None, + ): + """ + Copies metadata, info, and samples in this file to another file. Parameters ---------- @@ -772,12 +830,13 @@ def copy(self, other, ignore=None, parameters=None, parameter_names=None, ------- InferenceFile The open file handler to other. + """ if not isinstance(other, h5py.File): # check that we're not trying to overwrite this file if other == self.name: - raise IOError("destination is the same as this file") - other = self.__class__(other, 'w') + raise OSError("destination is the same as this file") + other = self.__class__(other, "w") # metadata self.copy_metadata(other) # info @@ -788,10 +847,13 @@ def copy(self, other, ignore=None, parameters=None, parameter_names=None, self.copy_info(other, ignore=ignore) # samples if self.samples_group not in ignore: - self.copy_samples(other, parameters=parameters, - parameter_names=parameter_names, - read_args=read_args, - write_args=write_args) + self.copy_samples( + other, + parameters=parameters, + parameter_names=parameter_names, + read_args=read_args, + write_args=write_args, + ) # if any down selection was done, re-set the default # thin-start/interval/end p = tuple(self[self.samples_group].keys())[0] @@ -799,14 +861,15 @@ def copy(self, other, ignore=None, parameters=None, parameter_names=None, p = tuple(other[other.samples_group].keys())[0] other_shape = other[other.samples_group][p].shape if my_shape != other_shape: - other.attrs['thin_start'] = 0 - other.attrs['thin_interval'] = 1 - other.attrs['thin_end'] = None + other.attrs["thin_start"] = 0 + other.attrs["thin_interval"] = 1 + other.attrs["thin_end"] = None return other @classmethod def write_kwargs_to_attrs(cls, attrs, **kwargs): - r"""Writes the given keywords to the given ``attrs``. + r""" + Writes the given keywords to the given ``attrs``. If any keyword argument points to a dict, the keyword will point to a list of the dict's keys. Each key is then written to the attrs with its @@ -818,6 +881,7 @@ def write_kwargs_to_attrs(cls, attrs, **kwargs): The ``attrs`` of an hdf file or a group in an hdf file. \**kwargs : The keywords to write. + """ for arg, val in kwargs.items(): if val is None: @@ -830,7 +894,8 @@ def write_kwargs_to_attrs(cls, attrs, **kwargs): attrs[str(arg)] = val def write_data(self, name, data, path=None, append=False): - """Convenience function to write data. + """ + Convenience function to write data. Given ``data`` is written as a dataset with ``name`` in ``path``. If the dataset or path do not exist yet, the dataset and path will @@ -858,9 +923,10 @@ def write_data(self, name, data, path=None, append=False): the data, it must be resizable along this dimension. If ``False`` (the default) what is in the file will be overwritten, and the given data must have the same shape. + """ if path is None: - path = '/' + path = "/" try: group = self[path] except KeyError: @@ -870,8 +936,7 @@ def write_data(self, name, data, path=None, append=False): if isinstance(data, dict): # call myself for each key, value pair in the dictionary for key, val in data.items(): - self.write_data(key, val, path='/'.join([path, name]), - append=append) + self.write_data(key, val, path="/".join([path, name]), append=append) # if appending, we need to resize the data on disk, or, if it doesn't # exist yet, create a dataset that is resizable along the last # dimension @@ -885,15 +950,20 @@ def write_data(self, name, data, path=None, append=False): ndata = dshape[-1] try: startidx = group[name].shape[-1] - group[name].resize(dshape[-1]+group[name].shape[-1], - axis=len(group[name].shape)-1) + group[name].resize( + dshape[-1] + group[name].shape[-1], axis=len(group[name].shape) - 1 + ) except KeyError: # dataset doesn't exist yet - group.create_dataset(name, dshape, - maxshape=tuple(list(dshape)[:-1]+[None]), - dtype=data.dtype, fletcher32=True) + group.create_dataset( + name, + dshape, + maxshape=tuple(list(dshape)[:-1] + [None]), + dtype=data.dtype, + fletcher32=True, + ) startidx = 0 - group[name][..., startidx:startidx+ndata] = data[..., :] + group[name][..., startidx : startidx + ndata] = data[..., :] else: try: group[name][()] = data @@ -902,7 +972,8 @@ def write_data(self, name, data, path=None, append=False): group[name] = data def write_config_file(self, cp): - """Writes the given config file parser. + """ + Writes the given config file parser. File is stored as a pickled buffer array to ``config_parser/{index}``, where ``{index}`` is an integer corresponding to the number of config @@ -913,6 +984,7 @@ def write_config_file(self, cp): ---------- cp : ConfigParser Config parser to save. + """ # get the index of the last saved file try: @@ -931,7 +1003,8 @@ def write_config_file(self, cp): dump_state(out, self, path=self.config_group, dsetname=str(index)) def read_config_file(self, return_cp=True, index=-1): - """Reads the config file that was used. + """ + Reads the config file that was used. A ``ValueError`` is raised if no config files have been saved, or if the requested index larger than the number of stored config files. @@ -952,6 +1025,7 @@ def read_config_file(self, return_cp=True, index=-1): ------- WorkflowConfigParser or StringIO : The parsed config file. + """ # get the stored indices try: @@ -970,7 +1044,8 @@ def read_config_file(self, return_cp=True, index=-1): return cf def read_data(self, group=None): - """Loads the data stored in the file as a FrequencySeries. + """ + Loads the data stored in the file as a FrequencySeries. Only works for models that store data as a frequency series in ``data/DET/stilde``. A ``KeyError`` will be raised if the model used @@ -986,22 +1061,24 @@ def read_data(self, group=None): ------- dict : Dictionary of detector name -> FrequencySeries. + """ - fmt = '{}/{}/stilde' - if group is None or group == '/': + fmt = "{}/{}/stilde" + if group is None or group == "/": path = self.data_group else: - path = '/'.join([group, self.data_group]) + path = "/".join([group, self.data_group]) data = {} for det in self[path].keys(): group = self[fmt.format(path, det)] data[det] = FrequencySeries( - group[()], delta_f=group.attrs['delta_f'], - epoch=group.attrs['epoch']) + group[()], delta_f=group.attrs["delta_f"], epoch=group.attrs["epoch"] + ) return data def read_psds(self, group=None): - """Loads the PSDs stored in the file as a FrequencySeries. + """ + Loads the PSDs stored in the file as a FrequencySeries. Only works for models that store PSDs in ``data/DET/psds/0``. A ``KeyError`` will be raised if the model used @@ -1017,15 +1094,15 @@ def read_psds(self, group=None): ------- dict : Dictionary of detector name -> FrequencySeries. + """ - fmt = '{}/{}/psds/0' - if group is None or group == '/': + fmt = "{}/{}/psds/0" + if group is None or group == "/": path = self.data_group else: - path = '/'.join([group, self.data_group]) + path = "/".join([group, self.data_group]) psds = {} for det in self[path].keys(): group = self[fmt.format(path, det)] - psds[det] = FrequencySeries( - group[()], delta_f=group.attrs['delta_f']) + psds[det] = FrequencySeries(group[()], delta_f=group.attrs["delta_f"]) return psds diff --git a/pycbc/inference/io/base_mcmc.py b/pycbc/inference/io/base_mcmc.py index 057b2976a77..a7679bdd5e2 100644 --- a/pycbc/inference/io/base_mcmc.py +++ b/pycbc/inference/io/base_mcmc.py @@ -21,24 +21,27 @@ # # ============================================================================= # -"""Provides I/O that is specific to MCMC samplers. -""" +"""Provides I/O that is specific to MCMC samplers.""" +import argparse import numpy -import argparse -class CommonMCMCMetadataIO(object): - """Provides functions for reading/writing MCMC metadata to file. +class CommonMCMCMetadataIO: + """ + Provides functions for reading/writing MCMC metadata to file. The functions here are common to both standard MCMC (in which chains are independent) and ensemble MCMC (in which chains/walkers share information). """ + def write_resume_point(self): - """Keeps a list of the number of iterations that were in a file when a - run was resumed from a checkpoint.""" + """ + Keeps a list of the number of iterations that were in a file when a + run was resumed from a checkpoint. + """ try: resume_pts = self.attrs["resume_points"].tolist() except KeyError: @@ -52,37 +55,40 @@ def write_resume_point(self): def write_niterations(self, niterations): """Writes the given number of iterations to the sampler group.""" - self[self.sampler_group].attrs['niterations'] = niterations + self[self.sampler_group].attrs["niterations"] = niterations @property def niterations(self): """Returns the number of iterations the sampler was run for.""" - return self[self.sampler_group].attrs['niterations'] + return self[self.sampler_group].attrs["niterations"] @property def nwalkers(self): - """Returns the number of walkers used by the sampler. + """ + Returns the number of walkers used by the sampler. Alias of ``nchains``. """ try: - return self[self.sampler_group].attrs['nwalkers'] + return self[self.sampler_group].attrs["nwalkers"] except KeyError: - return self[self.sampler_group].attrs['nchains'] + return self[self.sampler_group].attrs["nchains"] @property def nchains(self): - """Returns the number of chains used by the sampler. + """ + Returns the number of chains used by the sampler. Alias of ``nwalkers``. """ try: - return self[self.sampler_group].attrs['nchains'] + return self[self.sampler_group].attrs["nchains"] except KeyError: - return self[self.sampler_group].attrs['nwalkers'] + return self[self.sampler_group].attrs["nwalkers"] def _thin_data(self, group, params, thin_interval): - """Thins data on disk by the given interval. + """ + Thins data on disk by the given interval. This makes no effort to record the thinning interval that is applied. @@ -94,11 +100,16 @@ def _thin_data(self, group, params, thin_interval): The list of dataset names to thin. thin_interval : int The interval to thin the samples on disk by. + """ - samples = self.read_raw_samples(params, thin_start=0, - thin_interval=thin_interval, - thin_end=None, flatten=False, - group=group) + samples = self.read_raw_samples( + params, + thin_start=0, + thin_interval=thin_interval, + thin_end=None, + flatten=False, + group=group, + ) # now resize and write the data back to disk fpgroup = self[group] for param in params: @@ -109,7 +120,8 @@ def _thin_data(self, group, params, thin_interval): fpgroup[param][:] = data def thin(self, thin_interval): - """Thins the samples on disk to the given thinning interval. + """ + Thins the samples on disk to the given thinning interval. The interval must be a multiple of the file's current ``thinned_by``. @@ -117,13 +129,15 @@ def thin(self, thin_interval): ---------- thin_interval : int The interval the samples on disk should be thinned by. + """ # get the new interval to thin by new_interval = thin_interval / self.thinned_by if new_interval % 1: - raise ValueError("thin interval ({}) must be a multiple of the " - "current thinned_by ({})" - .format(thin_interval, self.thinned_by)) + raise ValueError( + f"thin interval ({thin_interval}) must be a multiple of the " + f"current thinned_by ({self.thinned_by})" + ) new_interval = int(new_interval) # now thin the data on disk params = list(self[self.samples_group].keys()) @@ -133,29 +147,32 @@ def thin(self, thin_interval): @property def thinned_by(self): - """Returns interval samples have been thinned by on disk. + """ + Returns interval samples have been thinned by on disk. This looks for ``thinned_by`` in the samples group attrs. If none is found, will just return 1. """ try: - thinned_by = self.attrs['thinned_by'] + thinned_by = self.attrs["thinned_by"] except KeyError: thinned_by = 1 return thinned_by @thinned_by.setter def thinned_by(self, thinned_by): - """Sets the thinned_by attribute. + """ + Sets the thinned_by attribute. This is the interval that samples have been thinned by on disk. The given value is written to ``self[self.samples_group].attrs['thinned_by']``. """ - self.attrs['thinned_by'] = int(thinned_by) + self.attrs["thinned_by"] = int(thinned_by) def last_iteration(self, parameter=None, group=None): - """Returns the iteration of the last sample of the given parameter. + """ + Returns the iteration of the last sample of the given parameter. Parameters ---------- @@ -165,6 +182,7 @@ def last_iteration(self, parameter=None, group=None): group : str, optional The name of the group to get the last iteration from. Default is the ``samples_group``. + """ if group is None: group = self.samples_group @@ -188,39 +206,42 @@ def iterations(self, parameter): def write_sampler_metadata(self, sampler): """Writes the sampler's metadata.""" - self.attrs['sampler'] = sampler.name + self.attrs["sampler"] = sampler.name try: - self[self.sampler_group].attrs['nchains'] = sampler.nchains + self[self.sampler_group].attrs["nchains"] = sampler.nchains except ValueError: - self[self.sampler_group].attrs['nwalkers'] = sampler.nwalkers + self[self.sampler_group].attrs["nwalkers"] = sampler.nwalkers # write the model's metadata sampler.model.write_metadata(self) @property def is_burned_in(self): - """Returns whether or not chains are burned in. + """ + Returns whether or not chains are burned in. Raises a ``ValueError`` if no burn in tests were done. """ try: - return self[self.sampler_group]['is_burned_in'][()] + return self[self.sampler_group]["is_burned_in"][()] except KeyError: raise ValueError("No burn in tests were performed") @property def burn_in_iteration(self): - """Returns the burn in iteration of all the chains. + """ + Returns the burn in iteration of all the chains. Raises a ``ValueError`` if no burn in tests were done. """ try: - return self[self.sampler_group]['burn_in_iteration'][()] + return self[self.sampler_group]["burn_in_iteration"][()] except KeyError: raise ValueError("No burn in tests were performed") @property def burn_in_index(self): - """Returns the burn in index. + """ + Returns the burn in index. This is the burn in iteration divided by the file's ``thinned_by``. Requires the class that this is used with has a ``burn_in_iteration`` @@ -230,19 +251,21 @@ def burn_in_index(self): @property def act(self): - """The autocorrelation time (ACT). + """ + The autocorrelation time (ACT). This is the ACL times the file's thinned by. Raises a ``ValueError`` if the ACT has not been calculated. """ try: - return self[self.sampler_group]['act'][()] + return self[self.sampler_group]["act"][()] except KeyError: raise ValueError("ACT has not been calculated") @act.setter def act(self, act): - """Writes the autocorrelation time(s). + """ + Writes the autocorrelation time(s). ACT(s) are written to the ``sample_group`` as a dataset with name ``act``. @@ -251,13 +274,15 @@ def act(self, act): ---------- act : array or int ACT(s) to write. + """ # pylint: disable=no-member - self.write_data('act', act, path=self.sampler_group) + self.write_data("act", act, path=self.sampler_group) @property def raw_acts(self): - """Dictionary of parameter names -> raw autocorrelation time(s). + """ + Dictionary of parameter names -> raw autocorrelation time(s). Depending on the sampler, the autocorrelation times may be floats, or [ntemps x] [nchains x] arrays. @@ -265,7 +290,7 @@ def raw_acts(self): Raises a ``ValueError`` is no raw acts have been set. """ try: - group = self[self.sampler_group]['raw_acts'] + group = self[self.sampler_group]["raw_acts"] except KeyError: raise ValueError("ACTs have not been calculated") acts = {} @@ -275,7 +300,8 @@ def raw_acts(self): @raw_acts.setter def raw_acts(self, acts): - """Writes the raw autocorrelation times. + """ + Writes the raw autocorrelation times. The ACT of each parameter is saved to ``[sampler_group]/raw_acts/{param}']``. Works for all types of MCMC @@ -285,14 +311,16 @@ def raw_acts(self, acts): ---------- acts : dict A dictionary of ACTs keyed by the parameter. + """ - path = self.sampler_group + '/raw_acts' + path = self.sampler_group + "/raw_acts" for param in acts: self.write_data(param, acts[param], path=path) @property def acl(self): - """The autocorrelation length (ACL) of the samples. + """ + The autocorrelation length (ACL) of the samples. This is the autocorrelation time (ACT) divided by the file's ``thinned_by`` attribute. Raises a ``ValueError`` if the ACT has not @@ -302,7 +330,8 @@ def acl(self): @acl.setter def acl(self, acl): - """Sets the autocorrelation length (ACL) of the samples. + """ + Sets the autocorrelation length (ACL) of the samples. This will convert the given value(s) to autocorrelation time(s) and save to the ``act`` attribute; see that attribute for details. @@ -311,7 +340,8 @@ def acl(self, acl): @property def raw_acls(self): - """Dictionary of parameter names -> raw autocorrelation length(s). + """ + Dictionary of parameter names -> raw autocorrelation length(s). Depending on the sampler, the autocorrelation lengths may be floats, or [ntemps x] [nchains x] arrays. @@ -324,7 +354,8 @@ def raw_acls(self): @raw_acls.setter def raw_acls(self, acls): - """Sets the raw autocorrelation lengths. + """ + Sets the raw autocorrelation lengths. The given ACLs are converted to autocorrelation times (ACTs) and saved to the ``raw_acts`` attribute; see that attribute for details. @@ -333,19 +364,21 @@ def raw_acls(self, acls): ---------- acls : dict A dictionary of ACLs keyed by the parameter. + """ self.raw_acts = {p: acls[p] * self.thinned_by for p in acls} def _update_sampler_history(self): - """Writes the number of iterations, effective number of samples, + """ + Writes the number of iterations, effective number of samples, autocorrelation times, and burn-in iteration to the history. """ - path = '/'.join([self.sampler_group, 'checkpoint_history']) + path = "/".join([self.sampler_group, "checkpoint_history"]) # write the current number of iterations - self.write_data('niterations', self.niterations, path=path, - append=True) - self.write_data('effective_nsamples', self.effective_nsamples, - path=path, append=True) + self.write_data("niterations", self.niterations, path=path, append=True) + self.write_data( + "effective_nsamples", self.effective_nsamples, path=path, append=True + ) # write the act: we'll make sure that this is 2D, so that the acts # can be appened along the last dimension try: @@ -354,8 +387,8 @@ def _update_sampler_history(self): # no acts were calculate act = None if act is not None: - act = act.reshape(tuple(list(act.shape)+[1])) - self.write_data('act', act, path=path, append=True) + act = act.reshape(tuple(list(act.shape) + [1])) + self.write_data("act", act, path=path, append=True) # write the burn in iteration in the same way try: burn_in = self.burn_in_iteration @@ -363,13 +396,13 @@ def _update_sampler_history(self): # no burn in tests were done burn_in = None if burn_in is not None: - burn_in = burn_in.reshape(tuple(list(burn_in.shape)+[1])) - self.write_data('burn_in_iteration', burn_in, path=path, - append=True) + burn_in = burn_in.reshape(tuple(list(burn_in.shape) + [1])) + self.write_data("burn_in_iteration", burn_in, path=path, append=True) @staticmethod def extra_args_parser(parser=None, skip_args=None, **kwargs): - r"""Create a parser to parse sampler-specific arguments for loading + r""" + Create a parser to parse sampler-specific arguments for loading samples. Parameters @@ -392,63 +425,85 @@ def extra_args_parser(parser=None, skip_args=None, **kwargs): An argument parser with th extra arguments added. actions : list of argparse.Action A list of the actions that were added. + """ if parser is None: parser = argparse.ArgumentParser(**kwargs) elif kwargs: - raise ValueError("No other keyword arguments should be provded if " - "a parser is provided.") + raise ValueError( + "No other keyword arguments should be provded if a parser is provided." + ) if skip_args is None: skip_args = [] actions = [] - if 'thin-start' not in skip_args: + if "thin-start" not in skip_args: act = parser.add_argument( - "--thin-start", type=int, default=None, + "--thin-start", + type=int, + default=None, help="Sample number to start collecting samples. If " - "none provided, will use the input file's `thin_start` " - "attribute.") + "none provided, will use the input file's `thin_start` " + "attribute.", + ) actions.append(act) - if 'thin-interval' not in skip_args: + if "thin-interval" not in skip_args: act = parser.add_argument( - "--thin-interval", type=int, default=None, + "--thin-interval", + type=int, + default=None, help="Interval to use for thinning samples. If none provided, " - "will use the input file's `thin_interval` attribute.") + "will use the input file's `thin_interval` attribute.", + ) actions.append(act) - if 'thin-end' not in skip_args: + if "thin-end" not in skip_args: act = parser.add_argument( - "--thin-end", type=int, default=None, + "--thin-end", + type=int, + default=None, help="Sample number to stop collecting samples. If " - "none provided, will use the input file's `thin_end` " - "attribute.") + "none provided, will use the input file's `thin_end` " + "attribute.", + ) actions.append(act) - if 'iteration' not in skip_args: + if "iteration" not in skip_args: act = parser.add_argument( - "--iteration", type=int, default=None, + "--iteration", + type=int, + default=None, help="Only retrieve the given iteration. To load " - "the last n-th sampe use -n, e.g., -1 will " - "load the last iteration. This overrides " - "the thin-start/interval/end options.") + "the last n-th sampe use -n, e.g., -1 will " + "load the last iteration. This overrides " + "the thin-start/interval/end options.", + ) actions.append(act) - if 'walkers' not in skip_args and 'chains' not in skip_args: + if "walkers" not in skip_args and "chains" not in skip_args: act = parser.add_argument( - "--walkers", "--chains", type=int, nargs="+", default=None, + "--walkers", + "--chains", + type=int, + nargs="+", + default=None, help="Only retrieve samples from the listed " - "walkers/chains. Default is to retrieve from all " - "walkers/chains.") + "walkers/chains. Default is to retrieve from all " + "walkers/chains.", + ) actions.append(act) return parser, actions -class MCMCMetadataIO(object): - """Provides functions for reading/writing metadata to file for MCMCs in +class MCMCMetadataIO: + """ + Provides functions for reading/writing metadata to file for MCMCs in which all chains are independent of each other. Overrides the ``BaseInference`` file's ``thin_start`` and ``thin_interval`` attributes. Instead of integers, these return arrays. """ + @property def thin_start(self): - """Returns the default thin start to use for reading samples. + """ + Returns the default thin start to use for reading samples. If burn-in tests were done, this will return the burn-in index of every chain that has burned in. The start index for chains that have not @@ -462,8 +517,9 @@ def thin_start(self): # replace any that have not been burned in with the number # of iterations; this will cause those chains to not return # any samples - thin_start[~self.is_burned_in] = \ - int(numpy.ceil(self.niterations/self.thinned_by)) + thin_start[~self.is_burned_in] = int( + numpy.ceil(self.niterations / self.thinned_by) + ) return thin_start except ValueError: # no burn in, just return array of zeros @@ -471,7 +527,8 @@ def thin_start(self): @property def thin_interval(self): - """Returns the default thin interval to use for reading samples. + """ + Returns the default thin interval to use for reading samples. If a finite ACL exists in the file, will return that. Otherwise, returns 1. @@ -485,13 +542,16 @@ def thin_interval(self): return numpy.ceil(acl).astype(int) -class EnsembleMCMCMetadataIO(object): - """Provides functions for reading/writing metadata to file for ensemble +class EnsembleMCMCMetadataIO: + """ + Provides functions for reading/writing metadata to file for ensemble MCMCs. """ + @property def thin_start(self): - """Returns the default thin start to use for reading samples. + """ + Returns the default thin start to use for reading samples. If burn-in tests were done, returns the burn in index. Otherwise, returns 0. @@ -504,7 +564,8 @@ def thin_start(self): @property def thin_interval(self): - """Returns the default thin interval to use for reading samples. + """ + Returns the default thin interval to use for reading samples. If a finite ACL exists in the file, will return that. Otherwise, returns 1. @@ -520,9 +581,11 @@ def thin_interval(self): return acl -def write_samples(fp, samples, parameters=None, last_iteration=None, - samples_group=None, thin_by=None): - """Writes samples to the given file. +def write_samples( + fp, samples, parameters=None, last_iteration=None, samples_group=None, thin_by=None +): + """ + Writes samples to the given file. This works for both standard MCMC and ensemble MCMC samplers without parallel tempering. @@ -541,7 +604,7 @@ def write_samples(fp, samples, parameters=None, last_iteration=None, sample, then none of the samples will be written. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to write files to. Must be an instance of BaseInferenceFile with CommonMCMCMetadataIO methods added. @@ -564,21 +627,22 @@ def write_samples(fp, samples, parameters=None, last_iteration=None, Override the ``thinned_by`` attribute in the file with the given value. **Only set this if you are using this function to write something other than inference samples!** + """ nwalkers, nsamples = list(samples.values())[0].shape - assert all(p.shape == (nwalkers, nsamples) - for p in samples.values()), ( - "all samples must have the same shape") + assert all(p.shape == (nwalkers, nsamples) for p in samples.values()), ( + "all samples must have the same shape" + ) if samples_group is None: samples_group = fp.samples_group if parameters is None: parameters = samples.keys() # thin the samples - samples = thin_samples_for_writing(fp, samples, parameters, - last_iteration, samples_group, - thin_by=thin_by) + samples = thin_samples_for_writing( + fp, samples, parameters, last_iteration, samples_group, thin_by=thin_by + ) # loop over number of dimensions - group = samples_group + '/{name}' + group = samples_group + "/{name}" for param in parameters: dataset_name = group.format(name=param) data = samples[param] @@ -597,22 +661,33 @@ def write_samples(fp, samples, parameters=None, last_iteration=None, # dataset doesn't exist yet istart = 0 istop = istart + data.shape[1] - fp.create_dataset(dataset_name, (nwalkers, istop), - maxshape=(nwalkers, None), - dtype=data.dtype, - fletcher32=True) + fp.create_dataset( + dataset_name, + (nwalkers, istop), + maxshape=(nwalkers, None), + dtype=data.dtype, + fletcher32=True, + ) fp[dataset_name][:, istart:istop] = data -def ensemble_read_raw_samples(fp, fields, thin_start=None, - thin_interval=None, thin_end=None, - iteration=None, walkers=None, flatten=True, - group=None): - """Base function for reading samples from ensemble MCMC files without +def ensemble_read_raw_samples( + fp, + fields, + thin_start=None, + thin_interval=None, + thin_end=None, + iteration=None, + walkers=None, + flatten=True, + group=None, +): + """ + Base function for reading samples from ensemble MCMC files without parallel tempering. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to write files to. Must be an instance of BaseInferenceFile with EnsembleMCMCMetadataIO methods added. @@ -643,18 +718,18 @@ def ensemble_read_raw_samples(fp, fields, thin_start=None, ------- dict A dictionary of field name -> numpy array pairs. + """ if isinstance(fields, str): fields = [fields] # walkers to load widx, nwalkers = _ensemble_get_walker_index(fp, walkers) # get the slice to use - get_index = _ensemble_get_index(fp, thin_start, thin_interval, thin_end, - iteration) + get_index = _ensemble_get_index(fp, thin_start, thin_interval, thin_end, iteration) # load if group is None: group = fp.samples_group - group = group + '/{name}' + group = group + "/{name}" arrays = {} for name in fields: arr = fp[group.format(name=name)][widx, get_index] @@ -669,7 +744,8 @@ def ensemble_read_raw_samples(fp, fields, thin_start=None, def _ensemble_get_walker_index(fp, walkers=None): - """Convenience function to determine which walkers to load. + """ + Convenience function to determine which walkers to load. Parameters ---------- @@ -685,6 +761,7 @@ def _ensemble_get_walker_index(fp, walkers=None): The walker indices to load. nwalkers : int The number of walkers that will be loaded. + """ if walkers is not None: widx = numpy.zeros(fp.nwalkers, dtype=bool) @@ -696,12 +773,14 @@ def _ensemble_get_walker_index(fp, walkers=None): return widx, nwalkers -def _ensemble_get_index(fp, thin_start=None, thin_interval=None, thin_end=None, - iteration=None): - """Determines the sample indices to retrieve for an ensemble MCMC. +def _ensemble_get_index( + fp, thin_start=None, thin_interval=None, thin_end=None, iteration=None +): + """ + Determines the sample indices to retrieve for an ensemble MCMC. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to write files to. Must be an instance of BaseInferenceFile with EnsembleMCMCMetadataIO methods added. @@ -721,6 +800,7 @@ def _ensemble_get_index(fp, thin_start=None, thin_interval=None, thin_end=None, ------- slice or int The indices to retrieve. + """ if iteration is not None: get_index = int(iteration) @@ -731,19 +811,21 @@ def _ensemble_get_index(fp, thin_start=None, thin_interval=None, thin_end=None, thin_interval = fp.thin_interval if thin_end is None: thin_end = fp.thin_end - get_index = fp.get_slice(thin_start=thin_start, - thin_interval=thin_interval, - thin_end=thin_end) + get_index = fp.get_slice( + thin_start=thin_start, thin_interval=thin_interval, thin_end=thin_end + ) return get_index -def _get_index(fp, chains, thin_start=None, thin_interval=None, thin_end=None, - iteration=None): - """Determines the sample indices to retrieve for an MCMC with independent +def _get_index( + fp, chains, thin_start=None, thin_interval=None, thin_end=None, iteration=None +): + """ + Determines the sample indices to retrieve for an MCMC with independent chains. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to read samples from. Must be an instance of BaseInferenceFile with EnsembleMCMCMetadataIO methods added. @@ -778,27 +860,32 @@ def _get_index(fp, chains, thin_start=None, thin_interval=None, thin_end=None, ------- get_index : list of slice or int The indices to retrieve. + """ nchains = len(chains) # convenience function to get the right thin start/interval/end if iteration is not None: - get_index = [int(iteration)]*nchains + get_index = [int(iteration)] * nchains else: # get the slice arguments thin_start = _format_slice_arg(thin_start, fp.thin_start, chains) - thin_interval = _format_slice_arg(thin_interval, fp.thin_interval, - chains) + thin_interval = _format_slice_arg(thin_interval, fp.thin_interval, chains) thin_end = _format_slice_arg(thin_end, fp.thin_end, chains) # the slices to use for each chain - get_index = [fp.get_slice(thin_start=thin_start[ci], - thin_interval=thin_interval[ci], - thin_end=thin_end[ci]) - for ci in range(nchains)] + get_index = [ + fp.get_slice( + thin_start=thin_start[ci], + thin_interval=thin_interval[ci], + thin_end=thin_end[ci], + ) + for ci in range(nchains) + ] return get_index def _format_slice_arg(value, default, chains): - """Formats a start/interval/end argument for picking out chains. + """ + Formats a start/interval/end argument for picking out chains. Parameters ---------- @@ -818,11 +905,12 @@ def _format_slice_arg(value, default, chains): array Array giving the value to use for each chain in ``chains``. The array will have the same length as ``chains``. + """ if value is None and default is None: # no value provided, and default is None, just return Nones with the # same length as chains - value = [None]*len(chains) + value = [None] * len(chains) elif value is None: # use the default, with the desired values extracted value = default[chains] @@ -832,15 +920,18 @@ def _format_slice_arg(value, default, chains): elif len(value) != len(chains): # a list of values was provided, but the length does not match the # chains, raise an error - raise ValueError("Number of requested thin-start/interval/end values " - "({}) does not match number of requested chains ({})" - .format(len(value), len(chains))) + raise ValueError( + "Number of requested thin-start/interval/end values " + f"({len(value)}) does not match number of requested chains ({len(chains)})" + ) return value -def thin_samples_for_writing(fp, samples, parameters, last_iteration, - group, thin_by=None): - """Thins samples for writing to disk. +def thin_samples_for_writing( + fp, samples, parameters, last_iteration, group, thin_by=None +): + """ + Thins samples for writing to disk. The thinning interval to use is determined by the given file handler's ``thinned_by`` attribute. If that attribute is 1, just returns the samples. @@ -873,14 +964,16 @@ def thin_samples_for_writing(fp, samples, parameters, last_iteration, ------- dict : Dictionary of the thinned samples to write. + """ if thin_by is None: thin_by = fp.thinned_by if thin_by > 1: if last_iteration is None: - raise ValueError("File's thinned_by attribute is > 1 ({}), " - "but last_iteration not provided." - .format(thin_by)) + raise ValueError( + f"File's thinned_by attribute is > 1 ({thin_by}), " + "but last_iteration not provided." + ) thinned_samples = {} for param in parameters: data = samples[param] @@ -892,8 +985,12 @@ def thin_samples_for_writing(fp, samples, parameters, last_iteration, # sample in samples. Subtracting the latter from the former - 1 # (-1 to convert from iteration to index) therefore gives the index # in the samples data to start using samples. - thin_start = fp.last_iteration(param, group) + thin_by \ - - (last_iteration - nsamples) - 1 + thin_start = ( + fp.last_iteration(param, group) + + thin_by + - (last_iteration - nsamples) + - 1 + ) thinned_samples[param] = data[..., thin_start::thin_by] else: thinned_samples = samples @@ -901,7 +998,8 @@ def thin_samples_for_writing(fp, samples, parameters, last_iteration, def nsamples_in_chain(start_iter, interval, niterations): - """Calculates the number of samples in an MCMC chain given a thinning + """ + Calculates the number of samples in an MCMC chain given a thinning start, end, and interval. This function will work with either python scalars, or numpy arrays. @@ -922,6 +1020,7 @@ def nsamples_in_chain(start_iter, interval, niterations): ------- num_samples : (array of) numpy.int The number of samples in a chain, >= 0. + """ # this is written in a slightly wonky way so that it will work with either # python scalars or numpy arrays; it is equivalent to: @@ -931,9 +1030,9 @@ def nsamples_in_chain(start_iter, interval, niterations): # count = max(niterations - start_iter, 0) slt0 = start_iter < 0 sgt0 = start_iter >= 0 - count = slt0*abs(start_iter) + sgt0*(niterations - start_iter) + count = slt0 * abs(start_iter) + sgt0 * (niterations - start_iter) # ensure count is in [0, niterations] cgtn = count > niterations cok = (count >= 0) & (count <= niterations) - count = cgtn*niterations + cok*count + count = cgtn * niterations + cok * count return numpy.ceil(count / interval).astype(int) diff --git a/pycbc/inference/io/base_multitemper.py b/pycbc/inference/io/base_multitemper.py index 572391d4dd3..d02d212bdbe 100644 --- a/pycbc/inference/io/base_multitemper.py +++ b/pycbc/inference/io/base_multitemper.py @@ -21,38 +21,45 @@ # # ============================================================================= # -"""Provides I/O support for multi-tempered sampler. -""" +"""Provides I/O support for multi-tempered sampler.""" import argparse + import numpy -from .base_mcmc import (CommonMCMCMetadataIO, thin_samples_for_writing, - _ensemble_get_index, _ensemble_get_walker_index, - _get_index) + +from .base_mcmc import ( + CommonMCMCMetadataIO, + _ensemble_get_index, + _ensemble_get_walker_index, + _get_index, + thin_samples_for_writing, +) + class ParseTempsArg(argparse.Action): - """Argparse action that will parse temps argument. + """ + Argparse action that will parse temps argument. If the provided argument is 'all', sets 'all' in the namespace dest. If a a sequence of numbers are provided, converts those numbers to ints before saving to the namespace. """ - def __init__(self, type=str, **kwargs): # pylint: disable=redefined-builtin + + def __init__(self, type=str, **kwargs): # pylint: disable=redefined-builtin # check that type is string if type != str: raise ValueError("the type for this action must be a string") - super(ParseTempsArg, self).__init__(type=type, **kwargs) + super().__init__(type=type, **kwargs) def __call__(self, parser, namespace, values, option_string=None): singlearg = isinstance(values, str) if singlearg: values = [values] - if values[0] == 'all': + if values[0] == "all": # check that only a single value was provided if len(values) > 1: - raise ValueError("if provide 'all', should not specify any " - "other temps") - temps = 'all' + raise ValueError("if provide 'all', should not specify any other temps") + temps = "all" else: temps = [] for val in values: @@ -67,43 +74,49 @@ def __call__(self, parser, namespace, values, option_string=None): class CommonMultiTemperedMetadataIO(CommonMCMCMetadataIO): - """Adds support for reading/writing multi-tempered metadata to + """ + Adds support for reading/writing multi-tempered metadata to :py:class:`~pycbc.inference.io.base_mcmc.CommonMCMCMetadatIO`. """ + @property def ntemps(self): """Returns the number of temperatures used by the sampler.""" - return self[self.sampler_group].attrs['ntemps'] + return self[self.sampler_group].attrs["ntemps"] def write_sampler_metadata(self, sampler): - """Adds writing ntemps to file. - """ - super(CommonMultiTemperedMetadataIO, self).write_sampler_metadata( - sampler) + """Adds writing ntemps to file.""" + super().write_sampler_metadata(sampler) self[self.sampler_group].attrs["ntemps"] = sampler.ntemps @staticmethod def extra_args_parser(parser=None, skip_args=None, **kwargs): - """Adds --temps to MCMCIO parser. - """ + """Adds --temps to MCMCIO parser.""" if skip_args is None: skip_args = [] parser, actions = CommonMCMCMetadataIO.extra_args_parser( - parser=parser, skip_args=skip_args, **kwargs) - if 'temps' not in skip_args: + parser=parser, skip_args=skip_args, **kwargs + ) + if "temps" not in skip_args: act = parser.add_argument( - "--temps", nargs="+", default=0, action=ParseTempsArg, + "--temps", + nargs="+", + default=0, + action=ParseTempsArg, help="Get the given temperatures. May provide either a " - "sequence of integers specifying the temperatures to " - "plot, or 'all' for all temperatures. Default is to only " - "plot the coldest (= 0) temperature chain.") + "sequence of integers specifying the temperatures to " + "plot, or 'all' for all temperatures. Default is to only " + "plot the coldest (= 0) temperature chain.", + ) actions.append(act) return parser, actions -def write_samples(fp, samples, parameters=None, last_iteration=None, - samples_group=None, thin_by=None): - """Writes samples to the given file. +def write_samples( + fp, samples, parameters=None, last_iteration=None, samples_group=None, thin_by=None +): + """ + Writes samples to the given file. This works both for standard MCMC and ensemble MCMC samplers with parallel tempering. @@ -113,7 +126,7 @@ def write_samples(fp, samples, parameters=None, last_iteration=None, ``ntemps x nwalkers x niterations`` array. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to write files to. Must be an instance of BaseInferenceFile with CommonMultiTemperedMetadataIO methods added. @@ -134,21 +147,22 @@ def write_samples(fp, samples, parameters=None, last_iteration=None, Override the ``thinned_by`` attribute in the file with the given value. **Only set this if you are using this function to write something other than inference samples!** + """ ntemps, nwalkers, niterations = tuple(samples.values())[0].shape - assert all(p.shape == (ntemps, nwalkers, niterations) - for p in samples.values()), ( - "all samples must have the same shape") + assert all(p.shape == (ntemps, nwalkers, niterations) for p in samples.values()), ( + "all samples must have the same shape" + ) if samples_group is None: samples_group = fp.samples_group if parameters is None: parameters = list(samples.keys()) # thin the samples - samples = thin_samples_for_writing(fp, samples, parameters, - last_iteration, samples_group, - thin_by=thin_by) + samples = thin_samples_for_writing( + fp, samples, parameters, last_iteration, samples_group, thin_by=thin_by + ) # loop over number of dimensions - group = samples_group + '/{name}' + group = samples_group + "/{name}" for param in parameters: dataset_name = group.format(name=param) data = samples[param] @@ -167,18 +181,30 @@ def write_samples(fp, samples, parameters=None, last_iteration=None, # dataset doesn't exist yet istart = 0 istop = istart + data.shape[2] - fp.create_dataset(dataset_name, (ntemps, nwalkers, istop), - maxshape=(ntemps, nwalkers, None), - dtype=data.dtype, - fletcher32=True) + fp.create_dataset( + dataset_name, + (ntemps, nwalkers, istop), + maxshape=(ntemps, nwalkers, None), + dtype=data.dtype, + fletcher32=True, + ) fp[dataset_name][:, :, istart:istop] = data -def read_raw_samples(fp, fields, - thin_start=None, thin_interval=None, thin_end=None, - iteration=None, temps='all', chains=None, - flatten=True, group=None): - """Base function for reading samples from a collection of independent +def read_raw_samples( + fp, + fields, + thin_start=None, + thin_interval=None, + thin_end=None, + iteration=None, + temps="all", + chains=None, + flatten=True, + group=None, +): + """ + Base function for reading samples from a collection of independent MCMC chains file with parallel tempering. This may collect differing numbering of samples from each chains, @@ -189,7 +215,7 @@ def read_raw_samples(fp, fields, ``numpy.nan``. If flattened, the NaNs are removed prior to returning. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to read samples from. Must be an instance of BaseInferenceFile with CommonMultiTemperedMetadataIO methods added. @@ -238,19 +264,19 @@ def read_raw_samples(fp, fields, ------- dict A dictionary of field name -> numpy array pairs. + """ if isinstance(fields, str): fields = [fields] if group is None: group = fp.samples_group - group = group + '/{name}' + group = group + "/{name}" # chains to load if chains is None: chains = numpy.arange(fp.nchains) elif not isinstance(chains, (list, numpy.ndarray)): chains = numpy.array([chains]).astype(int) - get_index = _get_index(fp, chains, thin_start, thin_interval, thin_end, - iteration) + get_index = _get_index(fp, chains, thin_start, thin_interval, thin_end, iteration) # load the samples arrays = {} for name in fields: @@ -269,7 +295,7 @@ def read_raw_samples(fp, fields, continue if isinstance(idx, (int, numpy.int_)): # make sure the last dimension corresponds to iteration - thisarr = thisarr.reshape(list(thisarr.shape)+[1]) + thisarr = thisarr.reshape(list(thisarr.shape) + [1]) # pull out the temperatures we need if selecttemps: thisarr = thisarr[temps, ...] @@ -278,11 +304,12 @@ def read_raw_samples(fp, fields, alist.append(thisarr) maxiters = max(maxiters, thisarr.shape[-1]) # stack into a single array - arr = numpy.full((ntemps, len(chains), maxiters), numpy.nan, - dtype=fp[dset].dtype) + arr = numpy.full( + (ntemps, len(chains), maxiters), numpy.nan, dtype=fp[dset].dtype + ) for ii, thisarr in enumerate(alist): if thisarr is not None: - arr[:, ii, :thisarr.shape[-1]] = thisarr + arr[:, ii, : thisarr.shape[-1]] = thisarr if flatten: # flatten and remove nans arr = arr.flatten() @@ -291,15 +318,24 @@ def read_raw_samples(fp, fields, return arrays -def ensemble_read_raw_samples(fp, fields, thin_start=None, - thin_interval=None, thin_end=None, - iteration=None, temps='all', walkers=None, - flatten=True, group=None): - """Base function for reading samples from ensemble MCMC file with +def ensemble_read_raw_samples( + fp, + fields, + thin_start=None, + thin_interval=None, + thin_end=None, + iteration=None, + temps="all", + walkers=None, + flatten=True, + group=None, +): + """ + Base function for reading samples from ensemble MCMC file with parallel tempering. Parameters - ----------- + ---------- fp : BaseInferenceFile Open file handler to write files to. Must be an instance of BaseInferenceFile with CommonMultiTemperedMetadataIO methods added. @@ -334,18 +370,18 @@ def ensemble_read_raw_samples(fp, fields, thin_start=None, ------- dict A dictionary of field name -> numpy array pairs. + """ if isinstance(fields, str): fields = [fields] # walkers to load widx, nwalkers = _ensemble_get_walker_index(fp, walkers) # get the slice to use - get_index = _ensemble_get_index(fp, thin_start, thin_interval, thin_end, - iteration) + get_index = _ensemble_get_index(fp, thin_start, thin_interval, thin_end, iteration) # load if group is None: group = fp.samples_group - group = group + '/{name}' + group = group + "/{name}" arrays = {} for name in fields: dset = group.format(name=name) @@ -365,10 +401,11 @@ def ensemble_read_raw_samples(fp, fields, thin_start=None, def _get_temps_index(temps, fp, dataset): - """Convenience function to determine which temperatures to load. + """ + Convenience function to determine which temperatures to load. Parameters - ----------- + ---------- temps : 'all' or (list of) int The temperature index (or list of indices) to retrieve. To retrieve all temperates pass 'all', or a list of all of the temperatures. @@ -387,8 +424,9 @@ def _get_temps_index(temps, fp, dataset): array after it is loaded from the file. ntemps : int The number of temperatures that will be loaded. + """ - if temps == 'all': + if temps == "all": # all temperatures were requested; just need to know how many ntemps = fp[dataset].shape[0] tidx = slice(None, None) diff --git a/pycbc/inference/io/base_nested_sampler.py b/pycbc/inference/io/base_nested_sampler.py index c09f7de5013..f3155b2ee9e 100644 --- a/pycbc/inference/io/base_nested_sampler.py +++ b/pycbc/inference/io/base_nested_sampler.py @@ -21,8 +21,8 @@ # # ============================================================================= # -"""Provides IO for the dynesty sampler. -""" +"""Provides IO for the dynesty sampler.""" + from .base_sampler import BaseSamplerFile from .posterior import read_raw_samples_from_file, write_samples_to_file @@ -30,7 +30,7 @@ class BaseNestedSamplerFile(BaseSamplerFile): """Class to handle file IO for the nested samplers cpnest and dynesty.""" - name = 'base_nest_file' + name = "base_nest_file" def read_raw_samples(self, fields, **kwargs): return read_raw_samples_from_file(self, fields, **kwargs) @@ -42,35 +42,37 @@ def write_niterations(self, niterations): """ Writes the given number of iterations to the sampler group. """ - self[self.sampler_group].attrs['niterations'] = niterations + self[self.sampler_group].attrs["niterations"] = niterations def write_sampler_metadata(self, sampler): """ Adds writing betas to MultiTemperedMCMCIO. """ - self.attrs['sampler'] = sampler.name + self.attrs["sampler"] = sampler.name if self.sampler_group not in self.keys(): # create the sampler group self.create_group(self.sampler_group) - self[self.sampler_group].attrs['nlivepoints'] = sampler.nlive + self[self.sampler_group].attrs["nlivepoints"] = sampler.nlive # write the model's metadata sampler.model.write_metadata(self) def write_samples(self, samples, parameters=None): - """Writes samples to the given file. + """ + Writes samples to the given file. Results are written to ``samples_group/{vararg}``, where ``{vararg}`` is the name of a model params. The samples are written as an array of length ``niterations``. Parameters - ----------- + ---------- samples : dict The samples to write. Each array in the dictionary should have length niterations. parameters : list, optional Only write the specified parameters to the file. If None, will write all of the keys in the ``samples`` dict. + """ # since we're just writing a posterior use # PosteriorFile's write_samples diff --git a/pycbc/inference/io/base_sampler.py b/pycbc/inference/io/base_sampler.py index b6b1611c65c..2fec288bc1f 100644 --- a/pycbc/inference/io/base_sampler.py +++ b/pycbc/inference/io/base_sampler.py @@ -15,22 +15,23 @@ """Provides abstract base class for all samplers.""" - import time -from abc import (ABCMeta, abstractmethod) +from abc import ABCMeta, abstractmethod from .base_hdf import BaseInferenceFile class BaseSamplerFile(BaseInferenceFile, metaclass=ABCMeta): - """Base HDF class for all samplers. + """ + Base HDF class for all samplers. This adds abstract methods ``write_resume_point`` and ``write_sampler_metadata`` to :py:class:`BaseInferenceFile`. """ def write_run_start_time(self): - """Writes the current (UNIX) time to the file. + """ + Writes the current (UNIX) time to the file. Times are stored as a list in the file's ``attrs``, with name ``run_start_time``. If the attrbute already exists, the current time @@ -46,44 +47,44 @@ def write_run_start_time(self): @property def run_start_time(self): - """The (UNIX) time pycbc inference began running. + """ + The (UNIX) time pycbc inference began running. If the run resumed from a checkpoint, the time the last checkpoint started is reported. """ - return self.attrs['run_start_time'][-1] + return self.attrs["run_start_time"][-1] def write_run_end_time(self): - """"Writes the curent (UNIX) time as the ``run_end_time`` attribute. - """ + """ "Writes the curent (UNIX) time as the ``run_end_time`` attribute.""" self.attrs["run_end_time"] = time.time() @property def run_end_time(self): - """The (UNIX) time pycbc inference finished. - """ + """The (UNIX) time pycbc inference finished.""" return self.attrs["run_end_time"] @abstractmethod def write_resume_point(self): - """Should write the point that a sampler starts up. + """ + Should write the point that a sampler starts up. How the resume point is indexed is up to the sampler. For example, MCMC samplers use the number of iterations that are stored in the checkpoint file. """ - pass @abstractmethod def write_sampler_metadata(self, sampler): - """This should write the given sampler's metadata to the file. + """ + This should write the given sampler's metadata to the file. This should also include the model's metadata. """ - pass def update_checkpoint_history(self): - """Writes a copy of relevant metadata to the file's checkpoint history. + """ + Writes a copy of relevant metadata to the file's checkpoint history. All data are written to ``sampler_info/checkpoint_history``. If the group does not exist yet, it will be created. @@ -92,7 +93,7 @@ def update_checkpoint_history(self): checkpoint to the file. It will also call :py:func:`_update_sampler_history` to write sampler-specific history. """ - path = '/'.join([self.sampler_group, 'checkpoint_history']) + path = "/".join([self.sampler_group, "checkpoint_history"]) try: history = self[path] except KeyError: @@ -101,10 +102,9 @@ def update_checkpoint_history(self): history = self[path] # write the checkpoint time current_time = time.time() - self.write_data('checkpoint_time', current_time, path=path, - append=True) + self.write_data("checkpoint_time", current_time, path=path, append=True) # get the amount of time since the last checkpoint - checkpoint_times = history['checkpoint_time'][()] + checkpoint_times = history["checkpoint_time"][()] if len(checkpoint_times) == 1: # this is the first checkpoint, get the run time for comparison lasttime = self.run_start_time @@ -112,24 +112,25 @@ def update_checkpoint_history(self): lasttime = checkpoint_times[-2] # if a resume happened since the last checkpoint, use the resume # time instad - if lasttime < self.run_start_time: - lasttime = self.run_start_time - self.write_data('checkpoint_dt', current_time-lasttime, path=path, - append=True) + lasttime = max(lasttime, self.run_start_time) + self.write_data( + "checkpoint_dt", current_time - lasttime, path=path, append=True + ) # write any sampler-specific history self._update_sampler_history() def _update_sampler_history(self): - """Writes sampler-specific history to the file. + """ + Writes sampler-specific history to the file. This function does nothing. Classes that inherit from it may override it to add any extra information they would like written. This is called by :py:func:`update_checkpoint_history`. """ - pass def validate(self): - """Runs a validation test. + """ + Runs a validation test. This checks that a samples group exist, and that there are more than one sample stored to it. @@ -138,9 +139,10 @@ def validate(self): ------- bool : Whether or not the file is valid as a checkpoint file. + """ try: - group = '{}/{}'.format(self.samples_group, self.variable_params[0]) + group = f"{self.samples_group}/{self.variable_params[0]}" checkpoint_valid = self[group].size != 0 except KeyError: checkpoint_valid = False diff --git a/pycbc/inference/io/cpnest.py b/pycbc/inference/io/cpnest.py index 781c0aada9a..bad6b23f637 100644 --- a/pycbc/inference/io/cpnest.py +++ b/pycbc/inference/io/cpnest.py @@ -21,12 +21,12 @@ # # ============================================================================= # -"""Provides IO for the emcee sampler. -""" +"""Provides IO for the emcee sampler.""" + from .base_nested_sampler import BaseNestedSamplerFile class CPNestFile(BaseNestedSamplerFile): """Class to handle file IO for the ``cpnest`` sampler.""" - name = 'cpnest_file' + name = "cpnest_file" diff --git a/pycbc/inference/io/dynesty.py b/pycbc/inference/io/dynesty.py index f00db506e01..0da23889174 100644 --- a/pycbc/inference/io/dynesty.py +++ b/pycbc/inference/io/dynesty.py @@ -21,23 +21,25 @@ # # ============================================================================= # -"""Provides IO for the dynesty sampler. -""" +"""Provides IO for the dynesty sampler.""" import argparse + import numpy -from pycbc.io.hdf import (dump_state, load_state) + +from pycbc.io.hdf import dump_state, load_state from .base_nested_sampler import BaseNestedSamplerFile -from .posterior import write_samples_to_file, read_raw_samples_from_file +from .posterior import read_raw_samples_from_file, write_samples_to_file + -class CommonNestedMetadataIO(object): - """Provides functions for reading/writing dynesty metadata to file. - """ +class CommonNestedMetadataIO: + """Provides functions for reading/writing dynesty metadata to file.""" @staticmethod def extra_args_parser(parser=None, skip_args=None, **kwargs): - r"""Create a parser to parse sampler-specific arguments for loading + r""" + Create a parser to parse sampler-specific arguments for loading samples. Parameters @@ -60,57 +62,63 @@ def extra_args_parser(parser=None, skip_args=None, **kwargs): An argument parser with th extra arguments added. actions : list of argparse.Action A list of the actions that were added. + """ if parser is None: parser = argparse.ArgumentParser(**kwargs) elif kwargs: - raise ValueError("No other keyword arguments should be provded if " - "a parser is provided.") + raise ValueError( + "No other keyword arguments should be provded if a parser is provided." + ) if skip_args is None: skip_args = [] actions = [] - if 'raw_samples' not in skip_args: + if "raw_samples" not in skip_args: act = parser.add_argument( - "--raw-samples", action='store_true', default=False, + "--raw-samples", + action="store_true", + default=False, help="Extract raw samples rather than a posterior. " - "Raw samples are the unweighted samples obtained from " - "the nested sampler. Default value is False, which means " - "raw samples are weighted by the log-weight array " - "obtained from the sampler, giving an estimate of the " - "posterior.") + "Raw samples are the unweighted samples obtained from " + "the nested sampler. Default value is False, which means " + "raw samples are weighted by the log-weight array " + "obtained from the sampler, giving an estimate of the " + "posterior.", + ) actions.append(act) - if 'seed' not in skip_args: + if "seed" not in skip_args: act = parser.add_argument( - "--seed", type=int, default=0, + "--seed", + type=int, + default=0, help="Set the random-number seed used for extracting the " - "posterior samples. This is needed because the " - "unweighted samples are randomly shuffled to produce " - "a posterior. Default is 0. Ignored if raw-samples are " - "extracted instead.") + "posterior samples. This is needed because the " + "unweighted samples are randomly shuffled to produce " + "a posterior. Default is 0. Ignored if raw-samples are " + "extracted instead.", + ) return parser, actions def write_pickled_data_into_checkpoint_file(self, state): - """Dump the sampler state into checkpoint file - """ - if 'sampler_info/saved_state' not in self: - self.create_group('sampler_info/saved_state') - dump_state(state, self, path='sampler_info/saved_state') + """Dump the sampler state into checkpoint file""" + if "sampler_info/saved_state" not in self: + self.create_group("sampler_info/saved_state") + dump_state(state, self, path="sampler_info/saved_state") def read_pickled_data_from_checkpoint_file(self): - """Load the sampler state (pickled) from checkpoint file - """ - return load_state(self, path='sampler_info/saved_state') + """Load the sampler state (pickled) from checkpoint file""" + return load_state(self, path="sampler_info/saved_state") def write_raw_samples(self, data, parameters=None): - """Write the nested samples to the file - """ - if 'samples' not in self: - self.create_group('samples') - write_samples_to_file(self, data, parameters=parameters, - group='samples') + """Write the nested samples to the file""" + if "samples" not in self: + self.create_group("samples") + write_samples_to_file(self, data, parameters=parameters, group="samples") + def validate(self): - """Runs a validation test. + """ + Runs a validation test. This checks that a samples group exist, and that pickeled data can be loaded. @@ -118,10 +126,11 @@ def validate(self): ------- bool : Whether or not the file is valid as a checkpoint file. + """ try: - if 'sampler_info/saved_state' in self: - load_state(self, path='sampler_info/saved_state') + if "sampler_info/saved_state" in self: + load_state(self, path="sampler_info/saved_state") checkpoint_valid = True except KeyError: checkpoint_valid = False @@ -131,10 +140,11 @@ def validate(self): class DynestyFile(CommonNestedMetadataIO, BaseNestedSamplerFile): """Class to handle file IO for the ``dynesty`` sampler.""" - name = 'dynesty_file' + name = "dynesty_file" def read_raw_samples(self, fields, raw_samples=False, seed=0): - """Reads samples from a dynesty file and constructs a posterior. + """ + Reads samples from a dynesty file and constructs a posterior. Parameters ---------- @@ -153,12 +163,14 @@ def read_raw_samples(self, fields, raw_samples=False, seed=0): ------- dict : Dictionary of parameter names -> samples. + """ samples = read_raw_samples_from_file(self, fields) - logwt = read_raw_samples_from_file(self, ['logwt'])['logwt'] - loglikelihood = read_raw_samples_from_file( - self, ['loglikelihood'])['loglikelihood'] - logz = self.attrs.get('log_evidence') + logwt = read_raw_samples_from_file(self, ["logwt"])["logwt"] + loglikelihood = read_raw_samples_from_file(self, ["loglikelihood"])[ + "loglikelihood" + ] + logz = self.attrs.get("log_evidence") if not raw_samples: weights = numpy.exp(logwt - logz) N = len(weights) @@ -180,9 +192,8 @@ def read_raw_samples(self, fields, raw_samples=False, seed=0): # Py27: delete this after we drop python 2.7 support rng = numpy.random.RandomState(seed) rng.shuffle(idx) - post = {'loglikelihood': loglikelihood[idx]} + post = {"loglikelihood": loglikelihood[idx]} for i, param in enumerate(fields): post[param] = samples[param][idx] return post - else: - return samples + return samples diff --git a/pycbc/inference/io/emcee.py b/pycbc/inference/io/emcee.py index c793e29fc26..a84fd42448f 100644 --- a/pycbc/inference/io/emcee.py +++ b/pycbc/inference/io/emcee.py @@ -21,22 +21,27 @@ # # ============================================================================= # -"""Provides IO for the emcee sampler. -""" +"""Provides IO for the emcee sampler.""" + import numpy +from .base_mcmc import ( + CommonMCMCMetadataIO, + EnsembleMCMCMetadataIO, + ensemble_read_raw_samples, + write_samples, +) from .base_sampler import BaseSamplerFile -from .base_mcmc import (EnsembleMCMCMetadataIO, CommonMCMCMetadataIO, - write_samples, ensemble_read_raw_samples) class EmceeFile(EnsembleMCMCMetadataIO, CommonMCMCMetadataIO, BaseSamplerFile): """Class to handle file IO for the ``emcee`` sampler.""" - name = 'emcee_file' + name = "emcee_file" def write_samples(self, samples, **kwargs): - r"""Writes samples to the given file. + r""" + Writes samples to the given file. Calls :py:func:`base_mcmc.write_samples`. See that function for details. @@ -49,17 +54,19 @@ def write_samples(self, samples, **kwargs): \**kwargs : All other keyword arguments are passed to :py:func:`base_mcmc.write_samples`. + """ write_samples(self, samples, **kwargs) def read_raw_samples(self, fields, **kwargs): - r"""Base function for reading samples. + r""" + Base function for reading samples. Calls :py:func:`base_mcmc.ensemble_read_raw_samples`. See that function for details. Parameters - ----------- + ---------- fields : list The list of field names to retrieve. \**kwargs : @@ -70,14 +77,16 @@ def read_raw_samples(self, fields, **kwargs): ------- dict A dictionary of field name -> numpy array pairs. + """ return ensemble_read_raw_samples(self, fields, **kwargs) def read_acceptance_fraction(self, walkers=None): - """Reads the acceptance fraction. + """ + Reads the acceptance fraction. Parameters - ----------- + ---------- walkers : (list of) int, optional The walker index (or a list of indices) to retrieve. If None, samples from all walkers will be obtained. @@ -86,8 +95,9 @@ def read_acceptance_fraction(self, walkers=None): ------- array Array of acceptance fractions with shape (requested walkers,). + """ - group = self.sampler_group + '/acceptance_fraction' + group = self.sampler_group + "/acceptance_fraction" if walkers is None: wmask = numpy.ones(self.nwalkers, dtype=bool) else: @@ -96,15 +106,17 @@ def read_acceptance_fraction(self, walkers=None): return self[group][wmask] def write_acceptance_fraction(self, acceptance_fraction): - """Write acceptance_fraction data to file. Results are written to + """ + Write acceptance_fraction data to file. Results are written to the ``[sampler_group]/acceptance_fraction``. Parameters - ----------- + ---------- acceptance_fraction : numpy.ndarray Array of acceptance fractions to write. + """ - group = self.sampler_group + '/acceptance_fraction' + group = self.sampler_group + "/acceptance_fraction" try: self[group][:] = acceptance_fraction except KeyError: diff --git a/pycbc/inference/io/emcee_pt.py b/pycbc/inference/io/emcee_pt.py index 78e0063a59b..9240f7c2952 100644 --- a/pycbc/inference/io/emcee_pt.py +++ b/pycbc/inference/io/emcee_pt.py @@ -14,24 +14,25 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Provides I/O support for emcee_pt. -""" - +"""Provides I/O support for emcee_pt.""" import numpy -from .base_sampler import BaseSamplerFile from .base_mcmc import EnsembleMCMCMetadataIO -from .base_multitemper import (CommonMultiTemperedMetadataIO, - write_samples, - ensemble_read_raw_samples) +from .base_multitemper import ( + CommonMultiTemperedMetadataIO, + ensemble_read_raw_samples, + write_samples, +) +from .base_sampler import BaseSamplerFile -class EmceePTFile(EnsembleMCMCMetadataIO, CommonMultiTemperedMetadataIO, - BaseSamplerFile): +class EmceePTFile( + EnsembleMCMCMetadataIO, CommonMultiTemperedMetadataIO, BaseSamplerFile +): """Class to handle file IO for the ``emcee`` sampler.""" - name = 'emcee_pt_file' + name = "emcee_pt_file" @property def betas(self): @@ -39,7 +40,8 @@ def betas(self): return self[self.sampler_group].attrs["betas"] def write_samples(self, samples, **kwargs): - r"""Writes samples to the given file. + r""" + Writes samples to the given file. Calls :py:func:`base_multitemper.write_samples`. See that function for details. @@ -52,17 +54,19 @@ def write_samples(self, samples, **kwargs): \**kwargs : All other keyword arguments are passed to :py:func:`base_multitemper.write_samples`. + """ write_samples(self, samples, **kwargs) def read_raw_samples(self, fields, **kwargs): - r"""Base function for reading samples. + r""" + Base function for reading samples. Calls :py:func:`base_multitemper.ensemble_read_raw_samples`. See that function for details. Parameters - ----------- + ---------- fields : list The list of field names to retrieve. \**kwargs : @@ -73,20 +77,21 @@ def read_raw_samples(self, fields, **kwargs): ------- dict A dictionary of field name -> numpy array pairs. + """ return ensemble_read_raw_samples(self, fields, **kwargs) def write_sampler_metadata(self, sampler): - """Adds writing betas to MultiTemperedMCMCIO. - """ - super(EmceePTFile, self).write_sampler_metadata(sampler) + """Adds writing betas to MultiTemperedMCMCIO.""" + super().write_sampler_metadata(sampler) self[self.sampler_group].attrs["betas"] = sampler.betas def read_acceptance_fraction(self, temps=None, walkers=None): - """Reads the acceptance fraction. + """ + Reads the acceptance fraction. Parameters - ----------- + ---------- temps : (list of) int, optional The temperature index (or a list of indices) to retrieve. If None, acfs from all temperatures and all walkers will be retrieved. @@ -99,8 +104,9 @@ def read_acceptance_fraction(self, temps=None, walkers=None): array Array of acceptance fractions with shape (requested temps, requested walkers). + """ - group = self.sampler_group + '/acceptance_fraction' + group = self.sampler_group + "/acceptance_fraction" if walkers is None: wmask = numpy.ones(self.nwalkers, dtype=bool) else: @@ -114,21 +120,24 @@ def read_acceptance_fraction(self, temps=None, walkers=None): return self[group][:][numpy.ix_(tmask, wmask)] def write_acceptance_fraction(self, acceptance_fraction): - """Write acceptance_fraction data to file. + """ + Write acceptance_fraction data to file. Results are written to ``[sampler_group]/acceptance_fraction``; the resulting dataset has shape (ntemps, nwalkers). Parameters - ----------- + ---------- acceptance_fraction : numpy.ndarray Array of acceptance fractions to write. Must have shape ntemps x nwalkers. + """ # check assert acceptance_fraction.shape == (self.ntemps, self.nwalkers), ( - "acceptance fraction must have shape ntemps x nwalker") - group = self.sampler_group + '/acceptance_fraction' + "acceptance fraction must have shape ntemps x nwalker" + ) + group = self.sampler_group + "/acceptance_fraction" try: self[group][:] = acceptance_fraction except KeyError: diff --git a/pycbc/inference/io/epsie.py b/pycbc/inference/io/epsie.py index 92302bd7b3d..92ab2fa7f1e 100644 --- a/pycbc/inference/io/epsie.py +++ b/pycbc/inference/io/epsie.py @@ -13,26 +13,26 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides IO classes for epsie samplers. -""" +"""This module provides IO classes for epsie samplers.""" +from pickle import UnpicklingError import numpy -from pickle import UnpicklingError from epsie import load_state -from .base_sampler import BaseSamplerFile from .base_mcmc import MCMCMetadataIO -from .base_multitemper import (CommonMultiTemperedMetadataIO, - write_samples, - read_raw_samples) +from .base_multitemper import ( + CommonMultiTemperedMetadataIO, + read_raw_samples, + write_samples, +) +from .base_sampler import BaseSamplerFile -class EpsieFile(MCMCMetadataIO, CommonMultiTemperedMetadataIO, - BaseSamplerFile): +class EpsieFile(MCMCMetadataIO, CommonMultiTemperedMetadataIO, BaseSamplerFile): """Class to handle IO for Epsie's parallel-tempered sampler.""" - name = 'epsie_file' + name = "epsie_file" @property def nchains(self): @@ -42,23 +42,23 @@ def nchains(self): @property def betas(self): """The betas that were used.""" - return self[self.sampler_group]['betas'][()] + return self[self.sampler_group]["betas"][()] @property def swap_interval(self): """The interval that temperature swaps occurred at.""" - return self[self.sampler_group].attrs['swap_interval'] + return self[self.sampler_group].attrs["swap_interval"] @swap_interval.setter def swap_interval(self, swap_interval): """Stores the swap interval to the sampler group's attrs.""" - self[self.sampler_group].attrs['swap_interval'] = swap_interval + self[self.sampler_group].attrs["swap_interval"] = swap_interval @property def seed(self): """The sampler's seed.""" # convert seed from str back to int (see setter below for reason) - return int(self[self.sampler_group].attrs['seed']) + return int(self[self.sampler_group].attrs["seed"]) @seed.setter def seed(self, seed): @@ -66,17 +66,17 @@ def seed(self, seed): # epsie uses the numpy's new random generators, which use long integers # for seeds. hdf5 doesn't know how to handle long integers, so we'll # store it as a string - self[self.sampler_group].attrs['seed'] = str(seed) + self[self.sampler_group].attrs["seed"] = str(seed) def write_sampler_metadata(self, sampler): - """Adds writing seed and betas to MultiTemperedMCMCIO. - """ - super(EpsieFile, self).write_sampler_metadata(sampler) + """Adds writing seed and betas to MultiTemperedMCMCIO.""" + super().write_sampler_metadata(sampler) self.seed = sampler.seed self.write_data("betas", sampler.betas, path=self.sampler_group) def thin(self, thin_interval): - """Thins the samples on disk to the given thinning interval. + """ + Thins the samples on disk to the given thinning interval. Also thins the acceptance ratio and the temperature data, both of which are stored in the ``sampler_info`` group. @@ -87,22 +87,20 @@ def thin(self, thin_interval): # what the current thinned by is. new_interval = thin_interval // self.thinned_by # now thin the samples - super(EpsieFile, self).thin(thin_interval) + super().thin(thin_interval) # thin the acceptance ratio - self._thin_data(self.sampler_group, ['acceptance_ratio'], - new_interval) + self._thin_data(self.sampler_group, ["acceptance_ratio"], new_interval) # thin the temperature swaps; since these may not happen every # iteration, the thin interval we use for these is different - ts_group = '/'.join([self.sampler_group, 'temperature_swaps']) + ts_group = "/".join([self.sampler_group, "temperature_swaps"]) ts_thin_interval = new_interval // self.swap_interval if ts_thin_interval > 1: - self._thin_data(ts_group, ['swap_index'], - ts_thin_interval) - self._thin_data(ts_group, ['acceptance_ratio'], - ts_thin_interval) + self._thin_data(ts_group, ["swap_index"], ts_thin_interval) + self._thin_data(ts_group, ["acceptance_ratio"], ts_thin_interval) def write_samples(self, samples, **kwargs): - r"""Writes samples to the given file. + r""" + Writes samples to the given file. Calls :py:func:`base_multitemper.write_samples`. See that function for details. @@ -115,17 +113,19 @@ def write_samples(self, samples, **kwargs): \**kwargs : All other keyword arguments are passed to :py:func:`base_multitemper.write_samples`. + """ write_samples(self, samples, **kwargs) def read_raw_samples(self, fields, **kwargs): - r"""Base function for reading samples. + r""" + Base function for reading samples. Calls :py:func:`base_multitemper.read_raw_samples`. See that function for details. Parameters - ----------- + ---------- fields : list The list of field names to retrieve. \**kwargs : @@ -136,30 +136,36 @@ def read_raw_samples(self, fields, **kwargs): ------- dict A dictionary of field name -> numpy array pairs. + """ return read_raw_samples(self, fields, **kwargs) def write_acceptance_ratio(self, acceptance_ratio, last_iteration=None): - """Writes the acceptance ratios to the sampler info group. + """ + Writes the acceptance ratios to the sampler info group. Parameters ---------- acceptance_ratio : array The acceptance ratios to write. Should have shape ``ntemps x nchains x niterations``. + """ # we'll use the write_samples machinery to write the acceptance ratios - self.write_samples({'acceptance_ratio': acceptance_ratio}, - last_iteration=last_iteration, - samples_group=self.sampler_group) + self.write_samples( + {"acceptance_ratio": acceptance_ratio}, + last_iteration=last_iteration, + samples_group=self.sampler_group, + ) def read_acceptance_ratio(self, temps=None, chains=None): - """Reads the acceptance ratios. + """ + Reads the acceptance ratios. Ratios larger than 1 are set back to 1 before returning. Parameters - ----------- + ---------- temps : (list of) int, optional The temperature index (or a list of indices) to retrieve. If None, acceptance ratios from all temperatures and all chains will be @@ -173,8 +179,9 @@ def read_acceptance_ratio(self, temps=None, chains=None): array Array of acceptance ratios with shape (requested temps, requested chains, niterations). + """ - group = self.sampler_group + '/acceptance_ratio' + group = self.sampler_group + "/acceptance_ratio" if chains is None: wmask = numpy.ones(self.nchains, dtype=bool) else: @@ -187,17 +194,18 @@ def read_acceptance_ratio(self, temps=None, chains=None): tmask[temps] = True all_ratios = self[group][:] # make sure values > 1 are set back to 1 - all_ratios[all_ratios > 1] = 1. + all_ratios[all_ratios > 1] = 1.0 return all_ratios[numpy.ix_(tmask, wmask)] def read_acceptance_rate(self, temps=None, chains=None): - """Reads the acceptance rate. + """ + Reads the acceptance rate. This calls :py:func:`read_acceptance_ratio`, then averages the ratios over all iterations to get the average rate. Parameters - ----------- + ---------- temps : (list of) int, optional The temperature index (or a list of indices) to retrieve. If None, acceptance rates from all temperatures and all chains will be @@ -211,6 +219,7 @@ def read_acceptance_rate(self, temps=None, chains=None): array Array of acceptance ratios with shape (requested temps, requested chains). + """ all_ratios = self.read_acceptance_ratio(temps, chains) # average over the number of iterations @@ -218,13 +227,14 @@ def read_acceptance_rate(self, temps=None, chains=None): return all_ratios def read_acceptance_fraction(self, temps=None, walkers=None): - """Alias for :py:func:`read_acceptance_rate`. - """ + """Alias for :py:func:`read_acceptance_rate`.""" return self.read_acceptance_rate(temps=temps, chains=walkers) - def write_temperature_data(self, swap_index, acceptance_ratio, - swap_interval, last_iteration): - """Writes temperature swaps and acceptance ratios. + def write_temperature_data( + self, swap_index, acceptance_ratio, swap_interval, last_iteration + ): + """ + Writes temperature swaps and acceptance ratios. Parameters ---------- @@ -239,9 +249,10 @@ def write_temperature_data(self, swap_index, acceptance_ratio, The number of iterations between temperature swaps. last_iteration : int The iteration of the last sample. + """ self.swap_interval = swap_interval - group = '/'.join([self.sampler_group, 'temperature_swaps']) + group = "/".join([self.sampler_group, "temperature_swaps"]) # we'll use the write_samples machinery to write the acceptance ratios; # if temperature swaps didn't happen every iteration, then a smaller # thinning interval than what is used for the samples should be used @@ -251,16 +262,22 @@ def write_temperature_data(self, swap_index, acceptance_ratio, last_iteration = last_iteration // swap_interval # we need to write the two arrays separately, since they have different # dimensions in temperature - self.write_samples({'swap_index': swap_index}, - last_iteration=last_iteration, - samples_group=group, thin_by=thin_by) - self.write_samples({'acceptance_ratio': acceptance_ratio}, - last_iteration=last_iteration, - samples_group=group, thin_by=thin_by) + self.write_samples( + {"swap_index": swap_index}, + last_iteration=last_iteration, + samples_group=group, + thin_by=thin_by, + ) + self.write_samples( + {"acceptance_ratio": acceptance_ratio}, + last_iteration=last_iteration, + samples_group=group, + thin_by=thin_by, + ) def validate(self): """Adds attemp to load checkpoint to validation test.""" - valid = super(EpsieFile, self).validate() + valid = super().validate() # try to load the checkpoint if valid: try: @@ -276,10 +293,11 @@ def _get_optional_args(args, opts, err_on_missing=False, **kwargs): # need this to make sure options called "walkers" are renamed to # "chains" parsed = BaseSamplerFile._get_optional_args( - args, opts, err_on_missing=err_on_missing, **kwargs) + args, opts, err_on_missing=err_on_missing, **kwargs + ) try: - chains = parsed.pop('walkers') - parsed['chains'] = chains + chains = parsed.pop("walkers") + parsed["chains"] = chains except KeyError: pass return parsed diff --git a/pycbc/inference/io/multinest.py b/pycbc/inference/io/multinest.py index 85ba9a70e62..5a3b72870dc 100644 --- a/pycbc/inference/io/multinest.py +++ b/pycbc/inference/io/multinest.py @@ -14,9 +14,7 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Provides I/O support for multinest. -""" - +"""Provides I/O support for multinest.""" from .base_sampler import BaseSamplerFile @@ -24,10 +22,11 @@ class MultinestFile(BaseSamplerFile): """Class to handle file IO for the ``multinest`` sampler.""" - name = 'multinest_file' + name = "multinest_file" def write_samples(self, samples, parameters=None): - """Writes samples to the given file. + """ + Writes samples to the given file. Results are written to ``samples_group/{vararg}``, where ``{vararg}`` is the name of a model params. The samples are written as an @@ -41,11 +40,13 @@ def write_samples(self, samples, parameters=None): parameters : list, optional Only write the specified parameters to the file. If None, will write all of the keys in the ``samples`` dict. + """ niterations = len(tuple(samples.values())[0]) assert all(len(p) == niterations for p in samples.values()), ( - "all samples must have the same shape") - group = self.samples_group + '/{name}' + "all samples must have the same shape" + ) + group = self.samples_group + "/{name}" if parameters is None: parameters = samples.keys() # loop over number of dimensions @@ -58,14 +59,18 @@ def write_samples(self, samples, parameters=None): self[dataset_name].resize(niterations, axis=0) except KeyError: # dataset doesn't exist yet - self.create_dataset(dataset_name, (niterations,), - maxshape=(None,), - dtype=samples[param].dtype, - fletcher32=True) + self.create_dataset( + dataset_name, + (niterations,), + maxshape=(None,), + dtype=samples[param].dtype, + fletcher32=True, + ) self[dataset_name][:] = samples[param] def write_logevidence(self, lnz, dlnz, importance_lnz, importance_dlnz): - """Writes the given log evidence and its error. + """ + Writes the given log evidence and its error. Results are saved to file's 'log_evidence' and 'dlog_evidence' attributes, as well as the importance-weighted versions of these @@ -81,18 +86,19 @@ def write_logevidence(self, lnz, dlnz, importance_lnz, importance_dlnz): The importance-weighted log of the evidence. importance_dlnz : float, optional The error in the importance-weighted estimate of the log evidence. + """ - self.attrs['log_evidence'] = lnz - self.attrs['dlog_evidence'] = dlnz + self.attrs["log_evidence"] = lnz + self.attrs["dlog_evidence"] = dlnz if all([e is not None for e in [importance_lnz, importance_dlnz]]): - self.attrs['importance_log_evidence'] = importance_lnz - self.attrs['importance_dlog_evidence'] = importance_dlnz + self.attrs["importance_log_evidence"] = importance_lnz + self.attrs["importance_dlog_evidence"] = importance_dlnz def read_raw_samples(self, fields, iteration=None): if isinstance(fields, str): fields = [fields] # load - group = self.samples_group + '/{name}' + group = self.samples_group + "/{name}" arrays = {} for name in fields: if iteration is not None: @@ -103,8 +109,10 @@ def read_raw_samples(self, fields, iteration=None): return arrays def write_resume_point(self): - """Keeps a list of the number of iterations that were in a file when a - run was resumed from a checkpoint.""" + """ + Keeps a list of the number of iterations that were in a file when a + run was resumed from a checkpoint. + """ try: resume_pts = self.attrs["resume_points"].tolist() except KeyError: @@ -119,18 +127,18 @@ def write_resume_point(self): @property def niterations(self): """Returns the number of iterations the sampler was run for.""" - return self[self.sampler_group].attrs['niterations'] + return self[self.sampler_group].attrs["niterations"] def write_niterations(self, niterations): """Writes the given number of iterations to the sampler group.""" - self[self.sampler_group].attrs['niterations'] = niterations + self[self.sampler_group].attrs["niterations"] = niterations def write_sampler_metadata(self, sampler): """Writes the sampler's metadata.""" - self.attrs['sampler'] = sampler.name + self.attrs["sampler"] = sampler.name if self.sampler_group not in self.keys(): # create the sampler group self.create_group(self.sampler_group) - self[self.sampler_group].attrs['nlivepoints'] = sampler.nlivepoints + self[self.sampler_group].attrs["nlivepoints"] = sampler.nlivepoints # write the model's metadata sampler.model.write_metadata(self) diff --git a/pycbc/inference/io/nessai.py b/pycbc/inference/io/nessai.py index 86c1bcfba41..eb1e5454d16 100644 --- a/pycbc/inference/io/nessai.py +++ b/pycbc/inference/io/nessai.py @@ -1,10 +1,10 @@ """Provides IO for the nessai sampler""" + import numpy from .base_nested_sampler import BaseNestedSamplerFile - -from .posterior import read_raw_samples_from_file from .dynesty import CommonNestedMetadataIO +from .posterior import read_raw_samples_from_file class NessaiFile(CommonNestedMetadataIO, BaseNestedSamplerFile): @@ -13,7 +13,8 @@ class NessaiFile(CommonNestedMetadataIO, BaseNestedSamplerFile): name = "nessai_file" def read_raw_samples(self, fields, raw_samples=False, seed=0): - """Reads samples from a nessai file and constructs a posterior. + """ + Reads samples from a nessai file and constructs a posterior. Using rejection sampling to resample the nested samples @@ -30,11 +31,13 @@ def read_raw_samples(self, fields, raw_samples=False, seed=0): ------- dict : Dictionary of parameter fields -> samples. + """ samples = read_raw_samples_from_file(self, fields) - logwt = read_raw_samples_from_file(self, ['logwt'])['logwt'] - loglikelihood = read_raw_samples_from_file( - self, ['loglikelihood'])['loglikelihood'] + logwt = read_raw_samples_from_file(self, ["logwt"])["logwt"] + loglikelihood = read_raw_samples_from_file(self, ["loglikelihood"])[ + "loglikelihood" + ] if not raw_samples: n_samples = len(logwt) # Rejection sample @@ -42,7 +45,7 @@ def read_raw_samples(self, fields, raw_samples=False, seed=0): logwt -= logwt.max() logu = numpy.log(rng.random(n_samples)) keep = logwt > logu - post = {'loglikelihood': loglikelihood[keep]} + post = {"loglikelihood": loglikelihood[keep]} for param in fields: post[param] = samples[param][keep] return post diff --git a/pycbc/inference/io/posterior.py b/pycbc/inference/io/posterior.py index 596690495b2..75f52f53d74 100644 --- a/pycbc/inference/io/posterior.py +++ b/pycbc/inference/io/posterior.py @@ -21,8 +21,7 @@ # # ============================================================================= # -"""Provides simplified standard format just for posterior data -""" +"""Provides simplified standard format just for posterior data""" from .base_hdf import BaseInferenceFile @@ -30,7 +29,7 @@ class PosteriorFile(BaseInferenceFile): """Class to handle file IO for the simplified Posterior file.""" - name = 'posterior_file' + name = "posterior_file" def read_raw_samples(self, fields, **kwargs): return read_raw_samples_from_file(self, fields, **kwargs) @@ -53,14 +52,15 @@ def read_raw_samples_from_file(fp, fields, **kwargs): def write_samples_to_file(fp, samples, parameters=None, group=None): - """Writes samples to the given file. + """ + Writes samples to the given file. Results are written to ``samples_group/{vararg}``, where ``{vararg}`` is the name of a model params. The samples are written as an array of length ``niterations``. Parameters - ----------- + ---------- fp : self Pass the 'self' from BaseInferenceFile class. samples : dict @@ -69,19 +69,20 @@ def write_samples_to_file(fp, samples, parameters=None, group=None): parameters : list, optional Only write the specified parameters to the file. If None, will write all of the keys in the ``samples`` dict. - """ + + """ # check data dimensions; we'll just use the first array in samples arr = list(samples.values())[0] if not arr.ndim == 1: raise ValueError("samples must be 1D arrays") niterations = arr.size - assert all(len(p) == niterations - for p in samples.values()), ( - "all samples must have the same shape") + assert all(len(p) == niterations for p in samples.values()), ( + "all samples must have the same shape" + ) if group is not None: - group = group + '/{name}' + group = group + "/{name}" else: - group = fp.samples_group + '/{name}' + group = fp.samples_group + "/{name}" if parameters is None: parameters = samples.keys() # loop over number of dimensions @@ -94,8 +95,11 @@ def write_samples_to_file(fp, samples, parameters=None, group=None): fp[dataset_name].resize(niterations, axis=0) except KeyError: # dataset doesn't exist yet - fp.create_dataset(dataset_name, (niterations,), - maxshape=(None,), - dtype=samples[param].dtype, - fletcher32=True) + fp.create_dataset( + dataset_name, + (niterations,), + maxshape=(None,), + dtype=samples[param].dtype, + fletcher32=True, + ) fp[dataset_name][:] = samples[param] diff --git a/pycbc/inference/io/ptemcee.py b/pycbc/inference/io/ptemcee.py index 578f9952160..c9fb27ad958 100644 --- a/pycbc/inference/io/ptemcee.py +++ b/pycbc/inference/io/ptemcee.py @@ -14,32 +14,38 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Provides I/O support for ptemcee. -""" +"""Provides I/O support for ptemcee.""" - -from .base_sampler import BaseSamplerFile from . import base_mcmc from .base_mcmc import EnsembleMCMCMetadataIO -from .base_multitemper import (CommonMultiTemperedMetadataIO, - write_samples, - ensemble_read_raw_samples) +from .base_multitemper import ( + CommonMultiTemperedMetadataIO, + ensemble_read_raw_samples, + write_samples, +) +from .base_sampler import BaseSamplerFile -class PTEmceeFile(EnsembleMCMCMetadataIO, CommonMultiTemperedMetadataIO, - BaseSamplerFile): +class PTEmceeFile( + EnsembleMCMCMetadataIO, CommonMultiTemperedMetadataIO, BaseSamplerFile +): """Class to handle file IO for the ``ptemcee`` sampler.""" - name = 'ptemcee_file' + name = "ptemcee_file" # attributes for setting up an ensemble from file - _ensemble_attrs = ['jumps_proposed', 'jumps_accepted', 'swaps_proposed', - 'swaps_accepted', 'logP', 'logl'] + _ensemble_attrs = [ + "jumps_proposed", + "jumps_accepted", + "swaps_proposed", + "swaps_accepted", + "logP", + "logl", + ] def write_sampler_metadata(self, sampler): - """Adds writing ptemcee-specific metadata to MultiTemperedMCMCIO. - """ - super(PTEmceeFile, self).write_sampler_metadata(sampler) + """Adds writing ptemcee-specific metadata to MultiTemperedMCMCIO.""" + super().write_sampler_metadata(sampler) group = self[self.sampler_group] group.attrs["starting_betas"] = sampler.starting_betas group.attrs["adaptive"] = sampler.adaptive @@ -53,23 +59,29 @@ def starting_betas(self): return self[self.sampler_group].attrs["starting_betas"] def write_betas(self, betas, last_iteration=None): - """Writes the betas to sampler group. + """ + Writes the betas to sampler group. As the betas may change with iterations, this writes the betas as a ntemps x niterations array to the file. """ # we'll use the single temperature write_samples to write the betas, # so that we get the thinning settings - base_mcmc.write_samples(self, {'betas': betas}, - last_iteration=last_iteration, - samples_group=self.sampler_group) - - def read_betas(self, thin_start=None, thin_interval=None, thin_end=None, - iteration=None): - """Reads betas from the file. + base_mcmc.write_samples( + self, + {"betas": betas}, + last_iteration=last_iteration, + samples_group=self.sampler_group, + ) + + def read_betas( + self, thin_start=None, thin_interval=None, thin_end=None, iteration=None + ): + """ + Reads betas from the file. Parameters - ----------- + ---------- thin_start : int, optional Start reading from the given iteration. Default is to start from the first iteration. @@ -86,21 +98,27 @@ def read_betas(self, thin_start=None, thin_interval=None, thin_end=None, ------- array A ntemps x niterations array of the betas. + """ - slc = base_mcmc._ensemble_get_index(self, thin_start=thin_start, - thin_interval=thin_interval, - thin_end=thin_end, - iteration=iteration) - betas = self[self.sampler_group]['betas'][:] + slc = base_mcmc._ensemble_get_index( + self, + thin_start=thin_start, + thin_interval=thin_interval, + thin_end=thin_end, + iteration=iteration, + ) + betas = self[self.sampler_group]["betas"][:] return betas[:, slc] def write_ensemble_attrs(self, ensemble): - """Writes ensemble attributes necessary to restart from checkpoint. + """ + Writes ensemble attributes necessary to restart from checkpoint. Parameters ---------- ensemble : ptemcee.Ensemble The ensemble to write attributes for. + """ group = self[self.sampler_group] for attr in self._ensemble_attrs: @@ -111,18 +129,21 @@ def write_ensemble_attrs(self, ensemble): group[attr] = vals def read_ensemble_attrs(self): - """Reads ensemble attributes from the file. + """ + Reads ensemble attributes from the file. Returns ------- dict : Dictionary of the ensemble attributes. + """ group = self[self.sampler_group] return {attr: group[attr][:] for attr in self._ensemble_attrs} def write_samples(self, samples, **kwargs): - r"""Writes samples to the given file. + r""" + Writes samples to the given file. Calls :py:func:`base_multitemper.write_samples`. See that function for details. @@ -135,11 +156,13 @@ def write_samples(self, samples, **kwargs): \**kwargs : All other keyword arguments are passed to :py:func:`base_multitemper.write_samples`. + """ write_samples(self, samples, **kwargs) def read_raw_samples(self, fields, **kwargs): - r"""Base function for reading samples. + r""" + Base function for reading samples. Calls :py:func:`base_multitemper.ensemble_read_raw_samples`. See that function for details. @@ -156,5 +179,6 @@ def read_raw_samples(self, fields, **kwargs): ------- dict A dictionary of field name -> numpy array pairs. + """ return ensemble_read_raw_samples(self, fields, **kwargs) diff --git a/pycbc/inference/io/snowline.py b/pycbc/inference/io/snowline.py index 4aa9edcbf13..0ebc8955bf9 100644 --- a/pycbc/inference/io/snowline.py +++ b/pycbc/inference/io/snowline.py @@ -21,12 +21,12 @@ # # ============================================================================= # -"""Provides IO for the snowline sampler. -""" +"""Provides IO for the snowline sampler.""" + from .posterior import PosteriorFile class SnowlineFile(PosteriorFile): """Class to handle file IO for the ``snowline`` sampler.""" - name = 'snowline_file' + name = "snowline_file" diff --git a/pycbc/inference/io/txt.py b/pycbc/inference/io/txt.py index b8714706755..cf55930242f 100644 --- a/pycbc/inference/io/txt.py +++ b/pycbc/inference/io/txt.py @@ -12,26 +12,30 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This modules defines functions for reading and samples that the +""" +This modules defines functions for reading and samples that the inference samplers generate and are stored in an ASCII TXT file. """ import numpy -class InferenceTXTFile(object): - """ A class that has extra functions for handling reading the samples +class InferenceTXTFile: + """ + A class that has extra functions for handling reading the samples from posterior-only TXT files. Parameters - ----------- + ---------- path : str The path to the TXT file. mode : {None, str} The mode to open the file. Only accepts "r" or "rb" for reading. delimiter : str Delimiter to use for TXT file. Default is space-delimited. + """ + name = "txt" comments = "" delimiter = " " @@ -46,10 +50,11 @@ def __init__(self, path, mode=None, delimiter=None): @classmethod def write(cls, output_file, samples, labels, delimiter=None): - """ Writes a text file with samples. + """ + Writes a text file with samples. Parameters - ----------- + ---------- output_file : str The path of the file to write. samples : FieldArray @@ -58,9 +63,14 @@ def write(cls, output_file, samples, labels, delimiter=None): A list of strings to include as header in TXT file. delimiter : str Delimiter to use in TXT file. + """ delimiter = delimiter if delimiter is not None else cls.delimiter header = delimiter.join(labels) - numpy.savetxt(output_file, samples, - comments=cls.comments, header=header, - delimiter=delimiter) + numpy.savetxt( + output_file, + samples, + comments=cls.comments, + header=header, + delimiter=delimiter, + ) diff --git a/pycbc/inference/io/ultranest.py b/pycbc/inference/io/ultranest.py index 56b98fd9f66..e5189a9433f 100644 --- a/pycbc/inference/io/ultranest.py +++ b/pycbc/inference/io/ultranest.py @@ -21,12 +21,12 @@ # # ============================================================================= # -"""Provides IO for the ultranest sampler. -""" +"""Provides IO for the ultranest sampler.""" + from .base_nested_sampler import BaseNestedSamplerFile class UltranestFile(BaseNestedSamplerFile): """Class to handle file IO for the ``ultranest`` sampler.""" - name = 'ultranest_file' + name = "ultranest_file" diff --git a/pycbc/inference/jump/__init__.py b/pycbc/inference/jump/__init__.py index f55975f6cfe..069dfcacd39 100644 --- a/pycbc/inference/jump/__init__.py +++ b/pycbc/inference/jump/__init__.py @@ -14,15 +14,19 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. """Provides custom jump proposals for samplers.""" -from .normal import (EpsieNormal, EpsieAdaptiveNormal, EpsieATAdaptiveNormal) -from .bounded_normal import (EpsieBoundedNormal, EpsieAdaptiveBoundedNormal, - EpsieATAdaptiveBoundedNormal) -from .angular import (EpsieAngular, EpsieAdaptiveAngular, - EpsieATAdaptiveAngular) -from .discrete import (EpsieNormalDiscrete, EpsieBoundedDiscrete, - EpsieAdaptiveNormalDiscrete, - EpsieAdaptiveBoundedDiscrete) - +from .angular import EpsieAdaptiveAngular, EpsieAngular, EpsieATAdaptiveAngular +from .bounded_normal import ( + EpsieAdaptiveBoundedNormal, + EpsieATAdaptiveBoundedNormal, + EpsieBoundedNormal, +) +from .discrete import ( + EpsieAdaptiveBoundedDiscrete, + EpsieAdaptiveNormalDiscrete, + EpsieBoundedDiscrete, + EpsieNormalDiscrete, +) +from .normal import EpsieAdaptiveNormal, EpsieATAdaptiveNormal, EpsieNormal epsie_proposals = { EpsieNormal.name: EpsieNormal, @@ -41,8 +45,9 @@ } -def epsie_proposals_from_config(cp, section='jump_proposal'): - """Loads epsie jump proposals from the given config file. +def epsie_proposals_from_config(cp, section="jump_proposal"): + """ + Loads epsie jump proposals from the given config file. This loads jump proposals from sub-sections starting with ``section`` (default is 'jump_proposal'). The tag part of the sub-sections' headers @@ -73,6 +78,7 @@ def epsie_proposals_from_config(cp, section='jump_proposal'): ------- list : List of the proposal instances. + """ tags = cp.get_subsections(section) proposals = [] diff --git a/pycbc/inference/jump/angular.py b/pycbc/inference/jump/angular.py index f6932d28fe7..8865a07e0a4 100644 --- a/pycbc/inference/jump/angular.py +++ b/pycbc/inference/jump/angular.py @@ -17,8 +17,11 @@ from epsie import proposals as epsie_proposals -from .normal import (epsie_from_config, epsie_adaptive_from_config, - epsie_at_adaptive_from_config) +from .normal import ( + epsie_adaptive_from_config, + epsie_at_adaptive_from_config, + epsie_from_config, +) class EpsieAngular(epsie_proposals.Angular): @@ -26,7 +29,8 @@ class EpsieAngular(epsie_proposals.Angular): @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_from_config` with ``cls`` set to :py:class:`epsie.proposals.Angular` and ``with_boundaries`` set @@ -52,6 +56,7 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.Angular`: An angular proposal for use with ``epsie`` samplers. + """ return epsie_from_config(cls, cp, section, tag, with_boundaries=False) @@ -61,7 +66,8 @@ class EpsieAdaptiveAngular(epsie_proposals.AdaptiveAngular): @classmethod def from_config(cls, cp, section, tag): - r"""Loads a proposal from a config file. + r""" + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveBoundedNormal` and @@ -89,9 +95,9 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveAngular`: An adaptive angular proposal for use with ``epsie`` samplers. + """ - return epsie_adaptive_from_config(cls, cp, section, tag, - with_boundaries=False) + return epsie_adaptive_from_config(cls, cp, section, tag, with_boundaries=False) class EpsieATAdaptiveAngular(epsie_proposals.ATAdaptiveAngular): @@ -99,7 +105,8 @@ class EpsieATAdaptiveAngular(epsie_proposals.ATAdaptiveAngular): @classmethod def from_config(cls, cp, section, tag): - r"""Loads a proposal from a config file. + r""" + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveBoundedNormal` and @@ -127,6 +134,8 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveAngularProposal`: An adaptive angular proposal for use with ``epsie`` samplers. + """ - return epsie_at_adaptive_from_config(cls, cp, section, tag, - with_boundaries=False) + return epsie_at_adaptive_from_config( + cls, cp, section, tag, with_boundaries=False + ) diff --git a/pycbc/inference/jump/bounded_normal.py b/pycbc/inference/jump/bounded_normal.py index 2e6f257086a..110bfe66d3d 100644 --- a/pycbc/inference/jump/bounded_normal.py +++ b/pycbc/inference/jump/bounded_normal.py @@ -17,8 +17,11 @@ from epsie import proposals as epsie_proposals -from .normal import (epsie_from_config, epsie_adaptive_from_config, - epsie_at_adaptive_from_config) +from .normal import ( + epsie_adaptive_from_config, + epsie_at_adaptive_from_config, + epsie_from_config, +) class EpsieBoundedNormal(epsie_proposals.BoundedNormal): @@ -26,7 +29,8 @@ class EpsieBoundedNormal(epsie_proposals.BoundedNormal): @classmethod def from_config(cls, cp, section, tag): - r"""Loads a proposal from a config file. + r""" + Loads a proposal from a config file. This calls :py:func:`epsie_from_config` with ``cls`` set to :py:class:`epsie.proposals.BoundedNormal` and ``with_boundaries`` set @@ -55,6 +59,7 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.BoundedNormal`: A bounded normal proposal for use with ``epsie`` samplers. + """ return epsie_from_config(cls, cp, section, tag, with_boundaries=True) @@ -64,7 +69,8 @@ class EpsieAdaptiveBoundedNormal(epsie_proposals.AdaptiveBoundedNormal): @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveBoundedNormal`. See that function @@ -92,6 +98,7 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveBoundedNormal`: An adaptive normal proposal for use with ``epsie`` samplers. + """ return epsie_adaptive_from_config(cls, cp, section, tag) @@ -101,7 +108,8 @@ class EpsieATAdaptiveBoundedNormal(epsie_proposals.ATAdaptiveBoundedNormal): @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveBoundedProposal`. See that function @@ -128,6 +136,8 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveBoundedProposal`: An adaptive bounded proposal for use with ``epsie`` samplers. + """ - return epsie_at_adaptive_from_config(cls, cp, section, tag, - with_boundaries=True) + return epsie_at_adaptive_from_config( + cls, cp, section, tag, with_boundaries=True + ) diff --git a/pycbc/inference/jump/discrete.py b/pycbc/inference/jump/discrete.py index bc4164666d5..044646157e5 100644 --- a/pycbc/inference/jump/discrete.py +++ b/pycbc/inference/jump/discrete.py @@ -17,7 +17,7 @@ from epsie import proposals as epsie_proposals -from .normal import (epsie_from_config, epsie_adaptive_from_config) +from .normal import epsie_adaptive_from_config, epsie_from_config class EpsieNormalDiscrete(epsie_proposals.NormalDiscrete): @@ -25,7 +25,8 @@ class EpsieNormalDiscrete(epsie_proposals.NormalDiscrete): @classmethod def from_config(cls, cp, section, tag): - r"""Loads a proposal from a config file. + r""" + Loads a proposal from a config file. This calls :py:func:`epsie_from_config` with ``cls`` set to :py:class:`epsie.proposals.NormalDiscrete` and ``with_boundaries`` set @@ -50,6 +51,7 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.BoundedDiscrete`: A bounded discrete proposal for use with ``epsie`` samplers. + """ return epsie_from_config(cls, cp, section, tag, with_boundaries=False) @@ -59,7 +61,8 @@ class EpsieBoundedDiscrete(epsie_proposals.BoundedDiscrete): @classmethod def from_config(cls, cp, section, tag): - r"""Loads a proposal from a config file. + r""" + Loads a proposal from a config file. This calls :py:func:`epsie_from_config` with ``cls`` set to :py:class:`epsie.proposals.BoundedDiscrete` and ``with_boundaries`` set @@ -86,17 +89,21 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.BoundedDiscrete`: A bounded discrete proposal for use with ``epsie`` samplers. + """ return epsie_from_config(cls, cp, section, tag, with_boundaries=True) class EpsieAdaptiveNormalDiscrete(epsie_proposals.AdaptiveNormalDiscrete): - """Adds ``from_config`` method to epsie's adaptive bounded discrete - proposal.""" + """ + Adds ``from_config`` method to epsie's adaptive bounded discrete + proposal. + """ @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveNormalDiscrete`. See that function @@ -124,18 +131,23 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveBoundedDiscrete`: An adaptive normal proposal for use with ``epsie`` samplers. + """ - return epsie_adaptive_from_config(cls, cp, section, tag, - boundary_arg_name='prior_widths') + return epsie_adaptive_from_config( + cls, cp, section, tag, boundary_arg_name="prior_widths" + ) class EpsieAdaptiveBoundedDiscrete(epsie_proposals.AdaptiveBoundedDiscrete): - """Adds ``from_config`` method to epsie's adaptive bounded discrete - proposal.""" + """ + Adds ``from_config`` method to epsie's adaptive bounded discrete + proposal. + """ @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveBoundedDiscrete`. See that function @@ -163,5 +175,6 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveBoundedDiscrete`: An adaptive normal proposal for use with ``epsie`` samplers. + """ return epsie_adaptive_from_config(cls, cp, section, tag) diff --git a/pycbc/inference/jump/normal.py b/pycbc/inference/jump/normal.py index 31ebdc8d5b3..8a2b0133ed3 100644 --- a/pycbc/inference/jump/normal.py +++ b/pycbc/inference/jump/normal.py @@ -15,9 +15,7 @@ """Jump proposals that use a normal distribution.""" - import numpy - from epsie import proposals as epsie_proposals from epsie.proposals import Boundaries @@ -29,7 +27,8 @@ class EpsieNormal(epsie_proposals.Normal): @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_from_config` with ``cls`` set to :py:class:`epsie.proposals.Normal` and ``with_boundaries`` set to @@ -55,6 +54,7 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.Normal`: A normal proposal for use with ``epsie`` samplers. + """ return epsie_from_config(cls, cp, section, tag, with_boundaries=False) @@ -64,7 +64,8 @@ class EpsieAdaptiveNormal(epsie_proposals.AdaptiveNormal): @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_adaptive_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveNormal`. See that function for @@ -94,9 +95,11 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveNormal`: An adaptive normal proposal for use with ``epsie`` samplers. + """ - return epsie_adaptive_from_config(cls, cp, section, tag, - boundary_arg_name='prior_widths') + return epsie_adaptive_from_config( + cls, cp, section, tag, boundary_arg_name="prior_widths" + ) class EpsieATAdaptiveNormal(epsie_proposals.ATAdaptiveNormal): @@ -104,7 +107,8 @@ class EpsieATAdaptiveNormal(epsie_proposals.ATAdaptiveNormal): @classmethod def from_config(cls, cp, section, tag): - """Loads a proposal from a config file. + """ + Loads a proposal from a config file. This calls :py:func:`epsie_from_config` with ``cls`` set to :py:class:`epsie.proposals.AdaptiveProposal` and ``with_boundaries`` @@ -131,13 +135,16 @@ def from_config(cls, cp, section, tag): ------- :py:class:`epsie.proposals.AdaptiveProposal`: An adaptive proposal for use with ``epsie`` samplers. + """ - return epsie_at_adaptive_from_config(cls, cp, section, tag, - with_boundaries=False) + return epsie_at_adaptive_from_config( + cls, cp, section, tag, with_boundaries=False + ) def epsie_from_config(cls, cp, section, tag, with_boundaries=False): - r"""Generic function for loading epsie proposals from a config file. + r""" + Generic function for loading epsie proposals from a config file. This should be used for proposals that are not adaptive. @@ -180,36 +187,40 @@ def epsie_from_config(cls, cp, section, tag, with_boundaries=False): ------- cls : The class initialized with the options read from the config file. + """ # check that the name matches assert cp.get_opt_tag(section, "name", tag) == cls.name, ( - "name in specified section must match mine") - params, opts = load_opts(cp, section, tag, skip=['name']) - args = {'parameters': params} + "name in specified section must match mine" + ) + params, opts = load_opts(cp, section, tag, skip=["name"]) + args = {"parameters": params} if with_boundaries: boundaries = get_param_boundaries(params, opts) - args['boundaries'] = boundaries - if 'discrete' in cls.name.split('_'): - args.update({'successive': - get_epsie_discrete_successive_settings(params, opts)}) + args["boundaries"] = boundaries + if "discrete" in cls.name.split("_"): + args.update( + {"successive": get_epsie_discrete_successive_settings(params, opts)} + ) # if there are any options left, assume they are for setting the variance if opts: cov = get_variance(params, opts) elif with_boundaries: - cov = numpy.array([abs(boundaries[p])/10. for p in params])**2. + cov = numpy.array([abs(boundaries[p]) / 10.0 for p in params]) ** 2.0 else: cov = None - args['cov'] = cov + args["cov"] = cov # no other options should remain if opts: - raise ValueError("unrecognized options {}" - .format(', '.join(opts.keys()))) + raise ValueError("unrecognized options {}".format(", ".join(opts.keys()))) return cls(**args) -def epsie_adaptive_from_config(cls, cp, section, tag, with_boundaries=True, - boundary_arg_name='boundaries'): - """Generic function for loading adaptive epsie proposals from a config +def epsie_adaptive_from_config( + cls, cp, section, tag, with_boundaries=True, boundary_arg_name="boundaries" +): + """ + Generic function for loading adaptive epsie proposals from a config file. The section that is read should have the format ``[{section}-{tag}]``, @@ -265,36 +276,41 @@ def epsie_adaptive_from_config(cls, cp, section, tag, with_boundaries=True, ------- cls : The class initialized with the options read from the config file. + """ # check that the name matches assert cp.get_opt_tag(section, "name", tag) == cls.name, ( - "name in specified section must match mine") - params, opts = load_opts(cp, section, tag, skip=['name']) - args = {'parameters': params} + "name in specified section must match mine" + ) + params, opts = load_opts(cp, section, tag, skip=["name"]) + args = {"parameters": params} # get the bounds if with_boundaries: args[boundary_arg_name] = get_param_boundaries(params, opts) - if 'discrete' in cls.name.split('_'): - args.update({'successive': - get_epsie_discrete_successive_settings(params, opts)}) + if "discrete" in cls.name.split("_"): + args.update( + {"successive": get_epsie_discrete_successive_settings(params, opts)} + ) # get the adaptation parameters args.update(get_epsie_adaptation_settings(opts)) # if there are any other options, assume they are for setting the # initial standard deviation if opts: var = get_variance(params, opts) - args['initial_std'] = var**0.5 + args["initial_std"] = var**0.5 # at this point, there should be no options left if opts: - raise ValueError('unrecognized options {} in section {}' - .format(', '.join(opts.keys()), - '-'.join([section, tag]))) + raise ValueError( + "unrecognized options {} in section {}".format( + ", ".join(opts.keys()), "-".join([section, tag]) + ) + ) return cls(**args) -def epsie_at_adaptive_from_config(cls, cp, section, tag, - with_boundaries=False): - """Generic function for loading AT Adaptive Normal proposals from a config +def epsie_at_adaptive_from_config(cls, cp, section, tag, with_boundaries=False): + """ + Generic function for loading AT Adaptive Normal proposals from a config file. The section that is read should have the format ``[{section}-{tag}]``, @@ -346,35 +362,38 @@ def epsie_at_adaptive_from_config(cls, cp, section, tag, ------- cls : The class initialized with the options read from the config file. + """ # check that the name matches assert cp.get_opt_tag(section, "name", tag) == cls.name, ( - "name in specified section must match mine") - params, opts = load_opts(cp, section, tag, skip=['name']) - args = {'parameters': params} + "name in specified section must match mine" + ) + params, opts = load_opts(cp, section, tag, skip=["name"]) + args = {"parameters": params} # get the bounds if with_boundaries: - args['boundaries'] = get_param_boundaries(params, opts) - if 'discrete' in cls.name.split('_'): - args.update({'successive': - get_epsie_discrete_successive_settings(params, opts)}) + args["boundaries"] = get_param_boundaries(params, opts) + if "discrete" in cls.name.split("_"): + args.update( + {"successive": get_epsie_discrete_successive_settings(params, opts)} + ) # get the adaptation parameters args.update(get_epsie_adaptation_settings(opts, cls.name)) # bounded and angular adaptive proposals support diagonal-only - diagonal = opts.pop('diagonal', None) - if not any(p in cls.name.split('_') for p in ['bounded', 'angular']): - args.update({'diagonal': diagonal is not None}) - componentwise = opts.pop('componentwise', None) + diagonal = opts.pop("diagonal", None) + if not any(p in cls.name.split("_") for p in ["bounded", "angular"]): + args.update({"diagonal": diagonal is not None}) + componentwise = opts.pop("componentwise", None) if componentwise is not None: - args.update({'componentwise': True}) + args.update({"componentwise": True}) if opts: - raise ValueError("unrecognized options {}" - .format(', '.join(opts.keys()))) + raise ValueError("unrecognized options {}".format(", ".join(opts.keys()))) return cls(**args) def load_opts(cp, section, tag, skip=None): - """Loads config options for jump proposals. + """ + Loads config options for jump proposals. All `-` in option names are converted to `_` before returning. @@ -396,19 +415,24 @@ def load_opts(cp, section, tag, skip=None): List of parameter names the jump proposal is for. opts : dict Dictionary of option names -> values, where all values are strings. + """ if skip is None: skip = [] params = tag.split(VARARGS_DELIM) # get options - readsection = '-'.join([section, tag]) - opts = {opt.replace('-', '_'): cp.get(readsection, opt) - for opt in cp.options(readsection) if opt not in skip} + readsection = "-".join([section, tag]) + opts = { + opt.replace("-", "_"): cp.get(readsection, opt) + for opt in cp.options(readsection) + if opt not in skip + } return params, opts -def get_variance(params, opts, default=1.): - """Gets variance for jump proposals from the dictionary of options. +def get_variance(params, opts, default=1.0): + """ + Gets variance for jump proposals from the dictionary of options. This looks for ``var_{param}`` for every parameter listed in ``params``. If found, the argument is popped from the given ``opts`` dictionary. If not @@ -430,15 +454,18 @@ def get_variance(params, opts, default=1.): numpy.array Array of variances to use. Order is the same as the parameter names given in ``params``. + """ - varfmt = 'var_{}' - cov = numpy.array([float(opts.pop(varfmt.format(param), default)) - for param in params]) + varfmt = "var_{}" + cov = numpy.array( + [float(opts.pop(varfmt.format(param), default)) for param in params] + ) return cov def get_param_boundaries(params, opts): - """Gets parameter boundaries for jump proposals. + """ + Gets parameter boundaries for jump proposals. The syntax for the options should be ``(min|max)_{param} = value``. Both a minimum and maximum should be provided for every parameter in ``params``. @@ -460,23 +487,27 @@ def get_param_boundaries(params, opts): ------- dict : Dictionary of parameter names -> :py:class:`epsie.proposals.Boundaries` + """ boundaries = {} for param in params: - minbound = opts.pop('min_{}'.format(param), None) + minbound = opts.pop(f"min_{param}", None) if minbound is None: - raise ValueError("Must provide a minimum bound for {p}." - "Syntax is min_{p} = val".format(p=param)) - maxbound = opts.pop('max_{}'.format(param), None) + raise ValueError( + f"Must provide a minimum bound for {param}.Syntax is min_{param} = val" + ) + maxbound = opts.pop(f"max_{param}", None) if maxbound is None: - raise ValueError("Must provide a maximum bound for {p}." - "Syntax is max_{p} = val".format(p=param)) + raise ValueError( + f"Must provide a maximum bound for {param}.Syntax is max_{param} = val" + ) boundaries[param] = Boundaries((float(minbound), float(maxbound))) return boundaries def get_epsie_adaptation_settings(opts, name=None): - """Get settings for Epsie adaptive proposals from a config file. + """ + Get settings for Epsie adaptive proposals from a config file. This requires that ``adaptation_duration`` is in the given dictionary. It will also look for ``adaptation_decay``, ``start_iteration``, and @@ -495,32 +526,34 @@ def get_epsie_adaptation_settings(opts, name=None): ------- dict : Dictionary of argument name -> values. + """ args = {} - adaptation_duration = opts.pop('adaptation_duration', None) + adaptation_duration = opts.pop("adaptation_duration", None) if adaptation_duration is None: if name is not None: - if all(p in name.split('_') for p in ['at', 'adaptive']): - args.update({'adaptation_duration': None}) + if all(p in name.split("_") for p in ["at", "adaptive"]): + args.update({"adaptation_duration": None}) else: raise ValueError("No adaptation_duration specified") else: - args.update({'adaptation_duration': int(adaptation_duration)}) + args.update({"adaptation_duration": int(adaptation_duration)}) # optional args - adaptation_decay = opts.pop('adaptation_decay', None) + adaptation_decay = opts.pop("adaptation_decay", None) if adaptation_decay is not None: - args.update({'adaptation_decay': int(adaptation_decay)}) - start_iteration = opts.pop('start_iteration', None) + args.update({"adaptation_decay": int(adaptation_decay)}) + start_iteration = opts.pop("start_iteration", None) if start_iteration is not None: - args.update({'start_iteration': int(start_iteration)}) - target_rate = opts.pop('target_rate', None) + args.update({"start_iteration": int(start_iteration)}) + target_rate = opts.pop("target_rate", None) if target_rate is not None: - args.update({'target_rate': float(target_rate)}) + args.update({"target_rate": float(target_rate)}) return args def get_epsie_discrete_successive_settings(params, opts): - """Get settings for Epsie successive discrete proposal successive jumps + """ + Get settings for Epsie successive discrete proposal successive jumps from a config file. If ``successive`` is not defined for a parameter then assumes successive @@ -546,9 +579,11 @@ def get_epsie_discrete_successive_settings(params, opts): ------- dict : Dictionary of parameter names -> bools + """ successive = {} for param in params: successive.update( - {param: opts.pop('successive_{}'.format(param), None) is not None}) + {param: opts.pop(f"successive_{param}", None) is not None} + ) return successive diff --git a/pycbc/inference/models/__init__.py b/pycbc/inference/models/__init__.py index ae0eba1a4e1..e4d895b4127 100644 --- a/pycbc/inference/models/__init__.py +++ b/pycbc/inference/models/__init__.py @@ -19,27 +19,39 @@ assuming various noise models. """ - import logging from importlib.metadata import entry_points +from .analytic import ( + TestEggbox, + TestNormal, + TestPosterior, + TestPrior, + TestRosenbrock, + TestVolcano, +) from .base import BaseModel from .base_data import BaseDataModel -from .analytic import (TestEggbox, TestNormal, TestRosenbrock, TestVolcano, - TestPrior, TestPosterior) +from .brute_marg import BruteLISASkyModesMarginalize, BruteParallelGaussianMarginalize +from .gated_gaussian_noise import ( + GatedGaussianMargPhase, + GatedGaussianMargPol, + GatedGaussianNoise, +) from .gaussian_noise import GaussianNoise -from .marginalized_gaussian_noise import MarginalizedPhaseGaussianNoise -from .marginalized_gaussian_noise import MarginalizedPolarization -from .marginalized_gaussian_noise import MarginalizedHMPolPhase -from .marginalized_gaussian_noise import MarginalizedTime -from .brute_marg import BruteParallelGaussianMarginalize -from .brute_marg import BruteLISASkyModesMarginalize -from .gated_gaussian_noise import (GatedGaussianNoise, GatedGaussianMargPol, - GatedGaussianMargPhase) -from .single_template import SingleTemplate +from .hierarchical import ( + HierarchicalModel, + JointPrimaryMarginalizedModel, + MultiSignalModel, +) +from .marginalized_gaussian_noise import ( + MarginalizedHMPolPhase, + MarginalizedPhaseGaussianNoise, + MarginalizedPolarization, + MarginalizedTime, +) from .relbin import Relative, RelativeTime, RelativeTimeDom -from .hierarchical import (HierarchicalModel, MultiSignalModel, - JointPrimaryMarginalizedModel) +from .single_template import SingleTemplate # Used to manage a model instance across multiple cores or MPI _global_instance = None @@ -51,17 +63,19 @@ def _call_global_model(*args, **kwds): def _call_global_model_logprior(*args, **kwds): - """Private function for a calling global's logprior. + """ + Private function for a calling global's logprior. This is needed for samplers that use a separate function for the logprior, like ``emcee_pt``. """ # pylint:disable=not-callable - return _global_instance(*args, callstat='logprior', **kwds) + return _global_instance(*args, callstat="logprior", **kwds) -class CallModel(object): - """Wrapper class for calling models from a sampler. +class CallModel: + """ + Wrapper class for calling models from a sampler. This class can be called like a function, with the parameter values to evaluate provided as a list in the same order as the model's @@ -125,7 +139,8 @@ def __getattr__(self, attr): return getattr(self.model, attr) def __call__(self, param_values, callstat=None, return_all_stats=None): - """Updates the model with the given parameter values, then calls the + """ + Updates the model with the given parameter values, then calls the call function. Parameters @@ -149,6 +164,7 @@ def __call__(self, param_values, callstat=None, return_all_stats=None): param values. Any stat that has not be calculated is set to ``numpy.nan``. This is only returned if ``return_all_stats`` is set to ``True``. + """ if callstat is None: callstat = self.callstat @@ -159,12 +175,12 @@ def __call__(self, param_values, callstat=None, return_all_stats=None): val = getattr(self.model, callstat) if return_all_stats: return val, self.model.get_current_stats() - else: - return val + return val def read_from_config(cp, **kwargs): - r"""Initializes a model from the given config file. + r""" + Initializes a model from the given config file. The section must have a ``name`` argument. The name argument corresponds to the name of the class to initialize. @@ -181,72 +197,82 @@ def read_from_config(cp, **kwargs): ------- cls The initialized model. + """ # use the name to get the distribution name = cp.get("model", "name") return get_model(name).from_config(cp, **kwargs) -_models = {_cls.name: _cls for _cls in ( - TestEggbox, - TestNormal, - TestRosenbrock, - TestVolcano, - TestPosterior, - TestPrior, - GaussianNoise, - MarginalizedPhaseGaussianNoise, - MarginalizedPolarization, - MarginalizedHMPolPhase, - MarginalizedTime, - BruteParallelGaussianMarginalize, - BruteLISASkyModesMarginalize, - GatedGaussianNoise, - GatedGaussianMargPol, - GatedGaussianMargPhase, - SingleTemplate, - Relative, - RelativeTime, - HierarchicalModel, - MultiSignalModel, - RelativeTimeDom, - JointPrimaryMarginalizedModel, -)} +_models = { + _cls.name: _cls + for _cls in ( + TestEggbox, + TestNormal, + TestRosenbrock, + TestVolcano, + TestPosterior, + TestPrior, + GaussianNoise, + MarginalizedPhaseGaussianNoise, + MarginalizedPolarization, + MarginalizedHMPolPhase, + MarginalizedTime, + BruteParallelGaussianMarginalize, + BruteLISASkyModesMarginalize, + GatedGaussianNoise, + GatedGaussianMargPol, + GatedGaussianMargPhase, + SingleTemplate, + Relative, + RelativeTime, + HierarchicalModel, + MultiSignalModel, + RelativeTimeDom, + JointPrimaryMarginalizedModel, + ) +} class _ModelManager(dict): - """Sub-classes dictionary to manage the collection of available models. + """ + Sub-classes dictionary to manage the collection of available models. The first time this is called, any plugin models that are available will be added to the dictionary before returning. """ + def __init__(self, *args, **kwargs): self.retrieve_plugins = True super().__init__(*args, **kwargs) def add_model(self, model): - """Adds a model to the dictionary. + """ + Adds a model to the dictionary. If the given model has the same name as a model already in the dictionary, the original model will be overridden. A warning will be printed in that case. """ if super().__contains__(model.name): - logging.warning("Custom model %s will override a model of the " - "same name. If you don't want this, change the " - "model's name attribute and restart.", model.name) + logging.warning( + "Custom model %s will override a model of the " + "same name. If you don't want this, change the " + "model's name attribute and restart.", + model.name, + ) self[model.name] = model def add_plugins(self): - """Adds any plugin models that are available. + """ + Adds any plugin models that are available. This will only add the plugins if ``self.retrieve_plugins = True``. After this runs, ``self.retrieve_plugins`` is set to ``False``, so that subsequent calls to this will no re-add models. """ - if self.retrieve_plugins: - for plugin in entry_points(group='pycbc.inference.models'): + for plugin in entry_points(group="pycbc.inference.models"): self.add_model(plugin.load()) self.retrieve_plugins = False @@ -312,7 +338,8 @@ def __delitem__(self, *args, **kwargs): def get_models(): - """Returns the dictionary of current models. + """ + Returns the dictionary of current models. Ensures that plugins are added to the dictionary first. """ @@ -321,7 +348,8 @@ def get_models(): def get_model(model_name): - """Retrieve the given model. + """ + Retrieve the given model. Parameters ---------- @@ -332,6 +360,7 @@ def get_model(model_name): ------- model : The requested model. + """ return get_models()[model_name] @@ -342,7 +371,8 @@ def available_models(): def register_model(model): - """Makes a custom model available to PyCBC. + """ + Makes a custom model available to PyCBC. The provided model will be added to the dictionary of models that PyCBC knows about, using the model's ``name`` attribute. If the ``name`` is the @@ -354,5 +384,6 @@ def register_model(model): The model to use. The model should be a sub-class of :py:class:`BaseModel ` to ensure it has the correct API for use within ``pycbc_inference``. + """ get_models().add_model(model) diff --git a/pycbc/inference/models/analytic.py b/pycbc/inference/models/analytic.py index 9b1886f4280..49f50ad1d8c 100644 --- a/pycbc/inference/models/analytic.py +++ b/pycbc/inference/models/analytic.py @@ -19,6 +19,7 @@ """ import logging + import numpy import numpy.random from scipy import stats @@ -27,7 +28,8 @@ class TestNormal(BaseModel): - r"""The test distribution is an multi-variate normal distribution. + r""" + The test distribution is an multi-variate normal distribution. The number of dimensions is set by the number of ``variable_params`` that are passed. For details on the distribution used, see @@ -64,31 +66,32 @@ class TestNormal(BaseModel): {'logjacobian': 0.0, 'loglikelihood': -1.8628770664093453, 'logprior': 0.0} """ + name = "test_normal" def __init__(self, variable_params, mean=None, cov=None, **kwargs): # set up base likelihood parameters - super(TestNormal, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) # store the pdf if mean is None: - mean = [0.]*len(variable_params) + mean = [0.0] * len(variable_params) if cov is None: - cov = [1.]*len(variable_params) + cov = [1.0] * len(variable_params) self._dist = stats.multivariate_normal(mean=mean, cov=cov) # check that the dimension is correct if self._dist.dim != len(variable_params): - raise ValueError("dimension mis-match between variable_params and " - "mean and/or cov") + raise ValueError( + "dimension mis-match between variable_params and mean and/or cov" + ) def _loglikelihood(self): - """Returns the log pdf of the multivariate normal. - """ - return self._dist.logpdf([self.current_params[p] - for p in self.variable_params]) + """Returns the log pdf of the multivariate normal.""" + return self._dist.logpdf([self.current_params[p] for p in self.variable_params]) class TestEggbox(BaseModel): - r"""The test distribution is an 'eggbox' function: + r""" + The test distribution is an 'eggbox' function: .. math:: @@ -106,21 +109,26 @@ class TestEggbox(BaseModel): All other keyword arguments are passed to ``BaseModel``. """ + name = "test_eggbox" def __init__(self, variable_params, **kwargs): # set up base likelihood parameters - super(TestEggbox, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) def _loglikelihood(self): - """Returns the log pdf of the eggbox function. - """ - return (2 + numpy.prod(numpy.cos([ - self.current_params[p]/2. for p in self.variable_params]))) ** 5 + """Returns the log pdf of the eggbox function.""" + return ( + 2 + + numpy.prod( + numpy.cos([self.current_params[p] / 2.0 for p in self.variable_params]) + ) + ) ** 5 class TestRosenbrock(BaseModel): - r"""The test distribution is the Rosenbrock function: + r""" + The test distribution is the Rosenbrock function: .. math:: @@ -138,24 +146,25 @@ class TestRosenbrock(BaseModel): All other keyword arguments are passed to ``BaseModel``. """ + name = "test_rosenbrock" def __init__(self, variable_params, **kwargs): # set up base likelihood parameters - super(TestRosenbrock, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) def _loglikelihood(self): - """Returns the log pdf of the Rosenbrock function. - """ + """Returns the log pdf of the Rosenbrock function.""" logl = 0 p = [self.current_params[p] for p in self.variable_params] for i in range(len(p) - 1): - logl -= ((1 - p[i])**2 + 100 * (p[i+1] - p[i]**2)**2) + logl -= (1 - p[i]) ** 2 + 100 * (p[i + 1] - p[i] ** 2) ** 2 return logl class TestVolcano(BaseModel): - r"""The test distribution is a two-dimensional 'volcano' function: + r""" + The test distribution is a two-dimensional 'volcano' function: .. math:: \Theta = @@ -171,30 +180,35 @@ class TestVolcano(BaseModel): All other keyword arguments are passed to ``BaseModel``. """ + name = "test_volcano" def __init__(self, variable_params, **kwargs): # set up base likelihood parameters - super(TestVolcano, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) # make sure there are exactly two variable args if len(self.variable_params) != 2: - raise ValueError("TestVolcano distribution requires exactly " - "two variable args") + raise ValueError( + "TestVolcano distribution requires exactly two variable args" + ) def _loglikelihood(self): - """Returns the log pdf of the 2D volcano function. - """ + """Returns the log pdf of the 2D volcano function.""" p = [self.current_params[p] for p in self.variable_params] - r = numpy.sqrt(p[0]**2 + p[1]**2) + r = numpy.sqrt(p[0] ** 2 + p[1] ** 2) mu, sigma = 5.0, 2.0 return 25 * ( - numpy.exp(-r/35) + 1 / (sigma * numpy.sqrt(2 * numpy.pi)) * - numpy.exp(-0.5 * ((r - mu) / sigma) ** 2)) + numpy.exp(-r / 35) + + 1 + / (sigma * numpy.sqrt(2 * numpy.pi)) + * numpy.exp(-0.5 * ((r - mu) / sigma) ** 2) + ) class TestPrior(BaseModel): - r"""Uses the prior as the test distribution. + r""" + Uses the prior as the test distribution. Parameters ---------- @@ -204,20 +218,21 @@ class TestPrior(BaseModel): All other keyword arguments are passed to ``BaseModel``. """ + name = "test_prior" def __init__(self, variable_params, **kwargs): # set up base likelihood parameters - super(TestPrior, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) def _loglikelihood(self): - """Returns zero. - """ - return 0. + """Returns zero.""" + return 0.0 class TestPosterior(BaseModel): - r"""Build a test posterior from a set of samples using a kde + r""" + Build a test posterior from a set of samples using a kde Parameters ---------- @@ -232,15 +247,17 @@ class TestPosterior(BaseModel): All other keyword arguments are passed to ``BaseModel``. """ + name = "test_posterior" def __init__(self, variable_params, posterior_file, nsamples, **kwargs): - super(TestPosterior, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) from pycbc.inference.io import loadfile # avoid cyclic import - logging.info('loading test posterior model') + + logging.info("loading test posterior model") inf_file = loadfile(posterior_file) - logging.info('reading samples') + logging.info("reading samples") samples = inf_file.read_samples(variable_params) samples = numpy.array([samples[v] for v in variable_params]) @@ -249,13 +266,12 @@ def __init__(self, variable_params, posterior_file, nsamples, **kwargs): idx = numpy.random.choice(idx, size=int(nsamples), replace=False) samples = samples[:, idx] - logging.info('making kde with %s samples', samples.shape[-1]) + logging.info("making kde with %s samples", samples.shape[-1]) self.kde = stats.gaussian_kde(samples) - logging.info('done initializing test posterior model') + logging.info("done initializing test posterior model") def _loglikelihood(self): - """Returns the log pdf of the test posterior kde - """ + """Returns the log pdf of the test posterior kde""" p = numpy.array([self.current_params[p] for p in self.variable_params]) logpost = self.kde.logpdf(p) return float(logpost[0]) diff --git a/pycbc/inference/models/base.py b/pycbc/inference/models/base.py index 6d112b9df27..d6fbffa1583 100644 --- a/pycbc/inference/models/base.py +++ b/pycbc/inference/models/base.py @@ -22,16 +22,16 @@ # ============================================================================= # -"""Base class for models. -""" +"""Base class for models.""" -import numpy import logging -from abc import (ABCMeta, abstractmethod) +from abc import ABCMeta, abstractmethod from configparser import NoSectionError -from pycbc import (transforms, distributions) -from pycbc.io import FieldArray +import numpy + +from pycbc import distributions, transforms +from pycbc.io import FieldArray # # ============================================================================= @@ -42,18 +42,18 @@ # -class _NoPrior(object): - """Dummy class to just return 0 if no prior is given to a model. - """ +class _NoPrior: + """Dummy class to just return 0 if no prior is given to a model.""" + @staticmethod def apply_boundary_conditions(**params): return params def __call__(self, **params): - return 0. + return 0.0 -class ModelStats(object): +class ModelStats: """Class to hold model's current stat values.""" @property @@ -62,7 +62,8 @@ def statnames(self): return list(self.__dict__.keys()) def getstats(self, names, default=numpy.nan): - """Get the requested stats as a tuple. + """ + Get the requested stats as a tuple. If a requested stat is not an attribute (implying it hasn't been stored), then the default value is returned for that stat. @@ -79,11 +80,13 @@ def getstats(self, names, default=numpy.nan): ------- tuple A tuple of the requested stats. + """ return tuple(getattr(self, n, default) for n in names) def getstatsdict(self, names, default=numpy.nan): - """Get the requested stats as a dictionary. + """ + Get the requested stats as a dictionary. If a requested stat is not an attribute (implying it hasn't been stored), then the default value is returned for that stat. @@ -100,23 +103,28 @@ def getstatsdict(self, names, default=numpy.nan): ------- dict A dictionary of the requested stats. + """ return dict(zip(names, self.getstats(names, default=default))) -class SamplingTransforms(object): - """Provides methods for transforming between sampling parameter space and +class SamplingTransforms: + """ + Provides methods for transforming between sampling parameter space and model parameter space. """ - def __init__(self, variable_params, sampling_params, - replace_parameters, sampling_transforms): + def __init__( + self, variable_params, sampling_params, replace_parameters, sampling_transforms + ): assert len(replace_parameters) == len(sampling_params), ( "number of sampling parameters must be the " - "same as the number of replace parameters") + "same as the number of replace parameters" + ) # pull out the replaced parameters - self.sampling_params = [arg for arg in variable_params - if arg not in replace_parameters] + self.sampling_params = [ + arg for arg in variable_params if arg not in replace_parameters + ] # add the sampling parameters self.sampling_params += sampling_params # sort to make sure we have a consistent order @@ -124,7 +132,8 @@ def __init__(self, variable_params, sampling_params, self.sampling_transforms = sampling_transforms def logjacobian(self, **params): - r"""Returns the log of the jacobian needed to transform pdfs in the + r""" + Returns the log of the jacobian needed to transform pdfs in the ``variable_params`` parameter space to the ``sampling_params`` parameter space. @@ -160,12 +169,19 @@ def logjacobian(self, **params): ------- float : The value of the jacobian. + """ - return numpy.log(abs(transforms.compute_jacobian( - params, self.sampling_transforms, inverse=True))) + return numpy.log( + abs( + transforms.compute_jacobian( + params, self.sampling_transforms, inverse=True + ) + ) + ) def apply(self, samples, inverse=False): - """Applies the sampling transforms to the given samples. + """ + Applies the sampling transforms to the given samples. Parameters ---------- @@ -179,13 +195,16 @@ def apply(self, samples, inverse=False): ------- dict or FieldArray The transformed samples, along with the original samples. + """ - return transforms.apply_transforms(samples, self.sampling_transforms, - inverse=inverse) + return transforms.apply_transforms( + samples, self.sampling_transforms, inverse=inverse + ) @classmethod def from_config(cls, cp, variable_params): - """Gets sampling transforms specified in a config file. + """ + Gets sampling transforms specified in a config file. Sampling parameters and the parameters they replace are read from the ``sampling_params`` section, if it exists. Sampling transforms are @@ -206,26 +225,31 @@ def from_config(cls, cp, variable_params): ------- SamplingTransforms A sampling transforms class. + """ # Check if a sampling_params section is provided try: - sampling_params, replace_parameters = \ - read_sampling_params_from_config(cp) + sampling_params, replace_parameters = read_sampling_params_from_config(cp) except NoSectionError as e: logging.warning("No sampling_params section read from config file") raise e # get sampling transformations sampling_transforms = transforms.read_transforms_from_config( - cp, 'sampling_transforms') - logging.info("Sampling in {} in place of {}".format( - ', '.join(sampling_params), ', '.join(replace_parameters))) - return cls(variable_params, sampling_params, - replace_parameters, sampling_transforms) - - -def read_sampling_params_from_config(cp, section_group=None, - section='sampling_params'): - """Reads sampling parameters from the given config file. + cp, "sampling_transforms" + ) + logging.info( + "Sampling in {} in place of {}".format( + ", ".join(sampling_params), ", ".join(replace_parameters) + ) + ) + return cls( + variable_params, sampling_params, replace_parameters, sampling_transforms + ) + + +def read_sampling_params_from_config(cp, section_group=None, section="sampling_params"): + """ + Reads sampling parameters from the given config file. Parameters are read from the `[({section_group}_){section}]` section. The options should list the variable args to transform; the parameters they @@ -261,18 +285,19 @@ def read_sampling_params_from_config(cp, section_group=None, The list of sampling parameters to use instead. replaced_params : list The list of variable args to replace in the sampler. + """ if section_group is not None: - section_prefix = '{}_'.format(section_group) + section_prefix = f"{section_group}_" else: - section_prefix = '' + section_prefix = "" section = section_prefix + section replaced_params = set() sampling_params = set() for args in cp.options(section): map_args = cp.get(section, args) - sampling_params.update(set(map(str.strip, map_args.split(',')))) - replaced_params.update(set(map(str.strip, args.split(',')))) + sampling_params.update(set(map(str.strip, map_args.split(",")))) + replaced_params.update(set(map(str.strip, args.split(",")))) return sorted(sampling_params), sorted(replaced_params) @@ -286,7 +311,8 @@ def read_sampling_params_from_config(cp, section_group=None, class BaseModel(metaclass=ABCMeta): - r"""Base class for all models. + r""" + Base class for all models. Given some model :math:`h` with parameters :math:`\Theta`, Bayes Theorem states that the probability of observing parameter values :math:`\vartheta` @@ -341,11 +367,20 @@ class BaseModel(metaclass=ABCMeta): the likelihood is most easily defined in. Since these are used solely for converting parameters, and not for rescaling the parameter space, a Jacobian is not required for these transforms. + """ + name = None - def __init__(self, variable_params, static_params=None, prior=None, - sampling_transforms=None, waveform_transforms=None, **kwargs): + def __init__( + self, + variable_params, + static_params=None, + prior=None, + sampling_transforms=None, + waveform_transforms=None, + **kwargs, + ): # store variable and static args self.variable_params = variable_params self.static_params = static_params @@ -353,8 +388,7 @@ def __init__(self, variable_params, static_params=None, prior=None, if prior is None: self.prior_distribution = _NoPrior() elif set(prior.variable_args) != set(variable_params): - raise ValueError("variable params of prior and model must be the " - "same") + raise ValueError("variable params of prior and model must be the same") else: self.prior_distribution = prior # store transforms @@ -391,7 +425,8 @@ def static_params(self, static_params): @property def sampling_params(self): - """Returns the sampling parameters. + """ + Returns the sampling parameters. If ``sampling_transforms`` is None, this is the same as the ``variable_params``. @@ -403,7 +438,8 @@ def sampling_params(self): return sampling_params def update(self, **params): - """Updates the current parameter positions and resets stats. + """ + Updates the current parameter positions and resets stats. If any sampling transforms are specified, they are applied to the params before being stored. @@ -417,18 +453,20 @@ def update(self, **params): @property def current_params(self): if self._current_params is None: - raise ValueError("no parameters values currently stored; " - "run update to add some") + raise ValueError( + "no parameters values currently stored; run update to add some" + ) return self._current_params @property def default_stats(self): """The stats that ``get_current_stats`` returns by default.""" - return ['logjacobian', 'logprior', 'loglikelihood'] + self._extra_stats + return ["logjacobian", "logprior", "loglikelihood"] + self._extra_stats @property def _extra_stats(self): - """Allows child classes to add more stats to the default stats. + """ + Allows child classes to add more stats to the default stats. This returns an empty list; classes that inherit should override this property if they want to add extra stats. @@ -436,7 +474,8 @@ def _extra_stats(self): return [] def get_current_stats(self, names=None): - """Return one or more of the current stats as a tuple. + """ + Return one or more of the current stats as a tuple. This function does no computation. It only returns what has already been calculated. If a stat hasn't been calculated, it will be returned @@ -453,6 +492,7 @@ def get_current_stats(self, names=None): tuple : The current values of the requested stats, as a tuple. The order of the stats is the same as the names. + """ if names is None: names = self.default_stats @@ -460,7 +500,8 @@ def get_current_stats(self, names=None): @property def current_stats(self): - """Return the ``default_stats`` as a dict. + """ + Return the ``default_stats`` as a dict. This does no computation. It only returns what has already been calculated. If a stat hasn't been calculated, it will be returned @@ -470,11 +511,13 @@ def current_stats(self): ------- dict : Dictionary of stat names -> current stat values. + """ return self._current_stats.getstatsdict(self.default_stats) def _trytoget(self, statname, fallback, apply_transforms=False, **kwargs): - r"""Helper function to get a stat from ``_current_stats``. + r""" + Helper function to get a stat from ``_current_stats``. If the statistic hasn't been calculated, ``_current_stats`` will raise an ``AttributeError``. In that case, the ``fallback`` function will @@ -497,6 +540,7 @@ def _trytoget(self, statname, fallback, apply_transforms=False, **kwargs): ------- float : The value of the property. + """ try: return getattr(self._current_stats, statname) @@ -504,32 +548,36 @@ def _trytoget(self, statname, fallback, apply_transforms=False, **kwargs): # apply waveform transforms if requested if apply_transforms and self.waveform_transforms is not None: self._current_params = transforms.apply_transforms( - self._current_params, self.waveform_transforms, - inverse=False) + self._current_params, self.waveform_transforms, inverse=False + ) val = fallback(**kwargs) setattr(self._current_stats, statname, val) return val @property def loglikelihood(self): - """The log likelihood at the current parameters. + """ + The log likelihood at the current parameters. This will initially try to return the ``current_stats.loglikelihood``. If that raises an ``AttributeError``, will call `_loglikelihood`` to calculate it and store it to ``current_stats``. """ - return self._trytoget('loglikelihood', self._loglikelihood, - apply_transforms=True) + return self._trytoget( + "loglikelihood", self._loglikelihood, apply_transforms=True + ) @abstractmethod def _loglikelihood(self): - """Low-level function that calculates the log likelihood of the current - params.""" - pass + """ + Low-level function that calculates the log likelihood of the current + params. + """ @property def logjacobian(self): - r"""The log jacobian of the sampling transforms at the current postion. + r""" + The log jacobian of the sampling transforms at the current postion. If no sampling transforms were provided, will just return 0. @@ -543,22 +591,22 @@ def logjacobian(self): ------- float : The value of the jacobian. + """ - return self._trytoget('logjacobian', self._logjacobian) + return self._trytoget("logjacobian", self._logjacobian) def _logjacobian(self): """Calculates the logjacobian of the current parameters.""" if self.sampling_transforms is None: - logj = 0. + logj = 0.0 else: - logj = self.sampling_transforms.logjacobian( - **self.current_params) + logj = self.sampling_transforms.logjacobian(**self.current_params) return logj @property def logprior(self): """Returns the log prior at the current parameters.""" - return self._trytoget('logprior', self._logprior) + return self._trytoget("logprior", self._logprior) def _logprior(self): """Calculates the log prior at the current parameters.""" @@ -570,7 +618,8 @@ def _logprior(self): @property def logposterior(self): - """Returns the log of the posterior of the current parameter values. + """ + Returns the log of the posterior of the current parameter values. The logprior is calculated first. If the logprior returns ``-inf`` (possibly indicating a non-physical point), then the ``loglikelihood`` @@ -579,11 +628,11 @@ def logposterior(self): logp = self.logprior if logp == -numpy.inf: return logp - else: - return logp + self.loglikelihood + return logp + self.loglikelihood def prior_rvs(self, size=1, prior=None): - """Returns random variates drawn from the prior. + """ + Returns random variates drawn from the prior. If the ``sampling_params`` are different from the ``variable_params``, the variates are transformed to the `sampling_params` parameter space @@ -600,6 +649,7 @@ def prior_rvs(self, size=1, prior=None): ------- FieldArray A field array of the random values. + """ # draw values from the prior if prior is None: @@ -609,13 +659,15 @@ def prior_rvs(self, size=1, prior=None): if self.sampling_transforms is not None: ptrans = self.sampling_transforms.apply(p0) # pull out the sampling args - p0 = FieldArray.from_arrays([ptrans[arg] - for arg in self.sampling_params], - names=self.sampling_params) + p0 = FieldArray.from_arrays( + [ptrans[arg] for arg in self.sampling_params], + names=self.sampling_params, + ) return p0 def _transform_params(self, **params): - r"""Applies sampling transforms and boundary conditions to parameters. + r""" + Applies sampling transforms and boundary conditions to parameters. Parameters ---------- @@ -626,6 +678,7 @@ def _transform_params(self, **params): ------- dict A dictionary of the transformed parameters. + """ # apply inverse transforms to go from sampling parameters to # variable args @@ -640,7 +693,8 @@ def _transform_params(self, **params): # @staticmethod def extra_args_from_config(cp, section, skip_args=None, dtypes=None): - """Gets any additional keyword in the given config file. + """ + Gets any additional keyword in the given config file. Parameters ---------- @@ -660,14 +714,14 @@ def extra_args_from_config(cp, section, skip_args=None, dtypes=None): ------- dict Dictionary of keyword arguments read from the config file. + """ kwargs = {} if dtypes is None: dtypes = {} if skip_args is None: skip_args = [] - read_args = [opt for opt in cp.options(section) - if opt not in skip_args] + read_args = [opt for opt in cp.options(section) if opt not in skip_args] for opt in read_args: val = cp.get(section, opt) # try to cast the value if a datatype was specified for this opt @@ -679,9 +733,11 @@ def extra_args_from_config(cp, section, skip_args=None, dtypes=None): return kwargs @staticmethod - def prior_from_config(cp, variable_params, static_params, prior_section, - constraint_section): - """Gets arguments and keyword arguments from a config file. + def prior_from_config( + cp, variable_params, static_params, prior_section, constraint_section + ): + """ + Gets arguments and keyword arguments from a config file. Parameters ---------- @@ -700,18 +756,22 @@ def prior_from_config(cp, variable_params, static_params, prior_section, ------- pycbc.distributions.JointDistribution The prior. + """ # get prior distribution for each variable parameter logging.info("Setting up priors for each parameter") dists = distributions.read_distributions_from_config(cp, prior_section) constraints = distributions.read_constraints_from_config( - cp, constraint_section, static_args=static_params) - return distributions.JointDistribution(variable_params, *dists, - constraints=constraints) + cp, constraint_section, static_args=static_params + ) + return distributions.JointDistribution( + variable_params, *dists, constraints=constraints + ) @classmethod def _init_args_from_config(cls, cp): - """Helper function for loading parameters. + """ + Helper function for loading parameters. This retrieves the prior, variable parameters, static parameterss, constraints, sampling transforms, and waveform transforms @@ -729,55 +789,65 @@ def _init_args_from_config(cls, cp): ``static_params``, ``prior``, and ``sampling_transforms``. If waveform transforms are in the config file, will also have ``waveform_transforms``. + """ section = "model" prior_section = "prior" - vparams_section = 'variable_params' - sparams_section = 'static_params' - constraint_section = 'constraint' + vparams_section = "variable_params" + sparams_section = "static_params" + constraint_section = "constraint" # check that the name exists and matches - name = cp.get(section, 'name') + name = cp.get(section, "name") if name != cls.name: - raise ValueError("section's {} name does not match mine {}".format( - name, cls.name)) + raise ValueError( + f"section's {name} name does not match mine {cls.name}" + ) # get model parameters variable_params, static_params = distributions.read_params_from_config( - cp, prior_section=prior_section, vargs_section=vparams_section, - sargs_section=sparams_section) + cp, + prior_section=prior_section, + vargs_section=vparams_section, + sargs_section=sparams_section, + ) # get prior prior = cls.prior_from_config( - cp, variable_params, static_params, prior_section, - constraint_section) - args = {'variable_params': variable_params, - 'static_params': static_params, - 'prior': prior} + cp, variable_params, static_params, prior_section, constraint_section + ) + args = { + "variable_params": variable_params, + "static_params": static_params, + "prior": prior, + } # try to load sampling transforms try: - sampling_transforms = SamplingTransforms.from_config( - cp, variable_params) + sampling_transforms = SamplingTransforms.from_config(cp, variable_params) except NoSectionError: sampling_transforms = None - args['sampling_transforms'] = sampling_transforms + args["sampling_transforms"] = sampling_transforms # get any waveform transforms - if any(cp.get_subsections('waveform_transforms')): + if any(cp.get_subsections("waveform_transforms")): logging.info("Loading waveform transforms") waveform_transforms = transforms.read_transforms_from_config( - cp, 'waveform_transforms') - args['waveform_transforms'] = waveform_transforms + cp, "waveform_transforms" + ) + args["waveform_transforms"] = waveform_transforms else: waveform_transforms = [] # safety check for spins # we won't do this if the following exists in the config file ignore = "no_err_on_missing_cartesian_spins" - check_for_cartesian_spins(1, variable_params, static_params, - waveform_transforms, cp, ignore) - check_for_cartesian_spins(2, variable_params, static_params, - waveform_transforms, cp, ignore) + check_for_cartesian_spins( + 1, variable_params, static_params, waveform_transforms, cp, ignore + ) + check_for_cartesian_spins( + 2, variable_params, static_params, waveform_transforms, cp, ignore + ) return args @classmethod def from_config(cls, cp, **kwargs): - r"""Initializes an instance of this class from the given config file. + r""" + Initializes an instance of this class from the given config file. Parameters ---------- @@ -786,16 +856,17 @@ def from_config(cls, cp, **kwargs): \**kwargs : All additional keyword arguments are passed to the class. Any provided keyword will over ride what is in the config file. + """ args = cls._init_args_from_config(cp) # get any other keyword arguments provided in the model section - args.update(cls.extra_args_from_config(cp, "model", - skip_args=['name'])) + args.update(cls.extra_args_from_config(cp, "model", skip_args=["name"])) args.update(kwargs) return cls(**args) def write_metadata(self, fp, group=None): - """Writes metadata to the given file handler. + """ + Writes metadata to the given file handler. Parameters ---------- @@ -805,17 +876,20 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ attrs = fp.getattrs(group=group) - attrs['model'] = self.name - attrs['variable_params'] = list(map(str, self.variable_params)) - attrs['sampling_params'] = list(map(str, self.sampling_params)) + attrs["model"] = self.name + attrs["variable_params"] = list(map(str, self.variable_params)) + attrs["sampling_params"] = list(map(str, self.sampling_params)) fp.write_kwargs_to_attrs(attrs, static_params=self.static_params) -def check_for_cartesian_spins(which, variable_params, static_params, - waveform_transforms, cp, ignore): - """Checks that if any spin parameters exist, cartesian spins also exist. +def check_for_cartesian_spins( + which, variable_params, static_params, waveform_transforms, cp, ignore +): + """ + Checks that if any spin parameters exist, cartesian spins also exist. This looks for parameters starting with ``spinN`` in the variable and static params, where ``N`` is either 1 or 2 (specified by the ``which`` @@ -842,11 +916,14 @@ def check_for_cartesian_spins(which, variable_params, static_params, ignore : str The section to check for in the config file. If the section is present in the config file, the check will not be done. + """ # don't do this check if the config file has the ignore section if cp.has_section(ignore): - logging.info("[{}] found in config file; not performing check for " - "cartesian spin{} parameters".format(ignore, which)) + logging.info( + f"[{ignore}] found in config file; not performing check for " + f"cartesian spin{which} parameters" + ) return errmsg = ( "Spin parameters {sp} found in variable/static " @@ -871,17 +948,21 @@ def check_for_cartesian_spins(which, variable_params, static_params, "(e.g., you are using a custom waveform or model) add\n\n" "[{ignore}]\n\n" "to your config file as an empty section and rerun. This check will " - "not be performed in that case.") + "not be performed in that case." + ) allparams = set(variable_params) | set(static_params.keys()) - spinparams = set(p for p in allparams - if p.startswith('spin{}'.format(which))) + spinparams = set(p for p in allparams if p.startswith(f"spin{which}")) if any(spinparams): - cartspins = set('spin{}{}'.format(which, coord) - for coord in ['x', 'y', 'z']) + cartspins = set(f"spin{which}{coord}" for coord in ["x", "y", "z"]) # add any parameters to all params that will be output by waveform # transforms allparams = allparams.union(*[t.outputs for t in waveform_transforms]) if not any(allparams & cartspins): - raise ValueError(errmsg.format(sp=', '.join(spinparams), - cp=', '.join(cartspins), - n=which, ignore=ignore)) + raise ValueError( + errmsg.format( + sp=", ".join(spinparams), + cp=", ".join(cartspins), + n=which, + ignore=ignore, + ) + ) diff --git a/pycbc/inference/models/base_data.py b/pycbc/inference/models/base_data.py index 1de12be1cc9..9439c0b758a 100644 --- a/pycbc/inference/models/base_data.py +++ b/pycbc/inference/models/base_data.py @@ -22,16 +22,18 @@ # ============================================================================= # -"""Base classes for mofdels with data. -""" +"""Base classes for mofdels with data.""" + +from abc import ABCMeta, abstractmethod import numpy -from abc import (ABCMeta, abstractmethod) + from .base import BaseModel class BaseDataModel(BaseModel, metaclass=ABCMeta): - r"""Base class for models that require data and a waveform generator. + r""" + Base class for models that require data and a waveform generator. This adds propeties for the log of the likelihood that the data contain noise, ``lognl``, and the log likelihood ratio ``loglr``. @@ -63,17 +65,26 @@ class BaseDataModel(BaseModel, metaclass=ABCMeta): See ``BaseModel`` for additional attributes and properties. + """ - def __init__(self, variable_params, data, recalibration=None, gates=None, - injection_file=None, no_save_data=False, **kwargs): + def __init__( + self, + variable_params, + data, + recalibration=None, + gates=None, + injection_file=None, + no_save_data=False, + **kwargs, + ): self._data = None self.data = data self.recalibration = recalibration self.no_save_data = no_save_data self.gates = gates self.injection_file = injection_file - super(BaseDataModel, self).__init__(variable_params, **kwargs) + super().__init__(variable_params, **kwargs) @property def data(self): @@ -88,26 +99,27 @@ def data(self, data): @property def _extra_stats(self): """Adds ``loglr`` and ``lognl`` to the ``default_stats``.""" - return ['loglr', 'lognl'] + return ["loglr", "lognl"] @property def lognl(self): - """The log likelihood of the model assuming the data is noise. + """ + The log likelihood of the model assuming the data is noise. This will initially try to return the ``current_stats.lognl``. If that raises an ``AttributeError``, will call `_lognl`` to calculate it and store it to ``current_stats``. """ - return self._trytoget('lognl', self._lognl) + return self._trytoget("lognl", self._lognl) @abstractmethod def _lognl(self): """Low-level function that calculates the lognl.""" - pass @property def loglr(self): - """The log likelihood ratio at the current parameters, + """ + The log likelihood ratio at the current parameters, or the inner product and if set the flag `self.return_sh_hh` to be True. @@ -115,16 +127,16 @@ def loglr(self): If that raises an ``AttributeError``, will call `_loglr`` to calculate it and store it to ``current_stats``. """ - return self._trytoget('loglr', self._loglr, apply_transforms=True) + return self._trytoget("loglr", self._loglr, apply_transforms=True) @abstractmethod def _loglr(self): """Low-level function that calculates the loglr.""" - pass @property def logplr(self): - """Returns the log of the prior-weighted likelihood ratio at the + """ + Returns the log of the prior-weighted likelihood ratio at the current parameter values. The logprior is calculated first. If the logprior returns ``-inf`` @@ -134,8 +146,7 @@ def logplr(self): logp = self.logprior if logp == -numpy.inf: return logp - else: - return logp + self.loglr + return logp + self.loglr @property def detectors(self): @@ -143,7 +154,8 @@ def detectors(self): return list(self._data.keys()) def write_metadata(self, fp, group=None): - """Adds data to the metadata that's written. + """ + Adds data to the metadata that's written. Parameters ---------- @@ -153,6 +165,7 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ super().write_metadata(fp, group=group) if not self.no_save_data: diff --git a/pycbc/inference/models/brute_marg.py b/pycbc/inference/models/brute_marg.py index 0b419838dd6..481ca612711 100644 --- a/pycbc/inference/models/brute_marg.py +++ b/pycbc/inference/models/brute_marg.py @@ -13,21 +13,26 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides model classes that do brute force marginalization +""" +This module provides model classes that do brute force marginalization using at the likelihood level. """ -import math + import logging +import math + import numpy +from scipy.special import logsumexp from pycbc.pool import BroadcastPool as Pool -from scipy.special import logsumexp from .gaussian_noise import BaseGaussianNoise from .tools import draw_sample _model = None -class likelihood_wrapper(object): + + +class likelihood_wrapper: def __init__(self, model): global _model _model = model @@ -38,17 +43,17 @@ def __call__(self, params): loglr = _model.loglr return loglr, _model.current_stats + class BruteParallelGaussianMarginalize(BaseGaussianNoise): name = "brute_parallel_gaussian_marginalize" - def __init__(self, variable_params, - cores=10, - base_model=None, - marginalize_phase=None, - **kwds): + def __init__( + self, variable_params, cores=10, base_model=None, marginalize_phase=None, **kwds + ): super().__init__(variable_params, **kwds) from pycbc.inference.models import models + self.model = models[base_model](variable_params, **kwds) self.call = likelihood_wrapper(self.model) @@ -65,9 +70,9 @@ def __init__(self, variable_params, @property def _extra_stats(self): stats = self.model._extra_stats - stats.append('maxl_phase') - if 'maxl_loglr' not in stats: - stats.append('maxl_loglr') + stats.append("maxl_phase") + if "maxl_loglr" not in stats: + stats.append("maxl_loglr") return stats def _loglr(self): @@ -75,15 +80,15 @@ def _loglr(self): params = [] for p in self.phase: pref = self.current_params.copy() - pref['coa_phase'] = p + pref["coa_phase"] = p params.append(pref) vals = list(self.pool.map(self.call, params)) loglr = numpy.array([v[0] for v in vals]) # get the maxl values - if 'maxl_loglr' not in self.model._extra_stats: + if "maxl_loglr" not in self.model._extra_stats: maxl_loglrs = loglr else: - maxl_loglrs = numpy.array([v[1]['maxl_loglr'] for v in vals]) + maxl_loglrs = numpy.array([v[1]["maxl_loglr"] for v in vals]) maxidx = maxl_loglrs.argmax() maxstats = vals[maxidx][1] maxphase = self.phase[maxidx] @@ -99,18 +104,17 @@ def _loglr(self): class BruteLISASkyModesMarginalize(BaseGaussianNoise): name = "brute_lisa_sky_modes_marginalize" - def __init__(self, variable_params, - cores=1, - loop_polarization=False, - base_model=None, - **kwds): + def __init__( + self, variable_params, cores=1, loop_polarization=False, base_model=None, **kwds + ): super().__init__(variable_params, **kwds) from pycbc.inference.models import models - kwds.update(models[base_model].extra_args_from_config( - kwds['config_object'], - "model", - skip_args=[]) + + kwds.update( + models[base_model].extra_args_from_config( + kwds["config_object"], "model", skip_args=[] + ) ) self.model = models[base_model](variable_params, **kwds) @@ -156,21 +160,20 @@ def _loglr(self): max_llr_idx = loglr.argmax() max_llr = loglr[max_llr_idx] marg_lrfac = sum([math.exp(llr - max_llr) for llr in loglr]) - marg_llr = max_llr + math.log(marg_lrfac/self.num_sky_modes) + marg_llr = max_llr + math.log(marg_lrfac / self.num_sky_modes) # set the stats for sym_num in range(self.num_sky_modes): - setattr(self._current_stats, f'llr_mode_{sym_num}', loglr[sym_num]) + setattr(self._current_stats, f"llr_mode_{sym_num}", loglr[sym_num]) return marg_llr def _apply_sky_point_rotation(self, pref, sky_num): - """ Apply the sky point rotation for mode sky_num to parameters pref - """ - lambdal = pref['eclipticlongitude'] - beta = pref['eclipticlatitude'] - psi = pref['polarization'] - inc = pref['inclination'] + """Apply the sky point rotation for mode sky_num to parameters pref""" + lambdal = pref["eclipticlongitude"] + beta = pref["eclipticlatitude"] + psi = pref["polarization"] + inc = pref["inclination"] pol_num = sky_num // 8 sky_num = sky_num % 8 @@ -179,51 +182,51 @@ def _apply_sky_point_rotation(self, pref, sky_num): # Apply latitude symmetry mode if lat_num: - beta = - beta + beta = -beta inc = numpy.pi - inc psi = numpy.pi - psi # Apply longitudonal symmetry mode - lambdal = (lambdal + long_num * 0.5 * numpy.pi) % (2*numpy.pi) - psi = (psi + long_num * 0.5 * numpy.pi) % (2*numpy.pi) + lambdal = (lambdal + long_num * 0.5 * numpy.pi) % (2 * numpy.pi) + psi = (psi + long_num * 0.5 * numpy.pi) % (2 * numpy.pi) # Apply additional polarization mode (shouldn't be needed) if pol_num: - psi = psi + (math.pi / 2.) + psi = psi + (math.pi / 2.0) - pref['eclipticlongitude'] = lambdal - pref['eclipticlatitude'] = beta - pref['polarization'] = psi - pref['inclination'] = inc + pref["eclipticlongitude"] = lambdal + pref["eclipticlatitude"] = beta + pref["polarization"] = psi + pref["inclination"] = inc @classmethod def from_config(cls, cp, **kwargs): - kwargs['config_object'] = cp - return super(BruteLISASkyModesMarginalize, cls).from_config( - cp, - **kwargs - ) + kwargs["config_object"] = cp + return super().from_config(cp, **kwargs) def reconstruct(self, seed=None): - """ Reconstruct a point from unwrapping the 8-fold sky symmetry - """ + """Reconstruct a point from unwrapping the 8-fold sky symmetry""" if seed: numpy.random.seed(seed) rec = {} - logging.info('Reconstruct LISA sky mode symmetry') + logging.info("Reconstruct LISA sky mode symmetry") self.reconstruct_sky_points = True loglr = self.loglr xl = draw_sample(loglr) - logging.info('Found point %d', xl) + logging.info("Found point %d", xl) # Undo rotations pref = self.current_params.copy() self._apply_sky_point_rotation(pref, xl) - for val in ['polarization', 'eclipticlongitude', 'eclipticlatitude', - 'inclination']: + for val in [ + "polarization", + "eclipticlongitude", + "eclipticlatitude", + "inclination", + ]: rec[val] = pref[val] - rec['loglr'] = loglr[xl] - rec['loglikelihood'] = self.lognl + rec['loglr'] + rec["loglr"] = loglr[xl] + rec["loglikelihood"] = self.lognl + rec["loglr"] self.reconstruct_sky_points = False return self.model.reconstruct(seed=seed, rec=rec) diff --git a/pycbc/inference/models/data_utils.py b/pycbc/inference/models/data_utils.py index 58f94228844..12e71dc30b6 100644 --- a/pycbc/inference/models/data_utils.py +++ b/pycbc/inference/models/data_utils.py @@ -13,31 +13,40 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Utilities for loading data for models. -""" +"""Utilities for loading data for models.""" import logging from argparse import ArgumentParser from time import sleep + import numpy + try: from mpi4py import MPI except ImportError: MPI = None +from pycbc import dq, strain +from pycbc.psd import ( + from_cli_multi_ifos as psd_from_cli_multi_ifos, +) +from pycbc.psd import ( + insert_psd_option_group_multi_ifo, + verify_psd_options_multi_ifo, +) +from pycbc.strain import ( + apply_gates_to_fd, + apply_gates_to_td, + gates_from_cli, + psd_gates_from_cli, + verify_strain_options_multi_ifo, +) from pycbc.types import MultiDetOptionAction -from pycbc.psd import (insert_psd_option_group_multi_ifo, - from_cli_multi_ifos as psd_from_cli_multi_ifos, - verify_psd_options_multi_ifo) -from pycbc import strain -from pycbc.strain import (gates_from_cli, psd_gates_from_cli, - apply_gates_to_td, apply_gates_to_fd, - verify_strain_options_multi_ifo) -from pycbc import dq def strain_from_cli_multi_ifos(*args, **kwargs): - """Wrapper around strain.from_cli_multi_ifos that tries a few times before + """ + Wrapper around strain.from_cli_multi_ifos that tries a few times before quiting. When running in a parallel environment, multiple concurrent queries to the @@ -64,85 +73,137 @@ def strain_from_cli_multi_ifos(*args, **kwargs): # ============================================================================= # class NoValidDataError(Exception): - """This should be raised if a continous segment of valid data could not be + """ + This should be raised if a continous segment of valid data could not be found. """ - pass + def create_data_parser(): """Creates an argument parser for loading GW data.""" parser = ArgumentParser() # add data options - parser.add_argument("--instruments", type=str, nargs="+", required=True, - help="Instruments to analyze, eg. H1 L1.") - parser.add_argument("--trigger-time", type=float, default=0., - help="Reference GPS time (at geocenter) from which " - "the (anlaysis|psd)-(start|end)-time options are " - "measured. The integer seconds will be used. " - "Default is 0; i.e., if not provided, " - "the analysis and psd times should be in GPS " - "seconds.") - parser.add_argument("--analysis-start-time", type=int, required=True, - nargs='+', action=MultiDetOptionAction, - metavar='IFO:TIME', - help="The start time to use for the analysis, " - "measured with respect to the trigger-time. " - "If psd-inverse-length is provided, the given " - "start time will be padded by half that length " - "to account for wrap-around effects.") - parser.add_argument("--analysis-end-time", type=int, required=True, - nargs='+', action=MultiDetOptionAction, - metavar='IFO:TIME', - help="The end time to use for the analysis, " - "measured with respect to the trigger-time. " - "If psd-inverse-length is provided, the given " - "end time will be padded by half that length " - "to account for wrap-around effects.") - parser.add_argument("--psd-start-time", type=int, default=None, - nargs='+', action=MultiDetOptionAction, - metavar='IFO:TIME', - help="Start time to use for PSD estimation, measured " - "with respect to the trigger-time.") - parser.add_argument("--psd-end-time", type=int, default=None, - nargs='+', action=MultiDetOptionAction, - metavar='IFO:TIME', - help="End time to use for PSD estimation, measured " - "with respect to the trigger-time.") - parser.add_argument("--data-conditioning-low-freq", type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:FLOW', dest="low_frequency_cutoff", - help="Low frequency cutoff of the data. Needed for " - "PSD estimation and when creating fake strain. " - "If not provided, will use the model's " - "low-frequency-cutoff.") + parser.add_argument( + "--instruments", + type=str, + nargs="+", + required=True, + help="Instruments to analyze, eg. H1 L1.", + ) + parser.add_argument( + "--trigger-time", + type=float, + default=0.0, + help="Reference GPS time (at geocenter) from which " + "the (anlaysis|psd)-(start|end)-time options are " + "measured. The integer seconds will be used. " + "Default is 0; i.e., if not provided, " + "the analysis and psd times should be in GPS " + "seconds.", + ) + parser.add_argument( + "--analysis-start-time", + type=int, + required=True, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="The start time to use for the analysis, " + "measured with respect to the trigger-time. " + "If psd-inverse-length is provided, the given " + "start time will be padded by half that length " + "to account for wrap-around effects.", + ) + parser.add_argument( + "--analysis-end-time", + type=int, + required=True, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="The end time to use for the analysis, " + "measured with respect to the trigger-time. " + "If psd-inverse-length is provided, the given " + "end time will be padded by half that length " + "to account for wrap-around effects.", + ) + parser.add_argument( + "--psd-start-time", + type=int, + default=None, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="Start time to use for PSD estimation, measured " + "with respect to the trigger-time.", + ) + parser.add_argument( + "--psd-end-time", + type=int, + default=None, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="End time to use for PSD estimation, measured " + "with respect to the trigger-time.", + ) + parser.add_argument( + "--data-conditioning-low-freq", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FLOW", + dest="low_frequency_cutoff", + help="Low frequency cutoff of the data. Needed for " + "PSD estimation and when creating fake strain. " + "If not provided, will use the model's " + "low-frequency-cutoff.", + ) insert_psd_option_group_multi_ifo(parser) strain.insert_strain_option_group_multi_ifo(parser, gps_times=False) strain.add_gate_option_group(parser) # add arguments for dq - dqgroup = parser.add_argument_group("Options for quering data quality " - "(DQ)") - dqgroup.add_argument('--dq-segment-name', default='DATA', - help='The status flag to query for data quality. ' - 'Default is "DATA".') - dqgroup.add_argument('--dq-source', choices=['any', 'GWOSC', 'dqsegdb'], - default='any', - help='Where to look for DQ information. If "any" ' - '(the default) will first try GWOSC, then ' - 'dqsegdb.') - dqgroup.add_argument('--dq-server', default='https://segments.ligo.org', - help='The server to use for dqsegdb.') - dqgroup.add_argument('--veto-definer', default=None, - help='Path to a veto definer file that defines ' - 'groups of flags, which themselves define a set ' - 'of DQ segments.') + dqgroup = parser.add_argument_group("Options for quering data quality (DQ)") + dqgroup.add_argument( + "--dq-segment-name", + default="DATA", + help='The status flag to query for data quality. Default is "DATA".', + ) + dqgroup.add_argument( + "--dq-source", + choices=["any", "GWOSC", "dqsegdb"], + default="any", + help='Where to look for DQ information. If "any" ' + "(the default) will first try GWOSC, then " + "dqsegdb.", + ) + dqgroup.add_argument( + "--dq-server", + default="https://segments.ligo.org", + help="The server to use for dqsegdb.", + ) + dqgroup.add_argument( + "--veto-definer", + default=None, + help="Path to a veto definer file that defines " + "groups of flags, which themselves define a set " + "of DQ segments.", + ) return parser -def check_validtimes(detector, gps_start, gps_end, shift_to_valid=False, - max_shift=None, segment_name='DATA', - **kwargs): - r"""Checks DQ server to see if the given times are in a valid segment. +def check_validtimes( + detector, + gps_start, + gps_end, + shift_to_valid=False, + max_shift=None, + segment_name="DATA", + **kwargs, +): + r""" + Checks DQ server to see if the given times are in a valid segment. If the ``shift_to_valid`` flag is provided, the times will be shifted left or right to try to find a continous valid block nearby. The shifting starts @@ -183,6 +244,7 @@ def check_validtimes(detector, gps_start, gps_end, shift_to_valid=False, use_end : int The end time to use. If ``shift_to_valid`` is True, this may differ from the given GPS end time. + """ # expand the times checked encase we need to shift if max_shift is None: @@ -192,13 +254,14 @@ def check_validtimes(detector, gps_start, gps_end, shift_to_valid=False, # if we're running in an mpi enviornment and we're not the parent process, # we'll wait before quering the segment database. This will result in # getting the segments from the cache, so as not to overload the database - if MPI is not None and (MPI.COMM_WORLD.Get_size() > 1 and - MPI.COMM_WORLD.Get_rank() != 0): + if MPI is not None and ( + MPI.COMM_WORLD.Get_size() > 1 and MPI.COMM_WORLD.Get_rank() != 0 + ): # we'll wait for 2 minutes sleep(120) - validsegs = dq.query_flag(detector, segment_name, check_start, - check_end, cache=True, - **kwargs) + validsegs = dq.query_flag( + detector, segment_name, check_start, check_end, cache=True, **kwargs + ) use_start = gps_start use_end = gps_end # shift if necessary @@ -215,15 +278,22 @@ def check_validtimes(detector, gps_start, gps_end, shift_to_valid=False, shiftsize += 1 # check that we have a valid range if (use_start, use_end) not in validsegs: - raise NoValidDataError("Could not find a continous valid segment in " - "in detector {}".format(detector)) + raise NoValidDataError( + f"Could not find a continous valid segment in in detector {detector}" + ) return use_start, use_end -def detectors_with_valid_data(detectors, gps_start_times, gps_end_times, - pad_data=None, err_on_missing_detectors=False, - **kwargs): - r"""Determines which detectors have valid data. +def detectors_with_valid_data( + detectors, + gps_start_times, + gps_end_times, + pad_data=None, + err_on_missing_detectors=False, + **kwargs, +): + r""" + Determines which detectors have valid data. Parameters ---------- @@ -257,32 +327,37 @@ def detectors_with_valid_data(detectors, gps_start_times, gps_end_times, no valid times could be found for a detector (and ``err_on_missing_detectors`` is False), it will not be included in the returned dictionary. + """ if pad_data is None: - pad_data = {det: 0 for det in detectors} + pad_data = dict.fromkeys(detectors, 0) dets_with_data = {} for det in detectors: - logging.info("Checking that %s has valid data in the requested " - "segment", det) + logging.info("Checking that %s has valid data in the requested segment", det) try: pad = pad_data[det] - start, end = check_validtimes(det, gps_start_times[det]-pad, - gps_end_times[det]+pad, - **kwargs) - dets_with_data[det] = (start+pad, end-pad) + start, end = check_validtimes( + det, gps_start_times[det] - pad, gps_end_times[det] + pad, **kwargs + ) + dets_with_data[det] = (start + pad, end - pad) except NoValidDataError as e: if err_on_missing_detectors: raise e - logging.warning("WARNING: Detector %s will not be used in " - "the analysis, as it does not have " - "continuous valid data that spans the " - "segment [%d, %d).", det, gps_start_times[det]-pad, - gps_end_times[det]+pad) + logging.warning( + "WARNING: Detector %s will not be used in " + "the analysis, as it does not have " + "continuous valid data that spans the " + "segment [%d, %d).", + det, + gps_start_times[det] - pad, + gps_end_times[det] + pad, + ) return dets_with_data def check_for_nans(strain_dict): - """Checks if any data in a dictionary of strains has NaNs. + """ + Checks if any data in a dictionary of strains has NaNs. If any NaNs are found, a ``ValueError`` is raised. @@ -291,14 +366,16 @@ def check_for_nans(strain_dict): strain_dict : dict Dictionary of detectors -> :py:class:`pycbc.types.timeseries.TimeSeries`. + """ for det, ts in strain_dict.items(): if numpy.isnan(ts.numpy()).any(): - raise ValueError("NaN found in strain from {}".format(det)) + raise ValueError(f"NaN found in strain from {det}") def data_opts_from_config(cp, section, filter_flow): - """Loads data options from a section in a config file. + """ + Loads data options from a section in a config file. Parameters ---------- @@ -319,6 +396,7 @@ def data_opts_from_config(cp, section, filter_flow): opts : parsed argparse.ArgumentParser An argument parser namespace that was constructed as if the options were specified on the command line. + """ # convert the section options into a command-line options optstr = cp.section_to_cli(section) @@ -336,10 +414,13 @@ def data_opts_from_config(cp, section, filter_flow): gps_end[det] += opts.trigger_time if opts.psd_inverse_length[det] is not None: pad = int(numpy.ceil(opts.psd_inverse_length[det] / 2)) - logging.info("Padding %s analysis start and end times by %d " - "(= psd-inverse-length/2) seconds to " - "account for PSD wrap around effects.", - det, pad) + logging.info( + "Padding %s analysis start and end times by %d " + "(= psd-inverse-length/2) seconds to " + "account for PSD wrap around effects.", + det, + pad, + ) else: pad = 0 gps_start[det] -= pad @@ -354,15 +435,21 @@ def data_opts_from_config(cp, section, filter_flow): low_freq_cutoff = filter_flow.copy() if opts.low_frequency_cutoff: # add in any missing detectors - low_freq_cutoff.update({det: opts.low_frequency_cutoff[det] - for det in opts.instruments - if opts.low_frequency_cutoff[det] is not None}) + low_freq_cutoff.update( + { + det: opts.low_frequency_cutoff[det] + for det in opts.instruments + if opts.low_frequency_cutoff[det] is not None + } + ) # make sure the data conditioning low frequency cutoff is < than # the matched filter cutoff if any(low_freq_cutoff[det] > filter_flow[det] for det in filter_flow): - raise ValueError("data conditioning low frequency cutoff must " - "be less than the filter low frequency " - "cutoff") + raise ValueError( + "data conditioning low frequency cutoff must " + "be less than the filter low frequency " + "cutoff" + ) opts.low_frequency_cutoff = low_freq_cutoff # verify options are sane @@ -371,10 +458,14 @@ def data_opts_from_config(cp, section, filter_flow): return opts -def data_from_cli(opts, check_for_valid_times=False, - shift_psd_times_to_valid=False, - err_on_missing_detectors=False): - """Loads the data needed for a model from the given command-line options. +def data_from_cli( + opts, + check_for_valid_times=False, + shift_psd_times_to_valid=False, + err_on_missing_detectors=False, +): + """ + Loads the data needed for a model from the given command-line options. Gates specifed on the command line are also applied. @@ -401,6 +492,7 @@ def data_from_cli(opts, check_for_valid_times=False, If ``opts.psd_(start|end)_time`` were set, a dctionary of detectors -> time series data to use for PSD estimation. Otherwise, ``None``. + """ # get gates to apply gates = gates_from_cli(opts) @@ -412,17 +504,21 @@ def data_from_cli(opts, check_for_valid_times=False, # validate times if check_for_valid_times: dets_with_data = detectors_with_valid_data( - instruments, opts.gps_start_time, opts.gps_end_time, + instruments, + opts.gps_start_time, + opts.gps_end_time, pad_data=opts.pad_data, err_on_missing_detectors=err_on_missing_detectors, shift_to_valid=False, - segment_name=opts.dq_segment_name, source=opts.dq_source, - server=opts.dq_server, veto_definer=opts.veto_definer) + segment_name=opts.dq_segment_name, + source=opts.dq_source, + server=opts.dq_server, + veto_definer=opts.veto_definer, + ) # reset instruments to only be those with valid data instruments = list(dets_with_data.keys()) - strain_dict = strain_from_cli_multi_ifos(opts, instruments, - precision="double") + strain_dict = strain_from_cli_multi_ifos(opts, instruments, precision="double") # apply gates if not waiting to overwhiten if not opts.gate_overwhitened: strain_dict = apply_gates_to_td(strain_dict, gates) @@ -433,32 +529,38 @@ def data_from_cli(opts, check_for_valid_times=False, # get strain time series to use for PSD estimation # if user has not given the PSD time options then use same data as analysis if opts.psd_start_time and opts.psd_end_time: - logging.info("Will generate a different time series for PSD " - "estimation") + logging.info("Will generate a different time series for PSD estimation") if check_for_valid_times: psd_times = detectors_with_valid_data( - instruments, opts.psd_start_time, opts.psd_end_time, + instruments, + opts.psd_start_time, + opts.psd_end_time, pad_data=opts.pad_data, err_on_missing_detectors=err_on_missing_detectors, shift_to_valid=shift_psd_times_to_valid, - segment_name=opts.dq_segment_name, source=opts.dq_source, - server=opts.dq_server, veto_definer=opts.veto_definer) + segment_name=opts.dq_segment_name, + source=opts.dq_source, + server=opts.dq_server, + veto_definer=opts.veto_definer, + ) # remove detectors from the strain dict that did not have valid # times for PSD estimation - for det in set(strain_dict.keys())-set(psd_times.keys()): + for det in set(strain_dict.keys()) - set(psd_times.keys()): _ = strain_dict.pop(det) # reset instruments to only be those with valid data instruments = list(psd_times.keys()) else: - psd_times = {det: (opts.psd_start_time[det], - opts.psd_end_time[det]) - for det in instruments} + psd_times = { + det: (opts.psd_start_time[det], opts.psd_end_time[det]) + for det in instruments + } psd_strain_dict = {} for det, (psd_start, psd_end) in psd_times.items(): opts.gps_start_time = psd_start opts.gps_end_time = psd_end psd_strain_dict.update( - strain_from_cli_multi_ifos(opts, [det], precision="double")) + strain_from_cli_multi_ifos(opts, [det], precision="double") + ) # apply any gates logging.info("Applying gates to PSD data") psd_strain_dict = apply_gates_to_td(psd_strain_dict, psd_gates) @@ -471,14 +573,16 @@ def data_from_cli(opts, check_for_valid_times=False, # check that we have data left to analyze if instruments == []: - raise NoValidDataError("No valid data could be found in any of the " - "requested instruments.") + raise NoValidDataError( + "No valid data could be found in any of the requested instruments." + ) return strain_dict, psd_strain_dict def fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict=None): - """Converts a dictionary of time series to the frequency domain, and gets + """ + Converts a dictionary of time series to the frequency domain, and gets the PSDs. Parameters @@ -500,6 +604,7 @@ def fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict=None): Dictionary of detectors -> frequency series data. psd_dict : dict Dictionary of detectors -> frequency-domain PSDs. + """ # FFT strain and save each of the length of the FFT, delta_f, and # low frequency cutoff to a dict @@ -522,15 +627,21 @@ def fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict=None): # get PSD as frequency series psd_dict = psd_from_cli_multi_ifos( - opts, length_dict, delta_f_dict, lfs, - list(psd_strain_dict.keys()), strain_dict=psd_strain_dict, - precision="double") + opts, + length_dict, + delta_f_dict, + lfs, + list(psd_strain_dict.keys()), + strain_dict=psd_strain_dict, + precision="double", + ) return stilde_dict, psd_dict def gate_overwhitened_data(stilde_dict, psd_dict, gates): - """Applies gates to overwhitened data. + """ + Applies gates to overwhitened data. Parameters ---------- @@ -546,6 +657,7 @@ def gate_overwhitened_data(stilde_dict, psd_dict, gates): dict : Dictionary of detectors -> frequency series data with the gates applied after overwhitening. The returned data is not overwhitened. + """ logging.info("Applying gates to overwhitened data") # overwhiten the data diff --git a/pycbc/inference/models/gated_gaussian_noise.py b/pycbc/inference/models/gated_gaussian_noise.py index f7306041096..18d35d9039e 100644 --- a/pycbc/inference/models/gated_gaussian_noise.py +++ b/pycbc/inference/models/gated_gaussian_noise.py @@ -13,41 +13,58 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides model classes that assume the noise is Gaussian and +""" +This module provides model classes that assume the noise is Gaussian and introduces a gate to remove given times from the data, using the inpainting method to fill the removed part such that it does not enter the likelihood. """ -from abc import abstractmethod import logging +import warnings +from abc import abstractmethod + import numpy import scipy from scipy import special -import warnings -from pycbc.types import FrequencySeries -from pycbc.detector import Detector -from pycbc.pnutils import hybrid_meco_frequency from pycbc import types -from pycbc.waveform.utils import time_from_frequencyseries -from pycbc.waveform import generator +from pycbc.detector import Detector from pycbc.filter import highpass +from pycbc.pnutils import hybrid_meco_frequency from pycbc.strain.gate import invert_covariance -from .gaussian_noise import (BaseGaussianNoise, create_waveform_generator, - catch_waveform_error) +from pycbc.types import FrequencySeries +from pycbc.waveform import generator +from pycbc.waveform.utils import time_from_frequencyseries + from .base_data import BaseDataModel from .data_utils import fd_data_from_strain_dict +from .gaussian_noise import ( + BaseGaussianNoise, + catch_waveform_error, + create_waveform_generator, +) class BaseGatedGaussian(BaseGaussianNoise): - r"""Base model for gated gaussian. + r""" + Base model for gated gaussian. Provides additional routines for applying a time-domain gate to data. See :py:class:`GatedGaussianNoise` for more details. """ - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, highpass_waveforms=False, **kwargs): + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + highpass_waveforms=False, + **kwargs, + ): # we'll want the time-domain data, so store that self._td_data = {} # cache the overwhitened data @@ -61,13 +78,12 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, self._gatetimes = {} self._det_lognls = {} # cache condition number calculations - self.check_condition_number = bool(kwargs.get('check-condition-number', - False)) + self.check_condition_number = bool(kwargs.get("check-condition-number", False)) self._cond = {} # cache inpainting options - paint_method = kwargs.get('paint_method') + paint_method = kwargs.get("paint_method") if paint_method is None: - paint_method = kwargs.get('paint-method', 'toeplitz') + paint_method = kwargs.get("paint-method", "toeplitz") self.paint_method = paint_method self._cov_matrices = {} # cache samples and linear regression for determinant extrapolation @@ -76,18 +92,23 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, # highpass waveforms with the given frequency self.highpass_waveforms = highpass_waveforms if self.highpass_waveforms: - logging.info("Will highpass waveforms at %f Hz", - highpass_waveforms) + logging.info("Will highpass waveforms at %f Hz", highpass_waveforms) # set up the boiler-plate attributes super().__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) @classmethod - def from_config(cls, cp, data_section='data', data=None, psds=None, - **kwargs): - """Adds additional keyword arguments based on config file. + def from_config(cls, cp, data_section="data", data=None, psds=None, **kwargs): + """ + Adds additional keyword arguments based on config file. Additional keyword arguments are: @@ -96,15 +117,18 @@ def from_config(cls, cp, data_section='data', data=None, psds=None, Also forces ``invpsd-trunc-low-freq-fill-value`` to ``fmin`` if not specified. """ - if cp.has_option(data_section, 'strain-high-pass') and \ - 'highpass_waveforms' not in kwargs: - kwargs['highpass_waveforms'] = float(cp.get(data_section, - 'strain-high-pass')) - if not cp.has_option(data_section, 'invpsd-trunc-low-freq-fill-value'): - cp.set(data_section, 'invpsd-trunc-low-freq-fill-value', 'fmin') - return super().from_config(cp, data_section=data_section, - data=data, psds=psds, - **kwargs) + if ( + cp.has_option(data_section, "strain-high-pass") + and "highpass_waveforms" not in kwargs + ): + kwargs["highpass_waveforms"] = float( + cp.get(data_section, "strain-high-pass") + ) + if not cp.has_option(data_section, "invpsd-trunc-low-freq-fill-value"): + cp.set(data_section, "invpsd-trunc-low-freq-fill-value", "fmin") + return super().from_config( + cp, data_section=data_section, data=data, psds=psds, **kwargs + ) @BaseDataModel.data.setter def data(self, data): @@ -120,7 +144,8 @@ def td_data(self): @BaseGaussianNoise.psds.setter def psds(self, psds): - """Sets the psds, and calculates the weight and norm from them. + """ + Sets the psds, and calculates the weight and norm from them. The data and the low and high frequency cutoffs must be set first. """ # check that the data has been set @@ -142,14 +167,13 @@ def psds(self, psds): for det, d in self._data.items(): if psds is None: # No psd means assume white PSD - p = FrequencySeries(numpy.ones(int(self._N/2+1)), - delta_f=d.delta_f) + p = FrequencySeries(numpy.ones(int(self._N / 2 + 1)), delta_f=d.delta_f) else: # copy for storage p = psds[det].copy() self._psds[det] = p # we'll store the weight to apply to the inner product - invp = 1./p + invp = 1.0 / p self._invpsds[det] = invp self._invasds[det] = invp**0.5 # store the autocorrelation function and covariance matrix for @@ -163,7 +187,8 @@ def psds(self, psds): self._overwhitened_data = self.whiten(self.data, 2, inplace=False) def _set_covfit(self, det): - """Sets the fit function for estimating the covariance determinant. + """ + Sets the fit function for estimating the covariance determinant. This must be called after the PSDs have been set, otherwise a ValueError will be raised. @@ -173,14 +198,14 @@ def _set_covfit(self, det): except KeyError: raise ValueError("No psd set for detector %s" % det) Rss = self._Rss[det] - cov = scipy.linalg.toeplitz(Rss/2) # full covariance matrix + cov = scipy.linalg.toeplitz(Rss / 2) # full covariance matrix samples, fit = self.logdet_fit(cov, p) self._cov_samples[det] = samples self._cov_regressions[det] = fit - return def logdet_fit(self, cov, p): - """Construct a linear regression from a sample of truncated covariance + """ + Construct a linear regression from a sample of truncated covariance matrices. Returns the sample points used for linear fit generation as well as the @@ -194,22 +219,23 @@ def logdet_fit(self, cov, p): s = cov.shape[0] max_size = 8192 if s > max_size: - sample_sizes = [s, max_size, max_size//2, max_size//4] + sample_sizes = [s, max_size, max_size // 2, max_size // 4] else: - sample_sizes = [s, s//2, s//4, s//8] + sample_sizes = [s, s // 2, s // 4, s // 8] for i in sample_sizes: # calculate logdet of the full matrix using circulant eigenvalue # approximation if i == s: - ld = 2*numpy.log(p/(2*p.delta_t)).sum() + ld = 2 * numpy.log(p / (2 * p.delta_t)).sum() sample_dets.append(ld) # generate three more sample matrices using exact calculations else: gate_size = s - i - start = (s - gate_size)//2 + start = (s - gate_size) // 2 end = start + gate_size - tc = numpy.delete(numpy.delete(cov, slice(start, end), 0), - slice(start, end), 1) + tc = numpy.delete( + numpy.delete(cov, slice(start, end), 0), slice(start, end), 1 + ) ld = numpy.linalg.slogdet(tc)[1] sample_dets.append(ld) # generate a linear regression using the four points (size, logdet) @@ -219,7 +245,8 @@ def logdet_fit(self, cov, p): @BaseGaussianNoise.normalize.setter def normalize(self, normalize): - """Clears the current stats if the normalization state is changed. + """ + Clears the current stats if the normalization state is changed. If normalize is set to True, the fit to the covariance determinant will be calculated if it hasn't yet and PSDs are set. @@ -234,8 +261,7 @@ def normalize(self, normalize): self._set_covfit(det) def gate_indices(self, det): - """Calculate the indices corresponding to start and end of gate. - """ + """Calculate the indices corresponding to start and end of gate.""" # get time series start and delta_t ts = self.td_data[det] # get gate start and length from get_gate_times @@ -247,7 +273,8 @@ def gate_indices(self, det): return lindex, rindex def det_lognorm(self, det, start_index=None, end_index=None): - """Calculate the normalization term from the truncated covariance + """ + Calculate the normalization term from the truncated covariance matrix. Determinant is estimated using a linear fit to logdet vs truncated @@ -265,29 +292,32 @@ def det_lognorm(self, det, start_index=None, end_index=None): # call the linear regression m, b = self._cov_regressions[det] # extrapolate from linear fit - ld = m*trunc_size + b + ld = m * trunc_size + b # full normalization term: - lognorm = -0.5*(numpy.log(2*numpy.pi)*trunc_size + ld) + lognorm = -0.5 * (numpy.log(2 * numpy.pi) * trunc_size + ld) # cache the result self._lognorm[(det, start_index, end_index)] = lognorm return lognorm def _nowaveform_handler(self): - """Convenience function to set logl values if no waveform generated. - """ + """Convenience function to set logl values if no waveform generated.""" return -numpy.inf def _loglr(self): - r"""Computes the log likelihood ratio. + r""" + Computes the log likelihood ratio. + Returns ------- float The value of the log likelihood ratio evaluated at the given point. + """ return self._loglikelihood() - self._lognl() def whiten(self, data, whiten, inplace=False): - """Whitens the given data. + """ + Whitens the given data. Parameters ---------- @@ -306,6 +336,7 @@ def whiten(self, data, whiten, inplace=False): dict : Dictionary of FrequencySeries after the requested whitening has been applied. + """ if not inplace: data = {det: d.copy() for det, d in data.items()} @@ -321,18 +352,19 @@ def whiten(self, data, whiten, inplace=False): return data def invert_covariance(self, det): - """Get the uninverted covariance matrix for the model's inverse PSDs. + """ + Get the uninverted covariance matrix for the model's inverse PSDs. Once the inverse matrix is calculated for a given gate time in this detector, store to cache; future calls of this function will pull from that cache instead. """ # don't bother with covariance matrix if we're using toeplitz solver - if self.paint_method == 'toeplitz': + if self.paint_method == "toeplitz": return None # check if there are cache results for this gate length lindex, rindex = self.gate_indices(det) try: - cov_matrices = self._cov_matrices[int(rindex-lindex)] + cov_matrices = self._cov_matrices[int(rindex - lindex)] except KeyError: cov_matrices = {} # check if this det has a precalculated matrix for this gate length @@ -344,12 +376,13 @@ def invert_covariance(self, det): invmat = invert_covariance(invpsd, lindex, rindex) cov_matrices[det] = invmat # cache results - self._cov_matrices[int(rindex-lindex)] = cov_matrices + self._cov_matrices[int(rindex - lindex)] = cov_matrices return invmat @abstractmethod def get_waveforms(self): - """The waveforms generated using the current parameters. + """ + The waveforms generated using the current parameters. If the waveforms haven't been generated yet, they will be generated, resized to the same length as the data, and cached. If the @@ -360,32 +393,36 @@ def get_waveforms(self): ------- dict : Dictionary of detector names -> waveforms + """ - pass @abstractmethod def get_gated_waveforms(self): - """Generates and gates waveforms using the current parameters. + """ + Generates and gates waveforms using the current parameters. Returns ------- dict : Dictionary of detector names -> FrequencySeries. + """ - pass def get_data(self): - """Return a copy of the data. + """ + Return a copy of the data. Returns ------- dict : Dictionary of detector names -> FrequencySeries. + """ return {det: d.copy() for det, d in self.data.items()} def get_gated_data(self): - """Return a copy of the gated data. + """ + Return a copy of the gated data. The gated data will be cached for faster retrieval. @@ -393,6 +430,7 @@ def get_gated_data(self): ------- dict : Dictionary of detector names -> FrequencySeries. + """ gate_times = self.get_gate_times() out = {} @@ -410,11 +448,15 @@ def get_gated_data(self): # doesn't exist yet, or the gate times changed cache.clear() invmat = self.invert_covariance(det) - d = d.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=True, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) + d = d.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=True, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) dtilde = d.to_frequencyseries() # save for next time cache[gatestartdelay, dgatedelay] = dtilde @@ -422,7 +464,8 @@ def get_gated_data(self): return out def get_gate_times(self): - """Gets the time to apply a gate based on the current sky position. + """ + Gets the time to apply a gate based on the current sky position. If the parameter ``gatefunc`` is set to ``'hmeco'``, the gate times will be calculated based on the hybrid MECO of the given set of @@ -438,19 +481,20 @@ def get_gate_times(self): ------- dict : Dictionary of detector names -> (gate start, gate width) + """ params = self.current_params try: - gatefunc = self.current_params['gatefunc'] + gatefunc = self.current_params["gatefunc"] except KeyError: gatefunc = None - if gatefunc == 'hmeco': + if gatefunc == "hmeco": return self.get_gate_times_hmeco() - gatestart = params['t_gate_start'] - gateend = params['t_gate_end'] + gatestart = params["t_gate_start"] + gateend = params["t_gate_end"] # we'll need the sky location for determining time shifts - ra = self.current_params['ra'] - dec = self.current_params['dec'] + ra = self.current_params["ra"] + dec = self.current_params["dec"] # try to get from cache try: gatetimes = self._gatetimes[gatestart, gateend, ra, dec] @@ -462,13 +506,15 @@ def get_gate_times(self): # check if the inpainting is numerically stable if self.check_condition_number: for det, d in self.td_data.items(): - lindex, rindex = d.get_gate_indices(gatetimes[det][0], - gatetimes[det][1]/2.) + lindex, rindex = d.get_gate_indices( + gatetimes[det][0], gatetimes[det][1] / 2.0 + ) self.condition_number(det, lindex, rindex) return gatetimes def _get_gate_times(self, gatestart, gateend, ra, dec): - """Calculates the gate times in each detector. + """ + Calculates the gate times in each detector. Parameters ---------- @@ -485,13 +531,14 @@ def _get_gate_times(self, gatestart, gateend, ra, dec): ------- dict : Dictionary of detector names -> (gate start, gate width) + """ gatetimes = {} for det in self._invpsds: thisdet = Detector(det) # account for the time delay between the waveforms of the # different detectors - refdet = self.current_params.get('tc_ref_frame', 'geocentric') + refdet = self.current_params.get("tc_ref_frame", "geocentric") gatestartdelay = thisdet.arrival_time(gatestart, ra, dec, refdet) gateenddelay = thisdet.arrival_time(gateend, ra, dec, refdet) dgatedelay = gateenddelay - gatestartdelay @@ -499,46 +546,49 @@ def _get_gate_times(self, gatestart, gateend, ra, dec): return gatetimes def get_gate_times_hmeco(self): - """Gets the time to apply a gate based on the current sky position. + """ + Gets the time to apply a gate based on the current sky position. Returns ------- dict : Dictionary of detector names -> (gate start, gate width) + """ # generate the template waveform wfs = self.get_waveforms() # get waveform parameters params = self.current_params - spin1 = params['spin1z'] - spin2 = params['spin2z'] + spin1 = params["spin1z"] + spin2 = params["spin2z"] # gate input for ringdown analysis which consideres a start time # and an end time - dgate = params['gate_window'] - meco_f = hybrid_meco_frequency(params['mass1'], params['mass2'], spin1, - spin2) + dgate = params["gate_window"] + meco_f = hybrid_meco_frequency(params["mass1"], params["mass2"], spin1, spin2) # figure out the gate times gatetimes = {} for det, h in wfs.items(): invpsd = self._invpsds[det] h.resize(len(invpsd)) ht = h.to_timeseries() - f_low = int((self._f_lower[det]+1)/h.delta_f) + f_low = int((self._f_lower[det] + 1) / h.delta_f) sample_freqs = h.sample_frequencies[f_low:].numpy() f_idx = numpy.where(sample_freqs <= meco_f)[0][-1] # find time corresponding to meco frequency t_from_freq = time_from_frequencyseries( - h[f_low:], sample_frequencies=sample_freqs) + h[f_low:], sample_frequencies=sample_freqs + ) if t_from_freq[f_idx] > 0: gatestartdelay = t_from_freq[f_idx] + float(t_from_freq.epoch) else: gatestartdelay = t_from_freq[f_idx] + ht.sample_times[-1] - gatestartdelay = min(gatestartdelay, params['t_gate_start']) + gatestartdelay = min(gatestartdelay, params["t_gate_start"]) gatetimes[det] = (gatestartdelay, dgate) return gatetimes def condition_number(self, det, lindex, rindex): - """Calculate the condition number associated with the inverse + """ + Calculate the condition number associated with the inverse covariance matrix used to gate and inpaint. Throws a warning if the condition number is greater than 1e16. @@ -556,6 +606,7 @@ def condition_number(self, det, lindex, rindex): float : The condition number of the inverse covariance matrix constructed from the inverse PSD with the given gate length. + """ gate_idx_len = int(rindex - lindex) if gate_idx_len not in self._cond.keys(): @@ -565,8 +616,8 @@ def condition_number(self, det, lindex, rindex): if det not in conds.keys(): # construct the matrix invpsd = self._invpsds[det] - tdfilter = invpsd.astype('complex').to_timeseries() * invpsd.delta_t - mat = scipy.linalg.toeplitz(tdfilter[:rindex-lindex]) + tdfilter = invpsd.astype("complex").to_timeseries() * invpsd.delta_t + mat = scipy.linalg.toeplitz(tdfilter[: rindex - lindex]) rcond = numpy.linalg.cond(mat) # cache the value conds[det] = rcond @@ -575,15 +626,16 @@ def condition_number(self, det, lindex, rindex): # pull from cache rcond = self._cond[gate_idx_len][det] if rcond >= 1e16: - warnings.warn(f'Condition number of inverse covariance matrix is ' - f'{rcond}; inpainting may be numerically unstable') + warnings.warn( + f"Condition number of inverse covariance matrix is " + f"{rcond}; inpainting may be numerically unstable" + ) return rcond def _lognl(self): - """Calculates the log of the noise likelihood. - """ + """Calculates the log of the noise likelihood.""" # clear variables - lognl = 0. + lognl = 0.0 self._det_lognls.clear() # get the times of the gates gate_times = self.get_gate_times() @@ -598,11 +650,15 @@ def _lognl(self): # gate the data data = self.td_data[det] invmat = self.invert_covariance(det) - gated_dt = data.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=True, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) + gated_dt = data.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=True, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) # convert to the frequency series gated_d = gated_dt.to_frequencyseries() # overwhiten @@ -610,7 +666,7 @@ def _lognl(self): d = self.data[det] # inner product ip = 4 * invpsd.delta_f * d[slc].inner(gated_d[slc]).real # - dd = norm - 0.5*ip + dd = norm - 0.5 * ip # store self._det_lognls[det] = dd lognl += dd @@ -618,13 +674,14 @@ def _lognl(self): def det_lognl(self, det): # make sure lognl has been called - _ = self._trytoget('lognl', self._lognl) + _ = self._trytoget("lognl", self._lognl) # the det_lognls dict should have been updated, so can call it now return self._det_lognls[det] @staticmethod def _fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict): - """Wrapper around :py:func:`data_utils.fd_data_from_strain_dict`. + """ + Wrapper around :py:func:`data_utils.fd_data_from_strain_dict`. Ensures that if the PSD is estimated from data, the inverse spectrum truncation uses a Hann window. Sets the low frequency cutoff for the @@ -633,21 +690,21 @@ def _fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict): if opts.psd_inverse_length and opts.invpsd_trunc_method is None: # make sure invpsd truncation is set to hanning logging.info("Using Hann window to truncate inverse PSD") - opts.invpsd_trunc_method = 'hann' + opts.invpsd_trunc_method = "hann" # set low frequency cutoff for PSDs if opts.psd_low_frequency_cutoff is None: opts.psd_low_frequency_cutoff = {} for d, lfs in opts.low_frequency_cutoff.items(): if d not in opts.psd_low_frequency_cutoff: # set to half the model's likelihood cutoffs - logging.info(f"Setting low frequency cutoff of {d} PSD to " - f"{lfs/2.}") - opts.psd_low_frequency_cutoff[d] = lfs/2. + logging.info(f"Setting low frequency cutoff of {d} PSD to {lfs / 2.0}") + opts.psd_low_frequency_cutoff[d] = lfs / 2.0 out = fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict) return out def write_metadata(self, fp, group=None): - """Adds writing the psds, and analyzed detectors. + """ + Adds writing the psds, and analyzed detectors. The analyzed detectors, their analysis segments, and the segments used for psd estimation are written as @@ -664,36 +721,36 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ BaseDataModel.write_metadata(self, fp, group=group) attrs = fp.getattrs(group=group) # write the analyzed detectors and times - attrs['analyzed_detectors'] = self.detectors + attrs["analyzed_detectors"] = self.detectors # store fitting values here for det, data in self.data.items(): - key = '{}_analysis_segment'.format(det) + key = f"{det}_analysis_segment" attrs[key] = [float(data.start_time), float(data.end_time)] # store covariance determinant extrapolation info (checkpoint) if self.normalize: - attrs['{}_cov_sample'.format(det)] = self._cov_samples[det] - attrs['{}_cov_regression'.format(det)] = \ - self._cov_regressions[det] + attrs[f"{det}_cov_sample"] = self._cov_samples[det] + attrs[f"{det}_cov_regression"] = self._cov_regressions[det] if self._psds is not None and not self.no_save_data: fp.write_psd(self._psds, group=group) # write the times used for psd estimation (if they were provided) for det in self.psd_segments: - key = '{}_psd_segment'.format(det) + key = f"{det}_psd_segment" attrs[key] = list(map(float, self.psd_segments[det])) # save the frequency cutoffs for det in self.detectors: - attrs['{}_likelihood_low_freq'.format(det)] = self._f_lower[det] + attrs[f"{det}_likelihood_low_freq"] = self._f_lower[det] if self._f_upper[det] is not None: - attrs['{}_likelihood_high_freq'.format(det)] = \ - self._f_upper[det] + attrs[f"{det}_likelihood_high_freq"] = self._f_upper[det] class GatedGaussianNoise(BaseGatedGaussian): - r"""Model that applies a time domain gate, assuming stationary Gaussian + r""" + Model that applies a time domain gate, assuming stationary Gaussian noise. The gate start and end times are set by providing ``t_gate_start`` and @@ -710,22 +767,40 @@ class GatedGaussianNoise(BaseGatedGaussian): use this model for fixed gate times. """ - name = 'gated_gaussian_noise' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, **kwargs): + name = "gated_gaussian_noise" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + **kwargs, + ): # set up the boiler-plate attributes super().__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) # create the waveform generator self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, - gates=self.gates, **self.static_params) + gates=self.gates, + **self.static_params, + ) @property def _extra_stats(self): @@ -734,7 +809,8 @@ def _extra_stats(self): @catch_waveform_error def _loglikelihood(self): - r"""Computes the log likelihood after removing the power within the + r""" + Computes the log likelihood after removing the power within the given time window, .. math:: @@ -747,12 +823,13 @@ def _loglikelihood(self): ------- float The value of the log likelihood. + """ # generate the template waveform wfs = self.get_waveforms() # get the times of the gates gate_times = self.get_gate_times() - logl = 0. + logl = 0.0 for det, h in wfs.items(): invpsd = self._invpsds[det] start_index, end_index = self.gate_indices(det) @@ -767,45 +844,52 @@ def _loglikelihood(self): res = data - ht rtilde = res.to_frequencyseries() invmat = self.invert_covariance(det) - gated_res = res.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=True, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) + gated_res = res.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=True, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) gated_rtilde = gated_res.to_frequencyseries() # overwhiten gated_rtilde *= invpsd rr = 4 * invpsd.delta_f * rtilde[slc].inner(gated_rtilde[slc]).real - logl += norm - 0.5*rr + logl += norm - 0.5 * rr return float(logl) @property def _extra_stats(self): - """Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each - detector.""" - return ['loglr', 'maxl_phase'] + ['{}_optimal_snrsq'.format(det) for det in self._data] + """ + Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each + detector. + """ + return ["loglr", "maxl_phase"] + [ + f"{det}_optimal_snrsq" for det in self._data + ] def _nowaveform_loglr(self): - """Convenience function to set loglr values if no waveform generated. - """ - setattr(self._current_stats, 'loglikelihood', -numpy.inf) + """Convenience function to set loglr values if no waveform generated.""" + self._current_stats.loglikelihood = -numpy.inf # maxl phase doesn't exist, so set it to nan - setattr(self._current_stats, 'maxl_phase', numpy.nan) + self._current_stats.maxl_phase = numpy.nan for det in self._data: # snr can't be < 0 by definition, so return 0 - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), 0.) + setattr(self._current_stats, f"{det}_optimal_snrsq", 0.0) return -numpy.inf @property def multi_signal_support(self): - """ The list of classes that this model supports in a multi-signal + """ + The list of classes that this model supports in a multi-signal likelihood """ return [type(self)] def multi_loglikelihood(self, models): - """ Calculate a multi-model (signal) likelihood - """ + """Calculate a multi-model (signal) likelihood""" # Generate the waveforms for each submodel wfs = [] for m in models + [self]: @@ -833,8 +917,8 @@ def get_waveforms(self): # apply high pass if self.highpass_waveforms: h = highpass( - h.to_timeseries(), - frequency=self.highpass_waveforms).to_frequencyseries() + h.to_timeseries(), frequency=self.highpass_waveforms + ).to_frequencyseries() wfs[det] = h self._current_wfs = wfs return self._current_wfs @@ -849,18 +933,23 @@ def get_gated_waveforms(self): gate_times = self.get_gate_times() gatestartdelay, dgatedelay = gate_times[det] invmat = self.invert_covariance(det) - ht = ht.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=False, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) + ht = ht.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=False, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) h = ht.to_frequencyseries() out[det] = h return out class GatedGaussianMargPol(BaseGatedGaussian): - r"""Gated gaussian model with numerical marginalization over polarization. + r""" + Gated gaussian model with numerical marginalization over polarization. This implements the GatedGaussian likelihood with an explicit numerical marginalization over polarization angle. This is accomplished using @@ -869,28 +958,45 @@ class GatedGaussianMargPol(BaseGatedGaussian): The 'polarization_samples' argument can be passed to set an alternate number of integration points. """ - name = 'gated_gaussian_margpol' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, - polarization_samples=1000, **kwargs): + name = "gated_gaussian_margpol" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + polarization_samples=1000, + **kwargs, + ): # set up the boiler-plate attributes super().__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) # the polarization parameters self.polarization_samples = polarization_samples - self.pol = numpy.linspace(0, 2*numpy.pi, self.polarization_samples) + self.pol = numpy.linspace(0, 2 * numpy.pi, self.polarization_samples) self.dets = {} # create the waveform generator self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameTwoPolGenerator, - **self.static_params) + **self.static_params, + ) def get_waveforms(self): if self._current_wfs is not None: @@ -904,11 +1010,11 @@ def get_waveforms(self): # apply high pass if self.highpass_waveforms: hp = highpass( - hp.to_timeseries(), - frequency=self.highpass_waveforms).to_frequencyseries() + hp.to_timeseries(), frequency=self.highpass_waveforms + ).to_frequencyseries() hc = highpass( - hc.to_timeseries(), - frequency=self.highpass_waveforms).to_frequencyseries() + hc.to_timeseries(), frequency=self.highpass_waveforms + ).to_frequencyseries() wfs[det] = (hp, hc) self._current_wfs = wfs return self._current_wfs @@ -925,34 +1031,40 @@ def get_gated_waveforms(self): for h in wfs[det]: ht = h.to_timeseries() invmat = self.invert_covariance(det) - ht = ht.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=False, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) + ht = ht.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=False, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) h = ht.to_frequencyseries() pols.append(h) out[det] = tuple(pols) return out def get_gate_times_hmeco(self): - """Gets the time to apply a gate based on the current sky position. + """ + Gets the time to apply a gate based on the current sky position. + Returns ------- dict : Dictionary of detector names -> (gate start, gate width) + """ # generate the template waveform wfs = self.get_waveforms() # get waveform parameters params = self.current_params - spin1 = params['spin1z'] - spin2 = params['spin2z'] + spin1 = params["spin1z"] + spin2 = params["spin2z"] # gate input for ringdown analysis which consideres a start time # and an end time - dgate = params['gate_window'] - meco_f = hybrid_meco_frequency(params['mass1'], params['mass2'], spin1, - spin2) + dgate = params["gate_window"] + meco_f = hybrid_meco_frequency(params["mass1"], params["mass2"], spin1, spin2) # figure out the gate times gatetimes = {} # for now only calculating time from plus polarization; should be all @@ -961,28 +1073,30 @@ def get_gate_times_hmeco(self): invpsd = self._invpsds[det] hp.resize(len(invpsd)) ht = hp.to_timeseries() - f_low = int((self._f_lower[det]+1)/hp.delta_f) + f_low = int((self._f_lower[det] + 1) / hp.delta_f) sample_freqs = hp.sample_frequencies[f_low:].numpy() f_idx = numpy.where(sample_freqs <= meco_f)[0][-1] # find time corresponding to meco frequency t_from_freq = time_from_frequencyseries( - hp[f_low:], sample_frequencies=sample_freqs) + hp[f_low:], sample_frequencies=sample_freqs + ) if t_from_freq[f_idx] > 0: gatestartdelay = t_from_freq[f_idx] + float(t_from_freq.epoch) else: gatestartdelay = t_from_freq[f_idx] + ht.sample_times[-1] - gatestartdelay = min(gatestartdelay, params['t_gate_start']) + gatestartdelay = min(gatestartdelay, params["t_gate_start"]) gatetimes[det] = (gatestartdelay, dgate) return gatetimes @property def _extra_stats(self): """Adds the maxL polarization and corresponding likelihood.""" - return ['maxl_polarization', 'maxl_logl'] + return ["maxl_polarization", "maxl_logl"] @catch_waveform_error def _loglikelihood(self): - r"""Computes the log likelihood after removing the power within the + r""" + Computes the log likelihood after removing the power within the given time window, .. math:: @@ -995,6 +1109,7 @@ def _loglikelihood(self): ------- float The value of the log likelihood. + """ # generate the template waveform wfs = self.get_waveforms() @@ -1002,12 +1117,12 @@ def _loglikelihood(self): gated_wfs = self.get_gated_waveforms() gated_data = self.get_gated_data() # cycle over - loglr = 0. - lognl = 0. - refframe = self.current_params.get('tc_ref_frame', 'geocentric') - ref_tc = self.current_params['tc'] - ra = self.current_params['ra'] - dec = self.current_params['dec'] + loglr = 0.0 + lognl = 0.0 + refframe = self.current_params.get("tc_ref_frame", "geocentric") + ref_tc = self.current_params["tc"] + ra = self.current_params["ra"] + dec = self.current_params["dec"] for det, (hp, hc) in wfs.items(): # get the antenna patterns if det not in self.dets: @@ -1029,8 +1144,8 @@ def _loglikelihood(self): d = self._overwhitened_data[det] # overwhiten the hp and hc invpsd = self._invpsds[det] - hp = hp*invpsd - hc = hc*invpsd + hp = hp * invpsd + hc = hc * invpsd # get the various gated inner products hpd = hp[slc].inner(gated_d[slc]).real # hcd = hc[slc].inner(gated_d[slc]).real # @@ -1044,34 +1159,34 @@ def _loglikelihood(self): # since the antenna patterns are real, # /2 + /2 = fp*(/2 + /2) # + fc*(/2 + /2) - hd = fp*(hpd + dhp) + fc*(hcd + dhc) + hd = fp * (hpd + dhp) + fc * (hcd + dhc) # /2 = /2 # = fp*fp*/2 + fc*fc*/2 # + fp*fc*/2 + fc*fp*/2 - hh = fp*fp*hphp + fc*fc*hchc + fp*fc*(hphc + hchp) + hh = fp * fp * hphp + fc * fc * hchc + fp * fc * (hphc + hchp) # sum up; note that the factor is 2df instead of 4df to account # for the factor of 1/2 - loglr += norm + 2*invpsd.delta_f*(hd - hh) + loglr += norm + 2 * invpsd.delta_f * (hd - hh) lognl += -2 * invpsd.delta_f * dd # store the maxl polarization idx = loglr.argmax() - setattr(self._current_stats, 'maxl_polarization', self.pol[idx]) - setattr(self._current_stats, 'maxl_logl', loglr[idx] + lognl) + self._current_stats.maxl_polarization = self.pol[idx] + self._current_stats.maxl_logl = loglr[idx] + lognl # compute the marginalized log likelihood marglogl = special.logsumexp(loglr) + lognl - numpy.log(len(self.pol)) return float(marglogl) @property def multi_signal_support(self): - """ The list of classes that this model supports in a multi-signal + """ + The list of classes that this model supports in a multi-signal likelihood """ return [type(self)] @catch_waveform_error def multi_loglikelihood(self, models): - """ Calculate a multi-model (signal) likelihood - """ + """Calculate a multi-model (signal) likelihood""" # Generate the waveforms for each submodel wfs = [] for m in models + [self]: @@ -1089,18 +1204,21 @@ def multi_loglikelihood(self, models): wfs[det][0] = wfs[det][0].copy().resize(mlen) wfs[det][1] = wfs[det][1].copy().resize(mlen) # combine waveforms - combine[det] = (sum([x[det][0] for x in wfs]), sum([x[det][1] - for x in wfs])) + combine[det] = ( + sum([x[det][0] for x in wfs]), + sum([x[det][1] for x in wfs]), + ) self._current_wfs = combine return self._loglikelihood() class GatedGaussianMargPhase(BaseGatedGaussian): - r"""Gated Gaussian noise model that analytically marginalizes over the + r""" + Gated Gaussian noise model that analytically marginalizes over the phase of a signal. - The phase to be marginalized over is specified by the user using the + The phase to be marginalized over is specified by the user using the `ref_phase` argument. If a model consists of multiple modes each with their own phase, only the reference phase is marginalized over. All phases must be specified with the `phase_names` argument. This can be passed as a list @@ -1116,56 +1234,78 @@ class GatedGaussianMargPhase(BaseGatedGaussian): respectively. The number of integration points can be controlled via the `phase_samples` argument. """ - name = 'gated_gaussian_margphase' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, - phase_samples=500000, phase_names=None, - ref_phase=None, **kwargs): + name = "gated_gaussian_margphase" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + phase_samples=500000, + phase_names=None, + ref_phase=None, + **kwargs, + ): # set up the boiler-plate attributes super().__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) self.det_names = list(self.data.keys()) self.dets = {} # phase marginalization parameters self.phase_samples = int(phase_samples) - self.phases = numpy.linspace(0, 2*numpy.pi, self.phase_samples) + self.phases = numpy.linspace(0, 2 * numpy.pi, self.phase_samples) if ref_phase is None: - raise KeyError('ref_phase is set to None. Please specify the ' - 'name of the phase parameter to marginalize ' - 'over') + raise KeyError( + "ref_phase is set to None. Please specify the " + "name of the phase parameter to marginalize " + "over" + ) self.ref_phase = ref_phase if phase_names is None: - logging.warning('No phase_names provided. Assuming single mode ' - f'specified by ref_phase {ref_phase}') + logging.warning( + "No phase_names provided. Assuming single mode " + f"specified by ref_phase {ref_phase}" + ) self.phase_names = [ref_phase] elif type(phase_names) == list: self.phase_names = phase_names elif type(phase_names) == str: - self.phase_names = phase_names.split(' ') + self.phase_names = phase_names.split(" ") else: - raise TypeError('Unrecognized format for phase_names arg. Accepts ' - 'string, list, or None') + raise TypeError( + "Unrecognized format for phase_names arg. Accepts string, list, or None" + ) # create the waveform generator self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameTwoPhaseGenerator, - **self.static_params) + **self.static_params, + ) def get_waveforms(self): - r"""Generate the waveforms. - """ + r"""Generate the waveforms.""" if self._current_wfs is None: params = self.current_params # generate the cosine and sine terms - wfs = self.waveform_generator.generate(phases=self.phase_names, - ref_phase=self.ref_phase, - **params) + wfs = self.waveform_generator.generate( + phases=self.phase_names, ref_phase=self.ref_phase, **params + ) for det, (hc, hs) in wfs.items(): # make the same length as the data hc.resize(len(self.data[det])) @@ -1173,18 +1313,17 @@ def get_waveforms(self): # apply high pass if self.highpass_waveforms: hc = highpass( - hc.to_timeseries(), - frequency=self.highpass_waveforms).to_frequencyseries() + hc.to_timeseries(), frequency=self.highpass_waveforms + ).to_frequencyseries() hs = highpass( - hs.to_timeseries(), - frequency=self.highpass_waveforms).to_frequencyseries() + hs.to_timeseries(), frequency=self.highpass_waveforms + ).to_frequencyseries() wfs[det] = (hc, hs) self._current_wfs = wfs return self._current_wfs def get_gated_waveforms(self): - r"""Generate the gated waveforms. - """ + r"""Generate the gated waveforms.""" wfs = self.get_waveforms() out = {} # apply the gate @@ -1195,16 +1334,24 @@ def get_gated_waveforms(self): gate_times = self.get_gate_times() gatestartdelay, dgatedelay = gate_times[det] invmat = self.invert_covariance(det) - hct = hct.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=False, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) - hst = hst.gate(gatestartdelay + dgatedelay/2, - window=dgatedelay/2, copy=False, - invpsd=invpsd, method='paint', - paint_method=self.paint_method, - paint_invmat=invmat) + hct = hct.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=False, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) + hst = hst.gate( + gatestartdelay + dgatedelay / 2, + window=dgatedelay / 2, + copy=False, + invpsd=invpsd, + method="paint", + paint_method=self.paint_method, + paint_invmat=invmat, + ) hc = hct.to_frequencyseries() hs = hst.to_frequencyseries() out[det] = (hc, hs) @@ -1213,12 +1360,11 @@ def get_gated_waveforms(self): @property def _extra_stats(self): """Adds the maxL phase and corresponding likelihood.""" - return ['maxl_phase', 'maxl_logl'] + return ["maxl_phase", "maxl_logl"] @catch_waveform_error def _loglikelihood(self): - r"""Computes the log likelihood. - """ + r"""Computes the log likelihood.""" # get waveforms wfs = self.get_waveforms() gated_wfs = self.get_gated_waveforms() @@ -1226,16 +1372,16 @@ def _loglikelihood(self): data = self.get_data() gated_data = self.get_gated_data() # cycle over all detectors - norm = 0. - hchc = 0. - hchs = 0. - hshc = 0. - hshs = 0. - dhc = 0. - dhs = 0. - hcd = 0. - hsd = 0. - dd = 0. + norm = 0.0 + hchc = 0.0 + hchs = 0.0 + hshc = 0.0 + hshs = 0.0 + dhc = 0.0 + dhs = 0.0 + hcd = 0.0 + hsd = 0.0 + dd = 0.0 for det in self.det_names: if det not in self.dets: self.dets[det] = Detector(det) @@ -1270,30 +1416,32 @@ def _loglikelihood(self): # numerical marginalization over phases cphi = numpy.cos(self.phases) sphi = numpy.sin(self.phases) - hh = cphi*cphi*hchc + sphi*sphi*hshs + cphi*sphi*(hchs+hshc) - dh = cphi*dhc + sphi*dhs - hd = cphi*hcd + sphi*hsd - loglr = -(hh-dh-hd) + hh = cphi * cphi * hchc + sphi * sphi * hshs + cphi * sphi * (hchs + hshc) + dh = cphi * dhc + sphi * dhs + hd = cphi * hcd + sphi * hsd + loglr = -(hh - dh - hd) lognl = -dd # get the maxL phase maxlidx = loglr.argmax() - setattr(self._current_stats, 'maxl_phase', self.phases[maxlidx]) - setattr(self._current_stats, 'maxl_logl', loglr[maxlidx] + lognl + norm) + self._current_stats.maxl_phase = self.phases[maxlidx] + self._current_stats.maxl_logl = loglr[maxlidx] + lognl + norm # get the marginalized log likelihood ratio - marglogl = special.logsumexp(loglr) + lognl + norm - numpy.log(self.phase_samples) + marglogl = ( + special.logsumexp(loglr) + lognl + norm - numpy.log(self.phase_samples) + ) return marglogl @property def multi_signal_support(self): - """ The list of classes that this model supports in a multi-signal + """ + The list of classes that this model supports in a multi-signal likelihood """ return [type(self)] @catch_waveform_error def multi_loglikelihood(self, models): - """ Calculate a multi-model (signal) likelihood - """ + """Calculate a multi-model (signal) likelihood""" # Generate the waveforms for each submodel wfs = [] for m in models + [self]: diff --git a/pycbc/inference/models/gaussian_noise.py b/pycbc/inference/models/gaussian_noise.py index b06b93a4fa8..2db7115bfd4 100644 --- a/pycbc/inference/models/gaussian_noise.py +++ b/pycbc/inference/models/gaussian_noise.py @@ -13,33 +13,37 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides model classes that assume the noise is Gaussian. -""" +"""This module provides model classes that assume the noise is Gaussian.""" import logging import shlex from abc import ABCMeta from functools import wraps + import numpy from pycbc import filter as pyfilter -from pycbc.waveform import (NoWaveformError, FailedWaveformError) -from pycbc.waveform import generator -from pycbc.types import FrequencySeries -from pycbc.strain import gates_from_cli -from pycbc.strain.calibration import Recalibrate from pycbc.inject import InjectionSet from pycbc.io import FieldArray +from pycbc.strain import gates_from_cli +from pycbc.strain.calibration import Recalibrate +from pycbc.types import FrequencySeries from pycbc.types.optparse import MultiDetOptionAction +from pycbc.waveform import FailedWaveformError, NoWaveformError, generator from .base import ModelStats from .base_data import BaseDataModel -from .data_utils import (data_opts_from_config, data_from_cli, - fd_data_from_strain_dict, gate_overwhitened_data) +from .data_utils import ( + data_from_cli, + data_opts_from_config, + fd_data_from_strain_dict, + gate_overwhitened_data, +) def catch_waveform_error(method): - """Decorator that will catch no waveform errors. + """ + Decorator that will catch no waveform errors. This can be added to a method in an inference model. The decorator will call the model's `_nowaveform_return` method if either of the following @@ -51,6 +55,7 @@ def catch_waveform_error(method): This requires the model to have a `_nowaveform_handler` method. """ + # the functools.wroaps decorator preserves the original method's name # and docstring @wraps(method) @@ -69,11 +74,13 @@ def method_wrapper(self, *args, **kwargs): else: raise e return retval + return method_wrapper class BaseGaussianNoise(BaseDataModel, metaclass=ABCMeta): - r"""Model for analyzing GW data with assuming a wide-sense stationary + r""" + Model for analyzing GW data with assuming a wide-sense stationary Gaussian noise model. This model will load gravitational wave data and calculate the log noise @@ -132,18 +139,30 @@ class BaseGaussianNoise(BaseDataModel, metaclass=ABCMeta): fail (i.e., they raise a ``FailedWaveformError``) will be treated as points with zero likelihood. Otherwise, such points will cause the model to raise a ``FailedWaveformError``. + """ - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, ignore_failed_waveforms=False, - no_save_data=False, - **kwargs): + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + ignore_failed_waveforms=False, + no_save_data=False, + **kwargs, + ): # set up the boiler-plate attributes - super(BaseGaussianNoise, self).__init__(variable_params, data, - static_params=static_params, - no_save_data=no_save_data, - **kwargs) + super().__init__( + variable_params, + data, + static_params=static_params, + no_save_data=no_save_data, + **kwargs, + ) self.ignore_failed_waveforms = ignore_failed_waveforms self.no_save_data = no_save_data # check if low frequency cutoff has been provided for every IFO with @@ -159,7 +178,8 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, "every detector in the `[model]` section, where " "`{DETECTOR} is the name of the detector," "or provide a single low-frequency-cutoff option" - "which will be used for all detectors") + "which will be used for all detectors" + ) # check that the data sets all have the same delta fs and delta ts dts = numpy.array([d.delta_t for d in self.data.values()]) @@ -170,16 +190,17 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, self.all_ifodata_same_rate_length = False logging.info( "You are using different data segment lengths or " - "sampling rates for different IFOs") + "sampling rates for different IFOs" + ) # store the number of samples in the time domain self._N = {} - for (det, d) in self._data.items(): - self._N[det] = int(1./(d.delta_f*d.delta_t)) + for det, d in self._data.items(): + self._N[det] = int(1.0 / (d.delta_f * d.delta_t)) # set lower/upper frequency cutoff if high_frequency_cutoff is None: - high_frequency_cutoff = {ifo: None for ifo in self.data} + high_frequency_cutoff = dict.fromkeys(self.data) self._f_upper = high_frequency_cutoff self._f_lower = low_frequency_cutoff @@ -187,10 +208,10 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, self._kmin = {} self._kmax = {} - for (det, d) in self._data.items(): - kmin, kmax = pyfilter.get_cutoff_indices(self._f_lower[det], - self._f_upper[det], - d.delta_f, self._N[det]) + for det, d in self._data.items(): + kmin, kmax = pyfilter.get_cutoff_indices( + self._f_lower[det], self._f_upper[det], d.delta_f, self._N[det] + ) self._kmin[det] = kmin self._kmax[det] = kmax @@ -228,7 +249,8 @@ def low_frequency_cutoff(self): @property def kmin(self): - """Dictionary of starting indices for the inner product. + """ + Dictionary of starting indices for the inner product. This is determined from the lower frequency cutoff and the ``delta_f`` of the data using @@ -238,7 +260,8 @@ def kmin(self): @property def kmax(self): - """Dictionary of ending indices for the inner product. + """ + Dictionary of ending indices for the inner product. This is determined from the high frequency cutoff and the ``delta_f`` of the data using @@ -250,7 +273,8 @@ def kmax(self): @property def psds(self): - """Dictionary of detectors -> PSD frequency series. + """ + Dictionary of detectors -> PSD frequency series. If no PSD was provided for a detector, this will just be a frequency series of ones. @@ -259,7 +283,8 @@ def psds(self): @psds.setter def psds(self, psds): - """Sets the psds, and calculates the weight and norm from them. + """ + Sets the psds, and calculates the weight and norm from them. The data and the low and high frequency cutoffs must be set first. """ @@ -280,8 +305,9 @@ def psds(self, psds): for det, d in self._data.items(): if psds is None: # No psd means assume white PSD - p = FrequencySeries(numpy.ones(int(self._N[det]/2+1)), - delta_f=d.delta_f) + p = FrequencySeries( + numpy.ones(int(self._N[det] / 2 + 1)), delta_f=d.delta_f + ) else: # copy for storage p = psds[det].copy() @@ -291,7 +317,7 @@ def psds(self, psds): kmin = self._kmin[det] kmax = self._kmax[det] invp = FrequencySeries(numpy.zeros(len(p)), delta_f=p.delta_f) - invp[kmin:kmax] = 1./p[kmin:kmax] + invp[kmin:kmax] = 1.0 / p[kmin:kmax] self._invpsds[det] = invp self._weight[det] = numpy.sqrt(4 * invp.delta_f * invp) self._whitened_data[det] = d.copy() @@ -301,7 +327,8 @@ def psds(self, psds): @property def psd_segments(self): - """Dictionary giving times used for PSD estimation for each detector. + """ + Dictionary giving times used for PSD estimation for each detector. If a detector's PSD was not estimated from data, or the segment wasn't provided, that detector will not be in the dictionary. @@ -309,7 +336,8 @@ def psd_segments(self): return self._psd_segments def set_psd_segments(self, psds): - """Sets the PSD segments from a dictionary of PSDs. + """ + Sets the PSD segments from a dictionary of PSDs. This attempts to get the PSD segment from a ``psd_segment`` attribute of each detector's PSD frequency series. If that attribute isn't set, @@ -321,6 +349,7 @@ def set_psd_segments(self, psds): Dictionary of detector name -> PSD frequency series. The segment used for each PSD will try to be retrieved from the PSD's ``.psd_segment`` attribute. + """ for det, p in psds.items(): try: @@ -330,7 +359,8 @@ def set_psd_segments(self, psds): @property def weight(self): - r"""Dictionary of detectors -> frequency series of inner-product + r""" + Dictionary of detectors -> frequency series of inner-product weights. The weights are :math:`\sqrt{4 \Delta f / S_n(f)}`. This is set when @@ -340,7 +370,8 @@ def weight(self): @property def whitened_data(self): - r"""Dictionary of detectors -> whitened data frequency series. + r""" + Dictionary of detectors -> whitened data frequency series. The whitened data is the data multiplied by the inner-product weight. Note that this includes the :math:`\sqrt{4 \Delta f}` factor. This @@ -349,12 +380,13 @@ def whitened_data(self): return self._whitened_data def det_lognorm(self, det): - """The log of the likelihood normalization in the given detector. + """ + The log of the likelihood normalization in the given detector. If ``self.normalize`` is False, will just return 0. """ if not self.normalize: - return 0. + return 0.0 try: return self._lognorm[det] except KeyError: @@ -363,21 +395,21 @@ def det_lognorm(self, det): dt = self._whitened_data[det].delta_t kmin = self._kmin[det] kmax = self._kmax[det] - lognorm = -float(self._N[det]*numpy.log(numpy.pi*self._N[det]*dt)/2. - + numpy.log(p[kmin:kmax]).sum()) + lognorm = -float( + self._N[det] * numpy.log(numpy.pi * self._N[det] * dt) / 2.0 + + numpy.log(p[kmin:kmax]).sum() + ) self._lognorm[det] = lognorm return self._lognorm[det] @property def normalize(self): - """Determines if the loglikelihood includes the normalization term. - """ + """Determines if the loglikelihood includes the normalization term.""" return self._normalize @normalize.setter def normalize(self, normalize): - """Clears the current stats if the normalization state is changed. - """ + """Clears the current stats if the normalization state is changed.""" if normalize != self._normalize: self._current_stats = ModelStats() self._lognorm.clear() @@ -390,7 +422,8 @@ def lognorm(self): return sum(self.det_lognorm(det) for det in self._data) def det_lognl(self, det): - r"""Returns the log likelihood of the noise in the given detector: + r""" + Returns the log likelihood of the noise in the given detector: .. math:: @@ -407,6 +440,7 @@ def det_lognl(self, det): ------- float : The log likelihood of the noise in the requested detector. + """ try: return self._det_lognls[det] @@ -421,7 +455,8 @@ def det_lognl(self, det): return self._det_lognls[det] def _lognl(self): - """Computes the log likelihood assuming the data is noise. + """ + Computes the log likelihood assuming the data is noise. Since this is a constant for Gaussian noise, this is only computed once then stored. @@ -435,7 +470,8 @@ def update(self, **params): self._current_wfs = None def _loglikelihood(self): - r"""Computes the log likelihood of the paramaters, + r""" + Computes the log likelihood of the paramaters, .. math:: @@ -448,13 +484,15 @@ def _loglikelihood(self): ------- float The value of the log likelihood evaluated at the given point. + """ # since the loglr has fewer terms, we'll call that, then just add # back the noise term that canceled in the log likelihood ratio return self.loglr + self.lognl def write_metadata(self, fp, group=None): - """Adds writing the psds, analyzed detectors, and lognl. + """ + Adds writing the psds, analyzed detectors, and lognl. The analyzed detectors, their analysis segments, and the segments used for psd estimation are written as @@ -478,39 +516,39 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ super().write_metadata(fp, group=group) attrs = fp.getattrs(group=group) # write the analyzed detectors and times - attrs['analyzed_detectors'] = self.detectors + attrs["analyzed_detectors"] = self.detectors for det, data in self.data.items(): - key = '{}_analysis_segment'.format(det) + key = f"{det}_analysis_segment" attrs[key] = [float(data.start_time), float(data.end_time)] if self._psds is not None and not self.no_save_data: fp.write_psd(self._psds, group=group) # write the times used for psd estimation (if they were provided) for det in self.psd_segments: - key = '{}_psd_segment'.format(det) + key = f"{det}_psd_segment" attrs[key] = list(map(float, self.psd_segments[det])) # save the frequency cutoffs for det in self.detectors: - attrs['{}_likelihood_low_freq'.format(det)] = self._f_lower[det] + attrs[f"{det}_likelihood_low_freq"] = self._f_lower[det] if self._f_upper[det] is not None: - attrs['{}_likelihood_high_freq'.format(det)] = \ - self._f_upper[det] + attrs[f"{det}_likelihood_high_freq"] = self._f_upper[det] # write the lognl to the samples group attrs sampattrs = fp.getattrs(group=fp.samples_group) # if a group is specified, prepend the lognl names with it - if group is None or group == '/': - prefix = '' + if group is None or group == "/": + prefix = "" else: - prefix = group.replace('/', '__') - if not prefix.endswith('__'): - prefix += '__' - sampattrs['{}lognl'.format(prefix)] = self.lognl + prefix = group.replace("/", "__") + if not prefix.endswith("__"): + prefix += "__" + sampattrs[f"{prefix}lognl"] = self.lognl # also save the lognl in each detector for det in self.detectors: - sampattrs['{}{}_lognl'.format(prefix, det)] = self.det_lognl(det) + sampattrs[f"{prefix}{det}_lognl"] = self.det_lognl(det) @staticmethod def _fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict): @@ -518,7 +556,8 @@ def _fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict): return fd_data_from_strain_dict(opts, strain_dict, psd_strain_dict) def _nowaveform_handler(self): - """Method that gets called if a NoWaveformError or FailedWaveformError + """ + Method that gets called if a NoWaveformError or FailedWaveformError is raised. See the :py:func:catch_waveform_error decorator for details. Here, this will just raise a NotImplementedError, since how this should @@ -527,12 +566,13 @@ def _nowaveform_handler(self): """ raise NotImplementedError( f"A waveform could not be generated, but this model does not know " - f"how to handle that. The parameters were: {self.current_params}.") + f"how to handle that. The parameters were: {self.current_params}." + ) @classmethod - def from_config(cls, cp, data_section='data', data=None, psds=None, - **kwargs): - r"""Initializes an instance of this class from the given config file. + def from_config(cls, cp, data_section="data", data=None, psds=None, **kwargs): + r""" + Initializes an instance of this class from the given config file. In addition to ``[model]``, a ``data_section`` (default ``[data]``) must be in the configuration file. The data section specifies settings @@ -582,10 +622,11 @@ def from_config(cls, cp, data_section='data', data=None, psds=None, \**kwargs : All additional keyword arguments are passed to the class. Any provided keyword will override what is in the config file. + """ # get the injection file, to replace any FROM_INJECTION settings - if 'injection-file' in cp.options('data'): - injection_file = cp.get('data', 'injection-file') + if "injection-file" in cp.options("data"): + injection_file = cp.get("data", "injection-file") else: injection_file = None # update any values that are to be retrieved from the injection @@ -593,85 +634,95 @@ def from_config(cls, cp, data_section='data', data=None, psds=None, get_values_from_injection(cp, injection_file, update_cp=True) args = cls._init_args_from_config(cp) # add the injection file - args['injection_file'] = injection_file + args["injection_file"] = injection_file # check if normalize is set - if cp.has_option('model', 'normalize'): - args['normalize'] = True - if cp.has_option('model', 'ignore-failed-waveforms'): - args['ignore_failed_waveforms'] = True - if cp.has_option('model', 'det-frame-waveform'): - args['det_frame_waveform'] = True - if cp.has_option('model', 'no-save-data'): - args['no_save_data'] = True + if cp.has_option("model", "normalize"): + args["normalize"] = True + if cp.has_option("model", "ignore-failed-waveforms"): + args["ignore_failed_waveforms"] = True + if cp.has_option("model", "det-frame-waveform"): + args["det_frame_waveform"] = True + if cp.has_option("model", "no-save-data"): + args["no_save_data"] = True # set inverse spectrum to truncate if not set - if cp.has_option(data_section, 'invpsd-trunc-which-spectrum'): - cp.set(data_section, 'invpsd-trunc-which-spectrum', 'invpsd') + if cp.has_option(data_section, "invpsd-trunc-which-spectrum"): + cp.set(data_section, "invpsd-trunc-which-spectrum", "invpsd") # get any other keyword arguments provided in the model section ignore_args = [ - 'name', - 'normalize', - 'ignore-failed-waveforms', - 'no-save-data', - 'det-frame-waveform' + "name", + "normalize", + "ignore-failed-waveforms", + "no-save-data", + "det-frame-waveform", ] for option in cp.options("model"): if option in ("low-frequency-cutoff", "high-frequency-cutoff"): ignore_args.append(option) - name = option.replace('-', '_') - args[name] = cp.get_cli_option('model', name, - nargs='+', type=float, - action=MultiDetOptionAction) + name = option.replace("-", "_") + args[name] = cp.get_cli_option( + "model", name, nargs="+", type=float, action=MultiDetOptionAction + ) - if 'low_frequency_cutoff' not in args: - raise ValueError("low-frequency-cutoff must be provided in the" - " model section, but is not found!") + if "low_frequency_cutoff" not in args: + raise ValueError( + "low-frequency-cutoff must be provided in the" + " model section, but is not found!" + ) # data args - bool_args = ['check-for-valid-times', 'shift-psd-times-to-valid', - 'err-on-missing-detectors'] - data_args = {arg.replace('-', '_'): True for arg in bool_args - if cp.has_option('model', arg)} + bool_args = [ + "check-for-valid-times", + "shift-psd-times-to-valid", + "err-on-missing-detectors", + ] + data_args = { + arg.replace("-", "_"): True + for arg in bool_args + if cp.has_option("model", arg) + } ignore_args += bool_args # load the data - opts = data_opts_from_config(cp, data_section, - args['low_frequency_cutoff']) + opts = data_opts_from_config(cp, data_section, args["low_frequency_cutoff"]) if data is None or psds is None: strain_dict, psd_strain_dict = data_from_cli(opts, **data_args) # convert to frequency domain and get psds stilde_dict, psds = cls._fd_data_from_strain_dict( - opts, strain_dict, psd_strain_dict) + opts, strain_dict, psd_strain_dict + ) # save the psd data segments if the psd was estimated from data if opts.psd_estimation: _tdict = psd_strain_dict or strain_dict for det in psds: - psds[det].psd_segment = (_tdict[det].start_time, - _tdict[det].end_time) + psds[det].psd_segment = ( + _tdict[det].start_time, + _tdict[det].end_time, + ) # gate overwhitened if desired if opts.gate_overwhitened and opts.gate is not None: - stilde_dict = gate_overwhitened_data( - stilde_dict, psds, opts.gate) + stilde_dict = gate_overwhitened_data(stilde_dict, psds, opts.gate) data = stilde_dict - args.update({'data': data, 'psds': psds}) + args.update({"data": data, "psds": psds}) # any extra args - args.update(cls.extra_args_from_config(cp, "model", - skip_args=ignore_args)) + args.update(cls.extra_args_from_config(cp, "model", skip_args=ignore_args)) # get ifo-specific instances of calibration model - if cp.has_section('calibration'): + if cp.has_section("calibration"): logging.info("Initializing calibration model") recalib = { - ifo: Recalibrate.from_config(cp, ifo, section='calibration') - for ifo in opts.instruments} - args['recalibration'] = recalib + ifo: Recalibrate.from_config(cp, ifo, section="calibration") + for ifo in opts.instruments + } + args["recalibration"] = recalib # get gates for templates gates = gates_from_cli(opts) if gates: - args['gates'] = gates + args["gates"] = gates args.update(kwargs) return cls(**args) class GaussianNoise(BaseGaussianNoise): - r"""Model that assumes data is stationary Gaussian noise. + r""" + Model that assumes data is stationary Gaussian noise. With Gaussian noise the log likelihood functions for signal :math:`\log p(d|\Theta, h)` and for noise :math:`\log p(d|n)` are given by: @@ -880,16 +931,32 @@ class GaussianNoise(BaseGaussianNoise): logprior: 0.92 """ - name = 'gaussian_noise' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, det_frame_waveform=False, **kwargs): + name = "gaussian_noise" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + det_frame_waveform=False, + **kwargs, + ): # set up the boiler-plate attributes - super(GaussianNoise, self).__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + super().__init__( + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) # Determine if all data have the same sampling rate and segment length if det_frame_waveform: generator_class = generator.FDomainDirectDetFrameGenerator @@ -898,52 +965,60 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, if self.all_ifodata_same_rate_length: # create a waveform generator for all ifos self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, generator_class=generator_class, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, - gates=self.gates, **self.static_params) + gates=self.gates, + **self.static_params, + ) else: # create a waveform generator for each ifo respestively self.waveform_generator = {} for det in self.data: self.waveform_generator[det] = create_waveform_generator( - self.variable_params, {det: self.data[det]}, + self.variable_params, + {det: self.data[det]}, generator_class=generator_class, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, - gates=self.gates, **self.static_params) + gates=self.gates, + **self.static_params, + ) @property def _extra_stats(self): - """Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each - detector.""" - return ['loglr'] + \ - ['{}_cplx_loglr'.format(det) for det in self._data] + \ - ['{}_optimal_snrsq'.format(det) for det in self._data] + """ + Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each + detector. + """ + return ( + ["loglr"] + + [f"{det}_cplx_loglr" for det in self._data] + + [f"{det}_optimal_snrsq" for det in self._data] + ) def _nowaveform_handler(self): - """Convenience function to set loglr values if no waveform generated. - """ + """Convenience function to set loglr values if no waveform generated.""" for det in self._data: - setattr(self._current_stats, 'loglikelihood', -numpy.inf) - setattr(self._current_stats, '{}_cplx_loglr'.format(det), - -numpy.inf) + self._current_stats.loglikelihood = -numpy.inf + setattr(self._current_stats, f"{det}_cplx_loglr", -numpy.inf) # snr can't be < 0 by definition, so return 0 - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), 0.) + setattr(self._current_stats, f"{det}_optimal_snrsq", 0.0) return -numpy.inf @property def multi_signal_support(self): - """ The list of classes that this model supports in a multi-signal + """ + The list of classes that this model supports in a multi-signal likelihood """ return [type(self)] @catch_waveform_error def multi_loglikelihood(self, models): - """ Calculate a multi-model (signal) likelihood - """ + """Calculate a multi-model (signal) likelihood""" # Generate the waveforms for each submodel wfs = [] for m in models + [self]: @@ -962,7 +1037,8 @@ def multi_loglikelihood(self, models): return loglr + self.lognl def get_waveforms(self): - """The waveforms generated using the current parameters. + """ + The waveforms generated using the current parameters. If the waveforms haven't been generated yet, they will be generated. @@ -970,6 +1046,7 @@ def get_waveforms(self): ------- dict : Dictionary of detector names -> FrequencySeries. + """ if self._current_wfs is None: params = self.current_params @@ -984,7 +1061,8 @@ def get_waveforms(self): @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, .. math:: @@ -998,9 +1076,10 @@ def _loglr(self): ------- float The value of the log likelihood ratio. + """ wfs = self.get_waveforms() - lr = 0. + lr = 0.0 for det, h in wfs.items(): # the kmax of the waveforms may be different than internal kmax kmax = min(len(h), self._kmax[det]) @@ -1008,20 +1087,19 @@ def _loglr(self): # if the waveform terminates before the filtering low frequency # cutoff, then the loglr is just 0 for this detector cplx_hd = 0j - hh = 0. + hh = 0.0 else: slc = slice(self._kmin[det], kmax) # whiten the waveform - h[self._kmin[det]:kmax] *= self._weight[det][slc] + h[self._kmin[det] : kmax] *= self._weight[det][slc] # the inner products cplx_hd = h[slc].inner(self._whitened_data[det][slc]) # hh = h[slc].inner(h[slc]).real # < h, h> cplx_loglr = cplx_hd - 0.5 * hh # store - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), hh) - setattr(self._current_stats, '{}_cplx_loglr'.format(det), - cplx_loglr) + setattr(self._current_stats, f"{det}_optimal_snrsq", hh) + setattr(self._current_stats, f"{det}_cplx_loglr", cplx_loglr) lr += cplx_loglr.real # also store the loglikelihood, to ensure it is populated in the # current stats even if loglikelihood is never called @@ -1029,7 +1107,8 @@ def _loglr(self): return float(lr) def det_cplx_loglr(self, det): - """Returns the complex log likelihood ratio in the given detector. + """ + Returns the complex log likelihood ratio in the given detector. Parameters ---------- @@ -1040,18 +1119,20 @@ def det_cplx_loglr(self, det): ------- complex float : The complex log likelihood ratio. + """ # try to get it from current stats try: - return getattr(self._current_stats, '{}_cplx_loglr'.format(det)) + return getattr(self._current_stats, f"{det}_cplx_loglr") except AttributeError: # hasn't been calculated yet; call loglr to do so self._loglr() # now try returning again - return getattr(self._current_stats, '{}_cplx_loglr'.format(det)) + return getattr(self._current_stats, f"{det}_cplx_loglr") def det_optimal_snrsq(self, det): - """Returns the opitmal SNR squared in the given detector. + """ + Returns the opitmal SNR squared in the given detector. Parameters ---------- @@ -1062,15 +1143,16 @@ def det_optimal_snrsq(self, det): ------- float : The opimtal SNR squared. + """ # try to get it from current stats try: - return getattr(self._current_stats, '{}_optimal_snrsq'.format(det)) + return getattr(self._current_stats, f"{det}_optimal_snrsq") except AttributeError: # hasn't been calculated yet; call loglr to do so self._loglr() # now try returning again - return getattr(self._current_stats, '{}_optimal_snrsq'.format(det)) + return getattr(self._current_stats, f"{det}_optimal_snrsq") # @@ -1083,7 +1165,8 @@ def det_optimal_snrsq(self, det): def get_values_from_injection(cp, injection_file, update_cp=True): - """Replaces all FROM_INJECTION values in a config file with the + """ + Replaces all FROM_INJECTION values in a config file with the corresponding value from the injection. This looks for any options that start with ``FROM_INJECTION[:ARG]`` in @@ -1132,8 +1215,9 @@ def get_values_from_injection(cp, injection_file, update_cp=True): list The parameters that were replaced, as a tuple of section name, option, value. + """ - lookfor = 'FROM_INJECTION' + lookfor = "FROM_INJECTION" # figure out what parameters need to be set replace_params = [] for sec in cp.sections(): @@ -1144,7 +1228,7 @@ def get_values_from_injection(cp, injection_file, update_cp=True): for ii, subval in enumerate(splitvals): if subval.startswith(lookfor): # determine what we should retrieve from the injection - subval = subval.split(':', 1) + subval = subval.split(":", 1) if len(subval) == 1: subval = opt else: @@ -1155,15 +1239,17 @@ def get_values_from_injection(cp, injection_file, update_cp=True): if replace_params: # check that we have an injection file if injection_file is None: - raise ValueError("One or values are set to {}, but no injection " - "file provided".format(lookfor)) + raise ValueError( + f"One or values are set to {lookfor}, but no injection file provided" + ) # load the injection file inj = InjectionSet(injection_file).table.view(type=FieldArray) # make sure there's only one injection provided if inj.size > 1: - raise ValueError("One or more values are set to {}, but more than " - "one injection exists in the injection file." - .format(lookfor)) + raise ValueError( + f"One or more values are set to {lookfor}, but more than " + "one injection exists in the injection file." + ) # get the injection values to replace for ii, (sec, opt, splitvals, replace_this) in enumerate(replace_params): # replace the value in the shlex-splitted string with the value @@ -1180,24 +1266,29 @@ def get_values_from_injection(cp, injection_file, update_cp=True): # following can just be replaced by: # replace_val = shlex.join(splitvals) for jj, arg in enumerate(splitvals): - if ' ' in arg: + if " " in arg: arg = "'" + arg + "'" splitvals[jj] = arg - replace_val = ' '.join(splitvals) + replace_val = " ".join(splitvals) replace_params[ii] = (sec, opt, replace_val) # replace in the config file if update_cp: - for (sec, opt, replace_val) in replace_params: + for sec, opt, replace_val in replace_params: cp.set(sec, opt, replace_val) return replace_params def create_waveform_generator( - variable_params, data, waveform_transforms=None, - recalibration=None, gates=None, - generator_class=generator.FDomainDetFrameGenerator, - **static_params): - r"""Creates a waveform generator for use with a model. + variable_params, + data, + waveform_transforms=None, + recalibration=None, + gates=None, + generator_class=generator.FDomainDetFrameGenerator, + **static_params, +): + r""" + Creates a waveform generator for use with a model. Parameters ---------- @@ -1228,24 +1319,24 @@ def create_waveform_generator( ------- pycbc.waveform.FDomainDetFrameGenerator A waveform generator for frequency domain generation. + """ # the waveform generator will get the variable_params + the output # of the waveform transforms, so we'll add them to the list of # parameters if waveform_transforms is not None: - wfoutputs = set.union(*[t.outputs - for t in waveform_transforms]) + wfoutputs = set.union(*[t.outputs for t in waveform_transforms]) else: wfoutputs = set() variable_params = list(variable_params) + list(wfoutputs) # figure out what generator to use based on the approximant try: - approximant = static_params['approximant'] + approximant = static_params["approximant"] except KeyError: raise ValueError("no approximant provided in the static args") - dm = static_params.get('preferred_domain', None) - if isinstance(dm, str) and dm.lower() == 'none': + dm = static_params.get("preferred_domain") + if isinstance(dm, str) and dm.lower() == "none": dm = None gen_function = generator_class.select_rframe_generator(approximant, dm) @@ -1257,15 +1348,21 @@ def create_waveform_generator( delta_f = d.delta_f delta_t = d.delta_t start_time = d.start_time - else: - if not all([d.delta_f == delta_f, d.delta_t == delta_t, - d.start_time == start_time]): - raise ValueError("data must all have the same delta_t, " - "delta_f, and start_time") + elif not all( + [d.delta_f == delta_f, d.delta_t == delta_t, d.start_time == start_time] + ): + raise ValueError( + "data must all have the same delta_t, delta_f, and start_time" + ) waveform_generator = generator_class( - gen_function, epoch=start_time, - variable_args=variable_params, detectors=list(data.keys()), - delta_f=delta_f, delta_t=delta_t, - recalib=recalibration, gates=gates, - **static_params) + gen_function, + epoch=start_time, + variable_args=variable_params, + detectors=list(data.keys()), + delta_f=delta_f, + delta_t=delta_t, + recalib=recalibration, + gates=gates, + **static_params, + ) return waveform_generator diff --git a/pycbc/inference/models/hierarchical.py b/pycbc/inference/models/hierarchical.py index 1a15caccc71..a0b0e45d2dc 100644 --- a/pycbc/inference/models/hierarchical.py +++ b/pycbc/inference/models/hierarchical.py @@ -25,11 +25,14 @@ """Hierarchical model definitions.""" -import shlex import logging +import shlex + import numpy + from pycbc import transforms from pycbc.workflow import WorkflowConfigParser + from .base import BaseModel # @@ -42,7 +45,8 @@ class HierarchicalModel(BaseModel): - r"""Model that is a combination of other models. + r""" + Model that is a combination of other models. Sub-models are treated as being independent of each other, although they can share parameters. In other words, the hierarchical likelihood is: @@ -79,8 +83,10 @@ class HierarchicalModel(BaseModel): \**kwargs : All other keyword arguments are passed to :py:class:`BaseModel `. + """ - name = 'hierarchical' + + name = "hierarchical" def __init__(self, variable_params, submodels, **kwargs): # sub models is assumed to be a dict of model labels -> model instances @@ -92,11 +98,11 @@ def __init__(self, variable_params, submodels, **kwargs): # add any parameters created by waveform transforms if self.waveform_transforms is not None: derived_params = set() - derived_params.update(*[t.outputs - for t in self.waveform_transforms]) + derived_params.update(*[t.outputs for t in self.waveform_transforms]) # convert to hierarchical params - derived_params = map_params(hpiter(derived_params, - list(self.submodels.keys()))) + derived_params = map_params( + hpiter(derived_params, list(self.submodels.keys())) + ) for lbl, pset in derived_params.items(): self.param_map[lbl].update(pset) # make sure the static parameters of all submodels are set correctly @@ -107,11 +113,18 @@ def __init__(self, variable_params, submodels, **kwargs): self.extra_stats_map = {} self.__extra_stats = [] for lbl, model in self.submodels.items(): - model.static_params = {p.subname: self.static_params[p.fullname] - for p in self.static_param_map[lbl]} - self.extra_stats_map.update(map_params([ - HierarchicalParam.from_subname(lbl, p) - for p in model._extra_stats+['loglikelihood']])) + model.static_params = { + p.subname: self.static_params[p.fullname] + for p in self.static_param_map[lbl] + } + self.extra_stats_map.update( + map_params( + [ + HierarchicalParam.from_subname(lbl, p) + for p in model._extra_stats + ["loglikelihood"] + ] + ) + ) self.__extra_stats += self.extra_stats_map[lbl] # also make sure the model's sampling transforms and waveform # transforms are not set, as these are handled by the hierarchical @@ -120,19 +133,24 @@ def __init__(self, variable_params, submodels, **kwargs): # transform with prefix on the submodel's level if self.name != "joint_primary_marginalized": if model.sampling_transforms is not None: - raise ValueError("Model {} has sampling transforms " - "set; in a hierarchical analysis, " - "these are handled by the " - "hierarchical model".format(lbl)) + raise ValueError( + f"Model {lbl} has sampling transforms " + "set; in a hierarchical analysis, " + "these are handled by the " + "hierarchical model" + ) if model.waveform_transforms is not None: - raise ValueError("Model {} has waveform transforms " - "set; in a hierarchical analysis, " - "these are handled by the " - "hierarchical model".format(lbl)) + raise ValueError( + f"Model {lbl} has waveform transforms " + "set; in a hierarchical analysis, " + "these are handled by the " + "hierarchical model" + ) @property def hvariable_params(self): - """The variable params as a tuple of :py:class:`HierarchicalParam` + """ + The variable params as a tuple of :py:class:`HierarchicalParam` instances. """ return self._variable_params @@ -148,12 +166,14 @@ def variable_params(self, variable_params): # as HierarchicalParam instances if isinstance(variable_params, str): variable_params = [variable_params] - self._variable_params = tuple(HierarchicalParam(p, self.submodels) - for p in variable_params) + self._variable_params = tuple( + HierarchicalParam(p, self.submodels) for p in variable_params + ) @property def hstatic_params(self): - """The static params with :py:class:`HierarchicalParam` instances used + """ + The static params with :py:class:`HierarchicalParam` instances used as dictionary keys. """ return self._static_params @@ -167,8 +187,10 @@ def static_params(self): def static_params(self, static_params): if static_params is None: static_params = {} - self._static_params = {HierarchicalParam(p, self.submodels): val - for p, val in static_params.items()} + self._static_params = { + HierarchicalParam(p, self.submodels): val + for p, val in static_params.items() + } @property def _extra_stats(self): @@ -181,13 +203,17 @@ def _hextra_stats(self): def _loglikelihood(self): # takes the sum of the constitutent models' loglikelihoods - logl = 0. + logl = 0.0 for lbl, model in self.submodels.items(): # update the model with the current params. This is done here # instead of in `update` because waveform transforms are not # applied until the loglikelihood function is called - model.update(**{p.subname: self.current_params[p.fullname] - for p in self.param_map[lbl]}) + model.update( + **{ + p.subname: self.current_params[p.fullname] + for p in self.param_map[lbl] + } + ) # now get the loglikelihood from the model sublogl = model.loglikelihood # store the extra stats @@ -199,7 +225,8 @@ def _loglikelihood(self): return logl def write_metadata(self, fp, group=None): - """Adds data to the metadata that's written. + """ + Adds data to the metadata that's written. Parameters ---------- @@ -215,24 +242,25 @@ def write_metadata(self, fp, group=None): super().write_metadata(fp, group=group) # write information about each submodel into a different group for # each one - if group is None or group == '/': - prefix = '' + if group is None or group == "/": + prefix = "" else: - prefix = group+'/' + prefix = group + "/" for lbl, model in self.submodels.items(): - model.write_metadata(fp, group=prefix+lbl) + model.write_metadata(fp, group=prefix + lbl) # if all submodels support it, write a combined lognl parameter try: sampattrs = fp.getattrs(group=fp.samples_group) lognl = [self.submodels[k].lognl for k in self.submodels] - sampattrs['{}lognl'.format(prefix)] = sum(lognl) + sampattrs[f"{prefix}lognl"] = sum(lognl) except (AttributeError, ValueError): pass @classmethod def from_config(cls, cp, **kwargs): - r"""Initializes an instance of this class from the given config file. + r""" + Initializes an instance of this class from the given config file. Sub-models are initialized before initializing this class. The model section must have a ``submodels`` argument that lists the names of all @@ -289,25 +317,25 @@ def from_config(cls, cp, **kwargs): Preferencing a keyword argument by ``{submodel}__`` will send the parameter as a keyword argument to the specified submodel's ``from_config`` method. + """ # we need the read from config function from the init; to prevent # circular imports, we import it here from pycbc.inference.models import read_from_config + # get the submodels - submodel_lbls = shlex.split(cp.get('model', 'submodels')) + submodel_lbls = shlex.split(cp.get("model", "submodels")) # sort parameters by model - vparam_map = map_params(hpiter(cp.options('variable_params'), - submodel_lbls)) - sparam_map = map_params(hpiter(cp.options('static_params'), - submodel_lbls)) + vparam_map = map_params(hpiter(cp.options("variable_params"), submodel_lbls)) + sparam_map = map_params(hpiter(cp.options("static_params"), submodel_lbls)) # we'll need any waveform transforms for the initializing sub-models, # as the underlying models will receive the output of those transforms - if any(cp.get_subsections('waveform_transforms')): + if any(cp.get_subsections("waveform_transforms")): waveform_transforms = transforms.read_transforms_from_config( - cp, 'waveform_transforms') - wfoutputs = set.union(*[t.outputs - for t in waveform_transforms]) + cp, "waveform_transforms" + ) + wfoutputs = set.union(*[t.outputs for t in waveform_transforms]) wfparam_map = map_params(hpiter(wfoutputs, submodel_lbls)) else: wfparam_map = {lbl: [] for lbl in submodel_lbls} @@ -322,56 +350,61 @@ def from_config(cls, cp, **kwargs): # include the [model] section for that model) copy_sections = [ HierarchicalParam(sec, submodel_lbls) - for sec in cp.sections() if lbl in - sec.split('-')[0].split(HierarchicalParam.delim, 1)[0]] + for sec in cp.sections() + if lbl in sec.split("-")[0].split(HierarchicalParam.delim, 1)[0] + ] for sec in copy_sections: # check that the user isn't trying to set variable or static # params for the model (we won't worry about waveform or # sampling transforms here, since that is checked for in the # __init__) - if sec.subname in ['variable_params', 'static_params']: - raise ValueError("Section {} found in the config file; " - "[variable_params] and [static_params] " - "sections should not include model " - "labels. To specify parameters unique to " - "one or more sub-models, prepend the " - "individual parameter names with the " - "model label. See HierarchicalParam for " - "details.".format(sec)) + if sec.subname in ["variable_params", "static_params"]: + raise ValueError( + f"Section {sec} found in the config file; " + "[variable_params] and [static_params] " + "sections should not include model " + "labels. To specify parameters unique to " + "one or more sub-models, prepend the " + "individual parameter names with the " + "model label. See HierarchicalParam for " + "details." + ) subcp.add_section(sec.subname) for opt, val in cp.items(sec): subcp.set(sec.subname, opt, val) # set the static params - subcp.add_section('static_params') + subcp.add_section("static_params") for param in sparam_map[lbl]: - subcp.set('static_params', param.subname, - cp.get('static_params', param.fullname)) + subcp.set( + "static_params", + param.subname, + cp.get("static_params", param.fullname), + ) # set the variable params: for now we'll just set all the # variable params as static params # so that the model doesn't raise an error looking for # prior sections. We'll then manually set the variable # params after the model is initialized - subcp.add_section('variable_params') + subcp.add_section("variable_params") for param in vparam_map[lbl]: - subcp.set('static_params', param.subname, 'REPLACE') + subcp.set("static_params", param.subname, "REPLACE") # add the outputs from the waveform transforms for param in wfparam_map[lbl]: - subcp.set('static_params', param.subname, 'REPLACE') + subcp.set("static_params", param.subname, "REPLACE") # extra any kwargs to pass subkwargs = {} for p, kwarg in list(kwargs.items()): - if p.startswith(lbl+'__'): + if p.startswith(lbl + "__"): val = kwargs.pop(p) - subkwargs[p.replace(lbl+'__', '', 1)] = val + subkwargs[p.replace(lbl + "__", "", 1)] = val # initialize submodel = read_from_config(subcp, **subkwargs) # move the static params back to variable for p in vparam_map[lbl]: submodel.static_params.pop(p.subname) - submodel.variable_params = tuple(p.subname - for p in vparam_map[lbl]) + submodel.variable_params = tuple(p.subname for p in vparam_map[lbl]) # remove the waveform transform parameters for p in wfparam_map[lbl]: submodel.static_params.pop(p.subname) @@ -384,7 +417,8 @@ def from_config(cls, cp, **kwargs): class HierarchicalParam(str): - """Sub-class of str for hierarchical parameter names. + """ + Sub-class of str for hierarchical parameter names. This adds attributes that keep track of the model label(s) the parameter is associated with, along with the name that is passed to the models. @@ -424,9 +458,11 @@ class HierarchicalParam(str): subname : str The name of the parameter without the model labels prepended to it. For example, ``e1_e2__foo`` yields ``foo``. + """ - delim = '__' - model_delim = '_' + + delim = "__" + model_delim = "_" def __new__(cls, fullname, possible_models): fullname = str(fullname) @@ -439,13 +475,13 @@ def __new__(cls, fullname, possible_models): @classmethod def from_subname(cls, model_label, subname): - """Creates a HierarchicalParam from the given subname and model label. - """ + """Creates a HierarchicalParam from the given subname and model label.""" return cls(cls.delim.join([model_label, subname]), set([model_label])) @classmethod def parse(cls, fullname, possible_models): - """Parses the full parameter name into the models the parameter is + """ + Parses the full parameter name into the models the parameter is associated with and the parameter name that is passed to the models. Parameters @@ -463,6 +499,7 @@ def parse(cls, fullname, possible_models): subp : str Parameter name that is passed to the models. This is the parameter name with the model label(s) stripped from it. + """ # make sure possible models is a set possible_models = set(possible_models) @@ -478,14 +515,17 @@ def parse(cls, fullname, possible_models): # make sure the given labels are in the list of possible models unknown = models - possible_models if any(unknown): - raise ValueError('unrecognized model label(s) {} present in ' - 'parameter {}'.format(', '.join(unknown), - fullname)) + raise ValueError( + "unrecognized model label(s) {} present in parameter {}".format( + ", ".join(unknown), fullname + ) + ) return models, subp def hpiter(params, possible_models): - """Turns a list of parameter strings into a list of HierarchicalParams. + """ + Turns a list of parameter strings into a list of HierarchicalParams. Parameters ---------- @@ -498,12 +538,14 @@ def hpiter(params, possible_models): ------- iterator : Iterator of :py:class:`HierarchicalParam` instances. + """ return map(lambda x: HierarchicalParam(x, possible_models), params) def map_params(params): - """Creates a map of models -> parameters. + """ + Creates a map of models -> parameters. Parameters ---------- @@ -514,6 +556,7 @@ def map_params(params): ------- dict : Dictionary of model labels -> associated parameters. + """ param_map = {} for p in params: @@ -526,7 +569,8 @@ def map_params(params): class MultiSignalModel(HierarchicalModel): - """ Model for multiple signals which share data + """ + Model for multiple signals which share data Sub models are treated as if the signals overlap in data. This requires constituent models to implement a specific method to handle this case. @@ -539,7 +583,8 @@ class MultiSignalModel(HierarchicalModel): configuration files is the same. The primary model is used to determine the noise terms , which by default will be the first model used. """ - name = 'multi_signal' + + name = "multi_signal" def __init__(self, variable_params, submodels, **kwargs): super().__init__(variable_params, submodels, **kwargs) @@ -551,27 +596,30 @@ def __init__(self, variable_params, submodels, **kwargs): model = self.submodels[lbl] ctypes.add(type(model)) - if hasattr(model, 'multi_signal_support'): + if hasattr(model, "multi_signal_support"): support[lbl] = set(model.multi_signal_support) # pick the primary model if it supports the set of constituent models for lbl in support: if ctypes <= support[lbl]: self.primary_model = lbl - logging.info('MultiSignalModel: PrimaryModel == %s', lbl) + logging.info("MultiSignalModel: PrimaryModel == %s", lbl) break else: # Oh, no, we don't support this combo! - raise RuntimeError("It looks like the combination of models, {}," - "for the MultiSignal model isn't supported by" - "any of the constituent models.".format(ctypes)) + raise RuntimeError( + f"It looks like the combination of models, {ctypes}," + "for the MultiSignal model isn't supported by" + "any of the constituent models." + ) self.other_models = self.submodels.copy() self.other_models.pop(self.primary_model) self.other_models = list(self.other_models.values()) def write_metadata(self, fp, group=None): - """Adds metadata to the output files + """ + Adds metadata to the output files Parameters ---------- @@ -581,27 +629,32 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ super().write_metadata(fp, group=group) sampattrs = fp.getattrs(group=fp.samples_group) # if a group is specified, prepend the lognl names with it - if group is None or group == '/': - prefix = '' + if group is None or group == "/": + prefix = "" else: - prefix = group.replace('/', '__') - if not prefix.endswith('__'): - prefix += '__' + prefix = group.replace("/", "__") + if not prefix.endswith("__"): + prefix += "__" try: model = self.submodels[self.primary_model] - sampattrs['{}lognl'.format(prefix)] = model.lognl + sampattrs[f"{prefix}lognl"] = model.lognl except AttributeError: pass def _loglikelihood(self): for lbl, model in self.submodels.items(): # Update the parameters of each - model.update(**{p.subname: self.current_params[p.fullname] - for p in self.param_map[lbl]}) + model.update( + **{ + p.subname: self.current_params[p.fullname] + for p in self.param_map[lbl] + } + ) # Calculate the combined loglikelihood p = self.primary_model @@ -616,7 +669,8 @@ def _loglikelihood(self): class JointPrimaryMarginalizedModel(HierarchicalModel): - """This likelihood model can be used for cases when one of the submodels + """ + This likelihood model can be used for cases when one of the submodels can be marginalized to accelerate the total likelihood. This likelihood model also allows for further acceleration of other models during marginalization, if some extrinsic parameters can be tightly constrained @@ -625,25 +679,29 @@ class JointPrimaryMarginalizedModel(HierarchicalModel): multiband observation, SOBHB signals' (tc, ra, dec) can be tightly constrained by 3G network, so this model is also useful for this case. """ - name = 'joint_primary_marginalized' + + name = "joint_primary_marginalized" def __init__(self, variable_params, submodels, **kwargs): super().__init__(variable_params, submodels, **kwargs) # assume the ground-based submodel as the primary model - self.primary_model = self.submodels[kwargs['primary_lbl'][0]] - self.primary_lbl = kwargs['primary_lbl'][0] + self.primary_model = self.submodels[kwargs["primary_lbl"][0]] + self.primary_lbl = kwargs["primary_lbl"][0] self.other_models = self.submodels.copy() - self.other_models.pop(kwargs['primary_lbl'][0]) + self.other_models.pop(kwargs["primary_lbl"][0]) self.other_models = list(self.other_models.values()) # determine whether to accelerate total_loglr from .tools import str_to_bool - self.static_margin_params_in_other_models = \ - str_to_bool(kwargs['static_margin_params_in_other_models'][0]) + + self.static_margin_params_in_other_models = str_to_bool( + kwargs["static_margin_params_in_other_models"][0] + ) def write_metadata(self, fp, group=None): - """Adds metadata to the output files + """ + Adds metadata to the output files Parameters ---------- @@ -653,25 +711,26 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ super().write_metadata(fp, group=group) sampattrs = fp.getattrs(group=fp.samples_group) # if a group is specified, prepend the lognl names with it - if group is None or group == '/': - prefix = '' + if group is None or group == "/": + prefix = "" else: - prefix = group.replace('/', '__') - if not prefix.endswith('__'): - prefix += '__' + prefix = group.replace("/", "__") + if not prefix.endswith("__"): + prefix += "__" try: for lbl, model in self.submodels.items(): - sampattrs['{}lognl'.format(prefix + '%s__' % lbl) - ] = model.lognl + sampattrs["{}lognl".format(prefix + "%s__" % lbl)] = model.lognl except AttributeError: pass def total_loglr(self): - r"""Computes the total log likelihood ratio, + r""" + Computes the total log likelihood ratio, .. math:: @@ -685,6 +744,7 @@ def total_loglr(self): ------- float The value of the log likelihood ratio. + """ # calculate = - 2 + up to a constant @@ -692,8 +752,7 @@ def total_loglr(self): sh_primary, hh_primary = self.primary_model.loglr self.primary_model.return_sh_hh = False # set logr, otherwise it will store (sh, hh) - setattr(self.primary_model._current_stats, 'loglr', - self.primary_model.marginalize_loglr(sh_primary, hh_primary)) + self.primary_model._current_stats.loglr = self.primary_model.marginalize_loglr(sh_primary, hh_primary) if isinstance(sh_primary, numpy.ndarray): nums = len(sh_primary) else: @@ -709,13 +768,12 @@ def total_loglr(self): # inclination has a very strong degeneracy, change of inclination # will change best match distance, so change the amplitude of # waveform. Using SNR will cancel out the effect of amplitude.err - i_max_extrinsic = numpy.argmax( - numpy.abs(sh_primary) / hh_primary**0.5) + i_max_extrinsic = numpy.argmax(numpy.abs(sh_primary) / hh_primary**0.5) for p in self.primary_model.marginalized_params_name: - if isinstance(self.primary_model.current_params[p], - numpy.ndarray): - margin_params[p] = \ - self.primary_model.current_params[p][i_max_extrinsic] + if isinstance(self.primary_model.current_params[p], numpy.ndarray): + margin_params[p] = self.primary_model.current_params[p][ + i_max_extrinsic + ] else: margin_params[p] = self.primary_model.current_params[p] else: @@ -737,8 +795,11 @@ def total_loglr(self): if not self.static_margin_params_in_other_models: for i in range(nums): current_params_other.update( - {key: value[i] if isinstance(value, numpy.ndarray) - else value for key, value in margin_params.items()}) + { + key: value[i] if isinstance(value, numpy.ndarray) else value + for key, value in margin_params.items() + } + ) other_model.update(**current_params_other) other_model.return_sh_hh = True sh_other, hh_other = other_model.loglr @@ -746,20 +807,21 @@ def total_loglr(self): hh_others[i] += hh_other other_model.return_sh_hh = False # set logr, otherwise it will store (sh, hh) - setattr(other_model._current_stats, 'loglr', - other_model.marginalize_loglr(sh_other, hh_other)) + other_model._current_stats.loglr = other_model.marginalize_loglr(sh_other, hh_other) else: # use one margin point set to approximate all the others current_params_other.update( - {key: value[0] if isinstance(value, numpy.ndarray) - else value for key, value in margin_params.items()}) + { + key: value[0] if isinstance(value, numpy.ndarray) else value + for key, value in margin_params.items() + } + ) other_model.update(**current_params_other) other_model.return_sh_hh = True sh_other, hh_other = other_model.loglr other_model.return_sh_hh = False # set logr, otherwise it will store (sh, hh) - setattr(other_model._current_stats, 'loglr', - other_model.marginalize_loglr(sh_other, hh_other)) + other_model._current_stats.loglr = other_model.marginalize_loglr(sh_other, hh_other) sh_others += sh_other hh_others += hh_other @@ -784,13 +846,17 @@ def others_lognl(self): return total_others_lognl def update_all_models(self, **params): - """This update method is also useful for loglr checking, + """ + This update method is also useful for loglr checking, the original update method in base module can't update - parameters in submodels correctly in loglr checking.""" + parameters in submodels correctly in loglr checking. + """ for lbl, model in self.submodels.items(): if self.param_map != {}: - p = {params.subname: self.current_params[params.fullname] - for params in self.param_map[lbl]} + p = { + params.subname: self.current_params[params.fullname] + for params in self.param_map[lbl] + } else: # dummy sampler doesn't have real variables, # which means self.param_map is {} @@ -802,8 +868,7 @@ def _loglikelihood(self): self.update_all_models() # calculate the combined loglikelihood - logl = self.total_loglr() + self.primary_model.lognl + \ - self.others_lognl() + logl = self.total_loglr() + self.primary_model.lognl + self.others_lognl() # store any extra stats from the submodels for lbl, model in self.submodels.items(): @@ -814,7 +879,8 @@ def _loglikelihood(self): @classmethod def from_config(cls, cp, **kwargs): - r"""Initializes an instance of this class from the given config file. + r""" + Initializes an instance of this class from the given config file. For more details, see `from_config` in `HierarchicalModel`. Parameters @@ -824,23 +890,24 @@ def from_config(cls, cp, **kwargs): \**kwargs : All additional keyword arguments are passed to the class. Any provided keyword will override what is in the config file. + """ # we need the read from config function from the init; to prevent # circular imports, we import it here from pycbc.inference.models import read_from_config + # get the submodels - kwargs['primary_lbl'] = shlex.split(cp.get('model', 'primary_model')) - kwargs['others_lbls'] = shlex.split(cp.get('model', 'other_models')) - submodel_lbls = kwargs['primary_lbl'] + kwargs['others_lbls'] + kwargs["primary_lbl"] = shlex.split(cp.get("model", "primary_model")) + kwargs["others_lbls"] = shlex.split(cp.get("model", "other_models")) + submodel_lbls = kwargs["primary_lbl"] + kwargs["others_lbls"] # sort parameters by model - vparam_map = map_params(hpiter(cp.options('variable_params'), - submodel_lbls)) - sparam_map = map_params(hpiter(cp.options('static_params'), - submodel_lbls)) + vparam_map = map_params(hpiter(cp.options("variable_params"), submodel_lbls)) + sparam_map = map_params(hpiter(cp.options("static_params"), submodel_lbls)) # get the acceleration label - kwargs['static_margin_params_in_other_models'] = shlex.split( - cp.get('model', 'static_margin_params_in_other_models')) + kwargs["static_margin_params_in_other_models"] = shlex.split( + cp.get("model", "static_margin_params_in_other_models") + ) # we'll need any waveform transforms for the initializing sub-models, # as the underlying models will receive the output of those transforms @@ -848,11 +915,11 @@ def from_config(cls, cp, **kwargs): # if `waveform_transforms` section doesn't have the prefix of # sub-model's name, then add this `waveform_transforms` section # into top level, if not, add it into sub-models' config - if any(cp.get_subsections('waveform_transforms')): + if any(cp.get_subsections("waveform_transforms")): waveform_transforms = transforms.read_transforms_from_config( - cp, 'waveform_transforms') - wfoutputs = set.union(*[t.outputs - for t in waveform_transforms]) + cp, "waveform_transforms" + ) + wfoutputs = set.union(*[t.outputs for t in waveform_transforms]) wfparam_map = map_params(hpiter(wfoutputs, submodel_lbls)) else: wfparam_map = {lbl: [] for lbl in submodel_lbls} @@ -867,30 +934,36 @@ def from_config(cls, cp, **kwargs): # include the [model] section for that model) copy_sections = [ HierarchicalParam(sec, submodel_lbls) - for sec in cp.sections() if lbl in - sec.split('-')[0].split(HierarchicalParam.delim, 1)[0]] + for sec in cp.sections() + if lbl in sec.split("-")[0].split(HierarchicalParam.delim, 1)[0] + ] for sec in copy_sections: # check that the user isn't trying to set variable or static # params for the model (we won't worry about waveform or # sampling transforms here, since that is checked for in the # __init__) - if sec.subname in ['variable_params', 'static_params']: - raise ValueError("Section {} found in the config file; " - "[variable_params] and [static_params] " - "sections should not include model " - "labels. To specify parameters unique to " - "one or more sub-models, prepend the " - "individual parameter names with the " - "model label. See HierarchicalParam for " - "details.".format(sec)) + if sec.subname in ["variable_params", "static_params"]: + raise ValueError( + f"Section {sec} found in the config file; " + "[variable_params] and [static_params] " + "sections should not include model " + "labels. To specify parameters unique to " + "one or more sub-models, prepend the " + "individual parameter names with the " + "model label. See HierarchicalParam for " + "details." + ) subcp.add_section(sec.subname) for opt, val in cp.items(sec): subcp.set(sec.subname, opt, val) # set the static params - subcp.add_section('static_params') + subcp.add_section("static_params") for param in sparam_map[lbl]: - subcp.set('static_params', param.subname, - cp.get('static_params', param.fullname)) + subcp.set( + "static_params", + param.subname, + cp.get("static_params", param.fullname), + ) # set the variable params: different from the standard # hierarchical model, in this JointPrimaryMarginalizedModel model, @@ -899,45 +972,47 @@ def from_config(cls, cp, **kwargs): # the primary model needs to do marginalization, so we must set # variable_params and prior section before initializing it. - subcp.add_section('variable_params') + subcp.add_section("variable_params") for param in vparam_map[lbl]: - if lbl in kwargs['primary_lbl']: + if lbl in kwargs["primary_lbl"]: # set variable_params for the primary model - subcp.set('variable_params', param.subname, - cp.get('variable_params', param.fullname)) + subcp.set( + "variable_params", + param.subname, + cp.get("variable_params", param.fullname), + ) else: # all variable_params in other models will come # from the primary model during sampling - subcp.set('static_params', param.subname, 'REPLACE') + subcp.set("static_params", param.subname, "REPLACE") for section in cp.sections(): # the primary model needs prior of marginlized parameters - if 'prior-' in section and lbl in kwargs['primary_lbl']: - prior_section = '%s' % section + if "prior-" in section and lbl in kwargs["primary_lbl"]: + prior_section = "%s" % section subcp[prior_section] = cp[prior_section] # similar to the standard hierarchical model, # add the outputs from the waveform transforms if sub-model # doesn't need marginalization - if lbl not in kwargs['primary_lbl']: + if lbl not in kwargs["primary_lbl"]: for param in wfparam_map[lbl]: - subcp.set('static_params', param.subname, 'REPLACE') + subcp.set("static_params", param.subname, "REPLACE") # save the vitual config file to disk for later check - with open('%s.ini' % lbl, 'w', encoding='utf-8') as file: + with open("%s.ini" % lbl, "w", encoding="utf-8") as file: subcp.write(file) # initialize submodel = read_from_config(subcp) - if lbl not in kwargs['primary_lbl']: + if lbl not in kwargs["primary_lbl"]: # similar to the standard hierarchical model, # move the static params back to variable if sub-model # doesn't need marginalization for p in vparam_map[lbl]: submodel.static_params.pop(p.subname) - submodel.variable_params = tuple(p.subname - for p in vparam_map[lbl]) + submodel.variable_params = tuple(p.subname for p in vparam_map[lbl]) # similar to the standard hierarchical model, # remove the waveform transform parameters if sub-model # doesn't need marginalization @@ -950,33 +1025,34 @@ def from_config(cls, cp, **kwargs): # `variable_params` and `prior` sections # here we ignore `coa_phase`, because if it's been marginalized, # it will not be listed in `variable_params` and `prior` sections - primary_model = submodels[kwargs['primary_lbl'][0]] + primary_model = submodels[kwargs["primary_lbl"][0]] marginalized_params = primary_model.marginalized_params_name.copy() for p in primary_model.static_params.keys(): - p_full = '%s__%s' % (kwargs['primary_lbl'][0], p) - if p_full not in cp['static_params']: - cp['static_params'][p_full] = "%s" % \ - primary_model.static_params[p] + p_full = "%s__%s" % (kwargs["primary_lbl"][0], p) + if p_full not in cp["static_params"]: + cp["static_params"][p_full] = "%s" % primary_model.static_params[p] for section in cp.sections(): - if 'prior-' in section: - p = section.split('-')[-1] + if "prior-" in section: + p = section.split("-")[-1] if p in marginalized_params: - cp['variable_params'].pop(p) + cp["variable_params"].pop(p) cp.pop(section) # save the vitual config file to disk for later check - with open('internal_top.ini', 'w', encoding='utf-8') as file: + with open("internal_top.ini", "w", encoding="utf-8") as file: cp.write(file) # now load the model logging.info("Loading joint_primary_marginalized model") return super(HierarchicalModel, cls).from_config( - cp, submodels=submodels, **kwargs) + cp, submodels=submodels, **kwargs + ) def reconstruct(self, rec=None, seed=None): - """ Reconstruct marginalized parameters by using the primary + """ + Reconstruct marginalized parameters by using the primary model's reconstruct method, total_loglr, and others_lognl. """ if seed: @@ -990,13 +1066,12 @@ def get_loglr(): # the top-level model if self.waveform_transforms is not None: self._current_params = transforms.apply_transforms( - self._current_params, self.waveform_transforms, - inverse=False) + self._current_params, self.waveform_transforms, inverse=False + ) self.update_all_models(**rec) return self.total_loglr() - rec = self.primary_model.reconstruct( - rec=rec, seed=seed, set_loglr=get_loglr) + rec = self.primary_model.reconstruct(rec=rec, seed=seed, set_loglr=get_loglr) # the primary model's reconstruct doesn't know lognl in other models - rec['loglikelihood'] += self.others_lognl() + rec["loglikelihood"] += self.others_lognl() return rec diff --git a/pycbc/inference/models/marginalized_gaussian_noise.py b/pycbc/inference/models/marginalized_gaussian_noise.py index f720a039a9b..07a7028cc3a 100644 --- a/pycbc/inference/models/marginalized_gaussian_noise.py +++ b/pycbc/inference/models/marginalized_gaussian_noise.py @@ -13,26 +13,33 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides model classes that assume the noise is Gaussian and +""" +This module provides model classes that assume the noise is Gaussian and allows for the likelihood to be marginalized over phase and/or time and/or distance. """ import itertools import logging + import numpy from scipy import special -from pycbc.waveform import generator from pycbc.detector import Detector -from .gaussian_noise import (BaseGaussianNoise, - create_waveform_generator, - GaussianNoise, catch_waveform_error) -from .tools import marginalize_likelihood, DistMarg +from pycbc.waveform import generator + +from .gaussian_noise import ( + BaseGaussianNoise, + GaussianNoise, + catch_waveform_error, + create_waveform_generator, +) +from .tools import DistMarg, marginalize_likelihood class MarginalizedPhaseGaussianNoise(GaussianNoise): - r"""The likelihood is analytically marginalized over phase. + r""" + The likelihood is analytically marginalized over phase. This class can be used with signal models that can be written as: @@ -110,47 +117,67 @@ class MarginalizedPhaseGaussianNoise(GaussianNoise): p(\Theta)\exp\left[\frac{1}{2}\sum_i\left( \left - \left \right)\right] """ - name = 'marginalized_phase' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - static_params=None, **kwargs): + name = "marginalized_phase" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + static_params=None, + **kwargs, + ): # set up the boiler-plate attributes - super(MarginalizedPhaseGaussianNoise, self).__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + super().__init__( + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) @property def _extra_stats(self): - """Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each - detector.""" - return ['loglr', 'maxl_phase'] + \ - ['{}_optimal_snrsq'.format(det) for det in self._data] + """ + Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each + detector. + """ + return ["loglr", "maxl_phase"] + [ + f"{det}_optimal_snrsq" for det in self._data + ] def _nowaveform_handler(self): - """Convenience function to set loglr values if no waveform generated. - """ - setattr(self._current_stats, 'loglikelihood', -numpy.inf) + """Convenience function to set loglr values if no waveform generated.""" + self._current_stats.loglikelihood = -numpy.inf # maxl phase doesn't exist, so set it to nan - setattr(self._current_stats, 'maxl_phase', numpy.nan) + self._current_stats.maxl_phase = numpy.nan for det in self._data: # snr can't be < 0 by definition, so return 0 - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), 0.) + setattr(self._current_stats, f"{det}_optimal_snrsq", 0.0) return -numpy.inf @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, .. math:: \log \mathcal{L}(\Theta) = I_0 \left(\left|\sum_i O(h^0_i, d_i)\right|\right) - \frac{1}{2}\left, at the current point in parameter space :math:`\Theta`. + Returns ------- float The value of the log likelihood ratio evaluated at the given point. + """ params = self.current_params if self.all_ifodata_same_rate_length: @@ -159,7 +186,7 @@ def _loglr(self): wfs = {} for det in self.data: wfs.update(self.waveform_generator[det].generate(**params)) - hh = 0. + hh = 0.0 hd = 0j for det, h in wfs.items(): # the kmax of the waveforms may be different than internal kmax @@ -167,19 +194,18 @@ def _loglr(self): if self._kmin[det] >= kmax: # if the waveform terminates before the filtering low frequency # cutoff, then the loglr is just 0 for this detector - hh_i = 0. + hh_i = 0.0 hd_i = 0j else: # whiten the waveform - h[self._kmin[det]:kmax] *= \ - self._weight[det][self._kmin[det]:kmax] + h[self._kmin[det] : kmax] *= self._weight[det][self._kmin[det] : kmax] # calculate inner products - hh_i = h[self._kmin[det]:kmax].inner( - h[self._kmin[det]:kmax]).real - hd_i = h[self._kmin[det]:kmax].inner( - self._whitened_data[det][self._kmin[det]:kmax]) + hh_i = h[self._kmin[det] : kmax].inner(h[self._kmin[det] : kmax]).real + hd_i = h[self._kmin[det] : kmax].inner( + self._whitened_data[det][self._kmin[det] : kmax] + ) # store - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), hh_i) + setattr(self._current_stats, f"{det}_optimal_snrsq", hh_i) hh += hh_i hd += hd_i self._current_stats.maxl_phase = numpy.angle(hd) @@ -187,72 +213,87 @@ def _loglr(self): class MarginalizedTime(DistMarg, BaseGaussianNoise): - r""" This likelihood numerically marginalizes over time + r""" + This likelihood numerically marginalizes over time This likelihood is optimized for marginalizing over time, but can also handle marginalization over polarization, phase (where appropriate), and sky location. The time series is interpolated using a quadratic apparoximation for sub-sample times. """ - name = 'marginalized_time' - def __init__(self, variable_params, - data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - sample_rate=None, - **kwargs): + name = "marginalized_time" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + sample_rate=None, + **kwargs, + ): # the flag used in `_loglr` self.return_sh_hh = False self.sample_rate = float(sample_rate) if sample_rate is not None else None self.kwargs = kwargs - variable_params, kwargs = self.setup_marginalization( - variable_params, - **kwargs) + variable_params, kwargs = self.setup_marginalization(variable_params, **kwargs) # set up the boiler-plate attributes - super(MarginalizedTime, self).__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - **kwargs) + super().__init__( + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + **kwargs, + ) # Determine if all data have the same sampling rate and segment length if self.all_ifodata_same_rate_length: # create a waveform generator for all ifos self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameTwoPolNoRespGenerator, - gates=self.gates, **kwargs['static_params']) + gates=self.gates, + **kwargs["static_params"], + ) else: # create a waveform generator for each ifo respectively self.waveform_generator = {} for det in self.data: self.waveform_generator[det] = create_waveform_generator( - self.variable_params, {det: self.data[det]}, + self.variable_params, + {det: self.data[det]}, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameTwoPolNoRespGenerator, - gates=self.gates, **kwargs['static_params']) + gates=self.gates, + **kwargs["static_params"], + ) self.dets = {} if sample_rate is not None: for ifo in self.data: if self.sample_rate < self.data[ifo].sample_rate: - raise ValueError("Model sample rate was set less than the" - " data. ") - logging.info("Using %s sample rate for marginalization", - sample_rate) + raise ValueError("Model sample rate was set less than the data. ") + logging.info("Using %s sample rate for marginalization", sample_rate) def _nowaveform_handler(self): - """Convenience function to set loglr values if no waveform generated. - """ + """Convenience function to set loglr values if no waveform generated.""" return -numpy.inf @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, or inner product and if `self.return_sh_hh` is True. .. math:: @@ -267,6 +308,7 @@ def _loglr(self): ------- float The value of the log likelihood ratio. + """ from pycbc.filter import matched_filter_core @@ -277,7 +319,7 @@ def _loglr(self): wfs = {} for det in self.data: wfs.update(self.waveform_generator[det].generate(**params)) - sh_total = hh_total = 0. + sh_total = hh_total = 0.0 snr_estimate = {} cplx_hpd = {} cplx_hcd = {} @@ -290,71 +332,69 @@ def _loglr(self): slc = slice(self._kmin[det], kmax) # whiten both polarizations - hp[self._kmin[det]:kmax] *= self._weight[det][slc] - hc[self._kmin[det]:kmax] *= self._weight[det][slc] + hp[self._kmin[det] : kmax] *= self._weight[det][slc] + hc[self._kmin[det] : kmax] *= self._weight[det][slc] # Use a higher sample rate if requested if self.sample_rate is not None: - tlen = int(round(self.sample_rate * - self.whitened_data[det].duration)) + tlen = int(round(self.sample_rate * self.whitened_data[det].duration)) flen = tlen // 2 + 1 else: flen = len(self._whitened_data[det]) - + hp.resize(flen) hc.resize(flen) self._whitened_data[det].resize(flen) cplx_hpd[det], _, _ = matched_filter_core( - hp, - self._whitened_data[det], - low_frequency_cutoff=self._f_lower[det], - high_frequency_cutoff=self._f_upper[det], - h_norm=1) + hp, + self._whitened_data[det], + low_frequency_cutoff=self._f_lower[det], + high_frequency_cutoff=self._f_upper[det], + h_norm=1, + ) cplx_hcd[det], _, _ = matched_filter_core( - hc, - self._whitened_data[det], - low_frequency_cutoff=self._f_lower[det], - high_frequency_cutoff=self._f_upper[det], - h_norm=1) + hc, + self._whitened_data[det], + low_frequency_cutoff=self._f_lower[det], + high_frequency_cutoff=self._f_upper[det], + h_norm=1, + ) hphp[det] = hp[slc].inner(hp[slc]).real hchc[det] = hc[slc].inner(hc[slc]).real hphc[det] = hp[slc].inner(hc[slc]).real - snr_proxy = ((cplx_hpd[det] / hphp[det] ** 0.5).squared_norm() + - (cplx_hcd[det] / hchc[det] ** 0.5).squared_norm()) + snr_proxy = (cplx_hpd[det] / hphp[det] ** 0.5).squared_norm() + ( + cplx_hcd[det] / hchc[det] ** 0.5 + ).squared_norm() snr_estimate[det] = (0.5 * snr_proxy) ** 0.5 self.draw_ifos(snr_estimate, log=False, **self.kwargs) self.snr_draw(snrs=snr_estimate) - - refframe = params.get('tc_ref_frame', 'geocentric') - ra = params['ra'] - dec = params['dec'] - ref_tc = params['tc'] + + refframe = params.get("tc_ref_frame", "geocentric") + ra = params["ra"] + dec = params["dec"] + ref_tc = params["tc"] for det in wfs: if det not in self.dets: self.dets[det] = Detector(det) tc = self.dets[det].arrival_time(ref_tc, ra, dec, refframe) if self.precalc_antenna_factors: fp, fc, dt = self.get_precalc_antenna_factors(det) - pol_phase = numpy.exp(-2.0j * params['polarization']) + pol_phase = numpy.exp(-2.0j * params["polarization"]) f = (fp + 1.0j * fc) * pol_phase fp = f.real fc = f.imag else: fp, fc = self.dets[det].antenna_pattern( - ra, dec, - params['polarization'], tc) + ra, dec, params["polarization"], tc + ) - cplx_hd = fp * cplx_hpd[det].at_time(tc, - interpolate='quadratic') - cplx_hd += fc * cplx_hcd[det].at_time(tc, - interpolate='quadratic') - hh = (fp * fp * hphp[det] + - fc * fc * hchc[det] + - 2.0 * fp * fc * hphc[det]) + cplx_hd = fp * cplx_hpd[det].at_time(tc, interpolate="quadratic") + cplx_hd += fc * cplx_hcd[det].at_time(tc, interpolate="quadratic") + hh = fp * fp * hphp[det] + fc * fc * hchc[det] + 2.0 * fp * fc * hphc[det] sh_total += cplx_hd hh_total += hh @@ -368,7 +408,8 @@ def _loglr(self): class MarginalizedPolarization(DistMarg, BaseGaussianNoise): - r""" This likelihood numerically marginalizes over polarization angle + r""" + This likelihood numerically marginalizes over polarization angle This class implements the Gaussian likelihood with an explicit numerical marginalization over polarization angle. This is accomplished using @@ -377,67 +418,87 @@ class MarginalizedPolarization(DistMarg, BaseGaussianNoise): The 'polarization_samples' argument can be passed to set an alternate number of integration points. """ - name = 'marginalized_polarization' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - polarization_samples=1000, - **kwargs): + name = "marginalized_polarization" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + polarization_samples=1000, + **kwargs, + ): variable_params, kwargs = self.setup_marginalization( - variable_params, - polarization_samples=polarization_samples, - **kwargs) + variable_params, polarization_samples=polarization_samples, **kwargs + ) # set up the boiler-plate attributes - super(MarginalizedPolarization, self).__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - **kwargs) + super().__init__( + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + **kwargs, + ) # Determine if all data have the same sampling rate and segment length if self.all_ifodata_same_rate_length: # create a waveform generator for all ifos self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameTwoPolGenerator, - gates=self.gates, **kwargs['static_params']) + gates=self.gates, + **kwargs["static_params"], + ) else: # create a waveform generator for each ifo respectively self.waveform_generator = {} for det in self.data: self.waveform_generator[det] = create_waveform_generator( - self.variable_params, {det: self.data[det]}, + self.variable_params, + {det: self.data[det]}, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameTwoPolGenerator, - gates=self.gates, **kwargs['static_params']) + gates=self.gates, + **kwargs["static_params"], + ) self.dets = {} @property def _extra_stats(self): - """Adds ``loglr``, ``maxl_polarization``, and the ``optimal_snrsq`` in + """ + Adds ``loglr``, ``maxl_polarization``, and the ``optimal_snrsq`` in each detector. """ - return ['loglr', 'maxl_polarization', 'maxl_loglr'] + \ - ['{}_optimal_snrsq'.format(det) for det in self._data] + return ["loglr", "maxl_polarization", "maxl_loglr"] + [ + f"{det}_optimal_snrsq" for det in self._data + ] def _nowaveform_handler(self): - """Convenience function to set loglr values if no waveform generated. - """ - setattr(self._current_stats, 'loglr', -numpy.inf) + """Convenience function to set loglr values if no waveform generated.""" + self._current_stats.loglr = -numpy.inf # maxl phase doesn't exist, so set it to nan - setattr(self._current_stats, 'maxl_polarization', numpy.nan) + self._current_stats.maxl_polarization = numpy.nan for det in self._data: # snr can't be < 0 by definition, so return 0 - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), 0.) + setattr(self._current_stats, f"{det}_optimal_snrsq", 0.0) return -numpy.inf @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, .. math:: @@ -451,6 +512,7 @@ def _loglr(self): ------- float The value of the log likelihood ratio. + """ params = self.current_params if self.all_ifodata_same_rate_length: @@ -460,25 +522,24 @@ def _loglr(self): for det in self.data: wfs.update(self.waveform_generator[det].generate(**params)) - lr = sh_total = hh_total = 0. - refframe = params.get('tc_ref_frame', 'geocentric') - ra = params['ra'] - dec = params['dec'] - ref_tc = params['tc'] + lr = sh_total = hh_total = 0.0 + refframe = params.get("tc_ref_frame", "geocentric") + ra = params["ra"] + dec = params["dec"] + ref_tc = params["tc"] for det, (hp, hc) in wfs.items(): if det not in self.dets: self.dets[det] = Detector(det) tc = self.dets[det].arrival_time(ref_tc, ra, dec, refframe) - fp, fc = self.dets[det].antenna_pattern(ra, dec, - params['polarization'], tc) + fp, fc = self.dets[det].antenna_pattern(ra, dec, params["polarization"], tc) # the kmax of the waveforms may be different than internal kmax kmax = min(max(len(hp), len(hc)), self._kmax[det]) slc = slice(self._kmin[det], kmax) # whiten both polarizations - hp[self._kmin[det]:kmax] *= self._weight[det][slc] - hc[self._kmin[det]:kmax] *= self._weight[det][slc] + hp[self._kmin[det] : kmax] *= self._weight[det][slc] + hc[self._kmin[det] : kmax] *= self._weight[det][slc] # h = fp * hp + hc * hc # = fp * + fc * @@ -501,30 +562,27 @@ def _loglr(self): hh = fp * fp * hphp + fc * fc * hchc + fp * fc * (hphc + hchp) # store - setattr(self._current_stats, '{}_optimal_snrsq'.format(det), hh) + setattr(self._current_stats, f"{det}_optimal_snrsq", hh) sh_total += cplx_hd hh_total += hh - lr, idx, maxl = self.marginalize_loglr(sh_total, hh_total, - return_peak=True) + lr, idx, maxl = self.marginalize_loglr(sh_total, hh_total, return_peak=True) # store the maxl polarization - setattr(self._current_stats, - 'maxl_polarization', - params['polarization'][idx]) - setattr(self._current_stats, 'maxl_loglr', maxl) + self._current_stats.maxl_polarization = params["polarization"][idx] + self._current_stats.maxl_loglr = maxl # just store the maxl optimal snrsq for det in wfs: - p = '{}_optimal_snrsq'.format(det) - setattr(self._current_stats, p, - getattr(self._current_stats, p)[idx]) + p = f"{det}_optimal_snrsq" + setattr(self._current_stats, p, getattr(self._current_stats, p)[idx]) return lr class MarginalizedHMPolPhase(BaseGaussianNoise): - r"""Numerically marginalizes waveforms with higher modes over polarization + r""" + Numerically marginalizes waveforms with higher modes over polarization `and` phase. This class implements the Gaussian likelihood with an explicit numerical @@ -577,27 +635,45 @@ class MarginalizedHMPolPhase(BaseGaussianNoise): `. """ - name = 'marginalized_hmpolphase' - def __init__(self, variable_params, data, low_frequency_cutoff, psds=None, - high_frequency_cutoff=None, normalize=False, - polarization_samples=100, - coa_phase_samples=100, - static_params=None, **kwargs): + name = "marginalized_hmpolphase" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + psds=None, + high_frequency_cutoff=None, + normalize=False, + polarization_samples=100, + coa_phase_samples=100, + static_params=None, + **kwargs, + ): # set up the boiler-plate attributes - super(MarginalizedHMPolPhase, self).__init__( - variable_params, data, low_frequency_cutoff, psds=psds, - high_frequency_cutoff=high_frequency_cutoff, normalize=normalize, - static_params=static_params, **kwargs) + super().__init__( + variable_params, + data, + low_frequency_cutoff, + psds=psds, + high_frequency_cutoff=high_frequency_cutoff, + normalize=normalize, + static_params=static_params, + **kwargs, + ) # create the waveform generator self.waveform_generator = create_waveform_generator( - self.variable_params, self.data, + self.variable_params, + self.data, waveform_transforms=self.waveform_transforms, recalibration=self.recalibration, generator_class=generator.FDomainDetFrameModesGenerator, - gates=self.gates, **self.static_params) - pol = numpy.linspace(0, 2*numpy.pi, polarization_samples) - phase = numpy.linspace(0, 2*numpy.pi, coa_phase_samples) + gates=self.gates, + **self.static_params, + ) + pol = numpy.linspace(0, 2 * numpy.pi, polarization_samples) + phase = numpy.linspace(0, 2 * numpy.pi, coa_phase_samples) # remap to every combination of the parameters # this gets every combination by mappin them to an NxM grid # one needs to be transposed so that they run allong opposite @@ -622,21 +698,23 @@ def phase_fac(self, m): @property def _extra_stats(self): - """Adds ``maxl_polarization`` and the ``maxl_phase`` - """ - return ['maxl_polarization', 'maxl_phase', ] + """Adds ``maxl_polarization`` and the ``maxl_phase``""" + return [ + "maxl_polarization", + "maxl_phase", + ] def _nowaveform_handler(self): - """Convenience function to set loglr values if no waveform generated. - """ + """Convenience function to set loglr values if no waveform generated.""" # maxl phase doesn't exist, so set it to nan - setattr(self._current_stats, 'maxl_polarization', numpy.nan) - setattr(self._current_stats, 'maxl_phase', numpy.nan) + self._current_stats.maxl_polarization = numpy.nan + self._current_stats.maxl_phase = numpy.nan return -numpy.inf @catch_waveform_error def _loglr(self, return_unmarginalized=False): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, .. math:: @@ -650,6 +728,7 @@ def _loglr(self, return_unmarginalized=False): ------- float The value of the log likelihood ratio. + """ params = self.current_params wfs = self.waveform_generator.generate(**params) @@ -660,13 +739,13 @@ def _loglr(self, return_unmarginalized=False): # * fp/fc need not be calculated except where polarization is different # * may be possible to simplify this by making smarter use of real/imag # --------------------------------------------------------------------- - lr = 0. + lr = 0.0 hds = {} hhs = {} - refframe = params.get('tc_ref_frame', 'geocentric') - ra = params['ra'] - dec = params['dec'] - ref_tc = params['tc'] + refframe = params.get("tc_ref_frame", "geocentric") + ra = params["ra"] + dec = params["dec"] + ref_tc = params["tc"] for det, modes in wfs.items(): if det not in self.dets: self.dets[det] = Detector(det) @@ -687,8 +766,8 @@ def _loglr(self, return_unmarginalized=False): # the kmax of the waveforms may be different than internal kmax kmax = min(max(len(ulm), len(vlm)), self._kmax[det]) slc = slice(self._kmin[det], kmax) - ulm[self._kmin[det]:kmax] *= self._weight[det][slc] - vlm[self._kmin[det]:kmax] *= self._weight[det][slc] + ulm[self._kmin[det] : kmax] *= self._weight[det][slc] + vlm[self._kmin[det] : kmax] *= self._weight[det][slc] # the inner products # @@ -698,12 +777,14 @@ def _loglr(self, return_unmarginalized=False): # add inclination, and pack into a complex number import lal + glm = lal.SpinWeightedSphericalHarmonic( - params['inclination'], 0, -2, l, m).real + params["inclination"], 0, -2, l, m + ).real if m not in zetas: zetas[m] = 0j - zetas[m] += glm * (ulmd + 1j*vlmd) + zetas[m] += glm * (ulmd + 1j * vlmd) # Get condense set of the parts of the waveform that only diff # by m, this is used next to help calculate @@ -738,11 +819,11 @@ def _loglr(self, return_unmarginalized=False): rs_m[m, mprime] = sr_m[mprime, m] sr_m[m, mprime] = rs_m[mprime, m] # now apply the phase to all the common ms - hpd = 0. - hcd = 0. - hphp = 0. - hchc = 0. - hphc = 0. + hpd = 0.0 + hcd = 0.0 + hphp = 0.0 + hchc = 0.0 + hphc = 0.0 for m, zeta in zetas.items(): phase_coeff = self.phase_fac(m) @@ -767,20 +848,26 @@ def _loglr(self, return_unmarginalized=False): rs = rs_m[m, mprime] sr = sr_m[m, mprime] # - hphp += rr * cosm * cosmprime \ - + ss * sinm * sinmprime \ - - rs * cosm * sinmprime \ + hphp += ( + rr * cosm * cosmprime + + ss * sinm * sinmprime + - rs * cosm * sinmprime - sr * sinm * cosmprime + ) # - hchc += rr * sinm * sinmprime \ - + ss * cosm * cosmprime \ - + rs * sinm * cosmprime \ + hchc += ( + rr * sinm * sinmprime + + ss * cosm * cosmprime + + rs * sinm * cosmprime + sr * cosm * sinmprime + ) # - hphc += -rr * cosm * sinmprime \ - + ss * sinm * cosmprime \ - + sr * sinm * sinmprime \ + hphc += ( + -rr * cosm * sinmprime + + ss * sinm * cosmprime + + sr * sinm * sinmprime - rs * cosm * cosmprime + ) # Now apply the polarizations and calculate the loglr # We have h = Fp * hp + Fc * hc @@ -803,6 +890,6 @@ def _loglr(self, return_unmarginalized=False): # store the maxl values idx = lr.argmax() - setattr(self._current_stats, 'maxl_polarization', self.pol[idx]) - setattr(self._current_stats, 'maxl_phase', self.phase[idx]) + self._current_stats.maxl_polarization = self.pol[idx] + self._current_stats.maxl_phase = self.phase[idx] return float(lr_total) diff --git a/pycbc/inference/models/relbin.py b/pycbc/inference/models/relbin.py index 37d41fa354c..f781b47ee2a 100644 --- a/pycbc/inference/models/relbin.py +++ b/pycbc/inference/models/relbin.py @@ -21,40 +21,56 @@ # # ============================================================================= # -"""This module provides model classes and functions for implementing +""" +This module provides model classes and functions for implementing a relative binning likelihood for parameter estimation. """ - +import itertools import logging + import numpy -import itertools from scipy.interpolate import interp1d -from pycbc.waveform import (get_fd_waveform_sequence, - get_fd_det_waveform_sequence, fd_det_sequence) from pycbc.detector import Detector from pycbc.types import Array, TimeSeries - -from .gaussian_noise import (BaseGaussianNoise, catch_waveform_error) -from pycbc.waveform import FailedWaveformError -from .relbin_cpu import (likelihood_parts, likelihood_parts_v, - likelihood_parts_multi, likelihood_parts_multi_v, - likelihood_parts_det, likelihood_parts_det_multi, - likelihood_parts_vector, - likelihood_parts_v_pol, - likelihood_parts_v_time, - likelihood_parts_v_pol_time, - likelihood_parts_vectorp, snr_predictor, - likelihood_parts_vectort, - snr_predictor_dom) +from pycbc.waveform import ( + FailedWaveformError, + fd_det_sequence, + get_fd_det_waveform_sequence, + get_fd_waveform_sequence, +) + +from .gaussian_noise import BaseGaussianNoise, catch_waveform_error +from .relbin_cpu import ( + likelihood_parts, + likelihood_parts_det, + likelihood_parts_det_multi, + likelihood_parts_multi, + likelihood_parts_multi_v, + likelihood_parts_v, + likelihood_parts_v_pol, + likelihood_parts_v_pol_time, + likelihood_parts_v_time, + likelihood_parts_vector, + likelihood_parts_vectorp, + likelihood_parts_vectort, + snr_predictor, + snr_predictor_dom, +) from .tools import DistMarg -def setup_bins(f_full, f_lo, f_hi, chi=1.0, - eps=0.1, gammas=None, - ): - """Construct frequency bins for use in a relative likelihood +def setup_bins( + f_full, + f_lo, + f_hi, + chi=1.0, + eps=0.1, + gammas=None, +): + """ + Construct frequency bins for use in a relative likelihood model. For details, see [Barak, Dai & Venumadhav 2018]. Parameters @@ -83,6 +99,7 @@ def setup_bins(f_full, f_lo, f_hi, chi=1.0, Bin edge frequencies. fbin_ind : numpy.array of ints Indices of bin edges in full frequency array. + """ f = numpy.linspace(f_lo, f_hi, 10000) # f^ga power law index @@ -92,17 +109,15 @@ def setup_bins(f_full, f_lo, f_hi, chi=1.0, else numpy.array([-5.0 / 3, -2.0 / 3, 1.0, 5.0 / 3, 7.0 / 3]) ) logging.info("Using powerlaw indices: %s", ga) - dalp = chi * 2.0 * numpy.pi / numpy.absolute((f_lo ** ga) - (f_hi ** ga)) + dalp = chi * 2.0 * numpy.pi / numpy.absolute((f_lo**ga) - (f_hi**ga)) dphi = numpy.sum( - numpy.array([numpy.sign(g) * d * (f ** g) for g, d in zip(ga, dalp)]), + numpy.array([numpy.sign(g) * d * (f**g) for g, d in zip(ga, dalp)]), axis=0, ) dphi_diff = dphi - dphi[0] # now construct frequency bins nbin = int(dphi_diff[-1] / eps) - dphi2f = interp1d( - dphi_diff, f, kind="slinear", bounds_error=False, fill_value=0.0 - ) + dphi2f = interp1d(dphi_diff, f, kind="slinear", bounds_error=False, fill_value=0.0) dphi_grid = numpy.linspace(dphi_diff[0], dphi_diff[-1], nbin + 1) # frequency grid points fbin = dphi2f(dphi_grid) @@ -115,7 +130,7 @@ def setup_bins(f_full, f_lo, f_hi, chi=1.0, curr_idx = len(f_full) - 1 else: abs1 = abs(f_full[idx_f_full] - fbin[idx_fbin]) - abs2 = abs(f_full[idx_f_full-1] - fbin[idx_fbin]) + abs2 = abs(f_full[idx_f_full - 1] - fbin[idx_fbin]) if abs1 > abs2: curr_idx = idx_f_full - 1 else: @@ -126,7 +141,8 @@ def setup_bins(f_full, f_lo, f_hi, chi=1.0, class Relative(DistMarg, BaseGaussianNoise): - r"""Model that assumes the likelihood in a region around the peak + r""" + Model that assumes the likelihood in a region around the peak is slowly varying such that a linear approximation can be made, and likelihoods can be calculated at a coarser frequency resolution. For more details on the implementation, see https://arxiv.org/abs/1806.08792. @@ -168,7 +184,9 @@ class Relative(DistMarg, BaseGaussianNoise): \**kwargs : All other keyword arguments are passed to :py:class:`BaseGaussianNoise`. + """ + name = "relative" def __init__( @@ -182,22 +200,21 @@ def __init__( earth_rotation=False, earth_rotation_mode=2, marginalize_phase=True, - **kwargs + **kwargs, ): variable_params, kwargs = self.setup_marginalization( - variable_params, - marginalize_phase=marginalize_phase, - **kwargs) + variable_params, marginalize_phase=marginalize_phase, **kwargs + ) - super(Relative, self).__init__( + super().__init__( variable_params, data, low_frequency_cutoff, **kwargs ) # If the waveform needs us to apply the detector response, # set flag to true (most cases for ground-based observatories). self.still_needs_det_response = False - if self.static_params['approximant'] in fd_det_sequence: + if self.static_params["approximant"] in fd_det_sequence: self.still_needs_det_response = True # reference waveform and bin edges @@ -217,8 +234,8 @@ def __init__( self.return_sh_hh = False for k in self.static_params: - if self.fid_params[k] == 'REPLACE': - self.fid_params.pop(k) + if self.fid_params[k] == "REPLACE": + self.fid_params.pop(k) for ifo in data: # store data and frequencies @@ -232,22 +249,25 @@ def __init__( f_hi = self.kmax[ifo] * self.df[ifo] logging.info( "%s: Generating fiducial waveform from %s to %s Hz", - ifo, f_lo, f_hi, + ifo, + f_lo, + f_hi, ) # prune low frequency samples to avoid waveform errors fpoints = Array(self.f[ifo].astype(numpy.float64)) - fpoints = fpoints[self.kmin[ifo]:self.kmax[ifo]+1] + fpoints = fpoints[self.kmin[ifo] : self.kmax[ifo] + 1] if self.still_needs_det_response: - wave = get_fd_det_waveform_sequence(ifos=ifo, - sample_points=fpoints, - **self.fid_params) + wave = get_fd_det_waveform_sequence( + ifos=ifo, sample_points=fpoints, **self.fid_params + ) curr_wav = wave[ifo] - self.ta[ifo] = 0. + self.ta[ifo] = 0.0 else: - fid_hp, fid_hc = get_fd_waveform_sequence(sample_points=fpoints, - **self.fid_params) + fid_hp, fid_hc = get_fd_waveform_sequence( + sample_points=fpoints, **self.fid_params + ) # Apply detector response if not handled by # the waveform generator self.det[ifo] = Detector(ifo) @@ -258,9 +278,12 @@ def __init__( ) self.ta[ifo] = self.fid_params["tc"] + dt fp, fc = self.det[ifo].antenna_pattern( - self.fid_params["ra"], self.fid_params["dec"], - self.fid_params["polarization"], self.fid_params["tc"]) - curr_wav = (fid_hp * fp + fid_hc * fc) + self.fid_params["ra"], + self.fid_params["dec"], + self.fid_params["polarization"], + self.fid_params["tc"], + ) + curr_wav = fid_hp * fp + fid_hc * fc # check for zeros at low and high frequencies # make sure only nonzero samples are included in bins @@ -271,7 +294,9 @@ def __init__( logging.info( "WARNING! Fiducial waveform starts above " "low-frequency-cutoff, initial bin frequency " - "will be %s Hz", f_lo) + "will be %s Hz", + f_lo, + ) numzeros_hi = list(curr_wav[::-1] != 0j).index(True) if numzeros_hi > 0: new_kmax = self.kmax[ifo] - numzeros_hi @@ -279,7 +304,9 @@ def __init__( logging.info( "WARNING! Fiducial waveform terminates below " "high-frequency-cutoff, final bin frequency " - "will be %s Hz", f_hi) + "will be %s Hz", + f_hi, + ) self.ta[ifo] -= self.end_time[ifo] curr_wav.resize(len(self.f[ifo])) @@ -288,13 +315,16 @@ def __init__( # We'll apply this to the data, in lieu of the ref waveform # This makes it easier to compare target signal to reference later tshift = numpy.exp(-2.0j * numpy.pi * self.f[ifo] * self.ta[ifo]) - self.h00[ifo] = numpy.array(curr_wav) # * tshift + self.h00[ifo] = numpy.array(curr_wav) # * tshift data_shifted = self.data[ifo] * numpy.conjugate(tshift) logging.info("Computing frequency bins") fbin_ind = setup_bins( - f_full=self.f[ifo], f_lo=f_lo, f_hi=f_hi, - gammas=gammas, eps=float(epsilon), + f_full=self.f[ifo], + f_lo=f_lo, + f_hi=f_hi, + gammas=gammas, + eps=float(epsilon), ) logging.info("Using %s bins for this model", len(fbin_ind)) @@ -302,17 +332,13 @@ def __init__( self.edges[ifo] = fbin_ind self.init_from_frequencies(data_shifted, self.h00, fbin_ind, ifo) self.antenna_time[ifo] = self.setup_antenna( - earth_rotation, - int(earth_rotation_mode), - self.fedges[ifo]) + earth_rotation, int(earth_rotation_mode), self.fedges[ifo] + ) self.combine_layout() def init_from_frequencies(self, data, h00, fbin_ind, ifo): bins = numpy.array( - [ - (fbin_ind[i], fbin_ind[i + 1]) - for i in range(len(fbin_ind) - 1) - ] + [(fbin_ind[i], fbin_ind[i + 1]) for i in range(len(fbin_ind) - 1)] ) # store low res copy of fiducial waveform @@ -378,59 +404,57 @@ def likelihood_function(self): if self.marginalize_vector_params: p = self.current_params - vmarg = set(k for k in self.marginalize_vector_params - if not numpy.isscalar(p[k])) + vmarg = set( + k for k in self.marginalize_vector_params if not numpy.isscalar(p[k]) + ) if self.earth_rotation: - if set(['tc', 'polarization']).issubset(vmarg): - self.lformat = 'earth_time_pol' + if set(["tc", "polarization"]).issubset(vmarg): + self.lformat = "earth_time_pol" return likelihood_parts_v_pol_time - elif set(['polarization']).issubset(vmarg): - self.lformat = 'earth_pol' + if set(["polarization"]).issubset(vmarg): + self.lformat = "earth_pol" return likelihood_parts_v_pol - elif set(['tc']).issubset(vmarg): - self.lformat = 'earth_time' + if set(["tc"]).issubset(vmarg): + self.lformat = "earth_time" return likelihood_parts_v_time - else: - if set(['ra', 'dec', 'tc']).issubset(vmarg): - return likelihood_parts_vector - elif set(['tc', 'polarization']).issubset(vmarg): - return likelihood_parts_vector - elif set(['tc']).issubset(vmarg): - return likelihood_parts_vectort - elif set(['polarization']).issubset(vmarg): - return likelihood_parts_vectorp + elif set(["ra", "dec", "tc"]).issubset(vmarg) or set(["tc", "polarization"]).issubset(vmarg): + return likelihood_parts_vector + elif set(["tc"]).issubset(vmarg): + return likelihood_parts_vectort + elif set(["polarization"]).issubset(vmarg): + return likelihood_parts_vectorp return self.lik def summary_product(self, h1, h2, bins, ifo): - """ Calculate the summary values for the inner product - """ + """Calculate the summary values for the inner product """ # calculate coefficients h12 = numpy.conjugate(h1) * h2 / self.psds[ifo] # constant terms - a0 = numpy.array([ - 4.0 * self.df[ifo] * h12[l:h].sum() - for l, h in bins - ]) + a0 = numpy.array([4.0 * self.df[ifo] * h12[l:h].sum() for l, h in bins]) # linear terms - a1 = numpy.array([ - 4.0 / (h - l) * - (h12[l:h] * (self.f[ifo][l:h] - self.f[ifo][l])).sum() - for l, h in bins]) + a1 = numpy.array( + [ + 4.0 / (h - l) * (h12[l:h] * (self.f[ifo][l:h] - self.f[ifo][l])).sum() + for l, h in bins + ] + ) return a0, a1 def get_waveforms(self, params): - """ Get the waveform polarizations for each ifo - """ + """Get the waveform polarizations for each ifo""" if self.still_needs_det_response: wfs = {} for ifo in self.data: - wfs.update(get_fd_det_waveform_sequence( - ifos=ifo, sample_points=self.fedges[ifo], **params)) + wfs.update( + get_fd_det_waveform_sequence( + ifos=ifo, sample_points=self.fedges[ifo], **params + ) + ) return wfs wfs = [] @@ -446,21 +470,25 @@ def get_waveforms(self, params): @property def multi_signal_support(self): - """ The list of classes that this model supports in a multi-signal + """ + The list of classes that this model supports in a multi-signal likelihood """ # Check if this model *can* be included in a multi-signal model. # All marginalizations must currently be disabled to work! - if (self.marginalize_vector_params or - self.marginalize_distance or - self.marginalize_phase): - logging.info("Cannot use single template model inside of" - "multi_signal if marginalizations are enabled") + if ( + self.marginalize_vector_params + or self.marginalize_distance + or self.marginalize_phase + ): + logging.info( + "Cannot use single template model inside of" + "multi_signal if marginalizations are enabled" + ) return [type(self)] def calculate_hihjs(self, models): - """ Pre-calculate the hihj inner products on a grid - """ + """Pre-calculate the hihj inner products on a grid""" self.hihj = {} for m1, m2 in itertools.combinations(models, 2): self.hihj[(m1, m2)] = {} @@ -476,23 +504,21 @@ def calculate_hihjs(self, models): edge = edge[keep] fedge = m1.f[ifo][edge] - bins = numpy.array([ - (edge[i], edge[i + 1]) - for i in range(len(edge) - 1) - ]) + bins = numpy.array( + [(edge[i], edge[i + 1]) for i in range(len(edge) - 1)] + ) a0, a1 = self.summary_product(h1, h2, bins, ifo) self.hihj[(m1, m2)][ifo] = a0, a1, fedge def multi_loglikelihood(self, models): - """ Calculate a multi-model (signal) likelihood - """ + """Calculate a multi-model (signal) likelihood""" models = [self] + models loglr = 0 # handle sum[ - 0.5 ] for m in models: loglr += m.loglr - if not hasattr(self, 'hihj'): + if not hasattr(self, "hihj"): self.calculate_hihjs(models) if self.still_needs_det_response: @@ -503,11 +529,10 @@ def multi_loglikelihood(self, models): dtc, channel, h00 = m1._current_wf_parts[det] dtc2, channel2, h002 = m2._current_wf_parts[det] - c1c2 = self.mlik(fedge, - dtc, channel, h00, - dtc2, channel2, h002, - a0, a1) - loglr += - c1c2.real # This is -0.5 * re( + ) + c1c2 = self.mlik( + fedge, dtc, channel, h00, dtc2, channel2, h002, a0, a1 + ) + loglr += -c1c2.real # This is -0.5 * re( + ) else: # finally add in the lognl term from this model for m1, m2 in itertools.combinations(models, 2): @@ -517,16 +542,30 @@ def multi_loglikelihood(self, models): fp, fc, dtc, hp, hc, h00 = m1._current_wf_parts[det] fp2, fc2, dtc2, hp2, hc2, h002 = m2._current_wf_parts[det] - h1h2 = self.mlik(fedge, - fp, fc, dtc, hp, hc, h00, - fp2, fc2, dtc2, hp2, hc2, h002, - a0, a1) - loglr += - h1h2.real # This is -0.5 * re( + ) + h1h2 = self.mlik( + fedge, + fp, + fc, + dtc, + hp, + hc, + h00, + fp2, + fc2, + dtc2, + hp2, + hc2, + h002, + a0, + a1, + ) + loglr += -h1h2.real # This is -0.5 * re( + ) return loglr + self.lognl @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, or inner product and if `self.return_sh_hh` is True. .. math:: @@ -544,6 +583,7 @@ def _loglr(self): or tuple The inner product (, ). + """ # get model params p = self.current_params @@ -552,7 +592,7 @@ def _loglr(self): norm = 0.0 filt = 0j self._current_wf_parts = {} - pol_phase = numpy.exp(-2.0j * p['polarization']) + pol_phase = numpy.exp(-2.0j * p["polarization"]) for ifo in self.data: freqs = self.fedges[ifo] @@ -565,34 +605,59 @@ def _loglr(self): # with detector response. Otherwise, skip detector response. if self.still_needs_det_response: - dtc = 0. + dtc = 0.0 channel = wfs[ifo].numpy() - filter_i, norm_i = lik(freqs, dtc, channel, h00, - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) + filter_i, norm_i = lik( + freqs, + dtc, + channel, + h00, + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) self._current_wf_parts[ifo] = (dtc, channel, h00) else: hp, hc = wfs[ifo] det = self.det[ifo] - fp, fc = det.antenna_pattern(p["ra"], p["dec"], - 0.0, times) + fp, fc = det.antenna_pattern(p["ra"], p["dec"], 0.0, times) dt = det.time_delay_from_earth_center(p["ra"], p["dec"], times) dtc = p["tc"] + dt - end_time - self.ta[ifo] - if self.lformat == 'earth_pol': - filter_i, norm_i = lik(freqs, fp, fc, dtc, pol_phase, - hp, hc, h00, - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) + if self.lformat == "earth_pol": + filter_i, norm_i = lik( + freqs, + fp, + fc, + dtc, + pol_phase, + hp, + hc, + h00, + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) else: f = (fp + 1.0j * fc) * pol_phase fp = f.real.copy() fc = f.imag.copy() - filter_i, norm_i = lik(freqs, fp, fc, dtc, - hp, hc, h00, - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) + filter_i, norm_i = lik( + freqs, + fp, + fc, + dtc, + hp, + hc, + h00, + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) self._current_wf_parts[ifo] = (fp, fc, dtc, hp, hc, h00) filt += filter_i @@ -606,18 +671,22 @@ def _loglr(self): return results def _nowaveform_handler(self): - """Returns -inf for loglr if no waveform generated. + """ + Returns -inf for loglr if no waveform generated. If `return_sh_hh` is set to True, a FailedWaveformError will be raised. """ if self.return_sh_hh: - raise FailedWaveformError("Waveform failed to generate and " - "return_sh_hh set to True! I don't know " - "what to return in this case.") + raise FailedWaveformError( + "Waveform failed to generate and " + "return_sh_hh set to True! I don't know " + "what to return in this case." + ) return -numpy.inf def write_metadata(self, fp, group=None): - """Adds writing the fiducial parameters and epsilon to file's attrs. + """ + Adds writing the fiducial parameters and epsilon to file's attrs. Parameters ---------- @@ -627,6 +696,7 @@ def write_metadata(self, fp, group=None): If provided, the metadata will be written to the attrs specified by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is written to the top-level attrs (``fp.attrs``). + """ super().write_metadata(fp, group=group) if group is None: @@ -634,17 +704,18 @@ def write_metadata(self, fp, group=None): else: attrs = fp[group].attrs for p, v in self.fid_params.items(): - attrs["{}_ref".format(p)] = v + attrs[f"{p}_ref"] = v def max_curvature_from_reference(self): - """ Return the maximum change in slope between frequency bins + """ + Return the maximum change in slope between frequency bins relative to the reference waveform. """ dmax = 0 for ifo in self.data: r = self.wf_ret[ifo][0] / self.h00_sparse[ifo] d = abs(numpy.diff(r / abs(r).min(), n=2)).max() - dmax = d if dmax < d else dmax + dmax = max(dmax, d) return dmax @staticmethod @@ -681,59 +752,66 @@ def extra_args_from_config(cp, section, skip_args=None, dtypes=None): "inclination": 0.0, "polarization": numpy.pi, } - fid_params.update( - {p: opt_params[p] for p in opt_params if p not in fid_params} - ) + fid_params.update({p: opt_params[p] for p in opt_params if p not in fid_params}) args.update({"fiducial_params": fid_params, "gammas": gammas}) return args class RelativeTime(Relative): - """ Heterodyne likelihood optimized for time marginalization. In addition + """ + Heterodyne likelihood optimized for time marginalization. In addition it supports phase (dominant-mode), sky location, and polarization marginalization. """ + name = "relative_time" - def __init__(self, *args, - sample_rate=4096, - **kwargs): - super(RelativeTime, self).__init__(*args, **kwargs) + def __init__(self, *args, sample_rate=4096, **kwargs): + super().__init__(*args, **kwargs) self.sample_rate = float(sample_rate) self.setup_peak_lock(sample_rate=self.sample_rate, **kwargs) self.draw_ifos(self.ref_snr, **kwargs) @property def ref_snr(self): - if not hasattr(self, '_ref_snr'): - wfs = {ifo: (self.h00_sparse[ifo], - self.h00_sparse[ifo]) for ifo in self.h00_sparse} + if not hasattr(self, "_ref_snr"): + wfs = { + ifo: (self.h00_sparse[ifo], self.h00_sparse[ifo]) + for ifo in self.h00_sparse + } self._ref_snr = self.get_snr(wfs) return self._ref_snr def get_snr(self, wfs): - """ Return hp/hc maximized SNR time series - """ + """Return hp/hc maximized SNR time series""" delta_t = 1.0 / self.sample_rate snrs = {} for ifo in wfs: sdat = self.sdat[ifo] dtc = self.tstart[ifo] - self.end_time[ifo] - self.ta[ifo] - snr = snr_predictor(self.fedges[ifo], - dtc - delta_t * 2.0, delta_t, - self.num_samples[ifo] + 4, - wfs[ifo][0], wfs[ifo][1], - self.h00_sparse[ifo], - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) - snrs[ifo] = TimeSeries(snr, delta_t=delta_t, - epoch=self.tstart[ifo] - delta_t * 2.0) + snr = snr_predictor( + self.fedges[ifo], + dtc - delta_t * 2.0, + delta_t, + self.num_samples[ifo] + 4, + wfs[ifo][0], + wfs[ifo][1], + self.h00_sparse[ifo], + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) + snrs[ifo] = TimeSeries( + snr, delta_t=delta_t, epoch=self.tstart[ifo] - delta_t * 2.0 + ) return snrs @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, .. math:: @@ -747,6 +825,7 @@ def _loglr(self): ------- float The value of the log likelihood ratio. + """ # get model params p = self.current_params @@ -754,7 +833,7 @@ def _loglr(self): lik = self.likelihood_function norm = 0.0 filt = 0j - pol_phase = numpy.exp(-2.0j * p['polarization']) + pol_phase = numpy.exp(-2.0j * p["polarization"]) self.snr_draw(wfs) p = self.current_params @@ -768,53 +847,80 @@ def _loglr(self): hp, hc = wfs[ifo] det = self.det[ifo] - fp, fc = det.antenna_pattern(p["ra"], p["dec"], - 0, times) + fp, fc = det.antenna_pattern(p["ra"], p["dec"], 0, times) times = det.time_delay_from_earth_center(p["ra"], p["dec"], times) dtc = p["tc"] - end_time - self.ta[ifo] - if self.lformat == 'earth_time_pol': + if self.lformat == "earth_time_pol": filter_i, norm_i = lik( - freqs, fp, fc, times, dtc, pol_phase, - hp, hc, h00, - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) + freqs, + fp, + fc, + times, + dtc, + pol_phase, + hp, + hc, + h00, + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) else: f = (fp + 1.0j * fc) * pol_phase fp = f.real.copy() fc = f.imag.copy() - if self.lformat == 'earth_time': + if self.lformat == "earth_time": filter_i, norm_i = lik( - freqs, fp, fc, times, dtc, - hp, hc, h00, - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) + freqs, + fp, + fc, + times, + dtc, + hp, + hc, + h00, + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) else: - filter_i, norm_i = lik(freqs, fp, fc, times + dtc, - hp, hc, h00, - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) + filter_i, norm_i = lik( + freqs, + fp, + fc, + times + dtc, + hp, + hc, + h00, + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) filt += filter_i norm += norm_i loglr = self.marginalize_loglr(filt, norm) return loglr def _nowaveform_handler(self): - """Sets loglr values if no waveform generated. - """ + """Sets loglr values if no waveform generated.""" return -numpy.inf class RelativeTimeDom(RelativeTime): - """ Heterodyne likelihood optimized for time marginalization and only + """ + Heterodyne likelihood optimized for time marginalization and only dominant-mode waveforms. This enables the ability to do inclination marginalization in addition to the other forms supportedy by RelativeTime. """ + name = "relative_time_dom" def get_snr(self, wfs): - """ Return hp/hc maximized SNR time series - """ + """Return hp/hc maximized SNR time series""" delta_t = 1.0 / self.sample_rate snrs = {} self.sh = {} @@ -823,17 +929,24 @@ def get_snr(self, wfs): sdat = self.sdat[ifo] dtc = self.tstart[ifo] - self.end_time[ifo] - self.ta[ifo] - sh, hh = snr_predictor_dom(self.fedges[ifo], - dtc - delta_t * 2.0, delta_t, - self.num_samples[ifo] + 4, - wfs[ifo][0], - self.h00_sparse[ifo], - sdat['a0'], sdat['a1'], - sdat['b0'], sdat['b1']) - snr = TimeSeries(abs(sh[2:-2]) / hh ** 0.5, delta_t=delta_t, - epoch=self.tstart[ifo]) - self.sh[ifo] = TimeSeries(sh, delta_t=delta_t, - epoch=self.tstart[ifo] - delta_t * 2.0) + sh, hh = snr_predictor_dom( + self.fedges[ifo], + dtc - delta_t * 2.0, + delta_t, + self.num_samples[ifo] + 4, + wfs[ifo][0], + self.h00_sparse[ifo], + sdat["a0"], + sdat["a1"], + sdat["b0"], + sdat["b1"], + ) + snr = TimeSeries( + abs(sh[2:-2]) / hh**0.5, delta_t=delta_t, epoch=self.tstart[ifo] + ) + self.sh[ifo] = TimeSeries( + sh, delta_t=delta_t, epoch=self.tstart[ifo] - delta_t * 2.0 + ) self.hh[ifo] = hh snrs[ifo] = snr @@ -841,7 +954,8 @@ def get_snr(self, wfs): @catch_waveform_error def _loglr(self): - r"""Computes the log likelihood ratio, + r""" + Computes the log likelihood ratio, or inner product and if `self.return_sh_hh` is True. .. math:: @@ -859,18 +973,19 @@ def _loglr(self): or tuple The inner product (, ). + """ # calculate = - 2 + up to a constant p = self.current_params p2 = p.copy() - p2.pop('inclination') + p2.pop("inclination") wfs = self.get_waveforms(p2) sh_total = hh_total = 0 - ic = numpy.cos(p['inclination']) + ic = numpy.cos(p["inclination"]) ip = 0.5 * (1.0 + ic * ic) - pol_phase = numpy.exp(-2.0j * p['polarization']) + pol_phase = numpy.exp(-2.0j * p["polarization"]) snrs = self.get_snr(wfs) self.snr_draw(snrs=snrs) @@ -879,19 +994,16 @@ def _loglr(self): if self.precalc_antenna_factors: fp, fc, dt = self.get_precalc_antenna_factors(ifo) else: - dt = self.det[ifo].time_delay_from_earth_center(p['ra'], - p['dec'], - p['tc']) - fp, fc = self.det[ifo].antenna_pattern(p['ra'], p['dec'], - 0, p['tc']) - dts = p['tc'] + dt + dt = self.det[ifo].time_delay_from_earth_center( + p["ra"], p["dec"], p["tc"] + ) + fp, fc = self.det[ifo].antenna_pattern(p["ra"], p["dec"], 0, p["tc"]) + dts = p["tc"] + dt f = (fp + 1.0j * fc) * pol_phase # Note, this includes complex conjugation already # as our stored inner products were hp* x data - htf = (f.real * ip + 1.0j * f.imag * ic) - sh = self.sh[ifo].at_time(dts, - interpolate='quadratic', - extrapolate=0.0j) + htf = f.real * ip + 1.0j * f.imag * ic + sh = self.sh[ifo].at_time(dts, interpolate="quadratic", extrapolate=0.0j) sh_total += sh * htf hh_total += self.hh[ifo] * abs(htf) ** 2.0 @@ -903,8 +1015,7 @@ def _loglr(self): return results def _nowaveform_handler(self): - """Sets loglr values if no waveform generated. - """ + """Sets loglr values if no waveform generated.""" loglr = sh_total = hh_total = -numpy.inf if self.return_sh_hh: results = (sh_total, hh_total) diff --git a/pycbc/inference/models/single_template.py b/pycbc/inference/models/single_template.py index ae7afa136b4..4dab83ac371 100644 --- a/pycbc/inference/models/single_template.py +++ b/pycbc/inference/models/single_template.py @@ -13,23 +13,24 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides model classes that assume the noise is Gaussian. -""" +"""This module provides model classes that assume the noise is Gaussian.""" +import itertools import logging + import numpy -import itertools from pycbc import filter as pyfilter -from pycbc.waveform import get_fd_waveform from pycbc.detector import Detector +from pycbc.waveform import get_fd_waveform from .gaussian_noise import BaseGaussianNoise from .tools import DistMarg class SingleTemplate(DistMarg, BaseGaussianNoise): - r"""Model that assumes we know all the intrinsic parameters. + r""" + Model that assumes we know all the intrinsic parameters. This model assumes we know all the intrinsic parameters, and are only maximizing over the extrinsic ones. We also assume a dominant mode waveform @@ -56,19 +57,26 @@ class SingleTemplate(DistMarg, BaseGaussianNoise): \**kwargs : All other keyword arguments are passed to :py:class:`BaseGaussianNoise`; see that class for details. + """ - name = 'single_template' - def __init__(self, variable_params, data, low_frequency_cutoff, - sample_rate=32768, - marginalize_phase=True, - **kwargs): + name = "single_template" + + def __init__( + self, + variable_params, + data, + low_frequency_cutoff, + sample_rate=32768, + marginalize_phase=True, + **kwargs, + ): variable_params, kwargs = self.setup_marginalization( - variable_params, - marginalize_phase=marginalize_phase, - **kwargs) - super(SingleTemplate, self).__init__( - variable_params, data, low_frequency_cutoff, **kwargs) + variable_params, marginalize_phase=marginalize_phase, **kwargs + ) + super().__init__( + variable_params, data, low_frequency_cutoff, **kwargs + ) sample_rate = float(sample_rate) @@ -77,12 +85,12 @@ def __init__(self, variable_params, data, low_frequency_cutoff, self.df = df p = self.static_params.copy() for k in self.static_params: - if p[k] == 'REPLACE': + if p[k] == "REPLACE": p.pop(k) - if 'distance' in p: - _ = p.pop('distance') - if 'inclination' in p: - _ = p.pop('inclination') + if "distance" in p: + _ = p.pop("distance") + if "inclination" in p: + _ = p.pop("inclination") hp, _ = get_fd_waveform(delta_f=df, distance=1, inclination=0, **p) @@ -102,18 +110,22 @@ def __init__(self, variable_params, data, low_frequency_cutoff, self.data[ifo].resize(flen) self.det[ifo] = Detector(ifo) snr, _, norm = pyfilter.matched_filter_core( - hp, self.data[ifo], + hp, + self.data[ifo], psd=self.psds[ifo], low_frequency_cutoff=flow, - high_frequency_cutoff=fhigh) + high_frequency_cutoff=fhigh, + ) self.sh[ifo] = 4 * df * snr self.snr[ifo] = snr * norm self.hh[ifo] = pyfilter.sigmasq( - hp, psd=self.psds[ifo], + hp, + psd=self.psds[ifo], low_frequency_cutoff=flow, - high_frequency_cutoff=fhigh) + high_frequency_cutoff=fhigh, + ) self.waveform = hp self.htfs = {} # Waveform phase / distance transformation factors @@ -121,28 +133,30 @@ def __init__(self, variable_params, data, low_frequency_cutoff, # Retrict to analyzing around peaks if chosen and choose what # ifos to draw from - self.setup_peak_lock(snrs=self.snr, - sample_rate=sample_rate, - **kwargs) + self.setup_peak_lock(snrs=self.snr, sample_rate=sample_rate, **kwargs) self.draw_ifos(self.snr) @property def multi_signal_support(self): - """ The list of classes that this model supports in a multi-signal + """ + The list of classes that this model supports in a multi-signal likelihood """ # Check if this model *can* be included in a multi-signal model. # All marginalizations must currently be disabled to work! - if (self.marginalize_vector_params or - self.marginalize_distance or - self.marginalize_phase): - logging.info("Cannot use single template model inside of" - "multi_signal if marginalizations are enabled") + if ( + self.marginalize_vector_params + or self.marginalize_distance + or self.marginalize_phase + ): + logging.info( + "Cannot use single template model inside of" + "multi_signal if marginalizations are enabled" + ) return [type(self)] def calculate_hihjs(self, models): - """ Pre-calculate the hihj inner products on a grid - """ + """Pre-calculate the hihj inner products on a grid""" self.hihj = {} for m1, m2 in itertools.combinations(models, 2): self.hihj[(m1, m2)] = {} @@ -152,22 +166,23 @@ def calculate_hihjs(self, models): flow = self.kmin[ifo] * self.df fhigh = self.kmax[ifo] * self.df h1h2, _, _ = pyfilter.matched_filter_core( - h1, h2, - psd=self.psds[ifo], - low_frequency_cutoff=flow, - high_frequency_cutoff=fhigh) + h1, + h2, + psd=self.psds[ifo], + low_frequency_cutoff=flow, + high_frequency_cutoff=fhigh, + ) self.hihj[(m1, m2)][ifo] = 4 * self.df * h1h2 def multi_loglikelihood(self, models): - """ Calculate a multi-model (signal) likelihood - """ + """Calculate a multi-model (signal) likelihood""" models = [self] + models loglr = 0 # handle sum[ - 0.5 ] for m in models: loglr += m.loglr - if not hasattr(self, 'hihj'): + if not hasattr(self, "hihj"): self.calculate_hihjs(models) # finally add in the lognl term from this model @@ -180,46 +195,46 @@ def multi_loglikelihood(self, models): h1h2 = hihj_vec.at_time(dt, nearest_sample=True) h1h2 *= m1.htfs[det] * m2.htfs[det].conj() - loglr += - h1h2.real # This is -0.5 * re( + ) + loglr += -h1h2.real # This is -0.5 * re( + ) return loglr + self.lognl def _loglr(self): - r"""Computes the log likelihood ratio + r""" + Computes the log likelihood ratio Returns ------- float The value of the log likelihood ratio. + """ # calculate = - 2 + up to a constant p = self.current_params phase = 1 - if 'coa_phase' in p: - phase = numpy.exp(-1.0j * 2 * p['coa_phase']) + if "coa_phase" in p: + phase = numpy.exp(-1.0j * 2 * p["coa_phase"]) sh_total = hh_total = 0 - ic = numpy.cos(p['inclination']) + ic = numpy.cos(p["inclination"]) ip = 0.5 * (1.0 + ic * ic) - pol_phase = numpy.exp(-2.0j * p['polarization']) + pol_phase = numpy.exp(-2.0j * p["polarization"]) self.snr_draw(snrs=self.snr) for ifo in self.sh: - dt = self.det[ifo].time_delay_from_earth_center(p['ra'], p['dec'], - p['tc']) - self.dts[ifo] = p['tc'] + dt + dt = self.det[ifo].time_delay_from_earth_center(p["ra"], p["dec"], p["tc"]) + self.dts[ifo] = p["tc"] + dt - fp, fc = self.det[ifo].antenna_pattern(p['ra'], p['dec'], - 0, p['tc']) + fp, fc = self.det[ifo].antenna_pattern(p["ra"], p["dec"], 0, p["tc"]) f = (fp + 1.0j * fc) * pol_phase # Note, this includes complex conjugation already # as our stored inner products were hp* x data - htf = (f.real * ip + 1.0j * f.imag * ic) / p['distance'] * phase + htf = (f.real * ip + 1.0j * f.imag * ic) / p["distance"] * phase self.htfs[ifo] = htf - sh = self.sh[ifo].at_time(self.dts[ifo], interpolate='quadratic') + sh = self.sh[ifo].at_time(self.dts[ifo], interpolate="quadratic") sh_total += sh * htf hh_total += self.hh[ifo] * abs(htf) ** 2.0 diff --git a/pycbc/inference/models/tools.py b/pycbc/inference/models/tools.py index 03d931da0f1..bfad63a4aa1 100644 --- a/pycbc/inference/models/tools.py +++ b/pycbc/inference/models/tools.py @@ -1,5 +1,4 @@ -""" Common utility functions for calculation of likelihoods -""" +"""Common utility functions for calculation of likelihoods""" import logging import warnings @@ -8,35 +7,32 @@ import numpy import numpy.random import tqdm - -from scipy.special import logsumexp, i0e from scipy.interpolate import RectBivariateSpline, interp1d -from pycbc.distributions import JointDistribution +from scipy.special import i0e, logsumexp from pycbc.detector import Detector - +from pycbc.distributions import JointDistribution # Earth radius in seconds EARTH_RADIUS = 0.031 def str_to_tuple(sval, ftype): - """ Convenience parsing to convert str to tuple""" + """Convenience parsing to convert str to tuple""" if sval is None: return () - return tuple(ftype(x.strip(' ')) for x in sval.split(',')) + return tuple(ftype(x.strip(" ")) for x in sval.split(",")) def str_to_bool(sval): - """ Ensure value is a bool if it can be converted """ + """Ensure value is a bool if it can be converted""" if isinstance(sval, str): return strtobool(sval) return sval def draw_sample(loglr, size=None): - """ Draw a random index from a 1-d vector with loglr weights - """ + """Draw a random index from a 1-d vector with loglr weights""" if size: x = numpy.random.uniform(size=size) else: @@ -48,23 +44,26 @@ def draw_sample(loglr, size=None): return xl -class DistMarg(): +class DistMarg: """Help class to add bookkeeping for likelihood marginalization""" - def setup_marginalization(self, - variable_params, - marginalize_phase=False, - marginalize_distance=False, - marginalize_distance_param='distance', - marginalize_distance_samples=int(1e4), - marginalize_distance_interpolator=False, - marginalize_distance_snr_range=None, - marginalize_distance_density=None, - marginalize_vector_params=None, - marginalize_vector_samples=1e3, - marginalize_sky_initial_samples=1e6, - **kwargs): - """ Setup the model for use with distance marginalization + def setup_marginalization( + self, + variable_params, + marginalize_phase=False, + marginalize_distance=False, + marginalize_distance_param="distance", + marginalize_distance_samples=int(1e4), + marginalize_distance_interpolator=False, + marginalize_distance_snr_range=None, + marginalize_distance_density=None, + marginalize_vector_params=None, + marginalize_vector_samples=1e3, + marginalize_sky_initial_samples=1e6, + **kwargs, + ): + """ + Setup the model for use with distance marginalization This function sets up precalculations for distance / phase marginalization. For distance margininalization it modifies the @@ -99,18 +98,17 @@ def setup_marginalization(self, kwags: dict The keyword arguments to the model initialization, may be modified from the original set by this function. + """ + def pop_prior(param): variable_params.remove(param) - old_prior = kwargs['prior'] - - dists = [d for d in old_prior.distributions - if param not in d.params] - dprior = [d for d in old_prior.distributions - if param in d.params][0] - prior = JointDistribution(variable_params, - *dists, **old_prior.kwargs) - kwargs['prior'] = prior + old_prior = kwargs["prior"] + + dists = [d for d in old_prior.distributions if param not in d.params] + dprior = [d for d in old_prior.distributions if param in d.params][0] + prior = JointDistribution(variable_params, *dists, **old_prior.kwargs) + kwargs["prior"] = prior return dprior self.reconstruct_phase = False @@ -123,22 +121,26 @@ def pop_prior(param): self.marginalized_vector_priors = {} self.vsamples = int(marginalize_vector_samples) - self.marginalize_sky_initial_samples = \ - int(float(marginalize_sky_initial_samples)) + self.marginalize_sky_initial_samples = int( + float(marginalize_sky_initial_samples) + ) for param in str_to_tuple(marginalize_vector_params, str): - logging.info('Marginalizing over %s, %s points from prior', - param, self.vsamples) + logging.info( + "Marginalizing over %s, %s points from prior", param, self.vsamples + ) self.marginalized_vector_priors[param] = pop_prior(param) # Remove in the future, backwards compatibility - if 'polarization_samples' in kwargs: - warnings.warn("use marginalize_vector_samples rather " - "than 'polarization_samples'", DeprecationWarning) + if "polarization_samples" in kwargs: + warnings.warn( + "use marginalize_vector_samples rather than 'polarization_samples'", + DeprecationWarning, + ) pol_uniform = numpy.linspace(0, numpy.pi * 2.0, self.vsamples) - self.marginalize_vector_params['polarization'] = pol_uniform - self.vsamples = int(kwargs['polarization_samples']) - kwargs.pop('polarization_samples') + self.marginalize_vector_params["polarization"] = pol_uniform + self.vsamples = int(kwargs["polarization_samples"]) + kwargs.pop("polarization_samples") self.reset_vector_params() @@ -153,57 +155,57 @@ def pop_prior(param): return variable_params, kwargs if isinstance(marginalize_distance_snr_range, str): - marginalize_distance_snr_range = \ - str_to_tuple(marginalize_distance_snr_range, float) + marginalize_distance_snr_range = str_to_tuple( + marginalize_distance_snr_range, float + ) if isinstance(marginalize_distance_density, str): - marginalize_distance_density = \ - str_to_tuple(marginalize_distance_density, int) + marginalize_distance_density = str_to_tuple( + marginalize_distance_density, int + ) - logging.info('Marginalizing over distance') + logging.info("Marginalizing over distance") # Take distance out of the variable params since we'll handle it # manually now dprior = pop_prior(marginalize_distance_param) - if len(dprior.params) != 1 or not hasattr(dprior, 'bounds'): - raise ValueError('Distance Marginalization requires a ' - 'univariate and bounded prior') + if len(dprior.params) != 1 or not hasattr(dprior, "bounds"): + raise ValueError( + "Distance Marginalization requires a univariate and bounded prior" + ) # Set up distance prior vector and samples # (1) prior is using distance - if dprior.params[0] == 'distance': - logging.info("Prior is directly on distance, setting up " - "%s grid weights", marginalize_distance_samples) - dmin, dmax = dprior.bounds['distance'] - dist_locs = numpy.linspace(dmin, dmax, - int(marginalize_distance_samples)) + if dprior.params[0] == "distance": + logging.info( + "Prior is directly on distance, setting up %s grid weights", + marginalize_distance_samples, + ) + dmin, dmax = dprior.bounds["distance"] + dist_locs = numpy.linspace(dmin, dmax, int(marginalize_distance_samples)) dist_weights = [dprior.pdf(distance=l) for l in dist_locs] dist_weights = numpy.array(dist_weights) # (2) prior is univariate and can be converted to distance - elif marginalize_distance_param != 'distance': - waveform_transforms = kwargs['waveform_transforms'] + elif marginalize_distance_param != "distance": + waveform_transforms = kwargs["waveform_transforms"] pname = dprior.params[0] - logging.info("Settings up transform, prior is in terms of" - " %s", pname) - wtrans = [d for d in waveform_transforms - if 'distance' not in d.outputs] + logging.info("Settings up transform, prior is in terms of %s", pname) + wtrans = [d for d in waveform_transforms if "distance" not in d.outputs] if len(wtrans) == 0: wtrans = None - kwargs['waveform_transforms'] = wtrans - dtrans = [d for d in waveform_transforms - if 'distance' in d.outputs][0] + kwargs["waveform_transforms"] = wtrans + dtrans = [d for d in waveform_transforms if "distance" in d.outputs][0] v = dprior.rvs(int(1e8)) - d = dtrans.transform({pname: v[pname]})['distance'] + d = dtrans.transform({pname: v[pname]})["distance"] d.sort() - cdf = numpy.arange(1, len(d)+1) / len(d) + cdf = numpy.arange(1, len(d) + 1) / len(d) i = interp1d(d, cdf) dmin, dmax = d.min(), d.max() - logging.info('Distance range %s-%s', dmin, dmax) - x = numpy.linspace(dmin, dmax, - int(marginalize_distance_samples) + 1) + logging.info("Distance range %s-%s", dmin, dmax) + x = numpy.linspace(dmin, dmax, int(marginalize_distance_samples) + 1) xl, xr = x[:-1], x[1:] dist_locs = 0.5 * (xr + xl) dist_weights = i(xr) - i(xl) @@ -219,39 +221,42 @@ def pop_prior(param): if str_to_bool(marginalize_distance_interpolator): setup_args = {} if marginalize_distance_snr_range: - setup_args['snr_range'] = marginalize_distance_snr_range + setup_args["snr_range"] = marginalize_distance_snr_range if marginalize_distance_density: - setup_args['density'] = marginalize_distance_density - i = setup_distance_marg_interpolant(self.distance_marginalization, - phase=self.marginalize_phase, - **setup_args) + setup_args["density"] = marginalize_distance_density + i = setup_distance_marg_interpolant( + self.distance_marginalization, + phase=self.marginalize_phase, + **setup_args, + ) self.distance_interpolator = i - kwargs['static_params']['distance'] = dist_ref + kwargs["static_params"]["distance"] = dist_ref # Save marginalized parameters' name into one place, # coa_phase will be a static param if been marginalized if marginalize_distance: - self.marginalized_params_name =\ - list(self.marginalize_vector_params.keys()) +\ - [marginalize_distance_param] + self.marginalized_params_name = list( + self.marginalize_vector_params.keys() + ) + [marginalize_distance_param] return variable_params, kwargs def reset_vector_params(self): - """ Redraw vector params from their priors - """ + """Redraw vector params from their priors""" for param in self.marginalized_vector_priors: vprior = self.marginalized_vector_priors[param] values = vprior.rvs(self.vsamples)[param] self.marginalize_vector_params[param] = values - def marginalize_loglr(self, sh_total, hh_total, - skip_vector=False, return_peak=False): - """ Return the marginal likelihood + def marginalize_loglr( + self, sh_total, hh_total, skip_vector=False, return_peak=False + ): + """ + Return the marginal likelihood Parameters - ----------- + ---------- sh_total: float or ndarray The total inner product summed over detectors hh_total: float or ndarray @@ -259,6 +264,7 @@ def marginalize_loglr(self, sh_total, hh_total, skip_vector: bool, False If true, and input is a vector, do not marginalize over that vector, instead return the likelihood values as a vector. + """ interpolator = self.distance_interpolator return_complex = False @@ -277,31 +283,32 @@ def marginalize_loglr(self, sh_total, hh_total, skip_vector = True return_complex = True - return marginalize_likelihood(sh_total, hh_total, - logw=self.marginalize_vector_weights, - phase=self.marginalize_phase, - interpolator=interpolator, - distance=distance, - skip_vector=skip_vector, - return_complex=return_complex, - return_peak=return_peak) + return marginalize_likelihood( + sh_total, + hh_total, + logw=self.marginalize_vector_weights, + phase=self.marginalize_phase, + interpolator=interpolator, + distance=distance, + skip_vector=skip_vector, + return_complex=return_complex, + return_peak=return_peak, + ) def premarg_draw(self): - """ Choose random samples from prechosen set""" - + """Choose random samples from prechosen set""" # Update the current proposed times and the marginalization values - logw = self.premarg['logw_partial'] + logw = self.premarg["logw_partial"] if self.vsamples == len(logw): choice = slice(None, None) else: - choice = numpy.random.choice(len(logw), size=self.vsamples, - replace=False) + choice = numpy.random.choice(len(logw), size=self.vsamples, replace=False) for k in self.snr_params: self.marginalize_vector_params[k] = self.premarg[k][choice] self._current_params.update(self.marginalize_vector_params) - self.sample_idx = self.premarg['sample_idx'][choice] + self.sample_idx = self.premarg["sample_idx"][choice] # Update the importance weights for each vector sample logw = self.marginalize_vector_weights + logw[choice] @@ -309,28 +316,32 @@ def premarg_draw(self): return self.marginalize_vector_params def snr_draw(self, wfs=None, snrs=None, size=None): - """ Improve the monte-carlo vector marginalization using the SNR time + """ + Improve the monte-carlo vector marginalization using the SNR time series of each detector """ try: p = self.current_params - set_scalar = numpy.isscalar(p['tc']) + set_scalar = numpy.isscalar(p["tc"]) except: set_scalar = False if not set_scalar: - if hasattr(self, 'premarg'): + if hasattr(self, "premarg"): return self.premarg_draw() if snrs is None: snrs = self.get_snr(wfs) - if ('tc' in self.marginalized_vector_priors and - not ('ra' in self.marginalized_vector_priors - or 'dec' in self.marginalized_vector_priors)): + if "tc" in self.marginalized_vector_priors and not ( + "ra" in self.marginalized_vector_priors + or "dec" in self.marginalized_vector_priors + ): return self.draw_times(snrs, size=size) - elif ('tc' in self.marginalized_vector_priors and - 'ra' in self.marginalized_vector_priors and - 'dec' in self.marginalized_vector_priors): + if ( + "tc" in self.marginalized_vector_priors + and "ra" in self.marginalized_vector_priors + and "dec" in self.marginalized_vector_priors + ): return self.draw_sky_times(snrs, size=size) else: # OK, we couldn't do anything with the requested monte-carlo @@ -339,30 +350,32 @@ def snr_draw(self, wfs=None, snrs=None, size=None): return None def draw_times(self, snrs, size=None): - """ Draw times consistent with the incoherent network SNR + """ + Draw times consistent with the incoherent network SNR Parameters ---------- snrs: dist of TimeSeries + """ - if not hasattr(self, 'tinfo'): + if not hasattr(self, "tinfo"): # determine the rough time offsets for this sky location - tcprior = self.marginalized_vector_priors['tc'] - tcmin, tcmax = tcprior.bounds['tc'] + tcprior = self.marginalized_vector_priors["tc"] + tcmin, tcmax = tcprior.bounds["tc"] tcave = (tcmax + tcmin) / 2.0 ifos = list(snrs.keys()) - if hasattr(self, 'keep_ifos'): + if hasattr(self, "keep_ifos"): ifos = self.keep_ifos d = {ifo: Detector(ifo, reference_time=tcave) for ifo in ifos} self.tinfo = tcmin, tcmax, tcave, ifos, d - self.snr_params = ['tc'] + self.snr_params = ["tc"] tcmin, tcmax, tcave, ifos, d = self.tinfo vsamples = size if size is not None else self.vsamples # Determine the weights for the valid time range - ra = self._current_params['ra'] - dec = self._current_params['dec'] + ra = self._current_params["ra"] + dec = self._current_params["dec"] # Determine the common valid time range iref = ifos[0] @@ -375,7 +388,7 @@ def draw_times(self, snrs, size=None): delt = snrs[iref].delta_t tmin = tcmin + dt - delt tmax = tcmax + dt + delt - if hasattr(self, 'tstart'): + if hasattr(self, "tstart"): tmin = self.tstart[iref] tmax = self.tend[iref] @@ -395,16 +408,16 @@ def draw_times(self, snrs, size=None): start = max(starts) end = min(ends) if end <= start: - return + return None # get the weights - snr = snrs[iref].time_slice(start, end, mode='nearest') + snr = snrs[iref].time_slice(start, end, mode="nearest") logweight = snr.squared_norm().numpy() for ifo in ifos[1:]: idel = idels[ifo] - snrv = snrs[ifo].time_slice(snr.start_time + idel, - snr.end_time + idel, - mode='nearest') + snrv = snrs[ifo].time_slice( + snr.start_time + idel, snr.end_time + idel, mode="nearest" + ) logweight += snrv.squared_norm().numpy() logweight /= 2.0 logweight -= logsumexp(logweight) # Normalize to PDF @@ -413,16 +426,14 @@ def draw_times(self, snrs, size=None): # Draw first which time sample tci = draw_sample(logweight, size=vsamples) # Second draw a subsample size offset so that all times are covered - tct = numpy.random.uniform(-snr.delta_t / 2.0, - snr.delta_t / 2.0, - size=vsamples) + tct = numpy.random.uniform(-snr.delta_t / 2.0, snr.delta_t / 2.0, size=vsamples) tc = tct + tci * snr.delta_t + float(snr.start_time) - dt # Update the current proposed times and the marginalization values # assumes uniform prior! - logw = - logweight[tci] + numpy.log(1.0 / len(logweight)) - self.marginalize_vector_params['tc'] = tc - self.marginalize_vector_params['logw_partial'] = logw + logw = -logweight[tci] + numpy.log(1.0 / len(logweight)) + self.marginalize_vector_params["tc"] = tc + self.marginalize_vector_params["logw_partial"] = logw if self._current_params is not None: # Update the importance weights for each vector sample @@ -432,39 +443,39 @@ def draw_times(self, snrs, size=None): return self.marginalize_vector_params def draw_sky_times(self, snrs, size=None): - """ Draw ra, dec, and tc together using SNR timeseries to determine + """ + Draw ra, dec, and tc together using SNR timeseries to determine monte-carlo weights. """ # First setup # precalculate dense sky grid and make dict and or array of the results ifos = list(snrs.keys()) - if hasattr(self, 'keep_ifos'): + if hasattr(self, "keep_ifos"): ifos = self.keep_ifos - ikey = ''.join(ifos) + ikey = "".join(ifos) vsamples = size if size is not None else self.vsamples # No good SNR peaks, go with prior draw if len(ifos) == 0: - self.marginalize_vector_params['logw_partial'] = numpy.zeros(vsamples) - return + self.marginalize_vector_params["logw_partial"] = numpy.zeros(vsamples) + return None def make_init(): - self.snr_params = ['tc', 'ra', 'dec'] + self.snr_params = ["tc", "ra", "dec"] size = self.marginalize_sky_initial_samples - logging.info('drawing samples: %s', size) - ra = self.marginalized_vector_priors['ra'].rvs(size=size)['ra'] - dec = self.marginalized_vector_priors['dec'].rvs(size=size)['dec'] - tcmin, tcmax = self.marginalized_vector_priors['tc'].bounds['tc'] + logging.info("drawing samples: %s", size) + ra = self.marginalized_vector_priors["ra"].rvs(size=size)["ra"] + dec = self.marginalized_vector_priors["dec"].rvs(size=size)["dec"] + tcmin, tcmax = self.marginalized_vector_priors["tc"].bounds["tc"] tcave = (tcmax + tcmin) / 2.0 d = {ifo: Detector(ifo, reference_time=tcave) for ifo in self.data} # What data structure to hold times? Dict of offset -> list? - logging.info('sorting into time delay dict') + logging.info("sorting into time delay dict") dts = [] for i in range(len(ifos) - 1): - dt = d[ifos[0]].time_delay_from_detector(d[ifos[i+1]], - ra, dec, tcave) + dt = d[ifos[0]].time_delay_from_detector(d[ifos[i + 1]], ra, dec, tcave) dt = numpy.rint(dt / snrs[ifos[0]].delta_t) dts.append(dt) @@ -487,11 +498,11 @@ def make_init(): return dmap, tcmin, tcmax, fp, fc, ra, dec, dtc, bin_prior - if not hasattr(self, 'tinfo'): + if not hasattr(self, "tinfo"): self.tinfo = {} if ikey not in self.tinfo: - logging.info('pregenerating sky pointings') + logging.info("pregenerating sky pointings") self.tinfo[ikey] = make_init() dmap, tcmin, tcmax, fp, fc, ra, dec, dtc, bin_prior = self.tinfo[ikey] @@ -506,13 +517,13 @@ def make_init(): for ifo in ifos: snr = snrs[ifo] tmin, tmax = tcmin - EARTH_RADIUS, tcmax + EARTH_RADIUS - if hasattr(self, 'tstart'): + if hasattr(self, "tstart"): tmin = self.tstart[ifo] tmax = self.tend[ifo] start = max(tmin, snr.start_time + snr.delta_t) end = min(tmax, snr.end_time - snr.delta_t * 2) - snr = snr.time_slice(start, end, mode='nearest') + snr = snr.time_slice(start, end, mode="nearest") w = snr.squared_norm().numpy() / 2.0 i = draw_sample(w, size=vsamples) @@ -546,8 +557,8 @@ def make_init(): # If we had really poor efficiency at finding a point, we should # give up and just use the original random draws if len(ix) < 0.05 * vsamples: - self.marginalize_vector_params['logw_partial'] = numpy.zeros(vsamples) - return + self.marginalize_vector_params["logw_partial"] = numpy.zeros(vsamples) + return None # fill back to fixed size with repeat samples # sample order is random, so this should be OK statistically @@ -564,9 +575,7 @@ def make_init(): wi = numpy.resize(numpy.array(wi), vsamples) # Second draw a subsample size offset so that all times are covered - tct = numpy.random.uniform(-snr.delta_t / 2.0, - snr.delta_t / 2.0, - size=len(ti)) + tct = numpy.random.uniform(-snr.delta_t / 2.0, snr.delta_t / 2.0, size=len(ti)) tc = tct + iref[ti] * snr.delta_t + float(sref.start_time) - dtc[ifos[0]] @@ -575,10 +584,10 @@ def make_init(): # factor at the moment. logw_sky = -mcweight[ti] + numpy.log(wi) - numpy.log(resize_factor) - self.marginalize_vector_params['tc'] = tc - self.marginalize_vector_params['ra'] = ra - self.marginalize_vector_params['dec'] = dec - self.marginalize_vector_params['logw_partial'] = logw_sky + self.marginalize_vector_params["tc"] = tc + self.marginalize_vector_params["ra"] = ra + self.marginalize_vector_params["dec"] = dec + self.marginalize_vector_params["logw_partial"] = logw_sky if self._current_params is not None: # Update the importance weights for each vector sample @@ -588,19 +597,22 @@ def make_init(): return self.marginalize_vector_params def get_precalc_antenna_factors(self, ifo): - """ Get the antenna factors for marginalized samples if they exist """ + """Get the antenna factors for marginalized samples if they exist""" ix = self.sample_idx fp, fc, dtc = self.precalc_antenna_factors return fp[ifo][ix], fc[ifo][ix], dtc[ifo][ix] - def setup_peak_lock(self, - sample_rate=4096, - snrs=None, - peak_lock_snr=None, - peak_lock_ratio=1e4, - peak_lock_region=4, - **kwargs): - """ Determine where to constrain marginalization based on + def setup_peak_lock( + self, + sample_rate=4096, + snrs=None, + peak_lock_snr=None, + peak_lock_ratio=1e4, + peak_lock_region=4, + **kwargs, + ): + """ + Determine where to constrain marginalization based on the observed reference SNR peaks. Parameters @@ -618,20 +630,20 @@ def setup_peak_lock(self, peak_lock_region: int Number of samples to inclue beyond the strict region determined by the relative likelihood - """ - if 'tc' not in self.marginalized_vector_priors: + """ + if "tc" not in self.marginalized_vector_priors: return - tcmin, tcmax = self.marginalized_vector_priors['tc'].bounds['tc'] + tcmin, tcmax = self.marginalized_vector_priors["tc"].bounds["tc"] tstart = tcmin - EARTH_RADIUS tmax = tcmax - tcmin + EARTH_RADIUS * 2.0 num_samples = int(tmax * sample_rate) - self.tstart = {ifo: tstart for ifo in self.data} - self.num_samples = {ifo: num_samples for ifo in self.data} + self.tstart = dict.fromkeys(self.data, tstart) + self.num_samples = dict.fromkeys(self.data, num_samples) if snrs is None: - if not hasattr(self, 'ref_snr'): + if not hasattr(self, "ref_snr"): raise ValueError("Model didn't have a reference SNR!") snrs = self.ref_snr @@ -645,16 +657,17 @@ def setup_peak_lock(self, for ifo in snrs: s = max(tstart, snrs[ifo].start_time) e = min(tstart + tmax, snrs[ifo].end_time) - z = snrs[ifo].time_slice(s, e, mode='nearest') + z = snrs[ifo].time_slice(s, e, mode="nearest") peak_snr, imax = z.abs_max_loc() times = z.sample_times peak_time = times[imax] - logging.info('%s: Max Ref SNR Peak of %s at %s', - ifo, peak_snr, peak_time) + logging.info( + "%s: Max Ref SNR Peak of %s at %s", ifo, peak_snr, peak_time + ) if peak_snr > peak_lock_snr: - target = peak_snr ** 2.0 / 2.0 - numpy.log(peak_lock_ratio) + target = peak_snr**2.0 / 2.0 - numpy.log(peak_lock_ratio) target = (target * 2.0) ** 0.5 region = numpy.where(abs(z) > target)[0] @@ -681,25 +694,36 @@ def setup_peak_lock(self, self.tstart[ifo] = ts self.num_samples[ifo] = int((te - ts) * sample_rate) + 1 - logging.info('%s: use region %s-%s, %s points', - ifo, ts, te, self.num_samples[ifo]) + logging.info( + "%s: use region %s-%s, %s points", + ifo, + ts, + te, + self.num_samples[ifo], + ) self.tend = self.tstart.copy() for ifo in snrs: self.tend[ifo] += self.num_samples[ifo] / sample_rate - def draw_ifos(self, snrs, peak_snr_threshold=4.0, log=True, - precalculate_marginalization_points=False, - **kwargs): - """ Helper utility to determine which ifos we should use based on the + def draw_ifos( + self, + snrs, + peak_snr_threshold=4.0, + log=True, + precalculate_marginalization_points=False, + **kwargs, + ): + """ + Helper utility to determine which ifos we should use based on the reference SNR time series. """ - if 'tc' not in self.marginalized_vector_priors: - return + if "tc" not in self.marginalized_vector_priors: + return None peak_snr_threshold = float(peak_snr_threshold) - tcmin, tcmax = self.marginalized_vector_priors['tc'].bounds['tc'] + tcmin, tcmax = self.marginalized_vector_priors["tc"].bounds["tc"] ifos = list(snrs.keys()) keep_ifos = [] psnrs = [] @@ -707,29 +731,33 @@ def draw_ifos(self, snrs, peak_snr_threshold=4.0, log=True, snr = snrs[ifo] start = max(tcmin - EARTH_RADIUS, snr.start_time) end = min(tcmax + EARTH_RADIUS, snr.end_time) - snr = snr.time_slice(start, end, mode='nearest') + snr = snr.time_slice(start, end, mode="nearest") psnr = abs(snr).max() if psnr > peak_snr_threshold: keep_ifos.append(ifo) psnrs.append(psnr) if log: - logging.info("Ifos used for SNR based draws:" - " %s, snrs: %s, peak_snr_threshold=%s", - keep_ifos, psnrs, peak_snr_threshold) + logging.info( + "Ifos used for SNR based draws: %s, snrs: %s, peak_snr_threshold=%s", + keep_ifos, + psnrs, + peak_snr_threshold, + ) self.keep_ifos = keep_ifos if precalculate_marginalization_points: num_points = int(float(precalculate_marginalization_points)) self.premarg = self.snr_draw(size=num_points, snrs=snrs).copy() - self.premarg['sample_idx'] = self.sample_idx + self.premarg["sample_idx"] = self.sample_idx return keep_ifos @property def current_params(self): - """ The current parameters + """ + The current parameters If a parameter has been vector marginalized, the likelihood should expect an array for the given parameter. This allows transparent @@ -739,11 +767,12 @@ def current_params(self): for k in self.marginalize_vector_params: if k not in params: params[k] = self.marginalize_vector_params[k] - self.marginalize_vector_weights = - numpy.log(self.vsamples) + self.marginalize_vector_weights = -numpy.log(self.vsamples) return params def reconstruct(self, rec=None, seed=None, set_loglr=None): - """ Reconstruct the distance or vectored marginalized parameter + """ + Reconstruct the distance or vectored marginalized parameter of this class. """ if seed: @@ -753,6 +782,7 @@ def reconstruct(self, rec=None, seed=None, set_loglr=None): rec = {} if set_loglr is None: + def get_loglr(): p = self.current_params.copy() p.update(rec) @@ -762,7 +792,7 @@ def get_loglr(): get_loglr = set_loglr if self.marginalize_vector_params: - logging.debug('Reconstruct vector') + logging.debug("Reconstruct vector") self.reconstruct_vector = True self.reset_vector_params() loglr = get_loglr() @@ -772,36 +802,36 @@ def get_loglr(): self.reconstruct_vector = False if self.distance_marginalization: - logging.debug('Reconstruct distance') + logging.debug("Reconstruct distance") # call likelihood to get vector output self.reconstruct_distance = True _, weights = self.distance_marginalization loglr = get_loglr() xl = draw_sample(loglr + numpy.log(weights)) - rec['distance'] = self.dist_locs[xl] + rec["distance"] = self.dist_locs[xl] self.reconstruct_distance = False if self.marginalize_phase: - logging.debug('Reconstruct phase') + logging.debug("Reconstruct phase") self.reconstruct_phase = True s, h = get_loglr() - phasev = numpy.linspace(0, numpy.pi*2.0, int(1e4)) + phasev = numpy.linspace(0, numpy.pi * 2.0, int(1e4)) # This assumes that the template was conjugated in inner products loglr = (numpy.exp(-2.0j * phasev) * s).real + h xl = draw_sample(loglr) - rec['coa_phase'] = phasev[xl] + rec["coa_phase"] = phasev[xl] self.reconstruct_phase = False - rec['loglr'] = loglr[xl] - rec['loglikelihood'] = self.lognl + rec['loglr'] + rec["loglr"] = loglr[xl] + rec["loglikelihood"] = self.lognl + rec["loglr"] return rec -def setup_distance_marg_interpolant(dist_marg, - phase=False, - snr_range=(1, 50), - density=(1000, 1000)): - """ Create the interpolant for distance marginalization +def setup_distance_marg_interpolant( + dist_marg, phase=False, snr_range=(1, 50), density=(1000, 1000) +): + """ + Create the interpolant for distance marginalization Parameters ---------- @@ -819,6 +849,7 @@ def setup_distance_marg_interpolant(dist_marg, interp: function Function which returns the precalculated likelihood for a given inner product sh/hh. + """ dist_rescale, _ = dist_marg logging.info("Interpolator valid for SNRs in %s", snr_range) @@ -828,21 +859,21 @@ def setup_distance_marg_interpolant(dist_marg, snr_min, snr_max = snr_range smax = dist_rescale.max() smin = dist_rescale.min() - shr_max = snr_max ** 2.0 / smin - hhr_max = snr_max ** 2.0 / smin / smin + shr_max = snr_max**2.0 / smin + hhr_max = snr_max**2.0 / smin / smin - shr_min = snr_min ** 2.0 / smax - hhr_min = snr_min ** 2.0 / smax / smax + shr_min = snr_min**2.0 / smax + hhr_min = snr_min**2.0 / smax / smax shr = numpy.geomspace(shr_min, shr_max, density[0]) hhr = numpy.geomspace(hhr_min, hhr_max, density[1]) lvals = numpy.zeros((len(shr), len(hhr))) - logging.info('Setup up likelihood interpolator') + logging.info("Setup up likelihood interpolator") for i, sh in enumerate(tqdm.tqdm(shr)): for j, hh in enumerate(hhr): - lvals[i, j] = marginalize_likelihood(sh, hh, - distance=dist_marg, - phase=phase) + lvals[i, j] = marginalize_likelihood( + sh, hh, distance=dist_marg, phase=phase + ) interp = RectBivariateSpline(shr, hhr, lvals) def interp_wrapper(x, y, bounds_check=True): @@ -859,19 +890,23 @@ def interp_wrapper(x, y, bounds_check=True): if k is not None: v[k] = -numpy.inf return v + return interp_wrapper -def marginalize_likelihood(sh, hh, - logw=None, - phase=False, - distance=False, - skip_vector=False, - interpolator=None, - return_peak=False, - return_complex=False, - ): - """ Return the marginalized likelihood. +def marginalize_likelihood( + sh, + hh, + logw=None, + phase=False, + distance=False, + skip_vector=False, + interpolator=None, + return_peak=False, + return_complex=False, +): + """ + Return the marginalized likelihood. Apply various marginalizations to the data, including phase, distance, and brute-force vector marginalizations. Several options relate @@ -910,10 +945,12 @@ def marginalize_likelihood(sh, hh, ------- loglr: float The marginalized loglikehood ratio + """ if distance and not interpolator and not numpy.isscalar(sh): - raise ValueError("Cannot do vector marginalization " - "and distance at the same time") + raise ValueError( + "Cannot do vector marginalization and distance at the same time" + ) if logw is None: if isinstance(hh, float): @@ -940,7 +977,7 @@ def marginalize_likelihood(sh, hh, # brute force distance path dist_rescale, dist_weights = distance sh = sh * dist_rescale - hh = hh * dist_rescale ** 2.0 + hh = hh * dist_rescale**2.0 logw = numpy.log(dist_weights) if return_complex: diff --git a/pycbc/inference/option_utils.py b/pycbc/inference/option_utils.py index 49c7934bf67..435742f1395 100644 --- a/pycbc/inference/option_utils.py +++ b/pycbc/inference/option_utils.py @@ -14,10 +14,10 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module contains standard options used for inference-related programs. -""" +"""This module contains standard options used for inference-related programs.""" import argparse + from pycbc import waveform # ----------------------------------------------------------------------------- @@ -28,7 +28,8 @@ class ParseLabelArg(argparse.Action): - """Argparse action that will parse arguments that can accept labels. + """ + Argparse action that will parse arguments that can accept labels. This assumes that the values set on the command line for its assigned argument are strings formatted like ``PARAM[:LABEL]``. When the arguments @@ -44,15 +45,14 @@ class ParseLabelArg(argparse.Action): This action can work on arguments that have ``nargs != 0`` and ``type`` set to ``str``. """ - def __init__(self, type=str, nargs=None, - **kwargs): # pylint: disable=redefined-builtin + + def __init__(self, type=str, nargs=None, **kwargs): # pylint: disable=redefined-builtin # check that type is string if type != str: raise ValueError("the type for this action must be a string") if nargs == 0: raise ValueError("nargs must not be 0 for this action") - super(ParseLabelArg, self).__init__(type=type, nargs=nargs, - **kwargs) + super().__init__(type=type, nargs=nargs, **kwargs) def __call__(self, parser, namespace, values, option_string=None): singlearg = isinstance(values, str) @@ -61,7 +61,7 @@ def __call__(self, parser, namespace, values, option_string=None): params = [] labels = {} for param in values: - psplit = param.split(':') + psplit = param.split(":") if len(psplit) == 2: param, label = psplit else: @@ -72,11 +72,12 @@ def __call__(self, parser, namespace, values, option_string=None): if singlearg: params = params[0] setattr(namespace, self.dest, params) - setattr(namespace, '{}_labels'.format(self.dest), labels) + setattr(namespace, f"{self.dest}_labels", labels) class ParseParametersArg(ParseLabelArg): - """Argparse action that will parse parameters and labels from an opton. + """ + Argparse action that will parse parameters and labels from an opton. Does the same as ``ParseLabelArg``, with the additional functionality that if ``LABEL`` is a known parameter in ``pycbc.waveform.parameters``, then @@ -117,12 +118,15 @@ class ParseParametersArg(ParseLabelArg): string was used in the labels dictionary. Since ``ni`` and ``bar`` (the label for ``z-arg``) are not recognized parameters, they were just used as-is in the labels dictionaries. + """ + def __call__(self, parser, namespace, values, option_string=None): - super(ParseParametersArg, self).__call__(parser, namespace, values, - option_string=option_string) + super().__call__( + parser, namespace, values, option_string=option_string + ) # try to replace the labels with a label from waveform.parameters - labels = getattr(namespace, '{}_labels'.format(self.dest)) + labels = getattr(namespace, f"{self.dest}_labels") for param, label in labels.items(): try: label = getattr(waveform.parameters, label).label @@ -132,103 +136,157 @@ def __call__(self, parser, namespace, values, option_string=None): def add_injsamples_map_opt(parser): - """Adds option to parser to specify a mapping between injection parameters + """ + Adds option to parser to specify a mapping between injection parameters an sample parameters. """ - parser.add_argument('--injection-samples-map', nargs='+', - metavar='INJECTION_PARAM:SAMPLES_PARAM', - help='Rename/apply functions to the injection ' - 'parameters and name them the same as one of the ' - 'parameters in samples. This can be used if the ' - 'injection parameters are not the same as the ' - 'samples parameters. INJECTION_PARAM may be a ' - 'function of the injection parameters; ' - 'SAMPLES_PARAM must a name of one of the ' - 'parameters in the samples group.') + parser.add_argument( + "--injection-samples-map", + nargs="+", + metavar="INJECTION_PARAM:SAMPLES_PARAM", + help="Rename/apply functions to the injection " + "parameters and name them the same as one of the " + "parameters in samples. This can be used if the " + "injection parameters are not the same as the " + "samples parameters. INJECTION_PARAM may be a " + "function of the injection parameters; " + "SAMPLES_PARAM must a name of one of the " + "parameters in the samples group.", + ) def add_plot_posterior_option_group(parser): - """Adds the options needed to configure plots of posterior results. + """ + Adds the options needed to configure plots of posterior results. Parameters ---------- parser : object ArgumentParser instance. + """ - pgroup = parser.add_argument_group("Options for what plots to create and " - "their formats.") - pgroup.add_argument('--plot-marginal', action='store_true', default=False, - help="Plot 1D marginalized distributions on the " - "diagonal axes.") - pgroup.add_argument('--marginal-percentiles', nargs='+', default=None, - type=float, - help="Percentiles to draw lines at on the 1D " - "histograms.") - pgroup.add_argument('--no-marginal-lines', action='store_true', - default=False, - help="Do not add vertical lines in the 1D marginal " - "plots showing the marginal percentiles.") - pgroup.add_argument('--no-marginal-titles', action='store_true', - default=False, - help="Do not add titles giving the 1D credible range " - "over the 1D marginal plots.") - pgroup.add_argument("--plot-scatter", action='store_true', default=False, - help="Plot each sample point as a scatter plot.") - pgroup.add_argument("--plot-density", action="store_true", default=False, - help="Plot the posterior density as a color map.") - pgroup.add_argument("--plot-contours", action="store_true", default=False, - help="Draw contours showing the 50th and 90th " - "percentile confidence regions.") - pgroup.add_argument('--contour-percentiles', nargs='+', default=None, - type=float, - help="Percentiles to draw contours if different " - "than 50th and 90th.") + pgroup = parser.add_argument_group( + "Options for what plots to create and their formats." + ) + pgroup.add_argument( + "--plot-marginal", + action="store_true", + default=False, + help="Plot 1D marginalized distributions on the diagonal axes.", + ) + pgroup.add_argument( + "--marginal-percentiles", + nargs="+", + default=None, + type=float, + help="Percentiles to draw lines at on the 1D histograms.", + ) + pgroup.add_argument( + "--no-marginal-lines", + action="store_true", + default=False, + help="Do not add vertical lines in the 1D marginal " + "plots showing the marginal percentiles.", + ) + pgroup.add_argument( + "--no-marginal-titles", + action="store_true", + default=False, + help="Do not add titles giving the 1D credible range " + "over the 1D marginal plots.", + ) + pgroup.add_argument( + "--plot-scatter", + action="store_true", + default=False, + help="Plot each sample point as a scatter plot.", + ) + pgroup.add_argument( + "--plot-density", + action="store_true", + default=False, + help="Plot the posterior density as a color map.", + ) + pgroup.add_argument( + "--plot-contours", + action="store_true", + default=False, + help="Draw contours showing the 50th and 90th percentile confidence regions.", + ) + pgroup.add_argument( + "--contour-percentiles", + nargs="+", + default=None, + type=float, + help="Percentiles to draw contours if different than 50th and 90th.", + ) # add mins, maxs options - pgroup.add_argument('--mins', nargs='+', metavar='PARAM:VAL', default=[], - help="Specify minimum parameter values to plot. This " - "should be done by specifying the parameter name " - "followed by the value. Parameter names must be " - "the same as the PARAM argument in --parameters " - "(or, if no parameters are provided, the same as " - "the parameter name specified in the variable " - "args in the input file. If none provided, " - "the smallest parameter value in the posterior " - "will be used.") - pgroup.add_argument('--maxs', nargs='+', metavar='PARAM:VAL', default=[], - help="Same as mins, but for the maximum values to " - "plot.") + pgroup.add_argument( + "--mins", + nargs="+", + metavar="PARAM:VAL", + default=[], + help="Specify minimum parameter values to plot. This " + "should be done by specifying the parameter name " + "followed by the value. Parameter names must be " + "the same as the PARAM argument in --parameters " + "(or, if no parameters are provided, the same as " + "the parameter name specified in the variable " + "args in the input file. If none provided, " + "the smallest parameter value in the posterior " + "will be used.", + ) + pgroup.add_argument( + "--maxs", + nargs="+", + metavar="PARAM:VAL", + default=[], + help="Same as mins, but for the maximum values to plot.", + ) # add expected parameters options - pgroup.add_argument('--expected-parameters', nargs='+', - metavar='PARAM:VAL', - default=[], - help="Specify expected parameter values to plot. If " - "provided, a cross will be plotted in each axis " - "that an expected parameter is provided. " - "Parameter names must be " - "the same as the PARAM argument in --parameters " - "(or, if no parameters are provided, the same as " - "the parameter name specified in the variable " - "args in the input file.") - pgroup.add_argument('--expected-parameters-color', default='r', - help="What to color the expected-parameters cross. " - "Default is red.") - pgroup.add_argument('--plot-injection-parameters', action='store_true', - default=False, - help="Get the expected parameters from the injection " - "in the input file. There must be only a single " - "injection in the file to work. Any values " - "specified by expected-parameters will override " - "the values obtained for the injection.") - pgroup.add_argument('--pick-injection-by-time', action='store_true', - default=False, - help="In the case of multiple injections, pick one" - " for plotting based on its proximity in time.") + pgroup.add_argument( + "--expected-parameters", + nargs="+", + metavar="PARAM:VAL", + default=[], + help="Specify expected parameter values to plot. If " + "provided, a cross will be plotted in each axis " + "that an expected parameter is provided. " + "Parameter names must be " + "the same as the PARAM argument in --parameters " + "(or, if no parameters are provided, the same as " + "the parameter name specified in the variable " + "args in the input file.", + ) + pgroup.add_argument( + "--expected-parameters-color", + default="r", + help="What to color the expected-parameters cross. Default is red.", + ) + pgroup.add_argument( + "--plot-injection-parameters", + action="store_true", + default=False, + help="Get the expected parameters from the injection " + "in the input file. There must be only a single " + "injection in the file to work. Any values " + "specified by expected-parameters will override " + "the values obtained for the injection.", + ) + pgroup.add_argument( + "--pick-injection-by-time", + action="store_true", + default=False, + help="In the case of multiple injections, pick one" + " for plotting based on its proximity in time.", + ) add_injsamples_map_opt(pgroup) return pgroup def plot_ranges_from_cli(opts): - """Parses the mins and maxs arguments from the `plot_posterior` option + """ + Parses the mins and maxs arguments from the `plot_posterior` option group. Parameters @@ -246,16 +304,17 @@ def plot_ranges_from_cli(opts): Dictionary of parameter name -> specified maxs. Only parameters that were specified in the --mins option will be included; if no parameters were provided, will return an empty dictionary. + """ mins = {} for x in opts.mins: - x = x.split(':') + x = x.split(":") if len(x) != 2: raise ValueError("option --mins not specified correctly; see help") mins[x[0]] = float(x[1]) maxs = {} for x in opts.maxs: - x = x.split(':') + x = x.split(":") if len(x) != 2: raise ValueError("option --maxs not specified correctly; see help") maxs[x[0]] = float(x[1]) @@ -263,7 +322,8 @@ def plot_ranges_from_cli(opts): def expected_parameters_from_cli(opts): - """Parses the --expected-parameters arguments from the `plot_posterior` + """ + Parses the --expected-parameters arguments from the `plot_posterior` option group. Parameters @@ -277,85 +337,124 @@ def expected_parameters_from_cli(opts): Dictionary of parameter name -> expected value. Only parameters that were specified in the --expected-parameters option will be included; if no parameters were provided, will return an empty dictionary. + """ expected = {} for x in opts.expected_parameters: - x = x.split(':') + x = x.split(":") if len(x) != 2: - raise ValueError("option --expected-paramters not specified " - "correctly; see help") + raise ValueError( + "option --expected-paramters not specified correctly; see help" + ) expected[x[0]] = float(x[1]) return expected def add_scatter_option_group(parser): - """Adds the options needed to configure scatter plots. + """ + Adds the options needed to configure scatter plots. Parameters ---------- parser : object ArgumentParser instance. + """ - scatter_group = parser.add_argument_group("Options for configuring the " - "scatter plot.") + scatter_group = parser.add_argument_group( + "Options for configuring the scatter plot." + ) scatter_group.add_argument( - '--z-arg', type=str, default=None, action=ParseParametersArg, - help='What to color the scatter points by. Syntax is the same as the ' - 'parameters option.') + "--z-arg", + type=str, + default=None, + action=ParseParametersArg, + help="What to color the scatter points by. Syntax is the same as the " + "parameters option.", + ) scatter_group.add_argument( - "--vmin", type=float, help="Minimum value for the colorbar.") + "--vmin", type=float, help="Minimum value for the colorbar." + ) scatter_group.add_argument( - "--vmax", type=float, help="Maximum value for the colorbar.") + "--vmax", type=float, help="Maximum value for the colorbar." + ) scatter_group.add_argument( - "--scatter-cmap", type=str, default='plasma', - help="Specify the colormap to use for points. Default is plasma.") + "--scatter-cmap", + type=str, + default="plasma", + help="Specify the colormap to use for points. Default is plasma.", + ) return scatter_group def add_density_option_group(parser): - """Adds the options needed to configure contours and density colour map. + """ + Adds the options needed to configure contours and density colour map. Parameters ---------- parser : object ArgumentParser instance. + """ - density_group = parser.add_argument_group("Options for configuring the " - "contours and density color map") + density_group = parser.add_argument_group( + "Options for configuring the contours and density color map" + ) density_group.add_argument( - "--density-cmap", type=str, default='viridis', - help="Specify the colormap to use for the density. " - "Default is viridis.") + "--density-cmap", + type=str, + default="viridis", + help="Specify the colormap to use for the density. Default is viridis.", + ) density_group.add_argument( - "--contour-color", type=str, default=None, + "--contour-color", + type=str, + default=None, help="Specify the color to use for the contour lines. Default is " - "white for density plots and black for scatter plots.") + "white for density plots and black for scatter plots.", + ) density_group.add_argument( - "--contour-linestyles", type=str, default=None, nargs="+", + "--contour-linestyles", + type=str, + default=None, + nargs="+", help="Specify the linestyles to use for the contour lines. Defaut " - "is solid for all.") + "is solid for all.", + ) density_group.add_argument( - "--no-contour-labels", action="store_true", default=False, - help="Don't put labels on the contours.") + "--no-contour-labels", + action="store_true", + default=False, + help="Don't put labels on the contours.", + ) density_group.add_argument( - '--use-kombine-kde', default=False, action="store_true", + "--use-kombine-kde", + default=False, + action="store_true", help="Use kombine's clustered KDE for determining 2D marginal " - "contours and density instead of scipy's gaussian_kde (the " - "default). This is better at distinguishing bimodal " - "distributions, but is much slower than the default. For speed, " - "suggest setting --kde-args 'max_samples:20000' or smaller if " - "using this. Requires kombine to be installed.") + "contours and density instead of scipy's gaussian_kde (the " + "default). This is better at distinguishing bimodal " + "distributions, but is much slower than the default. For speed, " + "suggest setting --kde-args 'max_samples:20000' or smaller if " + "using this. Requires kombine to be installed.", + ) density_group.add_argument( - '--max-kde-samples', type=int, default=None, + "--max-kde-samples", + type=int, + default=None, help="Limit the number of samples used for KDE construction to the " - "given value. This can substantially speed up plot generation " - "(particularly when plotting multiple parameters). Suggested " - "values: 5000 to 10000.") + "given value. This can substantially speed up plot generation " + "(particularly when plotting multiple parameters). Suggested " + "values: 5000 to 10000.", + ) density_group.add_argument( - '--kde-args', metavar="ARG:VALUE", nargs='+', default=None, + "--kde-args", + metavar="ARG:VALUE", + nargs="+", + default=None, help="Pass the given argrument, value pairs to the KDE function " - "(either scipy's or kombine's) when setting it up.") + "(either scipy's or kombine's) when setting it up.", + ) return density_group diff --git a/pycbc/inference/sampler/__init__.py b/pycbc/inference/sampler/__init__.py index 094a2298747..7d422a86b03 100644 --- a/pycbc/inference/sampler/__init__.py +++ b/pycbc/inference/sampler/__init__.py @@ -20,27 +20,31 @@ import logging # pylint: disable=unused-import -from .base import (initial_dist_from_config, create_new_output_file) -from .multinest import MultinestSampler -from .ultranest import UltranestSampler +from .base import create_new_output_file, initial_dist_from_config from .dummy import DummySampler +from .games import GameSampler +from .multinest import MultinestSampler from .refine import RefineSampler from .snowline import SnowlineSampler -from .games import GameSampler +from .ultranest import UltranestSampler # list of available samplers -samplers = {cls.name: cls for cls in ( - MultinestSampler, - UltranestSampler, - DummySampler, - RefineSampler, - SnowlineSampler, - GameSampler, -)} +samplers = { + cls.name: cls + for cls in ( + MultinestSampler, + UltranestSampler, + DummySampler, + RefineSampler, + SnowlineSampler, + GameSampler, + ) +} try: from .emcee import EmceeEnsembleSampler from .emcee_pt import EmceePTSampler + samplers[EmceeEnsembleSampler.name] = EmceeEnsembleSampler samplers[EmceePTSampler.name] = EmceePTSampler except ImportError: @@ -48,37 +52,43 @@ try: from .epsie import EpsieSampler + samplers[EpsieSampler.name] = EpsieSampler except ImportError: pass try: from .ptemcee import PTEmceeSampler + samplers[PTEmceeSampler.name] = PTEmceeSampler except ImportError: pass try: from .cpnest import CPNestSampler + samplers[CPNestSampler.name] = CPNestSampler except ImportError: pass try: from .dynesty import DynestySampler + samplers[DynestySampler.name] = DynestySampler except ImportError: pass try: from .nessai import NessaiSampler + samplers[NessaiSampler.name] = NessaiSampler except ImportError: pass def load_from_config(cp, model, **kwargs): - """Loads a sampler from the given config file. + """ + Loads a sampler from the given config file. This looks for a name in the section ``[sampler]`` to determine which sampler class to load. That sampler's ``from_config`` is then called. @@ -97,12 +107,13 @@ def load_from_config(cp, model, **kwargs): ------- sampler : The initialized sampler. + """ if len(model.variable_params) == 0: - logging.info('No variable params, so assuming Dummy Sampler') + logging.info("No variable params, so assuming Dummy Sampler") return DummySampler.from_config(cp, model, **kwargs) - name = cp.get('sampler', 'name') + name = cp.get("sampler", "name") try: return samplers[name].from_config(cp, model, **kwargs) except KeyError: diff --git a/pycbc/inference/sampler/base.py b/pycbc/inference/sampler/base.py index 6ec5b7bce82..a82be30bc0e 100644 --- a/pycbc/inference/sampler/base.py +++ b/pycbc/inference/sampler/base.py @@ -25,10 +25,9 @@ Defines the base sampler class to be inherited by all samplers. """ - -from abc import ABCMeta, abstractmethod, abstractproperty -import shutil import logging +import shutil +from abc import ABCMeta, abstractmethod, abstractproperty from six import add_metaclass @@ -45,8 +44,9 @@ @add_metaclass(ABCMeta) -class BaseSampler(object): - """Abstract base class for all inference samplers. +class BaseSampler: + """ + Abstract base class for all inference samplers. All sampler classes must inherit from this class and implement its abstract methods. @@ -55,7 +55,9 @@ class BaseSampler(object): ---------- model : Model An instance of a model from ``pycbc.inference.models``. + """ + name = None def __init__(self, model): @@ -67,83 +69,75 @@ def __init__(self, model): # @classmethod <--uncomment when we move to python 3.3 @abstractmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): - """This should initialize the sampler given a config file. - """ - pass + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): + """This should initialize the sampler given a config file.""" @property def variable_params(self): - """Returns the parameters varied in the model. - """ + """Returns the parameters varied in the model.""" return self.model.variable_params @property def sampling_params(self): - """Returns the sampling params used by the model. - """ + """Returns the sampling params used by the model.""" return self.model.sampling_params @property def static_params(self): - """Returns the model's fixed parameters. - """ + """Returns the model's fixed parameters.""" return self.model.static_params @abstractproperty def samples(self): - """A dict mapping variable_params to arrays of samples currently + """ + A dict mapping variable_params to arrays of samples currently in memory. The dictionary may also contain sampling_params. The sample arrays may have any shape, and may or may not be thinned. """ - pass @abstractproperty def model_stats(self): - """A dict mapping model's metadata fields to arrays of values for + """ + A dict mapping model's metadata fields to arrays of values for each sample in ``raw_samples``. The arrays may have any shape, and may or may not be thinned. """ - pass @abstractmethod def run(self): - """This function should run the sampler. + """ + This function should run the sampler. Any checkpointing should be done internally in this function. """ - pass @abstractproperty def io(self): - """A class that inherits from ``BaseInferenceFile`` to handle IO with + """ + A class that inherits from ``BaseInferenceFile`` to handle IO with an hdf file. This should be a class, not an instance of class, so that the sampler can initialize it when needed. """ - pass @abstractmethod def checkpoint(self): - """The sampler must have a checkpoint method for dumping raw samples + """ + The sampler must have a checkpoint method for dumping raw samples and stats to the file type defined by ``io``. """ - pass @abstractmethod def finalize(self): """Do any finalization to the samples file before exiting.""" - pass @abstractmethod def resume_from_checkpoint(self): - """Resume the sampler from the output file. - """ - pass + """Resume the sampler from the output file.""" + # # ============================================================================= @@ -155,7 +149,8 @@ def resume_from_checkpoint(self): def setup_output(sampler, output_file, check_nsamples=True, validate=True): - r"""Sets up the sampler's checkpoint and output files. + r""" + Sets up the sampler's checkpoint and output files. The checkpoint file has the same name as the output file, but with ``.checkpoint`` appended to the name. A backup file will also be @@ -167,17 +162,18 @@ def setup_output(sampler, output_file, check_nsamples=True, validate=True): Sampler output_file : str Name of the output file. + """ # check for backup file(s) - checkpoint_file = output_file + '.checkpoint' - backup_file = output_file + '.bkup' + checkpoint_file = output_file + ".checkpoint" + backup_file = output_file + ".bkup" # check if we have a good checkpoint and/or backup file logging.info("Looking for checkpoint file") checkpoint_valid = False if validate: - checkpoint_valid = validate_checkpoint_files(checkpoint_file, - backup_file, - check_nsamples) + checkpoint_valid = validate_checkpoint_files( + checkpoint_file, backup_file, check_nsamples + ) # Create a new file if the checkpoint doesn't exist, or if it is # corrupted sampler.new_checkpoint = False # keeps track if this is a new file or not @@ -201,7 +197,8 @@ def setup_output(sampler, output_file, check_nsamples=True, validate=True): def create_new_output_file(sampler, filename, **kwargs): - r"""Creates a new output file. + r""" + Creates a new output file. Parameters ---------- @@ -212,8 +209,9 @@ def create_new_output_file(sampler, filename, **kwargs): \**kwargs : All other keyword arguments are passed through to the file's ``write_metadata`` function. + """ - logging.info("Creating file {}".format(filename)) + logging.info(f"Creating file {filename}") with sampler.io(filename, "w") as fp: # create the samples group and sampler info group fp.create_group(fp.samples_group) @@ -223,7 +221,8 @@ def create_new_output_file(sampler, filename, **kwargs): def initial_dist_from_config(cp, variable_params, static_params=None): - r"""Loads a distribution for the sampler start from the given config file. + r""" + Loads a distribution for the sampler start from the given config file. A distribution will only be loaded if the config file has a [initial-\*] section(s). @@ -243,18 +242,21 @@ def initial_dist_from_config(cp, variable_params, static_params=None): JointDistribution or None : The initial distribution. If no [initial-\*] section found in the config file, will just return None. + """ if len(cp.get_subsections("initial")): - logging.info("Using a different distribution for the starting points " - "than the prior.") + logging.info( + "Using a different distribution for the starting points than the prior." + ) initial_dists = distributions.read_distributions_from_config( - cp, section="initial") + cp, section="initial" + ) constraints = distributions.read_constraints_from_config( - cp, constraint_section="initial_constraint", - static_args=static_params) + cp, constraint_section="initial_constraint", static_args=static_params + ) init_dist = distributions.JointDistribution( - variable_params, *initial_dists, - **{"constraints": constraints}) + variable_params, *initial_dists, constraints=constraints + ) else: init_dist = None return init_dist diff --git a/pycbc/inference/sampler/base_cube.py b/pycbc/inference/sampler/base_cube.py index 6b397208e12..16906cbbc61 100644 --- a/pycbc/inference/sampler/base_cube.py +++ b/pycbc/inference/sampler/base_cube.py @@ -25,6 +25,7 @@ Common utilities for samplers that rely on transforming between a unit cube and the prior space. This is typical of many nested sampling algorithms. """ + import numpy from .. import models @@ -39,10 +40,8 @@ def call_global_logprior(cube): def setup_calls(model, loglikelihood_function=None, copy_prior=False): - """ Configure calls for MPI support - """ - model_call = CubeModel(model, loglikelihood_function, - copy_prior=copy_prior) + """Configure calls for MPI support""" + model_call = CubeModel(model, loglikelihood_function, copy_prior=copy_prior) # these are used to help paralleize over multiple cores / MPI models._global_instance = model_call @@ -51,13 +50,15 @@ def setup_calls(model, loglikelihood_function=None, copy_prior=False): return log_likelihood_call, prior_call -class CubeModel(object): - """ Class for making PyCBC Inference 'model class' +class CubeModel: + """ + Class for making PyCBC Inference 'model class' Parameters ---------- model : inference.BaseModel instance A model instance from pycbc. + """ def __init__(self, model, loglikelihood_function=None, copy_prior=False): @@ -65,13 +66,13 @@ def __init__(self, model, loglikelihood_function=None, copy_prior=False): raise ValueError("Ultranest or dynesty do not support sampling transforms") self.model = model if loglikelihood_function is None: - loglikelihood_function = 'loglikelihood' + loglikelihood_function = "loglikelihood" self.loglikelihood_function = loglikelihood_function self.copy_prior = copy_prior def log_likelihood(self, cube): """ - returns log likelihood function + Returns log likelihood function """ params = dict(zip(self.model.sampling_params, cube)) self.model.update(**params) @@ -81,7 +82,7 @@ def log_likelihood(self, cube): def prior_transform(self, cube): """ - prior transform function for ultranest sampler + Prior transform function for ultranest sampler It takes unit cube as input parameter and apply prior transforms """ diff --git a/pycbc/inference/sampler/base_mcmc.py b/pycbc/inference/sampler/base_mcmc.py index e69ad727dc8..db49f74bf18 100644 --- a/pycbc/inference/sampler/base_mcmc.py +++ b/pycbc/inference/sampler/base_mcmc.py @@ -23,16 +23,16 @@ # """Provides constructor classes and convenience functions for MCMC samplers.""" -import logging -from abc import (ABCMeta, abstractmethod, abstractproperty) - import configparser as ConfigParser +import logging +from abc import ABCMeta, abstractmethod, abstractproperty import numpy from pycbc.filter import autocorrelation -from pycbc.inference.io import (validate_checkpoint_files, loadfile) +from pycbc.inference.io import loadfile, validate_checkpoint_files from pycbc.inference.io.base_mcmc import nsamples_in_chain + from .base import initial_dist_from_config # @@ -45,7 +45,8 @@ def raw_samples_to_dict(sampler, raw_samples): - """Convenience function for converting ND array to a dict of samples. + """ + Convenience function for converting ND array to a dict of samples. The samples are assumed to have dimension ``[sampler.base_shape x] niterations x len(sampler.sampling_params)``. @@ -64,23 +65,22 @@ def raw_samples_to_dict(sampler, raw_samples): sampling params are not the same as the variable params, they will also be included. Each array will have shape ``[sampler.base_shape x] niterations``. + """ sampling_params = sampler.sampling_params # convert to dictionary - samples = {param: raw_samples[..., ii] for - ii, param in enumerate(sampling_params)} + samples = {param: raw_samples[..., ii] for ii, param in enumerate(sampling_params)} # apply boundary conditions - samples = sampler.model.prior_distribution.apply_boundary_conditions( - **samples) + samples = sampler.model.prior_distribution.apply_boundary_conditions(**samples) # apply transforms to go to model's variable params space if sampler.model.sampling_transforms is not None: - samples = sampler.model.sampling_transforms.apply( - samples, inverse=True) + samples = sampler.model.sampling_transforms.apply(samples, inverse=True) return samples def blob_data_to_dict(stat_names, blobs): - """Converts list of "blobs" to a dictionary of model stats. + """ + Converts list of "blobs" to a dictionary of model stats. Samplers like ``emcee`` store the extra tuple returned by ``CallModel`` to a list called blobs. This is a list of lists of tuples with shape @@ -100,12 +100,14 @@ def blob_data_to_dict(stat_names, blobs): dict : A dictionary mapping the model's ``default_stats`` to arrays of values. Each array will have shape ``nwalkers x niterations``. + """ # get the dtypes of each of the stats; we'll just take this from the # first iteration and walker dtypes = [type(val) for val in blobs[0][0]] assert len(stat_names) == len(dtypes), ( - "number of stat names must match length of tuples in the blobs") + "number of stat names must match length of tuples in the blobs" + ) # convert to an array; to ensure that we get the dtypes correct, we'll # cast to a structured array raw_stats = numpy.array(blobs, dtype=list(zip(stat_names, dtypes))) @@ -116,7 +118,8 @@ def blob_data_to_dict(stat_names, blobs): def get_optional_arg_from_config(cp, section, arg, dtype=str): - """Convenience function to retrieve an optional argument from a config + """ + Convenience function to retrieve an optional argument from a config file. Parameters @@ -135,6 +138,7 @@ def get_optional_arg_from_config(cp, section, arg, dtype=str): ------- val : None or str If the argument is present, the value. Otherwise, None. + """ if cp.has_option(section, arg): val = dtype(cp.get(section, arg)) @@ -152,8 +156,9 @@ def get_optional_arg_from_config(cp, section, arg, dtype=str): # -class BaseMCMC(object, metaclass=ABCMeta): - r"""Abstract base class that provides methods common to MCMCs. +class BaseMCMC(metaclass=ABCMeta): + r""" + Abstract base class that provides methods common to MCMCs. This is not a sampler class itself. Sampler classes can inherit from this along with ``BaseSampler``. @@ -185,6 +190,7 @@ class BaseMCMC(object, metaclass=ABCMeta): [`classmethod`] Should compute the autocorrelation length using the given filename. Also allows for other keyword arguments. """ + _lastclear = None # the iteration when samples were cleared from memory _itercounter = None # the number of iterations since the last clear _pos = None @@ -201,14 +207,14 @@ class BaseMCMC(object, metaclass=ABCMeta): @abstractproperty def base_shape(self): - """What shape the sampler's samples arrays are in, excluding + """ + What shape the sampler's samples arrays are in, excluding the iterations dimension. For example, if a sampler uses 20 chains and 3 temperatures, this would be ``(3, 20)``. If a sampler only uses a single walker and no temperatures this would be ``()``. """ - pass @property def nchains(self): @@ -261,7 +267,8 @@ def thin_interval(self): @thin_interval.setter def thin_interval(self, interval): - """Sets the thin interval to use. + """ + Sets the thin interval to use. If ``None`` provided, will default to 1. """ @@ -286,21 +293,26 @@ def max_samples_per_chain(self, n): if n is not None: n = int(n) if n < self.thin_safety_factor: - raise ValueError("max samples per chain must be >= {}" - .format(self.thin_safety_factor)) + raise ValueError( + f"max samples per chain must be >= {self.thin_safety_factor}" + ) # also check that this is consistent with the target number of # effective samples if self.target_eff_nsamples is not None: - target_samps_per_chain = int(numpy.ceil( - self.target_eff_nsamples / self.nchains)) + target_samps_per_chain = int( + numpy.ceil(self.target_eff_nsamples / self.nchains) + ) if n <= target_samps_per_chain: - raise ValueError("max samples per chain must be > target " - "effective number of samples per walker " - "({})".format(target_samps_per_chain)) + raise ValueError( + "max samples per chain must be > target " + "effective number of samples per walker " + f"({target_samps_per_chain})" + ) self._max_samples_per_chain = n def get_thin_interval(self): - """Gets the thin interval to use. + """ + Gets the thin interval to use. If ``max_samples_per_chain`` is set, this will figure out what thin interval is needed to satisfy that criteria. In that case, the thin @@ -310,7 +322,7 @@ def get_thin_interval(self): # the extra factor of 2 is to account for the fact that the thin # interval will need to be at least twice as large as a previously # used interval - thinfactor = 2*(self.niterations // self.max_samples_per_chain) + thinfactor = 2 * (self.niterations // self.max_samples_per_chain) # make sure it's at least 1 thinfactor = max(thinfactor, 1) # make the new interval is a multiple of the previous, to ensure @@ -318,36 +330,36 @@ def get_thin_interval(self): if thinfactor < self.thin_interval: thin_interval = self.thin_interval else: - thin_interval = (thinfactor // self.thin_interval) * \ - self.thin_interval + thin_interval = (thinfactor // self.thin_interval) * self.thin_interval else: thin_interval = self.thin_interval return thin_interval def set_target(self, niterations=None, eff_nsamples=None): - """Sets the target niterations/nsamples for the sampler. + """ + Sets the target niterations/nsamples for the sampler. One or the other must be provided, not both. """ if niterations is None and eff_nsamples is None: - raise ValueError("Must provide a target niterations or " - "eff_nsamples") + raise ValueError("Must provide a target niterations or eff_nsamples") if niterations is not None and eff_nsamples is not None: - raise ValueError("Must provide a target niterations or " - "eff_nsamples, not both") - self._target_niterations = int(niterations) \ - if niterations is not None else None - self._target_eff_nsamples = int(eff_nsamples) \ - if eff_nsamples is not None else None + raise ValueError( + "Must provide a target niterations or eff_nsamples, not both" + ) + self._target_niterations = int(niterations) if niterations is not None else None + self._target_eff_nsamples = ( + int(eff_nsamples) if eff_nsamples is not None else None + ) @abstractmethod def clear_samples(self): """A method to clear samples from memory.""" - pass @property def pos(self): - """A dictionary of the current walker positions. + """ + A dictionary of the current walker positions. If the sampler hasn't been run yet, returns p0. """ @@ -355,13 +367,15 @@ def pos(self): if pos is None: return self.p0 # convert to dict - pos = {param: self._pos[..., k] - for (k, param) in enumerate(self.sampling_params)} + pos = { + param: self._pos[..., k] for (k, param) in enumerate(self.sampling_params) + } return pos @property def p0(self): - """A dictionary of the initial position of the chains. + """ + A dictionary of the initial position of the chains. This is set by using ``set_p0``. If not set yet, a ``ValueError`` is raised when the attribute is accessed. @@ -369,12 +383,12 @@ def p0(self): if self._p0 is None: raise ValueError("initial positions not set; run set_p0") # convert to dict - p0 = {param: self._p0[..., k] - for (k, param) in enumerate(self.sampling_params)} + p0 = {param: self._p0[..., k] for (k, param) in enumerate(self.sampling_params)} return p0 def set_p0(self, samples_file=None, prior=None): - """Sets the initial position of the chains. + """ + Sets the initial position of the chains. Parameters ---------- @@ -389,18 +403,20 @@ def set_p0(self, samples_file=None, prior=None): ------- p0 : dict A dictionary maping sampling params to the starting positions. + """ # if samples are given then use those as initial positions if samples_file is not None: - with self.io(samples_file, 'r') as fp: - samples = fp.read_samples(self.variable_params, - iteration=-1, flatten=False) + with self.io(samples_file, "r") as fp: + samples = fp.read_samples( + self.variable_params, iteration=-1, flatten=False + ) # remove the (length 1) niterations dimension samples = samples[..., 0] # make sure we have the same shape assert samples.shape == self.base_shape, ( - "samples in file {} have shape {}, but I have shape {}". - format(samples_file, samples.shape, self.base_shape)) + f"samples in file {samples_file} have shape {samples.shape}, but I have shape {self.base_shape}" + ) # transform to sampling parameter space if self.model.sampling_transforms is not None: samples = self.model.sampling_transforms.apply(samples) @@ -408,10 +424,11 @@ def set_p0(self, samples_file=None, prior=None): else: nsamples = numpy.prod(self.base_shape) samples = self.model.prior_rvs(size=nsamples, prior=prior).reshape( - self.base_shape) + self.base_shape + ) # store as ND array with shape [base_shape] x nparams ndim = len(self.variable_params) - p0 = numpy.ones(list(self.base_shape)+[ndim]) + p0 = numpy.ones(list(self.base_shape) + [ndim]) for i, param in enumerate(self.sampling_params): p0[..., i] = samples[param] self._p0 = p0 @@ -419,26 +436,23 @@ def set_p0(self, samples_file=None, prior=None): @abstractmethod def set_state_from_file(self, filename): - """Sets the state of the sampler to the instance saved in a file. - """ - pass + """Sets the state of the sampler to the instance saved in a file.""" def set_start_from_config(self, cp): - """Sets the initial state of the sampler from config file - """ - if cp.has_option('sampler', 'start-file'): - start_file = cp.get('sampler', 'start-file') + """Sets the initial state of the sampler from config file""" + if cp.has_option("sampler", "start-file"): + start_file = cp.get("sampler", "start-file") logging.info("Using file %s for initial positions", start_file) init_prior = None else: start_file = None init_prior = initial_dist_from_config( - cp, self.variable_params, self.static_params) + cp, self.variable_params, self.static_params + ) self.set_p0(samples_file=start_file, prior=init_prior) def resume_from_checkpoint(self): - """Resume the sampler from the checkpoint file - """ + """Resume the sampler from the checkpoint file""" with self.io(self.checkpoint_file, "r") as fp: self._lastclear = fp.niterations self.set_p0(samples_file=self.checkpoint_file) @@ -447,8 +461,10 @@ def resume_from_checkpoint(self): def run(self): """Runs the sampler.""" if self.target_eff_nsamples and self.checkpoint_interval is None: - raise ValueError("A checkpoint interval must be set if " - "targetting an effective number of samples") + raise ValueError( + "A checkpoint interval must be set if " + "targetting an effective number of samples" + ) # get the starting number of samples: # "nsamples" keeps track of the number of samples we've obtained (if # target_eff_nsamples is not None, this is the effective number of @@ -470,8 +486,10 @@ def run(self): target_nsamples = self.nchains * self.target_niterations nsamples = self._lastclear * self.nchains else: - raise ValueError("must set either target_eff_nsamples or " - "target_niterations; see set_target") + raise ValueError( + "must set either target_eff_nsamples or " + "target_niterations; see set_target" + ) self._itercounter = 0 # figure out the interval to use iterinterval = self.checkpoint_interval @@ -481,12 +499,14 @@ def run(self): while nsamples < target_nsamples: # adjust the interval if we would go past the number of iterations if self.target_niterations is not None and ( - self.niterations + iterinterval > self.target_niterations): + self.niterations + iterinterval > self.target_niterations + ): iterinterval = self.target_niterations - self.niterations # run sampler and set initial values to None so that sampler # picks up from where it left off next call - logging.info("Running sampler for {} to {} iterations".format( - self.niterations, self.niterations + iterinterval)) + logging.info( + f"Running sampler for {self.niterations} to {self.niterations + iterinterval} iterations" + ) # run the underlying sampler for the desired interval self.run_mcmc(iterinterval) # update the itercounter @@ -496,8 +516,7 @@ def run(self): # update nsamples for next loop if self.target_eff_nsamples is not None: nsamples = self.effective_nsamples - logging.info("Have {} effective samples post burn in".format( - nsamples)) + logging.info(f"Have {nsamples} effective samples post burn in") else: nsamples += iterinterval * self.nchains @@ -512,20 +531,18 @@ def set_burn_in(self, burn_in): @abstractmethod def effective_nsamples(self): - """The effective number of samples post burn-in that the sampler has + """ + The effective number of samples post burn-in that the sampler has acquired so far. """ - pass @abstractmethod def run_mcmc(self, niterations): """Run the MCMC for the given number of iterations.""" - pass @abstractmethod def write_results(self, filename): """Should write all samples currently in memory to the given file.""" - pass def checkpoint(self): """Dumps current samples to the checkpoint file.""" @@ -548,12 +565,16 @@ def checkpoint(self): thin_interval = fp.thinned_by elif thin_interval > fp.thinned_by: # we need to thin the samples on disk - logging.info("Thinning samples in %s by a factor " - "of %i", fn, int(thin_interval)) + logging.info( + "Thinning samples in %s by a factor of %i", + fn, + int(thin_interval), + ) fp.thin(thin_interval) fp_lastiter = fp.last_iteration() - logging.info("Writing samples to %s with thin interval %i", fn, - thin_interval) + logging.info( + "Writing samples to %s with thin interval %i", fn, thin_interval + ) self.write_results(fn) # update the running thin interval self.thin_interval = thin_interval @@ -590,14 +611,14 @@ def checkpoint(self): # check validity logging.info("Validating checkpoint and backup files") checkpoint_valid = validate_checkpoint_files( - self.checkpoint_file, self.backup_file) + self.checkpoint_file, self.backup_file + ) if not checkpoint_valid: - raise IOError("error writing to checkpoint file") - elif self.checkpoint_signal: + raise OSError("error writing to checkpoint file") + if self.checkpoint_signal: # kill myself with the specified signal - logging.info("Exiting with SIG{}".format(self.checkpoint_signal)) - kill_cmd="os.kill(os.getpid(), signal.SIG{})".format( - self.checkpoint_signal) + logging.info(f"Exiting with SIG{self.checkpoint_signal}") + kill_cmd = f"os.kill(os.getpid(), signal.SIG{self.checkpoint_signal})" exec(kill_cmd) # clear the in-memory chain to save memory logging.info("Clearing samples from memory") @@ -605,7 +626,8 @@ def checkpoint(self): @staticmethod def checkpoint_from_config(cp, section): - """Gets the checkpoint interval from the given config file. + """ + Gets the checkpoint interval from the given config file. This looks for 'checkpoint-interval' in the section. @@ -620,13 +642,16 @@ def checkpoint_from_config(cp, section): ------ int or None : The checkpoint interval, if it is in the section. Otherw + """ - return get_optional_arg_from_config(cp, section, 'checkpoint-interval', - dtype=int) + return get_optional_arg_from_config( + cp, section, "checkpoint-interval", dtype=int + ) @staticmethod def ckpt_signal_from_config(cp, section): - """Gets the checkpoint signal from the given config file. + """ + Gets the checkpoint signal from the given config file. This looks for 'checkpoint-signal' in the section. @@ -641,12 +666,13 @@ def ckpt_signal_from_config(cp, section): ------ int or None : The checkpoint interval, if it is in the section. Otherw + """ - return get_optional_arg_from_config(cp, section, 'checkpoint-signal', - dtype=str) + return get_optional_arg_from_config(cp, section, "checkpoint-signal", dtype=str) def set_target_from_config(self, cp, section): - """Sets the target using the given config file. + """ + Sets the target using the given config file. This looks for ``niterations`` to set the ``target_niterations``, and ``effective-nsamples`` to set the ``target_eff_nsamples``. @@ -657,6 +683,7 @@ def set_target_from_config(self, cp, section): Open config parser to retrieve the argument from. section : str Name of the section to retrieve from. + """ if cp.has_option(section, "niterations"): niterations = int(cp.get(section, "niterations")) @@ -669,7 +696,8 @@ def set_target_from_config(self, cp, section): self.set_target(niterations=niterations, eff_nsamples=nsamples) def set_burn_in_from_config(self, cp): - """Sets the burn in class from the given config file. + """ + Sets the burn in class from the given config file. If no burn-in section exists in the file, then this just set the burn-in class to None. @@ -681,8 +709,7 @@ def set_burn_in_from_config(self, cp): self.set_burn_in(bit) def set_thin_interval_from_config(self, cp, section): - """Sets thinning options from the given config file. - """ + """Sets thinning options from the given config file.""" if cp.has_option(section, "thin-interval"): thin_interval = int(cp.get(section, "thin-interval")) logging.info("Will thin samples using interval %i", thin_interval) @@ -690,25 +717,28 @@ def set_thin_interval_from_config(self, cp, section): thin_interval = None if cp.has_option(section, "max-samples-per-chain"): max_samps_per_chain = int(cp.get(section, "max-samples-per-chain")) - logging.info("Setting max samples per chain to %i", - max_samps_per_chain) + logging.info("Setting max samples per chain to %i", max_samps_per_chain) else: max_samps_per_chain = None # check for consistency if thin_interval is not None and max_samps_per_chain is not None: - raise ValueError("provide either thin-interval or " - "max-samples-per-chain, not both") + raise ValueError( + "provide either thin-interval or max-samples-per-chain, not both" + ) # check that the thin interval is < then the checkpoint interval - if thin_interval is not None and self.checkpoint_interval is not None \ - and thin_interval >= self.checkpoint_interval: - raise ValueError("thin interval must be less than the checkpoint " - "interval") + if ( + thin_interval is not None + and self.checkpoint_interval is not None + and thin_interval >= self.checkpoint_interval + ): + raise ValueError("thin interval must be less than the checkpoint interval") self.thin_interval = thin_interval self.max_samples_per_chain = max_samps_per_chain @property def raw_acls(self): - """Dictionary of parameter names -> autocorrelation lengths. + """ + Dictionary of parameter names -> autocorrelation lengths. Depending on the sampler, the ACLs may be an integer, or an arrray of values per chain and/or per temperature. @@ -724,28 +754,29 @@ def raw_acls(self, acls): @abstractmethod def acl(self): - """The autocorrelation length. + """ + The autocorrelation length. This method should convert the raw ACLs into an integer or array that can be used to extract independent samples from a chain. """ - pass @property def raw_acts(self): - """Dictionary of parameter names -> autocorrelation time(s). + """ + Dictionary of parameter names -> autocorrelation time(s). Returns ``None`` if no ACLs have been calculated. """ acls = self.raw_acls if acls is None: return None - return {p: acl * self.thin_interval - for (p, acl) in acls.items()} + return {p: acl * self.thin_interval for (p, acl) in acls.items()} @property def act(self): - """The autocorrelation time(s). + """ + The autocorrelation time(s). The autocorrelation time is defined as the autocorrelation length times the ``thin_interval``. It gives the number of iterations between @@ -761,23 +792,26 @@ def act(self): @abstractmethod def compute_acf(cls, filename, **kwargs): - """A method to compute the autocorrelation function of samples in the - given file.""" - pass + """ + A method to compute the autocorrelation function of samples in the + given file. + """ @abstractmethod def compute_acl(cls, filename, **kwargs): - """A method to compute the autocorrelation length of samples in the - given file.""" - pass + """ + A method to compute the autocorrelation length of samples in the + given file. + """ -class EnsembleSupport(object): +class EnsembleSupport: """Adds support for ensemble MCMC samplers.""" @property def nwalkers(self): - """The number of walkers used. + """ + The number of walkers used. Alias of ``nchains``. """ @@ -791,7 +825,8 @@ def nwalkers(self, value): @property def acl(self): - """The autocorrelation length of the ensemble. + """ + The autocorrelation length of the ensemble. This is calculated by taking the maximum over all of the ``raw_acls``. This works for both single and parallel-tempered ensemble samplers. @@ -805,7 +840,8 @@ def acl(self): @property def effective_nsamples(self): - """The effective number of samples post burn-in that the sampler has + """ + The effective number of samples post burn-in that the sampler has acquired so far. """ if self.burn_in is not None and not self.burn_in.is_burned_in: @@ -834,16 +870,23 @@ def effective_nsamples(self): # -def ensemble_compute_acf(filename, start_index=None, end_index=None, - per_walker=False, walkers=None, parameters=None): - """Computes the autocorrleation function for an ensemble MCMC. +def ensemble_compute_acf( + filename, + start_index=None, + end_index=None, + per_walker=False, + walkers=None, + parameters=None, +): + """ + Computes the autocorrleation function for an ensemble MCMC. By default, parameter values are averaged over all walkers at each iteration. The ACF is then calculated over the averaged chain. An ACF per-walker will be returned instead if ``per_walker=True``. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACFs for. start_index : int, optional @@ -868,9 +911,10 @@ def ensemble_compute_acf(filename, start_index=None, end_index=None, Dictionary of arrays giving the ACFs for each parameter. If ``per-walker`` is True, the arrays will have shape ``nwalkers x niterations``. + """ acfs = {} - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: if parameters is None: parameters = fp.variable_params if isinstance(parameters, str): @@ -881,26 +925,34 @@ def ensemble_compute_acf(filename, start_index=None, end_index=None, if walkers is None: walkers = numpy.arange(fp.nwalkers) arrays = [ - ensemble_compute_acf(filename, start_index=start_index, - end_index=end_index, - per_walker=False, walkers=ii, - parameters=param)[param] - for ii in walkers] + ensemble_compute_acf( + filename, + start_index=start_index, + end_index=end_index, + per_walker=False, + walkers=ii, + parameters=param, + )[param] + for ii in walkers + ] acfs[param] = numpy.vstack(arrays) else: samples = fp.read_raw_samples( - param, thin_start=start_index, thin_interval=1, - thin_end=end_index, walkers=walkers, - flatten=False)[param] + param, + thin_start=start_index, + thin_interval=1, + thin_end=end_index, + walkers=walkers, + flatten=False, + )[param] samples = samples.mean(axis=0) - acfs[param] = autocorrelation.calculate_acf( - samples).numpy() + acfs[param] = autocorrelation.calculate_acf(samples).numpy() return acfs -def ensemble_compute_acl(filename, start_index=None, end_index=None, - min_nsamples=10): - """Computes the autocorrleation length for an ensemble MCMC. +def ensemble_compute_acl(filename, start_index=None, end_index=None, min_nsamples=10): + """ + Computes the autocorrleation length for an ensemble MCMC. Parameter values are averaged over all walkers at each iteration. The ACL is then calculated over the averaged chain. If an ACL cannot @@ -908,7 +960,7 @@ def ensemble_compute_acl(filename, start_index=None, end_index=None, to ``inf``. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACLs for. start_index : int, optional @@ -927,13 +979,18 @@ def ensemble_compute_acl(filename, start_index=None, end_index=None, ------- dict A dictionary giving the ACL for each parameter. + """ acls = {} - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: for param in fp.variable_params: samples = fp.read_raw_samples( - param, thin_start=start_index, thin_interval=1, - thin_end=end_index, flatten=False)[param] + param, + thin_start=start_index, + thin_interval=1, + thin_end=end_index, + flatten=False, + )[param] samples = samples.mean(axis=0) # if < min number of samples, just set to inf if samples.size < min_nsamples: @@ -944,5 +1001,5 @@ def ensemble_compute_acl(filename, start_index=None, end_index=None, acl = numpy.inf acls[param] = acl maxacl = numpy.array(list(acls.values())).max() - logging.info("ACT: %s", str(maxacl*fp.thinned_by)) + logging.info("ACT: %s", str(maxacl * fp.thinned_by)) return acls diff --git a/pycbc/inference/sampler/base_multitemper.py b/pycbc/inference/sampler/base_multitemper.py index 84dfeb135f9..164c384e20c 100644 --- a/pycbc/inference/sampler/base_multitemper.py +++ b/pycbc/inference/sampler/base_multitemper.py @@ -21,20 +21,23 @@ # # ============================================================================= # -"""Provides constructor classes provide support for parallel tempered MCMC -samplers.""" - +""" +Provides constructor classes provide support for parallel tempered MCMC +samplers. +""" import logging -import numpy + import h5py +import numpy + from pycbc.filter import autocorrelation from pycbc.inference.io import loadfile -class MultiTemperedSupport(object): - """Provides methods for supporting multi-tempered samplers. - """ +class MultiTemperedSupport: + """Provides methods for supporting multi-tempered samplers.""" + _ntemps = None @property @@ -44,7 +47,8 @@ def ntemps(self): @staticmethod def betas_from_config(cp, section): - """Loads number of temperatures or betas from a config file. + """ + Loads number of temperatures or betas from a config file. This looks in the given section for: @@ -71,15 +75,17 @@ def betas_from_config(cp, section): betas : array The array of betas to use, if a inverse-temperatures-file was provided. + """ - if cp.has_option(section, "ntemps") and \ - cp.has_option(section, "inverse-temperatures-file"): - raise ValueError("Must specify either ntemps or " - "inverse-temperatures-file, not both.") + if cp.has_option(section, "ntemps") and cp.has_option( + section, "inverse-temperatures-file" + ): + raise ValueError( + "Must specify either ntemps or inverse-temperatures-file, not both." + ) if cp.has_option(section, "inverse-temperatures-file"): # get the path of the file containing inverse temperatures values. - inverse_temperatures_file = cp.get(section, - "inverse-temperatures-file") + inverse_temperatures_file = cp.get(section, "inverse-temperatures-file") betas = read_betas_from_hdf(inverse_temperatures_file) ntemps = betas.shape[0] else: @@ -90,12 +96,11 @@ def betas_from_config(cp, section): def read_betas_from_hdf(filename): - """Loads inverse temperatures from the given file. - """ + """Loads inverse temperatures from the given file.""" # get the path of the file containing inverse temperatures values. with h5py.File(filename, "r") as fp: try: - betas = numpy.array(fp.attrs['betas']) + betas = numpy.array(fp.attrs["betas"]) # betas must be in decending order betas = numpy.sort(betas)[::-1] except KeyError: @@ -112,13 +117,15 @@ def read_betas_from_hdf(filename): # -def compute_acf(filename, start_index=None, end_index=None, - chains=None, parameters=None, temps=None): - """Computes the autocorrleation function for independent MCMC chains with +def compute_acf( + filename, start_index=None, end_index=None, chains=None, parameters=None, temps=None +): + """ + Computes the autocorrleation function for independent MCMC chains with parallel tempering. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACFs for. start_index : int, optional @@ -145,9 +152,10 @@ def compute_acf(filename, start_index=None, end_index=None, dict : Dictionary parameter name -> ACF arrays. The arrays have shape ``ntemps x nchains x niterations``. + """ acfs = {} - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: if parameters is None: parameters = fp.variable_params if isinstance(parameters, str): @@ -161,8 +169,13 @@ def compute_acf(filename, start_index=None, end_index=None, subsubacfs = [] for ci in chains: samples = fp.read_raw_samples( - param, thin_start=start_index, thin_interval=1, - thin_end=end_index, chains=ci, temps=tk)[param] + param, + thin_start=start_index, + thin_interval=1, + thin_end=end_index, + chains=ci, + temps=tk, + )[param] thisacf = autocorrelation.calculate_acf(samples).numpy() subsubacfs.append(thisacf) # stack the chains @@ -172,15 +185,15 @@ def compute_acf(filename, start_index=None, end_index=None, return acfs -def compute_acl(filename, start_index=None, end_index=None, - min_nsamples=10): - """Computes the autocorrleation length for independent MCMC chains with +def compute_acl(filename, start_index=None, end_index=None, min_nsamples=10): + """ + Computes the autocorrleation length for independent MCMC chains with parallel tempering. ACLs are calculated separately for each chain. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACLs for. start_index : {None, int} @@ -200,7 +213,9 @@ def compute_acl(filename, start_index=None, end_index=None, dict A dictionary of ntemps x nchains arrays of the ACLs of each parameter. + """ + # following is a convenience function to calculate the acl for each chain # defined here so that we can use map for this below def _getacl(si): @@ -213,15 +228,21 @@ def _getacl(si): if acl <= 0: acl = numpy.inf return acl + acls = {} - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: tidx = numpy.arange(fp.ntemps) for param in fp.variable_params: these_acls = numpy.zeros((fp.ntemps, fp.nchains)) for tk in tidx: samples = fp.read_raw_samples( - param, thin_start=start_index, thin_interval=1, - thin_end=end_index, temps=tk, flatten=False)[param] + param, + thin_start=start_index, + thin_interval=1, + thin_end=end_index, + temps=tk, + flatten=False, + )[param] # flatten out the temperature samples = samples[0, ...] # samples now has shape nchains x maxiters @@ -231,17 +252,20 @@ def _getacl(si): these_acls[tk, :] = list(map(_getacl, samples)) acls[param] = these_acls # report the mean ACL: take the max over the temps and parameters - act = acl_from_raw_acls(acls)*fp.thinned_by + act = acl_from_raw_acls(acls) * fp.thinned_by finite = act[numpy.isfinite(act)] - logging.info("ACTs: min %s, mean (of finite) %s, max %s", - str(act.min()), - str(finite.mean() if finite.size > 0 else numpy.inf), - str(act.max())) + logging.info( + "ACTs: min %s, mean (of finite) %s, max %s", + str(act.min()), + str(finite.mean() if finite.size > 0 else numpy.inf), + str(act.max()), + ) return acls def acl_from_raw_acls(acls): - """Calculates the ACL for one or more chains from a dictionary of ACLs. + """ + Calculates the ACL for one or more chains from a dictionary of ACLs. This is for parallel tempered MCMCs in which the chains are independent of each other. @@ -258,14 +282,22 @@ def acl_from_raw_acls(acls): ------- array The ACL of each chain. + """ return numpy.array(list(acls.values())).max(axis=0).max(axis=0) -def ensemble_compute_acf(filename, start_index=None, end_index=None, - per_walker=False, walkers=None, parameters=None, - temps=None): - """Computes the autocorrleation function for a parallel tempered, ensemble +def ensemble_compute_acf( + filename, + start_index=None, + end_index=None, + per_walker=False, + walkers=None, + parameters=None, + temps=None, +): + """ + Computes the autocorrleation function for a parallel tempered, ensemble MCMC. By default, parameter values are averaged over all walkers at each @@ -305,9 +337,10 @@ def ensemble_compute_acf(filename, start_index=None, end_index=None, ``per-walker`` is True, the arrays will have shape ``ntemps x nwalkers x niterations``. Otherwise, the returned array will have shape ``ntemps x niterations``. + """ acfs = {} - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: if parameters is None: parameters = fp.variable_params if isinstance(parameters, str): @@ -320,14 +353,18 @@ def ensemble_compute_acf(filename, start_index=None, end_index=None, # just call myself with a single walker if walkers is None: walkers = numpy.arange(fp.nwalkers) - arrays = [ensemble_compute_acf(filename, - start_index=start_index, - end_index=end_index, - per_walker=False, - walkers=ii, - parameters=param, - temps=tk)[param][0, :] - for ii in walkers] + arrays = [ + ensemble_compute_acf( + filename, + start_index=start_index, + end_index=end_index, + per_walker=False, + walkers=ii, + parameters=param, + temps=tk, + )[param][0, :] + for ii in walkers + ] # we'll stack all of the walker arrays to make a single # nwalkers x niterations array; when these are stacked # below, we'll get a ntemps x nwalkers x niterations @@ -335,30 +372,34 @@ def ensemble_compute_acf(filename, start_index=None, end_index=None, subacfs.append(numpy.vstack(arrays)) else: samples = fp.read_raw_samples( - param, thin_start=start_index, - thin_interval=1, thin_end=end_index, - walkers=walkers, temps=tk, flatten=False)[param] + param, + thin_start=start_index, + thin_interval=1, + thin_end=end_index, + walkers=walkers, + temps=tk, + flatten=False, + )[param] # contract the walker dimension using the mean, and # flatten the (length 1) temp dimension samples = samples.mean(axis=1)[0, :] - thisacf = autocorrelation.calculate_acf( - samples).numpy() + thisacf = autocorrelation.calculate_acf(samples).numpy() subacfs.append(thisacf) # stack the temperatures acfs[param] = numpy.stack(subacfs) return acfs -def ensemble_compute_acl(filename, start_index=None, end_index=None, - min_nsamples=10): - """Computes the autocorrleation length for a parallel tempered, ensemble +def ensemble_compute_acl(filename, start_index=None, end_index=None, min_nsamples=10): + """ + Computes the autocorrleation length for a parallel tempered, ensemble MCMC. Parameter values are averaged over all walkers at each iteration and temperature. The ACL is then calculated over the averaged chain. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACLs for. start_index : int, optional @@ -377,9 +418,10 @@ def ensemble_compute_acl(filename, start_index=None, end_index=None, ------- dict A dictionary of ntemps-long arrays of the ACLs of each parameter. + """ acls = {} - with loadfile(filename, 'r') as fp: + with loadfile(filename, "r") as fp: if end_index is None: end_index = fp.niterations tidx = numpy.arange(fp.ntemps) @@ -387,8 +429,13 @@ def ensemble_compute_acl(filename, start_index=None, end_index=None, these_acls = numpy.zeros(fp.ntemps) for tk in tidx: samples = fp.read_raw_samples( - param, thin_start=start_index, thin_interval=1, - thin_end=end_index, temps=tk, flatten=False)[param] + param, + thin_start=start_index, + thin_interval=1, + thin_end=end_index, + temps=tk, + flatten=False, + )[param] # contract the walker dimension using the mean, and flatten # the (length 1) temp dimension samples = samples.mean(axis=1)[0, :] @@ -401,16 +448,15 @@ def ensemble_compute_acl(filename, start_index=None, end_index=None, these_acls[tk] = acl acls[param] = these_acls maxacl = numpy.array(list(acls.values())).max() - logging.info("ACT: %s", str(maxacl*fp.thinned_by)) + logging.info("ACT: %s", str(maxacl * fp.thinned_by)) return acls def _get_temps_idx(fp, temps): - """Gets the indices of temperatures to load for computing ACF. - """ + """Gets the indices of temperatures to load for computing ACF.""" if isinstance(temps, int): temps = [temps] - elif temps == 'all': + elif temps == "all": temps = numpy.arange(fp.ntemps) elif temps is None: temps = [0] diff --git a/pycbc/inference/sampler/cpnest.py b/pycbc/inference/sampler/cpnest.py index 71244f0a24d..3c40b1f19f5 100644 --- a/pycbc/inference/sampler/cpnest.py +++ b/pycbc/inference/sampler/cpnest.py @@ -26,17 +26,17 @@ packages for parameter estimation. """ - +import array import logging import os -import array + import cpnest import cpnest.model as cpm -from pycbc.inference.io import (CPNestFile, validate_checkpoint_files) -from .base import (BaseSampler, setup_output) -from .base_mcmc import get_optional_arg_from_config +from pycbc.inference.io import CPNestFile, validate_checkpoint_files +from .base import BaseSampler, setup_output +from .base_mcmc import get_optional_arg_from_config # # ============================================================================= @@ -46,8 +46,10 @@ # ============================================================================= # + class CPNestSampler(BaseSampler): - """This class is used to construct an CPNest sampler from the cpnest + """ + This class is used to construct an CPNest sampler from the cpnest package by John Veitch. Parameters @@ -60,12 +62,21 @@ class CPNestSampler(BaseSampler): A provider of a map function that allows a function call to be run over multiple sets of arguments and possibly maps them to cores/nodes/etc. + """ + name = "cpnest" _io = CPNestFile - def __init__(self, model, nlive, maxmcmc=1000, nthreads=1, verbose=1, - loglikelihood_function=None): + def __init__( + self, + model, + nlive, + maxmcmc=1000, + nthreads=1, + verbose=1, + loglikelihood_function=None, + ): self.model = model self.nlive = nlive self.maxmcmc = maxmcmc @@ -83,11 +94,15 @@ def __init__(self, model, nlive, maxmcmc=1000, nthreads=1, verbose=1, def run(self): out_dir = os.path.dirname(os.path.abspath(self.checkpoint_file)) if self._sampler is None: - self._sampler = cpnest.CPNest(self.model_call, verbose=1, - output=out_dir, - nthreads=self.nthreads, - nlive=self.nlive, - maxmcmc=self.maxmcmc, resume=True) + self._sampler = cpnest.CPNest( + self.model_call, + verbose=1, + output=out_dir, + nthreads=self.nthreads, + nlive=self.nlive, + maxmcmc=self.maxmcmc, + resume=True, + ) res = self._sampler.run() @property @@ -99,25 +114,31 @@ def niterations(self): return len(tuple(self.samples.values())[0]) @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): """ Loads the sampler from the given config file. """ section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) # get the number of live points to use nlive = int(cp.get(section, "nlive")) maxmcmc = int(cp.get(section, "maxmcmc")) nthreads = int(cp.get(section, "nthreads")) verbose = int(cp.get(section, "verbose")) - loglikelihood_function = \ - get_optional_arg_from_config(cp, section, 'loglikelihood-function') - obj = cls(model, nlive=nlive, maxmcmc=maxmcmc, nthreads=nthreads, - verbose=verbose, - loglikelihood_function=loglikelihood_function) + loglikelihood_function = get_optional_arg_from_config( + cp, section, "loglikelihood-function" + ) + obj = cls( + model, + nlive=nlive, + maxmcmc=maxmcmc, + nthreads=nthreads, + verbose=verbose, + loglikelihood_function=loglikelihood_function, + ) setup_output(obj, output_file, check_nsamples=False) if not obj.new_checkpoint: @@ -130,7 +151,7 @@ def checkpoint(self): def finalize(self): logz = self._sampler.NS.logZ dlogz = 0.1 #######FIXME!!!!!############### - logging.info("log Z, dlog Z: {}, {}".format(logz, dlogz)) + logging.info(f"log Z, dlog Z: {logz}, {dlogz}") for fn in [self.checkpoint_file]: with self.io(fn, "a") as fp: fp.write_logevidence(logz, dlogz) @@ -139,44 +160,48 @@ def finalize(self): self.write_results(fn) logging.info("Validating checkpoint and backup files") checkpoint_valid = validate_checkpoint_files( - self.checkpoint_file, self.backup_file, check_nsamples=False) + self.checkpoint_file, self.backup_file, check_nsamples=False + ) if not checkpoint_valid: - raise IOError("error writing to checkpoint file") + raise OSError("error writing to checkpoint file") @property def model_stats(self): - logl = self._sampler.posterior_samples['logL'] - logp = self._sampler.posterior_samples['logPrior'] - return {'loglikelihood': logl, 'logprior': logp} + logl = self._sampler.posterior_samples["logL"] + logp = self._sampler.posterior_samples["logPrior"] + return {"loglikelihood": logl, "logprior": logp} @property def samples(self): - samples_dict = {p: self._sampler.posterior_samples[p] for p in - self.posterior_samples.dtype.names} + samples_dict = { + p: self._sampler.posterior_samples[p] + for p in self.posterior_samples.dtype.names + } return samples_dict - def set_initial_conditions(self, initial_distribution=None, - samples_file=None): - """Sets up the starting point for the sampler. + def set_initial_conditions(self, initial_distribution=None, samples_file=None): + """ + Sets up the starting point for the sampler. Should also set the sampler's random state. """ - pass def resume_from_checkpoint(self): pass def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters - ----------- + ---------- filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples fp.write_samples(self.samples, self.model.variable_params) # write stats @@ -210,7 +235,9 @@ class CPNestModel(cpm.Model): ---------- model : inference.BaseModel instance A model instance from pycbc. + """ + def __init__(self, model, loglikelihood_function=None): if model.sampling_transforms is not None: raise ValueError("CPNest does not support sampling transforms") @@ -218,7 +245,7 @@ def __init__(self, model, loglikelihood_function=None): self.names = list(model.sampling_params) # set up lohlikelihood_function if loglikelihood_function is None: - loglikelihood_function = 'loglikelihood' + loglikelihood_function = "loglikelihood" self.loglikelihood_function = loglikelihood_function bounds = {} for dist in model.prior_distribution.distributions: @@ -227,10 +254,12 @@ def __init__(self, model, loglikelihood_function=None): def new_point(self): point = self.model.prior_rvs() - return cpm.LivePoint(list(self.model.sampling_params), - array.array('d', [point[p][0] for p in self.model.sampling_params])) + return cpm.LivePoint( + list(self.model.sampling_params), + array.array("d", [point[p][0] for p in self.model.sampling_params]), + ) - def log_prior(self,xx): + def log_prior(self, xx): self.model.update(**xx) return self.model.logprior diff --git a/pycbc/inference/sampler/dummy.py b/pycbc/inference/sampler/dummy.py index 460ce865f32..55ee9852a99 100644 --- a/pycbc/inference/sampler/dummy.py +++ b/pycbc/inference/sampler/dummy.py @@ -1,18 +1,20 @@ -""" Dummy class when no actual sampling is needed, but we may want to do +""" +Dummy class when no actual sampling is needed, but we may want to do some reconstruction supported by the likelihood model. """ import numpy -from pycbc.inference.io import PosteriorFile from pycbc.inference import models +from pycbc.inference.io import PosteriorFile from pycbc.pool import choose_pool -from .base import (BaseSampler, setup_output) +from .base import BaseSampler, setup_output def call_reconstruct(iteration): - """ Accessor to update the global model and call its reconstruction + """ + Accessor to update the global model and call its reconstruction routine. """ models._global_instance.update() @@ -20,17 +22,21 @@ def call_reconstruct(iteration): class DummySampler(BaseSampler): - """Dummy sampler for not doing sampling + """ + Dummy sampler for not doing sampling Parameters ---------- model : Model An instance of a model from ``pycbc.inference.models``. + """ - name = 'dummy' - def __init__(self, model, *args, nprocesses=1, use_mpi=False, - num_samples=1000, **kwargs): + name = "dummy" + + def __init__( + self, model, *args, nprocesses=1, use_mpi=False, num_samples=1000, **kwargs + ): super().__init__(model, *args) models._global_instance = model @@ -40,18 +46,17 @@ def __init__(self, model, *args, nprocesses=1, use_mpi=False, self.meta = {} @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): - """This should initialize the sampler given a config file. - """ - kwargs = {k: cp.get('sampler', k) for k in cp.options('sampler')} + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): + """This should initialize the sampler given a config file.""" + kwargs = {k: cp.get("sampler", k) for k in cp.options("sampler")} obj = cls(model, nprocesses=nprocesses, use_mpi=use_mpi, **kwargs) setup_output(obj, output_file, check_nsamples=False, validate=False) return obj @property def samples(self): - """A dict mapping variable_params to arrays of samples currently + """ + A dict mapping variable_params to arrays of samples currently in memory. The dictionary may also contain sampling_params. The sample arrays may have any shape, and may or may not be thinned. @@ -63,10 +68,8 @@ def model_stats(self): pass def run(self): - samples = self.pool.map(call_reconstruct, - range(self.num_samples)) - self._samples = {k: numpy.array([x[k] for x in samples]) - for k in samples[0]} + samples = self.pool.map(call_reconstruct, range(self.num_samples)) + self._samples = {k: numpy.array([x[k] for x in samples]) for k in samples[0]} def finalize(self): with self.io(self.checkpoint_file, "a") as fp: @@ -78,7 +81,8 @@ def finalize(self): @property def io(self): - """A class that inherits from ``BaseInferenceFile`` to handle IO with + """ + A class that inherits from ``BaseInferenceFile`` to handle IO with an hdf file. This should be a class, not an instance of class, so that the sampler diff --git a/pycbc/inference/sampler/dynesty.py b/pycbc/inference/sampler/dynesty.py index 33edce0db51..16c9fbbdf2b 100644 --- a/pycbc/inference/sampler/dynesty.py +++ b/pycbc/inference/sampler/dynesty.py @@ -28,15 +28,18 @@ import logging import time + +import dynesty +import dynesty.dynesty +import dynesty.nestedsamplers import numpy -import dynesty, dynesty.dynesty, dynesty.nestedsamplers + +from pycbc.inference.io import DynestyFile, loadfile, validate_checkpoint_files from pycbc.pool import choose_pool -from pycbc.inference.io import (DynestyFile, validate_checkpoint_files, - loadfile) -from .base import (BaseSampler, setup_output) -from .base_mcmc import get_optional_arg_from_config -from .base_cube import setup_calls +from .base import BaseSampler, setup_output +from .base_cube import setup_calls +from .base_mcmc import get_optional_arg_from_config # # ============================================================================= @@ -46,8 +49,10 @@ # ============================================================================= # + class DynestySampler(BaseSampler): - """This class is used to construct an Dynesty sampler from the dynesty + """ + This class is used to construct an Dynesty sampler from the dynesty package. Parameters @@ -60,25 +65,33 @@ class DynestySampler(BaseSampler): A provider of a map function that allows a function call to be run over multiple sets of arguments and possibly maps them to cores/nodes/etc. + """ + name = "dynesty" _io = DynestyFile - def __init__(self, model, nlive, nprocesses=1, - checkpoint_time_interval=None, maxcall=None, - loglikelihood_function=None, use_mpi=False, - no_save_state=False, - run_kwds=None, - extra_kwds=None, - internal_kwds=None, - **kwargs): + def __init__( + self, + model, + nlive, + nprocesses=1, + checkpoint_time_interval=None, + maxcall=None, + loglikelihood_function=None, + use_mpi=False, + no_save_state=False, + run_kwds=None, + extra_kwds=None, + internal_kwds=None, + **kwargs, + ): self.model = model self.no_save_state = no_save_state log_likelihood_call, prior_call = setup_calls( - model, - loglikelihood_function=loglikelihood_function, - copy_prior=True) + model, loglikelihood_function=loglikelihood_function, copy_prior=True + ) # Set up the pool self.pool = choose_pool(mpi=use_mpi, processes=nprocesses) @@ -97,9 +110,12 @@ def __init__(self, model, nlive, nprocesses=1, self.run_with_checkpoint = True if self.maxcall is None: self.maxcall = 5000 * self.pool.size - logging.info("Checkpointing enabled, will verify every %s calls" - " and try to checkpoint every %s seconds", - self.maxcall, self.checkpoint_time_interval) + logging.info( + "Checkpointing enabled, will verify every %s calls" + " and try to checkpoint every %s seconds", + self.maxcall, + self.checkpoint_time_interval, + ) else: self.run_with_checkpoint = False @@ -108,7 +124,7 @@ def __init__(self, model, nlive, nprocesses=1, cyclic = self.model.prior_distribution.cyclic for i, param in enumerate(self.variable_params): if param in cyclic: - logging.info('Param: %s will be cyclic', param) + logging.info("Param: %s will be cyclic", param) periodic.append(i) if len(periodic) == 0: @@ -126,44 +142,60 @@ def __init__(self, model, nlive, nprocesses=1, if len(reflective) == 0: reflective = None - if 'sample' in extra_kwds: - if 'rwalk2' in extra_kwds['sample']: + if "sample" in extra_kwds: + if "rwalk2" in extra_kwds["sample"]: dynesty.dynesty._SAMPLING["rwalk"] = sample_rwalk_mod dynesty.nestedsamplers._SAMPLING["rwalk"] = sample_rwalk_mod - extra_kwds['sample'] = 'rwalk' + extra_kwds["sample"] = "rwalk" if self.nlive < 0: # Interpret a negative input value for the number of live points # (which is clearly an invalid input in all senses) # as the desire to dynamically determine that number - self._sampler = dynesty.DynamicNestedSampler(log_likelihood_call, - prior_call, self.ndim, - pool=self.pool, - reflective=reflective, - periodic=periodic, - **extra_kwds) + self._sampler = dynesty.DynamicNestedSampler( + log_likelihood_call, + prior_call, + self.ndim, + pool=self.pool, + reflective=reflective, + periodic=periodic, + **extra_kwds, + ) self.run_with_checkpoint = False - logging.info("Checkpointing not currently supported with" - "DYNAMIC nested sampler") + logging.info( + "Checkpointing not currently supported withDYNAMIC nested sampler" + ) else: - self._sampler = dynesty.NestedSampler(log_likelihood_call, - prior_call, self.ndim, - nlive=self.nlive, - reflective=reflective, - periodic=periodic, - pool=self.pool, **extra_kwds) + self._sampler = dynesty.NestedSampler( + log_likelihood_call, + prior_call, + self.ndim, + nlive=self.nlive, + reflective=reflective, + periodic=periodic, + pool=self.pool, + **extra_kwds, + ) self._sampler.kwargs.update(internal_kwds) # properties of the internal sampler which should not be pickled - self.no_pickle = ['loglikelihood', - 'prior_transform', - 'propose_point', - 'update_proposal', - '_UPDATE', '_PROPOSE', - 'evolve_point', 'use_pool', 'queue_size', - 'use_pool_ptform', 'use_pool_logl', - 'use_pool_evolve', 'use_pool_update', - 'pool', 'M'] + self.no_pickle = [ + "loglikelihood", + "prior_transform", + "propose_point", + "update_proposal", + "_UPDATE", + "_PROPOSE", + "evolve_point", + "use_pool", + "queue_size", + "use_pool_ptform", + "use_pool_logl", + "use_pool_evolve", + "use_pool_update", + "pool", + "M", + ] def run(self): diff_niter = 1 @@ -172,7 +204,7 @@ def run(self): t0 = time.time() it = self._sampler.it - logging.info('Starting from iteration: %s', it) + logging.info("Starting from iteration: %s", it) while diff_niter != 0: self._sampler.run_nested(maxcall=self.maxcall, **self.run_kwds) @@ -181,7 +213,7 @@ def run(self): logging.info("Checking if we should checkpoint: %.2f s", delta_t) if delta_t >= self.checkpoint_time_interval: - logging.info('Checkpointing N={}'.format(n_checkpointing)) + logging.info(f"Checkpointing N={n_checkpointing}") self.checkpoint() n_checkpointing += 1 t0 = time.time() @@ -198,9 +230,17 @@ def niterations(self): return len(tuple(self.samples.values())[0]) @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False, loglikelihood_function=None): - """Loads the sampler from the given config file. Many options are + def from_config( + cls, + cp, + model, + output_file=None, + nprocesses=1, + use_mpi=False, + loglikelihood_function=None, + ): + """ + Loads the sampler from the given config file. Many options are directly passed to the underlying dynesty sampler, see the official dynesty documentation for more details on these. @@ -269,81 +309,91 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, ------- DynestySampler : The sampler instance. + """ section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) # get the number of live points to use nlive = int(cp.get(section, "nlive")) - loglikelihood_function = \ - get_optional_arg_from_config(cp, section, 'loglikelihood-function') + loglikelihood_function = get_optional_arg_from_config( + cp, section, "loglikelihood-function" + ) - no_save_state = cp.has_option(section, 'no-save-state') + no_save_state = cp.has_option(section, "no-save-state") # optional run_nested arguments for dynesty - rargs = {'maxiter': int, - 'dlogz': float, - 'logl_max': float, - 'n_effective': int, - } + rargs = { + "maxiter": int, + "dlogz": float, + "logl_max": float, + "n_effective": int, + } # optional arguments for dynesty - cargs = {'bound': str, - 'bootstrap': int, - 'enlarge': float, - 'update_interval': float, - 'sample': str, - 'first_update_min_ncall': int, - 'first_update_min_eff': float, - 'walks': int, - } + cargs = { + "bound": str, + "bootstrap": int, + "enlarge": float, + "update_interval": float, + "sample": str, + "first_update_min_ncall": int, + "first_update_min_eff": float, + "walks": int, + } # optional arguments that must be set internally internal_args = { - 'maxmcmc': int, - 'nact': int, - } + "maxmcmc": int, + "nact": int, + } extra = {} run_extra = {} internal_extra = {} - for args, argt in [(extra, cargs), - (run_extra, rargs), - (internal_extra, internal_args), - ]: + for args, argt in [ + (extra, cargs), + (run_extra, rargs), + (internal_extra, internal_args), + ]: for karg in argt: if cp.has_option(section, karg): args[karg] = argt[karg](cp.get(section, karg)) - #This arg needs to be a dict + # This arg needs to be a dict first_update = {} - if 'first_update_min_ncall' in extra: - first_update['min_ncall'] = extra.pop('first_update_min_ncall') - logging.info('First update: min_ncall:%s', - first_update['min_ncall']) - if 'first_update_min_eff' in extra: - first_update['min_eff'] = extra.pop('first_update_min_eff') - logging.info('First update: min_eff:%s', first_update['min_eff']) - extra['first_update'] = first_update + if "first_update_min_ncall" in extra: + first_update["min_ncall"] = extra.pop("first_update_min_ncall") + logging.info("First update: min_ncall:%s", first_update["min_ncall"]) + if "first_update_min_eff" in extra: + first_update["min_eff"] = extra.pop("first_update_min_eff") + logging.info("First update: min_eff:%s", first_update["min_eff"]) + extra["first_update"] = first_update # populate options for checkpointing checkpoint_time_interval = None maxcall = None - if cp.has_option(section, 'checkpoint_time_interval'): - ck_time = float(cp.get(section, 'checkpoint_time_interval')) + if cp.has_option(section, "checkpoint_time_interval"): + ck_time = float(cp.get(section, "checkpoint_time_interval")) checkpoint_time_interval = ck_time - if cp.has_option(section, 'maxcall'): - maxcall = int(cp.get(section, 'maxcall')) - - obj = cls(model, nlive=nlive, nprocesses=nprocesses, - loglikelihood_function=loglikelihood_function, - checkpoint_time_interval=checkpoint_time_interval, - maxcall=maxcall, - no_save_state=no_save_state, - use_mpi=use_mpi, run_kwds=run_extra, - extra_kwds=extra, - internal_kwds=internal_extra,) + if cp.has_option(section, "maxcall"): + maxcall = int(cp.get(section, "maxcall")) + + obj = cls( + model, + nlive=nlive, + nprocesses=nprocesses, + loglikelihood_function=loglikelihood_function, + checkpoint_time_interval=checkpoint_time_interval, + maxcall=maxcall, + no_save_state=no_save_state, + use_mpi=use_mpi, + run_kwds=run_extra, + extra_kwds=extra, + internal_kwds=internal_extra, + ) setup_output(obj, output_file, check_nsamples=False) if not obj.new_checkpoint: @@ -351,8 +401,7 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, return obj def checkpoint(self): - """Checkpoint function for dynesty sampler - """ + """Checkpoint function for dynesty sampler""" # Dynesty has its own __getstate__ which deletes # random state information and the pool saved = {} @@ -376,7 +425,7 @@ def checkpoint(self): def resume_from_checkpoint(self): try: - with loadfile(self.checkpoint_file, 'r') as fp: + with loadfile(self.checkpoint_file, "r") as fp: sampler = fp.read_pickled_data_from_checkpoint_file() for key in sampler.__dict__: @@ -385,27 +434,24 @@ def resume_from_checkpoint(self): setattr(self._sampler, key, value) self.set_state_from_file(self.checkpoint_file) - logging.info("Found valid checkpoint file: %s", - self.checkpoint_file) + logging.info("Found valid checkpoint file: %s", self.checkpoint_file) except Exception as e: print(e) logging.info("Failed to load checkpoint file") def set_state_from_file(self, filename): - """Sets the state of the sampler back to the instance saved in a file. - """ - with self.io(filename, 'r') as fp: + """Sets the state of the sampler back to the instance saved in a file.""" + with self.io(filename, "r") as fp: state = fp.read_random_state() # Dynesty handles most randomeness through rstate which is # pickled along with the class instance numpy.random.set_state(state) def finalize(self): - """Finalze and write it to the results file - """ + """Finalze and write it to the results file""" logz = self._sampler.results.logz[-1:][0] dlogz = self._sampler.results.logzerr[-1:][0] - logging.info("log Z, dlog Z: {}, {}".format(logz, dlogz)) + logging.info(f"log Z, dlog Z: {logz}, {dlogz}") if self.no_save_state: self.write_results(self.checkpoint_file) @@ -413,42 +459,43 @@ def finalize(self): self.checkpoint() logging.info("Validating checkpoint and backup files") checkpoint_valid = validate_checkpoint_files( - self.checkpoint_file, self.backup_file, check_nsamples=False) + self.checkpoint_file, self.backup_file, check_nsamples=False + ) if not checkpoint_valid: - raise IOError("error writing to checkpoint file") + raise OSError("error writing to checkpoint file") @property def samples(self): - """Returns raw nested samples - """ + """Returns raw nested samples""" results = self._sampler.results samples = results.samples nest_samp = {} for i, param in enumerate(self.variable_params): nest_samp[param] = samples[:, i] - nest_samp['logwt'] = results.logwt - nest_samp['loglikelihood'] = results.logl + nest_samp["logwt"] = results.logwt + nest_samp["loglikelihood"] = results.logl return nest_samp - def set_initial_conditions(self, initial_distribution=None, - samples_file=None): - """Sets up the starting point for the sampler. + def set_initial_conditions(self, initial_distribution=None, samples_file=None): + """ + Sets up the starting point for the sampler. Should also set the sampler's random state. """ - pass def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters - ----------- + ---------- filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # Write nested samples fp.write_raw_samples(self.samples) @@ -464,7 +511,7 @@ def model_stats(self): @property def logz(self): """ - return bayesian evidence estimated by + Return bayesian evidence estimated by dynesty sampler """ return self._sampler.results.logz[-1:][0] @@ -472,45 +519,45 @@ def logz(self): @property def logz_err(self): """ - return error in bayesian evidence estimated by + Return error in bayesian evidence estimated by dynesty sampler """ return self._sampler.results.logzerr[-1:][0] def sample_rwalk_mod(args): - """ Modified version of dynesty.sampling.sample_rwalk + """ + Modified version of dynesty.sampling.sample_rwalk - Adapted from version used in bilby/dynesty + Adapted from version used in bilby/dynesty """ try: # dynesty <= 1.1 - from dynesty.utils import unitcheck, reflect + from dynesty.utils import reflect, unitcheck # Unzipping. - (u, loglstar, axes, scale, - prior_transform, loglikelihood, kwargs) = args + (u, loglstar, axes, scale, prior_transform, loglikelihood, kwargs) = args except ImportError: # dynest >= 1.2 - from dynesty.utils import unitcheck, apply_reflect as reflect + from dynesty.utils import apply_reflect as reflect + from dynesty.utils import unitcheck - (u, loglstar, axes, scale, - prior_transform, loglikelihood, _, kwargs) = args + (u, loglstar, axes, scale, prior_transform, loglikelihood, _, kwargs) = args rstate = numpy.random # Bounds - nonbounded = kwargs.get('nonbounded', None) - periodic = kwargs.get('periodic', None) - reflective = kwargs.get('reflective', None) + nonbounded = kwargs.get("nonbounded", None) + periodic = kwargs.get("periodic", None) + reflective = kwargs.get("reflective", None) # Setup. n = len(u) - walks = kwargs.get('walks', 10 * n) # minimum number of steps - maxmcmc = kwargs.get('maxmcmc', 2000) # Maximum number of steps - nact = kwargs.get('nact', 5) # Number of ACT - old_act = kwargs.get('old_act', walks) + walks = kwargs.get("walks", 10 * n) # minimum number of steps + maxmcmc = kwargs.get("maxmcmc", 2000) # Maximum number of steps + nact = kwargs.get("nact", 5) # Number of ACT + old_act = kwargs.get("old_act", walks) # Initialize internal variables accept = 0 @@ -580,30 +627,31 @@ def sample_rwalk_mod(args): if accept + reject > walks: act = estimate_nmcmc( accept_ratio=accept / (accept + reject + nfail), - old_act=old_act, maxmcmc=maxmcmc) + old_act=old_act, + maxmcmc=maxmcmc, + ) # If we've taken too many likelihood evaluations then break if accept + reject > maxmcmc: logging.warning( - "Hit maximum number of walks {} with accept={}, reject={}, " - "and nfail={} try increasing maxmcmc" - .format(maxmcmc, accept, reject, nfail)) + f"Hit maximum number of walks {maxmcmc} with accept={accept}, reject={reject}, " + f"and nfail={nfail} try increasing maxmcmc" + ) break # If the act is finite, pick randomly from within the chain - if numpy.isfinite(act) and int(.5 * nact * act) < len(u_list): - idx = numpy.random.randint(int(.5 * nact * act), len(u_list)) + if numpy.isfinite(act) and int(0.5 * nact * act) < len(u_list): + idx = numpy.random.randint(int(0.5 * nact * act), len(u_list)) u = u_list[idx] v = v_list[idx] logl = logl_list[idx] else: - logging.debug("Unable to find a new point using walk: " - "returning a random point") + logging.debug("Unable to find a new point using walk: returning a random point") u = numpy.random.uniform(size=n) v = prior_transform(u) logl = loglikelihood(v) - blob = {'accept': accept, 'reject': reject, 'fail': nfail, 'scale': scale} + blob = {"accept": accept, "reject": reject, "fail": nfail, "scale": scale} kwargs["old_act"] = act ncall = accept + reject @@ -611,7 +659,8 @@ def sample_rwalk_mod(args): def estimate_nmcmc(accept_ratio, old_act, maxmcmc, safety=5, tau=None): - """Estimate autocorrelation length of chain using acceptance fraction + """ + Estimate autocorrelation length of chain using acceptance fraction Using ACL = (2/acc) - 1 multiplied by a safety margin. Code adapated from CPNest: @@ -631,6 +680,7 @@ def estimate_nmcmc(accept_ratio, old_act, maxmcmc, safety=5, tau=None): A safety factor applied in the calculation tau: int (optional) The ACT, if given, otherwise estimated. + """ if tau is None: tau = maxmcmc / safety @@ -638,10 +688,8 @@ def estimate_nmcmc(accept_ratio, old_act, maxmcmc, safety=5, tau=None): if accept_ratio == 0.0: Nmcmc_exact = (1 + 1 / tau) * old_act else: - Nmcmc_exact = ( - (1. - 1. / tau) * old_act + - (safety / tau) * (2. / accept_ratio - 1.) + Nmcmc_exact = (1.0 - 1.0 / tau) * old_act + (safety / tau) * ( + 2.0 / accept_ratio - 1.0 ) Nmcmc_exact = float(min(Nmcmc_exact, maxmcmc)) return max(safety, int(Nmcmc_exact)) - diff --git a/pycbc/inference/sampler/emcee.py b/pycbc/inference/sampler/emcee.py index 7409887b81d..72c9ad17ae8 100644 --- a/pycbc/inference/sampler/emcee.py +++ b/pycbc/inference/sampler/emcee.py @@ -26,20 +26,24 @@ packages for parameter estimation. """ - -import numpy import emcee -from pycbc.pool import choose_pool +import numpy -from .base import (BaseSampler, setup_output) -from .base_mcmc import (BaseMCMC, EnsembleSupport, - ensemble_compute_acf, ensemble_compute_acl, - raw_samples_to_dict, - blob_data_to_dict, get_optional_arg_from_config) -from ..burn_in import EnsembleMCMCBurnInTests from pycbc.inference.io import EmceeFile -from .. import models +from pycbc.pool import choose_pool +from .. import models +from ..burn_in import EnsembleMCMCBurnInTests +from .base import BaseSampler, setup_output +from .base_mcmc import ( + BaseMCMC, + EnsembleSupport, + blob_data_to_dict, + ensemble_compute_acf, + ensemble_compute_acl, + get_optional_arg_from_config, + raw_samples_to_dict, +) # # ============================================================================= @@ -49,12 +53,13 @@ # ============================================================================= # -if emcee.__version__ >= '3.0.0': +if emcee.__version__ >= "3.0.0": raise ImportError class EmceeEnsembleSampler(EnsembleSupport, BaseMCMC, BaseSampler): - """This class is used to construct an MCMC sampler from the emcee + """ + This class is used to construct an MCMC sampler from the emcee package's EnsembleSampler. Parameters @@ -67,19 +72,28 @@ class EmceeEnsembleSampler(EnsembleSupport, BaseMCMC, BaseSampler): A provider of a map function that allows a function call to be run over multiple sets of arguments and possibly maps them to cores/nodes/etc. + """ + name = "emcee" _io = EmceeFile burn_in_class = EnsembleMCMCBurnInTests - def __init__(self, model, nwalkers, - checkpoint_interval=None, checkpoint_signal=None, - logpost_function=None, nprocesses=1, use_mpi=False): + def __init__( + self, + model, + nwalkers, + checkpoint_interval=None, + checkpoint_signal=None, + logpost_function=None, + nprocesses=1, + use_mpi=False, + ): self.model = model # create a wrapper for calling the model if logpost_function is None: - logpost_function = 'logposterior' + logpost_function = "logposterior" model_call = models.CallModel(model, logpost_function) # these are used to help paralleize over multiple cores / MPI @@ -90,8 +104,7 @@ def __init__(self, model, nwalkers, # set up emcee self.nwalkers = nwalkers ndim = len(model.variable_params) - self._sampler = emcee.EnsembleSampler(nwalkers, ndim, model_call, - pool=pool) + self._sampler = emcee.EnsembleSampler(nwalkers, ndim, model_call, pool=pool) # emcee uses it's own internal random number generator; we'll set it # to have the same state as the numpy generator rstate = numpy.random.get_state() @@ -109,7 +122,8 @@ def base_shape(self): @property def samples(self): - """A dict mapping ``variable_params`` to arrays of samples currently + """ + A dict mapping ``variable_params`` to arrays of samples currently in memory. The arrays have shape ``nwalkers x niterations``. @@ -121,7 +135,8 @@ def samples(self): @property def model_stats(self): - """A dict mapping the model's ``default_stats`` to arrays of values. + """ + A dict mapping the model's ``default_stats`` to arrays of values. The returned array has shape ``nwalkers x niterations``. """ @@ -129,8 +144,7 @@ def model_stats(self): return blob_data_to_dict(stats, self._sampler.blobs) def clear_samples(self): - """Clears the samples and stats from memory. - """ + """Clears the samples and stats from memory.""" # store the iteration that the clear is occuring on self._lastclear = self.niterations self._itercounter = 0 @@ -139,9 +153,8 @@ def clear_samples(self): self._sampler.clear_blobs() def set_state_from_file(self, filename): - """Sets the state of the sampler back to the instance saved in a file. - """ - with self.io(filename, 'r') as fp: + """Sets the state of the sampler back to the instance saved in a file.""" + with self.io(filename, "r") as fp: rstate = fp.read_random_state() # set the numpy random state numpy.random.set_state(rstate) @@ -149,12 +162,14 @@ def set_state_from_file(self, filename): self._sampler.random_state = rstate def run_mcmc(self, niterations): - """Advance the ensemble for a number of samples. + """ + Advance the ensemble for a number of samples. Parameters ---------- niterations : int Number of iterations to run the sampler for. + """ pos = self._pos if pos is None: @@ -165,36 +180,41 @@ def run_mcmc(self, niterations): self._pos = p def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters - ----------- + ---------- filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples - fp.write_samples(self.samples, - parameters=self.model.variable_params, - last_iteration=self.niterations) + fp.write_samples( + self.samples, + parameters=self.model.variable_params, + last_iteration=self.niterations, + ) # write stats - fp.write_samples(self.model_stats, - last_iteration=self.niterations) + fp.write_samples(self.model_stats, last_iteration=self.niterations) # write accpetance fp.write_acceptance_fraction(self._sampler.acceptance_fraction) # write random state fp.write_random_state(state=self._sampler.random_state) def finalize(self): - """All data is written by the last checkpoint in the run method, so - this just passes.""" - pass + """ + All data is written by the last checkpoint in the run method, so + this just passes. + """ @staticmethod def compute_acf(filename, **kwargs): - r"""Computes the autocorrelation function. + r""" + Computes the autocorrelation function. Calls :py:func:`base_mcmc.ensemble_compute_acf`; see that function for details. @@ -213,18 +233,20 @@ def compute_acf(filename, **kwargs): Dictionary of arrays giving the ACFs for each parameter. If ``per-walker`` is True, the arrays will have shape ``nwalkers x niterations``. + """ return ensemble_compute_acf(filename, **kwargs) @staticmethod def compute_acl(filename, **kwargs): - r"""Computes the autocorrelation length. + r""" + Computes the autocorrelation length. Calls :py:func:`base_mcmc.ensemble_compute_acl`; see that function for details. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACLs for. \**kwargs : @@ -235,29 +257,34 @@ def compute_acl(filename, **kwargs): ------- dict A dictionary giving the ACL for each parameter. + """ return ensemble_compute_acl(filename, **kwargs) @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): """Loads the sampler from the given config file.""" section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) # get the number of walkers to use nwalkers = int(cp.get(section, "nwalkers")) # get the checkpoint interval, if it's specified checkpoint_interval = cls.checkpoint_from_config(cp, section) checkpoint_signal = cls.ckpt_signal_from_config(cp, section) # get the logpost function - lnpost = get_optional_arg_from_config(cp, section, 'logpost-function') - obj = cls(model, nwalkers, - checkpoint_interval=checkpoint_interval, - checkpoint_signal=checkpoint_signal, - logpost_function=lnpost, nprocesses=nprocesses, - use_mpi=use_mpi) + lnpost = get_optional_arg_from_config(cp, section, "logpost-function") + obj = cls( + model, + nwalkers, + checkpoint_interval=checkpoint_interval, + checkpoint_signal=checkpoint_signal, + logpost_function=lnpost, + nprocesses=nprocesses, + use_mpi=use_mpi, + ) # set target obj.set_target_from_config(cp, section) # add burn-in if it's specified diff --git a/pycbc/inference/sampler/emcee_pt.py b/pycbc/inference/sampler/emcee_pt.py index 45338bd6a5b..054c3bbb85f 100644 --- a/pycbc/inference/sampler/emcee_pt.py +++ b/pycbc/inference/sampler/emcee_pt.py @@ -20,32 +20,39 @@ """ import logging -import numpy + import emcee +import numpy +from pycbc.inference.io import EmceePTFile from pycbc.pool import choose_pool -from .base import (BaseSampler, setup_output) -from .base_mcmc import (BaseMCMC, EnsembleSupport, raw_samples_to_dict, - get_optional_arg_from_config) -from .base_multitemper import (MultiTemperedSupport, - ensemble_compute_acf, ensemble_compute_acl) -from ..burn_in import EnsembleMultiTemperedMCMCBurnInTests -from pycbc.inference.io import EmceePTFile from .. import models - +from ..burn_in import EnsembleMultiTemperedMCMCBurnInTests +from .base import BaseSampler, setup_output +from .base_mcmc import ( + BaseMCMC, + EnsembleSupport, + get_optional_arg_from_config, + raw_samples_to_dict, +) +from .base_multitemper import ( + MultiTemperedSupport, + ensemble_compute_acf, + ensemble_compute_acl, +) # This is a hack that will allow us to continue using emcee's abandoned # PTSampler, which relied on `numpy.float`, until the end of time. -numpy.float = float +float = float -if emcee.__version__ >= '3.0.0': +if emcee.__version__ >= "3.0.0": raise ImportError -class EmceePTSampler(MultiTemperedSupport, EnsembleSupport, BaseMCMC, - BaseSampler): - """This class is used to construct a parallel-tempered MCMC sampler from +class EmceePTSampler(MultiTemperedSupport, EnsembleSupport, BaseMCMC, BaseSampler): + """ + This class is used to construct a parallel-tempered MCMC sampler from the emcee package's PTSampler. Parameters @@ -70,25 +77,36 @@ class EmceePTSampler(MultiTemperedSupport, EnsembleSupport, BaseMCMC, use_mpi : bool, optional Use MPI for parallelization. Default (False) will use python's multiprocessing. + """ + name = "emcee_pt" _io = EmceePTFile burn_in_class = EnsembleMultiTemperedMCMCBurnInTests - def __init__(self, model, ntemps, nwalkers, betas=None, - checkpoint_interval=None, checkpoint_signal=None, - loglikelihood_function=None, - nprocesses=1, use_mpi=False): + def __init__( + self, + model, + ntemps, + nwalkers, + betas=None, + checkpoint_interval=None, + checkpoint_signal=None, + loglikelihood_function=None, + nprocesses=1, + use_mpi=False, + ): self.model = model # create a wrapper for calling the model if loglikelihood_function is None: - loglikelihood_function = 'loglikelihood' + loglikelihood_function = "loglikelihood" # frustratingly, emcee_pt does not support blob data, so we have to # turn it off - model_call = models.CallModel(model, loglikelihood_function, - return_all_stats=False) + model_call = models.CallModel( + model, loglikelihood_function, return_all_stats=False + ) # these are used to help paralleize over multiple cores / MPI models._global_instance = model_call @@ -99,9 +117,9 @@ def __init__(self, model, ntemps, nwalkers, betas=None, # construct the sampler: PTSampler needs the likelihood and prior # functions separately ndim = len(model.variable_params) - self._sampler = emcee.PTSampler(ntemps, nwalkers, ndim, - model_call, prior_call, pool=self.pool, - betas=betas) + self._sampler = emcee.PTSampler( + ntemps, nwalkers, ndim, model_call, prior_call, pool=self.pool, betas=betas + ) self.nwalkers = nwalkers self._ntemps = ntemps self._checkpoint_interval = checkpoint_interval @@ -113,7 +131,10 @@ def io(self): @property def base_shape(self): - return (self.ntemps, self.nwalkers,) + return ( + self.ntemps, + self.nwalkers, + ) @property def betas(self): @@ -121,7 +142,8 @@ def betas(self): @staticmethod def compute_acf(filename, **kwargs): - r"""Computes the autocorrelation function. + r""" + Computes the autocorrelation function. Calls :py:func:`base_multitemper.ensemble_compute_acf`; see that function for details. @@ -141,18 +163,20 @@ def compute_acf(filename, **kwargs): ``per-walker=True`` is passed as a keyword argument, the arrays will have shape ``ntemps x nwalkers x niterations``. Otherwise, the returned array will have shape ``ntemps x niterations``. + """ return ensemble_compute_acf(filename, **kwargs) @staticmethod def compute_acl(filename, **kwargs): - r"""Computes the autocorrelation length. + r""" + Computes the autocorrelation length. Calls :py:func:`base_multitemper.ensemble_compute_acl`; see that function for details. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACLs for. \**kwargs : @@ -163,13 +187,14 @@ def compute_acl(filename, **kwargs): ------- dict A dictionary of ntemps-long arrays of the ACLs of each parameter. + """ return ensemble_compute_acl(filename, **kwargs) @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): - """Loads the sampler from the given config file. + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): + """ + Loads the sampler from the given config file. The following options are retrieved in the ``[sampler]`` section: @@ -234,11 +259,13 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, ------- EmceePTSampler : The sampler instance. + """ section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) # get the number of walkers to use nwalkers = int(cp.get(section, "nwalkers")) # get the temps/betas @@ -247,12 +274,18 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, checkpoint_interval = cls.checkpoint_from_config(cp, section) checkpoint_signal = cls.ckpt_signal_from_config(cp, section) # get the loglikelihood function - logl = get_optional_arg_from_config(cp, section, 'logl-function') - obj = cls(model, ntemps, nwalkers, betas=betas, - checkpoint_interval=checkpoint_interval, - checkpoint_signal=checkpoint_signal, - loglikelihood_function=logl, nprocesses=nprocesses, - use_mpi=use_mpi) + logl = get_optional_arg_from_config(cp, section, "logl-function") + obj = cls( + model, + ntemps, + nwalkers, + betas=betas, + checkpoint_interval=checkpoint_interval, + checkpoint_signal=checkpoint_signal, + loglikelihood_function=logl, + nprocesses=nprocesses, + use_mpi=use_mpi, + ) # set target obj.set_target_from_config(cp, section) # add burn-in if it's specified @@ -269,7 +302,8 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, @property def samples(self): - """A dict mapping ``variable_params`` to arrays of samples currently + """ + A dict mapping ``variable_params`` to arrays of samples currently in memory. The arrays have shape ``ntemps x nwalkers x niterations``. @@ -281,7 +315,8 @@ def samples(self): @property def model_stats(self): - """Returns the log likelihood ratio and log prior as a dict of arrays. + """ + Returns the log likelihood ratio and log prior as a dict of arrays. The returned array has shape ntemps x nwalkers x niterations. @@ -303,12 +338,10 @@ def model_stats(self): # get prior from posterior logp = self._sampler.lnprobability - logl logjacobian = numpy.zeros(logp.shape) - return {'loglikelihood': logl, 'logprior': logp, - 'logjacobian': logjacobian} + return {"loglikelihood": logl, "logprior": logp, "logjacobian": logjacobian} def clear_samples(self): - """Clears the chain and blobs from memory. - """ + """Clears the chain and blobs from memory.""" # store the iteration that the clear is occuring on self._lastclear = self.niterations self._itercounter = 0 @@ -316,20 +349,21 @@ def clear_samples(self): self._sampler.reset() def set_state_from_file(self, filename): - """Sets the state of the sampler back to the instance saved in a file. - """ - with self.io(filename, 'r') as fp: + """Sets the state of the sampler back to the instance saved in a file.""" + with self.io(filename, "r") as fp: rstate = fp.read_random_state() # set the numpy random state numpy.random.set_state(rstate) def run_mcmc(self, niterations): - """Advance the ensemble for a number of samples. + """ + Advance the ensemble for a number of samples. Parameters ---------- niterations : int Number of samples to get from sampler. + """ pos = self._pos if pos is None: @@ -340,20 +374,24 @@ def run_mcmc(self, niterations): self._pos = p def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters - ----------- + ---------- filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples - fp.write_samples(self.samples, - parameters=self.model.variable_params, - last_iteration=self.niterations) + fp.write_samples( + self.samples, + parameters=self.model.variable_params, + last_iteration=self.niterations, + ) # write stats fp.write_samples(self.model_stats, last_iteration=self.niterations) # write accpetance @@ -362,9 +400,11 @@ def write_results(self, filename): fp.write_random_state() @classmethod - def calculate_logevidence(cls, filename, thin_start=None, thin_end=None, - thin_interval=None): - """Calculates the log evidence from the given file using ``emcee_pt``'s + def calculate_logevidence( + cls, filename, thin_start=None, thin_end=None, thin_interval=None + ): + """ + Calculates the log evidence from the given file using ``emcee_pt``'s thermodynamic integration. Parameters @@ -390,14 +430,18 @@ def calculate_logevidence(cls, filename, thin_start=None, thin_end=None, The estimate of log of the evidence. dlnZ : float The error on the estimate. + """ - with cls._io(filename, 'r') as fp: - logls = fp.read_raw_samples(['loglikelihood'], - thin_start=thin_start, - thin_interval=thin_interval, - thin_end=thin_end, - temps='all', flatten=False) - logls = logls['loglikelihood'] + with cls._io(filename, "r") as fp: + logls = fp.read_raw_samples( + ["loglikelihood"], + thin_start=thin_start, + thin_interval=thin_interval, + thin_end=thin_end, + temps="all", + flatten=False, + ) + logls = logls["loglikelihood"] # we need the betas that were used betas = fp.betas # annoyingly, theromdynaimc integration in PTSampler is an instance @@ -405,23 +449,24 @@ def calculate_logevidence(cls, filename, thin_start=None, thin_end=None, ntemps = fp.ntemps nwalkers = fp.nwalkers ndim = len(fp.variable_params) - dummy_sampler = emcee.PTSampler(ntemps, nwalkers, ndim, None, - None, betas=betas) + dummy_sampler = emcee.PTSampler(ntemps, nwalkers, ndim, None, None, betas=betas) return dummy_sampler.thermodynamic_integration_log_evidence( - logls=logls, fburnin=0.) + logls=logls, fburnin=0.0 + ) def _correctjacobian(self, samples): - """Corrects the log jacobian values stored on disk. + """ + Corrects the log jacobian values stored on disk. Parameters ---------- samples : dict Dictionary of the samples. + """ # flatten samples for evaluating orig_shape = list(samples.values())[0].shape - flattened_samples = {p: arr.ravel() - for p, arr in list(samples.items())} + flattened_samples = {p: arr.ravel() for p, arr in list(samples.items())} # convert to a list of tuples so we can use map function params = list(flattened_samples.keys()) size = flattened_samples[params[0]].size @@ -434,7 +479,8 @@ def _correctjacobian(self, samples): return logj.reshape(orig_shape) def finalize(self): - """Calculates the log evidence and writes to the checkpoint file. + """ + Calculates the log evidence and writes to the checkpoint file. If sampling transforms were used, this also corrects the jacobian stored on disk. @@ -445,27 +491,34 @@ def finalize(self): if self.model.sampling_transforms is not None: # fix the lobjacobian values stored on disk logging.info("Correcting logjacobian values on disk") - with self.io(self.checkpoint_file, 'r') as fp: - samples = fp.read_raw_samples(self.variable_params, - thin_start=0, - thin_interval=1, thin_end=None, - temps='all', flatten=False) + with self.io(self.checkpoint_file, "r") as fp: + samples = fp.read_raw_samples( + self.variable_params, + thin_start=0, + thin_interval=1, + thin_end=None, + temps="all", + flatten=False, + ) logjacobian = self._correctjacobian(samples) # write them back out for fn in [self.checkpoint_file, self.backup_file]: with self.io(fn, "a") as fp: - fp[fp.samples_group]['logjacobian'][()] = logjacobian + fp[fp.samples_group]["logjacobian"][()] = logjacobian logging.info("Calculating log evidence") # get the thinning settings - with self.io(self.checkpoint_file, 'r') as fp: + with self.io(self.checkpoint_file, "r") as fp: thin_start = fp.thin_start thin_interval = fp.thin_interval thin_end = fp.thin_end # calculate logz, dlogz = self.calculate_logevidence( - self.checkpoint_file, thin_start=thin_start, thin_end=thin_end, - thin_interval=thin_interval) - logging.info("log Z, dlog Z: {}, {}".format(logz, dlogz)) + self.checkpoint_file, + thin_start=thin_start, + thin_end=thin_end, + thin_interval=thin_interval, + ) + logging.info(f"log Z, dlog Z: {logz}, {dlogz}") # write to both the checkpoint and backup for fn in [self.checkpoint_file, self.backup_file]: with self.io(fn, "a") as fp: diff --git a/pycbc/inference/sampler/epsie.py b/pycbc/inference/sampler/epsie.py index b9125627167..3c0b0ac38c1 100644 --- a/pycbc/inference/sampler/epsie.py +++ b/pycbc/inference/sampler/epsie.py @@ -13,33 +13,34 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module provides classes for interacting with epsie samplers. -""" - - -import numpy +"""This module provides classes for interacting with epsie samplers.""" import epsie -from epsie.samplers import ParallelTemperedSampler +import numpy # we'll use emcee_pt's default beta ladder for temperature levels from emcee.ptsampler import default_beta_ladder +from epsie.samplers import ParallelTemperedSampler from pycbc.pool import choose_pool -from .base import (BaseSampler, setup_output) -from .base_mcmc import (BaseMCMC, get_optional_arg_from_config, - nsamples_in_chain) -from .base_multitemper import (MultiTemperedSupport, compute_acf, compute_acl, - acl_from_raw_acls) +from .. import models from ..burn_in import MultiTemperedMCMCBurnInTests -from ..jump import epsie_proposals_from_config from ..io import EpsieFile -from .. import models +from ..jump import epsie_proposals_from_config +from .base import BaseSampler, setup_output +from .base_mcmc import BaseMCMC, get_optional_arg_from_config, nsamples_in_chain +from .base_multitemper import ( + MultiTemperedSupport, + acl_from_raw_acls, + compute_acf, + compute_acl, +) class EpsieSampler(MultiTemperedSupport, BaseMCMC, BaseSampler): - """Constructs an MCMC sampler using epsie's parallel-tempered sampler. + """ + Constructs an MCMC sampler using epsie's parallel-tempered sampler. Parameters ---------- @@ -84,23 +85,34 @@ class EpsieSampler(MultiTemperedSupport, BaseMCMC, BaseSampler): use_mpi : bool, optional Use MPI for parallelization. Default (False) will use python's multiprocessing. + """ + name = "epsie" _io = EpsieFile burn_in_class = MultiTemperedMCMCBurnInTests - def __init__(self, model, nchains, ntemps=None, betas=None, - proposals=None, default_proposal=None, - default_proposal_args=None, seed=None, - swap_interval=1, - checkpoint_interval=None, checkpoint_signal=None, - loglikelihood_function=None, - nprocesses=1, use_mpi=False): + def __init__( + self, + model, + nchains, + ntemps=None, + betas=None, + proposals=None, + default_proposal=None, + default_proposal_args=None, + seed=None, + swap_interval=1, + checkpoint_interval=None, + checkpoint_signal=None, + loglikelihood_function=None, + nprocesses=1, + use_mpi=False, + ): # create the betas if not provided if betas is None: - betas = default_beta_ladder(len(model.variable_params), - ntemps=ntemps) + betas = default_beta_ladder(len(model.variable_params), ntemps=ntemps) self.model = model # create a wrapper for calling the model model_call = _EpsieCallModel(model, loglikelihood_function) @@ -114,11 +126,17 @@ def __init__(self, model, nchains, ntemps=None, betas=None, # initialize the sampler self._sampler = ParallelTemperedSampler( - model.sampling_params, model_call, nchains, betas=betas, + model.sampling_params, + model_call, + nchains, + betas=betas, swap_interval=swap_interval, - proposals=proposals, default_proposal=default_proposal, + proposals=proposals, + default_proposal=default_proposal, default_proposal_args=default_proposal_args, - seed=seed, pool=pool) + seed=seed, + pool=pool, + ) # set other parameters self.nchains = nchains self._ntemps = ntemps @@ -131,7 +149,10 @@ def io(self): @property def base_shape(self): - return (self.ntemps, self.nchains,) + return ( + self.ntemps, + self.nchains, + ) @property def betas(self): @@ -140,7 +161,8 @@ def betas(self): @property def seed(self): - """The seed used for epsie's random bit generator. + """ + The seed used for epsie's random bit generator. This is not the same as the seed used for the prior distributions. """ @@ -153,7 +175,8 @@ def swap_interval(self): @staticmethod def compute_acf(filename, **kwargs): - r"""Computes the autocorrelation function. + r""" + Computes the autocorrelation function. Calls :py:func:`base_multitemper.compute_acf`; see that function for details. @@ -171,18 +194,20 @@ def compute_acf(filename, **kwargs): dict : Dictionary of arrays giving the ACFs for each parameter. The arrays will have shape ``ntemps x nchains x niterations``. + """ return compute_acf(filename, **kwargs) @staticmethod def compute_acl(filename, **kwargs): - r"""Computes the autocorrelation length. + r""" + Computes the autocorrelation length. Calls :py:func:`base_multitemper.compute_acl`; see that function for details. Parameters - ----------- + ---------- filename : str Name of a samples file to compute ACLs for. \**kwargs : @@ -193,18 +218,19 @@ def compute_acl(filename, **kwargs): ------- dict A dictionary of ntemps-long arrays of the ACLs of each parameter. + """ return compute_acl(filename, **kwargs) @property def acl(self): # pylint: disable=invalid-overridden-method - """The autocorrelation lengths of the chains. - """ + """The autocorrelation lengths of the chains.""" return acl_from_raw_acls(self.raw_acls) @property def effective_nsamples(self): # pylint: disable=invalid-overridden-method - """The effective number of samples post burn-in that the sampler has + """ + The effective number of samples post burn-in that the sampler has acquired so far. """ act = self.act @@ -224,7 +250,8 @@ def effective_nsamples(self): # pylint: disable=invalid-overridden-method @property def samples(self): - """A dict mapping ``variable_params`` to arrays of samples currently + """ + A dict mapping ``variable_params`` to arrays of samples currently in memory. The arrays have shape ``ntemps x nchains x niterations``. @@ -233,25 +260,23 @@ def samples(self): """ samples = epsie.array2dict(self._sampler.positions) # apply boundary conditions - samples = self.model.prior_distribution.apply_boundary_conditions( - **samples) + samples = self.model.prior_distribution.apply_boundary_conditions(**samples) # apply transforms to go to model's variable params space if self.model.sampling_transforms is not None: - samples = self.model.sampling_transforms.apply( - samples, inverse=True) + samples = self.model.sampling_transforms.apply(samples, inverse=True) return samples @property def model_stats(self): - """A dict mapping the model's ``default_stats`` to arrays of values. + """ + A dict mapping the model's ``default_stats`` to arrays of values. The arrays have shape ``ntemps x nchains x niterations``. """ return epsie.array2dict(self._sampler.blobs) def clear_samples(self): - """Clears the chain and blobs from memory. - """ + """Clears the chain and blobs from memory.""" # store the iteration that the clear is occuring on self._lastclear = self.niterations self._itercounter = 0 @@ -259,12 +284,10 @@ def clear_samples(self): self._sampler.clear() def set_state_from_file(self, filename): - """Sets the state of the sampler back to the instance saved in a file. - """ - with self.io(filename, 'r') as fp: + """Sets the state of the sampler back to the instance saved in a file.""" + with self.io(filename, "r") as fp: # get the numpy state - numpy_rstate_group = '/'.join([fp.sampler_group, - 'numpy_random_state']) + numpy_rstate_group = "/".join([fp.sampler_group, "numpy_random_state"]) rstate = fp.read_random_state(group=numpy_rstate_group) # set the sampler state for epsie self._sampler.set_state_from_checkpoint(fp, path=fp.sampler_group) @@ -272,8 +295,7 @@ def set_state_from_file(self, filename): numpy.random.set_state(rstate) def set_p0(self, samples_file=None, prior=None): - p0 = super(EpsieSampler, self).set_p0(samples_file=samples_file, - prior=prior) + p0 = super().set_p0(samples_file=samples_file, prior=prior) self._sampler.start_position = p0 @property @@ -284,46 +306,55 @@ def pos(self): return self._sampler.current_positions def run_mcmc(self, niterations): - """Advance the chains for a number of iterations. + """ + Advance the chains for a number of iterations. Parameters ---------- niterations : int Number of samples to get from sampler. + """ self._sampler.run(niterations) def write_results(self, filename): - """Writes samples, model stats, acceptance ratios, and random state + """ + Writes samples, model stats, acceptance ratios, and random state to the given file. Parameters - ----------- + ---------- filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples - fp.write_samples(self.samples, - parameters=self.model.variable_params, - last_iteration=self.niterations) + fp.write_samples( + self.samples, + parameters=self.model.variable_params, + last_iteration=self.niterations, + ) # write stats fp.write_samples(self.model_stats, last_iteration=self.niterations) # write accpetance ratio acceptance = self._sampler.acceptance - fp.write_acceptance_ratio(acceptance['acceptance_ratio'], - last_iteration=self.niterations) + fp.write_acceptance_ratio( + acceptance["acceptance_ratio"], last_iteration=self.niterations + ) # write temperature data if self.ntemps > 1: temp_ar = self._sampler.temperature_acceptance temp_swaps = self._sampler.temperature_swaps - fp.write_temperature_data(temp_swaps, temp_ar, - self.swap_interval, - last_iteration=self.niterations) + fp.write_temperature_data( + temp_swaps, + temp_ar, + self.swap_interval, + last_iteration=self.niterations, + ) # write numpy's global state (for the distributions) - numpy_rstate_group = '/'.join([fp.sampler_group, - 'numpy_random_state']) + numpy_rstate_group = "/".join([fp.sampler_group, "numpy_random_state"]) fp.write_random_state(group=numpy_rstate_group) # write the sampler's state self._sampler.checkpoint(fp, path=fp.sampler_group) @@ -332,9 +363,9 @@ def finalize(self): pass @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): - """Loads the sampler from the given config file. + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): + """ + Loads the sampler from the given config file. The following options are retrieved in the ``[sampler]`` section: @@ -417,44 +448,55 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, ------- EpsiePTSampler : The sampler instance. + """ section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) nchains = int(cp.get(section, "nchains")) - seed = get_optional_arg_from_config(cp, section, 'seed', dtype=int) + seed = get_optional_arg_from_config(cp, section, "seed", dtype=int) ntemps, betas = cls.betas_from_config(cp, section) # get the swap interval - swap_interval = get_optional_arg_from_config(cp, section, - 'swap-interval', - dtype=int) + swap_interval = get_optional_arg_from_config( + cp, section, "swap-interval", dtype=int + ) if swap_interval is None: swap_interval = 1 # get the checkpoint interval, if it's specified checkpoint_interval = cls.checkpoint_from_config(cp, section) checkpoint_signal = cls.ckpt_signal_from_config(cp, section) # get the loglikelihood function - logl = get_optional_arg_from_config(cp, section, 'logl-function') + logl = get_optional_arg_from_config(cp, section, "logl-function") # get the proposals proposals = epsie_proposals_from_config(cp) # check that all of the sampling parameters have a specified # proposal sampling_params = set(model.sampling_params) - proposal_params = set(param for prop in proposals - for param in prop.parameters) + proposal_params = set(param for prop in proposals for param in prop.parameters) missing = sampling_params - proposal_params if missing: - raise ValueError("Missing jump proposals for sampling parameters " - "{}".format(', '.join(missing))) + raise ValueError( + "Missing jump proposals for sampling parameters {}".format( + ", ".join(missing) + ) + ) # initialize - obj = cls(model, nchains, - ntemps=ntemps, betas=betas, proposals=proposals, - swap_interval=swap_interval, seed=seed, - checkpoint_interval=checkpoint_interval, - checkpoint_signal=checkpoint_signal, - loglikelihood_function=logl, - nprocesses=nprocesses, use_mpi=use_mpi) + obj = cls( + model, + nchains, + ntemps=ntemps, + betas=betas, + proposals=proposals, + swap_interval=swap_interval, + seed=seed, + checkpoint_interval=checkpoint_interval, + checkpoint_signal=checkpoint_signal, + loglikelihood_function=logl, + nprocesses=nprocesses, + use_mpi=use_mpi, + ) # set target obj.set_target_from_config(cp, section) # add burn-in if it's specified @@ -470,8 +512,9 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, return obj -class _EpsieCallModel(object): - """Model wrapper for epsie. +class _EpsieCallModel: + """ + Model wrapper for epsie. Allows model to be called like a function. Returns the loglikelihood function, logprior, and the model's default stats. @@ -480,7 +523,7 @@ class _EpsieCallModel(object): def __init__(self, model, loglikelihood_function=None): self.model = model if loglikelihood_function is None: - loglikelihood_function = 'loglikelihood' + loglikelihood_function = "loglikelihood" self.loglikelihood_function = loglikelihood_function def __call__(self, **kwargs): diff --git a/pycbc/inference/sampler/games.py b/pycbc/inference/sampler/games.py index 3a5931fbd8f..93bcae2d195 100644 --- a/pycbc/inference/sampler/games.py +++ b/pycbc/inference/sampler/games.py @@ -1,22 +1,25 @@ -""" Direct monte carlo sampling using pregenerated mapping files that +""" +Direct monte carlo sampling using pregenerated mapping files that encode the intrinsic parameter space. """ + import logging -import tqdm + import h5py import numpy import numpy.random +import tqdm from scipy.special import logsumexp -from pycbc.io import FieldArray from pycbc.inference import models +from pycbc.io import FieldArray from pycbc.pool import choose_pool + from .dummy import DummySampler def call_likelihood(params): - """ Accessor to update the global model - """ + """Accessor to update the global model""" models._global_instance.update(**params) return models._global_instance.loglikelihood @@ -26,7 +29,8 @@ class OutOfSamples(Exception): class GameSampler(DummySampler): - """Direct importance sampling using a preconstructed parameter space + """ + Direct importance sampling using a preconstructed parameter space mapping file. Parameters @@ -42,15 +46,23 @@ class GameSampler(DummySampler): Try to use this many likelihood calls in each round of the analysis. rounds: int The number of iterations to use before terminated. + """ - name = 'games' - - def __init__(self, model, *args, nprocesses=1, use_mpi=False, - mapfile=None, - loglr_region=25, - target_likelihood_calls=1e5, - rounds=1, - **kwargs): + + name = "games" + + def __init__( + self, + model, + *args, + nprocesses=1, + use_mpi=False, + mapfile=None, + loglr_region=25, + target_likelihood_calls=1e5, + rounds=1, + **kwargs, + ): super().__init__(model, *args) self.meta = {} @@ -68,7 +80,7 @@ def __init__(self, model, *args, nprocesses=1, use_mpi=False, self.loglr_region = float(loglr_region) def draw_samples_from_bin(self, i, size): - """ Get samples from the binned prior space """ + """Get samples from the binned prior space""" if i not in self.draw: self.draw[i] = numpy.arange(0, len(self.dmap[i])) @@ -81,12 +93,13 @@ def draw_samples_from_bin(self, i, size): if size > 0: remain = len(self.draw[i]) - logging.info('Drew %i, %i remains in bin %i', size, remain, i) + logging.info("Drew %i, %i remains in bin %i", size, remain, i) return self.dmap[i][selected] def sample_round(self, bin_weight, node_idx, lengths): - """ Sample from the posterior using pre-binned sets of points and + """ + Sample from the posterior using pre-binned sets of points and the weighting factor of each bin. bin_weight: Array @@ -116,17 +129,16 @@ def sample_round(self, bin_weight, node_idx, lengths): drawweight = bin_weight / drawcount total_draw = drawcount.sum() - logging.info('...drawn random points within bins') + logging.info("...drawn random points within bins") psamp = FieldArray(total_draw, dtype=self.dtype) pweight = numpy.zeros(total_draw, dtype=float) bin_id = numpy.zeros(total_draw, dtype=int) j = 0 for i, (c, w) in enumerate(zip(drawcount, drawweight)): bdraw = self.draw_samples_from_bin(node_idx[i], c) - psamp[j:j+len(bdraw)] = FieldArray.from_records(bdraw, - dtype=self.dtype) - pweight[j:j+len(bdraw)] = numpy.log(bin_weight[i]) - numpy.log(w) - bin_id[j:j+len(bdraw)] = i + psamp[j : j + len(bdraw)] = FieldArray.from_records(bdraw, dtype=self.dtype) + pweight[j : j + len(bdraw)] = numpy.log(bin_weight[i]) - numpy.log(w) + bin_id[j : j + len(bdraw)] = i j += len(bdraw) logging.info("Possible unique values %s", lengths.sum()) @@ -140,8 +152,9 @@ def sample_round(self, bin_weight, node_idx, lengths): pset = {p: psamp[p][i] for p in self.model.variable_params} args.append(pset) - loglr_samp = list(tqdm.tqdm(self.pool.imap(call_likelihood, args), - total=len(args))) + loglr_samp = list( + tqdm.tqdm(self.pool.imap(call_likelihood, args), total=len(args)) + ) loglr_samp = numpy.array(loglr_samp) # Calculate the weights from the actual likelihood relative to the @@ -152,28 +165,30 @@ def sample_round(self, bin_weight, node_idx, lengths): return psamp, loglr_samp, weight2, bin_id def run(self): - """ Produce posterior samples """ - logging.info('Retrieving params of parameter space nodes') - with h5py.File(self.mapfile, 'r') as mapfile: - bparams = {p: mapfile['bank'][p][:] for p in self.variable_params} + """Produce posterior samples""" + logging.info("Retrieving params of parameter space nodes") + with h5py.File(self.mapfile, "r") as mapfile: + bparams = {p: mapfile["bank"][p][:] for p in self.variable_params} num_nodes = len(bparams[list(bparams.keys())[0]]) - lengths = numpy.array([len(mapfile['map'][str(x)]) - for x in range(num_nodes)]) - self.dtype = mapfile['map']['0'].dtype + lengths = numpy.array( + [len(mapfile["map"][str(x)]) for x in range(num_nodes)] + ) + self.dtype = mapfile["map"]["0"].dtype - logging.info('Calculating likelihood at nodes') + logging.info("Calculating likelihood at nodes") args = [] for i in range(num_nodes): pset = {p: bparams[p][i] for p in self.model.variable_params} args.append(pset) - node_loglrs = list(tqdm.tqdm(self.pool.imap(call_likelihood, args), - total=len(args))) + node_loglrs = list( + tqdm.tqdm(self.pool.imap(call_likelihood, args), total=len(args)) + ) node_loglrs = numpy.array(node_loglrs) loglr_bound = node_loglrs[~numpy.isnan(node_loglrs)].max() loglr_bound -= self.loglr_region - logging.info('Drawing proposal samples from node regions') + logging.info("Drawing proposal samples from node regions") logw = node_loglrs + numpy.log(lengths) passed = (node_loglrs > loglr_bound) & ~numpy.isnan(node_loglrs) passed = numpy.where(passed)[0] @@ -182,9 +197,9 @@ def run(self): weight = numpy.exp(logw2) logging.info("...reading template bins") - with h5py.File(self.mapfile, 'r') as mapfile: + with h5py.File(self.mapfile, "r") as mapfile: for i in passed: - self.dmap[i] = mapfile['map'][str(i)][:] + self.dmap[i] = mapfile["map"][str(i)][:] # Sample from posterior psamp = None @@ -196,9 +211,9 @@ def run(self): for i in range(self.rounds): try: - psamp_v, loglr_samp_v, weight2_v, bin_id = \ - self.sample_round(weight / weight.sum(), - passed, lengths[passed]) + psamp_v, loglr_samp_v, weight2_v, bin_id = self.sample_round( + weight / weight.sum(), passed, lengths[passed] + ) except OutOfSamples: logging.info("No more samples to draw from") break @@ -217,22 +232,22 @@ def run(self): weight2 = numpy.concatenate([weight2_v, weight2]) bin_ids = numpy.concatenate([bin_id, bin_ids]) - ess = 1.0 / ((weight2/weight2.sum()) ** 2.0).sum() + ess = 1.0 / ((weight2 / weight2.sum()) ** 2.0).sum() logging.info("ESS = %s", ess) # Prepare the equally weighted output samples - self.meta['ncalls'] = len(weight2) - self.meta['ess'] = ess + self.meta["ncalls"] = len(weight2) + self.meta["ess"] = ess weight2 /= weight2.sum() - draw2 = numpy.random.choice(len(psamp), size=int(ess * 1), - replace=True, p=weight2) - logging.info("Unique values second draw %s", - len(numpy.unique(psamp[draw2]))) + draw2 = numpy.random.choice( + len(psamp), size=int(ess * 1), replace=True, p=weight2 + ) + logging.info("Unique values second draw %s", len(numpy.unique(psamp[draw2]))) fsamp = FieldArray(len(draw2), dtype=self.dtype) for i, v in enumerate(draw2): fsamp[i] = psamp[v] self._samples = {p: fsamp[p] for p in self.model.variable_params} - self._samples['loglikelihood'] = loglr_samp[draw2] + self._samples["loglikelihood"] = loglr_samp[draw2] diff --git a/pycbc/inference/sampler/multinest.py b/pycbc/inference/sampler/multinest.py index db88ad57b47..a83d0f720ca 100644 --- a/pycbc/inference/sampler/multinest.py +++ b/pycbc/inference/sampler/multinest.py @@ -26,18 +26,18 @@ packages for parameter estimation. """ - import logging import sys + import numpy -from pycbc.inference.io import (MultinestFile, validate_checkpoint_files) from pycbc.distributions import read_constraints_from_config +from pycbc.inference.io import MultinestFile, validate_checkpoint_files from pycbc.pool import is_main_process from pycbc.transforms import apply_transforms -from .base import (BaseSampler, setup_output) -from .base_mcmc import get_optional_arg_from_config +from .base import BaseSampler, setup_output +from .base_mcmc import get_optional_arg_from_config # # ============================================================================= @@ -47,8 +47,10 @@ # ============================================================================= # + class MultinestSampler(BaseSampler): - """This class is used to construct a nested sampler from + """ + This class is used to construct a nested sampler from the Multinest package. Parameters @@ -57,25 +59,34 @@ class MultinestSampler(BaseSampler): A model from ``pycbc.inference.models``. nlivepoints : int Number of live points to use in sampler. + """ + name = "multinest" _io = MultinestFile - def __init__(self, model, nlivepoints, checkpoint_interval=1000, - importance_nested_sampling=False, - evidence_tolerance=0.1, sampling_efficiency=0.01, - constraints=None): + def __init__( + self, + model, + nlivepoints, + checkpoint_interval=1000, + importance_nested_sampling=False, + evidence_tolerance=0.1, + sampling_efficiency=0.01, + constraints=None, + ): try: loglevel = logging.getLogger().getEffectiveLevel() logging.getLogger().setLevel(logging.WARNING) from pymultinest import Analyzer, run + self.run_multinest = run self.analyzer = Analyzer logging.getLogger().setLevel(loglevel) except ImportError: raise ImportError("pymultinest is not installed.") - super(MultinestSampler, self).__init__(model) + super().__init__(model) self._constraints = constraints self._nlivepoints = nlivepoints @@ -99,8 +110,7 @@ def io(self): @property def niterations(self): - """Get the current number of iterations. - """ + """Get the current number of iterations.""" itercount = self._itercount if itercount is None: itercount = 0 @@ -108,55 +118,54 @@ def niterations(self): @property def checkpoint_interval(self): - """Get the number of iterations between checkpoints. - """ + """Get the number of iterations between checkpoints.""" return self._checkpoint_interval @property def nlivepoints(self): - """Get the number of live points used in sampling. - """ + """Get the number of live points used in sampling.""" return self._nlivepoints @property def logz(self): - """Get the current estimate of the log evidence. - """ + """Get the current estimate of the log evidence.""" return self._logz @property def dlogz(self): - """Get the current error estimate of the log evidence. - """ + """Get the current error estimate of the log evidence.""" return self._dlogz @property def importance_logz(self): - """Get the current importance weighted estimate of the log + """ + Get the current importance weighted estimate of the log evidence. """ return self._importance_logz @property def importance_dlogz(self): - """Get the current error estimate of the importance + """ + Get the current error estimate of the importance weighted log evidence. """ return self._importance_dlogz @property def samples(self): - """A dict mapping ``variable_params`` to arrays of samples currently + """ + A dict mapping ``variable_params`` to arrays of samples currently in memory. """ - samples_dict = {p: self._samples[:, i] for i, p in - enumerate(self.model.variable_params)} + samples_dict = { + p: self._samples[:, i] for i, p in enumerate(self.model.variable_params) + } return samples_dict @property def model_stats(self): - """A dict mapping the model's ``default_stats`` to arrays of values. - """ + """A dict mapping the model's ``default_stats`` to arrays of values.""" stats = [] for sample in self._samples: params = dict(zip(self.model.variable_params, sample)) @@ -169,29 +178,29 @@ def model_stats(self): return {s: stats[:, i] for i, s in enumerate(self.model.default_stats)} def get_posterior_samples(self): - """Read posterior samples from ASCII output file created by + """ + Read posterior samples from ASCII output file created by multinest. """ - post_file = self.backup_file[:-9]+'-post_equal_weights.dat' + post_file = self.backup_file[:-9] + "-post_equal_weights.dat" return numpy.loadtxt(post_file, ndmin=2) def check_if_finished(self): - """Estimate remaining evidence to see if desired evidence-tolerance + """ + Estimate remaining evidence to see if desired evidence-tolerance stopping criterion has been reached. """ - resume_file = self.backup_file[:-9] + '-resume.dat' - current_vol, _, _ = numpy.loadtxt( - resume_file, skiprows=6, unpack=True) + resume_file = self.backup_file[:-9] + "-resume.dat" + current_vol, _, _ = numpy.loadtxt(resume_file, skiprows=6, unpack=True) maxloglike = max(self.get_posterior_samples()[:, -1]) - logz_remain = numpy.exp(maxloglike + - numpy.log(current_vol) - self.logz) + logz_remain = numpy.exp(maxloglike + numpy.log(current_vol) - self.logz) logging.info("Estimate of remaining logZ is %s", logz_remain) done = logz_remain < self._ztol return done - def set_initial_conditions(self, initial_distribution=None, - samples_file=None): - """Sets the initial starting point for the sampler. + def set_initial_conditions(self, initial_distribution=None, samples_file=None): + """ + Sets the initial starting point for the sampler. If a starting samples file is provided, will also load the random state from it. @@ -201,14 +210,11 @@ def set_initial_conditions(self, initial_distribution=None, self.set_state_from_file(samples_file) def resume_from_checkpoint(self): - """Resume sampler from checkpoint - """ - pass + """Resume sampler from checkpoint""" def set_state_from_file(self, filename): - """Sets the state of the sampler back to the instance saved in a file. - """ - with self.io(filename, 'r') as f_p: + """Sets the state of the sampler back to the instance saved in a file.""" + with self.io(filename, "r") as f_p: rstate = f_p.read_random_state() # set the numpy random state numpy.random.set_state(rstate) @@ -216,8 +222,7 @@ def set_state_from_file(self, filename): self._random_state = rstate def loglikelihood(self, cube, *extra_args): - """Log likelihood evaluator that gets passed to multinest. - """ + """Log likelihood evaluator that gets passed to multinest.""" params = {p: v for p, v in zip(self.model.variable_params, cube)} # apply transforms if self.model.sampling_transforms is not None: @@ -225,14 +230,16 @@ def loglikelihood(self, cube, *extra_args): if self.model.waveform_transforms is not None: params = apply_transforms(params, self.model.waveform_transforms) # apply constraints - if (self._constraints is not None and - not all([c(params) for c in self._constraints])): + if self._constraints is not None and not all( + [c(params) for c in self._constraints] + ): return -numpy.inf self.model.update(**params) return self.model.loglikelihood def transform_prior(self, cube, *extra_args): - """Transforms the unit hypercube that multinest makes its draws + """ + Transforms the unit hypercube that multinest makes its draws from, into the prior space defined in the config file. """ dict_cube = dict(zip(self.model.variable_params, cube)) @@ -242,7 +249,8 @@ def transform_prior(self, cube, *extra_args): return cube def run(self): - """Runs the sampler until the specified evidence tolerance + """ + Runs the sampler until the specified evidence tolerance is reached. """ if self.new_checkpoint: @@ -251,35 +259,43 @@ def run(self): self.set_initial_conditions(samples_file=self.checkpoint_file) with self.io(self.checkpoint_file, "r") as f_p: self._itercount = f_p.niterations - outputfiles_basename = self.backup_file[:-9] + '-' - analyzer = self.analyzer(self._ndim, - outputfiles_basename=outputfiles_basename) + outputfiles_basename = self.backup_file[:-9] + "-" + analyzer = self.analyzer(self._ndim, outputfiles_basename=outputfiles_basename) iterinterval = self.checkpoint_interval done = False while not done: - logging.info("Running sampler for %s to %s iterations", - self.niterations, self.niterations + iterinterval) + logging.info( + "Running sampler for %s to %s iterations", + self.niterations, + self.niterations + iterinterval, + ) # run multinest - self.run_multinest(self.loglikelihood, self.transform_prior, - self._ndim, n_live_points=self.nlivepoints, - evidence_tolerance=self._ztol, - sampling_efficiency=self._eff, - importance_nested_sampling=self._ins, - max_iter=iterinterval, - n_iter_before_update=iterinterval, - seed=numpy.random.randint(0, 1e6), - outputfiles_basename=outputfiles_basename, - multimodal=False, verbose=True) + self.run_multinest( + self.loglikelihood, + self.transform_prior, + self._ndim, + n_live_points=self.nlivepoints, + evidence_tolerance=self._ztol, + sampling_efficiency=self._eff, + importance_nested_sampling=self._ins, + max_iter=iterinterval, + n_iter_before_update=iterinterval, + seed=numpy.random.randint(0, 1e6), + outputfiles_basename=outputfiles_basename, + multimodal=False, + verbose=True, + ) # parse results from multinest output files nest_stats = analyzer.get_mode_stats() self._logz = nest_stats["nested sampling global log-evidence"] - self._dlogz = nest_stats[ - "nested sampling global log-evidence error"] + self._dlogz = nest_stats["nested sampling global log-evidence error"] if self._ins: self._importance_logz = nest_stats[ - "nested importance sampling global log-evidence"] + "nested importance sampling global log-evidence" + ] self._importance_dlogz = nest_stats[ - "nested importance sampling global log-evidence error"] + "nested importance sampling global log-evidence error" + ] self._samples = self.get_posterior_samples()[:, :-1] logging.info("Have %s posterior samples", self._samples.shape[0]) # update the itercounter @@ -296,24 +312,26 @@ def run(self): sys.exit() def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters - ----------- + ---------- filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as f_p: + with self.io(filename, "a") as f_p: # write samples f_p.write_samples(self.samples, self.model.variable_params) # write stats f_p.write_samples(self.model_stats) # write evidence - f_p.write_logevidence(self.logz, self.dlogz, - self.importance_logz, - self.importance_dlogz) + f_p.write_logevidence( + self.logz, self.dlogz, self.importance_logz, self.importance_dlogz + ) # write random state (use default numpy.random_state) f_p.write_random_state() @@ -326,12 +344,14 @@ def checkpoint(self): f_p.write_niterations(self.niterations) logging.info("Validating checkpoint and backup files") checkpoint_valid = validate_checkpoint_files( - self.checkpoint_file, self.backup_file, check_nsamples=False) + self.checkpoint_file, self.backup_file, check_nsamples=False + ) if not checkpoint_valid: - raise IOError("error writing to checkpoint file") + raise OSError("error writing to checkpoint file") def setup_output(self, output_file): - """Sets up the sampler's checkpoint and output files. + """ + Sets up the sampler's checkpoint and output files. The checkpoint file has the same name as the output file, but with ``.checkpoint`` appended to the name. A backup file will also be @@ -343,56 +363,65 @@ def setup_output(self, output_file): Sampler output_file : str Name of the output file. + """ if self.is_main_process: setup_output(self, output_file) else: # child processes just store filenames - checkpoint_file = output_file + '.checkpoint' - backup_file = output_file + '.bkup' + checkpoint_file = output_file + ".checkpoint" + backup_file = output_file + ".bkup" self.checkpoint_file = checkpoint_file self.backup_file = backup_file self.checkpoint_valid = True self.new_checkpoint = True def finalize(self): - """All data is written by the last checkpoint in the run method, so - this just passes.""" - pass + """ + All data is written by the last checkpoint in the run method, so + this just passes. + """ @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): """Loads the sampler from the given config file.""" section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) # get the number of live points to use nlivepoints = int(cp.get(section, "nlivepoints")) # get the checkpoint interval, if it's specified checkpoint = get_optional_arg_from_config( - cp, section, 'checkpoint-interval', dtype=int) + cp, section, "checkpoint-interval", dtype=int + ) # get the evidence tolerance, if specified - ztol = get_optional_arg_from_config(cp, section, 'evidence-tolerance', - dtype=float) + ztol = get_optional_arg_from_config( + cp, section, "evidence-tolerance", dtype=float + ) # get the sampling efficiency, if specified - eff = get_optional_arg_from_config(cp, section, 'sampling-efficiency', - dtype=float) + eff = get_optional_arg_from_config( + cp, section, "sampling-efficiency", dtype=float + ) # get importance nested sampling setting, if specified - ins = get_optional_arg_from_config(cp, section, - 'importance-nested-sampling', - dtype=bool) + ins = get_optional_arg_from_config( + cp, section, "importance-nested-sampling", dtype=bool + ) # get constraints since we can't use the joint prior distribution constraints = read_constraints_from_config(cp) # build optional kwarg dict - kwarg_names = ['evidence_tolerance', 'sampling_efficiency', - 'importance_nested_sampling', - 'checkpoint_interval'] - optional_kwargs = {k: v for k, v in - zip(kwarg_names, [ztol, eff, ins, checkpoint]) if - v is not None} - obj = cls(model, nlivepoints, constraints=constraints, - **optional_kwargs) + kwarg_names = [ + "evidence_tolerance", + "sampling_efficiency", + "importance_nested_sampling", + "checkpoint_interval", + ] + optional_kwargs = { + k: v + for k, v in zip(kwarg_names, [ztol, eff, ins, checkpoint]) + if v is not None + } + obj = cls(model, nlivepoints, constraints=constraints, **optional_kwargs) obj.setup_output(output_file) return obj diff --git a/pycbc/inference/sampler/nessai.py b/pycbc/inference/sampler/nessai.py index a5aebac2ece..adf06140dab 100644 --- a/pycbc/inference/sampler/nessai.py +++ b/pycbc/inference/sampler/nessai.py @@ -4,22 +4,23 @@ Documentation for nessai: https://nessai.readthedocs.io/en/latest/ """ + import ast import logging import os import nessai.flowsampler -import nessai.model import nessai.livepoint +import nessai.model import nessai.utils.multiprocessing import nessai.utils.settings import numpy import numpy.lib.recfunctions as rfn +from ...pool import choose_pool +from ..io import NessaiFile, loadfile from .base import BaseSampler, setup_output from .base_mcmc import get_optional_arg_from_config -from ..io import NessaiFile, loadfile -from ...pool import choose_pool class NessaiSampler(BaseSampler): @@ -105,12 +106,8 @@ def run(self, **kwargs): if kwargs is not None: logging.info("Updating keyword arguments with %s", kwargs) - extra_kwds.update( - {k: v for k, v in kwargs.items() if k in default_kwds} - ) - run_kwds.update( - {k: v for k, v in kwargs.items() if k in default_run_kwds} - ) + extra_kwds.update({k: v for k, v in kwargs.items() if k in default_kwds}) + run_kwds.update({k: v for k, v in kwargs.items() if k in default_run_kwds}) if self._sampler is None: logging.info("Initialising nessai FlowSampler") @@ -130,7 +127,8 @@ def run(self, **kwargs): @staticmethod def get_default_kwds(importance_nested_sampler=False): - """Return lists of all allowed keyword arguments for nessai. + """ + Return lists of all allowed keyword arguments for nessai. Returns ------- @@ -138,6 +136,7 @@ def get_default_kwds(importance_nested_sampler=False): List of keyword arguments that can be passed to FlowSampler run_kwds: list List of keyword arguments that can be passed to FlowSampler.run + """ return nessai.utils.settings.get_all_kwargs( importance_nested_sampler=importance_nested_sampler, @@ -145,17 +144,15 @@ def get_default_kwds(importance_nested_sampler=False): ) @classmethod - def from_config( - cls, cp, model, output_file=None, nprocesses=1, use_mpi=False - ): + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): """ Loads the sampler from the given config file. """ section = "sampler" # check name - assert ( - cp.get(section, "name") == cls.name - ), "name in section [sampler] must match mine" + assert cp.get(section, "name") == cls.name, ( + "name in section [sampler] must match mine" + ) if cp.has_option(section, "importance_nested_sampler"): importance_nested_sampler = cp.get( @@ -171,9 +168,7 @@ def from_config( "Importance nested sampler is not currently supported" ) - default_kwds, default_run_kwds = cls.get_default_kwds( - importance_nested_sampler - ) + default_kwds, default_run_kwds = cls.get_default_kwds(importance_nested_sampler) # Keyword arguments the user cannot configure via the config remove_kwds = [ @@ -210,15 +205,11 @@ def from_config( # Specified kwds ignore_kwds = {"nlive", "name"} invalid_kwds = ( - cp[section].keys() - - set().union(kwds.keys(), run_kwds.keys()) - - ignore_kwds + cp[section].keys() - set().union(kwds.keys(), run_kwds.keys()) - ignore_kwds ) if invalid_kwds: - raise RuntimeError( - f"Config contains unknown options: {invalid_kwds}" - ) + raise RuntimeError(f"Config contains unknown options: {invalid_kwds}") logging.info("nessai keyword arguments: %s", kwds) logging.info("nessai run keyword arguments: %s", run_kwds) @@ -246,13 +237,15 @@ def set_initial_conditions( initial_distribution=None, samples_file=None, ): - """Sets up the starting point for the sampler. + """ + Sets up the starting point for the sampler. This is not used for nessai. """ def checkpoint_callback(self, state): - """Callback for checkpointing. + """ + Callback for checkpointing. This will be called periodically by nessai. """ @@ -270,9 +263,7 @@ def resume_from_checkpoint(self): try: with loadfile(self.checkpoint_file, "r") as fp: self.resume_data = fp.read_pickled_data_from_checkpoint_file() - logging.info( - "Found valid checkpoint file: %s", self.checkpoint_file - ) + logging.info("Found valid checkpoint file: %s", self.checkpoint_file) except Exception as e: logging.info("Failed to load checkpoint file with error: %s", e) @@ -284,7 +275,8 @@ def finalize(self): self.checkpoint() def write_results(self, filename): - """Write the results to a given file. + """ + Write the results to a given file. Writes the nested samples, log-evidence and log-evidence error. """ @@ -297,7 +289,8 @@ def write_results(self, filename): class NessaiModel(nessai.model.Model): - """Wrapper for PyCBC Inference model class for use with nessai. + """ + Wrapper for PyCBC Inference model class for use with nessai. Parameters ---------- @@ -305,6 +298,7 @@ class NessaiModel(nessai.model.Model): A model instance from PyCBC. loglikelihood_function : str Name of the log-likelihood method to call. + """ def __init__(self, model, loglikelihood_function=None): @@ -320,11 +314,7 @@ def __init__(self, model, loglikelihood_function=None): bounds = {} for dist in model.prior_distribution.distributions: bounds.update( - **{ - k: [v.min, v.max] - for k, v in dist.bounds.items() - if k in self.names - } + **{k: [v.min, v.max] for k, v in dist.bounds.items() if k in self.names} ) self.bounds = bounds # Prior and likelihood are not vectorised diff --git a/pycbc/inference/sampler/ptemcee.py b/pycbc/inference/sampler/ptemcee.py index 535a91c9ed9..56846a0b61b 100644 --- a/pycbc/inference/sampler/ptemcee.py +++ b/pycbc/inference/sampler/ptemcee.py @@ -19,25 +19,34 @@ packages for parameter estimation. """ - +import logging import shlex + import numpy import ptemcee -import logging -from pycbc.pool import choose_pool -from .base import (BaseSampler, setup_output) -from .base_mcmc import (BaseMCMC, EnsembleSupport, raw_samples_to_dict, - get_optional_arg_from_config) -from .base_multitemper import (read_betas_from_hdf, - ensemble_compute_acf, ensemble_compute_acl) -from ..burn_in import EnsembleMultiTemperedMCMCBurnInTests from pycbc.inference.io import PTEmceeFile +from pycbc.pool import choose_pool + from .. import models +from ..burn_in import EnsembleMultiTemperedMCMCBurnInTests +from .base import BaseSampler, setup_output +from .base_mcmc import ( + BaseMCMC, + EnsembleSupport, + get_optional_arg_from_config, + raw_samples_to_dict, +) +from .base_multitemper import ( + ensemble_compute_acf, + ensemble_compute_acl, + read_betas_from_hdf, +) class PTEmceeSampler(EnsembleSupport, BaseMCMC, BaseSampler): - """This class is used to construct the parallel-tempered ptemcee sampler. + """ + This class is used to construct the parallel-tempered ptemcee sampler. Parameters ---------- @@ -76,17 +85,30 @@ class PTEmceeSampler(EnsembleSupport, BaseMCMC, BaseSampler): use_mpi : bool, optional Use MPI for parallelization. Default (False) will use python's multiprocessing. + """ + name = "ptemcee" _io = PTEmceeFile burn_in_class = EnsembleMultiTemperedMCMCBurnInTests - def __init__(self, model, nwalkers, ntemps=None, Tmax=None, betas=None, - adaptive=False, adaptation_lag=None, adaptation_time=None, - scale_factor=None, - loglikelihood_function=None, - checkpoint_interval=None, checkpoint_signal=None, - nprocesses=1, use_mpi=False): + def __init__( + self, + model, + nwalkers, + ntemps=None, + Tmax=None, + betas=None, + adaptive=False, + adaptation_lag=None, + adaptation_time=None, + scale_factor=None, + loglikelihood_function=None, + checkpoint_interval=None, + checkpoint_signal=None, + nprocesses=1, + use_mpi=False, + ): self.model = model ndim = len(model.variable_params) @@ -98,30 +120,36 @@ def __init__(self, model, nwalkers, ntemps=None, Tmax=None, betas=None, # construct the keyword arguments to pass; if a kwarg is None, we # won't pass it, resulting in ptemcee's defaults being used kwargs = {} - kwargs['adaptive'] = adaptive - kwargs['betas'] = betas + kwargs["adaptive"] = adaptive + kwargs["betas"] = betas if adaptation_lag is not None: - kwargs['adaptation_lag'] = adaptation_lag + kwargs["adaptation_lag"] = adaptation_lag if adaptation_time is not None: - kwargs['adaptation_time'] = adaptation_time + kwargs["adaptation_time"] = adaptation_time if scale_factor is not None: - kwargs['scale_factor'] = scale_factor + kwargs["scale_factor"] = scale_factor # create a wrapper for calling the model if loglikelihood_function is None: - loglikelihood_function = 'loglikelihood' + loglikelihood_function = "loglikelihood" # frustratingly, ptemcee does not support blob data, so we have to # turn it off - model_call = models.CallModel(model, loglikelihood_function, - return_all_stats=False) + model_call = models.CallModel( + model, loglikelihood_function, return_all_stats=False + ) # these are used to help paralleize over multiple cores / MPI models._global_instance = model_call model_call = models._call_global_model prior_call = models._call_global_model_logprior self.pool = choose_pool(mpi=use_mpi, processes=nprocesses) # construct the sampler - self._sampler = ptemcee.Sampler(nwalkers=nwalkers, ndim=ndim, - logl=model_call, logp=prior_call, - mapper=self.pool.map, **kwargs) + self._sampler = ptemcee.Sampler( + nwalkers=nwalkers, + ndim=ndim, + logl=model_call, + logp=prior_call, + mapper=self.pool.map, + **kwargs, + ) self.nwalkers = nwalkers self._ntemps = ntemps self._checkpoint_interval = checkpoint_interval @@ -141,7 +169,10 @@ def ntemps(self): @property def base_shape(self): - return (self.ntemps, self.nwalkers,) + return ( + self.ntemps, + self.nwalkers, + ) @property def betas(self): @@ -178,7 +209,8 @@ def scale_factor(self): @property def ensemble(self): - """Returns the current ptemcee ensemble. + """ + Returns the current ptemcee ensemble. The ensemble stores the current location of and temperatures of the walkers. If the ensemble hasn't been setup yet, will set one up @@ -205,7 +237,8 @@ def _pos(self): @property def chain(self): - """The current chain of samples in memory. + """ + The current chain of samples in memory. The chain is returned as a :py:mod:`ptemcee.chain.Chain` instance. If no chain has been created yet (``_chain`` is None), then will create a new chain using the current ``ensemble``. @@ -216,8 +249,7 @@ def chain(self): return self._chain def clear_samples(self): - """Clears the chain and blobs from memory. - """ + """Clears the chain and blobs from memory.""" # store the iteration that the clear is occuring on self._lastclear = self.niterations self._itercounter = 0 @@ -228,7 +260,8 @@ def clear_samples(self): @property def samples(self): - """A dict mapping ``variable_params`` to arrays of samples currently + """ + A dict mapping ``variable_params`` to arrays of samples currently in memory. The arrays have shape ``ntemps x nwalkers x niterations``. """ @@ -239,7 +272,8 @@ def samples(self): @property def model_stats(self): - """Returns the log likelihood ratio and log prior as a dict of arrays. + """ + Returns the log likelihood ratio and log prior as a dict of arrays. The returned array has shape ntemps x nwalkers x niterations. @@ -262,13 +296,11 @@ def model_stats(self): logl = self._chain.logl.transpose((1, 2, 0)) logp = self._chain.logP.transpose((1, 2, 0)) logjacobian = numpy.zeros(logp.shape) - return {'loglikelihood': logl, 'logprior': logp, - 'logjacobian': logjacobian} + return {"loglikelihood": logl, "logprior": logp, "logjacobian": logjacobian} def set_state_from_file(self, filename): - """Sets the state of the sampler back to the instance saved in a file. - """ - with self.io(filename, 'r') as fp: + """Sets the state of the sampler back to the instance saved in a file.""" + with self.io(filename, "r") as fp: rstate = fp.read_random_state() # set the numpy random state numpy.random.set_state(rstate) @@ -280,19 +312,23 @@ def set_state_from_file(self, filename): ensemble.time = fp.niterations def run_mcmc(self, niterations): - """Advance the ensemble for a number of samples. + """ + Advance the ensemble for a number of samples. Parameters ---------- niterations : int Number of samples to get from sampler. + """ self.chain.run(niterations) @classmethod - def calculate_logevidence(cls, filename, thin_start=None, thin_end=None, - thin_interval=None): - """Calculates the log evidence from the given file. + def calculate_logevidence( + cls, filename, thin_start=None, thin_end=None, thin_interval=None + ): + """ + Calculates the log evidence from the given file. This uses ``ptemcee``'s thermodynamic integration. Parameters @@ -318,18 +354,22 @@ def calculate_logevidence(cls, filename, thin_start=None, thin_end=None, The estimate of log of the evidence. dlnZ : float The error on the estimate. + """ - with cls._io(filename, 'r') as fp: - logls = fp.read_raw_samples(['loglikelihood'], - thin_start=thin_start, - thin_interval=thin_interval, - thin_end=thin_end, - temps='all', flatten=False) - logls = logls['loglikelihood'] + with cls._io(filename, "r") as fp: + logls = fp.read_raw_samples( + ["loglikelihood"], + thin_start=thin_start, + thin_interval=thin_interval, + thin_end=thin_end, + temps="all", + flatten=False, + ) + logls = logls["loglikelihood"] # we need the betas that were used - betas = fp.read_betas(thin_start=thin_start, - thin_interval=thin_interval, - thin_end=thin_end) + betas = fp.read_betas( + thin_start=thin_start, thin_interval=thin_interval, thin_end=thin_end + ) # we'll separate betas out by their unique temperatures # there's probably a faster way to do this... mean_logls = [] @@ -346,11 +386,13 @@ def calculate_logevidence(cls, filename, thin_start=None, thin_end=None, mean_logls.append(loglsti[:, getiters].mean()) unique_betas.append(ubti[ii]) return ptemcee.util.thermodynamic_integration_log_evidence( - numpy.array(unique_betas), numpy.array(mean_logls)) + numpy.array(unique_betas), numpy.array(mean_logls) + ) @staticmethod def compute_acf(filename, **kwargs): - r"""Computes the autocorrelation function. + r""" + Computes the autocorrelation function. Calls :py:func:`base_multitemper.ensemble_compute_acf`; see that function for details. @@ -370,12 +412,14 @@ def compute_acf(filename, **kwargs): ``per-walker=True`` is passed as a keyword argument, the arrays will have shape ``ntemps x nwalkers x niterations``. Otherwise, the returned array will have shape ``ntemps x niterations``. + """ return ensemble_compute_acf(filename, **kwargs) @staticmethod def compute_acl(filename, **kwargs): - r"""Computes the autocorrelation length. + r""" + Computes the autocorrelation length. Calls :py:func:`base_multitemper.ensemble_compute_acl`; see that function for details. @@ -392,13 +436,14 @@ def compute_acl(filename, **kwargs): ------- dict A dictionary of ntemps-long arrays of the ACLs of each parameter. + """ return ensemble_compute_acl(filename, **kwargs) @classmethod - def from_config(cls, cp, model, output_file=None, nprocesses=1, - use_mpi=False): - """Loads the sampler from the given config file. + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): + """ + Loads the sampler from the given config file. The following options are retrieved in the ``[sampler]`` section: @@ -481,11 +526,13 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, ------- EmceePTSampler : The sampler instance. + """ section = "sampler" # check name assert cp.get(section, "name") == cls.name, ( - "name in section [sampler] must match mine") + "name in section [sampler] must match mine" + ) # get the number of walkers to use nwalkers = int(cp.get(section, "nwalkers")) # get the checkpoint interval, if it's specified @@ -493,51 +540,56 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, checkpoint_signal = cls.ckpt_signal_from_config(cp, section) optargs = {} # get the temperature level settings - ntemps = get_optional_arg_from_config(cp, section, 'ntemps', int) + ntemps = get_optional_arg_from_config(cp, section, "ntemps", int) if ntemps is not None: - optargs['ntemps'] = ntemps - tmax = get_optional_arg_from_config(cp, section, 'tmax', float) + optargs["ntemps"] = ntemps + tmax = get_optional_arg_from_config(cp, section, "tmax", float) if tmax is not None: - optargs['Tmax'] = tmax - betas = get_optional_arg_from_config(cp, section, 'betas') + optargs["Tmax"] = tmax + betas = get_optional_arg_from_config(cp, section, "betas") if betas is not None: # convert to list sorted in descencding order betas = numpy.sort(list(map(float, shlex.split(betas))))[::-1] - optargs['betas'] = betas - betas_file = get_optional_arg_from_config(cp, section, 'betas-file') + optargs["betas"] = betas + betas_file = get_optional_arg_from_config(cp, section, "betas-file") if betas_file is not None: - optargs['betas'] = read_betas_from_hdf(betas_file) + optargs["betas"] = read_betas_from_hdf(betas_file) # check for consistency if betas is not None and betas_file is not None: raise ValueError("provide either betas or betas-file, not both") - if 'betas' in optargs and (ntemps is not None or tmax is not None): - raise ValueError("provide either ntemps/tmax or betas/betas-file, " - "not both") + if "betas" in optargs and (ntemps is not None or tmax is not None): + raise ValueError("provide either ntemps/tmax or betas/betas-file, not both") # adaptation parameters - adaptive = get_optional_arg_from_config(cp, section, 'adaptive') + adaptive = get_optional_arg_from_config(cp, section, "adaptive") if adaptive is not None: - optargs['adaptive'] = True + optargs["adaptive"] = True else: - optargs['adaptive'] = False - adaptation_lag = get_optional_arg_from_config(cp, section, - 'adaptation-lag', int) + optargs["adaptive"] = False + adaptation_lag = get_optional_arg_from_config( + cp, section, "adaptation-lag", int + ) if adaptation_lag is not None: - optargs['adaptation_lag'] = adaptation_lag - adaptation_time = get_optional_arg_from_config(cp, section, - 'adaptation-time', int) + optargs["adaptation_lag"] = adaptation_lag + adaptation_time = get_optional_arg_from_config( + cp, section, "adaptation-time", int + ) if adaptation_time is not None: - optargs['adaptation_time'] = adaptation_time - scale_factor = get_optional_arg_from_config(cp, section, - 'scale-factor', float) + optargs["adaptation_time"] = adaptation_time + scale_factor = get_optional_arg_from_config(cp, section, "scale-factor", float) if scale_factor is not None: - optargs['scale_factor'] = scale_factor + optargs["scale_factor"] = scale_factor # get the loglikelihood function - logl = get_optional_arg_from_config(cp, section, 'logl-function') - obj = cls(model, nwalkers, - checkpoint_interval=checkpoint_interval, - checkpoint_signal=checkpoint_signal, - loglikelihood_function=logl, nprocesses=nprocesses, - use_mpi=use_mpi, **optargs) + logl = get_optional_arg_from_config(cp, section, "logl-function") + obj = cls( + model, + nwalkers, + checkpoint_interval=checkpoint_interval, + checkpoint_signal=checkpoint_signal, + loglikelihood_function=logl, + nprocesses=nprocesses, + use_mpi=use_mpi, + **optargs, + ) # set target obj.set_target_from_config(cp, section) # add burn-in if it's specified @@ -553,7 +605,8 @@ def from_config(cls, cp, model, output_file=None, nprocesses=1, return obj def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters @@ -561,12 +614,15 @@ def write_results(self, filename): filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples - fp.write_samples(self.samples, - parameters=self.model.variable_params, - last_iteration=self.niterations) + fp.write_samples( + self.samples, + parameters=self.model.variable_params, + last_iteration=self.niterations, + ) # write stats fp.write_samples(self.model_stats, last_iteration=self.niterations) # write random state @@ -579,17 +635,18 @@ def write_results(self, filename): fp.write_ensemble_attrs(self.ensemble) def _correctjacobian(self, samples): - """Corrects the log jacobian values stored on disk. + """ + Corrects the log jacobian values stored on disk. Parameters ---------- samples : dict Dictionary of the samples. + """ # flatten samples for evaluating orig_shape = list(samples.values())[0].shape - flattened_samples = {p: arr.ravel() - for p, arr in list(samples.items())} + flattened_samples = {p: arr.ravel() for p, arr in list(samples.items())} # convert to a list of tuples so we can use map function params = list(flattened_samples.keys()) size = flattened_samples[params[0]].size @@ -602,7 +659,8 @@ def _correctjacobian(self, samples): return logj.reshape(orig_shape) def finalize(self): - """Calculates the log evidence and writes to the checkpoint file. + """ + Calculates the log evidence and writes to the checkpoint file. If sampling transforms were used, this also corrects the jacobian stored on disk. @@ -613,27 +671,34 @@ def finalize(self): if self.model.sampling_transforms is not None: # fix the lobjacobian values stored on disk logging.info("Correcting logjacobian values on disk") - with self.io(self.checkpoint_file, 'r') as fp: - samples = fp.read_raw_samples(self.variable_params, - thin_start=0, - thin_interval=1, thin_end=None, - temps='all', flatten=False) + with self.io(self.checkpoint_file, "r") as fp: + samples = fp.read_raw_samples( + self.variable_params, + thin_start=0, + thin_interval=1, + thin_end=None, + temps="all", + flatten=False, + ) logjacobian = self._correctjacobian(samples) # write them back out for fn in [self.checkpoint_file, self.backup_file]: with self.io(fn, "a") as fp: - fp[fp.samples_group]['logjacobian'][()] = logjacobian + fp[fp.samples_group]["logjacobian"][()] = logjacobian logging.info("Calculating log evidence") # get the thinning settings - with self.io(self.checkpoint_file, 'r') as fp: + with self.io(self.checkpoint_file, "r") as fp: thin_start = fp.thin_start thin_interval = fp.thin_interval thin_end = fp.thin_end # calculate logz, dlogz = self.calculate_logevidence( - self.checkpoint_file, thin_start=thin_start, thin_end=thin_end, - thin_interval=thin_interval) - logging.info("log Z, dlog Z: {}, {}".format(logz, dlogz)) + self.checkpoint_file, + thin_start=thin_start, + thin_end=thin_end, + thin_interval=thin_interval, + ) + logging.info(f"log Z, dlog Z: {logz}, {dlogz}") # write to both the checkpoint and backup for fn in [self.checkpoint_file, self.backup_file]: with self.io(fn, "a") as fp: diff --git a/pycbc/inference/sampler/refine.py b/pycbc/inference/sampler/refine.py index 0139452294e..1b2d9fa1afd 100644 --- a/pycbc/inference/sampler/refine.py +++ b/pycbc/inference/sampler/refine.py @@ -1,19 +1,18 @@ -""" Sampler that uses kde refinement of an existing posterior estimate. -""" +"""Sampler that uses kde refinement of an existing posterior estimate.""" import logging + import numpy import numpy.random - from scipy.special import logsumexp -from scipy.stats import gaussian_kde from scipy.stats import entropy as sentropy +from scipy.stats import gaussian_kde from pycbc.inference import models -from pycbc.pool import choose_pool from pycbc.inference.io import loadfile +from pycbc.pool import choose_pool -from .base import setup_output, initial_dist_from_config +from .base import initial_dist_from_config, setup_output from .dummy import DummySampler @@ -51,7 +50,8 @@ def resample_equal(samples, logwt, seed=0): class RefineSampler(DummySampler): - """Sampler for kde drawn refinement of existing posterior estimate + """ + Sampler for kde drawn refinement of existing posterior estimate Parameters ---------- @@ -70,6 +70,7 @@ class RefineSampler(DummySampler): The target evidence difference between iterative kde updates kde: scipy.stats.gaussian_kde The inital kde to use. + """ name = "refine" @@ -90,7 +91,7 @@ def __init__( kde=None, update_groups=None, max_kde_samples=int(5e4), - **kwargs + **kwargs, ): super().__init__(model, *args) @@ -176,8 +177,7 @@ def converged(self, step, kde_new, factor, logp): frac_offbase = (logp < logp.max() - 5.0).sum() / len(logp) logging.info( - "%s: dlogz_iter=%.4f," - "dlogz_half=%.4f, entropy=%.4f offbase fraction=%.4f", + "%s: dlogz_iter=%.4f,dlogz_half=%.4f, entropy=%.4f offbase fraction=%.4f", step, dlogz, dlogz2, @@ -191,13 +191,10 @@ def converged(self, step, kde_new, factor, logp): and frac_offbase < self.offbase_fraction ): return True - else: - return False + return False @classmethod - def from_config( - cls, cp, model, output_file=None, nprocesses=1, use_mpi=False - ): + def from_config(cls, cp, model, output_file=None, nprocesses=1, use_mpi=False): """This should initialize the sampler given a config file.""" kwargs = {k: cp.get("sampler", k) for k in cp.options("sampler")} obj = cls(model, nprocesses=nprocesses, use_mpi=use_mpi, **kwargs) @@ -290,9 +287,7 @@ def run(self): gsample = self.group_kde.resample(int(1e5)) gsample = [ - gsample[i, :] - for i, k in enumerate(self.vparam) - if k in param_group + gsample[i, :] for i, k in enumerate(self.vparam) if k in param_group ] self.kde = gaussian_kde(numpy.array(gsample)) self.fixed_samples = self.group_kde.resample(1) @@ -307,9 +302,7 @@ def run(self): ) if total_samples is not None: - total_samples = numpy.concatenate( - [total_samples, ksamples], axis=1 - ) + total_samples = numpy.concatenate([total_samples, ksamples], axis=1) total_logp = numpy.concatenate([total_logp, logp]) total_logw = numpy.concatenate([total_logw, logw]) total_logl = numpy.concatenate([total_logl, logl]) @@ -321,9 +314,7 @@ def run(self): logging.info("setting up next kde iteration..") ntotal_logw = total_logw - logsumexp(total_logw) - kde_new = gaussian_kde( - total_samples, weights=numpy.exp(ntotal_logw) - ) + kde_new = gaussian_kde(total_samples, weights=numpy.exp(ntotal_logw)) if self.converged(r, kde_new, total_logl + total_logw, logp): break diff --git a/pycbc/inference/sampler/snowline.py b/pycbc/inference/sampler/snowline.py index 6a54825879a..34751dae8ae 100644 --- a/pycbc/inference/sampler/snowline.py +++ b/pycbc/inference/sampler/snowline.py @@ -26,15 +26,15 @@ packages for parameter estimation. """ -import sys import logging +import sys from pycbc.inference.io.snowline import SnowlineFile from pycbc.io.hdf import dump_state from pycbc.pool import use_mpi -from .base import (BaseSampler, setup_output) -from .base_cube import setup_calls +from .base import BaseSampler, setup_output +from .base_cube import setup_calls # # ============================================================================= @@ -44,15 +44,19 @@ # ============================================================================= # + class SnowlineSampler(BaseSampler): - """This class is used to construct an Snowline sampler from the snowline + """ + This class is used to construct an Snowline sampler from the snowline package. Parameters ---------- model : model A model from ``pycbc.inference.models`` + """ + name = "snowline" _io = SnowlineFile @@ -60,12 +64,12 @@ def __init__(self, model, **kwargs): super().__init__(model) import snowline + log_likelihood_call, prior_call = setup_calls(model, copy_prior=True) self._sampler = snowline.ReactiveImportanceSampler( - list(self.model.variable_params), - log_likelihood_call, - transform=prior_call) + list(self.model.variable_params), log_likelihood_call, transform=prior_call + ) do_mpi, _, rank = use_mpi() self.main = (not do_mpi) or (rank == 0) @@ -85,7 +89,7 @@ def io(self): @property def niterations(self): - return self.result['niter'] + return self.result["niter"] @classmethod def from_config(cls, cp, model, output_file=None, **kwds): @@ -93,15 +97,16 @@ def from_config(cls, cp, model, output_file=None, **kwds): Loads the sampler from the given config file. """ skeys = {} - opts = {'num_global_samples': int, - 'num_gauss_samples': int, - 'max_ncalls': int, - 'min_ess': int, - 'max_improvement_loops': int - } + opts = { + "num_global_samples": int, + "num_gauss_samples": int, + "max_ncalls": int, + "min_ess": int, + "max_improvement_loops": int, + } for opt_name in opts: - if cp.has_option('sampler', opt_name): - value = cp.get('sampler', opt_name) + if cp.has_option("sampler", opt_name): + value = cp.get("sampler", opt_name) skeys[opt_name] = opts[opt_name](value) inst = cls(model, **skeys) @@ -111,12 +116,10 @@ def from_config(cls, cp, model, output_file=None, **kwds): return inst def checkpoint(self): - """ There is currently no checkpointing implemented""" - pass + """There is currently no checkpointing implemented""" def resume_from_checkpoint(self): - """ There is currently no checkpointing implemented""" - pass + """There is currently no checkpointing implemented""" def finalize(self): logging.info("Writing samples to files") @@ -129,13 +132,14 @@ def model_stats(self): @property def samples(self): - samples = self.result['samples'] + samples = self.result["samples"] params = list(self.model.variable_params) samples_dict = {p: samples[:, i] for i, p in enumerate(params)} return samples_dict def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters @@ -143,26 +147,23 @@ def write_results(self, filename): filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples fp.write_samples(self.samples, self.samples.keys()) # write log evidence fp.write_logevidence(self.logz, self.logz_err) # write full results - dump_state(self.result, fp, - path='sampler_info', - dsetname='presult') + dump_state(self.result, fp, path="sampler_info", dsetname="presult") @property def logz(self): - """Return bayesian evidence estimated by snowline sampler. - """ - return self.result['logz'] + """Return bayesian evidence estimated by snowline sampler.""" + return self.result["logz"] @property def logz_err(self): - """Return error in bayesian evidence estimated by snowline sampler. - """ - return self.result['logzerr'] + """Return error in bayesian evidence estimated by snowline sampler.""" + return self.result["logzerr"] diff --git a/pycbc/inference/sampler/ultranest.py b/pycbc/inference/sampler/ultranest.py index f365af3c642..ed61ab191d0 100644 --- a/pycbc/inference/sampler/ultranest.py +++ b/pycbc/inference/sampler/ultranest.py @@ -26,16 +26,17 @@ packages for parameter estimation. """ -import sys import logging +import sys + import numpy from pycbc.inference.io.ultranest import UltranestFile from pycbc.io.hdf import dump_state from pycbc.pool import use_mpi -from .base import (BaseSampler, setup_output) -from .base_cube import setup_calls +from .base import BaseSampler, setup_output +from .base_cube import setup_calls # # ============================================================================= @@ -45,8 +46,10 @@ # ============================================================================= # + class UltranestSampler(BaseSampler): - """This class is used to construct an Ultranest sampler from the ultranest + """ + This class is used to construct an Ultranest sampler from the ultranest package. Parameters @@ -58,17 +61,19 @@ class UltranestSampler(BaseSampler): stepsampling : bool If false, uses rejection sampling. If true, uses hit-and-run sampler, which scales better with dimensionality. + """ + name = "ultranest" _io = UltranestFile - def __init__(self, model, log_dir=None, - stepsampling=False, - enable_plots=False, - **kwargs): - super(UltranestSampler, self).__init__(model) + def __init__( + self, model, log_dir=None, stepsampling=False, enable_plots=False, **kwargs + ): + super().__init__(model) import ultranest + log_likelihood_call, prior_call = setup_calls(model, copy_prior=True) # Check for cyclic boundaries @@ -76,7 +81,7 @@ def __init__(self, model, log_dir=None, cyclic = self.model.prior_distribution.cyclic for param in self.variable_params: if param in cyclic: - logging.info('Param: %s will be cyclic', param) + logging.info("Param: %s will be cyclic", param) periodic.append(True) else: periodic.append(False) @@ -84,15 +89,18 @@ def __init__(self, model, log_dir=None, self._sampler = ultranest.ReactiveNestedSampler( list(self.model.variable_params), log_likelihood_call, - prior_call, log_dir=log_dir, + prior_call, + log_dir=log_dir, wrapped_params=periodic, - resume=True) + resume=True, + ) if stepsampling: import ultranest.stepsampler + self._sampler.stepsampler = ultranest.stepsampler.RegionBallSliceSampler( - nsteps=100, adaptive_nsteps='move-distance', - region_filter=True) + nsteps=100, adaptive_nsteps="move-distance", region_filter=True + ) self.enable_plots = enable_plots self.nlive = 0 @@ -118,7 +126,7 @@ def io(self): @property def niterations(self): - return self.result['niter'] + return self.result["niter"] @classmethod def from_config(cls, cp, model, output_file=None, **kwds): @@ -126,26 +134,28 @@ def from_config(cls, cp, model, output_file=None, **kwds): Loads the sampler from the given config file. """ skeys = {} - opts = {'update_interval_iter_fraction': float, - 'update_interval_ncall': int, - 'log_interval': int, - 'show_status': bool, - 'dlogz': float, - 'dKL': float, - 'frac_remain': float, - 'Lepsilon': float, - 'min_ess': int, - 'max_iters': int, - 'max_ncalls': int, - 'log_dir': str, - 'stepsampling': bool, - 'enable_plots': bool, - 'max_num_improvement_loops': int, - 'min_num_live_points': int, - 'cluster_num_live_points:': int} + opts = { + "update_interval_iter_fraction": float, + "update_interval_ncall": int, + "log_interval": int, + "show_status": bool, + "dlogz": float, + "dKL": float, + "frac_remain": float, + "Lepsilon": float, + "min_ess": int, + "max_iters": int, + "max_ncalls": int, + "log_dir": str, + "stepsampling": bool, + "enable_plots": bool, + "max_num_improvement_loops": int, + "min_num_live_points": int, + "cluster_num_live_points:": int, + } for opt_name in opts: - if cp.has_option('sampler', opt_name): - value = cp.get('sampler', opt_name) + if cp.has_option("sampler", opt_name): + value = cp.get("sampler", opt_name) skeys[opt_name] = opts[opt_name](value) inst = cls(model, **skeys) @@ -176,22 +186,23 @@ def samples(self): # we'll do the resampling ourselves so we can pick up # additional parameters try: # Remove me on next ultranest release - wsamples = self.result['weighted_samples']['v'] - weights = self.result['weighted_samples']['w'] - logl = self.result['weighted_samples']['L'] + wsamples = self.result["weighted_samples"]["v"] + weights = self.result["weighted_samples"]["w"] + logl = self.result["weighted_samples"]["L"] except KeyError: - wsamples = self.result['weighted_samples']['points'] - weights = self.result['weighted_samples']['weights'] - logl = self.result['weighted_samples']['logl'] + wsamples = self.result["weighted_samples"]["points"] + weights = self.result["weighted_samples"]["weights"] + logl = self.result["weighted_samples"]["logl"] wsamples = numpy.column_stack((wsamples, logl)) - params = list(self.model.variable_params) + ['loglikelihood'] + params = list(self.model.variable_params) + ["loglikelihood"] samples = resample_equal(wsamples, weights / weights.sum()) samples_dict = {p: samples[:, i] for i, p in enumerate(params)} return samples_dict def write_results(self, filename): - """Writes samples, model stats, acceptance fraction, and random state + """ + Writes samples, model stats, acceptance fraction, and random state to the given file. Parameters @@ -199,26 +210,23 @@ def write_results(self, filename): filename : str The file to write to. The file is opened using the ``io`` class in an an append state. + """ - with self.io(filename, 'a') as fp: + with self.io(filename, "a") as fp: # write samples fp.write_samples(self.samples, self.samples.keys()) # write log evidence fp.write_logevidence(self.logz, self.logz_err) # write full ultranest formatted results - dump_state(self.result, fp, - path='sampler_info', - dsetname='presult') + dump_state(self.result, fp, path="sampler_info", dsetname="presult") @property def logz(self): - """Return bayesian evidence estimated by ultranest sampler. - """ - return self.result['logz'] + """Return bayesian evidence estimated by ultranest sampler.""" + return self.result["logz"] @property def logz_err(self): - """Return error in bayesian evidence estimated by ultranest sampler. - """ - return self.result['logzerr'] + """Return error in bayesian evidence estimated by ultranest sampler.""" + return self.result["logzerr"] diff --git a/pycbc/inject/__init__.py b/pycbc/inject/__init__.py index 21b9e339a13..4577c25c41b 100644 --- a/pycbc/inject/__init__.py +++ b/pycbc/inject/__init__.py @@ -1,2 +1,2 @@ -from pycbc.inject.injfilterrejector import * from pycbc.inject.inject import * +from pycbc.inject.injfilterrejector import * diff --git a/pycbc/inject/inject.py b/pycbc/inject/inject.py index f63d10df011..c9afd4f1ecc 100644 --- a/pycbc/inject/inject.py +++ b/pycbc/inject/inject.py @@ -25,31 +25,35 @@ # """This module provides utilities for injecting signals into data""" -import os -import numpy as np import copy import logging +import os from abc import ABCMeta, abstractmethod import lal -from igwn_ligolw import utils as ligolw_utils, ligolw, lsctables +import numpy as np +from igwn_ligolw import ligolw, lsctables +from igwn_ligolw import utils as ligolw_utils -from pycbc import waveform, frame, libutils -from pycbc.opt import LimitedSizeDict -from pycbc.waveform import (get_td_waveform, fd_det, - get_td_det_waveform_from_fd_det) -from pycbc.waveform import utils as wfutils -from pycbc.waveform import ringdown_td_approximants -from pycbc.types import float64, float32, TimeSeries, load_timeseries -from pycbc.detector import Detector +import pycbc.io +from pycbc import frame, libutils, waveform from pycbc.conversions import tau0_from_mass1_mass2 +from pycbc.detector import Detector from pycbc.filter import resample_to_delta_t -import pycbc.io from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.opt import LimitedSizeDict +from pycbc.types import TimeSeries, float32, float64, load_timeseries +from pycbc.waveform import ( + fd_det, + get_td_det_waveform_from_fd_det, + get_td_waveform, + ringdown_td_approximants, +) +from pycbc.waveform import utils as wfutils -logger = logging.getLogger('pycbc.inject.inject') +logger = logging.getLogger("pycbc.inject.inject") -sim = libutils.import_optional('lalsimulation') +sim = libutils.import_optional("lalsimulation") injection_func_map = { np.dtype(float32): lambda *args: sim.SimAddInjectionREAL4TimeSeries(*args), @@ -59,10 +63,11 @@ # Map parameter names used in pycbc to names used in the sim_inspiral # table, if they are different sim_inspiral_map = { - 'ra': 'longitude', - 'dec': 'latitude', - 'approximant': 'waveform', - } + "ra": "longitude", + "dec": "latitude", + "approximant": "waveform", +} + def set_sim_data(inj, field, data): """Sets data of a SimInspiral instance.""" @@ -71,18 +76,19 @@ def set_sim_data(inj, field, data): except KeyError: sim_field = field # for tc, map to geocentric times - if sim_field == 'tc': + if sim_field == "tc": inj.geocent_end_time = int(data) - inj.geocent_end_time_ns = int(1e9*(data % 1)) + inj.geocent_end_time_ns = int(1e9 * (data % 1)) # for spin1 and spin2 we need data to be an array - if sim_field in ['spin1', 'spin2']: + if sim_field in ["spin1", "spin2"]: setattr(inj, sim_field, [0, 0, data]) else: setattr(inj, sim_field, data) def projector(detector_name, inj, hp, hc, distance_scale=1): - """ Use the injection row to project the polarizations into the + """ + Use the injection row to project the polarizations into the detector frame """ detector = Detector(detector_name) @@ -104,31 +110,42 @@ def projector(detector_name, inj, hp, hc, distance_scale=1): # taper the polarizations try: - hp_tapered = hp.taper_timeseries(location=inj.taper, - tapermethod=inj.get('taper_method', 'lal'), - taper_window=inj.get('taper_window')) - hc_tapered = hc.taper_timeseries(location=inj.taper, - tapermethod=inj.get('taper_method', 'lal'), - taper_window=inj.get('taper_window')) + hp_tapered = hp.taper_timeseries( + location=inj.taper, + tapermethod=inj.get("taper_method", "lal"), + taper_window=inj.get("taper_window"), + ) + hc_tapered = hc.taper_timeseries( + location=inj.taper, + tapermethod=inj.get("taper_method", "lal"), + taper_window=inj.get("taper_window"), + ) except AttributeError: hp_tapered = hp hc_tapered = hc - projection_method = 'lal' - if hasattr(inj, 'detector_projection_method'): + projection_method = "lal" + if hasattr(inj, "detector_projection_method"): projection_method = inj.detector_projection_method - logger.info('Injecting at %s, method is %s', tc, projection_method) + logger.info("Injecting at %s, method is %s", tc, projection_method) # compute the detector response and add it to the strain - signal = detector.project_wave(hp_tapered, hc_tapered, - ra, dec, inj.polarization, - method=projection_method, - reference_time=tc,) + signal = detector.project_wave( + hp_tapered, + hc_tapered, + ra, + dec, + inj.polarization, + method=projection_method, + reference_time=tc, + ) return signal + def legacy_approximant_name(apx): - """Convert the old style xml approximant name to a name + """ + Convert the old style xml approximant name to a name and phase_order. Alex: I hate this function. Please delete this when we use Collin's new tables. """ @@ -142,9 +159,9 @@ def legacy_approximant_name(apx): return name, order -class _XMLInjectionSet(object): - - """Manages sets of injections: reads injections from LIGOLW XML files +class _XMLInjectionSet: + """ + Manages sets of injections: reads injections from LIGOLW XML files and injects them into time series. Parameters @@ -157,11 +174,13 @@ class _XMLInjectionSet(object): ---------- indoc table + """ def __init__(self, sim_file, **kwds): self.indoc = ligolw_utils.load_filename( - sim_file, False, contenthandler=LIGOLWContentHandler) + sim_file, False, contenthandler=LIGOLWContentHandler + ) self.table = lsctables.SimInspiralTable.get_table(self.indoc) self.extra_args = kwds @@ -176,7 +195,8 @@ def apply( injection_sample_rate=None, generate_injections=True, ): - """Add injections (as seen by a particular detector) to a time series. + """ + Add injections (as seen by a particular detector) to a time series. Parameters ---------- @@ -214,6 +234,7 @@ def apply( ------ TypeError For invalid types of `strain`. + """ if not generate_injections: raise NotImplementedError( @@ -222,8 +243,9 @@ def apply( ) if strain.dtype not in (float32, float64): - raise TypeError("Strain dtype must be float32 or float64, not " \ - + str(strain.dtype)) + raise TypeError( + "Strain dtype must be float32 or float64, not " + str(strain.dtype) + ) lalstrain = strain.lal() earth_travel_time = lal.REARTH_SI / lal.C_SI @@ -239,8 +261,9 @@ def apply( injections = self.table if simulation_ids: - injections = [inj for inj in injections \ - if inj.simulation_id in simulation_ids] + injections = [ + inj for inj in injections if inj.simulation_id in simulation_ids + ] injection_parameters = [] for inj in injections: f_l = inj.f_lower if f_lower is None else f_lower @@ -253,9 +276,10 @@ def apply( start_time = inj.time_geocent - 2 * (inj_length + 1) if end_time < t0 or start_time > t1: continue - signal = self.make_strain_from_inj_object(inj, delta_t, - detector_name, f_lower=f_l, distance_scale=distance_scale) - signal = resample_to_delta_t(signal, strain.delta_t, method='ldas') + signal = self.make_strain_from_inj_object( + inj, delta_t, detector_name, f_lower=f_l, distance_scale=distance_scale + ) + signal = resample_to_delta_t(signal, strain.delta_t, method="ldas") if float(signal.start_time) > t1: continue @@ -276,13 +300,15 @@ def apply( inj_filter_rejector.injection_params = injected return injected - def make_strain_from_inj_object(self, inj, delta_t, detector_name, - f_lower=None, distance_scale=1): - """Make a h(t) strain time-series from an injection object as read from + def make_strain_from_inj_object( + self, inj, delta_t, detector_name, f_lower=None, distance_scale=1 + ): + """ + Make a h(t) strain time-series from an injection object as read from a sim_inspiral table, for example. Parameters - ----------- + ---------- inj : injection object The injection object to turn into a strain h(t). delta_t : float @@ -297,9 +323,10 @@ def make_strain_from_inj_object(self, inj, delta_t, detector_name, no scaling. Returns - -------- + ------- signal : float h(t) corresponding to the injection. + """ f_l = inj.f_lower if f_lower is None else f_lower @@ -307,12 +334,15 @@ def make_strain_from_inj_object(self, inj, delta_t, detector_name, # compute the waveform time series hp, hc = get_td_waveform( - inj, approximant=name, delta_t=delta_t, + inj, + approximant=name, + delta_t=delta_t, phase_order=phase_order, - f_lower=f_l, distance=inj.distance, - **self.extra_args) - return projector(detector_name, - inj, hp, hc, distance_scale=distance_scale) + f_lower=f_l, + distance=inj.distance, + **self.extra_args, + ) + return projector(detector_name, inj, hp, hc, distance_scale=distance_scale) def end_times(self): """Return the end times of all injections""" @@ -320,7 +350,8 @@ def end_times(self): @staticmethod def write(filename, samples, write_params=None, static_args=None): - """Writes the injection samples to the given xml. + """ + Writes the injection samples to the given xml. Parameters ---------- @@ -334,6 +365,7 @@ def write(filename, samples, write_params=None, static_args=None): static_args : dict, optional Dictionary mapping static parameter names to values. These are written to the ``attrs``. + """ xmldoc = ligolw.Document() xmldoc.appendChild(ligolw.LIGO_LW()) @@ -352,17 +384,18 @@ def write(filename, samples, write_params=None, static_args=None): data = samples[ii][field] set_sim_data(sim, field, data) # set any static args - for (field, value) in static_args.items(): + for field, value in static_args.items(): set_sim_data(sim, field, value) simtable.append(sim) - ligolw_utils.write_filename(xmldoc, filename, compress='auto') + ligolw_utils.write_filename(xmldoc, filename, compress="auto") # ----------------------------------------------------------------------------- class _HDFInjectionSet(metaclass=ABCMeta): - r"""Manages sets of injections: reads injections from hdf files + r""" + Manages sets of injections: reads injections from hdf files and injects them into time series. Parameters @@ -381,6 +414,7 @@ class _HDFInjectionSet(metaclass=ABCMeta): required_params : tuple Parameter names that must exist in the injection HDF file in order to create an injection of that type. + """ _tableclass = pycbc.io.FieldArray @@ -389,7 +423,7 @@ class _HDFInjectionSet(metaclass=ABCMeta): def __init__(self, sim_file, hdf_group=None, **kwds): # open the file - fp = pycbc.io.HFile(sim_file, 'r') + fp = pycbc.io.HFile(sim_file, "r") group = fp if hdf_group is None else fp[hdf_group] # get parameters parameters = list(group.keys()) @@ -398,8 +432,8 @@ def __init__(self, sim_file, hdf_group=None, **kwds): # make sure Numpy S strings are loaded as strings and not bytestrings # (which could mess with approximant names, for example) for k in injvals: - if injvals[k].dtype.kind == 'S': - injvals[k] = injvals[k].astype('U') + if injvals[k].dtype.kind == "S": + injvals[k] = injvals[k].astype("U") # if there were no variable args, then we only have a single injection if len(parameters) == 0: numinj = 1 @@ -408,7 +442,7 @@ def __init__(self, sim_file, hdf_group=None, **kwds): # add any static args in the file try: # ensure parameter names are string types - self.static_args = group.attrs['static_args'].astype('U') + self.static_args = group.attrs["static_args"].astype("U") except KeyError: self.static_args = [] parameters.extend(self.static_args) @@ -427,14 +461,17 @@ def __init__(self, sim_file, hdf_group=None, **kwds): # times arr = np.repeat(val, numinj) # make sure any byte strings are stored as strings instead - if arr.dtype.char == 'S': - arr = arr.astype('U') + if arr.dtype.char == "S": + arr = arr.astype("U") injvals[param] = arr # make sure required parameters are provided missing = set(self.required_params) - set(injvals.keys()) if missing: - raise ValueError("required parameter(s) {} not found in the given " - "injection file".format(', '.join(missing))) + raise ValueError( + "required parameter(s) {} not found in the given injection file".format( + ", ".join(missing) + ) + ) # initialize the table self.table = self._tableclass.from_kwargs(**injvals) # save the extra arguments @@ -442,33 +479,35 @@ def __init__(self, sim_file, hdf_group=None, **kwds): fp.close() @abstractmethod - def apply(self, strain, detector_name, distance_scale=1, - simulation_ids=None, inj_filter_rejector=None, - **kwargs): + def apply( + self, + strain, + detector_name, + distance_scale=1, + simulation_ids=None, + inj_filter_rejector=None, + **kwargs, + ): """Adds injections to a detector's time series.""" - pass @abstractmethod - def make_strain_from_inj_object(self, inj, delta_t, detector_name, - distance_scale=1, **kwargs): - """Make a h(t) strain time-series from an injection object. - """ - pass + def make_strain_from_inj_object( + self, inj, delta_t, detector_name, distance_scale=1, **kwargs + ): + """Make a h(t) strain time-series from an injection object.""" @abstractmethod def end_times(self): """Return the end times of all injections""" - pass @abstractmethod def supported_approximants(self): """Return a list of the supported approximants.""" - pass @classmethod - def write(cls, filename, samples, write_params=None, static_args=None, - **metadata): - r"""Writes the injection samples to the given hdf file. + def write(cls, filename, samples, write_params=None, static_args=None, **metadata): + r""" + Writes the injection samples to the given hdf file. Parameters ---------- @@ -484,13 +523,14 @@ def write(cls, filename, samples, write_params=None, static_args=None, written to the ``attrs``. \**metadata : All other keyword arguments will be written to the file's attrs. + """ - with pycbc.io.HFile(filename, 'w') as fp: + with pycbc.io.HFile(filename, "w") as fp: # write metadata if static_args is None: static_args = {} fp.attrs["static_args"] = list(map(str, static_args.keys())) - fp.attrs['injtype'] = cls.injtype + fp.attrs["injtype"] = cls.injtype for key, val in metadata.items(): fp.attrs[key] = val if write_params is None: @@ -508,18 +548,18 @@ def write(cls, filename, samples, write_params=None, static_args=None, except TypeError as e: # can get this in python 3 if the val was a numpy.str_ type # we'll try again as a string type - if samples[field].dtype.char == 'U': - fp[field] = samples[field].astype('S') + if samples[field].dtype.char == "U": + fp[field] = samples[field].astype("S") else: raise e class CBCHDFInjectionSet(_HDFInjectionSet): - """Manages CBC injections. - """ + """Manages CBC injections.""" + _tableclass = pycbc.io.WaveformArray - injtype = 'cbc' - required_params = ('tc',) + injtype = "cbc" + required_params = ("tc",) def apply( self, @@ -532,7 +572,8 @@ def apply( injection_sample_rate=None, generate_injections=True, ): - """Add injections (as seen by a particular detector) to a time series. + """ + Add injections (as seen by a particular detector) to a time series. Parameters ---------- @@ -569,6 +610,7 @@ def apply( ------ TypeError For invalid types of `strain`. + """ # There is a case where injections are not added to the data, but # still need to be generated for the match threshold option in the @@ -582,18 +624,17 @@ def apply( "Be very sure that the injection file exactly matches " "how the injections were pregenerated!" ) - logging.warn(warn_msg) - must_make_injections = ( - generate_injections or ( - inj_filter_rejector and inj_filter_rejector.match_threshold - ) + logging.warning(warn_msg) + must_make_injections = generate_injections or ( + inj_filter_rejector and inj_filter_rejector.match_threshold ) if strain.dtype not in (float32, float64): - raise TypeError("Strain dtype must be float32 or float64, not " \ - + str(strain.dtype)) + raise TypeError( + "Strain dtype must be float32 or float64, not " + str(strain.dtype) + ) - if self.table[0]['approximant'] in fd_det: + if self.table[0]["approximant"] in fd_det: t0 = float(strain.start_time) t1 = float(strain.end_time) else: @@ -601,7 +642,6 @@ def apply( t0 = float(strain.start_time) - earth_travel_time t1 = float(strain.end_time) + earth_travel_time - if generate_injections: lalstrain = strain.lal() @@ -636,18 +676,13 @@ def apply( f_lower=f_l, distance_scale=distance_scale, ) - signal = resample_to_delta_t( - signal, - strain.delta_t, - method='ldas' - ) + signal = resample_to_delta_t(signal, strain.delta_t, method="ldas") if float(signal.start_time) > t1: continue signal = signal.astype(strain.dtype) injected_ids.append(ii) if inj_filter_rejector is not None: - inj_filter_rejector.generate_short_inj_from_inj(signal, ii) - + inj_filter_rejector.generate_short_inj_from_inj(signal, ii) if generate_injections: signal_lal = signal.lal() @@ -659,7 +694,7 @@ def apply( injected = copy.copy(self) injected.table = injections[np.array(injected_ids).astype(int)] if inj_filter_rejector is not None: - if hasattr(inj_filter_rejector, 'injected'): + if hasattr(inj_filter_rejector, "injected"): prev_p = inj_filter_rejector.injection_params prev_id = inj_filter_rejector.injection_ids injected = np.concatenate([prev_p, injected]) @@ -669,12 +704,14 @@ def apply( inj_filter_rejector.injection_ids = injected_ids return injected - def make_strain_from_inj_object(self, inj, delta_t, detector_name, - f_lower=None, distance_scale=1): - """Make a h(t) strain time-series from an injection object. + def make_strain_from_inj_object( + self, inj, delta_t, detector_name, f_lower=None, distance_scale=1 + ): + """ + Make a h(t) strain time-series from an injection object. Parameters - ----------- + ---------- inj : injection object The injection object to turn into a strain h(t). Can be any object which has waveform parameters as attributes, such as an @@ -691,26 +728,29 @@ def make_strain_from_inj_object(self, inj, delta_t, detector_name, no scaling. Returns - -------- + ------- signal : float h(t) corresponding to the injection. + """ if f_lower is None: f_l = inj.f_lower else: f_l = f_lower - if inj['approximant'] in fd_det: + if inj["approximant"] in fd_det: strain = get_td_det_waveform_from_fd_det( - inj, delta_t=delta_t, f_lower=f_l, - ifos=detector_name, **self.extra_args)[detector_name] + inj, delta_t=delta_t, f_lower=f_l, ifos=detector_name, **self.extra_args + )[detector_name] strain /= distance_scale else: # compute the waveform time series - hp, hc = get_td_waveform(inj, delta_t=delta_t, f_lower=f_l, - **self.extra_args) - strain = projector(detector_name, - inj, hp, hc, distance_scale=distance_scale) + hp, hc = get_td_waveform( + inj, delta_t=delta_t, f_lower=f_l, **self.extra_args + ) + strain = projector( + detector_name, inj, hp, hc, distance_scale=distance_scale + ) return strain def end_times(self): @@ -728,11 +768,13 @@ def supported_approximants(): class RingdownHDFInjectionSet(_HDFInjectionSet): - """Manages a ringdown injection: reads injection from hdf file + """ + Manages a ringdown injection: reads injection from hdf file and injects it into time series. """ - injtype = 'ringdown' - required_params = ('tc',) + + injtype = "ringdown" + required_params = ("tc",) def apply( self, @@ -744,7 +786,8 @@ def apply( injection_sample_rate=None, generate_injections=True, ): - """Add injection (as seen by a particular detector) to a time series. + """ + Add injection (as seen by a particular detector) to a time series. Parameters ---------- @@ -781,10 +824,12 @@ def apply( Or if generate_injections is not True. TypeError For invalid types of `strain`. + """ if inj_filter_rejector is not None: - raise NotImplementedError("Ringdown injections do not support " - "inj_filter_rejector") + raise NotImplementedError( + "Ringdown injections do not support inj_filter_rejector" + ) if not generate_injections: raise NotImplementedError( @@ -793,8 +838,9 @@ def apply( ) if strain.dtype not in (float32, float64): - raise TypeError("Strain dtype must be float32 or float64, not " \ - + str(strain.dtype)) + raise TypeError( + "Strain dtype must be float32 or float64, not " + str(strain.dtype) + ) lalstrain = strain.lal() @@ -811,22 +857,24 @@ def apply( for ii in range(injections.size): injection = injections[ii] signal = self.make_strain_from_inj_object( - injection, delta_t, detector_name, - distance_scale=distance_scale) - signal = resample_to_delta_t(signal, strain.delta_t, method='ldas') + injection, delta_t, detector_name, distance_scale=distance_scale + ) + signal = resample_to_delta_t(signal, strain.delta_t, method="ldas") signal = signal.astype(strain.dtype) signal_lal = signal.lal() add_injection(lalstrain, signal_lal, None) strain.data[:] = lalstrain.data.data[:] - def make_strain_from_inj_object(self, inj, delta_t, detector_name, - distance_scale=1): - """Make a h(t) strain time-series from an injection object as read from + def make_strain_from_inj_object( + self, inj, delta_t, detector_name, distance_scale=1 + ): + """ + Make a h(t) strain time-series from an injection object as read from an hdf file. Parameters - ----------- + ---------- inj : injection object The injection object to turn into a strain h(t). delta_t : float @@ -838,18 +886,20 @@ def make_strain_from_inj_object(self, inj, delta_t, detector_name, is no scaling. Returns - -------- + ------- signal : float h(t) corresponding to the injection. + """ # compute the waveform time series - hp, hc = ringdown_td_approximants[inj['approximant']]( - inj, delta_t=delta_t, **self.extra_args) - return projector(detector_name, - inj, hp, hc, distance_scale=distance_scale) + hp, hc = ringdown_td_approximants[inj["approximant"]]( + inj, delta_t=delta_t, **self.extra_args + ) + return projector(detector_name, inj, hp, hc, distance_scale=distance_scale) def end_times(self): - """Return the approximate end times of all injections. + """ + Return the approximate end times of all injections. Currently, this just assumes all ringdowns are 2 seconds long. """ @@ -864,7 +914,8 @@ def supported_approximants(): class IncoherentFromFileHDFInjectionSet(_HDFInjectionSet): - """Manages injecting an arbitrary time series loaded from a file. + """ + Manages injecting an arbitrary time series loaded from a file. The injections must have the following attributes set: @@ -921,22 +972,25 @@ class IncoherentFromFileHDFInjectionSet(_HDFInjectionSet): In order to use with ``pycbc_create_injections``, set the ``approximant`` name to ``'incoherent_from_file'``. """ - injtype = 'incoherent_from_file' - required_params = ('filename', 'ref_point') + + injtype = "incoherent_from_file" + required_params = ("filename", "ref_point") _buffersize = 10 _buffer = None _rtbuffer = None def end_times(self): - raise NotImplementedError("IncoherentFromFile times cannot be " - "determined without loading time series") + raise NotImplementedError( + "IncoherentFromFile times cannot be determined without loading time series" + ) @staticmethod def supported_approximants(): - return ['incoherent_from_file'] + return ["incoherent_from_file"] def loadts(self, inj): - """Loads an injection time series. + """ + Loads an injection time series. After the first time a time series is loaded it will be added to an internal buffer for faster in case another injection uses the same @@ -950,13 +1004,12 @@ def loadts(self, inj): except KeyError: pass # not in buffer, so load - if inj.filename.endswith('.gwf'): + if inj.filename.endswith(".gwf"): try: channel = inj.channel except AttributeError as _err: # Py3.XX: uncomment the "from _err" when we drop 2.7 - raise ValueError("Must provide a channel for " - "frame files") #from _err + raise ValueError("Must provide a channel for frame files") # from _err ts = frame.read_frame(inj.filename, channel) else: ts = load_timeseries(inj.filename) @@ -965,15 +1018,17 @@ def loadts(self, inj): return ts def set_ref_time(self, inj, ts): - """Sets t=0 of the given time series based on what the given + """ + Sets t=0 of the given time series based on what the given injection's ``ref_point`` is. """ try: ref_point = inj.ref_point except AttributeError as _err: # Py3.XX: uncomment the "from _err" when we drop 2.7 - raise ValueError("Must provide a ref_point for {} injections" - .format(self.injtype)) #from _err + raise ValueError( + f"Must provide a ref_point for {self.injtype} injections" + ) # from _err # try to get from buffer if self._rtbuffer is None: self._rtbuffer = LimitedSizeDict(size_limit=self._buffersize) @@ -981,24 +1036,24 @@ def set_ref_time(self, inj, ts): reftime = self._rtbuffer[inj.filename, ref_point] except KeyError: if ref_point == "start": - reftime = 0. + reftime = 0.0 elif ref_point == "end": - reftime = -len(ts)*ts.delta_t + reftime = -len(ts) * ts.delta_t elif ref_point == "center": - reftime = -len(ts)*ts.delta_t/2. + reftime = -len(ts) * ts.delta_t / 2.0 elif ref_point == "absmax": - reftime = -ts.abs_arg_max()*ts.delta_t + reftime = -ts.abs_arg_max() * ts.delta_t elif isinstance(ref_point, (float, int)): reftime = -float(ref_point) else: - raise ValueError("Unrecognized ref_point {} provided" - .format(ref_point)) + raise ValueError(f"Unrecognized ref_point {ref_point} provided") self._rtbuffer[inj.filename, ref_point] = reftime ts._epoch = reftime @staticmethod def slice_and_taper(inj, ts): - """Slices and tapers a timeseries based on the injection settings. + """ + Slices and tapers a timeseries based on the injection settings. This assumes that ``set_ref_time`` has been applied to the timeseries first. A copy of the time series will be returned even if no slicing @@ -1019,15 +1074,15 @@ def slice_and_taper(inj, ts): except AttributeError: twidth = 0 if twidth: - ts = wfutils.td_taper(ts, ts.start_time, ts.start_time+twidth, - side='left') + ts = wfutils.td_taper( + ts, ts.start_time, ts.start_time + twidth, side="left" + ) try: twidth = inj.right_taper_width except AttributeError: twidth = 0 if twidth: - ts = wfutils.td_taper(ts, ts.end_time-twidth, ts.end_time, - side='right') + ts = wfutils.td_taper(ts, ts.end_time - twidth, ts.end_time, side="right") return ts def apply( @@ -1037,11 +1092,12 @@ def apply( distance_scale=1, injection_sample_rate=None, inj_filter_rejector=None, - generate_injections=True + generate_injections=True, ): if inj_filter_rejector is not None: - raise NotImplementedError("IncoherentFromFile injections do not " - "support inj_filter_rejector") + raise NotImplementedError( + "IncoherentFromFile injections do not support inj_filter_rejector" + ) if not generate_injections: raise NotImplementedError( @@ -1050,7 +1106,7 @@ def apply( ) if injection_sample_rate is not None: - delta_t = 1./injection_sample_rate + delta_t = 1.0 / injection_sample_rate else: delta_t = strain.delta_t injections = self.table @@ -1066,7 +1122,7 @@ def apply( self.set_ref_time(inj, ts) # determine if we inject or not based on the times try: - injtime = inj['{}_gps_time'.format(detector_name).lower()] + injtime = inj[f"{detector_name}_gps_time".lower()] except ValueError: injtime = -np.inf if np.isnan(injtime): @@ -1074,43 +1130,41 @@ def apply( injtime = -np.inf start_time = injtime + ts.start_time end_time = injtime + ts.end_time - inject = (start_time < strain.end_time and - end_time > strain.start_time) + inject = start_time < strain.end_time and end_time > strain.start_time if inject: ts = self.make_strain_from_inj_object( - inj, delta_t, detector_name, - distance_scale=distance_scale, ts=ts) + inj, delta_t, detector_name, distance_scale=distance_scale, ts=ts + ) if ts.delta_t != strain.delta_t: - ts = resample_to_delta_t(ts, strain.delta_t, method='ldas') + ts = resample_to_delta_t(ts, strain.delta_t, method="ldas") strain.inject(ts, copy=False) - def make_strain_from_inj_object(self, inj, delta_t, detector_name, - distance_scale=1, ts=None): + def make_strain_from_inj_object( + self, inj, delta_t, detector_name, distance_scale=1, ts=None + ): if ts is None: ts = load_timeseries(inj.filename) self.set_ref_time(inj, ts) # slice and taper ts = self.slice_and_taper(inj, ts) # shift reference to the detector time - ts._epoch += inj['{}_gps_time'.format(detector_name).lower()] + ts._epoch += inj[f"{detector_name}_gps_time".lower()] # resample - ts = resample_to_delta_t(ts, delta_t, method='ldas') + ts = resample_to_delta_t(ts, delta_t, method="ldas") # apply any phase shift try: - phase_shift = inj[ - '{}_phase_shift'.format(detector_name).lower()] + phase_shift = inj[f"{detector_name}_phase_shift".lower()] except ValueError: phase_shift = 0 if phase_shift: fs = ts.to_frequencyseries() - fs *= np.exp(1j*phase_shift) + fs *= np.exp(1j * phase_shift) ts = fs.to_timeseries() # apply any scaling try: - amp_scale = inj[ - '{}_amp_scale'.format(detector_name).lower()] + amp_scale = inj[f"{detector_name}_amp_scale".lower()] except ValueError: - amp_scale = 1. + amp_scale = 1.0 amp_scale /= distance_scale ts *= amp_scale return ts @@ -1119,13 +1173,13 @@ def make_strain_from_inj_object(self, inj, delta_t, detector_name, hdfinjtypes = { CBCHDFInjectionSet.injtype: CBCHDFInjectionSet, RingdownHDFInjectionSet.injtype: RingdownHDFInjectionSet, - IncoherentFromFileHDFInjectionSet.injtype: - IncoherentFromFileHDFInjectionSet, + IncoherentFromFileHDFInjectionSet.injtype: IncoherentFromFileHDFInjectionSet, } def get_hdf_injtype(sim_file): - """Gets the HDFInjectionSet class to use with the given file. + """ + Gets the HDFInjectionSet class to use with the given file. This looks for the ``injtype`` in the given file's top level ``attrs``. If that attribute isn't set, will default to :py:class:`CBCHDFInjectionSet`. @@ -1139,10 +1193,11 @@ def get_hdf_injtype(sim_file): ------- HDFInjectionSet : The type of HDFInjectionSet to use. + """ - with pycbc.io.HFile(sim_file, 'r') as fp: + with pycbc.io.HFile(sim_file, "r") as fp: try: - ftype = fp.attrs['injtype'] + ftype = fp.attrs["injtype"] except KeyError: ftype = CBCHDFInjectionSet.injtype try: @@ -1159,7 +1214,8 @@ def get_hdf_injtype(sim_file): def hdf_injtype_from_approximant(approximant): - """Gets the HDFInjectionSet class to use with the given approximant. + """ + Gets the HDFInjectionSet class to use with the given approximant. Parameters ---------- @@ -1170,6 +1226,7 @@ def hdf_injtype_from_approximant(approximant): ------- HDFInjectionSet : The type of HDFInjectionSet to use. + """ retcls = None for cls in hdfinjtypes.values(): @@ -1177,13 +1234,15 @@ def hdf_injtype_from_approximant(approximant): retcls = cls if retcls is None: # none were found, raise an error - raise ValueError("Injection file type unknown for approximant {}" - .format(approximant)) + raise ValueError( + f"Injection file type unknown for approximant {approximant}" + ) return retcls -class InjectionSet(object): - r"""Manages sets of injections and injects them into time series. +class InjectionSet: + r""" + Manages sets of injections and injects them into time series. Injections are read from either LIGOLW XML files or HDF files. @@ -1199,11 +1258,12 @@ class InjectionSet(object): Attributes ---------- table + """ def __init__(self, sim_file, **kwds): ext = os.path.basename(sim_file) - if ext.endswith(('.xml', '.xml.gz', '.xmlgz')): + if ext.endswith((".xml", ".xml.gz", ".xmlgz")): self._injhandler = _XMLInjectionSet(sim_file, **kwds) self.indoc = self._injhandler.indoc else: @@ -1212,14 +1272,15 @@ def __init__(self, sim_file, **kwds): self.table = self._injhandler.table self.extra_args = self._injhandler.extra_args self.apply = self._injhandler.apply - self.make_strain_from_inj_object = \ - self._injhandler.make_strain_from_inj_object + self.make_strain_from_inj_object = self._injhandler.make_strain_from_inj_object self.end_times = self._injhandler.end_times @staticmethod - def write(filename, samples, write_params=None, static_args=None, - injtype=None, **metadata): - r"""Writes the injection samples to the given hdf file. + def write( + filename, samples, write_params=None, static_args=None, injtype=None, **metadata + ): + r""" + Writes the injection samples to the given hdf file. Parameters ---------- @@ -1239,39 +1300,39 @@ def write(filename, samples, write_params=None, static_args=None, the ``static_args``, followed by the ``samples``. \**metadata : All other keyword arguments will be written to the file's attrs. + """ # DELETE the following "if" once xml is dropped ext = os.path.basename(filename) - if ext.endswith(('.xml', '.xml.gz', '.xmlgz')): - _XMLInjectionSet.write(filename, samples, write_params, - static_args) + if ext.endswith((".xml", ".xml.gz", ".xmlgz")): + _XMLInjectionSet.write(filename, samples, write_params, static_args) else: # try determine the injtype if it isn't given if injtype is None: - if static_args is not None and 'approximant' in static_args: - injcls = hdf_injtype_from_approximant( - static_args['approximant']) - elif 'approximant' in samples.fieldnames: - apprxs = np.unique(samples['approximant']) + if static_args is not None and "approximant" in static_args: + injcls = hdf_injtype_from_approximant(static_args["approximant"]) + elif "approximant" in samples.fieldnames: + apprxs = np.unique(samples["approximant"]) # make sure they all correspond to the same injection type injcls = [hdf_injtype_from_approximant(a) for a in apprxs] if not all(c == injcls[0] for c in injcls): - raise ValueError("injections must all be of the same " - "type") + raise ValueError("injections must all be of the same type") injcls = injcls[0] else: - raise ValueError("Could not find an approximant in the " - "static args or samples to determine the " - "injection type. Please specify an " - "injtype instead.") + raise ValueError( + "Could not find an approximant in the " + "static args or samples to determine the " + "injection type. Please specify an " + "injtype instead." + ) else: injcls = hdfinjtypes[injtype] - injcls.write(filename, samples, write_params, static_args, - **metadata) + injcls.write(filename, samples, write_params, static_args, **metadata) @staticmethod def from_cli(opt): - """Return an instance of InjectionSet configured as specified + """ + Return an instance of InjectionSet configured as specified on the command line. """ if opt.injection_file is None: @@ -1279,14 +1340,15 @@ def from_cli(opt): kwa = {} if opt.injection_f_ref is not None: - kwa['f_ref'] = opt.injection_f_ref + kwa["f_ref"] = opt.injection_f_ref if opt.injection_f_final is not None: - kwa['f_final'] = opt.injection_f_final + kwa["f_final"] = opt.injection_f_final return InjectionSet(opt.injection_file, **kwa) -class SGBurstInjectionSet(object): - """Manages sets of sine-Gaussian burst injections: reads injections +class SGBurstInjectionSet: + """ + Manages sets of sine-Gaussian burst injections: reads injections from LIGOLW XML files and injects them into time series. Parameters @@ -1299,16 +1361,19 @@ class SGBurstInjectionSet(object): ---------- indoc table + """ def __init__(self, sim_file, **kwds): self.indoc = ligolw_utils.load_filename( - sim_file, False, contenthandler=LIGOLWContentHandler) + sim_file, False, contenthandler=LIGOLWContentHandler + ) self.table = lsctables.SimBurstTable.get_table(self.indoc) self.extra_args = kwds def apply(self, strain, detector_name, f_lower=None, distance_scale=1): - """Add injections (as seen by a particular detector) to a time series. + """ + Add injections (as seen by a particular detector) to a time series. Parameters ---------- @@ -1331,13 +1396,15 @@ def apply(self, strain, detector_name, f_lower=None, distance_scale=1): ------ TypeError For invalid types of `strain`. + """ if strain.dtype not in (float32, float64): - raise TypeError("Strain dtype must be float32 or float64, not " \ - + str(strain.dtype)) + raise TypeError( + "Strain dtype must be float32 or float64, not " + str(strain.dtype) + ) lalstrain = strain.lal() - #detector = Detector(detector_name) + # detector = Detector(detector_name) earth_travel_time = lal.REARTH_SI / lal.C_SI t0 = float(strain.start_time) - earth_travel_time t1 = float(strain.end_time) + earth_travel_time @@ -1348,7 +1415,7 @@ def apply(self, strain, detector_name, f_lower=None, distance_scale=1): for inj in self.table: # roughly estimate if the injection may overlap with the segment end_time = inj.time_geocent - #CHECK: This is a hack (10.0s); replace with an accurate estimate + # CHECK: This is a hack (10.0s); replace with an accurate estimate inj_length = 10.0 eccentricity = 0.0 polarization = 0.0 @@ -1357,9 +1424,14 @@ def apply(self, strain, detector_name, f_lower=None, distance_scale=1): continue # compute the waveform time series - hp, hc = sim.SimBurstSineGaussian(float(inj.q), - float(inj.frequency),float(inj.hrss),float(eccentricity), - float(polarization),float(strain.delta_t)) + hp, hc = sim.SimBurstSineGaussian( + float(inj.q), + float(inj.frequency), + float(inj.hrss), + float(eccentricity), + float(polarization), + float(strain.delta_t), + ) hp = TimeSeries(hp.data.data[:], delta_t=hp.deltaT, epoch=hp.epoch) hc = TimeSeries(hc.data.data[:], delta_t=hc.deltaT, epoch=hc.epoch) hp._epoch += float(end_time) @@ -1369,9 +1441,11 @@ def apply(self, strain, detector_name, f_lower=None, distance_scale=1): # compute the detector response, taper it if requested # and add it to the strain - strain = strain.taper_timeseries(location=inj.taper, - tapermethod=inj.get('taper_method', 'lal'), - taper_window=inj.get('taper_window')) + strain = strain.taper_timeseries( + location=inj.taper, + tapermethod=inj.get("taper_method", "lal"), + taper_window=inj.get("taper_window"), + ) signal_lal = hp.astype(strain.dtype).lal() add_injection(lalstrain, signal_lal, None) diff --git a/pycbc/inject/injfilterrejector.py b/pycbc/inject/injfilterrejector.py index 7c2a9163da1..dd8760dd8c3 100644 --- a/pycbc/inject/injfilterrejector.py +++ b/pycbc/inject/injfilterrejector.py @@ -27,15 +27,12 @@ """ import numpy as np -from igwn_segments import segment -from igwn_segments import segmentlist +from igwn_segments import segment, segmentlist + from pycbc import DYN_RANGE_FAC from pycbc.filter import match -from pycbc.pnutils import nearest_larger_binary_number -from pycbc.pnutils import mass1_mass2_to_tau0_tau3 -from pycbc.types import FrequencySeries, zeros -from pycbc.types import MultiDetOptionAction -from pycbc.types import positive_float +from pycbc.pnutils import mass1_mass2_to_tau0_tau3, nearest_larger_binary_number +from pycbc.types import FrequencySeries, MultiDetOptionAction, positive_float, zeros _injfilterrejector_group_help = ( "Options that, if injections are present in " @@ -102,86 +99,141 @@ def insert_injfilterrejector_option_group(parser): """Add options for injfilterrejector to executable.""" - injfilterrejector_group = \ - parser.add_argument_group(_injfilterrejector_group_help) + injfilterrejector_group = parser.add_argument_group(_injfilterrejector_group_help) curr_arg = "--injection-filter-rejector-chirp-time-window" - injfilterrejector_group.add_argument(curr_arg, type=float, default=None, - help=_injfilterer_cthresh_help) + injfilterrejector_group.add_argument( + curr_arg, type=float, default=None, help=_injfilterer_cthresh_help + ) curr_arg = "--injection-filter-rejector-match-threshold" - injfilterrejector_group.add_argument(curr_arg, type=float, default=None, - help=_injfilterer_mthresh_help) + injfilterrejector_group.add_argument( + curr_arg, type=float, default=None, help=_injfilterer_mthresh_help + ) curr_arg = "--injection-filter-rejector-coarsematch-deltaf" - injfilterrejector_group.add_argument(curr_arg, type=float, default=1., - help=_injfilterer_deltaf_help) + injfilterrejector_group.add_argument( + curr_arg, type=float, default=1.0, help=_injfilterer_deltaf_help + ) curr_arg = "--injection-filter-rejector-coarsematch-fmax" - injfilterrejector_group.add_argument(curr_arg, type=float, default=256., - help=_injfilterer_fmax_help) + injfilterrejector_group.add_argument( + curr_arg, type=float, default=256.0, help=_injfilterer_fmax_help + ) curr_arg = "--injection-filter-rejector-seg-buffer" - injfilterrejector_group.add_argument(curr_arg, type=int, default=10, - help=_injfilterer_buffer_help) + injfilterrejector_group.add_argument( + curr_arg, type=int, default=10, help=_injfilterer_buffer_help + ) curr_arg = "--injection-filter-rejector-f-lower" - injfilterrejector_group.add_argument(curr_arg, type=int, default=None, - help=_injfilterer_flower_help) + injfilterrejector_group.add_argument( + curr_arg, type=int, default=None, help=_injfilterer_flower_help + ) curr_arg = "--injection-filter-rejector-trigger-window" - injfilterrejector_group.add_argument(curr_arg, type=positive_float, - default=None, - help=_injfilterer_trwindow_help) + injfilterrejector_group.add_argument( + curr_arg, type=positive_float, default=None, help=_injfilterer_trwindow_help + ) def insert_injfilterrejector_option_group_multi_ifo(parser): """Add options for injfilterrejector to executable.""" - injfilterrejector_group = \ - parser.add_argument_group(_injfilterrejector_group_help) + injfilterrejector_group = parser.add_argument_group(_injfilterrejector_group_help) curr_arg = "--injection-filter-rejector-chirp-time-window" injfilterrejector_group.add_argument( - curr_arg, type=float, default=None, nargs='+', metavar='IFO:VALUE', - action=MultiDetOptionAction, help=_injfilterer_cthresh_help) + curr_arg, + type=float, + default=None, + nargs="+", + metavar="IFO:VALUE", + action=MultiDetOptionAction, + help=_injfilterer_cthresh_help, + ) curr_arg = "--injection-filter-rejector-match-threshold" injfilterrejector_group.add_argument( - curr_arg, type=float, default=None, nargs='+', metavar='IFO:VALUE', - action=MultiDetOptionAction, help=_injfilterer_mthresh_help) + curr_arg, + type=float, + default=None, + nargs="+", + metavar="IFO:VALUE", + action=MultiDetOptionAction, + help=_injfilterer_mthresh_help, + ) curr_arg = "--injection-filter-rejector-coarsematch-deltaf" injfilterrejector_group.add_argument( - curr_arg, type=float, default=1., nargs='+', metavar='IFO:VALUE', - action=MultiDetOptionAction, help=_injfilterer_deltaf_help) + curr_arg, + type=float, + default=1.0, + nargs="+", + metavar="IFO:VALUE", + action=MultiDetOptionAction, + help=_injfilterer_deltaf_help, + ) curr_arg = "--injection-filter-rejector-coarsematch-fmax" injfilterrejector_group.add_argument( - curr_arg, type=float, default=256., nargs='+', metavar='IFO:VALUE', - action=MultiDetOptionAction, help=_injfilterer_fmax_help) + curr_arg, + type=float, + default=256.0, + nargs="+", + metavar="IFO:VALUE", + action=MultiDetOptionAction, + help=_injfilterer_fmax_help, + ) curr_arg = "--injection-filter-rejector-seg-buffer" injfilterrejector_group.add_argument( - curr_arg, type=int, default=10, nargs='+', metavar='IFO:VALUE', - action=MultiDetOptionAction, help=_injfilterer_buffer_help) + curr_arg, + type=int, + default=10, + nargs="+", + metavar="IFO:VALUE", + action=MultiDetOptionAction, + help=_injfilterer_buffer_help, + ) curr_arg = "--injection-filter-rejector-f-lower" injfilterrejector_group.add_argument( - curr_arg, type=int, default=None, help=_injfilterer_flower_help, - metavar='IFO:VALUE', action=MultiDetOptionAction, nargs='+') + curr_arg, + type=int, + default=None, + help=_injfilterer_flower_help, + metavar="IFO:VALUE", + action=MultiDetOptionAction, + nargs="+", + ) curr_arg = "--injection-filter-rejector-trigger-window" injfilterrejector_group.add_argument( - curr_arg, type=positive_float, default=None, + curr_arg, + type=positive_float, + default=None, help=_injfilterer_trwindow_help, - metavar='IFO:VALUE', action=MultiDetOptionAction, nargs='+') + metavar="IFO:VALUE", + action=MultiDetOptionAction, + nargs="+", + ) -class InjFilterRejector(object): - """Class for holding parameters for using injection/template pre-filtering. +class InjFilterRejector: + """ + Class for holding parameters for using injection/template pre-filtering. This class is responsible for identifying where a matched-filter operation between templates and data is unncessary because the injections contained in the data will not match well with the given template. """ - def __init__(self, injection_file, chirp_time_window, - match_threshold, f_lower, coarsematch_deltaf=1., - coarsematch_fmax=256, seg_buffer=10, inj_trigger_window=None): + def __init__( + self, + injection_file, + chirp_time_window, + match_threshold, + f_lower, + coarsematch_deltaf=1.0, + coarsematch_fmax=256, + seg_buffer=10, + inj_trigger_window=None, + ): """Initialise InjFilterRejector instance.""" # Determine if InjFilterRejector is to be enabled if ( - injection_file is None or injection_file == 'False' or - ( - chirp_time_window is None and - match_threshold is None and - inj_trigger_window is None + injection_file is None + or injection_file == "False" + or ( + chirp_time_window is None + and match_threshold is None + and inj_trigger_window is None ) ): self.enabled = False @@ -196,7 +248,7 @@ def __init__(self, injection_file, chirp_time_window, self.seg_buffer = seg_buffer self.f_lower = f_lower self.inj_trigger_window = inj_trigger_window - assert(self.f_lower is not None) + assert self.f_lower is not None # Variables for storing arrays (reduced injections, memory # for templates, reduced PSDs ...) @@ -211,8 +263,7 @@ def __init__(self, injection_file, chirp_time_window, def from_cli(cls, opt): """Create an InjFilterRejector instance from command-line options.""" injection_file = opt.injection_file - chirp_time_window = \ - opt.injection_filter_rejector_chirp_time_window + chirp_time_window = opt.injection_filter_rejector_chirp_time_window match_threshold = opt.injection_filter_rejector_match_threshold coarsematch_deltaf = opt.injection_filter_rejector_coarsematch_deltaf coarsematch_fmax = opt.injection_filter_rejector_coarsematch_fmax @@ -226,20 +277,24 @@ def from_cli(cls, opt): # leave for future work, or if this is being used in another # code which doesn't have --low-frequency-cutoff f_lower = opt.low_frequency_cutoff - return cls(injection_file, chirp_time_window, match_threshold, - f_lower, coarsematch_deltaf=coarsematch_deltaf, - coarsematch_fmax=coarsematch_fmax, - seg_buffer=seg_buffer, inj_trigger_window=trig_window) + return cls( + injection_file, + chirp_time_window, + match_threshold, + f_lower, + coarsematch_deltaf=coarsematch_deltaf, + coarsematch_fmax=coarsematch_fmax, + seg_buffer=seg_buffer, + inj_trigger_window=trig_window, + ) @classmethod def from_cli_single_ifo(cls, opt, ifo): """Create an InjFilterRejector instance from command-line options.""" injection_file = opt.injection_file[ifo] - chirp_time_window = \ - opt.injection_filter_rejector_chirp_time_window[ifo] + chirp_time_window = opt.injection_filter_rejector_chirp_time_window[ifo] match_threshold = opt.injection_filter_rejector_match_threshold[ifo] - coarsematch_deltaf = \ - opt.injection_filter_rejector_coarsematch_deltaf[ifo] + coarsematch_deltaf = opt.injection_filter_rejector_coarsematch_deltaf[ifo] coarsematch_fmax = opt.injection_filter_rejector_coarsematch_fmax[ifo] seg_buffer = opt.injection_filter_rejector_seg_buffer[ifo] trig_window = opt.injection_filter_rejector_trigger_window[ifo] @@ -251,10 +306,16 @@ def from_cli_single_ifo(cls, opt, ifo): # leave for future work, or if this is being used in another # code which doesn't have --low-frequency-cutoff f_lower = opt.low_frequency_cutoff - return cls(injection_file, chirp_time_window, - match_threshold, f_lower, - coarsematch_deltaf, coarsematch_fmax, - seg_buffer=seg_buffer, inj_trigger_window=trig_window) + return cls( + injection_file, + chirp_time_window, + match_threshold, + f_lower, + coarsematch_deltaf, + coarsematch_fmax, + seg_buffer=seg_buffer, + inj_trigger_window=trig_window, + ) @classmethod def from_cli_multi_ifos(cls, opt, ifos): @@ -286,7 +347,7 @@ def precompute_injection_intervals(self): # Create individual intervals starts = inj_times - window ends = inj_times + window - intervals = segmentlist([segment(s,e) for s,e in zip(starts, ends)]) + intervals = segmentlist([segment(s, e) for s, e in zip(starts, ends)]) intervals.coalesce() return np.array(intervals) @@ -298,7 +359,7 @@ def find_indices_in_injection_intervals(self, trig_times): the injection intervals. """ if not self.enabled or self.inj_trigger_window is None: - return slice(None) # Pythonic way to say "take all triggers" + return slice(None) # Pythonic way to say "take all triggers" if self._injection_intervals is None: self._injection_intervals = self.precompute_injection_intervals() @@ -340,7 +401,7 @@ def generate_short_inj_from_inj(self, inj_waveform, simulation_id): curr_length = len(inj_waveform) new_length = int(nearest_larger_binary_number(curr_length)) # Don't want length less than 1/delta_f - while new_length * inj_waveform.delta_t < 1./self.coarsematch_deltaf: + while new_length * inj_waveform.delta_t < 1.0 / self.coarsematch_deltaf: new_length = new_length * 2 inj_waveform.resize(new_length) inj_tilde = inj_waveform.to_frequencyseries() @@ -352,15 +413,17 @@ def generate_short_inj_from_inj(self, inj_waveform, simulation_id): # 16384 Hz ... It is only a problem of injection sample rate # gives a lower Nyquist than the trunc_f_max. If this error is # ever raised one could consider zero-padding the injection. - assert(new_freq_len <= len(inj_tilde)) - df_ratio = int(self.coarsematch_deltaf/delta_f) + assert new_freq_len <= len(inj_tilde) + df_ratio = int(self.coarsematch_deltaf / delta_f) inj_tilde_np = inj_tilde_np[:new_freq_len:df_ratio] - new_inj = FrequencySeries(inj_tilde_np, dtype=np.complex64, - delta_f=self.coarsematch_deltaf) + new_inj = FrequencySeries( + inj_tilde_np, dtype=np.complex64, delta_f=self.coarsematch_deltaf + ) self.short_injections[simulation_id] = new_inj def template_segment_checker(self, bank, t_num, segment): - """Test if injections in segment are worth filtering with template. + """ + Test if injections in segment are worth filtering with template. Using the current template, current segment, and injections within that segment. Test if the injections and sufficiently "similar" to any of @@ -376,12 +439,13 @@ def template_segment_checker(self, bank, t_num, segment): templates. Parameters - ----------- + ---------- FIXME Returns - -------- + ------- FIXME + """ if not self.enabled: # If disabled, always filter (ie. return True) @@ -393,23 +457,22 @@ def template_segment_checker(self, bank, t_num, segment): # Chirp time test if self.chirp_time_window is not None: - m1 = bank.table[t_num]['mass1'] - m2 = bank.table[t_num]['mass2'] + m1 = bank.table[t_num]["mass1"] + m2 = bank.table[t_num]["mass2"] tau0_temp, _ = mass1_mass2_to_tau0_tau3(m1, m2, self.f_lower) for inj in self.injection_params.table: if isinstance(inj, np.record): # hdf format file - end_time = inj['tc'] + end_time = inj["tc"] else: # must be an xml file originally - end_time = inj.geocent_end_time + \ - 1E-9 * inj.geocent_end_time_ns + end_time = inj.geocent_end_time + 1e-9 * inj.geocent_end_time_ns - if not(seg_start_time <= end_time <= seg_end_time): + if not (seg_start_time <= end_time <= seg_end_time): continue - tau0_inj, _ = \ - mass1_mass2_to_tau0_tau3(inj.mass1, inj.mass2, - self.f_lower) + tau0_inj, _ = mass1_mass2_to_tau0_tau3( + inj.mass1, inj.mass2, self.f_lower + ) tau_diff = abs(tau0_temp - tau0_inj) if tau_diff <= self.chirp_time_window: break @@ -421,8 +484,7 @@ def template_segment_checker(self, bank, t_num, segment): if self.match_threshold: if self._short_template_mem is None: # Set the memory for the short templates - wav_len = 1 + int(self.coarsematch_fmax / - self.coarsematch_deltaf) + wav_len = 1 + int(self.coarsematch_fmax / self.coarsematch_deltaf) self._short_template_mem = zeros(wav_len, dtype=np.complex64) # Set the current short PSD to red_psd @@ -434,23 +496,26 @@ def template_segment_checker(self, bank, t_num, segment): step_size = int(self.coarsematch_deltaf / segment.psd.delta_f) max_idx = int(self.coarsematch_fmax / segment.psd.delta_f) + 1 red_psd_data = curr_psd[:max_idx:step_size] - red_psd = FrequencySeries(red_psd_data, #copy=False, - delta_f=self.coarsematch_deltaf) + red_psd = FrequencySeries( + red_psd_data, # copy=False, + delta_f=self.coarsematch_deltaf, + ) self._short_psd_storage[id(curr_psd)] = red_psd # Set htilde to be the current short template if not t_num == self._short_template_id: # Set the memory for the short templates if unset if self._short_template_mem is None: - wav_len = 1 + int(self.coarsematch_fmax / - self.coarsematch_deltaf) - self._short_template_mem = zeros(wav_len, - dtype=np.complex64) + wav_len = 1 + int(self.coarsematch_fmax / self.coarsematch_deltaf) + self._short_template_mem = zeros(wav_len, dtype=np.complex64) # Generate short waveform htilde = bank.generate_with_delta_f_and_max_freq( - t_num, self.coarsematch_fmax, self.coarsematch_deltaf, + t_num, + self.coarsematch_fmax, + self.coarsematch_deltaf, low_frequency_cutoff=bank.table[t_num].f_lower, - cached_mem=self._short_template_mem) + cached_mem=self._short_template_mem, + ) self._short_template_id = t_num self._short_template_wav = htilde else: @@ -459,19 +524,19 @@ def template_segment_checker(self, bank, t_num, segment): for ii, inj in enumerate(self.injection_params.table): if isinstance(inj, np.record): # hdf format file - end_time = inj['tc'] + end_time = inj["tc"] sim_id = self.injection_ids[ii] else: # must be an xml file originally - end_time = inj.geocent_end_time + \ - 1E-9 * inj.geocent_end_time_ns + end_time = inj.geocent_end_time + 1e-9 * inj.geocent_end_time_ns sim_id = inj.simulation_id - if not(seg_start_time < end_time < seg_end_time): + if not (seg_start_time < end_time < seg_end_time): continue curr_inj = self.short_injections[sim_id] - o, _ = match(htilde, curr_inj, psd=red_psd, - low_frequency_cutoff=self.f_lower) + o, _ = match( + htilde, curr_inj, psd=red_psd, low_frequency_cutoff=self.f_lower + ) if o > self.match_threshold: break else: diff --git a/pycbc/io/__init__.py b/pycbc/io/__init__.py index 81ba4884156..283718f569e 100644 --- a/pycbc/io/__init__.py +++ b/pycbc/io/__init__.py @@ -1,21 +1,26 @@ -import os -import logging -from astropy.utils.data import download_file import hashlib +import logging +import os from urllib.parse import urlparse + +from astropy.utils.data import download_file + +from .gracedb import * from .hdf import * from .record import * -from .gracedb import * -logger = logging.getLogger('pycbc.io') +logger = logging.getLogger("pycbc.io") # Backup URL in case GWOSC fails base_backup_url = "https://raw.githubusercontent.com/gwastro/pycbc_data/master/{}" -base_lfs_backup_url = "https://media.githubusercontent.com/media/gwastro/pycbc_data/master/{}" +base_lfs_backup_url = ( + "https://media.githubusercontent.com/media/gwastro/pycbc_data/master/{}" +) def get_file(url, retry=5, **args): - """ Retrieve file with retry upon failure + """ + Retrieve file with retry upon failure Uses the astropy download_file but adds a retry feature for flaky connections. See astropy for full options @@ -26,24 +31,20 @@ def get_file(url, retry=5, **args): # If this is in GitHub Actions we divert the URLs to a backup path if "gwosc.org/" in url: basename = os.path.basename(urlparse(url).path) - if basename.endswith('hdf5') or basename.endswith('gwf'): + if basename.endswith("hdf5") or basename.endswith("gwf"): # Just download file directly from backup new_url = base_lfs_backup_url.format(basename) logger.warning( - "Redirecting GWOSC URL %s to backup url %s", - url, - new_url + "Redirecting GWOSC URL %s to backup url %s", url, new_url ) url = new_url else: cleaned_url = url.strip().lower() - hash_object = hashlib.md5(cleaned_url.encode('utf-8')) + hash_object = hashlib.md5(cleaned_url.encode("utf-8")) hh = hash_object.hexdigest() - new_url = base_backup_url.format(hh + '.json') + new_url = base_backup_url.format(hh + ".json") logger.warning( - "Redirecting GWOSC URL %s to backup url %s", - url, - new_url + "Redirecting GWOSC URL %s to backup url %s", url, new_url ) url = new_url while True: diff --git a/pycbc/io/gracedb.py b/pycbc/io/gracedb.py index 0e546a4ea37..3df3ff72856 100644 --- a/pycbc/io/gracedb.py +++ b/pycbc/io/gracedb.py @@ -2,32 +2,30 @@ Class and function for use in dealing with GraceDB uploads """ +import copy +import json import logging import os -import numpy -import json -import copy from multiprocessing.dummy import threading import lal -from igwn_ligolw import ligolw -from igwn_ligolw import lsctables +import numpy +from igwn_ligolw import ligolw, lsctables from igwn_ligolw import utils as ligolw_utils import pycbc +from pycbc import constants, pnutils from pycbc import version as pycbc_version -from pycbc import pnutils, constants from pycbc.io.ligolw import ( - return_empty_sngl, create_process_table, make_psd_xmldoc, - snr_series_to_xml + return_empty_sngl, + snr_series_to_xml, ) -from pycbc.results import generate_asd_plot, generate_snr_plot -from pycbc.results import source_color from pycbc.mchirp_area import calc_probabilities +from pycbc.results import generate_asd_plot, generate_snr_plot, source_color -logger = logging.getLogger('pycbc.io.gracedb') +logger = logging.getLogger("pycbc.io.gracedb") def _single_value(value): @@ -35,12 +33,12 @@ def _single_value(value): return numpy.asarray(value).item() -class CandidateForGraceDB(object): - """This class provides an interface for uploading candidates to GraceDB. - """ +class CandidateForGraceDB: + """This class provides an interface for uploading candidates to GraceDB.""" def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): - """Initialize a representation of a zerolag candidate for upload to + """ + Initialize a representation of a zerolag candidate for upload to GraceDB. Parameters @@ -76,33 +74,33 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): mc_area_args: dict of dicts, optional Dictionary providing arguments to be used in source probability estimation with `pycbc/mchirp_area.py`. + """ self.coinc_results = coinc_results - self.psds = kwargs['psds'] + self.psds = kwargs["psds"] self.basename = None - if kwargs.get('gracedb'): - self.gracedb = kwargs['gracedb'] + if kwargs.get("gracedb"): + self.gracedb = kwargs["gracedb"] # Determine if the candidate should be marked as HWINJ - self.is_hardware_injection = ('HWINJ' in coinc_results - and coinc_results['HWINJ']) + self.is_hardware_injection = "HWINJ" in coinc_results and coinc_results["HWINJ"] # We may need to apply a time offset for premerger search self.time_offset = 0 - rtoff = f'foreground/{ifos[0]}/time_offset' + rtoff = f"foreground/{ifos[0]}/time_offset" if rtoff in coinc_results: self.time_offset = coinc_results[rtoff] # Check for ifos with SNR peaks in coinc_results - self.et_ifos = [i for i in ifos if f'foreground/{i}/end_time' in - coinc_results] + self.et_ifos = [i for i in ifos if f"foreground/{i}/end_time" in coinc_results] - if 'skyloc_data' in kwargs: - sld = kwargs['skyloc_data'] - assert len({sld[ifo]['snr_series'].delta_t for ifo in sld}) == 1, \ - "delta_t for all ifos do not match" + if "skyloc_data" in kwargs: + sld = kwargs["skyloc_data"] + assert len({sld[ifo]["snr_series"].delta_t for ifo in sld}) == 1, ( + "delta_t for all ifos do not match" + ) snr_ifos = sld.keys() # Ifos with SNR time series calculated - self.snr_series = {ifo: sld[ifo]['snr_series'] for ifo in snr_ifos} + self.snr_series = {ifo: sld[ifo]["snr_series"] for ifo in snr_ifos} # Extra ifos have SNR time series but not sngl inspiral triggers for ifo in snr_ifos: @@ -117,8 +115,9 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): outdoc = ligolw.Document() outdoc.appendChild(ligolw.LIGO_LW()) - proc_id = create_process_table(outdoc, program_name='pycbc', - detectors=snr_ifos).process_id + proc_id = create_process_table( + outdoc, program_name="pycbc", detectors=snr_ifos + ).process_id # Set up coinc_definer table coinc_def_table = lsctables.CoincDefTable.new() @@ -137,15 +136,14 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): coinc_event_row = lsctables.Coinc() coinc_event_row.coinc_def_id = coinc_def_id coinc_event_row.nevents = len(snr_ifos) - coinc_event_row.instruments = ','.join(snr_ifos) + coinc_event_row.instruments = ",".join(snr_ifos) coinc_event_row.time_slide_id = lsctables.TimeSlideID(0) coinc_event_row.process_id = proc_id coinc_event_row.coinc_event_id = coinc_id - if 'foreground/stat' in coinc_results: - coinc_event_row.likelihood = _single_value( - coinc_results['foreground/stat']) + if "foreground/stat" in coinc_results: + coinc_event_row.likelihood = _single_value(coinc_results["foreground/stat"]) else: - coinc_event_row.likelihood = 0. + coinc_event_row.likelihood = 0.0 coinc_event_table.append(coinc_event_row) outdoc.childNodes[0].appendChild(coinc_event_table) @@ -161,11 +159,12 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): sngl.event_id = lsctables.SnglInspiralID(sngl_id) sngl.process_id = proc_id sngl.ifo = ifo - names = [n.split('/')[-1] for n in coinc_results - if f'foreground/{ifo}' in n] + names = [ + n.split("/")[-1] for n in coinc_results if f"foreground/{ifo}" in n + ] for name in names: - val = coinc_results[f'foreground/{ifo}/{name}'] - if name == 'end_time': + val = coinc_results[f"foreground/{ifo}/{name}"] + if name == "end_time": val += self.time_offset sngl.end = lal.LIGOTimeGPS(val) else: @@ -176,20 +175,22 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): pass if sngl.mass1 and sngl.mass2: sngl.mtotal, sngl.eta = pnutils.mass1_mass2_to_mtotal_eta( - sngl.mass1, sngl.mass2) + sngl.mass1, sngl.mass2 + ) sngl.mchirp, _ = pnutils.mass1_mass2_to_mchirp_eta( - sngl.mass1, sngl.mass2) + sngl.mass1, sngl.mass2 + ) sngl_populated = sngl if sngl.snr: - sngl.eff_distance = sngl.sigmasq ** 0.5 / sngl.snr - network_snrsq += sngl.snr ** 2.0 - if 'channel_names' in kwargs and ifo in kwargs['channel_names']: - sngl.channel = kwargs['channel_names'][ifo] + sngl.eff_distance = sngl.sigmasq**0.5 / sngl.snr + network_snrsq += sngl.snr**2.0 + if "channel_names" in kwargs and ifo in kwargs["channel_names"]: + sngl.channel = kwargs["channel_names"][ifo] sngl_inspiral_table.append(sngl) # Set up coinc_map entry coinc_map_row = lsctables.CoincMap() - coinc_map_row.table_name = 'sngl_inspiral' + coinc_map_row.table_name = "sngl_inspiral" coinc_map_row.coinc_event_id = coinc_id coinc_map_row.event_id = sngl.event_id coinc_event_map_table.append(coinc_map_row) @@ -198,10 +199,12 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): snr_series_to_xml(self.snr_series[ifo], outdoc, sngl.event_id) # Set merger time to the mean of trigger peaks over coinc_results ifos - self.merger_time = \ - numpy.mean([coinc_results[f'foreground/{ifo}/end_time'] for ifo in - self.et_ifos]) \ + self.merger_time = ( + numpy.mean( + [coinc_results[f"foreground/{ifo}/end_time"] for ifo in self.et_ifos] + ) + self.time_offset + ) outdoc.childNodes[0].appendChild(coinc_event_map_table) outdoc.childNodes[0].appendChild(sngl_inspiral_table) @@ -210,16 +213,16 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): coinc_inspiral_table = lsctables.CoincInspiralTable.new() coinc_inspiral_row = lsctables.CoincInspiral() # This seems to be used as FAP, which should not be in gracedb - coinc_inspiral_row.false_alarm_rate = 0. - coinc_inspiral_row.minimum_duration = 0. + coinc_inspiral_row.false_alarm_rate = 0.0 + coinc_inspiral_row.minimum_duration = 0.0 coinc_inspiral_row.instruments = tuple(snr_ifos) coinc_inspiral_row.coinc_event_id = coinc_id coinc_inspiral_row.mchirp = sngl_populated.mchirp coinc_inspiral_row.mass = sngl_populated.mtotal coinc_inspiral_row.end_time = sngl_populated.end_time coinc_inspiral_row.end_time_ns = sngl_populated.end_time_ns - coinc_inspiral_row.snr = network_snrsq ** 0.5 - far = 1.0 / (constants.YRJUL_SI * coinc_results['foreground/ifar']) + coinc_inspiral_row.snr = network_snrsq**0.5 + far = 1.0 / (constants.YRJUL_SI * coinc_results["foreground/ifar"]) coinc_inspiral_row.combined_far = far coinc_inspiral_table.append(coinc_inspiral_row) outdoc.childNodes[0].appendChild(coinc_inspiral_table) @@ -227,74 +230,81 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): # Append the PSDs psds_lal = {} for ifo, psd in self.psds.items(): - kmin = int(kwargs['low_frequency_cutoff'] / psd.delta_f) + kmin = int(kwargs["low_frequency_cutoff"] / psd.delta_f) fseries = lal.CreateREAL8FrequencySeries( - "psd", psd.epoch, kwargs['low_frequency_cutoff'], psd.delta_f, - lal.StrainUnit**2 / lal.HertzUnit, len(psd) - kmin) + "psd", + psd.epoch, + kwargs["low_frequency_cutoff"], + psd.delta_f, + lal.StrainUnit**2 / lal.HertzUnit, + len(psd) - kmin, + ) fseries.data.data = ( - psd.numpy()[kmin:].astype(numpy.float64) - / pycbc.DYN_RANGE_FAC ** 2.0 + psd.numpy()[kmin:].astype(numpy.float64) / pycbc.DYN_RANGE_FAC**2.0 ) psds_lal[ifo] = fseries make_psd_xmldoc(psds_lal, outdoc) # P astro calculation - if 'padata' in kwargs: - if 'p_terr' in kwargs: + if "padata" in kwargs: + if "p_terr" in kwargs: raise RuntimeError( "Both p_astro calculation data and a " "previously calculated p_terr value were provided, this " "doesn't make sense!" ) - assert len(coinc_ifos) < 3, \ - f"p_astro can't handle {coinc_ifos} coinc ifos!" + assert len(coinc_ifos) < 3, f"p_astro can't handle {coinc_ifos} coinc ifos!" trigger_data = { - 'mass1': sngl_populated.mass1, - 'mass2': sngl_populated.mass2, - 'spin1z': sngl_populated.spin1z, - 'spin2z': sngl_populated.spin2z, - 'network_snr': network_snrsq ** 0.5, - 'far': far, - 'triggered': coinc_ifos, + "mass1": sngl_populated.mass1, + "mass2": sngl_populated.mass2, + "spin1z": sngl_populated.spin1z, + "spin2z": sngl_populated.spin2z, + "network_snr": network_snrsq**0.5, + "far": far, + "triggered": coinc_ifos, # Consider all ifos potentially relevant to detection, # ignore those that only contribute to sky loc - 'sensitive': self.et_ifos} + "sensitive": self.et_ifos, + } horizons = {i: self.psds[i].dist for i in self.et_ifos} - self.p_astro, self.p_terr = \ - kwargs['padata'].do_pastro_calc(trigger_data, horizons) - elif 'p_terr' in kwargs: - self.p_astro, self.p_terr = 1 - kwargs['p_terr'], kwargs['p_terr'] + self.p_astro, self.p_terr = kwargs["padata"].do_pastro_calc( + trigger_data, horizons + ) + elif "p_terr" in kwargs: + self.p_astro, self.p_terr = 1 - kwargs["p_terr"], kwargs["p_terr"] else: self.p_astro, self.p_terr = None, None # Source probabilities and hasmassgap estimation self.probabilities = None self.hasmassgap = None - if 'mc_area_args' in kwargs: + if "mc_area_args" in kwargs: eff_distances = [sngl.eff_distance for sngl in sngl_inspiral_table] self.probabilities = calc_probabilities( coinc_inspiral_row.mchirp, coinc_inspiral_row.snr, min(eff_distances), - kwargs['mc_area_args'] + kwargs["mc_area_args"], ) - if 'embright_mg_max' in kwargs['mc_area_args']: - hasmg_args = copy.deepcopy(kwargs['mc_area_args']) - hasmg_args['mass_gap'] = True - hasmg_args['mass_bdary']['gap_max'] = \ - kwargs['mc_area_args']['embright_mg_max'] + if "embright_mg_max" in kwargs["mc_area_args"]: + hasmg_args = copy.deepcopy(kwargs["mc_area_args"]) + hasmg_args["mass_gap"] = True + hasmg_args["mass_bdary"]["gap_max"] = kwargs["mc_area_args"][ + "embright_mg_max" + ] self.hasmassgap = calc_probabilities( coinc_inspiral_row.mchirp, coinc_inspiral_row.snr, min(eff_distances), - hasmg_args - )['Mass Gap'] + hasmg_args, + )["Mass Gap"] # Combine p astro and source probs if self.p_astro is not None and self.probabilities is not None: - self.astro_probs = {cl: pr * self.p_astro for - cl, pr in self.probabilities.items()} - self.astro_probs['Terrestrial'] = self.p_terr + self.astro_probs = { + cl: pr * self.p_astro for cl, pr in self.probabilities.items() + } + self.astro_probs["Terrestrial"] = self.p_terr else: self.astro_probs = None @@ -302,57 +312,65 @@ def __init__(self, coinc_ifos, ifos, coinc_results, **kwargs): self.time = sngl_populated.end def save(self, fname): - """Write a file representing this candidate in a LIGOLW XML format + """ + Write a file representing this candidate in a LIGOLW XML format compatible with GraceDB. Parameters ---------- fname: str Name of file to write to disk. + """ kwargs = {} if threading.current_thread() is not threading.main_thread(): # avoid an error due to no ability to do signal handling in threads - kwargs['trap_signals'] = None - ligolw_utils.write_filename(self.outdoc, fname, \ - compress='auto', **kwargs) + kwargs["trap_signals"] = None + ligolw_utils.write_filename(self.outdoc, fname, compress="auto", **kwargs) save_dir = os.path.dirname(fname) # Save EMBright properties info as json if self.hasmassgap is not None: - self.embright_file = os.path.join(save_dir, 'pycbc.em_bright.json') - with open(self.embright_file, 'w') as embrightf: - json.dump({'HasMassGap': self.hasmassgap}, embrightf) - logger.info('EM Bright file saved as %s', self.embright_file) + self.embright_file = os.path.join(save_dir, "pycbc.em_bright.json") + with open(self.embright_file, "w") as embrightf: + json.dump({"HasMassGap": self.hasmassgap}, embrightf) + logger.info("EM Bright file saved as %s", self.embright_file) # Save multi-cpt p astro as json if self.astro_probs is not None: - self.multipa_file = os.path.join(save_dir, 'pycbc.p_astro.json') - with open(self.multipa_file, 'w') as multipaf: + self.multipa_file = os.path.join(save_dir, "pycbc.p_astro.json") + with open(self.multipa_file, "w") as multipaf: json.dump(self.astro_probs, multipaf) - logger.info('Multi p_astro file saved as %s', self.multipa_file) + logger.info("Multi p_astro file saved as %s", self.multipa_file) # Save source probabilities in a json file if self.probabilities is not None: - self.prob_file = os.path.join(save_dir, 'src_probs.json') - with open(self.prob_file, 'w') as probf: + self.prob_file = os.path.join(save_dir, "src_probs.json") + with open(self.prob_file, "w") as probf: json.dump(self.probabilities, probf) - logger.info('Source probabilities file saved as %s', self.prob_file) + logger.info("Source probabilities file saved as %s", self.prob_file) # Don't save any other files! return # Save p astro / p terr as json if self.p_astro is not None: - self.pastro_file = os.path.join(save_dir, 'pa_pterr.json') - with open(self.pastro_file, 'w') as pastrof: - json.dump({'p_astro': self.p_astro, 'p_terr': self.p_terr}, - pastrof) - logger.info('P_astro file saved as %s', self.pastro_file) - - def upload(self, fname, gracedb_server=None, testing=True, - extra_strings=None, search='AllSky', labels=None, - **kwargs): - """Upload this candidate to GraceDB, and annotate it with a few useful + self.pastro_file = os.path.join(save_dir, "pa_pterr.json") + with open(self.pastro_file, "w") as pastrof: + json.dump({"p_astro": self.p_astro, "p_terr": self.p_terr}, pastrof) + logger.info("P_astro file saved as %s", self.pastro_file) + + def upload( + self, + fname, + gracedb_server=None, + testing=True, + extra_strings=None, + search="AllSky", + labels=None, + **kwargs, + ): + """ + Upload this candidate to GraceDB, and annotate it with a few useful plots and comments. Parameters @@ -372,16 +390,18 @@ def upload(self, fname, gracedb_server=None, testing=True, kwargs: named keyword arguments Extra keyword arguments to be passed to GraceDB upload function. Can also overwrite reload_certificate and reload_buffer here. + """ from matplotlib import pyplot as plt - if fname.endswith('.xml.gz'): - self.basename = fname.replace('.xml.gz', '') - elif fname.endswith('.xml'): - self.basename = fname.replace('.xml', '') + if fname.endswith(".xml.gz"): + self.basename = fname.replace(".xml.gz", "") + elif fname.endswith(".xml"): + self.basename = fname.replace(".xml", "") else: - raise ValueError("Upload filename must end in .xml or .xml.gz, got" - " %s" % fname) + raise ValueError( + "Upload filename must end in .xml or .xml.gz, got %s" % fname + ) # First make sure the event is saved on disk # as GraceDB operations can fail later @@ -389,181 +409,175 @@ def upload(self, fname, gracedb_server=None, testing=True, # hardware injections need to be maked with the INJ tag if self.is_hardware_injection: - labels = (labels or []) + ['INJ'] + labels = (labels or []) + ["INJ"] # connect to GraceDB if we are not reusing a connection - if not hasattr(self, 'gracedb'): - logger.info('Connecting to GraceDB') - gdbargs = {'reload_certificate': True, 'reload_buffer': 300} + if not hasattr(self, "gracedb"): + logger.info("Connecting to GraceDB") + gdbargs = {"reload_certificate": True, "reload_buffer": 300} if kwargs is not None: gdbargs.update(kwargs) if gracedb_server: - gdbargs['service_url'] = gracedb_server + gdbargs["service_url"] = gracedb_server try: from ligo.gracedb.rest import GraceDb + self.gracedb = GraceDb(**gdbargs) except Exception as exc: - logger.error('Failed to create GraceDB client') + logger.error("Failed to create GraceDB client") logger.error(exc) # create GraceDB event - logger.info('Uploading %s to GraceDB', fname) - group = 'Test' if testing else 'CBC' + logger.info("Uploading %s to GraceDB", fname) + group = "Test" if testing else "CBC" gid = None try: response = self.gracedb.create_event( - group, - "pycbc", - fname, - search=search, - labels=labels + group, "pycbc", fname, search=search, labels=labels ) gid = response.json()["graceid"] logger.info("Uploaded event %s", gid) except Exception as exc: - logger.error('Failed to create GraceDB event') + logger.error("Failed to create GraceDB event") logger.error(str(exc)) # Upload em_bright properties JSON if self.hasmassgap is not None and gid is not None: try: self.gracedb.write_log( - gid, 'EM Bright properties JSON file upload', + gid, + "EM Bright properties JSON file upload", filename=self.embright_file, - tag_name=['em_bright'] + tag_name=["em_bright"], ) - logger.info('Uploaded em_bright properties for %s', gid) + logger.info("Uploaded em_bright properties for %s", gid) except Exception as exc: - logger.error( - 'Failed to upload em_bright properties file ' - 'for %s', - gid - ) + logger.error("Failed to upload em_bright properties file for %s", gid) logger.error(str(exc)) # Upload multi-cpt p_astro JSON if self.astro_probs is not None and gid is not None: try: self.gracedb.write_log( - gid, 'Multi-component p_astro JSON file upload', + gid, + "Multi-component p_astro JSON file upload", filename=self.multipa_file, - tag_name=['p_astro'], - label='PASTRO_READY' + tag_name=["p_astro"], + label="PASTRO_READY", ) - logger.info('Uploaded multi p_astro for %s', gid) + logger.info("Uploaded multi p_astro for %s", gid) except Exception as exc: - logger.error( - 'Failed to upload multi p_astro file for %s', - gid - ) + logger.error("Failed to upload multi p_astro file for %s", gid) logger.error(str(exc)) # If there is p_astro but no probabilities, upload p_astro JSON - if hasattr(self, 'pastro_file') and gid is not None: + if hasattr(self, "pastro_file") and gid is not None: try: self.gracedb.write_log( - gid, '2-component p_astro JSON file upload', + gid, + "2-component p_astro JSON file upload", filename=self.pastro_file, - tag_name=['sig_info'] + tag_name=["sig_info"], ) - logger.info('Uploaded p_astro for %s', gid) + logger.info("Uploaded p_astro for %s", gid) except Exception as exc: - logger.error('Failed to upload p_astro file for %s', gid) + logger.error("Failed to upload p_astro file for %s", gid) logger.error(str(exc)) # plot the SNR timeseries and noise PSDs if self.snr_series is not None: - snr_series_fname = self.basename + '.hdf' - snr_series_plot_fname = self.basename + '_snr.png' - asd_series_plot_fname = self.basename + '_asd.png' + snr_series_fname = self.basename + ".hdf" + snr_series_plot_fname = self.basename + "_snr.png" + asd_series_plot_fname = self.basename + "_asd.png" triggers = { - ifo: (self.coinc_results[f'foreground/{ifo}/end_time'] - + self.time_offset, - self.coinc_results[f'foreground/{ifo}/snr']) + ifo: ( + self.coinc_results[f"foreground/{ifo}/end_time"] + self.time_offset, + self.coinc_results[f"foreground/{ifo}/snr"], + ) for ifo in self.et_ifos - } + } ref_time = int(self.merger_time) - generate_snr_plot(self.snr_series, snr_series_plot_fname, - triggers, ref_time) + generate_snr_plot( + self.snr_series, snr_series_plot_fname, triggers, ref_time + ) generate_asd_plot(self.psds, asd_series_plot_fname) for ifo in sorted(self.snr_series): curr_snrs = self.snr_series[ifo] - curr_snrs.save(snr_series_fname, group='%s/snr' % ifo) + curr_snrs.save(snr_series_fname, group="%s/snr" % ifo) # Additionally save the PSDs into the snr_series file for ifo in sorted(self.psds): # Undo dynamic range factor curr_psd = self.psds[ifo].astype(numpy.float64) - curr_psd /= pycbc.DYN_RANGE_FAC ** 2.0 - curr_psd.save(snr_series_fname, group='%s/psd' % ifo) + curr_psd /= pycbc.DYN_RANGE_FAC**2.0 + curr_psd.save(snr_series_fname, group="%s/psd" % ifo) # Upload SNR series in HDF format and plots if self.snr_series is not None and gid is not None: try: self.gracedb.write_log( - gid, 'SNR timeseries HDF file upload', - filename=snr_series_fname + gid, "SNR timeseries HDF file upload", filename=snr_series_fname ) self.gracedb.write_log( - gid, 'SNR timeseries plot upload', + gid, + "SNR timeseries plot upload", filename=snr_series_plot_fname, - tag_name=['background'], - displayName=['SNR timeseries'] + tag_name=["background"], + displayName=["SNR timeseries"], ) self.gracedb.write_log( - gid, 'ASD plot upload', + gid, + "ASD plot upload", filename=asd_series_plot_fname, - tag_name=['psd'], displayName=['ASDs'] + tag_name=["psd"], + displayName=["ASDs"], ) except Exception as exc: - logger.error( - 'Failed to upload SNR timeseries and ASD for %s', - gid - ) + logger.error("Failed to upload SNR timeseries and ASD for %s", gid) logger.error(str(exc)) # If 'self.prob_file' exists, make pie plot and do uploads. # The pie plot only shows relative astrophysical source # probabilities, not p_astro vs p_terrestrial - if hasattr(self, 'prob_file'): - self.prob_plotf = self.prob_file.replace('.json', '.png') + if hasattr(self, "prob_file"): + self.prob_plotf = self.prob_file.replace(".json", ".png") # Don't try to plot zero probabilities - prob_plot = {k: v for (k, v) in self.probabilities.items() - if v != 0.0} + prob_plot = {k: v for (k, v) in self.probabilities.items() if v != 0.0} labels, sizes = zip(*prob_plot.items()) colors = [source_color(label) for label in labels] fig, ax = plt.subplots() - ax.pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%', - textprops={'fontsize': 15}) - ax.axis('equal') + ax.pie( + sizes, + labels=labels, + colors=colors, + autopct="%1.1f%%", + textprops={"fontsize": 15}, + ) + ax.axis("equal") fig.savefig(self.prob_plotf) plt.close() if gid is not None: try: self.gracedb.write_log( gid, - 'Source probabilities JSON file upload', + "Source probabilities JSON file upload", filename=self.prob_file, - tag_name=['pe'] + tag_name=["pe"], ) - logger.info('Uploaded source probabilities for %s', gid) + logger.info("Uploaded source probabilities for %s", gid) self.gracedb.write_log( gid, - 'Source probabilities plot upload', + "Source probabilities plot upload", filename=self.prob_plotf, - tag_name=['pe'] - ) - logger.info( - 'Uploaded source probabilities pie chart for %s', - gid + tag_name=["pe"], ) + logger.info("Uploaded source probabilities pie chart for %s", gid) except Exception as exc: logger.error( - 'Failed to upload source probability results for %s', - gid + "Failed to upload source probability results for %s", gid ) logger.error(str(exc)) @@ -572,14 +586,13 @@ def upload(self, fname, gracedb_server=None, testing=True, # Add code version info gracedb_tag_with_version(self.gracedb, gid) # Add any annotations to the event log - for text in (extra_strings or []): - self.gracedb.write_log( - gid, text, tag_name=['analyst_comments']) + for text in extra_strings or []: + self.gracedb.write_log(gid, text, tag_name=["analyst_comments"]) except Exception as exc: logger.error( - 'Something failed during annotation of analyst ' - 'comments for event %s on GraceDB.', - fname + "Something failed during annotation of analyst " + "comments for event %s on GraceDB.", + fname, ) logger.error(str(exc)) @@ -587,17 +600,17 @@ def upload(self, fname, gracedb_server=None, testing=True, def gracedb_tag_with_version(gracedb, event_id): - """Add a GraceDB log entry reporting PyCBC's version and install location. - """ - version_str = 'Using PyCBC version {}{} at {}' + """Add a GraceDB log entry reporting PyCBC's version and install location.""" + version_str = "Using PyCBC version {}{} at {}" version_str = version_str.format( pycbc_version.version, - ' (release)' if pycbc_version.release else '', - os.path.dirname(pycbc.__file__) + " (release)" if pycbc_version.release else "", + os.path.dirname(pycbc.__file__), ) gracedb.write_log(event_id, version_str) __all__ = [ - 'CandidateForGraceDB', 'gracedb_tag_with_version', + "CandidateForGraceDB", + "gracedb_tag_with_version", ] diff --git a/pycbc/io/hdf.py b/pycbc/io/hdf.py index 4d4d048b6e7..b3c8da4b2a6 100644 --- a/pycbc/io/hdf.py +++ b/pycbc/io/hdf.py @@ -2,35 +2,31 @@ Convenience classes for accessing hdf5 trigger files """ -import h5py -import numpy as np -import logging import inspect +import logging import pickle - -from itertools import chain from io import BytesIO -from lal import LIGOTimeGPS +from itertools import chain -from igwn_ligolw import ligolw -from igwn_ligolw import lsctables +import h5py +import numpy as np +from igwn_ligolw import ligolw, lsctables from igwn_ligolw import utils as ligolw_utils +from lal import LIGOTimeGPS +from pycbc import conversions, events, pnutils +from pycbc.events import mean_if_greater_than_zero, ranking, veto from pycbc.io.ligolw import ( - return_search_summary, + create_process_table, return_empty_sngl, - create_process_table + return_search_summary, ) -from pycbc import events, conversions, pnutils -from pycbc.events import ranking, veto -from pycbc.events import mean_if_greater_than_zero -logger = logging.getLogger('pycbc.io.hdf') +logger = logging.getLogger("pycbc.io.hdf") class HGroup(h5py.Group): - """ Low level extensions to the h5py group object - """ + """Low level extensions to the h5py group object""" def create_group(self, *args, **kwargs): """ @@ -44,19 +40,16 @@ def create_dataset(self, name, *args, **kwds): Wrapper around h5py's create_dataset so that checksums are used """ # Do not allow callers to override including fletcher32 - if 'fletcher32' in kwds: - del kwds['fletcher32'] + kwds.pop("fletcher32", None) try: return super().create_dataset(name, *args, **kwds, fletcher32=True) except (ValueError, TypeError): logging.debug( - "Could not create a dataset with fletcher32, falling back " - "to default" + "Could not create a dataset with fletcher32, falling back to default" ) return super().create_dataset(name, *args, **kwds) - def __getitem__(self, name): """ Ensures that when accessing subgroups (e.g., g['subgroup']), @@ -70,7 +63,7 @@ def __getitem__(self, name): # Only wrap group-like objects. Datasets and other non-group # objects should be returned unchanged. Use presence of # mapping-like API (keys) as a quick proxy for group-like. - if isinstance(obj, h5py.Group) and hasattr(obj, 'keys'): + if isinstance(obj, h5py.Group) and hasattr(obj, "keys"): return HGroup(obj.id) return obj @@ -84,11 +77,20 @@ def parent(self): class HFile(HGroup, h5py.File): - """ Low level extensions to the capabilities of reading an hdf5 File - """ - def select(self, fcn, *args, chunksize=10**6, derived=None, group='', - return_data=True, premask=None): - """ Return arrays from an hdf5 file that satisfy the given function + """Low level extensions to the capabilities of reading an hdf5 File""" + + def select( + self, + fcn, + *args, + chunksize=10**6, + derived=None, + group="", + return_data=True, + premask=None, + ): + """ + Return arrays from an hdf5 file that satisfy the given function Parameters ---------- @@ -135,8 +137,8 @@ def select(self, fcn, *args, chunksize=10**6, derived=None, group='', >>> f = HFile(filename) >>> snr = f.select(lambda snr: snr > 6, 'H1/snr') - """ + """ # Required datasets are the arguments requested and datasets given # for any derived functions derived = derived if derived is not None else {} @@ -152,11 +154,13 @@ def select(self, fcn, *args, chunksize=10**6, derived=None, group='', refs = {} size = None for ds in dsets: - refs[ds] = self[group + '/' + ds] + refs[ds] = self[group + "/" + ds] if (size is not None) and (refs[ds].size != size): - raise RuntimeError(f"Dataset {ds} is {self[ds].size} " - "entries long, which does not match " - f"previous input datasets ({size}).") + raise RuntimeError( + f"Dataset {ds} is {self[ds].size} " + "entries long, which does not match " + f"previous input datasets ({size})." + ) size = refs[ds].size # Apply any pre-masks @@ -175,8 +179,10 @@ def select(self, fcn, *args, chunksize=10**6, derived=None, group='', if not mask.size == size: # You get here if you are using a boolean premask which # isn't the same size as the arrays - raise RuntimeError(f"Using premask of size {mask.size} which " - f"does not match the input datasets ({size}).") + raise RuntimeError( + f"Using premask of size {mask.size} which " + f"does not match the input datasets ({size})." + ) # datasets being returned (possibly) data = {} @@ -187,7 +193,7 @@ def select(self, fcn, *args, chunksize=10**6, derived=None, group='', # Loop through the chunks: i = 0 while i < size: - r = i + chunksize if i + chunksize < size else size + r = min(size, i + chunksize) if not any(mask[i:r]): # Nothing allowed through the mask in this chunk @@ -201,8 +207,7 @@ def select(self, fcn, *args, chunksize=10**6, derived=None, group='', submask = np.flatnonzero(mask[i:r]) # Read each chunk's worth of data - partial_data = {arg: refs[arg][i:r][mask[i:r]] - for arg in dsets} + partial_data = {arg: refs[arg][i:r][mask[i:r]] for arg in dsets} partial = [] for a in args: if a in derived.keys(): @@ -227,24 +232,26 @@ def select(self, fcn, *args, chunksize=10**6, derived=None, group='', i += chunksize if return_data: - return_tuple = tuple(np.concatenate(data[arg]) - for arg in args) + return_tuple = tuple(np.concatenate(data[arg]) for arg in args) else: return_tuple = None return indices.astype(np.uint64), return_tuple -class DictArray(object): - """ Utility for organizing sets of arrays of equal length. +class DictArray: + """ + Utility for organizing sets of arrays of equal length. Manages a dictionary of arrays of equal length. This can also be instantiated with a set of hdf5 files and the key values. The full data is always in memory and all operations create new instances of the DictArray. """ + def __init__(self, data=None, files=None, groups=None): - """ Create a DictArray + """ + Create a DictArray Parameters ---------- @@ -254,17 +261,21 @@ def __init__(self, data=None, files=None, groups=None): List of hdf5 file filenames. Incompatibile with the `data` option. groups: list of strings List of keys into each file. Required by the files option. + """ # Check that input fits with how the DictArray is set up if data and files: - raise RuntimeError('DictArray can only have data or files as ' - 'input, not both.') + raise RuntimeError( + "DictArray can only have data or files as input, not both." + ) if data is None and files is None: - raise RuntimeError('DictArray needs either data or files at' - 'initialization. To set up an empty instance' - 'use DictArray(data={})') + raise RuntimeError( + "DictArray needs either data or files at" + "initialization. To set up an empty instance" + "use DictArray(data={})" + ) if files and not groups: - raise RuntimeError('If files are given then need groups.') + raise RuntimeError("If files are given then need groups.") self.data = data self.groups = groups @@ -295,8 +306,10 @@ def __len__(self): def __add__(self, other): if self.data == {}: - logger.debug('Adding data to a DictArray instance which ' - 'was initialized with an empty dict') + logger.debug( + "Adding data to a DictArray instance which " + "was initialized with an empty dict" + ) return self._return(data=other) data = {} @@ -304,12 +317,11 @@ def __add__(self, other): try: data[k] = np.concatenate([self.data[k], other.data[k]]) except KeyError: - logger.info('%s does not exist in other data', k) + logger.info("%s does not exist in other data", k) return self._return(data=data) def select(self, idx): - """ Return a new DictArray containing only the indexed values - """ + """Return a new DictArray containing only the indexed values""" data = {} for k in self.data: # Make sure each entry is an array (not a scalar) @@ -317,8 +329,7 @@ def select(self, idx): return self._return(data=data) def remove(self, idx): - """ Return a new DictArray that does not contain the indexed values - """ + """Return a new DictArray that does not contain the indexed values""" data = {} for k in self.data: data[k] = np.delete(self.data[k], np.array(idx, dtype=int)) @@ -330,91 +341,114 @@ def save(self, outname): f.attrs[k] = self.attrs[k] for k in self.data: - f.create_dataset(k, data=self.data[k], - compression='gzip', - compression_opts=9, - shuffle=True) + f.create_dataset( + k, + data=self.data[k], + compression="gzip", + compression_opts=9, + shuffle=True, + ) f.close() class StatmapData(DictArray): - def __init__(self, data=None, seg=None, attrs=None, files=None, - groups=('stat', 'time1', 'time2', 'trigger_id1', - 'trigger_id2', 'template_id', 'decimation_factor', - 'timeslide_id')): - super(StatmapData, self).__init__(data=data, files=files, - groups=groups) + def __init__( + self, + data=None, + seg=None, + attrs=None, + files=None, + groups=( + "stat", + "time1", + "time2", + "trigger_id1", + "trigger_id2", + "template_id", + "decimation_factor", + "timeslide_id", + ), + ): + super().__init__(data=data, files=files, groups=groups) if data: - self.seg=seg - self.attrs=attrs + self.seg = seg + self.attrs = attrs elif files: f = HFile(files[0], "r") - self.seg = f['segments'] + self.seg = f["segments"] self.attrs = f.attrs def _return(self, data): return self.__class__(data=data, attrs=self.attrs, seg=self.seg) def cluster(self, window): - """ Cluster the dict array, assuming it has the relevant Coinc colums, + """ + Cluster the dict array, assuming it has the relevant Coinc colums, time1, time2, stat, and timeslide_id """ # If no events, do nothing if len(self.time1) == 0 or len(self.time2) == 0: return self from pycbc.events import cluster_coincs - interval = self.attrs['timeslide_interval'] - cid = cluster_coincs(self.stat, self.time1, self.time2, - self.timeslide_id, interval, window) + + interval = self.attrs["timeslide_interval"] + cid = cluster_coincs( + self.stat, self.time1, self.time2, self.timeslide_id, interval, window + ) return self.select(cid) def save(self, outname): - super(StatmapData, self).save(outname) + super().save(outname) with HFile(outname, "w") as f: for key in self.seg.keys(): - f['segments/%s/start' % key] = self.seg[key]['start'][:] - f['segments/%s/end' % key] = self.seg[key]['end'][:] + f["segments/%s/start" % key] = self.seg[key]["start"][:] + f["segments/%s/end" % key] = self.seg[key]["end"][:] class MultiifoStatmapData(StatmapData): - def __init__(self, data=None, seg=None, attrs=None, - files=None, ifos=None): - groups = ['decimation_factor', 'stat', 'template_id', 'timeslide_id'] + def __init__(self, data=None, seg=None, attrs=None, files=None, ifos=None): + groups = ["decimation_factor", "stat", "template_id", "timeslide_id"] for ifo in ifos: - groups += ['%s/time' % ifo] - groups += ['%s/trigger_id' % ifo] + groups += ["%s/time" % ifo] + groups += ["%s/trigger_id" % ifo] - super(MultiifoStatmapData, self).__init__(data=data, files=files, - groups=groups, attrs=attrs, - seg=seg) + super().__init__( + data=data, files=files, groups=groups, attrs=attrs, seg=seg + ) def _return(self, data): - ifolist = self.attrs['ifos'].split(' ') - return self.__class__(data=data, attrs=self.attrs, seg=self.seg, - ifos=ifolist) + ifolist = self.attrs["ifos"].split(" ") + return self.__class__(data=data, attrs=self.attrs, seg=self.seg, ifos=ifolist) def cluster(self, window): - """ Cluster the dict array, assuming it has the relevant Coinc colums, + """ + Cluster the dict array, assuming it has the relevant Coinc colums, time1, time2, stat, and timeslide_id """ # If no events, do nothing - pivot_ifo = self.attrs['pivot'] - fixed_ifo = self.attrs['fixed'] - if len(self.data['%s/time' % pivot_ifo]) == 0 or len(self.data['%s/time' % fixed_ifo]) == 0: + pivot_ifo = self.attrs["pivot"] + fixed_ifo = self.attrs["fixed"] + if ( + len(self.data["%s/time" % pivot_ifo]) == 0 + or len(self.data["%s/time" % fixed_ifo]) == 0 + ): return self from pycbc.events import cluster_coincs - interval = self.attrs['timeslide_interval'] - cid = cluster_coincs(self.stat, - self.data['%s/time' % pivot_ifo], - self.data['%s/time' % fixed_ifo], - self.timeslide_id, - interval, - window) + + interval = self.attrs["timeslide_interval"] + cid = cluster_coincs( + self.stat, + self.data["%s/time" % pivot_ifo], + self.data["%s/time" % fixed_ifo], + self.timeslide_id, + interval, + window, + ) return self.select(cid) -class FileData(object): +class FileData: def __init__(self, fname, group=None, columnlist=None, filter_func=None): """ Parameters @@ -426,20 +460,21 @@ def __init__(self, fname, group=None, columnlist=None, filter_func=None): filter_func : string String should evaluate to a Boolean expression using attributes of the class instance derived from columns: ex. 'self.snr < 6.5' + """ - if not fname: raise RuntimeError("Didn't get a file!") + if not fname: + raise RuntimeError("Didn't get a file!") self.fname = fname self.h5file = HFile(fname, "r") if group is None: if len(self.h5file.keys()) == 1: - group, = self.h5file.keys() + (group,) = self.h5file.keys() else: raise RuntimeError("Didn't get a group!") self.group_key = group self.group = self.h5file[group] - self.columns = columnlist if columnlist is not None \ - else list(self.group.keys()) + self.columns = columnlist if columnlist is not None else list(self.group.keys()) self.filter_func = filter_func self._mask = None @@ -455,18 +490,18 @@ def mask(self): ------- array of Boolean True for dataset indices to be returned by the get_column method + """ if self.filter_func is None: raise RuntimeError("Can't get a mask without a filter function!") - else: - # only evaluate if no previous calculation was done - if self._mask is None: - # get required columns into the namespace as numpy arrays - for column in self.columns: - if column in self.filter_func: - setattr(self, column, self.group[column][:]) - self._mask = eval(self.filter_func) - return self._mask + # only evaluate if no previous calculation was done + if self._mask is None: + # get required columns into the namespace as numpy arrays + for column in self.columns: + if column in self.filter_func: + setattr(self, column, self.group[column][:]) + self._mask = eval(self.filter_func) + return self._mask def get_column(self, col): """ @@ -482,6 +517,7 @@ def get_column(self, col): ------- numpy array Values from the dataset, filtered if requested + """ # catch corner case with an empty file (group with no datasets) if not len(self.group.keys()): @@ -489,12 +525,10 @@ def get_column(self, col): vals = self.group[col] if self.filter_func: return vals[self.mask] - else: - return vals[:] + return vals[:] -class DataFromFiles(object): - +class DataFromFiles: def __init__(self, filelist, group=None, columnlist=None, filter_func=None): self.files = filelist self.group = group @@ -515,27 +549,43 @@ def get_column(self, col): numpy array Values from the dataset, filtered if requested and concatenated in order of file list + """ - logger.info('getting %s', col) + logger.info("getting %s", col) vals = [] for f in self.files: - d = FileData(f, group=self.group, columnlist=self.columns, - filter_func=self.filter_func) + d = FileData( + f, + group=self.group, + columnlist=self.columns, + filter_func=self.filter_func, + ) vals.append(d.get_column(col)) # Close each file since h5py has an upper limit on the number of # open file objects (approx. 1000) d.close() - logger.info('- got %i values', sum(len(v) for v in vals)) + logger.info("- got %i values", sum(len(v) for v in vals)) return np.concatenate(vals) -class SingleDetTriggers(object): +class SingleDetTriggers: """ Provides easy access to the parameters of single-detector CBC triggers. """ - def __init__(self, trig_file, detector, bank_file=None, veto_file=None, - segment_name=None, premask=None, filter_rank=None, - filter_threshold=None, chunksize=10**6, filter_func=None): + + def __init__( + self, + trig_file, + detector, + bank_file=None, + veto_file=None, + segment_name=None, + premask=None, + filter_rank=None, + filter_threshold=None, + chunksize=10**6, + filter_func=None, + ): """ Create a SingleDetTriggers instance @@ -569,16 +619,17 @@ def __init__(self, trig_file, detector, bank_file=None, veto_file=None, chunksize : int , default 10**6 Size of chunks to read in for the filter_rank / threshold. + """ - logger.info('Loading triggers') - self.trigs_f = HFile(trig_file, 'r') + logger.info("Loading triggers") + self.trigs_f = HFile(trig_file, "r") self.trigs = self.trigs_f[detector] - self.ntriggers = self.trigs['end_time'].size + self.ntriggers = self.trigs["end_time"].size self.ifo = detector # convenience attributes self.detector = detector if bank_file: - logger.info('Loading bank') - self.bank = HFile(bank_file, 'r') + logger.info("Loading bank") + self.bank = HFile(bank_file, "r") else: # empty dict in place of non-existent hdf file self.bank = {} @@ -595,18 +646,21 @@ def __init__(self, trig_file, detector, bank_file=None, veto_file=None, if filter_rank: assert filter_threshold is not None - logger.info("Applying threshold of %.3f on %s", - filter_threshold, filter_rank) - fcn_dsets = (ranking.sngls_ranking_function_dict[filter_rank], - ranking.reqd_datasets[filter_rank]) + logger.info( + "Applying threshold of %.3f on %s", filter_threshold, filter_rank + ) + fcn_dsets = ( + ranking.sngls_ranking_function_dict[filter_rank], + ranking.reqd_datasets[filter_rank], + ) idx, _ = self.trigs_f.select( - lambda rank: rank > filter_threshold, - filter_rank, - derived={filter_rank: fcn_dsets}, - return_data=False, - premask=self.mask, - group=detector, - chunksize=chunksize, + lambda rank: rank > filter_threshold, + filter_rank, + derived={filter_rank: fcn_dsets}, + return_data=False, + premask=self.mask, + group=detector, + chunksize=chunksize, ) logger.info("%d triggers remain", idx.size) # If self.mask already has values, need to take these into account: @@ -615,42 +669,49 @@ def __init__(self, trig_file, detector, bank_file=None, veto_file=None, if filter_func: # Apply a filter on the triggers which is _not_ a ranking statistic for rank_str in ranking.sngls_ranking_function_dict.keys(): - if f'self.{rank_str}' in filter_func: - logger.warning('Supplying the ranking (%s) in ' - 'filter_func is inefficient, suggest to ' - 'use filter_rank instead.', rank_str) - logger.info('Setting up filter function') + if f"self.{rank_str}" in filter_func: + logger.warning( + "Supplying the ranking (%s) in " + "filter_func is inefficient, suggest to " + "use filter_rank instead.", + rank_str, + ) + logger.info("Setting up filter function") for c in self.trigs.keys(): if c in filter_func: - setattr(self, '_'+c, self.trigs[c][:]) + setattr(self, "_" + c, self.trigs[c][:]) for c in self.bank.keys(): if c in filter_func: # get template parameters corresponding to triggers - setattr(self, '_'+c, - np.array(self.bank[c])[self.trigs['template_id'][:]]) + setattr( + self, + "_" + c, + np.array(self.bank[c])[self.trigs["template_id"][:]], + ) - filter_mask = eval(filter_func.replace('self.', 'self._')) + filter_mask = eval(filter_func.replace("self.", "self._")) # remove the dummy attributes for c in chain(self.trigs.keys(), self.bank.keys()): - if c in filter_func: delattr(self, '_'+c) + if c in filter_func: + delattr(self, "_" + c) self.apply_mask(filter_mask) - logger.info('%i triggers remain after cut on %s', - sum(self.mask), filter_func) + logger.info( + "%i triggers remain after cut on %s", sum(self.mask), filter_func + ) if veto_file: - logger.info('Applying veto segments') + logger.info("Applying veto segments") # veto_mask is an array of indices into the trigger arrays # giving the surviving triggers - logger.info('%i triggers before vetoes', self.mask_size) + logger.info("%i triggers before vetoes", self.mask_size) veto_mask, _ = events.veto.indices_outside_segments( - self.end_time, [veto_file], - ifo=detector, segment_name=segment_name) + self.end_time, [veto_file], ifo=detector, segment_name=segment_name + ) # Update mask accordingly self.apply_mask(veto_mask) - logger.info('%i triggers remain after vetoes', - self.mask_size) + logger.info("%i triggers remain after vetoes", self.mask_size) def __getitem__(self, key): # Is key in the TRIGGER_MERGE file? @@ -664,44 +725,44 @@ def __getitem__(self, key): self.checkbank(key) return self.bank[key][:][self.template_id] except (RuntimeError, KeyError) as exc: - err_msg = "Cannot find {} in input files".format(key) + err_msg = f"Cannot find {key} in input files" raise ValueError(err_msg) from exc def checkbank(self, param): if self.bank == {}: - return RuntimeError("Can't get %s values without a bank file" - % param) + return RuntimeError("Can't get %s values without a bank file" % param) def trig_dict(self): """Returns dict of the masked trigger values""" mtrigs = {} for k in self.trigs: - if len(self.trigs[k]) == len(self.trigs['end_time']): + if len(self.trigs[k]) == len(self.trigs["end_time"]): if self.mask is not None: mtrigs[k] = self.trigs[k][self.mask] else: mtrigs[k] = self.trigs[k][:] - mtrigs['ifo'] = self.ifo + mtrigs["ifo"] = self.ifo return mtrigs @classmethod def get_param_names(cls): """Returns a list of plottable CBC parameter variables""" - return [m[0] for m in inspect.getmembers(cls) \ - if type(m[1]) == property] + return [m[0] for m in inspect.getmembers(cls) if type(m[1]) == property] def apply_mask(self, logic_mask): - """Apply a mask over the top of the current mask + """ + Apply a mask over the top of the current mask Parameters ---------- logic_mask : boolean array or numpy array of indices + """ if self.mask is None: self.mask = np.zeros(self.ntriggers, dtype=bool) self.mask[logic_mask] = True - elif hasattr(self.mask, 'dtype') and (self.mask.dtype == 'bool'): - if hasattr(logic_mask, 'dtype') and (logic_mask.dtype == 'bool'): + elif hasattr(self.mask, "dtype") and (self.mask.dtype == "bool"): + if hasattr(logic_mask, "dtype") and (logic_mask.dtype == "bool"): # So both new and old masks are boolean, numpy slice assignment # can be used directly, with no additional memory. self.mask[self.mask] = logic_mask @@ -716,11 +777,13 @@ def apply_mask(self, logic_mask): self.mask = list(np.array(self.mask)[logic_mask]) def and_masks(self, logic_mask): - """Apply a mask to be combined as a logical and with the current mask. + """ + Apply a mask to be combined as a logical and with the current mask. Parameters ---------- logic_mask : boolean array or numpy array/list of indices + """ if self.mask_size == self.ntriggers: # No mask exists, just update to use the given mask @@ -728,12 +791,12 @@ def and_masks(self, logic_mask): return # Use intersection of the indices of True values in the masks - if hasattr(logic_mask, 'dtype') and (logic_mask.dtype == 'bool'): + if hasattr(logic_mask, "dtype") and (logic_mask.dtype == "bool"): new_indices = np.flatnonzero(logic_mask) else: new_indices = np.array(logic_mask) - if hasattr(self.mask, 'dtype') and (self.mask.dtype == 'bool'): + if hasattr(self.mask, "dtype") and (self.mask.dtype == "bool"): orig_indices = np.flatnonzero(self.mask) else: orig_indices = np.array(self.mask) @@ -742,19 +805,21 @@ def and_masks(self, logic_mask): and_indices = np.intersect1d(new_indices, orig_indices) self.mask[and_indices.astype(np.uint64)] = True - def mask_to_n_loudest_clustered_events(self, rank_method, - statistic_threshold=None, - n_loudest=10, - cluster_window=10, - ): - """Edits the mask property of the class to point to the N loudest + def mask_to_n_loudest_clustered_events( + self, + rank_method, + statistic_threshold=None, + n_loudest=10, + cluster_window=10, + ): + """ + Edits the mask property of the class to point to the N loudest single detector events as ranked by ranking statistic. Events are clustered so that no more than 1 event within +/- cluster_window will be considered. Can apply a threshold on the statistic using statistic_threshold """ - sds = rank_method.single(self.trig_dict()) stat = rank_method.rank_stat_single( (self.ifo, sds), @@ -776,8 +841,7 @@ def mask_to_n_loudest_clustered_events(self, rank_method, if len(stat) == 0: logger.warning("No triggers after thresholding") return - else: - logger.info("%d triggers after thresholding", len(stat)) + logger.info("%d triggers after thresholding", len(stat)) index = stat.argsort()[::-1] new_times = [] @@ -814,72 +878,72 @@ def mask_size(self): @property def template_id(self): - return self.get_column('template_id').astype(int) + return self.get_column("template_id").astype(int) @property def mass1(self): - self.checkbank('mass1') - return self.bank['mass1'][:][self.template_id] + self.checkbank("mass1") + return self.bank["mass1"][:][self.template_id] @property def mass2(self): - self.checkbank('mass2') - return self.bank['mass2'][:][self.template_id] + self.checkbank("mass2") + return self.bank["mass2"][:][self.template_id] @property def spin1z(self): - self.checkbank('spin1z') - return self.bank['spin1z'][:][self.template_id] + self.checkbank("spin1z") + return self.bank["spin1z"][:][self.template_id] @property def spin2z(self): - self.checkbank('spin2z') - return self.bank['spin2z'][:][self.template_id] + self.checkbank("spin2z") + return self.bank["spin2z"][:][self.template_id] @property def spin2x(self): - self.checkbank('spin2x') - return self.bank['spin2x'][:][self.template_id] + self.checkbank("spin2x") + return self.bank["spin2x"][:][self.template_id] @property def spin2y(self): - self.checkbank('spin2y') - return self.bank['spin2y'][:][self.template_id] + self.checkbank("spin2y") + return self.bank["spin2y"][:][self.template_id] @property def spin1x(self): - self.checkbank('spin1x') - return self.bank['spin1x'][:][self.template_id] + self.checkbank("spin1x") + return self.bank["spin1x"][:][self.template_id] @property def spin1y(self): - self.checkbank('spin1y') - return self.bank['spin1y'][:][self.template_id] + self.checkbank("spin1y") + return self.bank["spin1y"][:][self.template_id] @property def inclination(self): - self.checkbank('inclination') - return self.bank['inclination'][:][self.template_id] + self.checkbank("inclination") + return self.bank["inclination"][:][self.template_id] @property def eccentricity(self): - self.checkbank('eccentricity') - return self.bank['eccentricity'][:][self.template_id] + self.checkbank("eccentricity") + return self.bank["eccentricity"][:][self.template_id] @property def rel_anomaly(self): - self.checkbank('rel_anomaly') - return self.bank['rel_anomaly'][:][self.template_id] + self.checkbank("rel_anomaly") + return self.bank["rel_anomaly"][:][self.template_id] @property def f_lower(self): - self.checkbank('f_lower') - return self.bank['f_lower'][:][self.template_id] + self.checkbank("f_lower") + return self.bank["f_lower"][:][self.template_id] @property def approximant(self): - self.checkbank('approximant') - return self.bank['approximant'][:][self.template_id] + self.checkbank("approximant") + return self.bank["approximant"][:][self.template_id] @property def mtotal(self): @@ -896,50 +960,50 @@ def eta(self): @property def effective_spin(self): # FIXME assumes aligned spins - return conversions.chi_eff(self.mass1, self.mass2, - self.spin1z, self.spin2z) + return conversions.chi_eff(self.mass1, self.mass2, self.spin1z, self.spin2z) # IMPROVEME: would like to have a way to access all get_freq and/or # other pnutils.* names rather than hard-coding each one # - eg make this part of a fancy interface to the bank file ? @property def f_seobnrv2_peak(self): - return pnutils.get_freq('fSEOBNRv2Peak', self.mass1, self.mass2, - self.spin1z, self.spin2z) + return pnutils.get_freq( + "fSEOBNRv2Peak", self.mass1, self.mass2, self.spin1z, self.spin2z + ) @property def f_seobnrv4_peak(self): - return pnutils.get_freq('fSEOBNRv4Peak', self.mass1, self.mass2, - self.spin1z, self.spin2z) + return pnutils.get_freq( + "fSEOBNRv4Peak", self.mass1, self.mass2, self.spin1z, self.spin2z + ) @property def end_time(self): - return self.get_column('end_time') + return self.get_column("end_time") @property def template_duration(self): - return self.get_column('template_duration') + return self.get_column("template_duration") @property def snr(self): - return self.get_column('snr') + return self.get_column("snr") @property def sgchisq(self): - return self.get_column('sg_chisq') + return self.get_column("sg_chisq") @property def u_vals(self): - return self.get_column('u_vals') + return self.get_column("u_vals") @property def rchisq(self): - return self.get_column('chisq') \ - / (self.get_column('chisq_dof') * 2 - 2) + return self.get_column("chisq") / (self.get_column("chisq_dof") * 2 - 2) @property def psd_var_val(self): - return self.get_column('psd_var_val') + return self.get_column("psd_var_val") @property def newsnr(self): @@ -951,13 +1015,15 @@ def newsnr_sgveto(self): @property def newsnr_sgveto_psdvar(self): - return ranking.newsnr_sgveto_psdvar(self.snr, self.rchisq, - self.sgchisq, self.psd_var_val) + return ranking.newsnr_sgveto_psdvar( + self.snr, self.rchisq, self.sgchisq, self.psd_var_val + ) @property def newsnr_sgveto_psdvar_threshold(self): - return ranking.newsnr_sgveto_psdvar_threshold(self.snr, self.rchisq, - self.sgchisq, self.psd_var_val) + return ranking.newsnr_sgveto_psdvar_threshold( + self.snr, self.rchisq, self.sgchisq, self.psd_var_val + ) def get_ranking(self, rank_name, **kwargs): return ranking.get_sngls_ranking_from_trigs(self, rank_name, **kwargs) @@ -972,30 +1038,33 @@ def get_column(self, cname): # If the mask accesses few enough elements then directly use it # This can be slower than reading in all the elements if most of them # will be read. - if isinstance(self.mask, list) or \ - self.mask_size < (self.ntriggers * MFRAC): + if isinstance(self.mask, list) or self.mask_size < (self.ntriggers * MFRAC): return self.trigs[cname][self.mask] # We have a lot of elements to read so we resort to readin the entire # array before masking. - elif self.mask is not None: + if self.mask is not None: return self.trigs[cname][:][self.mask] - else: - return self.trigs[cname][:] + return self.trigs[cname][:] -class ForegroundTriggers(object): - +class ForegroundTriggers: # Injection files are expected to only have 'exclusive' IFAR/FAP values, # should use has_inc=False for these. - def __init__(self, coinc_file, bank_file, sngl_files=None, n_loudest=None, - group='foreground', has_inc=True): + def __init__( + self, + coinc_file, + bank_file, + sngl_files=None, + n_loudest=None, + group="foreground", + has_inc=True, + ): self.coinc_file = FileData(coinc_file, group=group) - if 'ifos' in self.coinc_file.h5file.attrs: - self.ifos = self.coinc_file.h5file.attrs['ifos'].split(' ') + if "ifos" in self.coinc_file.h5file.attrs: + self.ifos = self.coinc_file.h5file.attrs["ifos"].split(" ") else: - raise ValueError("File doesn't have an 'ifos' attribute!", - coinc_file) + raise ValueError("File doesn't have an 'ifos' attribute!", coinc_file) self.sngl_files = {} if sngl_files is not None: for sngl_file in sngl_files: @@ -1004,13 +1073,17 @@ def __init__(self, coinc_file, bank_file, sngl_files=None, n_loudest=None, self.sngl_files[curr_ifo] = curr_dat if not all([ifo in self.sngl_files.keys() for ifo in self.ifos]): - print("sngl_files: {}".format(sngl_files)) - print("self.ifos: {}".format(self.ifos)) - raise RuntimeError("IFOs in statmap file not all represented " - "by single-detector trigger files.") + print(f"sngl_files: {sngl_files}") + print(f"self.ifos: {self.ifos}") + raise RuntimeError( + "IFOs in statmap file not all represented " + "by single-detector trigger files." + ) if not sorted(self.sngl_files.keys()) == sorted(self.ifos): - logger.warning("WARNING: Single-detector trigger files " - "given for IFOs not in the statmap file") + logger.warning( + "WARNING: Single-detector trigger files " + "given for IFOs not in the statmap file" + ) self.bank_file = HFile(bank_file, "r") self.n_loudest = n_loudest @@ -1026,24 +1099,26 @@ def sort_arr(self): if self._sort_arr is None: if self._inclusive: try: - ifar = self.coinc_file.get_column('ifar') + ifar = self.coinc_file.get_column("ifar") except KeyError: - logger.warning("WARNING: Can't find inclusive IFAR!" - "Using exclusive IFAR instead ...") - ifar = self.coinc_file.get_column('ifar_exc') + logger.warning( + "WARNING: Can't find inclusive IFAR!" + "Using exclusive IFAR instead ..." + ) + ifar = self.coinc_file.get_column("ifar_exc") self._inclusive = False else: - ifar = self.coinc_file.get_column('ifar_exc') + ifar = self.coinc_file.get_column("ifar_exc") sorting = ifar.argsort()[::-1] if self.n_loudest: - sorting = sorting[:self.n_loudest] + sorting = sorting[: self.n_loudest] self._sort_arr = sorting return self._sort_arr @property def template_id(self): if self._template_id is None: - template_id = self.get_coincfile_array('template_id') + template_id = self.get_coincfile_array("template_id") self._template_id = template_id.astype(int) return self._template_id @@ -1054,7 +1129,7 @@ def trig_id(self): self._trig_ids = {} for ifo in self.ifos: - self._trig_ids[ifo] = self.get_coincfile_array(ifo + '/trigger_id') + self._trig_ids[ifo] = self.get_coincfile_array(ifo + "/trigger_id") return self._trig_ids def get_coincfile_array(self, variable): @@ -1090,10 +1165,7 @@ def get_snglfile_array_dict(self, variable): needed_data = dataset[mask] # Get order and duplicate information back that was lost in # the boolean mask assignment - _, order_duplicate_index = np.unique( - tid, - return_inverse=True - ) + _, order_duplicate_index = np.unique(tid, return_inverse=True) curr = needed_data[order_duplicate_index] except IndexError: if len(self.trig_id[ifo]) == 0: @@ -1107,16 +1179,15 @@ def get_snglfile_array_dict(self, variable): def get_active_segments(self): self.active_segments = {} for ifo in self.ifos: - starts = self.sngl_files[ifo].get_column('search/start_time') - ends = self.sngl_files[ifo].get_column('search/end_time') - self.active_segments[ifo] = veto.start_end_to_segments(starts, - ends) + starts = self.sngl_files[ifo].get_column("search/start_time") + ends = self.sngl_files[ifo].get_column("search/end_time") + self.active_segments[ifo] = veto.start_end_to_segments(starts, ends) def get_end_time(self): - times_gen = (self.get_coincfile_array('{}/time'.format(ifo)) - for ifo in self.ifos) - ref_times = np.array([mean_if_greater_than_zero(t)[0] - for t in zip(*times_gen)]) + times_gen = ( + self.get_coincfile_array(f"{ifo}/time") for ifo in self.ifos + ) + ref_times = np.array([mean_if_greater_than_zero(t)[0] for t in zip(*times_gen)]) return ref_times def get_ifos(self): @@ -1126,17 +1197,16 @@ def get_ifos(self): ifos_list List of lists of ifo names involved in each foreground event. Ifos will be listed in the same order as self.ifos + """ # Ian thinks this could be coded more simply and efficiently # Note also that effectively the same thing is done as part of the # to_coinc_hdf_object method ifo_or_minus = [] for ifo in self.ifos: - ifo_trigs = np.where(self.get_coincfile_array(ifo + '/time') < 0, - '-', ifo) + ifo_trigs = np.where(self.get_coincfile_array(ifo + "/time") < 0, "-", ifo) ifo_or_minus.append(ifo_trigs) - ifos_list = [list(trig[trig != '-']) - for trig in iter(np.array(ifo_or_minus).T)] + ifos_list = [list(trig[trig != "-"]) for trig in iter(np.array(ifo_or_minus).T)] return ifos_list def to_coinc_xml_object(self, file_name): @@ -1144,32 +1214,27 @@ def to_coinc_xml_object(self, file_name): outdoc.appendChild(ligolw.LIGO_LW()) ifos = sorted(self.sngl_files) - proc_table = create_process_table( - outdoc, - program_name='pycbc', - detectors=ifos - ) + proc_table = create_process_table(outdoc, program_name="pycbc", detectors=ifos) proc_id = proc_table.process_id search_summ_table = lsctables.SearchSummaryTable.new() coinc_h5file = self.coinc_file.h5file try: - start_time = coinc_h5file['segments']['coinc']['start'][:].min() - end_time = coinc_h5file['segments']['coinc']['end'][:].max() + start_time = coinc_h5file["segments"]["coinc"]["start"][:].min() + end_time = coinc_h5file["segments"]["coinc"]["end"][:].max() except KeyError: start_times = [] end_times = [] - for ifo_comb in coinc_h5file['segments']: - if ifo_comb == 'foreground_veto': + for ifo_comb in coinc_h5file["segments"]: + if ifo_comb == "foreground_veto": continue - seg_group = coinc_h5file['segments'][ifo_comb] - start_times.append(seg_group['start'][:].min()) - end_times.append(seg_group['end'][:].max()) + seg_group = coinc_h5file["segments"][ifo_comb] + start_times.append(seg_group["start"][:].min()) + end_times.append(seg_group["end"][:].max()) start_time = min(start_times) end_time = max(end_times) num_trigs = len(self.sort_arr) - search_summary = return_search_summary(start_time, end_time, - num_trigs, ifos) + search_summary = return_search_summary(start_time, end_time, num_trigs, ifos) search_summ_table.append(search_summary) outdoc.childNodes[0].appendChild(search_summ_table) @@ -1194,29 +1259,37 @@ def to_coinc_xml_object(self, file_name): coinc_def_id = lsctables.CoincDefID(0) coinc_def_row = lsctables.CoincDef() coinc_def_row.search = "inspiral" - coinc_def_row.description = \ - "sngl_inspiral<-->sngl_inspiral coincidences" + coinc_def_row.description = "sngl_inspiral<-->sngl_inspiral coincidences" coinc_def_row.coinc_def_id = coinc_def_id coinc_def_row.search_coinc_type = 0 coinc_def_table.append(coinc_def_row) - bank_col_names = ['mass1', 'mass2', 'spin1z', 'spin2z'] + bank_col_names = ["mass1", "mass2", "spin1z", "spin2z"] bank_col_vals = {} for name in bank_col_names: bank_col_vals[name] = self.get_bankfile_array(name) - coinc_event_names = ['ifar', 'time', 'fap', 'stat'] + coinc_event_names = ["ifar", "time", "fap", "stat"] coinc_event_vals = {} for name in coinc_event_names: - if name == 'time': + if name == "time": coinc_event_vals[name] = self.get_end_time() else: coinc_event_vals[name] = self.get_coincfile_array(name) - sngl_col_names = ['snr', 'chisq', 'chisq_dof', 'bank_chisq', - 'bank_chisq_dof', 'cont_chisq', 'cont_chisq_dof', - 'end_time', 'template_duration', 'coa_phase', - 'sigmasq'] + sngl_col_names = [ + "snr", + "chisq", + "chisq_dof", + "bank_chisq", + "bank_chisq_dof", + "cont_chisq", + "cont_chisq_dof", + "end_time", + "template_duration", + "coa_phase", + "sigmasq", + ] sngl_col_vals = {} for name in sngl_col_names: sngl_col_vals[name] = self.get_snglfile_array_dict(name) @@ -1234,7 +1307,7 @@ def to_coinc_xml_object(self, file_name): for ifo in ifos: # If this ifo is not participating in this coincidence then # ignore it and move on. - if not sngl_col_vals['snr'][ifo][1][idx]: + if not sngl_col_vals["snr"][ifo][1][idx]: continue triggered_ifos += [ifo] event_id = lsctables.SnglInspiralID(sngl_event_count) @@ -1242,14 +1315,14 @@ def to_coinc_xml_object(self, file_name): sngl = return_empty_sngl() sngl.event_id = event_id sngl.ifo = ifo - net_snrsq += sngl_col_vals['snr'][ifo][0][idx]**2 + net_snrsq += sngl_col_vals["snr"][ifo][0][idx] ** 2 for name in sngl_col_names: val = sngl_col_vals[name][ifo][0][idx] - if name == 'end_time': + if name == "end_time": sngl.end = LIGOTimeGPS(val) - elif name == 'chisq': + elif name == "chisq": # Use reduced chisquared to be consistent with Live - dof = 2. * sngl_col_vals['chisq_dof'][ifo][0][idx] - 2. + dof = 2.0 * sngl_col_vals["chisq_dof"][ifo][0][idx] - 2.0 sngl.chisq = val / dof else: setattr(sngl, name, val) @@ -1257,10 +1330,12 @@ def to_coinc_xml_object(self, file_name): val = bank_col_vals[name][idx] setattr(sngl, name, val) sngl.mtotal, sngl.eta = pnutils.mass1_mass2_to_mtotal_eta( - sngl.mass1, sngl.mass2) + sngl.mass1, sngl.mass2 + ) sngl.mchirp, _ = pnutils.mass1_mass2_to_mchirp_eta( - sngl.mass1, sngl.mass2) - sngl.eff_distance = (sngl.sigmasq)**0.5 / sngl.snr + sngl.mass1, sngl.mass2 + ) + sngl.eff_distance = (sngl.sigmasq) ** 0.5 / sngl.snr # If exact match is not used, get masses from single triggers sngl_mchirps += [sngl.mchirp] sngl_mtots += [sngl.mtotal] @@ -1269,7 +1344,7 @@ def to_coinc_xml_object(self, file_name): # Set up coinc_map entry coinc_map_row = lsctables.CoincMap() - coinc_map_row.table_name = 'sngl_inspiral' + coinc_map_row.table_name = "sngl_inspiral" coinc_map_row.coinc_event_id = coinc_id coinc_map_row.event_id = event_id coinc_event_map_table.append(coinc_map_row) @@ -1285,7 +1360,7 @@ def to_coinc_xml_object(self, file_name): coinc_event_row.nevents = len(triggered_ifos) # NB, `coinc_event_row.instruments = triggered_ifos does not give a # correct result with ligo.lw 1.7.1 - coinc_event_row.instruments = ','.join(sorted(triggered_ifos)) + coinc_event_row.instruments = ",".join(sorted(triggered_ifos)) coinc_inspiral_row.instruments = triggered_ifos coinc_event_row.time_slide_id = time_slide_id coinc_event_row.process_id = proc_id @@ -1293,15 +1368,16 @@ def to_coinc_xml_object(self, file_name): coinc_inspiral_row.coinc_event_id = coinc_id coinc_inspiral_row.mchirp = sngl_combined_mchirp coinc_inspiral_row.mass = sngl_combined_mtot - coinc_inspiral_row.end = LIGOTimeGPS(coinc_event_vals['time'][idx]) + coinc_inspiral_row.end = LIGOTimeGPS(coinc_event_vals["time"][idx]) coinc_inspiral_row.snr = net_snrsq**0.5 - coinc_inspiral_row.false_alarm_rate = coinc_event_vals['fap'][idx] - coinc_inspiral_row.combined_far = 1./coinc_event_vals['ifar'][idx] + coinc_inspiral_row.false_alarm_rate = coinc_event_vals["fap"][idx] + coinc_inspiral_row.combined_far = 1.0 / coinc_event_vals["ifar"][idx] # Transform to Hz - coinc_inspiral_row.combined_far = \ - conversions.sec_to_year(coinc_inspiral_row.combined_far) - coinc_event_row.likelihood = coinc_event_vals['stat'][idx] - coinc_inspiral_row.minimum_duration = 0. + coinc_inspiral_row.combined_far = conversions.sec_to_year( + coinc_inspiral_row.combined_far + ) + coinc_event_row.likelihood = coinc_event_vals["stat"][idx] + coinc_inspiral_row.minimum_duration = 0.0 coinc_event_table.append(coinc_event_row) coinc_inspiral_table.append(coinc_inspiral_row) @@ -1315,55 +1391,53 @@ def to_coinc_xml_object(self, file_name): ligolw_utils.write_filename(outdoc, file_name) def to_coinc_hdf_object(self, file_name): - ofd = HFile(file_name,'w') + ofd = HFile(file_name, "w") # Some fields are special cases logger.info("Outputting search results") time = self.get_end_time() # time will be used later to determine active ifos - ofd['time'] = time + ofd["time"] = time if self._inclusive: - ofd['ifar'] = self.get_coincfile_array('ifar') - ofd['p_value'] = self.get_coincfile_array('fap') + ofd["ifar"] = self.get_coincfile_array("ifar") + ofd["p_value"] = self.get_coincfile_array("fap") - ofd['ifar_exclusive'] = self.get_coincfile_array('ifar_exc') - ofd['p_value_exclusive'] = self.get_coincfile_array('fap_exc') + ofd["ifar_exclusive"] = self.get_coincfile_array("ifar_exc") + ofd["p_value_exclusive"] = self.get_coincfile_array("fap_exc") # Coinc fields - for field in ['stat']: + for field in ["stat"]: ofd[field] = self.get_coincfile_array(field) logger.info("Outputting template information") # Bank fields - for field in ['mass1','mass2','spin1z','spin2z']: + for field in ["mass1", "mass2", "spin1z", "spin2z"]: ofd[field] = self.get_bankfile_array(field) - mass1 = self.get_bankfile_array('mass1') - mass2 = self.get_bankfile_array('mass2') - ofd['chirp_mass'], _ = pnutils.mass1_mass2_to_mchirp_eta(mass1, mass2) + mass1 = self.get_bankfile_array("mass1") + mass2 = self.get_bankfile_array("mass2") + ofd["chirp_mass"], _ = pnutils.mass1_mass2_to_mchirp_eta(mass1, mass2) logger.info("Outputting single-trigger information") logger.info("reduced chisquared") - chisq_vals_valid = self.get_snglfile_array_dict('chisq') - chisq_dof_vals_valid = self.get_snglfile_array_dict('chisq_dof') + chisq_vals_valid = self.get_snglfile_array_dict("chisq") + chisq_dof_vals_valid = self.get_snglfile_array_dict("chisq_dof") for ifo in self.ifos: chisq_vals = chisq_vals_valid[ifo][0] chisq_valid = chisq_vals_valid[ifo][1] chisq_dof_vals = chisq_dof_vals_valid[ifo][0] - rchisq = chisq_vals / (2. * chisq_dof_vals - 2.) - rchisq[np.logical_not(chisq_valid)] = -1. - ofd[ifo + '_chisq'] = rchisq + rchisq = chisq_vals / (2.0 * chisq_dof_vals - 2.0) + rchisq[np.logical_not(chisq_valid)] = -1.0 + ofd[ifo + "_chisq"] = rchisq # Single-detector fields - for field in ['sg_chisq', 'end_time', 'sigmasq', - 'psd_var_val']: + for field in ["sg_chisq", "end_time", "sigmasq", "psd_var_val"]: logger.info(field) try: vals_valid = self.get_snglfile_array_dict(field) except KeyError: - logger.info("%s is not present in the " - "single-detector files", field) + logger.info("%s is not present in the single-detector files", field) for ifo in self.ifos: # Some of the values will not be valid for all IFOs, @@ -1371,93 +1445,103 @@ def to_coinc_hdf_object(self, file_name): # tells us this, and we set the values to -1 vals = vals_valid[ifo][0] valid = vals_valid[ifo][1] - vals[np.logical_not(valid)] = -1. - ofd[f'{ifo}_{field}'] = vals + vals[np.logical_not(valid)] = -1.0 + ofd[f"{ifo}_{field}"] = vals - snr_vals_valid = self.get_snglfile_array_dict('snr') + snr_vals_valid = self.get_snglfile_array_dict("snr") network_snr_sq = np.zeros_like(snr_vals_valid[self.ifos[0]][0]) for ifo in self.ifos: vals = snr_vals_valid[ifo][0] valid = snr_vals_valid[ifo][1] - vals[np.logical_not(valid)] = -1. - ofd[ifo + '_snr'] = vals + vals[np.logical_not(valid)] = -1.0 + ofd[ifo + "_snr"] = vals network_snr_sq[valid] += vals[valid] ** 2.0 - ofd['network_snr'] = np.sqrt(network_snr_sq) + ofd["network_snr"] = np.sqrt(network_snr_sq) logger.info("Triggered detectors") # Create a n_ifos by n_events matrix, with the ifo letter if the # event contains a trigger from the ifo, empty string if not - triggered_matrix = [[ifo[0] if v else '' - for v in snr_vals_valid[ifo][1]] - for ifo in self.ifos] + triggered_matrix = [ + [ifo[0] if v else "" for v in snr_vals_valid[ifo][1]] for ifo in self.ifos + ] # Combine the ifo letters to make a single string per event - triggered_detectors = [''.join(triggered).encode('ascii') - for triggered in zip(*triggered_matrix)] - ofd.create_dataset('trig', data=triggered_detectors, - dtype=' 0 or gveto_after < 0: - raise ValueError("Gating veto window values must be negative " - "before gates and positive after gates.") + raise ValueError( + "Gating veto window values must be negative " + "before gates and positive after gates." + ) if not (gveto_before == 0 and gveto_after == 0): - autogate_times = np.unique( - self.file[self.ifo + '/gating/auto/time'][:]) - if self.ifo + '/gating/file' in self.file: - detgate_times = self.file[self.ifo + '/gating/file/time'][:] + autogate_times = np.unique(self.file[self.ifo + "/gating/auto/time"][:]) + if self.ifo + "/gating/file" in self.file: + detgate_times = self.file[self.ifo + "/gating/file/time"][:] else: detgate_times = [] gate_times = np.concatenate((autogate_times, detgate_times)) gating_veto_segs = veto.start_end_to_segments( - gate_times + gveto_before, - gate_times + gveto_after + gate_times + gveto_before, gate_times + gveto_after ).coalesce() self.segs = (self.segs - gating_veto_segs).coalesce() self.valid = veto.segments_to_start_end(self.segs) def get_data(self, col, num): - """Get a column of data for template with id 'num'. + """ + Get a column of data for template with id 'num'. Parameters ---------- @@ -1470,12 +1554,14 @@ def get_data(self, col, num): ------- data: numpy.ndarray The requested column of data + """ - ref = self.file['%s/%s_template' % (self.ifo, col)][num] - return self.file['%s/%s' % (self.ifo, col)][ref] + ref = self.file["%s/%s_template" % (self.ifo, col)][num] + return self.file["%s/%s" % (self.ifo, col)][ref] def set_template(self, num): - """Set the active template to read from. + """ + Set the active template to read from. Parameters ---------- @@ -1486,23 +1572,23 @@ def set_template(self, num): ------- trigger_id: numpy.ndarray The indices of this templates triggers. + """ self.template_num = num - times = self.get_data('end_time', num) + times = self.get_data("end_time", num) # Determine which of these template's triggers are kept after # applying vetoes if self.valid: - self.keep = veto.indices_within_times(times, self.valid[0], - self.valid[1]) -# logger.info('applying vetoes') + self.keep = veto.indices_within_times(times, self.valid[0], self.valid[1]) + # logger.info('applying vetoes') else: self.keep = np.arange(0, len(times)) if self.bank != {}: self.param = {} - if 'parameters' in self.bank.attrs: - for col in self.bank.attrs['parameters']: + if "parameters" in self.bank.attrs: + for col in self.bank.attrs["parameters"]: self.param[col] = self.bank[col][self.template_num] else: for col in self.bank: @@ -1511,12 +1597,12 @@ def set_template(self, num): # Calculate the trigger id by adding the relative offset in self.keep # to the absolute beginning index of this templates triggers stored # in 'template_boundaries' - trigger_id = self.keep + \ - self.file['%s/template_boundaries' % self.ifo][num] + trigger_id = self.keep + self.file["%s/template_boundaries" % self.ifo][num] return trigger_id def __getitem__(self, col): - """ Return the column of data for current active template after + """ + Return the column of data for current active template after applying vetoes Parameters @@ -1528,17 +1614,29 @@ def __getitem__(self, col): ------- data: numpy.ndarray The requested column of data + """ if self.template_num is None: - raise ValueError('You must call set_template to first pick the ' - 'template to read data from') + raise ValueError( + "You must call set_template to first pick the " + "template to read data from" + ) data = self.get_data(col, self.template_num) data = data[self.keep] if self.valid else data return data -chisq_choices = ['traditional', 'cont', 'bank', 'max_cont_trad', 'sg', - 'max_bank_cont', 'max_bank_trad', 'max_bank_cont_trad'] +chisq_choices = [ + "traditional", + "cont", + "bank", + "max_cont_trad", + "sg", + "max_bank_cont", + "max_bank_trad", + "max_bank_cont_trad", +] + def get_chisq_from_file_choice(hdfile, chisq_choice): """ @@ -1546,58 +1644,62 @@ def get_chisq_from_file_choice(hdfile, chisq_choice): Parameters ---------- - hdfile: HDF file object, or dictionary, or ReadByTemplate object + hdfile: HDF file object, or dictionary, or ReadByTemplate object or SingleDetTriggers object The object to retrieve the chi-squared values from. - chisq_choice: str + chisq_choice: str The choice of chi-squared values to retrieve. Returns ------- chisq: numpy.ndarray The reduced chi-squared values based on the specified choice. + """ # Get the reduced chi-squared values - if chisq_choice in ['traditional','max_cont_trad', 'max_bank_trad', - 'max_bank_cont_trad']: - trad_chisq = hdfile['chisq'][:] + if chisq_choice in [ + "traditional", + "max_cont_trad", + "max_bank_trad", + "max_bank_cont_trad", + ]: + trad_chisq = hdfile["chisq"][:] # We now need to handle the case where chisq is not actually calculated # 0 is used as a sentinel value - trad_chisq_dof = hdfile['chisq_dof'][:] + trad_chisq_dof = hdfile["chisq_dof"][:] red_trad_chisq = trad_chisq / (trad_chisq_dof * 2 - 2) - if chisq_choice in ['cont', 'max_cont_trad', 'max_bank_cont', - 'max_bank_cont_trad']: - cont_chisq = hdfile['cont_chisq'][:] - cont_chisq_dof = hdfile['cont_chisq_dof'][:] + if chisq_choice in ["cont", "max_cont_trad", "max_bank_cont", "max_bank_cont_trad"]: + cont_chisq = hdfile["cont_chisq"][:] + cont_chisq_dof = hdfile["cont_chisq_dof"][:] red_cont_chisq = cont_chisq / cont_chisq_dof - if chisq_choice in ['bank', 'max_bank_cont', 'max_bank_trad', - 'max_bank_cont_trad']: - bank_chisq = hdfile['bank_chisq'][:] - bank_chisq_dof = hdfile['bank_chisq_dof'][:] + if chisq_choice in ["bank", "max_bank_cont", "max_bank_trad", "max_bank_cont_trad"]: + bank_chisq = hdfile["bank_chisq"][:] + bank_chisq_dof = hdfile["bank_chisq_dof"][:] red_bank_chisq = bank_chisq / bank_chisq_dof # return the corresponding reduced chi-squared values depending on the choice - if chisq_choice == 'sg': - chisq = hdfile['sg_chisq'][:] - elif chisq_choice == 'traditional': + if chisq_choice == "sg": + chisq = hdfile["sg_chisq"][:] + elif chisq_choice == "traditional": chisq = red_trad_chisq - elif chisq_choice == 'cont': + elif chisq_choice == "cont": chisq = red_cont_chisq - elif chisq_choice == 'bank': + elif chisq_choice == "bank": chisq = red_bank_chisq - elif chisq_choice == 'max_cont_trad': + elif chisq_choice == "max_cont_trad": chisq = np.maximum(red_trad_chisq, red_cont_chisq) - elif chisq_choice == 'max_bank_cont': + elif chisq_choice == "max_bank_cont": chisq = np.maximum(red_bank_chisq, red_cont_chisq) - elif chisq_choice == 'max_bank_trad': + elif chisq_choice == "max_bank_trad": chisq = np.maximum(red_bank_chisq, red_trad_chisq) - elif chisq_choice == 'max_bank_cont_trad': + elif chisq_choice == "max_bank_cont_trad": chisq = np.maximum(np.maximum(red_bank_chisq, red_cont_chisq), red_trad_chisq) else: err_msg = "Do not recognize --chisq-choice %s" % chisq_choice raise ValueError(err_msg) return chisq + def save_dict_to_hdf5(dic, filename): """ Parameters @@ -1606,9 +1708,11 @@ def save_dict_to_hdf5(dic, filename): python dictionary to be converted to hdf5 format filename: desired name of hdf5 file + """ - with HFile(filename, 'w') as h5file: - recursively_save_dict_contents_to_group(h5file, '/', dic) + with HFile(filename, "w") as h5file: + recursively_save_dict_contents_to_group(h5file, "/", dic) + def recursively_save_dict_contents_to_group(h5file, path, dic): """ @@ -1620,15 +1724,19 @@ def recursively_save_dict_contents_to_group(h5file, path, dic): path within h5py file to saved dictionary dic: python dictionary to be converted to hdf5 format + """ for key, item in dic.items(): - if isinstance(item, (np.ndarray, np.int64, np.float64, str, int, float, - bytes, tuple, list)): + if isinstance( + item, + (np.ndarray, np.int64, np.float64, str, int, float, bytes, tuple, list), + ): h5file[path + str(key)] = item elif isinstance(item, dict): - recursively_save_dict_contents_to_group(h5file, path + key + '/', item) + recursively_save_dict_contents_to_group(h5file, path + key + "/", item) else: - raise ValueError('Cannot save %s type' % type(item)) + raise ValueError("Cannot save %s type" % type(item)) + def load_hdf5_to_dict(h5file, path): """ @@ -1643,47 +1751,53 @@ def load_hdf5_to_dict(h5file, path): ------- dic: dictionary with hdf5 file group content + """ dic = {} for key, item in h5file[path].items(): if isinstance(item, h5py.Dataset): dic[key] = item[()] elif isinstance(item, h5py.Group): - dic[key] = load_hdf5_to_dict(h5file, path + key + '/') + dic[key] = load_hdf5_to_dict(h5file, path + key + "/") else: - raise ValueError('Cannot load %s type' % type(item)) + raise ValueError("Cannot load %s type" % type(item)) return dic + def combine_and_copy(f, files, group): - """ Combine the same column from multiple files and save to a third""" + """Combine the same column from multiple files and save to a third""" # ensure that the files input is stable for iteration order assert isinstance(files, (list, tuple)) - f[group] = np.concatenate([fi[group][:] if group in fi else \ - np.array([], dtype=np.uint32) for fi in files]) + f[group] = np.concatenate( + [fi[group][:] if group in fi else np.array([], dtype=np.uint32) for fi in files] + ) + def name_all_datasets(files): assert isinstance(files, (list, tuple)) datasets = [] for fi in files: - datasets += get_all_subkeys(fi, '/') + datasets += get_all_subkeys(fi, "/") return set(datasets) + def get_all_subkeys(grp, key): subkey_list = [] subkey_start = key - if key == '': + if key == "": grpk = grp else: grpk = grp[key] for sk in grpk.keys(): - path = subkey_start + '/' + sk + path = subkey_start + "/" + sk if isinstance(grp[path], h5py.Dataset): - subkey_list.append(path.lstrip('/')) + subkey_list.append(path.lstrip("/")) else: subkey_list += get_all_subkeys(grp, path) # returns an empty list if there is no dataset or subgroup within the group return subkey_list + # # ============================================================================= # @@ -1693,9 +1807,11 @@ def get_all_subkeys(grp, key): # -def dump_state(state, fp, path=None, dsetname='state', - protocol=pickle.HIGHEST_PROTOCOL): - """Dumps the given state to an hdf5 file handler. +def dump_state( + state, fp, path=None, dsetname="state", protocol=pickle.HIGHEST_PROTOCOL +): + """ + Dumps the given state to an hdf5 file handler. The state is stored as a raw binary array to ``{path}/{dsetname}`` in the given hdf5 file handler. If a dataset with the same name and path is @@ -1719,14 +1835,16 @@ def dump_state(state, fp, path=None, dsetname='state', protocol : int, optional The protocol version to use for pickling. See the :py:mod:`pickle` module for more details. + """ memfp = BytesIO() pickle.dump(state, memfp, protocol=protocol) dump_pickle_to_hdf(memfp, fp, path=path, dsetname=dsetname) -def dump_pickle_to_hdf(memfp, fp, path=None, dsetname='state'): - """Dumps pickled data to an hdf5 file object. +def dump_pickle_to_hdf(memfp, fp, path=None, dsetname="state"): + """ + Dumps pickled data to an hdf5 file object. Parameters ---------- @@ -1740,21 +1858,24 @@ def dump_pickle_to_hdf(memfp, fp, path=None, dsetname='state'): dsetname : str, optional The name of the dataset to store the binary array to. Default is ``state``. + """ memfp.seek(0) - bdata = np.frombuffer(memfp.read(), dtype='S1') + bdata = np.frombuffer(memfp.read(), dtype="S1") if path is not None: - dsetname = path + '/' + dsetname + dsetname = path + "/" + dsetname if dsetname not in fp: - fp.create_dataset(dsetname, shape=bdata.shape, maxshape=(None,), - dtype=bdata.dtype) + fp.create_dataset( + dsetname, shape=bdata.shape, maxshape=(None,), dtype=bdata.dtype + ) elif bdata.size != fp[dsetname].shape[0]: fp[dsetname].resize((bdata.size,)) fp[dsetname][:] = bdata -def load_state(fp, path=None, dsetname='state'): - """Loads a sampler state from the given hdf5 file object. +def load_state(fp, path=None, dsetname="state"): + """ + Loads a sampler state from the given hdf5 file object. The sampler state is expected to be stored as a raw bytes array which can be loaded by pickle. @@ -1769,6 +1890,7 @@ def load_state(fp, path=None, dsetname='state'): dsetname : str, optional The name of the dataset that the state data is stored to. Default is ``state``. + """ if path is not None: fp = fp[path] @@ -1776,10 +1898,25 @@ def load_state(fp, path=None, dsetname='state'): return pickle.load(BytesIO(bdata)) -__all__ = ('HFile', 'DictArray', 'StatmapData', 'MultiifoStatmapData', - 'FileData', 'DataFromFiles', 'SingleDetTriggers', - 'ForegroundTriggers', 'ReadByTemplate', 'chisq_choices', - 'get_chisq_from_file_choice', 'save_dict_to_hdf5', - 'recursively_save_dict_contents_to_group', 'load_hdf5_to_dict', - 'combine_and_copy', 'name_all_datasets', 'get_all_subkeys', - 'dump_state', 'dump_pickle_to_hdf', 'load_state') +__all__ = ( + "DataFromFiles", + "DictArray", + "FileData", + "ForegroundTriggers", + "HFile", + "MultiifoStatmapData", + "ReadByTemplate", + "SingleDetTriggers", + "StatmapData", + "chisq_choices", + "combine_and_copy", + "dump_pickle_to_hdf", + "dump_state", + "get_all_subkeys", + "get_chisq_from_file_choice", + "load_hdf5_to_dict", + "load_state", + "name_all_datasets", + "recursively_save_dict_contents_to_group", + "save_dict_to_hdf5", +) diff --git a/pycbc/io/ligolw.py b/pycbc/io/ligolw.py index aec0a54c07b..0792dfdf0ea 100644 --- a/pycbc/io/ligolw.py +++ b/pycbc/io/ligolw.py @@ -17,44 +17,46 @@ import os import sys + import numpy -from igwn_ligolw import lsctables -from igwn_ligolw import ligolw +from igwn_ligolw import ligolw, lsctables from igwn_ligolw.ligolw import LIGOLWContentHandler as OrigLIGOLWContentHandler from igwn_ligolw.lsctables import TableByName from igwn_ligolw.types import FormatFunc, FromPyType, ToPyType -import pycbc.version as pycbc_version +import pycbc.version as pycbc_version __all__ = ( - 'default_null_value', - 'return_empty_sngl', - 'return_search_summary', - 'create_process_table', - 'legacy_row_id_converter', - 'get_table_columns', - 'LIGOLWContentHandler' + "LIGOLWContentHandler", + "create_process_table", + "default_null_value", + "get_table_columns", + "legacy_row_id_converter", + "return_empty_sngl", + "return_search_summary", ) ROWID_PYTYPE = int ROWID_TYPE = FromPyType[ROWID_PYTYPE] ROWID_FORMATFUNC = FormatFunc[ROWID_TYPE] -IDTypes = set([u"ilwd:char", u"ilwd:char_u"]) +IDTypes = set(["ilwd:char", "ilwd:char_u"]) def default_null_value(col_name, col_type): """ Associate a sensible "null" default value to a given LIGOLW column type. """ - if col_type in ['real_4', 'real_8']: - return 0. - if col_type in ['int_4s', 'int_8s']: + if col_type in ["real_4", "real_8"]: + return 0.0 + if col_type in ["int_4s", "int_8s"]: # this case includes row IDs return 0 - if col_type == 'lstring': - return '' - raise NotImplementedError(('Do not know how to initialize column ' - '{} of type {}').format(col_name, col_type)) + if col_type == "lstring": + return "" + raise NotImplementedError( + f"Do not know how to initialize column {col_name} of type {col_type}" + ) + def return_empty_sngl(nones=False): """ @@ -71,11 +73,11 @@ def return_empty_sngl(nones=False): If True, just set all columns to None. Returns - -------- + ------- lsctables.SnglInspiral The "empty" SnglInspiral object. - """ + """ sngl = lsctables.SnglInspiral() cols = lsctables.SnglInspiralTable.validcolumns for entry in cols: @@ -84,6 +86,7 @@ def return_empty_sngl(nones=False): setattr(sngl, col_name, value) return sngl + def return_search_summary(start_time=0, end_time=0, nevents=0, ifos=None): """ Function to create a SearchSummary object where all columns are populated @@ -96,9 +99,10 @@ def return_search_summary(start_time=0, end_time=0, nevents=0, ifos=None): It then populates columns if given them as options. Returns - -------- + ------- lsctables.SeachSummary The "empty" SearchSummary object. + """ if ifos is None: ifos = [] @@ -128,12 +132,14 @@ def return_search_summary(start_time=0, end_time=0, nevents=0, ifos=None): return search_summary -def create_process_table(document, program_name=None, detectors=None, - comment=None, options=None): - """Create a LIGOLW process table with sane defaults, add it to a LIGOLW + +def create_process_table( + document, program_name=None, detectors=None, comment=None, options=None +): + """ + Create a LIGOLW process table with sane defaults, add it to a LIGOLW document, and return it. """ - if program_name is None: program_name = os.path.basename(sys.argv[0]) if options is None: @@ -155,15 +161,17 @@ def create_process_table(document, program_name=None, detectors=None, program_name, opts, version=pycbc_version.version, - cvs_repository='pycbc/'+pycbc_version.git_branch, + cvs_repository="pycbc/" + pycbc_version.git_branch, cvs_entry_time=cvs_entry_time, instruments=detectors, - comment=comment + comment=comment, ) return process + def legacy_row_id_converter(ContentHandler): - """Convert from old-style to new-style row IDs on the fly. + """ + Convert from old-style to new-style row IDs on the fly. This is loosely adapted from :func:`ligo.lw.utils.ilwd.strip_ilwdchar`. @@ -172,24 +180,26 @@ def legacy_row_id_converter(ContentHandler): When building a ContentHandler, this must be the _outermost_ decorator, outside of :func:`ligo.lw.lsctables.use_in`, :func:`ligo.lw.param.use_in`, or :func:`ligo.lw.table.use_in`. + """ - def endElementNS(self, uri_localname, qname, - __orig_endElementNS=ContentHandler.endElementNS): + def endElementNS( + self, uri_localname, qname, __orig_endElementNS=ContentHandler.endElementNS + ): """Convert values of elements from ilwdchar to int.""" if isinstance(self.current, ligolw.Param) and self.current.Type in IDTypes: old_type = ToPyType[self.current.Type] old_val = str(old_type(self.current.pcdata)) - new_value = ROWID_PYTYPE(old_val.split(":")[-1]) + new_value = ROWID_PYTYPE(old_val.rsplit(":", maxsplit=1)[-1]) self.current.Type = ROWID_TYPE self.current.pcdata = ROWID_FORMATFUNC(new_value) __orig_endElementNS(self, uri_localname, qname) remapped = {} - def startColumn(self, parent, attrs, - __orig_startColumn=ContentHandler.startColumn): - """Convert types in elements from ilwdchar to int. + def startColumn(self, parent, attrs, __orig_startColumn=ContentHandler.startColumn): + """ + Convert types in elements from ilwdchar to int. Notes ----- @@ -215,18 +225,18 @@ def converter(old_value): validcolumns = TableByName[parent.Name].validcolumns if result.Name not in validcolumns: stripped_column_to_valid_column = { - ligolw.Column.ColumnName(name): name - for name in validcolumns + ligolw.Column.ColumnName(name): name for name in validcolumns } if result.Name in stripped_column_to_valid_column: result.setAttribute( - 'Name', stripped_column_to_valid_column[result.Name]) + "Name", stripped_column_to_valid_column[result.Name] + ) return result - def startStream(self, parent, attrs, - __orig_startStream=ContentHandler.startStream): - """Convert values in table elements from ilwdchar to int. + def startStream(self, parent, attrs, __orig_startStream=ContentHandler.startStream): + """ + Convert values in table elements from ilwdchar to int. Notes ----- @@ -241,11 +251,16 @@ def startStream(self, parent, attrs, # FIXME: convert loadcolumns attributes to sets to # avoid the conversion. loadcolumns &= set(parent.loadcolumns) - result._tokenizer.set_types([ - (remapped.pop((id(parent), colname), pytype) - if colname in loadcolumns else None) - for pytype, colname - in zip(parent.columnpytypes, parent.columnnames)]) + result._tokenizer.set_types( + [ + ( + remapped.pop((id(parent), colname), pytype) + if colname in loadcolumns + else None + ) + for pytype, colname in zip(parent.columnpytypes, parent.columnnames) + ] + ) return result ContentHandler.endElementNS = endElementNS @@ -254,29 +269,28 @@ def startStream(self, parent, attrs, return ContentHandler + def _build_series(series, dim_names, comment, delta_name, delta_unit): Attributes = ligolw.sax.xmlreader.AttributesImpl - elem = ligolw.LIGO_LW( - Attributes({'Name': str(series.__class__.__name__)})) + elem = ligolw.LIGO_LW(Attributes({"Name": str(series.__class__.__name__)})) if comment is not None: elem.appendChild(ligolw.Comment()).pcdata = comment - elem.appendChild(ligolw.Time.from_gps(series.epoch, 'epoch')) - elem.appendChild(ligolw.Param.from_pyvalue('f0', series.f0, unit='s^-1')) + elem.appendChild(ligolw.Time.from_gps(series.epoch, "epoch")) + elem.appendChild(ligolw.Param.from_pyvalue("f0", series.f0, unit="s^-1")) delta = getattr(series, delta_name) if numpy.iscomplexobj(series.data.data): - data = numpy.vstack(( - numpy.arange(len(series.data.data)) * delta, - series.data.data.real, - series.data.data.imag - )) + data = numpy.vstack( + ( + numpy.arange(len(series.data.data)) * delta, + series.data.data.real, + series.data.data.imag, + ) + ) else: - data = numpy.vstack(( - numpy.arange(len(series.data.data)) * delta, - series.data.data - )) - a = ligolw.Array.build( - series.name, data, dim_names=dim_names, encoding='base64' - ) + data = numpy.vstack( + (numpy.arange(len(series.data.data)) * delta, series.data.data) + ) + a = ligolw.Array.build(series.name, data, dim_names=dim_names, encoding="base64") a.Unit = str(series.sampleUnits) dim0 = a.getElementsByTagName(ligolw.Dim.tagName)[0] dim0.Unit = delta_unit @@ -285,50 +299,47 @@ def _build_series(series, dim_names, comment, delta_name, delta_unit): elem.appendChild(a) return elem + def make_psd_xmldoc(psddict, xmldoc=None): - """Add a set of PSDs to a LIGOLW XML document. If the document is not + """ + Add a set of PSDs to a LIGOLW XML document. If the document is not given, a new one is created first. """ xmldoc = ligolw.Document() if xmldoc is None else xmldoc.childNodes[0] # the PSDs must be children of a LIGO_LW with name "psd" - root_name = 'psd' + root_name = "psd" Attributes = ligolw.sax.xmlreader.AttributesImpl - lw = xmldoc.appendChild( - ligolw.LIGO_LW(Attributes({'Name': root_name}))) + lw = xmldoc.appendChild(ligolw.LIGO_LW(Attributes({"Name": root_name}))) for instrument, psd in psddict.items(): xmlseries = _build_series( - psd, - ('Frequency,Real', 'Frequency'), - None, - 'deltaF', - 's^-1' + psd, ("Frequency,Real", "Frequency"), None, "deltaF", "s^-1" ) fs = lw.appendChild(xmlseries) - fs.appendChild(ligolw.Param.from_pyvalue('instrument', instrument)) + fs.appendChild(ligolw.Param.from_pyvalue("instrument", instrument)) return xmldoc + def snr_series_to_xml(snr_series, document, sngl_inspiral_id): - """Save an SNR time series into an XML document, in a format compatible + """ + Save an SNR time series into an XML document, in a format compatible with BAYESTAR. """ snr_lal = snr_series.lal() - snr_lal.name = 'snr' - snr_lal.sampleUnits = '' + snr_lal.name = "snr" + snr_lal.sampleUnits = "" snr_xml = _build_series( - snr_lal, - ('Time', 'Time,Real,Imaginary'), - None, - 'deltaT', - 's' + snr_lal, ("Time", "Time,Real,Imaginary"), None, "deltaT", "s" ) snr_node = document.childNodes[-1].appendChild(snr_xml) - eid_param = ligolw.Param.from_pyvalue('event_id', sngl_inspiral_id) + eid_param = ligolw.Param.from_pyvalue("event_id", sngl_inspiral_id) snr_node.appendChild(eid_param) + def get_table_columns(table): - """Return a list of columns that are present in the given table, in a + """ + Return a list of columns that are present in the given table, in a format that can be passed to `igwn_ligolw.Table.new()`. The split on ":" is needed for columns like `process:process_id`, which @@ -338,7 +349,7 @@ def get_table_columns(table): """ columns = [] for col in table.validcolumns: - att = col.split(':')[-1] + att = col.split(":")[-1] if att in table.columnnames: columns.append(col) return columns @@ -346,4 +357,4 @@ def get_table_columns(table): @legacy_row_id_converter class LIGOLWContentHandler(OrigLIGOLWContentHandler): - "Dummy class needed for loading LIGOLW files" + """Dummy class needed for loading LIGOLW files""" diff --git a/pycbc/io/live.py b/pycbc/io/live.py index 60cf82e0605..6fad6ca0320 100644 --- a/pycbc/io/live.py +++ b/pycbc/io/live.py @@ -1,13 +1,13 @@ +import datetime import logging import os import pathlib -import datetime + import numpy from pycbc.time import gps_to_utc_datetime - -logger = logging.getLogger('pycbc.io.live') +logger = logging.getLogger("pycbc.io.live") def maximum_string(numbers): @@ -21,6 +21,7 @@ def maximum_string(numbers): A list of integers from which to determine the longest common string prefix. E.g. '12345', '12346', '12356' returns '123' + """ # The max length of the number will be the integer above log10 # of the biggest number @@ -42,6 +43,7 @@ def filter_file(filename, start_time, end_time): """ Indicate whether the filename indicates that the file is within the start and end times + Parameters ---------- filename : string @@ -58,11 +60,12 @@ def filter_file(filename, start_time, end_time): ------- boolean Does any of the file lie within the start/end times + """ # FIX ME eventually - this uses the gps time and duration from the filename # Is there a better way? (i.e. trigger gps times in the file or # add an attribute) - fend = filename.split('-')[-2:] + fend = filename.split("-")[-2:] file_start = float(fend[0]) duration = float(fend[1][:-4]) @@ -74,56 +77,59 @@ def add_live_trigger_selection_options(parser): Add options required for obtaining the right set of PyCBC live triggers into an argument parser """ - finding_group = parser.add_argument_group('Trigger Finding') + finding_group = parser.add_argument_group("Trigger Finding") finding_group.add_argument( "--trigger-directory", metavar="PATH", required=True, help="Directory containing trigger files, directory " - "can contain subdirectories. Required." + "can contain subdirectories. Required.", ) finding_group.add_argument( "--gps-start-time", type=int, required=True, - help="Start time of the analysis. Integer, required" + help="Start time of the analysis. Integer, required", ) finding_group.add_argument( "--gps-end-time", type=int, required=True, - help="End time of the analysis. Integer, required" + help="End time of the analysis. Integer, required", ) finding_group.add_argument( "--date-directories", action="store_true", - help="Indicate if the trigger files are stored in " - "directories by date." + help="Indicate if the trigger files are stored in directories by date.", ) default_dd_format = "%Y_%m_%d" finding_group.add_argument( "--date-directory-format", default=default_dd_format, help="Format of date, see datetime strftime " - "documentation for details. Default: " - "%%Y_%%m_%%d" + "documentation for details. Default: " + "%%Y_%%m_%%d", ) finding_group.add_argument( "--file-identifier", default="H1L1V1-Live", help="String required in filename to be considered for " - "analysis. Default: 'H1L1V1-Live'." + "analysis. Default: 'H1L1V1-Live'.", ) -def find_trigger_files(directory, gps_start_time, gps_end_time, - id_string='*', date_directories=False, - date_directory_format="%Y_%m_%d"): +def find_trigger_files( + directory, + gps_start_time, + gps_end_time, + id_string="*", + date_directories=False, + date_directory_format="%Y_%m_%d", +): """ Find a list of PyCBC live trigger files which are between the gps start and end times given """ - # Find the string at the start of the gps time which will match all # files in this range - this helps to cut which ones we need to # compare later @@ -132,7 +138,7 @@ def find_trigger_files(directory, gps_start_time, gps_end_time, # ** means recursive, so for large directories, this is expensive. # It is not too bad if date_directories is set, as we don't waste time # in directories where there cant be any files. - glob_string = f'**/*{id_string}*{num_match}*.hdf' + glob_string = f"**/*{id_string}*{num_match}*.hdf" if date_directories: # convert the GPS times into dates, and only use the directories # of those dates to search @@ -154,8 +160,9 @@ def find_trigger_files(directory, gps_start_time, gps_end_time, matching_files = [f.as_posix() for f in matching_files_gen] # Is the file in the time window? - matching_files = [f for f in matching_files - if filter_file(f, gps_start_time, gps_end_time)] + matching_files = [ + f for f in matching_files if filter_file(f, gps_start_time, gps_end_time) + ] return sorted(matching_files) @@ -171,12 +178,12 @@ def find_trigger_files_from_cli(args): args.gps_end_time, id_string=args.file_identifier, date_directories=args.date_directories, - date_directory_format=args.date_directory_format + date_directory_format=args.date_directory_format, ) __all__ = [ - 'add_live_trigger_selection_options', - 'find_trigger_files', - 'find_trigger_files_from_cli', + "add_live_trigger_selection_options", + "find_trigger_files", + "find_trigger_files_from_cli", ] diff --git a/pycbc/io/record.py b/pycbc/io/record.py index 1a1220786c5..0088194bfc6 100644 --- a/pycbc/io/record.py +++ b/pycbc/io/record.py @@ -28,15 +28,22 @@ waves. """ -import types, re, copy, numpy, inspect +import copy +import inspect +import re +import types + +import numpy from igwn_ligolw import types as ligolw_types -from pycbc import coordinates, conversions, cosmology + +from pycbc import conversions, coordinates, cosmology from pycbc.population import population_models from pycbc.waveform import parameters # what functions are given to the eval in FieldArray's __getitem__: -_numpy_function_lib = {_x: _y for _x,_y in numpy.__dict__.items() - if isinstance(_y, (numpy.ufunc, float))} +_numpy_function_lib = { + _x: _y for _x, _y in numpy.__dict__.items() if isinstance(_y, (numpy.ufunc, float)) +} # # ============================================================================= @@ -47,9 +54,13 @@ # # add ligolw_types to numpy sctypeDict # but don't include bindings that numpy already defines -numpy.sctypeDict.update({_k: _val - for (_k, _val) in ligolw_types.ToNumPyType.items() - if _k not in numpy.sctypeDict}) +numpy.sctypeDict.update( + { + _k: _val + for (_k, _val) in ligolw_types.ToNumPyType.items() + if _k not in numpy.sctypeDict + } +) # Annoyingly, numpy has no way to store NaNs in an integer field to indicate # the equivalent of None. This can be problematic for fields that store ids: @@ -58,23 +69,26 @@ # we define here an integer to indicate 'id not set'. ID_NOT_SET = -1 EMPTY_OBJECT = None -VIRTUALFIELD_DTYPE = 'VIRTUAL' +VIRTUALFIELD_DTYPE = "VIRTUAL" + def set_default_empty(array): if array.dtype.names is None: # scalar dtype, just set - if array.dtype.str[1] == 'i': + if array.dtype.str[1] == "i": # integer, set to ID_NOT_SET array[:] = ID_NOT_SET - elif array.dtype.str[1] == 'O': + elif array.dtype.str[1] == "O": # object, set to EMPTY_OBJECT array[:] = EMPTY_OBJECT else: for name in array.dtype.names: set_default_empty(array[name]) + def default_empty(shape, dtype): - """Numpy's empty array can have random values in it. To prevent that, we + """ + Numpy's empty array can have random values in it. To prevent that, we define here a default emtpy array. This default empty is a numpy.zeros array, except that objects are set to None, and all ints to ID_NOT_SET. """ @@ -82,15 +96,18 @@ def default_empty(shape, dtype): set_default_empty(default) return default + # set default data types _default_types_status = { - 'default_strlen': 50, - 'ilwd_as_int': True, - 'lstring_as_obj': False + "default_strlen": 50, + "ilwd_as_int": True, + "lstring_as_obj": False, } + def lstring_as_obj(true_or_false=None): - """Toggles whether lstrings should be treated as strings or as objects. + """ + Toggles whether lstrings should be treated as strings or as objects. When FieldArrays is first loaded, the default is True. Parameters @@ -119,36 +136,44 @@ def lstring_as_obj(true_or_false=None): FieldArray([('0.0',), ('0.0',), ('0.0',), ('0.0',), ('0.0',), ('0.0',), ('0.0',), ('0.0',), ('0.0',), ('0.0',)], dtype=[('foo', 'S50')]) + """ if true_or_false is not None: - _default_types_status['lstring_as_obj'] = true_or_false + _default_types_status["lstring_as_obj"] = true_or_false # update the sctypeDict - numpy.sctypeDict[u'lstring'] = numpy.object_ \ - if _default_types_status['lstring_as_obj'] \ - else 'S%i' % _default_types_status['default_strlen'] - return _default_types_status['lstring_as_obj'] + numpy.sctypeDict["lstring"] = ( + numpy.object_ + if _default_types_status["lstring_as_obj"] + else "S%i" % _default_types_status["default_strlen"] + ) + return _default_types_status["lstring_as_obj"] + def ilwd_as_int(true_or_false=None): - """Similar to lstring_as_obj, sets whether or not ilwd:chars should be + """ + Similar to lstring_as_obj, sets whether or not ilwd:chars should be treated as strings or as ints. Default is True. """ if true_or_false is not None: - _default_types_status['ilwd_as_int'] = true_or_false - numpy.sctypeDict[u'ilwd:char'] = int \ - if _default_types_status['ilwd_as_int'] \ - else 'S%i' % default_strlen - return _default_types_status['ilwd_as_int'] + _default_types_status["ilwd_as_int"] = true_or_false + numpy.sctypeDict["ilwd:char"] = ( + int if _default_types_status["ilwd_as_int"] else "S%i" % default_strlen + ) + return _default_types_status["ilwd_as_int"] + def default_strlen(strlen=None): - """Sets the default string length for lstring and ilwd:char, if they are + """ + Sets the default string length for lstring and ilwd:char, if they are treated as strings. Default is 50. """ if strlen is not None: - _default_types_status['default_strlen'] = strlen + _default_types_status["default_strlen"] = strlen # update the sctypeDicts as needed - lstring_as_obj(_default_types_status['lstring_as_obj']) - ilwd_as_int(_default_types_status['ilwd_as_int']) - return _default_types_status['default_strlen'] + lstring_as_obj(_default_types_status["lstring_as_obj"]) + ilwd_as_int(_default_types_status["ilwd_as_int"]) + return _default_types_status["default_strlen"] + # set the defaults lstring_as_obj(True) @@ -169,40 +194,52 @@ def default_strlen(strlen=None): # # this parser will pull out sufields as separate identifiers from their parent # field; e.g., foo.bar --> ['foo', 'bar'] -_pyparser = re.compile(r'(?P[\w_][\w\d_]*)') +_pyparser = re.compile(r"(?P[\w_][\w\d_]*)") # this parser treats subfields as one identifier with their parent field; # e.g., foo.bar --> ['foo.bar'] -_fieldparser = re.compile(r'(?P[\w_][.\w\d_]*)') +_fieldparser = re.compile(r"(?P[\w_][.\w\d_]*)") + + def get_vars_from_arg(arg): - """Given a python string, gets the names of any identifiers use in it. + """ + Given a python string, gets the names of any identifiers use in it. For example, if ``arg = '3*narf/foo.bar'``, this will return ``set(['narf', 'foo', 'bar'])``. """ return set(_pyparser.findall(arg)) + def get_fields_from_arg(arg): - """Given a python string, gets FieldArray field names used in it. This + """ + Given a python string, gets FieldArray field names used in it. This differs from get_vars_from_arg in that any identifier with a '.' in it will be treated as one identifier. For example, if ``arg = '3*narf/foo.bar'``, this will return ``set(['narf', 'foo.bar'])``. """ return set(_fieldparser.findall(arg)) + # this parser looks for fields inside a class method function. This is done by # looking for variables that start with self.{x} or self["{x}"]; e.g., # self.a.b*3 + self.c, self['a.b']*3 + self.c, self.a.b*3 + self["c"], all # return set('a.b', 'c'). _instfieldparser = re.compile( - r'''self(?:\.|(?:\[['"]))(?P[\w_][.\w\d_]*)''') + r"""self(?:\.|(?:\[['"]))(?P[\w_][.\w\d_]*)""" +) + + def get_instance_fields_from_arg(arg): - """Given a python string definining a method function on an instance of an + """ + Given a python string definining a method function on an instance of an FieldArray, returns the field names used in it. This differs from get_fields_from_arg in that it looks for variables that start with 'self'. """ return set(_instfieldparser.findall(arg)) + def get_needed_fieldnames(arr, names): - """Given a FieldArray-like array and a list of names, determines what + """ + Given a FieldArray-like array and a list of names, determines what fields are needed from the array so that using the names does not result in an error. @@ -221,6 +258,7 @@ def get_needed_fieldnames(arr, names): ------- set The set of the fields needed to evaluate the names. + """ fieldnames = set([]) # we'll need the class that the array is an instance of to evaluate some @@ -266,13 +304,14 @@ def get_needed_fieldnames(arr, names): def get_dtype_descr(dtype): - """Numpy's ``dtype.descr`` will return empty void fields if a dtype has + """ + Numpy's ``dtype.descr`` will return empty void fields if a dtype has offsets specified. This function tries to fix that by not including fields that have no names and are void types. """ dts = [] for dt in dtype.descr: - if (dt[0] == '' and dt[1][1] == 'V'): + if dt[0] == "" and dt[1][1] == "V": continue # Downstream codes (numpy, etc) can't handle metadata in dtype @@ -284,7 +323,8 @@ def get_dtype_descr(dtype): def combine_fields(dtypes): - """Combines the fields in the list of given dtypes into a single dtype. + """ + Combines the fields in the list of given dtypes into a single dtype. Parameters ---------- @@ -295,13 +335,13 @@ def combine_fields(dtypes): ------- numpy.dtype A new dtype combining the fields in the list of dtypes. + """ if not isinstance(dtypes, list): dtypes = [dtypes] # Note: incase any of the dtypes have offsets, we won't include any fields # that have no names and are void - new_dt = numpy.dtype([dt for dtype in dtypes \ - for dt in get_dtype_descr(dtype)]) + new_dt = numpy.dtype([dt for dtype in dtypes for dt in get_dtype_descr(dtype)]) return new_dt @@ -309,12 +349,15 @@ def _ensure_array_list(arrays): """Ensures that every element in a list is an instance of a numpy array.""" # Note: the isinstance test is needed below so that instances of FieldArray # are not converted to numpy arrays - return [numpy.array(arr, ndmin=1) if not isinstance(arr, numpy.ndarray) - else arr for arr in arrays] + return [ + numpy.array(arr, ndmin=1) if not isinstance(arr, numpy.ndarray) else arr + for arr in arrays + ] def merge_arrays(merge_list, names=None, flatten=True, outtype=None): - """Merges the given arrays into a single array. The arrays must all have + """ + Merges the given arrays into a single array. The arrays must all have the same shape. If one or more of the given arrays has multiple fields, all of the fields will be included as separate fields in the new array. @@ -344,17 +387,20 @@ def merge_arrays(merge_list, names=None, flatten=True, outtype=None): new array : {numpy.ndarray | outtype} A new array with all of the fields in all of the arrays merged into a single array. + """ # make sure everything in merge_list is an array merge_list = _ensure_array_list(merge_list) if not all(merge_list[0].shape == arr.shape for arr in merge_list): - raise ValueError("all of the arrays in merge_list must have the " + - "same shape") + raise ValueError( + "all of the arrays in merge_list must have the " + "same shape" + ) if flatten: new_dt = combine_fields([arr.dtype for arr in merge_list]) else: - new_dt = numpy.dtype([('f%i' %ii, arr.dtype.descr) \ - for ii,arr in enumerate(merge_list)]) + new_dt = numpy.dtype( + [("f%i" % ii, arr.dtype.descr) for ii, arr in enumerate(merge_list)] + ) new_arr = merge_list[0].__class__(merge_list[0].shape, dtype=new_dt) # ii is a counter to keep track of which fields from the new array # go with which arrays in merge list @@ -375,8 +421,10 @@ def merge_arrays(merge_list, names=None, flatten=True, outtype=None): new_arr = new_arr.view(type=outtype) return new_arr + def add_fields(input_array, arrays, names=None, assubarray=False): - """Adds the given array(s) as new field(s) to the given input array. + """ + Adds the given array(s) as new field(s) to the given input array. Returns a new instance of the input_array with the new fields added. Parameters @@ -402,6 +450,7 @@ def add_fields(input_array, arrays, names=None, assubarray=False): ------- new_array : new instance of `input_array` A copy of the `input_array` with the desired fields added. + """ if not isinstance(arrays, list): arrays = [arrays] @@ -413,17 +462,18 @@ def add_fields(input_array, arrays, names=None, assubarray=False): names = [names] # check if any names are subarray names; if so, we have to add them # separately - subarray_names = [name for name in names if len(name.split('.')) > 1] + subarray_names = [name for name in names if len(name.split(".")) > 1] else: subarray_names = [] if any(subarray_names): - subarrays = [arrays[ii] for ii,name in enumerate(names) \ - if name in subarray_names] + subarrays = [ + arrays[ii] for ii, name in enumerate(names) if name in subarray_names + ] # group together by subarray groups = {} - for name,arr in zip(subarray_names, subarrays): - key = name.split('.')[0] - subkey = '.'.join(name.split('.')[1:]) + for name, arr in zip(subarray_names, subarrays): + key = name.split(".")[0] + subkey = ".".join(name.split(".")[1:]) try: groups[key].append((subkey, arr)) except KeyError: @@ -440,20 +490,21 @@ def add_fields(input_array, arrays, names=None, assubarray=False): # get the data new_subarray = input_array[group_name] # add the new fields to the subarray - new_subarray = add_fields(new_subarray, thisdict.values(), - thisdict.keys()) + new_subarray = add_fields( + new_subarray, thisdict.values(), thisdict.keys() + ) # remove the original from the input array input_array = input_array.without_fields(group_name) else: new_subarray = thisdict.values() # add the new subarray to input_array as a subarray - input_array = add_fields(input_array, new_subarray, - names=group_name, assubarray=True) + input_array = add_fields( + input_array, new_subarray, names=group_name, assubarray=True + ) # set the subarray names input_array[group_name].dtype.names = thisdict.keys() # remove the subarray names from names - keep_idx = [ii for ii,name in enumerate(names) \ - if name not in subarray_names] + keep_idx = [ii for ii, name in enumerate(names) if name not in subarray_names] names = [names[ii] for ii in keep_idx] # if there's nothing left, just return if names == []: @@ -465,16 +516,16 @@ def add_fields(input_array, arrays, names=None, assubarray=False): if len(arrays) > 1: arrays = [merge_arrays(arrays, flatten=True)] # now merge all the fields as a single subarray - merged_arr = numpy.empty(len(arrays[0]), - dtype=[('f0', arrays[0].dtype.descr)]) - merged_arr['f0'] = arrays[0] + merged_arr = numpy.empty(len(arrays[0]), dtype=[("f0", arrays[0].dtype.descr)]) + merged_arr["f0"] = arrays[0] arrays = [merged_arr] merge_list = [input_array] + arrays if names is not None: names = list(input_array.dtype.names) + names # merge into a single array - return merge_arrays(merge_list, names=names, flatten=True, - outtype=type(input_array)) + return merge_arrays( + merge_list, names=names, flatten=True, outtype=type(input_array) + ) # @@ -487,11 +538,13 @@ def add_fields(input_array, arrays, names=None, assubarray=False): # We'll include functions in various pycbc modules in FieldArray's function # library. All modules used must have an __all__ list defined. -_modules_for_functionlib = [conversions, coordinates, cosmology, - population_models] -_fieldarray_functionlib = {_funcname : getattr(_mod, _funcname) - for _mod in _modules_for_functionlib - for _funcname in getattr(_mod, '__all__')} +_modules_for_functionlib = [conversions, coordinates, cosmology, population_models] +_fieldarray_functionlib = { + _funcname: getattr(_mod, _funcname) + for _mod in _modules_for_functionlib + for _funcname in _mod.__all__ +} + class FieldArray(numpy.recarray): """ @@ -762,19 +815,17 @@ def bar(self): arrays. """ + _virtualfields = [] _functionlib = _fieldarray_functionlib - __persistent_attributes__ = ['name', '_virtualfields', '_functionlib'] + __persistent_attributes__ = ["name", "_virtualfields", "_functionlib"] def __new__(cls, shape, name=None, zero=True, **kwargs): - """Initializes a new empty array. - """ - obj = super(FieldArray, cls).__new__(cls, shape, **kwargs).view( - type=cls) + """Initializes a new empty array.""" + obj = super().__new__(cls, shape, **kwargs).view(type=cls) obj.name = name - obj.__persistent_attributes__ = [a - for a in cls.__persistent_attributes__] - obj._functionlib = {f: func for f,func in cls._functionlib.items()} + obj.__persistent_attributes__ = [a for a in cls.__persistent_attributes__] + obj._functionlib = {f: func for f, func in cls._functionlib.items()} obj._virtualfields = [f for f in cls._virtualfields] # zero out the array if desired if zero: @@ -783,7 +834,8 @@ def __new__(cls, shape, name=None, zero=True, **kwargs): return obj def __array_finalize__(self, obj): - """Default values are set here. + """ + Default values are set here. See for details. @@ -797,22 +849,24 @@ def __array_finalize__(self, obj): pass def __copy_attributes__(self, other, default=None): - """Copies the values of all of the attributes listed in + """ + Copies the values of all of the attributes listed in `self.__persistent_attributes__` to other. """ - [setattr(other, attr, copy.deepcopy(getattr(self, attr, default))) \ - for attr in self.__persistent_attributes__] + [ + setattr(other, attr, copy.deepcopy(getattr(self, attr, default))) + for attr in self.__persistent_attributes__ + ] def __getattribute__(self, attr, no_fallback=False): - """Allows fields to be accessed as attributes. - """ + """Allows fields to be accessed as attributes.""" # first try to get the attribute try: return numpy.ndarray.__getattribute__(self, attr) except AttributeError as e: # don't try getitem, which might get back here if no_fallback: - raise(e) + raise (e) # might be a field, try to retrive it using getitem if attr in self.fields: @@ -821,28 +875,28 @@ def __getattribute__(self, attr, no_fallback=False): raise AttributeError(e) def __setitem__(self, item, values): - """Wrap's recarray's setitem to allow attribute-like indexing when + """ + Wrap's recarray's setitem to allow attribute-like indexing when setting values. """ if type(item) is int and type(values) is numpy.ndarray: # numpy >=1.14 only accepts tuples values = tuple(values) try: - return super(FieldArray, self).__setitem__(item, values) + return super().__setitem__(item, values) except ValueError: # we'll get a ValueError if a subarray is being referenced using # '.'; so we'll try to parse it out here - fields = item.split('.') + fields = item.split(".") if len(fields) > 1: for field in fields[:-1]: self = self[field] item = fields[-1] # now try again - return super(FieldArray, self).__setitem__(item, values) + return super().__setitem__(item, values) def __getbaseitem__(self, item): - """Gets an item assuming item is either an index or a fieldname. - """ + """Gets an item assuming item is either an index or a fieldname.""" # We cast to a ndarray to avoid calling array_finalize, which can be # slow out = self.view(numpy.ndarray)[item] @@ -851,27 +905,26 @@ def __getbaseitem__(self, item): return out # if there are fields, but only a single entry, we'd just get a # record by casting to self, so just cast immediately to recarray - elif out.ndim == 0: + if out.ndim == 0: return out.view(numpy.recarray) # otherwise, cast back to an instance of self - else: - return out.view(type(self)) + return out.view(type(self)) def __getsubitem__(self, item): - """Gets a subfield using `field.subfield` notation. - """ + """Gets a subfield using `field.subfield` notation.""" try: return self.__getbaseitem__(item) except ValueError as err: - subitems = item.split('.') + subitems = item.split(".") if len(subitems) > 1: - return self.__getbaseitem__(subitems[0] - ).__getsubitem__('.'.join(subitems[1:])) - else: - raise ValueError(err) + return self.__getbaseitem__(subitems[0]).__getsubitem__( + ".".join(subitems[1:]) + ) + raise ValueError(err) def __getitem__(self, item): - """Wraps recarray's `__getitem__` so that math functions on fields and + """ + Wraps recarray's `__getitem__` so that math functions on fields and attributes can be retrieved. Any function in numpy's library may be used. """ @@ -881,11 +934,11 @@ def __getitem__(self, item): # # arg isn't a simple argument of row, so we'll have to eval it # - if not hasattr(self, '_code_cache'): + if not hasattr(self, "_code_cache"): self._code_cache = {} if item not in self._code_cache: - code = compile(item, '', 'eval') + code = compile(item, "", "eval") # get the function library item_dict = dict(_numpy_function_lib.items()) @@ -934,8 +987,9 @@ def __contains__(self, field): """Returns True if the given field name is in self's fields.""" return field in self.fields - def sort(self, axis=-1, kind='quicksort', order=None): - """Sort an array, in-place. + def sort(self, axis=-1, kind="quicksort", order=None): + """ + Sort an array, in-place. This function extends the standard numpy record array in-place sort to allow the basic use of Field array virtual fields. Only a single @@ -952,6 +1006,7 @@ def sort(self, axis=-1, kind='quicksort', order=None): When `a` is an array with fields defined, this argument specifies which fields to compare first, second, etc. Not all fields need be specified. + """ try: numpy.recarray.sort(self, axis=axis, kind=kind, order=order) @@ -961,7 +1016,8 @@ def sort(self, axis=-1, kind='quicksort', order=None): self[:] = self[numpy.argsort(self[order])] def addattr(self, attrname, value=None, persistent=True): - """Adds an attribute to self. If persistent is True, the attribute will + """ + Adds an attribute to self. If persistent is True, the attribute will be made a persistent attribute. Persistent attributes are copied whenever a view or copy of this array is created. Otherwise, new views or copies of this will not have the attribute. @@ -972,17 +1028,19 @@ def addattr(self, attrname, value=None, persistent=True): self.__persistent_attributes__.append(attrname) def add_methods(self, names, methods): - """Adds the given method(s) as instance method(s) of self. The + """ + Adds the given method(s) as instance method(s) of self. The method(s) must take `self` as a first argument. """ if isinstance(names, str): names = [names] methods = [methods] - for name,method in zip(names, methods): + for name, method in zip(names, methods): setattr(self, name, types.MethodType(method, self)) def add_properties(self, names, methods): - """Returns a view of self with the given methods added as properties. + """ + Returns a view of self with the given methods added as properties. From: . """ @@ -991,12 +1049,13 @@ def add_properties(self, names, methods): if isinstance(names, str): names = [names] methods = [methods] - for name,method in zip(names, methods): + for name, method in zip(names, methods): setattr(cls, name, property(method)) return self.view(type=cls) def add_virtualfields(self, names, methods): - """Returns a view of this array with the given methods added as virtual + """ + Returns a view of this array with the given methods added as virtual fields. Specifically, the given methods are added using add_properties and their names are added to the list of virtual fields. Virtual fields are properties that are assumed to operate on one or more of self's @@ -1012,7 +1071,8 @@ def add_virtualfields(self, names, methods): return out def add_functions(self, names, functions): - """Adds the given functions to the function library. + """ + Adds the given functions to the function library. Functions are added to this instance of the array; all copies of and slices of this array will also have the new functions included. @@ -1023,17 +1083,20 @@ def add_functions(self, names, functions): Name or list of names of the functions. functions : (list of) function(s) The function(s) to call. + """ if isinstance(names, str): names = [names] functions = [functions] if len(functions) != len(names): - raise ValueError("number of provided names must be same as number " - "of functions") + raise ValueError( + "number of provided names must be same as number of functions" + ) self._functionlib.update(dict(zip(names, functions))) def del_functions(self, names): - """Removes the specified function names from the function library. + """ + Removes the specified function names from the function library. Functions are removed from this instance of the array; all copies and slices of this array will also have the functions removed. @@ -1042,6 +1105,7 @@ def del_functions(self, names): ---------- names : (list of) string(s) Name or list of names of the functions to remove. + """ if isinstance(names, str): names = [names] @@ -1050,7 +1114,8 @@ def del_functions(self, names): @classmethod def from_arrays(cls, arrays, name=None, **kwargs): - """Creates a new instance of self from the given (list of) array(s). + """ + Creates a new instance of self from the given (list of) array(s). This is done by calling numpy.rec.fromarrays on the given arrays with the given kwargs. The type of the returned array is cast to this class, and the name (if provided) is set. @@ -1069,6 +1134,7 @@ def from_arrays(cls, arrays, name=None, **kwargs): array : instance of this class An array that is an instance of this class in which the field data is from the given array(s). + """ obj = numpy.rec.fromarrays(arrays, **kwargs).view(type=cls) obj.name = name @@ -1076,7 +1142,8 @@ def from_arrays(cls, arrays, name=None, **kwargs): @classmethod def from_records(cls, records, name=None, **kwargs): - """Creates a new instance of self from the given (list of) record(s). + """ + Creates a new instance of self from the given (list of) record(s). A "record" is a tuple in which each element is the value of one field in the resulting record array. This is done by calling @@ -1100,15 +1167,16 @@ def from_records(cls, records, name=None, **kwargs): array : instance of this class An array that is an instance of this class in which the field data is from the given record(s). + """ - obj = numpy.rec.fromrecords(records, **kwargs).view( - type=cls) + obj = numpy.rec.fromrecords(records, **kwargs).view(type=cls) obj.name = name return obj @classmethod def from_kwargs(cls, **kwargs): - """Creates a new instance of self from the given keyword arguments. + """ + Creates a new instance of self from the given keyword arguments. Each argument will correspond to a field in the returned array, with the name of the field given by the keyword, and the value(s) whatever the keyword was set to. Each keyword may be set to a single value or @@ -1128,10 +1196,11 @@ def from_kwargs(cls, **kwargs): >>> a = FieldArray.from_kwargs(mass1=1.1, mass2=2.) >>> a.mass1, a.mass2 (array([ 1.1]), array([ 2.])) + """ arrays = [] names = [] - for p,vals in kwargs.items(): + for p, vals in kwargs.items(): if not isinstance(vals, numpy.ndarray): if not isinstance(vals, list): vals = [vals] @@ -1140,10 +1209,10 @@ def from_kwargs(cls, **kwargs): names.append(p) return cls.from_arrays(arrays, names=names) - @classmethod def from_ligolw_table(cls, table, columns=None, cast_to_dtypes=None): - """Converts the given ligolw table into an FieldArray. The `tableName` + """ + Converts the given ligolw table into an FieldArray. The `tableName` attribute is copied to the array's `name`. Parameters @@ -1165,8 +1234,9 @@ def from_ligolw_table(cls, table, columns=None, cast_to_dtypes=None): ------- array : FieldArray The input table as an FieldArray. + """ - name = table.tableName.split(':')[0] + name = table.tableName.split(":")[0] if columns is None: # get all the columns columns = table.validcolumns @@ -1182,23 +1252,26 @@ def from_ligolw_table(cls, table, columns=None, cast_to_dtypes=None): else: dtype = list(columns.items()) # get the values - if _default_types_status['ilwd_as_int']: + if _default_types_status["ilwd_as_int"]: # columns like `process:process_id` have corresponding attributes # with names that are only the part after the colon, so we split - input_array = \ - [tuple(getattr(row, col.split(':')[-1]) if dt != 'ilwd:char' - else int(getattr(row, col)) - for col,dt in columns.items()) - for row in table] + input_array = [ + tuple( + getattr(row, col.split(":")[-1]) + if dt != "ilwd:char" + else int(getattr(row, col)) + for col, dt in columns.items() + ) + for row in table + ] else: - input_array = \ - [tuple(getattr(row, col) for col in columns) for row in table] + input_array = [tuple(getattr(row, col) for col in columns) for row in table] # return the values as an instance of cls - return cls.from_records(input_array, dtype=dtype, - name=name) + return cls.from_records(input_array, dtype=dtype, name=name) def to_array(self, fields=None, axis=0): - """Returns an `numpy.ndarray` of self in which the fields are included + """ + Returns an `numpy.ndarray` of self in which the fields are included as an extra dimension. Parameters @@ -1217,6 +1290,7 @@ def to_array(self, fields=None, axis=0): ------- numpy.ndarray The desired fields as a numpy array. + """ if fields is None: fields = self.fieldnames @@ -1226,15 +1300,15 @@ def to_array(self, fields=None, axis=0): @property def fieldnames(self): - """Returns a tuple listing the field names in self. Equivalent to + """ + Returns a tuple listing the field names in self. Equivalent to `array.dtype.names`, where `array` is self. """ return self.dtype.names @property def virtualfields(self): - """Returns a tuple listing the names of virtual fields in self. - """ + """Returns a tuple listing the names of virtual fields in self.""" if self._virtualfields is None: vfs = tuple() else: @@ -1243,20 +1317,24 @@ def virtualfields(self): @property def functionlib(self): - """Returns the library of functions that are available when calling + """ + Returns the library of functions that are available when calling items. """ return self._functionlib @property def fields(self): - """Returns a tuple listing the names of fields and virtual fields in - self.""" + """ + Returns a tuple listing the names of fields and virtual fields in + self. + """ return tuple(list(self.fieldnames) + list(self.virtualfields)) @property def aliases(self): - """Returns a dictionary of the aliases, or "titles", of the field names + """ + Returns a dictionary of the aliases, or "titles", of the field names in self. An alias can be specified by passing a tuple in the name part of the dtype. For example, if an array is created with ``dtype=[(('foo', 'bar'), float)]``, the array will have a field @@ -1296,13 +1374,15 @@ def add_fields(self, arrays, names=None, assubarray=False): ------- new_array : new instance of this array A copy of this array with the desired fields added. + """ newself = add_fields(self, arrays, names=names, assubarray=assubarray) self.__copy_attributes__(newself) return newself def parse_boolargs(self, args): - """Returns an array populated by given values, with the indices of + """ + Returns an array populated by given values, with the indices of those values dependent on given boolen tests on self. The given `args` should be a list of tuples, with the first element the @@ -1409,11 +1489,12 @@ def parse_boolargs(self, args): out = numpy.zeros(self.size, dtype=outdtype) mask = numpy.zeros(self.size, dtype=bool) leftovers = numpy.ones(self.size, dtype=bool) - for ii,(boolarg,val) in enumerate(zip(bool_args, return_vals)): - if boolarg is None or boolarg == '' or boolarg.lower() == 'else': - if ii+1 != len(bool_args): - raise ValueError("only the last item may not provide " - "any boolean arguments") + for ii, (boolarg, val) in enumerate(zip(bool_args, return_vals)): + if boolarg is None or boolarg == "" or boolarg.lower() == "else": + if ii + 1 != len(bool_args): + raise ValueError( + "only the last item may not provide any boolean arguments" + ) mask = leftovers else: mask = leftovers & self[boolarg] @@ -1422,7 +1503,8 @@ def parse_boolargs(self, args): return out, numpy.where(leftovers)[0] def append(self, other): - """Appends another array to this array. + """ + Appends another array to this array. The returned array will have all of the class methods and virutal fields of this array, including any that were added using `add_method` @@ -1447,6 +1529,7 @@ def append(self, other): An array with others values appended to this array's values. The returned array is an instance of the same class as this array, including all methods and virtual fields. + """ try: return numpy.append(self, other).view(type=self.__class__) @@ -1454,13 +1537,15 @@ def append(self, other): # see if the dtype error was due to string fields having different # lengths; if so, we'll make the joint field the larger of the # two - str_fields = [name for name in self.fieldnames - if _isstring(self.dtype[name])] + str_fields = [ + name for name in self.fieldnames if _isstring(self.dtype[name]) + ] # get the larger of the two new_strlens = dict( - [[name, - max(self.dtype[name].itemsize, other.dtype[name].itemsize)] - for name in str_fields] + [ + [name, max(self.dtype[name].itemsize, other.dtype[name].itemsize)] + for name in str_fields + ] ) # cast both to the new string lengths new_dt = [] @@ -1470,14 +1555,14 @@ def append(self, other): dt = (name, self.dtype[name].type, new_strlens[name]) new_dt.append(dt) new_dt = numpy.dtype(new_dt) - return numpy.append( - self.astype(new_dt), - other.astype(new_dt) - ).view(type=self.__class__) + return numpy.append(self.astype(new_dt), other.astype(new_dt)).view( + type=self.__class__ + ) @classmethod def parse_parameters(cls, parameters, possible_fields): - """Parses a list of parameters to get the list of fields needed in + """ + Parses a list of parameters to get the list of fields needed in order to evaluate those parameters. Parameters @@ -1493,43 +1578,45 @@ def parse_parameters(cls, parameters, possible_fields): list : The list of names of the fields that are needed in order to evaluate the given parameters. + """ if isinstance(possible_fields, str): possible_fields = [possible_fields] possible_fields = list(map(str, possible_fields)) # we'll just use float as the dtype, as we just need this for names - arr = cls(1, dtype=list(zip(possible_fields, - len(possible_fields)*[float]))) + arr = cls(1, dtype=list(zip(possible_fields, len(possible_fields) * [float]))) # try to perserve order return list(get_needed_fieldnames(arr, parameters)) + def _isstring(dtype): - """Given a numpy dtype, determines whether it is a string. Returns True + """ + Given a numpy dtype, determines whether it is a string. Returns True if the dtype is string or unicode. """ return dtype.type == numpy.str_ or dtype.type == numpy.bytes_ def aliases_from_fields(fields): - """Given a dictionary of fields, will return a dictionary mapping the + """ + Given a dictionary of fields, will return a dictionary mapping the aliases to the names. """ return dict(c for c in fields if isinstance(c, tuple)) def fields_from_names(fields, names=None): - """Given a dictionary of fields and a list of names, will return a + """ + Given a dictionary of fields and a list of names, will return a dictionary consisting of the fields specified by names. Names can be either the names of fields, or their aliases. """ - if names is None: return fields if isinstance(names, str): names = [names] aliases_to_names = aliases_from_fields(fields) - names_to_aliases = dict(zip(aliases_to_names.values(), - aliases_to_names.keys())) + names_to_aliases = dict(zip(aliases_to_names.values(), aliases_to_names.keys())) outfields = {} for name in names: try: @@ -1540,7 +1627,7 @@ def fields_from_names(fields, names=None): elif name in names_to_aliases: key = (names_to_aliases[name], name) else: - raise KeyError('default fields has no field %s' % name) + raise KeyError("default fields has no field %s" % name) outfields[key] = fields[key] return outfields @@ -1553,6 +1640,7 @@ def fields_from_names(fields, names=None): # ============================================================================= # + class _FieldArrayWithDefaults(FieldArray): """ Subclasses FieldArray, adding class attribute ``_staticfields``, and @@ -1593,9 +1681,11 @@ class method ``default_fields``. The ``_staticfields`` should be a """ _staticfields = {} + @classmethod def default_fields(cls, include_virtual=True, **kwargs): - """The default fields and their dtypes. By default, this returns + """ + The default fields and their dtypes. By default, this returns whatever the class's ``_staticfields`` and ``_virtualfields`` is set to as a dictionary of fieldname, dtype (the dtype of virtualfields is given by VIRTUALFIELD_DTYPE). This function should be overridden by @@ -1605,24 +1695,24 @@ def default_fields(cls, include_virtual=True, **kwargs): """ output = cls._staticfields.copy() if include_virtual: - output.update({name: VIRTUALFIELD_DTYPE - for name in cls._virtualfields}) + output.update(dict.fromkeys(cls._virtualfields, VIRTUALFIELD_DTYPE)) return output - def __new__(cls, shape, name=None, additional_fields=None, - field_kwargs=None, **kwargs): - """The ``additional_fields`` should be specified in the same way as + def __new__( + cls, shape, name=None, additional_fields=None, field_kwargs=None, **kwargs + ): + """ + The ``additional_fields`` should be specified in the same way as ``dtype`` is normally given to FieldArray. The ``field_kwargs`` are passed to the class's default_fields method as keyword arguments. """ if field_kwargs is None: field_kwargs = {} - if 'names' in kwargs and 'dtype' in kwargs: + if "names" in kwargs and "dtype" in kwargs: raise ValueError("Please provide names or dtype, not both") - default_fields = cls.default_fields(include_virtual=False, - **field_kwargs) - if 'names' in kwargs: - names = kwargs.pop('names') + default_fields = cls.default_fields(include_virtual=False, **field_kwargs) + if "names" in kwargs: + names = kwargs.pop("names") if isinstance(names, str): names = [names] # evaluate the names to figure out what base fields are needed @@ -1631,23 +1721,23 @@ def __new__(cls, shape, name=None, additional_fields=None, # block of code is skipped) arr = cls(1, field_kwargs=field_kwargs) # try to perserve order - sortdict = dict([[nm, ii] for ii,nm in enumerate(names)]) + sortdict = dict([[nm, ii] for ii, nm in enumerate(names)]) names = list(get_needed_fieldnames(arr, names)) - names.sort(key=lambda x: sortdict[x] if x in sortdict - else len(names)) + names.sort(key=lambda x: sortdict[x] if x in sortdict else len(names)) # add the fields as the dtype argument for initializing - kwargs['dtype'] = [(fld, default_fields[fld]) for fld in names] - if 'dtype' not in kwargs: - kwargs['dtype'] = list(default_fields.items()) + kwargs["dtype"] = [(fld, default_fields[fld]) for fld in names] + if "dtype" not in kwargs: + kwargs["dtype"] = list(default_fields.items()) # add the additional fields if additional_fields is not None: if not isinstance(additional_fields, list): additional_fields = [additional_fields] - if not isinstance(kwargs['dtype'], list): - kwargs['dtype'] = [kwargs['dtype']] - kwargs['dtype'] += additional_fields - return super(_FieldArrayWithDefaults, cls).__new__(cls, shape, - name=name, **kwargs) + if not isinstance(kwargs["dtype"], list): + kwargs["dtype"] = [kwargs["dtype"]] + kwargs["dtype"] += additional_fields + return super().__new__( + cls, shape, name=name, **kwargs + ) def add_default_fields(self, names, **kwargs): """ @@ -1665,6 +1755,7 @@ def add_default_fields(self, names, **kwargs): ------- new array : instance of this array A copy of this array with the field added. + """ if isinstance(names, str): names = [names] @@ -1672,21 +1763,21 @@ def add_default_fields(self, names, **kwargs): # parse out any virtual fields arr = self.__class__(1, field_kwargs=kwargs) # try to perserve order - sortdict = dict([[nm, ii] for ii,nm in enumerate(names)]) + sortdict = dict([[nm, ii] for ii, nm in enumerate(names)]) names = list(get_needed_fieldnames(arr, names)) - names.sort(key=lambda x: sortdict[x] if x in sortdict - else len(names)) + names.sort(key=lambda x: sortdict[x] if x in sortdict else len(names)) fields = [(name, default_fields[name]) for name in names] arrays = [] names = [] - for name,dt in fields: + for name, dt in fields: arrays.append(default_empty(self.size, dtype=[(name, dt)])) names.append(name) return self.add_fields(arrays, names) @classmethod def parse_parameters(cls, parameters, possible_fields=None): - """Parses a list of parameters to get the list of fields needed in + """ + Parses a list of parameters to get the list of fields needed in order to evaluate those parameters. Parameters @@ -1704,16 +1795,19 @@ def parse_parameters(cls, parameters, possible_fields=None): list : The list of names of the fields that are needed in order to evaluate the given parameters. + """ if possible_fields is not None: # make sure field names are strings and not unicode - possible_fields = dict([[f, dt] - for f,dt in possible_fields.items()]) + possible_fields = dict([[f, dt] for f, dt in possible_fields.items()]) + class ModifiedArray(cls): _staticfields = possible_fields + cls = ModifiedArray return cls(1, names=parameters).fieldnames + # # ============================================================================= # @@ -1722,6 +1816,7 @@ class ModifiedArray(cls): # ============================================================================= # + class WaveformArray(_FieldArrayWithDefaults): """ A FieldArray with some default fields and properties commonly used @@ -1795,18 +1890,33 @@ class WaveformArray(_FieldArrayWithDefaults): '$d_L$ (Mpc)' """ - _staticfields = (parameters.cbc_intrinsic_params + - parameters.extrinsic_params).dtype_dict + + _staticfields = ( + parameters.cbc_intrinsic_params + parameters.extrinsic_params + ).dtype_dict _virtualfields = [ - parameters.mchirp, parameters.eta, parameters.mtotal, - parameters.q, parameters.primary_mass, parameters.secondary_mass, + parameters.mchirp, + parameters.eta, + parameters.mtotal, + parameters.q, + parameters.primary_mass, + parameters.secondary_mass, parameters.chi_eff, - parameters.spin_px, parameters.spin_py, parameters.spin_pz, - parameters.spin_sx, parameters.spin_sy, parameters.spin_sz, - parameters.spin1_a, parameters.spin1_azimuthal, parameters.spin1_polar, - parameters.spin2_a, parameters.spin2_azimuthal, parameters.spin2_polar, - parameters.remnant_mass] + parameters.spin_px, + parameters.spin_py, + parameters.spin_pz, + parameters.spin_sx, + parameters.spin_sy, + parameters.spin_sz, + parameters.spin1_a, + parameters.spin1_azimuthal, + parameters.spin1_polar, + parameters.spin2_a, + parameters.spin2_azimuthal, + parameters.spin2_polar, + parameters.remnant_mass, + ] @property def primary_mass(self): @@ -1841,89 +1951,98 @@ def mchirp(self): @property def chi_eff(self): """Returns the effective spin.""" - return conversions.chi_eff(self.mass1, self.mass2, self.spin1z, - self.spin2z) + return conversions.chi_eff(self.mass1, self.mass2, self.spin1z, self.spin2z) @property def spin_px(self): """Returns the x-component of the spin of the primary mass.""" - return conversions.primary_spin(self.mass1, self.mass2, self.spin1x, - self.spin2x) + return conversions.primary_spin( + self.mass1, self.mass2, self.spin1x, self.spin2x + ) @property def spin_py(self): """Returns the y-component of the spin of the primary mass.""" - return conversions.primary_spin(self.mass1, self.mass2, self.spin1y, - self.spin2y) + return conversions.primary_spin( + self.mass1, self.mass2, self.spin1y, self.spin2y + ) @property def spin_pz(self): """Returns the z-component of the spin of the primary mass.""" - return conversions.primary_spin(self.mass1, self.mass2, self.spin1z, - self.spin2z) + return conversions.primary_spin( + self.mass1, self.mass2, self.spin1z, self.spin2z + ) @property def spin_sx(self): """Returns the x-component of the spin of the secondary mass.""" - return conversions.secondary_spin(self.mass1, self.mass2, self.spin1x, - self.spin2x) + return conversions.secondary_spin( + self.mass1, self.mass2, self.spin1x, self.spin2x + ) @property def spin_sy(self): """Returns the y-component of the spin of the secondary mass.""" - return conversions.secondary_spin(self.mass1, self.mass2, self.spin1y, - self.spin2y) + return conversions.secondary_spin( + self.mass1, self.mass2, self.spin1y, self.spin2y + ) @property def spin_sz(self): """Returns the z-component of the spin of the secondary mass.""" - return conversions.secondary_spin(self.mass1, self.mass2, self.spin1z, - self.spin2z) + return conversions.secondary_spin( + self.mass1, self.mass2, self.spin1z, self.spin2z + ) @property def spin1_a(self): """Returns the dimensionless spin magnitude of mass 1.""" return coordinates.cartesian_to_spherical_rho( - self.spin1x, self.spin1y, self.spin1z) + self.spin1x, self.spin1y, self.spin1z + ) @property def spin1_azimuthal(self): """Returns the azimuthal spin angle of mass 1.""" - return coordinates.cartesian_to_spherical_azimuthal( - self.spin1x, self.spin1y) + return coordinates.cartesian_to_spherical_azimuthal(self.spin1x, self.spin1y) @property def spin1_polar(self): """Returns the polar spin angle of mass 1.""" return coordinates.cartesian_to_spherical_polar( - self.spin1x, self.spin1y, self.spin1z) + self.spin1x, self.spin1y, self.spin1z + ) @property def spin2_a(self): """Returns the dimensionless spin magnitude of mass 2.""" return coordinates.cartesian_to_spherical_rho( - self.spin1x, self.spin1y, self.spin1z) + self.spin1x, self.spin1y, self.spin1z + ) @property def spin2_azimuthal(self): """Returns the azimuthal spin angle of mass 2.""" - return coordinates.cartesian_to_spherical_azimuthal( - self.spin2x, self.spin2y) + return coordinates.cartesian_to_spherical_azimuthal(self.spin2x, self.spin2y) @property def spin2_polar(self): """Returns the polar spin angle of mass 2.""" return coordinates.cartesian_to_spherical_polar( - self.spin2x, self.spin2y, self.spin2z) + self.spin2x, self.spin2y, self.spin2z + ) @property def remnant_mass(self): """Returns the remnant mass for an NS-BH binary.""" return conversions.remnant_mass_from_mass1_mass2_cartesian_spin_eos( - self.mass1, self.mass2, - spin1x=self.spin1x, - spin1y=self.spin1y, - spin1z=self.spin1z) + self.mass1, + self.mass2, + spin1x=self.spin1x, + spin1y=self.spin1y, + spin1z=self.spin1z, + ) -__all__ = ['FieldArray', 'WaveformArray'] +__all__ = ["FieldArray", "WaveformArray"] diff --git a/pycbc/libutils.py b/pycbc/libutils.py index 6671ebf94c6..f805f071336 100644 --- a/pycbc/libutils.py +++ b/pycbc/libutils.py @@ -20,19 +20,19 @@ according to the paths that pkg-config specifies. """ +import ctypes +import fnmatch import importlib -import logging import inspect +import logging import os -import fnmatch -import ctypes -import sys import subprocess -from ctypes.util import find_library +import sys from collections import deque +from ctypes.util import find_library from subprocess import getoutput -logger = logging.getLogger('pycbc.libutils') +logger = logging.getLogger("pycbc.libutils") # Be careful setting the mode for opening libraries! Some libraries (e.g. # libgomp) seem to require the DEFAULT_MODE is used. Others (e.g. FFTW when @@ -42,32 +42,38 @@ def pkg_config(pkg_libraries): - """Use pkg-config to query for the location of libraries, library directories, - and header directories + """ + Use pkg-config to query for the location of libraries, library directories, + and header directories - Arguments: - pkg_libries(list): A list of packages as strings + Arguments: + pkg_libries(list): A list of packages as strings + + Returns + ------- + libraries(list), library_dirs(list), include_dirs(list) - Returns: - libraries(list), library_dirs(list), include_dirs(list) """ - libraries=[] - library_dirs=[] - include_dirs=[] + libraries = [] + library_dirs = [] + include_dirs = [] # Check that we have the packages for pkg in pkg_libraries: - if os.system('pkg-config --exists %s 2>/dev/null' % pkg) == 0: + if os.system("pkg-config --exists %s 2>/dev/null" % pkg) == 0: pass else: - print("Could not find library {0}".format(pkg)) + print(f"Could not find library {pkg}") sys.exit(1) # Get the pck-config flags - if len(pkg_libraries)>0 : + if len(pkg_libraries) > 0: # PKG_CONFIG_ALLOW_SYSTEM_CFLAGS explicitly lists system paths. # On system-wide LAL installs, this is needed for swig to find lalswig.i - for token in getoutput("PKG_CONFIG_ALLOW_SYSTEM_CFLAGS=1 pkg-config --libs --cflags %s" % ' '.join(pkg_libraries)).split(): + for token in getoutput( + "PKG_CONFIG_ALLOW_SYSTEM_CFLAGS=1 pkg-config --libs --cflags %s" + % " ".join(pkg_libraries) + ).split(): if token.startswith("-l"): libraries.append(token[2:]) elif token.startswith("-L"): @@ -77,9 +83,9 @@ def pkg_config(pkg_libraries): return libraries, library_dirs, include_dirs + def pkg_config_header_strings(pkg_libraries): - """ Returns a list of header strings that could be passed to a compiler - """ + """Returns a list of header strings that could be passed to a compiler""" _, _, header_dirs = pkg_config(pkg_libraries) header_strings = [] @@ -89,8 +95,10 @@ def pkg_config_header_strings(pkg_libraries): return header_strings + def pkg_config_check_exists(package): - return (os.system('pkg-config --exists {0} 2>/dev/null'.format(package)) == 0) + return os.system(f"pkg-config --exists {package} 2>/dev/null") == 0 + def pkg_config_libdirs(packages): """ @@ -99,7 +107,6 @@ def pkg_config_libdirs(packages): that the package may be found in the standard system locations, irrespective of pkg-config. """ - # don't try calling pkg-config if NO_PKGCONFIG is set in environment if os.environ.get("NO_PKGCONFIG", None): return [] @@ -110,24 +117,30 @@ def pkg_config_libdirs(packages): subprocess.check_call(["pkg-config", "--version"], stdout=FNULL) except: print( - "PyCBC.libutils: pkg-config call failed, " - "setting NO_PKGCONFIG=1", + "PyCBC.libutils: pkg-config call failed, setting NO_PKGCONFIG=1", file=sys.stderr, ) - os.environ['NO_PKGCONFIG'] = "1" + os.environ["NO_PKGCONFIG"] = "1" return [] # First, check that we can call pkg-config on each package in the list for pkg in packages: if not pkg_config_check_exists(pkg): - raise ValueError("Package {0} cannot be found on the pkg-config search path".format(pkg)) + raise ValueError( + f"Package {pkg} cannot be found on the pkg-config search path" + ) libdirs = [] - for token in getoutput("PKG_CONFIG_ALLOW_SYSTEM_LIBS=1 pkg-config --libs-only-L {0}".format(' '.join(packages))).split(): + for token in getoutput( + "PKG_CONFIG_ALLOW_SYSTEM_LIBS=1 pkg-config --libs-only-L {0}".format( + " ".join(packages) + ) + ).split(): if token.startswith("-L"): libdirs.append(token[2:]) return libdirs + def get_libpath_from_dirlist(libname, dirs): """ This function tries to find the architecture-independent library given by libname in the first @@ -140,27 +153,30 @@ def get_libpath_from_dirlist(libname, dirs): None is returned. """ dirqueue = deque(dirs) - while (len(dirqueue) > 0): + while len(dirqueue) > 0: nextdir = dirqueue.popleft() possible = [] # Our directory might be no good, so try/except try: for libfile in os.listdir(nextdir): - if fnmatch.fnmatch(libfile,'lib'+libname+'.so*') or \ - fnmatch.fnmatch(libfile,'lib'+libname+'.dylib*') or \ - fnmatch.fnmatch(libfile,'lib'+libname+'.*.dylib*') or \ - fnmatch.fnmatch(libfile,libname+'.dll') or \ - fnmatch.fnmatch(libfile,'cyg'+libname+'-*.dll'): + if ( + fnmatch.fnmatch(libfile, "lib" + libname + ".so*") + or fnmatch.fnmatch(libfile, "lib" + libname + ".dylib*") + or fnmatch.fnmatch(libfile, "lib" + libname + ".*.dylib*") + or fnmatch.fnmatch(libfile, libname + ".dll") + or fnmatch.fnmatch(libfile, "cyg" + libname + "-*.dll") + ): possible.append(libfile) except OSError: pass # There might be more than one library found, we want the highest-numbered - if (len(possible) > 0): + if len(possible) > 0: possible.sort() - return os.path.join(nextdir,possible[-1]) + return os.path.join(nextdir, possible[-1]) # If we get here, we didn't find it... return None + def get_ctypes_library(libname, packages, mode=DEFAULT_RTLD_MODE): """ This function takes a library name, specified in architecture-independent fashion (i.e. @@ -195,14 +211,14 @@ def get_ctypes_library(libname, packages, mode=DEFAULT_RTLD_MODE): if fullpath is None: # We got nothin' return None - else: - if mode is None: - return ctypes.CDLL(fullpath) - else: - return ctypes.CDLL(fullpath, mode=mode) + if mode is None: + return ctypes.CDLL(fullpath) + return ctypes.CDLL(fullpath, mode=mode) + def import_optional(library_name): - """ Try to import library but and return stub if not found + """ + Try to import library but and return stub if not found Parameters ---------- @@ -214,18 +230,19 @@ def import_optional(library_name): library: library or stub Either returns the library if importing is sucessful or it returns a stub which raises an import error and message when accessed. + """ try: return importlib.import_module(library_name) except ImportError: # module wasn't found so let's return a stub instead to inform # the user what has happened when they try to use related functions - class no_module(object): + class no_module: def __init__(self, library): self.library = library def __getattribute__(self, attr): - if attr == 'library': + if attr == "library": return super().__getattribute__(attr) lib = self.library @@ -233,11 +250,12 @@ def __getattribute__(self, attr): curframe = inspect.currentframe() calframe = inspect.getouterframes(curframe, 2) fun = calframe[1][3] - msg =""" The function {} tried to access - '{}' of library '{}', however, - '{}' is not currently installed. To enable this - functionality install '{}' (e.g. through pip + msg = f""" The function {fun} tried to access + '{attr}' of library '{lib}', however, + '{lib}' is not currently installed. To enable this + functionality install '{lib}' (e.g. through pip / conda / system packages / source). - """.format(fun, attr, lib, lib, lib) + """ raise ImportError(inspect.cleandoc(msg)) + return no_module(library_name) diff --git a/pycbc/live/__init__.py b/pycbc/live/__init__.py index 5a3a40c901a..bc627ccf7d4 100644 --- a/pycbc/live/__init__.py +++ b/pycbc/live/__init__.py @@ -2,6 +2,6 @@ This packages contains modules to help with pycbc live running """ -from .snr_optimizer import * from .significance_fits import * +from .snr_optimizer import * from .supervision import * diff --git a/pycbc/live/significance_fits.py b/pycbc/live/significance_fits.py index 88a53f736d6..104e69f58c8 100644 --- a/pycbc/live/significance_fits.py +++ b/pycbc/live/significance_fits.py @@ -3,10 +3,11 @@ """ import logging + import h5py import numpy -logger = logging.getLogger('pycbc.live.single_fits') +logger = logging.getLogger("pycbc.live.single_fits") def add_live_significance_trigger_pruning_options(parser): @@ -17,20 +18,18 @@ def add_live_significance_trigger_pruning_options(parser): pruning_group.add_argument( "--prune-loudest", type=int, - help="Maximum number of loudest trigger clusters to " - "remove from each bin." + help="Maximum number of loudest trigger clusters to remove from each bin.", ) pruning_group.add_argument( "--prune-window", type=float, help="Window (seconds) either side of the --prune-loudest " - "loudest triggers in each duration bin to remove." + "loudest triggers in each duration bin to remove.", ) pruning_group.add_argument( "--prune-stat-threshold", type=float, - help="Minimum statistic value to consider a " - "trigger for pruning." + help="Minimum statistic value to consider a trigger for pruning.", ) @@ -39,12 +38,13 @@ def verify_live_significance_trigger_pruning_options(args, parser): Verify options used for pruning in live singles significance fits """ # Pruning options are mutually required or not needed - prune_options = [args.prune_loudest, args.prune_window, - args.prune_stat_threshold] + prune_options = [args.prune_loudest, args.prune_window, args.prune_stat_threshold] if any(prune_options) and not all(prune_options): - parser.error("Require all or none of --prune-loudest, " - "--prune-window and --prune-stat-threshold") + parser.error( + "Require all or none of --prune-loudest, " + "--prune-window and --prune-stat-threshold" + ) def add_live_significance_duration_bin_options(parser): @@ -52,46 +52,45 @@ def add_live_significance_duration_bin_options(parser): Add options used to calculate duration bin edges in live singles significance fits """ - durbin_group = parser.add_argument_group('Duration Bins') + durbin_group = parser.add_argument_group("Duration Bins") durbin_group.add_argument( "--duration-bin-edges", - nargs='+', + nargs="+", type=float, help="Durations to use for bin edges. " - "Use if specifying exact bin edges, " - "Not compatible with --duration-bin-start, " - "--duration-bin-end and --num-duration-bins" + "Use if specifying exact bin edges, " + "Not compatible with --duration-bin-start, " + "--duration-bin-end and --num-duration-bins", ) durbin_group.add_argument( "--duration-bin-start", type=float, help="Shortest duration to use for duration bins." - "Not compatible with --duration-bins, requires " - "--duration-bin-end and --num-duration-bins." + "Not compatible with --duration-bins, requires " + "--duration-bin-end and --num-duration-bins.", ) durbin_group.add_argument( - "--duration-bin-end", type=float, - help="Longest duration to use for duration bins." + "--duration-bin-end", + type=float, + help="Longest duration to use for duration bins.", ) durbin_group.add_argument( "--duration-from-bank", - help="Path to the template bank file to get max/min " - "durations from." + help="Path to the template bank file to get max/min durations from.", ) durbin_group.add_argument( "--num-duration-bins", type=int, - help="How many template duration bins to split the bank " - "into before fitting." + help="How many template duration bins to split the bank into before fitting.", ) durbin_group.add_argument( "--duration-bin-spacing", - choices=['linear', 'log'], - default='log', + choices=["linear", "log"], + default="log", help="How to set spacing for bank split " - "if using --num-duration-bins and " - "--duration-bin-start + --duration-bin-end " - "or --duration-from-bank." + "if using --num-duration-bins and " + "--duration-bin-start + --duration-bin-end " + "or --duration-from-bank.", ) @@ -102,30 +101,41 @@ def verify_live_significance_duration_bin_options(args, parser): """ # Check the bin options if args.duration_bin_edges: - if (args.duration_bin_start or args.duration_bin_end or - args.duration_from_bank or args.num_duration_bins): - parser.error("Cannot use --duration-bin-edges with " - "--duration-bin-start, --duration-bin-end, " - "--duration-from-bank or --num-duration-bins.") + if ( + args.duration_bin_start + or args.duration_bin_end + or args.duration_from_bank + or args.num_duration_bins + ): + parser.error( + "Cannot use --duration-bin-edges with " + "--duration-bin-start, --duration-bin-end, " + "--duration-from-bank or --num-duration-bins." + ) else: if not args.num_duration_bins: - parser.error("--num-duration-bins must be set if not using " - "--duration-bin-edges.") - if not ((args.duration_bin_start and args.duration_bin_end) or - args.duration_from_bank): - parser.error("--duration-bin-start & --duration-bin-end or " - "--duration-from-bank must be set if not using " - "--duration-bin-edges.") - if args.duration_bin_end and \ - args.duration_bin_end <= args.duration_bin_start: - parser.error("--duration-bin-end must be greater than " - "--duration-bin-start, got " - f"{args.duration_bin_end} and {args.duration_bin_start}") + parser.error( + "--num-duration-bins must be set if not using --duration-bin-edges." + ) + if not ( + (args.duration_bin_start and args.duration_bin_end) + or args.duration_from_bank + ): + parser.error( + "--duration-bin-start & --duration-bin-end or " + "--duration-from-bank must be set if not using " + "--duration-bin-edges." + ) + if args.duration_bin_end and args.duration_bin_end <= args.duration_bin_start: + parser.error( + "--duration-bin-end must be greater than " + "--duration-bin-start, got " + f"{args.duration_bin_end} and {args.duration_bin_start}" + ) def duration_bins_from_cli(args): - """Create the duration bins from CLI options. - """ + """Create the duration bins from CLI options.""" if args.duration_bin_edges: # direct bin specification return numpy.array(args.duration_bin_edges) @@ -134,28 +144,22 @@ def duration_bins_from_cli(args): max_dur = args.duration_bin_end if args.duration_from_bank: # read min/max duration directly from the bank itself - with h5py.File(args.duration_from_bank, 'r') as bank_file: - temp_durs = bank_file['template_duration'][:] + with h5py.File(args.duration_from_bank, "r") as bank_file: + temp_durs = bank_file["template_duration"][:] min_dur, max_dur = min(temp_durs), max(temp_durs) - if args.duration_bin_spacing == 'log': + if args.duration_bin_spacing == "log": return numpy.logspace( - numpy.log10(min_dur), - numpy.log10(max_dur), - args.num_duration_bins + 1 - ) - if args.duration_bin_spacing == 'linear': - return numpy.linspace( - min_dur, - max_dur, - args.num_duration_bins + 1 + numpy.log10(min_dur), numpy.log10(max_dur), args.num_duration_bins + 1 ) + if args.duration_bin_spacing == "linear": + return numpy.linspace(min_dur, max_dur, args.num_duration_bins + 1) raise RuntimeError("Invalid duration bin specification") __all__ = [ - 'add_live_significance_trigger_pruning_options', - 'verify_live_significance_trigger_pruning_options', - 'add_live_significance_duration_bin_options', - 'verify_live_significance_duration_bin_options', - 'duration_bins_from_cli', + "add_live_significance_duration_bin_options", + "add_live_significance_trigger_pruning_options", + "duration_bins_from_cli", + "verify_live_significance_duration_bin_options", + "verify_live_significance_trigger_pruning_options", ] diff --git a/pycbc/live/snr_optimizer.py b/pycbc/live/snr_optimizer.py index b139fe7e36d..2898c82455a 100644 --- a/pycbc/live/snr_optimizer.py +++ b/pycbc/live/snr_optimizer.py @@ -21,26 +21,27 @@ # # ============================================================================= # -""" This module contains functions for optimizing the signal-to-noise ratio +""" +This module contains functions for optimizing the signal-to-noise ratio of triggers produced by PyCBC Live. Also contained within this module are the command line arguments required and options group for the SNR optimization. This module is primarily used in the pycbc_optimize_snr program. """ -import time import logging +import time import types + import numpy from scipy.optimize import differential_evolution, shgo -from pycbc import ( - DYN_RANGE_FAC, waveform -) -from pycbc.types import zeros + +import pycbc.conversions as cv import pycbc.waveform.bank +from pycbc import DYN_RANGE_FAC, waveform from pycbc.filter import matched_filter_core -import pycbc.conversions as cv +from pycbc.types import zeros -logger = logging.getLogger('pycbc.live.snr_optimizer') +logger = logging.getLogger("pycbc.live.snr_optimizer") try: import pyswarms as ps @@ -52,7 +53,7 @@ MIN_CPT_MASS = 0.99 # Set a large maximum total mass -MAX_MTOTAL = 500. +MAX_MTOTAL = 500.0 Nfeval = 0 start_time = time.time() @@ -67,8 +68,19 @@ def callback_func(Xi, convergence=0): Nfeval += 1 -def compute_network_snr_core(v, data, coinc_times, ifos, flen, approximant, - flow, f_end, delta_f, sample_rate, raise_err=False): +def compute_network_snr_core( + v, + data, + coinc_times, + ifos, + flen, + approximant, + flow, + f_end, + delta_f, + sample_rate, + raise_err=False, +): """ Compute network SNR as a function over mchirp, eta, and two aligned spin components, stored in that order in the sequence v. @@ -120,11 +132,18 @@ def compute_network_snr_core(v, data, coinc_times, ifos, flen, approximant, try: htilde = waveform.get_waveform_filter( - zeros(flen, dtype=numpy.complex64), - approximant=approximant, - mass1=mass1, mass2=mass2, spin1z=v[2], spin2z=v[3], - f_lower=flow, f_final=f_end, delta_f=delta_f, - delta_t=1./sample_rate, distance=distance) + zeros(flen, dtype=numpy.complex64), + approximant=approximant, + mass1=mass1, + mass2=mass2, + spin1z=v[2], + spin2z=v[3], + f_lower=flow, + f_final=f_end, + delta_f=delta_f, + delta_t=1.0 / sample_rate, + distance=distance, + ) except RuntimeError: if raise_err: raise @@ -132,14 +151,12 @@ def compute_network_snr_core(v, data, coinc_times, ifos, flen, approximant, # due to the choice of parameters and carry on return -numpy.inf, {} - if not hasattr(htilde, 'params'): - htilde.params = dict(mass1=mass1, mass2=mass2, - spin1z=v[2], spin2z=v[3]) - if not hasattr(htilde, 'end_idx'): + if not hasattr(htilde, "params"): + htilde.params = dict(mass1=mass1, mass2=mass2, spin1z=v[2], spin2z=v[3]) + if not hasattr(htilde, "end_idx"): htilde.end_idx = int(f_end / htilde.delta_f) htilde.approximant = approximant - htilde.sigmasq = types.MethodType(pycbc.waveform.bank.sigma_cached, - htilde) + htilde.sigmasq = types.MethodType(pycbc.waveform.bank.sigma_cached, htilde) htilde.min_f_lower = flow htilde.end_frequency = f_end htilde.f_lower = flow @@ -147,58 +164,49 @@ def compute_network_snr_core(v, data, coinc_times, ifos, flen, approximant, snr_series_dict = {} for ifo in ifos: sigmasq = htilde.sigmasq(data[ifo].psd) - snr, _, norm = matched_filter_core(htilde, data[ifo], - h_norm=sigmasq) + snr, _, norm = matched_filter_core(htilde, data[ifo], h_norm=sigmasq) duration = 0.095 half_dur_samples = int(snr.sample_rate * duration / 2) - onsource_idx = (float(coinc_times[ifo] - snr.start_time) * - snr.sample_rate) + onsource_idx = float(coinc_times[ifo] - snr.start_time) * snr.sample_rate onsource_idx = int(round(onsource_idx)) - onsource_slice = slice(onsource_idx - half_dur_samples, - onsource_idx + half_dur_samples + 1) + onsource_slice = slice( + onsource_idx - half_dur_samples, onsource_idx + half_dur_samples + 1 + ) snr_series = snr[onsource_slice] * norm snr_series_dict[ifo] = snr * norm - snr_series_dict['sigmasq_' + ifo] = sigmasq - network_snrsq += max(abs(snr_series._data)) ** 2. + snr_series_dict["sigmasq_" + ifo] = sigmasq + network_snrsq += max(abs(snr_series._data)) ** 2.0 - return network_snrsq ** 0.5, snr_series_dict + return network_snrsq**0.5, snr_series_dict def compute_minus_network_snr(v, *argv): if len(argv) == 1: argv = argv[0] nsnr, _ = compute_network_snr_core(v, *argv) - logger.debug('snr: %s', nsnr) + logger.debug("snr: %s", nsnr) return -nsnr def compute_minus_network_snr_pso(v, *argv, **kwargs): - argv = kwargs['args'] - nsnr_array = numpy.array([ - compute_network_snr_core(v_i, *argv)[0] - for v_i in v]) + argv = kwargs["args"] + nsnr_array = numpy.array([compute_network_snr_core(v_i, *argv)[0] for v_i in v]) return -nsnr_array def optimize_di(bounds, cli_args, extra_args, initial_point): # Convert from dict to array with parameters in a given order - bounds = numpy.array([ - bounds['mchirp'], - bounds['eta'], - bounds['spin1z'], - bounds['spin2z'] - ]) + bounds = numpy.array( + [bounds["mchirp"], bounds["eta"], bounds["spin1z"], bounds["spin2z"]] + ) # Initialize the population with random values within specified bounds population = numpy.random.uniform( - bounds[:, 0], - bounds[:, 1], - size=(int(cli_args.snr_opt_di_popsize), len(bounds)) + bounds[:, 0], bounds[:, 1], size=(int(cli_args.snr_opt_di_popsize), len(bounds)) ) if cli_args.snr_opt_include_candidate: # add the initial point to the population - population = numpy.concatenate((population[:-1], - initial_point)) - logger.debug('Initial population: %s', population) + population = numpy.concatenate((population[:-1], initial_point)) + logger.debug("Initial population: %s", population) results = differential_evolution( compute_minus_network_snr, @@ -210,81 +218,77 @@ def optimize_di(bounds, cli_args, extra_args, initial_point): recombination=0.7, callback=callback_func, args=extra_args, - init=population + init=population, ) return results.x -def optimize_shgo(bounds, cli_args, extra_args, initial_point): # pylint: disable=unused-argument - bounds = [ - bounds['mchirp'], - bounds['eta'], - bounds['spin1z'], - bounds['spin2z'] - ] +def optimize_shgo(bounds, cli_args, extra_args, initial_point): # pylint: disable=unused-argument + bounds = [bounds["mchirp"], bounds["eta"], bounds["spin1z"], bounds["spin2z"]] results = shgo( compute_minus_network_snr, bounds=bounds, args=extra_args, iters=cli_args.snr_opt_shgo_iters, n=cli_args.snr_opt_shgo_samples, - sampling_method="sobol" + sampling_method="sobol", ) return results.x def optimize_pso(bounds, cli_args, extra_args, initial_point): options = { - 'c1': cli_args.snr_opt_pso_c1, - 'c2': cli_args.snr_opt_pso_c2, - 'w': cli_args.snr_opt_pso_w + "c1": cli_args.snr_opt_pso_c1, + "c2": cli_args.snr_opt_pso_c2, + "w": cli_args.snr_opt_pso_w, } - min_bounds = numpy.array([ - bounds['mchirp'][0], - bounds['eta'][0], - bounds['spin1z'][0], - bounds['spin2z'][0] - ]) - max_bounds = numpy.array([ - bounds['mchirp'][1], - bounds['eta'][1], - bounds['spin1z'][1], - bounds['spin2z'][1] - ]) + min_bounds = numpy.array( + [ + bounds["mchirp"][0], + bounds["eta"][0], + bounds["spin1z"][0], + bounds["spin2z"][0], + ] + ) + max_bounds = numpy.array( + [ + bounds["mchirp"][1], + bounds["eta"][1], + bounds["spin1z"][1], + bounds["spin2z"][1], + ] + ) # Initialize the population with random values within specified bounds population = numpy.random.uniform( - min_bounds, - max_bounds, - size=(int(cli_args.snr_opt_pso_particles), len(bounds)) + min_bounds, max_bounds, size=(int(cli_args.snr_opt_pso_particles), len(bounds)) ) if cli_args.snr_opt_include_candidate: # add the initial point to the population - population = numpy.concatenate((population[:-1], - initial_point)) - logger.debug('Initial population: %s', population) + population = numpy.concatenate((population[:-1], initial_point)) + logger.debug("Initial population: %s", population) optimizer = ps.single.GlobalBestPSO( n_particles=int(cli_args.snr_opt_pso_particles), dimensions=4, options=options, bounds=(min_bounds, max_bounds), - init_pos=population + init_pos=population, ) _, results = optimizer.optimize( compute_minus_network_snr_pso, iters=int(cli_args.snr_opt_pso_iters), n_processes=cli_args.cores, - args=extra_args + args=extra_args, ) return results optimize_funcs = { - 'differential_evolution': optimize_di, - 'shgo': optimize_shgo, - 'pso': optimize_pso + "differential_evolution": optimize_di, + "shgo": optimize_shgo, + "pso": optimize_pso, } # The following sets the default values of the options, but allows us to check @@ -294,62 +298,78 @@ def optimize_pso(bounds, cli_args, extra_args, initial_point): # message and default value option_dict = { - 'differential_evolution': { - 'maxiter': ('The maximum number of generations over which the entire ' - 'population is evolved.', 50), - 'popsize': ('A multiplier for setting the total population size.', - 100), + "differential_evolution": { + "maxiter": ( + "The maximum number of generations over which the entire " + "population is evolved.", + 50, + ), + "popsize": ("A multiplier for setting the total population size.", 100), }, - 'shgo': { - 'samples': ('Number of sampling points used in the construction of ' - 'the simplicial complex.', 76), - 'iters': ('Number of iterations used in the construction of the ' - 'simplicial complex.', 3), + "shgo": { + "samples": ( + "Number of sampling points used in the construction of " + "the simplicial complex.", + 76, + ), + "iters": ( + "Number of iterations used in the construction of the simplicial complex.", + 3, + ), + }, + "pso": { + "iters": ("Number of iterations used in the particle swarm optimization.", 5), + "particles": ("Number of particles used in the swarm.", 250), + "c1": ("The hyperparameter c1: the cognitive parameter.", 0.5), + "c2": ("The hyperparameter c2: the social parameter.", 2.0), + "w": ("The hyperparameter w: the inertia parameter.", 0.01), }, - 'pso': { - 'iters': ('Number of iterations used in the particle swarm ' - 'optimization.', 5), - 'particles': ('Number of particles used in the swarm.', 250), - 'c1': ('The hyperparameter c1: the cognitive parameter.', 0.5), - 'c2': ('The hyperparameter c2: the social parameter.', 2.0), - 'w': ('The hyperparameter w: the inertia parameter.', 0.01), - } } + def insert_snr_optimizer_options(parser): - opt_opt_group = parser.add_argument_group("SNR optimizer configuration " - "options.") + opt_opt_group = parser.add_argument_group("SNR optimizer configuration options.") # Option to choose which optimizer to use: optimizer_choices = sorted(list(option_dict.keys())) - opt_opt_group.add_argument('--snr-opt-method', - default='differential_evolution', + opt_opt_group.add_argument( + "--snr-opt-method", + default="differential_evolution", choices=optimizer_choices, - help='SNR Optimizer choices: ' + ', '.join(optimizer_choices)) + help="SNR Optimizer choices: " + ", ".join(optimizer_choices), + ) # Add the generic options - opt_opt_group.add_argument('--snr-opt-include-candidate', - action='store_true', - help='Include parameters of the candidate event in the initialized ' - 'array for the optimizer. Only relevant for --snr-opt-method pso ' - 'or differential_evolution') - opt_opt_group.add_argument('--snr-opt-seed', - default='42', - help='Seed to supply to the random generation of initial array to ' - 'pass to the optimizer. Only relevant for --snr-opt-method pso ' - 'or differential_evolution. Set to ''random'' for a random seed') + opt_opt_group.add_argument( + "--snr-opt-include-candidate", + action="store_true", + help="Include parameters of the candidate event in the initialized " + "array for the optimizer. Only relevant for --snr-opt-method pso " + "or differential_evolution", + ) + opt_opt_group.add_argument( + "--snr-opt-seed", + default="42", + help="Seed to supply to the random generation of initial array to " + "pass to the optimizer. Only relevant for --snr-opt-method pso " + "or differential_evolution. Set to " + "random" + " for a random seed", + ) # For each optimizer, add the possible options for optimizer, option_subdict in option_dict.items(): - optimizer_name = optimizer.replace('_', '-') - if optimizer_name == 'differential-evolution': - optimizer_name = 'di' + optimizer_name = optimizer.replace("_", "-") + if optimizer_name == "differential-evolution": + optimizer_name = "di" for opt_name, opt_help_default in option_subdict.items(): option_name = f"--snr-opt-{optimizer_name}-{opt_name}" - opt_opt_group.add_argument(option_name, + opt_opt_group.add_argument( + option_name, type=float, - help=f'Only relevant for --snr-opt-method {optimizer}: ' + - opt_help_default[0] + - f' Default = {opt_help_default[1]}') + help=f"Only relevant for --snr-opt-method {optimizer}: " + + opt_help_default[0] + + f" Default = {opt_help_default[1]}", + ) def check_snr_optimizer_options(args, parser): @@ -357,29 +377,37 @@ def check_snr_optimizer_options(args, parser): Deal with default options and required parameters given optimizer option """ options = {} - options['differential_evolution'] = [args.snr_opt_di_maxiter, - args.snr_opt_di_popsize] - options['shgo'] = [args.snr_opt_shgo_samples, args.snr_opt_shgo_iters] - options['pso'] = [args.snr_opt_pso_iters, args.snr_opt_pso_particles, - args.snr_opt_pso_c1, args.snr_opt_pso_c2, - args.snr_opt_pso_w] + options["differential_evolution"] = [ + args.snr_opt_di_maxiter, + args.snr_opt_di_popsize, + ] + options["shgo"] = [args.snr_opt_shgo_samples, args.snr_opt_shgo_iters] + options["pso"] = [ + args.snr_opt_pso_iters, + args.snr_opt_pso_particles, + args.snr_opt_pso_c1, + args.snr_opt_pso_c2, + args.snr_opt_pso_w, + ] - if args.snr_opt_method == 'pso' and ps is None: - parser.error('You need to install pyswarms to use the pso optimizer.') + if args.snr_opt_method == "pso" and ps is None: + parser.error("You need to install pyswarms to use the pso optimizer.") # Check all the options are suitable for the chosen optimizer - for k in options.keys(): + for k in options: if args.snr_opt_method == k: continue if any(options[k]): - parser.error("Argument has been supplied which is not suitable " + - f"for the optimizer given ({args.snr_opt_method})") + parser.error( + "Argument has been supplied which is not suitable " + f"for the optimizer given ({args.snr_opt_method})" + ) # Give the arguments the default values according to the dictionary - optimizer_name = args.snr_opt_method.replace('_', '-') - if optimizer_name == 'differential-evolution': - optimizer_name = 'di' + optimizer_name = args.snr_opt_method.replace("_", "-") + if optimizer_name == "differential-evolution": + optimizer_name = "di" for key, value in option_dict[args.snr_opt_method].items(): - key_name = f'snr_opt_{optimizer_name}_{key}' + key_name = f"snr_opt_{optimizer_name}_{key}" if not getattr(args, key_name): setattr(args, key_name, value[1]) diff --git a/pycbc/live/supervision.py b/pycbc/live/supervision.py index 858dc6c782b..8e3f3059520 100644 --- a/pycbc/live/supervision.py +++ b/pycbc/live/supervision.py @@ -22,15 +22,16 @@ """ import logging +import os import subprocess import time -import os from datetime import datetime + from dateutil.relativedelta import relativedelta import pycbc -logger = logging.getLogger('pycbc.live.supervision') +logger = logging.getLogger("pycbc.live.supervision") def symlink(target, link_name): @@ -43,9 +44,9 @@ def symlink(target, link_name): link_name = os.path.abspath(link_name) logger.info("Linking %s to %s", target, link_name) try: - subprocess.run(['ln', '-sf', target, link_name], check=True) + subprocess.run(["ln", "-sf", target, link_name], check=True) except subprocess.CalledProcessError as sub_err: - logging.error("Could not link %s to %s", target, link_name) + logging.exception("Could not link %s to %s", target, link_name) raise sub_err @@ -55,8 +56,8 @@ def dict_to_args(opts_dict): """ dargs = [] for option, value in opts_dict.items(): - dargs.append('--' + option.strip()) - if value == '': + dargs.append("--" + option.strip()) + if value == "": # option is a flag, do nothing continue if len(value.split()) > 1: @@ -74,21 +75,18 @@ def mail_volunteers_error(controls, mail_body_lines, subject): Email a list of people, defined by mail-volunteers-file To be used for errors or unusual occurences """ - with open(controls['mail-volunteers-file'], 'r') as mail_volunteers_file: - volunteers = [volunteer.strip() for volunteer in - mail_volunteers_file.readlines()] - logger.info("Emailing %s with warnings", ' '.join(volunteers)) - mail_command = [ - 'mail', - '-s', - subject - ] + with open(controls["mail-volunteers-file"]) as mail_volunteers_file: + volunteers = [ + volunteer.strip() for volunteer in mail_volunteers_file + ] + logger.info("Emailing %s with warnings", " ".join(volunteers)) + mail_command = ["mail", "-s", subject] mail_command += volunteers - mail_body = '\n'.join(mail_body_lines) + mail_body = "\n".join(mail_body_lines) try: subprocess.run(mail_command, input=mail_body, text=True, check=True) except subprocess.CalledProcessError as sub_err: - logging.error("Could not send mail on error") + logging.exception("Could not send mail on error") raise sub_err @@ -97,29 +95,28 @@ def run_and_error(command_arguments, controls): Wrapper around subprocess.run to catch errors and send emails if required """ logger.info("Running %s", " ".join(command_arguments)) - command_output = subprocess.run( - command_arguments, - capture_output=True - ) + command_output = subprocess.run(command_arguments, capture_output=True) if command_output.returncode: - error_contents = [' '.join(command_arguments), '\n', - command_output.stderr.decode()] - if 'mail-volunteers-file' in controls: + error_contents = [ + " ".join(command_arguments), + "\n", + command_output.stderr.decode(), + ] + if "mail-volunteers-file" in controls: mail_volunteers_error( controls, error_contents, - f"PyCBC live could not run {command_arguments[0]}" + f"PyCBC live could not run {command_arguments[0]}", ) err_msg = f"Could not run {command_arguments[0]}:\n" - err_msg += ' '.join(error_contents) + err_msg += " ".join(error_contents) raise subprocess.SubprocessError(err_msg) def wait_for_utc_time(target_str): - """Wait until the UTC time is as given by `target_str`, in HH:MM:SS format. - """ - target_hour, target_minute, target_second = map(int, target_str.split(':')) + """Wait until the UTC time is as given by `target_str`, in HH:MM:SS format.""" + target_hour, target_minute, target_second = map(int, target_str.split(":")) now = datetime.utcnow() # for today's target, take now and replace the time target_today = now + relativedelta( @@ -131,7 +128,7 @@ def wait_for_utc_time(target_str): ) next_target = target_today if now <= target_today else target_tomorrow sleep_seconds = (next_target - now).total_seconds() - logger.info('Waiting %.0f s', sleep_seconds) + logger.info("Waiting %.0f s", sleep_seconds) time.sleep(sleep_seconds) @@ -139,16 +136,10 @@ def ensure_directories(control_values, day_str): """ Ensure that the required directories exist """ - output_dir = os.path.join( - control_values['output-directory'], - day_str - ) + output_dir = os.path.join(control_values["output-directory"], day_str) pycbc.makedir(output_dir) - if 'public-dir' in control_values: + if "public-dir" in control_values: # The public directory wil be in subdirectories for the year, month, # day, e.g. 2024_04_12 will be in 2024/04/12. - public_dir = os.path.join( - control_values['public-dir'], - *day_str.split('_') - ) + public_dir = os.path.join(control_values["public-dir"], *day_str.split("_")) pycbc.makedir(public_dir) diff --git a/pycbc/mchirp_area.py b/pycbc/mchirp_area.py index eeca5036759..d3bd0c255eb 100644 --- a/pycbc/mchirp_area.py +++ b/pycbc/mchirp_area.py @@ -2,151 +2,189 @@ # Initial code by A. Curiel Barroso, August 2019 # Modified by V. Villa-Ortega, January 2020, March 2021 -"""Functions to compute the area corresponding to different CBC on the m1 & m2 +""" +Functions to compute the area corresponding to different CBC on the m1 & m2 plane when given a central mchirp value and uncertainty. It also includes a function that calculates the source frame when given the detector frame mass and redshift. """ -import math import logging -import numpy as np +import math -from scipy.integrate import quad +import numpy as np from astropy.cosmology import FlatLambdaCDM +from scipy.integrate import quad -from pycbc.cosmology import _redshift from pycbc.conversions import mass2_from_mchirp_mass1 as m2mcm1 +from pycbc.cosmology import _redshift -logger = logging.getLogger('pycbc.mchirp_area') +logger = logging.getLogger("pycbc.mchirp_area") def insert_args(parser): - mchirp_group = parser.add_argument_group("Arguments for estimating the " - "source probabilities of a " - "candidate event using the snr, " - "mchirp, and effective distance.") - mchirp_group.add_argument('--src-class-mass-limits', type=float, nargs=3, - metavar=('MIN_M2', 'MAX_NS', 'MAX_M1'), - default=[1.0, 3.0, 45.0], - help="Minimum and maximum values for the mass " - "of the binary components and maximum mass " - "of a neutron star, used as limits " - "when computing the area corresponding" - "to different CBC sources.") - mchirp_group.add_argument('--src-class-mass-gap-max', type=float, - metavar=('MAX_GAP'), - help="Upper limit of the mass gap, corresponding" - " to the minimum mass of a black hole. " - "Used as limit of integration of the " - "different CBC regions when considering " - "the MassGap category.") - mchirp_group.add_argument('--src-class-mchirp-to-delta', type=float, - metavar='m0', required=True, - help='Coefficient to estimate the value of the ' - 'mchirp uncertainty by mchirp_delta = ' - 'm0 * mchirp.') - mchirp_group.add_argument('--src-class-eff-to-lum-distance', type=float, - metavar='a0', required=True, - help='Coefficient to estimate the value of the ' - 'luminosity distance from the minimum ' - 'eff distance by D_lum = a0 * min(D_eff).') - mchirp_group.add_argument('--src-class-lum-distance-to-delta', type=float, - nargs=2, metavar=('b0', 'b1'), required=True, - help='Coefficients to estimate the value of the ' - 'uncertainty on the luminosity distance ' - 'from the estimated luminosity distance and' - ' the coinc snr by delta_lum = D_lum * ' - 'exp(b0) * coinc_snr ** b1.') - mchirp_group.add_argument('--src-class-mass-gap-separate', - action='store_true', - help='Gives separate probabilities for each kind' - ' of mass gap CBC sources: GNS, GG, BHG.') - mchirp_group.add_argument('--src-class-lal-cosmology', - action='store_true', - help='Uses the Planck15 cosmology defined in ' - 'lalsuite instead of the astropy Planck15 ' - 'default model.') + mchirp_group = parser.add_argument_group( + "Arguments for estimating the " + "source probabilities of a " + "candidate event using the snr, " + "mchirp, and effective distance." + ) + mchirp_group.add_argument( + "--src-class-mass-limits", + type=float, + nargs=3, + metavar=("MIN_M2", "MAX_NS", "MAX_M1"), + default=[1.0, 3.0, 45.0], + help="Minimum and maximum values for the mass " + "of the binary components and maximum mass " + "of a neutron star, used as limits " + "when computing the area corresponding" + "to different CBC sources.", + ) + mchirp_group.add_argument( + "--src-class-mass-gap-max", + type=float, + metavar=("MAX_GAP"), + help="Upper limit of the mass gap, corresponding" + " to the minimum mass of a black hole. " + "Used as limit of integration of the " + "different CBC regions when considering " + "the MassGap category.", + ) + mchirp_group.add_argument( + "--src-class-mchirp-to-delta", + type=float, + metavar="m0", + required=True, + help="Coefficient to estimate the value of the " + "mchirp uncertainty by mchirp_delta = " + "m0 * mchirp.", + ) + mchirp_group.add_argument( + "--src-class-eff-to-lum-distance", + type=float, + metavar="a0", + required=True, + help="Coefficient to estimate the value of the " + "luminosity distance from the minimum " + "eff distance by D_lum = a0 * min(D_eff).", + ) + mchirp_group.add_argument( + "--src-class-lum-distance-to-delta", + type=float, + nargs=2, + metavar=("b0", "b1"), + required=True, + help="Coefficients to estimate the value of the " + "uncertainty on the luminosity distance " + "from the estimated luminosity distance and" + " the coinc snr by delta_lum = D_lum * " + "exp(b0) * coinc_snr ** b1.", + ) + mchirp_group.add_argument( + "--src-class-mass-gap-separate", + action="store_true", + help="Gives separate probabilities for each kind" + " of mass gap CBC sources: GNS, GG, BHG.", + ) + mchirp_group.add_argument( + "--src-class-lal-cosmology", + action="store_true", + help="Uses the Planck15 cosmology defined in " + "lalsuite instead of the astropy Planck15 " + "default model.", + ) def from_cli(args, parser): mass_limits_sorted = sorted(args.src_class_mass_limits) if args.src_class_mass_gap_max: if args.src_class_mass_gap_max < mass_limits_sorted[1]: - parser.error('MAX_GAP value cannot be lower than MAX_NS limit') - return {'mass_limits': - {'max_m1': mass_limits_sorted[2], - 'min_m2': mass_limits_sorted[0]}, - 'mass_bdary': - {'ns_max': mass_limits_sorted[1], - 'gap_max': args.src_class_mass_gap_max}, - 'estimation_coeff': - {'a0': args.src_class_eff_to_lum_distance, - 'b0': args.src_class_lum_distance_to_delta[0], - 'b1': args.src_class_lum_distance_to_delta[1], - 'm0': args.src_class_mchirp_to_delta}, - 'mass_gap': True, - 'mass_gap_separate': args.src_class_mass_gap_separate, - 'lal_cosmology': args.src_class_lal_cosmology} - return {'mass_limits': - {'max_m1': mass_limits_sorted[2], - 'min_m2': mass_limits_sorted[0]}, - 'mass_bdary': - {'ns_max': mass_limits_sorted[1], - 'gap_max': mass_limits_sorted[1]}, - 'estimation_coeff': - {'a0': args.src_class_eff_to_lum_distance, - 'b0': args.src_class_lum_distance_to_delta[0], - 'b1': args.src_class_lum_distance_to_delta[1], - 'm0': args.src_class_mchirp_to_delta}, - 'mass_gap': False, - 'mass_gap_separate': args.src_class_mass_gap_separate, - 'lal_cosmology': args.src_class_lal_cosmology} + parser.error("MAX_GAP value cannot be lower than MAX_NS limit") + return { + "mass_limits": { + "max_m1": mass_limits_sorted[2], + "min_m2": mass_limits_sorted[0], + }, + "mass_bdary": { + "ns_max": mass_limits_sorted[1], + "gap_max": args.src_class_mass_gap_max, + }, + "estimation_coeff": { + "a0": args.src_class_eff_to_lum_distance, + "b0": args.src_class_lum_distance_to_delta[0], + "b1": args.src_class_lum_distance_to_delta[1], + "m0": args.src_class_mchirp_to_delta, + }, + "mass_gap": True, + "mass_gap_separate": args.src_class_mass_gap_separate, + "lal_cosmology": args.src_class_lal_cosmology, + } + return { + "mass_limits": { + "max_m1": mass_limits_sorted[2], + "min_m2": mass_limits_sorted[0], + }, + "mass_bdary": { + "ns_max": mass_limits_sorted[1], + "gap_max": mass_limits_sorted[1], + }, + "estimation_coeff": { + "a0": args.src_class_eff_to_lum_distance, + "b0": args.src_class_lum_distance_to_delta[0], + "b1": args.src_class_lum_distance_to_delta[1], + "m0": args.src_class_mchirp_to_delta, + }, + "mass_gap": False, + "mass_gap_separate": args.src_class_mass_gap_separate, + "lal_cosmology": args.src_class_lal_cosmology, + } def redshift_estimation(distance, distance_std, lal_cosmology): - """Takes values of distance and its uncertainty and returns a - dictionary with estimates of the redshift and its uncertainty. - If the argument 'lal_cosmology' is True, it uses Planck15 cosmology - model as defined in lalsuite instead of the astropy default. - Constants for lal_cosmology taken from Planck15_lal_cosmology() in - https://git.ligo.org/lscsoft/pesummary/-/blob/master/pesummary/gw/ - cosmology.py. + """ + Takes values of distance and its uncertainty and returns a + dictionary with estimates of the redshift and its uncertainty. + If the argument 'lal_cosmology' is True, it uses Planck15 cosmology + model as defined in lalsuite instead of the astropy default. + Constants for lal_cosmology taken from Planck15_lal_cosmology() in + https://git.ligo.org/lscsoft/pesummary/-/blob/master/pesummary/gw/ + cosmology.py. """ if lal_cosmology: cosmology = FlatLambdaCDM(H0=67.90, Om0=0.3065) else: cosmology = None z_estimation = _redshift(distance, cosmology=cosmology) - z_est_max = _redshift((distance + distance_std), - cosmology=cosmology) - z_est_min = _redshift((distance - distance_std), - cosmology=cosmology) + z_est_max = _redshift((distance + distance_std), cosmology=cosmology) + z_est_min = _redshift((distance - distance_std), cosmology=cosmology) z_std_estimation = 0.5 * (z_est_max - z_est_min) - z = {'central': z_estimation, 'delta': z_std_estimation} + z = {"central": z_estimation, "delta": z_std_estimation} return z def src_mass_from_z_det_mass(z, del_z, mdet, del_mdet): - """Takes values of redshift, redshift uncertainty, detector mass and its + """ + Takes values of redshift, redshift uncertainty, detector mass and its uncertainty and computes the source mass and its uncertainty. """ - msrc = mdet / (1. + z) - del_msrc = msrc * ((del_mdet / mdet) ** 2. - + (del_z / (1. + z)) ** 2.) ** 0.5 + msrc = mdet / (1.0 + z) + del_msrc = msrc * ((del_mdet / mdet) ** 2.0 + (del_z / (1.0 + z)) ** 2.0) ** 0.5 return (msrc, del_msrc) def intmc(mc, x_min, x_max): - """Returns the integral of m2 over m1 between x_min and x_max, - assuming that mchirp is fixed. + """ + Returns the integral of m2 over m1 between x_min and x_max, + assuming that mchirp is fixed. """ integral = quad(lambda x, mc: m2mcm1(mc, x), x_min, x_max, args=mc) return integral[0] def get_area(trig_mc, lim_h1, lim_h2, lim_v1, lim_v2): - """Returns the area under the chirp mass contour in a region of the m1-m2 + """ + Returns the area under the chirp mass contour in a region of the m1-m2 plane (m1 > m2). Parameters @@ -162,15 +200,16 @@ def get_area(trig_mc, lim_h1, lim_h2, lim_v1, lim_v2): Returns ------- area : float + """ mc_max = trig_mc[0] + trig_mc[1] mc_min = trig_mc[0] - trig_mc[1] # The points where the equal mass line and a chirp mass # curve intersect is m1 = m2 = 2**0.2 * mchirp - mi_max = (2.**0.2) * mc_max - mi_min = (2.**0.2) * mc_min + mi_max = (2.0**0.2) * mc_max + mi_min = (2.0**0.2) * mc_min - if lim_h1 == 'diagonal': + if lim_h1 == "diagonal": max_h1 = mi_max min_h1 = mi_min fun_sup = lambda x: x @@ -196,32 +235,27 @@ def get_area(trig_mc, lim_h1, lim_h2, lim_v1, lim_v2): return area -def calc_areas( - trig_mc_det, - mass_limits, - mass_bdary, - z, - mass_gap, - mass_gap_separate): - """Computes the area inside the lines of the second component mass as a +def calc_areas(trig_mc_det, mass_limits, mass_bdary, z, mass_gap, mass_gap_separate): + """ + Computes the area inside the lines of the second component mass as a function of the first component mass for the two extreme values of mchirp: mchirp +/- mchirp_uncertainty, for each region of the source classifying diagram. """ - trig_mc = src_mass_from_z_det_mass(z["central"], z["delta"], - trig_mc_det["central"], - trig_mc_det["delta"]) + trig_mc = src_mass_from_z_det_mass( + z["central"], z["delta"], trig_mc_det["central"], trig_mc_det["delta"] + ) m2_min = mass_limits["min_m2"] m1_max = mass_limits["max_m1"] ns_max = mass_bdary["ns_max"] gap_max = mass_bdary["gap_max"] - abbh = get_area(trig_mc, 'diagonal', gap_max, gap_max, m1_max) + abbh = get_area(trig_mc, "diagonal", gap_max, gap_max, m1_max) abhg = get_area(trig_mc, gap_max, ns_max, gap_max, m1_max) ansbh = get_area(trig_mc, ns_max, m2_min, gap_max, m1_max) - agg = get_area(trig_mc, 'diagonal', ns_max, ns_max, gap_max) + agg = get_area(trig_mc, "diagonal", ns_max, ns_max, gap_max) agns = get_area(trig_mc, ns_max, m2_min, ns_max, gap_max) - abns = get_area(trig_mc, 'diagonal', m2_min, m2_min, ns_max) + abns = get_area(trig_mc, "diagonal", m2_min, m2_min, ns_max) if mass_gap: if mass_gap_separate: @@ -231,73 +265,76 @@ def calc_areas( "NSBH": ansbh, "GG": agg, "BHG": abhg, - "BBH": abbh - } - return { - "BNS": abns, - "NSBH": ansbh, - "BBH": abbh, - "Mass Gap": agns + agg + abhg + "BBH": abbh, } - return { - "BNS": abns, - "NSBH": ansbh, - "BBH": abbh - } + return {"BNS": abns, "NSBH": ansbh, "BBH": abbh, "Mass Gap": agns + agg + abhg} + return {"BNS": abns, "NSBH": ansbh, "BBH": abbh} def calc_probabilities(mchirp, snr, eff_distance, src_args): - """Computes the different probabilities that a candidate event belongs to - each CBC source category taking as arguments the chirp mass, the - coincident SNR and the effective distance, and estimating the - chirp mass uncertainty, the luminosity distance (and its uncertainty) - and the redshift (and its uncertainty). Probability is estimated to be - directly proportional to the area of the corresponding CBC region. """ - mass_limits = src_args['mass_limits'] - mass_bdary = src_args['mass_bdary'] - coeff = src_args['estimation_coeff'] - trig_mc_det = {'central': mchirp, 'delta': mchirp * coeff['m0']} - dist_estimation = coeff['a0'] * eff_distance - dist_std_estimation = (dist_estimation * math.exp(coeff['b0']) * - snr ** coeff['b1']) - z = redshift_estimation(dist_estimation, dist_std_estimation, - src_args['lal_cosmology']) - mass_gap = src_args['mass_gap'] - mass_gap_separate = src_args['mass_gap_separate'] + Computes the different probabilities that a candidate event belongs to + each CBC source category taking as arguments the chirp mass, the + coincident SNR and the effective distance, and estimating the + chirp mass uncertainty, the luminosity distance (and its uncertainty) + and the redshift (and its uncertainty). Probability is estimated to be + directly proportional to the area of the corresponding CBC region. + """ + mass_limits = src_args["mass_limits"] + mass_bdary = src_args["mass_bdary"] + coeff = src_args["estimation_coeff"] + trig_mc_det = {"central": mchirp, "delta": mchirp * coeff["m0"]} + dist_estimation = coeff["a0"] * eff_distance + dist_std_estimation = dist_estimation * math.exp(coeff["b0"]) * snr ** coeff["b1"] + z = redshift_estimation( + dist_estimation, dist_std_estimation, src_args["lal_cosmology"] + ) + mass_gap = src_args["mass_gap"] + mass_gap_separate = src_args["mass_gap_separate"] # If the mchirp is greater than the mchirp corresponding to two masses # equal to the maximum mass, the probability for BBH is 100%. # If it is less than the mchirp corresponding to two masses equal to the # minimum mass, the probability for BNS is 100%. - mc_max = mass_limits['max_m1'] / (2 ** 0.2) - mc_min = mass_limits['min_m2'] / (2 ** 0.2) + mc_max = mass_limits["max_m1"] / (2**0.2) + mc_min = mass_limits["min_m2"] / (2**0.2) - if trig_mc_det['central'] > mc_max * (1 + z['central']): + if trig_mc_det["central"] > mc_max * (1 + z["central"]): if mass_gap: if mass_gap_separate: - probabilities = {"BNS": 0.0, "GNS": 0.0, "NSBH": 0.0, - "GG": 0.0, "BHG": 0.0, "BBH": 1.0} + probabilities = { + "BNS": 0.0, + "GNS": 0.0, + "NSBH": 0.0, + "GG": 0.0, + "BHG": 0.0, + "BBH": 1.0, + } else: - probabilities = {"BNS": 0.0, "NSBH": 0.0, "BBH": 1.0, - "Mass Gap": 0.0} + probabilities = {"BNS": 0.0, "NSBH": 0.0, "BBH": 1.0, "Mass Gap": 0.0} else: probabilities = {"BNS": 0.0, "NSBH": 0.0, "BBH": 1.0} - elif trig_mc_det['central'] < mc_min * (1 + z['central']): + elif trig_mc_det["central"] < mc_min * (1 + z["central"]): if mass_gap: if mass_gap_separate: - probabilities = {"BNS": 1.0, "GNS": 0.0, "NSBH": 0.0, - "GG": 0.0, "BHG": 0.0, "BBH": 0.0} + probabilities = { + "BNS": 1.0, + "GNS": 0.0, + "NSBH": 0.0, + "GG": 0.0, + "BHG": 0.0, + "BBH": 0.0, + } else: - probabilities = {"BNS": 1.0, "NSBH": 0.0, "BBH": 0.0, - "Mass Gap": 0.0} + probabilities = {"BNS": 1.0, "NSBH": 0.0, "BBH": 0.0, "Mass Gap": 0.0} else: probabilities = {"BNS": 1.0, "NSBH": 0.0, "BBH": 0.0} else: - areas = calc_areas(trig_mc_det, mass_limits, mass_bdary, z, - mass_gap, mass_gap_separate) + areas = calc_areas( + trig_mc_det, mass_limits, mass_bdary, z, mass_gap, mass_gap_separate + ) total_area = sum(areas.values()) probabilities = {key: areas[key] / total_area for key in areas} diff --git a/pycbc/neutron_stars/__init__.py b/pycbc/neutron_stars/__init__.py index 33c02d9529d..2afe85148af 100644 --- a/pycbc/neutron_stars/__init__.py +++ b/pycbc/neutron_stars/__init__.py @@ -1,9 +1,11 @@ import os.path + # Setup the directory with the NS equilibrium sequence(s) -NS_DATA_DIRECTORY = os.path.join( - os.path.dirname(__file__), 'ns_data') +NS_DATA_DIRECTORY = os.path.join(os.path.dirname(__file__), "ns_data") NS_SEQUENCES = [ - f.replace('equil_', '').replace('.dat', '') - for f in os.listdir(NS_DATA_DIRECTORY) if f.endswith('.dat')] + f.replace("equil_", "").replace(".dat", "") + for f in os.listdir(NS_DATA_DIRECTORY) + if f.endswith(".dat") +] from pycbc.neutron_stars.eos_utils import * from pycbc.neutron_stars.pg_isso_solver import * diff --git a/pycbc/neutron_stars/eos_utils.py b/pycbc/neutron_stars/eos_utils.py index 90661beb8d2..54b31ebaaed 100644 --- a/pycbc/neutron_stars/eos_utils.py +++ b/pycbc/neutron_stars/eos_utils.py @@ -17,19 +17,23 @@ """ Utility functions for handling NS equations of state """ + import os.path + import numpy as np from scipy.interpolate import interp1d -from . import NS_SEQUENCES, NS_DATA_DIRECTORY -from .pg_isso_solver import PG_ISSO_solver from pycbc.libutils import import_optional + +from . import NS_DATA_DIRECTORY, NS_SEQUENCES +from .pg_isso_solver import PG_ISSO_solver + # Imports needed if we implement the lalsimulation EOS interface # from pycbc.constants import ( # MSUN_SI, G_SI, C_SI # ) -lalsim = import_optional('lalsimulation') +lalsim = import_optional("lalsimulation") def load_ns_sequence(eos_name): @@ -39,13 +43,13 @@ def load_ns_sequence(eos_name): File format is: grav mass (Msun), baryonic mass (Msun), compactness Parameters - ----------- + ---------- eos_name : string NS equation of state label ('2H' is the only supported choice at the moment) Returns - ---------- + ------- ns_sequence : numpy.array contains the sequence data in the form NS gravitational mass (in solar masses), NS baryonic mass (in solar @@ -54,15 +58,16 @@ def load_ns_sequence(eos_name): the maximum NS gravitational mass (in solar masses) in the sequence (this is the mass of the most massive stable NS) + """ - ns_sequence_file = os.path.join( - NS_DATA_DIRECTORY, 'equil_{}.dat'.format(eos_name)) + ns_sequence_file = os.path.join(NS_DATA_DIRECTORY, f"equil_{eos_name}.dat") if eos_name not in NS_SEQUENCES: raise NotImplementedError( - f'{eos_name} does not have an implemented NS sequence file! ' - f'To implement, the file {ns_sequence_file} must exist and ' - 'contain: NS gravitational mass (in solar masses), NS baryonic ' - 'mass (in solar masses), NS compactness (dimensionless)') + f"{eos_name} does not have an implemented NS sequence file! " + f"To implement, the file {ns_sequence_file} must exist and " + "contain: NS gravitational mass (in solar masses), NS baryonic " + "mass (in solar masses), NS compactness (dimensionless)" + ) ns_sequence = np.loadtxt(ns_sequence_file) max_ns_g_mass = max(ns_sequence[:, 0]) return (ns_sequence, max_ns_g_mass) @@ -74,7 +79,7 @@ def interp_grav_mass_to_baryon_mass(ns_g_mass, ns_sequence, extrapolate=False): mass and an NS equilibrium sequence (in solar masses). Parameters - ----------- + ---------- ns_g_mass : float NS gravitational mass (in solar masses) ns_sequence : numpy.array @@ -86,8 +91,9 @@ def interp_grav_mass_to_baryon_mass(ns_g_mass, ns_sequence, extrapolate=False): Default is False (so ValueError is raised for ns_g_mass out of bounds) Returns - ---------- + ------- float + """ x = ns_sequence[:, 0] y = ns_sequence[:, 1] @@ -102,7 +108,7 @@ def interp_grav_mass_to_compactness(ns_g_mass, ns_sequence, extrapolate=False): its gravitational mass and an NS equilibrium sequence. Parameters - ----------- + ---------- ns_g_mass : float NS gravitational mass (in solar masses) ns_sequence : numpy.array @@ -114,8 +120,9 @@ def interp_grav_mass_to_compactness(ns_g_mass, ns_sequence, extrapolate=False): Default is False (so ValueError is raised for ns_g_mass out of bounds) Returns - ---------- + ------- float + """ x = ns_sequence[:, 0] y = ns_sequence[:, 2] @@ -125,7 +132,8 @@ def interp_grav_mass_to_compactness(ns_g_mass, ns_sequence, extrapolate=False): def initialize_eos(ns_mass, eos, extrapolate=False): - """Load an equation of state and return the compactness and baryonic + """ + Load an equation of state and return the compactness and baryonic mass for a given neutron star mass Parameters @@ -146,6 +154,7 @@ def initialize_eos(ns_mass, eos, extrapolate=False): Compactness parameter of the neutron star. ns_b_mass : float Baryonic mass of the neutron star. + """ if isinstance(ns_mass, np.ndarray): input_is_array = True @@ -155,36 +164,39 @@ def initialize_eos(ns_mass, eos, extrapolate=False): try: if any(ns_mass > ns_max) and input_is_array: raise ValueError( - f'Maximum NS mass for {eos} is {ns_max}, received masses ' - f'up to {max(ns_mass[ns_mass > ns_max])}') + f"Maximum NS mass for {eos} is {ns_max}, received masses " + f"up to {max(ns_mass[ns_mass > ns_max])}" + ) except TypeError: if ns_mass > ns_max and not input_is_array: raise ValueError( - f'Maximum NS mass for {eos} is {ns_max}, received ' - f'{ns_mass}') + f"Maximum NS mass for {eos} is {ns_max}, received {ns_mass}" + ) # Interpolate NS compactness and rest mass ns_compactness = interp_grav_mass_to_compactness( - ns_mass, ns_seq, extrapolate=extrapolate) + ns_mass, ns_seq, extrapolate=extrapolate + ) ns_b_mass = interp_grav_mass_to_baryon_mass( - ns_mass, ns_seq, extrapolate=extrapolate) + ns_mass, ns_seq, extrapolate=extrapolate + ) elif eos in lalsim.SimNeutronStarEOSNames: - #from pycbc.constants import MSUN_SI, G_SI, C_SI - #eos_obj = lalsim.SimNeutronStarEOSByName(eos) - #eos_fam = lalsim.CreateSimNeutronStarFamily(eos_obj) - #r_ns = lalsim.SimNeutronStarRadius(ns_mass * MSUN_SI, eos_obj) - #ns_compactness = G_SI * ns_mass * MSUN_SI / (r_ns * C_SI**2) - raise NotImplementedError( - 'LALSimulation EOS interface not yet implemented!') + # from pycbc.constants import MSUN_SI, G_SI, C_SI + # eos_obj = lalsim.SimNeutronStarEOSByName(eos) + # eos_fam = lalsim.CreateSimNeutronStarFamily(eos_obj) + # r_ns = lalsim.SimNeutronStarRadius(ns_mass * MSUN_SI, eos_obj) + # ns_compactness = G_SI * ns_mass * MSUN_SI / (r_ns * C_SI**2) + raise NotImplementedError("LALSimulation EOS interface not yet implemented!") else: raise NotImplementedError( - f'{eos} is not implemented! Available are: ' - f'{NS_SEQUENCES + list(lalsim.SimNeutronStarEOSNames)}') + f"{eos} is not implemented! Available are: " + f"{NS_SEQUENCES + list(lalsim.SimNeutronStarEOSNames)}" + ) return (ns_compactness, ns_b_mass) -def foucart18( - eta, ns_compactness, ns_b_mass, bh_spin_mag, bh_spin_pol): - """Function that determines the remnant disk mass of an NS-BH system +def foucart18(eta, ns_compactness, ns_b_mass, bh_spin_mag, bh_spin_pol): + """ + Function that determines the remnant disk mass of an NS-BH system using the fit to numerical-relativity results discussed in `Foucart, Hinderer & Nissanke, PRD 98, 081501(R) (2018)`_. @@ -204,16 +216,17 @@ def foucart18( Dimensionless spin magnitude of the BH. bh_spin_pol : {float, array} The tilt angle of the BH spin. + """ isso = PG_ISSO_solver(bh_spin_mag, bh_spin_pol) # Fit parameters and tidal correction alpha = 0.406 - beta = 0.139 + beta = 0.139 gamma = 0.255 delta = 1.761 fit = ( - alpha / eta ** (1/3) * (1 - 2 * ns_compactness) + alpha / eta ** (1 / 3) * (1 - 2 * ns_compactness) - beta * ns_compactness / eta * isso + gamma - ) + ) return ns_b_mass * np.where(fit > 0.0, fit, 0.0) ** delta diff --git a/pycbc/neutron_stars/pg_isso_solver.py b/pycbc/neutron_stars/pg_isso_solver.py index 5774e019b16..42add73731c 100644 --- a/pycbc/neutron_stars/pg_isso_solver.py +++ b/pycbc/neutron_stars/pg_isso_solver.py @@ -28,7 +28,8 @@ def ISCO_solution(chi, incl): - r"""Analytic solution of the innermost + r""" + Analytic solution of the innermost stable circular orbit (ISCO) for the Kerr metric. ..See eq. (2.21) of @@ -36,7 +37,7 @@ def ISCO_solution(chi, incl): https://articles.adsabs.harvard.edu/pdf/1972ApJ...178..347B Parameters - ----------- + ---------- chi: float the BH dimensionless spin parameter incl: float @@ -44,8 +45,9 @@ def ISCO_solution(chi, incl): momentum in radians Returns - ---------- + ------- float + """ chi2 = chi * chi sgn = np.sign(np.cos(incl)) @@ -55,7 +57,8 @@ def ISCO_solution(chi, incl): def ISSO_eq_at_pole(r, chi): - r"""Polynomial that enables the calculation of the Kerr polar + r""" + Polynomial that enables the calculation of the Kerr polar (:math:`\iota = \pm \pi / 2`) innermost stable spherical orbit (ISSO) radius via the roots of @@ -78,15 +81,17 @@ def ISSO_eq_at_pole(r, chi): Returns ------- float + """ chi2 = chi * chi - return ( - r**3 * (r**2 * (r - 6) + chi2 * (3 * r + 4)) - + chi2 * chi2 * (3 * r * (r - 2) + chi2)) + return r**3 * (r**2 * (r - 6) + chi2 * (3 * r + 4)) + chi2 * chi2 * ( + 3 * r * (r - 2) + chi2 + ) def ISSO_eq_at_pole_dr(r, chi): - """Partial derivative of :func:`ISSO_eq_at_pole` with respect to r. + """ + Partial derivative of :func:`ISSO_eq_at_pole` with respect to r. Parameters ---------- @@ -98,17 +103,19 @@ def ISSO_eq_at_pole_dr(r, chi): Returns ------- float + """ chi2 = chi * chi twlvchi2 = 12 * chi2 sxchi4 = 6 * chi2 * chi2 return ( - 6 * r**5 - 30 * r**4 + twlvchi2 * r**3 + twlvchi2 * r**2 + sxchi4 * r - + sxchi4) + 6 * r**5 - 30 * r**4 + twlvchi2 * r**3 + twlvchi2 * r**2 + sxchi4 * r + sxchi4 + ) def ISSO_eq_at_pole_dr2(r, chi): - """Double partial derivative of :func:`ISSO_eq_at_pole` with + """ + Double partial derivative of :func:`ISSO_eq_at_pole` with respect to r. Parameters @@ -121,15 +128,15 @@ def ISSO_eq_at_pole_dr2(r, chi): Returns ------- float + """ chi2 = chi * chi - return ( - 30 * r**4 - 120 * r**3 + 36 * chi2 * r**2 + 24 * chi2 * r - + 6 * chi2 * chi2) + return 30 * r**4 - 120 * r**3 + 36 * chi2 * r**2 + 24 * chi2 * r + 6 * chi2 * chi2 def PG_ISSO_eq(r, chi, incl): - r"""Polynomial that enables the calculation of a generic innermost + r""" + Polynomial that enables the calculation of a generic innermost stable spherical orbit (ISSO) radius via the roots in :math:`r` of .. math:: @@ -172,6 +179,7 @@ def PG_ISSO_eq(r, chi, incl): Returns ------- float + """ chi2 = chi * chi chi4 = chi2 * chi2 @@ -179,25 +187,22 @@ def PG_ISSO_eq(r, chi, incl): r4 = r2 * r2 three_r = 3 * r r_minus_2 = r - 2 - sin_incl2 = (np.sin(incl))**2 - - X = ( - chi2 * ( - chi2 * (3 * chi2 + 4 * r * (2 * r - 3)) - + r2 * (15 * r * (r - 4) + 28)) - - 6 * r4 * (r2 - 4)) - Y = ( - chi4 * (chi4 + r2 * (7 * r * (three_r - 4) + 36)) - + 6 * r * r_minus_2 * ( - chi4 * chi2 + 2 * r2 * r * ( - chi2 * (three_r + 2) + 3 * r2 * r_minus_2))) - Z = (r * (r - 6))**2 - chi2 * (2 * r * (3 * r + 14) - 9 * chi2) + sin_incl2 = (np.sin(incl)) ** 2 + + X = chi2 * ( + chi2 * (3 * chi2 + 4 * r * (2 * r - 3)) + r2 * (15 * r * (r - 4) + 28) + ) - 6 * r4 * (r2 - 4) + Y = chi4 * (chi4 + r2 * (7 * r * (three_r - 4) + 36)) + 6 * r * r_minus_2 * ( + chi4 * chi2 + 2 * r2 * r * (chi2 * (three_r + 2) + 3 * r2 * r_minus_2) + ) + Z = (r * (r - 6)) ** 2 - chi2 * (2 * r * (3 * r + 14) - 9 * chi2) return r4 * r4 * Z + chi2 * sin_incl2 * (chi2 * sin_incl2 * Y - 2 * r4 * X) def PG_ISSO_eq_dr(r, chi, incl): - """Partial derivative of :func:`PG_ISSO_eq` with respect to r. + """ + Partial derivative of :func:`PG_ISSO_eq` with respect to r. Parameters ---------- @@ -212,6 +217,7 @@ def PG_ISSO_eq_dr(r, chi, incl): Returns ------- float + """ sini = np.sin(incl) sin2i = sini * sini @@ -222,23 +228,33 @@ def PG_ISSO_eq_dr(r, chi, incl): chi8 = chi4 * chi4 chi10 = chi6 * chi4 return ( - 12 * r**11 - 132 * r**10 - + r**9 * (120 * chi2 * sin2i - 60 * chi2 + 360) - r**8 * 252 * chi2 - + 8 * r**7 * ( - 36 * chi4 * sin4i - 30 * chi4 * sin2i + 9 * chi4 - - 48 * chi2 * sin2i) + 12 * r**11 + - 132 * r**10 + + r**9 * (120 * chi2 * sin2i - 60 * chi2 + 360) + - r**8 * 252 * chi2 + + 8 + * r**7 + * (36 * chi4 * sin4i - 30 * chi4 * sin2i + 9 * chi4 - 48 * chi2 * sin2i) + 7 * r**6 * (120 * chi4 * sin2i - 144 * chi4 * sin4i) - + 6 * r**5 * ( - 36 * chi6 * sin4i - 16 * chi6 * sin2i + 144 * chi4 * sin4i - - 56 * chi4 * sin2i) + + 6 + * r**5 + * ( + 36 * chi6 * sin4i + - 16 * chi6 * sin2i + + 144 * chi4 * sin4i + - 56 * chi4 * sin2i + ) + r**4 * (120 * chi6 * sin2i - 240 * chi6 * sin4i) + r**3 * (84 * chi8 * sin4i - 24 * chi8 * sin2i - 192 * chi6 * sin4i) - 84 * r**2 * chi8 * sin4i - + r * (12 * chi10 * sin4i + 72 * chi8 * sin4i) - 12 * chi10 * sin4i) + + r * (12 * chi10 * sin4i + 72 * chi8 * sin4i) + - 12 * chi10 * sin4i + ) def PG_ISSO_eq_dr2(r, chi, incl): - """Second partial derivative of :func:`PG_ISSO_eq` with respect to + """ + Second partial derivative of :func:`PG_ISSO_eq` with respect to r. Parameters @@ -254,6 +270,7 @@ def PG_ISSO_eq_dr2(r, chi, incl): Returns ------- float + """ sini = np.sin(incl) sin2i = sini * sini @@ -263,24 +280,33 @@ def PG_ISSO_eq_dr2(r, chi, incl): chi6 = chi4 * chi2 chi8 = chi4 * chi4 return ( - 132 * r**10 - 1320 * r**9 - + 90 * r**8 * (12 * chi2 * sin2i - 6 * chi2 + 36) - 2016 * chi2 * r**7 - + 56 * r**6 * ( - 36 * chi4 * sin4i - 30 * chi4 * sin2i + 9 * chi4 - - 48 * chi2 * sin2i) + 132 * r**10 + - 1320 * r**9 + + 90 * r**8 * (12 * chi2 * sin2i - 6 * chi2 + 36) + - 2016 * chi2 * r**7 + + 56 + * r**6 + * (36 * chi4 * sin4i - 30 * chi4 * sin2i + 9 * chi4 - 48 * chi2 * sin2i) + 42 * r**5 * (120 * chi4 * sin2i - 144 * chi4 * sin4i) - + 30 * r**4 * ( - 36 * chi6 * sin4i - 16 * chi6 * sin2i + 144 * chi4 * sin4i - - 56 * chi4 * sin2i) + + 30 + * r**4 + * ( + 36 * chi6 * sin4i + - 16 * chi6 * sin2i + + 144 * chi4 * sin4i + - 56 * chi4 * sin2i + ) + r**3 * (480 * chi6 * sin2i - 960 * chi6 * sin4i) - + r**2 * ( - 252 * chi8 * sin4i - 72 * chi8 * sin2i - 576 * chi6 * sin4i) + + r**2 * (252 * chi8 * sin4i - 72 * chi8 * sin2i - 576 * chi6 * sin4i) - r * 168 * chi8 * sin4i - + 12 * chi8 * chi2 * sin4i + 72 * chi8 * sin4i) + + 12 * chi8 * chi2 * sin4i + + 72 * chi8 * sin4i + ) def PG_ISSO_solver(chi, incl): - """Function that determines the radius of the innermost stable + """ + Function that determines the radius of the innermost stable spherical orbit (ISSO) for a Kerr BH and a generic inclination angle between the BH spin and the orbital angular momentum. This function finds the appropriate root of :func:`PG_ISSO_eq`. @@ -297,6 +323,7 @@ def PG_ISSO_solver(chi, incl): ------- solution: array the radius of the orbit in BH mass units + """ # Auxiliary variables if np.isscalar(chi): @@ -314,13 +341,20 @@ def PG_ISSO_solver(chi, incl): # Initial guess is based on the extrema of the polar ISSO radius equation, # that are: r=6 (chi=1) and r=1+sqrt(3)+sqrt(3+sqrt(12))=5.274... (chi=0) initial_guess = [5.27451056440629 if c > 0.5 else 6 for c in chi] - rISSO_at_pole_limit = np.array([ - root_scalar( - ISSO_eq_at_pole, x0=g0, fprime=ISSO_eq_at_pole_dr, - fprime2=ISSO_eq_at_pole_dr2, args=(c)).root - for g0, c in zip(initial_guess, chi)]) + rISSO_at_pole_limit = np.array( + [ + root_scalar( + ISSO_eq_at_pole, + x0=g0, + fprime=ISSO_eq_at_pole_dr, + fprime2=ISSO_eq_at_pole_dr2, + args=(c), + ).root + for g0, c in zip(initial_guess, chi) + ] + ) # If the inclination is pi/2, just output the ISSO radius at the pole(s) - polar = np.isclose(incl, 0.5*np.pi) + polar = np.isclose(incl, 0.5 * np.pi) if all(polar): return rISSO_at_pole_limit @@ -328,24 +362,45 @@ def PG_ISSO_solver(chi, incl): initial_hi = np.maximum(rISCO_limit, rISSO_at_pole_limit) initial_lo = np.minimum(rISCO_limit, rISSO_at_pole_limit) brackets = [ - (bl, bh) if (c != 1 and PG_ISSO_eq(bl, c, inc) * - PG_ISSO_eq(bh, c, inc) < 0) else None - for bl, bh, c, inc in zip(initial_lo, initial_hi, chi, incl)] - solution = np.array([ - root_scalar( - PG_ISSO_eq, x0=g0, fprime=PG_ISSO_eq_dr, bracket=bracket, - fprime2=PG_ISSO_eq_dr2, args=(c, inc), xtol=1e-12).root - for g0, bracket, c, inc in zip(initial_hi, brackets, chi, incl)]) + (bl, bh) + if (c != 1 and PG_ISSO_eq(bl, c, inc) * PG_ISSO_eq(bh, c, inc) < 0) + else None + for bl, bh, c, inc in zip(initial_lo, initial_hi, chi, incl) + ] + solution = np.array( + [ + root_scalar( + PG_ISSO_eq, + x0=g0, + fprime=PG_ISSO_eq_dr, + bracket=bracket, + fprime2=PG_ISSO_eq_dr2, + args=(c, inc), + xtol=1e-12, + ).root + for g0, bracket, c, inc in zip(initial_hi, brackets, chi, incl) + ] + ) oob = (solution < 1) | (solution > 9) if any(oob): - solution = np.array([ - root_scalar( - PG_ISSO_eq, x0=g0, fprime=PG_ISSO_eq_dr, bracket=bracket, - fprime2=PG_ISSO_eq_dr2, args=(c, inc)).root - if ob else sol for g0, bracket, c, inc, ob, sol - in zip(initial_lo, brackets, chi, incl, oob, solution) - ]) + solution = np.array( + [ + root_scalar( + PG_ISSO_eq, + x0=g0, + fprime=PG_ISSO_eq_dr, + bracket=bracket, + fprime2=PG_ISSO_eq_dr2, + args=(c, inc), + ).root + if ob + else sol + for g0, bracket, c, inc, ob, sol in zip( + initial_lo, brackets, chi, incl, oob, solution + ) + ] + ) oob = (solution < 1) | (solution > 9) if any(oob): - raise RuntimeError('Unable to obtain some solutions!') + raise RuntimeError("Unable to obtain some solutions!") return solution diff --git a/pycbc/noise/__init__.py b/pycbc/noise/__init__.py index 76e0b193156..ce33b7e7e76 100644 --- a/pycbc/noise/__init__.py +++ b/pycbc/noise/__init__.py @@ -1 +1 @@ -from .gaussian import noise_from_psd, noise_from_string, frequency_noise_from_psd # noqa +from .gaussian import noise_from_psd, noise_from_string, frequency_noise_from_psd # noqa diff --git a/pycbc/noise/gaussian.py b/pycbc/noise/gaussian.py index 7082a1be9c9..0996a166547 100644 --- a/pycbc/noise/gaussian.py +++ b/pycbc/noise/gaussian.py @@ -23,20 +23,23 @@ # # ============================================================================= # -"""This module contains functions to generate gaussian noise colored with a +""" +This module contains functions to generate gaussian noise colored with a noise spectrum. """ -from pycbc import libutils -from pycbc.types import TimeSeries, zeros -from pycbc.types import complex_same_precision_as, FrequencySeries import lal import numpy.random -lalsimulation = libutils.import_optional('lalsimulation') +from pycbc import libutils +from pycbc.types import FrequencySeries, TimeSeries, complex_same_precision_as, zeros + +lalsimulation = libutils.import_optional("lalsimulation") + def frequency_noise_from_psd(psd, seed=None): - """ Create noise with a given psd. + """ + Create noise with a given psd. Return noise coloured with the given psd. The returned noise FrequencySeries has the same length and frequency step as the given psd. @@ -51,9 +54,10 @@ def frequency_noise_from_psd(psd, seed=None): the seed will not be reset. Returns - -------- + ------- noise : FrequencySeriesSeries A FrequencySeries containing gaussian noise colored by the given psd. + """ sigma = 0.5 * (psd / psd.delta_f) ** (0.5) if seed is not None: @@ -61,7 +65,7 @@ def frequency_noise_from_psd(psd, seed=None): sigma = sigma.numpy() dtype = complex_same_precision_as(psd) - not_zero = (sigma != 0) + not_zero = sigma != 0 sigma_red = sigma[not_zero] noise_re = numpy.random.normal(0, sigma_red) @@ -71,12 +75,12 @@ def frequency_noise_from_psd(psd, seed=None): noise = numpy.zeros(len(sigma), dtype=dtype) noise[not_zero] = noise_red - return FrequencySeries(noise, - delta_f=psd.delta_f, - dtype=dtype) + return FrequencySeries(noise, delta_f=psd.delta_f, dtype=dtype) + def noise_from_psd(length, delta_t, psd, seed=None): - """ Create noise with a given psd. + """ + Create noise with a given psd. Return noise with a given psd. Note that if unique noise is desired a unique seed should be provided. @@ -93,9 +97,10 @@ def noise_from_psd(length, delta_t, psd, seed=None): The seed to generate the noise. Returns - -------- + ------- noise : TimeSeries A TimeSeries containing gaussian noise colored by the given psd. + """ noise_ts = TimeSeries(zeros(length), delta_t=delta_t) @@ -104,34 +109,40 @@ def noise_from_psd(length, delta_t, psd, seed=None): randomness = lal.gsl_rng("ranlux", seed) - N = int (1.0 / delta_t / psd.delta_f) - n = N//2+1 - stride = N//2 + N = int(1.0 / delta_t / psd.delta_f) + n = N // 2 + 1 + stride = N // 2 if n > len(psd): raise ValueError("PSD not compatible with requested delta_t") psd = (psd[0:n]).lal() - psd.data.data[n-1] = 0 + psd.data.data[n - 1] = 0 psd.data.data[0] = 0 segment = TimeSeries(zeros(N), delta_t=delta_t).lal() length_generated = 0 lalsimulation.SimNoise(segment, 0, psd, randomness) - while (length_generated < length): + while length_generated < length: if (length_generated + stride) < length: - noise_ts.data[length_generated:length_generated+stride] = segment.data.data[0:stride] + noise_ts.data[length_generated : length_generated + stride] = ( + segment.data.data[0:stride] + ) else: - noise_ts.data[length_generated:length] = segment.data.data[0:length-length_generated] + noise_ts.data[length_generated:length] = segment.data.data[ + 0 : length - length_generated + ] length_generated += stride lalsimulation.SimNoise(segment, stride, psd, randomness) return noise_ts + def noise_from_string(psd_name, length, delta_t, seed=None, low_frequency_cutoff=10.0): - """ Create noise from an analytic PSD + """ + Create noise from an analytic PSD Return noise from the chosen PSD. Note that if unique noise is desired a unique seed should be provided. @@ -151,14 +162,15 @@ def noise_from_string(psd_name, length, delta_t, seed=None, low_frequency_cutoff The low frequency cutoff to pass to the PSD generation. Returns - -------- + ------- noise : TimeSeries A TimeSeries containing gaussian noise colored by the given psd. + """ import pycbc.psd # We just need enough resolution to resolve lines delta_f = 1.0 / 8 - flen = int(.5 / delta_t / delta_f) + 1 + flen = int(0.5 / delta_t / delta_f) + 1 psd = pycbc.psd.from_string(psd_name, flen, delta_f, low_frequency_cutoff) return noise_from_psd(int(length), delta_t, psd, seed=seed) diff --git a/pycbc/noise/reproduceable.py b/pycbc/noise/reproduceable.py index 03574da231a..799dec73929 100644 --- a/pycbc/noise/reproduceable.py +++ b/pycbc/noise/reproduceable.py @@ -21,15 +21,19 @@ # # ============================================================================= # -import numpy, pycbc.psd -from pycbc.types import TimeSeries, complex_same_precision_as +import numpy from numpy.random import RandomState +import pycbc.psd +from pycbc.types import TimeSeries, complex_same_precision_as + # This constant need to be constant to be able to recover identical results. BLOCK_SAMPLES = 1638400 + def block(seed, sample_rate): - """ Return block of normal random numbers + """ + Return block of normal random numbers Parameters ---------- @@ -39,17 +43,20 @@ def block(seed, sample_rate): Sets the variance of the white noise Returns - -------- + ------- noise : numpy.ndarray Array of random numbers + """ num = BLOCK_SAMPLES rng = RandomState(seed % 2**32) variance = sample_rate / 2 return rng.normal(size=num, scale=variance**0.5) + def normal(start, end, sample_rate=16384, seed=0): - """ Generate data with a white Gaussian (normal) distribution + """ + Generate data with a white Gaussian (normal) distribution Parameters ---------- @@ -64,9 +71,10 @@ def normal(start, end, sample_rate=16384, seed=0): The seed to generate the noise. Returns - -------- + ------- noise : TimeSeries A TimeSeries containing gaussian noise + """ # This is reproduceable because we used fixed seeds from known values block_dur = BLOCK_SAMPLES / sample_rate @@ -77,18 +85,26 @@ def normal(start, end, sample_rate=16384, seed=0): if end % block_dur == 0: e -= 1 - sv = RandomState(seed).randint(-2**50, 2**50) - data = numpy.concatenate([block(i + sv, sample_rate) - for i in numpy.arange(s, e + 1, 1)]) + sv = RandomState(seed).randint(-(2**50), 2**50) + data = numpy.concatenate( + [block(i + sv, sample_rate) for i in numpy.arange(s, e + 1, 1)] + ) ts = TimeSeries(data, delta_t=1.0 / sample_rate, epoch=(s * block_dur)) return ts.time_slice(start, end) -def colored_noise(psd, start_time, end_time, - seed=0, sample_rate=16384, - low_frequency_cutoff=1.0, - filter_duration=128, - scale=1.0): - """ Create noise from a PSD + +def colored_noise( + psd, + start_time, + end_time, + seed=0, + sample_rate=16384, + low_frequency_cutoff=1.0, + filter_duration=128, + scale=1.0, +): + """ + Create noise from a PSD Return noise from the chosen PSD. Note that if unique noise is desired a unique seed should be provided. @@ -112,9 +128,10 @@ def colored_noise(psd, start_time, end_time, The duration in seconds of the coloring filter Returns - -------- + ------- noise : TimeSeries A TimeSeries containing gaussian noise colored by the given psd. + """ psd = psd.copy() @@ -125,25 +142,26 @@ def colored_noise(psd, start_time, end_time, # Want to avoid zeroes in PSD. max_val = psd.max() for i in range(len(psd)): - if i >= (oldlen-1): + if i >= (oldlen - 1): psd.data[i] = psd[oldlen - 2] if psd[i] == 0: psd.data[i] = max_val fil_len = int(filter_duration * sample_rate) wn_dur = int(end_time - start_time) + 2 * filter_duration - if psd.delta_f >= 1. / (2.*filter_duration): + if psd.delta_f >= 1.0 / (2.0 * filter_duration): # If the PSD is short enough, this method is less memory intensive than # resizing and then calling inverse_spectrum_truncation - psd = pycbc.psd.interpolate(psd, 1.0 / (2. * filter_duration)) + psd = pycbc.psd.interpolate(psd, 1.0 / (2.0 * filter_duration)) # inverse_spectrum_truncation truncates the inverted PSD. To truncate # the non-inverted PSD we give it the inverted PSD to truncate and then # invert the output. - psd = 1. / pycbc.psd.inverse_spectrum_truncation( - 1./psd, - fil_len, - low_frequency_cutoff=low_frequency_cutoff, - trunc_method='hann') + psd = 1.0 / pycbc.psd.inverse_spectrum_truncation( + 1.0 / psd, + fil_len, + low_frequency_cutoff=low_frequency_cutoff, + trunc_method="hann", + ) psd = psd.astype(complex_same_precision_as(psd)) # Zero-pad the time-domain PSD to desired length. Zeroes must be added # in the middle, so some rolling between a resize is used. @@ -156,36 +174,45 @@ def colored_noise(psd, start_time, end_time, psd = psd.to_frequencyseries() else: psd = pycbc.psd.interpolate(psd, 1.0 / wn_dur) - psd = 1. / pycbc.psd.inverse_spectrum_truncation( - 1./psd, - fil_len, - low_frequency_cutoff=low_frequency_cutoff, - trunc_method='hann') + psd = 1.0 / pycbc.psd.inverse_spectrum_truncation( + 1.0 / psd, + fil_len, + low_frequency_cutoff=low_frequency_cutoff, + trunc_method="hann", + ) kmin = int(low_frequency_cutoff / psd.delta_f) psd[:kmin].clear() - asd = (psd.squared_norm())**0.25 + asd = (psd.squared_norm()) ** 0.25 del psd - white_noise = normal(start_time - filter_duration, - end_time + filter_duration, - seed=seed, - sample_rate=sample_rate) + white_noise = normal( + start_time - filter_duration, + end_time + filter_duration, + seed=seed, + sample_rate=sample_rate, + ) white_noise = white_noise.to_frequencyseries() # Here we color. Do not want to duplicate memory here though so use '*=' - white_noise *= asd*scale + white_noise *= asd * scale del asd - colored = white_noise.to_timeseries(delta_t=1.0/sample_rate) + colored = white_noise.to_timeseries(delta_t=1.0 / sample_rate) del white_noise return colored.time_slice(start_time, end_time) -def noise_from_string(psd_name, start_time, end_time, - seed=0, - sample_rate=16384, - low_frequency_cutoff=1.0, - filter_duration=128, - scale=1.0): - """ Create noise from an analytic PSD + +def noise_from_string( + psd_name, + start_time, + end_time, + seed=0, + sample_rate=16384, + low_frequency_cutoff=1.0, + filter_duration=128, + scale=1.0, +): + """ + Create noise from an analytic PSD Return noise from the chosen PSD. Note that if unique noise is desired a unique seed should be provided. @@ -209,16 +236,21 @@ def noise_from_string(psd_name, start_time, end_time, The duration in seconds of the coloring filter Returns - -------- + ------- noise : TimeSeries A TimeSeries containing gaussian noise colored by the given psd. + """ delta_f = 1.0 / filter_duration flen = int(sample_rate / delta_f) // 2 + 1 psd = pycbc.psd.from_string(psd_name, flen, delta_f, low_frequency_cutoff) - return colored_noise(psd, start_time, end_time, - seed=seed, - sample_rate=sample_rate, - low_frequency_cutoff=low_frequency_cutoff, - filter_duration=filter_duration, - scale=scale) + return colored_noise( + psd, + start_time, + end_time, + seed=seed, + sample_rate=sample_rate, + low_frequency_cutoff=low_frequency_cutoff, + filter_duration=filter_duration, + scale=scale, + ) diff --git a/pycbc/opt.py b/pycbc/opt.py index c4d7307ae5b..a3cb91f2cce 100644 --- a/pycbc/opt.py +++ b/pycbc/opt.py @@ -17,11 +17,13 @@ """ This module defines optimization flags and some optimized utilities. """ -import os, sys + import logging +import os +import sys from collections import OrderedDict -logger = logging.getLogger('pycbc.opt') +logger = logging.getLogger("pycbc.opt") def get_l2_cache_size(): @@ -32,6 +34,7 @@ def get_l2_cache_size(): ------- int or None The L2 cache size in bytes if the environment variable is set, None otherwise. + """ cache_size_str = os.environ.get("_PYCBC_L2_CACHE_SIZE", None) if cache_size_str is not None: @@ -47,21 +50,30 @@ def insert_optimization_option_group(parser): ---------- parser : object OptionParser instance + """ - optimization_group = parser.add_argument_group("Options for selecting " - "optimization-specific settings") + optimization_group = parser.add_argument_group( + "Options for selecting optimization-specific settings" + ) - optimization_group.add_argument("--cpu-affinity", help=""" + optimization_group.add_argument( + "--cpu-affinity", + help=""" A set of CPUs on which to run, specified in a format suitable - to pass to taskset.""") - optimization_group.add_argument("--cpu-affinity-from-env", help=""" + to pass to taskset.""", + ) + optimization_group.add_argument( + "--cpu-affinity-from-env", + help=""" The name of an enivornment variable containing a set of CPUs on which to run, specified in a format suitable - to pass to taskset.""") + to pass to taskset.""", + ) def verify_optimization_options(opt, parser): - """Parses the CLI options, verifies that they are consistent and + """ + Parses the CLI options, verifies that they are consistent and reasonable, and acts on them if they are Parameters @@ -71,16 +83,14 @@ def verify_optimization_options(opt, parser): required attributes parser : object OptionParser instance. - """ + """ # Pin to specified CPUs if requested requested_cpus = None if opt.cpu_affinity_from_env is not None: if opt.cpu_affinity is not None: - logger.error( - "Both --cpu_affinity_from_env and --cpu_affinity specified" - ) + logger.error("Both --cpu_affinity_from_env and --cpu_affinity specified") sys.exit(1) requested_cpus = os.environ.get(opt.cpu_affinity_from_env) @@ -89,15 +99,15 @@ def verify_optimization_options(opt, parser): logger.error( "CPU affinity requested from environment variable %s " "but this variable is not defined", - opt.cpu_affinity_from_env + opt.cpu_affinity_from_env, ) sys.exit(1) - if requested_cpus == '': + if requested_cpus == "": logger.error( "CPU affinity requested from environment variable %s " "but this variable is empty", - opt.cpu_affinity_from_env + opt.cpu_affinity_from_env, ) sys.exit(1) @@ -105,20 +115,20 @@ def verify_optimization_options(opt, parser): requested_cpus = opt.cpu_affinity if requested_cpus is not None: - command = 'taskset -pc %s %d' % (requested_cpus, os.getpid()) + command = "taskset -pc %s %d" % (requested_cpus, os.getpid()) retcode = os.system(command) if retcode != 0: logger.error( - 'taskset command <%s> failed with return code %d', - command, retcode + "taskset command <%s> failed with return code %d", command, retcode ) sys.exit(1) logger.info("Pinned to CPUs %s ", requested_cpus) + class LimitedSizeDict(OrderedDict): - """ Fixed sized dict for FIFO caching""" + """Fixed sized dict for FIFO caching""" def __init__(self, *args, **kwds): self.size_limit = kwds.pop("size_limit", None) diff --git a/pycbc/pnutils.py b/pycbc/pnutils.py index 0dea2310bdc..5394cf5fcbe 100644 --- a/pycbc/pnutils.py +++ b/pycbc/pnutils.py @@ -23,17 +23,18 @@ # # ============================================================================= # -"""This module contains convenience pN functions. This includes calculating conversions +""" +This module contains convenience pN functions. This includes calculating conversions between quantities. """ import logging -import numpy +import numpy from scipy.optimize import bisect, brentq, minimize from pycbc import conversions, libutils -from pycbc.constants import MSUN_SI, PI, MTSUN_SI, PC_SI +from pycbc.constants import MSUN_SI, MTSUN_SI, PC_SI, PI logger = logging.getLogger("pycbc.pnutils") @@ -111,7 +112,8 @@ def eta_mass1_to_mass2(eta, mass1, return_mass_heavier=False, force_real=True): def mchirp_q_to_mass1_mass2(mchirp, q): - """This function takes a value of mchirp and the mass ratio + """ + This function takes a value of mchirp and the mass ratio mass1/mass2 and returns the two component masses. The map from q to eta is @@ -127,14 +129,15 @@ def mchirp_q_to_mass1_mass2(mchirp, q): def A0(f_lower): - """used in calculating chirp times: see Cokelaer, arxiv.org:0706.4437 + """ + Used in calculating chirp times: see Cokelaer, arxiv.org:0706.4437 appendix 1, also lalinspiral/python/sbank/tau0tau3.py """ return conversions._a0(f_lower) def A3(f_lower): - """another parameter used for chirp times""" + """Another parameter used for chirp times""" return conversions._a3(f_lower) @@ -171,7 +174,7 @@ def get_beta_sigma_from_aligned_spins(eta, spin1z, spin2z): See . Parameters - ----------- + ---------- eta : float or numpy.array Symmetric mass ratio of the input system(s) spin1z : float or numpy.array @@ -180,7 +183,7 @@ def get_beta_sigma_from_aligned_spins(eta, spin1z, spin2z): Spin(s) parallel to the orbit of the smallest body(ies) Returns - -------- + ------- beta : float or numpy.array The 1.5PN spin combination sigma : float or numpy.array @@ -189,6 +192,7 @@ def get_beta_sigma_from_aligned_spins(eta, spin1z, spin2z): The 2.5PN spin combination chis : float or numpy.array (spin1z + spin2z) / 2. + """ chiS = 0.5 * (spin1z + spin2z) chiA = 0.5 * (spin1z - spin2z) @@ -238,6 +242,7 @@ def f_SchwarzISCO(M): ------- f : float or numpy.array Frequency in Hz + """ return conversions.f_schwarzchild_isco(M) @@ -259,6 +264,7 @@ def f_BKLISCO(m1, m2): ------- f : float or numpy.array Frequency in Hz + """ # q is defined to be in [0,1] for this formula q = numpy.minimum(m1 / m2, m2 / m1) @@ -279,6 +285,7 @@ def f_LightRing(M): ------- f : float or numpy.array Frequency in Hz + """ return 1.0 / (3.0 ** (1.5) * PI * M * MTSUN_SI) @@ -299,6 +306,7 @@ def f_ERD(M): ------- f : float or numpy.array Frequency in Hz + """ return 1.07 * 0.5326 / (2 * PI * 0.955 * M * MTSUN_SI) @@ -321,6 +329,7 @@ def f_FRD(m1, m2): ------- f : float or numpy.array Frequency in Hz + """ m_total, eta = mass1_mass2_to_mtotal_eta(m1, m2) tmp = (1.0 - 0.63 * (1.0 - 3.4641016 * eta + 2.9 * eta**2) ** (0.3)) / ( @@ -345,12 +354,14 @@ def f_LRD(m1, m2): ------- f : float or numpy.array Frequency in Hz + """ return 1.2 * f_FRD(m1, m2) def _get_freq(freqfunc, m1, m2, s1z, s2z): - """Wrapper of the LALSimulation function returning the frequency + """ + Wrapper of the LALSimulation function returning the frequency for a given frequency function and template parameters. Parameters @@ -370,6 +381,7 @@ def _get_freq(freqfunc, m1, m2, s1z, s2z): ------- f : float Frequency in Hz + """ return lalsim.SimInspiralGetFrequency( solar_mass_to_kg(m1), @@ -410,13 +422,15 @@ def get_freq(freqfunc, m1, m2, s1z, s2z): ------- f : float or numpy.array Frequency in Hz + """ lalsim_ffunc = getattr(lalsim, freqfunc) return _vec_get_freq(lalsim_ffunc, m1, m2, s1z, s2z) def _get_final_freq(approx, m1, m2, s1z, s2z): - """Wrapper of the LALSimulation function returning the final (highest) + """ + Wrapper of the LALSimulation function returning the final (highest) frequency for a given approximant an template parameters Parameters @@ -436,6 +450,7 @@ def _get_final_freq(approx, m1, m2, s1z, s2z): ------- f : float Frequency in Hz + """ return lalsim.SimInspiralGetFinalFreq( solar_mass_to_kg(m1), @@ -455,7 +470,8 @@ def _get_final_freq(approx, m1, m2, s1z, s2z): def get_final_freq(approx, m1, m2, s1z, s2z): - """Returns the final (highest) frequency for a given approximant using + """ + Returns the final (highest) frequency for a given approximant using given template parameters. NOTE: TaylorTx and TaylorFx are currently all given an ISCO cutoff !! @@ -477,6 +493,7 @@ def get_final_freq(approx, m1, m2, s1z, s2z): ------- f : float or numpy.array Frequency in Hz + """ # Unfortunately we need a few special cases (quite hacky in the case of # IMRPhenomXAS) because some useful approximants are not understood by @@ -569,6 +586,7 @@ def frequency_cutoff_from_name(name, m1, m2, s1z, s2z): ------- f : float or numpy.array Frequency in Hz + """ params = {"mass1": m1, "mass2": m2, "spin1z": s1z, "spin2z": s2z} return named_frequency_cutoffs[name](params) @@ -636,7 +654,8 @@ def get_inspiral_tf( pn_2order=7, approximant="TaylorF2", ): - """Compute the time-frequency evolution of an inspiral signal. + """ + Compute the time-frequency evolution of an inspiral signal. Return a tuple of time and frequency vectors tracking the evolution of an inspiral signal in the time-frequency plane. @@ -764,6 +783,7 @@ def meco_velocity(m1, m2, chi1, chi2): ------- v : float Velocity (dimensionless) + """ _, energy2, energy3, energy4, energy5, energy6 = _energy_coeffs(m1, m2, chi1, chi2) @@ -837,8 +857,7 @@ def dtdv_func(v): if dtdv_func(1.0) < 0.0: return bisect(dtdv_func, 0.05, 1.0) - else: - return 1.0 + return 1.0 def energy_coefficients(m1, m2, s1z=0, s2z=0, phase_order=-1, spin_order=-1): @@ -847,12 +866,12 @@ def energy_coefficients(m1, m2, s1z=0, s2z=0, phase_order=-1, spin_order=-1): implemented_spin_order = 7 if phase_order > implemented_phase_order: raise ValueError("pN coeffiecients of that order have not been implemented") - elif phase_order == -1: + if phase_order == -1: phase_order = implemented_phase_order if spin_order > implemented_spin_order: raise ValueError("pN coeffiecients of that order have not been implemented") - elif spin_order == -1: + if spin_order == -1: spin_order = implemented_spin_order qmdef1 = 1.0 @@ -985,12 +1004,12 @@ def kerr_lightring_velocity(chi): # If chi > 0.9996, the algorithm cannot solve the function if chi >= 0.9996: return brentq(kerr_lightring, 0, 0.8, args=(0.9996)) - else: - return brentq(kerr_lightring, 0, 0.8, args=(chi)) + return brentq(kerr_lightring, 0, 0.8, args=(chi)) def hybridEnergy(v, m1, m2, chi1, chi2, qm1, qm2): - """Return hybrid MECO energy. + """ + Return hybrid MECO energy. Return the hybrid energy [eq. (6)] whose minimum defines the hybrid MECO up to 3.5PN (including the 3PN spin-spin) @@ -1014,6 +1033,7 @@ def hybridEnergy(v, m1, m2, chi1, chi2, qm1, qm2): ------- h_E: float The hybrid energy as a function of v + """ pi_sq = numpy.pi**2 v2, v3, v4, v5, v6, v7 = v**2, v**3, v**4, v**5, v**6, v**7 @@ -1092,7 +1112,8 @@ def hybridEnergy(v, m1, m2, chi1, chi2, qm1, qm2): def hybrid_meco_velocity(m1, m2, chi1, chi2, qm1=None, qm2=None): - """Return the velocity of the hybrid MECO + """ + Return the velocity of the hybrid MECO Parameters ---------- @@ -1115,8 +1136,8 @@ def hybrid_meco_velocity(m1, m2, chi1, chi2, qm1=None, qm2=None): ------- v: float The velocity (dimensionless) of the hybrid MECO - """ + """ if qm1 is None: qm1 = 1 if qm2 is None: @@ -1132,7 +1153,8 @@ def hybrid_meco_velocity(m1, m2, chi1, chi2, qm1=None, qm2=None): def hybrid_meco_frequency(m1, m2, chi1, chi2, qm1=None, qm2=None): - """Return the frequency of the hybrid MECO + """ + Return the frequency of the hybrid MECO Parameters ---------- @@ -1155,6 +1177,7 @@ def hybrid_meco_frequency(m1, m2, chi1, chi2, qm1=None, qm2=None): ------- f: float The frequency (in Hz) of the hybrid MECO + """ if qm1 is None: qm1 = 1 @@ -1179,7 +1202,8 @@ def jframe_to_l0frame( spin2_polar=0.0, spin12_deltaphi=0.0, ): - """Converts J-frame parameters into L0 frame. + """ + Converts J-frame parameters into L0 frame. Parameters ---------- @@ -1236,6 +1260,7 @@ def jframe_to_l0frame( * spin2z : float The z component of the second binary component's dimensionless spin. + """ inclination, spin1x, spin1y, spin1z, spin2x, spin2y, spin2z = ( lalsim.SimInspiralTransformPrecessingNewInitialConditions( @@ -1277,7 +1302,8 @@ def l0frame_to_jframe( spin2y=0.0, spin2z=0.0, ): - """Converts L0-frame parameters to J-frame. + """ + Converts L0-frame parameters to J-frame. Parameters ---------- @@ -1334,6 +1360,7 @@ def l0frame_to_jframe( * spin12_deltaphi : float Difference between the azimuthal angles of the spin of the larger object (S1) and the spin of the smaller object (S2). + """ # Note: unlike other LALSimulation functions, this one takes masses in # solar masses diff --git a/pycbc/pool.py b/pycbc/pool.py index 1e3539eea1d..c15af38ed87 100644 --- a/pycbc/pool.py +++ b/pycbc/pool.py @@ -1,26 +1,28 @@ -""" Tools for creating pools of worker processes -""" -import multiprocessing.pool -import functools -from multiprocessing import TimeoutError, cpu_count, get_context -import types -import signal +"""Tools for creating pools of worker processes""" + import atexit +import functools import logging +import multiprocessing.pool +import signal +import types +from multiprocessing import TimeoutError, cpu_count, get_context + +logger = logging.getLogger("pycbc.pool") -logger = logging.getLogger('pycbc.pool') def is_main_process(): - """ Check if this is the main control process and may handle one time tasks - """ + """Check if this is the main control process and may handle one time tasks""" try: from mpi4py import MPI + comm = MPI.COMM_WORLD rank = comm.Get_rank() return rank == 0 except (ImportError, ValueError, RuntimeError): return True + # Allow the pool to be interrupted, need to disable the children processes # from intercepting the keyboard interrupt def _noint(init, *args): @@ -28,10 +30,13 @@ def _noint(init, *args): if init is not None: return init(*args) + _process_lock = None _numdone = None + + def _lockstep_fcn(values): - """ Wrapper to ensure that all processes execute together """ + """Wrapper to ensure that all processes execute together""" numrequired, fcn, args = values with _process_lock: _numdone.value += 1 @@ -42,34 +47,39 @@ def _lockstep_fcn(values): if _numdone.value == numrequired: return fcn(args) + def _shutdown_pool(p): p.terminate() p.join() + class BroadcastPool(multiprocessing.pool.Pool): - """ Multiprocessing pool with a broadcast method - """ - def __init__(self, processes=None, initializer=None, initargs=(), - context=None, **kwds): + """Multiprocessing pool with a broadcast method""" + + def __init__( + self, processes=None, initializer=None, initargs=(), context=None, **kwds + ): global _process_lock global _numdone _process_lock = multiprocessing.Lock() - _numdone = multiprocessing.Value('i', 0) + _numdone = multiprocessing.Value("i", 0) noint = functools.partial(_noint, initializer) # Default is fork to preserve child memory inheritance and # copy on write if context is None: context = get_context("fork") - super(BroadcastPool, self).__init__(processes, noint, initargs, - context=context, **kwds) + super().__init__( + processes, noint, initargs, context=context, **kwds + ) atexit.register(_shutdown_pool, self) def __len__(self): return len(self._pool) def broadcast(self, fcn, args): - """ Do a function call on every worker. + """ + Do a function call on every worker. Parameters ---------- @@ -77,13 +87,15 @@ def broadcast(self, fcn, args): Function to call. args: tuple The arguments for Pool.map + """ results = self.map(_lockstep_fcn, [(len(self), fcn, args)] * len(self)) _numdone.value = 0 return results def allmap(self, fcn, args): - """ Do a function call on every worker with different arguments + """ + Do a function call on every worker with different arguments Parameters ---------- @@ -91,14 +103,15 @@ def allmap(self, fcn, args): Function to call. args: tuple The arguments for Pool.map + """ - results = self.map(_lockstep_fcn, - [(len(self), fcn, arg) for arg in args]) + results = self.map(_lockstep_fcn, [(len(self), fcn, arg) for arg in args]) _numdone.value = 0 return results def map(self, func, items, chunksize=None): - """ Catch keyboard interrupts to allow the pool to exit cleanly. + """ + Catch keyboard interrupts to allow the pool to exit cleanly. Parameters ---------- @@ -108,6 +121,7 @@ def map(self, func, items, chunksize=None): Arguments to pass chunksize: int, Optional Number of calls for each process to handle at once + """ results = self.map_async(func, items, chunksize) while True: @@ -121,17 +135,17 @@ def map(self, func, items, chunksize=None): raise KeyboardInterrupt def close_pool(self): - """ Close the pool and remove the reference - """ + """Close the pool and remove the reference""" self.close() self.join() atexit.unregister(_shutdown_pool) + def _dummy_broadcast(self, f, args): self.map(f, [args] * self.size) -class SinglePool(object): +class SinglePool: def __init__(self, **_): pass @@ -144,31 +158,27 @@ def map(self, f, items): # This is single core, so imap and map # would not behave differently. This is defined # so that the general pool interfaces can use - # imap irrespective of the pool type. + # imap irrespective of the pool type. imap = map imap_unordered = map def close_pool(self): - ''' Dummy function to be consistent with BroadcastPool - ''' - pass + """Dummy function to be consistent with BroadcastPool""" + def use_mpi(require_mpi=False, log=True): - """ Get whether MPI is enabled and if so the current size and rank - """ + """Get whether MPI is enabled and if so the current size and rank""" use_mpi = False try: from mpi4py import MPI + comm = MPI.COMM_WORLD size = comm.Get_size() rank = comm.Get_rank() if size > 1: use_mpi = True if log: - logger.info( - 'Running under mpi with size: %s, rank: %s', - size, rank - ) + logger.info("Running under mpi with size: %s, rank: %s", size, rank) except ImportError as e: if require_mpi: print(e) @@ -179,7 +189,8 @@ def use_mpi(require_mpi=False, log=True): def choose_pool(processes, mpi=False, **kwargs): - """ Get processing pool. + """ + Get processing pool. Keyword arguments are passed to the pool constructor. """ @@ -187,22 +198,23 @@ def choose_pool(processes, mpi=False, **kwargs): if do_mpi: try: import schwimmbad - pool = schwimmbad.choose_pool(mpi=do_mpi, - processes=(size - 1), - **kwargs) + + pool = schwimmbad.choose_pool(mpi=do_mpi, processes=(size - 1), **kwargs) pool.broadcast = types.MethodType(_dummy_broadcast, pool) atexit.register(pool.close) if processes: - logger.info('NOTE: that for MPI process size determined by ' - 'MPI launch size, not the processes argument') + logger.info( + "NOTE: that for MPI process size determined by " + "MPI launch size, not the processes argument" + ) if do_mpi and not mpi: - logger.info('NOTE: using MPI as this process was launched' - 'under MPI') + logger.info("NOTE: using MPI as this process was launchedunder MPI") except ImportError: - raise ValueError("Failed to start up an MPI pool, " - "install mpi4py / schwimmbad") + raise ValueError( + "Failed to start up an MPI pool, install mpi4py / schwimmbad" + ) elif processes == 1: pool = SinglePool(**kwargs) else: @@ -214,4 +226,3 @@ def choose_pool(processes, mpi=False, **kwargs): if size: pool.size = size return pool - diff --git a/pycbc/population/__init__.py b/pycbc/population/__init__.py index 5f50880e721..1ee80546180 100644 --- a/pycbc/population/__init__.py +++ b/pycbc/population/__init__.py @@ -1,7 +1,7 @@ -from pycbc.population.rates_functions import * -from pycbc.population.scale_injections import * -from pycbc.population.population_models import * from pycbc.population.fgmc_functions import * from pycbc.population.fgmc_laguerre import * from pycbc.population.live_pastro import * from pycbc.population.live_pastro_utils import * +from pycbc.population.population_models import * +from pycbc.population.rates_functions import * +from pycbc.population.scale_injections import * diff --git a/pycbc/population/fgmc_functions.py b/pycbc/population/fgmc_functions.py index b5466adaa9f..ef6a070f16c 100644 --- a/pycbc/population/fgmc_functions.py +++ b/pycbc/population/fgmc_functions.py @@ -11,76 +11,83 @@ See https://dcc.ligo.org/LIGO-T2100060/public for technical explanations """ -from os.path import basename import bisect -from itertools import chain as it_chain, combinations as it_comb +from itertools import chain as it_chain +from itertools import combinations as it_comb +from os.path import basename + import numpy as np from pycbc import conversions as conv from pycbc import events -from pycbc.events.coinc import mean_if_greater_than_zero as coinc_meanigz from pycbc.events import triggers +from pycbc.events.coinc import mean_if_greater_than_zero as coinc_meanigz from pycbc.io.hdf import HFile def filter_bin_lo_hi(values, lo, hi): in_bin = np.sign((values - lo) * (hi - values)) if np.any(in_bin == 0): - raise RuntimeError('Edge case! Bin edges', lo, hi, - 'value(s)', values[in_bin == 0]) + raise RuntimeError( + "Edge case! Bin edges", lo, hi, "value(s)", values[in_bin == 0] + ) return in_bin == 1 def filter_tmplt_mchirp(bankf, lo_mchirp, hi_mchirp): with HFile(bankf) as bank: - mchirp = conv.mchirp_from_mass1_mass2(bank['mass1'][:], bank['mass2'][:]) + mchirp = conv.mchirp_from_mass1_mass2(bank["mass1"][:], bank["mass2"][:]) # Boolean over template id return filter_bin_lo_hi(mchirp, lo_mchirp, hi_mchirp) def read_full_data(fullf, rhomin, tmplt_filter=None): - """Read the zero- and time-lagged triggers identified by a specific - set of templates. - - Parameters - ---------- - fullf: - File that stores zerolag and slide triggers - bankf: - File with template mass/spin information - rhomin: float - Ranking statistic threshold - tmplt_filter: array of Booleans - Filter over the array of templates stored in bankf - - Returns - ------- - dictionary - containing foreground triggers and background information """ - with HFile(fullf, 'r') as full_data: + Read the zero- and time-lagged triggers identified by a specific + set of templates. + + Parameters + ---------- + fullf: + File that stores zerolag and slide triggers + bankf: + File with template mass/spin information + rhomin: float + Ranking statistic threshold + tmplt_filter: array of Booleans + Filter over the array of templates stored in bankf + + Returns + ------- + dictionary + containing foreground triggers and background information + + """ + with HFile(fullf, "r") as full_data: # apply template filter - tid_bkg = full_data['background_exc/template_id'][:] - tid_fg = full_data['foreground/template_id'][:] + tid_bkg = full_data["background_exc/template_id"][:] + tid_fg = full_data["foreground/template_id"][:] bkg_inbin = tmplt_filter[tid_bkg] # Boolean over bg events fg_inbin = tmplt_filter[tid_fg] # Boolean over fg events - zerolagstat = full_data['foreground/stat'][:][fg_inbin] - zerolagifar = full_data['foreground/ifar'][:][fg_inbin] + zerolagstat = full_data["foreground/stat"][:][fg_inbin] + zerolagifar = full_data["foreground/ifar"][:][fg_inbin] # arbitrarily choose time from one of the ifos - zerolagtime = full_data['foreground/time1'][:][fg_inbin] + zerolagtime = full_data["foreground/time1"][:][fg_inbin] - cstat_back_exc = full_data['background_exc/stat'][:][bkg_inbin] - dec_factors = full_data['background_exc/decimation_factor'][:][bkg_inbin] + cstat_back_exc = full_data["background_exc/stat"][:][bkg_inbin] + dec_factors = full_data["background_exc/decimation_factor"][:][bkg_inbin] # filter on stat value above = zerolagstat > rhomin back_above = cstat_back_exc > rhomin - return {'zerolagstat': zerolagstat[above], - 'zerolagifar': zerolagifar[above], - 'zerolagtime': zerolagtime[above], - 'dec_factors': dec_factors[back_above], - 'cstat_back_exc': cstat_back_exc[back_above], - 'file_name': fullf} + return { + "zerolagstat": zerolagstat[above], + "zerolagifar": zerolagifar[above], + "zerolagtime": zerolagtime[above], + "dec_factors": dec_factors[back_above], + "cstat_back_exc": cstat_back_exc[back_above], + "file_name": fullf, + } def read_full_data_mchirp(fullf, bankf, rhomin, mc_lo, mc_hi): @@ -94,9 +101,11 @@ def log_rho_bg(trigs, counts, bins): counts: background histogram bins: bin edges of the background histogram - Returns: + Returns + ------- log of background PDF at the zerolag statistic values, fractional uncertainty due to Poisson count (set to 100% for empty bins) + """ trigs = np.atleast_1d(trigs) if len(trigs) == 0: # corner case @@ -120,22 +129,23 @@ def log_rho_bg(trigs, counts, bins): # a bin that extends from the limits of the slide triggers out to the # loudest trigger. Fractional error is 100% log_rhos.append(-np.log(N) - np.log(np.max(trigs) - bins[-1])) - fracerr.append(1.) + fracerr.append(1.0) else: i = bisect.bisect(bins, t) - 1 # If there are no counts for a foreground trigger put a fictitious # count in the background bin if counts[i] == 0: counts[i] = 1 - log_rhos.append(np.log(counts[i]) - np.log(bins[i+1] - bins[i]) - - np.log(N)) + log_rhos.append( + np.log(counts[i]) - np.log(bins[i + 1] - bins[i]) - np.log(N) + ) fracerr.append(counts[i] ** -0.5) return np.array(log_rhos), np.array(fracerr) def log_rho_fg_analytic(trigs, rhomin): # PDF of a rho^-4 distribution defined above the threshold rhomin - return np.log(3.) + 3. * np.log(rhomin) - 4 * np.log(trigs) + return np.log(3.0) + 3.0 * np.log(rhomin) - 4 * np.log(trigs) def log_rho_fg(trigs, injstats, bins): @@ -144,9 +154,11 @@ def log_rho_fg(trigs, injstats, bins): injstats: injection event statistic values bins: histogram bins - Returns: + Returns + ------- log of signal PDF at the zerolag statistic values, fractional uncertainty from Poisson count + """ trigs = np.atleast_1d(trigs) if len(trigs) == 0: # corner case @@ -156,7 +168,7 @@ def log_rho_fg(trigs, injstats, bins): # allow 'very loud' triggers bmax = np.max(bins) if np.max(trigs) >= bmax: - print('Replacing stat values lying above highest bin') + print("Replacing stat values lying above highest bin") print(str(bmax)) trigs = np.where(trigs >= bmax, bmax - 1e-6, trigs) assert np.max(trigs) < np.max(bins) # check it worked @@ -164,7 +176,7 @@ def log_rho_fg(trigs, injstats, bins): counts, bins = np.histogram(injstats, bins) N = sum(counts) dens = counts / np.diff(bins) / float(N) - fracerr = counts ** -0.5 + fracerr = counts**-0.5 tinds = np.searchsorted(bins, trigs) - 1 return np.log(dens[tinds]), fracerr[tinds] @@ -173,23 +185,22 @@ def log_rho_fg(trigs, injstats, bins): def get_start_dur(path): fname = basename(path) # remove directory path # file name is IFOS-DESCRIPTION-START-DURATION.type - pieces = fname.split('.')[0].split('-') + pieces = fname.split(".")[0].split("-") return pieces[2], pieces[3] def in_coinc_time_incl(f, cstring, test_times): - """ filter to all times where coincs of type given by cstring exist - """ + """Filter to all times where coincs of type given by cstring exist""" in_time = np.zeros(len(test_times)) - starts = np.array(f['segments/%s/start' % cstring][:]) - ends = np.array(f['segments/%s/end' % cstring][:]) + starts = np.array(f["segments/%s/start" % cstring][:]) + ends = np.array(f["segments/%s/end" % cstring][:]) idx_within_segment = events.indices_within_times(test_times, starts, ends) in_time[idx_within_segment] = np.ones_like(idx_within_segment) return in_time # what to change for more/fewer ifos -_ifoset = ('H1', 'L1', 'V1') +_ifoset = ("H1", "L1", "V1") # all combinations of ifos with length mincount or more @@ -197,18 +208,19 @@ def in_coinc_time_incl(f, cstring, test_times): def alltimes(ifos, mincount=1): assert mincount <= len(ifos) assert len(set(ifos)) == len(ifos) # can't work with duplicates - return it_chain.from_iterable(it_comb(ifos, r) for r in - np.arange(mincount, len(ifos) + 1)) + return it_chain.from_iterable( + it_comb(ifos, r) for r in np.arange(mincount, len(ifos) + 1) + ) _alltimes = frozenset(alltimes(_ifoset, mincount=1)) -_alltimestring = frozenset([''.join(t) for t in _alltimes]) +_alltimestring = frozenset(["".join(t) for t in _alltimes]) _allctimes = frozenset(alltimes(_ifoset, mincount=2)) def ifos_from_combo(ct): # extract ifos in alphabetical order from a coinc time string - return sorted([ct[i:i + 2] for i in range(0, len(ct), 2)]) + return sorted([ct[i : i + 2] for i in range(0, len(ct), 2)]) def type_in_time(ct, cty): @@ -216,25 +228,32 @@ def type_in_time(ct, cty): return all(i in ct for i in cty) -class EventRate(object): - def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', - bin_lo=None, bin_hi=None): +class EventRate: + def __init__( + self, + args, + coinc_times, + coinc_types=None, + bin_param="mchirp", + bin_lo=None, + bin_hi=None, + ): """ coinc_times: iterable of strings indicating combinations of ifos operating coinc_types: list of strings indicating coinc event types to be considered """ # allow for single-ifo time although not supported in pipeline yet - if hasattr(args, 'min_ifos'): + if hasattr(args, "min_ifos"): self.mincount = args.min_ifos else: self.mincount = 2 - if hasattr(args, 'network') and sorted(args.network) != list(_ifoset): + if hasattr(args, "network") and sorted(args.network) != list(_ifoset): self.ifos = sorted(args.network) else: self.ifos = _ifoset self.allctimes = frozenset(alltimes(self.ifos, mincount=self.mincount)) - self.allctimestring = frozenset([''.join(t) for t in self.allctimes]) + self.allctimestring = frozenset(["".join(t) for t in self.allctimes]) for ct in coinc_times: assert ct in list(self.allctimestring) self.ctimes = coinc_times @@ -246,10 +265,9 @@ def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', # any coinc type must also be a time (?) for ty in coinc_types: assert ty in list(self.allctimes) - self.coinc_types = frozenset([''.join(t) for t in coinc_types]) + self.coinc_types = frozenset(["".join(t) for t in coinc_types]) if args.verbose: - print('Using', self.coinc_types, 'coincs in', - self.allctimestring, 'times') + print("Using", self.coinc_types, "coincs in", self.allctimestring, "times") self.args = args self.thr = self.args.stat_threshold self.bin_param = bin_param @@ -265,29 +283,28 @@ def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', def add_bank(self, bank_file): self.bank = bank_file - with HFile(self.bank, 'r') as b: - tids = np.arange(len(b['mass1'][:])) + with HFile(self.bank, "r") as b: + tids = np.arange(len(b["mass1"][:])) # tuples of m1, m2, s1z, s2z in template id order self.massspins = triggers.get_mass_spin(b, tids) def filter_templates(self): """ - calculate array of Booleans in template id order to filter events + Calculate array of Booleans in template id order to filter events """ assert self.massspins is not None assert self.lo is not None assert self.hi is not None if self.args.verbose: - print('Cutting on %s between %f - %f' % - (self.bin_param, self.lo, self.hi)) + print("Cutting on %s between %f - %f" % (self.bin_param, self.lo, self.hi)) self.tpars = triggers.get_param(self.bin_param, None, *self.massspins) self.in_bin = filter_bin_lo_hi(self.tpars, self.lo, self.hi) - def make_bins(self, maxval, choice='bg'): + def make_bins(self, maxval, choice="bg"): # allow options to be strings describing bin formulae as well as floats? try: - linbw = getattr(self.args, choice + '_bin_width') - logbw = getattr(self.args, choice + '_log_bin_width') + linbw = getattr(self.args, choice + "_bin_width") + logbw = getattr(self.args, choice + "_log_bin_width") except AttributeError: pass if linbw is not None: @@ -295,13 +312,15 @@ def make_bins(self, maxval, choice='bg'): bins = np.linspace(self.thr - 0.0001, maxval, n_bins + 1) elif logbw is not None: n_bins = int(np.log(maxval / self.thr) / float(logbw)) - bins = np.logspace(np.log10(self.thr) - 0.0001, np.log10(maxval), - n_bins + 1) + bins = np.logspace( + np.log10(self.thr) - 0.0001, np.log10(maxval), n_bins + 1 + ) else: - raise RuntimeError("Can't make bins without a %s bin width option!" - % choice) + raise RuntimeError( + "Can't make bins without a %s bin width option!" % choice + ) if self.args.verbose: - print(str(n_bins) + ' ' + choice + ' stat bins') + print(str(n_bins) + " " + choice + " stat bins") return bins def get_ctypes(self, tdict): @@ -310,14 +329,15 @@ def get_ctypes(self, tdict): cty = [] for ts in ifotimes: # if an ifo doesn't participate, time is sentinel value -1 - cty.append(''.join([i for i, t in zip(self.ifos, ts) if t > 0])) + cty.append("".join([i for i, t in zip(self.ifos, ts) if t > 0])) # return is array of coinc types strings return np.array(cty) def moreifotimes(self, ctstring): # get list of coinc times with more ifos than ctstring - allctime_moreifos = [ct for ct in self.allctimestring if - len(ct) > len(ctstring)] + allctime_moreifos = [ + ct for ct in self.allctimestring if len(ct) > len(ctstring) + ] # only return those when at least the same ifos are operating ret = [] ifos = ifos_from_combo(ctstring) @@ -327,8 +347,7 @@ def moreifotimes(self, ctstring): return ret def in_coinc_time_excl(self, f, cstring, test_times): - """ filter to all times where exactly the ifos in cstring are observing - """ + """Filter to all times where exactly the ifos in cstring are observing""" if len(cstring) == max(len(s) for s in self.allctimestring): # ctime string already uniquely specifies time return in_coinc_time_incl(f, cstring, test_times) @@ -343,26 +362,40 @@ def in_coinc_time_excl(self, f, cstring, test_times): return in_time def get_livetimes(self, fi): - with HFile(fi, 'r') as f: + with HFile(fi, "r") as f: for ct in self.ctimes: # 'inclusive' time when at least the ifos specified by ct are on - fgt = conv.sec_to_year(f[ct].attrs['foreground_time']) + fgt = conv.sec_to_year(f[ct].attrs["foreground_time"]) # index dict on chunk start time / coinc type self.incl_livetimes[(get_start_dur(fi)[0], ct)] = fgt # subtract times during which 1 more ifo was on, # ie subtract H1L1* time from H1L1; subtract H1* time from H1; etc for combo in self.moreifotimes(ct): if len(combo) == len(ct) + 2: - fgt -= conv.sec_to_year(f[combo].attrs['foreground_time']) + fgt -= conv.sec_to_year(f[combo].attrs["foreground_time"]) # index dict on chunk start time / coinc time self.livetimes[(get_start_dur(fi)[0], ct)] = fgt class ForegroundEvents(EventRate): - def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', - bin_lo=None, bin_hi=None): - EventRate.__init__(self, args, coinc_times, coinc_types=coinc_types, - bin_param=bin_param, bin_lo=bin_lo, bin_hi=bin_hi) + def __init__( + self, + args, + coinc_times, + coinc_types=None, + bin_param="mchirp", + bin_lo=None, + bin_hi=None, + ): + EventRate.__init__( + self, + args, + coinc_times, + coinc_types=coinc_types, + bin_param=bin_param, + bin_lo=bin_lo, + bin_hi=bin_hi, + ) self.thr = self.args.stat_threshold # set of arrays in parallel containing zerolag event properties self.starttimes = [] @@ -379,13 +412,13 @@ def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', def add_zerolag(self, full_file): start = get_start_dur(full_file)[0] self.starttimes.append(start) - with HFile(full_file, 'r') as f: + with HFile(full_file, "r") as f: # get stat values & threshold - _stats = f['foreground/stat'][:] + _stats = f["foreground/stat"][:] _keepstat = _stats > self.thr # get templates & apply filter - _tids = f['foreground/template_id'][:] + _tids = f["foreground/template_id"][:] # we need the template filter to have already been made assert self.in_bin is not None _keep = np.logical_and(_keepstat, self.in_bin[_tids]) @@ -394,39 +427,39 @@ def add_zerolag(self, full_file): # assign times and coinc types _times = {} for i in self.ifos: - _times[i] = f['foreground/' + i + '/time'][:][_keep] + _times[i] = f["foreground/" + i + "/time"][:][_keep] # if an ifo doesn't participate, time is sentinel value -1 # event time is mean of remaining positive GPS times - meantimes = np.array([coinc_meanigz(ts)[0] - for ts in zip(*_times.values())]) + meantimes = np.array([coinc_meanigz(ts)[0] for ts in zip(*_times.values())]) _ctype = self.get_ctypes(_times) if len(_ctype) == 0: if self.args.verbose: - print('No events in ' + start) + print("No events in " + start) return # filter events in_ctypes = np.array([cty in self.coinc_types for cty in _ctype]) meantimes = meantimes[in_ctypes] # get coinc time as strings # (strings may have different lengths) - _ctime = np.repeat(np.array([''], dtype=object), len(meantimes)) + _ctime = np.repeat(np.array([""], dtype=object), len(meantimes)) for ct in self.allctimestring: intime = self.in_coinc_time_excl(f, ct, meantimes) _ctime[intime == 1] = ct if self.args.verbose: - print('Got %i events in %s time' % (len(_ctime[intime == 1]), ct)) + print("Got %i events in %s time" % (len(_ctime[intime == 1]), ct)) # store self.stat = np.append(self.stat, _stats[_keep][in_ctypes]) try: # injection analyses only have 'ifar_exc', not 'ifar' - self.ifar = np.append(self.ifar, - f['foreground/ifar'][:][_keep][in_ctypes]) + self.ifar = np.append( + self.ifar, f["foreground/ifar"][:][_keep][in_ctypes] + ) except KeyError: - self.ifar = np.append(self.ifar, - f['foreground/ifar_exc'][:][_keep][in_ctypes]) + self.ifar = np.append( + self.ifar, f["foreground/ifar_exc"][:][_keep][in_ctypes] + ) self.gpst = np.append(self.gpst, meantimes) self.masspars = np.append(self.masspars, massp) - self.start = np.append(self.start, int(start) * - np.ones_like(meantimes)) + self.start = np.append(self.start, int(start) * np.ones_like(meantimes)) self.ctime = np.append(self.ctime, _ctime) self.ctype = np.append(self.ctype, _ctype[in_ctypes]) @@ -450,7 +483,7 @@ def get_bg_pdf(self, bg_rate): # store self.bg_pdf[_idx] = _pdf if self.args.verbose: - print('Found bg PDFs for ' + cty + ' coincs from ' + st) + print("Found bg PDFs for " + cty + " coincs from " + st) def get_sg_pdf(self, sg_rate): assert isinstance(sg_rate, SignalEventRate) @@ -471,15 +504,31 @@ def get_sg_pdf(self, sg_rate): # store self.sg_pdf[_idx] = _pdf if self.args.verbose: - print('Found sg PDFs for %s coincs in %s time from %s' % - (cty, ct, st)) + print( + "Found sg PDFs for %s coincs in %s time from %s" + % (cty, ct, st) + ) class BackgroundEventRate(EventRate): - def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', - bin_lo=None, bin_hi=None): - EventRate.__init__(self, args, coinc_times, coinc_types=coinc_types, - bin_param=bin_param, bin_lo=bin_lo, bin_hi=bin_hi) + def __init__( + self, + args, + coinc_times, + coinc_types=None, + bin_param="mchirp", + bin_lo=None, + bin_hi=None, + ): + EventRate.__init__( + self, + args, + coinc_times, + coinc_types=coinc_types, + bin_param=bin_param, + bin_lo=bin_lo, + bin_hi=bin_hi, + ) self.thr = self.args.stat_threshold # BG values in dict indexed on tuple (chunk start, coinc type) self.bg_vals = {} @@ -497,24 +546,24 @@ def add_background(self, full_file): start = get_start_dur(full_file)[0] self.get_livetimes(full_file) - with HFile(full_file, 'r') as ff: + with HFile(full_file, "r") as ff: # get stat values and threshold - _bgstat = ff['background_exc/stat'][:] + _bgstat = ff["background_exc/stat"][:] _keepstat = _bgstat > self.thr # get template ids and filter - _bgtid = ff['background_exc/template_id'][:] + _bgtid = ff["background_exc/template_id"][:] # need the template filter to have already been made assert self.in_bin is not None _keep = np.logical_and(_keepstat, self.in_bin[_bgtid]) _bgstat = _bgstat[_keep] - _bgdec = ff['background_exc/decimation_factor'][:][_keep] + _bgdec = ff["background_exc/decimation_factor"][:][_keep] # assign coinc types _times = {} for i in self.ifos: # NB times are time-shifted between ifos - _times[i] = ff['background_exc/' + i + '/time'][:][_keep] + _times[i] = ff["background_exc/" + i + "/time"][:][_keep] _ctype = self.get_ctypes(_times) for cty in self.coinc_types: self.bg_vals[(start, cty)] = _bgstat[_ctype == cty] @@ -522,43 +571,54 @@ def add_background(self, full_file): # get bg livetime for noise rate estimate # - convert to years self.bg_livetimes[(start, cty)] = conv.sec_to_year( - ff[cty].attrs['background_time_exc']) + ff[cty].attrs["background_time_exc"] + ) # make histogram - bins = self.make_bins(np.max(_bgstat[_ctype == cty]), 'bg') + bins = self.make_bins(np.max(_bgstat[_ctype == cty]), "bg") # hack to make larger bins for H1L1V1 - if cty == 'H1L1V1': + if cty == "H1L1V1": if self.args.verbose: - print('Halving bg bins for triple bg hist') + print("Halving bg bins for triple bg hist") bins = bins[::2].copy() # take every 2nd bin edge - self.bg_hist[(start, cty)] = \ - np.histogram(_bgstat[_ctype == cty], - weights=_bgdec[_ctype == cty], bins=bins) + self.bg_hist[(start, cty)] = np.histogram( + _bgstat[_ctype == cty], weights=_bgdec[_ctype == cty], bins=bins + ) # get expected number of bg events for this chunk and coinc type - self.exp_bg[(start, cty)] = _bgdec[_ctype == cty].sum() * \ - self.incl_livetimes[(start, cty)] / \ - self.bg_livetimes[(start, cty)] + self.exp_bg[(start, cty)] = ( + _bgdec[_ctype == cty].sum() + * self.incl_livetimes[(start, cty)] + / self.bg_livetimes[(start, cty)] + ) def plot_bg(self): from matplotlib import pyplot as plt + for chunk_type, hist in self.bg_hist.items(): - print('Plotting', chunk_type, 'background PDF ...') + print("Plotting", chunk_type, "background PDF ...") xplot = np.linspace(self.thr, self.args.plot_max_stat, 500) heights, bins = hist[0], hist[1] logpdf, _ = log_rho_bg(xplot, heights, bins) plt.plot(xplot, np.exp(logpdf)) # plot error bars at bin centres lpdf, fracerr = log_rho_bg(0.5 * (bins[:-1] + bins[1:]), heights, bins) - plt.errorbar(0.5 * (bins[:-1] + bins[1:]), np.exp(lpdf), - yerr=np.exp(lpdf) * fracerr, fmt='none') + plt.errorbar( + 0.5 * (bins[:-1] + bins[1:]), + np.exp(lpdf), + yerr=np.exp(lpdf) * fracerr, + fmt="none", + ) plt.semilogy() plt.grid(True) plt.xlim(xmax=self.args.plot_max_stat + 0.5) plt.ylim(ymin=0.7 * np.exp(logpdf.min())) - plt.xlabel('Ranking statistic') - plt.ylabel('Background PDF') - plt.savefig(self.args.plot_dir + '%s-bg_pdf-%s' % - (chunk_type[1], chunk_type[0]) + '.png') + plt.xlabel("Ranking statistic") + plt.ylabel("Background PDF") + plt.savefig( + self.args.plot_dir + + "%s-bg_pdf-%s" % (chunk_type[1], chunk_type[0]) + + ".png" + ) plt.close() def get_norms(self): @@ -573,18 +633,31 @@ def eval_pdf(self, chunk, ctime, ctype, statvals): # fraction of expected noise events in given chunk & coinc type frac_chunk_type = self.exp_bg[chunk_type] / self.norm # fraction of inj in specified chunk, coinc type *and* time - frac_in_time = self.livetimes[(chunk, ctime)] /\ - self.incl_livetimes[chunk_type] + frac_in_time = self.livetimes[(chunk, ctime)] / self.incl_livetimes[chunk_type] # unpack heights / bins from bg hist object local_pdfs, _ = log_rho_bg(statvals, *self.bg_hist[chunk_type]) return local_pdfs + np.log(frac_chunk_type * frac_in_time) class SignalEventRate(EventRate): - def __init__(self, args, coinc_times, coinc_types=None, bin_param='mchirp', - bin_lo=None, bin_hi=None): - EventRate.__init__(self, args, coinc_times, coinc_types=coinc_types, - bin_param=bin_param, bin_lo=bin_lo, bin_hi=bin_hi) + def __init__( + self, + args, + coinc_times, + coinc_types=None, + bin_param="mchirp", + bin_lo=None, + bin_hi=None, + ): + EventRate.__init__( + self, + args, + coinc_times, + coinc_types=coinc_types, + bin_param=bin_param, + bin_lo=bin_lo, + bin_hi=bin_hi, + ) self.thr = self.args.stat_threshold self.starts = [] # bookkeeping # for the moment roll all inj chunks together @@ -598,13 +671,13 @@ def add_injections(self, inj_file, fg_file): self.starts.append(get_start_dur(inj_file)[0]) self.get_livetimes(inj_file) - with HFile(inj_file, 'r') as jf: + with HFile(inj_file, "r") as jf: # get stat values and threshold - _injstat = jf['found_after_vetoes/stat'][:] + _injstat = jf["found_after_vetoes/stat"][:] _keepstat = _injstat > self.thr # get template ids and filter - _injtid = jf['found_after_vetoes/template_id'][:] + _injtid = jf["found_after_vetoes/template_id"][:] assert self.in_bin is not None _keep = np.logical_and(_keepstat, self.in_bin[_injtid]) _injstat = _injstat[_keep] @@ -612,32 +685,31 @@ def add_injections(self, inj_file, fg_file): # assign coinc types _times = {} for i in self.ifos: - _times[i] = jf['found_after_vetoes/' + i + '/time'][:][_keep] - meantimes = np.array([coinc_meanigz(ts)[0] - for ts in zip(*_times.values())]) + _times[i] = jf["found_after_vetoes/" + i + "/time"][:][_keep] + meantimes = np.array([coinc_meanigz(ts)[0] for ts in zip(*_times.values())]) _ctype = self.get_ctypes(_times) # get coinc time as strings # (strings may have different lengths) - _ctime = np.repeat(np.array([''], dtype=object), len(meantimes)) + _ctime = np.repeat(np.array([""], dtype=object), len(meantimes)) for ct in self.allctimestring: # get coinc time info from segments in fg file - intime = self.in_coinc_time_excl( - HFile(fg_file, 'r'), ct, meantimes) + intime = self.in_coinc_time_excl(HFile(fg_file, "r"), ct, meantimes) _ctime[intime == 1] = ct # do we need this? if self.args.verbose: - print('Got %i ' % (intime == 1).sum() + 'inj in %s time' % ct) + print("Got %i " % (intime == 1).sum() + "inj in %s time" % ct) # filter by coinc type and add to array for cty in self.coinc_types: if not type_in_time(ct, cty): continue my_vals = _injstat[np.logical_and(_ctype == cty, intime == 1)] if self.args.verbose: - print('%d ' % len(my_vals) + 'are %s coincs' % cty) + print("%d " % len(my_vals) + "are %s coincs" % cty) if (ct, cty) not in self.inj_vals: # initialize self.inj_vals[(ct, cty)] = np.array([]) if len(my_vals) > 0: - self.inj_vals[(ct, cty)] = \ - np.append(self.inj_vals[(ct, cty)], my_vals) + self.inj_vals[(ct, cty)] = np.append( + self.inj_vals[(ct, cty)], my_vals + ) del intime, my_vals def make_all_bins(self): @@ -648,37 +720,40 @@ def make_all_bins(self): vals = self.inj_vals[(ct, cty)] # get norm of fg histogram by taking bins out to max injection stat binmax = vals.max() * 1.01 - self.fg_bins[(ct, cty)] = self.make_bins(binmax, 'inj') + self.fg_bins[(ct, cty)] = self.make_bins(binmax, "inj") def plot_inj(self): from matplotlib import pyplot as plt + for ct in self.allctimestring: for cty in self.coinc_types: if not type_in_time(ct, cty): continue - print('Plotting ' + cty + ' signal PDF in ' + ct + ' time ...') + print("Plotting " + cty + " signal PDF in " + ct + " time ...") samples = self.inj_vals[(ct, cty)] bins = self.fg_bins[(ct, cty)] - xplot = np.logspace(np.log10(self.thr), - np.log10(samples.max()), 500) + xplot = np.logspace(np.log10(self.thr), np.log10(samples.max()), 500) logpdf, _ = log_rho_fg(xplot, samples, bins) plt.plot(xplot, np.exp(logpdf)) # plot error bars at bin centres - lpdf, fracerr = log_rho_fg(0.5 * (bins[:-1] + bins[1:]), - samples, bins) - plt.errorbar(0.5 * (bins[:-1] + bins[1:]), np.exp(lpdf), - yerr=np.exp(lpdf) * fracerr, fmt='none') + lpdf, fracerr = log_rho_fg(0.5 * (bins[:-1] + bins[1:]), samples, bins) + plt.errorbar( + 0.5 * (bins[:-1] + bins[1:]), + np.exp(lpdf), + yerr=np.exp(lpdf) * fracerr, + fmt="none", + ) plt.semilogy() plt.grid(True) # zoom in on the 'interesting' range - plt.xlim(xmin=self.thr, xmax=2. * self.args.plot_max_stat) + plt.xlim(xmin=self.thr, xmax=2.0 * self.args.plot_max_stat) plt.ylim(ymin=0.7 * np.exp(logpdf.min())) - plt.title(r'%i injs plotted, \# of bins %i' % - (len(samples), len(bins) - 1)) - plt.xlabel('Ranking statistic') - plt.ylabel('Signal PDF') - plt.savefig(self.args.plot_dir + '%s-fg_pdf-%s' % (ct, cty) - + '.png') + plt.title( + r"%i injs plotted, \# of bins %i" % (len(samples), len(bins) - 1) + ) + plt.xlabel("Ranking statistic") + plt.ylabel("Signal PDF") + plt.savefig(self.args.plot_dir + "%s-fg_pdf-%s" % (ct, cty) + ".png") plt.close() def get_norms(self): @@ -696,14 +771,28 @@ def eval_pdf(self, chunk, ctime, ctype, statvals): # total livetime for specified coinc time total_coinc_time = sum([self.livetimes[(ch, ctime)] for ch in self.starts]) # fraction of inj in specified chunk *and* coinc time/type - this_norm = frac_time_type * self.livetimes[(chunk, ctime)] / \ - total_coinc_time - local_pdfs, _ = log_rho_fg(statvals, self.inj_vals[time_type], - self.fg_bins[time_type]) + this_norm = frac_time_type * self.livetimes[(chunk, ctime)] / total_coinc_time + local_pdfs, _ = log_rho_fg( + statvals, self.inj_vals[time_type], self.fg_bins[time_type] + ) return local_pdfs + np.log(this_norm) -__all__ = ['filter_bin_lo_hi', 'filter_tmplt_mchirp', 'read_full_data', - 'read_full_data_mchirp', 'log_rho_bg', 'log_rho_fg_analytic', - 'log_rho_fg', 'get_start_dur', 'in_coinc_time_incl', 'alltimes', - 'ifos_from_combo', 'type_in_time', 'EventRate', 'ForegroundEvents', - 'BackgroundEventRate', 'SignalEventRate'] + +__all__ = [ + "BackgroundEventRate", + "EventRate", + "ForegroundEvents", + "SignalEventRate", + "alltimes", + "filter_bin_lo_hi", + "filter_tmplt_mchirp", + "get_start_dur", + "ifos_from_combo", + "in_coinc_time_incl", + "log_rho_bg", + "log_rho_fg", + "log_rho_fg_analytic", + "read_full_data", + "read_full_data_mchirp", + "type_in_time", +] diff --git a/pycbc/population/fgmc_laguerre.py b/pycbc/population/fgmc_laguerre.py index eeb339efe62..2584a915332 100644 --- a/pycbc/population/fgmc_laguerre.py +++ b/pycbc/population/fgmc_laguerre.py @@ -13,17 +13,21 @@ """ import numpy -import scipy.stats as sst -import scipy.special as ssp import scipy.integrate as sig import scipy.optimize as sop +import scipy.special as ssp +import scipy.stats as sst class augmented_rv_continuous(sst.rv_continuous): - - def __init__(self, unit='dimensionless', texunit=r'\mbox{dimensionless}', - texsymb=r'x', **kwargs): - ''' + def __init__( + self, + unit="dimensionless", + texunit=r"\mbox{dimensionless}", + texsymb=r"x", + **kwargs, + ): + """ Parameters ---------- unit : string, optional @@ -32,9 +36,9 @@ def __init__(self, unit='dimensionless', texunit=r'\mbox{dimensionless}', units of independent variable, in tex format texsymb : string, optional symbol of independent variable, in tex format - ''' - super(augmented_rv_continuous, self).__init__(**kwargs) + """ + super().__init__(**kwargs) self._hpd_interval_vec = numpy.vectorize(self._hpd_interval_scalar) self.unit = unit self.texunit = texunit @@ -54,7 +58,7 @@ def width(a): return a, b def hpd_interval(self, alpha): - ''' + """ Confidence interval of highest probability density. Parameters @@ -68,7 +72,8 @@ def hpd_interval(self, alpha): a, b : ndarray of float end-points of range that contain ``100 * alpha %`` of the rv's possible values. - ''' + + """ if isinstance(alpha, (float, numpy.number)): a, b = self._hpd_interval_scalar(alpha) else: @@ -77,15 +82,22 @@ def hpd_interval(self, alpha): class count_posterior(augmented_rv_continuous): - ''' + """ Count posterior distribution. - ''' - - def __init__(self, logbf, laguerre_n, Lambda0, prior=-0.5, - name='count posterior', unit='signals/experiment', - texunit=r'\mathrm{signals}/\mathrm{experiment}', - texsymb=r'\Lambda_1'): - ''' + """ + + def __init__( + self, + logbf, + laguerre_n, + Lambda0, + prior=-0.5, + name="count posterior", + unit="signals/experiment", + texunit=r"\mathrm{signals}/\mathrm{experiment}", + texsymb=r"\Lambda_1", + ): + """ Parameters ---------- logbf : array_like @@ -98,10 +110,11 @@ def __init__(self, logbf, laguerre_n, Lambda0, prior=-0.5, prior distribution power law of improper prior if float or count posterior distribution if count_posterior (default=-0.5: Jeffreys prior) - ''' - super(count_posterior, self).__init__(a=0.0, b=numpy.inf, name=name, - unit=unit, texunit=texunit, - texsymb=texsymb) + + """ + super().__init__( + a=0.0, b=numpy.inf, name=name, unit=unit, texunit=texunit, texsymb=texsymb + ) self.Lambda0 = Lambda0 # weighted Bayes factor self.k = numpy.exp(numpy.array(logbf)) / self.Lambda0 @@ -111,7 +124,7 @@ def __init__(self, logbf, laguerre_n, Lambda0, prior=-0.5, if prior == 0: self.prior = lambda x: 1.0 elif prior > 0: - self.prior = lambda x: x ** prior + self.prior = lambda x: x**prior else: # regularize at x = 0 self.prior = lambda x: (x + self.xtol) ** prior @@ -119,8 +132,9 @@ def __init__(self, logbf, laguerre_n, Lambda0, prior=-0.5, # pre-compute Gaussian-Generalized-Laguerre quadrature # abscissas and weights, along with pdf at these abscissas self.x, w = ssp.la_roots(laguerre_n, self.alpha) - self.p = numpy.array([ww * numpy.prod(1.0 + self.k * xx) - for xx, ww in zip(self.x, w)]) + self.p = numpy.array( + [ww * numpy.prod(1.0 + self.k * xx) for xx, ww in zip(self.x, w)] + ) self.norm = 1.0 / sum(self.p) self.p *= self.norm @@ -134,7 +148,7 @@ def _cdf(self, x): return sig.quad(self._pdf, 0.0, x) def expect(self, func): - ''' + """ Calculate expected value of a function with respect to the distribution. @@ -152,7 +166,8 @@ def expect(self, func): ------- expect : float The calculated expected value. - ''' + + """ # FIXME: not as feature rich as the expect method this overrides return sum(pp * func(xx) for xx, pp in zip(self.x, self.p)) @@ -160,18 +175,19 @@ def _munp(self, n): return self.expect(lambda x: x**n) def p_bg(self, logbf): - ''' + """ Calculate the false alarm probabilities of the events. Parameters ---------- logbf : array_like Logs of foreground over background probability ratios of events. - ''' + + """ # get weighted bayes factor k = numpy.exp(numpy.asarray(logbf)) / self.Lambda0 - P0 = numpy.dot(1./(1. + numpy.outer(k, self.x)), self.p) + P0 = numpy.dot(1.0 / (1.0 + numpy.outer(k, self.x)), self.p) if isinstance(k, (float, int, numpy.number)): return P0.item() if isinstance(k, numpy.ndarray) and k.ndim == 0: @@ -179,4 +195,5 @@ def p_bg(self, logbf): # except in special cases above, return array of values return P0 -__all__ = ['augmented_rv_continuous', 'count_posterior'] + +__all__ = ["augmented_rv_continuous", "count_posterior"] diff --git a/pycbc/population/fgmc_plots.py b/pycbc/population/fgmc_plots.py index 529ab59d487..a6214914199 100644 --- a/pycbc/population/fgmc_plots.py +++ b/pycbc/population/fgmc_plots.py @@ -6,8 +6,8 @@ # option) any later version. import json -import numpy +import numpy from matplotlib import figure from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas @@ -25,11 +25,11 @@ def plotodds(rankstats, p_b): # odds vs ranking stat fig, ax = plot_setup() ax.loglog() - ax.plot(rankstats, (1.0 - p_b) / p_b, 'k.') - ax.plot([rankstats.min(), rankstats.max()], [1.0, 1.0], 'c--') - ax.set_title(r'Foreground/Background Odds') - ax.set_xlabel(r'ranking statistic') - ax.set_ylabel(r'$P_1/P_0$') + ax.plot(rankstats, (1.0 - p_b) / p_b, "k.") + ax.plot([rankstats.min(), rankstats.max()], [1.0, 1.0], "c--") + ax.set_title(r"Foreground/Background Odds") + ax.set_xlabel(r"ranking statistic") + ax.set_ylabel(r"$P_1/P_0$") ax.set_xlim(0.99 * rankstats.min(), 1.2 * rankstats.max()) return fig @@ -38,10 +38,10 @@ def plotpbg(rankstats, p_b): # p_terr vs ranking stat fig, ax = plot_setup() ax.loglog() - ax.plot(rankstats, p_b, 'k.') - ax.set_title(r'Probability of background origin') - ax.set_xlabel(r'ranking statistic') - ax.set_ylabel(r'$P_0$') + ax.plot(rankstats, p_b, "k.") + ax.set_title(r"Probability of background origin") + ax.set_xlabel(r"ranking statistic") + ax.set_ylabel(r"$P_0$") ax.set_xlim(0.99 * rankstats.min(), 1.2 * rankstats.max()) return fig @@ -50,11 +50,11 @@ def plotoddsifar(ifar, p_b): # odds vs IFAR fig, ax = plot_setup() ax.loglog() - ax.plot(ifar, (1.0 - p_b) / p_b, 'k.') - ax.plot([ifar.min(), ifar.max()], [1.0, 1.0], 'c--') - ax.set_title(r'Foreground/Background Odds') - ax.set_xlabel(r'IFAR') - ax.set_ylabel(r'$P_1/P_0$') + ax.plot(ifar, (1.0 - p_b) / p_b, "k.") + ax.plot([ifar.min(), ifar.max()], [1.0, 1.0], "c--") + ax.set_title(r"Foreground/Background Odds") + ax.set_xlabel(r"IFAR") + ax.set_ylabel(r"$P_1/P_0$") ax.set_xlim(0.9 * ifar.min(), 1.1 * ifar.max()) return fig @@ -66,116 +66,130 @@ def plotfdr(p_b, ntop): p_b = numpy.sort(p_b)[:ntop] # cumulative probable noise/signal counts cum_false = p_b.cumsum() - cum_true = (1. - p_b).cumsum() + cum_true = (1.0 - p_b).cumsum() ax.semilogy() - ax.plot(p_b, cum_false / cum_true, 'b+') - ax.plot(p_b, 1. / (numpy.arange(len(p_b)) + 1), 'c--', label=r'1 noise event') + ax.plot(p_b, cum_false / cum_true, "b+") + ax.plot(p_b, 1.0 / (numpy.arange(len(p_b)) + 1), "c--", label=r"1 noise event") ax.legend() - ax.set_xlabel(r'$p_{\rm terr}$') - ax.set_ylabel(r'Cumulative $p_{\rm terr}$ / Cumulative $p_{\rm astro}$') - ax.set_xlim(0., 1.05 * p_b.max()) + ax.set_xlabel(r"$p_{\rm terr}$") + ax.set_ylabel(r"Cumulative $p_{\rm terr}$ / Cumulative $p_{\rm astro}$") + ax.set_xlim(0.0, 1.05 * p_b.max()) return fig def finalize_plot(fig, args, extensions, name, pltype, tag): # Helper function for extn in extensions: - filename = args.pldir + '_'.join(name.split()) + '_' + pltype + tag + extn + filename = args.pldir + "_".join(name.split()) + "_" + pltype + tag + extn if args.verbose: - print('writing %s ...' % filename) + print("writing %s ..." % filename) fig.savefig(filename) -def odds_summary(args, rankstats, ifars, p_b, ntop, times=None, mchirps=None, - name='events', plot_extensions=None): +def odds_summary( + args, + rankstats, + ifars, + p_b, + ntop, + times=None, + mchirps=None, + name="events", + plot_extensions=None, +): - print('\nSummary of Top %i %s' % (ntop, name.title())) + print("\nSummary of Top %i %s" % (ntop, name.title())) # do sort in reverse order - statsort = numpy.argsort(1. / numpy.array(rankstats)) + statsort = numpy.argsort(1.0 / numpy.array(rankstats)) topn = statsort[:ntop] # indices giving top n topgps = [] topstat = [] topifar = [] toppastro = [] for n, i in enumerate(topn): - gps = times[i] if times is not None else '' + gps = times[i] if times is not None else "" ifar = ifars[i] stat = rankstats[i] - mchirpstring = 'mchirp %.3F' % mchirps[i] if mchirps is not None else '' + mchirpstring = "mchirp %.3F" % mchirps[i] if mchirps is not None else "" topgps.append(gps) topstat.append(stat) topifar.append(ifar) - print('#%d event:' % (n + 1), str(gps), mchirpstring) - print(' rankstat = %-8.3f' % stat) - print(' IFAR = %.2f' % ifar) - print(' odds = %g' % ((1. - p_b[i]) / p_b[i])) - toppastro.append(1. - p_b[i]) + print("#%d event:" % (n + 1), str(gps), mchirpstring) + print(" rankstat = %-8.3f" % stat) + print(" IFAR = %.2f" % ifar) + print(" odds = %g" % ((1.0 - p_b[i]) / p_b[i])) + toppastro.append(1.0 - p_b[i]) if args.p_astro_txt is not None: - numpy.savetxt(args.p_astro_txt, - numpy.column_stack((topgps, topstat, topifar, toppastro)), - fmt=['%.3F', '%.2F', '%.2F', '%.5F'], - delimiter=',', - header='GPS seconds, stat, IFAR/yr, p_astro') - - if hasattr(args, 'json_tag') and args.json_tag is not None: + numpy.savetxt( + args.p_astro_txt, + numpy.column_stack((topgps, topstat, topifar, toppastro)), + fmt=["%.3F", "%.2F", "%.2F", "%.5F"], + delimiter=",", + header="GPS seconds, stat, IFAR/yr, p_astro", + ) + + if hasattr(args, "json_tag") and args.json_tag is not None: # save to catalog-style files def dump_json(gps, p_a, p_b): - jfile = args.plot_dir + 'H1L1V1-PYCBC_%s-%s-1.json' % \ - (args.json_tag, str(int(gps))) # truncate to integer GPS - with open(jfile, 'w') as jf: - json.dump({'Astro': p_a, 'Terrestrial': p_b}, jf) - if hasattr(args, 'json_min_ifar') and args.json_min_ifar is not None: + jfile = args.plot_dir + "H1L1V1-PYCBC_%s-%s-1.json" % ( + args.json_tag, + str(int(gps)), + ) # truncate to integer GPS + with open(jfile, "w") as jf: + json.dump({"Astro": p_a, "Terrestrial": p_b}, jf) + + if hasattr(args, "json_min_ifar") and args.json_min_ifar is not None: for g, ifar, pt in zip(times, ifars, p_b): if ifar < args.json_min_ifar: continue - dump_json(g, 1. - pt, pt) + dump_json(g, 1.0 - pt, pt) else: for g, pa in zip(topgps, toppastro): - dump_json(g, pa, 1. - pa) + dump_json(g, pa, 1.0 - pa) if plot_extensions is not None: - plottag = args.plot_tag or '' - if plottag != '': - plottag = '_' + plottag + plottag = args.plot_tag or "" + if plottag != "": + plottag = "_" + plottag fig = plotodds(rankstats, p_b) - finalize_plot(fig, args, plot_extensions, name, 'odds', plottag) + finalize_plot(fig, args, plot_extensions, name, "odds", plottag) fig = plotpbg(rankstats, p_b) - finalize_plot(fig, args, plot_extensions, name, 'pbg', plottag) + finalize_plot(fig, args, plot_extensions, name, "pbg", plottag) fig = plotoddsifar(ifars, p_b) - finalize_plot(fig, args, plot_extensions, name, 'ifarodds', plottag) + finalize_plot(fig, args, plot_extensions, name, "ifarodds", plottag) fig = plotfdr(p_b, ntop) - finalize_plot(fig, args, plot_extensions, name, 'fdr', plottag) + finalize_plot(fig, args, plot_extensions, name, "fdr", plottag) -def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style='linear'): +def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style="linear"): fig = figure.Figure() FigureCanvas(fig) ax = fig.gca() - name = rv.name if hasattr(rv, 'name') else None - symb = rv.texsymb if hasattr(rv, 'texsymb') else r'x' - unit = rv.texunit if hasattr(rv, 'texunit') else None + name = rv.name if hasattr(rv, "name") else None + symb = rv.texsymb if hasattr(rv, "texsymb") else r"x" + unit = rv.texunit if hasattr(rv, "texunit") else None - xlabel = r'$' + symb + '$' + xlabel = r"$" + symb + "$" if unit is not None: - xlabel += r' ($' + unit + r'$)' + xlabel += r" ($" + unit + r"$)" a, b = rv.interval(0.9999) - if style == 'loglog': + if style == "loglog": ax.loglog() - ylabel = r'$p(' + symb + r')$' + ylabel = r"$p(" + symb + r")$" space = lambda a, b: numpy.logspace(numpy.log10(a), numpy.log10(b), 100) func = numpy.vectorize(rv.pdf) xmin = a ymin = rv.pdf(b) - elif style == 'semilogx': + elif style == "semilogx": ax.semilogx() - ylabel = r'$' + symb + r'\,p(' + symb + r')$' + ylabel = r"$" + symb + r"\,p(" + symb + r")$" space = lambda a, b: numpy.logspace(numpy.log10(a), numpy.log10(b), 100) func = numpy.vectorize(lambda x: x * rv.pdf(x)) xmin = a @@ -183,7 +197,7 @@ def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style='lin else: # linear ax.yaxis.set_ticklabels([]) - ylabel = r'$p(' + symb + r')$' + ylabel = r"$p(" + symb + r")$" space = lambda a, b: numpy.linspace(a, b, 100) func = numpy.vectorize(rv.pdf) xmin = 0.0 @@ -192,7 +206,7 @@ def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style='lin x = space(a, b) y = func(x) - ax.plot(x, y, color='k', linestyle='-') + ax.plot(x, y, color="k", linestyle="-") if plot_lim is not None: xmin, xmax = plot_lim ax.set_xlim(xmin=xmin, xmax=xmax) @@ -201,7 +215,7 @@ def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style='lin ax.set_ylim(ymin=ymin) if y.max() < 2 and y.max() > 1: - ax.set_ylim(ymax=2.) + ax.set_ylim(ymax=2.0) if name is not None: ax.set_title(name.title()) @@ -209,7 +223,7 @@ def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style='lin ax.set_ylabel(ylabel) if middle is not None: - ax.plot([middle, middle], [ymin, func(middle)], 'k--') + ax.plot([middle, middle], [ymin, func(middle)], "k--") if credible_intervals is not None: # alpha : density of fill shading @@ -219,43 +233,57 @@ def plotdist(rv, plot_lim=None, middle=None, credible_intervals=None, style='lin hi = min(b, hi) x = space(lo, hi) y = func(x) - ax.fill_between(x, y, ymin, color='k', alpha=alpha) + ax.fill_between(x, y, ymin, color="k", alpha=alpha) return fig -def dist_summary(args, rv, plot_styles=('linear', 'loglog', 'semilogx'), - plot_extensions=None, middle=None, credible_intervals=None): +def dist_summary( + args, + rv, + plot_styles=("linear", "loglog", "semilogx"), + plot_extensions=None, + middle=None, + credible_intervals=None, +): - name = rv.name if hasattr(rv, 'name') else 'posterior' - unit = rv.unit if hasattr(rv, 'unit') else '' + name = rv.name if hasattr(rv, "name") else "posterior" + unit = rv.unit if hasattr(rv, "unit") else "" median = rv.median() - mode = rv.mode() if hasattr(rv, 'mode') else None + mode = rv.mode() if hasattr(rv, "mode") else None - print('Summary of ' + name.title()) - print('mean =', rv.mean(), unit) - print('median =', median, unit) + print("Summary of " + name.title()) + print("mean =", rv.mean(), unit) + print("median =", median, unit) if mode is not None: - print('mode =', mode, unit) - print('stddev =', rv.std(), unit) + print("mode =", mode, unit) + print("stddev =", rv.std(), unit) if credible_intervals is not None and len(credible_intervals) > 0: - print('equal-tailed credible intervals:') + print("equal-tailed credible intervals:") equal_tailed_credible_intervals = {} for cred in credible_intervals: lo, hi = rv.interval(cred) equal_tailed_credible_intervals[cred] = (lo, hi) - print('%g%%' % (cred * 100), 'credible interval =', '[%g, %g]' % - (lo, hi), unit) - - if hasattr(rv, 'hpd_interval'): - print('highest probability density credible intervals:') + print( + "%g%%" % (cred * 100), + "credible interval =", + "[%g, %g]" % (lo, hi), + unit, + ) + + if hasattr(rv, "hpd_interval"): + print("highest probability density credible intervals:") hpd_credible_intervals = {} for cred in credible_intervals: hpdlo, hpdhi = rv.hpd_interval(cred) hpd_credible_intervals[cred] = (hpdlo, hpdhi) - print('%g%%' % (cred * 100), 'credible interval =', '[%g, %g]' % - (hpdlo, hpdhi), unit) + print( + "%g%%" % (cred * 100), + "credible interval =", + "[%g, %g]" % (hpdlo, hpdhi), + unit, + ) else: hpd_credible_intervals = None @@ -263,7 +291,7 @@ def dist_summary(args, rv, plot_styles=('linear', 'loglog', 'semilogx'), credible_intervals = None intervals = None - if middle == 'mode' and mode is not None: + if middle == "mode" and mode is not None: middle = mode if credible_intervals is not None: # use hpd intervals with mode @@ -276,12 +304,17 @@ def dist_summary(args, rv, plot_styles=('linear', 'loglog', 'semilogx'), # plot distributions if plot_extensions is not None: - plottag = args.plot_tag or '' - if plottag != '': - plottag = '_' + plottag + plottag = args.plot_tag or "" + if plottag != "": + plottag = "_" + plottag for style in plot_styles: - fig = plotdist(rv, plot_lim=args.plot_limits, middle=middle, - credible_intervals=intervals, style=style) + fig = plotdist( + rv, + plot_lim=args.plot_limits, + middle=middle, + credible_intervals=intervals, + style=style, + ) finalize_plot(fig, args, plot_extensions, name, style, plottag) if credible_intervals is not None and len(credible_intervals) == 1: @@ -289,4 +322,5 @@ def dist_summary(args, rv, plot_styles=('linear', 'loglog', 'semilogx'), # keep codeclimate happy with explicit return statement return None -__all__ = ['plotdist', 'odds_summary', 'dist_summary'] + +__all__ = ["dist_summary", "odds_summary", "plotdist"] diff --git a/pycbc/population/live_pastro.py b/pycbc/population/live_pastro.py index c933cd4e32c..d9341425c70 100644 --- a/pycbc/population/live_pastro.py +++ b/pycbc/population/live_pastro.py @@ -1,15 +1,17 @@ import logging + import h5py import numpy -from pycbc.tmpltbank import bank_conversions as bankconv -from pycbc.events import triggers from pycbc import conversions as conv +from pycbc.events import triggers +from pycbc.tmpltbank import bank_conversions as bankconv + from . import fgmc_functions as fgmcfun -_s_per_yr = 1. / conv.sec_to_year(1.) +_s_per_yr = 1.0 / conv.sec_to_year(1.0) -logger = logging.getLogger('pycbc.population.live_pastro') +logger = logging.getLogger("pycbc.population.live_pastro") def check_template_param_bin_data(spec_json): @@ -22,16 +24,16 @@ def check_template_param_bin_data(spec_json): Returns ------- spec_json: dictionary + """ # Check the necessary data are present - assert 'param' in spec_json - assert 'bin_edges' in spec_json # should be a list of floats - assert 'sig_per_yr_binned' in spec_json # signal rate per bin (per year) + assert "param" in spec_json + assert "bin_edges" in spec_json # should be a list of floats + assert "sig_per_yr_binned" in spec_json # signal rate per bin (per year) # Do the lengths of bin arrays match? - assert len(spec_json['bin_edges']) == \ - len(spec_json['sig_per_yr_binned']) + 1 - assert 'ref_bns_horizon' in spec_json # float - assert 'netsnr_thresh' in spec_json # float + assert len(spec_json["bin_edges"]) == len(spec_json["sig_per_yr_binned"]) + 1 + assert "ref_bns_horizon" in spec_json # float + assert "netsnr_thresh" in spec_json # float return spec_json @@ -46,13 +48,14 @@ def check_template_param_bin_farlim_data(spec_json): Returns ------- spec_json: dictionary + """ # Standard template param bin checks check_template_param_bin_data(spec_json) # In addition, need limiting FAR and SNR values - assert 'limit_far' in spec_json - assert 'limit_snr' in spec_json + assert "limit_far" in spec_json + assert "limit_snr" in spec_json return spec_json @@ -70,16 +73,17 @@ def read_template_bank_param(spec_d, bankf): ------- bank_data: dictionary Template counts binned over specified param + """ - with h5py.File(bankf, 'r') as bank: + with h5py.File(bankf, "r") as bank: # All the templates - tids = numpy.arange(len(bank['mass1'])) + tids = numpy.arange(len(bank["mass1"])) # Get param vals - logger.info('Getting %s values from bank', spec_d['param']) - parvals = bankconv.get_bank_property(spec_d['param'], bank, tids) - counts, edges = numpy.histogram(parvals, bins=spec_d['bin_edges']) - bank_data = {'bin_edges': edges, 'tcounts': counts, 'num_t': counts.sum()} - logger.info('Binned template counts: %s', counts) + logger.info("Getting %s values from bank", spec_d["param"]) + parvals = bankconv.get_bank_property(spec_d["param"], bank, tids) + counts, edges = numpy.histogram(parvals, bins=spec_d["bin_edges"]) + bank_data = {"bin_edges": edges, "tcounts": counts, "num_t": counts.sum()} + logger.info("Binned template counts: %s", counts) return bank_data @@ -111,13 +115,14 @@ def trials_type(ntriggered, nactive): if ntriggered == 2 and nactive == 3: return 6 # All valid inputs are exhausted, throw an error - raise ValueError(f"I don't know what to do with {ntriggered} triggered and" - f" {nactive} active ifos!") + raise ValueError( + f"I don't know what to do with {ntriggered} triggered and" + f" {nactive} active ifos!" + ) def signal_pdf_from_snr(netsnr, thresh): - """ FGMC approximate signal distribution ~ SNR ** -4 - """ + """FGMC approximate signal distribution ~ SNR ** -4""" return numpy.exp(fgmcfun.log_rho_fg_analytic(netsnr, thresh)) @@ -127,9 +132,9 @@ def signal_rate_rescale(horizons, ref_dhor): to account for network sensitivity variation relative to a reference state """ # Combine sensitivities over ifos in a way analogous to network SNR - net_horizon = sum(hor ** 2. for hor in horizons.values()) ** 0.5 + net_horizon = sum(hor**2.0 for hor in horizons.values()) ** 0.5 # signal rate is proportional to horizon distance cubed - return net_horizon ** 3. / ref_dhor ** 3. + return net_horizon**3.0 / ref_dhor**3.0 def signal_rate_trig_type(horizons, sens_ifos, trig_ifos): @@ -140,17 +145,20 @@ def signal_rate_trig_type(horizons, sens_ifos, trig_ifos): # Single-ifo time if len(sens_ifos) == 1: assert len(trig_ifos) == 1 - return 1. + return 1.0 # Single trigger in multi-ifo time if len(trig_ifos) == 1: # Sensitive volume scales with horizon^3 # Suppress horizon by sqrt(2) wrt coincs - return (horizons[trig_ifos[0]] / 2**0.5) ** 3. /\ - sum([horizons[i] ** 3. for i in sens_ifos]) + return (horizons[trig_ifos[0]] / 2**0.5) ** 3.0 / sum( + [horizons[i] ** 3.0 for i in sens_ifos] + ) # Double coinc : volume determined by less sensitive ifo # Compare to 2nd most sensitive ifo over the observing network - return sorted([horizons[i] for i in trig_ifos])[0] ** 3. /\ - sorted([horizons[i] for i in sens_ifos])[-2] ** 3. + return ( + sorted([horizons[i] for i in trig_ifos])[0] ** 3.0 + / sorted([horizons[i] for i in sens_ifos])[-2] ** 3.0 + ) def template_param_bin_pa(padata, trdata, horizons): @@ -167,40 +175,38 @@ def template_param_bin_pa(padata, trdata, horizons): Returns ------- p_astro, p_terr: tuple of floats + """ - massspin = (trdata['mass1'], trdata['mass2'], - trdata['spin1z'], trdata['spin2z']) - trig_param = triggers.get_param(padata.spec['param'], None, *massspin) + massspin = (trdata["mass1"], trdata["mass2"], trdata["spin1z"], trdata["spin2z"]) + trig_param = triggers.get_param(padata.spec["param"], None, *massspin) # NB digitize gives '1' for first bin, '2' for second etc. - bind = numpy.digitize(trig_param, padata.bank['bin_edges']) - 1 - logger.debug('Trigger %s is in bin %i', padata.spec['param'], bind) + bind = numpy.digitize(trig_param, padata.bank["bin_edges"]) - 1 + logger.debug("Trigger %s is in bin %i", padata.spec["param"], bind) # Get noise rate density - if 'bg_fac' not in padata.spec: - expfac = 6. + if "bg_fac" not in padata.spec: + expfac = 6.0 else: - expfac = padata.spec['bg_fac'] + expfac = padata.spec["bg_fac"] # FAR is in Hz, therefore convert to rate per year (per SNR) - dnoise = noise_density_from_far(trdata['far'], expfac) * _s_per_yr - logger.debug('FAR %.3g, noise density per yr per SNR %.3g', - trdata['far'], dnoise) + dnoise = noise_density_from_far(trdata["far"], expfac) * _s_per_yr + logger.debug("FAR %.3g, noise density per yr per SNR %.3g", trdata["far"], dnoise) # Scale by fraction of templates in bin - dnoise *= padata.bank['tcounts'][bind] / padata.bank['num_t'] - logger.debug('Noise density in bin %.3g', dnoise) + dnoise *= padata.bank["tcounts"][bind] / padata.bank["num_t"] + logger.debug("Noise density in bin %.3g", dnoise) # Get signal rate density per year at given SNR - dsig = signal_pdf_from_snr(trdata['network_snr'], - padata.spec['netsnr_thresh']) - logger.debug('SNR %.3g, signal pdf %.3g', trdata['network_snr'], dsig) - dsig *= padata.spec['sig_per_yr_binned'][bind] - logger.debug('Signal density per yr per SNR in bin %.3g', dsig) + dsig = signal_pdf_from_snr(trdata["network_snr"], padata.spec["netsnr_thresh"]) + logger.debug("SNR %.3g, signal pdf %.3g", trdata["network_snr"], dsig) + dsig *= padata.spec["sig_per_yr_binned"][bind] + logger.debug("Signal density per yr per SNR in bin %.3g", dsig) # Scale by network sensitivity accounting for BNS horizon distances - dsig *= signal_rate_rescale(horizons, padata.spec['ref_bns_horizon']) - logger.debug('After horizon rescaling %.3g', dsig) + dsig *= signal_rate_rescale(horizons, padata.spec["ref_bns_horizon"]) + logger.debug("After horizon rescaling %.3g", dsig) p_astro = dsig / (dsig + dnoise) - logger.debug('p_astro %.4g', p_astro) + logger.debug("p_astro %.4g", p_astro) return p_astro, 1 - p_astro @@ -218,49 +224,47 @@ def template_param_bin_types_pa(padata, trdata, horizons): Returns ------- p_astro, p_terr: tuple of floats + """ - massspin = (trdata['mass1'], trdata['mass2'], - trdata['spin1z'], trdata['spin2z']) - trig_param = triggers.get_param(padata.spec['param'], None, *massspin) + massspin = (trdata["mass1"], trdata["mass2"], trdata["spin1z"], trdata["spin2z"]) + trig_param = triggers.get_param(padata.spec["param"], None, *massspin) # NB digitize gives '1' for first bin, '2' for second etc. - bind = numpy.digitize(trig_param, padata.bank['bin_edges']) - 1 - logger.debug('Trigger %s is in bin %i', padata.spec['param'], bind) + bind = numpy.digitize(trig_param, padata.bank["bin_edges"]) - 1 + logger.debug("Trigger %s is in bin %i", padata.spec["param"], bind) # Get noise rate density - if 'bg_fac' not in padata.spec: - expfac = 6. + if "bg_fac" not in padata.spec: + expfac = 6.0 else: - expfac = padata.spec['bg_fac'] + expfac = padata.spec["bg_fac"] # List of ifos over trigger threshold - tr_ifos = trdata['triggered'] + tr_ifos = trdata["triggered"] # FAR is in Hz, therefore convert to rate per year (per SNR) - dnoise = noise_density_from_far(trdata['far'], expfac) * _s_per_yr - logger.debug('FAR %.3g, noise density per yr per SNR %.3g', - trdata['far'], dnoise) + dnoise = noise_density_from_far(trdata["far"], expfac) * _s_per_yr + logger.debug("FAR %.3g, noise density per yr per SNR %.3g", trdata["far"], dnoise) # Scale by fraction of templates in bin - dnoise *= padata.bank['tcounts'][bind] / padata.bank['num_t'] - logger.debug('Noise density in bin %.3g', dnoise) + dnoise *= padata.bank["tcounts"][bind] / padata.bank["num_t"] + logger.debug("Noise density in bin %.3g", dnoise) # Back out trials factor to give noise density for triggered event type - dnoise /= float(trials_type(len(tr_ifos), len(trdata['sensitive']))) - logger.debug('Divide by previously applied trials factor: %.3g', dnoise) + dnoise /= float(trials_type(len(tr_ifos), len(trdata["sensitive"]))) + logger.debug("Divide by previously applied trials factor: %.3g", dnoise) # Get signal rate density per year at given SNR - dsig = signal_pdf_from_snr(trdata['network_snr'], - padata.spec['netsnr_thresh']) - logger.debug('SNR %.3g, signal pdf %.3g', trdata['network_snr'], dsig) - dsig *= padata.spec['sig_per_yr_binned'][bind] - logger.debug('Total signal density per yr per SNR in bin %.3g', dsig) + dsig = signal_pdf_from_snr(trdata["network_snr"], padata.spec["netsnr_thresh"]) + logger.debug("SNR %.3g, signal pdf %.3g", trdata["network_snr"], dsig) + dsig *= padata.spec["sig_per_yr_binned"][bind] + logger.debug("Total signal density per yr per SNR in bin %.3g", dsig) # Scale by network sensitivity accounting for BNS horizons - dsig *= signal_rate_rescale(horizons, padata.spec['ref_bns_horizon']) - logger.debug('After network horizon rescaling %.3g', dsig) + dsig *= signal_rate_rescale(horizons, padata.spec["ref_bns_horizon"]) + logger.debug("After network horizon rescaling %.3g", dsig) # Scale by relative signal rate in triggered ifos - dsig *= signal_rate_trig_type(horizons, trdata['sensitive'], tr_ifos) - logger.debug('After triggered ifo rate rescaling %.3g', dsig) + dsig *= signal_rate_trig_type(horizons, trdata["sensitive"], tr_ifos) + logger.debug("After triggered ifo rate rescaling %.3g", dsig) p_astro = dsig / (dsig + dnoise) - logger.debug('p_astro %.4g', p_astro) + logger.debug("p_astro %.4g", p_astro) return p_astro, 1 - p_astro @@ -278,6 +282,7 @@ def template_param_bin_types_farlim_pa(padata, trdata, horizons): Returns ------- p_astro, p_terr: tuple of floats + """ # If the network SNR and FAR indicate saturation of the FAR estimate, # set them to specified fixed values @@ -289,12 +294,12 @@ def template_param_bin_types_farlim_pa(padata, trdata, horizons): __all__ = [ "check_template_param_bin_data", - "read_template_bank_param", "noise_density_from_far", + "read_template_bank_param", "signal_pdf_from_snr", "signal_rate_rescale", "signal_rate_trig_type", "template_param_bin_pa", - "template_param_bin_types_pa", "template_param_bin_types_farlim_pa", + "template_param_bin_types_pa", ] diff --git a/pycbc/population/live_pastro_utils.py b/pycbc/population/live_pastro_utils.py index be9c9d5993f..40d0288fe30 100644 --- a/pycbc/population/live_pastro_utils.py +++ b/pycbc/population/live_pastro_utils.py @@ -1,12 +1,14 @@ -import logging import json +import logging + from . import live_pastro as livepa -logger = logging.getLogger('pycbc.population.live_pastro_utils') +logger = logging.getLogger("pycbc.population.live_pastro_utils") def insert_live_pastro_option_group(parser): - """ Add low-latency p astro options to the argparser object. + """ + Add low-latency p astro options to the argparser object. Parameters ---------- @@ -17,40 +19,40 @@ def insert_live_pastro_option_group(parser): ------- live_pastro_group : Argument group object - """ - live_pastro_group = parser.add_argument_group('Options for live p_astro') - live_pastro_group.add_argument('--p-astro-spec', - help='File containing information to set ' - 'up p_astro calculation') + """ + live_pastro_group = parser.add_argument_group("Options for live p_astro") + live_pastro_group.add_argument( + "--p-astro-spec", + help="File containing information to set up p_astro calculation", + ) return live_pastro_group # Choices of p astro calc method _check_spec = { - 'template_param_bins': livepa.check_template_param_bin_data, - 'template_param_bins_types': livepa.check_template_param_bin_data, - 'template_param_bins_types_farlim': - livepa.check_template_param_bin_farlim_data + "template_param_bins": livepa.check_template_param_bin_data, + "template_param_bins_types": livepa.check_template_param_bin_data, + "template_param_bins_types_farlim": livepa.check_template_param_bin_farlim_data, } _read_bank = { - 'template_param_bins': livepa.read_template_bank_param, - 'template_param_bins_types': livepa.read_template_bank_param, - 'template_param_bins_types_farlim': livepa.read_template_bank_param + "template_param_bins": livepa.read_template_bank_param, + "template_param_bins_types": livepa.read_template_bank_param, + "template_param_bins_types_farlim": livepa.read_template_bank_param, } _do_calc = { - 'template_param_bins': livepa.template_param_bin_pa, - 'template_param_bins_types': livepa.template_param_bin_types_pa, - 'template_param_bins_types_farlim': - livepa.template_param_bin_types_farlim_pa + "template_param_bins": livepa.template_param_bin_pa, + "template_param_bins_types": livepa.template_param_bin_types_pa, + "template_param_bins_types_farlim": livepa.template_param_bin_types_farlim_pa, } -class PAstroData(): - """ Class for managing live p_astro calculation persistent info """ +class PAstroData: + """Class for managing live p_astro calculation persistent info""" + def __init__(self, specfile, bank): """ Read in spec file and extract relevant info from bank @@ -61,6 +63,7 @@ def __init__(self, specfile, bank): Path to file giving method and static data used in calculation bank: str Path to hdf template bank file + """ if specfile is None: self.do = False @@ -70,11 +73,10 @@ def __init__(self, specfile, bank): with open(specfile) as specf: self.spec_json = json.load(specf) try: - self.method = self.spec_json['method'] + self.method = self.spec_json["method"] except KeyError as ke: - raise ValueError("Can't find 'method' in p_astro spec file!") \ - from ke - logger.info('Setting up p_astro data with method %s', self.method) + raise ValueError("Can't find 'method' in p_astro spec file!") from ke + logger.info("Setting up p_astro data with method %s", self.method) self.spec = _check_spec[self.method](self.spec_json) self.bank = _read_bank[self.method](self.spec, bank) @@ -84,37 +86,37 @@ def apply_significance_limits(self, trigger_data): set them to the fixed values given in the specification. """ # This only happens for double or triple events - if len(trigger_data['triggered']) == 1: + if len(trigger_data["triggered"]) == 1: return trigger_data - if len(trigger_data['triggered']) > 1: - farlim = self.spec['limit_far'] - snrlim = self.spec['limit_snr'] + if len(trigger_data["triggered"]) > 1: + farlim = self.spec["limit_far"] + snrlim = self.spec["limit_snr"] # Only do anything if FAR and SNR are beyond given limits - if trigger_data['far'] > farlim or \ - trigger_data['network_snr'] < snrlim: + if trigger_data["far"] > farlim or trigger_data["network_snr"] < snrlim: return trigger_data - logger.debug('Truncating FAR and SNR from %f, %f to %f, %f', - trigger_data['far'], trigger_data['network_snr'], - farlim, snrlim) - trigger_data['network_snr'] = snrlim - trigger_data['far'] = farlim + logger.debug( + "Truncating FAR and SNR from %f, %f to %f, %f", + trigger_data["far"], + trigger_data["network_snr"], + farlim, + snrlim, + ) + trigger_data["network_snr"] = snrlim + trigger_data["far"] = farlim return trigger_data - raise RuntimeError('Number of triggered ifos must be >0 !') + raise RuntimeError("Number of triggered ifos must be >0 !") def do_pastro_calc(self, trigger_data, horizons): - """ No-op, or call the despatch dictionary to evaluate p_astro """ + """No-op, or call the despatch dictionary to evaluate p_astro""" if not self.do: return None, None - logger.info('Computing p_astro') + logger.info("Computing p_astro") p_astro, p_terr = _do_calc[self.method](self, trigger_data, horizons) return p_astro, p_terr -__all__ = [ - "insert_live_pastro_option_group", - "PAstroData" -] +__all__ = ["PAstroData", "insert_live_pastro_option_group"] diff --git a/pycbc/population/population_models.py b/pycbc/population/population_models.py index 2df5ba87c77..7e349007990 100644 --- a/pycbc/population/population_models.py +++ b/pycbc/population/population_models.py @@ -27,16 +27,18 @@ """ from functools import partial + import numpy as np import scipy.integrate as scipy_integrate import scipy.interpolate as scipy_interpolate from astropy import units -from pycbc.cosmology import get_cosmology -from pycbc.cosmology import cosmological_quantity_from_redshift + +from pycbc.cosmology import cosmological_quantity_from_redshift, get_cosmology def sfr_grb_2008(z): - r""" The star formation rate (SFR) calibrated by high-z GRBs data. + r""" + The star formation rate (SFR) calibrated by high-z GRBs data. Parameters ---------- @@ -51,18 +53,22 @@ def sfr_grb_2008(z): Note ---- Please see Eq.(5) in for more details. - """ + """ rho_local = 0.02 # Msolar/yr/Mpc^3 eta = -10 - rho_z = rho_local*((1+z)**(3.4*eta) + ((1+z)/5000)**(-0.3*eta) + - ((1+z)/9)**(-3.5*eta))**(1./eta) + rho_z = rho_local * ( + (1 + z) ** (3.4 * eta) + + ((1 + z) / 5000) ** (-0.3 * eta) + + ((1 + z) / 9) ** (-3.5 * eta) + ) ** (1.0 / eta) return rho_z def sfr_madau_dickinson_2014(z, gamma=2.7, kappa=5.6, z_peak=1.9): - r""" The madau-dickinson 2014 star formation rate (SFR). + r""" + The madau-dickinson 2014 star formation rate (SFR). Parameters ---------- @@ -77,14 +83,15 @@ def sfr_madau_dickinson_2014(z, gamma=2.7, kappa=5.6, z_peak=1.9): Notes ----- Pease see Eq.(15) in for more details. - """ - rho_z = 0.015 * (1+z)**gamma / (1 + ((1+z)/(1+z_peak))**kappa) + """ + rho_z = 0.015 * (1 + z) ** gamma / (1 + ((1 + z) / (1 + z_peak)) ** kappa) return rho_z -def sfr_madau_fragos_2017(z, k_imf=0.66, mode='high'): - r""" The madau-fragos 2017 star formation rate (SFR), +def sfr_madau_fragos_2017(z, k_imf=0.66, mode="high"): + r""" + The madau-fragos 2017 star formation rate (SFR), which updates madau-dickinson 2014 SFR by better reproducing a number of recent 4 < z < 10 results. @@ -107,25 +114,26 @@ def sfr_madau_fragos_2017(z, k_imf=0.66, mode='high'): Notes ----- Pease see and for more details. - """ - if mode == 'low': + """ + if mode == "low": factor_a = 2.6 factor_b = 3.2 factor_c = 6.2 - elif mode == 'high': + elif mode == "high": factor_a = 2.7 factor_b = 3.0 factor_c = 5.35 else: raise ValueError("'mode' must choose from 'high' or 'low'.") - rho_z = k_imf * 0.015 * (1+z)**factor_a / (1 + ((1+z)/factor_b)**factor_c) + rho_z = k_imf * 0.015 * (1 + z) ** factor_a / (1 + ((1 + z) / factor_b) ** factor_c) return rho_z def diff_lookback_time(z, **kwargs): - r""" The derivative of lookback time t(z) + r""" + The derivative of lookback time t(z) with respect to redshit z. Parameters @@ -145,18 +153,21 @@ def diff_lookback_time(z, **kwargs): Notes ----- Pease see Eq.(A3) in for more details. + """ from sympy import sqrt cosmology = get_cosmology(**kwargs) - H0 = cosmology.H0.value * \ - (3.0856776E+19)**(-1)/(1/24/3600/365*1e-9) # Gyr^-1 - dt_dz = 1/H0/(1+z)/sqrt((cosmology.Ode0+cosmology.Om0*(1+z)**3)) + H0 = ( + cosmology.H0.value * (3.0856776e19) ** (-1) / (1 / 24 / 3600 / 365 * 1e-9) + ) # Gyr^-1 + dt_dz = 1 / H0 / (1 + z) / sqrt(cosmology.Ode0 + cosmology.Om0 * (1 + z) ** 3) return dt_dz def p_tau(tau, td_model="inverse"): - r""" The probability distribution of the time delay. + r""" + The probability distribution of the time delay. Parameters ---------- @@ -175,40 +186,51 @@ def p_tau(tau, td_model="inverse"): Notes ----- Pease see the Appendix in for more details. + """ - from sympy import sqrt, exp, log, Piecewise + from sympy import Piecewise, exp, log, sqrt if td_model == "log_normal": t_ln = 2.9 # Gyr sigma_ln = 0.2 - p_t = exp(-(log(tau)-log(t_ln))**2/(2*sigma_ln**2)) / \ - (sqrt(2*np.pi)*sigma_ln) + p_t = exp(-((log(tau) - log(t_ln)) ** 2) / (2 * sigma_ln**2)) / ( + sqrt(2 * np.pi) * sigma_ln + ) elif td_model == "gaussian": t_g = 2 # Gyr sigma_g = 0.3 - p_t = exp(-(tau-t_g)**2/(2*sigma_g**2)) / (sqrt(2*np.pi)*sigma_g) + p_t = exp(-((tau - t_g) ** 2) / (2 * sigma_g**2)) / (sqrt(2 * np.pi) * sigma_g) elif td_model == "power_law": alpha_t = 0.81 - p_t = tau**(-alpha_t) + p_t = tau ** (-alpha_t) elif td_model == "inverse": # make sure that there is a minimum and maximum time delay - td_min = 0.02 # Taken from Regimbau et al. https://journals.aps.org/prd/abstract/10.1103/PhysRevD.86.122001 - td_max = cosmological_quantity_from_redshift(0, 'age') - norm_const = 1/np.log(td_max/td_min) + td_min = 0.02 # Taken from Regimbau et al. https://journals.aps.org/prd/abstract/10.1103/PhysRevD.86.122001 + td_max = cosmological_quantity_from_redshift(0, "age") + norm_const = 1 / np.log(td_max / td_min) if isinstance(tau, (float, int)) or isinstance(tau, np.ndarray): - p_t = np.where((tau < td_min) | (tau > td_max), 0, norm_const * tau**(-0.999)) + p_t = np.where( + (tau < td_min) | (tau > td_max), 0, norm_const * tau ** (-0.999) + ) else: - p_t = Piecewise((0, tau < td_min), (0, tau > td_max), (norm_const * tau**(-0.999), True)) + p_t = Piecewise( + (0, tau < td_min), + (0, tau > td_max), + (norm_const * tau ** (-0.999), True), + ) else: - raise ValueError("'model' must choose from \ - ['log_normal', 'gaussian', 'power_law', 'inverse'].") + raise ValueError( + "'model' must choose from \ + ['log_normal', 'gaussian', 'power_law', 'inverse']." + ) return p_t def convolution_trans(sfr, diff_lookback_t, model_td, **kwargs): - r""" This function is used in a symbolic integral, which to calculate + r""" + This function is used in a symbolic integral, which to calculate the merger rate density of CBC sources. This function converts the convolution of the star formation rate SFR(tau) and the time delay probability P(tau) on the time delay 'tau' into the convolution on @@ -236,26 +258,31 @@ def convolution_trans(sfr, diff_lookback_t, model_td, **kwargs): Notes ----- Pease see Eq.(A2) in for more details. + """ from sympy import integrate, symbols - if model_td not in ['log_normal', 'gaussian', 'power_law', 'inverse']: - raise ValueError("'model_td' must choose from \ - ['log_normal', 'gaussian', 'power_law', 'inverse'].") + if model_td not in ["log_normal", "gaussian", "power_law", "inverse"]: + raise ValueError( + "'model_td' must choose from \ + ['log_normal', 'gaussian', 'power_law', 'inverse']." + ) # Fix the cosmology, set 'z/z_0' to be the only # parameter in the symbolic integration. diff_lookback_time_z = partial(diff_lookback_t, **kwargs) - z = symbols('z') - z_0 = symbols('z_0') + z = symbols("z") + z_0 = symbols("z_0") tau = integrate(diff_lookback_time_z(z), (z, z_0, z)) func = sfr(z) * p_tau(tau, model_td) * diff_lookback_time_z(z) return func -def merger_rate_density(sfr_func, td_model, rho_local, maxz=10.0, - npoints=10000, z_array=None, **kwargs): - r""" This function uses the symbolic integral to calculate +def merger_rate_density( + sfr_func, td_model, rho_local, maxz=10.0, npoints=10000, z_array=None, **kwargs +): + r""" + This function uses the symbolic integral to calculate the merger rate density of CBC sources. This function converts the convolution of the star formation rate SFR(tau) and the time delay probability P(tau) on the time delay 'tau' into the convolution on @@ -290,36 +317,38 @@ def merger_rate_density(sfr_func, td_model, rho_local, maxz=10.0, Notes ----- Pease see Eq.(A1), Eq.(A2) in for more details. + """ - from sympy import symbols, lambdify + from sympy import lambdify, symbols if z_array is None: z_array = np.linspace(0, maxz, npoints) - if td_model not in ['log_normal', 'gaussian', 'power_law', 'inverse']: - raise ValueError("'td_model' must choose from \ - ['log_normal', 'gaussian', 'power_law', 'inverse'].") + if td_model not in ["log_normal", "gaussian", "power_law", "inverse"]: + raise ValueError( + "'td_model' must choose from \ + ['log_normal', 'gaussian', 'power_law', 'inverse']." + ) - z = symbols('z') - z_0 = symbols('z_0') + z = symbols("z") + z_0 = symbols("z_0") f_z = np.zeros(len(z_array)) func_1 = convolution_trans( - sfr=sfr_func, diff_lookback_t=diff_lookback_time, - model_td=td_model, **kwargs) + sfr=sfr_func, diff_lookback_t=diff_lookback_time, model_td=td_model, **kwargs + ) for i in range(len(z_array)): - func_2 = lambdify(z, func_1.subs(z_0, z_array[i]), 'scipy') - f_z[i] = scipy_integrate.quad( - func_2, z_array[i], np.inf, epsabs=1.49e-3)[0] + func_2 = lambdify(z, func_1.subs(z_0, z_array[i]), "scipy") + f_z[i] = scipy_integrate.quad(func_2, z_array[i], np.inf, epsabs=1.49e-3)[0] - f_z = f_z/f_z[0]*rho_local # Normalize & Rescale + f_z = f_z / f_z[0] * rho_local # Normalize & Rescale rho_z = scipy_interpolate.interp1d(z_array, f_z) return rho_z -def coalescence_rate(rate_den, maxz=10.0, npoints=10000, - z_array=None, **kwargs): - r""" This function calculates the coalescence(merger) rate at the redshift z. +def coalescence_rate(rate_den, maxz=10.0, npoints=10000, z_array=None, **kwargs): + r""" + This function calculates the coalescence(merger) rate at the redshift z. Parameters ---------- @@ -346,8 +375,8 @@ def coalescence_rate(rate_den, maxz=10.0, npoints=10000, Notes ----- Pease see Eq.(1) in for more details. - """ + """ if z_array is None: z_array = np.linspace(0, maxz, npoints) @@ -355,17 +384,21 @@ def coalescence_rate(rate_den, maxz=10.0, npoints=10000, cosmology = get_cosmology(**kwargs) for z in z_array: - dr = cosmology.differential_comoving_volume(z) / (1+z) - dr_dz.append((dr*4*np.pi*units.sr*rate_den(z)*(units.Mpc)**(-3)).value) + dr = cosmology.differential_comoving_volume(z) / (1 + z) + dr_dz.append( + (dr * 4 * np.pi * units.sr * rate_den(z) * (units.Mpc) ** (-3)).value + ) coalescence_rate_interp = scipy_interpolate.interp1d( - z_array, dr_dz, fill_value='extrapolate') + z_array, dr_dz, fill_value="extrapolate" + ) return coalescence_rate_interp def total_rate_upto_redshift(z, merger_rate): - r"""Total rate of occurrences out to some redshift. + r""" + Total rate of occurrences out to some redshift. Parameters ---------- @@ -381,29 +414,32 @@ def total_rate_upto_redshift(z, merger_rate): rate: float or list The total rate of occurrences out to some redshift. In the unit of "yr^-1". - """ + """ if isinstance(z, (float, int)): total_rate = scipy_integrate.quad( - merger_rate, 0, z, - epsabs=2.00e-4, epsrel=2.00e-4, limit=1000)[0] + merger_rate, 0, z, epsabs=2.00e-4, epsrel=2.00e-4, limit=1000 + )[0] elif isinstance(z, (tuple, np.ndarray, list)): total_rate = [] for redshift in z: total_rate.append( scipy_integrate.quad( - merger_rate, 0, redshift, - epsabs=2.00e-4, epsrel=2.00e-4, limit=1000)[0] + merger_rate, 0, redshift, epsabs=2.00e-4, epsrel=2.00e-4, limit=1000 + )[0] ) else: - raise ValueError("'z' must be 'int', 'float', 'tuple', \ - 'numpy.ndarray' or 'list'.") + raise ValueError( + "'z' must be 'int', 'float', 'tuple', \ + 'numpy.ndarray' or 'list'." + ) return total_rate def average_time_between_signals(z_array, merger_rate): - r""" This function calculates the average time interval + r""" + This function calculates the average time interval of a certain type of CBC source. Parameters @@ -418,16 +454,16 @@ def average_time_between_signals(z_array, merger_rate): ------- average_time : float The average time interval (s). - """ - total_rate = total_rate_upto_redshift( - z_array[-1], merger_rate) # yr^-1 - average_time = 1./total_rate * 365*24*3600 # s + """ + total_rate = total_rate_upto_redshift(z_array[-1], merger_rate) # yr^-1 + average_time = 1.0 / total_rate * 365 * 24 * 3600 # s return average_time def norm_redshift_distribution(z_array, merger_rate): - r""" This function calculates the normalized redshift distribution + r""" + This function calculates the normalized redshift distribution of a certain type of CBC source. Parameters @@ -447,16 +483,16 @@ def norm_redshift_distribution(z_array, merger_rate): ----- The can be used as a population-informed prior for redshift and luminosity distance of CBC sources. - """ + """ lambd = average_time_between_signals(z_array, merger_rate) - norm_coalescence_rate = lambd/(365*24*3600) * merger_rate(z_array) + norm_coalescence_rate = lambd / (365 * 24 * 3600) * merger_rate(z_array) return norm_coalescence_rate -def distance_from_rate( - total_rate, merger_rate, maxz=10, npoints=10000, **kwargs): - r"""Returns the luminosity distance from the given total rate value. +def distance_from_rate(total_rate, merger_rate, maxz=10, npoints=10000, **kwargs): + r""" + Returns the luminosity distance from the given total rate value. Parameters ---------- @@ -490,20 +526,24 @@ def distance_from_rate( function to plot the curve and find the point where the curve starts to stay almost horizontal, then set `maxz` to the corresponding value and change `npoints` to a reasonable value. + """ cosmology = get_cosmology(**kwargs) - if not hasattr(merger_rate, 'dist_interp'): + if not hasattr(merger_rate, "dist_interp"): merger_rate.dist_interp = {} - if ((cosmology.name not in merger_rate.dist_interp) or - (len(merger_rate.dist_interp[cosmology.name].x) != npoints)): + if (cosmology.name not in merger_rate.dist_interp) or ( + len(merger_rate.dist_interp[cosmology.name].x) != npoints + ): + def rate_func(redshift): return total_rate_upto_redshift(redshift, merger_rate) z_array = np.linspace(0, maxz, npoints) dists = cosmological_quantity_from_redshift( - z_array, 'luminosity_distance', **kwargs) + z_array, "luminosity_distance", **kwargs + ) total_rates = rate_func(z_array) interp = scipy_interpolate.interp1d(total_rates, dists) merger_rate.dist_interp[cosmology.name] = interp @@ -514,8 +554,16 @@ def rate_func(redshift): return dl -__all__ = ['sfr_grb_2008', 'sfr_madau_dickinson_2014', - 'sfr_madau_fragos_2017', 'diff_lookback_time', - 'p_tau', 'merger_rate_density', 'coalescence_rate', - 'norm_redshift_distribution', 'total_rate_upto_redshift', - 'distance_from_rate', 'average_time_between_signals'] +__all__ = [ + "average_time_between_signals", + "coalescence_rate", + "diff_lookback_time", + "distance_from_rate", + "merger_rate_density", + "norm_redshift_distribution", + "p_tau", + "sfr_grb_2008", + "sfr_madau_dickinson_2014", + "sfr_madau_fragos_2017", + "total_rate_upto_redshift", +] diff --git a/pycbc/population/rates_functions.py b/pycbc/population/rates_functions.py index 91379263054..9835941fa45 100644 --- a/pycbc/population/rates_functions.py +++ b/pycbc/population/rates_functions.py @@ -3,42 +3,43 @@ """ import numpy as np +import scipy.stats as ss from numpy import log from scipy import integrate, optimize -import scipy.stats as ss from pycbc.conversions import mchirp_from_mass1_mass2 from pycbc.io.hdf import HFile def process_full_data(fname, rhomin, mass1, mass2, lo_mchirp, hi_mchirp): - """Read the zero-lag and time-lag triggers identified by templates in - a specified range of chirp mass. - - Parameters - ---------- - hdfile: - File that stores all the triggers - rhomin: float - Minimum value of SNR threhold (will need including ifar) - mass1: array - First mass of the waveform in the template bank - mass2: array - Second mass of the waveform in the template bank - lo_mchirp: float - Minimum chirp mass for the template - hi_mchirp: float - Maximum chirp mass for the template - - Returns - ------- - dictionary - containing foreground triggers and background information """ - with HFile(fname, 'r') as bulk: + Read the zero-lag and time-lag triggers identified by templates in + a specified range of chirp mass. + + Parameters + ---------- + hdfile: + File that stores all the triggers + rhomin: float + Minimum value of SNR threhold (will need including ifar) + mass1: array + First mass of the waveform in the template bank + mass2: array + Second mass of the waveform in the template bank + lo_mchirp: float + Minimum chirp mass for the template + hi_mchirp: float + Maximum chirp mass for the template + + Returns + ------- + dictionary + containing foreground triggers and background information - id_bkg = bulk['background_exc/template_id'][:] - id_fg = bulk['foreground/template_id'][:] + """ + with HFile(fname, "r") as bulk: + id_bkg = bulk["background_exc/template_id"][:] + id_fg = bulk["foreground/template_id"][:] mchirp_bkg = mchirp_from_mass1_mass2(mass1[id_bkg], mass2[id_bkg]) bound = np.sign((mchirp_bkg - lo_mchirp) * (hi_mchirp - mchirp_bkg)) @@ -47,48 +48,53 @@ def process_full_data(fname, rhomin, mass1, mass2, lo_mchirp, hi_mchirp): bound = np.sign((mchirp_fg - lo_mchirp) * (hi_mchirp - mchirp_fg)) idx_fg = np.where(bound == 1) - zerolagstat = bulk['foreground/stat'][:][idx_fg] - cstat_back_exc = bulk['background_exc/stat'][:][idx_bkg] - dec_factors = bulk['background_exc/decimation_factor'][:][idx_bkg] + zerolagstat = bulk["foreground/stat"][:][idx_fg] + cstat_back_exc = bulk["background_exc/stat"][:][idx_bkg] + dec_factors = bulk["background_exc/decimation_factor"][:][idx_bkg] - return {'zerolagstat': zerolagstat[zerolagstat > rhomin], - 'dec_factors': dec_factors[cstat_back_exc > rhomin], - 'cstat_back_exc': cstat_back_exc[cstat_back_exc > rhomin]} + return { + "zerolagstat": zerolagstat[zerolagstat > rhomin], + "dec_factors": dec_factors[cstat_back_exc > rhomin], + "cstat_back_exc": cstat_back_exc[cstat_back_exc > rhomin], + } def save_bkg_falloff(fname_statmap, fname_bank, path, rhomin, lo_mchirp, hi_mchirp): - ''' Read the STATMAP files to derive snr falloff for the background events. - Save the output to a txt file - Bank file is also provided to restrict triggers to BBH templates. - - Parameters - ---------- - fname_statmap: string - STATMAP file containing trigger information - fname_bank: string - File name of the template bank - path: string - Destination where txt file is saved - rhomin: float - Minimum value of SNR threhold (will need including ifar) - lo_mchirp: float - Minimum chirp mass for the template - hi_mchirp: float - Maximum chirp mass for template - ''' - - with HFile(fname_bank, 'r') as bulk: - mass1_bank = bulk['mass1'][:] - mass2_bank = bulk['mass2'][:] - full_data = process_full_data(fname_statmap, rhomin, - mass1_bank, mass2_bank, lo_mchirp, hi_mchirp) - - max_bg_stat = np.max(full_data['cstat_back_exc']) + """ + Read the STATMAP files to derive snr falloff for the background events. + Save the output to a txt file + Bank file is also provided to restrict triggers to BBH templates. + + Parameters + ---------- + fname_statmap: string + STATMAP file containing trigger information + fname_bank: string + File name of the template bank + path: string + Destination where txt file is saved + rhomin: float + Minimum value of SNR threhold (will need including ifar) + lo_mchirp: float + Minimum chirp mass for the template + hi_mchirp: float + Maximum chirp mass for template + + """ + with HFile(fname_bank, "r") as bulk: + mass1_bank = bulk["mass1"][:] + mass2_bank = bulk["mass2"][:] + full_data = process_full_data( + fname_statmap, rhomin, mass1_bank, mass2_bank, lo_mchirp, hi_mchirp + ) + + max_bg_stat = np.max(full_data["cstat_back_exc"]) bg_bins = np.linspace(rhomin, max_bg_stat, 76) - bg_counts = np.histogram(full_data['cstat_back_exc'], - weights=full_data['dec_factors'], bins=bg_bins)[0] + bg_counts = np.histogram( + full_data["cstat_back_exc"], weights=full_data["dec_factors"], bins=bg_bins + )[0] - zerolagstat = full_data['zerolagstat'] + zerolagstat = full_data["zerolagstat"] coincs = zerolagstat[zerolagstat >= rhomin] bkg = (bg_bins[:-1], bg_bins[1:], bg_counts) @@ -111,11 +117,10 @@ def log_rho_fgmc(t, injstats, bins): def fgmc(log_fg_ratios, mu_log_vt, sigma_log_vt, Rf, maxfg): - ''' + """ Function to fit the likelihood Fixme - ''' - - Lb = np.random.uniform(0., maxfg, len(Rf)) + """ + Lb = np.random.uniform(0.0, maxfg, len(Rf)) pquit = 0 while pquit < 0.1: @@ -133,7 +138,7 @@ def fgmc(log_fg_ratios, mu_log_vt, sigma_log_vt, Rf, maxfg): for lfr in log_fg_ratios: plR += np.logaddexp(lfr + log_Lf, log_Lb) - plR -= (Lf + Lb) + plR -= Lf + Lb plRn = plR - max(plR) idx = np.exp(plRn) > np.random.random(len(plRn)) @@ -146,28 +151,32 @@ def fgmc(log_fg_ratios, mu_log_vt, sigma_log_vt, Rf, maxfg): def _optm(x, alpha, mu, sigma): - '''Return probability density of skew-lognormal - See scipy.optimize.curve_fit - ''' + """ + Return probability density of skew-lognormal + See scipy.optimize.curve_fit + """ return ss.skewnorm.pdf(x, alpha, mu, sigma) def fit(R): - ''' Fit skew - lognormal to the rate samples achived from a prior analysis - Parameters - ---------- - R: array - Rate samples - Returns - ------- - ff[0]: float - The skewness - ff[1]: float - The mean - ff[2]: float - The standard deviation - ''' + """ + Fit skew - lognormal to the rate samples achived from a prior analysis + + Parameters + ---------- + R: array + Rate samples + + Returns + ------- + ff[0]: float + The skewness + ff[1]: float + The mean + ff[2]: float + The standard deviation + """ lR = np.log(R) mu_norm, sigma_norm = np.mean(lR), np.std(lR) @@ -177,36 +186,39 @@ def fit(R): # Initial guess has been taken as the mean and std-dev of the data # And a guess assuming small skewness - ff = optimize.curve_fit(_optm, xs, pxs, p0 = [0.1, mu_norm, sigma_norm])[0] + ff = optimize.curve_fit(_optm, xs, pxs, p0=[0.1, mu_norm, sigma_norm])[0] return ff[0], ff[1], ff[2] def skew_lognormal_samples(alpha, mu, sigma, minrp, maxrp): - ''' Returns a large number of Skew lognormal samples - Parameters - ---------- - alpha: float - Skewness of the distribution - mu: float - Mean of the distribution - sigma: float - Scale of the distribution - minrp: float - Minimum value for the samples - maxrp: float - Maximum value for the samples - Returns - ------- - Rfs: array - Large number of samples (may need fixing) - ''' + """ + Returns a large number of Skew lognormal samples + + Parameters + ---------- + alpha: float + Skewness of the distribution + mu: float + Mean of the distribution + sigma: float + Scale of the distribution + minrp: float + Minimum value for the samples + maxrp: float + Maximum value for the samples + + Returns + ------- + Rfs: array + Large number of samples (may need fixing) + """ nsamp = 100000000 lRu = np.random.uniform(minrp, maxrp, nsamp) plRu = ss.skewnorm.pdf(lRu, alpha, mu, sigma) rndn = np.random.random(nsamp) maxp = max(plRu) - idx = np.where(plRu/maxp > rndn) + idx = np.where(plRu / maxp > rndn) log_Rf = lRu[idx] Rfs = np.exp(log_Rf) @@ -215,33 +227,39 @@ def skew_lognormal_samples(alpha, mu, sigma, minrp, maxrp): # The flat in log and power-law mass distribution models # + # PDF for the two canonical models plus flat in mass model def prob_lnm(m1, m2, s1z, s2z, **kwargs): - ''' Return probability density for uniform in log - Parameters - ---------- - m1: array - Component masses 1 - m2: array - Component masses 2 - s1z: array - Aligned spin 1(Not in use currently) - s2z: - Aligned spin 2(Not in use currently) - **kwargs: string - Keyword arguments as model parameters - Returns - ------- - p_m1_m2: array - The probability density for m1, m2 pair - ''' - - min_mass = kwargs.get('min_mass', 5.) - max_mass = kwargs.get('max_mass', 95.) + """ + Return probability density for uniform in log + + Parameters + ---------- + m1: array + Component masses 1 + m2: array + Component masses 2 + s1z: array + Aligned spin 1(Not in use currently) + s2z: + Aligned spin 2(Not in use currently) + **kwargs: string + Keyword arguments as model parameters + + Returns + ------- + p_m1_m2: array + The probability density for m1, m2 pair + + """ + min_mass = kwargs.get("min_mass", 5.0) + max_mass = kwargs.get("max_mass", 95.0) max_mtotal = min_mass + max_mass m1, m2 = np.array(m1), np.array(m2) - C_lnm = integrate.quad(lambda x: (log(max_mtotal - x) - log(min_mass))/x, min_mass, max_mass)[0] + C_lnm = integrate.quad( + lambda x: (log(max_mtotal - x) - log(min_mass)) / x, min_mass, max_mass + )[0] xx = np.minimum(m1, m2) m1 = np.maximum(m1, m2) @@ -251,41 +269,43 @@ def prob_lnm(m1, m2, s1z, s2z, **kwargs): bound += np.sign(max_mass - m1) * np.sign(m2 - min_mass) idx = np.where(bound != 2) - p_m1_m2 = (1/C_lnm)*(1./m1)*(1./m2) + p_m1_m2 = (1 / C_lnm) * (1.0 / m1) * (1.0 / m2) p_m1_m2[idx] = 0 return p_m1_m2 def prob_imf(m1, m2, s1z, s2z, **kwargs): - ''' Return probability density for power-law - Parameters - ---------- - m1: array - Component masses 1 - m2: array - Component masses 2 - s1z: array - Aligned spin 1(Not in use currently) - s2z: - Aligned spin 2(Not in use currently) - **kwargs: string - Keyword arguments as model parameters - - Returns - ------- - p_m1_m2: array - the probability density for m1, m2 pair - ''' - - min_mass = kwargs.get('min_mass', 5.) - max_mass = kwargs.get('max_mass', 95.) - alpha = kwargs.get('alpha', -2.35) + """ + Return probability density for power-law + + Parameters + ---------- + m1: array + Component masses 1 + m2: array + Component masses 2 + s1z: array + Aligned spin 1(Not in use currently) + s2z: + Aligned spin 2(Not in use currently) + **kwargs: string + Keyword arguments as model parameters + + Returns + ------- + p_m1_m2: array + the probability density for m1, m2 pair + + """ + min_mass = kwargs.get("min_mass", 5.0) + max_mass = kwargs.get("max_mass", 95.0) + alpha = kwargs.get("alpha", -2.35) max_mtotal = min_mass + max_mass m1, m2 = np.array(m1), np.array(m2) - C_imf = max_mass**(alpha + 1)/(alpha + 1) - C_imf -= min_mass**(alpha + 1)/(alpha + 1) + C_imf = max_mass ** (alpha + 1) / (alpha + 1) + C_imf -= min_mass ** (alpha + 1) / (alpha + 1) xx = np.minimum(m1, m2) m1 = np.maximum(m1, m2) @@ -296,44 +316,46 @@ def prob_imf(m1, m2, s1z, s2z, **kwargs): idx = np.where(bound != 2) p_m1_m2 = np.zeros_like(m1) - idx = np.where(m1 <= max_mtotal/2.) - p_m1_m2[idx] = (1./C_imf) * m1[idx]**alpha /(m1[idx] - min_mass) - idx = np.where(m1 > max_mtotal/2.) - p_m1_m2[idx] = (1./C_imf) * m1[idx]**alpha /(max_mass - m1[idx]) + idx = np.where(m1 <= max_mtotal / 2.0) + p_m1_m2[idx] = (1.0 / C_imf) * m1[idx] ** alpha / (m1[idx] - min_mass) + idx = np.where(m1 > max_mtotal / 2.0) + p_m1_m2[idx] = (1.0 / C_imf) * m1[idx] ** alpha / (max_mass - m1[idx]) p_m1_m2[idx] = 0 - return p_m1_m2/2. + return p_m1_m2 / 2.0 def prob_flat(m1, m2, s1z, s2z, **kwargs): - ''' Return probability density for uniform in component mass - Parameters - ---------- - m1: array - Component masses 1 - m2: array - Component masses 2 - s1z: array - Aligned spin 1 (not in use currently) - s2z: - Aligned spin 2 (not in use currently) - **kwargs: string - Keyword arguments as model parameters - - Returns - ------- - p_m1_m2: array - the probability density for m1, m2 pair - ''' - - min_mass = kwargs.get('min_mass', 1.) - max_mass = kwargs.get('max_mass', 2.) + """ + Return probability density for uniform in component mass + + Parameters + ---------- + m1: array + Component masses 1 + m2: array + Component masses 2 + s1z: array + Aligned spin 1 (not in use currently) + s2z: + Aligned spin 2 (not in use currently) + **kwargs: string + Keyword arguments as model parameters + + Returns + ------- + p_m1_m2: array + the probability density for m1, m2 pair + + """ + min_mass = kwargs.get("min_mass", 1.0) + max_mass = kwargs.get("max_mass", 2.0) bound = np.sign(m1 - m2) bound += np.sign(max_mass - m1) * np.sign(m2 - min_mass) idx = np.where(bound != 2) - p_m1_m2 = 2. / (max_mass - min_mass)**2 + p_m1_m2 = 2.0 / (max_mass - min_mass) ** 2 p_m1_m2[idx] = 0 return p_m1_m2 @@ -341,32 +363,33 @@ def prob_flat(m1, m2, s1z, s2z, **kwargs): # Generate samples for the two canonical models plus flat in mass model def draw_imf_samples(**kwargs): - ''' Draw samples for power-law model - - Parameters - ---------- - **kwargs: string - Keyword arguments as model parameters and number of samples - - Returns - ------- - array - The first mass - array - The second mass - ''' - - alpha_salpeter = kwargs.get('alpha', -2.35) - nsamples = kwargs.get('nsamples', 1) - min_mass = kwargs.get('min_mass', 5.) - max_mass = kwargs.get('max_mass', 95.) + """ + Draw samples for power-law model + + Parameters + ---------- + **kwargs: string + Keyword arguments as model parameters and number of samples + + Returns + ------- + array + The first mass + array + The second mass + + """ + alpha_salpeter = kwargs.get("alpha", -2.35) + nsamples = kwargs.get("nsamples", 1) + min_mass = kwargs.get("min_mass", 5.0) + max_mass = kwargs.get("max_mass", 95.0) max_mtotal = min_mass + max_mass - a = (max_mass/min_mass)**(alpha_salpeter + 1.0) - 1.0 + a = (max_mass / min_mass) ** (alpha_salpeter + 1.0) - 1.0 beta = 1.0 / (alpha_salpeter + 1.0) - k = nsamples * int(1.5 + log(1 + 100./nsamples)) - aa = min_mass * (1.0 + a * np.random.random(k))**beta + k = nsamples * int(1.5 + log(1 + 100.0 / nsamples)) + aa = min_mass * (1.0 + a * np.random.random(k)) ** beta bb = np.random.uniform(min_mass, aa, k) idx = np.where(aa + bb < max_mtotal) @@ -376,30 +399,31 @@ def draw_imf_samples(**kwargs): def draw_lnm_samples(**kwargs): - ''' Draw samples for uniform-in-log model - - Parameters - ---------- - **kwargs: string - Keyword arguments as model parameters and number of samples - - Returns - ------- - array - The first mass - array - The second mass - ''' - - #PDF doesnt match with sampler - nsamples = kwargs.get('nsamples', 1) - min_mass = kwargs.get('min_mass', 5.) - max_mass = kwargs.get('max_mass', 95.) + """ + Draw samples for uniform-in-log model + + Parameters + ---------- + **kwargs: string + Keyword arguments as model parameters and number of samples + + Returns + ------- + array + The first mass + array + The second mass + + """ + # PDF doesnt match with sampler + nsamples = kwargs.get("nsamples", 1) + min_mass = kwargs.get("min_mass", 5.0) + max_mass = kwargs.get("max_mass", 95.0) max_mtotal = min_mass + max_mass lnmmin = log(min_mass) lnmmax = log(max_mass) - k = nsamples * int(1.5 + log(1 + 100./nsamples)) + k = nsamples * int(1.5 + log(1 + 100.0 / nsamples)) aa = np.exp(np.random.uniform(lnmmin, lnmmax, k)) bb = np.exp(np.random.uniform(lnmmin, lnmmax, k)) @@ -410,25 +434,26 @@ def draw_lnm_samples(**kwargs): def draw_flat_samples(**kwargs): - ''' Draw samples for uniform in mass - - Parameters - ---------- - **kwargs: string - Keyword arguments as model parameters and number of samples - - Returns - ------- - array - The first mass - array - The second mass - ''' - - #PDF doesnt match with sampler - nsamples = kwargs.get('nsamples', 1) - min_mass = kwargs.get('min_mass', 1.) - max_mass = kwargs.get('max_mass', 2.) + """ + Draw samples for uniform in mass + + Parameters + ---------- + **kwargs: string + Keyword arguments as model parameters and number of samples + + Returns + ------- + array + The first mass + array + The second mass + + """ + # PDF doesnt match with sampler + nsamples = kwargs.get("nsamples", 1) + min_mass = kwargs.get("min_mass", 1.0) + max_mass = kwargs.get("max_mass", 2.0) m1 = np.random.uniform(min_mass, max_mass, nsamples) m2 = np.random.uniform(min_mass, max_mass, nsamples) @@ -438,18 +463,20 @@ def draw_flat_samples(**kwargs): # Functions to generate chirp mass samples for the two canonical models def mchirp_sampler_lnm(**kwargs): - ''' Draw chirp mass samples for uniform-in-log model - - Parameters - ---------- - **kwargs: string - Keyword arguments as model parameters and number of samples - - Returns - ------- - mchirp-astro: array - The chirp mass samples for the population - ''' + """ + Draw chirp mass samples for uniform-in-log model + + Parameters + ---------- + **kwargs: string + Keyword arguments as model parameters and number of samples + + Returns + ------- + mchirp-astro: array + The chirp mass samples for the population + + """ m1, m2 = draw_lnm_samples(**kwargs) mchirp_astro = mchirp_from_mass1_mass2(m1, m2) @@ -457,18 +484,20 @@ def mchirp_sampler_lnm(**kwargs): def mchirp_sampler_imf(**kwargs): - ''' Draw chirp mass samples for power-law model - - Parameters - ---------- - **kwargs: string - Keyword arguments as model parameters and number of samples - - Returns - ------- - mchirp-astro: array - The chirp mass samples for the population - ''' + """ + Draw chirp mass samples for power-law model + + Parameters + ---------- + **kwargs: string + Keyword arguments as model parameters and number of samples + + Returns + ------- + mchirp-astro: array + The chirp mass samples for the population + + """ m1, m2 = draw_imf_samples(**kwargs) mchirp_astro = mchirp_from_mass1_mass2(m1, m2) @@ -476,18 +505,20 @@ def mchirp_sampler_imf(**kwargs): def mchirp_sampler_flat(**kwargs): - ''' Draw chirp mass samples for flat in mass model - - Parameters - ---------- - **kwargs: string - Keyword arguments as model parameters and number of samples - - Returns - ------- - mchirp-astro: array - The chirp mass samples for the population - ''' + """ + Draw chirp mass samples for flat in mass model + + Parameters + ---------- + **kwargs: string + Keyword arguments as model parameters and number of samples + + Returns + ------- + mchirp-astro: array + The chirp mass samples for the population + + """ m1, m2 = draw_flat_samples(**kwargs) mchirp_astro = mchirp_from_mass1_mass2(m1, m2) diff --git a/pycbc/population/scale_injections.py b/pycbc/population/scale_injections.py index 125dffe0036..5e8e6b461fc 100644 --- a/pycbc/population/scale_injections.py +++ b/pycbc/population/scale_injections.py @@ -1,323 +1,341 @@ +import copy + import numpy as np +from astropy.cosmology import WMAP9 as cosmo from numpy import log -import copy -from scipy.interpolate import interp1d from scipy.integrate import quad -from astropy.cosmology import WMAP9 as cosmo +from scipy.interpolate import interp1d from pycbc.conversions import mchirp_from_mass1_mass2 as m1m2tomch from pycbc.io.hdf import HFile -_mch_BNS = 1.4/2**.2 -_redshifts, _d_lum, _I = np.arange(0., 5., 0.01), [], [] -_save_params = ['mass1', 'mass2', 'spin1z', 'spin2z', 'spin1y', 'spin2y', - 'spin1x', 'spin2x', 'distance', 'end_time'] +_mch_BNS = 1.4 / 2**0.2 +_redshifts, _d_lum, _I = np.arange(0.0, 5.0, 0.01), [], [] +_save_params = [ + "mass1", + "mass2", + "spin1z", + "spin2z", + "spin1y", + "spin2y", + "spin1x", + "spin2x", + "distance", + "end_time", +] for zz in _redshifts: _d_lum.append(cosmo.luminosity_distance(zz).value) _dlum_interp = interp1d(_d_lum, _redshifts) + def read_injections(sim_files, m_dist, s_dist, d_dist): - ''' Read all the injections from the files in the provided folder. - The files must belong to individual set i.e. no files that combine - all the injections in a run. - Identify injection strategies and finds parameter boundaries. - Collect injection according to GPS. - - Parameters - ---------- - sim_files: list - List containign names of the simulation files - m_dist: list - The mass distribution used in the simulation runs - s_dist: list - The spin distribution used in the simulation runs - d_dist: list - The distance distribution used in the simulation runs - - Returns - ------- - injections: dictionary - Contains the organized information about the injections - ''' + """ + Read all the injections from the files in the provided folder. + The files must belong to individual set i.e. no files that combine + all the injections in a run. + Identify injection strategies and finds parameter boundaries. + Collect injection according to GPS. + + Parameters + ---------- + sim_files: list + List containign names of the simulation files + m_dist: list + The mass distribution used in the simulation runs + s_dist: list + The spin distribution used in the simulation runs + d_dist: list + The distance distribution used in the simulation runs + + Returns + ------- + injections: dictionary + Contains the organized information about the injections + """ injections = {} min_d, max_d = 1e12, 0 nf = len(sim_files) for i in range(nf): - key = str(i) injections[key] = process_injections(sim_files[i]) - injections[key]['file_name'] = sim_files[i] - injections[key]['m_dist'] = m_dist[i] - injections[key]['s_dist'] = s_dist[i] - injections[key]['d_dist'] = d_dist[i] + injections[key]["file_name"] = sim_files[i] + injections[key]["m_dist"] = m_dist[i] + injections[key]["s_dist"] = s_dist[i] + injections[key]["d_dist"] = d_dist[i] - mass1, mass2 = injections[key]['mass1'], injections[key]['mass2'] - distance = injections[key]['distance'] + mass1, mass2 = injections[key]["mass1"], injections[key]["mass2"] + distance = injections[key]["distance"] mchirp = m1m2tomch(mass1, mass2) - injections[key]['chirp_mass'] = mchirp - injections[key]['total_mass'] = mass1 + mass2 + injections[key]["chirp_mass"] = mchirp + injections[key]["total_mass"] = mass1 + mass2 - injections[key]['mtot_range'] = [min(mass1 + mass2), max(mass1 + mass2)] - injections[key]['m1_range'] = [min(mass1), max(mass1)] - injections[key]['m2_range'] = [min(mass2), max(mass2)] - injections[key]['d_range'] = [min(distance), max(distance)] + injections[key]["mtot_range"] = [min(mass1 + mass2), max(mass1 + mass2)] + injections[key]["m1_range"] = [min(mass1), max(mass1)] + injections[key]["m2_range"] = [min(mass2), max(mass2)] + injections[key]["d_range"] = [min(distance), max(distance)] min_d, max_d = min(min_d, min(distance)), max(max_d, max(distance)) - injections['z_range'] = [dlum_to_z(min_d), dlum_to_z(max_d)] + injections["z_range"] = [dlum_to_z(min_d), dlum_to_z(max_d)] return injections + def estimate_vt(injections, mchirp_sampler, model_pdf, **kwargs): - #Try including ifar threshold - '''Based on injection strategy and the desired astro model estimate the injected volume. - Scale injections and estimate sensitive volume. - - Parameters - ---------- - injections: dictionary - Dictionary obtained after reading injections from read_injections - mchirp_sampler: function - Sampler for producing chirp mass samples for the astro model. - model_pdf: function - The PDF for astro model in mass1-mass2-spin1z-spin2z space. - This is easily extendible to include precession - kwargs: key words - Inputs for thresholds and astrophysical models - - Returns - ------- - injection_chunks: dictionary - The input dictionary with VT and VT error included with the injections - ''' - - thr_var = kwargs.get('thr_var') - thr_val = kwargs.get('thr_val') - - nsamples = 1000000 #Used to calculate injected astro volume + # Try including ifar threshold + """ + Based on injection strategy and the desired astro model estimate the injected volume. + Scale injections and estimate sensitive volume. + + Parameters + ---------- + injections: dictionary + Dictionary obtained after reading injections from read_injections + mchirp_sampler: function + Sampler for producing chirp mass samples for the astro model. + model_pdf: function + The PDF for astro model in mass1-mass2-spin1z-spin2z space. + This is easily extendible to include precession + kwargs: key words + Inputs for thresholds and astrophysical models + + Returns + ------- + injection_chunks: dictionary + The input dictionary with VT and VT error included with the injections + + """ + thr_var = kwargs.get("thr_var") + thr_val = kwargs.get("thr_val") + + nsamples = 1000000 # Used to calculate injected astro volume injections = copy.deepcopy(injections) - min_z, max_z = injections['z_range'] - V = quad(contracted_dVdc, 0., max_z)[0] + min_z, max_z = injections["z_range"] + V = quad(contracted_dVdc, 0.0, max_z)[0] z_astro = astro_redshifts(min_z, max_z, nsamples) astro_lum_dist = cosmo.luminosity_distance(z_astro).value - mch_astro = np.array(mchirp_sampler(nsamples = nsamples, **kwargs)) - mch_astro_det = mch_astro * (1. + z_astro) + mch_astro = np.array(mchirp_sampler(nsamples=nsamples, **kwargs)) + mch_astro_det = mch_astro * (1.0 + z_astro) idx_within = np.zeros(nsamples) for key in injections.keys(): - - if key == 'z_range': + if key == "z_range": # This is repeated down again and is so continue - mchirp = injections[key]['chirp_mass'] - min_mchirp, max_mchirp = min(mchirp), max(mchirp) - distance = injections[key]['distance'] + mchirp = injections[key]["chirp_mass"] + min_mchirp, max_mchirp = min(mchirp), max(mchirp) + distance = injections[key]["distance"] - if injections[key]['d_dist'] == 'uniform': + if injections[key]["d_dist"] == "uniform": d_min, d_max = min(distance), max(distance) - elif injections[key]['d_dist'] == 'dchirp': - d_fid_min = min(distance / (mchirp/_mch_BNS)**(5/6.)) - d_fid_max = max(distance / (mchirp/_mch_BNS)**(5/6.)) + elif injections[key]["d_dist"] == "dchirp": + d_fid_min = min(distance / (mchirp / _mch_BNS) ** (5 / 6.0)) + d_fid_max = max(distance / (mchirp / _mch_BNS) ** (5 / 6.0)) - d_min = d_fid_min * (mch_astro_det/_mch_BNS)**(5/6.) - d_max = d_fid_max * (mch_astro_det/_mch_BNS)**(5/6.) + d_min = d_fid_min * (mch_astro_det / _mch_BNS) ** (5 / 6.0) + d_max = d_fid_max * (mch_astro_det / _mch_BNS) ** (5 / 6.0) - bound = np.sign((max_mchirp-mch_astro_det)*(mch_astro_det-min_mchirp)) - bound += np.sign((d_max - astro_lum_dist)*(astro_lum_dist - d_min)) + bound = np.sign((max_mchirp - mch_astro_det) * (mch_astro_det - min_mchirp)) + bound += np.sign((d_max - astro_lum_dist) * (astro_lum_dist - d_min)) idx = np.where(bound == 2) idx_within[idx] = 1 - inj_V0 = 4*np.pi*V*len(idx_within[idx_within == 1])/float(nsamples) - injections['inj_astro_vol'] = inj_V0 + inj_V0 = 4 * np.pi * V * len(idx_within[idx_within == 1]) / float(nsamples) + injections["inj_astro_vol"] = inj_V0 # Estimate the sensitive volume - z_range = injections['z_range'] - V_min = quad(contracted_dVdc, 0., z_range[0])[0] - V_max = quad(contracted_dVdc, 0., z_range[1])[0] + z_range = injections["z_range"] + V_min = quad(contracted_dVdc, 0.0, z_range[0])[0] + V_max = quad(contracted_dVdc, 0.0, z_range[1])[0] thr_falloff, i_inj, i_det, i_det_sq = [], 0, 0, 0 gps_min, gps_max = 1e15, 0 keys = injections.keys() for key in keys: - - if key == 'z_range' or key == 'inj_astro_vol': + if key == "z_range" or key == "inj_astro_vol": continue data = injections[key] - distance = data['distance'] - mass1, mass2 = data['mass1'], data['mass2'] - spin1z, spin2z = data['spin1z'], data['spin2z'] - mchirp = data['chirp_mass'] - gps_min = min(gps_min, min(data['end_time'])) - gps_max = max(gps_max, max(data['end_time'])) + distance = data["distance"] + mass1, mass2 = data["mass1"], data["mass2"] + spin1z, spin2z = data["spin1z"], data["spin2z"] + mchirp = data["chirp_mass"] + gps_min = min(gps_min, min(data["end_time"])) + gps_max = max(gps_max, max(data["end_time"])) z_inj = dlum_to_z(distance) - m1_sc, m2_sc = mass1/(1 + z_inj), mass2/(1 + z_inj) + m1_sc, m2_sc = mass1 / (1 + z_inj), mass2 / (1 + z_inj) p_out = model_pdf(m1_sc, m2_sc, spin1z, spin2z) p_out *= pdf_z_astro(z_inj, V_min, V_max) p_in = 0 J = cosmo.luminosity_distance(z_inj + 0.0005).value J -= cosmo.luminosity_distance(z_inj - 0.0005).value - J = abs(J)/0.001 # A quick way to get dD_l/dz + J = abs(J) / 0.001 # A quick way to get dD_l/dz # Sum probability of injections from j-th set for all the strategies for key2 in keys: - - if key2 == 'z_range' or key2 == 'inj_astro_vol': + if key2 == "z_range" or key2 == "inj_astro_vol": continue dt_j = injections[key2] - dist_j = dt_j['distance'] - m1_j, m2_j = dt_j['mass1'], dt_j['mass2'] - s1x_2, s2x_2 = dt_j['spin1x'], dt_j['spin2x'] - s1y_2, s2y_2 = dt_j['spin1y'], dt_j['spin2y'] - s1z_2, s2z_2 = dt_j['spin1z'], dt_j['spin2z'] + dist_j = dt_j["distance"] + m1_j, m2_j = dt_j["mass1"], dt_j["mass2"] + s1x_2, s2x_2 = dt_j["spin1x"], dt_j["spin2x"] + s1y_2, s2y_2 = dt_j["spin1y"], dt_j["spin2y"] + s1z_2, s2z_2 = dt_j["spin1z"], dt_j["spin2z"] s1 = np.sqrt(s1x_2**2 + s1y_2**2 + s1z_2**2) s2 = np.sqrt(s2x_2**2 + s2y_2**2 + s2z_2**2) - mch_j = dt_j['chirp_mass'] + mch_j = dt_j["chirp_mass"] - #Get probability density for injections in mass-distance space - if dt_j['m_dist'] == 'totalMass': + # Get probability density for injections in mass-distance space + if dt_j["m_dist"] == "totalMass": lomass, himass = min(min(m1_j), min(m2_j), max(max(m1_j), max(m2_j))) lomass_2, himass_2 = lomass, himass - elif dt_j['m_dist'] == 'componentMass' or dt_j['m_dist'] == 'log': + elif dt_j["m_dist"] == "componentMass" or dt_j["m_dist"] == "log": lomass, himass = min(m1_j), max(m1_j) lomass_2, himass_2 = min(m2_j), max(m2_j) - if dt_j['d_dist'] == 'dchirp': - l_dist = min(dist_j / (mch_j/_mch_BNS)**(5/6.)) - h_dist = max(dist_j / (mch_j/_mch_BNS)**(5/6.)) - elif dt_j['d_dist'] == 'uniform': + if dt_j["d_dist"] == "dchirp": + l_dist = min(dist_j / (mch_j / _mch_BNS) ** (5 / 6.0)) + h_dist = max(dist_j / (mch_j / _mch_BNS) ** (5 / 6.0)) + elif dt_j["d_dist"] == "uniform": l_dist, h_dist = min(dist_j), max(dist_j) - mdist = dt_j['m_dist'] - prob_mass = inj_mass_pdf(mdist, mass1, mass2, - lomass, himass, lomass_2, himass_2) + mdist = dt_j["m_dist"] + prob_mass = inj_mass_pdf( + mdist, mass1, mass2, lomass, himass, lomass_2, himass_2 + ) - ddist = dt_j['d_dist'] - prob_dist = inj_distance_pdf(ddist, distance, l_dist, - h_dist, mchirp) + ddist = dt_j["d_dist"] + prob_dist = inj_distance_pdf(ddist, distance, l_dist, h_dist, mchirp) hspin1, hspin2 = max(s1), max(s2) - prob_spin = inj_spin_pdf(dt_j['s_dist'], hspin1, spin1z) - prob_spin *= inj_spin_pdf(dt_j['s_dist'], hspin2, spin2z) + prob_spin = inj_spin_pdf(dt_j["s_dist"], hspin1, spin1z) + prob_spin *= inj_spin_pdf(dt_j["s_dist"], hspin2, spin2z) - p_in += prob_mass * prob_dist * prob_spin * J * (1 + z_inj)**2 + p_in += prob_mass * prob_dist * prob_spin * J * (1 + z_inj) ** 2 p_in[p_in == 0] = 1e12 - p_out_in = p_out/p_in + p_out_in = p_out / p_in i_inj += np.sum(p_out_in) i_det += np.sum((p_out_in)[data[thr_var] > thr_val]) - i_det_sq += np.sum((p_out_in)[data[thr_var] > thr_val]**2) + i_det_sq += np.sum((p_out_in)[data[thr_var] > thr_val] ** 2) idx_thr = np.where(data[thr_var] > thr_val) thrs = data[thr_var][idx_thr] - ratios = p_out_in[idx_thr]/max(p_out_in[idx_thr]) + ratios = p_out_in[idx_thr] / max(p_out_in[idx_thr]) rndn = np.random.uniform(0, 1, len(ratios)) idx_ratio = np.where(ratios > rndn) thr_falloff.append(thrs[idx_ratio]) - inj_V0 = injections['inj_astro_vol'] - injections['ninj'] = i_inj - injections['ndet'] = i_det - injections['ndetsq'] = i_det_sq - injections['VT'] = ((inj_V0*i_det/i_inj) * (gps_max - gps_min)/31557600) - injections['VT_err'] = injections['VT'] * np.sqrt(i_det_sq)/i_det - injections['thr_falloff'] = np.hstack(np.array(thr_falloff).flat) + inj_V0 = injections["inj_astro_vol"] + injections["ninj"] = i_inj + injections["ndet"] = i_det + injections["ndetsq"] = i_det_sq + injections["VT"] = (inj_V0 * i_det / i_inj) * (gps_max - gps_min) / 31557600 + injections["VT_err"] = injections["VT"] * np.sqrt(i_det_sq) / i_det + injections["thr_falloff"] = np.hstack(np.array(thr_falloff).flat) return injections + def process_injections(hdffile): - """Function to read in the injection file and - extract the found injections and all injections - - Parameters - ---------- - hdffile: hdf file - File for which injections are to be processed - - Returns - ------- - data: dictionary - Dictionary containing injection read from the input file + """ + Function to read in the injection file and + extract the found injections and all injections + + Parameters + ---------- + hdffile: hdf file + File for which injections are to be processed + + Returns + ------- + data: dictionary + Dictionary containing injection read from the input file + """ data = {} - with HFile(hdffile, 'r') as inp: - found_index = inp['found_after_vetoes/injection_index'][:] + with HFile(hdffile, "r") as inp: + found_index = inp["found_after_vetoes/injection_index"][:] for param in _save_params: - data[param] = inp['injections/'+param][:] + data[param] = inp["injections/" + param][:] ifar = np.zeros_like(data[_save_params[0]]) - ifar[found_index] = inp['found_after_vetoes/ifar'][:] + ifar[found_index] = inp["found_after_vetoes/ifar"][:] - data['ifar'] = ifar + data["ifar"] = ifar stat = np.zeros_like(data[_save_params[0]]) - stat[found_index] = inp['found_after_vetoes/stat'][:] + stat[found_index] = inp["found_after_vetoes/stat"][:] - data['stat'] = stat + data["stat"] = stat return data + def dlum_to_z(dl): - ''' Get the redshift for a luminosity distance + """ + Get the redshift for a luminosity distance - Parameters - ---------- - dl: array - The array of luminosity distances + Parameters + ---------- + dl: array + The array of luminosity distances - Returns - ------- - array - The redshift values corresponding to the luminosity distances - ''' + Returns + ------- + array + The redshift values corresponding to the luminosity distances + """ return _dlum_interp(dl) + def astro_redshifts(min_z, max_z, nsamples): - '''Sample the redshifts for sources, with redshift - independent rate, using standard cosmology - - Parameters - ---------- - min_z: float - Minimum redshift - max_z: float - Maximum redshift - nsamples: int - Number of samples - - Returns - ------- - z_astro: array - nsamples of redshift, between min_z, max_z, by standard cosmology - ''' + """ + Sample the redshifts for sources, with redshift + independent rate, using standard cosmology + + Parameters + ---------- + min_z: float + Minimum redshift + max_z: float + Maximum redshift + nsamples: int + Number of samples + + Returns + ------- + z_astro: array + nsamples of redshift, between min_z, max_z, by standard cosmology + """ dz, fac = 0.001, 3.0 # use interpolation instead of directly estimating all the pdfz for rndz - V = quad(contracted_dVdc, 0., max_z)[0] - zbins = np.arange(min_z, max_z + dz/2., dz) + V = quad(contracted_dVdc, 0.0, max_z)[0] + zbins = np.arange(min_z, max_z + dz / 2.0, dz) zcenter = (zbins[:-1] + zbins[1:]) / 2 - pdfz = cosmo.differential_comoving_volume(zcenter).value/(1+zcenter)/V + pdfz = cosmo.differential_comoving_volume(zcenter).value / (1 + zcenter) / V int_pdf = interp1d(zcenter, pdfz, bounds_error=False, fill_value=0) - rndz = np.random.uniform(min_z, max_z, int(fac*nsamples)) + rndz = np.random.uniform(min_z, max_z, int(fac * nsamples)) pdf_zs = int_pdf(rndz) maxpdf = max(pdf_zs) - rndn = np.random.uniform(0, 1, int(fac*nsamples)) * maxpdf + rndn = np.random.uniform(0, 1, int(fac * nsamples)) * maxpdf diff = pdf_zs - rndn idx = np.where(diff > 0) z_astro = rndz[idx] @@ -327,44 +345,49 @@ def astro_redshifts(min_z, max_z, nsamples): return z_astro + def pdf_z_astro(z, V_min, V_max): - ''' Get the probability density for the rate of events - at a redshift assuming standard cosmology - ''' - return contracted_dVdc(z)/(V_max - V_min) + """ + Get the probability density for the rate of events + at a redshift assuming standard cosmology + """ + return contracted_dVdc(z) / (V_max - V_min) + def contracted_dVdc(z): - #Return the time-dilated differential comoving volume - return cosmo.differential_comoving_volume(z).value/(1+z) + # Return the time-dilated differential comoving volume + return cosmo.differential_comoving_volume(z).value / (1 + z) + ##### Defining current standard strategies used for making injections ##### -def inj_mass_pdf(key, mass1, mass2, lomass, himass, lomass_2 = 0, himass_2 = 0): - - '''Estimate the probability density based on the injection strategy - - Parameters - ---------- - key: string - Injection strategy - mass1: array - First mass of the injections - mass2: array - Second mass of the injections - lomass: float - Lower value of the mass distributions - himass: float - higher value of the mass distribution - - Returns - ------- - pdf: array - Probability density of the injections - ''' +def inj_mass_pdf(key, mass1, mass2, lomass, himass, lomass_2=0, himass_2=0): + """ + Estimate the probability density based on the injection strategy + + Parameters + ---------- + key: string + Injection strategy + mass1: array + First mass of the injections + mass2: array + Second mass of the injections + lomass: float + Lower value of the mass distributions + himass: float + higher value of the mass distribution + + Returns + ------- + pdf: array + Probability density of the injections + + """ mass1, mass2 = np.array(mass1), np.array(mass2) - if key == 'totalMass': + if key == "totalMass": # Returns the PDF of mass when total mass is uniformly distributed. # Both the component masses have the same distribution for this case. @@ -374,15 +397,15 @@ def inj_mass_pdf(key, mass1, mass2, lomass, himass, lomass_2 = 0, himass_2 = 0): # himass: higher component mass bound = np.sign((lomass + himass) - (mass1 + mass2)) - bound += np.sign((himass - mass1)*(mass1 - lomass)) - bound += np.sign((himass - mass2)*(mass2 - lomass)) + bound += np.sign((himass - mass1) * (mass1 - lomass)) + bound += np.sign((himass - mass2) * (mass2 - lomass)) idx = np.where(bound != 3) - pdf = 1./(himass - lomass)/(mass1 + mass2 - 2 * lomass) + pdf = 1.0 / (himass - lomass) / (mass1 + mass2 - 2 * lomass) pdf[idx] = 0 return pdf - if key == 'componentMass': + if key == "componentMass": # Returns the PDF of mass when component mass is uniformly # distributed. Component masses are independent for this case. @@ -391,15 +414,15 @@ def inj_mass_pdf(key, mass1, mass2, lomass, himass, lomass_2 = 0, himass_2 = 0): # lomass: lower component mass # himass: higher component mass - bound = np.sign((himass - mass1)*(mass1 - lomass)) - bound += np.sign((himass_2 - mass2)*(mass2 - lomass_2)) + bound = np.sign((himass - mass1) * (mass1 - lomass)) + bound += np.sign((himass_2 - mass2) * (mass2 - lomass_2)) idx = np.where(bound != 2) pdf = np.ones_like(mass1) / (himass - lomass) / (himass_2 - lomass_2) pdf[idx] = 0 return pdf - if key == 'log': + if key == "log": # Returns the PDF of mass when component mass is uniform in log. # Component masses are independent for this case. @@ -408,29 +431,31 @@ def inj_mass_pdf(key, mass1, mass2, lomass, himass, lomass_2 = 0, himass_2 = 0): # lomass: lower component mass # himass: higher component mass - bound = np.sign((himass - mass1)*(mass1 - lomass)) - bound += np.sign((himass_2 - mass2)*(mass2 - lomass_2)) + bound = np.sign((himass - mass1) * (mass1 - lomass)) + bound += np.sign((himass_2 - mass2) * (mass2 - lomass_2)) idx = np.where(bound != 2) pdf = 1 / (log(himass) - log(lomass)) / (log(himass_2) - log(lomass_2)) - pdf /= (mass1 * mass2) + pdf /= mass1 * mass2 pdf[idx] = 0 return pdf + def inj_spin_pdf(key, high_spin, spinz): - ''' Estimate the probability density of the - injections for the spin distribution. - - Parameters - ---------- - key: string - Injections strategy - high_spin: float - Maximum spin used in the strategy - spinz: array - Spin of the injections (for one component) - ''' + """ + Estimate the probability density of the + injections for the spin distribution. + + Parameters + ---------- + key: string + Injections strategy + high_spin: float + Maximum spin used in the strategy + spinz: array + Spin of the injections (for one component) + """ # If the data comes from disable_spin simulation if spinz[0] == 0: return np.ones_like(spinz) @@ -440,69 +465,73 @@ def inj_spin_pdf(key, high_spin, spinz): bound = np.sign(np.absolute(high_spin) - np.absolute(spinz)) bound += np.sign(1 - np.absolute(spinz)) - if key == 'precessing': + if key == "precessing": # Returns the PDF of spins when total spin is # isotropically distributed. Both the component # masses have the same distribution for this case. - pdf = (np.log(high_spin - np.log(abs(spinz)))/high_spin/2) + pdf = np.log(high_spin - np.log(abs(spinz))) / high_spin / 2 idx = np.where(bound != 2) pdf[idx] = 0 return pdf - if key == 'aligned': + if key == "aligned": # Returns the PDF of mass when spins are aligned and uniformly # distributed. Component spins are independent for this case. - pdf = (np.ones_like(spinz) / 2 / high_spin) + pdf = np.ones_like(spinz) / 2 / high_spin idx = np.where(bound != 2) pdf[idx] = 0 return pdf - if key == 'disable_spin': + if key == "disable_spin": # Returns unit array pdf = np.ones_like(spinz) return pdf -def inj_distance_pdf(key, distance, low_dist, high_dist, mchirp = 1): - ''' Estimate the probability density of the - injections for the distance distribution. - - Parameters - ---------- - key: string - Injections strategy - distance: array - Array of distances - low_dist: float - Lower value of distance used in the injection strategy - high_dist: float - Higher value of distance used in the injection strategy - ''' +def inj_distance_pdf(key, distance, low_dist, high_dist, mchirp=1): + """ + Estimate the probability density of the + injections for the distance distribution. + + Parameters + ---------- + key: string + Injections strategy + distance: array + Array of distances + low_dist: float + Lower value of distance used in the injection strategy + high_dist: float + Higher value of distance used in the injection strategy + + """ distance = np.array(distance) - if key == 'uniform': + if key == "uniform": # Returns the PDF at a distance when # distance is uniformly distributed. - pdf = np.ones_like(distance)/(high_dist - low_dist) - bound = np.sign((high_dist - distance)*(distance - low_dist)) + pdf = np.ones_like(distance) / (high_dist - low_dist) + bound = np.sign((high_dist - distance) * (distance - low_dist)) idx = np.where(bound != 1) pdf[idx] = 0 return pdf - if key == 'dchirp': + if key == "dchirp": # Returns the PDF at a distance when distance is uniformly # distributed but scaled by the chirp mass - weight = (mchirp/_mch_BNS)**(5./6) + weight = (mchirp / _mch_BNS) ** (5.0 / 6) pdf = np.ones_like(distance) / weight / (high_dist - low_dist) - bound = np.sign((weight*high_dist - distance)*(distance - weight*low_dist)) + bound = np.sign( + (weight * high_dist - distance) * (distance - weight * low_dist) + ) idx = np.where(bound != 1) pdf[idx] = 0 return pdf diff --git a/pycbc/psd/__init__.py b/pycbc/psd/__init__.py index ce452d8cc10..978fba00277 100644 --- a/pycbc/psd/__init__.py +++ b/pycbc/psd/__init__.py @@ -14,22 +14,40 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import copy + import igwn_segments as segments -from pycbc.psd.read import * + from pycbc.psd.analytical import * from pycbc.psd.analytical_space import * from pycbc.psd.estimate import * +from pycbc.psd.read import * from pycbc.psd.variation import * -from pycbc.types import float32,float64 -from pycbc.types import MultiDetOptionAppendAction, MultiDetOptionAction -from pycbc.types import DictOptionAction, MultiDetDictOptionAction -from pycbc.types import copy_opts_for_single_ifo -from pycbc.types import required_opts, required_opts_multi_ifo -from pycbc.types import ensure_one_opt, ensure_one_opt_multi_ifo - -def from_cli(opt, length, delta_f, low_frequency_cutoff, - strain=None, dyn_range_factor=1, precision=None): - """Parses the CLI options related to the noise PSD and returns a +from pycbc.types import ( + DictOptionAction, + MultiDetDictOptionAction, + MultiDetOptionAction, + MultiDetOptionAppendAction, + copy_opts_for_single_ifo, + ensure_one_opt, + ensure_one_opt_multi_ifo, + float32, + float64, + required_opts, + required_opts_multi_ifo, +) + + +def from_cli( + opt, + length, + delta_f, + low_frequency_cutoff, + strain=None, + dyn_range_factor=1, + precision=None, +): + """ + Parses the CLI options related to the noise PSD and returns a FrequencySeries with the corresponding PSD. If necessary, the PSD is linearly interpolated to achieve the resolution specified in the CLI. @@ -66,62 +84,72 @@ def from_cli(opt, length, delta_f, low_frequency_cutoff, ------- psd : FrequencySeries The frequency series containing the PSD. + """ if opt.psd_low_frequency_cutoff is not None: f_low = opt.psd_low_frequency_cutoff else: f_low = low_frequency_cutoff - sample_rate = (length -1) * 2 * delta_f + sample_rate = (length - 1) * 2 * delta_f try: psd_estimation = opt.psd_estimation is not None except AttributeError: psd_estimation = False - exclusive_opts = [opt.psd_model, opt.psd_file, opt.asd_file, - psd_estimation] + exclusive_opts = [opt.psd_model, opt.psd_file, opt.asd_file, psd_estimation] if sum(map(bool, exclusive_opts)) != 1: err_msg = "You must specify exactly one of '--psd-file', " err_msg += "'--psd-model', '--asd-file', '--psd-estimation'" raise ValueError(err_msg) - if (opt.psd_model or opt.psd_file or opt.asd_file): + if opt.psd_model or opt.psd_file or opt.asd_file: # PSD from lalsimulation or file if opt.psd_model: - psd = from_string(opt.psd_model, length, delta_f, f_low, - **opt.psd_extra_args) + psd = from_string( + opt.psd_model, length, delta_f, f_low, **opt.psd_extra_args + ) elif opt.psd_file or opt.asd_file: if opt.asd_file: psd_file_name = opt.asd_file else: psd_file_name = opt.psd_file - if psd_file_name.endswith(('.dat', '.txt')): + if psd_file_name.endswith((".dat", ".txt")): is_asd_file = bool(opt.asd_file) - psd = from_txt(psd_file_name, length, - delta_f, f_low, is_asd_file=is_asd_file) + psd = from_txt( + psd_file_name, length, delta_f, f_low, is_asd_file=is_asd_file + ) elif opt.asd_file: raise ValueError( "ASD files must be in ASCII format (extension .dat or " f".txt). Got {psd_file_name} instead" ) - elif psd_file_name.endswith(('.xml', '.xml.gz')): - psd = from_xml(psd_file_name, length, delta_f, f_low, - ifo_string=opt.psd_file_xml_ifo_string, - root_name=opt.psd_file_xml_root_name) + elif psd_file_name.endswith((".xml", ".xml.gz")): + psd = from_xml( + psd_file_name, + length, + delta_f, + f_low, + ifo_string=opt.psd_file_xml_ifo_string, + root_name=opt.psd_file_xml_root_name, + ) # Set values < flow to the value at flow (if flow > 0) kmin = int(f_low / psd.delta_f) if kmin > 0: psd[0:kmin] = psd[kmin] - psd *= dyn_range_factor ** 2 + psd *= dyn_range_factor**2 elif psd_estimation: # estimate PSD from data - psd = welch(strain, avg_method=opt.psd_estimation, - seg_len=int(opt.psd_segment_length * sample_rate + 0.5), - seg_stride=int(opt.psd_segment_stride * sample_rate + 0.5), - num_segments=opt.psd_num_segments, - require_exact_data_fit=False) + psd = welch( + strain, + avg_method=opt.psd_estimation, + seg_len=int(opt.psd_segment_length * sample_rate + 0.5), + seg_stride=int(opt.psd_segment_stride * sample_rate + 0.5), + num_segments=opt.psd_num_segments, + require_exact_data_fit=False, + ) if delta_f != psd.delta_f: psd = interpolate(psd, delta_f, length) @@ -134,44 +162,52 @@ def from_cli(opt, length, delta_f, low_frequency_cutoff, try: which_spectrum = opt.invpsd_trunc_which_spectrum except AttributeError: - which_spectrum = 'invasd' + which_spectrum = "invasd" try: fill_value = opt.invpsd_trunc_low_freq_fill_value except AttributeError: - fill_value = 0. - psd = inverse_spectrum_truncation(psd, + fill_value = 0.0 + psd = inverse_spectrum_truncation( + psd, int(opt.psd_inverse_length * sample_rate), which_spectrum=which_spectrum, low_frequency_cutoff=f_low, low_frequency_fill_value=fill_value, - trunc_method=opt.invpsd_trunc_method) + trunc_method=opt.invpsd_trunc_method, + ) - if hasattr(opt, 'psd_output') and opt.psd_output: - (psd.astype(float64) / (dyn_range_factor ** 2)).save(opt.psd_output) + if hasattr(opt, "psd_output") and opt.psd_output: + (psd.astype(float64) / (dyn_range_factor**2)).save(opt.psd_output) if precision is None: return psd - elif precision == 'single': + if precision == "single": return psd.astype(float32) - elif precision == 'double': + if precision == "double": return psd.astype(float64) raise ValueError( "If provided, the precision kwarg must be either 'single' or " f"'double'. You provided {precision}" ) -def from_cli_single_ifo(opt, length, delta_f, low_frequency_cutoff, ifo, - **kwargs): + +def from_cli_single_ifo(opt, length, delta_f, low_frequency_cutoff, ifo, **kwargs): """ Get the PSD for a single ifo when using the multi-detector CLI """ single_det_opt = copy_opts_for_single_ifo(opt, ifo) - return from_cli(single_det_opt, length, delta_f, low_frequency_cutoff, - **kwargs) - -def from_cli_multi_ifos(opt, length_dict, delta_f_dict, - low_frequency_cutoff_dict, ifos, strain_dict=None, - **kwargs): + return from_cli(single_det_opt, length, delta_f, low_frequency_cutoff, **kwargs) + + +def from_cli_multi_ifos( + opt, + length_dict, + delta_f_dict, + low_frequency_cutoff_dict, + ifos, + strain_dict=None, + **kwargs, +): """ Get the PSD for all ifos when using the multi-detector CLI """ @@ -181,11 +217,18 @@ def from_cli_multi_ifos(opt, length_dict, delta_f_dict, strain = strain_dict[ifo] else: strain = None - psd[ifo] = from_cli_single_ifo(opt, length_dict[ifo], delta_f_dict[ifo], - low_frequency_cutoff_dict[ifo], ifo, - strain=strain, **kwargs) + psd[ifo] = from_cli_single_ifo( + opt, + length_dict[ifo], + delta_f_dict[ifo], + low_frequency_cutoff_dict[ifo], + ifo, + strain=strain, + **kwargs, + ) return psd + def insert_psd_option_group(parser, output=True, include_data_options=True): """ Adds the options used to call the pycbc.psd.from_cli function to an @@ -193,120 +236,181 @@ def insert_psd_option_group(parser, output=True, include_data_options=True): want to use these options in your code. Parameters - ----------- + ---------- parser : object OptionParser instance. + """ psd_options = parser.add_argument_group( - "Options to select the method of PSD generation", - "The options --psd-model, --psd-file, --asd-file, " - "and --psd-estimation are mutually exclusive.") - psd_options.add_argument("--psd-model", - help="Get PSD from given analytical model. ", - choices=get_psd_model_list()) - psd_options.add_argument("--psd-extra-args", - nargs='+', action=DictOptionAction, - metavar='PARAM:VALUE', default={}, type=float, - help="(optional) Extra arguments passed to " - "the PSD models.") - psd_options.add_argument("--psd-file", - help="Get PSD using given PSD ASCII file") - psd_options.add_argument("--asd-file", - help="Get PSD using given ASD ASCII file") - psd_options.add_argument("--psd-inverse-length", type=float, - help="(Optional) The maximum length of the " - "impulse response of the overwhitening " - "filter (s)") - psd_options.add_argument("--psd-low-frequency-cutoff", type=float, - help="(Optional) The low frequency cutoff for the " - "PSD. If not specified, the low frequency " - "cutoff of the matched filter/likelihood " - "integral will be used.") + "Options to select the method of PSD generation", + "The options --psd-model, --psd-file, --asd-file, " + "and --psd-estimation are mutually exclusive.", + ) + psd_options.add_argument( + "--psd-model", + help="Get PSD from given analytical model. ", + choices=get_psd_model_list(), + ) + psd_options.add_argument( + "--psd-extra-args", + nargs="+", + action=DictOptionAction, + metavar="PARAM:VALUE", + default={}, + type=float, + help="(optional) Extra arguments passed to the PSD models.", + ) + psd_options.add_argument("--psd-file", help="Get PSD using given PSD ASCII file") + psd_options.add_argument("--asd-file", help="Get PSD using given ASD ASCII file") + psd_options.add_argument( + "--psd-inverse-length", + type=float, + help="(Optional) The maximum length of the " + "impulse response of the overwhitening " + "filter (s)", + ) + psd_options.add_argument( + "--psd-low-frequency-cutoff", + type=float, + help="(Optional) The low frequency cutoff for the " + "PSD. If not specified, the low frequency " + "cutoff of the matched filter/likelihood " + "integral will be used.", + ) # Truncation options if specifying psd-inverse-length - psd_options.add_argument("--invpsd-trunc-method", choices=["hann"], - help="(Optional) What truncation method to use " - "when applying psd-inverse-length. If not " - "provided, a hard truncation will be used.") - psd_options.add_argument("--invpsd-trunc-which-spectrum", type=str, - default='invasd', choices=["invasd", "invpsd"], - help="(Optional) Which spectrum to use to perform " - "truncation when applying psd-inverse-length. " - "The default ('invasd') truncates the inverse " - "ASD, while 'invpsd' truncates the inverse " - "PSD.") - psd_options.add_argument("--invpsd-trunc-low-freq-fill-value", default=0., - help="(Optional) Value to set the inverse PSD to " - "for frequencies below the low frequency " - "cutoff when applying psd-inverse-length. " - "Default 0; accepts any float. Also accepts " - "'fmin'; this sets values below the cutoff " - "to the inverse PSD value specified by " - "`psd-low-frequency-cutoff` (or the matched " - "filter/likelihood integral otherwise).") + psd_options.add_argument( + "--invpsd-trunc-method", + choices=["hann"], + help="(Optional) What truncation method to use " + "when applying psd-inverse-length. If not " + "provided, a hard truncation will be used.", + ) + psd_options.add_argument( + "--invpsd-trunc-which-spectrum", + type=str, + default="invasd", + choices=["invasd", "invpsd"], + help="(Optional) Which spectrum to use to perform " + "truncation when applying psd-inverse-length. " + "The default ('invasd') truncates the inverse " + "ASD, while 'invpsd' truncates the inverse " + "PSD.", + ) + psd_options.add_argument( + "--invpsd-trunc-low-freq-fill-value", + default=0.0, + help="(Optional) Value to set the inverse PSD to " + "for frequencies below the low frequency " + "cutoff when applying psd-inverse-length. " + "Default 0; accepts any float. Also accepts " + "'fmin'; this sets values below the cutoff " + "to the inverse PSD value specified by " + "`psd-low-frequency-cutoff` (or the matched " + "filter/likelihood integral otherwise).", + ) # Options specific to XML PSD files - psd_options.add_argument("--psd-file-xml-ifo-string", - help="If using an XML PSD file, use the PSD in " - "the file's PSD dictionary with this " - "ifo string. If not given and only one " - "PSD present in the file return that, if " - "not given and multiple (or zero) PSDs " - "present an exception will be raised.") - psd_options.add_argument("--psd-file-xml-root-name", default='psd', - help="If given use this as the root name for " - "the PSD XML file. If this means nothing " - "to you, then it is probably safe to " - "ignore this option.") + psd_options.add_argument( + "--psd-file-xml-ifo-string", + help="If using an XML PSD file, use the PSD in " + "the file's PSD dictionary with this " + "ifo string. If not given and only one " + "PSD present in the file return that, if " + "not given and multiple (or zero) PSDs " + "present an exception will be raised.", + ) + psd_options.add_argument( + "--psd-file-xml-root-name", + default="psd", + help="If given use this as the root name for " + "the PSD XML file. If this means nothing " + "to you, then it is probably safe to " + "ignore this option.", + ) # Options for PSD variation - psd_options.add_argument("--psdvar-segment", type=float, - metavar="SECONDS", help="Length of segment " - "for mean square calculation of PSD variation.") - psd_options.add_argument("--psdvar-short-segment", type=float, - metavar="SECONDS", help="Length of short segment " - "for outliers removal in PSD variability " - "calculation.") - psd_options.add_argument("--psdvar-long-segment", type=float, - metavar="SECONDS", help="Length of long segment " - "when calculating the PSD variability.") - psd_options.add_argument("--psdvar-psd-duration", type=float, - metavar="SECONDS", help="Duration of short " - "segments for PSD estimation.") - psd_options.add_argument("--psdvar-psd-stride", type=float, - metavar="SECONDS", help="Separation between PSD " - "estimation segments.") - psd_options.add_argument("--psdvar-low-freq", type=float, metavar="HERTZ", - help="Minimum frequency to consider in strain " - "bandpass.") - psd_options.add_argument("--psdvar-high-freq", type=float, metavar="HERTZ", - help="Maximum frequency to consider in strain " - "bandpass.") - - if include_data_options : - psd_options.add_argument("--psd-estimation", - help="Measure PSD from the data, using " - "given average method.", - choices=["mean", "median", "median-mean"]) - psd_options.add_argument("--psd-segment-length", type=float, - help="(Required for --psd-estimation) The " - "segment length for PSD estimation (s)") - psd_options.add_argument("--psd-segment-stride", type=float, - help="(Required for --psd-estimation) " - "The separation between consecutive " - "segments (s)") - psd_options.add_argument("--psd-num-segments", type=int, - help="(Optional, used only with " - "--psd-estimation). If given, PSDs will " - "be estimated using only this number of " - "segments. If more data is given than " - "needed to make this number of segments " - "then excess data will not be used in " - "the PSD estimate. If not enough data " - "is given, the code will fail.") + psd_options.add_argument( + "--psdvar-segment", + type=float, + metavar="SECONDS", + help="Length of segment for mean square calculation of PSD variation.", + ) + psd_options.add_argument( + "--psdvar-short-segment", + type=float, + metavar="SECONDS", + help="Length of short segment " + "for outliers removal in PSD variability " + "calculation.", + ) + psd_options.add_argument( + "--psdvar-long-segment", + type=float, + metavar="SECONDS", + help="Length of long segment when calculating the PSD variability.", + ) + psd_options.add_argument( + "--psdvar-psd-duration", + type=float, + metavar="SECONDS", + help="Duration of short segments for PSD estimation.", + ) + psd_options.add_argument( + "--psdvar-psd-stride", + type=float, + metavar="SECONDS", + help="Separation between PSD estimation segments.", + ) + psd_options.add_argument( + "--psdvar-low-freq", + type=float, + metavar="HERTZ", + help="Minimum frequency to consider in strain bandpass.", + ) + psd_options.add_argument( + "--psdvar-high-freq", + type=float, + metavar="HERTZ", + help="Maximum frequency to consider in strain bandpass.", + ) + + if include_data_options: + psd_options.add_argument( + "--psd-estimation", + help="Measure PSD from the data, using given average method.", + choices=["mean", "median", "median-mean"], + ) + psd_options.add_argument( + "--psd-segment-length", + type=float, + help="(Required for --psd-estimation) The " + "segment length for PSD estimation (s)", + ) + psd_options.add_argument( + "--psd-segment-stride", + type=float, + help="(Required for --psd-estimation) " + "The separation between consecutive " + "segments (s)", + ) + psd_options.add_argument( + "--psd-num-segments", + type=int, + help="(Optional, used only with " + "--psd-estimation). If given, PSDs will " + "be estimated using only this number of " + "segments. If more data is given than " + "needed to make this number of segments " + "then excess data will not be used in " + "the PSD estimate. If not enough data " + "is given, the code will fail.", + ) if output: - psd_options.add_argument("--psd-output", - help="(Optional) Write PSD to specified file") + psd_options.add_argument( + "--psd-output", help="(Optional) Write PSD to specified file" + ) return psd_options + def insert_psd_option_group_multi_ifo(parser): """ Adds the options used to call the pycbc.psd.from_cli function to an @@ -314,119 +418,207 @@ def insert_psd_option_group_multi_ifo(parser): want to use these options in your code. Parameters - ----------- + ---------- parser : object OptionParser instance. + """ psd_options = parser.add_argument_group( - "Options to select the method of PSD generation", - "The options --psd-model, --psd-file, --asd-file, " - "and --psd-estimation are mutually exclusive.") - psd_options.add_argument("--psd-model", nargs="+", - action=MultiDetOptionAction, metavar='IFO:MODEL', - help="Get PSD from given analytical model. " - "Choose from %s" %(', '.join(get_psd_model_list()),)) - psd_options.add_argument("--psd-extra-args", - nargs='+', action=MultiDetDictOptionAction, - metavar='DETECTOR:PARAM:VALUE', default={}, - type=float, help="(optional) Extra arguments " - "passed to the PSD models.") - psd_options.add_argument("--psd-file", nargs="+", - action=MultiDetOptionAction, metavar='IFO:FILE', - help="Get PSD using given PSD ASCII file") - psd_options.add_argument("--asd-file", nargs="+", - action=MultiDetOptionAction, metavar='IFO:FILE', - help="Get PSD using given ASD ASCII file") - psd_options.add_argument("--psd-estimation", nargs="+", - action=MultiDetOptionAction, metavar='IFO:FILE', - help="Measure PSD from the data, using given " - "average method. Choose from " - "mean, median or median-mean.") - psd_options.add_argument("--psd-low-frequency-cutoff", nargs="+", type=float, - action=MultiDetOptionAction, metavar='IFO:FREQ', - help="(Optional) The low frequency cutoff for the " - "PSD. If not specified, the low frequency " - "cutoff of the matched filter/likelihood " - "integral will be used.") - psd_options.add_argument("--psd-segment-length", type=float, nargs="+", - action=MultiDetOptionAction, metavar='IFO:LENGTH', - help="(Required for --psd-estimation) The segment " - "length for PSD estimation (s)") - psd_options.add_argument("--psd-segment-stride", type=float, nargs="+", - action=MultiDetOptionAction, metavar='IFO:STRIDE', - help="(Required for --psd-estimation) The separation" - " between consecutive segments (s)") - psd_options.add_argument("--psd-num-segments", type=int, nargs="+", - action=MultiDetOptionAction, metavar='IFO:NUM', - help="(Optional, used only with --psd-estimation). " - "If given PSDs will be estimated using only " - "this number of segments. If more data is " - "given than needed to make this number of " - "segments than excess data will not be used in " - "the PSD estimate. If not enough data is given " - "the code will fail.") - psd_options.add_argument("--psd-inverse-length", type=float, nargs="+", - action=MultiDetOptionAction, metavar='IFO:LENGTH', - help="(Optional) The maximum length of the impulse" - " response of the overwhitening filter (s)") - psd_options.add_argument("--invpsd-trunc-method", choices=["hann"], - help="(Optional) What truncation method to use " - "when applying psd-inverse-length. If not " - "provided, a hard truncation will be used.") - psd_options.add_argument("--invpsd-trunc-which-spectrum", type=str, - default='invasd', choices=["invasd", "invpsd"], - help="(Optional) Which spectrum to use to perform " - "truncation when applying psd-inverse-length. " - "The default ('invasd') truncates the inverse " - "ASD, while 'invpsd' truncates the inverse " - "PSD.") - psd_options.add_argument("--invpsd-trunc-low-freq-fill-value", default=0., - action=MultiDetOptionAction, metavar='IFO:VALUE', - nargs="+", - help="(Optional) Value to set the inverse PSD to " - "for frequencies below the low frequency " - "cutoff when applying psd-inverse-length. " - "Default 0; accepts any float. Also accepts " - "'fmin'; this sets values below the cutoff " - "to the inverse PSD value specified by " - "`psd-low-frequency-cutoff` (or the matched " - "filter/likelihood integral otherwise).") - psd_options.add_argument("--psd-output", nargs="+", - action=MultiDetOptionAction, metavar='IFO:FILE', - help="(Optional) Write PSD to specified file") + "Options to select the method of PSD generation", + "The options --psd-model, --psd-file, --asd-file, " + "and --psd-estimation are mutually exclusive.", + ) + psd_options.add_argument( + "--psd-model", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:MODEL", + help="Get PSD from given analytical model. " + "Choose from %s" % (", ".join(get_psd_model_list()),), + ) + psd_options.add_argument( + "--psd-extra-args", + nargs="+", + action=MultiDetDictOptionAction, + metavar="DETECTOR:PARAM:VALUE", + default={}, + type=float, + help="(optional) Extra arguments passed to the PSD models.", + ) + psd_options.add_argument( + "--psd-file", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="Get PSD using given PSD ASCII file", + ) + psd_options.add_argument( + "--asd-file", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="Get PSD using given ASD ASCII file", + ) + psd_options.add_argument( + "--psd-estimation", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="Measure PSD from the data, using given " + "average method. Choose from " + "mean, median or median-mean.", + ) + psd_options.add_argument( + "--psd-low-frequency-cutoff", + nargs="+", + type=float, + action=MultiDetOptionAction, + metavar="IFO:FREQ", + help="(Optional) The low frequency cutoff for the " + "PSD. If not specified, the low frequency " + "cutoff of the matched filter/likelihood " + "integral will be used.", + ) + psd_options.add_argument( + "--psd-segment-length", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:LENGTH", + help="(Required for --psd-estimation) The segment " + "length for PSD estimation (s)", + ) + psd_options.add_argument( + "--psd-segment-stride", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:STRIDE", + help="(Required for --psd-estimation) The separation" + " between consecutive segments (s)", + ) + psd_options.add_argument( + "--psd-num-segments", + type=int, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:NUM", + help="(Optional, used only with --psd-estimation). " + "If given PSDs will be estimated using only " + "this number of segments. If more data is " + "given than needed to make this number of " + "segments than excess data will not be used in " + "the PSD estimate. If not enough data is given " + "the code will fail.", + ) + psd_options.add_argument( + "--psd-inverse-length", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:LENGTH", + help="(Optional) The maximum length of the impulse" + " response of the overwhitening filter (s)", + ) + psd_options.add_argument( + "--invpsd-trunc-method", + choices=["hann"], + help="(Optional) What truncation method to use " + "when applying psd-inverse-length. If not " + "provided, a hard truncation will be used.", + ) + psd_options.add_argument( + "--invpsd-trunc-which-spectrum", + type=str, + default="invasd", + choices=["invasd", "invpsd"], + help="(Optional) Which spectrum to use to perform " + "truncation when applying psd-inverse-length. " + "The default ('invasd') truncates the inverse " + "ASD, while 'invpsd' truncates the inverse " + "PSD.", + ) + psd_options.add_argument( + "--invpsd-trunc-low-freq-fill-value", + default=0.0, + action=MultiDetOptionAction, + metavar="IFO:VALUE", + nargs="+", + help="(Optional) Value to set the inverse PSD to " + "for frequencies below the low frequency " + "cutoff when applying psd-inverse-length. " + "Default 0; accepts any float. Also accepts " + "'fmin'; this sets values below the cutoff " + "to the inverse PSD value specified by " + "`psd-low-frequency-cutoff` (or the matched " + "filter/likelihood integral otherwise).", + ) + psd_options.add_argument( + "--psd-output", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="(Optional) Write PSD to specified file", + ) # Options for PSD variation - psd_options.add_argument("--psdvar-segment", type=float, - metavar="SECONDS", help="Length of segment " - "when calculating the PSD variability.") - psd_options.add_argument("--psdvar-short-segment", type=float, - metavar="SECONDS", help="Length of short segment " - "for outliers removal in PSD variability " - "calculation.") - psd_options.add_argument("--psdvar-long-segment", type=float, - metavar="SECONDS", help="Length of long segment " - "when calculating the PSD variability.") - psd_options.add_argument("--psdvar-psd-duration", type=float, - metavar="SECONDS", help="Duration of short " - "segments for PSD estimation.") - psd_options.add_argument("--psdvar-psd-stride", type=float, - metavar="SECONDS", help="Separation between PSD " - "estimation segments.") - psd_options.add_argument("--psdvar-low-freq", type=float, metavar="HERTZ", - help="Minimum frequency to consider in strain " - "bandpass.") - psd_options.add_argument("--psdvar-high-freq", type=float, metavar="HERTZ", - help="Maximum frequency to consider in strain " - "bandpass.") + psd_options.add_argument( + "--psdvar-segment", + type=float, + metavar="SECONDS", + help="Length of segment when calculating the PSD variability.", + ) + psd_options.add_argument( + "--psdvar-short-segment", + type=float, + metavar="SECONDS", + help="Length of short segment " + "for outliers removal in PSD variability " + "calculation.", + ) + psd_options.add_argument( + "--psdvar-long-segment", + type=float, + metavar="SECONDS", + help="Length of long segment when calculating the PSD variability.", + ) + psd_options.add_argument( + "--psdvar-psd-duration", + type=float, + metavar="SECONDS", + help="Duration of short segments for PSD estimation.", + ) + psd_options.add_argument( + "--psdvar-psd-stride", + type=float, + metavar="SECONDS", + help="Separation between PSD estimation segments.", + ) + psd_options.add_argument( + "--psdvar-low-freq", + type=float, + metavar="HERTZ", + help="Minimum frequency to consider in strain bandpass.", + ) + psd_options.add_argument( + "--psdvar-high-freq", + type=float, + metavar="HERTZ", + help="Maximum frequency to consider in strain bandpass.", + ) return psd_options + ensure_one_opt_groups = [] -ensure_one_opt_groups.append(['--psd-file', '--psd-model', - '--psd-estimation', '--asd-file']) +ensure_one_opt_groups.append( + ["--psd-file", "--psd-model", "--psd-estimation", "--asd-file"] +) + def verify_psd_options(opt, parser): - """Parses the CLI options and verifies that they are consistent and + """ + Parses the CLI options and verifies that they are consistent and reasonable. Parameters @@ -437,6 +629,7 @@ def verify_psd_options(opt, parser): psd_segment_length, psd_segment_stride, psd_inverse_length, psd_output). parser : object OptionParser instance. + """ try: psd_estimation = opt.psd_estimation is not None @@ -447,12 +640,17 @@ def verify_psd_options(opt, parser): ensure_one_opt(opt, parser, opt_group) if psd_estimation: - required_opts(opt, parser, - ['--psd-segment-stride', '--psd-segment-length'], - required_by = "--psd-estimation") + required_opts( + opt, + parser, + ["--psd-segment-stride", "--psd-segment-length"], + required_by="--psd-estimation", + ) + def verify_psd_options_multi_ifo(opt, parser, ifos): - """Parses the CLI options and verifies that they are consistent and + """ + Parses the CLI options and verifies that they are consistent and reasonable. Parameters @@ -463,24 +661,32 @@ def verify_psd_options_multi_ifo(opt, parser, ifos): psd_segment_length, psd_segment_stride, psd_inverse_length, psd_output). parser : object OptionParser instance. + """ for ifo in ifos: for opt_group in ensure_one_opt_groups: ensure_one_opt_multi_ifo(opt, parser, ifo, opt_group) if opt.psd_estimation[ifo]: - required_opts_multi_ifo(opt, parser, ifo, - ['--psd-segment-stride', '--psd-segment-length'], - required_by = "--psd-estimation") - -def generate_overlapping_psds(opt, gwstrain, flen, delta_f, flow, - dyn_range_factor=1., precision=None): - """Generate a set of overlapping PSDs to cover a stretch of data. This + required_opts_multi_ifo( + opt, + parser, + ifo, + ["--psd-segment-stride", "--psd-segment-length"], + required_by="--psd-estimation", + ) + + +def generate_overlapping_psds( + opt, gwstrain, flen, delta_f, flow, dyn_range_factor=1.0, precision=None +): + """ + Generate a set of overlapping PSDs to cover a stretch of data. This allows one to analyse a long stretch of data with PSD measurements that change with time. Parameters - ----------- + ---------- opt : object Result of parsing the CLI with OptionParser, or any object with the required attributes (psd_model, psd_file, asd_file, psd_estimation, @@ -505,17 +711,25 @@ def generate_overlapping_psds(opt, gwstrain, flen, delta_f, flow, not already in that precision. Returns - -------- + ------- psd_and_times : list of (start, end, PSD) tuples This is a list of tuples containing one entry for each PSD. The first and second entries (start, end) in each tuple represent the index range of the gwstrain data that was used to estimate that PSD. The third entry (psd) contains the PSD estimate between that interval. + """ if not opt.psd_estimation: - psd = from_cli(opt, flen, delta_f, flow, strain=gwstrain, - dyn_range_factor=dyn_range_factor, precision=precision) - psds_and_times = [ (0, len(gwstrain), psd) ] + psd = from_cli( + opt, + flen, + delta_f, + flow, + strain=gwstrain, + dyn_range_factor=dyn_range_factor, + precision=precision, + ) + psds_and_times = [(0, len(gwstrain), psd)] return psds_and_times # Figure out the data length used for PSD generation @@ -538,14 +752,14 @@ def generate_overlapping_psds(opt, gwstrain, flen, delta_f, flow, if input_data_len < psd_data_len: err_msg = "Input data length must be longer than data length needed " err_msg += "to estimate a PSD. You specified that a PSD should be " - err_msg += "estimated with %d seconds. " %(psd_data_len) - err_msg += "Input data length is %d seconds. " %(input_data_len) + err_msg += "estimated with %d seconds. " % (psd_data_len) + err_msg += "Input data length is %d seconds. " % (input_data_len) raise ValueError(err_msg) - elif input_data_len == psd_data_len: + if input_data_len == psd_data_len: num_psd_measurements = 1 psd_stride = 0 else: - num_psd_measurements = int(2 * (input_data_len-1) / psd_data_len) + num_psd_measurements = int(2 * (input_data_len - 1) / psd_data_len) psd_stride = int((input_data_len - psd_data_len) / num_psd_measurements) for idx in range(num_psd_measurements): @@ -556,18 +770,35 @@ def generate_overlapping_psds(opt, gwstrain, flen, delta_f, flow, start_idx = psd_stride * idx end_idx = psd_data_len + psd_stride * idx strain_part = gwstrain[start_idx:end_idx] - psd = from_cli(opt, flen, delta_f, flow, strain=strain_part, - dyn_range_factor=dyn_range_factor, precision=precision) - psds_and_times.append( (start_idx, end_idx, psd) ) + psd = from_cli( + opt, + flen, + delta_f, + flow, + strain=strain_part, + dyn_range_factor=dyn_range_factor, + precision=precision, + ) + psds_and_times.append((start_idx, end_idx, psd)) return psds_and_times -def associate_psds_to_segments(opt, fd_segments, gwstrain, flen, delta_f, flow, - dyn_range_factor=1., precision=None): - """Generate a set of overlapping PSDs covering the data in GWstrain. + +def associate_psds_to_segments( + opt, + fd_segments, + gwstrain, + flen, + delta_f, + flow, + dyn_range_factor=1.0, + precision=None, +): + """ + Generate a set of overlapping PSDs covering the data in GWstrain. Then associate these PSDs with the appropriate segment in strain_segments. Parameters - ----------- + ---------- opt : object Result of parsing the CLI with OptionParser, or any object with the required attributes (psd_model, psd_file, asd_file, psd_estimation, @@ -593,16 +824,24 @@ def associate_psds_to_segments(opt, fd_segments, gwstrain, flen, delta_f, flow, If 'single' the PSD will be converted to float32, if not already in that precision. If 'double' the PSD will be converted to float64, if not already in that precision. + """ - psds_and_times = generate_overlapping_psds(opt, gwstrain, flen, delta_f, - flow, dyn_range_factor=dyn_range_factor, - precision=precision) + psds_and_times = generate_overlapping_psds( + opt, + gwstrain, + flen, + delta_f, + flow, + dyn_range_factor=dyn_range_factor, + precision=precision, + ) for fd_segment in fd_segments: best_psd = None psd_overlap = 0 - inp_seg = segments.segment(fd_segment.seg_slice.start, - fd_segment.seg_slice.stop) + inp_seg = segments.segment( + fd_segment.seg_slice.start, fd_segment.seg_slice.stop + ) for start_idx, end_idx, psd in psds_and_times: psd_seg = segments.segment(start_idx, end_idx) if psd_seg.intersects(inp_seg): @@ -614,21 +853,46 @@ def associate_psds_to_segments(opt, fd_segments, gwstrain, flen, delta_f, flow, raise ValueError("No PSDs found intersecting segment!") fd_segment.psd = best_psd -def associate_psds_to_single_ifo_segments(opt, fd_segments, gwstrain, flen, - delta_f, flow, ifo, - dyn_range_factor=1., precision=None): + +def associate_psds_to_single_ifo_segments( + opt, + fd_segments, + gwstrain, + flen, + delta_f, + flow, + ifo, + dyn_range_factor=1.0, + precision=None, +): """ Associate PSDs to segments for a single ifo when using the multi-detector CLI """ single_det_opt = copy_opts_for_single_ifo(opt, ifo) - associate_psds_to_segments(single_det_opt, fd_segments, gwstrain, flen, - delta_f, flow, dyn_range_factor=dyn_range_factor, - precision=precision) + associate_psds_to_segments( + single_det_opt, + fd_segments, + gwstrain, + flen, + delta_f, + flow, + dyn_range_factor=dyn_range_factor, + precision=precision, + ) + -def associate_psds_to_multi_ifo_segments(opt, fd_segments, gwstrain, flen, - delta_f, flow, ifos, - dyn_range_factor=1., precision=None): +def associate_psds_to_multi_ifo_segments( + opt, + fd_segments, + gwstrain, + flen, + delta_f, + flow, + ifos, + dyn_range_factor=1.0, + precision=None, +): """ Associate PSDs to segments for all ifos when using the multi-detector CLI """ @@ -643,6 +907,14 @@ def associate_psds_to_multi_ifo_segments(opt, fd_segments, gwstrain, flen, else: segments = None - associate_psds_to_single_ifo_segments(opt, segments, strain, flen, - delta_f, flow, ifo, dyn_range_factor=dyn_range_factor, - precision=precision) + associate_psds_to_single_ifo_segments( + opt, + segments, + strain, + flen, + delta_f, + flow, + ifo, + dyn_range_factor=dyn_range_factor, + precision=precision, + ) diff --git a/pycbc/psd/analytical.py b/pycbc/psd/analytical.py index b216601754a..6b6cf79a420 100644 --- a/pycbc/psd/analytical.py +++ b/pycbc/psd/analytical.py @@ -14,37 +14,52 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Provides reference PSDs from LALSimulation and pycbc.psd.analytical_space. +""" +Provides reference PSDs from LALSimulation and pycbc.psd.analytical_space. More information about how to use these ground-based detectors' PSD can be found in the guide about :ref:`Analytic PSDs from lalsimulation`. For space-borne ones, see `pycbc.psd.analytical_space` module. """ + import numbers -from pycbc.types import FrequencySeries -from pycbc.psd.analytical_space import ( - analytical_psd_lisa_tdi_XYZ, analytical_psd_lisa_tdi_AE, - analytical_psd_lisa_tdi_T, sh_transformed_psd_lisa_tdi_XYZ, - analytical_psd_lisa_tdi_AE_confusion, - analytical_psd_tianqin_tdi_XYZ, analytical_psd_tianqin_tdi_AE, - analytical_psd_tianqin_tdi_T, analytical_psd_tianqin_tdi_AE_confusion, - analytical_psd_taiji_tdi_XYZ, analytical_psd_taiji_tdi_AE, - analytical_psd_taiji_tdi_T, analytical_psd_taiji_tdi_AE_confusion, - ) + import lal import numpy +from pycbc.psd.analytical_space import ( + analytical_psd_lisa_tdi_AE, + analytical_psd_lisa_tdi_AE_confusion, + analytical_psd_lisa_tdi_T, + analytical_psd_lisa_tdi_XYZ, + analytical_psd_taiji_tdi_AE, + analytical_psd_taiji_tdi_AE_confusion, + analytical_psd_taiji_tdi_T, + analytical_psd_taiji_tdi_XYZ, + analytical_psd_tianqin_tdi_AE, + analytical_psd_tianqin_tdi_AE_confusion, + analytical_psd_tianqin_tdi_T, + analytical_psd_tianqin_tdi_XYZ, + sh_transformed_psd_lisa_tdi_XYZ, +) +from pycbc.types import FrequencySeries + # build a list of usable PSD functions from lalsimulation -_name_prefix = 'SimNoisePSD' -_name_suffix = 'Ptr' -_name_blacklist = ('FromFile', 'MirrorTherm', 'Quantum', 'Seismic', 'Shot', 'SuspTherm') +_name_prefix = "SimNoisePSD" +_name_suffix = "Ptr" +_name_blacklist = ("FromFile", "MirrorTherm", "Quantum", "Seismic", "Shot", "SuspTherm") _psd_list = [] try: import lalsimulation + for _name in lalsimulation.__dict__: - if _name != _name_prefix and _name.startswith(_name_prefix) and not _name.endswith(_name_suffix): - _name = _name[len(_name_prefix):] + if ( + _name != _name_prefix + and _name.startswith(_name_prefix) + and not _name.endswith(_name_suffix) + ): + _name = _name[len(_name_prefix) :] if _name not in _name_blacklist: _psd_list.append(_name) except ImportError: @@ -54,42 +69,53 @@ # add functions wrapping lalsimulation PSDs for _name in _psd_list: - exec(""" + exec( + """ def %s(length, delta_f, low_freq_cutoff): \"\"\"Return a FrequencySeries containing the %s PSD from LALSimulation. \"\"\" return from_string("%s", length, delta_f, low_freq_cutoff) -""" % (_name, _name, _name)) +""" + % (_name, _name, _name) + ) + def get_psd_model_list(): - """ Returns a list of available reference PSD functions. + """ + Returns a list of available reference PSD functions. Returns ------- list Returns a list of names of reference PSD functions. + """ return get_lalsim_psd_list() + get_pycbc_psd_list() + def get_lalsim_psd_list(): - """Return a list of available reference PSD functions from LALSimulation. - """ + """Return a list of available reference PSD functions from LALSimulation.""" return _psd_list + def get_pycbc_psd_list(): - """ Return a list of available reference PSD functions coded in PyCBC. + """ + Return a list of available reference PSD functions coded in PyCBC. Returns ------- list Returns a list of names of all reference PSD functions coded in PyCBC. + """ pycbc_analytical_psd_list = pycbc_analytical_psds.keys() pycbc_analytical_psd_list = sorted(pycbc_analytical_psd_list) return pycbc_analytical_psd_list + def from_string(psd_name, length, delta_f, low_freq_cutoff, **kwargs): - """Generate a frequency series containing a LALSimulation or + """ + Generate a frequency series containing a LALSimulation or built-in space-borne detectors' PSD specified by name. Parameters @@ -110,26 +136,24 @@ def from_string(psd_name, length, delta_f, low_freq_cutoff, **kwargs): ------- psd : FrequencySeries The generated frequency series. - """ + """ # check if valid PSD model if psd_name not in get_psd_model_list(): - raise ValueError( - psd_name + ' not found among analytical PSD functions.' - ) + raise ValueError(psd_name + " not found among analytical PSD functions.") # make sure length has the right type for CreateREAL8FrequencySeries if not isinstance(length, numbers.Integral) or length <= 0: - raise TypeError('length must be a positive integer') + raise TypeError("length must be a positive integer") length = int(length) # if PSD model is in LALSimulation if psd_name in get_lalsim_psd_list(): lalseries = lal.CreateREAL8FrequencySeries( - '', lal.LIGOTimeGPS(0), 0, delta_f, lal.DimensionlessUnit, length) + "", lal.LIGOTimeGPS(0), 0, delta_f, lal.DimensionlessUnit, length + ) try: - func = lalsimulation.__dict__[ - _name_prefix + psd_name + _name_suffix] + func = lalsimulation.__dict__[_name_prefix + psd_name + _name_suffix] except KeyError: func = lalsimulation.__dict__[_name_prefix + psd_name] func(lalseries, low_freq_cutoff) @@ -148,8 +172,10 @@ def from_string(psd_name, length, delta_f, low_freq_cutoff, **kwargs): return psd + def flat_unity(length, delta_f, low_freq_cutoff): - """ Returns a FrequencySeries of ones above the low_frequency_cutoff. + """ + Returns a FrequencySeries of ones above the low_frequency_cutoff. Parameters ---------- @@ -164,29 +190,28 @@ def flat_unity(length, delta_f, low_freq_cutoff): ------- FrequencySeries Returns a FrequencySeries containing the unity PSD model. + """ fseries = FrequencySeries(numpy.ones(length), delta_f=delta_f) kmin = int(low_freq_cutoff / fseries.delta_f) fseries.data[:kmin] = 0 return fseries + # dict of analytical PSDs coded in PyCBC pycbc_analytical_psds = { - 'flat_unity' : flat_unity, - - 'analytical_psd_lisa_tdi_XYZ' : analytical_psd_lisa_tdi_XYZ, - 'analytical_psd_lisa_tdi_AE' : analytical_psd_lisa_tdi_AE, - 'analytical_psd_lisa_tdi_T' : analytical_psd_lisa_tdi_T, - 'sh_transformed_psd_lisa_tdi_XYZ' : sh_transformed_psd_lisa_tdi_XYZ, - 'analytical_psd_lisa_tdi_AE_confusion' : analytical_psd_lisa_tdi_AE_confusion, - - 'analytical_psd_tianqin_tdi_XYZ' : analytical_psd_tianqin_tdi_XYZ, - 'analytical_psd_tianqin_tdi_AE' : analytical_psd_tianqin_tdi_AE, - 'analytical_psd_tianqin_tdi_T' : analytical_psd_tianqin_tdi_T, - 'analytical_psd_tianqin_tdi_AE_confusion' : analytical_psd_tianqin_tdi_AE_confusion, - - 'analytical_psd_taiji_tdi_XYZ' : analytical_psd_taiji_tdi_XYZ, - 'analytical_psd_taiji_tdi_AE' : analytical_psd_taiji_tdi_AE, - 'analytical_psd_taiji_tdi_T' : analytical_psd_taiji_tdi_T, - 'analytical_psd_taiji_tdi_AE_confusion' : analytical_psd_taiji_tdi_AE_confusion, + "flat_unity": flat_unity, + "analytical_psd_lisa_tdi_XYZ": analytical_psd_lisa_tdi_XYZ, + "analytical_psd_lisa_tdi_AE": analytical_psd_lisa_tdi_AE, + "analytical_psd_lisa_tdi_T": analytical_psd_lisa_tdi_T, + "sh_transformed_psd_lisa_tdi_XYZ": sh_transformed_psd_lisa_tdi_XYZ, + "analytical_psd_lisa_tdi_AE_confusion": analytical_psd_lisa_tdi_AE_confusion, + "analytical_psd_tianqin_tdi_XYZ": analytical_psd_tianqin_tdi_XYZ, + "analytical_psd_tianqin_tdi_AE": analytical_psd_tianqin_tdi_AE, + "analytical_psd_tianqin_tdi_T": analytical_psd_tianqin_tdi_T, + "analytical_psd_tianqin_tdi_AE_confusion": analytical_psd_tianqin_tdi_AE_confusion, + "analytical_psd_taiji_tdi_XYZ": analytical_psd_taiji_tdi_XYZ, + "analytical_psd_taiji_tdi_AE": analytical_psd_taiji_tdi_AE, + "analytical_psd_taiji_tdi_T": analytical_psd_taiji_tdi_T, + "analytical_psd_taiji_tdi_AE_confusion": analytical_psd_taiji_tdi_AE_confusion, } diff --git a/pycbc/psd/analytical_space.py b/pycbc/psd/analytical_space.py index e3dd4efb174..9ca11a0fe22 100644 --- a/pycbc/psd/analytical_space.py +++ b/pycbc/psd/analytical_space.py @@ -31,13 +31,15 @@ """ import numpy as np -from scipy.interpolate import interp1d from astropy.constants import c +from scipy.interpolate import interp1d + from pycbc.psd.read import from_numpy_arrays def _psd_acc_noise(f, acc_noise_level=None): - """ The PSD of TDI-based space-borne GW + """ + The PSD of TDI-based space-borne GW detectors' acceleration noise. Note that this is suitable for LISA and Taiji, TianQin has a different form. @@ -53,19 +55,22 @@ def _psd_acc_noise(f, acc_noise_level=None): ------- s_acc_nu : float or numpy.array The PSD value or array for acceleration noise. + Notes ----- Please see Eq.(11-13) in for more details. + """ - s_acc = acc_noise_level**2 * (1+(4e-4/f)**2)*(1+(f/8e-3)**4) - s_acc_d = s_acc * (2*np.pi*f)**(-4) - s_acc_nu = (2*np.pi*f/c.value)**2 * s_acc_d + s_acc = acc_noise_level**2 * (1 + (4e-4 / f) ** 2) * (1 + (f / 8e-3) ** 4) + s_acc_d = s_acc * (2 * np.pi * f) ** (-4) + s_acc_nu = (2 * np.pi * f / c.value) ** 2 * s_acc_d return s_acc_nu def psd_lisa_acc_noise(f, acc_noise_level=3e-15): - """ The PSD of LISA's acceleration noise. + """ + The PSD of LISA's acceleration noise. Parameters ---------- @@ -78,9 +83,11 @@ def psd_lisa_acc_noise(f, acc_noise_level=3e-15): ------- s_acc_nu : float or numpy.array The PSD value or array for acceleration noise. + Notes ----- Please see Eq.(11-13) in for more details. + """ s_acc_nu = _psd_acc_noise(f, acc_noise_level) @@ -88,7 +95,8 @@ def psd_lisa_acc_noise(f, acc_noise_level=3e-15): def psd_tianqin_acc_noise(f, acc_noise_level=1e-15): - """ The PSD of TianQin's acceleration noise. + """ + The PSD of TianQin's acceleration noise. Parameters ---------- @@ -101,19 +109,22 @@ def psd_tianqin_acc_noise(f, acc_noise_level=1e-15): ------- s_acc_nu : float or numpy.array The PSD value or array for acceleration noise. + Notes ----- Please see Table(1) in <10.1088/0264-9381/33/3/035010> and that paper for more details. + """ - s_acc_d = acc_noise_level**2 * (2*np.pi*f)**(-4) * (1+1e-4/f) - s_acc_nu = (2*np.pi*f/c.value)**2 * s_acc_d + s_acc_d = acc_noise_level**2 * (2 * np.pi * f) ** (-4) * (1 + 1e-4 / f) + s_acc_nu = (2 * np.pi * f / c.value) ** 2 * s_acc_d return s_acc_nu def psd_taiji_acc_noise(f, acc_noise_level=3e-15): - """ The PSD of Taiji's acceleration noise. + """ + The PSD of Taiji's acceleration noise. Parameters ---------- @@ -126,9 +137,11 @@ def psd_taiji_acc_noise(f, acc_noise_level=3e-15): ------- s_acc_nu : float or numpy.array The PSD value or array for acceleration noise. + Notes ----- Please see Eq.(2) in <10.1103/PhysRevD.107.064021> for more details. + """ s_acc_nu = _psd_acc_noise(f, acc_noise_level) @@ -136,7 +149,8 @@ def psd_taiji_acc_noise(f, acc_noise_level=3e-15): def _psd_oms_noise(f, oms_noise_level=None): - """ The PSD of TDI-based space-borne GW detectors' OMS noise. + """ + The PSD of TDI-based space-borne GW detectors' OMS noise. Note that this is suitable for LISA and Taiji, TianQin has a different form. @@ -151,18 +165,21 @@ def _psd_oms_noise(f, oms_noise_level=None): ------- s_oms_nu : float or numpy.array The PSD value or array for OMS noise. + Notes ----- Please see Eq.(9-10) in for more details. + """ - s_oms_d = oms_noise_level**2 * (1+(2e-3/f)**4) - s_oms_nu = s_oms_d * (2*np.pi*f/c.value)**2 + s_oms_d = oms_noise_level**2 * (1 + (2e-3 / f) ** 4) + s_oms_nu = s_oms_d * (2 * np.pi * f / c.value) ** 2 return s_oms_nu def psd_lisa_oms_noise(f, oms_noise_level=15e-12): - """ The PSD of LISA's OMS noise. + """ + The PSD of LISA's OMS noise. Parameters ---------- @@ -175,9 +192,11 @@ def psd_lisa_oms_noise(f, oms_noise_level=15e-12): ------- s_oms_nu : float or numpy.array The PSD value or array for OMS noise. + Notes ----- Please see Eq.(9-10) in for more details. + """ s_oms_nu = _psd_oms_noise(f, oms_noise_level) @@ -185,7 +204,8 @@ def psd_lisa_oms_noise(f, oms_noise_level=15e-12): def psd_tianqin_oms_noise(f, oms_noise_level=1e-12): - """ The PSD of TianQin's OMS noise. + """ + The PSD of TianQin's OMS noise. Parameters ---------- @@ -198,19 +218,22 @@ def psd_tianqin_oms_noise(f, oms_noise_level=1e-12): ------- s_oms_nu : float or numpy.array The PSD value or array for OMS noise. + Notes ----- Please see Table(1) in <10.1088/0264-9381/33/3/035010> and that paper for more details. + """ s_oms_d = oms_noise_level**2 - s_oms_nu = s_oms_d * (2*np.pi*f/c.value)**2 + s_oms_nu = s_oms_d * (2 * np.pi * f / c.value) ** 2 return s_oms_nu def psd_taiji_oms_noise(f, oms_noise_level=8e-12): - """ The PSD of Taiji's OMS noise. + """ + The PSD of Taiji's OMS noise. Parameters ---------- @@ -223,9 +246,11 @@ def psd_taiji_oms_noise(f, oms_noise_level=8e-12): ------- s_oms_nu : float or numpy.array The PSD value or array for OMS noise. + Notes ----- Please see Eq.(1) in <10.1103/PhysRevD.107.064021> for more details. + """ s_oms_nu = _psd_oms_noise(f, oms_noise_level) @@ -233,7 +258,8 @@ def psd_taiji_oms_noise(f, oms_noise_level=8e-12): def lisa_psd_components(f, acc_noise_level=3e-15, oms_noise_level=15e-12): - """ The PSD of LISA's acceleration and OMS noise. + """ + The PSD of LISA's acceleration and OMS noise. Parameters ---------- @@ -248,6 +274,7 @@ def lisa_psd_components(f, acc_noise_level=3e-15, oms_noise_level=15e-12): ------- low_freq_component, high_freq_component : The PSD value or array for acceleration and OMS noise. + """ acc_noise_level = np.float64(acc_noise_level) oms_noise_level = np.float64(oms_noise_level) @@ -258,7 +285,8 @@ def lisa_psd_components(f, acc_noise_level=3e-15, oms_noise_level=15e-12): def tianqin_psd_components(f, acc_noise_level=1e-15, oms_noise_level=1e-12): - """ The PSD of TianQin's acceleration and OMS noise. + """ + The PSD of TianQin's acceleration and OMS noise. Parameters ---------- @@ -273,6 +301,7 @@ def tianqin_psd_components(f, acc_noise_level=1e-15, oms_noise_level=1e-12): ------- low_freq_component, high_freq_component : The PSD value or array for acceleration and OMS noise. + """ acc_noise_level = np.float64(acc_noise_level) oms_noise_level = np.float64(oms_noise_level) @@ -283,7 +312,8 @@ def tianqin_psd_components(f, acc_noise_level=1e-15, oms_noise_level=1e-12): def taiji_psd_components(f, acc_noise_level=3e-15, oms_noise_level=8e-12): - """ The PSD of Taiji's acceleration and OMS noise. + """ + The PSD of Taiji's acceleration and OMS noise. Parameters ---------- @@ -298,6 +328,7 @@ def taiji_psd_components(f, acc_noise_level=3e-15, oms_noise_level=8e-12): ------- low_freq_component, high_freq_component : The PSD value or array for acceleration and OMS noise. + """ acc_noise_level = np.float64(acc_noise_level) oms_noise_level = np.float64(oms_noise_level) @@ -308,7 +339,8 @@ def taiji_psd_components(f, acc_noise_level=3e-15, oms_noise_level=8e-12): def omega_length(f, len_arm=None): - """ The function to calculate 2*pi*f*arm_length. + """ + The function to calculate 2*pi*f*arm_length. Parameters ---------- @@ -322,15 +354,18 @@ def omega_length(f, len_arm=None): ------- omega_len : float or numpy.array The value of 2*pi*f*arm_length. + """ - omega_len = 2*np.pi*f * len_arm/c.value + omega_len = 2 * np.pi * f * len_arm / c.value return omega_len -def _analytical_psd_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm=None, psd_components=None, tdi=None): - """ The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for TDI-based +def _analytical_psd_tdi_XYZ( + length, delta_f, low_freq_cutoff, len_arm=None, psd_components=None, tdi=None +): + """ + The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for TDI-based space-borne GW detectors. Parameters @@ -352,30 +387,42 @@ def _analytical_psd_tdi_XYZ(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel). + Notes ----- Please see Eq.(19-20) in for more details. + """ len_arm = np.float64(len_arm) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) s_acc_nu, s_oms_nu = psd_components omega_len = omega_length(fr, len_arm) - psd = 16*(np.sin(omega_len))**2 * (s_oms_nu + - s_acc_nu*(3+np.cos(2*omega_len))) + psd = ( + 16 + * (np.sin(omega_len)) ** 2 + * (s_oms_nu + s_acc_nu * (3 + np.cos(2 * omega_len))) + ) if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 psd *= tdi2_factor fseries = from_numpy_arrays(fr, psd, length, delta_f, low_freq_cutoff) return fseries -def analytical_psd_lisa_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, tdi=None): - """ The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for LISA. +def analytical_psd_lisa_tdi_XYZ( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + tdi=None, +): + """ + The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for LISA. Parameters ---------- @@ -398,24 +445,32 @@ def analytical_psd_lisa_tdi_XYZ(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA. + Notes ----- Please see Eq.(19-20) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(lisa_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array(lisa_psd_components(fr, acc_noise_level, oms_noise_level)) + fseries = _analytical_psd_tdi_XYZ( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def analytical_psd_tianqin_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - acc_noise_level=1e-15, - oms_noise_level=1e-12, tdi=None): - """ The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for TianQin. +def analytical_psd_tianqin_tdi_XYZ( + length, + delta_f, + low_freq_cutoff, + len_arm=np.sqrt(3) * 1e8, + acc_noise_level=1e-15, + oms_noise_level=1e-12, + tdi=None, +): + """ + The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for TianQin. Parameters ---------- @@ -438,24 +493,35 @@ def analytical_psd_tianqin_tdi_XYZ(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for TianQin. + Notes ----- Please see Table(1) in <10.1088/0264-9381/33/3/035010> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(tianqin_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array( + tianqin_psd_components(fr, acc_noise_level, oms_noise_level) + ) + fseries = _analytical_psd_tdi_XYZ( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def analytical_psd_taiji_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm=3e9, acc_noise_level=3e-15, - oms_noise_level=8e-12, tdi=None): - """ The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for Taiji. +def analytical_psd_taiji_tdi_XYZ( + length, + delta_f, + low_freq_cutoff, + len_arm=3e9, + acc_noise_level=3e-15, + oms_noise_level=8e-12, + tdi=None, +): + """ + The TDI-1.5/2.0 analytical PSD (X,Y,Z channel) for Taiji. Parameters ---------- @@ -478,22 +544,28 @@ def analytical_psd_taiji_tdi_XYZ(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for Taiji. + Notes ----- Please see <10.1103/PhysRevD.107.064021> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(taiji_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array( + taiji_psd_components(fr, acc_noise_level, oms_noise_level) + ) + fseries = _analytical_psd_tdi_XYZ( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def _analytical_csd_tdi_XY(length, delta_f, low_freq_cutoff, - len_arm=None, psd_components=None, tdi=None): - """ The cross-spectrum density between TDI channel X and Y. +def _analytical_csd_tdi_XY( + length, delta_f, low_freq_cutoff, len_arm=None, psd_components=None, tdi=None +): + """ + The cross-spectrum density between TDI channel X and Y. Parameters ---------- @@ -514,30 +586,38 @@ def _analytical_csd_tdi_XY(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The CSD between TDI-1.5/2.0 channel X and Y. + Notes ----- Please see Eq.(56) in for more details. + """ len_arm = np.float64(len_arm) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) s_acc_nu, s_oms_nu = psd_components omega_len = omega_length(fr, len_arm) - csd = (-8*np.sin(omega_len)**2 * np.cos(omega_len) * - (s_oms_nu+4*s_acc_nu)) + csd = -8 * np.sin(omega_len) ** 2 * np.cos(omega_len) * (s_oms_nu + 4 * s_acc_nu) if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 csd *= tdi2_factor fseries = from_numpy_arrays(fr, csd, length, delta_f, low_freq_cutoff) return fseries -def analytical_csd_lisa_tdi_XY(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, tdi=None): - """ The cross-spectrum density between LISA's TDI channel X and Y. +def analytical_csd_lisa_tdi_XY( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + tdi=None, +): + """ + The cross-spectrum density between LISA's TDI channel X and Y. Parameters ---------- @@ -560,22 +640,26 @@ def analytical_csd_lisa_tdi_XY(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The CSD between LISA's TDI-1.5/2.0 channel X and Y. + Notes ----- Please see Eq.(56) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(lisa_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_csd_tdi_XY(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array(lisa_psd_components(fr, acc_noise_level, oms_noise_level)) + fseries = _analytical_csd_tdi_XY( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def _analytical_psd_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm=None, psd_components=None, tdi=None): - """ The PSD of TDI-1.5/2.0 channel A and E. +def _analytical_psd_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm=None, psd_components=None, tdi=None +): + """ + The PSD of TDI-1.5/2.0 channel A and E. Parameters ---------- @@ -596,31 +680,45 @@ def _analytical_psd_tdi_AE(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of TDI-1.5/2.0 channel A and E. + Notes ----- Please see Eq.(58) in for more details. + """ len_arm = np.float64(len_arm) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) s_acc_nu, s_oms_nu = psd_components omega_len = omega_length(fr, len_arm) - psd = (8*(np.sin(omega_len))**2 * - (4*(1+np.cos(omega_len)+np.cos(omega_len)**2)*s_acc_nu + - (2+np.cos(omega_len))*s_oms_nu)) + psd = ( + 8 + * (np.sin(omega_len)) ** 2 + * ( + 4 * (1 + np.cos(omega_len) + np.cos(omega_len) ** 2) * s_acc_nu + + (2 + np.cos(omega_len)) * s_oms_nu + ) + ) if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 psd *= tdi2_factor fseries = from_numpy_arrays(fr, psd, length, delta_f, low_freq_cutoff) return fseries -def analytical_psd_lisa_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, tdi=None): - """ The PSD of LISA's TDI-1.5/2.0 channel A and E. +def analytical_psd_lisa_tdi_AE( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + tdi=None, +): + """ + The PSD of LISA's TDI-1.5/2.0 channel A and E. Parameters ---------- @@ -643,24 +741,32 @@ def analytical_psd_lisa_tdi_AE(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of LISA's TDI-1.5/2.0 channel A and E. + Notes ----- Please see Eq.(58) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(lisa_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array(lisa_psd_components(fr, acc_noise_level, oms_noise_level)) + fseries = _analytical_psd_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def analytical_psd_tianqin_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - acc_noise_level=1e-15, - oms_noise_level=1e-12, tdi=None): - """ The PSD of TianQin's TDI-1.5/2.0 channel A and E. +def analytical_psd_tianqin_tdi_AE( + length, + delta_f, + low_freq_cutoff, + len_arm=np.sqrt(3) * 1e8, + acc_noise_level=1e-15, + oms_noise_level=1e-12, + tdi=None, +): + """ + The PSD of TianQin's TDI-1.5/2.0 channel A and E. Parameters ---------- @@ -683,24 +789,35 @@ def analytical_psd_tianqin_tdi_AE(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of TianQin's TDI-1.5/2.0 channel A and E. + Notes ----- Please see Table(1) in <10.1088/0264-9381/33/3/035010> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(tianqin_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array( + tianqin_psd_components(fr, acc_noise_level, oms_noise_level) + ) + fseries = _analytical_psd_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def analytical_psd_taiji_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm=3e9, acc_noise_level=3e-15, - oms_noise_level=8e-12, tdi=None): - """ The PSD of Taiji's TDI-1.5/2.0 channel A and E. +def analytical_psd_taiji_tdi_AE( + length, + delta_f, + low_freq_cutoff, + len_arm=3e9, + acc_noise_level=3e-15, + oms_noise_level=8e-12, + tdi=None, +): + """ + The PSD of Taiji's TDI-1.5/2.0 channel A and E. Parameters ---------- @@ -723,22 +840,28 @@ def analytical_psd_taiji_tdi_AE(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of Taiji's TDI-1.5/2.0 channel A and E. + Notes ----- Please see <10.1103/PhysRevD.107.064021> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(taiji_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array( + taiji_psd_components(fr, acc_noise_level, oms_noise_level) + ) + fseries = _analytical_psd_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def _analytical_psd_tdi_T(length, delta_f, low_freq_cutoff, - len_arm=None, psd_components=None, tdi=None): - """ The PSD of TDI-1.5/2.0 channel T. +def _analytical_psd_tdi_T( + length, delta_f, low_freq_cutoff, len_arm=None, psd_components=None, tdi=None +): + """ + The PSD of TDI-1.5/2.0 channel T. Parameters ---------- @@ -759,30 +882,43 @@ def _analytical_psd_tdi_T(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of TDI-1.5/2.0 channel T. + Notes ----- Please see Eq.(59) in for more details. + """ len_arm = np.float64(len_arm) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) s_acc_nu, s_oms_nu = psd_components omega_len = omega_length(fr, len_arm) - psd = (32*np.sin(omega_len)**2 * np.sin(omega_len/2)**2 * - (4*s_acc_nu*np.sin(omega_len/2)**2 + s_oms_nu)) + psd = ( + 32 + * np.sin(omega_len) ** 2 + * np.sin(omega_len / 2) ** 2 + * (4 * s_acc_nu * np.sin(omega_len / 2) ** 2 + s_oms_nu) + ) if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 psd *= tdi2_factor fseries = from_numpy_arrays(fr, psd, length, delta_f, low_freq_cutoff) return fseries -def analytical_psd_lisa_tdi_T(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, tdi=None): - """ The PSD of LISA's TDI-1.5/2.0 channel T. +def analytical_psd_lisa_tdi_T( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + tdi=None, +): + """ + The PSD of LISA's TDI-1.5/2.0 channel T. Parameters ---------- @@ -805,24 +941,32 @@ def analytical_psd_lisa_tdi_T(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of LISA's TDI-1.5/2.0 channel T. + Notes ----- Please see Eq.(59) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(lisa_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_T(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array(lisa_psd_components(fr, acc_noise_level, oms_noise_level)) + fseries = _analytical_psd_tdi_T( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def analytical_psd_tianqin_tdi_T(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - acc_noise_level=1e-15, - oms_noise_level=1e-12, tdi=None): - """ The PSD of TianQin's TDI-1.5/2.0 channel T. +def analytical_psd_tianqin_tdi_T( + length, + delta_f, + low_freq_cutoff, + len_arm=np.sqrt(3) * 1e8, + acc_noise_level=1e-15, + oms_noise_level=1e-12, + tdi=None, +): + """ + The PSD of TianQin's TDI-1.5/2.0 channel T. Parameters ---------- @@ -845,24 +989,35 @@ def analytical_psd_tianqin_tdi_T(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of TianQin's TDI-1.5/2.0 channel T. + Notes ----- Please see Table(1) in <10.1088/0264-9381/33/3/035010> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(tianqin_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_T(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array( + tianqin_psd_components(fr, acc_noise_level, oms_noise_level) + ) + fseries = _analytical_psd_tdi_T( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries -def analytical_psd_taiji_tdi_T(length, delta_f, low_freq_cutoff, - len_arm=3e9, acc_noise_level=3e-15, - oms_noise_level=8e-12, tdi=None): - """ The PSD of Taiji's TDI-1.5/2.0 channel T. +def analytical_psd_taiji_tdi_T( + length, + delta_f, + low_freq_cutoff, + len_arm=3e9, + acc_noise_level=3e-15, + oms_noise_level=8e-12, + tdi=None, +): + """ + The PSD of Taiji's TDI-1.5/2.0 channel T. Parameters ---------- @@ -885,21 +1040,26 @@ def analytical_psd_taiji_tdi_T(length, delta_f, low_freq_cutoff, ------- fseries : FrequencySeries The PSD of Taiji's TDI-1.5/2.0 channel T. + Notes ----- Please see <10.1103/PhysRevD.107.064021> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - psd_components = np.array(taiji_psd_components( - fr, acc_noise_level, oms_noise_level)) - fseries = _analytical_psd_tdi_T(length, delta_f, low_freq_cutoff, - len_arm, psd_components, tdi) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + psd_components = np.array( + taiji_psd_components(fr, acc_noise_level, oms_noise_level) + ) + fseries = _analytical_psd_tdi_T( + length, delta_f, low_freq_cutoff, len_arm, psd_components, tdi + ) return fseries def averaged_lisa_fplus_sq_numerical(f, len_arm=2.5e9): - """ A numerical fit for LISA's squared antenna response function, + """ + A numerical fit for LISA's squared antenna response function, averaged over sky and polarization angle. Parameters @@ -913,9 +1073,11 @@ def averaged_lisa_fplus_sq_numerical(f, len_arm=2.5e9): ------- fp_sq_numerical : float or numpy.array The sky and polarization angle averaged squared antenna response. + Notes ----- Please see Eq.(36) in for more details. + """ from astropy.utils.data import download_file @@ -928,15 +1090,15 @@ def averaged_lisa_fplus_sq_numerical(f, len_arm=2.5e9): # Padding the end. freqs = np.append(freqs, 2) fp_sq = np.append(fp_sq, 0.0012712348970728724) - fp_sq_interp = interp1d(freqs, fp_sq, kind='linear', - fill_value="extrapolate") - fp_sq_numerical = fp_sq_interp(f)/16 + fp_sq_interp = interp1d(freqs, fp_sq, kind="linear", fill_value="extrapolate") + fp_sq_numerical = fp_sq_interp(f) / 16 return fp_sq_numerical def averaged_fplus_sq_approximated(f, len_arm=None): - r""" A simplified fit for TDI-based space-borne GW detectors' + r""" + A simplified fit for TDI-based space-borne GW detectors' squared antenna response function, averaged over sky and polarization angle. @@ -955,17 +1117,20 @@ def averaged_fplus_sq_approximated(f, len_arm=None): ------- fp_sq_approx : float or numpy.array The sky and polarization angle averaged squared antenna response. + Notes ----- Please see Eq.(9) in <10.1088/1361-6382/ab1101> for more details. + """ - fp_sq_approx = (3./20.)*(1./(1.+0.6*omega_length(f, len_arm)**2)) + fp_sq_approx = (3.0 / 20.0) * (1.0 / (1.0 + 0.6 * omega_length(f, len_arm) ** 2)) return fp_sq_approx -def averaged_tianqin_fplus_sq_numerical(f, len_arm=np.sqrt(3)*1e8): - """ A numerical fit for TianQin's squared antenna response function, +def averaged_tianqin_fplus_sq_numerical(f, len_arm=np.sqrt(3) * 1e8): + """ + A numerical fit for TianQin's squared antenna response function, averaged over sky and polarization angle. Parameters @@ -979,24 +1144,35 @@ def averaged_tianqin_fplus_sq_numerical(f, len_arm=np.sqrt(3)*1e8): ------- fp_sq_numerical : float or numpy.array The sky and polarization angle averaged squared antenna response. + Notes ----- Please see Eq.(15-16) in <10.1103/PhysRevD.100.043003> for more details. + """ base = averaged_fplus_sq_approximated(f, len_arm) - a = [1, 1e-4, 2639e-4, 231/5*1e-4, -2093/1.25*1e-4, 2173e-5, - 2101e-6, 3027/2*1e-5, -42373/5*1e-6, 176087e-8, - -8023/5*1e-7, 5169e-9] + a = [ + 1, + 1e-4, + 2639e-4, + 231 / 5 * 1e-4, + -2093 / 1.25 * 1e-4, + 2173e-5, + 2101e-6, + 3027 / 2 * 1e-5, + -42373 / 5 * 1e-6, + 176087e-8, + -8023 / 5 * 1e-7, + 5169e-9, + ] omega_len = omega_length(f, len_arm) omega_len_low_f = omega_len[omega_len < 4.1] omega_len_high_f = omega_len[omega_len >= 4.1] - base_low_f = base[:len(omega_len_low_f)] - base_high_f = base[len(omega_len_low_f):] + base_low_f = base[: len(omega_len_low_f)] + base_high_f = base[len(omega_len_low_f) :] low_f_modulation = np.polyval(a[::-1], omega_len_low_f) - high_f_modulation = np.exp( - -0.322 * np.sin(2*omega_len_high_f-4.712) + 0.078 - ) + high_f_modulation = np.exp(-0.322 * np.sin(2 * omega_len_high_f - 4.712) + 0.078) low_f_result = np.multiply(base_low_f, low_f_modulation) high_f_result = np.multiply(base_high_f, high_f_modulation) fp_sq_numerical = np.concatenate((low_f_result, high_f_result)) @@ -1005,7 +1181,8 @@ def averaged_tianqin_fplus_sq_numerical(f, len_arm=np.sqrt(3)*1e8): def averaged_response_lisa_tdi(f, len_arm=2.5e9, tdi=None): - """ LISA's TDI-1.5/2.0 response function to GW, + """ + LISA's TDI-1.5/2.0 response function to GW, averaged over sky and polarization angle. Parameters @@ -1021,24 +1198,27 @@ def averaged_response_lisa_tdi(f, len_arm=2.5e9, tdi=None): ------- response_tdi : float or numpy.array The sky and polarization angle averaged TDI-1.5/2.0 response to GW. + Notes ----- Please see Eq.(39-40) in for more details. + """ omega_len = omega_length(f, len_arm) ave_fp2 = averaged_lisa_fplus_sq_numerical(f, len_arm) - response_tdi = (4*omega_len)**2 * np.sin(omega_len)**2 * ave_fp2 + response_tdi = (4 * omega_len) ** 2 * np.sin(omega_len) ** 2 * ave_fp2 if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 response_tdi *= tdi2_factor return response_tdi -def averaged_response_tianqin_tdi(f, len_arm=np.sqrt(3)*1e8, tdi=None): - """ TianQin's TDI-1.5/2.0 response function to GW, +def averaged_response_tianqin_tdi(f, len_arm=np.sqrt(3) * 1e8, tdi=None): + """ + TianQin's TDI-1.5/2.0 response function to GW, averaged over sky and polarization angle. Parameters @@ -1054,21 +1234,23 @@ def averaged_response_tianqin_tdi(f, len_arm=np.sqrt(3)*1e8, tdi=None): ------- response_tdi : float or numpy.array The sky and polarization angle averaged TDI-1.5/2.0 response to GW. + """ omega_len = omega_length(f, len_arm) ave_fp2 = averaged_tianqin_fplus_sq_numerical(f, len_arm) - response_tdi = (4*omega_len)**2 * np.sin(omega_len)**2 * ave_fp2 + response_tdi = (4 * omega_len) ** 2 * np.sin(omega_len) ** 2 * ave_fp2 if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 response_tdi *= tdi2_factor return response_tdi def averaged_response_taiji_tdi(f, len_arm=3e9, tdi=None): - """ Taiji's TDI-1.5/2.0 response function to GW, + """ + Taiji's TDI-1.5/2.0 response function to GW, averaged over sky and polarization angle. Parameters @@ -1084,24 +1266,30 @@ def averaged_response_taiji_tdi(f, len_arm=3e9, tdi=None): ------- response_tdi : float or numpy.array The sky and polarization angle averaged TDI-1.5/2.0 response to GW. + """ omega_len = omega_length(f, len_arm) ave_fp2 = averaged_fplus_sq_approximated(f, len_arm) - response_tdi = (4*omega_len)**2 * np.sin(omega_len)**2 * ave_fp2 + response_tdi = (4 * omega_len) ** 2 * np.sin(omega_len) ** 2 * ave_fp2 if str(tdi) not in ["1.5", "2.0"]: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") if str(tdi) == "2.0": - tdi2_factor = 4*(np.sin(2*omega_len))**2 + tdi2_factor = 4 * (np.sin(2 * omega_len)) ** 2 response_tdi *= tdi2_factor return response_tdi -def sensitivity_curve_lisa_semi_analytical(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, - acc_noise_level=3e-15, - oms_noise_level=15e-12): - """ The semi-analytical LISA's sensitivity curve (6-links), +def sensitivity_curve_lisa_semi_analytical( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, +): + """ + The semi-analytical LISA's sensitivity curve (6-links), averaged over sky and polarization angle. Parameters @@ -1124,31 +1312,37 @@ def sensitivity_curve_lisa_semi_analytical(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The sky and polarization angle averaged semi-analytical LISA's sensitivity curve (6-links). + Notes ----- Please see Eq.(42-43) in for more details. + """ len_arm = np.float64(len_arm) acc_noise_level = np.float64(acc_noise_level) oms_noise_level = np.float64(oms_noise_level) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) fp_sq = averaged_lisa_fplus_sq_numerical(fr, len_arm) - s_acc_nu, s_oms_nu = lisa_psd_components( - fr, acc_noise_level, oms_noise_level) + s_acc_nu, s_oms_nu = lisa_psd_components(fr, acc_noise_level, oms_noise_level) omega_len = omega_length(fr, len_arm) - sense_curve = ((s_oms_nu + s_acc_nu*(3+np.cos(2*omega_len))) / - (omega_len**2*fp_sq)) - fseries = from_numpy_arrays(fr, sense_curve/2, - length, delta_f, low_freq_cutoff) + sense_curve = (s_oms_nu + s_acc_nu * (3 + np.cos(2 * omega_len))) / ( + omega_len**2 * fp_sq + ) + fseries = from_numpy_arrays(fr, sense_curve / 2, length, delta_f, low_freq_cutoff) return fseries -def sensitivity_curve_tianqin_analytical(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - acc_noise_level=1e-15, - oms_noise_level=1e-12): - """ The analytical TianQin's sensitivity curve (6-links), +def sensitivity_curve_tianqin_analytical( + length, + delta_f, + low_freq_cutoff, + len_arm=np.sqrt(3) * 1e8, + acc_noise_level=1e-15, + oms_noise_level=1e-12, +): + """ + The analytical TianQin's sensitivity curve (6-links), averaged over sky and polarization angle. Parameters @@ -1171,27 +1365,33 @@ def sensitivity_curve_tianqin_analytical(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The sky and polarization angle averaged analytical TianQin's sensitivity curve (6-links). + """ len_arm = np.float64(len_arm) acc_noise_level = np.float64(acc_noise_level) oms_noise_level = np.float64(oms_noise_level) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) fp_sq = averaged_tianqin_fplus_sq_numerical(fr, len_arm) - s_acc_nu, s_oms_nu = tianqin_psd_components( - fr, acc_noise_level, oms_noise_level) + s_acc_nu, s_oms_nu = tianqin_psd_components(fr, acc_noise_level, oms_noise_level) omega_len = omega_length(fr, len_arm) - sense_curve = ((s_oms_nu + s_acc_nu*(3+np.cos(2*omega_len))) / - (omega_len**2*fp_sq)) - fseries = from_numpy_arrays(fr, sense_curve/2, - length, delta_f, low_freq_cutoff) + sense_curve = (s_oms_nu + s_acc_nu * (3 + np.cos(2 * omega_len))) / ( + omega_len**2 * fp_sq + ) + fseries = from_numpy_arrays(fr, sense_curve / 2, length, delta_f, low_freq_cutoff) return fseries -def sensitivity_curve_taiji_analytical(length, delta_f, low_freq_cutoff, - len_arm=3e9, acc_noise_level=3e-15, - oms_noise_level=8e-12): - """ The analytical Taiji's sensitivity curve (6-links), +def sensitivity_curve_taiji_analytical( + length, + delta_f, + low_freq_cutoff, + len_arm=3e9, + acc_noise_level=3e-15, + oms_noise_level=8e-12, +): + """ + The analytical Taiji's sensitivity curve (6-links), averaged over sky and polarization angle. Parameters @@ -1214,25 +1414,26 @@ def sensitivity_curve_taiji_analytical(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The sky and polarization angle averaged analytical Taiji's sensitivity curve (6-links). + """ len_arm = np.float64(len_arm) acc_noise_level = np.float64(acc_noise_level) oms_noise_level = np.float64(oms_noise_level) - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) fp_sq = averaged_fplus_sq_approximated(fr, len_arm) - s_acc_nu, s_oms_nu = taiji_psd_components( - fr, acc_noise_level, oms_noise_level) + s_acc_nu, s_oms_nu = taiji_psd_components(fr, acc_noise_level, oms_noise_level) omega_len = omega_length(fr, len_arm) - sense_curve = ((s_oms_nu + s_acc_nu*(3+np.cos(2*omega_len))) / - (omega_len**2*fp_sq)) - fseries = from_numpy_arrays(fr, sense_curve/2, - length, delta_f, low_freq_cutoff) + sense_curve = (s_oms_nu + s_acc_nu * (3 + np.cos(2 * omega_len))) / ( + omega_len**2 * fp_sq + ) + fseries = from_numpy_arrays(fr, sense_curve / 2, length, delta_f, low_freq_cutoff) return fseries def sensitivity_curve_lisa_SciRD(length, delta_f, low_freq_cutoff): - """ The analytical LISA's sensitivity curve in SciRD, + """ + The analytical LISA's sensitivity curve in SciRD, averaged over sky and polarization angle. Parameters @@ -1249,23 +1450,25 @@ def sensitivity_curve_lisa_SciRD(length, delta_f, low_freq_cutoff): fseries : FrequencySeries The sky and polarization angle averaged analytical LISA's sensitivity curve in SciRD. + Notes ----- Please see Eq.(114) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - s_I = 5.76e-48 * (1+(4e-4/fr)**2) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + s_I = 5.76e-48 * (1 + (4e-4 / fr) ** 2) s_II = 3.6e-41 - R = 1 + (fr/2.5e-2)**2 - sense_curve = 10/3 * (s_I/(2*np.pi*fr)**4+s_II) * R - fseries = from_numpy_arrays(fr, sense_curve, length, - delta_f, low_freq_cutoff) + R = 1 + (fr / 2.5e-2) ** 2 + sense_curve = 10 / 3 * (s_I / (2 * np.pi * fr) ** 4 + s_II) * R + fseries = from_numpy_arrays(fr, sense_curve, length, delta_f, low_freq_cutoff) return fseries def confusion_fit_lisa(length, delta_f, low_freq_cutoff, duration=1.0): - """ The LISA's sensitivity curve for Galactic confusion noise, + """ + The LISA's sensitivity curve for Galactic confusion noise, averaged over sky and polarization angle. No instrumental noise. Parameters @@ -1285,23 +1488,30 @@ def confusion_fit_lisa(length, delta_f, low_freq_cutoff, duration=1.0): The sky and polarization angle averaged LISA's sensitivity curve for Galactic confusion noise. No instrumental noise. + Notes ----- Please see Eq.(85-86) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) - f1 = 10**(-0.25*np.log10(duration)-2.7) - fk = 10**(-0.27*np.log10(duration)-2.47) - sh_confusion = (0.5*1.14e-44*fr**(-7/3)*np.exp(-(fr/f1)**1.8) * - (1.0+np.tanh((fk-fr)/0.31e-3))) - fseries = from_numpy_arrays(fr, sh_confusion, length, delta_f, - low_freq_cutoff) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) + f1 = 10 ** (-0.25 * np.log10(duration) - 2.7) + fk = 10 ** (-0.27 * np.log10(duration) - 2.47) + sh_confusion = ( + 0.5 + * 1.14e-44 + * fr ** (-7 / 3) + * np.exp(-((fr / f1) ** 1.8)) + * (1.0 + np.tanh((fk - fr) / 0.31e-3)) + ) + fseries = from_numpy_arrays(fr, sh_confusion, length, delta_f, low_freq_cutoff) return fseries def confusion_fit_tianqin(length, delta_f, low_freq_cutoff, duration=1.0): - """ The TianQin's sensitivity curve for Galactic confusion noise, + """ + The TianQin's sensitivity curve for Galactic confusion noise, averaged over sky and polarization angle. No instrumental noise. Only valid for 0.5 mHz < f < 10 mHz. Note that the results between 0.5, 1, 2, 4, and 5 years are extrapolated, might be non-physical. @@ -1323,12 +1533,14 @@ def confusion_fit_tianqin(length, delta_f, low_freq_cutoff, duration=1.0): The sky and polarization angle averaged TianQin's sensitivity curve for Galactic confusion noise. No instrumental noise. + Notes ----- Please see Table(II) in <10.1103/PhysRevD.102.063021> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) t_obs = [0.5, 1, 2, 4, 5] a0 = [-18.6, -18.6, -18.6, -18.6, -18.6] a1 = [-1.22, -1.13, -1.45, -1.43, -1.51] @@ -1337,41 +1549,47 @@ def confusion_fit_tianqin(length, delta_f, low_freq_cutoff, duration=1.0): a4 = [0.65, 4.05, -4.48, -0.15, -0.83] a5 = [3.6, -4.5, 10.8, -1.8, 13.2] a6 = [-4.6, -0.5, -9.4, -3.2, -19.1] - fit_a0 = interp1d(t_obs, a0, kind='cubic', fill_value="extrapolate") - fit_a1 = interp1d(t_obs, a1, kind='cubic', fill_value="extrapolate") - fit_a2 = interp1d(t_obs, a2, kind='cubic', fill_value="extrapolate") - fit_a3 = interp1d(t_obs, a3, kind='cubic', fill_value="extrapolate") - fit_a4 = interp1d(t_obs, a4, kind='cubic', fill_value="extrapolate") - fit_a5 = interp1d(t_obs, a5, kind='cubic', fill_value="extrapolate") - fit_a6 = interp1d(t_obs, a6, kind='cubic', fill_value="extrapolate") + fit_a0 = interp1d(t_obs, a0, kind="cubic", fill_value="extrapolate") + fit_a1 = interp1d(t_obs, a1, kind="cubic", fill_value="extrapolate") + fit_a2 = interp1d(t_obs, a2, kind="cubic", fill_value="extrapolate") + fit_a3 = interp1d(t_obs, a3, kind="cubic", fill_value="extrapolate") + fit_a4 = interp1d(t_obs, a4, kind="cubic", fill_value="extrapolate") + fit_a5 = interp1d(t_obs, a5, kind="cubic", fill_value="extrapolate") + fit_a6 = interp1d(t_obs, a6, kind="cubic", fill_value="extrapolate") if duration not in t_obs: - raise Warning("Note that the results between " + - "0.5, 1, 2, 4, and 5 years are extrapolated, " + - "might be non-physical.") + raise Warning( + "Note that the results between " + "0.5, 1, 2, 4, and 5 years are extrapolated, " + "might be non-physical." + ) # 10/3 is the factor for sky-average, the original fit in the paper # is not sky-averaged. - sh_confusion = 10./3 * np.power( - 10, - fit_a0(duration) + - fit_a1(duration) * np.log10(fr*1e3) + - fit_a2(duration) * np.log10(fr*1e3)**2 + - fit_a3(duration) * np.log10(fr*1e3)**3 + - fit_a4(duration) * np.log10(fr*1e3)**4 + - fit_a5(duration) * np.log10(fr*1e3)**5 + - fit_a6(duration) * np.log10(fr*1e3)**6 - )**2 + sh_confusion = ( + 10.0 + / 3 + * np.power( + 10, + fit_a0(duration) + + fit_a1(duration) * np.log10(fr * 1e3) + + fit_a2(duration) * np.log10(fr * 1e3) ** 2 + + fit_a3(duration) * np.log10(fr * 1e3) ** 3 + + fit_a4(duration) * np.log10(fr * 1e3) ** 4 + + fit_a5(duration) * np.log10(fr * 1e3) ** 5 + + fit_a6(duration) * np.log10(fr * 1e3) ** 6, + ) + ** 2 + ) # avoid the jump of values - sh_confusion[(fr > 3e-4) & (fr < 5e-4)] = \ - sh_confusion[(np.abs(fr - 5e-4)).argmin()] + sh_confusion[(fr > 3e-4) & (fr < 5e-4)] = sh_confusion[(np.abs(fr - 5e-4)).argmin()] sh_confusion[(fr < 3e-4) | (fr > 1e-2)] = 0 - fseries = from_numpy_arrays(fr, sh_confusion, length, delta_f, - low_freq_cutoff) + fseries = from_numpy_arrays(fr, sh_confusion, length, delta_f, low_freq_cutoff) return fseries def confusion_fit_taiji(length, delta_f, low_freq_cutoff, duration=1.0): - """ The Taiji's sensitivity curve for Galactic confusion noise, + """ + The Taiji's sensitivity curve for Galactic confusion noise, averaged over sky and polarization angle. No instrumental noise. Only valid for 0.1 mHz < f < 10 mHz. Note that the results between 0.5, 1, 2, and 4 years are extrapolated, might be non-physical. @@ -1393,12 +1611,14 @@ def confusion_fit_taiji(length, delta_f, low_freq_cutoff, duration=1.0): The sky and polarization angle averaged Taiji's sensitivity curve for Galactic confusion noise. No instrumental noise. + Notes ----- Please see Eq.(6) and Table(I) in <10.1103/PhysRevD.107.064021> for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) t_obs = [0.5, 1, 2, 4] a0 = [-85.3498, -85.4336, -85.3919, -85.5448] a1 = [-2.64899, -2.46276, -2.69735, -3.23671] @@ -1406,36 +1626,44 @@ def confusion_fit_taiji(length, delta_f, low_freq_cutoff, duration=1.0): a3 = [-0.478447, -0.884147, -1.15302, -1.14711] a4 = [-0.334821, -0.427176, -0.302761, 0.0325887] a5 = [0.0658353, 0.128666, 0.175521, 0.187854] - fit_a0 = interp1d(t_obs, a0, kind='cubic', fill_value="extrapolate") - fit_a1 = interp1d(t_obs, a1, kind='cubic', fill_value="extrapolate") - fit_a2 = interp1d(t_obs, a2, kind='cubic', fill_value="extrapolate") - fit_a3 = interp1d(t_obs, a3, kind='cubic', fill_value="extrapolate") - fit_a4 = interp1d(t_obs, a4, kind='cubic', fill_value="extrapolate") - fit_a5 = interp1d(t_obs, a5, kind='cubic', fill_value="extrapolate") + fit_a0 = interp1d(t_obs, a0, kind="cubic", fill_value="extrapolate") + fit_a1 = interp1d(t_obs, a1, kind="cubic", fill_value="extrapolate") + fit_a2 = interp1d(t_obs, a2, kind="cubic", fill_value="extrapolate") + fit_a3 = interp1d(t_obs, a3, kind="cubic", fill_value="extrapolate") + fit_a4 = interp1d(t_obs, a4, kind="cubic", fill_value="extrapolate") + fit_a5 = interp1d(t_obs, a5, kind="cubic", fill_value="extrapolate") if duration not in t_obs: - raise Warning("Note that the results between " + - "0.5, 1, 2, and 4 years are extrapolated, " + - "might be non-physical.") + raise Warning( + "Note that the results between " + "0.5, 1, 2, and 4 years are extrapolated, " + "might be non-physical." + ) sh_confusion = np.exp( - fit_a0(duration) + - fit_a1(duration) * np.log(fr*1e3) + - fit_a2(duration) * np.log(fr*1e3)**2 + - fit_a3(duration) * np.log(fr*1e3)**3 + - fit_a4(duration) * np.log(fr*1e3)**4 + - fit_a5(duration) * np.log(fr*1e3)**5 + fit_a0(duration) + + fit_a1(duration) * np.log(fr * 1e3) + + fit_a2(duration) * np.log(fr * 1e3) ** 2 + + fit_a3(duration) * np.log(fr * 1e3) ** 3 + + fit_a4(duration) * np.log(fr * 1e3) ** 4 + + fit_a5(duration) * np.log(fr * 1e3) ** 5 ) sh_confusion[(fr < 1e-4) | (fr > 1e-2)] = 0 - fseries = from_numpy_arrays(fr, sh_confusion, length, delta_f, - low_freq_cutoff) + fseries = from_numpy_arrays(fr, sh_confusion, length, delta_f, low_freq_cutoff) return fseries -def sensitivity_curve_lisa_confusion(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, - base_model="semi", duration=1.0): - """ The LISA's sensitivity curve with Galactic confusion noise, +def sensitivity_curve_lisa_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + base_model="semi", + duration=1.0, +): + """ + The LISA's sensitivity curve with Galactic confusion noise, averaged over sky and polarization angle. Parameters @@ -1462,35 +1690,45 @@ def sensitivity_curve_lisa_confusion(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The sky and polarization angle averaged LISA's sensitivity curve with Galactic confusion noise. + Notes ----- Please see Eq.(85-86) in for more details. + """ if base_model == "semi": base_curve = sensitivity_curve_lisa_semi_analytical( - length, delta_f, low_freq_cutoff, - len_arm, acc_noise_level, oms_noise_level) + length, delta_f, low_freq_cutoff, len_arm, acc_noise_level, oms_noise_level + ) elif base_model == "SciRD": - base_curve = sensitivity_curve_lisa_SciRD( - length, delta_f, low_freq_cutoff) + base_curve = sensitivity_curve_lisa_SciRD(length, delta_f, low_freq_cutoff) else: raise ValueError("Must choose from 'semi' or 'SciRD'.") if duration < 0 or duration > 10: raise ValueError("Must between 0 and 10.") - fseries_confusion = confusion_fit_lisa( - length, delta_f, low_freq_cutoff, duration) - fseries = from_numpy_arrays(base_curve.sample_frequencies, - base_curve+fseries_confusion, - length, delta_f, low_freq_cutoff) + fseries_confusion = confusion_fit_lisa(length, delta_f, low_freq_cutoff, duration) + fseries = from_numpy_arrays( + base_curve.sample_frequencies, + base_curve + fseries_confusion, + length, + delta_f, + low_freq_cutoff, + ) return fseries -def sensitivity_curve_tianqin_confusion(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - acc_noise_level=1e-15, - oms_noise_level=1e-12, duration=1.0): - """ The TianQin's sensitivity curve with Galactic confusion noise, +def sensitivity_curve_tianqin_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm=np.sqrt(3) * 1e8, + acc_noise_level=1e-15, + oms_noise_level=1e-12, + duration=1.0, +): + """ + The TianQin's sensitivity curve with Galactic confusion noise, averaged over sky and polarization angle. Parameters @@ -1515,25 +1753,38 @@ def sensitivity_curve_tianqin_confusion(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The sky and polarization angle averaged TianQin's sensitivity curve with Galactic confusion noise. + """ base_curve = sensitivity_curve_tianqin_analytical( - length, delta_f, low_freq_cutoff, - len_arm, acc_noise_level, oms_noise_level) + length, delta_f, low_freq_cutoff, len_arm, acc_noise_level, oms_noise_level + ) if duration < 0 or duration > 5: raise ValueError("Must between 0 and 5.") fseries_confusion = confusion_fit_tianqin( - length, delta_f, low_freq_cutoff, duration) - fseries = from_numpy_arrays(base_curve.sample_frequencies, - base_curve+fseries_confusion, - length, delta_f, low_freq_cutoff) + length, delta_f, low_freq_cutoff, duration + ) + fseries = from_numpy_arrays( + base_curve.sample_frequencies, + base_curve + fseries_confusion, + length, + delta_f, + low_freq_cutoff, + ) return fseries -def sensitivity_curve_taiji_confusion(length, delta_f, low_freq_cutoff, - len_arm=3e9, acc_noise_level=3e-15, - oms_noise_level=8e-12, duration=1.0): - """ The Taiji's sensitivity curve with Galactic confusion noise, +def sensitivity_curve_taiji_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm=3e9, + acc_noise_level=3e-15, + oms_noise_level=8e-12, + duration=1.0, +): + """ + The Taiji's sensitivity curve with Galactic confusion noise, averaged over sky and polarization angle. Parameters @@ -1558,27 +1809,38 @@ def sensitivity_curve_taiji_confusion(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The sky and polarization angle averaged Taiji's sensitivity curve with Galactic confusion noise. + """ base_curve = sensitivity_curve_taiji_analytical( - length, delta_f, low_freq_cutoff, - len_arm, acc_noise_level, oms_noise_level) + length, delta_f, low_freq_cutoff, len_arm, acc_noise_level, oms_noise_level + ) if duration < 0 or duration > 4: raise ValueError("Must between 0 and 4.") - fseries_confusion = confusion_fit_taiji( - length, delta_f, low_freq_cutoff, duration) - fseries = from_numpy_arrays(base_curve.sample_frequencies, - base_curve+fseries_confusion, - length, delta_f, low_freq_cutoff) + fseries_confusion = confusion_fit_taiji(length, delta_f, low_freq_cutoff, duration) + fseries = from_numpy_arrays( + base_curve.sample_frequencies, + base_curve + fseries_confusion, + length, + delta_f, + low_freq_cutoff, + ) return fseries -def sh_transformed_psd_lisa_tdi_XYZ(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, - base_model="semi", duration=1.0, - tdi=None): - """ The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA +def sh_transformed_psd_lisa_tdi_XYZ( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + base_model="semi", + duration=1.0, + tdi=None, +): + """ + The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA with Galactic confusion noise, transformed from LISA sensitivity curve. Parameters @@ -1607,32 +1869,43 @@ def sh_transformed_psd_lisa_tdi_XYZ(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA with Galactic confusion noise, transformed from LISA sensitivity curve. + Notes ----- Please see Eq.(7,41-43) in for more details. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) if str(tdi) in ["1.5", "2.0"]: response = averaged_response_lisa_tdi(fr, len_arm, tdi) else: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") - fseries_response = from_numpy_arrays(fr, np.array(response), - length, delta_f, low_freq_cutoff) - sh = sensitivity_curve_lisa_confusion(length, delta_f, low_freq_cutoff, - len_arm, acc_noise_level, - oms_noise_level, base_model, - duration) - psd = 2*sh.data * fseries_response.data - fseries = from_numpy_arrays(sh.sample_frequencies, psd, - length, delta_f, low_freq_cutoff) + fseries_response = from_numpy_arrays( + fr, np.array(response), length, delta_f, low_freq_cutoff + ) + sh = sensitivity_curve_lisa_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm, + acc_noise_level, + oms_noise_level, + base_model, + duration, + ) + psd = 2 * sh.data * fseries_response.data + fseries = from_numpy_arrays( + sh.sample_frequencies, psd, length, delta_f, low_freq_cutoff + ) return fseries -def semi_analytical_psd_lisa_confusion_noise(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, duration=1.0, - tdi=None): - """ The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA Galactic confusion noise, +def semi_analytical_psd_lisa_confusion_noise( + length, delta_f, low_freq_cutoff, len_arm=2.5e9, duration=1.0, tdi=None +): + """ + The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA Galactic confusion noise, no instrumental noise. Parameters @@ -1655,28 +1928,34 @@ def semi_analytical_psd_lisa_confusion_noise(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for LISA Galactic confusion noise, no instrumental noise. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) if str(tdi) in ["1.5", "2.0"]: response = averaged_response_lisa_tdi(fr, len_arm, tdi) else: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") - fseries_response = from_numpy_arrays(fr, np.array(response), - length, delta_f, low_freq_cutoff) - fseries_confusion = confusion_fit_lisa( - length, delta_f, low_freq_cutoff, duration) - psd_confusion = 2*fseries_confusion.data * fseries_response.data - fseries = from_numpy_arrays(fseries_confusion.sample_frequencies, - psd_confusion, length, delta_f, - low_freq_cutoff) + fseries_response = from_numpy_arrays( + fr, np.array(response), length, delta_f, low_freq_cutoff + ) + fseries_confusion = confusion_fit_lisa(length, delta_f, low_freq_cutoff, duration) + psd_confusion = 2 * fseries_confusion.data * fseries_response.data + fseries = from_numpy_arrays( + fseries_confusion.sample_frequencies, + psd_confusion, + length, + delta_f, + low_freq_cutoff, + ) return fseries -def analytical_psd_tianqin_confusion_noise(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - duration=1.0, tdi=None): - """ The TDI-1.5/2.0 PSD (X,Y,Z channel) for TianQin Galactic confusion +def analytical_psd_tianqin_confusion_noise( + length, delta_f, low_freq_cutoff, len_arm=np.sqrt(3) * 1e8, duration=1.0, tdi=None +): + """ + The TDI-1.5/2.0 PSD (X,Y,Z channel) for TianQin Galactic confusion noise, no instrumental noise. Parameters @@ -1699,28 +1978,36 @@ def analytical_psd_tianqin_confusion_noise(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for TianQin Galactic confusion noise, no instrumental noise. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) if str(tdi) in ["1.5", "2.0"]: response = averaged_response_tianqin_tdi(fr, len_arm, tdi) else: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") - fseries_response = from_numpy_arrays(fr, np.array(response), - length, delta_f, low_freq_cutoff) + fseries_response = from_numpy_arrays( + fr, np.array(response), length, delta_f, low_freq_cutoff + ) fseries_confusion = confusion_fit_tianqin( - length, delta_f, low_freq_cutoff, duration) - psd_confusion = 2*fseries_confusion.data * fseries_response.data - fseries = from_numpy_arrays(fseries_confusion.sample_frequencies, - psd_confusion, length, delta_f, - low_freq_cutoff) + length, delta_f, low_freq_cutoff, duration + ) + psd_confusion = 2 * fseries_confusion.data * fseries_response.data + fseries = from_numpy_arrays( + fseries_confusion.sample_frequencies, + psd_confusion, + length, + delta_f, + low_freq_cutoff, + ) return fseries -def analytical_psd_taiji_confusion_noise(length, delta_f, low_freq_cutoff, - len_arm=3e9, duration=1.0, - tdi=None): - """ The TDI-1.5/2.0 PSD (X,Y,Z channel) for Taiji Galactic confusion +def analytical_psd_taiji_confusion_noise( + length, delta_f, low_freq_cutoff, len_arm=3e9, duration=1.0, tdi=None +): + """ + The TDI-1.5/2.0 PSD (X,Y,Z channel) for Taiji Galactic confusion noise, no instrumental noise. Parameters @@ -1743,29 +2030,41 @@ def analytical_psd_taiji_confusion_noise(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (X,Y,Z channel) for Taiji Galactic confusion noise, no instrumental noise. + """ - fr = np.linspace(low_freq_cutoff, (length-1)*2*delta_f, length) + fr = np.linspace(low_freq_cutoff, (length - 1) * 2 * delta_f, length) if str(tdi) in ["1.5", "2.0"]: response = averaged_response_taiji_tdi(fr, len_arm, tdi) else: raise ValueError("The version of TDI, currently only for 1.5 or 2.0.") - fseries_response = from_numpy_arrays(fr, np.array(response), - length, delta_f, low_freq_cutoff) - fseries_confusion = confusion_fit_taiji( - length, delta_f, low_freq_cutoff, duration) - psd_confusion = 2*fseries_confusion.data * fseries_response.data - fseries = from_numpy_arrays(fseries_confusion.sample_frequencies, - psd_confusion, length, delta_f, - low_freq_cutoff) + fseries_response = from_numpy_arrays( + fr, np.array(response), length, delta_f, low_freq_cutoff + ) + fseries_confusion = confusion_fit_taiji(length, delta_f, low_freq_cutoff, duration) + psd_confusion = 2 * fseries_confusion.data * fseries_response.data + fseries = from_numpy_arrays( + fseries_confusion.sample_frequencies, + psd_confusion, + length, + delta_f, + low_freq_cutoff, + ) return fseries -def analytical_psd_lisa_tdi_AE_confusion(length, delta_f, low_freq_cutoff, - len_arm=2.5e9, acc_noise_level=3e-15, - oms_noise_level=15e-12, - duration=1.0, tdi=None): - """ The TDI-1.5/2.0 PSD (A,E channel) for LISA +def analytical_psd_lisa_tdi_AE_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm=2.5e9, + acc_noise_level=3e-15, + oms_noise_level=15e-12, + duration=1.0, + tdi=None, +): + """ + The TDI-1.5/2.0 PSD (A,E channel) for LISA with Galactic confusion noise. Parameters @@ -1792,13 +2091,14 @@ def analytical_psd_lisa_tdi_AE_confusion(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (A,E channel) for LISA with Galactic confusion noise. + """ - psd_AE = analytical_psd_lisa_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm, acc_noise_level, - oms_noise_level, tdi) + psd_AE = analytical_psd_lisa_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm, acc_noise_level, oms_noise_level, tdi + ) psd_X_confusion = semi_analytical_psd_lisa_confusion_noise( - length, delta_f, low_freq_cutoff, - len_arm, duration, tdi) + length, delta_f, low_freq_cutoff, len_arm, duration, tdi + ) # S_A = S_E = S_X - S_XY, confusion noise's contribution to # S_XY is -0.5 * psd_X_confusion, while for S_X is psd_X_confusion. # S_T = S_X + 2*S_XY, so S_T keeps the same. @@ -1807,12 +2107,18 @@ def analytical_psd_lisa_tdi_AE_confusion(length, delta_f, low_freq_cutoff, return fseries -def analytical_psd_tianqin_tdi_AE_confusion(length, delta_f, low_freq_cutoff, - len_arm=np.sqrt(3)*1e8, - acc_noise_level=1e-15, - oms_noise_level=1e-12, - duration=1.0, tdi=None): - """ The TDI-1.5/2.0 PSD (A,E channel) for TianQin +def analytical_psd_tianqin_tdi_AE_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm=np.sqrt(3) * 1e8, + acc_noise_level=1e-15, + oms_noise_level=1e-12, + duration=1.0, + tdi=None, +): + """ + The TDI-1.5/2.0 PSD (A,E channel) for TianQin with Galactic confusion noise. Parameters @@ -1839,14 +2145,14 @@ def analytical_psd_tianqin_tdi_AE_confusion(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (A,E channel) for TianQin with Galactic confusion noise. + """ - psd_AE = analytical_psd_tianqin_tdi_AE(length, delta_f, - low_freq_cutoff, - len_arm, acc_noise_level, - oms_noise_level, tdi) + psd_AE = analytical_psd_tianqin_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm, acc_noise_level, oms_noise_level, tdi + ) psd_X_confusion = analytical_psd_tianqin_confusion_noise( - length, delta_f, low_freq_cutoff, - len_arm, duration, tdi) + length, delta_f, low_freq_cutoff, len_arm, duration, tdi + ) # S_A = S_E = S_X - S_XY, confusion noise's contribution to # S_XY is -0.5 * psd_X_confusion, while for S_X is psd_X_confusion. # S_T = S_X + 2*S_XY, so S_T keeps the same. @@ -1855,11 +2161,18 @@ def analytical_psd_tianqin_tdi_AE_confusion(length, delta_f, low_freq_cutoff, return fseries -def analytical_psd_taiji_tdi_AE_confusion(length, delta_f, low_freq_cutoff, - len_arm=3e9, acc_noise_level=3e-15, - oms_noise_level=8e-12, - duration=1.0, tdi=None): - """ The TDI-1.5/2.0 PSD (A,E channel) for Taiji +def analytical_psd_taiji_tdi_AE_confusion( + length, + delta_f, + low_freq_cutoff, + len_arm=3e9, + acc_noise_level=3e-15, + oms_noise_level=8e-12, + duration=1.0, + tdi=None, +): + """ + The TDI-1.5/2.0 PSD (A,E channel) for Taiji with Galactic confusion noise. Parameters @@ -1886,13 +2199,14 @@ def analytical_psd_taiji_tdi_AE_confusion(length, delta_f, low_freq_cutoff, fseries : FrequencySeries The TDI-1.5/2.0 PSD (A,E channel) for Taiji with Galactic confusion noise. + """ - psd_AE = analytical_psd_taiji_tdi_AE(length, delta_f, low_freq_cutoff, - len_arm, acc_noise_level, - oms_noise_level, tdi) + psd_AE = analytical_psd_taiji_tdi_AE( + length, delta_f, low_freq_cutoff, len_arm, acc_noise_level, oms_noise_level, tdi + ) psd_X_confusion = analytical_psd_taiji_confusion_noise( - length, delta_f, low_freq_cutoff, - len_arm, duration, tdi) + length, delta_f, low_freq_cutoff, len_arm, duration, tdi + ) # S_A = S_E = S_X - S_XY, confusion noise's contribution to # S_XY is -0.5 * psd_X_confusion, while for S_X is psd_X_confusion. # S_T = S_X + 2*S_XY, so S_T keeps the same. diff --git a/pycbc/psd/estimate.py b/pycbc/psd/estimate.py index a667bacc371..3823642079e 100644 --- a/pycbc/psd/estimate.py +++ b/pycbc/psd/estimate.py @@ -13,13 +13,19 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Utilites to estimate PSDs from data. -""" +"""Utilites to estimate PSDs from data.""" import numpy -from pycbc.types import Array, FrequencySeries, TimeSeries, zeros -from pycbc.types import real_same_precision_as, complex_same_precision_as + from pycbc.fft import fft, ifft +from pycbc.types import ( + Array, + FrequencySeries, + TimeSeries, + complex_same_precision_as, + real_same_precision_as, + zeros, +) # Change to True in front-end if you want this function to use caching # This is a mostly-hidden optimization option that most users will not want @@ -32,8 +38,10 @@ WELCH_UNIQUE_ID = 438716587 INVSPECTRUNC_UNIQUE_ID = 100257896 + def median_bias(n): - """Calculate the bias of the median average PSD computed from `n` segments. + """ + Calculate the bias of the median average PSD computed from `n` segments. Parameters ---------- @@ -53,19 +61,29 @@ def median_bias(n): Notes ----- See arXiv:gr-qc/0509116 appendix B for details. + """ if type(n) is not int or n <= 0: - raise ValueError('n must be a positive integer') + raise ValueError("n must be a positive integer") if n >= 1000: return numpy.log(2) ans = 1 for i in range(1, (n - 1) // 2 + 1): - ans += 1.0 / (2*i + 1) - 1.0 / (2*i) + ans += 1.0 / (2 * i + 1) - 1.0 / (2 * i) return ans -def welch(timeseries, seg_len=4096, seg_stride=2048, window='hann', - avg_method='median', num_segments=None, require_exact_data_fit=False): - """PSD estimator based on Welch's method. + +def welch( + timeseries, + seg_len=4096, + seg_stride=2048, + window="hann", + avg_method="median", + num_segments=None, + require_exact_data_fit=False, +): + """ + PSD estimator based on Welch's method. Parameters ---------- @@ -96,27 +114,30 @@ def welch(timeseries, seg_len=4096, seg_stride=2048, window='hann', Notes ----- See arXiv:gr-qc/0509116 for details. + """ from pycbc.strain.strain import execute_cached_fft - window_map = { - 'hann': numpy.hanning - } + window_map = {"hann": numpy.hanning} # sanity checks if isinstance(window, numpy.ndarray) and window.size != seg_len: - raise ValueError('Invalid window: incorrect window length') + raise ValueError("Invalid window: incorrect window length") if not isinstance(window, numpy.ndarray) and window not in window_map: - raise ValueError('Invalid window: unknown window {!r}'.format(window)) - if avg_method not in ('mean', 'median', 'median-mean'): - raise ValueError('Invalid averaging method') - if type(seg_len) is not int or type(seg_stride) is not int \ - or seg_len <= 0 or seg_stride <= 0: - raise ValueError('Segment length and stride must be positive integers') - - if timeseries.precision == 'single': + raise ValueError(f"Invalid window: unknown window {window!r}") + if avg_method not in ("mean", "median", "median-mean"): + raise ValueError("Invalid averaging method") + if ( + type(seg_len) is not int + or type(seg_stride) is not int + or seg_len <= 0 + or seg_stride <= 0 + ): + raise ValueError("Segment length and stride must be positive integers") + + if timeseries.precision == "single": fs_dtype = numpy.complex64 - elif timeseries.precision == 'double': + elif timeseries.precision == "double": fs_dtype = numpy.complex128 num_samples = len(timeseries) @@ -141,19 +162,19 @@ def welch(timeseries, seg_len=4096, seg_stride=2048, window='hann', timeseries = timeseries[start:end] num_samples = len(timeseries) if data_len > num_samples: - err_msg = "I was asked to estimate a PSD on %d " %(data_len) + err_msg = "I was asked to estimate a PSD on %d " % (data_len) err_msg += "data samples. However the data provided only contains " - err_msg += "%d data samples." %(num_samples) + err_msg += "%d data samples." % (num_samples) if num_samples != (num_segments - 1) * seg_stride + seg_len: - raise ValueError('Incorrect choice of segmentation parameters') + raise ValueError("Incorrect choice of segmentation parameters") if not isinstance(window, numpy.ndarray): window = window_map[window](seg_len) w = Array(window.astype(timeseries.dtype)) # calculate psd of each segment - delta_f = 1. / timeseries.delta_t / seg_len + delta_f = 1.0 / timeseries.delta_t / seg_len if not USE_CACHING_FOR_WELCH_FFTS: segment_tilde = FrequencySeries( numpy.zeros(int(seg_len / 2 + 1)), @@ -170,11 +191,10 @@ def welch(timeseries, seg_len=4096, seg_stride=2048, window='hann', if not USE_CACHING_FOR_WELCH_FFTS: fft(segment * w, segment_tilde) else: - segment_tilde = execute_cached_fft(segment * w, - uid=WELCH_UNIQUE_ID) + segment_tilde = execute_cached_fft(segment * w, uid=WELCH_UNIQUE_ID) seg_psd = abs(segment_tilde * segment_tilde.conj()).numpy() - #halve the DC and Nyquist components to be consistent with TO10095 + # halve the DC and Nyquist components to be consistent with TO10095 seg_psd[0] /= 2 seg_psd[-1] /= 2 @@ -182,29 +202,35 @@ def welch(timeseries, seg_len=4096, seg_stride=2048, window='hann', segment_psds = numpy.array(segment_psds) - if avg_method == 'mean': + if avg_method == "mean": psd = numpy.mean(segment_psds, axis=0) - elif avg_method == 'median': + elif avg_method == "median": psd = numpy.median(segment_psds, axis=0) / median_bias(num_segments) - elif avg_method == 'median-mean': + elif avg_method == "median-mean": odd_psds = segment_psds[::2] even_psds = segment_psds[1::2] - odd_median = numpy.median(odd_psds, axis=0) / \ - median_bias(len(odd_psds)) - even_median = numpy.median(even_psds, axis=0) / \ - median_bias(len(even_psds)) + odd_median = numpy.median(odd_psds, axis=0) / median_bias(len(odd_psds)) + even_median = numpy.median(even_psds, axis=0) / median_bias(len(even_psds)) psd = (odd_median + even_median) / 2 w = w.numpy() - psd *= 2 * delta_f * seg_len / (w*w).sum() + psd *= 2 * delta_f * seg_len / (w * w).sum() + + return FrequencySeries( + psd, delta_f=delta_f, dtype=timeseries.dtype, epoch=timeseries.start_time + ) - return FrequencySeries(psd, delta_f=delta_f, dtype=timeseries.dtype, - epoch=timeseries.start_time) -def inverse_spectrum_truncation(psd, max_filter_len, which_spectrum='invasd', - low_frequency_cutoff=None, - low_frequency_fill_value=0., trunc_method=None): - """Modify a PSD such that the impulse response associated with its inverse +def inverse_spectrum_truncation( + psd, + max_filter_len, + which_spectrum="invasd", + low_frequency_cutoff=None, + low_frequency_fill_value=0.0, + trunc_method=None, +): + """ + Modify a PSD such that the impulse response associated with its inverse square root is no longer than `max_filter_len` time samples. In practice this corresponds to a coarse graining or smoothing of the PSD. @@ -218,7 +244,7 @@ def inverse_spectrum_truncation(psd, max_filter_len, which_spectrum='invasd', Which spectrum to truncate. If 'invasd' (default), apply truncation to the inverse ASD. If 'invpsd', apply to the inverse PSD. low_frequency_cutoff : {None, int} - Frequencies below `low_frequency_cutoff` are set to value specified by + Frequencies below `low_frequency_cutoff` are set to value specified by `low_frequency_fill_value`. low_frequency_fill_value : {float, 'fmin'} Value to set PSD to at frequencies below `low_frequency_cutoff`. @@ -227,7 +253,7 @@ def inverse_spectrum_truncation(psd, max_filter_len, which_spectrum='invasd', trunc_method : {None, 'hann'} Function used for truncating the time-domain filter. None produces a hard truncation at `max_filter_len`. - + Returns ------- @@ -242,62 +268,64 @@ def inverse_spectrum_truncation(psd, max_filter_len, which_spectrum='invasd', Notes ----- See arXiv:gr-qc/0509116 for details. + """ from pycbc.strain.strain import execute_cached_fft, execute_cached_ifft # sanity checks if type(max_filter_len) is not int or max_filter_len <= 0: - raise ValueError('max_filter_len must be a positive integer') - if low_frequency_cutoff is not None and \ - (low_frequency_cutoff < 0. or - low_frequency_cutoff > psd.sample_frequencies[-1]): - raise ValueError('low_frequency_cutoff must be within the bandwidth of ' - 'the PSD') + raise ValueError("max_filter_len must be a positive integer") + if low_frequency_cutoff is not None and ( + low_frequency_cutoff < 0.0 or low_frequency_cutoff > psd.sample_frequencies[-1] + ): + raise ValueError("low_frequency_cutoff must be within the bandwidth of the PSD") - N = (len(psd)-1)*2 + N = (len(psd) - 1) * 2 - inv_spectrum = FrequencySeries(zeros(len(psd)), delta_f=psd.delta_f, \ - dtype=complex_same_precision_as(psd)) + inv_spectrum = FrequencySeries( + zeros(len(psd)), delta_f=psd.delta_f, dtype=complex_same_precision_as(psd) + ) kmin = 1 if low_frequency_cutoff: kmin = int(low_frequency_cutoff / psd.delta_f) - + # set values below low frequency cutoff - if low_frequency_fill_value != 0.: - if low_frequency_fill_value == 'fmin': - low_frequency_fill_value = 1./psd[kmin] + if low_frequency_fill_value != 0.0: + if low_frequency_fill_value == "fmin": + low_frequency_fill_value = 1.0 / psd[kmin] inv_spectrum[:kmin] = float(low_frequency_fill_value) - inv_spectrum[kmin:N//2] = (1.0 / psd[kmin:N//2]) + inv_spectrum[kmin : N // 2] = 1.0 / psd[kmin : N // 2] # if truncating asd, take sqrt - if which_spectrum == 'invasd': - inv_spectrum[:N//2] = inv_spectrum[:N//2]**0.5 - elif which_spectrum != 'invpsd': - raise ValueError(f'Invalid which_spectrum input {which_spectrum}; ' - f'input must be either "invpsd" or "invasd"') + if which_spectrum == "invasd": + inv_spectrum[: N // 2] = inv_spectrum[: N // 2] ** 0.5 + elif which_spectrum != "invpsd": + raise ValueError( + f"Invalid which_spectrum input {which_spectrum}; " + f'input must be either "invpsd" or "invasd"' + ) if not USE_CACHING_FOR_INV_SPEC_TRUNC: q = TimeSeries( - numpy.zeros(N), - delta_t=(N / psd.delta_f), - dtype=real_same_precision_as(psd) + numpy.zeros(N), delta_t=(N / psd.delta_f), dtype=real_same_precision_as(psd) ) ifft(inv_spectrum, q) else: - q = execute_cached_ifft(inv_spectrum, copy_output=False, - uid=INVSPECTRUNC_UNIQUE_ID) + q = execute_cached_ifft( + inv_spectrum, copy_output=False, uid=INVSPECTRUNC_UNIQUE_ID + ) trunc_start = max_filter_len // 2 trunc_end = N - max_filter_len // 2 if trunc_end < trunc_start: - raise ValueError('Invalid value in inverse_spectrum_truncation') + raise ValueError("Invalid value in inverse_spectrum_truncation") - if trunc_method == 'hann': + if trunc_method == "hann": trunc_window = Array(numpy.hanning(max_filter_len), dtype=q.dtype) q[0:trunc_start] *= trunc_window[-trunc_start:] - q[trunc_end:N] *= trunc_window[0:max_filter_len//2] + q[trunc_end:N] *= trunc_window[0 : max_filter_len // 2] if trunc_start < trunc_end: q[trunc_start:trunc_end] = 0 @@ -305,20 +333,21 @@ def inverse_spectrum_truncation(psd, max_filter_len, which_spectrum='invasd', psd_trunc = FrequencySeries( numpy.zeros(len(psd)), delta_f=psd.delta_f, - dtype=complex_same_precision_as(psd) + dtype=complex_same_precision_as(psd), ) fft(q, psd_trunc) else: - psd_trunc = execute_cached_fft(q, copy_output=False, - uid=INVSPECTRUNC_UNIQUE_ID) - if which_spectrum == 'invasd': + psd_trunc = execute_cached_fft(q, copy_output=False, uid=INVSPECTRUNC_UNIQUE_ID) + if which_spectrum == "invasd": psd_trunc *= psd_trunc.conj() - psd_out = 1. / abs(psd_trunc) + psd_out = 1.0 / abs(psd_trunc) return psd_out + def interpolate(series, delta_f, length=None): - """Return a new PSD that has been interpolated to the desired delta_f. + """ + Return a new PSD that has been interpolated to the desired delta_f. Parameters ---------- @@ -336,18 +365,24 @@ def interpolate(series, delta_f, length=None): ------- interpolated series : FrequencySeries A new FrequencySeries that has been interpolated. + """ if length is None: - new_n = (len(series)-1) * series.delta_f / delta_f + 1 + new_n = (len(series) - 1) * series.delta_f / delta_f + 1 else: new_n = length samples = numpy.arange(0, numpy.rint(new_n)) * delta_f - interpolated_series = numpy.interp(samples, series.sample_frequencies.numpy(), series.numpy()) - return FrequencySeries(interpolated_series, epoch=series.epoch, - delta_f=delta_f, dtype=series.dtype) + interpolated_series = numpy.interp( + samples, series.sample_frequencies.numpy(), series.numpy() + ) + return FrequencySeries( + interpolated_series, epoch=series.epoch, delta_f=delta_f, dtype=series.dtype + ) + def bandlimited_interpolate(series, delta_f): - """Return a new PSD that has been interpolated to the desired delta_f. + """ + Return a new PSD that has been interpolated to the desired delta_f. Parameters ---------- @@ -360,8 +395,11 @@ def bandlimited_interpolate(series, delta_f): ------- interpolated series : FrequencySeries A new FrequencySeries that has been interpolated. + """ - series = FrequencySeries(series, dtype=complex_same_precision_as(series), delta_f=series.delta_f) + series = FrequencySeries( + series, dtype=complex_same_precision_as(series), delta_f=series.delta_f + ) N = (len(series) - 1) * 2 delta_t = 1.0 / series.delta_f / N @@ -369,15 +407,20 @@ def bandlimited_interpolate(series, delta_f): new_N = int(1.0 / (delta_t * delta_f)) new_n = new_N // 2 + 1 - series_in_time = TimeSeries(zeros(N), dtype=real_same_precision_as(series), delta_t=delta_t) + series_in_time = TimeSeries( + zeros(N), dtype=real_same_precision_as(series), delta_t=delta_t + ) ifft(series, series_in_time) - padded_series_in_time = TimeSeries(zeros(new_N), dtype=series_in_time.dtype, delta_t=delta_t) - padded_series_in_time[0:N//2] = series_in_time[0:N//2] - padded_series_in_time[new_N-N//2:new_N] = series_in_time[N//2:N] + padded_series_in_time = TimeSeries( + zeros(new_N), dtype=series_in_time.dtype, delta_t=delta_t + ) + padded_series_in_time[0 : N // 2] = series_in_time[0 : N // 2] + padded_series_in_time[new_N - N // 2 : new_N] = series_in_time[N // 2 : N] - interpolated_series = FrequencySeries(zeros(new_n), dtype=series.dtype, delta_f=delta_f) + interpolated_series = FrequencySeries( + zeros(new_n), dtype=series.dtype, delta_f=delta_f + ) fft(padded_series_in_time, interpolated_series) return interpolated_series - diff --git a/pycbc/psd/read.py b/pycbc/psd/read.py index 9417b0ee2a8..a917aee669b 100644 --- a/pycbc/psd/read.py +++ b/pycbc/psd/read.py @@ -13,18 +13,21 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Utilities to read PSDs from files. -""" +"""Utilities to read PSDs from files.""" import logging + import numpy import scipy.interpolate + from pycbc.types import FrequencySeries -logger = logging.getLogger('pycbc.psd.read') +logger = logging.getLogger("pycbc.psd.read") + def from_numpy_arrays(freq_data, noise_data, length, delta_f, low_freq_cutoff): - """Interpolate n PSD (as two 1-dimensional arrays of frequency and data) + """ + Interpolate n PSD (as two 1-dimensional arrays of frequency and data) to the desired length, delta_f and low frequency cutoff. Parameters @@ -44,49 +47,59 @@ def from_numpy_arrays(freq_data, noise_data, length, delta_f, low_freq_cutoff): ------- psd : FrequencySeries The generated frequency series. + """ # Only include points above the low frequency cutoff if freq_data[0] > low_freq_cutoff: raise ValueError( - f'Lowest frequency in input PSD data ({freq_data[0]} Hz) is ' - f'higher than requested low-frequency cutoff ({low_freq_cutoff} Hz)' + f"Lowest frequency in input PSD data ({freq_data[0]} Hz) is " + f"higher than requested low-frequency cutoff ({low_freq_cutoff} Hz)" ) kmin = int(low_freq_cutoff / delta_f) flow = kmin * delta_f - data_start = (0 if freq_data[0]==low_freq_cutoff else numpy.searchsorted(freq_data, flow) - 1) + data_start = ( + 0 + if freq_data[0] == low_freq_cutoff + else numpy.searchsorted(freq_data, flow) - 1 + ) data_start = max(0, data_start) # If the cutoff is exactly in the file, start there - if freq_data[data_start+1] == low_freq_cutoff: + if freq_data[data_start + 1] == low_freq_cutoff: data_start += 1 freq_data = freq_data[data_start:] noise_data = noise_data[data_start:] if (length - 1) * delta_f > freq_data[-1]: - logger.warning('Requested number of samples exceeds the highest ' - 'available frequency in the input data, ' - 'will use max available frequency instead. ' - '(requested %f Hz, available %f Hz)', - (length - 1) * delta_f, freq_data[-1]) - length = int(freq_data[-1]/delta_f + 1) + logger.warning( + "Requested number of samples exceeds the highest " + "available frequency in the input data, " + "will use max available frequency instead. " + "(requested %f Hz, available %f Hz)", + (length - 1) * delta_f, + freq_data[-1], + ) + length = int(freq_data[-1] / delta_f + 1) flog = numpy.log(freq_data) slog = numpy.log(noise_data) psd_interp = scipy.interpolate.interp1d( - flog, slog, fill_value=(slog[0], slog[-1]), bounds_error=False) + flog, slog, fill_value=(slog[0], slog[-1]), bounds_error=False + ) psd = numpy.zeros(length, dtype=numpy.float64) vals = numpy.log(numpy.arange(kmin, length) * delta_f) - psd[kmin:] = numpy.exp(psd_interp(vals)) + psd[kmin:] = numpy.exp(psd_interp(vals)) return FrequencySeries(psd, delta_f=delta_f) def from_txt(filename, length, delta_f, low_freq_cutoff, is_asd_file=True): - """Read an ASCII file containing one-sided ASD or PSD data and generate + """ + Read an ASCII file containing one-sided ASD or PSD data and generate a frequency series with the corresponding PSD. The ASD or PSD data is interpolated in order to match the desired frequency resolution. @@ -117,23 +130,25 @@ def from_txt(filename, length, delta_f, low_freq_cutoff, is_asd_file=True): ValueError If the ASCII file contains negative, infinite or NaN frequencies or amplitude densities. + """ file_data = numpy.loadtxt(filename) - if (file_data < 0).any() or \ - numpy.logical_not(numpy.isfinite(file_data)).any(): - raise ValueError('Invalid data in ' + filename) + if (file_data < 0).any() or numpy.logical_not(numpy.isfinite(file_data)).any(): + raise ValueError("Invalid data in " + filename) freq_data = file_data[:, 0] noise_data = file_data[:, 1] if is_asd_file: - noise_data = noise_data ** 2 + noise_data = noise_data**2 + + return from_numpy_arrays(freq_data, noise_data, length, delta_f, low_freq_cutoff) - return from_numpy_arrays(freq_data, noise_data, length, delta_f, - low_freq_cutoff) -def from_xml(filename, length, delta_f, low_freq_cutoff, ifo_string=None, - root_name='psd'): - """Read a LIGOLW XML file containing one-sided PSD data and generate +def from_xml( + filename, length, delta_f, low_freq_cutoff, ifo_string=None, root_name="psd" +): + """ + Read a LIGOLW XML file containing one-sided PSD data and generate a frequency series with the corresponding PSD. The data is interpolated in order to match the desired frequency resolution. @@ -159,14 +174,16 @@ def from_xml(filename, length, delta_f, low_freq_cutoff, ifo_string=None, ------- psd : FrequencySeries The generated frequency series. + """ import lal.series from igwn_ligolw import utils as ligolw_utils - with open(filename, 'rb') as fp: + with open(filename, "rb") as fp: ct_handler = lal.series.PSDContentHandler - xml_doc = ligolw_utils.load_fileobj(fp, compress='auto', - contenthandler=ct_handler) + xml_doc = ligolw_utils.load_fileobj( + fp, compress="auto", contenthandler=ct_handler + ) psd_dict = lal.series.read_psd_xmldoc(xml_doc, root_name=root_name) if ifo_string is not None: @@ -184,9 +201,9 @@ def from_xml(filename, length, delta_f, low_freq_cutoff, ifo_string=None, if f0 != 0.0: logger.warning( "XML PSD has non-zero start frequency f0=%.4f Hz; " - "applying offset to frequency axis.", f0 + "applying offset to frequency axis.", + f0, ) freq_data = f0 + numpy.arange(len(noise_data)) * psd_freq_series.deltaF - return from_numpy_arrays(freq_data, noise_data, length, delta_f, - low_freq_cutoff) + return from_numpy_arrays(freq_data, noise_data, length, delta_f, low_freq_cutoff) diff --git a/pycbc/psd/variation.py b/pycbc/psd/variation.py index 4f8dfb07c20..f923dee20fe 100644 --- a/pycbc/psd/variation.py +++ b/pycbc/psd/variation.py @@ -1,8 +1,8 @@ -""" PSD Variation """ +"""PSD Variation""" import numpy -from numpy.fft import rfft, irfft import scipy.signal as sig +from numpy.fft import irfft, rfft from scipy.interpolate import interp1d import pycbc.psd @@ -10,7 +10,8 @@ def create_full_filt(freqs, filt, plong, srate, psd_duration): - """Create a filter to convolve with strain data to find PSD variation. + """ + Create a filter to convolve with strain data to find PSD variation. Parameters ---------- @@ -29,26 +30,28 @@ def create_full_filt(freqs, filt, plong, srate, psd_duration): ------- full_filt : numpy.ndarray The full filter used to calculate PSD variation. - """ + """ # Make the weighting filter - bandpass, which weight by f^-7/6, # and whiten. The normalization is chosen so that the variance # will be one if this filter is applied to white noise which # already has a variance of one. - fweight = freqs ** (-7./6.) * filt / numpy.sqrt(plong) - fweight[0] = 0. - norm = (sum(abs(fweight) ** 2) / (len(fweight) - 1.)) ** -0.5 + fweight = freqs ** (-7.0 / 6.0) * filt / numpy.sqrt(plong) + fweight[0] = 0.0 + norm = (sum(abs(fweight) ** 2) / (len(fweight) - 1.0)) ** -0.5 fweight = norm * fweight - fwhiten = numpy.sqrt(2. / srate) / numpy.sqrt(plong) - fwhiten[0] = 0. + fwhiten = numpy.sqrt(2.0 / srate) / numpy.sqrt(plong) + fwhiten[0] = 0.0 full_filt = sig.windows.hann(int(psd_duration * srate)) * numpy.roll( - irfft(fwhiten * fweight), int(psd_duration / 2) * srate) + irfft(fwhiten * fweight), int(psd_duration / 2) * srate + ) return full_filt def mean_square(data, delta_t, srate, short_stride, stride): - """ Calculate mean square of given time series once per stride + """ + Calculate mean square of given time series once per stride First of all this function calculate the mean square of given time series once per short_stride. This is used to find and remove @@ -75,31 +78,40 @@ def mean_square(data, delta_t, srate, short_stride, stride): ------- m_s: List Mean square of given time series - """ + """ # Calculate mean square of data once per short stride and replace # outliers - short_ms = numpy.mean(data.reshape(-1, int(srate * short_stride)) ** 2, - axis=1) + short_ms = numpy.mean(data.reshape(-1, int(srate * short_stride)) ** 2, axis=1) # Define an array of averages that is used to substitute outliers ave = 0.5 * (short_ms[2:] + short_ms[:-2]) - outliers = short_ms[1:-1] > (2. * ave) + outliers = short_ms[1:-1] > (2.0 * ave) short_ms[1:-1][outliers] = ave[outliers] # Calculate mean square of data every step within a window equal to # stride seconds m_s = [] - inv_time = int(1. / short_stride) + inv_time = int(1.0 / short_stride) for index in range(int(delta_t - stride + 1)): - m_s.append(numpy.mean(short_ms[inv_time * index:inv_time * - int(index+stride)])) + m_s.append( + numpy.mean(short_ms[inv_time * index : inv_time * int(index + stride)]) + ) return m_s -def calc_filt_psd_variation(strain, segment, short_segment, psd_long_segment, - psd_duration, psd_stride, psd_avg_method, low_freq, - high_freq): - """ Calculates time series of PSD variability +def calc_filt_psd_variation( + strain, + segment, + short_segment, + psd_long_segment, + psd_duration, + psd_stride, + psd_avg_method, + low_freq, + high_freq, +): + """ + Calculates time series of PSD variability This function first splits the segment up into 512 second chunks. It then calculates the PSD over this 512 second. The PSD is used to @@ -140,11 +152,12 @@ def calc_filt_psd_variation(strain, segment, short_segment, psd_long_segment, ------- psd_var : TimeSeries Time series of the variability in the PSD estimation + """ # Calculate strain precision - if strain.precision == 'single': + if strain.precision == "single": fs_dtype = numpy.float32 - elif strain.precision == 'double': + elif strain.precision == "double": fs_dtype = numpy.float64 # Convert start and end times immediately to floats @@ -158,13 +171,14 @@ def calc_filt_psd_variation(strain, segment, short_segment, psd_long_segment, strain_crop = 8.0 # Find the times of the long segments - times_long = numpy.arange(start_time, end_time, - psd_long_segment - 2 * strain_crop - - segment + step) + times_long = numpy.arange( + start_time, end_time, psd_long_segment - 2 * strain_crop - segment + step + ) # Create a bandpass filter between low_freq and high_freq - filt = sig.firwin(4 * srate, [low_freq, high_freq], pass_zero=False, - window='hann', fs=srate) + filt = sig.firwin( + 4 * srate, [low_freq, high_freq], pass_zero=False, window="hann", fs=srate + ) filt.resize(int(psd_duration * srate)) # Fourier transform the filter and take the absolute value to get # rid of the phase. @@ -179,37 +193,42 @@ def calc_filt_psd_variation(strain, segment, short_segment, psd_long_segment, astrain, seg_len=int(psd_duration * strain.sample_rate), seg_stride=int(psd_stride * strain.sample_rate), - avg_method=psd_avg_method) + avg_method=psd_avg_method, + ) else: astrain = strain.time_slice(tlong, end_time) plong = pycbc.psd.welch( - strain.time_slice(end_time - psd_long_segment, - end_time), - seg_len=int(psd_duration * strain.sample_rate), - seg_stride=int(psd_stride * strain.sample_rate), - avg_method=psd_avg_method) + strain.time_slice(end_time - psd_long_segment, end_time), + seg_len=int(psd_duration * strain.sample_rate), + seg_stride=int(psd_stride * strain.sample_rate), + avg_method=psd_avg_method, + ) astrain = astrain.numpy() freqs = numpy.array(plong.sample_frequencies, dtype=fs_dtype) plong = plong.numpy() full_filt = create_full_filt(freqs, filt, plong, srate, psd_duration) # Convolve the filter with long segment of data - wstrain = sig.fftconvolve(astrain, full_filt, mode='same') - wstrain = wstrain[int(strain_crop * srate):-int(strain_crop * srate)] + wstrain = sig.fftconvolve(astrain, full_filt, mode="same") + wstrain = wstrain[int(strain_crop * srate) : -int(strain_crop * srate)] # compute the mean square of the chunk of data delta_t = len(wstrain) * strain.delta_t variation = mean_square(wstrain, delta_t, srate, short_segment, segment) psd_var_list.append(numpy.array(variation, dtype=wstrain.dtype)) # Package up the time series to return - psd_var = TimeSeries(numpy.concatenate(psd_var_list), delta_t=step, - epoch=start_time + strain_crop + segment) + psd_var = TimeSeries( + numpy.concatenate(psd_var_list), + delta_t=step, + epoch=start_time + strain_crop + segment, + ) return psd_var def find_trigger_value(psd_var, idx, start, sample_rate): - """ Find the PSD variation value at a particular time with the filter + """ + Find the PSD variation value at a particular time with the filter method. If the time is outside the timeseries bound, 1. is given. Parameters @@ -227,27 +246,27 @@ def find_trigger_value(psd_var, idx, start, sample_rate): ------- vals : Array PSD variation value at a particular time + """ # Find gps time of the trigger time = start + idx / sample_rate # Extract the PSD variation at trigger time through linear # interpolation - if not hasattr(psd_var, 'cached_psd_var_interpolant'): - psd_var.cached_psd_var_interpolant = \ - interp1d(psd_var.sample_times.numpy(), - psd_var.numpy(), - fill_value=1.0, - bounds_error=False) + if not hasattr(psd_var, "cached_psd_var_interpolant"): + psd_var.cached_psd_var_interpolant = interp1d( + psd_var.sample_times.numpy(), + psd_var.numpy(), + fill_value=1.0, + bounds_error=False, + ) vals = psd_var.cached_psd_var_interpolant(time) return vals -def live_create_filter(psd_estimated, - psd_duration, - sample_rate, - low_freq=20, - high_freq=480): +def live_create_filter( + psd_estimated, psd_duration, sample_rate, low_freq=20, high_freq=480 +): """ Create a filter to be used in the calculation of the psd variation for the PyCBC Live search. This filter combines a bandpass between a lower and @@ -278,13 +297,14 @@ def live_create_filter(psd_estimated, find the psd variation value. """ - # Create a bandpass filter between low_freq and high_freq once - filt = sig.firwin(4 * sample_rate, - [low_freq, high_freq], - pass_zero=False, - window='hann', - fs=sample_rate) + filt = sig.firwin( + 4 * sample_rate, + [low_freq, high_freq], + pass_zero=False, + window="hann", + fs=sample_rate, + ) filt.resize(int(psd_duration * sample_rate)) # Fourier transform the filter and take the absolute value to get @@ -299,11 +319,9 @@ def live_create_filter(psd_estimated, return full_filt -def live_calc_psd_variation(strain, - full_filt, - increment, - data_trim=2.0, - short_stride=0.25): +def live_calc_psd_variation( + strain, full_filt, increment, data_trim=2.0, short_stride=0.25 +): """ Calculate the psd variation in the PyCBC Live search. @@ -338,22 +356,24 @@ def live_calc_psd_variation(strain, sample_rate = int(strain.sample_rate) # Grab the last increments worth of data, plus padding for edge effects. - astrain = strain.time_slice(strain.end_time - increment - (data_trim * 3), - strain.end_time) + astrain = strain.time_slice( + strain.end_time - increment - (data_trim * 3), strain.end_time + ) # Convolve the data and the filter to produce the PSD variation timeseries, # then trim the beginning and end of the data to prevent edge effects. - wstrain = sig.fftconvolve(astrain, full_filt, mode='same') - wstrain = wstrain[int(data_trim * sample_rate):-int(data_trim * sample_rate)] + wstrain = sig.fftconvolve(astrain, full_filt, mode="same") + wstrain = wstrain[int(data_trim * sample_rate) : -int(data_trim * sample_rate)] # Create a PSD variation array by taking the mean square of the PSD # variation timeseries every short_stride short_ms = numpy.mean( - wstrain.reshape(-1, int(sample_rate * short_stride)) ** 2, axis=1) + wstrain.reshape(-1, int(sample_rate * short_stride)) ** 2, axis=1 + ) # Define an array of averages that is used to substitute outliers ave = 0.5 * (short_ms[2:] + short_ms[:-2]) - outliers = short_ms[1:-1] > (2. * ave) + outliers = short_ms[1:-1] > (2.0 * ave) short_ms[1:-1][outliers] = ave[outliers] # Calculate the PSD variation every second by a moving window average @@ -366,15 +386,14 @@ def live_calc_psd_variation(strain, m_s.append(numpy.mean(short_ms[start:end])) m_s = numpy.array(m_s, dtype=wstrain.dtype) - psd_var = TimeSeries(m_s, - delta_t=1.0, - epoch=strain.end_time - increment - (data_trim * 2)) + psd_var = TimeSeries( + m_s, delta_t=1.0, epoch=strain.end_time - increment - (data_trim * 2) + ) return psd_var -def live_find_var_value(triggers, - psd_var_timeseries): +def live_find_var_value(triggers, psd_var_timeseries): """ Extract the PSD variation values at trigger times by linear interpolation. @@ -390,14 +409,16 @@ def live_find_var_value(triggers, ------- psd_var_vals : numpy.ndarray Array of interpolated PSD variation values at trigger times. - """ + """ # Create the interpolator - interpolator = interp1d(psd_var_timeseries.sample_times.numpy(), - psd_var_timeseries.numpy(), - fill_value=1.0, - bounds_error=False) + interpolator = interp1d( + psd_var_timeseries.sample_times.numpy(), + psd_var_timeseries.numpy(), + fill_value=1.0, + bounds_error=False, + ) # Evaluate at the trigger times - psd_var_vals = interpolator(triggers['end_time']) + psd_var_vals = interpolator(triggers["end_time"]) return psd_var_vals diff --git a/pycbc/rate.py b/pycbc/rate.py index bf526506637..97883152ae7 100644 --- a/pycbc/rate.py +++ b/pycbc/rate.py @@ -1,21 +1,22 @@ -import numpy import bisect import logging +import numpy + from . import bin_utils -logger = logging.getLogger('pycbc.rate') +logger = logging.getLogger("pycbc.rate") def integral_element(mu, pdf): - ''' + """ Returns an array of elements of the integrand dP = p(mu) dmu for a density p(mu) defined at sample values mu ; samples need not be equally spaced. Uses a simple trapezium rule. Number of dP elements is 1 - (number of mu samples). - ''' + """ dmu = mu[1:] - mu[:-1] - bin_mean = (pdf[1:] + pdf[:-1]) / 2. + bin_mean = (pdf[1:] + pdf[:-1]) / 2.0 return dmu * bin_mean @@ -26,14 +27,18 @@ def normalize_pdf(mu, pofmu): arrays or lists of the same length. """ if min(pofmu) < 0: - raise ValueError("Probabilities cannot be negative, don't ask me to " - "normalize a function with negative values!") + raise ValueError( + "Probabilities cannot be negative, don't ask me to " + "normalize a function with negative values!" + ) if min(mu) < 0: - raise ValueError("Rates cannot be negative, don't ask me to " - "normalize a function over a negative domain!") + raise ValueError( + "Rates cannot be negative, don't ask me to " + "normalize a function over a negative domain!" + ) dp = integral_element(mu, pofmu) - return mu, pofmu/sum(dp) + return mu, pofmu / sum(dp) def compute_upper_limit(mu_in, post, alpha=0.9): @@ -79,10 +84,10 @@ def compute_lower_limit(mu_in, post, alpha=0.9): def confidence_interval_min_width(mu, post, alpha=0.9): - ''' + """ Returns the minimal-width confidence interval [mu_low, mu_high] of confidence level alpha for a posterior distribution post on the parameter mu. - ''' + """ if not 0 < alpha < 1: raise ValueError("Confidence level must be in (0,1).") @@ -105,32 +110,37 @@ def confidence_interval_min_width(mu, post, alpha=0.9): def hpd_coverage(mu, pdf, thresh): - ''' + """ Integrates a pdf over mu taking only bins where the mean over the bin is above a given threshold This gives the coverage of the HPD interval for the given threshold. - ''' + """ dp = integral_element(mu, pdf) - bin_mean = (pdf[1:] + pdf[:-1]) / 2. + bin_mean = (pdf[1:] + pdf[:-1]) / 2.0 return dp[bin_mean > thresh].sum() def hpd_threshold(mu_in, post, alpha, tol): - ''' + """ For a PDF post over samples mu_in, find a density threshold such that the region having higher density has coverage of at least alpha, and less than alpha plus a given tolerance. - ''' + """ norm_post = normalize_pdf(mu_in, post) # initialize bisection search p_minus = 0.0 p_plus = max(post) - while abs(hpd_coverage(mu_in, norm_post, p_minus) - - hpd_coverage(mu_in, norm_post, p_plus)) >= tol: - p_test = (p_minus + p_plus) / 2. + while ( + abs( + hpd_coverage(mu_in, norm_post, p_minus) + - hpd_coverage(mu_in, norm_post, p_plus) + ) + >= tol + ): + p_test = (p_minus + p_plus) / 2.0 if hpd_coverage(mu_in, post, p_test) >= alpha: # test value was too low or just right p_minus = p_test @@ -144,7 +154,7 @@ def hpd_threshold(mu_in, post, alpha, tol): def hpd_credible_interval(mu_in, post, alpha=0.9, tolerance=1e-3): - ''' + """ Returns the minimum and maximum rate values of the HPD (Highest Posterior Density) credible interval for a posterior post defined at the sample values mu_in. Samples need not be @@ -154,7 +164,7 @@ def hpd_credible_interval(mu_in, post, alpha=0.9, tolerance=1e-3): is multimodal and the correct interval is not contiguous; in this case will over-cover by including the whole range from minimum to maximum mu. - ''' + """ if alpha == 1: nonzero_samples = mu_in[post > 0] mu_low = numpy.min(nonzero_samples) @@ -181,40 +191,37 @@ def integrate_efficiency(dbins, eff, err=0, logbins=False): logd = numpy.log(dbins) dlogd = logd[1:] - logd[:-1] # use log midpoint of bins - dreps = numpy.exp((numpy.log(dbins[1:]) + numpy.log(dbins[:-1])) / 2.) - vol = numpy.sum(4.*numpy.pi * dreps**3. * eff * dlogd) + dreps = numpy.exp((numpy.log(dbins[1:]) + numpy.log(dbins[:-1])) / 2.0) + vol = numpy.sum(4.0 * numpy.pi * dreps**3.0 * eff * dlogd) # propagate errors in eff to errors in v - verr = numpy.sqrt( - numpy.sum((4.*numpy.pi * dreps**3. * err * dlogd)**2.) - ) + verr = numpy.sqrt(numpy.sum((4.0 * numpy.pi * dreps**3.0 * err * dlogd) ** 2.0)) else: dd = dbins[1:] - dbins[:-1] - dreps = (dbins[1:] + dbins[:-1]) / 2. - vol = numpy.sum(4. * numpy.pi * dreps**2. * eff * dd) + dreps = (dbins[1:] + dbins[:-1]) / 2.0 + vol = numpy.sum(4.0 * numpy.pi * dreps**2.0 * eff * dd) # propagate errors - verr = numpy.sqrt(numpy.sum((4.*numpy.pi * dreps**2. * err * dd)**2.)) + verr = numpy.sqrt(numpy.sum((4.0 * numpy.pi * dreps**2.0 * err * dd) ** 2.0)) return vol, verr def compute_efficiency(f_dist, m_dist, dbins): - ''' + """ Compute the efficiency as a function of distance for the given sets of found and missed injection distances. Note that injections that do not fit into any dbin get lost :( - ''' + """ efficiency = numpy.zeros(len(dbins) - 1) error = numpy.zeros(len(dbins) - 1) for j, dlow in enumerate(dbins[:-1]): dhigh = dbins[j + 1] found = numpy.sum((dlow <= f_dist) * (f_dist < dhigh)) missed = numpy.sum((dlow <= m_dist) * (m_dist < dhigh)) - if found+missed == 0: + if found + missed == 0: # avoid divide by 0 in empty bins - missed = 1. + missed = 1.0 efficiency[j] = float(found) / (found + missed) - error[j] = numpy.sqrt(efficiency[j] * (1 - efficiency[j]) / - (found + missed)) + error[j] = numpy.sqrt(efficiency[j] * (1 - efficiency[j]) / (found + missed)) return efficiency, error @@ -237,26 +244,32 @@ def mean_efficiency_volume(found, missed, dbins): def filter_injections_by_mass(injs, mbins, bin_num, bin_type, bin_num2=None): - ''' + """ For a given set of injections (sim_inspiral rows), return the subset of injections that fall within the given mass range. - ''' + """ if bin_type == "Mass1_Mass2": - m1bins = numpy.concatenate((mbins.lower()[0], - numpy.array([mbins.upper()[0][-1]]))) + m1bins = numpy.concatenate( + (mbins.lower()[0], numpy.array([mbins.upper()[0][-1]])) + ) m1lo = m1bins[bin_num] m1hi = m1bins[bin_num + 1] - m2bins = numpy.concatenate((mbins.lower()[1], - numpy.array([mbins.upper()[1][-1]]))) + m2bins = numpy.concatenate( + (mbins.lower()[1], numpy.array([mbins.upper()[1][-1]])) + ) m2lo = m2bins[bin_num2] m2hi = m2bins[bin_num2 + 1] - newinjs = [l for l in injs if - ((m1lo <= l.mass1 < m1hi and m2lo <= l.mass2 < m2hi) or - (m1lo <= l.mass2 < m1hi and m2lo <= l.mass1 < m2hi))] + newinjs = [ + l + for l in injs + if ( + (m1lo <= l.mass1 < m1hi and m2lo <= l.mass2 < m2hi) + or (m1lo <= l.mass2 < m1hi and m2lo <= l.mass1 < m2hi) + ) + ] return newinjs - mbins = numpy.concatenate((mbins.lower()[0], - numpy.array([mbins.upper()[0][-1]]))) + mbins = numpy.concatenate((mbins.lower()[0], numpy.array([mbins.upper()[0][-1]]))) mlow = mbins[bin_num] mhigh = mbins[bin_num + 1] if bin_type == "Chirp_Mass": @@ -269,15 +282,22 @@ def filter_injections_by_mass(injs, mbins, bin_num, bin_type, bin_num2=None): elif bin_type == "BNS_BBH": if bin_num in [0, 2]: # BNS/BBH case - newinjs = [l for l in injs if - (mlow <= l.mass1 < mhigh and mlow <= l.mass2 < mhigh)] + newinjs = [ + l for l in injs if (mlow <= l.mass1 < mhigh and mlow <= l.mass2 < mhigh) + ] else: # NSBH - newinjs = [l for l in injs if (mbins[0] <= l.mass1 < mbins[1] and - mbins[2] <= l.mass2 < mbins[3])] + newinjs = [ + l + for l in injs + if (mbins[0] <= l.mass1 < mbins[1] and mbins[2] <= l.mass2 < mbins[3]) + ] # BHNS - newinjs += [l for l in injs if (mbins[0] <= l.mass2 < mbins[1] and - mbins[2] <= l.mass1 < mbins[3])] + newinjs += [ + l + for l in injs + if (mbins[0] <= l.mass2 < mbins[1] and mbins[2] <= l.mass1 < mbins[3]) + ] return newinjs @@ -304,17 +324,16 @@ def compute_volume_vs_mass(found, missed, mass_bins, bin_type, dbins=None): if bin_type == "Mass1_Mass2": for j, mc1 in enumerate(mass_bins.centres()[0]): for k, mc2 in enumerate(mass_bins.centres()[1]): - newfound = filter_injections_by_mass( - found, mass_bins, j, bin_type, k) - newmissed = filter_injections_by_mass( - missed, mass_bins, j, bin_type, k) + newfound = filter_injections_by_mass(found, mass_bins, j, bin_type, k) + newmissed = filter_injections_by_mass(missed, mass_bins, j, bin_type, k) foundArray[(mc1, mc2)] = len(newfound) missedArray[(mc1, mc2)] = len(newmissed) # compute the volume using this injection set meaneff, efferr, meanvol, volerr = mean_efficiency_volume( - newfound, newmissed, dbins) + newfound, newmissed, dbins + ) effvmass.append(meaneff) errvmass.append(efferr) volArray[(mc1, mc2)] = meanvol @@ -323,20 +342,20 @@ def compute_volume_vs_mass(found, missed, mass_bins, bin_type, dbins=None): return volArray, vol2Array, foundArray, missedArray, effvmass, errvmass for j, mc in enumerate(mass_bins.centres()[0]): - # filter out injections not in this mass bin newfound = filter_injections_by_mass(found, mass_bins, j, bin_type) newmissed = filter_injections_by_mass(missed, mass_bins, j, bin_type) - foundArray[(mc, )] = len(newfound) - missedArray[(mc, )] = len(newmissed) + foundArray[(mc,)] = len(newfound) + missedArray[(mc,)] = len(newmissed) # compute the volume using this injection set meaneff, efferr, meanvol, volerr = mean_efficiency_volume( - newfound, newmissed, dbins) + newfound, newmissed, dbins + ) effvmass.append(meaneff) errvmass.append(efferr) - volArray[(mc, )] = meanvol - vol2Array[(mc, )] = volerr + volArray[(mc,)] = meanvol + vol2Array[(mc,)] = volerr return volArray, vol2Array, foundArray, missedArray, effvmass, errvmass diff --git a/pycbc/results/__init__.py b/pycbc/results/__init__.py index 63e0aca882b..ed3b54e1503 100644 --- a/pycbc/results/__init__.py +++ b/pycbc/results/__init__.py @@ -1,12 +1,12 @@ -from pycbc.results.table_utils import * -from pycbc.results.metadata import * -from pycbc.results.versioning import * from pycbc.results.color import * +from pycbc.results.dq import * +from pycbc.results.layout import * +from pycbc.results.metadata import * from pycbc.results.plot import * from pycbc.results.psd import * +from pycbc.results.pygrb_plotting_utils import * +from pycbc.results.pygrb_postprocessing_utils import * from pycbc.results.snr import * -from pycbc.results.layout import * -from pycbc.results.dq import * from pycbc.results.str_utils import * -from pycbc.results.pygrb_postprocessing_utils import * -from pycbc.results.pygrb_plotting_utils import * +from pycbc.results.table_utils import * +from pycbc.results.versioning import * diff --git a/pycbc/results/color.py b/pycbc/results/color.py index cd7585cba39..2fb64f5d161 100644 --- a/pycbc/results/color.py +++ b/pycbc/results/color.py @@ -1,23 +1,22 @@ -""" Utilities for managing matplotlib colors and mapping ifos to color -""" +"""Utilities for managing matplotlib colors and mapping ifos to color""" _ifo_color_map = { - 'G1': '#222222', # dark gray - 'K1': '#ffb200', # yellow/orange - 'H1': '#ee0000', # red - 'I1': '#b0dd8b', # light green - 'L1': '#4ba6ff', # blue - 'V1': '#9b59b6', # magenta/purple + "G1": "#222222", # dark gray + "K1": "#ffb200", # yellow/orange + "H1": "#ee0000", # red + "I1": "#b0dd8b", # light green + "L1": "#4ba6ff", # blue + "V1": "#9b59b6", # magenta/purple } _source_color_map = { - 'BNS': '#A2C8F5', # light blue - 'NSBH': '#FFB482', # light orange - 'BBH': '#FE9F9B', # light red - 'Mass Gap': '#8EE5A1', # light green - 'GNS': '#98D6CB', # turquoise - 'GG': '#79BB87', # green - 'BHG': '#C6C29E' # dark khaki + "BNS": "#A2C8F5", # light blue + "NSBH": "#FFB482", # light orange + "BBH": "#FE9F9B", # light red + "Mass Gap": "#8EE5A1", # light green + "GNS": "#98D6CB", # turquoise + "GG": "#79BB87", # green + "BHG": "#C6C29E", # dark khaki } diff --git a/pycbc/results/dq.py b/pycbc/results/dq.py index dcf8c3abd41..cdb5ae890a9 100644 --- a/pycbc/results/dq.py +++ b/pycbc/results/dq.py @@ -1,4 +1,4 @@ -'''This module contains utilities for following up search triggers''' +"""This module contains utilities for following up search triggers""" # JavaScript for searching the aLOG redirect_javascript = """""" -search_form_string="""""" @@ -30,7 +30,7 @@ 'https://alog.ligo-wa.caltech.edu/aLOG/includes/search.php?adminType=search'); return true;">aLOG""" -data_l1_string=""" +data_l1_string = """ Summary   @@ -40,7 +40,8 @@ def get_summary_page_link(ifo, utc_time): - """Return a string that links to the summary page and aLOG for this ifo + """ + Return a string that links to the summary page and aLOG for this ifo Parameters ---------- @@ -54,37 +55,38 @@ def get_summary_page_link(ifo, utc_time): ------- return_string : string String containing HTML for links to summary page and aLOG search + """ search_form = search_form_string - data = {'H1': data_h1_string, 'L1': data_l1_string} + data = {"H1": data_h1_string, "L1": data_l1_string} if ifo not in data: return ifo - else: - # support datetime/date objects or sequences (year, month, day) - try: - if hasattr(utc_time, 'year') and hasattr(utc_time, 'month') and hasattr(utc_time, 'day'): - year = int(utc_time.year) - month = int(utc_time.month) - day = int(utc_time.day) - else: - year = int(utc_time[0]) - month = int(utc_time[1]) - day = int(utc_time[2]) - except (AttributeError, TypeError, IndexError, ValueError) as e: - # Give a more informative error including the received value/type and - # the original exception to help debugging (for example when utc_time - # is a float and indexing utc_time[0] raises a TypeError). - raise TypeError( - "utc_time must be a datetime/date or a sequence (year, month, day); " - "got {} (type {}) - original error: {}".format( - utc_time, type(utc_time).__name__, e - ) - ) - - # alog format is day-month-year - alog_utc = '%02d-%02d-%4d' % (day, month, year) - # summary page is exactly the reverse - ext = '%4d%02d%02d' % (year, month, day) - return_string = search_form % (ifo.lower(), ifo.lower(), alog_utc, alog_utc) - return return_string + data[ifo] % ext + # support datetime/date objects or sequences (year, month, day) + try: + if ( + hasattr(utc_time, "year") + and hasattr(utc_time, "month") + and hasattr(utc_time, "day") + ): + year = int(utc_time.year) + month = int(utc_time.month) + day = int(utc_time.day) + else: + year = int(utc_time[0]) + month = int(utc_time[1]) + day = int(utc_time[2]) + except (AttributeError, TypeError, IndexError, ValueError) as e: + # Give a more informative error including the received value/type and + # the original exception to help debugging (for example when utc_time + # is a float and indexing utc_time[0] raises a TypeError). + raise TypeError( + "utc_time must be a datetime/date or a sequence (year, month, day); " + f"got {utc_time} (type {type(utc_time).__name__}) - original error: {e}" + ) + # alog format is day-month-year + alog_utc = "%02d-%02d-%4d" % (day, month, year) + # summary page is exactly the reverse + ext = "%4d%02d%02d" % (year, month, day) + return_string = search_form % (ifo.lower(), ifo.lower(), alog_utc, alog_utc) + return return_string + data[ifo] % ext diff --git a/pycbc/results/followup.py b/pycbc/results/followup.py index ccf522f1d1c..1c7162f98f8 100644 --- a/pycbc/results/followup.py +++ b/pycbc/results/followup.py @@ -21,25 +21,34 @@ # # ============================================================================= # -""" This module provides functions to generate followup plots and trigger +""" +This module provides functions to generate followup plots and trigger time series. """ -import numpy, matplotlib + # Only if a backend is not already set ... This should really *not* be done # here, but in the executables you should set matplotlib.use() # This matches the check that matplotlib does internally, but this *may* be # version dependenant. If this is a problem then remove this and control from # the executables directly. import sys -if 'matplotlib.backends' not in sys.modules: - matplotlib.use('agg') + +import matplotlib +import numpy + +if "matplotlib.backends" not in sys.modules: + matplotlib.use("agg") from matplotlib import pyplot as plt -import mpld3, mpld3.plugins +import mpld3 +import mpld3.plugins from igwn_segments import segment + from pycbc.io.hdf import HFile + def columns_from_file_list(file_list, columns, ifo, start, end): - """ Return columns of information stored in single detector trigger + """ + Return columns of information stored in single detector trigger files. Parameters @@ -60,15 +69,16 @@ def columns_from_file_list(file_list, columns, ifo, start, end): ------- trigger_dict : dict A dictionary of column vectors with column names as keys. + """ file_list = file_list.find_output_with_ifo(ifo) file_list = file_list.find_all_output_in_range(ifo, segment(start, end)) trig_dict = {} for trig_file in file_list: - f = HFile(trig_file.storage_path, 'r') + f = HFile(trig_file.storage_path, "r") - time = f['end_time'][:] + time = f["end_time"][:] pick = numpy.logical_and(time < end, time > start) pick_loc = numpy.where(pick)[0] @@ -79,53 +89,55 @@ def columns_from_file_list(file_list, columns, ifo, start, end): return trig_dict -ifo_color = {'H1': 'blue', 'L1':'red', 'V1':'green'} + +ifo_color = {"H1": "blue", "L1": "red", "V1": "green"} + def coinc_timeseries_plot(coinc_file, start, end): fig = plt.figure() - f = HFile(coinc_file, 'r') + f = HFile(coinc_file, "r") - stat1 = f['foreground/stat1'] - stat2 = f['foreground/stat2'] - time1 = f['foreground/time1'] - time2 = f['foreground/time2'] - ifo1 = f.attrs['detector_1'] - ifo2 = f.attrs['detector_2'] + stat1 = f["foreground/stat1"] + stat2 = f["foreground/stat2"] + time1 = f["foreground/time1"] + time2 = f["foreground/time2"] + ifo1 = f.attrs["detector_1"] + ifo2 = f.attrs["detector_2"] plt.scatter(time1, stat1, label=ifo1, color=ifo_color[ifo1]) plt.scatter(time2, stat2, label=ifo2, color=ifo_color[ifo2]) - fmt = '.12g' + fmt = ".12g" mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fmt=fmt)) plt.legend() - plt.xlabel('Time (s)') - plt.ylabel('NewSNR') + plt.xlabel("Time (s)") + plt.ylabel("NewSNR") plt.grid() return mpld3.fig_to_html(fig) + def trigger_timeseries_plot(file_list, ifos, start, end): fig = plt.figure() for ifo in ifos: - trigs = columns_from_file_list(file_list, - ['snr', 'end_time'], - ifo, start, end) + trigs = columns_from_file_list(file_list, ["snr", "end_time"], ifo, start, end) print(trigs) - plt.scatter(trigs['end_time'], trigs['snr'], label=ifo, - color=ifo_color[ifo]) + plt.scatter(trigs["end_time"], trigs["snr"], label=ifo, color=ifo_color[ifo]) - fmt = '.12g' + fmt = ".12g" mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fmt=fmt)) plt.legend() - plt.xlabel('Time (s)') - plt.ylabel('SNR') + plt.xlabel("Time (s)") + plt.ylabel("SNR") plt.grid() return mpld3.fig_to_html(fig) + def times_to_urls(times, window, tag): - base = '/../followup/%s/%s/%s' + base = "/../followup/%s/%s/%s" return times_to_links(times, window, tag, base=base) + def times_to_links(times, window, tag, base=None): if base is None: base = "followup" @@ -137,11 +149,10 @@ def times_to_links(times, window, tag, base=None): urls.append(base % (tag, start, end)) return urls + def get_gracedb_search_link(time): # Set up a search string for a 3s window around the coincidence - gdb_search_query = '%.0f+..+%.0f' % (numpy.floor(time) - 1, - numpy.ceil(time) + 1) - gdb_search_url = ('https://gracedb.ligo.org/search/?query=' - '{}&query_type=S'.format(gdb_search_query)) + gdb_search_query = "%.0f+..+%.0f" % (numpy.floor(time) - 1, numpy.ceil(time) + 1) + gdb_search_url = f"https://gracedb.ligo.org/search/?query={gdb_search_query}&query_type=S" gdb_search_link = 'Search' return gdb_search_link diff --git a/pycbc/results/layout.py b/pycbc/results/layout.py index 8450e2a4664..7620066d770 100644 --- a/pycbc/results/layout.py +++ b/pycbc/results/layout.py @@ -13,13 +13,15 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This module contains result page layout and numbering helper functions -""" +"""This module contains result page layout and numbering helper functions""" + import os.path from itertools import zip_longest -def two_column_layout(path, cols, unique='', **kwargs): - """ Make a well layout in a two column format + +def two_column_layout(path, cols, unique="", **kwargs): + """ + Make a well layout in a two column format Parameters ---------- @@ -31,13 +33,17 @@ def two_column_layout(path, cols, unique='', **kwargs): The format of the items on the well result section. Each tuple contains the two files that are shown in the left and right hand side of a row in the well.html page. + """ - path = os.path.join(os.getcwd(), path, 'well{}.html'.format(unique)) + path = os.path.join(os.getcwd(), path, f"well{unique}.html") from pycbc.results.render import render_workflow_html_template - render_workflow_html_template(path, 'two_column.html', cols, **kwargs) + + render_workflow_html_template(path, "two_column.html", cols, **kwargs) + def single_layout(path, files, **kwargs): - """ Make a well layout in single column format + """ + Make a well layout in single column format path: str Location to make the well html file @@ -46,14 +52,16 @@ def single_layout(path, files, **kwargs): """ two_column_layout(path, [(f,) for f in files], **kwargs) + def grouper(iterable, n, fillvalue=None): - """ Group items into chunks of n length - """ + """Group items into chunks of n length""" args = [iter(iterable)] * n return zip_longest(*args, fillvalue=fillvalue) + def group_layout(path, files, **kwargs): - """ Make a well layout in chunks of two from a list of files + """ + Make a well layout in chunks of two from a list of files path: str Location to make the well html file @@ -64,11 +72,13 @@ def group_layout(path, files, **kwargs): if len(files) > 0: two_column_layout(path, list(grouper(files, 2)), **kwargs) -class SectionNumber(object): - """ Class to help with numbering sections in an output page. - """ + +class SectionNumber: + """Class to help with numbering sections in an output page.""" + def __init__(self, base, secs): - """ Create section numbering instance + """ + Create section numbering instance Parameters ---------- @@ -76,6 +86,7 @@ def __init__(self, base, secs): The path of the of output html results directory secs: list of strings List of the subsections of the output html page + """ self.base = base self.secs = secs @@ -84,22 +95,23 @@ def __init__(self, base, secs): self.num = {} for num, sec in enumerate(secs): - self.name[sec] = '%s._%s' % (num + 1, sec) + self.name[sec] = "%s._%s" % (num + 1, sec) self.num[sec] = num self.count[sec] = 1 - def __getitem__ (self, path): - """ Return the path to use for the given subsection with numbering + def __getitem__(self, path): + """ + Return the path to use for the given subsection with numbering included. The numbering increments for each new subsection request. If a section is re-requested, it gets the original numbering. """ if path in self.name: name = self.name[path] else: - sec, subsec = path.split('/') + sec, subsec = path.split("/") subnum = self.count[sec] num = self.num[sec] - name = '%s/%s.%02d_%s' % (self.name[sec], num + 1, subnum, subsec) + name = "%s/%s.%02d_%s" % (self.name[sec], num + 1, subnum, subsec) self.count[sec] += 1 self.name[path] = name path = os.path.join(os.getcwd(), self.base, name) diff --git a/pycbc/results/metadata.py b/pycbc/results/metadata.py index fb211927713..5968ede5677 100644 --- a/pycbc/results/metadata.py +++ b/pycbc/results/metadata.py @@ -2,26 +2,31 @@ This Module contains generic utility functions for creating plots within PyCBC. """ -import os.path, pycbc.version + import configparser as ConfigParser +import os.path from html.parser import HTMLParser from xml.sax.saxutils import escape, unescape +import pycbc.version + escape_table = { - '"': """, - "'": "'", - "@": "@", - } + '"': """, + "'": "'", + "@": "@", +} unescape_table = { - "@" : "@", - } + "@": "@", +} for k, v in escape_table.items(): unescape_table[v] = k + def html_escape(text): - """ Sanitize text for html parsing """ + """Sanitize text for html parsing""" return escape(text, escape_table) + class MetaParser(HTMLParser): def __init__(self): self.metadata = {} @@ -31,45 +36,48 @@ def handle_data(self, data): pass def handle_starttag(self, tag, attrs): - attr= {} + attr = {} for key, value in attrs: attr[key] = value - if tag == 'div' and 'class' in attr and attr['class'] == 'pycbc-meta': - self.metadata[attr['key']] = unescape(attr['value'], unescape_table) + if tag == "div" and "class" in attr and attr["class"] == "pycbc-meta": + self.metadata[attr["key"]] = unescape(attr["value"], unescape_table) def save_html_with_metadata(fig, filename, fig_kwds, kwds): - """ Save a html output to file with metadata """ + """Save a html output to file with metadata""" if isinstance(fig, str): text = fig else: from mpld3 import fig_to_html + text = fig_to_html(fig, **fig_kwds) - f = open(filename, 'w') + f = open(filename, "w") for key, value in kwds.items(): value = escape(value, escape_table) - line = "
" % (str(key), value) + line = '
' % (str(key), value) f.write(line) f.write(text) + def load_html_metadata(filename): - """ Get metadata from html file """ + """Get metadata from html file""" parser = MetaParser() - data = open(filename, 'r').read() + data = open(filename).read() - if 'pycbc-meta' in data: + if "pycbc-meta" in data: print("LOADING HTML FILE %s" % filename) parser.feed(data) cp = ConfigParser.ConfigParser(parser.metadata) cp.add_section(os.path.basename(filename)) return cp + def save_png_with_metadata(fig, filename, fig_kwds, kwds): - """ Save a matplotlib figure to a png with metadata - """ + """Save a matplotlib figure to a png with metadata""" from PIL import Image, PngImagePlugin + fig.savefig(filename, **fig_kwds) im = Image.open(filename) @@ -80,39 +88,46 @@ def save_png_with_metadata(fig, filename, fig_kwds, kwds): im.save(filename, "png", pnginfo=meta) + def save_pdf_with_metadata(fig, filename, fig_kwds, kwds): - """Save a matplotlib figure to a PDF file with metadata. - """ + """Save a matplotlib figure to a PDF file with metadata.""" # https://stackoverflow.com/a/17462125 from matplotlib.backends.backend_pdf import PdfPages with PdfPages(filename) as pdfp: - fig.savefig(pdfp, format='pdf', **fig_kwds) + fig.savefig(pdfp, format="pdf", **fig_kwds) metadata = pdfp.infodict() for key in kwds: - if str(key).lower() == 'title': + if str(key).lower() == "title": # map the title to the official PDF keyword (capitalized) - metadata['Title'] = str(kwds[key]) + metadata["Title"] = str(kwds[key]) else: metadata[str(key)] = str(kwds[key]) + def load_png_metadata(filename): from PIL import Image + data = Image.open(filename).info cp = ConfigParser.ConfigParser(data) cp.add_section(os.path.basename(filename)) return cp -_metadata_saver = {'.png': save_png_with_metadata, - '.html': save_html_with_metadata, - '.pdf': save_pdf_with_metadata, - } -_metadata_loader = {'.png': load_png_metadata, - '.html': load_html_metadata, - } + +_metadata_saver = { + ".png": save_png_with_metadata, + ".html": save_html_with_metadata, + ".pdf": save_pdf_with_metadata, +} +_metadata_loader = { + ".png": load_png_metadata, + ".html": load_html_metadata, +} + def save_fig_with_metadata(fig, filename, fig_kwds=None, **kwds): - """ Save plot to file with metadata included. Kewords translate to metadata + """ + Save plot to file with metadata included. Kewords translate to metadata that is stored directly in the plot file. Limited format types available. Parameters @@ -121,19 +136,24 @@ def save_fig_with_metadata(fig, filename, fig_kwds=None, **kwds): The matplotlib figure to save to the file filename: str Name of file to store the plot. + """ if fig_kwds is None: fig_kwds = {} try: extension = os.path.splitext(filename)[1] - kwds['version'] = pycbc.version.git_verbose_msg + kwds["version"] = pycbc.version.git_verbose_msg _metadata_saver[extension](fig, filename, fig_kwds, kwds) except KeyError: - raise TypeError('Cannot save file %s with metadata, extension %s not ' - 'supported at this time' % (filename, extension)) + raise TypeError( + "Cannot save file %s with metadata, extension %s not " + "supported at this time" % (filename, extension) + ) + def load_metadata_from_file(filename): - """ Load the plot related metadata saved in a file + """ + Load the plot related metadata saved in a file Parameters ---------- @@ -144,10 +164,13 @@ def load_metadata_from_file(filename): ------- cp: ConfigParser A configparser object containing the metadata + """ try: extension = os.path.splitext(filename)[1] return _metadata_loader[extension](filename) except KeyError: - raise TypeError('Cannot read metadata from file %s, extension %s not ' - 'supported at this time' % (filename, extension)) + raise TypeError( + "Cannot read metadata from file %s, extension %s not " + "supported at this time" % (filename, extension) + ) diff --git a/pycbc/results/mpld3_utils.py b/pycbc/results/mpld3_utils.py index 6bad980f15e..7bab48c7311 100644 --- a/pycbc/results/mpld3_utils.py +++ b/pycbc/results/mpld3_utils.py @@ -1,6 +1,9 @@ -""" This module provides functionality to extend mpld3 -""" -import mpld3, mpld3.plugins, mpld3.utils +"""This module provides functionality to extend mpld3""" + +import mpld3 +import mpld3.plugins +import mpld3.utils + class ClickLink(mpld3.plugins.PluginBase): """Plugin for following a link on click""" @@ -28,11 +31,14 @@ class ClickLink(mpld3.plugins.PluginBase): ); } """ + def __init__(self, points, links): - self.dict_ = {"type": "clicklink", - "id": mpld3.utils.get_id(points), - "links": links, - } + self.dict_ = { + "type": "clicklink", + "id": mpld3.utils.get_id(points), + "links": links, + } + class MPLSlide(mpld3.plugins.PluginBase): JAVASCRIPT = """ @@ -102,20 +108,22 @@ class MPLSlide(mpld3.plugins.PluginBase): if (this.props.enabled) this.fig.enable_zoom(); else this.fig.disable_zoom(); }; """ + def __init__(self, button=True, enabled=None): if enabled is None: enabled = not button - self.dict_ = {"type": "zoom", - "button": button, - "enabled": enabled} + self.dict_ = {"type": "zoom", "button": button, "enabled": enabled} + class Tooltip(mpld3.plugins.PointHTMLTooltip): JAVASCRIPT = "" - def __init__(self, points, labels=None, - hoffset=0, voffset=10, css=None): - super(Tooltip, self).__init__(points, labels, hoffset, voffset, "") + + def __init__(self, points, labels=None, hoffset=0, voffset=10, css=None): + super().__init__(points, labels, hoffset, voffset, "") + class LineTooltip(mpld3.plugins.LineHTMLTooltip): JAVASCRIPT = "" + def __init__(self, line, label=None, hoffset=0, voffset=10, css=None): - super(LineTooltip, self).__init__(line, label, hoffset, voffset, "") + super().__init__(line, label, hoffset, voffset, "") diff --git a/pycbc/results/plot.py b/pycbc/results/plot.py index 6832b5a1eab..91a9aee3dbe 100644 --- a/pycbc/results/plot.py +++ b/pycbc/results/plot.py @@ -1,26 +1,26 @@ -""" Plotting utilities and premade plot configurations -""" +"""Plotting utilities and premade plot configurations""" + def hist_overflow(val, val_max, **kwds): - """ Make a histogram with an overflow bar above val_max """ + """Make a histogram with an overflow bar above val_max""" from matplotlib import pyplot as plt - overflow = len(val[val>=val_max]) - plt.hist(val[val= val_max]) + plt.hist(val[val < val_max], **kwds) - if 'color' in kwds: - color = kwds['color'] + if "color" in kwds: + color = kwds["color"] else: color = None if overflow > 0: - rect = plt.bar(val_max+0.05, overflow, .5, color=color)[0] - plt.text(rect.get_x(), - 1.10*rect.get_height(), '%s+' % val_max) + rect = plt.bar(val_max + 0.05, overflow, 0.5, color=color)[0] + plt.text(rect.get_x(), 1.10 * rect.get_height(), "%s+" % val_max) def add_style_opt_to_parser(parser, default=None): - """Adds an option to set the matplotlib style to a parser. + """ + Adds an option to set the matplotlib style to a parser. Parameters ---------- @@ -29,20 +29,27 @@ def add_style_opt_to_parser(parser, default=None): default : str, optional The default style to use. Default, None, will result in the default matplotlib style to be used. + """ from matplotlib import pyplot - parser.add_argument('--mpl-style', default=default, - choices=['default']+pyplot.style.available+['xkcd'], - help='Set the matplotlib style to use.') + + parser.add_argument( + "--mpl-style", + default=default, + choices=["default"] + pyplot.style.available + ["xkcd"], + help="Set the matplotlib style to use.", + ) def set_style_from_cli(opts): - """Uses the mpl-style option to set the style for plots. + """ + Uses the mpl-style option to set the style for plots. Note: This will change the global rcParams. """ from matplotlib import pyplot - if opts.mpl_style == 'xkcd': + + if opts.mpl_style == "xkcd": # this is treated differently for some reason pyplot.xkcd() elif opts.mpl_style is not None: diff --git a/pycbc/results/psd.py b/pycbc/results/psd.py index 3a69c9e7d2d..82f3e3658b8 100644 --- a/pycbc/results/psd.py +++ b/pycbc/results/psd.py @@ -26,11 +26,12 @@ """ Module to generate PSD figures """ -from pycbc.results import ifo_color + from pycbc import DYN_RANGE_FAC +from pycbc.results import ifo_color -def generate_asd_plot(psddict, output_filename, f_min=10.): +def generate_asd_plot(psddict, output_filename, f_min=10.0): """ Generate an ASD plot as used for upload to GraceDB. @@ -49,28 +50,27 @@ def generate_asd_plot(psddict, output_filename, f_min=10.): Returns ------- None + """ from matplotlib import pyplot as plt + asd_fig, asd_ax = plt.subplots(1) - asd_min = [1E-24] # Default minimum to plot + asd_min = [1e-24] # Default minimum to plot for ifo in sorted(psddict.keys()): curr_psd = psddict[ifo] freqs = curr_psd.sample_frequencies - physical = (freqs >= f_min) # Ignore lower frequencies + physical = freqs >= f_min # Ignore lower frequencies asd_to_plot = curr_psd[physical] ** 0.5 / DYN_RANGE_FAC asd_min.append(min(asd_to_plot)) - asd_ax.loglog(freqs[physical], - asd_to_plot, - c=ifo_color(ifo), - label=ifo) + asd_ax.loglog(freqs[physical], asd_to_plot, c=ifo_color(ifo), label=ifo) asd_ax.grid(True) asd_ax.legend() asd_ax.set_xlim([f_min, 1300]) - asd_ax.set_ylim([min(asd_min), 1E-20]) - asd_ax.set_xlabel('Frequency (Hz)') - asd_ax.set_ylabel('ASD') + asd_ax.set_ylim([min(asd_min), 1e-20]) + asd_ax.set_xlabel("Frequency (Hz)") + asd_ax.set_ylabel("ASD") asd_fig.savefig(output_filename) diff --git a/pycbc/results/pygrb_plotting_utils.py b/pycbc/results/pygrb_plotting_utils.py index abf1faea53d..19eaf2d7d87 100644 --- a/pycbc/results/pygrb_plotting_utils.py +++ b/pycbc/results/pygrb_plotting_utils.py @@ -24,8 +24,10 @@ """ import copy -import numpy + import igwn_segments as segments +import numpy + from pycbc.results import save_fig_with_metadata @@ -34,17 +36,16 @@ # ============================================================================= def contour_plotter(axis, snr_vals, contours, colors, vert_spike=False): """Plot contours in a scatter plot where SNR is on the horizontal axis""" - for i, _ in enumerate(contours): plot_vals_x = [] plot_vals_y = [] if vert_spike: for j, _ in enumerate(snr_vals): # Workaround to ensure vertical spike is shown on veto plots - if contours[i][j] > 1E-15 and not plot_vals_x: + if contours[i][j] > 1e-15 and not plot_vals_x: plot_vals_x.append(snr_vals[j]) plot_vals_y.append(0.1) - if contours[i][j] > 1E-15 and plot_vals_x: + if contours[i][j] > 1e-15 and plot_vals_x: plot_vals_x.append(snr_vals[j]) plot_vals_y.append(contours[i][j]) else: @@ -57,31 +58,44 @@ def contour_plotter(axis, snr_vals, contours, colors, vert_spike=False): # Functions used in executables # + # ============================================================================= # Plot trigger time and offsource extent over segments # Courtesy of Alex Dietz # ============================================================================= -def make_grb_segments_plot(wkflow, science_segs, trigger_time, trigger_name, - out_dir, coherent_seg=None, fail_criterion=None): +def make_grb_segments_plot( + wkflow, + science_segs, + trigger_time, + trigger_name, + out_dir, + coherent_seg=None, + fail_criterion=None, +): """Plot trigger time and offsource extent over segments""" - import matplotlib.pyplot as plt - from matplotlib.patches import Rectangle from matplotlib.lines import Line2D + from matplotlib.patches import Rectangle + from pycbc.results.color import ifo_color ifos = wkflow.ifos if len(sum(science_segs.values(), [])) == 0: - extent = segments.segment(int(wkflow.cp.get("workflow", "start-time")), - int(wkflow.cp.get("workflow", "end-time"))) + extent = segments.segment( + int(wkflow.cp.get("workflow", "start-time")), + int(wkflow.cp.get("workflow", "end-time")), + ) else: - pltpad = [science_segs.extent_all()[1] - trigger_time, - trigger_time - science_segs.extent_all()[0]] - extent = segments.segmentlist([science_segs.extent_all(), - segments.segment(trigger_time - - pltpad[0], - trigger_time - + pltpad[1])]).extent() + pltpad = [ + science_segs.extent_all()[1] - trigger_time, + trigger_time - science_segs.extent_all()[0], + ] + extent = segments.segmentlist( + [ + science_segs.extent_all(), + segments.segment(trigger_time - pltpad[0], trigger_time + pltpad[1]), + ] + ).extent() ifo_colors = {} for ifo in ifos: @@ -93,38 +107,68 @@ def make_grb_segments_plot(wkflow, science_segs, trigger_time, trigger_name, fig, subs = plt.subplots(len(ifos), sharey=True) if len(ifos) == 1: subs = [subs] - plt.xticks(rotation=20, ha='right') + plt.xticks(rotation=20, ha="right") for sub, ifo in zip(subs, ifos): for seg in science_segs[ifo]: - sub.add_patch(Rectangle((seg[0], 0.1), abs(seg), 0.8, - facecolor=ifo_colors[ifo], - edgecolor='none')) + sub.add_patch( + Rectangle( + (seg[0], 0.1), + abs(seg), + 0.8, + facecolor=ifo_colors[ifo], + edgecolor="none", + ) + ) if coherent_seg: - if len(science_segs[ifo]) > 0 and \ - coherent_seg in science_segs[ifo]: - sub.plot([trigger_time, trigger_time], [0, 1], '-', - c='orange') - sub.add_patch(Rectangle((coherent_seg[0], 0), - abs(coherent_seg), 1, alpha=0.5, - facecolor='orange', edgecolor='none')) + if len(science_segs[ifo]) > 0 and coherent_seg in science_segs[ifo]: + sub.plot([trigger_time, trigger_time], [0, 1], "-", c="orange") + sub.add_patch( + Rectangle( + (coherent_seg[0], 0), + abs(coherent_seg), + 1, + alpha=0.5, + facecolor="orange", + edgecolor="none", + ) + ) else: - sub.plot([trigger_time, trigger_time], [0, 1], ':', - c='orange') - sub.plot([coherent_seg[0], coherent_seg[0]], [0, 1], '--', - c='orange', alpha=0.5) - sub.plot([coherent_seg[1], coherent_seg[1]], [0, 1], '--', - c='orange', alpha=0.5) + sub.plot([trigger_time, trigger_time], [0, 1], ":", c="orange") + sub.plot( + [coherent_seg[0], coherent_seg[0]], + [0, 1], + "--", + c="orange", + alpha=0.5, + ) + sub.plot( + [coherent_seg[1], coherent_seg[1]], + [0, 1], + "--", + c="orange", + alpha=0.5, + ) else: - sub.plot([trigger_time, trigger_time], [0, 1], ':k') + sub.plot([trigger_time, trigger_time], [0, 1], ":k") if fail_criterion: if len(science_segs[ifo]) > 0: - style_str = '--' + style_str = "--" else: - style_str = '-' - sub.plot([fail_criterion[0], fail_criterion[0]], [0, 1], style_str, - c='black', alpha=0.5) - sub.plot([fail_criterion[1], fail_criterion[1]], [0, 1], style_str, - c='black', alpha=0.5) + style_str = "-" + sub.plot( + [fail_criterion[0], fail_criterion[0]], + [0, 1], + style_str, + c="black", + alpha=0.5, + ) + sub.plot( + [fail_criterion[1], fail_criterion[1]], + [0, 1], + style_str, + c="black", + alpha=0.5, + ) sub.set_frame_on(False) sub.set_yticks([]) @@ -142,17 +186,18 @@ def make_grb_segments_plot(wkflow, science_segs, trigger_time, trigger_name, xmin, xmax = fig.axes[-1].get_xaxis().get_view_interval() ymin, _ = fig.axes[-1].get_yaxis().get_view_interval() - fig.axes[-1].add_artist(Line2D((xmin, xmax), (ymin, ymin), color='black', - linewidth=2)) - fig.axes[-1].set_xlabel('GPS Time') + fig.axes[-1].add_artist( + Line2D((xmin, xmax), (ymin, ymin), color="black", linewidth=2) + ) + fig.axes[-1].set_xlabel("GPS Time") - fig.axes[0].set_title('Science Segments for GRB%s' % trigger_name) + fig.axes[0].set_title("Science Segments for GRB%s" % trigger_name) plt.tight_layout() fig.subplots_adjust(hspace=0) - plot_name = 'GRB%s_segments.png' % trigger_name - plot_url = 'file://localhost%s/%s' % (out_dir, plot_name) - fig.savefig('%s/%s' % (out_dir, plot_name)) + plot_name = "GRB%s_segments.png" % trigger_name + plot_url = "file://localhost%s/%s" % (out_dir, plot_name) + fig.savefig("%s/%s" % (out_dir, plot_name)) return [ifos, plot_name, extent, plot_url] @@ -162,7 +207,6 @@ def make_grb_segments_plot(wkflow, science_segs, trigger_time, trigger_name, # ============================================================================= def axis_max_value(trig_values, inj_values, inj_file): """Deterime the maximum of a quantity in the trigger and injection data""" - axis_max = trig_values.max() if inj_file and inj_values.size and inj_values.max() > axis_max: axis_max = inj_values.max() @@ -175,7 +219,6 @@ def axis_max_value(trig_values, inj_values, inj_file): # ============================================================================= def axis_min_value(trig_values, inj_values, inj_file): """Deterime the minimum of a quantity in the trigger and injection data""" - axis_min = trig_values.min() if inj_file and inj_values.size and inj_values.min() < axis_min: axis_min = inj_values.min() @@ -186,9 +229,19 @@ def axis_min_value(trig_values, inj_values, inj_file): # ============================================================================= # Master plotting function: fits all plotting needs in for PyGRB results # ============================================================================= -def pygrb_plotter(trigs, injs, xlabel, ylabel, opts, - snr_vals=None, conts=None, shade_cont_value=None, - colors=None, vert_spike=False, cmd=None): +def pygrb_plotter( + trigs, + injs, + xlabel, + ylabel, + opts, + snr_vals=None, + conts=None, + shade_cont_value=None, + colors=None, + vert_spike=False, + cmd=None, +): """Master function to plot PyGRB results""" from matplotlib import pyplot as plt @@ -197,9 +250,9 @@ def pygrb_plotter(trigs, injs, xlabel, ylabel, opts, cax = fig.gca() # Plot trigger-related and (if present) injection-related quantities cax_plotter = cax.loglog if opts.use_logs else cax.plot - cax_plotter(trigs[0], trigs[1], 'bx') + cax_plotter(trigs[0], trigs[1], "bx") if not (injs[0] is None and injs[1] is None): - cax_plotter(injs[0], injs[1], 'r+') + cax_plotter(injs[0], injs[1], "r+") cax.grid() # Plot contours if conts is not None: @@ -211,19 +264,19 @@ def pygrb_plotter(trigs, injs, xlabel, ylabel, opts, polyy = copy.deepcopy(conts[shade_cont_value]) polyx = numpy.append(polyx, [max(snr_vals), min(snr_vals)]) polyy = numpy.append(polyy, [limy, limy]) - cax.fill(polyx, polyy, color='#dddddd') + cax.fill(polyx, polyy, color="#dddddd") # Axes: labels and limits cax.set_xlabel(xlabel) cax.set_ylabel(ylabel) if opts.x_lims: - x_lims = map(float, opts.x_lims.split(',')) + x_lims = map(float, opts.x_lims.split(",")) cax.set_xlim(x_lims) if opts.y_lims: - y_lims = map(float, opts.y_lims.split(',')) + y_lims = map(float, opts.y_lims.split(",")) cax.set_ylim(y_lims) # Wrap up plt.tight_layout() - save_fig_with_metadata(fig, opts.output_file, cmd=cmd, - title=opts.plot_title, - caption=opts.plot_caption) + save_fig_with_metadata( + fig, opts.output_file, cmd=cmd, title=opts.plot_title, caption=opts.plot_caption + ) plt.close() diff --git a/pycbc/results/pygrb_postprocessing_utils.py b/pycbc/results/pygrb_postprocessing_utils.py index 1428af3c5a7..3db6a4f7a88 100644 --- a/pycbc/results/pygrb_postprocessing_utils.py +++ b/pycbc/results/pygrb_postprocessing_utils.py @@ -23,21 +23,22 @@ Module to generate PyGRB figures: scatter plots and timeseries. """ -import logging import argparse import copy -import numpy +import logging + import h5py import igwn_segments as segments - -from scipy import stats +import numpy from igwn_segments.utils import fromsegwizard -from pycbc.events.coherent import reweightedsnr_cut -from pycbc.events import veto +from scipy import stats + from pycbc import add_common_pycbc_options +from pycbc.events import veto +from pycbc.events.coherent import reweightedsnr_cut from pycbc.io.hdf import HFile -logger = logging.getLogger('pycbc.results.pygrb_postprocessing_utils') +logger = logging.getLogger("pycbc.results.pygrb_postprocessing_utils") # ============================================================================= @@ -49,50 +50,68 @@ # ============================================================================= def pygrb_initialize_plot_parser(description=None): """Sets up a basic argument parser object for PyGRB plotting scripts""" - formatter_class = argparse.ArgumentDefaultsHelpFormatter - parser = argparse.ArgumentParser(description=description, - formatter_class=formatter_class) + parser = argparse.ArgumentParser( + description=description, formatter_class=formatter_class + ) add_common_pycbc_options(parser) - parser.add_argument("-o", "--output-file", default=None, - help="Output file.") - parser.add_argument("--x-lims", default=None, - help="Comma separated minimum and maximum values " - "for the horizontal axis. When using negative " - "values an equal sign after --x-lims is necessary.") - parser.add_argument("--y-lims", default=None, - help="Comma separated minimum and maximum values " - "for the vertical axis. When using negative values " - "an equal sign after --y-lims is necessary.") - parser.add_argument("--use-logs", default=False, action="store_true", - help="Produce a log-log plot") - parser.add_argument("-i", "--ifo", default=None, help="IFO used for IFO " - "specific plots") - parser.add_argument("-a", "--seg-files", nargs="+", - default=[], help="The location of the buffer, " - "onsource and offsource txt segment files.") - parser.add_argument("-V", "--veto-file", - help="The location of the xml veto file.") - parser.add_argument('--plot-title', default=None, - help="If provided, use the given string as the plot " - "title.") - parser.add_argument('--plot-caption', default=None, - help="If provided, use the given string as the plot " - "caption") + parser.add_argument("-o", "--output-file", default=None, help="Output file.") + parser.add_argument( + "--x-lims", + default=None, + help="Comma separated minimum and maximum values " + "for the horizontal axis. When using negative " + "values an equal sign after --x-lims is necessary.", + ) + parser.add_argument( + "--y-lims", + default=None, + help="Comma separated minimum and maximum values " + "for the vertical axis. When using negative values " + "an equal sign after --y-lims is necessary.", + ) + parser.add_argument( + "--use-logs", default=False, action="store_true", help="Produce a log-log plot" + ) + parser.add_argument( + "-i", "--ifo", default=None, help="IFO used for IFO specific plots" + ) + parser.add_argument( + "-a", + "--seg-files", + nargs="+", + default=[], + help="The location of the buffer, onsource and offsource txt segment files.", + ) + parser.add_argument("-V", "--veto-file", help="The location of the xml veto file.") + parser.add_argument( + "--plot-title", + default=None, + help="If provided, use the given string as the plot title.", + ) + parser.add_argument( + "--plot-caption", + default=None, + help="If provided, use the given string as the plot caption", + ) return parser def pygrb_add_slide_opts(parser): """Add to parser object arguments related to short timeslides""" - parser.add_argument("--slide-id", type=str, default='0', - help="Select a specific slide or set to all to plot " - "results from all short slides.") + parser.add_argument( + "--slide-id", + type=str, + default="0", + help="Select a specific slide or set to all to plot " + "results from all short slides.", + ) def slide_opts_helper(args): """ - This function overwrites the types of input slide_id information - when loading data in postprocessing scripts. + This function overwrites the types of input slide_id information + when loading data in postprocessing scripts. """ if args.slide_id.isdigit(): args.slide_id = int(args.slide_id) @@ -106,70 +125,123 @@ def pygrb_add_injmc_opts(parser): """Add to parser object the arguments used for Monte-Carlo on distance.""" if parser is None: parser = argparse.ArgumentParser() - parser.add_argument("-M", "--num-mc-injections", - type=int, default=100, help="Number of Monte " - "Carlo injection simulations to perform.") - parser.add_argument("-S", "--seed", type=int, - default=1234, help="Seed to initialize Monte Carlo.") - parser.add_argument("-U", "--upper-inj-dist", - type=float, default=1000, help="The upper distance " - "of the injections in Mpc, if used.") - parser.add_argument("-L", "--lower-inj-dist", - type=float, default=0, help="The lower distance of " - "the injections in Mpc, if used.") - parser.add_argument("-n", "--num-bins", type=int, - default=0, help="The number of bins used to " - "calculate injection efficiency.") - parser.add_argument("-w", "--waveform-error", - type=float, default=0, help="The standard deviation " - "to use when calculating the waveform error.") + parser.add_argument( + "-M", + "--num-mc-injections", + type=int, + default=100, + help="Number of Monte Carlo injection simulations to perform.", + ) + parser.add_argument( + "-S", "--seed", type=int, default=1234, help="Seed to initialize Monte Carlo." + ) + parser.add_argument( + "-U", + "--upper-inj-dist", + type=float, + default=1000, + help="The upper distance of the injections in Mpc, if used.", + ) + parser.add_argument( + "-L", + "--lower-inj-dist", + type=float, + default=0, + help="The lower distance of the injections in Mpc, if used.", + ) + parser.add_argument( + "-n", + "--num-bins", + type=int, + default=0, + help="The number of bins used to calculate injection efficiency.", + ) + parser.add_argument( + "-w", + "--waveform-error", + type=float, + default=0, + help="The standard deviation to use when calculating the waveform error.", + ) for ifo in ["g1", "h1", "k1", "l1", "v1"]: - parser.add_argument(f"--{ifo}-cal-error", type=float, - default=0, help="The standard deviation to use " - f"when calculating the {ifo.upper()} " - "calibration amplitude error.") - parser.add_argument(f"--{ifo}-dc-cal-error", - type=float, default=1.0, help="The scaling " - "factor to use when calculating the " - f"{ifo.upper()} calibration amplitude error.") + parser.add_argument( + f"--{ifo}-cal-error", + type=float, + default=0, + help="The standard deviation to use " + f"when calculating the {ifo.upper()} " + "calibration amplitude error.", + ) + parser.add_argument( + f"--{ifo}-dc-cal-error", + type=float, + default=1.0, + help="The scaling " + "factor to use when calculating the " + f"{ifo.upper()} calibration amplitude error.", + ) def pygrb_add_bestnr_opts(parser): """Add to the parser object the arguments used for BestNR calculation""" if parser is None: parser = argparse.ArgumentParser() - parser.add_argument("-Q", "--chisq-index", type=float, - default=6.0, help="chisq_index for newSNR " - "calculation (default: 6)") - parser.add_argument("-N", "--chisq-nhigh", type=float, - default=2.0, help="chisq_nhigh for newSNR " - "calculation (default: 2") + parser.add_argument( + "-Q", + "--chisq-index", + type=float, + default=6.0, + help="chisq_index for newSNR calculation (default: 6)", + ) + parser.add_argument( + "-N", + "--chisq-nhigh", + type=float, + default=2.0, + help="chisq_nhigh for newSNR calculation (default: 2", + ) def pygrb_add_null_snr_opts(parser): - """Add to the parser object the arguments used for null SNR calculation - and null SNR cut.""" - parser.add_argument("-A", "--null-snr-threshold", - default=5.25, - type=float, - help="Null SNR threshold for null SNR cut " - "(default: 5.25)") - parser.add_argument("-T", "--null-grad-thresh", type=float, - default=20., help="Threshold above which to " - "increase the values of the null SNR cut") - parser.add_argument("-D", "--null-grad-val", type=float, - default=0.2, help="Rate the null SNR cut will " - "increase above the threshold") + """ + Add to the parser object the arguments used for null SNR calculation + and null SNR cut. + """ + parser.add_argument( + "-A", + "--null-snr-threshold", + default=5.25, + type=float, + help="Null SNR threshold for null SNR cut (default: 5.25)", + ) + parser.add_argument( + "-T", + "--null-grad-thresh", + type=float, + default=20.0, + help="Threshold above which to increase the values of the null SNR cut", + ) + parser.add_argument( + "-D", + "--null-grad-val", + type=float, + default=0.2, + help="Rate the null SNR cut will increase above the threshold", + ) def pygrb_add_bestnr_cut_opt(parser): """Add to the parser object an argument to place a threshold on BestNR.""" if parser is None: parser = argparse.ArgumentParser() - parser.add_argument("--newsnr-threshold", type=float, metavar='THRESHOLD', - default=0., - help="Cut triggers with NewSNR less than THRESHOLD. " - "Default 0: all events are considered.") + parser.add_argument( + "--newsnr-threshold", + type=float, + metavar="THRESHOLD", + default=0.0, + help="Cut triggers with NewSNR less than THRESHOLD. " + "Default 0: all events are considered.", + ) # ============================================================================= @@ -178,7 +250,6 @@ def pygrb_add_bestnr_cut_opt(parser): # An underscore at the name start flags a function called only in this file def _read_seg_files(seg_files): """Read segments txt files""" - if len(seg_files) != 3: err_msg = "The location of three segment files is necessary." err_msg += "[bufferSeg.txt, offSourceSeg.txt, onSourceSeg.txt]" @@ -189,9 +260,9 @@ def _read_seg_files(seg_files): keys = ["buffer", "off", "on"] for key, seg_file in zip(keys, seg_files): - segs = fromsegwizard(open(seg_file, 'r')) + segs = fromsegwizard(open(seg_file)) if len(segs) > 1: - err_msg = 'More than one segment, an error has occured.' + err_msg = "More than one segment, an error has occured." raise RuntimeError(err_msg) times[key] = segs[0] @@ -203,7 +274,6 @@ def _read_seg_files(seg_files): # ============================================================================= def _extract_vetoes(veto_file, ifos, offsource): """Extracts the veto segments from the veto File""" - clean_segs = {} vetoes = segments.segmentlistdict() @@ -223,8 +293,7 @@ def _extract_vetoes(veto_file, ifos, offsource): for ifo in ifos: for v in vetoes[ifo]: v_span = v[1] - v[0] - logging.info("%ds of data vetoed at GPS time %d", - v_span, v[0]) + logging.info("%ds of data vetoed at GPS time %d", v_span, v[0]) return vetoes @@ -234,7 +303,6 @@ def _extract_vetoes(veto_file, ifos, offsource): # ============================================================================= def _slide_vetoes(vetoes, slide_dict_or_list, slide_id, ifos): """Build a dictionary (indexed by ifo) of time-slid vetoes""" - # Copy vetoes if vetoes: slid_vetoes = copy.deepcopy(vetoes) @@ -250,13 +318,12 @@ def _slide_vetoes(vetoes, slide_dict_or_list, slide_id, ifos): # ============================================================================= # Recursive function to reach all datasets in an HDF file handle # ============================================================================= -def _dataset_iterator(g, prefix=''): +def _dataset_iterator(g, prefix=""): """Reach all datasets in an HDF file handle""" - for key, item in g.items(): # Avoid slash as first character - pref = prefix[1:] if prefix.startswith('/') else prefix - path = pref + '/' + key + pref = prefix.removeprefix("/") + path = pref + "/" + key if isinstance(item, h5py.Dataset): yield (path, item) elif isinstance(item, h5py.Group): @@ -266,37 +333,44 @@ def _dataset_iterator(g, prefix=''): # ============================================================================= # Function to load trigger/injection data # ============================================================================= -def load_data(input_file, ifos, rw_snr_threshold=None, data_tag=None, - slide_id=None): - """Load data from a trigger/injection PyGRB output file, returning a +def load_data(input_file, ifos, rw_snr_threshold=None, data_tag=None, slide_id=None): + """ + Load data from a trigger/injection PyGRB output file, returning a dictionary. If the input_file is None, None is returned. data_tag enables logging information about the number of triggers/injections found, so the - user should not set it to 'trigs'/'injs' when processing the onsource.""" - + user should not set it to 'trigs'/'injs' when processing the onsource. + """ if not input_file: return None - trigs = HFile(input_file, 'r') - rw_snr = trigs['network/reweighted_snr'][:] \ - if 'network/reweighted_snr' in trigs.keys() else numpy.array([]) - net_ids = trigs['network/event_id'][:] \ - if 'network/event_id' in trigs.keys() \ + trigs = HFile(input_file, "r") + rw_snr = ( + trigs["network/reweighted_snr"][:] + if "network/reweighted_snr" in trigs.keys() + else numpy.array([]) + ) + net_ids = ( + trigs["network/event_id"][:] + if "network/event_id" in trigs.keys() else numpy.array([], dtype=numpy.int64) + ) # Output the number of items loaded only upon a request by the user who is # expected not to set data_tag to 'trigs'or 'injs' when processing the # onsource - if data_tag == 'trigs': + if data_tag == "trigs": logging.info("%d triggers loaded.", len(rw_snr)) - elif data_tag == 'injs': + elif data_tag == "injs": logging.info("%d injections loaded.", len(rw_snr)) else: logging.info("Loading triggers.") ifo_ids = {} for ifo in ifos: - ifo_ids[ifo] = trigs[ifo+'/event_id'][:] \ - if ifo+'/event_id' in trigs.keys() \ + ifo_ids[ifo] = ( + trigs[ifo + "/event_id"][:] + if ifo + "/event_id" in trigs.keys() else numpy.array([], dtype=numpy.int64) + ) trigs.close() # Apply the reweighted SNR cut on the reweighted SNR @@ -308,9 +382,9 @@ def load_data(input_file, ifos, rw_snr_threshold=None, data_tag=None, # Output the number of items surviging vetoes with the same logic as above msg = "" - if data_tag == 'trigs': + if data_tag == "trigs": msg += f"{sum(above_thresh)} triggers " - elif data_tag == 'injs': + elif data_tag == "injs": msg = f"{sum(above_thresh)} injections " if msg: msg += f"surviving reweighted SNR cut at {rw_snr_threshold}." @@ -326,10 +400,7 @@ def load_data(input_file, ifos, rw_snr_threshold=None, data_tag=None, # 2. Do not assume ifo_ids are sorted and produce a sorting tracker sorter = numpy.argsort(input_ids) # 3. Find out where the target_ids are in the sorted ifo_ids[ifo] - insert_indices = numpy.searchsorted( - input_ids, - target_ids, - sorter=sorter) + insert_indices = numpy.searchsorted(input_ids, target_ids, sorter=sorter) # 4. Safety checks: # a. searchsorted returns len(arr) when the element is not in the arr # b. ensure that all target_ids were recovered in ifo_ids[ifo] @@ -347,35 +418,35 @@ def load_data(input_file, ifos, rw_snr_threshold=None, data_tag=None, mask = numpy.full(sum(above_thresh), True) # When necessary and possible, update the mask so it selects a given # slide id - if slide_id is not None and 'network/slide_id' in trigs.keys(): - mask = numpy.where(trigs['network/slide_id'][above_thresh] == - slide_id)[0] - for (path, dset) in _dataset_iterator(trigs): + if slide_id is not None and "network/slide_id" in trigs.keys(): + mask = numpy.where(trigs["network/slide_id"][above_thresh] == slide_id)[0] + for path, dset in _dataset_iterator(trigs): # The dataset contains search or missed injections information, # not properties of triggers or found injections: just copy it - if 'search' in path or 'missed' in path or 'gating' in path: + if "search" in path or "missed" in path or "gating" in path: trigs_dict[path] = dset[:] # The dataset is trig/inj info at an IFO: # cut with the correct index elif path[:2] in ifos: ifo = path[:2] if ifo_ids_above_thresh_locations[ifo].size != 0: - trigs_dict[path] = \ - dset[:][ifo_ids_above_thresh_locations[ifo]] + trigs_dict[path] = dset[:][ifo_ids_above_thresh_locations[ifo]] else: trigs_dict[path] = numpy.array([]) # The dataset is trig/inj network info: cut it before copying else: trigs_dict[path] = dset[above_thresh] - if 'network/slide_id' in trigs.keys(): + if "network/slide_id" in trigs.keys(): # The slide selection is applied to datasets that contain # properties of surviving triggers. These datasets are # identified knowing that each trigger has a slide id, so # they must have as many entries as the 'network/slide_id' # dataset once triggers below threshold are removed from it. - if trigs_dict[path].size == \ - trigs['network/slide_id'][above_thresh].size: + if ( + trigs_dict[path].size + == trigs["network/slide_id"][above_thresh].size + ): trigs_dict[path] = trigs_dict[path][mask] return trigs_dict @@ -384,9 +455,11 @@ def load_data(input_file, ifos, rw_snr_threshold=None, data_tag=None, # ============================================================================= # Function to apply vetoes to found injections # ============================================================================= -def apply_vetoes_to_found_injs(found_missed_file, found_injs, ifos, - veto_file=None, keys=None): - """Separate injections surviving vetoes from vetoed injections. +def apply_vetoes_to_found_injs( + found_missed_file, found_injs, ifos, veto_file=None, keys=None +): + """ + Separate injections surviving vetoes from vetoed injections. Parameters ---------- @@ -409,40 +482,33 @@ def apply_vetoes_to_found_injs(found_missed_file, found_injs, ifos, found_idx: numpy.array of indices of surviving injections veto_idx: numpy.array of indices of vetoed injections - """ - keep_keys = keys if keys else found_injs.keys() + """ + keep_keys = keys or found_injs.keys() - if not found_missed_file or ifos[0]+'/end_time' not in found_injs.keys(): - t_id_key = 'network/template_id' + if not found_missed_file or ifos[0] + "/end_time" not in found_injs.keys(): + t_id_key = "network/template_id" if t_id_key not in keep_keys: - keep_keys = list(keep_keys+[t_id_key]) + keep_keys = list(keep_keys + [t_id_key]) empty_dict = dict.fromkeys(keep_keys, numpy.array([])) - empty_dict[t_id_key] = \ - empty_dict[t_id_key].astype(dtype=numpy.int64) + empty_dict[t_id_key] = empty_dict[t_id_key].astype(dtype=numpy.int64) return (empty_dict, empty_dict, None, None) - found_idx = numpy.arange(len(found_injs[ifos[0]+'/end_time'][:])) + found_idx = numpy.arange(len(found_injs[ifos[0] + "/end_time"][:])) veto_idx = numpy.array([], dtype=numpy.int64) if veto_file: logging.info("Applying data vetoes to found injections...") for ifo in ifos: - inj_time = found_injs[ifo+'/end_time'][:] - segs = veto.select_segments_by_definer(veto_file, - segment_name=None, - ifo=ifo) + inj_time = found_injs[ifo + "/end_time"][:] + segs = veto.select_segments_by_definer( + veto_file, segment_name=None, ifo=ifo + ) if len(segs) > 0: - idx, _ = veto.indices_outside_segments(inj_time, - [veto_file], - ifo, - None) + idx, _ = veto.indices_outside_segments(inj_time, [veto_file], ifo, None) veto_idx = numpy.append(veto_idx, idx) logging.info("%d injections vetoed due to %s.", len(idx), ifo) - idx, _ = veto.indices_within_segments(inj_time, - [veto_file], - ifo, - None) + idx, _ = veto.indices_within_segments(inj_time, [veto_file], ifo, None) found_idx = numpy.intersect1d(found_idx, idx) veto_idx = numpy.unique(veto_idx) @@ -452,7 +518,7 @@ def apply_vetoes_to_found_injs(found_missed_file, found_injs, ifos, found_after_vetoes = {} missed_after_vetoes = {} for key in keep_keys: - if key == 'network/coincident_snr': + if key == "network/coincident_snr": found_injs[key] = get_coinc_snr(found_injs) if isinstance(found_injs[key], numpy.ndarray): found_after_vetoes[key] = found_injs[key][found_idx] @@ -467,44 +533,45 @@ def apply_vetoes_to_found_injs(found_missed_file, found_injs, ifos, # * Function to calculate the antenna distance factor # ============================================================================= def _get_antenna_single_response(antenna, ra, dec, geocent_time): - """Returns the antenna response F+^2 + Fx^2 of an IFO (passed as pycbc - Detector type) at a given sky location and time.""" - + """ + Returns the antenna response F+^2 + Fx^2 of an IFO (passed as pycbc + Detector type) at a given sky location and time. + """ fp, fc = antenna.antenna_pattern(ra, dec, 0, geocent_time) return fp**2 + fc**2 # Vectorize the function above on all but the first argument -get_antenna_responses = numpy.vectorize(_get_antenna_single_response, - otypes=[float]) +get_antenna_responses = numpy.vectorize(_get_antenna_single_response, otypes=[float]) get_antenna_responses.excluded.add(0) def get_antenna_dist_factor(antenna, ra, dec, geocent_time, inc=0.0): - """Returns the antenna factors (defined as eq. 4.3 on page 57 of + """ + Returns the antenna factors (defined as eq. 4.3 on page 57 of Duncan Brown's Ph.D.) for an IFO (passed as pycbc Detector type) at - a given sky location and time.""" - + a given sky location and time. + """ fp, fc = antenna.antenna_pattern(ra, dec, 0, geocent_time) - return numpy.sqrt(fp ** 2 * (1 + numpy.cos(inc)) ** 2 / 4 + fc ** 2) + return numpy.sqrt(fp**2 * (1 + numpy.cos(inc)) ** 2 / 4 + fc**2) # ============================================================================= # Construct sorted triggers from trials # ============================================================================= def sort_trigs(trial_dict, trigs, slide_dict, seg_dict): - """Constructs sorted triggers from a trials dictionary for the slides - requested via slide_dict.""" - + """ + Constructs sorted triggers from a trials dictionary for the slides + requested via slide_dict. + """ sorted_trigs = {} # Begin by sorting the triggers into each slide for slide_id in slide_dict: sorted_trigs[slide_id] = [] - for slide_id, event_id in zip(trigs['network/slide_id'], - trigs['network/event_id']): + for slide_id, event_id in zip(trigs["network/slide_id"], trigs["network/event_id"]): if slide_id in slide_dict: sorted_trigs[slide_id].append(event_id) @@ -514,8 +581,8 @@ def sort_trigs(trial_dict, trigs, slide_dict, seg_dict): # Check the triggers are all in the analysed segment lists for event_id in sorted_trigs[slide_id]: - index = numpy.flatnonzero(trigs['network/event_id'] == event_id)[0] - end_time = trigs['network/end_time_gc'][index] + index = numpy.flatnonzero(trigs["network/event_id"] == event_id)[0] + end_time = trigs["network/end_time_gc"][index] if end_time not in curr_seg_list: # This can be raised if the trigger is on the segment boundary, # so check if the trigger is within 1/100 of a second within @@ -531,17 +598,19 @@ def sort_trigs(trial_dict, trigs, slide_dict, seg_dict): # Keep triggers that are in trial_dict num_trigs_before = len(sorted_trigs[slide_id]) - sorted_trigs[slide_id] = [event_id for event_id in - sorted_trigs[slide_id] - if trigs['network/end_time_gc'][ - trigs['network/event_id'] == event_id][0] - in trial_dict[slide_id]] + sorted_trigs[slide_id] = [ + event_id + for event_id in sorted_trigs[slide_id] + if trigs["network/end_time_gc"][trigs["network/event_id"] == event_id][0] + in trial_dict[slide_id] + ] # Check that the number of triggers has not increased after vetoes - assert len(sorted_trigs[slide_id]) <= num_trigs_before, \ - f"Slide {slide_id} has {num_trigs_before} triggers before the "\ - f"trials dictionary was used and {len(sorted_trigs[slide_id])} "\ + assert len(sorted_trigs[slide_id]) <= num_trigs_before, ( + f"Slide {slide_id} has {num_trigs_before} triggers before the " + f"trials dictionary was used and {len(sorted_trigs[slide_id])} " "after. This should not happen." + ) # END OF CHECK # return sorted_trigs @@ -551,9 +620,10 @@ def sort_trigs(trial_dict, trigs, slide_dict, seg_dict): # Extract trigger properties and store them as dictionaries # ============================================================================= def extract_trig_properties(trial_dict, trigs, slide_dict, seg_dict, keys): - """Extract and store as dictionaries specific keys of time-slid - triggers (trigs) compatibly with the trials dictionary (trial_dict)""" - + """ + Extract and store as dictionaries specific keys of time-slid + triggers (trigs) compatibly with the trials dictionary (trial_dict) + """ # Sort the triggers into each slide sorted_trigs = sort_trigs(trial_dict, trigs, slide_dict, seg_dict) n_surviving_trigs = sum(len(i) for i in sorted_trigs.values()) @@ -567,12 +637,14 @@ def extract_trig_properties(trial_dict, trigs, slide_dict, seg_dict, keys): for slide_id in slide_dict: slide_trigs = sorted_trigs[slide_id] - indices = numpy.nonzero( - numpy.isin(trigs['network/event_id'], slide_trigs))[0] + indices = numpy.nonzero(numpy.isin(trigs["network/event_id"], slide_trigs))[0] for key in keys: if slide_trigs: - found_trigs[key][slide_id] = get_coinc_snr(trigs)[indices] \ - if key == 'network/coincident_snr' else trigs[key][indices] + found_trigs[key][slide_id] = ( + get_coinc_snr(trigs)[indices] + if key == "network/coincident_snr" + else trigs[key][indices] + ) else: found_trigs[key][slide_id] = numpy.asarray([]) @@ -586,15 +658,14 @@ def extract_trig_properties(trial_dict, trigs, slide_dict, seg_dict, keys): # ============================================================================= def extract_ifos(trig_file, ifo=None): """Extracts IFOs from hdf file and checks for presence of a specific IFO""" - # Load hdf file - hdf_file = HFile(trig_file, 'r') + hdf_file = HFile(trig_file, "r") # Extract IFOs ifos = sorted(list(hdf_file.keys())) # Remove unwanted keys from key list to reduce it to the ifos - for key in ['network', 'found', 'missed']: + for key in ["network", "found", "missed"]: if key in ifos: ifos.remove(key) @@ -612,17 +683,16 @@ def extract_ifos(trig_file, ifo=None): def load_time_slides(hdf_file_path): """Loads timeslides from PyGRB output file as a dictionary""" logging.info("Loading timeslides.") - hdf_file = HFile(hdf_file_path, 'r') + hdf_file = HFile(hdf_file_path, "r") ifos = extract_ifos(hdf_file_path) - ids = numpy.arange(len(hdf_file[f'{ifos[0]}/search/time_slides'])) + ids = numpy.arange(len(hdf_file[f"{ifos[0]}/search/time_slides"])) time_slide_dict = { - slide_id: { - ifo: hdf_file[f'{ifo}/search/time_slides'][slide_id] - for ifo in ifos} - for slide_id in ids} + slide_id: {ifo: hdf_file[f"{ifo}/search/time_slides"][slide_id] for ifo in ifos} + for slide_id in ids + } # Check time_slide_ids are ordered correctly. - if not (numpy.all(ids[1:] == numpy.array(ids[:-1])+1) and ids[0] == 0): + if not (numpy.all(ids[1:] == numpy.array(ids[:-1]) + 1) and ids[0] == 0): err_msg = "time_slide_ids list should start at zero and increase by " err_msg += "one for every element" raise RuntimeError(err_msg) @@ -644,20 +714,23 @@ def load_segment_dict(hdf_file_path): Loads the segment dictionary with the format {slide_id: segmentlist(segments analyzed)} """ - logging.info("Loading segments.") # Long time slides will require mapping between slides and segments - hdf_file = HFile(hdf_file_path, 'r') + hdf_file = HFile(hdf_file_path, "r") ifos = extract_ifos(hdf_file_path) # Get slide IDs - slide_ids = numpy.arange(len(hdf_file[f'{ifos[0]}/search/time_slides'])) + slide_ids = numpy.arange(len(hdf_file[f"{ifos[0]}/search/time_slides"])) # Get segment start/end times - seg_starts = hdf_file['network/search/segments/start_times'][:] - seg_ends = hdf_file['network/search/segments/end_times'][:] + seg_starts = hdf_file["network/search/segments/start_times"][:] + seg_ends = hdf_file["network/search/segments/end_times"][:] # Write list of segments - seg_list = segments.segmentlist([segments.segment(seg_start, seg_ends[i]) - for i, seg_start in enumerate(seg_starts)]) + seg_list = segments.segmentlist( + [ + segments.segment(seg_start, seg_ends[i]) + for i, seg_start in enumerate(seg_starts) + ] + ) # Write segment_dict in proper format # At the moment of this comment, there is only one segment @@ -669,10 +742,10 @@ def load_segment_dict(hdf_file_path): # ============================================================================= # Construct the trials from the timeslides, segments, and vetoes # ============================================================================= -def construct_trials(seg_files, seg_dict, ifos, slide_dict, veto_file, - hide_onsource=True): +def construct_trials( + seg_files, seg_dict, ifos, slide_dict, veto_file, hide_onsource=True +): """Constructs trials from segments, timeslides, and vetoes""" - logging.info("Constructing trials.") trial_dict = {} @@ -680,10 +753,10 @@ def construct_trials(seg_files, seg_dict, ifos, slide_dict, veto_file, segs = _read_seg_files(seg_files) # Separate segments - trial_time = abs(segs['on']) + trial_time = abs(segs["on"]) # Determine the veto segments - vetoes = _extract_vetoes(veto_file, ifos, segs['off']) + vetoes = _extract_vetoes(veto_file, ifos, segs["off"]) # Slide vetoes over trials: this can only *reduce* the analysis time for slide_id in slide_dict: @@ -694,10 +767,12 @@ def construct_trials(seg_files, seg_dict, ifos, slide_dict, veto_file, if hide_onsource: for ifo in ifos: slide_offset = slide_dict[slide_id][ifo] - seg_buffer.append(segments.segment(segs['buffer'][0] - - slide_offset, - segs['buffer'][1] - - slide_offset)) + seg_buffer.append( + segments.segment( + segs["buffer"][0] - slide_offset, + segs["buffer"][1] - slide_offset, + ) + ) seg_buffer.coalesce() # Construct the ifo-indexed dictionary of slid veteoes @@ -708,15 +783,17 @@ def construct_trials(seg_files, seg_dict, ifos, slide_dict, veto_file, for curr_seg in curr_seg_list: iter_int = 1 while 1: - trial_end = curr_seg[0] + trial_time*iter_int + trial_end = curr_seg[0] + trial_time * iter_int if trial_end > curr_seg[1]: break - curr_trial = segments.segment(trial_end - trial_time, - trial_end) + curr_trial = segments.segment(trial_end - trial_time, trial_end) if not seg_buffer.intersects_segment(curr_trial): - intersect = numpy.any([slid_vetoes[ifo]. - intersects_segment(curr_trial) - for ifo in ifos]) + intersect = numpy.any( + [ + slid_vetoes[ifo].intersects_segment(curr_trial) + for ifo in ifos + ] + ) if not intersect: trial_dict[slide_id].append(curr_trial) @@ -733,7 +810,6 @@ def construct_trials(seg_files, seg_dict, ifos, slide_dict, veto_file, # ============================================================================= def sort_stat(time_veto_max_stat): """Sort a dictionary of loudest SNRs/BestNRs""" - full_time_veto_max_stat = list(time_veto_max_stat.values()) full_time_veto_max_stat = numpy.concatenate(full_time_veto_max_stat) full_time_veto_max_stat.sort() @@ -746,17 +822,19 @@ def sort_stat(time_veto_max_stat): # ============================================================================= def max_median_stat(slide_dict, time_veto_max_stat, trig_stat, total_trials): """Return maximum and median of trig_stat and sorted time_veto_max_stat""" - - max_stat = max(trig_stat[slide_id].max() if trig_stat[slide_id].size - else 0 for slide_id in slide_dict) + max_stat = max( + trig_stat[slide_id].max() if trig_stat[slide_id].size else 0 + for slide_id in slide_dict + ) full_time_veto_max_stat = sort_stat(time_veto_max_stat) if total_trials % 2: median_stat = full_time_veto_max_stat[(total_trials - 1) // 2] else: - median_stat = numpy.mean((full_time_veto_max_stat) - [total_trials//2 - 1: total_trials//2 + 1]) + median_stat = numpy.mean( + (full_time_veto_max_stat)[total_trials // 2 - 1 : total_trials // 2 + 1] + ) return max_stat, median_stat, full_time_veto_max_stat @@ -766,21 +844,20 @@ def max_median_stat(slide_dict, time_veto_max_stat, trig_stat, total_trials): # ============================================================================= def mc_cal_wf_errs(num_mc_injs, inj_dists, cal_err, wf_err, max_dc_cal_err): """Includes calibration and waveform errors by running an MC""" - # The efficiency calculations include calibration and waveform # errors incorporated by running over each injection num_mc_injs times, # where each time we draw a random value of distance. num_injs = len(inj_dists) - inj_dist_mc = numpy.ndarray((num_mc_injs+1, num_injs)) + inj_dist_mc = numpy.ndarray((num_mc_injs + 1, num_injs)) inj_dist_mc[0, :] = inj_dists for i in range(num_mc_injs): cal_dist_red = stats.norm.rvs(size=num_injs) * cal_err wf_dist_red = numpy.abs(stats.norm.rvs(size=num_injs) * wf_err) - inj_dist_mc[i+1, :] = inj_dists / (max_dc_cal_err * - (1 + cal_dist_red) * - (1 + wf_dist_red)) + inj_dist_mc[i + 1, :] = inj_dists / ( + max_dc_cal_err * (1 + cal_dist_red) * (1 + wf_dist_red) + ) return inj_dist_mc @@ -789,13 +866,14 @@ def mc_cal_wf_errs(num_mc_injs, inj_dists, cal_err, wf_err, max_dc_cal_err): # Function to calculate the coincident SNR # ============================================================================= def get_coinc_snr(trigs_or_injs): - """ Calculate coincident SNR using coherent and null SNRs""" - + """Calculate coincident SNR using coherent and null SNRs""" coinc_snr = numpy.array([]) - if 'network/coherent_snr' in trigs_or_injs.keys() and \ - 'network/null_snr' in trigs_or_injs.keys(): - coh_snr_sq = numpy.square(trigs_or_injs['network/coherent_snr'][:]) - null_snr_sq = numpy.square(trigs_or_injs['network/null_snr'][:]) + if ( + "network/coherent_snr" in trigs_or_injs.keys() + and "network/null_snr" in trigs_or_injs.keys() + ): + coh_snr_sq = numpy.square(trigs_or_injs["network/coherent_snr"][:]) + null_snr_sq = numpy.square(trigs_or_injs["network/null_snr"][:]) coinc_snr = numpy.sqrt(coh_snr_sq + null_snr_sq) return coinc_snr @@ -806,17 +884,19 @@ def template_hash_to_id(trigger_file, bank_path): This function converts the template hashes from a trigger file into 'template_id's that represent indices of the templates within the bank. + Parameters ---------- trigger_file: HFile object for trigger file bank_file: filepath for template bank + """ - ifos = [k for k in trigger_file.keys() if k != 'network'] - if ifos[0]+'/template_hash' not in trigger_file.keys(): + ifos = [k for k in trigger_file.keys() if k != "network"] + if ifos[0] + "/template_hash" not in trigger_file.keys(): return numpy.array([], dtype=int) with HFile(bank_path, "r") as bank: - hashes = bank['template_hash'][:] - trig_hashes = trigger_file[f'{ifos[0]}/template_hash'][:] + hashes = bank["template_hash"][:] + trig_hashes = trigger_file[f"{ifos[0]}/template_hash"][:] trig_ids = numpy.zeros(trig_hashes.shape[0], dtype=int) for idx, t_hash in enumerate(hashes): matches = numpy.where(trig_hashes == t_hash) diff --git a/pycbc/results/render.py b/pycbc/results/render.py index 3f58c4bf6fb..3258add5a93 100644 --- a/pycbc/results/render.py +++ b/pycbc/results/render.py @@ -14,47 +14,49 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -import os.path, types import codecs - +import os.path +import types from configparser import ConfigParser -from jinja2 import Environment, FileSystemLoader from xml.sax.saxutils import unescape +from jinja2 import Environment, FileSystemLoader + import pycbc.results from pycbc.results import unescape_table from pycbc.results.metadata import save_html_with_metadata from pycbc.workflow.core import SegFile, makedir + def render_workflow_html_template(filename, subtemplate, filelists, **kwargs): - """ Writes a template given inputs from the workflow generator. Takes + """ + Writes a template given inputs from the workflow generator. Takes a list of tuples. Each tuple is a pycbc File object. Also the name of the subtemplate to render and the filename of the output. """ - dirnam = os.path.dirname(filename) makedir(dirnam) try: - filenames = [f.name for filelist in filelists for f in filelist if f is not None] + filenames = [ + f.name for filelist in filelists for f in filelist if f is not None + ] except TypeError: filenames = [] # render subtemplate - subtemplate_dir = pycbc.results.__path__[0] + '/templates/wells' + subtemplate_dir = pycbc.results.__path__[0] + "/templates/wells" env = Environment(loader=FileSystemLoader(subtemplate_dir)) env.globals.update(get_embedded_config=get_embedded_config) env.globals.update(path_exists=os.path.exists) env.globals.update(len=len) subtemplate = env.get_template(subtemplate) - context = {'filelists' : filelists, - 'dir' : dirnam} + context = {"filelists": filelists, "dir": dirnam} context.update(kwargs) output = subtemplate.render(context) # save as html page - kwds = {'render-function' : 'render_tmplt', - 'filenames' : ','.join(filenames)} + kwds = {"render-function": "render_tmplt", "filenames": ",".join(filenames)} kwds.update(kwargs) for key in kwds: @@ -62,12 +64,14 @@ def render_workflow_html_template(filename, subtemplate, filelists, **kwargs): save_html_with_metadata(str(output), filename, None, kwds) + def get_embedded_config(filename): - """ Attempt to load config data attached to file - """ + """Attempt to load config data attached to file""" + def check_option(self, section, name): - return (self.has_section(section) and - (self.has_option(section, name) or (name in self.defaults()))) + return self.has_section(section) and ( + self.has_option(section, name) or (name in self.defaults()) + ) try: cp = pycbc.results.load_metadata_from_file(filename) @@ -78,18 +82,17 @@ def check_option(self, section, name): return cp -def setup_template_render(path, config_path): - """ This function is the gateway for rendering a template for a file. - """ +def setup_template_render(path, config_path): + """This function is the gateway for rendering a template for a file.""" # initialization cp = get_embedded_config(path) - output = '' + output = "" filename = os.path.basename(path) # use meta-data if not empty for rendering - if cp.has_option(filename, 'render-function'): - render_function_name = cp.get(filename, 'render-function') + if cp.has_option(filename, "render-function"): + render_function_name = cp.get(filename, "render-function") render_function = eval(render_function_name) output = render_function(path, cp) @@ -98,8 +101,8 @@ def setup_template_render(path, config_path): cp.read(config_path) # render template - if cp.has_option(filename, 'render-function'): - render_function_name = cp.get(filename, 'render-function') + if cp.has_option(filename, "render-function"): + render_function_name = cp.get(filename, "render-function") render_function = eval(render_function_name) output = render_function(path, cp) else: @@ -112,22 +115,23 @@ def setup_template_render(path, config_path): return output + def render_default(path, cp): - """ This is the default function that will render a template to a string of HTML. The + """ + This is the default function that will render a template to a string of HTML. The string will be for a drop-down tab that contains a link to the file. If the file extension requires information to be read, then that is passed to the content variable (eg. a segmentlistdict). """ - # define filename and slug from path filename = os.path.basename(path) - slug = filename.replace('.', '_') + slug = filename.replace(".", "_") # initializations content = None - if path.endswith('.xml') or path.endswith('.xml.gz'): + if path.endswith(".xml") or path.endswith(".xml.gz"): # segment or veto file return a segmentslistdict instance try: wf_file = SegFile.from_segment_xml(path) @@ -135,111 +139,101 @@ def render_default(path, cp): # for now I just coalesce. wf_file.return_union_seglist() except Exception as e: - print('No segment table found in %s : %s' % (path, e)) + print("No segment table found in %s : %s" % (path, e)) # render template - template_dir = pycbc.results.__path__[0] + '/templates/files' + template_dir = pycbc.results.__path__[0] + "/templates/files" env = Environment(loader=FileSystemLoader(template_dir)) env.globals.update(abs=abs) env.globals.update(open=open) env.globals.update(path_exists=os.path.exists) - template = env.get_template('file_default.html') - context = {'path' : path, - 'filename' : filename, - 'slug' : slug, - 'cp' : cp, - 'content' : content} + template = env.get_template("file_default.html") + context = { + "path": path, + "filename": filename, + "slug": slug, + "cp": cp, + "content": content, + } output = template.render(context) return output -def render_glitchgram(path, cp): - """ Render a glitchgram file template. - """ +def render_glitchgram(path, cp): + """Render a glitchgram file template.""" # define filename and slug from path filename = os.path.basename(path) - slug = filename.replace('.', '_') + slug = filename.replace(".", "_") # render template - template_dir = pycbc.results.__path__[0] + '/templates/files' + template_dir = pycbc.results.__path__[0] + "/templates/files" env = Environment(loader=FileSystemLoader(template_dir)) env.globals.update(abs=abs) - template = env.get_template(cp.get(filename, 'template')) - context = {'filename' : filename, - 'slug' : slug, - 'cp' : cp} + template = env.get_template(cp.get(filename, "template")) + context = {"filename": filename, "slug": slug, "cp": cp} output = template.render(context) return output -def render_text(path, cp): - """ Render a file as text. - """ +def render_text(path, cp): + """Render a file as text.""" # define filename and slug from path filename = os.path.basename(path) - slug = filename.replace('.', '_') + slug = filename.replace(".", "_") # initializations content = None # read file as a string - with codecs.open(path, 'r', encoding='utf-8', errors='replace') as fp: + with codecs.open(path, "r", encoding="utf-8", errors="replace") as fp: content = fp.read() # replace all the escaped characters content = unescape(content, unescape_table) # render template - template_dir = pycbc.results.__path__[0] + '/templates/files' + template_dir = pycbc.results.__path__[0] + "/templates/files" env = Environment(loader=FileSystemLoader(template_dir)) env.globals.update(abs=abs) env.globals.update(path_exists=os.path.exists) - template = env.get_template('file_pre.html') - context = {'filename' : filename, - 'slug' : slug, - 'cp' : cp, - 'content' : content} + template = env.get_template("file_pre.html") + context = {"filename": filename, "slug": slug, "cp": cp, "content": content} output = template.render(context) return output + def render_ignore(path, cp): - """ Does not render anything. - """ + """Does not render anything.""" + return "" - return '' def render_tmplt(path, cp): - """ Render a file as text. - """ - + """Render a file as text.""" # define filename and slug from path filename = os.path.basename(path) - slug = filename.replace('.', '_') + slug = filename.replace(".", "_") # initializations content = None # read file as a string - with open(path, 'r') as fp: + with open(path) as fp: content = fp.read() # replace all the escaped characters content = unescape(content, unescape_table) # render template - template_dir = '/'.join(path.split('/')[:-1]) + template_dir = "/".join(path.split("/")[:-1]) env = Environment(loader=FileSystemLoader(template_dir)) env.globals.update(setup_template_render=setup_template_render) env.globals.update(get_embedded_config=get_embedded_config) env.globals.update(path_exists=os.path.exists) template = env.get_template(filename) - context = {'filename' : filename, - 'slug' : slug, - 'cp' : cp, - 'content' : content} + context = {"filename": filename, "slug": slug, "cp": cp, "content": content} output = template.render(context) return output diff --git a/pycbc/results/scatter_histograms.py b/pycbc/results/scatter_histograms.py index ad3e5cc09f9..7d607e5b8e0 100644 --- a/pycbc/results/scatter_histograms.py +++ b/pycbc/results/scatter_histograms.py @@ -29,29 +29,29 @@ import itertools import sys +import matplotlib import numpy - import scipy.stats -import matplotlib - # Only if a backend is not already set ... This should really *not* be done # here, but in the executables you should set matplotlib.use() # This matches the check that matplotlib does internally, but this *may* be # version dependenant. If this is a problem then remove this and control from # the executables directly. -if 'matplotlib.backends' not in sys.modules: # nopep8 - matplotlib.use('agg') +if "matplotlib.backends" not in sys.modules: # nopep8 + matplotlib.use("agg") -from matplotlib import (offsetbox, pyplot, gridspec, colors) +from matplotlib import colors, gridspec, offsetbox, pyplot -from pycbc.results import str_utils from pycbc.io import FieldArray +from pycbc.results import str_utils -def create_axes_grid(parameters, labels=None, height_ratios=None, - width_ratios=None, no_diagonals=False): - """Given a list of parameters, creates a figure with an axis for +def create_axes_grid( + parameters, labels=None, height_ratios=None, width_ratios=None, no_diagonals=False +): + """ + Given a list of parameters, creates a figure with an axis for every possible combination of the parameters. Parameters @@ -77,6 +77,7 @@ def create_axes_grid(parameters, labels=None, height_ratios=None, A dictionary mapping the parameter combinations to the axis and their location in the subplots grid; i.e., the key, values are: `{('param1', 'param2'): (pyplot.axes, row index, column index)}` + """ if labels is None: labels = {p: p for p in parameters} @@ -89,12 +90,17 @@ def create_axes_grid(parameters, labels=None, height_ratios=None, if ndim < 3: fsize = (8, 7) else: - fsize = (ndim*3 - 1, ndim*3 - 2) + fsize = (ndim * 3 - 1, ndim * 3 - 2) fig = pyplot.figure(figsize=fsize) # create the axis grid - gs = gridspec.GridSpec(ndim, ndim, width_ratios=width_ratios, - height_ratios=height_ratios, - wspace=0.05, hspace=0.05) + gs = gridspec.GridSpec( + ndim, + ndim, + width_ratios=width_ratios, + height_ratios=height_ratios, + wspace=0.05, + hspace=0.05, + ) # create grid of axis numbers to easily create axes in the right locations axes = numpy.arange(ndim**2).reshape((ndim, ndim)) @@ -113,40 +119,42 @@ def create_axes_grid(parameters, labels=None, height_ratios=None, # map to a parameter index px = parameters[ncolumn] if no_diagonals: - py = parameters[nrow+1] + py = parameters[nrow + 1] else: py = parameters[nrow] if (px, py) in combos: axis_dict[px, py] = (ax, nrow, ncolumn) # x labels only on bottom if nrow + 1 == ndim: - ax.set_xlabel('{}'.format(labels[px]), fontsize=18) + ax.set_xlabel(f"{labels[px]}", fontsize=18) else: pyplot.setp(ax.get_xticklabels(), visible=False) ax.xaxis.offsetText.set_visible(False) # y labels only on left if ncolumn == 0: - ax.set_ylabel('{}'.format(labels[py]), fontsize=18) + ax.set_ylabel(f"{labels[py]}", fontsize=18) else: pyplot.setp(ax.get_yticklabels(), visible=False) ax.yaxis.offsetText.set_visible(False) else: # make non-used axes invisible - ax.axis('off') + ax.axis("off") return fig, axis_dict def get_scale_fac(fig, fiducial_width=8, fiducial_height=7): - """Gets a factor to scale fonts by for the given figure. The scale + """ + Gets a factor to scale fonts by for the given figure. The scale factor is relative to a figure with dimensions (`fiducial_width`, `fiducial_height`). """ width, height = fig.get_size_inches() - return (width*height/(fiducial_width*fiducial_height))**0.5 + return (width * height / (fiducial_width * fiducial_height)) ** 0.5 def construct_kde(samples_array, use_kombine=False, kdeargs=None): - """Constructs a KDE from the given samples. + """ + Constructs a KDE from the given samples. Parameters ---------- @@ -165,6 +173,7 @@ def construct_kde(samples_array, use_kombine=False, kdeargs=None): ------- kde : scipy.stats.gaussian_kde The KDE. + """ # make sure samples are randomly sorted numpy.random.seed(0) @@ -174,7 +183,7 @@ def construct_kde(samples_array, use_kombine=False, kdeargs=None): kdeargs = {} else: kdeargs = kdeargs.copy() - max_nsamples = kdeargs.pop('max_kde_samples', None) + max_nsamples = kdeargs.pop("max_kde_samples", None) samples_array = samples_array[:max_nsamples] if use_kombine: try: @@ -191,15 +200,29 @@ def construct_kde(samples_array, use_kombine=False, kdeargs=None): return kde -def create_density_plot(xparam, yparam, samples, plot_density=True, - plot_contours=True, percentiles=None, cmap='viridis', - contour_color=None, label_contours=True, - contour_linestyles=None, - xmin=None, xmax=None, - ymin=None, ymax=None, exclude_region=None, - fig=None, ax=None, use_kombine=False, - kdeargs=None): - """Computes and plots posterior density and confidence intervals using the +def create_density_plot( + xparam, + yparam, + samples, + plot_density=True, + plot_contours=True, + percentiles=None, + cmap="viridis", + contour_color=None, + label_contours=True, + contour_linestyles=None, + xmin=None, + xmax=None, + ymin=None, + ymax=None, + exclude_region=None, + fig=None, + ax=None, + use_kombine=False, + kdeargs=None, +): + """ + Computes and plots posterior density and confidence intervals using the given samples. Parameters @@ -259,10 +282,11 @@ def create_density_plot(xparam, yparam, samples, plot_density=True, The figure the plot was made on. ax : pyplot.axes The axes the plot was drawn on. + """ if percentiles is None: - percentiles = numpy.array([50., 90.]) - percentiles = 100. - numpy.array(percentiles) + percentiles = numpy.array([50.0, 90.0]) + percentiles = 100.0 - numpy.array(percentiles) percentiles.sort() if ax is None and fig is None: @@ -287,8 +311,9 @@ def create_density_plot(xparam, yparam, samples, plot_density=True, ymax = ysamples.max() npts = 100 X, Y = numpy.mgrid[ - xmin:xmax:complex(0, npts), # pylint:disable=invalid-slice-index - ymin:ymax:complex(0, npts)] # pylint:disable=invalid-slice-index + xmin : xmax : complex(0, npts), # pylint:disable=invalid-slice-index + ymin : ymax : complex(0, npts), + ] # pylint:disable=invalid-slice-index pos = numpy.vstack([X.ravel(), Y.ravel()]) if use_kombine: Z = numpy.exp(kde(pos.T).reshape(X.shape)) @@ -301,13 +326,18 @@ def create_density_plot(xparam, yparam, samples, plot_density=True, # convert X,Y to a single FieldArray so we can use it's ability to # evaluate strings farr = FieldArray.from_kwargs(**{xparam: X, yparam: Y}) - Z[farr[exclude_region]] = 0. + Z[farr[exclude_region]] = 0.0 if plot_density: - ax.imshow(numpy.rot90(Z), extent=[xmin, xmax, ymin, ymax], - aspect='auto', cmap=cmap, zorder=1) + ax.imshow( + numpy.rot90(Z), + extent=[xmin, xmax, ymin, ymax], + aspect="auto", + cmap=cmap, + zorder=1, + ) if contour_color is None: - contour_color = 'w' + contour_color = "w" if plot_contours: # compute the percentile values @@ -316,17 +346,25 @@ def create_density_plot(xparam, yparam, samples, plot_density=True, resamps = numpy.exp(resamps) s = numpy.percentile(resamps, percentiles) if contour_color is None: - contour_color = 'k' + contour_color = "k" # make linewidths thicker if not plotting density for clarity if plot_density: lw = 1 else: lw = 2 - ct = ax.contour(X, Y, Z, s, colors=contour_color, linewidths=lw, - linestyles=contour_linestyles, zorder=3) + ct = ax.contour( + X, + Y, + Z, + s, + colors=contour_color, + linewidths=lw, + linestyles=contour_linestyles, + zorder=3, + ) # label contours if label_contours: - lbls = ['{p}%'.format(p=int(p)) for p in (100. - percentiles)] + lbls = [f"{int(p)}%" for p in (100.0 - percentiles)] fmt = dict(zip(ct.levels, lbls)) fs = 12 ax.clabel(ct, ct.levels, inline=True, fmt=fmt, fontsize=fs) @@ -334,13 +372,26 @@ def create_density_plot(xparam, yparam, samples, plot_density=True, return fig, ax -def create_marginalized_hist(ax, values, label, percentiles=None, - color='k', fillcolor='gray', linecolor='navy', - linestyle='-', plot_marginal_lines=True, - title=True, expected_value=None, - expected_color='red', rotated=False, - plot_min=None, plot_max=None, log_scale=False): - """Plots a 1D marginalized histogram of the given param from the given +def create_marginalized_hist( + ax, + values, + label, + percentiles=None, + color="k", + fillcolor="gray", + linecolor="navy", + linestyle="-", + plot_marginal_lines=True, + title=True, + expected_value=None, + expected_color="red", + rotated=False, + plot_min=None, + plot_max=None, + log_scale=False, +): + """ + Plots a 1D marginalized histogram of the given param from the given samples. Parameters @@ -384,21 +435,20 @@ def create_marginalized_hist(ax, values, label, percentiles=None, Factor to scale the default font sizes by. Default is 1 (no scaling). log_scale : boolean Should the histogram bins be logarithmically spaced + """ if fillcolor is None: - htype = 'step' - fillcolor = 'none' + htype = "step" + fillcolor = "none" else: - htype = 'stepfilled' + htype = "stepfilled" if rotated: - orientation = 'horizontal' + orientation = "horizontal" else: - orientation = 'vertical' + orientation = "vertical" if log_scale: bins = numpy.logspace( - numpy.log10(numpy.nanmin(values)), - numpy.log10(numpy.nanmax(values)), - 50 + numpy.log10(numpy.nanmin(values)), numpy.log10(numpy.nanmax(values)), 50 ) else: bins = numpy.linspace( @@ -406,11 +456,19 @@ def create_marginalized_hist(ax, values, label, percentiles=None, numpy.nanmax(values), 50, ) - ax.hist(values, bins=bins, histtype=htype, orientation=orientation, - facecolor=fillcolor, edgecolor=color, ls=linestyle, lw=2, - density=True) + ax.hist( + values, + bins=bins, + histtype=htype, + orientation=orientation, + facecolor=fillcolor, + edgecolor=color, + ls=linestyle, + lw=2, + density=True, + ) if percentiles is None: - percentiles = [5., 50., 95.] + percentiles = [5.0, 50.0, 95.0] if len(percentiles) > 0: plotp = numpy.percentile(values, percentiles) else: @@ -418,9 +476,9 @@ def create_marginalized_hist(ax, values, label, percentiles=None, if plot_marginal_lines: for val in plotp: if rotated: - ax.axhline(y=val, ls='dashed', color=linecolor, lw=2, zorder=3) + ax.axhline(y=val, ls="dashed", color=linecolor, lw=2, zorder=3) else: - ax.axvline(x=val, ls='dashed', color=linecolor, lw=2, zorder=3) + ax.axvline(x=val, ls="dashed", color=linecolor, lw=2, zorder=3) # plot expected if expected_value is not None: if rotated: @@ -431,7 +489,7 @@ def create_marginalized_hist(ax, values, label, percentiles=None, if len(percentiles) > 0: minp = min(percentiles) maxp = max(percentiles) - medp = (maxp + minp) / 2. + medp = (maxp + minp) / 2.0 else: minp = 5 medp = 50 @@ -441,13 +499,11 @@ def create_marginalized_hist(ax, values, label, percentiles=None, values_max = numpy.percentile(values, maxp) negerror = values_med - values_min poserror = values_max - values_med - fmt = '${0}$'.format(str_utils.format_value( - values_med, negerror, plus_error=poserror)) + fmt = f"${str_utils.format_value(values_med, negerror, plus_error=poserror)}$" if rotated: ax.yaxis.set_label_position("right") # sets colored title for marginal histogram - set_marginal_histogram_title(ax, fmt, color, - label=label, rotated=rotated) + set_marginal_histogram_title(ax, fmt, color, label=label, rotated=rotated) else: # sets colored title for marginal histogram set_marginal_histogram_title(ax, fmt, color, label=label) @@ -456,7 +512,7 @@ def create_marginalized_hist(ax, values, label, percentiles=None, # Remove x-ticks ax.set_xticks([]) # turn off x-labels - ax.set_xlabel('') + ax.set_xlabel("") # set limits ymin, ymax = ax.get_ylim() if plot_min is not None: @@ -468,7 +524,7 @@ def create_marginalized_hist(ax, values, label, percentiles=None, # Remove y-ticks ax.set_yticks([]) # turn off y-label - ax.set_ylabel('') + ax.set_ylabel("") # set limits xmin, xmax = ax.get_xlim() if plot_min is not None: @@ -479,7 +535,8 @@ def create_marginalized_hist(ax, values, label, percentiles=None, def set_marginal_histogram_title(ax, fmt, color, label=None, rotated=False): - """ Sets the title of the marginal histograms. + """ + Sets the title of the marginal histograms. Parameters ---------- @@ -493,8 +550,8 @@ def set_marginal_histogram_title(ax, fmt, color, label=None, rotated=False): If title does not exist, then include label at beginning of the string. rotated : bool If `True` then rotate the text 270 degrees for sideways title. - """ + """ # get rotation angle of the title rotation = 270 if rotated else 0 @@ -512,66 +569,91 @@ def set_marginal_histogram_title(ax, fmt, color, label=None, rotated=False): # if no title exists if not hasattr(ax, "title_boxes"): - # create a text box - title = "{} = {}".format(label, fmt) + title = f"{label} = {fmt}" tbox1 = offsetbox.TextArea( - title, - textprops=dict(color=color, size=15, rotation=rotation, - ha='left', va='bottom')) + title, + textprops=dict( + color=color, size=15, rotation=rotation, ha="left", va="bottom" + ), + ) # save a list of text boxes as attribute for later ax.title_boxes = [tbox1] # pack text boxes - ybox = packer_class(children=ax.title_boxes, - align="bottom", pad=0, sep=5) + ybox = packer_class(children=ax.title_boxes, align="bottom", pad=0, sep=5) # else append existing title else: - # delete old title ax.title_anchor.remove() # add new text box to list tbox1 = offsetbox.TextArea( - " {}".format(fmt), - textprops=dict(color=color, size=15, rotation=rotation, - ha='left', va='bottom')) + f" {fmt}", + textprops=dict( + color=color, size=15, rotation=rotation, ha="left", va="bottom" + ), + ) ax.title_boxes = ax.title_boxes + [tbox1] # pack text boxes - ybox = packer_class(children=ax.title_boxes, - align="bottom", pad=0, sep=5) + ybox = packer_class(children=ax.title_boxes, align="bottom", pad=0, sep=5) # add new title and keep reference to instance as an attribute anchored_ybox = offsetbox.AnchoredOffsetbox( - loc=2, child=ybox, pad=0., - frameon=False, bbox_to_anchor=(xscale, yscale), - bbox_transform=ax.transAxes, borderpad=0.) + loc=2, + child=ybox, + pad=0.0, + frameon=False, + bbox_to_anchor=(xscale, yscale), + bbox_transform=ax.transAxes, + borderpad=0.0, + ) ax.title_anchor = ax.add_artist(anchored_ybox) -def create_multidim_plot(parameters, samples, labels=None, - mins=None, maxs=None, expected_parameters=None, - expected_parameters_color='r', - plot_marginal=True, plot_scatter=True, - plot_maxl=False, - plot_marginal_lines=True, - marginal_percentiles=None, contour_percentiles=None, - marginal_title=True, marginal_linestyle='-', - zvals=None, show_colorbar=True, cbar_label=None, - vmin=None, vmax=None, scatter_cmap='plasma', - scatter_log_cmap=False, log_parameters=None, - plot_density=False, plot_contours=True, - density_cmap='viridis', - contour_color=None, label_contours=True, - contour_linestyles=None, - hist_color='black', - line_color=None, fill_color='gray', - use_kombine=False, kdeargs=None, - fig=None, axis_dict=None): - """Generate a figure with several plots and histograms. +def create_multidim_plot( + parameters, + samples, + labels=None, + mins=None, + maxs=None, + expected_parameters=None, + expected_parameters_color="r", + plot_marginal=True, + plot_scatter=True, + plot_maxl=False, + plot_marginal_lines=True, + marginal_percentiles=None, + contour_percentiles=None, + marginal_title=True, + marginal_linestyle="-", + zvals=None, + show_colorbar=True, + cbar_label=None, + vmin=None, + vmax=None, + scatter_cmap="plasma", + scatter_log_cmap=False, + log_parameters=None, + plot_density=False, + plot_contours=True, + density_cmap="viridis", + contour_color=None, + label_contours=True, + contour_linestyles=None, + hist_color="black", + line_color=None, + fill_color="gray", + use_kombine=False, + kdeargs=None, + fig=None, + axis_dict=None, +): + """ + Generate a figure with several plots and histograms. Parameters ---------- @@ -668,6 +750,7 @@ def create_multidim_plot(parameters, samples, labels=None, A dictionary mapping the parameter combinations to the axis and their location in the subplots grid; i.e., the key, values are: `{('param1', 'param2'): (pyplot.axes, row index, column index)}` + """ if labels is None: labels = {p: p for p in parameters} @@ -692,27 +775,30 @@ def create_multidim_plot(parameters, samples, labels=None, zvals = zvals[sort_indices] samples = samples[sort_indices] if contour_color is None: - contour_color = 'k' + contour_color = "k" elif show_colorbar: raise ValueError("must provide z values to create a colorbar") else: # just make all scatter points same color - zvals = 'gray' + zvals = "gray" if plot_contours and contour_color is None: - contour_color = 'navy' + contour_color = "navy" if plot_maxl: # make sure loglikelihood is provide - if 'loglikelihood' not in samples.fieldnames: + if "loglikelihood" not in samples.fieldnames: raise ValueError("plot-maxl requires loglikelihood") - maxidx = samples['loglikelihood'].argmax() + maxidx = samples["loglikelihood"].argmax() # create the axis grid if fig is None and axis_dict is None: fig, axis_dict = create_axes_grid( - parameters, labels=labels, - width_ratios=width_ratios, height_ratios=height_ratios, - no_diagonals=not plot_marginal) + parameters, + labels=labels, + width_ratios=width_ratios, + height_ratios=height_ratios, + no_diagonals=not plot_marginal, + ) # convert samples to a dictionary to avoid re-computing derived parameters # every time they are needed @@ -744,7 +830,7 @@ def create_multidim_plot(parameters, samples, labels=None, ax, _, _ = axis_dict[param, param] # if only plotting 2 parameters and on the second parameter, # rotate the marginal plot - rotated = nparams == 2 and pi == nparams-1 + rotated = nparams == 2 and pi == nparams - 1 # see if there are expected values if expected_parameters is not None: try: @@ -754,15 +840,23 @@ def create_multidim_plot(parameters, samples, labels=None, else: expected_value = None create_marginalized_hist( - ax, samples[param], label=labels[param], - color=hist_color, fillcolor=fill_color, + ax, + samples[param], + label=labels[param], + color=hist_color, + fillcolor=fill_color, log_scale=param in log_parameters, plot_marginal_lines=plot_marginal_lines, - linestyle=marginal_linestyle, linecolor=line_color, - title=marginal_title, expected_value=expected_value, + linestyle=marginal_linestyle, + linecolor=line_color, + title=marginal_title, + expected_value=expected_value, expected_color=expected_parameters_color, - rotated=rotated, plot_min=mins[param], plot_max=maxs[param], - percentiles=marginal_percentiles) + rotated=rotated, + plot_min=mins[param], + plot_max=maxs[param], + percentiles=marginal_percentiles, + ) # Off-diagonals... for px, py in axis_dict: @@ -773,50 +867,75 @@ def create_multidim_plot(parameters, samples, labels=None, if plot_density: alpha = 0.3 else: - alpha = 1. + alpha = 1.0 if scatter_log_cmap: cmap_norm = colors.LogNorm(vmin=vmin, vmax=vmax) else: cmap_norm = colors.Normalize(vmin=vmin, vmax=vmax) - plt = ax.scatter(x=samples[px], y=samples[py], c=zvals, s=5, - edgecolors='none', norm=cmap_norm, - cmap=scatter_cmap, alpha=alpha, zorder=2) + plt = ax.scatter( + x=samples[px], + y=samples[py], + c=zvals, + s=5, + edgecolors="none", + norm=cmap_norm, + cmap=scatter_cmap, + alpha=alpha, + zorder=2, + ) if plot_contours or plot_density: # Exclude out-of-bound regions # this is a bit kludgy; should probably figure out a better # solution to eventually allow for more than just m_p m_s - if (px == 'm_p' and py == 'm_s') or (py == 'm_p' and px == 'm_s'): - exclude_region = 'm_s > m_p' + if (px == "m_p" and py == "m_s") or (py == "m_p" and px == "m_s"): + exclude_region = "m_s > m_p" else: exclude_region = None create_density_plot( - px, py, samples, plot_density=plot_density, - plot_contours=plot_contours, cmap=density_cmap, + px, + py, + samples, + plot_density=plot_density, + plot_contours=plot_contours, + cmap=density_cmap, percentiles=contour_percentiles, - contour_color=contour_color, label_contours=label_contours, + contour_color=contour_color, + label_contours=label_contours, contour_linestyles=contour_linestyles, - xmin=mins[px], xmax=maxs[px], - ymin=mins[py], ymax=maxs[py], - exclude_region=exclude_region, ax=ax, - use_kombine=use_kombine, kdeargs=kdeargs) + xmin=mins[px], + xmax=maxs[px], + ymin=mins[py], + ymax=maxs[py], + exclude_region=exclude_region, + ax=ax, + use_kombine=use_kombine, + kdeargs=kdeargs, + ) if plot_maxl: maxlx = samples[px][maxidx] maxly = samples[py][maxidx] - ax.scatter(maxlx, maxly, marker='x', s=20, c=contour_color, - zorder=5) + ax.scatter(maxlx, maxly, marker="x", s=20, c=contour_color, zorder=5) if expected_parameters is not None: try: - ax.axvline(expected_parameters[px], lw=1.5, - color=expected_parameters_color, zorder=5) + ax.axvline( + expected_parameters[px], + lw=1.5, + color=expected_parameters_color, + zorder=5, + ) except KeyError: pass try: - ax.axhline(expected_parameters[py], lw=1.5, - color=expected_parameters_color, zorder=5) + ax.axhline( + expected_parameters[py], + lw=1.5, + color=expected_parameters_color, + zorder=5, + ) except KeyError: pass @@ -827,8 +946,8 @@ def create_multidim_plot(parameters, samples, labels=None, if len(parameters) > 3: for px, py in axis_dict: ax, _, _ = axis_dict[px, py] - ax.set_xticks(reduce_ticks(ax, 'x', maxticks=3)) - ax.set_yticks(reduce_ticks(ax, 'y', maxticks=3)) + ax.set_xticks(reduce_ticks(ax, "x", maxticks=3)) + ax.set_yticks(reduce_ticks(ax, "y", maxticks=3)) if plot_scatter and show_colorbar: # compute font size based on fig size @@ -837,14 +956,15 @@ def create_multidim_plot(parameters, samples, labels=None, cbar_ax = fig.add_axes([0.9, 0.1, 0.03, 0.8]) cb = fig.colorbar(plt, cax=cbar_ax) if cbar_label is not None: - cb.set_label(cbar_label, fontsize=12*scale_fac) - cb.ax.tick_params(labelsize=8*scale_fac) + cb.set_label(cbar_label, fontsize=12 * scale_fac) + cb.ax.tick_params(labelsize=8 * scale_fac) return fig, axis_dict def remove_common_offset(arr): - """Given an array of data, removes a common offset > 1000, returning the + """ + Given an array of data, removes a common offset > 1000, returning the removed value. """ offset = 0 @@ -864,7 +984,8 @@ def remove_common_offset(arr): def reduce_ticks(ax, which, maxticks=3): - """Given a pyplot axis, resamples its `which`-axis ticks such that are at most + """ + Given a pyplot axis, resamples its `which`-axis ticks such that are at most `maxticks` left. Parameters @@ -880,12 +1001,13 @@ def reduce_ticks(ax, which, maxticks=3): ------- array An array of the selected ticks. + """ - ticks = getattr(ax, 'get_{}ticks'.format(which))() + ticks = getattr(ax, f"get_{which}ticks")() if len(ticks) > maxticks: # make sure the left/right value is not at the edge - minax, maxax = getattr(ax, 'get_{}lim'.format(which))() - dw = abs(maxax-minax)/10. + minax, maxax = getattr(ax, f"get_{which}lim")() + dw = abs(maxax - minax) / 10.0 start_idx, end_idx = 0, len(ticks) if ticks[0] < minax + dw: start_idx += 1 diff --git a/pycbc/results/snr.py b/pycbc/results/snr.py index e56c16104e9..9d3f402995f 100644 --- a/pycbc/results/snr.py +++ b/pycbc/results/snr.py @@ -26,7 +26,9 @@ """ Module to generate SNR figures """ + from matplotlib import pyplot as plt + from pycbc.results import ifo_color @@ -36,7 +38,6 @@ def generate_snr_plot(snrdict, output_filename, triggers, ref_time): Parameters ---------- - snrdict: dictionary A dictionary keyed on ifo containing the SNR TimeSeries objects @@ -51,21 +52,30 @@ def generate_snr_plot(snrdict, output_filename, triggers, ref_time): Returns ------- None + """ plt.figure() ref_time = int(ref_time) for ifo in sorted(snrdict): curr_snrs = snrdict[ifo] - plt.plot(curr_snrs.sample_times - ref_time, abs(curr_snrs), - c=ifo_color(ifo), label=ifo) + plt.plot( + curr_snrs.sample_times - ref_time, + abs(curr_snrs), + c=ifo_color(ifo), + label=ifo, + ) if ifo in triggers: - plt.plot(triggers[ifo][0] - ref_time, - triggers[ifo][1], marker='x', c=ifo_color(ifo)) + plt.plot( + triggers[ifo][0] - ref_time, + triggers[ifo][1], + marker="x", + c=ifo_color(ifo), + ) plt.legend() - plt.xlabel(f'GPS time from {ref_time:d} (s)') - plt.ylabel('SNR') + plt.xlabel(f"GPS time from {ref_time:d} (s)") + plt.ylabel("SNR") plt.savefig(output_filename) plt.close() diff --git a/pycbc/results/str_utils.py b/pycbc/results/str_utils.py index 08318a0307e..a009f108faa 100644 --- a/pycbc/results/str_utils.py +++ b/pycbc/results/str_utils.py @@ -38,61 +38,74 @@ def mathjax_html_header(): - """Standard header to use for html pages to display latex math. + """ + Standard header to use for html pages to display latex math. Returns ------- header: str The necessary html head needed to use latex on an html page. + """ return mjax_header + def drop_trailing_zeros(num): """ Drops the trailing zeros in a float that is printed. """ - txt = '%f' %(num) - txt = txt.rstrip('0') - if txt.endswith('.'): - txt = txt[:-1] + txt = "%f" % (num) + txt = txt.rstrip("0") + txt = txt.removesuffix(".") return txt + def get_signum(val, err, max_sig=numpy.inf): """ Given an error, returns a string for val formated to the appropriate number of significant figures. """ - coeff, pwr = ('%e' % err).split('e') - if pwr.startswith('-'): + coeff, pwr = ("%e" % err).split("e") + if pwr.startswith("-"): pwr = int(pwr[1:]) - if round(float(coeff)) == 10.: + if round(float(coeff)) == 10.0: pwr -= 1 pwr = min(pwr, max_sig) - tmplt = '%.' + str(pwr+1) + 'f' + tmplt = "%." + str(pwr + 1) + "f" return tmplt % val - else: - pwr = int(pwr[1:]) - if round(float(coeff)) == 10.: - pwr += 1 - # if the error is large, we can sometimes get 0; - # adjust the round until we don't get 0 (assuming the actual - # value isn't 0) - return_val = round(val, -pwr+1) - if val != 0.: - loop_count = 0 - max_recursion = 100 - while return_val == 0.: - pwr -= 1 - return_val = round(val, -pwr+1) - loop_count += 1 - if loop_count > max_recursion: - raise ValueError("Maximum recursion depth hit! Input " +\ - "values are: val = %f, err = %f" %(val, err)) - return drop_trailing_zeros(return_val) - -def format_value(value, error, plus_error=None, use_scientific_notation=3, - include_error=True, use_relative_error=False, ndecs=None): - r"""Given a numerical value and some bound on it, formats the number into a + pwr = int(pwr[1:]) + if round(float(coeff)) == 10.0: + pwr += 1 + # if the error is large, we can sometimes get 0; + # adjust the round until we don't get 0 (assuming the actual + # value isn't 0) + return_val = round(val, -pwr + 1) + if val != 0.0: + loop_count = 0 + max_recursion = 100 + while return_val == 0.0: + pwr -= 1 + return_val = round(val, -pwr + 1) + loop_count += 1 + if loop_count > max_recursion: + raise ValueError( + "Maximum recursion depth hit! Input " + + "values are: val = %f, err = %f" % (val, err) + ) + return drop_trailing_zeros(return_val) + + +def format_value( + value, + error, + plus_error=None, + use_scientific_notation=3, + include_error=True, + use_relative_error=False, + ndecs=None, +): + r""" + Given a numerical value and some bound on it, formats the number into a string such that the value is rounded to the nearest significant figure, which is determined by the error = abs(value-bound). @@ -179,7 +192,7 @@ def format_value(value, error, plus_error=None, use_scientific_notation=3, '3.928\\times 10^{-22}\\,^{+2.1\\%}_{-5.7\\%}' """ - minus_sign = '-' if value < 0. else '' + minus_sign = "-" if value < 0.0 else "" value = abs(value) minus_err = abs(error) if plus_error is None: @@ -187,64 +200,69 @@ def format_value(value, error, plus_error=None, use_scientific_notation=3, else: plus_err = abs(plus_error) error = min(minus_err, plus_err) - if value == 0. or abs(numpy.log10(value)) < use_scientific_notation: - conversion_factor = 0. + if value == 0.0 or abs(numpy.log10(value)) < use_scientific_notation: + conversion_factor = 0.0 else: conversion_factor = numpy.floor(numpy.log10(value)) - value = value * 10**(-conversion_factor) - error = error * 10**(-conversion_factor) - if conversion_factor == 0.: - powfactor = '' - elif conversion_factor == 1.: - powfactor = r'\times 10' + value = value * 10 ** (-conversion_factor) + error = error * 10 ** (-conversion_factor) + if conversion_factor == 0.0: + powfactor = "" + elif conversion_factor == 1.0: + powfactor = r"\times 10" else: - powfactor = r'\times 10^{%i}' %(int(conversion_factor)) + powfactor = r"\times 10^{%i}" % (int(conversion_factor)) if ndecs is not None: - decs = value * 10**(-ndecs) + decs = value * 10 ** (-ndecs) else: decs = error # now round the the appropriate number of sig figs valtxt = get_signum(value, decs) - valtxt = '{}{}'.format(minus_sign, valtxt) + valtxt = f"{minus_sign}{valtxt}" if include_error: if plus_error is None: errtxt = get_signum(error, error) - if use_relative_error and float(valtxt) != 0.: - relative_err = 100.*float(errtxt)/float(valtxt) + if use_relative_error and float(valtxt) != 0.0: + relative_err = 100.0 * float(errtxt) / float(valtxt) # we round the relative error to the nearest 1% using # get_signum; Note that if the relative error is < 1%, # get_signum will automatically increase the number of values # after the decimal until it gets to the first non-zero value - relative_err = get_signum(relative_err, 1.) - txt = r'%s %s \pm%s\%%' %(valtxt, powfactor, relative_err) + relative_err = get_signum(relative_err, 1.0) + txt = r"%s %s \pm%s\%%" % (valtxt, powfactor, relative_err) else: - txt = r'%s \pm %s%s' %(valtxt, errtxt, powfactor) + txt = r"%s \pm %s%s" % (valtxt, errtxt, powfactor) else: - plus_err = plus_err * 10**(-conversion_factor) - minus_err = minus_err * 10**(-conversion_factor) + plus_err = plus_err * 10 ** (-conversion_factor) + minus_err = minus_err * 10 ** (-conversion_factor) minus_err_txt = get_signum(minus_err, decs) plus_err_txt = get_signum(plus_err, decs) - if use_relative_error and float(valtxt) != 0.: + if use_relative_error and float(valtxt) != 0.0: # same as above, but with plus and minus rel_plus_err = get_signum( - 100.*float(plus_err_txt)/float(valtxt), 1.) + 100.0 * float(plus_err_txt) / float(valtxt), 1.0 + ) rel_minus_err = get_signum( - 100.*float(minus_err_txt)/float(valtxt), 1.) - txt = r'%s%s\,^{+%s\%%}_{-%s\%%}' %(valtxt, powfactor, - rel_plus_err, rel_minus_err) + 100.0 * float(minus_err_txt) / float(valtxt), 1.0 + ) + txt = r"%s%s\,^{+%s\%%}_{-%s\%%}" % ( + valtxt, + powfactor, + rel_plus_err, + rel_minus_err, + ) else: - txt = r'%s^{+%s}_{-%s}%s' %(valtxt, plus_err_txt, - minus_err_txt, powfactor) + txt = r"%s^{+%s}_{-%s}%s" % ( + valtxt, + plus_err_txt, + minus_err_txt, + powfactor, + ) else: - txt = r'%s%s' %(valtxt, powfactor) + txt = r"%s%s" % (valtxt, powfactor) return txt -__all__ = [ - "mathjax_html_header", - "drop_trailing_zeros", - "get_signum", - "format_value" -] +__all__ = ["drop_trailing_zeros", "format_value", "get_signum", "mathjax_html_header"] diff --git a/pycbc/results/table_utils.py b/pycbc/results/table_utils.py index e74ebde3a56..cbbce10c505 100644 --- a/pycbc/results/table_utils.py +++ b/pycbc/results/table_utils.py @@ -21,11 +21,12 @@ # # ============================================================================= # -""" This module provides functions to generate sortable html tables -""" -import mako.template -import uuid +"""This module provides functions to generate sortable html tables""" + import copy +import uuid + +import mako.template import numpy google_table_template = mako.template.Template(""" @@ -58,8 +59,10 @@
""") + def html_table(columns, names, page_size=None, format_strings=None): - """ Return an HTML table of this data. + """ + Return an HTML table of this data. Parameters ---------- @@ -76,36 +79,37 @@ def html_table(columns, names, page_size=None, format_strings=None): ------- html_table : str A str containing the html code to display a table of this data + """ if len(columns) != len(names): raise ValueError( - 'I need the same number of columns and names, ' - f'got {len(columns)} and {len(names)} instead' + "I need the same number of columns and names, " + f"got {len(columns)} and {len(names)} instead" ) if format_strings is not None and len(format_strings) != len(columns): raise ValueError( - 'I need the same number of columns and format strings, ' - f'got {len(columns)} and {len(names)} instead' + "I need the same number of columns and format strings, " + f"got {len(columns)} and {len(names)} instead" ) if len({len(column) for column in columns}) != 1: - raise ValueError('All columns must have the same length') + raise ValueError("All columns must have the same length") if page_size is None: - page = 'disable' + page = "disable" else: - page = 'enable' + page = "enable" div_id = uuid.uuid4() column_descriptions = [] for column, name in zip(columns, names): - if column.dtype.kind in 'iuf': + if column.dtype.kind in "iuf": # signed and unsigned integers and floats - ctype = 'number' + ctype = "number" else: # this comprises strings, bools, complex, void, etc # but we will convert all those to str in a moment - ctype = 'string' + ctype = "string" column_descriptions.append((ctype, name)) data = [] @@ -114,9 +118,9 @@ def html_table(columns, names, page_size=None, format_strings=None): # the explicit conversions here are to make sure the JS code # sees proper numbers and not something like 'np.float64(12)' for item, column in zip(row, columns): - if column.dtype.kind == 'f': + if column.dtype.kind == "f": row2.append(float(item)) - elif column.dtype.kind in 'iu': + elif column.dtype.kind in "iu": row2.append(int(item)) else: row2.append(str(item)) @@ -131,6 +135,7 @@ def html_table(columns, names, page_size=None, format_strings=None): format_strings=format_strings, ) + static_table_template = mako.template.Template(""" % for row in range(n_rows): @@ -166,8 +171,10 @@ def html_table(columns, names, page_size=None, format_strings=None):
""") + def static_table(data, titles=None, columns_max=None, row_labels=None): - """ Return an html table of this data + """ + Return an html table of this data Parameters ---------- @@ -185,6 +192,7 @@ def static_table(data, titles=None, columns_max=None, row_labels=None): ------- html_table : str A string containing the html table. + """ data = copy.deepcopy(data) titles = copy.deepcopy(titles) @@ -194,9 +202,7 @@ def static_table(data, titles=None, columns_max=None, row_labels=None): raise ValueError("titles and data lengths do not match") if row_labels is not None and not len(row_labels) == drows: - raise ValueError( - "row_labels must be the same number of rows supplied to data" - ) + raise ValueError("row_labels must be the same number of rows supplied to data") if columns_max is not None: n_rows = int(numpy.ceil(len(data[0]) / columns_max)) @@ -204,9 +210,9 @@ def static_table(data, titles=None, columns_max=None, row_labels=None): if len(data[0]) < n_rows * n_columns: # Pad the data and titles with empty strings n_missing = int(n_rows * n_columns - len(data[0])) - data = numpy.hstack((data, numpy.zeros((len(data), n_missing), dtype='U1'))) + data = numpy.hstack((data, numpy.zeros((len(data), n_missing), dtype="U1"))) if titles is not None: - titles += [' '] * n_missing + titles += [" "] * n_missing else: n_rows = 1 n_columns = len(data[0]) diff --git a/pycbc/results/versioning.py b/pycbc/results/versioning.py index ebc47a4130e..a340cc4d5b5 100644 --- a/pycbc/results/versioning.py +++ b/pycbc/results/versioning.py @@ -21,114 +21,120 @@ import pycbc.version from pycbc.libutils import import_optional -lal = import_optional('lal') -lalframe = import_optional('lalframe') -lalsimulation = import_optional('lalsimulation') +lal = import_optional("lal") +lalframe = import_optional("lalframe") +lalsimulation = import_optional("lalsimulation") -logger = logging.getLogger('pycbc.results.versioning') +logger = logging.getLogger("pycbc.results.versioning") + def get_library_version_info(): - """This will return a list of dictionaries containing versioning + """ + This will return a list of dictionaries containing versioning information about the various LIGO libraries that PyCBC will use in an - analysis run.""" + analysis run. + """ library_list = [] def add_info_new_version(info_dct, curr_module, extra_str): - vcs_object = getattr(curr_module, extra_str +'VCSInfo') - info_dct['ID'] = vcs_object.vcsId - info_dct['Status'] = vcs_object.vcsStatus - info_dct['Version'] = vcs_object.version - info_dct['Tag'] = vcs_object.vcsTag - info_dct['Author'] = vcs_object.vcsAuthor - info_dct['Branch'] = vcs_object.vcsBranch - info_dct['Committer'] = vcs_object.vcsCommitter - info_dct['Date'] = vcs_object.vcsDate + vcs_object = getattr(curr_module, extra_str + "VCSInfo") + info_dct["ID"] = vcs_object.vcsId + info_dct["Status"] = vcs_object.vcsStatus + info_dct["Version"] = vcs_object.version + info_dct["Tag"] = vcs_object.vcsTag + info_dct["Author"] = vcs_object.vcsAuthor + info_dct["Branch"] = vcs_object.vcsBranch + info_dct["Committer"] = vcs_object.vcsCommitter + info_dct["Date"] = vcs_object.vcsDate if lal is not None: lalinfo = {} - lalinfo['Name'] = 'LAL' + lalinfo["Name"] = "LAL" try: - lalinfo['ID'] = lal.VCSId - lalinfo['Status'] = lal.VCSStatus - lalinfo['Version'] = lal.VCSVersion - lalinfo['Tag'] = lal.VCSTag - lalinfo['Author'] = lal.VCSAuthor - lalinfo['Branch'] = lal.VCSBranch - lalinfo['Committer'] = lal.VCSCommitter - lalinfo['Date'] = lal.VCSDate + lalinfo["ID"] = lal.VCSId + lalinfo["Status"] = lal.VCSStatus + lalinfo["Version"] = lal.VCSVersion + lalinfo["Tag"] = lal.VCSTag + lalinfo["Author"] = lal.VCSAuthor + lalinfo["Branch"] = lal.VCSBranch + lalinfo["Committer"] = lal.VCSCommitter + lalinfo["Date"] = lal.VCSDate except AttributeError: - add_info_new_version(lalinfo, lal, '') + add_info_new_version(lalinfo, lal, "") library_list.append(lalinfo) if lalframe is not None: lalframeinfo = {} try: - lalframeinfo['Name'] = 'LALFrame' - lalframeinfo['ID'] = lalframe.FrameVCSId - lalframeinfo['Status'] = lalframe.FrameVCSStatus - lalframeinfo['Version'] = lalframe.FrameVCSVersion - lalframeinfo['Tag'] = lalframe.FrameVCSTag - lalframeinfo['Author'] = lalframe.FrameVCSAuthor - lalframeinfo['Branch'] = lalframe.FrameVCSBranch - lalframeinfo['Committer'] = lalframe.FrameVCSCommitter - lalframeinfo['Date'] = lalframe.FrameVCSDate + lalframeinfo["Name"] = "LALFrame" + lalframeinfo["ID"] = lalframe.FrameVCSId + lalframeinfo["Status"] = lalframe.FrameVCSStatus + lalframeinfo["Version"] = lalframe.FrameVCSVersion + lalframeinfo["Tag"] = lalframe.FrameVCSTag + lalframeinfo["Author"] = lalframe.FrameVCSAuthor + lalframeinfo["Branch"] = lalframe.FrameVCSBranch + lalframeinfo["Committer"] = lalframe.FrameVCSCommitter + lalframeinfo["Date"] = lalframe.FrameVCSDate except AttributeError: - add_info_new_version(lalframeinfo, lalframe, 'Frame') + add_info_new_version(lalframeinfo, lalframe, "Frame") library_list.append(lalframeinfo) if lalsimulation is not None: lalsimulationinfo = {} - lalsimulationinfo['Name'] = 'LALSimulation' + lalsimulationinfo["Name"] = "LALSimulation" try: - lalsimulationinfo['ID'] = lalsimulation.SimulationVCSId - lalsimulationinfo['Status'] = lalsimulation.SimulationVCSStatus - lalsimulationinfo['Version'] = lalsimulation.SimulationVCSVersion - lalsimulationinfo['Tag'] = lalsimulation.SimulationVCSTag - lalsimulationinfo['Author'] = lalsimulation.SimulationVCSAuthor - lalsimulationinfo['Branch'] = lalsimulation.SimulationVCSBranch - lalsimulationinfo['Committer'] = lalsimulation.SimulationVCSCommitter - lalsimulationinfo['Date'] = lalsimulation.SimulationVCSDate + lalsimulationinfo["ID"] = lalsimulation.SimulationVCSId + lalsimulationinfo["Status"] = lalsimulation.SimulationVCSStatus + lalsimulationinfo["Version"] = lalsimulation.SimulationVCSVersion + lalsimulationinfo["Tag"] = lalsimulation.SimulationVCSTag + lalsimulationinfo["Author"] = lalsimulation.SimulationVCSAuthor + lalsimulationinfo["Branch"] = lalsimulation.SimulationVCSBranch + lalsimulationinfo["Committer"] = lalsimulation.SimulationVCSCommitter + lalsimulationinfo["Date"] = lalsimulation.SimulationVCSDate except AttributeError: - add_info_new_version(lalsimulationinfo, lalsimulation, 'Simulation') + add_info_new_version(lalsimulationinfo, lalsimulation, "Simulation") library_list.append(lalsimulationinfo) pycbcinfo = {} - pycbcinfo['Name'] = 'PyCBC' - pycbcinfo['ID'] = pycbc.version.git_hash - pycbcinfo['Status'] = pycbc.version.git_status - pycbcinfo['Version'] = pycbc.version.version + pycbcinfo["Name"] = "PyCBC" + pycbcinfo["ID"] = pycbc.version.git_hash + pycbcinfo["Status"] = pycbc.version.git_status + pycbcinfo["Version"] = pycbc.version.version if pycbc.version.release: - pycbcinfo['Version'] += ' (release)' - pycbcinfo['Tag'] = pycbc.version.git_tag - pycbcinfo['Author'] = pycbc.version.git_author - pycbcinfo['Builder'] = pycbc.version.git_builder - pycbcinfo['Branch'] = pycbc.version.git_branch - pycbcinfo['Committer'] = pycbc.version.git_committer - pycbcinfo['Date'] = pycbc.version.git_build_date + pycbcinfo["Version"] += " (release)" + pycbcinfo["Tag"] = pycbc.version.git_tag + pycbcinfo["Author"] = pycbc.version.git_author + pycbcinfo["Builder"] = pycbc.version.git_builder + pycbcinfo["Branch"] = pycbc.version.git_branch + pycbcinfo["Committer"] = pycbc.version.git_committer + pycbcinfo["Date"] = pycbc.version.git_build_date library_list.append(pycbcinfo) return library_list + def get_code_version_numbers(executable_names, executable_files): - """Will extract the version information from the executables listed in + """ + Will extract the version information from the executables listed in the executable section of the supplied ConfigParser object. Returns - -------- + ------- dict A dictionary keyed by the executable name with values giving the version string for each executable. + """ code_version_dict = {} for exe_name, value in zip(executable_names, executable_files): value = urllib.parse.urlparse(value) logger.info("Getting version info for %s", exe_name) version_string = None - if value.scheme in ['gsiftp', 'http', 'https']: + if value.scheme in ["gsiftp", "http", "https"]: code_version_dict[exe_name] = "Using bundle downloaded from %s" % value - elif value.scheme == 'singularity': + elif value.scheme == "singularity": txt = ( "Executable run from a singularity image. See config file " "and site catalog for details of what image was used." @@ -137,8 +143,7 @@ def get_code_version_numbers(executable_names, executable_files): else: try: version_string = subprocess.check_output( - [value.path, '--version'], - stderr=subprocess.STDOUT + [value.path, "--version"], stderr=subprocess.STDOUT ).decode() except subprocess.CalledProcessError: version_string = "Executable fails on {} --version" diff --git a/pycbc/scheme.py b/pycbc/scheme.py index 9d27d146792..20a38732c85 100644 --- a/pycbc/scheme.py +++ b/pycbc/scheme.py @@ -26,33 +26,36 @@ This modules provides python contexts that set the default behavior for PyCBC objects. """ + +import logging import os -import pycbc from functools import wraps -import logging + +import pycbc + from .libutils import get_ctypes_library from .pool import use_mpi -logger = logging.getLogger('pycbc.scheme') +logger = logging.getLogger("pycbc.scheme") -class _SchemeManager(object): +class _SchemeManager: _single = None def __init__(self): if _SchemeManager._single is not None: raise RuntimeError("SchemeManager is a private class") - _SchemeManager._single= self + _SchemeManager._single = self - self.state= None - self._lock= False + self.state = None + self._lock = False def lock(self): - self._lock= True + self._lock = True def unlock(self): - self._lock= False + self._lock = False def shift_to(self, state): if self._lock is False: @@ -60,65 +63,82 @@ def shift_to(self, state): else: raise RuntimeError("The state is locked, cannot shift schemes") + # Create the global processing scheme manager mgr = _SchemeManager() DefaultScheme = None default_context = None -class Scheme(object): - """Context that sets PyCBC objects to use CPU processing. """ +class Scheme: + """Context that sets PyCBC objects to use CPU processing.""" + _single = None + def __init__(self): if DefaultScheme is type(self): return if Scheme._single is not None: raise RuntimeError("Only one processing scheme can be used") Scheme._single = True + def __enter__(self): mgr.shift_to(self) mgr.lock() + def __exit__(self, type, value, traceback): mgr.unlock() mgr.shift_to(default_context) + def __del__(self): if Scheme is not None: Scheme._single = None -_cuda_cleanup_list=[] + +_cuda_cleanup_list = [] + def register_clean_cuda(function): _cuda_cleanup_list.append(function) + def clean_cuda(context): - #Before cuda context is destroyed, all item destructions dependent on cuda + # Before cuda context is destroyed, all item destructions dependent on cuda # must take place. This calls all functions that have been registered # with _register_clean_cuda() in reverse order - #So the last one registered, is the first one cleaned + # So the last one registered, is the first one cleaned _cuda_cleanup_list.reverse() for func in _cuda_cleanup_list: func() context.pop() from pycuda.tools import clear_context_caches + clear_context_caches() + class CUDAScheme(Scheme): - """Context that sets PyCBC objects to use a CUDA processing scheme. """ + """Context that sets PyCBC objects to use a CUDA processing scheme.""" + def __init__(self, device_num=0): Scheme.__init__(self) if not pycbc.HAVE_CUDA: raise RuntimeError("Install PyCUDA to use CUDA processing") import pycuda.driver + pycuda.driver.init() self.device = pycuda.driver.Device(device_num) - self.context = self.device.make_context(flags=pycuda.driver.ctx_flags.SCHED_BLOCKING_SYNC) + self.context = self.device.make_context( + flags=pycuda.driver.ctx_flags.SCHED_BLOCKING_SYNC + ) import atexit - atexit.register(clean_cuda,self.context) + + atexit.register(clean_cuda, self.context) class CUPYScheme(Scheme): - """Scheme for using CUPY. + """ + Scheme for using CUPY. Supports using CUPY with MPI. If MPI is enabled, will use all available devices. The environment variable `CUDA_VISIBLE_DEVICES` can be used to @@ -129,9 +149,11 @@ class CUPYScheme(Scheme): device_num : int, optional The device number to use. If not provided, will use the default, 0. Should not be provided when using MPI to parallelize across devices. + """ + def __init__(self, device_num=None): - import cupy # Fail now if cupy is not there. + import cupy # Fail now if cupy is not there. import cupy.cuda do_mpi, _, rank = use_mpi(require_mpi=False, log=False) @@ -167,19 +189,21 @@ def __exit__(self, *args): class CPUScheme(Scheme): def __init__(self, num_threads=1): if isinstance(num_threads, int): - self.num_threads=num_threads - elif num_threads == 'env' and "PYCBC_NUM_THREADS" in os.environ: + self.num_threads = num_threads + elif num_threads == "env" and "PYCBC_NUM_THREADS" in os.environ: self.num_threads = int(os.environ["PYCBC_NUM_THREADS"]) else: import multiprocessing + self.num_threads = multiprocessing.cpu_count() self._libgomp = None def __enter__(self): Scheme.__enter__(self) try: - self._libgomp = get_ctypes_library("gomp", ['gomp'], - mode=ctypes.RTLD_GLOBAL) + self._libgomp = get_ctypes_library( + "gomp", ["gomp"], mode=ctypes.RTLD_GLOBAL + ) except: # Should we fail or give a warning if we cannot import # libgomp? Seems to work even for MKL scheme, but @@ -188,7 +212,7 @@ def __enter__(self): os.environ["OMP_NUM_THREADS"] = str(self.num_threads) if self._libgomp is not None: - self._libgomp.omp_set_num_threads( int(self.num_threads) ) + self._libgomp.omp_set_num_threads(int(self.num_threads)) def __exit__(self, type, value, traceback): os.environ["OMP_NUM_THREADS"] = "1" @@ -196,12 +220,14 @@ def __exit__(self, type, value, traceback): self._libgomp.omp_set_num_threads(1) Scheme.__exit__(self, type, value, traceback) + class MKLScheme(CPUScheme): def __init__(self, num_threads=1): CPUScheme.__init__(self, num_threads) if not pycbc.HAVE_MKL: raise RuntimeError("Can't find MKL libraries") + class NumpyScheme(CPUScheme): pass @@ -218,7 +244,7 @@ class NumpyScheme(CPUScheme): _default_scheme_prefix = os.getenv("PYCBC_SCHEME", "cpu") try: _default_scheme_class = _scheme_map[_default_scheme_prefix] -except KeyError as exc: +except KeyError: raise RuntimeError( "PYCBC_SCHEME={!r} not recognised, please select one of: {}".format( _default_scheme_prefix, @@ -226,17 +252,23 @@ class NumpyScheme(CPUScheme): ), ) + class DefaultScheme(_default_scheme_class): pass + default_context = DefaultScheme() mgr.state = default_context scheme_prefix[DefaultScheme] = _default_scheme_prefix + def current_prefix(): return scheme_prefix[type(mgr.state)] + _import_cache = {} + + def schemed(prefix): def scheming_function(func): @@ -248,8 +280,9 @@ def _scheming_function(*args, **kwds): exc_errors = [] for sch in mgr.state.__class__.__mro__[0:-2]: try: - backend = __import__(prefix + scheme_prefix[sch], - fromlist=[func.__name__]) + backend = __import__( + prefix + scheme_prefix[sch], fromlist=[func.__name__] + ) schemed_fn = getattr(backend, func.__name__) except (ImportError, AttributeError) as e: exc_errors += [e] @@ -262,25 +295,31 @@ def _scheming_function(*args, **kwds): return schemed_fn(*args, **kwds) - err = (f"Failed to find implementation of {func.__name__} " - f"for {current_prefix()} scheme. ") + err = ( + f"Failed to find implementation of {func.__name__} " + f"for {current_prefix()} scheme. " + ) for emsg in exc_errors: err += str(emsg) + " " raise RuntimeError(err) + return _scheming_function return scheming_function + def cpuonly(func): @wraps(func) def _cpuonly(*args, **kwds): if not issubclass(type(mgr.state), CPUScheme): - raise TypeError(fn.__name__ + - " can only be called from a CPU processing scheme.") - else: - return func(*args, **kwds) + raise TypeError( + fn.__name__ + " can only be called from a CPU processing scheme." + ) + return func(*args, **kwds) + return _cpuonly + def insert_processing_option_group(parser): """ Adds the options used to choose a processing scheme. This should be used @@ -290,31 +329,40 @@ def insert_processing_option_group(parser): ---------- parser : object OptionParser instance + """ - processing_group = parser.add_argument_group("Options for selecting the" - " processing scheme in this program.") - processing_group.add_argument("--processing-scheme", - help="The choice of processing scheme. " - "Choices are " + str(list(set(scheme_prefix.values()))) + - ". (optional for CPU scheme) The number of " - "execution threads " - "can be indicated by cpu:NUM_THREADS, " - "where NUM_THREADS " - "is an integer. The default is a single thread. " - "If the scheme is provided as cpu:env, the number " - "of threads can be provided by the PYCBC_NUM_THREADS " - "environment variable. If the environment variable " - "is not set, the number of threads matches the number " - "of logical cores. ", - default="cpu") - - processing_group.add_argument("--processing-device-id", - help="(optional) ID of GPU to use for accelerated " - "processing", - default=0, type=int) + processing_group = parser.add_argument_group( + "Options for selecting the processing scheme in this program." + ) + processing_group.add_argument( + "--processing-scheme", + help="The choice of processing scheme. " + "Choices are " + + str(list(set(scheme_prefix.values()))) + + ". (optional for CPU scheme) The number of " + "execution threads " + "can be indicated by cpu:NUM_THREADS, " + "where NUM_THREADS " + "is an integer. The default is a single thread. " + "If the scheme is provided as cpu:env, the number " + "of threads can be provided by the PYCBC_NUM_THREADS " + "environment variable. If the environment variable " + "is not set, the number of threads matches the number " + "of logical cores. ", + default="cpu", + ) + + processing_group.add_argument( + "--processing-device-id", + help="(optional) ID of GPU to use for accelerated processing", + default=0, + type=int, + ) + def from_cli(opt): - """Parses the command line options and returns a processing scheme. + """ + Parses the command line options and returns a processing scheme. Parameters ---------- @@ -326,8 +374,9 @@ def from_cli(opt): ------- ctx: Scheme Returns the requested processing scheme. + """ - scheme_str = opt.processing_scheme.split(':') + scheme_str = opt.processing_scheme.split(":") name = scheme_str[0] if name == "cuda": @@ -342,7 +391,7 @@ def from_cli(opt): else: ctx = MKLScheme() logger.info("Running with MKL support: %s threads" % ctx.num_threads) - elif name == 'cupy': + elif name == "cupy": logger.info("Running with CUPY support") ctx = CUPYScheme() else: @@ -356,8 +405,10 @@ def from_cli(opt): logger.info("Running with CPU support: %s threads" % ctx.num_threads) return ctx + def verify_processing_options(opt, parser): - """Parses the processing scheme options and verifies that they are + """ + Parses the processing scheme options and verifies that they are reasonable. @@ -368,16 +419,20 @@ def verify_processing_options(opt, parser): required attributes. parser : object OptionParser instance. + """ scheme_types = scheme_prefix.values() - if opt.processing_scheme.split(':')[0] not in scheme_types: + if opt.processing_scheme.split(":")[0] not in scheme_types: parser.error("(%s) is not a valid scheme type.") + class ChooseBySchemeDict(dict): - """ This class represents a dictionary whose purpose is to chose objects + """ + This class represents a dictionary whose purpose is to chose objects based on their processing scheme. The keys are intended to be processing schemes. """ + def __getitem__(self, scheme): for base in scheme.__mro__[0:-1]: try: @@ -385,4 +440,3 @@ def __getitem__(self, scheme): break except: pass - diff --git a/pycbc/sensitivity.py b/pycbc/sensitivity.py index d7232593608..397725be2ca 100644 --- a/pycbc/sensitivity.py +++ b/pycbc/sensitivity.py @@ -1,9 +1,11 @@ -""" This module contains utilities for calculating search sensitivity -""" -import numpy +"""This module contains utilities for calculating search sensitivity""" + import logging +import numpy + from pycbc.conversions import chirp_distance + from . import bin_utils # numpy renamed trapz to trapezoid in 2.0 and removed trapz in 2.x @@ -12,12 +14,12 @@ except ImportError: # numpy < 2.0 from numpy import trapz as trapezoid -logger = logging.getLogger('pycbc.sensitivity') +logger = logging.getLogger("pycbc.sensitivity") def compute_search_efficiency_in_bins( - found, total, ndbins, - sim_to_bins_function=lambda sim: (sim.distance,)): + found, total, ndbins, sim_to_bins_function=lambda sim: (sim.distance,) +): """ Calculate search efficiency in the given ndbins. @@ -37,7 +39,7 @@ def compute_search_efficiency_in_bins( eff = bin_utils.BinnedArray(bin_utils.NDBins(ndbins), array=bins.ratio()) # compute binomial uncertainties in each bin - err_arr = numpy.sqrt(eff.array * (1-eff.array)/bins.denominator.array) + err_arr = numpy.sqrt(eff.array * (1 - eff.array) / bins.denominator.array) err = bin_utils.BinnedArray(bin_utils.NDBins(ndbins), array=err_arr) return eff, err @@ -52,7 +54,8 @@ def compute_search_volume_in_bins(found, total, ndbins, sim_to_bins_function): sim_to_bins_function must maps an object to a tuple indexing the ndbins. """ eff, err = compute_search_efficiency_in_bins( - found, total, ndbins, sim_to_bins_function) + found, total, ndbins, sim_to_bins_function + ) dx = ndbins[0].upper() - ndbins[0].lower() r = ndbins[0].centres() @@ -61,18 +64,19 @@ def compute_search_volume_in_bins(found, total, ndbins, sim_to_bins_function): errors = bin_utils.BinnedArray(bin_utils.NDBins(ndbins[1:])) # integrate efficiency to obtain volume - vol.array = trapezoid(eff.array.T * 4. * numpy.pi * r**2, r, dx) + vol.array = trapezoid(eff.array.T * 4.0 * numpy.pi * r**2, r, dx) # propagate errors in eff to errors in V errors.array = numpy.sqrt( - ((4 * numpy.pi * r**2 * err.array.T * dx)**2).sum(axis=-1) + ((4 * numpy.pi * r**2 * err.array.T * dx) ** 2).sum(axis=-1) ) return vol, errors def volume_to_distance_with_errors(vol, vol_err): - """ Return the distance and standard deviation upper and lower bounds + """ + Return the distance and standard deviation upper and lower bounds Parameters ---------- @@ -86,16 +90,24 @@ def volume_to_distance_with_errors(vol, vol_err): elow: float """ - dist = (vol * 3.0/4.0/numpy.pi) ** (1.0/3.0) - ehigh = ((vol + vol_err) * 3.0/4.0/numpy.pi) ** (1.0/3.0) - dist + dist = (vol * 3.0 / 4.0 / numpy.pi) ** (1.0 / 3.0) + ehigh = ((vol + vol_err) * 3.0 / 4.0 / numpy.pi) ** (1.0 / 3.0) - dist delta = numpy.where(vol >= vol_err, vol - vol_err, 0) - elow = dist - (delta * 3.0/4.0/numpy.pi) ** (1.0/3.0) + elow = dist - (delta * 3.0 / 4.0 / numpy.pi) ** (1.0 / 3.0) return dist, ehigh, elow -def volume_montecarlo(found_d, missed_d, found_mchirp, missed_mchirp, - distribution_param, distribution, limits_param, - min_param=None, max_param=None): +def volume_montecarlo( + found_d, + missed_d, + found_mchirp, + missed_mchirp, + distribution_param, + distribution, + limits_param, + min_param=None, + max_param=None, +): """ Compute sensitive volume and standard error via direct Monte Carlo integral @@ -106,7 +118,7 @@ def volume_montecarlo(found_d, missed_d, found_mchirp, missed_mchirp, OR get that coded as a new function? Parameters - ----------- + ---------- found_d: numpy.ndarray The distances of found injections missed_d: numpy.ndarray @@ -136,63 +148,58 @@ def volume_montecarlo(found_d, missed_d, found_mchirp, missed_mchirp, the maximum actually injected value will be used Returns - -------- + ------- volume: float Volume estimate volume_error: float The standard error in the volume + """ - d_power = { - 'log' : 3., - 'uniform' : 2., - 'distancesquared' : 1., - 'volume' : 0. - }[distribution] + d_power = {"log": 3.0, "uniform": 2.0, "distancesquared": 1.0, "volume": 0.0}[ + distribution + ] mchirp_power = { - 'log' : 0., - 'uniform' : 5. / 6., - 'distancesquared' : 5. / 3., - 'volume' : 15. / 6. + "log": 0.0, + "uniform": 5.0 / 6.0, + "distancesquared": 5.0 / 3.0, + "volume": 15.0 / 6.0, }[distribution] # establish maximum physical distance: first for chirp distance distribution - if limits_param == 'chirp_distance': - mchirp_standard_bns = 1.4 * 2.**(-1. / 5.) + if limits_param == "chirp_distance": + mchirp_standard_bns = 1.4 * 2.0 ** (-1.0 / 5.0) all_mchirp = numpy.concatenate((found_mchirp, missed_mchirp)) max_mchirp = all_mchirp.max() if max_param is not None: # use largest injected mchirp to convert back to distance - max_distance = max_param * \ - (max_mchirp / mchirp_standard_bns)**(5. / 6.) + max_distance = max_param * (max_mchirp / mchirp_standard_bns) ** (5.0 / 6.0) else: max_distance = max(found_d.max(), missed_d.max()) - elif limits_param == 'distance': + elif limits_param == "distance": if max_param is not None: max_distance = max_param else: # if no max distance given, use max distance actually injected max_distance = max(found_d.max(), missed_d.max()) else: - raise NotImplementedError("%s is not a recognized parameter" - % limits_param) + raise NotImplementedError("%s is not a recognized parameter" % limits_param) # volume of sphere - montecarlo_vtot = (4. / 3.) * numpy.pi * max_distance**3. + montecarlo_vtot = (4.0 / 3.0) * numpy.pi * max_distance**3.0 # arrays of weights for the MC integral - if distribution_param == 'distance': - found_weights = found_d ** d_power - missed_weights = missed_d ** d_power - elif distribution_param == 'chirp_distance': + if distribution_param == "distance": + found_weights = found_d**d_power + missed_weights = missed_d**d_power + elif distribution_param == "chirp_distance": # weight by a power of mchirp to rescale injection density to the # target mass distribution - found_weights = found_d ** d_power * \ - found_mchirp ** mchirp_power - missed_weights = missed_d ** d_power * \ - missed_mchirp ** mchirp_power + found_weights = found_d**d_power * found_mchirp**mchirp_power + missed_weights = missed_d**d_power * missed_mchirp**mchirp_power else: - raise NotImplementedError("%s is not a recognized distance parameter" - % distribution_param) + raise NotImplementedError( + "%s is not a recognized distance parameter" % distribution_param + ) all_weights = numpy.concatenate((found_weights, missed_weights)) @@ -202,39 +209,40 @@ def volume_montecarlo(found_d, missed_d, found_mchirp, missed_mchirp, mc_weight_samples = numpy.concatenate((found_weights, 0 * missed_weights)) mc_sum = sum(mc_weight_samples) - if limits_param == 'distance': + if limits_param == "distance": mc_norm = sum(all_weights) - elif limits_param == 'chirp_distance': + elif limits_param == "chirp_distance": # if injections are made up to a maximum chirp distance, account for # extra missed injections that would occur when injecting up to # maximum physical distance : this works out to a 'chirp volume' factor - mc_norm = sum(all_weights * (max_mchirp / all_mchirp) ** (5. / 2.)) + mc_norm = sum(all_weights * (max_mchirp / all_mchirp) ** (5.0 / 2.0)) # take out a constant factor mc_prefactor = montecarlo_vtot / mc_norm # count the samples - if limits_param == 'distance': + if limits_param == "distance": Ninj = len(mc_weight_samples) - elif limits_param == 'chirp_distance': + elif limits_param == "chirp_distance": # find the total expected number after extending from maximum chirp # dist up to maximum physical distance - if distribution == 'log': + if distribution == "log": # only need minimum distance in this one case if min_param is not None: - min_distance = min_param * \ - (numpy.min(all_mchirp) / mchirp_standard_bns) ** (5. / 6.) + min_distance = min_param * ( + numpy.min(all_mchirp) / mchirp_standard_bns + ) ** (5.0 / 6.0) else: min_distance = min(numpy.min(found_d), numpy.min(missed_d)) logrange = numpy.log(max_distance / min_distance) - Ninj = len(mc_weight_samples) + (5. / 6.) * \ - sum(numpy.log(max_mchirp / all_mchirp) / logrange) + Ninj = len(mc_weight_samples) + (5.0 / 6.0) * sum( + numpy.log(max_mchirp / all_mchirp) / logrange + ) else: Ninj = sum((max_mchirp / all_mchirp) ** mchirp_power) # sample variance of efficiency: mean of the square - square of the mean - mc_sample_variance = sum(mc_weight_samples ** 2.) / Ninj - \ - (mc_sum / Ninj) ** 2. + mc_sample_variance = sum(mc_weight_samples**2.0) / Ninj - (mc_sum / Ninj) ** 2.0 # return MC integral and its standard deviation; variance of mc_sum scales # relative to sample variance by Ninj (Bienayme' rule) @@ -244,63 +252,88 @@ def volume_montecarlo(found_d, missed_d, found_mchirp, missed_mchirp, def chirp_volume_montecarlo( - found_d, missed_d, found_mchirp, missed_mchirp, - distribution_param, distribution, limits_param, min_param, max_param): - - assert distribution_param == 'chirp_distance' - assert limits_param == 'chirp_distance' + found_d, + missed_d, + found_mchirp, + missed_mchirp, + distribution_param, + distribution, + limits_param, + min_param, + max_param, +): + + assert distribution_param == "chirp_distance" + assert limits_param == "chirp_distance" found_dchirp = chirp_distance(found_d, found_mchirp) missed_dchirp = chirp_distance(missed_d, missed_mchirp) # treat chirp distances in MC volume estimate as physical distances - return volume_montecarlo(found_dchirp, missed_dchirp, found_mchirp, - missed_mchirp, 'distance', distribution, - 'distance', min_param, max_param) + return volume_montecarlo( + found_dchirp, + missed_dchirp, + found_mchirp, + missed_mchirp, + "distance", + distribution, + "distance", + min_param, + max_param, + ) def volume_binned_pylal(f_dist, m_dist, bins=15): - """ Compute the sensitive volume using a distance binned efficiency estimate + """ + Compute the sensitive volume using a distance binned efficiency estimate Parameters - ----------- + ---------- f_dist: numpy.ndarray The distances of found injections m_dist: numpy.ndarray The distances of missed injections Returns - -------- + ------- volume: float Volume estimate volume_error: float The standard error in the volume + """ + def sims_to_bin(sim): return (sim, 0) total = numpy.concatenate([f_dist, m_dist]) - ndbins = bin_utils.NDBins([bin_utils.LinearBins(min(total), max(total), bins), - bin_utils.LinearBins(0., 1, 1)]) + ndbins = bin_utils.NDBins( + [ + bin_utils.LinearBins(min(total), max(total), bins), + bin_utils.LinearBins(0.0, 1, 1), + ] + ) vol, verr = compute_search_volume_in_bins(f_dist, total, ndbins, sims_to_bin) return vol.array[0], verr.array[0] def volume_shell(f_dist, m_dist): - """ Compute the sensitive volume using sum over spherical shells. + """ + Compute the sensitive volume using sum over spherical shells. Parameters - ----------- + ---------- f_dist: numpy.ndarray The distances of found injections m_dist: numpy.ndarray The distances of missed injections Returns - -------- + ------- volume: float Volume estimate volume_error: float The standard error in the volume + """ f_dist.sort() m_dist.sort() @@ -314,13 +347,13 @@ def volume_shell(f_dist, m_dist): if i == len(distances) - 1: break - high = (distances[i+1] + distances[i]) / 2 + high = (distances[i + 1] + distances[i]) / 2 bin_width = high - low if dist_sorting[i] < len(f_dist): - vol += 4 * numpy.pi * distances[i]**2.0 * bin_width - vol_err += (4 * numpy.pi * distances[i]**2.0 * bin_width)**2.0 + vol += 4 * numpy.pi * distances[i] ** 2.0 * bin_width + vol_err += (4 * numpy.pi * distances[i] ** 2.0 * bin_width) ** 2.0 low = high - vol_err = vol_err ** 0.5 + vol_err = vol_err**0.5 return vol, vol_err diff --git a/pycbc/strain/__init__.py b/pycbc/strain/__init__.py index d3811e1df61..5621629c988 100644 --- a/pycbc/strain/__init__.py +++ b/pycbc/strain/__init__.py @@ -1,22 +1,31 @@ +from .gate import ( + add_gate_option_group, + apply_gates_to_fd, + apply_gates_to_td, + gates_from_cli, + psd_gates_from_cli, +) from .recalibrate import CubicSpline, PhysicalModel - -from .strain import detect_loud_glitches -from .strain import from_cli, from_cli_single_ifo, from_cli_multi_ifos -from .strain import insert_strain_option_group, insert_strain_option_group_multi_ifo -from .strain import verify_strain_options, verify_strain_options_multi_ifo -from .strain import gate_data, StrainSegments, StrainBuffer - -from .gate import add_gate_option_group, gates_from_cli -from .gate import apply_gates_to_td, apply_gates_to_fd, psd_gates_from_cli - -models = { - CubicSpline.name: CubicSpline, - PhysicalModel.name: PhysicalModel -} +from .strain import ( + StrainBuffer, + StrainSegments, + detect_loud_glitches, + from_cli, + from_cli_multi_ifos, + from_cli_single_ifo, + gate_data, + insert_strain_option_group, + insert_strain_option_group_multi_ifo, + verify_strain_options, + verify_strain_options_multi_ifo, +) + +models = {CubicSpline.name: CubicSpline, PhysicalModel.name: PhysicalModel} def read_model_from_config(cp, ifo, section="calibration"): - """Returns an instance of the calibration model specified in the + """ + Returns an instance of the calibration model specified in the given configuration file. Parameters @@ -32,8 +41,9 @@ def read_model_from_config(cp, ifo, section="calibration"): ------- instance An instance of the calibration model class. + """ - model = cp.get_opt_tag(section, "{}_model".format(ifo.lower()), None) + model = cp.get_opt_tag(section, f"{ifo.lower()}_model", None) recalibrator = models[model].from_config(cp, ifo.lower(), section) return recalibrator diff --git a/pycbc/strain/calibration.py b/pycbc/strain/calibration.py index d26fe696ba9..f73801b8295 100644 --- a/pycbc/strain/calibration.py +++ b/pycbc/strain/calibration.py @@ -13,12 +13,12 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Functions for adding calibration factors to waveform templates. -""" +"""Functions for adding calibration factors to waveform templates.""" + +from abc import ABCMeta, abstractmethod import numpy as np from scipy.interpolate import UnivariateSpline -from abc import (ABCMeta, abstractmethod) class Recalibrate(metaclass=ABCMeta): @@ -30,7 +30,8 @@ def __init__(self, ifo_name): @abstractmethod def apply_calibration(self, strain): - """Apply calibration model + """ + Apply calibration model This method should be overwritten by subclasses @@ -43,11 +44,13 @@ def apply_calibration(self, strain): ------ strain_adjusted : FrequencySeries The recalibrated strain. + """ return - def map_to_adjust(self, strain, prefix='recalib_', **params): - """Map an input dictionary of sampling parameters to the + def map_to_adjust(self, strain, prefix="recalib_", **params): + """ + Map an input dictionary of sampling parameters to the adjust_strain function by filtering the dictionary for the calibration parameters, then calling adjust_strain. @@ -64,11 +67,15 @@ def map_to_adjust(self, strain, prefix='recalib_', **params): ------ strain_adjusted : FrequencySeries The recalibrated strain. - """ - self.params.update({ - key[len(prefix):]: params[key] - for key in params if prefix in key and self.ifo_name in key}) + """ + self.params.update( + { + key[len(prefix) :]: params[key] + for key in params + if prefix in key and self.ifo_name in key + } + ) strain_adjusted = self.apply_calibration(strain) @@ -76,7 +83,8 @@ def map_to_adjust(self, strain, prefix='recalib_', **params): @classmethod def from_config(cls, cp, ifo, section): - """Read a config file to get calibration options and transfer + """ + Read a config file to get calibration options and transfer functions which will be used to intialize the model. Parameters @@ -93,21 +101,24 @@ def from_config(cls, cp, ifo, section): ------ instance An instance of the class. + """ all_params = dict(cp.items(section)) - params = {key[len(ifo)+1:]: all_params[key] - for key in all_params if ifo.lower() in key} - model = params.pop('model') - params['ifo_name'] = ifo.lower() + params = { + key[len(ifo) + 1 :]: all_params[key] + for key in all_params + if ifo.lower() in key + } + model = params.pop("model") + params["ifo_name"] = ifo.lower() return all_models[model](**params) class CubicSpline(Recalibrate): - name = 'cubic_spline' + name = "cubic_spline" - def __init__(self, minimum_frequency, maximum_frequency, n_points, - ifo_name): + def __init__(self, minimum_frequency, maximum_frequency, n_points, ifo_name): """ Cubic spline recalibration @@ -125,20 +136,22 @@ def __init__(self, minimum_frequency, maximum_frequency, n_points, maximum frequency of spline points n_points: int number of spline points + """ Recalibrate.__init__(self, ifo_name=ifo_name) minimum_frequency = float(minimum_frequency) maximum_frequency = float(maximum_frequency) n_points = int(n_points) if n_points < 4: - raise ValueError( - 'Use at least 4 spline points for calibration model') + raise ValueError("Use at least 4 spline points for calibration model") self.n_points = n_points - self.spline_points = np.logspace(np.log10(minimum_frequency), - np.log10(maximum_frequency), n_points) + self.spline_points = np.logspace( + np.log10(minimum_frequency), np.log10(maximum_frequency), n_points + ) def apply_calibration(self, strain): - """Apply calibration model + """ + Apply calibration model This applies cubic spline calibration to the strain. @@ -151,26 +164,30 @@ def apply_calibration(self, strain): ------ strain_adjusted : FrequencySeries The recalibrated strain. + """ - amplitude_parameters =\ - [self.params['amplitude_{}_{}'.format(self.ifo_name, ii)] - for ii in range(self.n_points)] - amplitude_spline = UnivariateSpline(self.spline_points, - amplitude_parameters) + amplitude_parameters = [ + self.params[f"amplitude_{self.ifo_name}_{ii}"] + for ii in range(self.n_points) + ] + amplitude_spline = UnivariateSpline(self.spline_points, amplitude_parameters) delta_amplitude = amplitude_spline(strain.sample_frequencies.numpy()) - phase_parameters =\ - [self.params['phase_{}_{}'.format(self.ifo_name, ii)] - for ii in range(self.n_points)] + phase_parameters = [ + self.params[f"phase_{self.ifo_name}_{ii}"] + for ii in range(self.n_points) + ] phase_spline = UnivariateSpline(self.spline_points, phase_parameters) delta_phase = phase_spline(strain.sample_frequencies.numpy()) - strain_adjusted = strain * (1.0 + delta_amplitude)\ - * (2.0 + 1j * delta_phase) / (2.0 - 1j * delta_phase) + strain_adjusted = ( + strain + * (1.0 + delta_amplitude) + * (2.0 + 1j * delta_phase) + / (2.0 - 1j * delta_phase) + ) return strain_adjusted -all_models = { - CubicSpline.name: CubicSpline -} +all_models = {CubicSpline.name: CubicSpline} diff --git a/pycbc/strain/gate.py b/pycbc/strain/gate.py index 4224d7be7fa..6e7df3e4fc7 100644 --- a/pycbc/strain/gate.py +++ b/pycbc/strain/gate.py @@ -13,15 +13,16 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Functions for applying gates to data. -""" +"""Functions for applying gates to data.""" from scipy import linalg + from . import strain def _gates_from_cli(opts, gate_opt): - """Parses the given `gate_opt` into something understandable by + """ + Parses the given `gate_opt` into something understandable by `strain.gate_data`. """ gates = {} @@ -29,13 +30,14 @@ def _gates_from_cli(opts, gate_opt): return gates for gate in getattr(opts, gate_opt): try: - ifo, central_time, half_dur, taper_dur = gate.split(':') + ifo, central_time, half_dur, taper_dur = gate.split(":") central_time = float(central_time) half_dur = float(half_dur) taper_dur = float(taper_dur) except ValueError: - raise ValueError("--gate {} not formatted correctly; ".format( - gate) + "see help") + raise ValueError( + f"--gate {gate} not formatted correctly; " + "see help" + ) try: gates[ifo].append((central_time, half_dur, taper_dur)) except KeyError: @@ -44,21 +46,24 @@ def _gates_from_cli(opts, gate_opt): def gates_from_cli(opts): - """Parses the --gate option into something understandable by + """ + Parses the --gate option into something understandable by `strain.gate_data`. """ - return _gates_from_cli(opts, 'gate') + return _gates_from_cli(opts, "gate") def psd_gates_from_cli(opts): - """Parses the --psd-gate option into something understandable by + """ + Parses the --psd-gate option into something understandable by `strain.gate_data`. """ - return _gates_from_cli(opts, 'psd_gate') + return _gates_from_cli(opts, "psd_gate") def apply_gates_to_td(strain_dict, gates): - """Applies the given dictionary of gates to the given dictionary of + """ + Applies the given dictionary of gates to the given dictionary of strain. Parameters @@ -74,6 +79,7 @@ def apply_gates_to_td(strain_dict, gates): ------- dict Dictionary of time-domain strain with the gates applied. + """ # copy data to new dictionary outdict = dict(strain_dict.items()) @@ -83,7 +89,8 @@ def apply_gates_to_td(strain_dict, gates): def apply_gates_to_fd(stilde_dict, gates): - """Applies the given dictionary of gates to the given dictionary of + """ + Applies the given dictionary of gates to the given dictionary of strain in the frequency domain. Gates are applied by IFFT-ing the strain data to the time domain, applying @@ -102,47 +109,61 @@ def apply_gates_to_fd(stilde_dict, gates): ------- dict Dictionary of frequency-domain strain with the gates applied. + """ # copy data to new dictionary outdict = dict(stilde_dict.items()) # create a time-domin strain dictionary to apply the gates to strain_dict = dict([[ifo, outdict[ifo].to_timeseries()] for ifo in gates]) # apply gates and fft back to the frequency domain - for ifo,d in apply_gates_to_td(strain_dict, gates).items(): + for ifo, d in apply_gates_to_td(strain_dict, gates).items(): outdict[ifo] = d.to_frequencyseries() return outdict def add_gate_option_group(parser): - """Adds the options needed to apply gates to data. + """ + Adds the options needed to apply gates to data. Parameters ---------- parser : object ArgumentParser instance. + """ gate_group = parser.add_argument_group("Options for gating data") - gate_group.add_argument("--gate", nargs="+", type=str, - metavar="IFO:CENTRALTIME:HALFDUR:TAPERDUR", - help="Apply one or more gates to the data before " - "filtering.") - gate_group.add_argument("--gate-overwhitened", action="store_true", - help="Overwhiten data first, then apply the " - "gates specified in --gate. Overwhitening " - "allows for sharper tapers to be used, " - "since lines are not blurred.") - gate_group.add_argument("--psd-gate", nargs="+", type=str, - metavar="IFO:CENTRALTIME:HALFDUR:TAPERDUR", - help="Apply one or more gates to the data used " - "for computing the PSD. Gates are applied " - "prior to FFT-ing the data for PSD " - "estimation.") + gate_group.add_argument( + "--gate", + nargs="+", + type=str, + metavar="IFO:CENTRALTIME:HALFDUR:TAPERDUR", + help="Apply one or more gates to the data before filtering.", + ) + gate_group.add_argument( + "--gate-overwhitened", + action="store_true", + help="Overwhiten data first, then apply the " + "gates specified in --gate. Overwhitening " + "allows for sharper tapers to be used, " + "since lines are not blurred.", + ) + gate_group.add_argument( + "--psd-gate", + nargs="+", + type=str, + metavar="IFO:CENTRALTIME:HALFDUR:TAPERDUR", + help="Apply one or more gates to the data used " + "for computing the PSD. Gates are applied " + "prior to FFT-ing the data for PSD " + "estimation.", + ) return gate_group def gate_and_paint(data, lindex, rindex, invpsd, copy=True): - """Gates and in-paints data using a Toeplitz solver. + """ + Gates and in-paints data using a Toeplitz solver. Parameters ---------- @@ -162,6 +183,7 @@ def gate_and_paint(data, lindex, rindex, invpsd, copy=True): ------- TimeSeries : The gated and in-painted time series. + """ # Uses the hole-filling method of # https://arxiv.org/pdf/1908.05644.pdf @@ -170,17 +192,21 @@ def gate_and_paint(data, lindex, rindex, invpsd, copy=True): data = data.copy() data[lindex:rindex] = 0 # get the over-whitened gated data - tdfilter = invpsd.astype('complex').to_timeseries() * invpsd.delta_t + tdfilter = invpsd.astype("complex").to_timeseries() * invpsd.delta_t owhgated_data = (data.to_frequencyseries() * invpsd).to_timeseries() # remove the projection into the null space - proj = linalg.solve_toeplitz(tdfilter[:(rindex - lindex)], - owhgated_data[lindex:rindex]) + proj = linalg.solve_toeplitz( + tdfilter[: (rindex - lindex)], owhgated_data[lindex:rindex] + ) data[lindex:rindex] -= proj return data + def invert_covariance(invpsd, lindex, rindex): - """Calculate the uninverted covariance matrix. + """ + Calculate the uninverted covariance matrix. + Parameters ---------- invpsd : FrequencySeries @@ -195,14 +221,17 @@ def invert_covariance(invpsd, lindex, rindex): array : The uninverted covariance matrix associated with the inverse PSD in the time window [lindex, rindex]. + """ - tdfilter = invpsd.astype('complex').to_timeseries() * invpsd.delta_t - mat = linalg.toeplitz(tdfilter[:(rindex-lindex)]) + tdfilter = invpsd.astype("complex").to_timeseries() * invpsd.delta_t + mat = linalg.toeplitz(tdfilter[: (rindex - lindex)]) invmat = linalg.inv(mat) return invmat + def gate_and_paint_matmul(data, lindex, rindex, invpsd, invmat=None, copy=True): - """Gates and in-paints data using explicit matrix multiplication. + """ + Gates and in-paints data using explicit matrix multiplication. Parameters ---------- @@ -219,11 +248,12 @@ def gate_and_paint_matmul(data, lindex, rindex, invpsd, invmat=None, copy=True): copy : bool, optional Copy the data before applying the gate. Otherwise, the gate will be applied in-place. Default is True. - + Returns ------- TimeSeries : The gated and in-painted time series. + """ if copy: data = data.copy() @@ -238,4 +268,4 @@ def gate_and_paint_matmul(data, lindex, rindex, invpsd, invmat=None, copy=True): # remove the projection into the null space proj = invmat @ owhgated_data[lindex:rindex] data[lindex:rindex] -= proj - return data \ No newline at end of file + return data diff --git a/pycbc/strain/lines.py b/pycbc/strain/lines.py index de27bcd1f73..888b5464da3 100644 --- a/pycbc/strain/lines.py +++ b/pycbc/strain/lines.py @@ -13,14 +13,16 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Functions for removing frequency lines from real data. -""" +"""Functions for removing frequency lines from real data.""" import numpy + from pycbc.types import TimeSeries, zeros + def complex_median(complex_list): - """ Get the median value of a list of complex numbers. + """ + Get the median value of a list of complex numbers. Parameters ---------- @@ -31,15 +33,16 @@ def complex_median(complex_list): ------- a + 1.j*b: complex number The median of the real and imaginary parts. + """ - median_real = numpy.median([complex_number.real - for complex_number in complex_list]) - median_imag = numpy.median([complex_number.imag - for complex_number in complex_list]) - return median_real + 1.j*median_imag + median_real = numpy.median([complex_number.real for complex_number in complex_list]) + median_imag = numpy.median([complex_number.imag for complex_number in complex_list]) + return median_real + 1.0j * median_imag + def avg_inner_product(data1, data2, bin_size): - """ Calculate the time-domain inner product averaged over bins. + """ + Calculate the time-domain inner product averaged over bins. Parameters ---------- @@ -59,16 +62,20 @@ def avg_inner_product(data1, data2, bin_size): The absolute value of the median of the inner product. phi: float The angle of the median of the inner product. + """ assert data1.duration == data2.duration assert data1.sample_rate == data2.sample_rate seglen = int(bin_size * data1.sample_rate) inner_prod = [] for idx in range(int(data1.duration / bin_size)): - start, end = idx * seglen, (idx+1) * seglen + start, end = idx * seglen, (idx + 1) * seglen norm = len(data1[start:end]) - bin_prod = 2 * sum(data1.data[start:end].real * - numpy.conjugate(data2.data[start:end])) / norm + bin_prod = ( + 2 + * sum(data1.data[start:end].real * numpy.conjugate(data2.data[start:end])) + / norm + ) inner_prod.append(bin_prod) # Get the median over all bins to avoid outliers due to the presence @@ -76,8 +83,10 @@ def avg_inner_product(data1, data2, bin_size): inner_median = complex_median(inner_prod) return inner_prod, numpy.abs(inner_median), numpy.angle(inner_median) + def line_model(freq, data, tref, amp=1, phi=0): - """ Simple time-domain model for a frequency line. + """ + Simple time-domain model for a frequency line. Parameters ---------- @@ -99,18 +108,22 @@ def line_model(freq, data, tref, amp=1, phi=0): data are complex to allow measuring the amplitude and phase of the corresponding frequency line in the strain data. For extraction, use only the real part of the data. + """ - freq_line = TimeSeries(zeros(len(data)), delta_t=data.delta_t, - epoch=data.start_time) + freq_line = TimeSeries( + zeros(len(data)), delta_t=data.delta_t, epoch=data.start_time + ) times = data.sample_times - float(tref) alpha = 2 * numpy.pi * freq * times + phi - freq_line.data = amp * numpy.exp(1.j * alpha) + freq_line.data = amp * numpy.exp(1.0j * alpha) return freq_line + def matching_line(freq, data, tref, bin_size=1): - """ Find the parameter of the line with frequency 'freq' in the data. + """ + Find the parameter of the line with frequency 'freq' in the data. Parameters ---------- @@ -128,15 +141,17 @@ def matching_line(freq, data, tref, bin_size=1): line_model: pycbc.types.TimeSeries A timeseries containing the frequency line with the amplitude and phase measured from the data. + """ template_line = line_model(freq, data, tref=tref) # Measure amplitude and phase of the line in the data - _, amp, phi = avg_inner_product(data, template_line, - bin_size=bin_size) + _, amp, phi = avg_inner_product(data, template_line, bin_size=bin_size) return line_model(freq, data, tref=tref, amp=amp, phi=phi) + def calibration_lines(freqs, data, tref=None): - """ Extract the calibration lines from strain data. + """ + Extract the calibration lines from strain data. Parameters ---------- @@ -151,18 +166,20 @@ def calibration_lines(freqs, data, tref=None): ------- data: pycbc.types.TimeSeries The strain data with the calibration lines removed. + """ if tref is None: tref = float(data.start_time) for freq in freqs: - measured_line = matching_line(freq, data, tref, - bin_size=data.duration) + measured_line = matching_line(freq, data, tref, bin_size=data.duration) data -= measured_line.data.real return data + def clean_data(freqs, data, chunk, avg_bin): - """ Extract time-varying (wandering) lines from strain data. + """ + Extract time-varying (wandering) lines from strain data. Parameters ---------- @@ -183,35 +200,40 @@ def clean_data(freqs, data, chunk, avg_bin): ------- data: pycbc.types.TimeSeries The strain data with the wandering lines removed. + """ if avg_bin >= chunk: - raise ValueError('The bin size for averaging the inner product ' - 'must be less than the chunk size.') + raise ValueError( + "The bin size for averaging the inner product " + "must be less than the chunk size." + ) if chunk >= data.duration: - raise ValueError('The chunk size must be less than the ' - 'data duration.') - steps = numpy.arange(0, int(data.duration/chunk)-0.5, 0.5) + raise ValueError("The chunk size must be less than the data duration.") + steps = numpy.arange(0, int(data.duration / chunk) - 0.5, 0.5) seglen = chunk * data.sample_rate tref = float(data.start_time) for freq in freqs: for step in steps: - start, end = int(step*seglen), int((step+1)*seglen) - chunk_line = matching_line(freq, data[start:end], - tref, bin_size=avg_bin) + start, end = int(step * seglen), int((step + 1) * seglen) + chunk_line = matching_line(freq, data[start:end], tref, bin_size=avg_bin) # Apply hann window on sides of chunk_line to smooth boundaries # and avoid discontinuities hann_window = numpy.hanning(len(chunk_line)) - apply_hann = TimeSeries(numpy.ones(len(chunk_line)), - delta_t=chunk_line.delta_t, - epoch=chunk_line.start_time) + apply_hann = TimeSeries( + numpy.ones(len(chunk_line)), + delta_t=chunk_line.delta_t, + epoch=chunk_line.start_time, + ) if step == 0: - apply_hann.data[len(hann_window)/2:] *= \ - hann_window[len(hann_window)/2:] + apply_hann.data[len(hann_window) / 2 :] *= hann_window[ + len(hann_window) / 2 : + ] elif step == steps[-1]: - apply_hann.data[:len(hann_window)/2] *= \ - hann_window[:len(hann_window)/2] + apply_hann.data[: len(hann_window) / 2] *= hann_window[ + : len(hann_window) / 2 + ] else: apply_hann.data *= hann_window chunk_line.data *= apply_hann.data diff --git a/pycbc/strain/recalibrate.py b/pycbc/strain/recalibrate.py index b1096f5530b..e86177ff5b5 100644 --- a/pycbc/strain/recalibrate.py +++ b/pycbc/strain/recalibrate.py @@ -1,5 +1,4 @@ -""" Classes and functions for adjusting strain data. -""" +"""Classes and functions for adjusting strain data.""" # Copyright (C) 2015 Ben Lackey, Christopher M. Biwer, # Daniel Finstad, Colm Talbot, Alex Nitz # @@ -17,18 +16,20 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -from abc import (ABCMeta, abstractmethod) +import glob +import os +from abc import ABCMeta, abstractmethod import numpy as np -import glob, os -from scipy.interpolate import UnivariateSpline -from pycbc.types import FrequencySeries +from scipy.interpolate import InterpolatedUnivariateSpline, UnivariateSpline + from pycbc.frame.gwosc import get_run -from scipy.interpolate import InterpolatedUnivariateSpline +from pycbc.types import FrequencySeries class Recalibrate(metaclass=ABCMeta): - """ Base class for modifying calibration """ + """Base class for modifying calibration""" + name = None def __init__(self, ifo_name): @@ -37,7 +38,8 @@ def __init__(self, ifo_name): @abstractmethod def apply_calibration(self, strain): - """Apply calibration model + """ + Apply calibration model This method should be overwritten by subclasses @@ -50,11 +52,13 @@ def apply_calibration(self, strain): ------ strain_adjusted : FrequencySeries The recalibrated strain. + """ return - def map_to_adjust(self, strain, prefix='recalib_', **params): - """Map an input dictionary of sampling parameters to the + def map_to_adjust(self, strain, prefix="recalib_", **params): + """ + Map an input dictionary of sampling parameters to the adjust_strain function by filtering the dictionary for the calibration parameters, then calling adjust_strain. @@ -71,11 +75,15 @@ def map_to_adjust(self, strain, prefix='recalib_', **params): ------ strain_adjusted : FrequencySeries The recalibrated strain. - """ - self.params.update({ - key[len(prefix):]: params[key] - for key in params if prefix in key and self.ifo_name in key}) + """ + self.params.update( + { + key[len(prefix) :]: params[key] + for key in params + if prefix in key and self.ifo_name in key + } + ) strain_adjusted = self.apply_calibration(strain) @@ -83,7 +91,8 @@ def map_to_adjust(self, strain, prefix='recalib_', **params): @classmethod def from_config(cls, cp, ifo, section): - """Read a config file to get calibration options and transfer + """ + Read a config file to get calibration options and transfer functions which will be used to intialize the model. Parameters @@ -100,19 +109,24 @@ def from_config(cls, cp, ifo, section): ------ instance An instance of the class. + """ all_params = dict(cp.items(section)) - params = {key[len(ifo)+1:]: all_params[key] - for key in all_params if ifo.lower() in key} + params = { + key[len(ifo) + 1 :]: all_params[key] + for key in all_params + if ifo.lower() in key + } params = {key: params[key] for key in params} - params.pop('model') - params['ifo_name'] = ifo.lower() + params.pop("model") + params["ifo_name"] = ifo.lower() return cls(**params) class CubicSpline(Recalibrate): - """Cubic spline recalibration + """ + Cubic spline recalibration see https://dcc.ligo.org/LIGO-T1400682/public @@ -128,24 +142,26 @@ class CubicSpline(Recalibrate): maximum frequency of spline points n_points: int number of spline points + """ - name = 'cubic_spline' - def __init__(self, minimum_frequency, maximum_frequency, n_points, - ifo_name): + name = "cubic_spline" + + def __init__(self, minimum_frequency, maximum_frequency, n_points, ifo_name): Recalibrate.__init__(self, ifo_name=ifo_name) minimum_frequency = float(minimum_frequency) maximum_frequency = float(maximum_frequency) n_points = int(n_points) if n_points < 4: - raise ValueError( - 'Use at least 4 spline points for calibration model') + raise ValueError("Use at least 4 spline points for calibration model") self.n_points = n_points - self.spline_points = np.logspace(np.log10(minimum_frequency), - np.log10(maximum_frequency), n_points) + self.spline_points = np.logspace( + np.log10(minimum_frequency), np.log10(maximum_frequency), n_points + ) def apply_calibration(self, strain): - """Apply calibration model + """ + Apply calibration model This applies cubic spline calibration to the strain. @@ -158,28 +174,35 @@ def apply_calibration(self, strain): ------ strain_adjusted : FrequencySeries The recalibrated strain. + """ - amplitude_parameters =\ - [self.params['amplitude_{}_{}'.format(self.ifo_name, ii)] - for ii in range(self.n_points)] - amplitude_spline = UnivariateSpline(self.spline_points, - amplitude_parameters) + amplitude_parameters = [ + self.params[f"amplitude_{self.ifo_name}_{ii}"] + for ii in range(self.n_points) + ] + amplitude_spline = UnivariateSpline(self.spline_points, amplitude_parameters) delta_amplitude = amplitude_spline(strain.sample_frequencies.numpy()) - phase_parameters =\ - [self.params['phase_{}_{}'.format(self.ifo_name, ii)] - for ii in range(self.n_points)] + phase_parameters = [ + self.params[f"phase_{self.ifo_name}_{ii}"] + for ii in range(self.n_points) + ] phase_spline = UnivariateSpline(self.spline_points, phase_parameters) delta_phase = phase_spline(strain.sample_frequencies.numpy()) - strain_adjusted = strain * (1.0 + delta_amplitude)\ - * (2.0 + 1j * delta_phase) / (2.0 - 1j * delta_phase) + strain_adjusted = ( + strain + * (1.0 + delta_amplitude) + * (2.0 + 1j * delta_phase) + / (2.0 - 1j * delta_phase) + ) return strain_adjusted -class PhysicalModel(object): - """ Class for adjusting time-varying calibration parameters of given +class PhysicalModel: + """ + Class for adjusting time-varying calibration parameters of given strain data. Parameters @@ -206,11 +229,22 @@ class PhysicalModel(object): qinv0 : float Initial inverse quality factor at t0 for the signal recycling cavity. + """ - name = 'physical_model' - def __init__(self, freq=None, fc0=None, c0=None, d0=None, - a_tst0=None, a_pu0=None, fs0=None, qinv0=None): + name = "physical_model" + + def __init__( + self, + freq=None, + fc0=None, + c0=None, + d0=None, + a_tst0=None, + a_pu0=None, + fs0=None, + qinv0=None, + ): self.freq = np.real(freq) self.c0 = c0 self.d0 = d0 @@ -221,8 +255,9 @@ def __init__(self, freq=None, fc0=None, c0=None, d0=None, self.qinv0 = float(qinv0) # initial detuning at time t0 - init_detuning = self.freq**2 / (self.freq**2 - 1.0j * self.freq * \ - self.fs0 * self.qinv0 + self.fs0**2) + init_detuning = self.freq**2 / ( + self.freq**2 - 1.0j * self.freq * self.fs0 * self.qinv0 + self.fs0**2 + ) # initial open loop gain self.g0 = self.c0 * self.d0 * (self.a_tst0 + self.a_pu0) @@ -231,10 +266,11 @@ def __init__(self, freq=None, fc0=None, c0=None, d0=None, self.r0 = (1.0 + self.g0) / self.c0 # residual of c0 after factoring out the coupled cavity pole fc0 - self.c_res = self.c0 * (1 + 1.0j * self.freq / self.fc0)/init_detuning + self.c_res = self.c0 * (1 + 1.0j * self.freq / self.fc0) / init_detuning def update_c(self, fs=None, qinv=None, fc=None, kappa_c=1.0): - """ Calculate the sensing function c(f,t) given the new parameters + """ + Calculate the sensing function c(f,t) given the new parameters kappa_c(t), kappa_a(t), f_c(t), fs, and qinv. Parameters @@ -252,15 +288,26 @@ def update_c(self, fs=None, qinv=None, fc=None, kappa_c=1.0): ------- c : numpy.array The new sensing function c(f,t). + """ - detuning_term = self.freq**2 / (self.freq**2 - 1.0j *self.freq*fs * \ - qinv + fs**2) - return self.c_res * kappa_c / (1 + 1.0j * self.freq/fc)*detuning_term - - def update_g(self, fs=None, qinv=None, fc=None, kappa_tst_re=1.0, - kappa_tst_im=0.0, kappa_pu_re=1.0, kappa_pu_im=0.0, - kappa_c=1.0): - """ Calculate the open loop gain g(f,t) given the new parameters + detuning_term = self.freq**2 / ( + self.freq**2 - 1.0j * self.freq * fs * qinv + fs**2 + ) + return self.c_res * kappa_c / (1 + 1.0j * self.freq / fc) * detuning_term + + def update_g( + self, + fs=None, + qinv=None, + fc=None, + kappa_tst_re=1.0, + kappa_tst_im=0.0, + kappa_pu_re=1.0, + kappa_pu_im=0.0, + kappa_c=1.0, + ): + """ + Calculate the open loop gain g(f,t) given the new parameters kappa_c(t), kappa_a(t), f_c(t), fs, and qinv. Parameters @@ -290,16 +337,26 @@ def update_g(self, fs=None, qinv=None, fc=None, kappa_tst_re=1.0, ------- g : numpy.array The new open loop gain g(f,t). + """ c = self.update_c(fs=fs, qinv=qinv, fc=fc, kappa_c=kappa_c) a_tst = self.a_tst0 * (kappa_tst_re + 1.0j * kappa_tst_im) a_pu = self.a_pu0 * (kappa_pu_re + 1.0j * kappa_pu_im) return c * self.d0 * (a_tst + a_pu) - def update_r(self, fs=None, qinv=None, fc=None, kappa_c=1.0, - kappa_tst_re=1.0, kappa_tst_im=0.0, kappa_pu_re=1.0, - kappa_pu_im=0.0): - """ Calculate the response function R(f,t) given the new parameters + def update_r( + self, + fs=None, + qinv=None, + fc=None, + kappa_c=1.0, + kappa_tst_re=1.0, + kappa_tst_im=0.0, + kappa_pu_re=1.0, + kappa_pu_im=0.0, + ): + """ + Calculate the response function R(f,t) given the new parameters kappa_c(t), kappa_a(t), f_c(t), fs, and qinv. Parameters @@ -329,18 +386,35 @@ def update_r(self, fs=None, qinv=None, fc=None, kappa_c=1.0, ------- r : numpy.array The new response function r(f,t). + """ c = self.update_c(fs=fs, qinv=qinv, fc=fc, kappa_c=kappa_c) - g = self.update_g(fs=fs, qinv=qinv, fc=fc, kappa_c=kappa_c, - kappa_tst_re=kappa_tst_re, - kappa_tst_im=kappa_tst_im, - kappa_pu_re=kappa_pu_re, kappa_pu_im=kappa_pu_im) + g = self.update_g( + fs=fs, + qinv=qinv, + fc=fc, + kappa_c=kappa_c, + kappa_tst_re=kappa_tst_re, + kappa_tst_im=kappa_tst_im, + kappa_pu_re=kappa_pu_re, + kappa_pu_im=kappa_pu_im, + ) return (1.0 + g) / c - def adjust_strain(self, strain, delta_fs=None, delta_qinv=None, - delta_fc=None, kappa_c=1.0, kappa_tst_re=1.0, - kappa_tst_im=0.0, kappa_pu_re=1.0, kappa_pu_im=0.0): - """Adjust the FrequencySeries strain by changing the time-dependent + def adjust_strain( + self, + strain, + delta_fs=None, + delta_qinv=None, + delta_fc=None, + kappa_c=1.0, + kappa_tst_re=1.0, + kappa_tst_im=0.0, + kappa_pu_re=1.0, + kappa_pu_im=0.0, + ): + """ + Adjust the FrequencySeries strain by changing the time-dependent calibration parameters kappa_c(t), kappa_a(t), f_c(t), fs, and qinv. Parameters @@ -372,17 +446,23 @@ def adjust_strain(self, strain, delta_fs=None, delta_qinv=None, ------- strain_adjusted : FrequencySeries The adjusted strain. + """ fc = self.fc0 + delta_fc if delta_fc else self.fc0 fs = self.fs0 + delta_fs if delta_fs else self.fs0 qinv = self.qinv0 + delta_qinv if delta_qinv else self.qinv0 # calculate adjusted response function - r_adjusted = self.update_r(fs=fs, qinv=qinv, fc=fc, kappa_c=kappa_c, - kappa_tst_re=kappa_tst_re, - kappa_tst_im=kappa_tst_im, - kappa_pu_re=kappa_pu_re, - kappa_pu_im=kappa_pu_im) + r_adjusted = self.update_r( + fs=fs, + qinv=qinv, + fc=fc, + kappa_c=kappa_c, + kappa_tst_re=kappa_tst_re, + kappa_tst_im=kappa_tst_im, + kappa_pu_re=kappa_pu_re, + kappa_pu_im=kappa_pu_im, + ) # calculate error function k = r_adjusted / self.r0 @@ -396,17 +476,17 @@ def adjust_strain(self, strain, delta_fs=None, delta_qinv=None, k_amp_off = UnivariateSpline(self.freq, k_amp, k=order, s=0) k_phase_off = UnivariateSpline(self.freq, k_phase, k=order, s=0) freq_even = strain.sample_frequencies.numpy() - k_even_sample = k_amp_off(freq_even) * \ - np.exp(1.0j * k_phase_off(freq_even)) - strain_adjusted = FrequencySeries(strain.numpy() * \ - k_even_sample, - delta_f=strain.delta_f) + k_even_sample = k_amp_off(freq_even) * np.exp(1.0j * k_phase_off(freq_even)) + strain_adjusted = FrequencySeries( + strain.numpy() * k_even_sample, delta_f=strain.delta_f + ) return strain_adjusted @classmethod def tf_from_file(cls, path, delimiter=" "): - """Convert the contents of a file with the columns + """ + Convert the contents of a file with the columns [freq, real(h), imag(h)] to a numpy.array with columns [freq, real(h)+j*imag(h)]. @@ -418,6 +498,7 @@ def tf_from_file(cls, path, delimiter=" "): Return ------ numpy.array + """ data = np.loadtxt(path, delimiter=delimiter) freq = data[:, 0] @@ -426,7 +507,8 @@ def tf_from_file(cls, path, delimiter=" "): @classmethod def from_config(cls, cp, ifo, section): - """Read a config file to get calibration options and transfer + """ + Read a config file to get calibration options and transfer functions which will be used to intialize the model. Parameters @@ -444,12 +526,12 @@ def from_config(cls, cp, ifo, section): ------ instance An instance of the Recalibrate class. + """ # read transfer functions tfs = [] tf_names = ["a-tst", "a-pu", "c", "d"] - for tag in ['-'.join([ifo, "transfer-function", name]) - for name in tf_names]: + for tag in ["-".join([ifo, "transfer-function", name]) for name in tf_names]: tf_path = cp.get_opt_tag(section, tag, None) tfs.append(cls.tf_from_file(tf_path)) a_tst0 = tfs[0][:, 1] @@ -460,21 +542,30 @@ def from_config(cls, cp, ifo, section): # if upper stage actuation is included, read that in and add it # to a_pu0 - uim_tag = '-'.join([ifo, 'transfer-function-a-uim']) + uim_tag = "-".join([ifo, "transfer-function-a-uim"]) if cp.has_option(section, uim_tag): tf_path = cp.get_opt_tag(section, uim_tag, None) a_pu0 += cls.tf_from_file(tf_path)[:, 1] # read fc0, fs0, and qinv0 - fc0 = cp.get_opt_tag(section, '-'.join([ifo, "fc0"]), None) - fs0 = cp.get_opt_tag(section, '-'.join([ifo, "fs0"]), None) - qinv0 = cp.get_opt_tag(section, '-'.join([ifo, "qinv0"]), None) - - return cls(freq=freq, fc0=fc0, c0=c0, d0=d0, a_tst0=a_tst0, - a_pu0=a_pu0, fs0=fs0, qinv0=qinv0) + fc0 = cp.get_opt_tag(section, "-".join([ifo, "fc0"]), None) + fs0 = cp.get_opt_tag(section, "-".join([ifo, "fs0"]), None) + qinv0 = cp.get_opt_tag(section, "-".join([ifo, "qinv0"]), None) + + return cls( + freq=freq, + fc0=fc0, + c0=c0, + d0=d0, + a_tst0=a_tst0, + a_pu0=a_pu0, + fs0=fs0, + qinv0=qinv0, + ) def map_to_adjust(self, strain, **params): - """Map an input dictionary of sampling parameters to the + """ + Map an input dictionary of sampling parameters to the adjust_strain function by filtering the dictionary for the calibration parameters, then calling adjust_strain. @@ -490,13 +581,21 @@ def map_to_adjust(self, strain, **params): ------ strain_adjusted : FrequencySeries The recalibrated strain. + """ # calibration param names - arg_names = ['delta_fs', 'delta_fc', 'delta_qinv', 'kappa_c', - 'kappa_tst_re', 'kappa_tst_im', 'kappa_pu_re', - 'kappa_pu_im'] + arg_names = [ + "delta_fs", + "delta_fc", + "delta_qinv", + "kappa_c", + "kappa_tst_re", + "kappa_tst_im", + "kappa_pu_re", + "kappa_pu_im", + ] # calibration param labels as they exist in config files - arg_labels = [''.join(['calib_', name]) for name in arg_names] + arg_labels = ["".join(["calib_", name]) for name in arg_names] # default values for calibration params default_values = [0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0] # make list of calibration param values @@ -507,18 +606,23 @@ def map_to_adjust(self, strain, **params): else: calib_args.append(val) # adjust the strain using calibration param values - strain_adjusted = self.adjust_strain(strain, delta_fs=calib_args[0], - delta_fc=calib_args[1], delta_qinv=calib_args[2], - kappa_c=calib_args[3], - kappa_tst_re=calib_args[4], - kappa_tst_im=calib_args[5], - kappa_pu_re=calib_args[6], - kappa_pu_im=calib_args[7]) + strain_adjusted = self.adjust_strain( + strain, + delta_fs=calib_args[0], + delta_fc=calib_args[1], + delta_qinv=calib_args[2], + kappa_c=calib_args[3], + kappa_tst_re=calib_args[4], + kappa_tst_im=calib_args[5], + kappa_pu_re=calib_args[6], + kappa_pu_im=calib_args[7], + ) return strain_adjusted -def read_calibration_envelop_file(calibration_file, correction_type, - minimum_frequency, maximum_frequency, n_nodes): +def read_calibration_envelop_file( + calibration_file, correction_type, minimum_frequency, maximum_frequency, n_nodes +): """ This function reads the calibration envelop file and provide arrays needed to construct cubic splines @@ -540,19 +644,20 @@ def read_calibration_envelop_file(calibration_file, correction_type, calibration_data = np.loadtxt(calibration_file).T log_frequency_array = np.log(calibration_data[0]) - log_nodes = np.linspace(np.log(minimum_frequency), - np.log(maximum_frequency), n_nodes) + log_nodes = np.linspace( + np.log(minimum_frequency), np.log(maximum_frequency), n_nodes + ) - if correction_type.lower()=='template': + if correction_type.lower() == "template": amplitude_median = calibration_data[1] - 1 phase_median = calibration_data[2] amplitude_sigma = abs(calibration_data[5] - calibration_data[3]) / 2 phase_sigma = abs(calibration_data[6] - calibration_data[4]) / 2 - elif correction_type.lower()=='data': - amplitude_median = 1/calibration_data[1] -1 + elif correction_type.lower() == "data": + amplitude_median = 1 / calibration_data[1] - 1 phase_median = -calibration_data[2] - amplitude_sigma = abs(1/calibration_data[3] - 1/calibration_data[5]) / 2 + amplitude_sigma = abs(1 / calibration_data[3] - 1 / calibration_data[5]) / 2 phase_sigma = abs(calibration_data[6] - calibration_data[4]) / 2 # Some sanity checks @@ -560,27 +665,41 @@ def read_calibration_envelop_file(calibration_file, correction_type, if maximum_frequency - calibration_data[0][-1] < 0.5: maximum_frequency = calibration_data[0][-1] else: - raise ValueError("Maximum frequency (=%d) for tc=%s is more than" - "maximum frequency in calibration envelop (=%.3g)"%(maximum_frequency,gps_time,calibration_data[0][-1])) + raise ValueError( + "Maximum frequency (=%d) for tc=%s is more than" + "maximum frequency in calibration envelop (=%.3g)" + % (maximum_frequency, gps_time, calibration_data[0][-1]) + ) else: pass if calibration_data[0][0] > minimum_frequency: - raise ValueError("Minimum frequency (=%d) for tc=%s is less than" - "minimum frequency in calibration envelop (=%.3g)"%(minimum_frequency,gps_time,calibration_data[0][0])) - else: - pass - - amplitude_median_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, amplitude_median)(log_nodes) - amplitude_sigma_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, amplitude_sigma)(log_nodes) - phase_median_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, phase_median)(log_nodes) - phase_sigma_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, phase_sigma)(log_nodes) - - return log_nodes, amplitude_median_nodes, amplitude_sigma_nodes, phase_median_nodes, phase_sigma_nodes + raise ValueError( + "Minimum frequency (=%d) for tc=%s is less than" + "minimum frequency in calibration envelop (=%.3g)" + % (minimum_frequency, gps_time, calibration_data[0][0]) + ) + + amplitude_median_nodes = InterpolatedUnivariateSpline( + log_frequency_array, amplitude_median + )(log_nodes) + amplitude_sigma_nodes = InterpolatedUnivariateSpline( + log_frequency_array, amplitude_sigma + )(log_nodes) + phase_median_nodes = InterpolatedUnivariateSpline( + log_frequency_array, phase_median + )(log_nodes) + phase_sigma_nodes = InterpolatedUnivariateSpline(log_frequency_array, phase_sigma)( + log_nodes + ) + + return ( + log_nodes, + amplitude_median_nodes, + amplitude_sigma_nodes, + phase_median_nodes, + phase_sigma_nodes, + ) def get_calibration_files(ifos, gps_time, calibration_file_path): @@ -597,36 +716,63 @@ def get_calibration_files(ifos, gps_time, calibration_file_path): Output: A dictionary of calibration envelop file path for each IFO. """ - - RUN_NAME=get_run(gps_time).split('_')[0] + RUN_NAME = get_run(gps_time).split("_")[0] get_run(gps_time) dict_calibration_file = {} for ifo in ifos: - if RUN_NAME=='O4a': - all_calibration_files = glob.glob('%s/%s_%s/*.txt'%(calibration_file_path, ifo, RUN_NAME)) + if RUN_NAME == "O4a": + all_calibration_files = glob.glob( + "%s/%s_%s/*.txt" % (calibration_file_path, ifo, RUN_NAME) + ) else: - all_calibration_files = glob.glob('%s/%s/*FinalResults.txt'%(calibration_file_path, ifo)) - if ifo !='V1': + all_calibration_files = glob.glob( + "%s/%s/*FinalResults.txt" % (calibration_file_path, ifo) + ) + if ifo != "V1": # H1 and L1 detector files are treated seperately compared to V1 detector # This list of GPS times can be coded in a better way!!! - if RUN_NAME=='O4a': - list_gpstimes = np.array([int(item.split('_')[-1].split('.')[0]) for item in all_calibration_files]) + if RUN_NAME == "O4a": + list_gpstimes = np.array( + [ + int(item.split("_")[-1].split(".")[0]) + for item in all_calibration_files + ] + ) else: - list_gpstimes = np.array([int(item.split('_')[-4]) for item in all_calibration_files]) + list_gpstimes = np.array( + [int(item.split("_")[-4]) for item in all_calibration_files] + ) dt_list = abs(list_gpstimes - gps_time) ifo_calibration_file = all_calibration_files[dt_list.argmin()] + elif RUN_NAME == "O2": + ifo_calibration_file = ( + "%s/calibration_envelops/V1/V_calibrationUncertaintyEnvelope_magnitude5p1percent_phase40mraddeg20microsecond.txt" + % os.getcwd() + ) + elif RUN_NAME == "O3a": + ifo_calibration_file = ( + "%s/calibration_envelops/V1/V_O3a_calibrationUncertaintyEnvelope_magnitude5percent_phase35milliradians10microseconds.txt" + % os.getcwd() + ) + elif RUN_NAME == "O3b": + ifo_calibration_file = ( + "%s/calibration_envelops/V1/V_O3b_calibrationUncertaintyEnvelope_magnitude5percent_phase35milliradians10microseconds.txt" + % os.getcwd() + ) else: - if RUN_NAME=='O2': - ifo_calibration_file = '%s/calibration_envelops/V1/V_calibrationUncertaintyEnvelope_magnitude5p1percent_phase40mraddeg20microsecond.txt'%os.getcwd() - elif RUN_NAME=='O3a': - ifo_calibration_file = '%s/calibration_envelops/V1/V_O3a_calibrationUncertaintyEnvelope_magnitude5percent_phase35milliradians10microseconds.txt'%os.getcwd() - elif RUN_NAME=='O3b': - ifo_calibration_file = '%s/calibration_envelops/V1/V_O3b_calibrationUncertaintyEnvelope_magnitude5percent_phase35milliradians10microseconds.txt'%os.getcwd() - else: - raise ValueError("Virgo GPS time is not in valid range") - - if RUN_NAME=='O4a': - dict_calibration_file[ifo] = '%s/%s_%s/%s'%(calibration_file_path, ifo,RUN_NAME,ifo_calibration_file.split('/')[-1]) + raise ValueError("Virgo GPS time is not in valid range") + + if RUN_NAME == "O4a": + dict_calibration_file[ifo] = "%s/%s_%s/%s" % ( + calibration_file_path, + ifo, + RUN_NAME, + ifo_calibration_file.split("/")[-1], + ) else: - dict_calibration_file[ifo] = '%s/%s/%s'%(calibration_file_path, ifo,ifo_calibration_file.split('/')[-1]) + dict_calibration_file[ifo] = "%s/%s/%s" % ( + calibration_file_path, + ifo, + ifo_calibration_file.split("/")[-1], + ) return dict_calibration_file diff --git a/pycbc/strain/strain.py b/pycbc/strain/strain.py index 77cbf5e5a89..b2fe18bd5b8 100644 --- a/pycbc/strain/strain.py +++ b/pycbc/strain/strain.py @@ -1,4 +1,4 @@ -#Copyright (C) 2013 Alex Nitz +# Copyright (C) 2013 Alex Nitz # # This program is free software; you can redistribute it and/or modify it # under the terms of the GNU General Public License as published by the @@ -16,36 +16,48 @@ """ This modules contains functions reading, generating, and segmenting strain data """ + import copy -import logging import functools -import numpy +import logging +import numpy from scipy.signal import kaiserord +import pycbc.events +import pycbc.filter +import pycbc.frame +import pycbc.psd import pycbc.types -from pycbc.types import TimeSeries, zeros -from pycbc.types import Array, FrequencySeries -from pycbc.types import MultiDetOptionAppendAction, MultiDetOptionAction -from pycbc.types import MultiDetOptionActionSpecial -from pycbc.types import DictOptionAction, MultiDetDictOptionAction -from pycbc.types import required_opts, required_opts_multi_ifo -from pycbc.types import ensure_one_opt, ensure_one_opt_multi_ifo -from pycbc.types import copy_opts_for_single_ifo, complex_same_precision_as -from pycbc.inject import InjectionSet, SGBurstInjectionSet -from pycbc.filter import resample_to_delta_t, lowpass, highpass, make_frequency_series +from pycbc.fft import FFT, IFFT +from pycbc.filter import highpass, lowpass, make_frequency_series, resample_to_delta_t from pycbc.filter.zpk import filter_zpk +from pycbc.inject import InjectionSet, SGBurstInjectionSet +from pycbc.types import ( + Array, + DictOptionAction, + FrequencySeries, + MultiDetDictOptionAction, + MultiDetOptionAction, + MultiDetOptionActionSpecial, + MultiDetOptionAppendAction, + TimeSeries, + complex_same_precision_as, + copy_opts_for_single_ifo, + ensure_one_opt, + ensure_one_opt_multi_ifo, + required_opts, + required_opts_multi_ifo, + zeros, +) from pycbc.waveform.spa_tmplt import spa_distance -import pycbc.psd -from pycbc.fft import FFT, IFFT -import pycbc.events -import pycbc.frame -import pycbc.filter -logger = logging.getLogger('pycbc.strain.strain') +logger = logging.getLogger("pycbc.strain.strain") + def next_power_of_2(n): - """Return the smallest integer power of 2 larger than the argument. + """ + Return the smallest integer power of 2 larger than the argument. Parameters ---------- @@ -56,14 +68,25 @@ def next_power_of_2(n): ------- m : int Smallest integer power of 2 larger than n. + """ return 1 << n.bit_length() -def detect_loud_glitches(strain, psd_duration=4., psd_stride=2., - psd_avg_method='median', low_freq_cutoff=30., - threshold=50., cluster_window=5., corrupt_time=4., - high_freq_cutoff=None, output_intermediates=False): - """Automatic identification of loud transients for gating purposes. + +def detect_loud_glitches( + strain, + psd_duration=4.0, + psd_stride=2.0, + psd_avg_method="median", + low_freq_cutoff=30.0, + threshold=50.0, + cluster_window=5.0, + corrupt_time=4.0, + high_freq_cutoff=None, + output_intermediates=False, +): + """ + Automatic identification of loud transients for gating purposes. This function first estimates the PSD of the input time series using the FindChirp Welch method. Then it whitens the time series using that @@ -97,11 +120,10 @@ def detect_loud_glitches(strain, psd_duration=4., psd_stride=2., frequency is used. output_intermediates : {bool, False} Save intermediate time series for debugging. - """ + """ if high_freq_cutoff: - strain = resample_to_delta_t(strain, 0.5 / high_freq_cutoff, - method='ldas') + strain = resample_to_delta_t(strain, 0.5 / high_freq_cutoff, method="ldas") else: strain = strain.copy() @@ -109,11 +131,12 @@ def detect_loud_glitches(strain, psd_duration=4., psd_stride=2., corrupt_length = int(corrupt_time * strain.sample_rate) w = numpy.arange(corrupt_length) / float(corrupt_length) strain[0:corrupt_length] *= pycbc.types.Array(w, dtype=strain.dtype) - strain[(len(strain) - corrupt_length):] *= \ - pycbc.types.Array(w[::-1], dtype=strain.dtype) + strain[(len(strain) - corrupt_length) :] *= pycbc.types.Array( + w[::-1], dtype=strain.dtype + ) if output_intermediates: - strain.save_to_wav('strain_conditioned.wav') + strain.save_to_wav("strain_conditioned.wav") # zero-pad strain to a power-of-2 length strain_pad_length = next_power_of_2(len(strain)) @@ -121,21 +144,28 @@ def detect_loud_glitches(strain, psd_duration=4., psd_stride=2., pad_end = pad_start + len(strain) pad_epoch = strain.start_time - pad_start / float(strain.sample_rate) strain_pad = pycbc.types.TimeSeries( - pycbc.types.zeros(strain_pad_length, dtype=strain.dtype), - delta_t=strain.delta_t, copy=False, epoch=pad_epoch) + pycbc.types.zeros(strain_pad_length, dtype=strain.dtype), + delta_t=strain.delta_t, + copy=False, + epoch=pad_epoch, + ) strain_pad[pad_start:pad_end] = strain[:] # estimate the PSD - psd = pycbc.psd.welch(strain[corrupt_length:(len(strain)-corrupt_length)], - seg_len=int(psd_duration * strain.sample_rate), - seg_stride=int(psd_stride * strain.sample_rate), - avg_method=psd_avg_method, - require_exact_data_fit=False) - psd = pycbc.psd.interpolate(psd, 1. / strain_pad.duration) + psd = pycbc.psd.welch( + strain[corrupt_length : (len(strain) - corrupt_length)], + seg_len=int(psd_duration * strain.sample_rate), + seg_stride=int(psd_stride * strain.sample_rate), + avg_method=psd_avg_method, + require_exact_data_fit=False, + ) + psd = pycbc.psd.interpolate(psd, 1.0 / strain_pad.duration) psd = pycbc.psd.inverse_spectrum_truncation( - psd, int(psd_duration * strain.sample_rate), - low_frequency_cutoff=low_freq_cutoff, - trunc_method='hann') + psd, + int(psd_duration * strain.sample_rate), + low_frequency_cutoff=low_freq_cutoff, + trunc_method="hann", + ) kmin = int(low_freq_cutoff / psd.delta_f) psd[0:kmin] = numpy.inf if high_freq_cutoff: @@ -148,38 +178,38 @@ def detect_loud_glitches(strain, psd_duration=4., psd_stride=2., if high_freq_cutoff: norm = high_freq_cutoff - low_freq_cutoff else: - norm = strain.sample_rate / 2. - low_freq_cutoff + norm = strain.sample_rate / 2.0 - low_freq_cutoff strain_tilde *= (psd * norm) ** (-0.5) strain_pad = strain_tilde.to_timeseries() if output_intermediates: - strain_pad[pad_start:pad_end].save_to_wav('strain_whitened.wav') + strain_pad[pad_start:pad_end].save_to_wav("strain_whitened.wav") mag = abs(strain_pad[pad_start:pad_end]) if output_intermediates: - mag.save('strain_whitened_mag.npy') + mag.save("strain_whitened_mag.npy") mag = mag.numpy() # remove strain corrupted by filters at the ends mag[0:corrupt_length] = 0 - mag[-1:-corrupt_length-1:-1] = 0 + mag[-1 : -corrupt_length - 1 : -1] = 0 # find peaks and their times indices = numpy.where(mag > threshold)[0] cluster_idx = pycbc.events.findchirp_cluster_over_window( - indices, numpy.array(mag[indices]), - int(cluster_window*strain.sample_rate)) - times = [idx * strain.delta_t + strain.start_time \ - for idx in indices[cluster_idx]] + indices, numpy.array(mag[indices]), int(cluster_window * strain.sample_rate) + ) + times = [idx * strain.delta_t + strain.start_time for idx in indices[cluster_idx]] return times -def from_cli(opt, dyn_range_fac=1, precision='single', - inj_filter_rejector=None): - """Parses the CLI options related to strain data reading and conditioning. + +def from_cli(opt, dyn_range_fac=1, precision="single", inj_filter_rejector=None): + """ + Parses the CLI options related to strain data reading and conditioning. Parameters ---------- @@ -201,6 +231,7 @@ def from_cli(opt, dyn_range_fac=1, precision='single', ------- strain : TimeSeries The time series containing the conditioned strain data. + """ gating_info = {} @@ -214,27 +245,34 @@ def from_cli(opt, dyn_range_fac=1, precision='single', logger.info("Reading Frames") - if hasattr(opt, 'frame_sieve') and opt.frame_sieve: + if hasattr(opt, "frame_sieve") and opt.frame_sieve: sieve = opt.frame_sieve else: sieve = None if opt.frame_type: strain = pycbc.frame.query_and_read_frame( - opt.frame_type, opt.channel_name, - start_time=opt.gps_start_time-opt.pad_data, - end_time=opt.gps_end_time+opt.pad_data, - sieve=sieve) + opt.frame_type, + opt.channel_name, + start_time=opt.gps_start_time - opt.pad_data, + end_time=opt.gps_end_time + opt.pad_data, + sieve=sieve, + ) elif opt.frame_files or opt.frame_cache: strain = pycbc.frame.read_frame( - frame_source, opt.channel_name, - start_time=opt.gps_start_time-opt.pad_data, - end_time=opt.gps_end_time+opt.pad_data, - sieve=sieve) + frame_source, + opt.channel_name, + start_time=opt.gps_start_time - opt.pad_data, + end_time=opt.gps_end_time + opt.pad_data, + sieve=sieve, + ) elif opt.hdf_store: - strain = pycbc.frame.read_store(opt.hdf_store, opt.channel_name, - opt.gps_start_time - opt.pad_data, - opt.gps_end_time + opt.pad_data) + strain = pycbc.frame.read_store( + opt.hdf_store, + opt.channel_name, + opt.gps_start_time - opt.pad_data, + opt.gps_end_time + opt.pad_data, + ) elif opt.fake_strain or opt.fake_strain_from_file: logger.info("Generating Fake Strain") @@ -247,30 +285,35 @@ def from_cli(opt, dyn_range_fac=1, precision='single', plen = round(opt.sample_rate / pdf) // 2 + 1 if opt.fake_strain_from_file: logger.info("Reading ASD from file") - strain_psd = pycbc.psd.from_txt(opt.fake_strain_from_file, - plen, pdf, - fake_flow, - is_asd_file=True) - elif opt.fake_strain != 'zeroNoise': + strain_psd = pycbc.psd.from_txt( + opt.fake_strain_from_file, plen, pdf, fake_flow, is_asd_file=True + ) + elif opt.fake_strain != "zeroNoise": logger.info("Making PSD for strain") - strain_psd = pycbc.psd.from_string(opt.fake_strain, plen, pdf, - fake_flow, **fake_extra_args) + strain_psd = pycbc.psd.from_string( + opt.fake_strain, plen, pdf, fake_flow, **fake_extra_args + ) - if opt.fake_strain == 'zeroNoise': + if opt.fake_strain == "zeroNoise": logger.info("Making zero-noise time series") - strain = TimeSeries(pycbc.types.zeros(duration * fake_rate), - delta_t=1.0 / fake_rate, - epoch=opt.gps_start_time - opt.pad_data) + strain = TimeSeries( + pycbc.types.zeros(duration * fake_rate), + delta_t=1.0 / fake_rate, + epoch=opt.gps_start_time - opt.pad_data, + ) else: logger.info("Making colored noise") from pycbc.noise.reproduceable import colored_noise - strain = colored_noise(strain_psd, - opt.gps_start_time - opt.pad_data, - opt.gps_end_time + opt.pad_data, - seed=opt.fake_strain_seed, - sample_rate=fake_rate, - low_frequency_cutoff=fake_flow, - filter_duration=1.0/pdf) + + strain = colored_noise( + strain_psd, + opt.gps_start_time - opt.pad_data, + opt.gps_end_time + opt.pad_data, + seed=opt.fake_strain_seed, + sample_rate=fake_rate, + low_frequency_cutoff=fake_flow, + filter_duration=1.0 / pdf, + ) if not strain.sample_rate_close(fake_rate): err_msg = "Actual sample rate of generated data does not match " @@ -283,8 +326,7 @@ def from_cli(opt, dyn_range_fac=1, precision='single', if opt.injection_frame_files or opt.injection_frame_type: if opt.injection_frame_files and opt.injection_frame_type: err_msg = ( - "You cannot supply both injection-frame-files and " - "injection-frame-type" + "You cannot supply both injection-frame-files and injection-frame-type" ) raise ValueError(err_msg) @@ -294,25 +336,26 @@ def from_cli(opt, dyn_range_fac=1, precision='single', injection_strain = pycbc.frame.query_and_read_frame( opt.injection_frame_type, opt.injection_channel_name, - start_time=opt.gps_start_time-opt.pad_data, - end_time=opt.gps_end_time+opt.pad_data, - sieve=None + start_time=opt.gps_start_time - opt.pad_data, + end_time=opt.gps_end_time + opt.pad_data, + sieve=None, ) else: injection_strain = pycbc.frame.read_frame( opt.injection_frame_files, opt.injection_channel_name, - start_time=opt.gps_start_time-opt.pad_data, - end_time=opt.gps_end_time+opt.pad_data, - sieve=None + start_time=opt.gps_start_time - opt.pad_data, + end_time=opt.gps_end_time + opt.pad_data, + sieve=None, ) strain.inject(injection_strain, copy=False) - if not opt.channel_name and (opt.injection_file \ - or opt.sgburst_injection_file): - raise ValueError('Please provide channel names with the format ' - 'ifo:channel (e.g. H1:CALIB-STRAIN) to inject ' - 'simulated signals into strain') + if not opt.channel_name and (opt.injection_file or opt.sgburst_injection_file): + raise ValueError( + "Please provide channel names with the format " + "ifo:channel (e.g. H1:CALIB-STRAIN) to inject " + "simulated signals into strain" + ) if opt.zpk_z and opt.zpk_p and opt.zpk_k: logger.info("Highpass Filtering") @@ -335,34 +378,37 @@ def from_cli(opt, dyn_range_fac=1, precision='single', if opt.sample_rate: logger.info("Resampling data") - strain = resample_to_delta_t(strain, - 1. / opt.sample_rate, - method='ldas') + strain = resample_to_delta_t(strain, 1.0 / opt.sample_rate, method="ldas") if injector is not None: logger.info("Applying injections") - injections = \ - injector.apply(strain, opt.channel_name.split(':')[0], - distance_scale=opt.injection_scale_factor, - injection_sample_rate=opt.injection_sample_rate, - inj_filter_rejector=inj_filter_rejector, - generate_injections=opt.generate_injections) + injections = injector.apply( + strain, + opt.channel_name.split(":")[0], + distance_scale=opt.injection_scale_factor, + injection_sample_rate=opt.injection_sample_rate, + inj_filter_rejector=inj_filter_rejector, + generate_injections=opt.generate_injections, + ) if opt.sgburst_injection_file: logger.info("Applying sine-Gaussian burst injections") injector = SGBurstInjectionSet(opt.sgburst_injection_file) - injector.apply(strain, opt.channel_name.split(':')[0], - distance_scale=opt.injection_scale_factor, - generate_injections=opt.generate_injections) + injector.apply( + strain, + opt.channel_name.split(":")[0], + distance_scale=opt.injection_scale_factor, + generate_injections=opt.generate_injections, + ) - if precision == 'single': + if precision == "single": logger.info("Converting to float32") strain = (strain * dyn_range_fac).astype(pycbc.types.float32) elif precision == "double": logger.info("Converting to float64") strain = (strain * dyn_range_fac).astype(pycbc.types.float64) else: - raise ValueError("Unrecognized precision {}".format(precision)) + raise ValueError(f"Unrecognized precision {precision}") if opt.gating_file is not None: logger.info("Gating times contained in gating file") @@ -370,35 +416,46 @@ def from_cli(opt, dyn_range_fac=1, precision='single', if len(gate_params.shape) == 1: gate_params = [gate_params] for gate_time, gate_window, gate_taper in gate_params: - strain = strain.gate(gate_time, window=gate_window, - method=opt.gating_method, - copy=False, - taper_width=gate_taper) - gating_info['file'] = \ - [gp for gp in gate_params \ - if (gp[0] + gp[1] + gp[2] >= strain.start_time) \ - and (gp[0] - gp[1] - gp[2] <= strain.end_time)] + strain = strain.gate( + gate_time, + window=gate_window, + method=opt.gating_method, + copy=False, + taper_width=gate_taper, + ) + gating_info["file"] = [ + gp + for gp in gate_params + if (gp[0] + gp[1] + gp[2] >= strain.start_time) + and (gp[0] - gp[1] - gp[2] <= strain.end_time) + ] if opt.autogating_threshold is not None: - gating_info['auto'] = [] + gating_info["auto"] = [] for _ in range(opt.autogating_max_iterations): glitch_times = detect_loud_glitches( - strain, threshold=opt.autogating_threshold, - cluster_window=opt.autogating_cluster, - low_freq_cutoff=opt.strain_high_pass, - corrupt_time=opt.pad_data + opt.autogating_pad) - gate_params = [[gt, opt.autogating_width, opt.autogating_taper] - for gt in glitch_times] - gating_info['auto'] += gate_params + strain, + threshold=opt.autogating_threshold, + cluster_window=opt.autogating_cluster, + low_freq_cutoff=opt.strain_high_pass, + corrupt_time=opt.pad_data + opt.autogating_pad, + ) + gate_params = [ + [gt, opt.autogating_width, opt.autogating_taper] for gt in glitch_times + ] + gating_info["auto"] += gate_params for gate_time, gate_window, gate_taper in gate_params: - strain = strain.gate(gate_time, window=gate_window, - method=opt.gating_method, - copy=False, - taper_width=gate_taper) + strain = strain.gate( + gate_time, + window=gate_window, + method=opt.gating_method, + copy=False, + taper_width=gate_taper, + ) if len(glitch_times) > 0: - logger.info('Autogating at %s', - ', '.join(['%.3f' % gt - for gt in glitch_times])) + logger.info( + "Autogating at %s", ", ".join(["%.3f" % gt for gt in glitch_times]) + ) else: break @@ -410,18 +467,20 @@ def from_cli(opt, dyn_range_fac=1, precision='single', logger.info("Lowpass Filtering") strain = lowpass(strain, frequency=opt.strain_low_pass) - if hasattr(opt, 'witness_frame_type') and opt.witness_frame_type: + if hasattr(opt, "witness_frame_type") and opt.witness_frame_type: stilde = strain.to_frequencyseries() from pycbc.io.hdf import HFile + tf_file = HFile(opt.witness_tf_file) for key in tf_file: - witness = pycbc.frame.query_and_read_frame(opt.witness_frame_type, - str(key), - start_time=strain.start_time, - end_time=strain.end_time) + witness = pycbc.frame.query_and_read_frame( + opt.witness_frame_type, + str(key), + start_time=strain.start_time, + end_time=strain.end_time, + ) witness = (witness * dyn_range_fac).astype(strain.dtype) - tf = pycbc.types.load_frequencyseries(opt.witness_tf_file, - group=key) + tf = pycbc.types.load_frequencyseries(opt.witness_tf_file, group=key) tf = tf.astype(stilde.dtype) flen = int(opt.witness_filter_length * strain.sample_rate) @@ -430,7 +489,7 @@ def from_cli(opt, dyn_range_fac=1, precision='single', tf_time = tf.to_timeseries() window = Array(numpy.hanning(flen * 2), dtype=strain.dtype) tf_time[0:flen] *= window[flen:] - tf_time[len(tf_time)-flen:] *= window[0:flen] + tf_time[len(tf_time) - flen :] *= window[0:flen] tf = tf_time.to_frequencyseries() kmax = min(len(tf), len(stilde) - 1) @@ -448,8 +507,8 @@ def from_cli(opt, dyn_range_fac=1, precision='single', logger.info("Tapering data") # Use auto-gating, a one-sided gate is a taper pd_taper_window = opt.taper_data - gate_params = [(strain.start_time, 0., pd_taper_window)] - gate_params.append((strain.end_time, 0., pd_taper_window)) + gate_params = [(strain.start_time, 0.0, pd_taper_window)] + gate_params.append((strain.end_time, 0.0, pd_taper_window)) gate_data(strain, gate_params) if injector is not None: @@ -458,13 +517,14 @@ def from_cli(opt, dyn_range_fac=1, precision='single', return strain + def from_cli_single_ifo(opt, ifo, inj_filter_rejector=None, **kwargs): """ Get the strain for a single ifo when using the multi-detector CLI """ single_det_opt = copy_opts_for_single_ifo(opt, ifo) - return from_cli(single_det_opt, - inj_filter_rejector=inj_filter_rejector, **kwargs) + return from_cli(single_det_opt, inj_filter_rejector=inj_filter_rejector, **kwargs) + def from_cli_multi_ifos(opt, ifos, inj_filter_rejector_dict=None, **kwargs): """ @@ -472,215 +532,342 @@ def from_cli_multi_ifos(opt, ifos, inj_filter_rejector_dict=None, **kwargs): """ strain = {} if inj_filter_rejector_dict is None: - inj_filter_rejector_dict = {ifo: None for ifo in ifos} + inj_filter_rejector_dict = dict.fromkeys(ifos) for ifo in ifos: - strain[ifo] = from_cli_single_ifo(opt, ifo, - inj_filter_rejector_dict[ifo], **kwargs) + strain[ifo] = from_cli_single_ifo( + opt, ifo, inj_filter_rejector_dict[ifo], **kwargs + ) return strain def insert_strain_option_group(parser, gps_times=True): - """ Add strain-related options to the optparser object. + """ + Add strain-related options to the optparser object. Adds the options used to call the pycbc.strain.from_cli function to an optparser as an OptionGroup. This should be used if you want to use these options in your code. Parameters - ----------- + ---------- parser : object OptionParser instance. gps_times : bool, optional Include ``--gps-start-time`` and ``--gps-end-time`` options. Default is True. - """ - data_reading_group = parser.add_argument_group("Options for obtaining h(t)", - "These options are used for generating h(t) either by " - "reading from a file or by generating it. This is only " - "needed if the PSD is to be estimated from the data, ie. " - " if the --psd-estimation option is given.") + """ + data_reading_group = parser.add_argument_group( + "Options for obtaining h(t)", + "These options are used for generating h(t) either by " + "reading from a file or by generating it. This is only " + "needed if the PSD is to be estimated from the data, ie. " + " if the --psd-estimation option is given.", + ) # Required options if gps_times: - data_reading_group.add_argument("--gps-start-time", - help="The gps start time of the data " - "(integer seconds)", type=int) - data_reading_group.add_argument("--gps-end-time", - help="The gps end time of the data " - "(integer seconds)", type=int) - - data_reading_group.add_argument("--strain-high-pass", type=float, - help="High pass frequency") - data_reading_group.add_argument("--strain-low-pass", type=float, - help="Low pass frequency") - data_reading_group.add_argument("--pad-data", default=8, - help="Extra padding to remove highpass corruption " - "(integer seconds, default 8)", type=int) - data_reading_group.add_argument("--taper-data", - help="Taper ends of data to zero using the supplied length as a " - "window (integer seconds)", type=int, default=0) - data_reading_group.add_argument("--sample-rate", type=float, - help="The sample rate to use for h(t) generation (integer Hz)") - data_reading_group.add_argument("--channel-name", type=str, - help="The channel containing the gravitational strain data") + data_reading_group.add_argument( + "--gps-start-time", + help="The gps start time of the data (integer seconds)", + type=int, + ) + data_reading_group.add_argument( + "--gps-end-time", + help="The gps end time of the data (integer seconds)", + type=int, + ) + + data_reading_group.add_argument( + "--strain-high-pass", type=float, help="High pass frequency" + ) + data_reading_group.add_argument( + "--strain-low-pass", type=float, help="Low pass frequency" + ) + data_reading_group.add_argument( + "--pad-data", + default=8, + help="Extra padding to remove highpass corruption (integer seconds, default 8)", + type=int, + ) + data_reading_group.add_argument( + "--taper-data", + help="Taper ends of data to zero using the supplied length as a " + "window (integer seconds)", + type=int, + default=0, + ) + data_reading_group.add_argument( + "--sample-rate", + type=float, + help="The sample rate to use for h(t) generation (integer Hz)", + ) + data_reading_group.add_argument( + "--channel-name", + type=str, + help="The channel containing the gravitational strain data", + ) # Read from cache file - data_reading_group.add_argument("--frame-cache", type=str, nargs="+", - help="Cache file containing the frame locations.") + data_reading_group.add_argument( + "--frame-cache", + type=str, + nargs="+", + help="Cache file containing the frame locations.", + ) # Read from frame files - data_reading_group.add_argument("--frame-files", - type=str, nargs="+", - help="list of frame files") + data_reading_group.add_argument( + "--frame-files", type=str, nargs="+", help="list of frame files" + ) # Read from hdf store file - data_reading_group.add_argument("--hdf-store", - type=str, - help="Store of time series data in hdf format") + data_reading_group.add_argument( + "--hdf-store", type=str, help="Store of time series data in hdf format" + ) # Use datafind to get frame files - data_reading_group.add_argument("--frame-type", - type=str, - metavar="S:TYPE", - help="(optional), replaces frame-files. Use datafind " - "to get the needed frame file(s) of this type " - "from site S.") + data_reading_group.add_argument( + "--frame-type", + type=str, + metavar="S:TYPE", + help="(optional), replaces frame-files. Use datafind " + "to get the needed frame file(s) of this type " + "from site S.", + ) # Filter frame files by URL - data_reading_group.add_argument("--frame-sieve", - type=str, - help="(optional), Only use frame files where the " - "URL matches the regular expression given.") + data_reading_group.add_argument( + "--frame-sieve", + type=str, + help="(optional), Only use frame files where the " + "URL matches the regular expression given.", + ) # Generate gaussian noise with given psd - data_reading_group.add_argument("--fake-strain", - help="Name of model PSD for generating fake gaussian noise.", - choices=pycbc.psd.get_psd_model_list() + ['zeroNoise']) - data_reading_group.add_argument("--fake-strain-extra-args", - nargs='+', action=DictOptionAction, - metavar='PARAM:VALUE', default={}, type=float, - help="(optional) Extra arguments passed to " - "the PSD models.") - data_reading_group.add_argument("--fake-strain-seed", type=int, default=0, - help="Seed value for the generation of fake colored" - " gaussian noise") - data_reading_group.add_argument("--fake-strain-from-file", - help="File containing ASD for generating fake noise from it.") - data_reading_group.add_argument("--fake-strain-flow", - default=1.0, type=float, - help="Low frequency cutoff of the fake strain") - data_reading_group.add_argument("--fake-strain-filter-duration", - default=128.0, type=float, - help="Duration in seconds of the fake data coloring filter") - data_reading_group.add_argument("--fake-strain-sample-rate", - default=16384, type=float, - help="Sample rate of the fake data generation") + data_reading_group.add_argument( + "--fake-strain", + help="Name of model PSD for generating fake gaussian noise.", + choices=pycbc.psd.get_psd_model_list() + ["zeroNoise"], + ) + data_reading_group.add_argument( + "--fake-strain-extra-args", + nargs="+", + action=DictOptionAction, + metavar="PARAM:VALUE", + default={}, + type=float, + help="(optional) Extra arguments passed to the PSD models.", + ) + data_reading_group.add_argument( + "--fake-strain-seed", + type=int, + default=0, + help="Seed value for the generation of fake colored gaussian noise", + ) + data_reading_group.add_argument( + "--fake-strain-from-file", + help="File containing ASD for generating fake noise from it.", + ) + data_reading_group.add_argument( + "--fake-strain-flow", + default=1.0, + type=float, + help="Low frequency cutoff of the fake strain", + ) + data_reading_group.add_argument( + "--fake-strain-filter-duration", + default=128.0, + type=float, + help="Duration in seconds of the fake data coloring filter", + ) + data_reading_group.add_argument( + "--fake-strain-sample-rate", + default=16384, + type=float, + help="Sample rate of the fake data generation", + ) # Injection options - data_reading_group.add_argument("--injection-file", type=str, - help="(optional) Injection file containing parameters" - " of CBC signals to be added to the strain") - data_reading_group.add_argument("--sgburst-injection-file", type=str, - help="(optional) Injection file containing parameters" - "of sine-Gaussian burst signals to add to the strain") - data_reading_group.add_argument("--do-not-inject-from-file", - action="store_false", dest="generate_injections", - default=True, - help="If this options are given, the injections in " - "injection-file or sgburst-injection-file are not " - "added into the data. This can be used for " - "debugging (ie. in minifollowups), or if using " - "injections from frame files where you need to " - "know injection parameters to allow the " - "injection optimization settings.") - data_reading_group.add_argument("--injection-scale-factor", type=float, - default=1, - help="Divide injections by this factor " - "before adding to the strain data") - data_reading_group.add_argument("--injection-sample-rate", type=float, - help="Sample rate to use for injections (integer Hz). " - "Typically similar to the strain data sample rate." - "If not provided, the strain sample rate will be " - "used") - data_reading_group.add_argument("--injection-f-ref", type=float, - help="Reference frequency in Hz for creating CBC " - "injections from an XML file") - data_reading_group.add_argument("--injection-f-final", type=float, - help="Override the f_final field of a CBC XML " - "injection file (frequency in Hz)") + data_reading_group.add_argument( + "--injection-file", + type=str, + help="(optional) Injection file containing parameters" + " of CBC signals to be added to the strain", + ) + data_reading_group.add_argument( + "--sgburst-injection-file", + type=str, + help="(optional) Injection file containing parameters" + "of sine-Gaussian burst signals to add to the strain", + ) + data_reading_group.add_argument( + "--do-not-inject-from-file", + action="store_false", + dest="generate_injections", + default=True, + help="If this options are given, the injections in " + "injection-file or sgburst-injection-file are not " + "added into the data. This can be used for " + "debugging (ie. in minifollowups), or if using " + "injections from frame files where you need to " + "know injection parameters to allow the " + "injection optimization settings.", + ) + data_reading_group.add_argument( + "--injection-scale-factor", + type=float, + default=1, + help="Divide injections by this factor before adding to the strain data", + ) + data_reading_group.add_argument( + "--injection-sample-rate", + type=float, + help="Sample rate to use for injections (integer Hz). " + "Typically similar to the strain data sample rate." + "If not provided, the strain sample rate will be " + "used", + ) + data_reading_group.add_argument( + "--injection-f-ref", + type=float, + help="Reference frequency in Hz for creating CBC injections from an XML file", + ) + data_reading_group.add_argument( + "--injection-f-final", + type=float, + help="Override the f_final field of a CBC XML injection file (frequency in Hz)", + ) # Options for getting injection from frame files - data_reading_group.add_argument("--injection-channel-name", type=str, - help="The channel containing the injection strain data") - data_reading_group.add_argument("--injection-frame-type", type=str, - help="We are going to add injections from frame files. " - "This will use datafind to get the needed frame " - "files of this type.") - data_reading_group.add_argument("--injection-frame-files", - type=str, nargs="+", - help="We are going to add injections from frame files. " - "This provides the list of frame files containing " - "injection strain.") + data_reading_group.add_argument( + "--injection-channel-name", + type=str, + help="The channel containing the injection strain data", + ) + data_reading_group.add_argument( + "--injection-frame-type", + type=str, + help="We are going to add injections from frame files. " + "This will use datafind to get the needed frame " + "files of this type.", + ) + data_reading_group.add_argument( + "--injection-frame-files", + type=str, + nargs="+", + help="We are going to add injections from frame files. " + "This provides the list of frame files containing " + "injection strain.", + ) # Gating options - data_reading_group.add_argument("--gating-file", type=str, - help="(optional) Text file of gating segments to apply." - " Format of each line is (all values in seconds):" - " gps_time zeros_half_width pad_half_width") - data_reading_group.add_argument('--autogating-threshold', type=float, - metavar='SIGMA', - help='If given, find and gate glitches ' - 'producing a deviation larger than ' - 'SIGMA in the whitened strain time ' - 'series.') - data_reading_group.add_argument('--autogating-max-iterations', type=int, - metavar='SIGMA', default=1, - help='If given, iteratively apply ' - 'autogating') - data_reading_group.add_argument('--autogating-cluster', type=float, - metavar='SECONDS', default=5., - help='Length of clustering window for ' - 'detecting glitches for autogating.') - data_reading_group.add_argument('--autogating-width', type=float, - metavar='SECONDS', default=0.25, - help='Half-width of the gating window.') - data_reading_group.add_argument('--autogating-taper', type=float, - metavar='SECONDS', default=0.25, - help='Taper the strain before and after ' - 'each gating window over a duration ' - 'of SECONDS.') - data_reading_group.add_argument('--autogating-pad', type=float, - metavar='SECONDS', default=16, - help='Ignore the given length of whitened ' - 'strain at the ends of a segment, to ' - 'avoid filters ringing.') - data_reading_group.add_argument('--gating-method', type=str, - default='taper', - help='Choose the method for gating. ' - 'Default: `taper`', - choices=['hard', 'taper', 'paint']) + data_reading_group.add_argument( + "--gating-file", + type=str, + help="(optional) Text file of gating segments to apply." + " Format of each line is (all values in seconds):" + " gps_time zeros_half_width pad_half_width", + ) + data_reading_group.add_argument( + "--autogating-threshold", + type=float, + metavar="SIGMA", + help="If given, find and gate glitches " + "producing a deviation larger than " + "SIGMA in the whitened strain time " + "series.", + ) + data_reading_group.add_argument( + "--autogating-max-iterations", + type=int, + metavar="SIGMA", + default=1, + help="If given, iteratively apply autogating", + ) + data_reading_group.add_argument( + "--autogating-cluster", + type=float, + metavar="SECONDS", + default=5.0, + help="Length of clustering window for detecting glitches for autogating.", + ) + data_reading_group.add_argument( + "--autogating-width", + type=float, + metavar="SECONDS", + default=0.25, + help="Half-width of the gating window.", + ) + data_reading_group.add_argument( + "--autogating-taper", + type=float, + metavar="SECONDS", + default=0.25, + help="Taper the strain before and after " + "each gating window over a duration " + "of SECONDS.", + ) + data_reading_group.add_argument( + "--autogating-pad", + type=float, + metavar="SECONDS", + default=16, + help="Ignore the given length of whitened " + "strain at the ends of a segment, to " + "avoid filters ringing.", + ) + data_reading_group.add_argument( + "--gating-method", + type=str, + default="taper", + help="Choose the method for gating. Default: `taper`", + choices=["hard", "taper", "paint"], + ) # Optional - data_reading_group.add_argument("--normalize-strain", type=float, - help="(optional) Divide frame data by constant.") - data_reading_group.add_argument("--zpk-z", type=float, nargs="+", - help="(optional) Zero-pole-gain (zpk) filter strain. " - "A list of zeros for transfer function") - data_reading_group.add_argument("--zpk-p", type=float, nargs="+", - help="(optional) Zero-pole-gain (zpk) filter strain. " - "A list of poles for transfer function") - data_reading_group.add_argument("--zpk-k", type=float, - help="(optional) Zero-pole-gain (zpk) filter strain. " - "Transfer function gain") + data_reading_group.add_argument( + "--normalize-strain", + type=float, + help="(optional) Divide frame data by constant.", + ) + data_reading_group.add_argument( + "--zpk-z", + type=float, + nargs="+", + help="(optional) Zero-pole-gain (zpk) filter strain. " + "A list of zeros for transfer function", + ) + data_reading_group.add_argument( + "--zpk-p", + type=float, + nargs="+", + help="(optional) Zero-pole-gain (zpk) filter strain. " + "A list of poles for transfer function", + ) + data_reading_group.add_argument( + "--zpk-k", + type=float, + help="(optional) Zero-pole-gain (zpk) filter strain. Transfer function gain", + ) # Options to apply to subtract noise from a witness channel and known # transfer function. - data_reading_group.add_argument("--witness-frame-type", type=str, - help="(optional), frame type which will be use to query the" - " witness channel data.") - data_reading_group.add_argument("--witness-tf-file", type=str, - help="an hdf file containing the transfer" - " functions and the associated channel names") - data_reading_group.add_argument("--witness-filter-length", type=float, - help="filter length in seconds for the transfer function") + data_reading_group.add_argument( + "--witness-frame-type", + type=str, + help="(optional), frame type which will be use to query the" + " witness channel data.", + ) + data_reading_group.add_argument( + "--witness-tf-file", + type=str, + help="an hdf file containing the transfer" + " functions and the associated channel names", + ) + data_reading_group.add_argument( + "--witness-filter-length", + type=float, + help="filter length in seconds for the transfer function", + ) return data_reading_group + # FIXME: This repeats almost all of the options above. Any nice way of reducing # this? def insert_strain_option_group_multi_ifo(parser, gps_times=True): @@ -690,264 +877,439 @@ def insert_strain_option_group_multi_ifo(parser, gps_times=True): want to use these options in your code. Parameters - ----------- + ---------- parser : object OptionParser instance. gps_times : bool, optional Include ``--gps-start-time`` and ``--gps-end-time`` options. Default is True. - """ - data_reading_group_multi = parser.add_argument_group("Options for obtaining" - " h(t)", - "These options are used for generating h(t) either by " - "reading from a file or by generating it. This is only " - "needed if the PSD is to be estimated from the data, ie. " - "if the --psd-estimation option is given. This group " - "supports reading from multiple ifos simultaneously.") + """ + data_reading_group_multi = parser.add_argument_group( + "Options for obtaining h(t)", + "These options are used for generating h(t) either by " + "reading from a file or by generating it. This is only " + "needed if the PSD is to be estimated from the data, ie. " + "if the --psd-estimation option is given. This group " + "supports reading from multiple ifos simultaneously.", + ) # Required options if gps_times: data_reading_group_multi.add_argument( - "--gps-start-time", nargs='+', action=MultiDetOptionAction, - metavar='IFO:TIME', type=int, - help="The gps start time of the data (integer seconds)") + "--gps-start-time", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + type=int, + help="The gps start time of the data (integer seconds)", + ) data_reading_group_multi.add_argument( - "--gps-end-time", nargs='+', action=MultiDetOptionAction, - metavar='IFO:TIME', type=int, - help="The gps end time of the data (integer seconds)") - - data_reading_group_multi.add_argument("--strain-high-pass", nargs='+', - action=MultiDetOptionAction, - type=float, metavar='IFO:FREQUENCY', - help="High pass frequency") - data_reading_group_multi.add_argument("--strain-low-pass", nargs='+', - action=MultiDetOptionAction, - type=float, metavar='IFO:FREQUENCY', - help="Low pass frequency") - data_reading_group_multi.add_argument("--pad-data", nargs='+', default=8, - action=MultiDetOptionAction, - type=int, metavar='IFO:LENGTH', - help="Extra padding to remove highpass corruption " - "(integer seconds, default 8)") - data_reading_group_multi.add_argument("--taper-data", nargs='+', - action=MultiDetOptionAction, - type=int, default=0, metavar='IFO:LENGTH', - help="Taper ends of data to zero using the " - "supplied length as a window (integer seconds)") - data_reading_group_multi.add_argument("--sample-rate", type=float, - nargs='+', - action=MultiDetOptionAction, metavar='IFO:RATE', - help="The sample rate to use for h(t) generation " - " (integer Hz).") - data_reading_group_multi.add_argument("--channel-name", type=str, nargs='+', - action=MultiDetOptionActionSpecial, - metavar='IFO:CHANNEL', - help="The channel containing the gravitational " - "strain data") + "--gps-end-time", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + type=int, + help="The gps end time of the data (integer seconds)", + ) + + data_reading_group_multi.add_argument( + "--strain-high-pass", + nargs="+", + action=MultiDetOptionAction, + type=float, + metavar="IFO:FREQUENCY", + help="High pass frequency", + ) + data_reading_group_multi.add_argument( + "--strain-low-pass", + nargs="+", + action=MultiDetOptionAction, + type=float, + metavar="IFO:FREQUENCY", + help="Low pass frequency", + ) + data_reading_group_multi.add_argument( + "--pad-data", + nargs="+", + default=8, + action=MultiDetOptionAction, + type=int, + metavar="IFO:LENGTH", + help="Extra padding to remove highpass corruption (integer seconds, default 8)", + ) + data_reading_group_multi.add_argument( + "--taper-data", + nargs="+", + action=MultiDetOptionAction, + type=int, + default=0, + metavar="IFO:LENGTH", + help="Taper ends of data to zero using the " + "supplied length as a window (integer seconds)", + ) + data_reading_group_multi.add_argument( + "--sample-rate", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:RATE", + help="The sample rate to use for h(t) generation (integer Hz).", + ) + data_reading_group_multi.add_argument( + "--channel-name", + type=str, + nargs="+", + action=MultiDetOptionActionSpecial, + metavar="IFO:CHANNEL", + help="The channel containing the gravitational strain data", + ) # Read from cache file - data_reading_group_multi.add_argument("--frame-cache", type=str, nargs="+", - action=MultiDetOptionAppendAction, - metavar='IFO:FRAME_CACHE', - help="Cache file containing the frame locations.") + data_reading_group_multi.add_argument( + "--frame-cache", + type=str, + nargs="+", + action=MultiDetOptionAppendAction, + metavar="IFO:FRAME_CACHE", + help="Cache file containing the frame locations.", + ) # Read from frame files - data_reading_group_multi.add_argument("--frame-files", type=str, nargs="+", - action=MultiDetOptionAppendAction, - metavar='IFO:FRAME_FILES', - help="list of frame files") + data_reading_group_multi.add_argument( + "--frame-files", + type=str, + nargs="+", + action=MultiDetOptionAppendAction, + metavar="IFO:FRAME_FILES", + help="list of frame files", + ) # Read from hdf store file - data_reading_group_multi.add_argument("--hdf-store", type=str, nargs='+', - action=MultiDetOptionAction, - metavar='IFO:HDF_STORE_FILE', - help="Store of time series data in hdf format") + data_reading_group_multi.add_argument( + "--hdf-store", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:HDF_STORE_FILE", + help="Store of time series data in hdf format", + ) # Use datafind to get frame files - data_reading_group_multi.add_argument("--frame-type", type=str, nargs="+", - action=MultiDetOptionActionSpecial, - metavar='IFO:FRAME_TYPE', - help="(optional) Replaces frame-files. " - "Use datafind to get the needed frame " - "file(s) of this type.") + data_reading_group_multi.add_argument( + "--frame-type", + type=str, + nargs="+", + action=MultiDetOptionActionSpecial, + metavar="IFO:FRAME_TYPE", + help="(optional) Replaces frame-files. " + "Use datafind to get the needed frame " + "file(s) of this type.", + ) # Filter frame files by URL - data_reading_group_multi.add_argument("--frame-sieve", type=str, nargs="+", - action=MultiDetOptionAction, - metavar='IFO:FRAME_SIEVE', - help="(optional), Only use frame files where the " - "URL matches the regular expression given.") + data_reading_group_multi.add_argument( + "--frame-sieve", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FRAME_SIEVE", + help="(optional), Only use frame files where the " + "URL matches the regular expression given.", + ) # Generate gaussian noise with given psd - data_reading_group_multi.add_argument("--fake-strain", type=str, nargs="+", - action=MultiDetOptionAction, metavar='IFO:CHOICE', - help="Name of model PSD for generating fake " - "gaussian noise. Choose from %s or zeroNoise" \ - %((', ').join(pycbc.psd.get_lalsim_psd_list()),) ) - data_reading_group_multi.add_argument("--fake-strain-extra-args", - nargs='+', action=MultiDetDictOptionAction, - metavar='DETECTOR:PARAM:VALUE', default={}, - type=float, help="(optional) Extra arguments " - "passed to the PSD models.") - data_reading_group_multi.add_argument("--fake-strain-seed", type=int, - default=0, nargs="+", action=MultiDetOptionAction, - metavar='IFO:SEED', - help="Seed value for the generation of fake " - "colored gaussian noise") - data_reading_group_multi.add_argument("--fake-strain-from-file", nargs="+", - action=MultiDetOptionAction, metavar='IFO:FILE', - help="File containing ASD for generating fake " - "noise from it.") - data_reading_group_multi.add_argument("--fake-strain-flow", - default=1.0, type=float, - nargs="+", action=MultiDetOptionAction, - help="Low frequency cutoff of the fake strain") - data_reading_group_multi.add_argument("--fake-strain-filter-duration", - default=128.0, type=float, - nargs="+", action=MultiDetOptionAction, - help="Duration in seconds of the fake data coloring filter") - data_reading_group_multi.add_argument("--fake-strain-sample-rate", - default=16384, type=float, - nargs="+", action=MultiDetOptionAction, - help="Sample rate of the fake data generation") + data_reading_group_multi.add_argument( + "--fake-strain", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:CHOICE", + help="Name of model PSD for generating fake " + "gaussian noise. Choose from %s or zeroNoise" + % ((", ").join(pycbc.psd.get_lalsim_psd_list()),), + ) + data_reading_group_multi.add_argument( + "--fake-strain-extra-args", + nargs="+", + action=MultiDetDictOptionAction, + metavar="DETECTOR:PARAM:VALUE", + default={}, + type=float, + help="(optional) Extra arguments passed to the PSD models.", + ) + data_reading_group_multi.add_argument( + "--fake-strain-seed", + type=int, + default=0, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:SEED", + help="Seed value for the generation of fake colored gaussian noise", + ) + data_reading_group_multi.add_argument( + "--fake-strain-from-file", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="File containing ASD for generating fake noise from it.", + ) + data_reading_group_multi.add_argument( + "--fake-strain-flow", + default=1.0, + type=float, + nargs="+", + action=MultiDetOptionAction, + help="Low frequency cutoff of the fake strain", + ) + data_reading_group_multi.add_argument( + "--fake-strain-filter-duration", + default=128.0, + type=float, + nargs="+", + action=MultiDetOptionAction, + help="Duration in seconds of the fake data coloring filter", + ) + data_reading_group_multi.add_argument( + "--fake-strain-sample-rate", + default=16384, + type=float, + nargs="+", + action=MultiDetOptionAction, + help="Sample rate of the fake data generation", + ) # Injection options - data_reading_group_multi.add_argument("--injection-file", type=str, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:FILE', - help="(optional) Injection file containing parameters" - "of CBC signals to be added to the strain") - data_reading_group_multi.add_argument("--do-not-inject-from-file", - default=True, action="store_false", - dest="generate_injections", - help="If this options are given, the injections in " - "injection-file or sgburst-injection-file are not " - "added into the data. This can be used for " - "debugging (ie. in minifollowups), or if using " - "injections from frame files where you need to " - "know injection parameters to allow the " - "injection optimization settings.") - - data_reading_group_multi.add_argument("--sgburst-injection-file", type=str, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:FILE', - help="(optional) Injection file containing parameters" - "of sine-Gaussian burst signals to add to the strain") - data_reading_group_multi.add_argument("--injection-scale-factor", - type=float, nargs="+", action=MultiDetOptionAction, - metavar="IFO:VAL", default=1., - help="Divide injections by this factor " - "before adding to the strain data") - data_reading_group_multi.add_argument("--injection-sample-rate", - type=float, nargs="+", action=MultiDetOptionAction, - metavar="IFO:VAL", - help="Sample rate to use for injections (integer Hz). " - "Typically similar to the strain data sample rate." - "If not provided, the strain sample rate will be " - "used") - - data_reading_group_multi.add_argument("--injection-f-ref", type=float, - action=MultiDetOptionAction, metavar='IFO:VALUE', - help="Reference frequency in Hz for creating CBC " - "injections from an XML file") - data_reading_group_multi.add_argument('--injection-f-final', type=float, - action=MultiDetOptionAction, metavar='IFO:VALUE', - help="Override the f_final field of a CBC XML " - "injection file (frequency in Hz)") + data_reading_group_multi.add_argument( + "--injection-file", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="(optional) Injection file containing parameters" + "of CBC signals to be added to the strain", + ) + data_reading_group_multi.add_argument( + "--do-not-inject-from-file", + default=True, + action="store_false", + dest="generate_injections", + help="If this options are given, the injections in " + "injection-file or sgburst-injection-file are not " + "added into the data. This can be used for " + "debugging (ie. in minifollowups), or if using " + "injections from frame files where you need to " + "know injection parameters to allow the " + "injection optimization settings.", + ) + + data_reading_group_multi.add_argument( + "--sgburst-injection-file", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="(optional) Injection file containing parameters" + "of sine-Gaussian burst signals to add to the strain", + ) + data_reading_group_multi.add_argument( + "--injection-scale-factor", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:VAL", + default=1.0, + help="Divide injections by this factor before adding to the strain data", + ) + data_reading_group_multi.add_argument( + "--injection-sample-rate", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:VAL", + help="Sample rate to use for injections (integer Hz). " + "Typically similar to the strain data sample rate." + "If not provided, the strain sample rate will be " + "used", + ) + + data_reading_group_multi.add_argument( + "--injection-f-ref", + type=float, + action=MultiDetOptionAction, + metavar="IFO:VALUE", + help="Reference frequency in Hz for creating CBC injections from an XML file", + ) + data_reading_group_multi.add_argument( + "--injection-f-final", + type=float, + action=MultiDetOptionAction, + metavar="IFO:VALUE", + help="Override the f_final field of a CBC XML injection file (frequency in Hz)", + ) # Options for getting injection from frame files - data_reading_group_multi.add_argument("--injection-channel-name", type=str, - action=MultiDetOptionAction, metavar='IFO:VALUE', - help="The channel containing the injection strain data") - data_reading_group_multi.add_argument("--injection-frame-type", type=str, - action=MultiDetOptionAction, metavar='IFO:VALUE', - help="We are going to add injections from frame files. " - "This will use datafind to get the needed frame " - "files of this type.") - data_reading_group_multi.add_argument("--injection-frame-files", - type=str, nargs="+", metavar='IFO:FRAME_FILES', - action=MultiDetOptionAppendAction, - help="We are going to add injections from frame files. " - "This provides the list of frame files containing " - "injection strain.") + data_reading_group_multi.add_argument( + "--injection-channel-name", + type=str, + action=MultiDetOptionAction, + metavar="IFO:VALUE", + help="The channel containing the injection strain data", + ) + data_reading_group_multi.add_argument( + "--injection-frame-type", + type=str, + action=MultiDetOptionAction, + metavar="IFO:VALUE", + help="We are going to add injections from frame files. " + "This will use datafind to get the needed frame " + "files of this type.", + ) + data_reading_group_multi.add_argument( + "--injection-frame-files", + type=str, + nargs="+", + metavar="IFO:FRAME_FILES", + action=MultiDetOptionAppendAction, + help="We are going to add injections from frame files. " + "This provides the list of frame files containing " + "injection strain.", + ) # Gating options - data_reading_group_multi.add_argument("--gating-file", nargs="+", - action=MultiDetOptionAction, - metavar='IFO:FILE', - help='(optional) Text file of gating segments to apply.' - ' Format of each line (units s) :' - ' gps_time zeros_half_width pad_half_width') - data_reading_group_multi.add_argument('--autogating-threshold', type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:SIGMA', - help='If given, find and gate glitches producing a ' - 'deviation larger than SIGMA in the whitened strain' - ' time series') - data_reading_group_multi.add_argument('--autogating-max-iterations', type=int, - metavar='SIGMA', default=1, - help='If given, iteratively apply ' - 'autogating') - data_reading_group_multi.add_argument('--autogating-cluster', type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:SECONDS', default=5., - help='Length of clustering window for ' - 'detecting glitches for autogating.') - data_reading_group_multi.add_argument('--autogating-width', type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:SECONDS', default=0.25, - help='Half-width of the gating window.') - data_reading_group_multi.add_argument('--autogating-taper', type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:SECONDS', default=0.25, - help='Taper the strain before and after ' - 'each gating window over a duration ' - 'of SECONDS.') - data_reading_group_multi.add_argument('--autogating-pad', type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:SECONDS', default=16, - help='Ignore the given length of whitened ' - 'strain at the ends of a segment, to ' - 'avoid filters ringing.') - data_reading_group_multi.add_argument('--gating-method', type=str, - nargs='+', action=MultiDetOptionAction, - default='taper', - help='Choose the method for gating. ' - 'Default: `taper`', - choices=['hard', 'taper', 'paint']) + data_reading_group_multi.add_argument( + "--gating-file", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", + help="(optional) Text file of gating segments to apply." + " Format of each line (units s) :" + " gps_time zeros_half_width pad_half_width", + ) + data_reading_group_multi.add_argument( + "--autogating-threshold", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:SIGMA", + help="If given, find and gate glitches producing a " + "deviation larger than SIGMA in the whitened strain" + " time series", + ) + data_reading_group_multi.add_argument( + "--autogating-max-iterations", + type=int, + metavar="SIGMA", + default=1, + help="If given, iteratively apply autogating", + ) + data_reading_group_multi.add_argument( + "--autogating-cluster", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:SECONDS", + default=5.0, + help="Length of clustering window for detecting glitches for autogating.", + ) + data_reading_group_multi.add_argument( + "--autogating-width", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:SECONDS", + default=0.25, + help="Half-width of the gating window.", + ) + data_reading_group_multi.add_argument( + "--autogating-taper", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:SECONDS", + default=0.25, + help="Taper the strain before and after " + "each gating window over a duration " + "of SECONDS.", + ) + data_reading_group_multi.add_argument( + "--autogating-pad", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:SECONDS", + default=16, + help="Ignore the given length of whitened " + "strain at the ends of a segment, to " + "avoid filters ringing.", + ) + data_reading_group_multi.add_argument( + "--gating-method", + type=str, + nargs="+", + action=MultiDetOptionAction, + default="taper", + help="Choose the method for gating. Default: `taper`", + choices=["hard", "taper", "paint"], + ) # Optional - data_reading_group_multi.add_argument("--normalize-strain", type=float, - nargs="+", action=MultiDetOptionAction, - metavar='IFO:VALUE', - help="(optional) Divide frame data by constant.") - data_reading_group_multi.add_argument("--zpk-z", type=float, - nargs="+", action=MultiDetOptionAppendAction, - metavar='IFO:VALUE', - help="(optional) Zero-pole-gain (zpk) filter strain. " - "A list of zeros for transfer function") - data_reading_group_multi.add_argument("--zpk-p", type=float, - nargs="+", action=MultiDetOptionAppendAction, - metavar='IFO:VALUE', - help="(optional) Zero-pole-gain (zpk) filter strain. " - "A list of poles for transfer function") - data_reading_group_multi.add_argument("--zpk-k", type=float, - nargs="+", action=MultiDetOptionAppendAction, - metavar='IFO:VALUE', - help="(optional) Zero-pole-gain (zpk) filter strain. " - "Transfer function gain") + data_reading_group_multi.add_argument( + "--normalize-strain", + type=float, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:VALUE", + help="(optional) Divide frame data by constant.", + ) + data_reading_group_multi.add_argument( + "--zpk-z", + type=float, + nargs="+", + action=MultiDetOptionAppendAction, + metavar="IFO:VALUE", + help="(optional) Zero-pole-gain (zpk) filter strain. " + "A list of zeros for transfer function", + ) + data_reading_group_multi.add_argument( + "--zpk-p", + type=float, + nargs="+", + action=MultiDetOptionAppendAction, + metavar="IFO:VALUE", + help="(optional) Zero-pole-gain (zpk) filter strain. " + "A list of poles for transfer function", + ) + data_reading_group_multi.add_argument( + "--zpk-k", + type=float, + nargs="+", + action=MultiDetOptionAppendAction, + metavar="IFO:VALUE", + help="(optional) Zero-pole-gain (zpk) filter strain. Transfer function gain", + ) return data_reading_group_multi ensure_one_opt_groups = [] -ensure_one_opt_groups.append(['--frame-cache','--fake-strain', - '--fake-strain-from-file', - '--frame-files', '--frame-type', - '--hdf-store']) - -required_opts_list = ['--gps-start-time', '--gps-end-time', - '--pad-data', '--sample-rate', - '--channel-name'] +ensure_one_opt_groups.append( + [ + "--frame-cache", + "--fake-strain", + "--fake-strain-from-file", + "--frame-files", + "--frame-type", + "--hdf-store", + ] +) + +required_opts_list = [ + "--gps-start-time", + "--gps-end-time", + "--pad-data", + "--sample-rate", + "--channel-name", +] def verify_strain_options(opts, parser): - """Sanity check provided strain arguments. + """ + Sanity check provided strain arguments. Parses the strain data CLI options and verifies that they are consistent and reasonable. @@ -961,6 +1323,7 @@ def verify_strain_options(opts, parser): fake-strain-seed). parser : object OptionParser instance. + """ for opt_group in ensure_one_opt_groups: ensure_one_opt(opts, parser, opt_group) @@ -968,7 +1331,8 @@ def verify_strain_options(opts, parser): def verify_strain_options_multi_ifo(opts, parser, ifos): - """Sanity check provided strain arguments. + """ + Sanity check provided strain arguments. Parses the strain data CLI options and verifies that they are consistent and reasonable. @@ -984,6 +1348,7 @@ def verify_strain_options_multi_ifo(opts, parser, ifos): OptionParser instance. ifos : list of strings List of ifos for which to verify options for + """ for ifo in ifos: for opt_group in ensure_one_opt_groups: @@ -992,7 +1357,8 @@ def verify_strain_options_multi_ifo(opts, parser, ifos): def gate_data(data, gate_params): - """Apply a set of gating windows to a time series. + """ + Apply a set of gating windows to a time series. Each gating window is defined by a central time, a given duration (centered on the given @@ -1015,44 +1381,57 @@ def gate_data(data, gate_params): ------- data: TimeSeries The gated time series. + """ + def inverted_tukey(M, n_pad): - midlen = M - 2*n_pad + midlen = M - 2 * n_pad if midlen < 0: raise ValueError("No zeros left after applying padding.") - padarr = 0.5*(1.+numpy.cos(numpy.pi*numpy.arange(n_pad)/n_pad)) - return numpy.concatenate((padarr,numpy.zeros(midlen),padarr[::-1])) + padarr = 0.5 * (1.0 + numpy.cos(numpy.pi * numpy.arange(n_pad) / n_pad)) + return numpy.concatenate((padarr, numpy.zeros(midlen), padarr[::-1])) - sample_rate = 1./data.delta_t + sample_rate = 1.0 / data.delta_t temp = data.data for glitch_time, glitch_width, pad_width in gate_params: t_start = glitch_time - glitch_width - pad_width - data.start_time t_end = glitch_time + glitch_width + pad_width - data.start_time - if t_start > data.duration or t_end < 0.: - continue # Skip gate segments that don't overlap - win_samples = int(2*sample_rate*(glitch_width+pad_width)) - pad_samples = int(sample_rate*pad_width) + if t_start > data.duration or t_end < 0.0: + continue # Skip gate segments that don't overlap + win_samples = int(2 * sample_rate * (glitch_width + pad_width)) + pad_samples = int(sample_rate * pad_width) window = inverted_tukey(win_samples, pad_samples) offset = int(t_start * sample_rate) idx1 = max(0, -offset) - idx2 = min(len(window), len(data)-offset) - temp[idx1+offset:idx2+offset] *= window[idx1:idx2] + idx2 = min(len(window), len(data) - offset) + temp[idx1 + offset : idx2 + offset] *= window[idx1:idx2] return data -class StrainSegments(object): - """ Class for managing manipulation of strain data for the purpose of - matched filtering. This includes methods for segmenting and - conditioning. +class StrainSegments: + """ + Class for managing manipulation of strain data for the purpose of + matched filtering. This includes methods for segmenting and + conditioning. """ - def __init__(self, strain, segment_length=None, segment_start_pad=0, - segment_end_pad=0, trigger_start=None, trigger_end=None, - filter_inj_only=False, injection_window=None, - allow_zero_padding=False): - """ Determine how to chop up the strain data into smaller segments - for analysis. + + def __init__( + self, + strain, + segment_length=None, + segment_start_pad=0, + segment_end_pad=0, + trigger_start=None, + trigger_end=None, + filter_inj_only=False, + injection_window=None, + allow_zero_padding=False, + ): + """ + Determine how to chop up the strain data into smaller segments + for analysis. """ self._fourier_segments = None self.strain = strain @@ -1081,8 +1460,8 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, min_start_time = int(strain.start_time) if trigger_start < min_start_time: err_msg = "Trigger start time must be within analysable " - err_msg += "window. Asked to start from %d " %(trigger_start) - err_msg += "but can only analyse from %d." %(min_start_time) + err_msg += "window. Asked to start from %d " % (trigger_start) + err_msg += "but can only analyse from %d." % (min_start_time) raise ValueError(err_msg) if not trigger_end: @@ -1094,25 +1473,23 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, max_end_time = int(strain.end_time) if trigger_end > max_end_time: err_msg = "Trigger end time must be within analysable " - err_msg += "window. Asked to end at %d " %(trigger_end) - err_msg += "but can only analyse to %d." %(max_end_time) + err_msg += "window. Asked to end at %d " % (trigger_end) + err_msg += "but can only analyse to %d." % (max_end_time) raise ValueError(err_msg) - throwaway_size = seg_start_pad + seg_end_pad seg_width = seg_len - throwaway_size # The amount of time we can actually analyze given the # amount of padding that is needed analyzable = trigger_end - trigger_start - data_start = (trigger_start - segment_start_pad) - \ - int(strain.start_time) + data_start = (trigger_start - segment_start_pad) - int(strain.start_time) data_end = trigger_end + segment_end_pad - int(strain.start_time) data_dur = data_end - data_start data_start = data_start * strain.sample_rate data_end = data_end * strain.sample_rate - #number of segments we need to analyze this data + # number of segments we need to analyze this data num_segs = int(numpy.ceil(float(analyzable) / float(seg_width))) # The offset we will use between segments @@ -1121,9 +1498,9 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, self.analyze_slices = [] # Determine how to chop up the strain into smaller segments - for nseg in range(num_segs-1): + for nseg in range(num_segs - 1): # boundaries for time slices into the strain - seg_start = int(data_start + (nseg*seg_offset) * strain.sample_rate) + seg_start = int(data_start + (nseg * seg_offset) * strain.sample_rate) seg_end = int(seg_start + seg_len * strain.sample_rate) seg_slice = slice(seg_start, seg_end) self.segment_slices.append(seg_slice) @@ -1140,7 +1517,7 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, seg_slice = slice(seg_start, seg_end) self.segment_slices.append(seg_slice) - remaining = (data_dur - ((num_segs - 1) * seg_offset + seg_start_pad)) + remaining = data_dur - ((num_segs - 1) * seg_offset + seg_start_pad) ana_start = int((seg_len - remaining) * strain.sample_rate) ana_end = int((seg_len - seg_end_pad) * strain.sample_rate) ana_slice = slice(ana_start, ana_end) @@ -1148,16 +1525,23 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, self.full_segment_slices = copy.deepcopy(self.segment_slices) - #Remove segments that are outside trig start and end + # Remove segments that are outside trig start and end segment_slices_red = [] analyze_slices_red = [] trig_start_idx = (trigger_start - int(strain.start_time)) * strain.sample_rate trig_end_idx = (trigger_end - int(strain.start_time)) * strain.sample_rate - if filter_inj_only and hasattr(strain, 'injections'): + if filter_inj_only and hasattr(strain, "injections"): end_times = strain.injections.end_times() - end_times = [time for time in end_times if float(time) < trigger_end and float(time) > trigger_start] - inj_idx = [(float(time) - float(strain.start_time)) * strain.sample_rate for time in end_times] + end_times = [ + time + for time in end_times + if float(time) < trigger_end and float(time) > trigger_start + ] + inj_idx = [ + (float(time) - float(strain.start_time)) * strain.sample_rate + for time in end_times + ] for seg, ana in zip(self.segment_slices, self.analyze_slices): start = ana.start @@ -1167,18 +1551,19 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, # adjust first segment if trig_start_idx > cum_start: - start += (trig_start_idx - cum_start) + start += trig_start_idx - cum_start # adjust last segment if trig_end_idx < cum_end: - stop -= (cum_end - trig_end_idx) + stop -= cum_end - trig_end_idx - if filter_inj_only and hasattr(strain, 'injections'): + if filter_inj_only and hasattr(strain, "injections"): analyze_this = False inj_window = strain.sample_rate * 8 for inj_id in inj_idx: - if inj_id < (cum_end + inj_window) and \ - inj_id > (cum_start - inj_window): + if inj_id < (cum_end + inj_window) and inj_id > ( + cum_start - inj_window + ): analyze_this = True if not analyze_this: @@ -1192,7 +1577,8 @@ def __init__(self, strain, segment_length=None, segment_start_pad=0, self.analyze_slices = analyze_slices_red def fourier_segments(self): - """ Return a list of the FFT'd segments. + """ + Return a list of the FFT'd segments. Return the list of FrequencySeries. Additional properties are added that describe the strain segment. The property 'analyze' is a slice corresponding to the portion of the time domain equivalent @@ -1207,11 +1593,11 @@ def fourier_segments(self): # Assume that we cannot have a case where we both zero-pad on # both sides elif seg_slice.start < 0: - strain_chunk = self.strain[:seg_slice.stop] + strain_chunk = self.strain[: seg_slice.stop] strain_chunk.prepend_zeros(-seg_slice.start) freq_seg = make_frequency_series(strain_chunk) elif seg_slice.stop > len(self.strain): - strain_chunk = self.strain[seg_slice.start:] + strain_chunk = self.strain[seg_slice.start :] strain_chunk.append_zeros(seg_slice.stop - len(self.strain)) freq_seg = make_frequency_series(strain_chunk) freq_seg.analyze = ana @@ -1223,116 +1609,182 @@ def fourier_segments(self): @classmethod def from_cli(cls, opt, strain): - """Calculate the segmentation of the strain data for analysis from + """ + Calculate the segmentation of the strain data for analysis from the command line options. """ - return cls(strain, segment_length=opt.segment_length, - segment_start_pad=opt.segment_start_pad, - segment_end_pad=opt.segment_end_pad, - trigger_start=opt.trig_start_time, - trigger_end=opt.trig_end_time, - filter_inj_only=opt.filter_inj_only, - injection_window=opt.injection_window, - allow_zero_padding=opt.allow_zero_padding) + return cls( + strain, + segment_length=opt.segment_length, + segment_start_pad=opt.segment_start_pad, + segment_end_pad=opt.segment_end_pad, + trigger_start=opt.trig_start_time, + trigger_end=opt.trig_end_time, + filter_inj_only=opt.filter_inj_only, + injection_window=opt.injection_window, + allow_zero_padding=opt.allow_zero_padding, + ) @classmethod def insert_segment_option_group(cls, parser): segment_group = parser.add_argument_group( - "Options for segmenting the strain", - "These options are used to determine how to " - "segment the strain into smaller chunks, " - "and for determining the portion of each to " - "analyze for triggers. ") - segment_group.add_argument("--trig-start-time", type=int, default=0, - help="(optional) The gps time to start recording triggers") - segment_group.add_argument("--trig-end-time", type=int, default=0, - help="(optional) The gps time to stop recording triggers") - segment_group.add_argument("--segment-length", type=int, - help="The length of each strain segment in seconds.") - segment_group.add_argument("--segment-start-pad", type=int, - help="The time in seconds to ignore of the " - "beginning of each segment in seconds. ") - segment_group.add_argument("--segment-end-pad", type=int, - help="The time in seconds to ignore at the " - "end of each segment in seconds.") - segment_group.add_argument("--allow-zero-padding", action='store_true', - help="Allow for zero padding of data to " - "analyze requested times, if needed.") + "Options for segmenting the strain", + "These options are used to determine how to " + "segment the strain into smaller chunks, " + "and for determining the portion of each to " + "analyze for triggers. ", + ) + segment_group.add_argument( + "--trig-start-time", + type=int, + default=0, + help="(optional) The gps time to start recording triggers", + ) + segment_group.add_argument( + "--trig-end-time", + type=int, + default=0, + help="(optional) The gps time to stop recording triggers", + ) + segment_group.add_argument( + "--segment-length", + type=int, + help="The length of each strain segment in seconds.", + ) + segment_group.add_argument( + "--segment-start-pad", + type=int, + help="The time in seconds to ignore of the " + "beginning of each segment in seconds. ", + ) + segment_group.add_argument( + "--segment-end-pad", + type=int, + help="The time in seconds to ignore at the end of each segment in seconds.", + ) + segment_group.add_argument( + "--allow-zero-padding", + action="store_true", + help="Allow for zero padding of data to " + "analyze requested times, if needed.", + ) # Injection optimization options - segment_group.add_argument("--filter-inj-only", action='store_true', - help="Analyze only segments that contain an injection.") - segment_group.add_argument("--injection-window", default=None, - type=float, help="""If using --filter-inj-only then + segment_group.add_argument( + "--filter-inj-only", + action="store_true", + help="Analyze only segments that contain an injection.", + ) + segment_group.add_argument( + "--injection-window", + default=None, + type=float, + help="""If using --filter-inj-only then only search for injections within +/- injection window of the injections's end time. This is useful to speed up a coherent search or a search where we initially filter at lower sample rate, and then filter at full rate where needed. NOTE: Reverts to full analysis if two injections are in the same - segment.""") - + segment.""", + ) @classmethod def from_cli_single_ifo(cls, opt, strain, ifo): - """Calculate the segmentation of the strain data for analysis from + """ + Calculate the segmentation of the strain data for analysis from the command line options. """ - return cls(strain, segment_length=opt.segment_length[ifo], - segment_start_pad=opt.segment_start_pad[ifo], - segment_end_pad=opt.segment_end_pad[ifo], - trigger_start=opt.trig_start_time[ifo], - trigger_end=opt.trig_end_time[ifo], - filter_inj_only=opt.filter_inj_only, - allow_zero_padding=opt.allow_zero_padding) + return cls( + strain, + segment_length=opt.segment_length[ifo], + segment_start_pad=opt.segment_start_pad[ifo], + segment_end_pad=opt.segment_end_pad[ifo], + trigger_start=opt.trig_start_time[ifo], + trigger_end=opt.trig_end_time[ifo], + filter_inj_only=opt.filter_inj_only, + allow_zero_padding=opt.allow_zero_padding, + ) @classmethod def from_cli_multi_ifos(cls, opt, strain_dict, ifos): - """Calculate the segmentation of the strain data for analysis from + """ + Calculate the segmentation of the strain data for analysis from the command line options. """ strain_segments = {} for ifo in ifos: - strain_segments[ifo] = cls.from_cli_single_ifo( - opt, strain_dict[ifo], ifo) + strain_segments[ifo] = cls.from_cli_single_ifo(opt, strain_dict[ifo], ifo) return strain_segments @classmethod def insert_segment_option_group_multi_ifo(cls, parser): segment_group = parser.add_argument_group( - "Options for segmenting the strain", - "These options are used to determine how to " - "segment the strain into smaller chunks, " - "and for determining the portion of each to " - "analyze for triggers. ") - segment_group.add_argument("--trig-start-time", type=int, default=0, - nargs='+', action=MultiDetOptionAction, metavar='IFO:TIME', - help="(optional) The gps time to start recording triggers") - segment_group.add_argument("--trig-end-time", type=int, default=0, - nargs='+', action=MultiDetOptionAction, metavar='IFO:TIME', - help="(optional) The gps time to stop recording triggers") - segment_group.add_argument("--segment-length", type=int, - nargs='+', action=MultiDetOptionAction, - metavar='IFO:LENGTH', - help="The length of each strain segment in seconds.") - segment_group.add_argument("--segment-start-pad", type=int, - nargs='+', action=MultiDetOptionAction, metavar='IFO:TIME', - help="The time in seconds to ignore of the " - "beginning of each segment in seconds. ") - segment_group.add_argument("--segment-end-pad", type=int, - nargs='+', action=MultiDetOptionAction, metavar='IFO:TIME', - help="The time in seconds to ignore at the " - "end of each segment in seconds.") - segment_group.add_argument("--allow-zero-padding", action='store_true', - help="Allow for zero padding of data to analyze " - "requested times, if needed.") - segment_group.add_argument("--filter-inj-only", action='store_true', - help="Analyze only segments that contain " - "an injection.") - - required_opts_list = ['--segment-length', - '--segment-start-pad', - '--segment-end-pad', - ] + "Options for segmenting the strain", + "These options are used to determine how to " + "segment the strain into smaller chunks, " + "and for determining the portion of each to " + "analyze for triggers. ", + ) + segment_group.add_argument( + "--trig-start-time", + type=int, + default=0, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="(optional) The gps time to start recording triggers", + ) + segment_group.add_argument( + "--trig-end-time", + type=int, + default=0, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="(optional) The gps time to stop recording triggers", + ) + segment_group.add_argument( + "--segment-length", + type=int, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:LENGTH", + help="The length of each strain segment in seconds.", + ) + segment_group.add_argument( + "--segment-start-pad", + type=int, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="The time in seconds to ignore of the " + "beginning of each segment in seconds. ", + ) + segment_group.add_argument( + "--segment-end-pad", + type=int, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:TIME", + help="The time in seconds to ignore at the end of each segment in seconds.", + ) + segment_group.add_argument( + "--allow-zero-padding", + action="store_true", + help="Allow for zero padding of data to analyze " + "requested times, if needed.", + ) + segment_group.add_argument( + "--filter-inj-only", + action="store_true", + help="Analyze only segments that contain an injection.", + ) + + required_opts_list = [ + "--segment-length", + "--segment-start-pad", + "--segment-end-pad", + ] @classmethod def verify_segment_options(cls, opt, parser): @@ -1346,13 +1798,10 @@ def verify_segment_options_multi_ifo(cls, opt, parser, ifos): @functools.lru_cache(maxsize=500) def create_memory_and_engine_for_class_based_fft( - npoints_time, - dtype, - delta_t=1, - ifft=False, - uid=0 + npoints_time, dtype, delta_t=1, ifft=False, uid=0 ): - """ Create memory and engine for class-based FFT/IFFT + """ + Create memory and engine for class-based FFT/IFFT Currently only supports R2C FFT / C2R IFFTs, but this could be expanded if use-cases arise. @@ -1374,24 +1823,15 @@ def create_memory_and_engine_for_class_based_fft( Provide a unique identifier. This is used to provide a separate set of memory in the cache, for instance if calling this from different codes. + """ npoints_freq = npoints_time // 2 + 1 delta_f_tmp = 1.0 / (npoints_time * delta_t) - vec = TimeSeries( - zeros( - npoints_time, - dtype=dtype - ), - delta_t=delta_t, - copy=False - ) + vec = TimeSeries(zeros(npoints_time, dtype=dtype), delta_t=delta_t, copy=False) vectilde = FrequencySeries( - zeros( - npoints_freq, - dtype=complex_same_precision_as(vec) - ), + zeros(npoints_freq, dtype=complex_same_precision_as(vec)), delta_f=delta_f_tmp, - copy=False + copy=False, ) if ifft: fft_class = IFFT(vectilde, vec) @@ -1405,12 +1845,14 @@ def create_memory_and_engine_for_class_based_fft( return invec, outvec, fft_class -def execute_cached_fft(invec_data, normalize_by_rate=True, ifft=False, - copy_output=True, uid=0): - """ Executes a cached FFT +def execute_cached_fft( + invec_data, normalize_by_rate=True, ifft=False, copy_output=True, uid=0 +): + """ + Executes a cached FFT Parameters - ----------- + ---------- invec_data : Array Array which will be used as input when fft_class is executed. normalize_by_rate : boolean (optional, default:False) @@ -1430,8 +1872,10 @@ def execute_cached_fft(invec_data, normalize_by_rate=True, ifft=False, Provide a unique identifier. This is used to provide a separate set of memory in the cache, for instance if calling this from different codes. + """ from pycbc.types import real_same_precision_as + if ifft: npoints_time = (len(invec_data) - 1) * 2 else: @@ -1449,11 +1893,7 @@ def execute_cached_fft(invec_data, normalize_by_rate=True, ifft=False, dtype = real_same_precision_as(invec_data) invec, outvec, fft_class = create_memory_and_engine_for_class_based_fft( - npoints_time, - dtype, - delta_t=delta_t, - ifft=ifft, - uid=uid + npoints_time, dtype, delta_t=delta_t, ifft=ifft, uid=uid ) if invec_data is not None: @@ -1474,10 +1914,11 @@ def execute_cached_fft(invec_data, normalize_by_rate=True, ifft=False, def execute_cached_ifft(*args, **kwargs): - """ Executes a cached IFFT + """ + Executes a cached IFFT Parameters - ----------- + ---------- invec_data : Array Array which will be used as input when fft_class is executed. normalize_by_rate : boolean (optional, default:False) @@ -1494,6 +1935,7 @@ def execute_cached_ifft(*args, **kwargs): Provide a unique identifier. This is used to provide a separate set of memory in the cache, for instance if calling this from different codes. + """ return execute_cached_fft(*args, **kwargs, ifft=True) @@ -1505,40 +1947,47 @@ def execute_cached_ifft(*args, **kwargs): STRAINBUFFER_UNIQUE_ID_2 = 778946541 STRAINBUFFER_UNIQUE_ID_3 = 665849947 + class StrainBuffer(pycbc.frame.DataBuffer): - def __init__(self, frame_src, channel_name, start_time, - max_buffer, - sample_rate, - low_frequency_cutoff=20, - highpass_frequency=15.0, - highpass_reduction=200.0, - highpass_bandwidth=5.0, - psd_samples=30, - psd_segment_length=4, - psd_inverse_length=3.5, - trim_padding=0.25, - autogating_threshold=None, - autogating_cluster=None, - autogating_pad=None, - autogating_width=None, - autogating_taper=None, - autogating_duration=None, - autogating_psd_segment_length=None, - autogating_psd_stride=None, - state_channel=None, - data_quality_channel=None, - idq_channel=None, - idq_state_channel=None, - idq_threshold=None, - dyn_range_fac=pycbc.DYN_RANGE_FAC, - psd_abort_difference=None, - psd_recalculate_difference=None, - force_update_cache=True, - increment_update_cache=None, - analyze_flags=None, - data_quality_flags=None, - dq_padding=0): - """ Class to produce overwhitened strain incrementally + def __init__( + self, + frame_src, + channel_name, + start_time, + max_buffer, + sample_rate, + low_frequency_cutoff=20, + highpass_frequency=15.0, + highpass_reduction=200.0, + highpass_bandwidth=5.0, + psd_samples=30, + psd_segment_length=4, + psd_inverse_length=3.5, + trim_padding=0.25, + autogating_threshold=None, + autogating_cluster=None, + autogating_pad=None, + autogating_width=None, + autogating_taper=None, + autogating_duration=None, + autogating_psd_segment_length=None, + autogating_psd_stride=None, + state_channel=None, + data_quality_channel=None, + idq_channel=None, + idq_state_channel=None, + idq_threshold=None, + dyn_range_fac=pycbc.DYN_RANGE_FAC, + psd_abort_difference=None, + psd_recalculate_difference=None, + force_update_cache=True, + increment_update_cache=None, + analyze_flags=None, + data_quality_flags=None, + dq_padding=0, + ): + """ + Class to produce overwhitened strain incrementally Parameters ---------- @@ -1621,11 +2070,16 @@ def __init__(self, frame_src, channel_name, start_time, is an alternate to the forced updated of the frame cache, and apptempts to predict the next frame file name without probing the filesystem. + """ - super(StrainBuffer, self).__init__(frame_src, channel_name, start_time, - max_buffer=max_buffer, - force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache) + super().__init__( + frame_src, + channel_name, + start_time, + max_buffer=max_buffer, + force_update_cache=force_update_cache, + increment_update_cache=increment_update_cache, + ) self.low_frequency_cutoff = low_frequency_cutoff @@ -1640,52 +2094,68 @@ def __init__(self, frame_src, channel_name, start_time, # State channel if state_channel is not None: valid_mask = pycbc.frame.flag_names_to_bitmask(self.analyze_flags) - logger.info('State channel %s interpreted as bitmask %s = good', - state_channel, bin(valid_mask)) + logger.info( + "State channel %s interpreted as bitmask %s = good", + state_channel, + bin(valid_mask), + ) self.state = pycbc.frame.StatusBuffer( frame_src, - state_channel, start_time, + state_channel, + start_time, max_buffer=max_buffer, valid_mask=valid_mask, force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache) + increment_update_cache=increment_update_cache, + ) # low latency dq channel if data_quality_channel is not None: - sb_kwargs = dict(max_buffer=max_buffer, - force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache) - if len(self.data_quality_flags) == 1 \ - and self.data_quality_flags[0] == 'veto_nonzero': - sb_kwargs['valid_on_zero'] = True - logger.info('DQ channel %s interpreted as zero = good', - data_quality_channel) + sb_kwargs = dict( + max_buffer=max_buffer, + force_update_cache=force_update_cache, + increment_update_cache=increment_update_cache, + ) + if ( + len(self.data_quality_flags) == 1 + and self.data_quality_flags[0] == "veto_nonzero" + ): + sb_kwargs["valid_on_zero"] = True + logger.info( + "DQ channel %s interpreted as zero = good", data_quality_channel + ) else: - sb_kwargs['valid_mask'] = pycbc.frame.flag_names_to_bitmask( - self.data_quality_flags) + sb_kwargs["valid_mask"] = pycbc.frame.flag_names_to_bitmask( + self.data_quality_flags + ) logger.info( - 'DQ channel %s interpreted as bitmask %s = good', + "DQ channel %s interpreted as bitmask %s = good", data_quality_channel, - bin(sb_kwargs['valid_mask']) + bin(sb_kwargs["valid_mask"]), ) - self.dq = pycbc.frame.StatusBuffer(frame_src, data_quality_channel, - start_time, **sb_kwargs) + self.dq = pycbc.frame.StatusBuffer( + frame_src, data_quality_channel, start_time, **sb_kwargs + ) if idq_channel is not None: if idq_state_channel is None: raise ValueError( - 'Each detector with an iDQ channel requires an iDQ state channel as well') + "Each detector with an iDQ channel requires an iDQ state channel as well" + ) if idq_threshold is None: raise ValueError( - 'If an iDQ channel is provided, a veto threshold must also be provided') - self.idq = pycbc.frame.iDQBuffer(frame_src, - idq_channel, - idq_state_channel, - idq_threshold, - start_time, - max_buffer=max_buffer, - force_update_cache=force_update_cache, - increment_update_cache=increment_update_cache) + "If an iDQ channel is provided, a veto threshold must also be provided" + ) + self.idq = pycbc.frame.iDQBuffer( + frame_src, + idq_channel, + idq_state_channel, + idq_threshold, + start_time, + max_buffer=max_buffer, + force_update_cache=force_update_cache, + increment_update_cache=increment_update_cache, + ) self.highpass_frequency = highpass_frequency self.highpass_reduction = highpass_reduction @@ -1713,27 +2183,30 @@ def __init__(self, frame_src, channel_name, start_time, self.psds = {} strain_len = int(max_buffer * self.sample_rate) - self.strain = TimeSeries(zeros(strain_len, dtype=numpy.float32), - delta_t=1.0/self.sample_rate, - epoch=start_time-max_buffer) + self.strain = TimeSeries( + zeros(strain_len, dtype=numpy.float32), + delta_t=1.0 / self.sample_rate, + epoch=start_time - max_buffer, + ) # Determine the total number of corrupted samples for highpass # and PSD over whitening - highpass_samples, self.beta = kaiserord(self.highpass_reduction, - self.highpass_bandwidth / self.raw_buffer.sample_rate * 2 * numpy.pi) - self.highpass_samples = int(highpass_samples / 2) - resample_corruption = 10 # If using the ldas method + highpass_samples, self.beta = kaiserord( + self.highpass_reduction, + self.highpass_bandwidth / self.raw_buffer.sample_rate * 2 * numpy.pi, + ) + self.highpass_samples = int(highpass_samples / 2) + resample_corruption = 10 # If using the ldas method self.factor = round(1.0 / self.raw_buffer.delta_t / self.sample_rate) self.corruption = self.highpass_samples // self.factor + resample_corruption - self.psd_corruption = self.psd_inverse_length * self.sample_rate + self.psd_corruption = self.psd_inverse_length * self.sample_rate self.total_corruption = self.corruption + self.psd_corruption # Determine how much padding is needed after removing the parts # associated with PSD over whitening and highpass filtering self.trim_padding = int(trim_padding * self.sample_rate) - if self.trim_padding > self.total_corruption: - self.trim_padding = self.total_corruption + self.trim_padding = min(self.trim_padding, self.total_corruption) self.psd_duration = (psd_samples - 1) // 2 * psd_segment_length @@ -1746,50 +2219,74 @@ def __init__(self, frame_src, channel_name, start_time, @property def start_time(self): - """ Return the start time of the current valid segment of data """ + """Return the start time of the current valid segment of data""" return self.end_time - self.blocksize @property def end_time(self): - """ Return the end time of the current valid segment of data """ - return float(self.strain.start_time + (len(self.strain) - self.total_corruption) / self.sample_rate) + """Return the end time of the current valid segment of data""" + return float( + self.strain.start_time + + (len(self.strain) - self.total_corruption) / self.sample_rate + ) def add_hard_count(self): - """ Reset the countdown timer, so that we don't analyze data long enough + """ + Reset the countdown timer, so that we don't analyze data long enough to generate a new PSD. """ - self.wait_duration = int(numpy.ceil(self.total_corruption / self.sample_rate + self.psd_duration)) + self.wait_duration = int( + numpy.ceil(self.total_corruption / self.sample_rate + self.psd_duration) + ) self.invalidate_psd() def invalidate_psd(self): - """ Make the current PSD invalid. A new one will be generated when - it is next required """ + """ + Make the current PSD invalid. A new one will be generated when + it is next required + """ self.psd = None self.psds = {} def recalculate_psd(self): - """ Recalculate the psd - """ - + """Recalculate the psd""" seg_len = int(self.sample_rate * self.psd_segment_length) e = len(self.strain) s = e - (self.psd_samples + 1) * seg_len // 2 - psd = pycbc.psd.welch(self.strain[s:e], seg_len=seg_len, seg_stride=seg_len//2) + psd = pycbc.psd.welch( + self.strain[s:e], seg_len=seg_len, seg_stride=seg_len // 2 + ) - psd.dist = spa_distance(psd, 1.4, 1.4, self.low_frequency_cutoff) * pycbc.DYN_RANGE_FAC + psd.dist = ( + spa_distance(psd, 1.4, 1.4, self.low_frequency_cutoff) * pycbc.DYN_RANGE_FAC + ) # If the new psd is similar to the old one, don't replace it if self.psd and self.psd_recalculate_difference: - if abs(self.psd.dist - psd.dist) / self.psd.dist < self.psd_recalculate_difference: - logger.info("Skipping recalculation of %s PSD, %s-%s", - self.detector, self.psd.dist, psd.dist) + if ( + abs(self.psd.dist - psd.dist) / self.psd.dist + < self.psd_recalculate_difference + ): + logger.info( + "Skipping recalculation of %s PSD, %s-%s", + self.detector, + self.psd.dist, + psd.dist, + ) return True # If the new psd is *really* different than the old one, return an error if self.psd and self.psd_abort_difference: - if abs(self.psd.dist - psd.dist) / self.psd.dist > self.psd_abort_difference: - logger.info("%s PSD is CRAZY, aborting!!!!, %s-%s", - self.detector, self.psd.dist, psd.dist) + if ( + abs(self.psd.dist - psd.dist) / self.psd.dist + > self.psd_abort_difference + ): + logger.info( + "%s PSD is CRAZY, aborting!!!!, %s-%s", + self.detector, + self.psd.dist, + psd.dist, + ) self.psd = psd self.psds = {} return False @@ -1801,7 +2298,8 @@ def recalculate_psd(self): return True def check_psd_dist(self, min_dist, max_dist): - """Check that the horizon distance of a detector is within a required + """ + Check that the horizon distance of a detector is within a required range. If so, return True, otherwise log a warning and return False. """ if self.psd is None: @@ -1817,12 +2315,13 @@ def check_psd_dist(self, min_dist, max_dist): self.detector, self.psd.dist, min_dist, - max_dist + max_dist, ) return good def overwhitened_data(self, delta_f): - """ Return overwhitened data + """ + Return overwhitened data Parameters ---------- @@ -1833,6 +2332,7 @@ def overwhitened_data(self, delta_f): ------- htilde: FrequencySeries Overwhited strain data + """ # we haven't already computed htilde for this delta_f if delta_f not in self.segments: @@ -1841,23 +2341,27 @@ def overwhitened_data(self, delta_f): s = int(e - buffer_length * self.sample_rate - self.reduced_pad * 2) # FFT the contents of self.strain[s:e] into fseries - fseries = execute_cached_fft(self.strain[s:e], - copy_output=False, - uid=STRAINBUFFER_UNIQUE_ID_1) - fseries._epoch = self.strain._epoch + s*self.strain.delta_t + fseries = execute_cached_fft( + self.strain[s:e], copy_output=False, uid=STRAINBUFFER_UNIQUE_ID_1 + ) + fseries._epoch = self.strain._epoch + s * self.strain.delta_t # we haven't calculated a resample psd for this delta_f if delta_f not in self.psds: psdt = pycbc.psd.interpolate(self.psd, fseries.delta_f) - psdt = pycbc.psd.inverse_spectrum_truncation(psdt, - int(self.sample_rate * self.psd_inverse_length), - low_frequency_cutoff=self.low_frequency_cutoff) + psdt = pycbc.psd.inverse_spectrum_truncation( + psdt, + int(self.sample_rate * self.psd_inverse_length), + low_frequency_cutoff=self.low_frequency_cutoff, + ) psdt._delta_f = fseries.delta_f psd = pycbc.psd.interpolate(self.psd, delta_f) - psd = pycbc.psd.inverse_spectrum_truncation(psd, - int(self.sample_rate * self.psd_inverse_length), - low_frequency_cutoff=self.low_frequency_cutoff) + psd = pycbc.psd.inverse_spectrum_truncation( + psd, + int(self.sample_rate * self.psd_inverse_length), + low_frequency_cutoff=self.low_frequency_cutoff, + ) psd.psdt = psdt self.psds[delta_f] = psd @@ -1866,26 +2370,30 @@ def overwhitened_data(self, delta_f): fseries /= psd.psdt # trim ends of strain - if self.reduced_pad != 0: + if self.reduced_pad != 0: # IFFT the contents of fseries into overwhite - overwhite = execute_cached_ifft(fseries, - copy_output=False, - uid=STRAINBUFFER_UNIQUE_ID_2) + overwhite = execute_cached_ifft( + fseries, copy_output=False, uid=STRAINBUFFER_UNIQUE_ID_2 + ) - overwhite2 = overwhite[self.reduced_pad:len(overwhite)-self.reduced_pad] + overwhite2 = overwhite[ + self.reduced_pad : len(overwhite) - self.reduced_pad + ] taper_window = self.trim_padding / 2.0 / overwhite.sample_rate - gate_params = [(overwhite2.start_time, 0., taper_window), - (overwhite2.end_time, 0., taper_window)] + gate_params = [ + (overwhite2.start_time, 0.0, taper_window), + (overwhite2.end_time, 0.0, taper_window), + ] gate_data(overwhite2, gate_params) # FFT the contents of overwhite2 into fseries_trimmed fseries_trimmed = execute_cached_fft( - overwhite2, - copy_output=True, - uid=STRAINBUFFER_UNIQUE_ID_3 + overwhite2, copy_output=True, uid=STRAINBUFFER_UNIQUE_ID_3 ) - fseries_trimmed.start_time = fseries.start_time + self.reduced_pad * self.strain.delta_t + fseries_trimmed.start_time = ( + fseries.start_time + self.reduced_pad * self.strain.delta_t + ) else: fseries_trimmed = fseries @@ -1896,37 +2404,43 @@ def overwhitened_data(self, delta_f): return stilde def near_hwinj(self): - """Check that the current set of triggers could be influenced by + """ + Check that the current set of triggers could be influenced by a hardware injection. """ if not self.state: return False - if not self.state.is_extent_valid(self.start_time, self.blocksize, pycbc.frame.NO_HWINJ): + if not self.state.is_extent_valid( + self.start_time, self.blocksize, pycbc.frame.NO_HWINJ + ): return True return False def null_advance_strain(self, blocksize): - """ Advance and insert zeros + """ + Advance and insert zeros Parameters ---------- blocksize: int The number of seconds to attempt to read from the channel + """ sample_step = int(blocksize * self.sample_rate) csize = sample_step + self.corruption * 2 self.strain.roll(-sample_step) # We should roll this off at some point too... - self.strain[len(self.strain) - csize + self.corruption:] = 0 + self.strain[len(self.strain) - csize + self.corruption :] = 0 self.strain.start_time += blocksize # The next time we need strain will need to be tapered self.taper_immediate_strain = True def advance(self, blocksize, timeout=10): - """Advanced buffer blocksize seconds. + """ + Advanced buffer blocksize seconds. Add blocksize seconds more to the buffer, push blocksize seconds from the beginning. @@ -1940,8 +2454,9 @@ def advance(self, blocksize, timeout=10): ------- status: boolean Returns True if this block is analyzable. + """ - ts = super(StrainBuffer, self).attempt_advance(blocksize, timeout=timeout) + ts = super().attempt_advance(blocksize, timeout=timeout) self.blocksize = blocksize self.gate_params = [] @@ -1970,8 +2485,7 @@ def advance(self, blocksize, timeout=10): self.dq.null_advance(blocksize) if self.idq: self.idq.null_advance(blocksize) - logger.info("%s time has invalid data, resetting buffer", - self.detector) + logger.info("%s time has invalid data, resetting buffer", self.detector) return False # Also advance the dq vector and idq timeseries in lockstep @@ -1991,27 +2505,29 @@ def advance(self, blocksize, timeout=10): start = len(self.raw_buffer) - csize * self.factor strain = self.raw_buffer[start:] - strain = pycbc.filter.highpass_fir(strain, self.highpass_frequency, - self.highpass_samples, - beta=self.beta) + strain = pycbc.filter.highpass_fir( + strain, self.highpass_frequency, self.highpass_samples, beta=self.beta + ) strain = (strain * self.dyn_range_fac).astype(numpy.float32) - strain = pycbc.filter.resample_to_delta_t(strain, - 1.0/self.sample_rate, method='ldas') + strain = pycbc.filter.resample_to_delta_t( + strain, 1.0 / self.sample_rate, method="ldas" + ) # remove corruption at beginning - strain = strain[self.corruption:] + strain = strain[self.corruption :] # taper beginning if needed if self.taper_immediate_strain: logger.info("Tapering start of %s strain block", self.detector) strain = gate_data( - strain, [(strain.start_time, 0., self.autogating_taper)]) + strain, [(strain.start_time, 0.0, self.autogating_taper)] + ) self.taper_immediate_strain = False # Stitch into continuous stream self.strain.roll(-sample_step) - self.strain[len(self.strain) - csize + self.corruption:] = strain[:] + self.strain[len(self.strain) - csize + self.corruption :] = strain[:] self.strain.start_time += blocksize # apply gating if needed @@ -2019,60 +2535,72 @@ def advance(self, blocksize, timeout=10): autogating_duration_length = self.autogating_duration * self.sample_rate autogating_start_sample = int(len(self.strain) - autogating_duration_length) glitch_times = detect_loud_glitches( - self.strain[autogating_start_sample:-self.corruption], - psd_duration=self.autogating_psd_segment_length, psd_stride=self.autogating_psd_stride, - threshold=self.autogating_threshold, - cluster_window=self.autogating_cluster, - low_freq_cutoff=self.highpass_frequency, - corrupt_time=self.autogating_pad) + self.strain[autogating_start_sample : -self.corruption], + psd_duration=self.autogating_psd_segment_length, + psd_stride=self.autogating_psd_stride, + threshold=self.autogating_threshold, + cluster_window=self.autogating_cluster, + low_freq_cutoff=self.highpass_frequency, + corrupt_time=self.autogating_pad, + ) if len(glitch_times) > 0: - logger.info('Autogating %s at %s', self.detector, - ', '.join(['%.3f' % gt for gt in glitch_times])) - self.gate_params = \ - [(gt, self.autogating_width, self.autogating_taper) - for gt in glitch_times] + logger.info( + "Autogating %s at %s", + self.detector, + ", ".join(["%.3f" % gt for gt in glitch_times]), + ) + self.gate_params = [ + (gt, self.autogating_width, self.autogating_taper) + for gt in glitch_times + ] self.strain = gate_data(self.strain, self.gate_params) - if self.psd is None and self.wait_duration <=0: + if self.psd is None and self.wait_duration <= 0: self.recalculate_psd() return self.wait_duration <= 0 @classmethod def from_cli(cls, ifo, args): - """Initialize a StrainBuffer object (data reader) for a particular + """ + Initialize a StrainBuffer object (data reader) for a particular detector. """ state_channel = analyze_flags = None - if args.state_channel and ifo in args.state_channel \ - and args.analyze_flags and ifo in args.analyze_flags: - state_channel = ':'.join([ifo, args.state_channel[ifo]]) - analyze_flags = args.analyze_flags[ifo].split(',') + if ( + args.state_channel + and ifo in args.state_channel + and args.analyze_flags + and ifo in args.analyze_flags + ): + state_channel = ":".join([ifo, args.state_channel[ifo]]) + analyze_flags = args.analyze_flags[ifo].split(",") dq_channel = dq_flags = None - if args.data_quality_channel and ifo in args.data_quality_channel \ - and args.data_quality_flags and ifo in args.data_quality_flags: - dq_channel = ':'.join([ifo, args.data_quality_channel[ifo]]) - dq_flags = args.data_quality_flags[ifo].split(',') + if ( + args.data_quality_channel + and ifo in args.data_quality_channel + and args.data_quality_flags + and ifo in args.data_quality_flags + ): + dq_channel = ":".join([ifo, args.data_quality_channel[ifo]]) + dq_flags = args.data_quality_flags[ifo].split(",") idq_channel = None if args.idq_channel and ifo in args.idq_channel: - idq_channel = ':'.join([ifo, args.idq_channel[ifo]]) + idq_channel = ":".join([ifo, args.idq_channel[ifo]]) idq_state_channel = None if args.idq_state_channel and ifo in args.idq_state_channel: - idq_state_channel = ':'.join([ifo, args.idq_state_channel[ifo]]) + idq_state_channel = ":".join([ifo, args.idq_state_channel[ifo]]) if args.frame_type: frame_src = pycbc.frame.frame_paths( - args.frame_type[ifo], - args.start_time, - args.end_time, - site=ifo[0] + args.frame_type[ifo], args.start_time, args.end_time, site=ifo[0] ) else: frame_src = [args.frame_src[ifo]] - strain_channel = ':'.join([ifo, args.channel_name[ifo]]) + strain_channel = ":".join([ifo, args.channel_name[ifo]]) return cls( frame_src, @@ -2107,5 +2635,5 @@ def from_cli(cls, ifo, args): increment_update_cache=args.increment_update_cache[ifo], analyze_flags=analyze_flags, data_quality_flags=dq_flags, - dq_padding=args.data_quality_padding + dq_padding=args.data_quality_padding, ) diff --git a/pycbc/time.py b/pycbc/time.py index c60514d2970..3b6958fcadf 100644 --- a/pycbc/time.py +++ b/pycbc/time.py @@ -2,9 +2,10 @@ Module to contain time conversions used in pycbc """ -from astropy.time import Time from datetime import timezone +from astropy.time import Time + def _ensure_utc_datetime(date): """ @@ -30,8 +31,9 @@ def gps_to_utc_datetime(gps): Returns ------- datetime + """ - dt = Time(gps, format='gps', scale='utc').to_datetime() + dt = Time(gps, format="gps", scale="utc").to_datetime() return _ensure_utc_datetime(dt) @@ -47,10 +49,11 @@ def datetime_to_str(date, format="%Y-%m-%d %H:%M:%S"): The format to use for the string. Default is yyyy-mm-dd HH:MM:SS. Supply as a format according to datetime documentation https://docs.python.org/3/library/datetime.html#strftime-strptime-behavior - + Returns ------- str + """ return date.strftime(format) @@ -68,6 +71,7 @@ def gps_to_utc_str(gps, format="%Y-%m-%d %H:%M:%S"): Returns ------- str + """ return datetime_to_str(gps_to_utc_datetime(gps), format=format) @@ -84,6 +88,7 @@ def strip_time_from_date(date): Returns ------- datetime + """ stripped = date.replace(hour=0, minute=0, second=0, microsecond=0) return _ensure_utc_datetime(stripped) @@ -101,10 +106,13 @@ def strip_time_from_gps(gps, format="%Y-%m-%d"): Returns ------- float, str + """ gps_datetime = gps_to_utc_datetime(gps) midnight_datetime = strip_time_from_date(gps_datetime) - return utc_datetime_to_gps(midnight_datetime), datetime_to_str(midnight_datetime, format=format) + return utc_datetime_to_gps(midnight_datetime), datetime_to_str( + midnight_datetime, format=format + ) def utc_datetime_to_gps(date): @@ -119,19 +127,21 @@ def utc_datetime_to_gps(date): Returns ------- float + """ date_utc = _ensure_utc_datetime(date) - return float(Time(date_utc, format='datetime', scale='utc').gps) + return float(Time(date_utc, format="datetime", scale="utc").gps) def gps_now(): - """Return the current GPS time as a float using Astropy. + """ + Return the current GPS time as a float using Astropy. Returns ------- float - """ + """ return float(Time.now().gps) @@ -148,17 +158,22 @@ def gmst_accurate(gps_time): Returns ------- float + """ - gmst = Time(gps_time, format='gps', scale='utc', - location=(0, 0)).sidereal_time('mean').rad + gmst = ( + Time(gps_time, format="gps", scale="utc", location=(0, 0)) + .sidereal_time("mean") + .rad + ) return gmst + __all__ = [ - 'gps_to_utc_datetime', - 'gps_to_utc_str', - 'strip_time_from_date', - 'strip_time_from_gps', - 'utc_datetime_to_gps', - 'gps_now', - 'gmst_accurate', -] \ No newline at end of file + "gmst_accurate", + "gps_now", + "gps_to_utc_datetime", + "gps_to_utc_str", + "strip_time_from_date", + "strip_time_from_gps", + "utc_datetime_to_gps", +] diff --git a/pycbc/tmpltbank/__init__.py b/pycbc/tmpltbank/__init__.py index 952cdfccb30..664f04a1cf3 100644 --- a/pycbc/tmpltbank/__init__.py +++ b/pycbc/tmpltbank/__init__.py @@ -1,8 +1,8 @@ +from pycbc.tmpltbank.bank_conversions import * +from pycbc.tmpltbank.brute_force_methods import * from pycbc.tmpltbank.calc_moments import * -from pycbc.tmpltbank.lambda_mapping import * from pycbc.tmpltbank.coord_utils import * +from pycbc.tmpltbank.lambda_mapping import * from pycbc.tmpltbank.lattice_utils import * -from pycbc.tmpltbank.brute_force_methods import * from pycbc.tmpltbank.option_utils import * from pycbc.tmpltbank.partitioned_bank import * -from pycbc.tmpltbank.bank_conversions import * diff --git a/pycbc/tmpltbank/bank_conversions.py b/pycbc/tmpltbank/bank_conversions.py index 51f66e7ebb3..e5cf8f42e54 100644 --- a/pycbc/tmpltbank/bank_conversions.py +++ b/pycbc/tmpltbank/bank_conversions.py @@ -27,42 +27,59 @@ """ import logging + import numpy as np from pycbc import conversions as conv from pycbc import pnutils -logger = logging.getLogger('pycbc.tmpltbank.bank_conversions') +logger = logging.getLogger("pycbc.tmpltbank.bank_conversions") # Convert from parameter name to helper function # some multiple names are used for the same function -conversion_options = ['mass1', 'mass2', 'spin1z', 'spin2z', 'duration', - 'template_duration', 'mtotal', 'total_mass', - 'q', 'invq', 'eta', 'chirp_mass', 'mchirp', - 'chieff', 'chi_eff', 'effective_spin', 'chi_a', - 'premerger_duration'] +conversion_options = [ + "mass1", + "mass2", + "spin1z", + "spin2z", + "duration", + "template_duration", + "mtotal", + "total_mass", + "q", + "invq", + "eta", + "chirp_mass", + "mchirp", + "chieff", + "chi_eff", + "effective_spin", + "chi_a", + "premerger_duration", +] mass_conversions = { - 'mtotal': conv.mtotal_from_mass1_mass2, - 'total_mass': conv.mtotal_from_mass1_mass2, - 'q': conv.q_from_mass1_mass2, - 'invq': conv.invq_from_mass1_mass2, - 'eta': conv.eta_from_mass1_mass2, - 'mchirp': conv.mchirp_from_mass1_mass2, - 'chirp_mass': conv.mchirp_from_mass1_mass2, + "mtotal": conv.mtotal_from_mass1_mass2, + "total_mass": conv.mtotal_from_mass1_mass2, + "q": conv.q_from_mass1_mass2, + "invq": conv.invq_from_mass1_mass2, + "eta": conv.eta_from_mass1_mass2, + "mchirp": conv.mchirp_from_mass1_mass2, + "chirp_mass": conv.mchirp_from_mass1_mass2, } spin_conversions = { - 'chieff': conv.chi_eff, - 'chi_eff': conv.chi_eff, - 'effective_spin': conv.chi_eff, - 'chi_a': conv.chi_a + "chieff": conv.chi_eff, + "chi_eff": conv.chi_eff, + "effective_spin": conv.chi_eff, + "chi_a": conv.chi_a, } def get_bank_property(parameter, bank, template_ids): - """ Get a specific value from a hdf file object in standard PyCBC + """ + Get a specific value from a hdf file object in standard PyCBC template bank format Parameters @@ -88,48 +105,46 @@ def get_bank_property(parameter, bank, template_ids): if parameter in bank: values = bank[parameter][:][template_ids] # Duration may be in the bank, but if not, we need to calculate - elif parameter.endswith('duration'): + elif parameter.endswith("duration"): fullband_req = False prem_required = False - if parameter != "premerger_duration" and 'template_duration' in bank: + if parameter != "premerger_duration" and "template_duration" in bank: # This statement should be the reached only if 'duration' # is given, but 'template_duration' is in the bank - fullband_dur = bank['template_duration'][:][template_ids] - elif parameter in ['template_duration', 'duration']: + fullband_dur = bank["template_duration"][:][template_ids] + elif parameter in ["template_duration", "duration"]: # Only calculate fullband/premerger durations if we need to fullband_req = True - if 'f_final' in bank: + if "f_final" in bank: prem_required = True elif parameter == "premerger_duration": prem_required = True # Set up the arguments for get_imr_duration - imr_args = ['mass1', 'mass2', 'spin1z', 'spin2z'] - if 'approximant' in bank: - kwargs = {'approximant': bank['approximant'][:][template_ids]} + imr_args = ["mass1", "mass2", "spin1z", "spin2z"] + if "approximant" in bank: + kwargs = {"approximant": bank["approximant"][:][template_ids]} else: kwargs = {} if fullband_req: # Unpack the appropriate arguments fullband_dur = pnutils.get_imr_duration( - *[bank[k][:][template_ids] - for k in imr_args + ['f_lower']], - **kwargs) + *[bank[k][:][template_ids] for k in imr_args + ["f_lower"]], **kwargs + ) - if prem_required and 'f_final' in bank: + if prem_required and "f_final" in bank: # If f_final is in the bank, then we need to calculate # the premerger time of the end of the template prem_dur = pnutils.get_imr_duration( - *[bank[k][:][template_ids] - for k in imr_args + ['f_final']], - **kwargs) + *[bank[k][:][template_ids] for k in imr_args + ["f_final"]], **kwargs + ) elif prem_required: # Pre-merger for bank without f_final is zero prem_dur = np.zeros_like(template_ids) # Now we decide what to return: - if parameter in ['template_duration', 'duration']: + if parameter in ["template_duration", "duration"]: values = fullband_dur if prem_required: values -= prem_dur @@ -137,19 +152,26 @@ def get_bank_property(parameter, bank, template_ids): values = prem_dur # Basic conversions - elif parameter in mass_conversions.keys(): - values = mass_conversions[parameter](bank['mass1'][:][template_ids], - bank['mass2'][:][template_ids]) - - elif parameter in spin_conversions.keys(): - values = spin_conversions[parameter](bank['mass1'][:][template_ids], - bank['mass2'][:][template_ids], - bank['spin1z'][:][template_ids], - bank['spin2z'][:][template_ids]) + elif parameter in mass_conversions: + values = mass_conversions[parameter]( + bank["mass1"][:][template_ids], bank["mass2"][:][template_ids] + ) + + elif parameter in spin_conversions: + values = spin_conversions[parameter]( + bank["mass1"][:][template_ids], + bank["mass2"][:][template_ids], + bank["spin1z"][:][template_ids], + bank["spin2z"][:][template_ids], + ) else: # parameter not in the current conversion parameter list - raise NotImplementedError("Bank conversion function " + parameter - + " not recognised: choose from '" + - "', '".join(conversion_options) + "'.") + raise NotImplementedError( + "Bank conversion function " + + parameter + + " not recognised: choose from '" + + "', '".join(conversion_options) + + "'." + ) return values diff --git a/pycbc/tmpltbank/bank_output_utils.py b/pycbc/tmpltbank/bank_output_utils.py index fee23894b0a..793038a389f 100644 --- a/pycbc/tmpltbank/bank_output_utils.py +++ b/pycbc/tmpltbank/bank_output_utils.py @@ -1,27 +1,30 @@ import logging -import numpy -from igwn_ligolw import ligolw, lsctables, utils as ligolw_utils +import numpy +from igwn_ligolw import ligolw, lsctables +from igwn_ligolw import utils as ligolw_utils from pycbc import pnutils -from pycbc.tmpltbank.lambda_mapping import ethinca_order_from_string +from pycbc.constants import GAMMA, MTSUN_SI, PI, TWOPI +from pycbc.io.hdf import HFile from pycbc.io.ligolw import ( - return_empty_sngl, return_search_summary, create_process_table + create_process_table, + return_empty_sngl, + return_search_summary, ) -from pycbc.io.hdf import HFile -from pycbc.constants import PI, MTSUN_SI, TWOPI, GAMMA - +from pycbc.tmpltbank.lambda_mapping import ethinca_order_from_string from pycbc.waveform import get_waveform_filter_length_in_time as gwflit -logger = logging.getLogger('pycbc.tmpltbank.bank_output_utils') +logger = logging.getLogger("pycbc.tmpltbank.bank_output_utils") + def convert_to_sngl_inspiral_table(params, proc_id): - ''' + """ Convert a list of m1,m2,spin1z,spin2z values into a basic sngl_inspiral table with mass and spin parameters populated and event IDs assigned Parameters - ----------- + ---------- params : iterable Each entry in the params iterable should be a sequence of [mass1, mass2, spin1z, spin2z] in that order @@ -29,12 +32,13 @@ def convert_to_sngl_inspiral_table(params, proc_id): Process ID to add to each row of the sngl_inspiral table Returns - ---------- + ------- SnglInspiralTable Bank of templates in SnglInspiralTable format - ''' + + """ sngl_inspiral_table = lsctables.SnglInspiralTable.new() - col_names = ['mass1','mass2','spin1z','spin2z'] + col_names = ["mass1", "mass2", "spin1z", "spin2z"] for values in params: tmplt = return_empty_sngl() @@ -43,17 +47,19 @@ def convert_to_sngl_inspiral_table(params, proc_id): for colname, value in zip(col_names, values): setattr(tmplt, colname, value) tmplt.mtotal, tmplt.eta = pnutils.mass1_mass2_to_mtotal_eta( - tmplt.mass1, tmplt.mass2) - tmplt.mchirp, _ = pnutils.mass1_mass2_to_mchirp_eta( - tmplt.mass1, tmplt.mass2) - tmplt.template_duration = 0 # FIXME + tmplt.mass1, tmplt.mass2 + ) + tmplt.mchirp, _ = pnutils.mass1_mass2_to_mchirp_eta(tmplt.mass1, tmplt.mass2) + tmplt.template_duration = 0 # FIXME tmplt.event_id = sngl_inspiral_table.get_next_id() sngl_inspiral_table.append(tmplt) return sngl_inspiral_table -def calculate_ethinca_metric_comps(metricParams, ethincaParams, mass1, mass2, - spin1z=0., spin2z=0., full_ethinca=True): + +def calculate_ethinca_metric_comps( + metricParams, ethincaParams, mass1, mass2, spin1z=0.0, spin2z=0.0, full_ethinca=True +): r""" Calculate the Gamma components needed to use the ethinca metric. At present this outputs the standard TaylorF2 metric over the end time @@ -64,7 +70,7 @@ def calculate_ethinca_metric_comps(metricParams, ethincaParams, mass1, mass2, bank layout options fLow and f0 (which must be the same as each other). Parameters - ----------- + ---------- metricParams : metricParameters instance Structure holding all the options for construction of the metric and the eigenvalues, eigenvectors and covariance matrix @@ -83,8 +89,9 @@ def calculate_ethinca_metric_comps(metricParams, ethincaParams, mass1, mass2, If True calculate the ethinca components in all 3 directions (mass1, mass2 and time). If False calculate only the time component (which is stored in Gamma0). + Returns - -------- + ------- fMax_theor : float Value of the upper frequency cutoff given by the template parameters and the cutoff formula requested. @@ -93,128 +100,153 @@ def calculate_ethinca_metric_comps(metricParams, ethincaParams, mass1, mass2, Array holding 6 independent metric components in (end_time, tau_0, tau_3) coordinates to be stored in the Gamma0-5 slots of a SnglInspiral object. + """ - if (float(spin1z) != 0. or float(spin2z) != 0.) and full_ethinca: - raise NotImplementedError("Ethinca cannot at present be calculated " - "for nonzero component spins!") + if (float(spin1z) != 0.0 or float(spin2z) != 0.0) and full_ethinca: + raise NotImplementedError( + "Ethinca cannot at present be calculated for nonzero component spins!" + ) f0 = metricParams.f0 if f0 != metricParams.fLow: - raise ValueError("If calculating ethinca the bank f0 value must be " - "equal to f-low!") - if ethincaParams.fLow is not None and ( - ethincaParams.fLow != metricParams.fLow): - raise NotImplementedError("An ethinca metric f-low different from the" - " bank metric f-low is not supported!") + raise ValueError( + "If calculating ethinca the bank f0 value must be equal to f-low!" + ) + if ethincaParams.fLow is not None and (ethincaParams.fLow != metricParams.fLow): + raise NotImplementedError( + "An ethinca metric f-low different from the" + " bank metric f-low is not supported!" + ) twicePNOrder = ethinca_order_from_string(ethincaParams.pnOrder) piFl = PI * f0 totalMass, eta = pnutils.mass1_mass2_to_mtotal_eta(mass1, mass2) totalMass = totalMass * MTSUN_SI - v0cube = totalMass*piFl - v0 = v0cube**(1./3.) + v0cube = totalMass * piFl + v0 = v0cube ** (1.0 / 3.0) # Get theoretical cutoff frequency and work out the closest # frequency for which moments were calculated fMax_theor = pnutils.frequency_cutoff_from_name( - ethincaParams.cutoff, mass1, mass2, spin1z, spin2z) - fMaxes = list(metricParams.moments['J4'].keys()) - fMaxIdx = abs(numpy.array(fMaxes,dtype=float) - fMax_theor).argmin() + ethincaParams.cutoff, mass1, mass2, spin1z, spin2z + ) + fMaxes = list(metricParams.moments["J4"].keys()) + fMaxIdx = abs(numpy.array(fMaxes, dtype=float) - fMax_theor).argmin() fMax = fMaxes[fMaxIdx] # Set the appropriate moments - Js = numpy.zeros([18,3],dtype=float) + Js = numpy.zeros([18, 3], dtype=float) for i in range(18): - Js[i,0] = metricParams.moments['J%d'%(i)][fMax] - Js[i,1] = metricParams.moments['log%d'%(i)][fMax] - Js[i,2] = metricParams.moments['loglog%d'%(i)][fMax] + Js[i, 0] = metricParams.moments["J%d" % (i)][fMax] + Js[i, 1] = metricParams.moments["log%d" % (i)][fMax] + Js[i, 2] = metricParams.moments["loglog%d" % (i)][fMax] # Compute the time-dependent metric term. two_pi_flower_sq = TWOPI * f0 * TWOPI * f0 - gammaVals = numpy.zeros([6],dtype=float) - gammaVals[0] = 0.5 * two_pi_flower_sq * \ - ( Js[(1,0)] - (Js[(4,0)]*Js[(4,0)]) ) + gammaVals = numpy.zeros([6], dtype=float) + gammaVals[0] = 0.5 * two_pi_flower_sq * (Js[(1, 0)] - (Js[(4, 0)] * Js[(4, 0)])) # If mass terms not required stop here if not full_ethinca: return fMax_theor, gammaVals # 3pN is a mess, so split it into pieces - a0 = 11583231236531/200286535680 - 5*PI*PI - 107*GAMMA/14 - a1 = (-15737765635/130056192 + 2255*PI*PI/512)*eta - a2 = (76055/73728)*eta*eta - a3 = (-127825/55296)*eta*eta*eta - alog = numpy.log(4*v0) # Log terms are tricky - be careful + a0 = 11583231236531 / 200286535680 - 5 * PI * PI - 107 * GAMMA / 14 + a1 = (-15737765635 / 130056192 + 2255 * PI * PI / 512) * eta + a2 = (76055 / 73728) * eta * eta + a3 = (-127825 / 55296) * eta * eta * eta + alog = numpy.log(4 * v0) # Log terms are tricky - be careful # Get the Psi coefficients - Psi = [{},{}] #Psi = numpy.zeros([2,8,2],dtype=float) - Psi[0][0,0] = 3/5 - Psi[0][2,0] = (743/756 + 11*eta/3)*v0*v0 - Psi[0][3,0] = 0. - Psi[0][4,0] = (-3058673/508032 + 5429*eta/504 + 617*eta*eta/24)\ - *v0cube*v0 - Psi[0][5,1] = (-7729*PI/126)*v0cube*v0*v0/3 - Psi[0][6,0] = (128/15)*(-3*a0 - a1 + a2 + 3*a3 + 107*(1+3*alog)/14)\ - *v0cube*v0cube - Psi[0][6,1] = (6848/35)*v0cube*v0cube/3 - Psi[0][7,0] = (-15419335/63504 - 75703*eta/756)*PI*v0cube*v0cube*v0 - - Psi[1][0,0] = 0. - Psi[1][2,0] = (3715/12096 - 55*eta/96)/PI/v0; - Psi[1][3,0] = -3/2 - Psi[1][4,0] = (15293365/4064256 - 27145*eta/16128 - 3085*eta*eta/384)\ - *v0/PI - Psi[1][5,1] = (193225/8064)*v0*v0/3 - Psi[1][6,0] = (4/PI)*(2*a0 + a1/3 - 4*a2/3 - 3*a3 -107*(1+6*alog)/42)\ - *v0cube - Psi[1][6,1] = (-428/PI/7)*v0cube/3 - Psi[1][7,0] = (77096675/1161216 + 378515*eta/24192 + 74045*eta*eta/8064)\ - *v0cube*v0 + Psi = [{}, {}] # Psi = numpy.zeros([2,8,2],dtype=float) + Psi[0][0, 0] = 3 / 5 + Psi[0][2, 0] = (743 / 756 + 11 * eta / 3) * v0 * v0 + Psi[0][3, 0] = 0.0 + Psi[0][4, 0] = ( + (-3058673 / 508032 + 5429 * eta / 504 + 617 * eta * eta / 24) * v0cube * v0 + ) + Psi[0][5, 1] = (-7729 * PI / 126) * v0cube * v0 * v0 / 3 + Psi[0][6, 0] = ( + (128 / 15) + * (-3 * a0 - a1 + a2 + 3 * a3 + 107 * (1 + 3 * alog) / 14) + * v0cube + * v0cube + ) + Psi[0][6, 1] = (6848 / 35) * v0cube * v0cube / 3 + Psi[0][7, 0] = (-15419335 / 63504 - 75703 * eta / 756) * PI * v0cube * v0cube * v0 + + Psi[1][0, 0] = 0.0 + Psi[1][2, 0] = (3715 / 12096 - 55 * eta / 96) / PI / v0 + Psi[1][3, 0] = -3 / 2 + Psi[1][4, 0] = ( + (15293365 / 4064256 - 27145 * eta / 16128 - 3085 * eta * eta / 384) * v0 / PI + ) + Psi[1][5, 1] = (193225 / 8064) * v0 * v0 / 3 + Psi[1][6, 0] = ( + (4 / PI) + * (2 * a0 + a1 / 3 - 4 * a2 / 3 - 3 * a3 - 107 * (1 + 6 * alog) / 42) + * v0cube + ) + Psi[1][6, 1] = (-428 / PI / 7) * v0cube / 3 + Psi[1][7, 0] = ( + (77096675 / 1161216 + 378515 * eta / 24192 + 74045 * eta * eta / 8064) + * v0cube + * v0 + ) # Set the appropriate moments - Js = numpy.zeros([18,3],dtype=float) + Js = numpy.zeros([18, 3], dtype=float) for i in range(18): - Js[i,0] = metricParams.moments['J%d'%(i)][fMax] - Js[i,1] = metricParams.moments['log%d'%(i)][fMax] - Js[i,2] = metricParams.moments['loglog%d'%(i)][fMax] + Js[i, 0] = metricParams.moments["J%d" % (i)][fMax] + Js[i, 1] = metricParams.moments["log%d" % (i)][fMax] + Js[i, 2] = metricParams.moments["loglog%d" % (i)][fMax] # Calculate the g matrix - PNterms = [(0,0),(2,0),(3,0),(4,0),(5,1),(6,0),(6,1),(7,0)] + PNterms = [(0, 0), (2, 0), (3, 0), (4, 0), (5, 1), (6, 0), (6, 1), (7, 0)] PNterms = [term for term in PNterms if term[0] <= twicePNOrder] # Now can compute the mass-dependent gamma values for m in [0, 1]: for k in PNterms: - gammaVals[1+m] += 0.5 * two_pi_flower_sq * Psi[m][k] * \ - ( Js[(9-k[0],k[1])] - - Js[(12-k[0],k[1])] * Js[(4,0)] ) - - g = numpy.zeros([2,2],dtype=float) - for (m,n) in [(0,0),(0,1),(1,1)]: + gammaVals[1 + m] += ( + 0.5 + * two_pi_flower_sq + * Psi[m][k] + * (Js[(9 - k[0], k[1])] - Js[(12 - k[0], k[1])] * Js[(4, 0)]) + ) + + g = numpy.zeros([2, 2], dtype=float) + for m, n in [(0, 0), (0, 1), (1, 1)]: for k in PNterms: for l in PNterms: - g[m,n] += Psi[m][k] * Psi[n][l] * \ - ( Js[(17-k[0]-l[0], k[1]+l[1])] - - Js[(12-k[0],k[1])] * Js[(12-l[0],l[1])] ) - g[m,n] = 0.5 * two_pi_flower_sq * g[m,n] - g[n,m] = g[m,n] - - gammaVals[3] = g[0,0] - gammaVals[4] = g[0,1] - gammaVals[5] = g[1,1] + g[m, n] += ( + Psi[m][k] + * Psi[n][l] + * ( + Js[(17 - k[0] - l[0], k[1] + l[1])] + - Js[(12 - k[0], k[1])] * Js[(12 - l[0], l[1])] + ) + ) + g[m, n] = 0.5 * two_pi_flower_sq * g[m, n] + g[n, m] = g[m, n] + + gammaVals[3] = g[0, 0] + gammaVals[4] = g[0, 1] + gammaVals[5] = g[1, 1] return fMax_theor, gammaVals -def output_sngl_inspiral_table(outputFile, tempBank, programName="", - optDict = None, outdoc=None, - **kwargs): # pylint:disable=unused-argument + +def output_sngl_inspiral_table( + outputFile, tempBank, programName="", optDict=None, outdoc=None, **kwargs +): # pylint:disable=unused-argument """ Function that converts the information produced by the various PyCBC bank generation codes into a valid LIGOLW XML file containing a sngl_inspiral table and outputs to file. Parameters - ----------- + ---------- outputFile : string Name of the file that the bank will be written to tempBank : iterable @@ -229,6 +261,7 @@ def output_sngl_inspiral_table(outputFile, tempBank, programName="", write to disk. If not given create a new document. kwargs : optional key-word arguments Allows unused options to be passed to this function (for modularity) + """ if optDict is None: optDict = {} @@ -238,33 +271,33 @@ def output_sngl_inspiral_table(outputFile, tempBank, programName="", # get IFO to put in search summary table ifos = [] - if 'channel_name' in optDict.keys(): - if optDict['channel_name'] is not None: - ifos = [optDict['channel_name'][0:2]] + if "channel_name" in optDict.keys(): + if optDict["channel_name"] is not None: + ifos = [optDict["channel_name"][0:2]] proc = create_process_table( - outdoc, - program_name=programName, - detectors=ifos, - options=optDict + outdoc, program_name=programName, detectors=ifos, options=optDict ) proc_id = proc.process_id sngl_inspiral_table = convert_to_sngl_inspiral_table(tempBank, proc_id) # set per-template low-frequency cutoff - if 'f_low_column' in optDict and 'f_low' in optDict and \ - optDict['f_low_column'] is not None: + if ( + "f_low_column" in optDict + and "f_low" in optDict + and optDict["f_low_column"] is not None + ): for sngl in sngl_inspiral_table: - setattr(sngl, optDict['f_low_column'], optDict['f_low']) + setattr(sngl, optDict["f_low_column"], optDict["f_low"]) outdoc.childNodes[0].appendChild(sngl_inspiral_table) # get times to put in search summary table start_time = 0 end_time = 0 - if 'gps_start_time' in optDict.keys() and 'gps_end_time' in optDict.keys(): - start_time = optDict['gps_start_time'] - end_time = optDict['gps_end_time'] + if "gps_start_time" in optDict.keys() and "gps_end_time" in optDict.keys(): + start_time = optDict["gps_start_time"] + end_time = optDict["gps_end_time"] # make search summary table search_summary_table = lsctables.SearchSummaryTable.new() @@ -278,15 +311,21 @@ def output_sngl_inspiral_table(outputFile, tempBank, programName="", ligolw_utils.write_filename(outdoc, outputFile) -def output_bank_to_hdf(outputFile, tempBank, optDict=None, programName='', - approximant=None, output_duration=False, - **kwargs): # pylint:disable=unused-argument +def output_bank_to_hdf( + outputFile, + tempBank, + optDict=None, + programName="", + approximant=None, + output_duration=False, + **kwargs, +): # pylint:disable=unused-argument """ Function that converts the information produced by the various PyCBC bank generation codes into a hdf5 file. Parameters - ----------- + ---------- outputFile : string Name of the file that the bank will be written to tempBank : iterable @@ -304,62 +343,55 @@ def output_bank_to_hdf(outputFile, tempBank, optDict=None, programName='', get_waveform_filter_length_in_time, to the file. kwargs : optional key-word arguments Allows unused options to be passed to this function (for modularity) + """ bank_dict = {} mass1, mass2, spin1z, spin2z = list(zip(*tempBank)) - bank_dict['mass1'] = mass1 - bank_dict['mass2'] = mass2 - bank_dict['spin1z'] = spin1z - bank_dict['spin2z'] = spin2z + bank_dict["mass1"] = mass1 + bank_dict["mass2"] = mass2 + bank_dict["spin1z"] = spin1z + bank_dict["spin2z"] = spin2z # Add other values to the bank dictionary as appropriate if optDict is not None: - bank_dict['f_lower'] = numpy.ones_like(mass1) * \ - optDict['f_low'] - argument_string = [f'{k}:{v}' for k, v in optDict.items()] + bank_dict["f_lower"] = numpy.ones_like(mass1) * optDict["f_low"] + argument_string = [f"{k}:{v}" for k, v in optDict.items()] - if optDict is not None and optDict['output_f_final']: - bank_dict['f_final'] = numpy.ones_like(mass1) * \ - optDict['f_upper'] + if optDict is not None and optDict["output_f_final"]: + bank_dict["f_final"] = numpy.ones_like(mass1) * optDict["f_upper"] if approximant: if not isinstance(approximant, bytes): appx = approximant.encode() - bank_dict['approximant'] = numpy.repeat(appx, len(mass1)) + bank_dict["approximant"] = numpy.repeat(appx, len(mass1)) if output_duration: - appx = approximant if approximant else 'SPAtmplt' + appx = approximant or "SPAtmplt" tmplt_durations = numpy.zeros_like(mass1) for i in range(len(mass1)): - wfrm_length = gwflit(appx, - mass1=mass1[i], - mass2=mass2[i], - f_lower=optDict['f_low'], - phase_order=7) + wfrm_length = gwflit( + appx, + mass1=mass1[i], + mass2=mass2[i], + f_lower=optDict["f_low"], + phase_order=7, + ) tmplt_durations[i] = wfrm_length - bank_dict['template_duration'] = tmplt_durations + bank_dict["template_duration"] = tmplt_durations - with HFile(outputFile, 'w') as bankf_out: - bankf_out.attrs['program'] = programName + with HFile(outputFile, "w") as bankf_out: + bankf_out.attrs["program"] = programName if optDict is not None: - bankf_out.attrs['arguments'] = argument_string + bankf_out.attrs["arguments"] = argument_string for k, v in bank_dict.items(): bankf_out[k] = v def output_bank_to_file(outputFile, tempBank, **kwargs): - if outputFile.endswith(('.xml','.xml.gz','.xmlgz')): - output_sngl_inspiral_table( - outputFile, - tempBank, - **kwargs - ) - elif outputFile.endswith(('.h5','.hdf','.hdf5')): - output_bank_to_hdf( - outputFile, - tempBank, - **kwargs - ) + if outputFile.endswith((".xml", ".xml.gz", ".xmlgz")): + output_sngl_inspiral_table(outputFile, tempBank, **kwargs) + elif outputFile.endswith((".h5", ".hdf", ".hdf5")): + output_bank_to_hdf(outputFile, tempBank, **kwargs) else: err_msg = f"Unrecognized extension for file {outputFile}." raise ValueError(err_msg) diff --git a/pycbc/tmpltbank/brute_force_methods.py b/pycbc/tmpltbank/brute_force_methods.py index 34ca473c5df..be6afcd13b3 100644 --- a/pycbc/tmpltbank/brute_force_methods.py +++ b/pycbc/tmpltbank/brute_force_methods.py @@ -1,21 +1,29 @@ import logging + import numpy from pycbc.tmpltbank.coord_utils import get_cov_params -logger = logging.getLogger('pycbc.tmpltbank.brute_force_methods') +logger = logging.getLogger("pycbc.tmpltbank.brute_force_methods") -def get_physical_covaried_masses(xis, bestMasses, bestXis, req_match, - massRangeParams, metricParams, fUpper, - giveUpThresh = 5000): +def get_physical_covaried_masses( + xis, + bestMasses, + bestXis, + req_match, + massRangeParams, + metricParams, + fUpper, + giveUpThresh=5000, +): """ This function takes the position of a point in the xi parameter space and iteratively finds a close point in the physical coordinate space (masses and spins). - + Parameters - ----------- + ---------- xis : list or array Desired position of the point in the xi space. If only N values are provided and the xi space's dimension is larger then it is assumed that @@ -48,7 +56,7 @@ def get_physical_covaried_masses(xis, bestMasses, bestXis, req_match, has been found after this it will give up. Returns - -------- + ------- mass1 : float The heavier mass of the obtained point. mass2 : float @@ -63,6 +71,7 @@ def get_physical_covaried_masses(xis, bestMasses, bestXis, req_match, The mismatch between the obtained point and the input xis. new_xis : list The position of the point in the xi space + """ # TUNABLE PARAMETERS GO HERE! # This states how far apart to scatter test points in the first proposal @@ -71,37 +80,46 @@ def get_physical_covaried_masses(xis, bestMasses, bestXis, req_match, # Set up xi_size = len(xis) scaleFactor = origScaleFactor - bestChirpmass = bestMasses[0] * (bestMasses[1])**(3./5.) + bestChirpmass = bestMasses[0] * (bestMasses[1]) ** (3.0 / 5.0) count = 0 unFixedCount = 0 currDist = 100000000000000000 - while(1): + while 1: # If we are a long way away we use larger jumps if count: if currDist > 1 and scaleFactor == origScaleFactor: - scaleFactor = origScaleFactor*10 + scaleFactor = origScaleFactor * 10 # Get a set of test points with mass -> xi mappings - totmass, eta, spin1z, spin2z, mass1, mass2, new_xis = \ - get_mass_distribution([bestChirpmass, bestMasses[1], bestMasses[2], - bestMasses[3]], - scaleFactor, massRangeParams, metricParams, - fUpper) - cDist = (new_xis[0] - xis[0])**2 - for j in range(1,xi_size): - cDist += (new_xis[j] - xis[j])**2 - if (cDist.min() < req_match): + totmass, eta, spin1z, spin2z, mass1, mass2, new_xis = get_mass_distribution( + [bestChirpmass, bestMasses[1], bestMasses[2], bestMasses[3]], + scaleFactor, + massRangeParams, + metricParams, + fUpper, + ) + cDist = (new_xis[0] - xis[0]) ** 2 + for j in range(1, xi_size): + cDist += (new_xis[j] - xis[j]) ** 2 + if cDist.min() < req_match: idx = cDist.argmin() scaleFactor = origScaleFactor new_xis_list = [new_xis[ldx][idx] for ldx in range(len(new_xis))] - return mass1[idx], mass2[idx], spin1z[idx], spin2z[idx], count, \ - cDist.min(), new_xis_list - if (cDist.min() < currDist): + return ( + mass1[idx], + mass2[idx], + spin1z[idx], + spin2z[idx], + count, + cDist.min(), + new_xis_list, + ) + if cDist.min() < currDist: idx = cDist.argmin() bestMasses[0] = totmass[idx] bestMasses[1] = eta[idx] bestMasses[2] = spin1z[idx] bestMasses[3] = spin2z[idx] - bestChirpmass = bestMasses[0] * (bestMasses[1])**(3./5.) + bestChirpmass = bestMasses[0] * (bestMasses[1]) ** (3.0 / 5.0) currDist = cDist.min() unFixedCount = 0 scaleFactor = origScaleFactor @@ -109,12 +127,19 @@ def get_physical_covaried_masses(xis, bestMasses, bestXis, req_match, unFixedCount += 1 if unFixedCount > giveUpThresh: # Stop at this point - diff = (bestMasses[0]*bestMasses[0] * (1-4*bestMasses[1]))**0.5 - mass1 = (bestMasses[0] + diff)/2. - mass2 = (bestMasses[0] - diff)/2. + diff = (bestMasses[0] * bestMasses[0] * (1 - 4 * bestMasses[1])) ** 0.5 + mass1 = (bestMasses[0] + diff) / 2.0 + mass2 = (bestMasses[0] - diff) / 2.0 new_xis_list = [new_xis[ldx][0] for ldx in range(len(new_xis))] - return mass1, mass2, bestMasses[2], bestMasses[3], count, \ - currDist, new_xis_list + return ( + mass1, + mass2, + bestMasses[2], + bestMasses[3], + count, + currDist, + new_xis_list, + ) if not unFixedCount % 100: scaleFactor *= 2 if scaleFactor > 64: @@ -122,17 +147,25 @@ def get_physical_covaried_masses(xis, bestMasses, bestXis, req_match, # Shouldn't be here! raise RuntimeError -def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, - metricParams, fUpper, - numJumpPoints=100, chirpMassJumpFac=0.0001, - etaJumpFac=0.01, spin1zJumpFac=0.01, - spin2zJumpFac=0.01): + +def get_mass_distribution( + bestMasses, + scaleFactor, + massRangeParams, + metricParams, + fUpper, + numJumpPoints=100, + chirpMassJumpFac=0.0001, + etaJumpFac=0.01, + spin1zJumpFac=0.01, + spin2zJumpFac=0.01, +): """ Given a set of masses, this function will create a set of points nearby in the mass space and map these to the xi space. Parameters - ----------- + ---------- bestMasses : list Contains [ChirpMass, eta, spin1z, spin2z]. Points will be placed around tjos @@ -167,8 +200,8 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, The jump points will be chosen with absolute variation in spin2z up to this multiplied by scaleFactor. - Returns - -------- + Returns + ------- Totmass : numpy.array Total mass of the resulting points Eta : numpy.array @@ -185,8 +218,9 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, Mass2 (mass of smaller body) of the resulting points new_xis : list of numpy.array Position of points in the xi coordinates + """ - # FIXME: It would be better if rejected values could be drawn from the + # FIXME: It would be better if rejected values could be drawn from the # full possible mass/spin distribution. However speed in this function is # a major factor and must be considered. bestChirpmass = bestMasses[0] @@ -195,21 +229,19 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, bestSpin2z = bestMasses[3] # Firstly choose a set of values for masses and spins - chirpmass = bestChirpmass * (1 - (numpy.random.random(numJumpPoints)-0.5) \ - * chirpMassJumpFac * scaleFactor ) + chirpmass = bestChirpmass * ( + 1 - (numpy.random.random(numJumpPoints) - 0.5) * chirpMassJumpFac * scaleFactor + ) etaRange = massRangeParams.maxEta - massRangeParams.minEta currJumpFac = etaJumpFac * scaleFactor - if currJumpFac > etaRange: - currJumpFac = etaRange - eta = bestEta * ( 1 - (numpy.random.random(numJumpPoints) - 0.5) \ - * currJumpFac) + currJumpFac = min(currJumpFac, etaRange) + eta = bestEta * (1 - (numpy.random.random(numJumpPoints) - 0.5) * currJumpFac) maxSpinMag = max(massRangeParams.maxNSSpinMag, massRangeParams.maxBHSpinMag) minSpinMag = min(massRangeParams.maxNSSpinMag, massRangeParams.maxBHSpinMag) # Note that these two are cranged by spinxzFac, *not* spinxzFac/spinxz currJumpFac = spin1zJumpFac * scaleFactor - if currJumpFac > maxSpinMag: - currJumpFac = maxSpinMag + currJumpFac = min(currJumpFac, maxSpinMag) # Actually set the new spin trial points if massRangeParams.nsbhFlag or (maxSpinMag == minSpinMag): @@ -220,10 +252,12 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, curr_spin_1z_jump_fac = massRangeParams.maxBHSpinMag if currJumpFac > massRangeParams.maxNSSpinMag: curr_spin_2z_jump_fac = massRangeParams.maxNSSpinMag - spin1z = bestSpin1z + ( (numpy.random.random(numJumpPoints) - 0.5) \ - * curr_spin_1z_jump_fac) - spin2z = bestSpin2z + ( (numpy.random.random(numJumpPoints) - 0.5) \ - * curr_spin_2z_jump_fac) + spin1z = bestSpin1z + ( + (numpy.random.random(numJumpPoints) - 0.5) * curr_spin_1z_jump_fac + ) + spin2z = bestSpin2z + ( + (numpy.random.random(numJumpPoints) - 0.5) * curr_spin_2z_jump_fac + ) else: # If maxNSSpinMag is very low (0) and maxBHSpinMag is high we can # find it hard to place any points. So mix these when @@ -237,21 +271,23 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, curr_spin_ns_jump_fac = massRangeParams.maxNSSpinMag spin1z = numpy.zeros(numJumpPoints, dtype=float) spin2z = numpy.zeros(numJumpPoints, dtype=float) - split_point = int(numJumpPoints/2) + split_point = int(numJumpPoints / 2) # So set the first half to be at least within the BH range and the # second half to be at least within the NS range - spin1z[:split_point] = bestSpin1z + \ - ( (numpy.random.random(split_point) - 0.5)\ - * curr_spin_bh_jump_fac) - spin1z[split_point:] = bestSpin1z + \ - ( (numpy.random.random(numJumpPoints-split_point) - 0.5)\ - * curr_spin_ns_jump_fac) - spin2z[:split_point] = bestSpin2z + \ - ( (numpy.random.random(split_point) - 0.5)\ - * curr_spin_bh_jump_fac) - spin2z[split_point:] = bestSpin2z + \ - ( (numpy.random.random(numJumpPoints-split_point) - 0.5)\ - * curr_spin_ns_jump_fac) + spin1z[:split_point] = bestSpin1z + ( + (numpy.random.random(split_point) - 0.5) * curr_spin_bh_jump_fac + ) + spin1z[split_point:] = bestSpin1z + ( + (numpy.random.random(numJumpPoints - split_point) - 0.5) + * curr_spin_ns_jump_fac + ) + spin2z[:split_point] = bestSpin2z + ( + (numpy.random.random(split_point) - 0.5) * curr_spin_bh_jump_fac + ) + spin2z[split_point:] = bestSpin2z + ( + (numpy.random.random(numJumpPoints - split_point) - 0.5) + * curr_spin_ns_jump_fac + ) # Point[0] is always set to the original point chirpmass[0] = bestChirpmass @@ -267,10 +303,10 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, eta[eta < 0.0001] = 0.0001 # Total mass, masses and mass diff - totmass = chirpmass / (eta**(3./5.)) - diff = (totmass*totmass * (1-4*eta))**0.5 - mass1 = (totmass + diff)/2. - mass2 = (totmass - diff)/2. + totmass = chirpmass / (eta ** (3.0 / 5.0)) + diff = (totmass * totmass * (1 - 4 * eta)) ** 0.5 + mass1 = (totmass + diff) / 2.0 + mass2 = (totmass - diff) / 2.0 # Check the validity of the spin values # Do the first spin @@ -285,10 +321,12 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, else: # Do have to consider masses boundary_mass = massRangeParams.ns_bh_boundary_mass - numploga1 = numpy.logical_and(mass1 >= boundary_mass, - abs(spin1z) <= massRangeParams.maxBHSpinMag) - numploga2 = numpy.logical_and(mass1 < boundary_mass, - abs(spin1z) <= massRangeParams.maxNSSpinMag) + numploga1 = numpy.logical_and( + mass1 >= boundary_mass, abs(spin1z) <= massRangeParams.maxBHSpinMag + ) + numploga2 = numpy.logical_and( + mass1 < boundary_mass, abs(spin1z) <= massRangeParams.maxNSSpinMag + ) numploga = numpy.logical_or(numploga1, numploga2) numploga = numpy.logical_not(numploga) spin1z[numploga] = 0 @@ -304,10 +342,12 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, else: # Do have to consider masses boundary_mass = massRangeParams.ns_bh_boundary_mass - numplogb1 = numpy.logical_and(mass2 >= boundary_mass, - abs(spin2z) <= massRangeParams.maxBHSpinMag) - numplogb2 = numpy.logical_and(mass2 < boundary_mass, - abs(spin2z) <= massRangeParams.maxNSSpinMag) + numplogb1 = numpy.logical_and( + mass2 >= boundary_mass, abs(spin2z) <= massRangeParams.maxBHSpinMag + ) + numplogb2 = numpy.logical_and( + mass2 < boundary_mass, abs(spin2z) <= massRangeParams.maxNSSpinMag + ) numplogb = numpy.logical_or(numplogb1, numplogb2) numplogb = numpy.logical_not(numplogb) spin2z[numplogb] = 0 @@ -321,20 +361,20 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, # larger than any thresholds used in the functions in brute_force_utils.py # and will always be rejected. An unphysical value cannot be used as it # would result in unphysical metric distances and cause failures. - totmass[mass1 < massRangeParams.minMass1*0.9999] = 0.0001 - totmass[mass1 > massRangeParams.maxMass1*1.0001] = 0.0001 - totmass[mass2 < massRangeParams.minMass2*0.9999] = 0.0001 - totmass[mass2 > massRangeParams.maxMass2*1.0001] = 0.0001 + totmass[mass1 < massRangeParams.minMass1 * 0.9999] = 0.0001 + totmass[mass1 > massRangeParams.maxMass1 * 1.0001] = 0.0001 + totmass[mass2 < massRangeParams.minMass2 * 0.9999] = 0.0001 + totmass[mass2 > massRangeParams.maxMass2 * 1.0001] = 0.0001 # There is some numerical error which can push this a bit higher. We do # *not* want to reject the initial guide point. This error comes from # Masses -> totmass, eta -> masses conversion, we will have points pushing # onto the boudaries of the space. - totmass[totmass > massRangeParams.maxTotMass*1.0001] = 0.0001 - totmass[totmass < massRangeParams.minTotMass*0.9999] = 0.0001 + totmass[totmass > massRangeParams.maxTotMass * 1.0001] = 0.0001 + totmass[totmass < massRangeParams.minTotMass * 0.9999] = 0.0001 if massRangeParams.max_chirp_mass: - totmass[chirpmass > massRangeParams.max_chirp_mass*1.0001] = 0.0001 + totmass[chirpmass > massRangeParams.max_chirp_mass * 1.0001] = 0.0001 if massRangeParams.min_chirp_mass: - totmass[chirpmass < massRangeParams.min_chirp_mass*0.9999] = 0.0001 + totmass[chirpmass < massRangeParams.min_chirp_mass * 0.9999] = 0.0001 if totmass[0] < 0.00011: raise ValueError("Cannot remove the guide point!") @@ -343,21 +383,29 @@ def get_mass_distribution(bestMasses, scaleFactor, massRangeParams, mass2[totmass < 0.00011] = 0.0001 # Then map to xis - new_xis = get_cov_params(mass1, mass2, spin1z, spin2z, - metricParams, fUpper) + new_xis = get_cov_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper) return totmass, eta, spin1z, spin2z, mass1, mass2, new_xis -def stack_xi_direction_brute(xis, bestMasses, bestXis, direction_num, - req_match, massRangeParams, metricParams, fUpper, - scaleFactor=0.8, numIterations=3000): + +def stack_xi_direction_brute( + xis, + bestMasses, + bestXis, + direction_num, + req_match, + massRangeParams, + metricParams, + fUpper, + scaleFactor=0.8, + numIterations=3000, +): """ This function is used to assess the depth of the xi_space in a specified dimension at a specified point in the higher dimensions. It does this by iteratively throwing points at the space to find maxima and minima. Parameters - ----------- - + ---------- xis : list or array Position in the xi space at which to assess the depth. This can be only a subset of the higher dimensions than that being sampled. @@ -391,38 +439,64 @@ def stack_xi_direction_brute(xis, bestMasses, bestXis, direction_num, numIterations : int, optional (default = 3000) The number of times to make calls to get_mass_distribution when assessing the maximum/minimum of this parameter space. Making this - smaller makes the code faster, but at the cost of accuracy. - + smaller makes the code faster, but at the cost of accuracy. + Returns - -------- + ------- xi_min : float The minimal value of the specified dimension at the specified point in parameter space. xi_max : float The maximal value of the specified dimension at the specified point in parameter space. - """ + """ # Find minimum - ximin = find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, \ - req_match, massRangeParams, metricParams, \ - fUpper, find_minimum=True, \ - scaleFactor=scaleFactor, \ - numIterations=numIterations) - + ximin = find_xi_extrema_brute( + xis, + bestMasses, + bestXis, + direction_num, + req_match, + massRangeParams, + metricParams, + fUpper, + find_minimum=True, + scaleFactor=scaleFactor, + numIterations=numIterations, + ) + # Find maximum - ximax = find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, \ - req_match, massRangeParams, metricParams, \ - fUpper, find_minimum=False, \ - scaleFactor=scaleFactor, \ - numIterations=numIterations) + ximax = find_xi_extrema_brute( + xis, + bestMasses, + bestXis, + direction_num, + req_match, + massRangeParams, + metricParams, + fUpper, + find_minimum=False, + scaleFactor=scaleFactor, + numIterations=numIterations, + ) return ximin, ximax -def find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, req_match, \ - massRangeParams, metricParams, fUpper, \ - find_minimum=False, scaleFactor=0.8, \ - numIterations=3000): + +def find_xi_extrema_brute( + xis, + bestMasses, + bestXis, + direction_num, + req_match, + massRangeParams, + metricParams, + fUpper, + find_minimum=False, + scaleFactor=0.8, + numIterations=3000, +): """ This function is used to find the largest or smallest value of the xi space in a specified @@ -430,8 +504,7 @@ def find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, req_match, \ iteratively throwing points at the space to find extrema. Parameters - ----------- - + ---------- xis : list or array Position in the xi space at which to assess the depth. This can be only a subset of the higher dimensions than that being sampled. @@ -468,18 +541,18 @@ def find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, req_match, \ numIterations : int, optional (default = 3000) The number of times to make calls to get_mass_distribution when assessing the maximum/minimum of this parameter space. Making this - smaller makes the code faster, but at the cost of accuracy. + smaller makes the code faster, but at the cost of accuracy. Returns - -------- + ------- xi_extent : float The extremal value of the specified dimension at the specified point in parameter space. - """ + """ # Setup xi_size = len(xis) - bestChirpmass = bestMasses[0] * (bestMasses[1])**(3./5.) + bestChirpmass = bestMasses[0] * (bestMasses[1]) ** (3.0 / 5.0) if find_minimum: xiextrema = 10000000000 else: @@ -487,14 +560,16 @@ def find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, req_match, \ for _ in range(numIterations): # Evaluate extrema of the xi direction specified - totmass, eta, spin1z, spin2z, _, _, new_xis = \ - get_mass_distribution([bestChirpmass,bestMasses[1],bestMasses[2], - bestMasses[3]], - scaleFactor, massRangeParams, metricParams, - fUpper) - cDist = (new_xis[0] - xis[0])**2 + totmass, eta, spin1z, spin2z, _, _, new_xis = get_mass_distribution( + [bestChirpmass, bestMasses[1], bestMasses[2], bestMasses[3]], + scaleFactor, + massRangeParams, + metricParams, + fUpper, + ) + cDist = (new_xis[0] - xis[0]) ** 2 for j in range(1, xi_size): - cDist += (new_xis[j] - xis[j])**2 + cDist += (new_xis[j] - xis[j]) ** 2 redCDist = cDist[cDist < req_match] if len(redCDist): if not find_minimum: @@ -505,13 +580,13 @@ def find_xi_extrema_brute(xis, bestMasses, bestXis, direction_num, req_match, \ new_xis[direction_num][cDist > req_match] = 10000000 currXiExtrema = (new_xis[direction_num]).min() idx = (new_xis[direction_num]).argmin() - if ( ((not find_minimum) and (currXiExtrema > xiextrema)) or \ - (find_minimum and (currXiExtrema < xiextrema)) ): + if ((not find_minimum) and (currXiExtrema > xiextrema)) or ( + find_minimum and (currXiExtrema < xiextrema) + ): xiextrema = currXiExtrema bestMasses[0] = totmass[idx] bestMasses[1] = eta[idx] bestMasses[2] = spin1z[idx] bestMasses[3] = spin2z[idx] - bestChirpmass = bestMasses[0] * (bestMasses[1])**(3./5.) + bestChirpmass = bestMasses[0] * (bestMasses[1]) ** (3.0 / 5.0) return xiextrema - diff --git a/pycbc/tmpltbank/calc_moments.py b/pycbc/tmpltbank/calc_moments.py index e9ef6ef6cef..ea665239b55 100644 --- a/pycbc/tmpltbank/calc_moments.py +++ b/pycbc/tmpltbank/calc_moments.py @@ -15,15 +15,17 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import logging + import numpy from pycbc.tmpltbank.lambda_mapping import generate_mapping -logger = logging.getLogger('pycbc.tmpltbank.calc_moments') +logger = logging.getLogger("pycbc.tmpltbank.calc_moments") -def determine_eigen_directions(metricParams, preserveMoments=False, - vary_fmax=False, vary_density=None): +def determine_eigen_directions( + metricParams, preserveMoments=False, vary_fmax=False, vary_density=None +): """ This function will calculate the coordinate transfomations that are needed to rotate from a coordinate system described by the various Lambda @@ -31,7 +33,7 @@ def determine_eigen_directions(metricParams, preserveMoments=False, metric is Cartesian. Parameters - ----------- + ---------- metricParams : metricParameters instance Structure holding all the options for construction of the metric. preserveMoments : boolean, optional (default False) @@ -51,7 +53,7 @@ def determine_eigen_directions(metricParams, preserveMoments=False, ranges as described for vary_fmax. Returns - -------- + ------- metricParams : metricParameters instance Structure holding all the options for construction of the metric. **THIS FUNCTION ONLY RETURNS THE CLASS** @@ -80,8 +82,8 @@ def determine_eigen_directions(metricParams, preserveMoments=False, contains the result of all the integrals used in computing the metrics above. It can be used for the ethinca components calculation, or other similar calculations. - """ + """ evals = {} evecs = {} metric = {} @@ -89,41 +91,41 @@ def determine_eigen_directions(metricParams, preserveMoments=False, # First step is to get the moments needed to calculate the metric if not (metricParams.moments and preserveMoments): - get_moments(metricParams, vary_fmax=vary_fmax, - vary_density=vary_density) + get_moments(metricParams, vary_fmax=vary_fmax, vary_density=vary_density) # What values are going to be in the moments # J7 is the normalization factor so it *MUST* be present - list = metricParams.moments['J7'].keys() + list = metricParams.moments["J7"].keys() # We start looping over every item in the list of metrics for item in list: # Here we convert the moments into a form easier to use here Js = {} - for i in range(-7,18): - Js[i] = metricParams.moments['J%d'%(i)][item] + for i in range(-7, 18): + Js[i] = metricParams.moments["J%d" % (i)][item] logJs = {} - for i in range(-1,18): - logJs[i] = metricParams.moments['log%d'%(i)][item] + for i in range(-1, 18): + logJs[i] = metricParams.moments["log%d" % (i)][item] loglogJs = {} - for i in range(-1,18): - loglogJs[i] = metricParams.moments['loglog%d'%(i)][item] + for i in range(-1, 18): + loglogJs[i] = metricParams.moments["loglog%d" % (i)][item] logloglogJs = {} - for i in range(-1,18): - logloglogJs[i] = metricParams.moments['logloglog%d'%(i)][item] + for i in range(-1, 18): + logloglogJs[i] = metricParams.moments["logloglog%d" % (i)][item] loglogloglogJs = {} - for i in range(-1,18): - loglogloglogJs[i] = metricParams.moments['loglogloglog%d'%(i)][item] + for i in range(-1, 18): + loglogloglogJs[i] = metricParams.moments["loglogloglog%d" % (i)][item] mapping = generate_mapping(metricParams.pnOrder) # Calculate the metric - gs, unmax_metric_curr = calculate_metric(Js, logJs, loglogJs, - logloglogJs, loglogloglogJs, mapping) + gs, unmax_metric_curr = calculate_metric( + Js, logJs, loglogJs, logloglogJs, loglogloglogJs, mapping + ) metric[item] = gs unmax_metric[item] = unmax_metric_curr @@ -136,12 +138,12 @@ def determine_eigen_directions(metricParams, preserveMoments=False, # Due to numerical imprecision the very small eigenvalues can # be negative. Make these positive. evals[item][i] = -evals[item][i] - if evecs[item][i,i] < 0: + if evecs[item][i, i] < 0: # We demand a convention that all diagonal terms in the matrix # of eigenvalues are positive. # This is done to help visualization of the spaces (increasing # mchirp always goes the same way) - evecs[item][:,i] = - evecs[item][:,i] + evecs[item][:, i] = -evecs[item][:, i] metricParams.evals = evals metricParams.evecs = evecs @@ -150,6 +152,7 @@ def determine_eigen_directions(metricParams, preserveMoments=False, return metricParams + def get_moments(metricParams, vary_fmax=False, vary_density=None): """ This function will calculate the various integrals (moments) that are @@ -157,7 +160,7 @@ def get_moments(metricParams, vary_fmax=False, vary_density=None): coincidence. Parameters - ----------- + ---------- metricParams : metricParameters instance Structure holding all the options for construction of the metric. vary_fmax : boolean, optional (default False) @@ -173,7 +176,7 @@ def get_moments(metricParams, vary_fmax=False, vary_density=None): ranges as described for vary_fmax. Returns - -------- + ------- None : None **THIS FUNCTION RETURNS NOTHING** The following will be **added** to the metricParams structure @@ -219,6 +222,7 @@ def get_moments(metricParams, vary_fmax=False, vary_density=None): The normalization factor can be obtained in moments['I7'][f_cutoff] + """ # NOTE: Unless the TaylorR2F4 metric is used the log^3 and log^4 terms are # not needed. As this calculation is not too slow compared to bank @@ -229,55 +233,106 @@ def get_moments(metricParams, vary_fmax=False, vary_density=None): new_f, new_amp = interpolate_psd(psd_f, psd_amp, metricParams.deltaF) # Need I7 first as this is the normalization factor - funct = lambda x,f0: 1 - I7 = calculate_moment(new_f, new_amp, metricParams.fLow, \ - metricParams.fUpper, metricParams.f0, funct,\ - vary_fmax=vary_fmax, vary_density=vary_density) + funct = lambda x, f0: 1 + I7 = calculate_moment( + new_f, + new_amp, + metricParams.fLow, + metricParams.fUpper, + metricParams.f0, + funct, + vary_fmax=vary_fmax, + vary_density=vary_density, + ) # Do all the J moments moments = {} - moments['I7'] = I7 - for i in range(-7,18): - funct = lambda x,f0: x**((-i+7)/3.) - moments['J%d' %(i)] = calculate_moment(new_f, new_amp, \ - metricParams.fLow, metricParams.fUpper, \ - metricParams.f0, funct, norm=I7, \ - vary_fmax=vary_fmax, vary_density=vary_density) + moments["I7"] = I7 + for i in range(-7, 18): + funct = lambda x, f0: x ** ((-i + 7) / 3.0) + moments["J%d" % (i)] = calculate_moment( + new_f, + new_amp, + metricParams.fLow, + metricParams.fUpper, + metricParams.f0, + funct, + norm=I7, + vary_fmax=vary_fmax, + vary_density=vary_density, + ) # Do the logx multiplied by some power terms - for i in range(-1,18): - funct = lambda x,f0: (numpy.log((x*f0)**(1./3.))) * x**((-i+7)/3.) - moments['log%d' %(i)] = calculate_moment(new_f, new_amp, \ - metricParams.fLow, metricParams.fUpper, \ - metricParams.f0, funct, norm=I7, \ - vary_fmax=vary_fmax, vary_density=vary_density) + for i in range(-1, 18): + funct = lambda x, f0: ( + (numpy.log((x * f0) ** (1.0 / 3.0))) * x ** ((-i + 7) / 3.0) + ) + moments["log%d" % (i)] = calculate_moment( + new_f, + new_amp, + metricParams.fLow, + metricParams.fUpper, + metricParams.f0, + funct, + norm=I7, + vary_fmax=vary_fmax, + vary_density=vary_density, + ) # Do the loglog term - for i in range(-1,18): - funct = lambda x,f0: (numpy.log((x*f0)**(1./3.)))**2 * x**((-i+7)/3.) - moments['loglog%d' %(i)] = calculate_moment(new_f, new_amp, \ - metricParams.fLow, metricParams.fUpper, \ - metricParams.f0, funct, norm=I7, \ - vary_fmax=vary_fmax, vary_density=vary_density) + for i in range(-1, 18): + funct = lambda x, f0: ( + (numpy.log((x * f0) ** (1.0 / 3.0))) ** 2 * x ** ((-i + 7) / 3.0) + ) + moments["loglog%d" % (i)] = calculate_moment( + new_f, + new_amp, + metricParams.fLow, + metricParams.fUpper, + metricParams.f0, + funct, + norm=I7, + vary_fmax=vary_fmax, + vary_density=vary_density, + ) # Do the logloglog term - for i in range(-1,18): - funct = lambda x,f0: (numpy.log((x*f0)**(1./3.)))**3 * x**((-i+7)/3.) - moments['logloglog%d' %(i)] = calculate_moment(new_f, new_amp, \ - metricParams.fLow, metricParams.fUpper, \ - metricParams.f0, funct, norm=I7, \ - vary_fmax=vary_fmax, vary_density=vary_density) + for i in range(-1, 18): + funct = lambda x, f0: ( + (numpy.log((x * f0) ** (1.0 / 3.0))) ** 3 * x ** ((-i + 7) / 3.0) + ) + moments["logloglog%d" % (i)] = calculate_moment( + new_f, + new_amp, + metricParams.fLow, + metricParams.fUpper, + metricParams.f0, + funct, + norm=I7, + vary_fmax=vary_fmax, + vary_density=vary_density, + ) # Do the logloglog term - for i in range(-1,18): - funct = lambda x,f0: (numpy.log((x*f0)**(1./3.)))**4 * x**((-i+7)/3.) - moments['loglogloglog%d' %(i)] = calculate_moment(new_f, new_amp, \ - metricParams.fLow, metricParams.fUpper, \ - metricParams.f0, funct, norm=I7, \ - vary_fmax=vary_fmax, vary_density=vary_density) + for i in range(-1, 18): + funct = lambda x, f0: ( + (numpy.log((x * f0) ** (1.0 / 3.0))) ** 4 * x ** ((-i + 7) / 3.0) + ) + moments["loglogloglog%d" % (i)] = calculate_moment( + new_f, + new_amp, + metricParams.fLow, + metricParams.fUpper, + metricParams.f0, + funct, + norm=I7, + vary_fmax=vary_fmax, + vary_density=vary_density, + ) metricParams.moments = moments + def interpolate_psd(psd_f, psd_amp, deltaF): """ Function to interpolate a PSD to a different value of deltaF. Uses linear @@ -293,11 +348,12 @@ def interpolate_psd(psd_f, psd_amp, deltaF): Value of deltaF to interpolate the PSD to. Returns - -------- + ------- new_psd_f : numpy.array Array of the frequencies contained within the interpolated PSD new_psd_amp : numpy.array Array of the interpolated PSD values at the frequencies in new_psd_f. + """ # In some cases this will be a no-op. I thought about removing this, but # this function can take unequally sampled PSDs and it is difficult to @@ -310,10 +366,10 @@ def interpolate_psd(psd_f, psd_amp, deltaF): for i in range(len(psd_f) - 1): f_low = psd_f[i] - f_high = psd_f[i+1] + f_high = psd_f[i + 1] amp_low = psd_amp[i] - amp_high = psd_amp[i+1] - while(1): + amp_high = psd_amp[i + 1] + while 1: if fcurr > f_high: break new_psd_f.append(fcurr) @@ -324,8 +380,9 @@ def interpolate_psd(psd_f, psd_amp, deltaF): return numpy.asarray(new_psd_f), numpy.asarray(new_psd_amp) -def calculate_moment(psd_f, psd_amp, fmin, fmax, f0, funct, - norm=None, vary_fmax=False, vary_density=None): +def calculate_moment( + psd_f, psd_amp, fmin, fmax, f0, funct, norm=None, vary_fmax=False, vary_density=None +): r""" Function for calculating one of the integrals used to construct a template bank placement metric. The integral calculated will be @@ -337,7 +394,7 @@ def calculate_moment(psd_f, psd_amp, fmin, fmax, f0, funct, chosen Parameters - ----------- + ---------- psd_f : numpy.array numpy array holding the set of evenly spaced frequencies used in the PSD psd_amp : numpy.array @@ -372,20 +429,21 @@ def calculate_moment(psd_f, psd_amp, fmin, fmax, f0, funct, ranges as described for vary_fmax. Returns - -------- + ------- moment : Dictionary of floats moment[f_cutoff] will store the value of the moment at the frequency cutoff given by f_cutoff. - """ + """ # Must ensure deltaF in psd_f is constant psd_x = psd_f / f0 deltax = psd_x[1] - psd_x[0] mask = numpy.logical_and(psd_f > fmin, psd_f < fmax) psdf_red = psd_f[mask] - comps_red = psd_x[mask] ** (-7./3.) * funct(psd_x[mask], f0) * deltax / \ - psd_amp[mask] + comps_red = ( + psd_x[mask] ** (-7.0 / 3.0) * funct(psd_x[mask], f0) * deltax / psd_amp[mask] + ) moment = {} moment[fmax] = comps_red.sum() if norm: @@ -397,14 +455,14 @@ def calculate_moment(psd_f, psd_amp, fmin, fmax, f0, funct, moment[t_fmax] = moment[t_fmax] / norm[t_fmax] return moment -def calculate_metric(Js, logJs, loglogJs, logloglogJs, loglogloglogJs, \ - mapping): + +def calculate_metric(Js, logJs, loglogJs, logloglogJs, loglogloglogJs, mapping): """ This function will take the various integrals calculated by get_moments and convert this into a metric for the appropriate parameter space. Parameters - ----------- + ---------- Js : Dictionary The list of (log^0 x) * x**(-i/3) integrals computed by get_moments() The index is Js[i] @@ -425,114 +483,126 @@ def calculate_metric(Js, logJs, loglogJs, logloglogJs, loglogloglogJs, \ space and map these to entries in the metric matrix. Returns - -------- + ------- metric : numpy.matrix The resulting metric. - """ + """ # How many dimensions in the parameter space? maxLen = len(mapping.keys()) - metric = numpy.zeros(shape=(maxLen,maxLen), dtype=float) - unmax_metric = numpy.zeros(shape=(maxLen+1,maxLen+1), dtype=float) + metric = numpy.zeros(shape=(maxLen, maxLen), dtype=float) + unmax_metric = numpy.zeros(shape=(maxLen + 1, maxLen + 1), dtype=float) for i in range(16): for j in range(16): - calculate_metric_comp(metric, unmax_metric, i, j, Js, - logJs, loglogJs, logloglogJs, - loglogloglogJs, mapping) + calculate_metric_comp( + metric, + unmax_metric, + i, + j, + Js, + logJs, + loglogJs, + logloglogJs, + loglogloglogJs, + mapping, + ) return metric, unmax_metric -def calculate_metric_comp(gs, unmax_metric, i, j, Js, logJs, loglogJs, - logloglogJs, loglogloglogJs, mapping): +def calculate_metric_comp( + gs, unmax_metric, i, j, Js, logJs, loglogJs, logloglogJs, loglogloglogJs, mapping +): """ Used to compute part of the metric. Only call this from within calculate_metric(). Please see the documentation for that function. """ # Time term in unmax_metric. Note that these terms are recomputed a bunch # of time, but this cost is insignificant compared to computing the moments - unmax_metric[-1,-1] = (Js[1] - Js[4]*Js[4]) + unmax_metric[-1, -1] = Js[1] - Js[4] * Js[4] # Normal terms - if 'Lambda%d'%i in mapping and 'Lambda%d'%j in mapping: - gammaij = Js[17-i-j] - Js[12-i]*Js[12-j] - gamma0i = (Js[9-i] - Js[4]*Js[12-i]) - gamma0j = (Js[9-j] - Js[4] * Js[12-j]) - gs[mapping['Lambda%d'%i],mapping['Lambda%d'%j]] = \ - 0.5 * (gammaij - gamma0i*gamma0j/(Js[1] - Js[4]*Js[4])) - unmax_metric[mapping['Lambda%d'%i], -1] = gamma0i - unmax_metric[-1, mapping['Lambda%d'%j]] = gamma0j - unmax_metric[mapping['Lambda%d'%i],mapping['Lambda%d'%j]] = gammaij + if "Lambda%d" % i in mapping and "Lambda%d" % j in mapping: + gammaij = Js[17 - i - j] - Js[12 - i] * Js[12 - j] + gamma0i = Js[9 - i] - Js[4] * Js[12 - i] + gamma0j = Js[9 - j] - Js[4] * Js[12 - j] + gs[mapping["Lambda%d" % i], mapping["Lambda%d" % j]] = 0.5 * ( + gammaij - gamma0i * gamma0j / (Js[1] - Js[4] * Js[4]) + ) + unmax_metric[mapping["Lambda%d" % i], -1] = gamma0i + unmax_metric[-1, mapping["Lambda%d" % j]] = gamma0j + unmax_metric[mapping["Lambda%d" % i], mapping["Lambda%d" % j]] = gammaij # Normal,log cross terms - if 'Lambda%d'%i in mapping and 'LogLambda%d'%j in mapping: - gammaij = logJs[17-i-j] - logJs[12-j] * Js[12-i] - gamma0i = (Js[9-i] - Js[4] * Js[12-i]) - gamma0j = logJs[9-j] - logJs[12-j] * Js[4] - gs[mapping['Lambda%d'%i],mapping['LogLambda%d'%j]] = \ - gs[mapping['LogLambda%d'%j],mapping['Lambda%d'%i]] = \ - 0.5 * (gammaij - gamma0i*gamma0j/(Js[1] - Js[4]*Js[4])) - unmax_metric[mapping['Lambda%d'%i], -1] = gamma0i - unmax_metric[-1, mapping['Lambda%d'%i]] = gamma0i - unmax_metric[-1, mapping['LogLambda%d'%j]] = gamma0j - unmax_metric[mapping['LogLambda%d'%j], -1] = gamma0j - unmax_metric[mapping['Lambda%d'%i],mapping['LogLambda%d'%j]] = gammaij - unmax_metric[mapping['LogLambda%d'%j],mapping['Lambda%d'%i]] = gammaij + if "Lambda%d" % i in mapping and "LogLambda%d" % j in mapping: + gammaij = logJs[17 - i - j] - logJs[12 - j] * Js[12 - i] + gamma0i = Js[9 - i] - Js[4] * Js[12 - i] + gamma0j = logJs[9 - j] - logJs[12 - j] * Js[4] + gs[mapping["Lambda%d" % i], mapping["LogLambda%d" % j]] = gs[ + mapping["LogLambda%d" % j], mapping["Lambda%d" % i] + ] = 0.5 * (gammaij - gamma0i * gamma0j / (Js[1] - Js[4] * Js[4])) + unmax_metric[mapping["Lambda%d" % i], -1] = gamma0i + unmax_metric[-1, mapping["Lambda%d" % i]] = gamma0i + unmax_metric[-1, mapping["LogLambda%d" % j]] = gamma0j + unmax_metric[mapping["LogLambda%d" % j], -1] = gamma0j + unmax_metric[mapping["Lambda%d" % i], mapping["LogLambda%d" % j]] = gammaij + unmax_metric[mapping["LogLambda%d" % j], mapping["Lambda%d" % i]] = gammaij # Log,log terms - if 'LogLambda%d'%i in mapping and 'LogLambda%d'%j in mapping: - gammaij = loglogJs[17-i-j] - logJs[12-j] * logJs[12-i] - gamma0i = (logJs[9-i] - Js[4] * logJs[12-i]) - gamma0j = logJs[9-j] - logJs[12-j] * Js[4] - gs[mapping['LogLambda%d'%i],mapping['LogLambda%d'%j]] = \ - 0.5 * (gammaij - gamma0i*gamma0j/(Js[1] - Js[4]*Js[4])) - unmax_metric[mapping['LogLambda%d'%i], -1] = gamma0i - unmax_metric[-1, mapping['LogLambda%d'%j]] = gamma0j - unmax_metric[mapping['LogLambda%d'%i],mapping['LogLambda%d'%j]] =\ - gammaij + if "LogLambda%d" % i in mapping and "LogLambda%d" % j in mapping: + gammaij = loglogJs[17 - i - j] - logJs[12 - j] * logJs[12 - i] + gamma0i = logJs[9 - i] - Js[4] * logJs[12 - i] + gamma0j = logJs[9 - j] - logJs[12 - j] * Js[4] + gs[mapping["LogLambda%d" % i], mapping["LogLambda%d" % j]] = 0.5 * ( + gammaij - gamma0i * gamma0j / (Js[1] - Js[4] * Js[4]) + ) + unmax_metric[mapping["LogLambda%d" % i], -1] = gamma0i + unmax_metric[-1, mapping["LogLambda%d" % j]] = gamma0j + unmax_metric[mapping["LogLambda%d" % i], mapping["LogLambda%d" % j]] = gammaij # Normal,loglog cross terms - if 'Lambda%d'%i in mapping and 'LogLogLambda%d'%j in mapping: - gammaij = loglogJs[17-i-j] - loglogJs[12-j] * Js[12-i] - gamma0i = (Js[9-i] - Js[4] * Js[12-i]) - gamma0j = loglogJs[9-j] - loglogJs[12-j] * Js[4] - gs[mapping['Lambda%d'%i],mapping['LogLogLambda%d'%j]] = \ - gs[mapping['LogLogLambda%d'%j],mapping['Lambda%d'%i]] = \ - 0.5 * (gammaij - gamma0i*gamma0j/(Js[1] - Js[4]*Js[4])) - unmax_metric[mapping['Lambda%d'%i], -1] = gamma0i - unmax_metric[-1, mapping['Lambda%d'%i]] = gamma0i - unmax_metric[-1, mapping['LogLogLambda%d'%j]] = gamma0j - unmax_metric[mapping['LogLogLambda%d'%j], -1] = gamma0j - unmax_metric[mapping['Lambda%d'%i],mapping['LogLogLambda%d'%j]] = \ - gammaij - unmax_metric[mapping['LogLogLambda%d'%j],mapping['Lambda%d'%i]] = \ - gammaij + if "Lambda%d" % i in mapping and "LogLogLambda%d" % j in mapping: + gammaij = loglogJs[17 - i - j] - loglogJs[12 - j] * Js[12 - i] + gamma0i = Js[9 - i] - Js[4] * Js[12 - i] + gamma0j = loglogJs[9 - j] - loglogJs[12 - j] * Js[4] + gs[mapping["Lambda%d" % i], mapping["LogLogLambda%d" % j]] = gs[ + mapping["LogLogLambda%d" % j], mapping["Lambda%d" % i] + ] = 0.5 * (gammaij - gamma0i * gamma0j / (Js[1] - Js[4] * Js[4])) + unmax_metric[mapping["Lambda%d" % i], -1] = gamma0i + unmax_metric[-1, mapping["Lambda%d" % i]] = gamma0i + unmax_metric[-1, mapping["LogLogLambda%d" % j]] = gamma0j + unmax_metric[mapping["LogLogLambda%d" % j], -1] = gamma0j + unmax_metric[mapping["Lambda%d" % i], mapping["LogLogLambda%d" % j]] = gammaij + unmax_metric[mapping["LogLogLambda%d" % j], mapping["Lambda%d" % i]] = gammaij # log,loglog cross terms - if 'LogLambda%d'%i in mapping and 'LogLogLambda%d'%j in mapping: - gammaij = logloglogJs[17-i-j] - loglogJs[12-j] * logJs[12-i] - gamma0i = (logJs[9-i] - Js[4] * logJs[12-i]) - gamma0j = loglogJs[9-j] - loglogJs[12-j] * Js[4] - gs[mapping['LogLambda%d'%i],mapping['LogLogLambda%d'%j]] = \ - gs[mapping['LogLogLambda%d'%j],mapping['LogLambda%d'%i]] = \ - 0.5 * (gammaij - gamma0i*gamma0j/(Js[1] - Js[4]*Js[4])) - unmax_metric[mapping['LogLambda%d'%i], -1] = gamma0i - unmax_metric[-1, mapping['LogLambda%d'%i]] = gamma0i - unmax_metric[-1, mapping['LogLogLambda%d'%j]] = gamma0j - unmax_metric[mapping['LogLogLambda%d'%j], -1] = gamma0j - unmax_metric[mapping['LogLambda%d'%i],mapping['LogLogLambda%d'%j]] = \ + if "LogLambda%d" % i in mapping and "LogLogLambda%d" % j in mapping: + gammaij = logloglogJs[17 - i - j] - loglogJs[12 - j] * logJs[12 - i] + gamma0i = logJs[9 - i] - Js[4] * logJs[12 - i] + gamma0j = loglogJs[9 - j] - loglogJs[12 - j] * Js[4] + gs[mapping["LogLambda%d" % i], mapping["LogLogLambda%d" % j]] = gs[ + mapping["LogLogLambda%d" % j], mapping["LogLambda%d" % i] + ] = 0.5 * (gammaij - gamma0i * gamma0j / (Js[1] - Js[4] * Js[4])) + unmax_metric[mapping["LogLambda%d" % i], -1] = gamma0i + unmax_metric[-1, mapping["LogLambda%d" % i]] = gamma0i + unmax_metric[-1, mapping["LogLogLambda%d" % j]] = gamma0j + unmax_metric[mapping["LogLogLambda%d" % j], -1] = gamma0j + unmax_metric[mapping["LogLambda%d" % i], mapping["LogLogLambda%d" % j]] = ( gammaij - unmax_metric[mapping['LogLogLambda%d'%j],mapping['LogLambda%d'%i]] = \ + ) + unmax_metric[mapping["LogLogLambda%d" % j], mapping["LogLambda%d" % i]] = ( gammaij + ) # Loglog,loglog terms - if 'LogLogLambda%d'%i in mapping and 'LogLogLambda%d'%j in mapping: - gammaij = loglogloglogJs[17-i-j] - loglogJs[12-j] * loglogJs[12-i] - gamma0i = (loglogJs[9-i] - Js[4] * loglogJs[12-i]) - gamma0j = loglogJs[9-j] - loglogJs[12-j] * Js[4] - gs[mapping['LogLogLambda%d'%i],mapping['LogLogLambda%d'%j]] = \ - 0.5 * (gammaij - gamma0i*gamma0j/(Js[1] - Js[4]*Js[4])) - unmax_metric[mapping['LogLogLambda%d'%i], -1] = gamma0i - unmax_metric[-1, mapping['LogLogLambda%d'%j]] = gamma0j - unmax_metric[mapping['LogLogLambda%d'%i],mapping['LogLogLambda%d'%j]] =\ + if "LogLogLambda%d" % i in mapping and "LogLogLambda%d" % j in mapping: + gammaij = loglogloglogJs[17 - i - j] - loglogJs[12 - j] * loglogJs[12 - i] + gamma0i = loglogJs[9 - i] - Js[4] * loglogJs[12 - i] + gamma0j = loglogJs[9 - j] - loglogJs[12 - j] * Js[4] + gs[mapping["LogLogLambda%d" % i], mapping["LogLogLambda%d" % j]] = 0.5 * ( + gammaij - gamma0i * gamma0j / (Js[1] - Js[4] * Js[4]) + ) + unmax_metric[mapping["LogLogLambda%d" % i], -1] = gamma0i + unmax_metric[-1, mapping["LogLogLambda%d" % j]] = gamma0j + unmax_metric[mapping["LogLogLambda%d" % i], mapping["LogLogLambda%d" % j]] = ( gammaij - + ) diff --git a/pycbc/tmpltbank/coord_utils.py b/pycbc/tmpltbank/coord_utils.py index 7eecdd3b4f7..8791ce7ea7d 100644 --- a/pycbc/tmpltbank/coord_utils.py +++ b/pycbc/tmpltbank/coord_utils.py @@ -15,18 +15,17 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import logging + import numpy -from pycbc.tmpltbank.lambda_mapping import get_chirp_params -from pycbc import conversions -from pycbc import pnutils +from pycbc import conversions, pnutils from pycbc.neutron_stars import load_ns_sequence +from pycbc.tmpltbank.lambda_mapping import get_chirp_params -logger = logging.getLogger('pycbc.tmpltbank.coord_utils') +logger = logging.getLogger("pycbc.tmpltbank.coord_utils") -def estimate_mass_range(numPoints, massRangeParams, metricParams, fUpper,\ - covary=True): +def estimate_mass_range(numPoints, massRangeParams, metricParams, fUpper, covary=True): """ This function will generate a large set of points with random masses and spins (using pycbc.tmpltbank.get_random_mass) and translate these points @@ -59,6 +58,7 @@ def estimate_mass_range(numPoints, massRangeParams, metricParams, fUpper,\ ------- xis : numpy.array A list of the positions of each point in the xi_i coordinate system. + """ vals_set = get_random_mass(numPoints, massRangeParams) mass1 = vals_set[0] @@ -66,14 +66,13 @@ def estimate_mass_range(numPoints, massRangeParams, metricParams, fUpper,\ spin1z = vals_set[2] spin2z = vals_set[3] if covary: - lambdas = get_cov_params(mass1, mass2, spin1z, spin2z, metricParams, - fUpper) + lambdas = get_cov_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper) else: - lambdas = get_conv_params(mass1, mass2, spin1z, spin2z, metricParams, - fUpper) + lambdas = get_conv_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper) return numpy.array(lambdas) + def get_random_mass_point_particles(numPoints, massRangeParams): """ This function will generate a large set of points within the chosen mass @@ -88,7 +87,7 @@ def get_random_mass_point_particles(numPoints, massRangeParams): Instance holding all the details of mass ranges and spin ranges. Returns - -------- + ------- mass1 : float Mass of heavier body. mass2 : float @@ -97,26 +96,30 @@ def get_random_mass_point_particles(numPoints, massRangeParams): Spin of body 1. spin2z : float Spin of body 2. - """ + """ # WARNING: We expect mass1 > mass2 ALWAYS # First we choose the total masses from a unifrom distribution in mass # to the -5/3. power. - mass = numpy.random.random(numPoints) * \ - (massRangeParams.minTotMass**(-5./3.) \ - - massRangeParams.maxTotMass**(-5./3.)) \ - + massRangeParams.maxTotMass**(-5./3.) - mass = mass**(-3./5.) + mass = numpy.random.random(numPoints) * ( + massRangeParams.minTotMass ** (-5.0 / 3.0) + - massRangeParams.maxTotMass ** (-5.0 / 3.0) + ) + massRangeParams.maxTotMass ** (-5.0 / 3.0) + mass = mass ** (-3.0 / 5.0) # Next we choose the mass ratios, this will take different limits based on # the value of total mass - maxmass2 = numpy.minimum(mass/2., massRangeParams.maxMass2) - minmass1 = numpy.maximum(massRangeParams.minMass1, mass/2.) - mineta = numpy.maximum(massRangeParams.minCompMass \ - * (mass-massRangeParams.minCompMass)/(mass*mass), \ - massRangeParams.maxCompMass \ - * (mass-massRangeParams.maxCompMass)/(mass*mass)) + maxmass2 = numpy.minimum(mass / 2.0, massRangeParams.maxMass2) + minmass1 = numpy.maximum(massRangeParams.minMass1, mass / 2.0) + mineta = numpy.maximum( + massRangeParams.minCompMass + * (mass - massRangeParams.minCompMass) + / (mass * mass), + massRangeParams.maxCompMass + * (mass - massRangeParams.maxCompMass) + / (mass * mass), + ) # Note that mineta is a numpy.array because mineta depends on the total # mass. Therefore this is not precomputed in the massRangeParams instance if massRangeParams.minEta: @@ -124,17 +127,17 @@ def get_random_mass_point_particles(numPoints, massRangeParams): # Eta also restricted by chirp mass restrictions if massRangeParams.min_chirp_mass: eta_val_at_min_chirp = massRangeParams.min_chirp_mass / mass - eta_val_at_min_chirp = eta_val_at_min_chirp**(5./3.) + eta_val_at_min_chirp = eta_val_at_min_chirp ** (5.0 / 3.0) mineta = numpy.maximum(mineta, eta_val_at_min_chirp) - maxeta = numpy.minimum(massRangeParams.maxEta, maxmass2 \ - * (mass - maxmass2) / (mass*mass)) - maxeta = numpy.minimum(maxeta, minmass1 \ - * (mass - minmass1) / (mass*mass)) + maxeta = numpy.minimum( + massRangeParams.maxEta, maxmass2 * (mass - maxmass2) / (mass * mass) + ) + maxeta = numpy.minimum(maxeta, minmass1 * (mass - minmass1) / (mass * mass)) # max eta also affected by chirp mass restrictions if massRangeParams.max_chirp_mass: eta_val_at_max_chirp = massRangeParams.max_chirp_mass / mass - eta_val_at_max_chirp = eta_val_at_max_chirp**(5./3.) + eta_val_at_max_chirp = eta_val_at_max_chirp ** (5.0 / 3.0) maxeta = numpy.minimum(maxeta, eta_val_at_max_chirp) if (maxeta < mineta).any(): @@ -143,49 +146,52 @@ def get_random_mass_point_particles(numPoints, massRangeParams): eta = numpy.random.random(numPoints) * (maxeta - mineta) + mineta # Also calculate the component masses; mass1 > mass2 - diff = (mass*mass * (1-4*eta))**0.5 - mass1 = (mass + diff)/2. - mass2 = (mass - diff)/2. + diff = (mass * mass * (1 - 4 * eta)) ** 0.5 + mass1 = (mass + diff) / 2.0 + mass2 = (mass - diff) / 2.0 # Check the masses are where we want them to be (allowing some floating # point rounding error). - if (mass1 > massRangeParams.maxMass1*1.001).any() \ - or (mass1 < massRangeParams.minMass1*0.999).any(): + if (mass1 > massRangeParams.maxMass1 * 1.001).any() or ( + mass1 < massRangeParams.minMass1 * 0.999 + ).any(): errMsg = "Mass1 is not within the specified mass range." raise ValueError(errMsg) - if (mass2 > massRangeParams.maxMass2*1.001).any() \ - or (mass2 < massRangeParams.minMass2*0.999).any(): + if (mass2 > massRangeParams.maxMass2 * 1.001).any() or ( + mass2 < massRangeParams.minMass2 * 0.999 + ).any(): errMsg = "Mass2 is not within the specified mass range." raise ValueError(errMsg) # Next up is the spins. First check if we have non-zero spins if massRangeParams.maxNSSpinMag == 0 and massRangeParams.maxBHSpinMag == 0: - spin1z = numpy.zeros(numPoints,dtype=float) - spin2z = numpy.zeros(numPoints,dtype=float) + spin1z = numpy.zeros(numPoints, dtype=float) + spin2z = numpy.zeros(numPoints, dtype=float) elif massRangeParams.nsbhFlag: # Spin 1 first mspin = numpy.zeros(len(mass1)) mspin += massRangeParams.maxBHSpinMag - spin1z = (2*numpy.random.random(numPoints) - 1) * mspin + spin1z = (2 * numpy.random.random(numPoints) - 1) * mspin # Then spin2 mspin = numpy.zeros(len(mass2)) mspin += massRangeParams.maxNSSpinMag - spin2z = (2*numpy.random.random(numPoints) - 1) * mspin - else: + spin2z = (2 * numpy.random.random(numPoints) - 1) * mspin + else: boundary_mass = massRangeParams.ns_bh_boundary_mass # Spin 1 first mspin = numpy.zeros(len(mass1)) mspin += massRangeParams.maxNSSpinMag mspin[mass1 > boundary_mass] = massRangeParams.maxBHSpinMag - spin1z = (2*numpy.random.random(numPoints) - 1) * mspin + spin1z = (2 * numpy.random.random(numPoints) - 1) * mspin # Then spin 2 mspin = numpy.zeros(len(mass2)) mspin += massRangeParams.maxNSSpinMag mspin[mass2 > boundary_mass] = massRangeParams.maxBHSpinMag - spin2z = (2*numpy.random.random(numPoints) - 1) * mspin + spin2z = (2 * numpy.random.random(numPoints) - 1) * mspin return mass1, mass2, spin1z, spin2z -def get_random_mass(numPoints, massRangeParams, eos='2H'): + +def get_random_mass(numPoints, massRangeParams, eos="2H"): """ This function will generate a large set of points within the chosen mass and spin space, and with the desired minimum remnant disk mass (this applies @@ -201,9 +207,9 @@ def get_random_mass(numPoints, massRangeParams, eos='2H'): Instance holding all the details of mass ranges and spin ranges. eos : string Name of equation of state of neutron star. - + Returns - -------- + ------- mass1 : float Mass of heavier body. mass2 : float @@ -212,16 +218,17 @@ def get_random_mass(numPoints, massRangeParams, eos='2H'): Spin of body 1. spin2z : float Spin of body 2. - """ + """ # WARNING: We expect mass1 > mass2 ALWAYS # Check if EM contraints are required, i.e. if the systems must produce # a minimum remnant disk mass. If this is not the case, proceed treating # the systems as point particle binaries if massRangeParams.remnant_mass_threshold is None: - mass1, mass2, spin1z, spin2z = \ - get_random_mass_point_particles(numPoints, massRangeParams) + mass1, mass2, spin1z, spin2z = get_random_mass_point_particles( + numPoints, massRangeParams + ) # otherwise, load EOS dependent data, generate the EM constraint # (i.e. compute the minimum symmetric mass ratio needed to # generate a given remnant disk mass as a function of the NS @@ -233,10 +240,10 @@ def get_random_mass(numPoints, massRangeParams, eos='2H'): boundary_mass = massRangeParams.ns_bh_boundary_mass if max_ns_g_mass < boundary_mass: warn_msg = "WARNING: " - warn_msg += "Option of ns-bh-boundary-mass is %s " %(boundary_mass) + warn_msg += "Option of ns-bh-boundary-mass is %s " % (boundary_mass) warn_msg += "which is higher than the maximum NS gravitational " warn_msg += "mass admitted by the EOS that was prescribed " - warn_msg += "(%s). " %(max_ns_g_mass) + warn_msg += "(%s). " % (max_ns_g_mass) warn_msg += "The code will proceed using the latter value " warn_msg += "as the boundary mass." logger.warning(warn_msg) @@ -255,9 +262,9 @@ def get_random_mass(numPoints, massRangeParams, eos='2H'): while numPointsFound < numPoints: # Generate the random points within the required mass # and spin cuts - mass1, mass2, spin1z, spin2z = \ - get_random_mass_point_particles(numPoints-numPointsFound, - massRangeParams) + mass1, mass2, spin1z, spin2z = get_random_mass_point_particles( + numPoints - numPointsFound, massRangeParams + ) # Now proceed with cutting out EM dim systems # Use a logical mask to track points that do not correspond to @@ -303,23 +310,27 @@ def get_random_mass(numPoints, massRangeParams, eos='2H'): spin1x=0.0, spin1y=0.0, spin1z=spin1z_nsbh, - eos=eos + eos=eos, + ) + mask_bright_nsbh[remnant > massRangeParams.remnant_mass_threshold] = ( + True ) - mask_bright_nsbh[remnant - > - massRangeParams.remnant_mass_threshold] = True # Keep only points that correspond to binaries that can produce an # EM counterpart (i.e., BNSs and EM-bright NSBHs) and add their # properties to the pile of accpeted points to output - mass1_out = numpy.concatenate((mass1_out, mass1_bns, - mass1_nsbh[mask_bright_nsbh])) - mass2_out = numpy.concatenate((mass2_out, mass2_bns, - mass2_nsbh[mask_bright_nsbh])) - spin1z_out = numpy.concatenate((spin1z_out, spin1z_bns, - spin1z_nsbh[mask_bright_nsbh])) - spin2z_out = numpy.concatenate((spin2z_out, spin2z_bns, - spin2z_nsbh[mask_bright_nsbh])) + mass1_out = numpy.concatenate( + (mass1_out, mass1_bns, mass1_nsbh[mask_bright_nsbh]) + ) + mass2_out = numpy.concatenate( + (mass2_out, mass2_bns, mass2_nsbh[mask_bright_nsbh]) + ) + spin1z_out = numpy.concatenate( + (spin1z_out, spin1z_bns, spin1z_nsbh[mask_bright_nsbh]) + ) + spin2z_out = numpy.concatenate( + (spin2z_out, spin2z_bns, spin2z_nsbh[mask_bright_nsbh]) + ) # Number of points that survived all cuts numPointsFound = len(mass1_out) @@ -332,15 +343,25 @@ def get_random_mass(numPoints, massRangeParams, eos='2H'): return mass1, mass2, spin1z, spin2z -def get_cov_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper, - lambda1=None, lambda2=None, quadparam1=None, - quadparam2=None): + +def get_cov_params( + mass1, + mass2, + spin1z, + spin2z, + metricParams, + fUpper, + lambda1=None, + lambda2=None, + quadparam1=None, + quadparam2=None, +): """ Function to convert between masses and spins and locations in the xi parameter space. Xi = Cartesian metric and rotated to principal components. Parameters - ----------- + ---------- mass1 : float Mass of heavier body. mass2 : float @@ -361,28 +382,47 @@ def get_cov_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper, the given value of fUpper) Returns - -------- + ------- xis : list of floats or numpy.arrays Position of the system(s) in the xi coordinate system - """ + """ # Do this by doing masses - > lambdas -> mus - mus = get_conv_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper, - lambda1=lambda1, lambda2=lambda2, - quadparam1=quadparam1, quadparam2=quadparam2) + mus = get_conv_params( + mass1, + mass2, + spin1z, + spin2z, + metricParams, + fUpper, + lambda1=lambda1, + lambda2=lambda2, + quadparam1=quadparam1, + quadparam2=quadparam2, + ) # and then mus -> xis xis = get_covaried_params(mus, metricParams.evecsCV[fUpper]) return xis -def get_conv_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper, - lambda1=None, lambda2=None, quadparam1=None, - quadparam2=None): + +def get_conv_params( + mass1, + mass2, + spin1z, + spin2z, + metricParams, + fUpper, + lambda1=None, + lambda2=None, + quadparam1=None, + quadparam2=None, +): """ Function to convert between masses and spins and locations in the mu parameter space. Mu = Cartesian metric, but not principal components. Parameters - ----------- + ---------- mass1 : float Mass of heavier body. mass2 : float @@ -402,27 +442,36 @@ def get_conv_params(mass1, mass2, spin1z, spin2z, metricParams, fUpper, the given value of fUpper) Returns - -------- + ------- mus : list of floats or numpy.arrays Position of the system(s) in the mu coordinate system - """ + """ # Do this by masses -> lambdas - lambdas = get_chirp_params(mass1, mass2, spin1z, spin2z, - metricParams.f0, metricParams.pnOrder, - lambda1=lambda1, lambda2=lambda2, - quadparam1=quadparam1, quadparam2=quadparam2) + lambdas = get_chirp_params( + mass1, + mass2, + spin1z, + spin2z, + metricParams.f0, + metricParams.pnOrder, + lambda1=lambda1, + lambda2=lambda2, + quadparam1=quadparam1, + quadparam2=quadparam2, + ) # and lambdas -> mus mus = get_mu_params(lambdas, metricParams, fUpper) return mus + def get_mu_params(lambdas, metricParams, fUpper): """ Function to rotate from the lambda coefficients into position in the mu coordinate system. Mu = Cartesian metric, but not principal components. Parameters - ----------- + ---------- lambdas : list of floats or numpy.arrays Position of the system(s) in the lambda coefficients metricParams : metricParameters instance @@ -436,15 +485,16 @@ def get_mu_params(lambdas, metricParams, fUpper): the given value of fUpper) Returns - -------- + ------- mus : list of floats or numpy.arrays Position of the system(s) in the mu coordinate system + """ lambdas = numpy.asarray(lambdas) # If original inputs were floats we need to make this a 2D array if len(lambdas.shape) == 1: resize_needed = True - lambdas = lambdas[:,None] + lambdas = lambdas[:, None] else: resize_needed = False @@ -454,20 +504,21 @@ def get_mu_params(lambdas, metricParams, fUpper): evecs = numpy.asarray(evecs) mus = ((lambdas.T).dot(evecs)).T - mus = mus * numpy.sqrt(evals)[:,None] + mus = mus * numpy.sqrt(evals)[:, None] if resize_needed: mus = numpy.ndarray.flatten(mus) return mus + def get_covaried_params(mus, evecsCV): """ Function to rotate from position(s) in the mu_i coordinate system into the position(s) in the xi_i coordinate system Parameters - ----------- + ---------- mus : list of floats or numpy.arrays Position of the system(s) in the mu coordinate system evecsCV : numpy.matrix @@ -475,15 +526,16 @@ def get_covaried_params(mus, evecsCV): coordinate system. Returns - -------- + ------- xis : list of floats or numpy.arrays Position of the system(s) in the xi coordinate system + """ mus = numpy.asarray(mus) # If original inputs were floats we need to make this a 2D array if len(mus.shape) == 1: resize_needed = True - mus = mus[:,None] + mus = mus[:, None] else: resize_needed = False @@ -494,13 +546,14 @@ def get_covaried_params(mus, evecsCV): return xis + def rotate_vector(evecs, old_vector, rescale_factor, index): """ Function to find the position of the system(s) in one of the xi_i or mu_i directions. Parameters - ----------- + ---------- evecs : numpy.matrix Matrix of the eigenvectors of the metric in lambda_i coordinates. Used to rotate to a Cartesian coordinate system. @@ -513,15 +566,17 @@ def rotate_vector(evecs, old_vector, rescale_factor, index): if we are going from mu_i -> xi_j, this will give j. Returns - -------- + ------- positions : float or numpy.array Position of the point(s) in the resulting coordinate. + """ temp = 0 for i in range(len(evecs)): - temp += (evecs[i,index] * rescale_factor) * old_vector[i] + temp += (evecs[i, index] * rescale_factor) * old_vector[i] return temp + def get_point_distance(point1, point2, metricParams, fUpper): """ Function to calculate the mismatch between two points, supplied in terms @@ -552,13 +607,14 @@ def get_point_distance(point1, point2, metricParams, fUpper): the given value of fUpper) Returns - -------- + ------- dist : float or numpy.array Distance between the point2 and all points in point1 xis1 : List of floats or numpy.arrays Position of the input point1(s) in the xi_i parameter space xis2 : List of floats Position of the input point2 in the xi_i parameter space + """ aMass1 = point1[0] aMass2 = point1[1] @@ -574,12 +630,13 @@ def get_point_distance(point1, point2, metricParams, fUpper): bXis = get_cov_params(bMass1, bMass2, bSpin1, bSpin2, metricParams, fUpper) - dist = (aXis[0] - bXis[0])**2 - for i in range(1,len(aXis)): - dist += (aXis[i] - bXis[i])**2 + dist = (aXis[0] - bXis[0]) ** 2 + for i in range(1, len(aXis)): + dist += (aXis[i] - bXis[i]) ** 2 return dist, aXis, bXis + def calc_point_dist(vsA, entryA): r""" This function is used to determine the distance between two points. @@ -594,19 +651,21 @@ def calc_point_dist(vsA, entryA): The minimal mismatch allowed between the points Returns - -------- + ------- val : float The metric distance between the two points. + """ chi_diffs = vsA - entryA - val = ((chi_diffs)*(chi_diffs)).sum() - return val + val = ((chi_diffs) * (chi_diffs)).sum() + return val + def test_point_dist(point_1_chis, point_2_chis, distance_threshold): r""" This function tests if the difference between two points in the chi parameter space is less than a distance threshold. Returns True if it is - and False if it is not. + and False if it is not. Parameters ---------- @@ -616,6 +675,7 @@ def test_point_dist(point_1_chis, point_2_chis, distance_threshold): An array of point 2's position in the \chi_i coordinate system distance_threshold : float The distance threshold to use. + """ return calc_point_dist(point_1_chis, point_2_chis) < distance_threshold @@ -651,44 +711,44 @@ def calc_point_dist_vary(mus1, fUpper1, mus2, fUpper2, fMap, norm_map, MMdistA): The minimal mismatch allowed between the points Returns - -------- + ------- Boolean True if the points have a mismatch < MMdistA False if the points have a mismatch > MMdistA + """ f_upper = min(fUpper1, fUpper2) f_other = max(fUpper1, fUpper2) idx = fMap[f_upper] vecs1 = mus1[idx] vecs2 = mus2[idx] - val = ((vecs1 - vecs2)*(vecs1 - vecs2)).sum() - if (val > MMdistA): + val = ((vecs1 - vecs2) * (vecs1 - vecs2)).sum() + if val > MMdistA: return False # Reduce match to account for normalization. norm_fac = norm_map[f_upper] / norm_map[f_other] - val = 1 - (1 - val)*norm_fac - return (val < MMdistA) + val = 1 - (1 - val) * norm_fac + return val < MMdistA def find_max_and_min_frequencies(name, mass_range_params, freqs): """ ADD DOCS """ - cutoff_fns = pnutils.named_frequency_cutoffs if name not in cutoff_fns.keys(): - err_msg = "%s not recognized as a valid cutoff frequency choice." %name + err_msg = "%s not recognized as a valid cutoff frequency choice." % name err_msg += "Recognized choices: " + " ".join(cutoff_fns.keys()) raise ValueError(err_msg) # Can I do this quickly? total_mass_approxs = { "SchwarzISCO": pnutils.f_SchwarzISCO, - "LightRing" : pnutils.f_LightRing, - "ERD" : pnutils.f_ERD + "LightRing": pnutils.f_LightRing, + "ERD": pnutils.f_ERD, } - - if name in total_mass_approxs.keys(): + + if name in total_mass_approxs: # This can be done quickly if the cutoff only depends on total mass # Assumes that lower total mass = higher cutoff frequency upper_f_cutoff = total_mass_approxs[name](mass_range_params.minTotMass) @@ -696,31 +756,30 @@ def find_max_and_min_frequencies(name, mass_range_params, freqs): else: # Do this numerically # FIXME: Is 1000000 the right choice? I think so, but just highlighting - mass1, mass2, spin1z, spin2z = \ - get_random_mass(1000000, mass_range_params) + mass1, mass2, spin1z, spin2z = get_random_mass(1000000, mass_range_params) mass_dict = {} - mass_dict['mass1'] = mass1 - mass_dict['mass2'] = mass2 - mass_dict['spin1z'] = spin1z - mass_dict['spin2z'] = spin2z + mass_dict["mass1"] = mass1 + mass_dict["mass2"] = mass2 + mass_dict["spin1z"] = spin1z + mass_dict["spin2z"] = spin2z tmp_freqs = cutoff_fns[name](mass_dict) upper_f_cutoff = tmp_freqs.max() lower_f_cutoff = tmp_freqs.min() - cutoffs = numpy.array([lower_f_cutoff,upper_f_cutoff]) + cutoffs = numpy.array([lower_f_cutoff, upper_f_cutoff]) if lower_f_cutoff < freqs.min(): warn_msg = "WARNING: " - warn_msg += "Lowest frequency cutoff is %s Hz " %(lower_f_cutoff,) + warn_msg += "Lowest frequency cutoff is %s Hz " % (lower_f_cutoff,) warn_msg += "which is lower than the lowest frequency calculated " - warn_msg += "for the metric: %s Hz. " %(freqs.min()) + warn_msg += "for the metric: %s Hz. " % (freqs.min()) warn_msg += "Distances for these waveforms will be calculated at " warn_msg += "the lowest available metric frequency." logger.warning(warn_msg) if upper_f_cutoff > freqs.max(): warn_msg = "WARNING: " - warn_msg += "Highest frequency cutoff is %s Hz " %(upper_f_cutoff,) + warn_msg += "Highest frequency cutoff is %s Hz " % (upper_f_cutoff,) warn_msg += "which is larger than the highest frequency calculated " - warn_msg += "for the metric: %s Hz. " %(freqs.max()) + warn_msg += "for the metric: %s Hz. " % (freqs.max()) warn_msg += "Distances for these waveforms will be calculated at " warn_msg += "the largest available metric frequency." logger.warning(warn_msg) @@ -750,25 +809,27 @@ def return_nearest_cutoff(name, mass_dict, freqs): ------- numpy.array The frequencies closest to the cutoff for each value of totmass. + """ # A bypass for the redundant case if len(freqs) == 1: - return numpy.zeros(len(mass_dict['m1']), dtype=float) + freqs[0] + return numpy.zeros(len(mass_dict["m1"]), dtype=float) + freqs[0] cutoff_fns = pnutils.named_frequency_cutoffs if name not in cutoff_fns.keys(): - err_msg = "%s not recognized as a valid cutoff frequency choice." %name + err_msg = "%s not recognized as a valid cutoff frequency choice." % name err_msg += "Recognized choices: " + " ".join(cutoff_fns.keys()) raise ValueError(err_msg) f_cutoff = cutoff_fns[name](mass_dict) return find_closest_calculated_frequencies(f_cutoff, freqs) + def find_closest_calculated_frequencies(input_freqs, metric_freqs): """ Given a value (or array) of input frequencies find the closest values in the list of frequencies calculated in the metric. Parameters - ----------- + ---------- input_freqs : numpy.array or float The frequency(ies) that you want to find the closest value in metric_freqs @@ -776,13 +837,14 @@ def find_closest_calculated_frequencies(input_freqs, metric_freqs): The list of frequencies calculated by the metric Returns - -------- + ------- output_freqs : numpy.array or float The list of closest values to input_freqs for which the metric was computed + """ try: - refEv = numpy.zeros(len(input_freqs),dtype=float) + refEv = numpy.zeros(len(input_freqs), dtype=float) except TypeError: refEv = numpy.zeros(1, dtype=float) input_freqs = numpy.array([input_freqs]) @@ -801,17 +863,17 @@ def find_closest_calculated_frequencies(input_freqs, metric_freqs): if i == 0: # If frequency is lower than halfway between the first two entries # use the first (lowest) value - logicArr = input_freqs < ((metric_freqs[0] + metric_freqs[1])/2.) - elif i == (len(metric_freqs)-1): + logicArr = input_freqs < ((metric_freqs[0] + metric_freqs[1]) / 2.0) + elif i == (len(metric_freqs) - 1): # If frequency is larger than halfway between the last two entries # use the last (highest) value - logicArr = input_freqs > ((metric_freqs[-2] + metric_freqs[-1])/2.) + logicArr = input_freqs > ((metric_freqs[-2] + metric_freqs[-1]) / 2.0) else: # For frequencies within the range in freqs, check which points # should use the frequency corresponding to index i. - logicArrA = input_freqs > ((metric_freqs[i-1] + metric_freqs[i])/2.) - logicArrB = input_freqs < ((metric_freqs[i] + metric_freqs[i+1])/2.) - logicArr = numpy.logical_and(logicArrA,logicArrB) + logicArrA = input_freqs > ((metric_freqs[i - 1] + metric_freqs[i]) / 2.0) + logicArrB = input_freqs < ((metric_freqs[i] + metric_freqs[i + 1]) / 2.0) + logicArr = numpy.logical_and(logicArrA, logicArrB) if logicArr.any(): refEv[logicArr] = metric_freqs[i] return refEv @@ -825,7 +887,7 @@ def outspiral_loop(N): a number of bins, but want to start in the center and work outwards. """ # Create a 2D lattice of all points - X,Y = numpy.meshgrid(numpy.arange(-N,N+1), numpy.arange(-N,N+1)) + X, Y = numpy.meshgrid(numpy.arange(-N, N + 1), numpy.arange(-N, N + 1)) # Flatten it X = numpy.ndarray.flatten(X) @@ -834,14 +896,14 @@ def outspiral_loop(N): # Force to an integer X = numpy.array(X, dtype=int) Y = numpy.array(Y, dtype=int) - + # Calculate distances - G = numpy.sqrt(X**2+Y**2) + G = numpy.sqrt(X**2 + Y**2) # Combine back into an array - out_arr = numpy.array([X,Y,G]) - + out_arr = numpy.array([X, Y, G]) + # And order correctly - sorted_out_arr = out_arr[:,out_arr[2].argsort()] + sorted_out_arr = out_arr[:, out_arr[2].argsort()] - return sorted_out_arr[:2,:].T + return sorted_out_arr[:2, :].T diff --git a/pycbc/tmpltbank/lambda_mapping.py b/pycbc/tmpltbank/lambda_mapping.py index 5b23219a56b..2b2feb6f3d2 100644 --- a/pycbc/tmpltbank/lambda_mapping.py +++ b/pycbc/tmpltbank/lambda_mapping.py @@ -13,24 +13,31 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -import re import logging -import numpy +import re +import numpy -from pycbc.constants import MTSUN_SI, PI import pycbc.libutils +from pycbc.constants import MTSUN_SI, PI -lal = pycbc.libutils.import_optional('lal') -lalsimulation = pycbc.libutils.import_optional('lalsimulation') +lal = pycbc.libutils.import_optional("lal") +lalsimulation = pycbc.libutils.import_optional("lalsimulation") -logger = logging.getLogger('pycbc.tmpltbank.lambda_mapping') +logger = logging.getLogger("pycbc.tmpltbank.lambda_mapping") # PLEASE ENSURE THESE ARE KEPT UP TO DATE WITH THE REST OF THIS FILE -pycbcValidTmpltbankOrders = ['zeroPN','onePN','onePointFivePN','twoPN',\ - 'twoPointFivePN','threePN','threePointFivePN'] - -pycbcValidOrdersHelpDescriptions=""" +pycbcValidTmpltbankOrders = [ + "zeroPN", + "onePN", + "onePointFivePN", + "twoPN", + "twoPointFivePN", + "threePN", + "threePointFivePN", +] + +pycbcValidOrdersHelpDescriptions = """ * zeroPN: Will only include the dominant term (proportional to chirp mass) * onePN: Will only the leading orbit term and first correction at 1PN * onePointFivePN: Will include orbit and spin terms to 1.5PN. @@ -58,47 +65,52 @@ def generate_mapping(order): A string containing a PN order. Valid values are given above. Returns - -------- + ------- mapping : dictionary A mapping between the active Lambda terms and index in the metric + """ mapping = {} - mapping['Lambda0'] = 0 - if order == 'zeroPN': + mapping["Lambda0"] = 0 + if order == "zeroPN": return mapping - mapping['Lambda2'] = 1 - if order == 'onePN': + mapping["Lambda2"] = 1 + if order == "onePN": return mapping - mapping['Lambda3'] = 2 - if order == 'onePointFivePN': + mapping["Lambda3"] = 2 + if order == "onePointFivePN": return mapping - mapping['Lambda4'] = 3 - if order == 'twoPN': + mapping["Lambda4"] = 3 + if order == "twoPN": return mapping - mapping['LogLambda5'] = 4 - if order == 'twoPointFivePN': + mapping["LogLambda5"] = 4 + if order == "twoPointFivePN": return mapping - mapping['Lambda6'] = 5 - mapping['LogLambda6'] = 6 - if order == 'threePN': + mapping["Lambda6"] = 5 + mapping["LogLambda6"] = 6 + if order == "threePN": return mapping - mapping['Lambda7'] = 7 - if order == 'threePointFivePN': + mapping["Lambda7"] = 7 + if order == "threePointFivePN": return mapping # For some as-of-yet unknown reason, the tidal terms are not giving correct # match estimates when enabled. So, for now, this order is commented out. - #if order == 'tidalTesting': + # if order == 'tidalTesting': # mapping['Lambda10'] = 8 # mapping['Lambda12'] = 9 # return mapping - raise ValueError("Order %s is not understood." %(order)) + raise ValueError("Order %s is not understood." % (order)) + # Override doc so the PN orders are added automatically to online docs -generate_mapping.__doc__ = \ - generate_mapping.__doc__.format(pycbcValidOrdersHelpDescriptions) +generate_mapping.__doc__ = generate_mapping.__doc__.format( + pycbcValidOrdersHelpDescriptions +) + def generate_inverse_mapping(order): - """Genereate a lambda entry -> PN order map. + """ + Genereate a lambda entry -> PN order map. This function will generate the opposite of generate mapping. So where generate_mapping gives dict[key] = item this will give @@ -111,36 +123,42 @@ def generate_inverse_mapping(order): A string containing a PN order. Valid values are given above. Returns - -------- + ------- mapping : dictionary An inverse mapping between the active Lambda terms and index in the metric + """ mapping = generate_mapping(order) inv_mapping = {} - for key,value in mapping.items(): + for key, value in mapping.items(): inv_mapping[value] = key return inv_mapping -generate_inverse_mapping.__doc__ = \ - generate_inverse_mapping.__doc__.format(pycbcValidOrdersHelpDescriptions) + +generate_inverse_mapping.__doc__ = generate_inverse_mapping.__doc__.format( + pycbcValidOrdersHelpDescriptions +) + def get_ethinca_orders(): """ Returns the dictionary mapping TaylorF2 PN order names to twice-PN orders (powers of v/c) """ - ethinca_orders = {"zeroPN" : 0, - "onePN" : 2, - "onePointFivePN" : 3, - "twoPN" : 4, - "twoPointFivePN" : 5, - "threePN" : 6, - "threePointFivePN" : 7 - } + ethinca_orders = { + "zeroPN": 0, + "onePN": 2, + "onePointFivePN": 3, + "twoPN": 4, + "twoPointFivePN": 5, + "threePN": 6, + "threePointFivePN": 7, + } return ethinca_orders + def ethinca_order_from_string(order): """ Returns the integer giving twice the post-Newtonian order @@ -153,16 +171,28 @@ def ethinca_order_from_string(order): Returns ------- int + """ if order in get_ethinca_orders().keys(): return get_ethinca_orders()[order] - else: raise ValueError("Order "+str(order)+" is not valid for ethinca" - "calculation! Valid orders: "+ - str(get_ethinca_orders().keys())) - -def get_chirp_params(mass1, mass2, spin1z, spin2z, f0, order, - quadparam1=None, quadparam2=None, lambda1=None, - lambda2=None): + raise ValueError( + "Order " + str(order) + " is not valid for ethinca" + "calculation! Valid orders: " + str(get_ethinca_orders().keys()) + ) + + +def get_chirp_params( + mass1, + mass2, + spin1z, + spin2z, + f0, + order, + quadparam1=None, + quadparam2=None, + lambda1=None, + lambda2=None, +): """ Take a set of masses and spins and convert to the various lambda coordinates that describe the orbital phase. Accepted PN orders are: @@ -190,11 +220,11 @@ def get_chirp_params(mass1, mass2, spin1z, spin2z, f0, order, spins to the lambda_i coordinate system. Valid orders given above. Returns - -------- + ------- lambdas : list of floats or numpy.arrays The lambda coordinates for the input system(s) - """ + """ # Determine whether array or single value input sngl_inp = False try: @@ -239,15 +269,22 @@ def get_chirp_params(mass1, mass2, spin1z, spin2z, f0, order, lambda2_v = lal.CreateREAL8Vector(len(mass1)) lambda2_v.data[:] = lambda2[:] dquadparam1_v = lal.CreateREAL8Vector(len(mass1)) - dquadparam1_v.data[:] = quadparam1[:] - 1. + dquadparam1_v.data[:] = quadparam1[:] - 1.0 dquadparam2_v = lal.CreateREAL8Vector(len(mass1)) - dquadparam2_v.data[:] = quadparam2[:] - 1. - - phasing_arr = lalsimulation.SimInspiralTaylorF2AlignedPhasingArray\ - (mass1_v, mass2_v, spin1z_v, spin2z_v, lambda1_v, lambda2_v, - dquadparam1_v, dquadparam2_v) - - vec_len = lalsimulation.PN_PHASING_SERIES_MAX_ORDER + 1; + dquadparam2_v.data[:] = quadparam2[:] - 1.0 + + phasing_arr = lalsimulation.SimInspiralTaylorF2AlignedPhasingArray( + mass1_v, + mass2_v, + spin1z_v, + spin2z_v, + lambda1_v, + lambda2_v, + dquadparam1_v, + dquadparam2_v, + ) + + vec_len = lalsimulation.PN_PHASING_SERIES_MAX_ORDER + 1 phasing_vs = numpy.zeros([num_points, vec_len]) phasing_vlogvs = numpy.zeros([num_points, vec_len]) phasing_vlogvsqs = numpy.zeros([num_points, vec_len]) @@ -255,49 +292,52 @@ def get_chirp_params(mass1, mass2, spin1z, spin2z, f0, order, lng = len(mass1) jmp = lng * vec_len for idx in range(vec_len): - phasing_vs[:,idx] = phasing_arr.data[lng*idx : lng*(idx+1)] - phasing_vlogvs[:,idx] = \ - phasing_arr.data[jmp + lng*idx : jmp + lng*(idx+1)] - phasing_vlogvsqs[:,idx] = \ - phasing_arr.data[2*jmp + lng*idx : 2*jmp + lng*(idx+1)] - - pim = PI * (mass1 + mass2)*MTSUN_SI + phasing_vs[:, idx] = phasing_arr.data[lng * idx : lng * (idx + 1)] + phasing_vlogvs[:, idx] = phasing_arr.data[ + jmp + lng * idx : jmp + lng * (idx + 1) + ] + phasing_vlogvsqs[:, idx] = phasing_arr.data[ + 2 * jmp + lng * idx : 2 * jmp + lng * (idx + 1) + ] + + pim = PI * (mass1 + mass2) * MTSUN_SI pmf = pim * f0 - pmf13 = pmf**(1./3.) - logpim13 = numpy.log((pim)**(1./3.)) + pmf13 = pmf ** (1.0 / 3.0) + logpim13 = numpy.log((pim) ** (1.0 / 3.0)) mapping = generate_inverse_mapping(order) lambdas = [] - lambda_str = '^Lambda([0-9]+)' - loglambda_str = '^LogLambda([0-9]+)' - logloglambda_str = '^LogLogLambda([0-9]+)' + lambda_str = "^Lambda([0-9]+)" + loglambda_str = "^LogLambda([0-9]+)" + logloglambda_str = "^LogLogLambda([0-9]+)" for idx in range(len(mapping.keys())): # RE magic engage! rematch = re.match(lambda_str, mapping[idx]) if rematch: pn_order = int(rematch.groups()[0]) - term = phasing_vs[:,pn_order] - term = term + logpim13 * phasing_vlogvs[:,pn_order] - lambdas.append(term * pmf13**(-5+pn_order)) + term = phasing_vs[:, pn_order] + term = term + logpim13 * phasing_vlogvs[:, pn_order] + lambdas.append(term * pmf13 ** (-5 + pn_order)) continue rematch = re.match(loglambda_str, mapping[idx]) if rematch: pn_order = int(rematch.groups()[0]) - lambdas.append((phasing_vlogvs[:,pn_order]) * pmf13**(-5+pn_order)) + lambdas.append((phasing_vlogvs[:, pn_order]) * pmf13 ** (-5 + pn_order)) continue rematch = re.match(logloglambda_str, mapping[idx]) if rematch: raise ValueError("LOGLOG terms are not implemented") - #pn_order = int(rematch.groups()[0]) - #lambdas.append(phasing_vlogvsqs[:,pn_order] * pmf13**(-5+pn_order)) - #continue - err_msg = "Failed to parse " + mapping[idx] + # pn_order = int(rematch.groups()[0]) + # lambdas.append(phasing_vlogvsqs[:,pn_order] * pmf13**(-5+pn_order)) + # continue + err_msg = "Failed to parse " + mapping[idx] raise ValueError(err_msg) if sngl_inp: return [l[0] for l in lambdas] - else: - return lambdas + return lambdas + -get_chirp_params.__doc__ = \ - get_chirp_params.__doc__.format(pycbcValidOrdersHelpDescriptions) +get_chirp_params.__doc__ = get_chirp_params.__doc__.format( + pycbcValidOrdersHelpDescriptions +) diff --git a/pycbc/tmpltbank/lattice_utils.py b/pycbc/tmpltbank/lattice_utils.py index 05992ab5531..3254784aa75 100644 --- a/pycbc/tmpltbank/lattice_utils.py +++ b/pycbc/tmpltbank/lattice_utils.py @@ -14,12 +14,13 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -import logging import copy +import logging + import numpy +logger = logging.getLogger("pycbc.tmpltbank.lattice_utils") -logger = logging.getLogger('pycbc.tmpltbank.lattice_utils') def generate_hexagonal_lattice(maxv1, minv1, maxv2, minv2, mindist): """ @@ -27,7 +28,7 @@ def generate_hexagonal_lattice(maxv1, minv1, maxv2, minv2, mindist): lattice. Parameters - ----------- + ---------- maxv1 : float Largest value in the 1st dimension to cover minv1 : float @@ -41,11 +42,12 @@ def generate_hexagonal_lattice(maxv1, minv1, maxv2, minv2, mindist): generated bank of points. Returns - -------- + ------- v1s : numpy.array Array of positions in the first dimension v2s : numpy.array Array of positions in the second dimension + """ if minv1 > maxv1: raise ValueError("Invalid input to function.") @@ -54,47 +56,47 @@ def generate_hexagonal_lattice(maxv1, minv1, maxv2, minv2, mindist): # Place first point v1s = [minv1] v2s = [minv2] - initPoint = [minv1,minv2] + initPoint = [minv1, minv2] # Place first line initLine = [initPoint] tmpv1 = minv1 - while (tmpv1 < maxv1): - tmpv1 = tmpv1 + (3 * mindist)**(0.5) - initLine.append([tmpv1,minv2]) + while tmpv1 < maxv1: + tmpv1 = tmpv1 + (3 * mindist) ** (0.5) + initLine.append([tmpv1, minv2]) v1s.append(tmpv1) v2s.append(minv2) initLine = numpy.array(initLine) initLine2 = copy.deepcopy(initLine) - initLine2[:,0] += 0.5 * (3*mindist)**0.5 - initLine2[:,1] += 1.5 * (mindist)**0.5 + initLine2[:, 0] += 0.5 * (3 * mindist) ** 0.5 + initLine2[:, 1] += 1.5 * (mindist) ** 0.5 for i in range(len(initLine2)): - v1s.append(initLine2[i,0]) - v2s.append(initLine2[i,1]) - tmpv2_1 = initLine[0,1] - tmpv2_2 = initLine2[0,1] + v1s.append(initLine2[i, 0]) + v2s.append(initLine2[i, 1]) + tmpv2_1 = initLine[0, 1] + tmpv2_2 = initLine2[0, 1] while tmpv2_1 < maxv2 and tmpv2_2 < maxv2: - tmpv2_1 = tmpv2_1 + 3.0 * (mindist)**0.5 - tmpv2_2 = tmpv2_2 + 3.0 * (mindist)**0.5 - initLine[:,1] = tmpv2_1 - initLine2[:,1] = tmpv2_2 + tmpv2_1 = tmpv2_1 + 3.0 * (mindist) ** 0.5 + tmpv2_2 = tmpv2_2 + 3.0 * (mindist) ** 0.5 + initLine[:, 1] = tmpv2_1 + initLine2[:, 1] = tmpv2_2 for i in range(len(initLine)): - v1s.append(initLine[i,0]) - v2s.append(initLine[i,1]) + v1s.append(initLine[i, 0]) + v2s.append(initLine[i, 1]) for i in range(len(initLine2)): - v1s.append(initLine2[i,0]) - v2s.append(initLine2[i,1]) + v1s.append(initLine2[i, 0]) + v2s.append(initLine2[i, 1]) v1s = numpy.array(v1s) v2s = numpy.array(v2s) return v1s, v2s -def generate_anstar_3d_lattice(maxv1, minv1, maxv2, minv2, maxv3, minv3, \ - mindist): + +def generate_anstar_3d_lattice(maxv1, minv1, maxv2, minv2, maxv3, minv3, mindist): """ This function calls into LAL routines to generate a 3-dimensional array of points using the An^* lattice. Parameters - ----------- + ---------- maxv1 : float Largest value in the 1st dimension to cover minv1 : float @@ -112,13 +114,14 @@ def generate_anstar_3d_lattice(maxv1, minv1, maxv2, minv2, maxv3, minv3, \ generated bank of points. Returns - -------- + ------- v1s : numpy.array Array of positions in the first dimension v2s : numpy.array Array of positions in the second dimension v3s : numpy.array Array of positions in the second dimension + """ # Lal/Lalpulsar are not a requirement for the rest of pycbc, so check if we have it # here in this function. @@ -136,16 +139,16 @@ def generate_anstar_3d_lattice(maxv1, minv1, maxv2, minv2, maxv3, minv3, \ lalpulsar.SetLatticeTilingConstantBound(tiling, 1, minv2, maxv2) lalpulsar.SetLatticeTilingConstantBound(tiling, 2, minv3, maxv3) # Make a 3x3 Euclidean lattice - a = lal.gsl_matrix(3,3) - a.data[0,0] = 1 - a.data[1,1] = 1 - a.data[2,2] = 1 + a = lal.gsl_matrix(3, 3) + a.data[0, 0] = 1 + a.data[1, 1] = 1 + a.data[2, 2] = 1 try: # old versions of lalpulsar used an enumeration lattice = lalpulsar.TILING_LATTICE_ANSTAR except AttributeError: # newer versions of lalpulsar use a string - lattice = 'An-star' + lattice = "An-star" lalpulsar.SetTilingLatticeAndMetric(tiling, lattice, a, mindist) try: iterator = lalpulsar.CreateLatticeTilingIterator(tiling, 3) @@ -158,9 +161,8 @@ def generate_anstar_3d_lattice(maxv1, minv1, maxv2, minv2, maxv3, minv3, \ vs2 = [] vs3 = [] curr_point = lal.gsl_vector(3) - while (lalpulsar.NextLatticeTilingPoint(iterator, curr_point) > 0): + while lalpulsar.NextLatticeTilingPoint(iterator, curr_point) > 0: vs1.append(curr_point.data[0]) vs2.append(curr_point.data[1]) vs3.append(curr_point.data[2]) return vs1, vs2, vs3 - diff --git a/pycbc/tmpltbank/option_utils.py b/pycbc/tmpltbank/option_utils.py index 935e00c9418..4b404be63d5 100644 --- a/pycbc/tmpltbank/option_utils.py +++ b/pycbc/tmpltbank/option_utils.py @@ -16,41 +16,47 @@ import argparse import logging +import os import textwrap + import numpy -import os -from pycbc.tmpltbank.lambda_mapping import get_ethinca_orders, pycbcValidOrdersHelpDescriptions from pycbc import pnutils from pycbc.neutron_stars import load_ns_sequence -from pycbc.types import positive_float, nonnegative_float +from pycbc.tmpltbank.lambda_mapping import ( + get_ethinca_orders, + pycbcValidOrdersHelpDescriptions, +) +from pycbc.types import nonnegative_float, positive_float -logger = logging.getLogger('pycbc.tmpltbank.option_utils') +logger = logging.getLogger("pycbc.tmpltbank.option_utils") class IndentedHelpFormatterWithNL(argparse.ArgumentDefaultsHelpFormatter): """ - This class taken from + This class taken from https://groups.google.com/forum/#!topic/comp.lang.python/bfbmtUGhW8I and is used to format the argparse help messages to deal with line breaking nicer. Specfically the pn-order help is large and looks crappy without this. This function is (C) Tim Chase """ + def format_description(self, description): """ No documentation """ - if not description: return "" + if not description: + return "" desc_width = self.width - self.current_indent - indent = " "*self.current_indent + indent = " " * self.current_indent # the above is still the same - bits = description.split('\n') + bits = description.split("\n") formatted_bits = [ - textwrap.fill(bit, - desc_width, - initial_indent=indent, - subsequent_indent=indent) - for bit in bits] + textwrap.fill( + bit, desc_width, initial_indent=indent, subsequent_indent=indent + ) + for bit in bits + ] result = "\n".join(formatted_bits) + "\n" return result @@ -78,7 +84,7 @@ def format_option(self, option): if len(opts) > opt_width: opts = "%*s%s\n" % (self.current_indent, "", opts) indent_first = self.help_position - else: # start help on same line as opts + else: # start help on same line as opts opts = "%*s%-*s " % (self.current_indent, "", opt_width, opts) indent_first = 0 result.append(opts) @@ -89,14 +95,15 @@ def format_option(self, option): for para in help_text.split("\n"): help_lines.extend(textwrap.wrap(para, self.help_width)) # Everything is the same after here - result.append("%*s%s\n" % ( - indent_first, "", help_lines[0])) - result.extend(["%*s%s\n" % (self.help_position, "", line) - for line in help_lines[1:]]) + result.append("%*s%s\n" % (indent_first, "", help_lines[0])) + result.extend( + ["%*s%s\n" % (self.help_position, "", line) for line in help_lines[1:]] + ) elif opts[-1] != "\n": result.append("\n") return "".join(result) + def get_options_from_group(option_group): """ Take an option group and return all the options that are defined in that @@ -107,10 +114,11 @@ def get_options_from_group(option_group): for option in option_list: option_strings = option.option_strings for string in option_strings: - if string.startswith('--'): + if string.startswith("--"): command_lines.append(string) return command_lines + def insert_base_bank_options(parser, match_req=True): """ Adds essential common options for template bank generation to an @@ -128,19 +136,30 @@ def match_type(s): return value parser.add_argument( - '-m', '--min-match', type=match_type, required=match_req, - help="Generate bank with specified minimum match. Required.") + "-m", + "--min-match", + type=match_type, + required=match_req, + help="Generate bank with specified minimum match. Required.", + ) parser.add_argument( - '-O', '--output-file', required=True, - help="Output file name. Required.") - parser.add_argument('--f-low-column', type=str, metavar='NAME', - help='If given, store the lower frequency cutoff into ' - 'column NAME of the single-inspiral table. ' - '(Requires an output file ending in .xml)') + "-O", "--output-file", required=True, help="Output file name. Required." + ) parser.add_argument( - '--output-f-final', action='store_true', default=False, - help="Include 'f_final' in the output hdf file." + "--f-low-column", + type=str, + metavar="NAME", + help="If given, store the lower frequency cutoff into " + "column NAME of the single-inspiral table. " + "(Requires an output file ending in .xml)", ) + parser.add_argument( + "--output-f-final", + action="store_true", + default=False, + help="Include 'f_final' in the output hdf file.", + ) + def insert_metric_calculation_options(parser): """ @@ -149,41 +168,66 @@ def insert_metric_calculation_options(parser): options in your code. """ metricOpts = parser.add_argument_group( - "Options related to calculating the parameter space metric") - metricOpts.add_argument("--pn-order", action="store", type=str, - required=True, - help="Determines the PN order to use. For a bank of " - "non-spinning templates, spin-related terms in the " - "metric will be zero. REQUIRED. " - "Choices: %s" %(pycbcValidOrdersHelpDescriptions)) - metricOpts.add_argument("--f0", action="store", type=positive_float, - default=70.,\ - help="f0 is used as a dynamic scaling factor when " - "calculating integrals used in metric construction. " - "I.e. instead of integrating F(f) we integrate F(f/f0) " - "then rescale by powers of f0. The default value 70Hz " - "should be fine for most applications. OPTIONAL. " - "UNITS=Hz. **WARNING: If the ethinca metric is to be " - "calculated, f0 must be set equal to f-low**") - metricOpts.add_argument("--f-low", action="store", type=positive_float, - required=True, - help="Lower frequency cutoff used in computing the " - "parameter space metric. REQUIRED. UNITS=Hz") - metricOpts.add_argument("--f-upper", action="store", type=positive_float, - required=True, - help="Upper frequency cutoff used in computing the " - "parameter space metric. REQUIRED. UNITS=Hz") - metricOpts.add_argument("--delta-f", action="store", type=positive_float, - required=True, - help="Frequency spacing used in computing the parameter " - r"space metric: integrals of the form \int F(f) df " - r"are approximated as \sum F(f) delta_f. REQUIRED. " - "UNITS=Hz") - metricOpts.add_argument("--write-metric", action="store_true", - default=False, help="If given write the metric components " - "to disk as they are calculated.") + "Options related to calculating the parameter space metric" + ) + metricOpts.add_argument( + "--pn-order", + action="store", + type=str, + required=True, + help="Determines the PN order to use. For a bank of " + "non-spinning templates, spin-related terms in the " + "metric will be zero. REQUIRED. " + "Choices: %s" % (pycbcValidOrdersHelpDescriptions), + ) + metricOpts.add_argument( + "--f0", + action="store", + type=positive_float, + default=70.0, + help="f0 is used as a dynamic scaling factor when " + "calculating integrals used in metric construction. " + "I.e. instead of integrating F(f) we integrate F(f/f0) " + "then rescale by powers of f0. The default value 70Hz " + "should be fine for most applications. OPTIONAL. " + "UNITS=Hz. **WARNING: If the ethinca metric is to be " + "calculated, f0 must be set equal to f-low**", + ) + metricOpts.add_argument( + "--f-low", + action="store", + type=positive_float, + required=True, + help="Lower frequency cutoff used in computing the " + "parameter space metric. REQUIRED. UNITS=Hz", + ) + metricOpts.add_argument( + "--f-upper", + action="store", + type=positive_float, + required=True, + help="Upper frequency cutoff used in computing the " + "parameter space metric. REQUIRED. UNITS=Hz", + ) + metricOpts.add_argument( + "--delta-f", + action="store", + type=positive_float, + required=True, + help="Frequency spacing used in computing the parameter " + r"space metric: integrals of the form \int F(f) df " + r"are approximated as \sum F(f) delta_f. REQUIRED. " + "UNITS=Hz", + ) + metricOpts.add_argument( + "--write-metric", + action="store_true", + default=False, + help="If given write the metric components to disk as they are calculated.", + ) return metricOpts + def verify_metric_calculation_options(opts, parser): """ Parses the metric calculation options given and verifies that they are @@ -195,11 +239,13 @@ def verify_metric_calculation_options(opts, parser): Result of parsing the input options with OptionParser parser : object The OptionParser instance. + """ if not opts.pn_order: parser.error("Must supply --pn-order") -class metricParameters(object): + +class metricParameters: """ This class holds all of the options that are parsed in the function insert_metric_calculation_options @@ -207,26 +253,27 @@ class metricParameters(object): from the __init__ function, providing directly the options normally provided on the command line. """ + _psd = None _metric = None _evals = None _evecs = None _evecsCV = None - def __init__(self, pnOrder, fLow, fUpper, deltaF, f0=70, - write_metric=False): + + def __init__(self, pnOrder, fLow, fUpper, deltaF, f0=70, write_metric=False): """ Initialize an instance of the metricParameters by providing all options directly. See the help message associated with any code that uses the metric options for more details of how to set each of these, e.g. pycbc_aligned_stoch_bank --help """ - self.pnOrder=pnOrder - self.fLow=fLow - self.fUpper=fUpper - self.deltaF=deltaF - self.f0=f0 - self._moments=None - self.write_metric=write_metric + self.pnOrder = pnOrder + self.fLow = fLow + self.fUpper = fUpper + self.deltaF = deltaF + self.f0 = f0 + self._moments = None + self.write_metric = write_metric @classmethod def from_argparse(cls, opts): @@ -238,8 +285,14 @@ def from_argparse(cls, opts): verify_metric_calculation_options have already been called before initializing the class. """ - return cls(opts.pn_order, opts.f_low, opts.f_upper, opts.delta_f,\ - f0=opts.f0, write_metric=opts.write_metric) + return cls( + opts.pn_order, + opts.f_low, + opts.f_upper, + opts.delta_f, + f0=opts.f0, + write_metric=opts.write_metric, + ) @property def psd(self): @@ -273,23 +326,23 @@ def moments(self): For the first entries the options are: moments['J%d' %(i)][f_cutoff] - This stores the integral of + This stores the integral of x**((-i)/3.) * delta X / PSD(x) moments['log%d' %(i)][f_cutoff] - This stores the integral of + This stores the integral of (numpy.log(x**(1./3.))) x**((-i)/3.) * delta X / PSD(x) moments['loglog%d' %(i)][f_cutoff] - This stores the integral of + This stores the integral of (numpy.log(x**(1./3.)))**2 x**((-i)/3.) * delta X / PSD(x) moments['loglog%d' %(i)][f_cutoff] - This stores the integral of + This stores the integral of (numpy.log(x**(1./3.)))**3 x**((-i)/3.) * delta X / PSD(x) moments['loglog%d' %(i)][f_cutoff] - This stores the integral of + This stores the integral of (numpy.log(x**(1./3.)))**4 x**((-i)/3.) * delta X / PSD(x) The second entry stores the frequency cutoff that was used when @@ -299,7 +352,7 @@ def moments(self): @moments.setter def moments(self, inMoments): - self._moments=inMoments + self._moments = inMoments @property def evals(self): @@ -326,8 +379,7 @@ def evals(self): def evals(self, inEvals): if self.write_metric: for frequency in inEvals.keys(): - numpy.savetxt("metric_evals_%d.dat" %(frequency), - inEvals[frequency]) + numpy.savetxt("metric_evals_%d.dat" % (frequency), inEvals[frequency]) self._evals = inEvals @property @@ -350,8 +402,7 @@ def evecs(self): def evecs(self, inEvecs): if self.write_metric: for frequency in inEvecs.keys(): - numpy.savetxt("metric_evecs_%d.dat" %(frequency), - inEvecs[frequency]) + numpy.savetxt("metric_evecs_%d.dat" % (frequency), inEvecs[frequency]) self._evecs = inEvecs @property @@ -373,8 +424,9 @@ def metric(self): def metric(self, inMetric): if self.write_metric: for frequency in inMetric.keys(): - numpy.savetxt("metric_components_%d.dat" %(frequency), - inMetric[frequency]) + numpy.savetxt( + "metric_components_%d.dat" % (frequency), inMetric[frequency] + ) self._metric = inMetric @property @@ -397,8 +449,9 @@ def time_unprojected_metric(self): def time_unprojected_metric(self, inMetric): if self.write_metric: for frequency in inMetric.keys(): - numpy.savetxt("metric_timeunprojected_%d.dat" %(frequency), - inMetric[frequency]) + numpy.savetxt( + "metric_timeunprojected_%d.dat" % (frequency), inMetric[frequency] + ) self._time_unprojected_metric = inMetric @property @@ -421,140 +474,216 @@ def evecsCV(self): def evecsCV(self, inEvecs): if self.write_metric: for frequency in inEvecs.keys(): - numpy.savetxt("covariance_evecs_%d.dat" %(frequency), - inEvecs[frequency]) + numpy.savetxt( + "covariance_evecs_%d.dat" % (frequency), inEvecs[frequency] + ) self._evecsCV = inEvecs -def insert_mass_range_option_group(parser,nonSpin=False): +def insert_mass_range_option_group(parser, nonSpin=False): """ Adds the options used to specify mass ranges in the bank generation codes to an argparser as an OptionGroup. This should be used if you want to use these options in your code. - + Parameters - ----------- + ---------- parser : object OptionParser instance. nonSpin : boolean, optional (default=False) If this is provided the spin-related options will not be added. + """ - massOpts = parser.add_argument_group("Options related to mass and spin " - "limits for bank generation") - massOpts.add_argument("--min-mass1", action="store", type=positive_float, - required=True, - help="Minimum mass1: must be >= min-mass2. " - "REQUIRED. UNITS=Solar mass") - massOpts.add_argument("--max-mass1", action="store", type=positive_float, - required=True, - help="Maximum mass1: must be >= max-mass2. " - "REQUIRED. UNITS=Solar mass") - massOpts.add_argument("--min-mass2", action="store", type=positive_float, - required=True, - help="Minimum mass2. REQUIRED. UNITS=Solar mass") - massOpts.add_argument("--max-mass2", action="store", type=positive_float, - required=True, - help="Maximum mass2. REQUIRED. UNITS=Solar mass") - massOpts.add_argument("--max-total-mass", action="store", - type=positive_float, default=None, - help="Maximum total mass. OPTIONAL, if not provided " - "the max total mass is determined by the component " - "masses. UNITS=Solar mass") - massOpts.add_argument("--min-total-mass", action="store", - type=positive_float, default=None, - help="Minimum total mass. OPTIONAL, if not provided the " - "min total mass is determined by the component masses." - " UNITS=Solar mass") - massOpts.add_argument("--max-chirp-mass", action="store", - type=positive_float, default=None, - help="Maximum chirp mass. OPTIONAL, if not provided the " - "max chirp mass is determined by the component masses." - " UNITS=Solar mass") - massOpts.add_argument("--min-chirp-mass", action="store", - type=positive_float, default=None, - help="Minimum total mass. OPTIONAL, if not provided the " - "min chirp mass is determined by the component masses." - " UNITS=Solar mass") - massOpts.add_argument("--max-eta", action="store", type=positive_float, - default=0.25, - help="Maximum symmetric mass ratio. OPTIONAL, no upper bound" - " on eta will be imposed if not provided. " - "UNITS=Solar mass.") - massOpts.add_argument("--min-eta", action="store", type=nonnegative_float, - default=0., - help="Minimum symmetric mass ratio. OPTIONAL, no lower bound" - " on eta will be imposed if not provided. " - "UNITS=Solar mass.") - massOpts.add_argument("--ns-eos", action="store", - default=None, - help="Select the EOS to be used for the NS when calculating " - "the remnant disk mass. Only 2H is currently supported. " - "OPTIONAL") - massOpts.add_argument("--remnant-mass-threshold", action="store", - type=nonnegative_float, default=None, - help="Setting this filters EM dim NS-BH binaries: if the " - "remnant disk mass does not exceed this value, the NS-BH " - "binary is dropped from the target parameter space. " - "When it is set to None (default value) the EM dim " - "filter is not activated. OPTIONAL") - massOpts.add_argument("--use-eos-max-ns-mass", action="store_true", default=False, - help="Cut the mass range of the smaller object to the maximum " - "mass allowed by EOS. " - "OPTIONAL") - massOpts.add_argument("--delta-bh-spin", action="store", - type=positive_float, default=None, - help="Grid spacing used for the BH spin z component when " - "generating the surface of the minumum minimum symmetric " - "mass ratio as a function of BH spin and NS mass required " - "to produce a remnant disk mass that exceeds the threshold " - "specificed in --remnant-mass-threshold. " - "OPTIONAL (0.1 by default) ") - massOpts.add_argument("--delta-ns-mass", action="store", - type=positive_float, default=None, - help="Grid spacing used for the NS mass when generating the " - "surface of the minumum minimum symmetric mass ratio as " - "a function of BH spin and NS mass required to produce " - "a remnant disk mass that exceeds the thrsehold specified " - "in --remnant-mass-threshold. " - "OPTIONAL (0.1 by default) ") + massOpts = parser.add_argument_group( + "Options related to mass and spin limits for bank generation" + ) + massOpts.add_argument( + "--min-mass1", + action="store", + type=positive_float, + required=True, + help="Minimum mass1: must be >= min-mass2. REQUIRED. UNITS=Solar mass", + ) + massOpts.add_argument( + "--max-mass1", + action="store", + type=positive_float, + required=True, + help="Maximum mass1: must be >= max-mass2. REQUIRED. UNITS=Solar mass", + ) + massOpts.add_argument( + "--min-mass2", + action="store", + type=positive_float, + required=True, + help="Minimum mass2. REQUIRED. UNITS=Solar mass", + ) + massOpts.add_argument( + "--max-mass2", + action="store", + type=positive_float, + required=True, + help="Maximum mass2. REQUIRED. UNITS=Solar mass", + ) + massOpts.add_argument( + "--max-total-mass", + action="store", + type=positive_float, + default=None, + help="Maximum total mass. OPTIONAL, if not provided " + "the max total mass is determined by the component " + "masses. UNITS=Solar mass", + ) + massOpts.add_argument( + "--min-total-mass", + action="store", + type=positive_float, + default=None, + help="Minimum total mass. OPTIONAL, if not provided the " + "min total mass is determined by the component masses." + " UNITS=Solar mass", + ) + massOpts.add_argument( + "--max-chirp-mass", + action="store", + type=positive_float, + default=None, + help="Maximum chirp mass. OPTIONAL, if not provided the " + "max chirp mass is determined by the component masses." + " UNITS=Solar mass", + ) + massOpts.add_argument( + "--min-chirp-mass", + action="store", + type=positive_float, + default=None, + help="Minimum total mass. OPTIONAL, if not provided the " + "min chirp mass is determined by the component masses." + " UNITS=Solar mass", + ) + massOpts.add_argument( + "--max-eta", + action="store", + type=positive_float, + default=0.25, + help="Maximum symmetric mass ratio. OPTIONAL, no upper bound" + " on eta will be imposed if not provided. " + "UNITS=Solar mass.", + ) + massOpts.add_argument( + "--min-eta", + action="store", + type=nonnegative_float, + default=0.0, + help="Minimum symmetric mass ratio. OPTIONAL, no lower bound" + " on eta will be imposed if not provided. " + "UNITS=Solar mass.", + ) + massOpts.add_argument( + "--ns-eos", + action="store", + default=None, + help="Select the EOS to be used for the NS when calculating " + "the remnant disk mass. Only 2H is currently supported. " + "OPTIONAL", + ) + massOpts.add_argument( + "--remnant-mass-threshold", + action="store", + type=nonnegative_float, + default=None, + help="Setting this filters EM dim NS-BH binaries: if the " + "remnant disk mass does not exceed this value, the NS-BH " + "binary is dropped from the target parameter space. " + "When it is set to None (default value) the EM dim " + "filter is not activated. OPTIONAL", + ) + massOpts.add_argument( + "--use-eos-max-ns-mass", + action="store_true", + default=False, + help="Cut the mass range of the smaller object to the maximum " + "mass allowed by EOS. " + "OPTIONAL", + ) + massOpts.add_argument( + "--delta-bh-spin", + action="store", + type=positive_float, + default=None, + help="Grid spacing used for the BH spin z component when " + "generating the surface of the minumum minimum symmetric " + "mass ratio as a function of BH spin and NS mass required " + "to produce a remnant disk mass that exceeds the threshold " + "specificed in --remnant-mass-threshold. " + "OPTIONAL (0.1 by default) ", + ) + massOpts.add_argument( + "--delta-ns-mass", + action="store", + type=positive_float, + default=None, + help="Grid spacing used for the NS mass when generating the " + "surface of the minumum minimum symmetric mass ratio as " + "a function of BH spin and NS mass required to produce " + "a remnant disk mass that exceeds the thrsehold specified " + "in --remnant-mass-threshold. " + "OPTIONAL (0.1 by default) ", + ) if nonSpin: parser.add_argument_group(massOpts) return massOpts - massOpts.add_argument("--max-ns-spin-mag", action="store", - type=nonnegative_float, default=None, - help="Maximum neutron star spin magnitude. Neutron stars " - "are defined as components lighter than the NS-BH " - "boundary (3 Msun by default). REQUIRED if min-mass2 " - "< ns-bh-boundary-mass") - massOpts.add_argument("--max-bh-spin-mag", action="store", - type=nonnegative_float, default=None, - help="Maximum black hole spin magnitude. Black holes are " - "defined as components at or above the NS-BH boundary " - "(3 Msun by default). REQUIRED if max-mass1 >= " - "ns-bh-boundary-mass") + massOpts.add_argument( + "--max-ns-spin-mag", + action="store", + type=nonnegative_float, + default=None, + help="Maximum neutron star spin magnitude. Neutron stars " + "are defined as components lighter than the NS-BH " + "boundary (3 Msun by default). REQUIRED if min-mass2 " + "< ns-bh-boundary-mass", + ) + massOpts.add_argument( + "--max-bh-spin-mag", + action="store", + type=nonnegative_float, + default=None, + help="Maximum black hole spin magnitude. Black holes are " + "defined as components at or above the NS-BH boundary " + "(3 Msun by default). REQUIRED if max-mass1 >= " + "ns-bh-boundary-mass", + ) # Mutually exclusive group prevents both options being set on command line # If --nsbh-flag is True then spinning bank generation must ignore the # default value of ns-bh-boundary-mass. action = massOpts.add_mutually_exclusive_group(required=False) - action.add_argument("--ns-bh-boundary-mass", action='store', - type=positive_float, - help="Mass boundary between neutron stars and black holes. " - "Components below this mass are considered neutron " - "stars and are subject to the neutron star spin limits. " - "Components at/above are subject to the black hole spin " - "limits. OPTIONAL, default=%f. UNITS=Solar mass" \ - % massRangeParameters.default_nsbh_boundary_mass) - action.add_argument("--nsbh-flag", action="store_true", default=False, - help="Set this flag if generating a bank that contains only " - "systems with 1 black hole and 1 neutron star. With " - "this flag set the heavier body will always be subject " - "to the black hole spin restriction and the lighter " - "to the neutron star spin restriction, regardless of " - "mass. OPTIONAL. If set, the value of " - "--ns-bh-boundary-mass will be ignored.") + action.add_argument( + "--ns-bh-boundary-mass", + action="store", + type=positive_float, + help="Mass boundary between neutron stars and black holes. " + "Components below this mass are considered neutron " + "stars and are subject to the neutron star spin limits. " + "Components at/above are subject to the black hole spin " + "limits. OPTIONAL, default=%f. UNITS=Solar mass" + % massRangeParameters.default_nsbh_boundary_mass, + ) + action.add_argument( + "--nsbh-flag", + action="store_true", + default=False, + help="Set this flag if generating a bank that contains only " + "systems with 1 black hole and 1 neutron star. With " + "this flag set the heavier body will always be subject " + "to the black hole spin restriction and the lighter " + "to the neutron star spin restriction, regardless of " + "mass. OPTIONAL. If set, the value of " + "--ns-bh-boundary-mass will be ignored.", + ) return massOpts + def verify_mass_range_options(opts, parser, nonSpin=False): """ Parses the metric calculation options given and verifies that they are @@ -568,6 +697,7 @@ def verify_mass_range_options(opts, parser, nonSpin=False): The OptionParser instance. nonSpin : boolean, optional (default=False) If this is provided the spin-related options will not be checked. + """ # Mass1 must be the heavier! if opts.min_mass1 < opts.min_mass2: @@ -575,26 +705,29 @@ def verify_mass_range_options(opts, parser, nonSpin=False): if opts.max_mass1 < opts.max_mass2: parser.error("max-mass1 cannot be less than max-mass2!") # If given are min/max total mass/chirp mass possible? - if opts.min_total_mass \ - and (opts.min_total_mass > opts.max_mass1 + opts.max_mass2): - err_msg = "Supplied minimum total mass %f " %(opts.min_total_mass,) + if opts.min_total_mass and (opts.min_total_mass > opts.max_mass1 + opts.max_mass2): + err_msg = "Supplied minimum total mass %f " % (opts.min_total_mass,) err_msg += "greater than the sum of the two max component masses " - err_msg += " %f and %f." %(opts.max_mass1,opts.max_mass2) + err_msg += " %f and %f." % (opts.max_mass1, opts.max_mass2) parser.error(err_msg) - if opts.max_total_mass \ - and (opts.max_total_mass < opts.min_mass1 + opts.min_mass2): - err_msg = "Supplied maximum total mass %f " %(opts.max_total_mass,) + if opts.max_total_mass and (opts.max_total_mass < opts.min_mass1 + opts.min_mass2): + err_msg = "Supplied maximum total mass %f " % (opts.max_total_mass,) err_msg += "smaller than the sum of the two min component masses " - err_msg += " %f and %f." %(opts.min_mass1,opts.min_mass2) + err_msg += " %f and %f." % (opts.min_mass1, opts.min_mass2) parser.error(err_msg) - if opts.max_total_mass and opts.min_total_mass \ - and (opts.max_total_mass < opts.min_total_mass): + if ( + opts.max_total_mass + and opts.min_total_mass + and (opts.max_total_mass < opts.min_total_mass) + ): parser.error("Min total mass must be larger than max total mass.") # Warn the user that his/her setup is such that EM dim NS-BH binaries # will not be targeted by the template bank that is being built. Also # inform him/her about the caveats involved in this. - if hasattr(opts, 'remnant_mass_threshold') \ - and opts.remnant_mass_threshold is not None: + if ( + hasattr(opts, "remnant_mass_threshold") + and opts.remnant_mass_threshold is not None + ): logger.info("""You have asked to exclude EM dim NS-BH systems from the target parameter space. The script will assume that m1 is the BH and m2 is the NS: make sure that your settings @@ -608,16 +741,18 @@ def verify_mass_range_options(opts, parser, nonSpin=False): # and inform user whether this will be read from file or generated. # This is the minumum eta as a function of BH spin and NS mass # required to produce an EM counterpart - if os.path.isfile('constraint_em_bright.npz'): + if os.path.isfile("constraint_em_bright.npz"): logger.info("""The constraint surface for EM bright binaries will be read in from constraint_em_bright.npz.""") # Assign min/max total mass from mass1, mass2 if not specified - if (not opts.min_total_mass) or \ - ((opts.min_mass1 + opts.min_mass2) > opts.min_total_mass): + if (not opts.min_total_mass) or ( + (opts.min_mass1 + opts.min_mass2) > opts.min_total_mass + ): opts.min_total_mass = opts.min_mass1 + opts.min_mass2 - if (not opts.max_total_mass) or \ - ((opts.max_mass1 + opts.max_mass2) < opts.max_total_mass): + if (not opts.max_total_mass) or ( + (opts.max_mass1 + opts.max_mass2) < opts.max_total_mass + ): opts.max_total_mass = opts.max_mass1 + opts.max_mass2 # It is vital that min and max total mass be set correctly. @@ -635,29 +770,34 @@ def verify_mass_range_options(opts, parser, nonSpin=False): # mass line within the parameter space, or it doesn't intersect # at all. # First let's get the masses at both of these possible points - m1_at_max_m2 = pnutils.mchirp_mass1_to_mass2(opts.min_chirp_mass, - opts.max_mass2) + m1_at_max_m2 = pnutils.mchirp_mass1_to_mass2( + opts.min_chirp_mass, opts.max_mass2 + ) if m1_at_max_m2 < opts.max_mass2: # Unphysical, remove m1_at_max_m2 = -1 - m2_at_min_m1 = pnutils.mchirp_mass1_to_mass2(opts.min_chirp_mass, - opts.min_mass1) + m2_at_min_m1 = pnutils.mchirp_mass1_to_mass2( + opts.min_chirp_mass, opts.min_mass1 + ) if m2_at_min_m1 > opts.min_mass1: # Unphysical, remove m2_at_min_m1 = -1 # Get the values on the equal mass line m1_at_equal_mass, m2_at_equal_mass = pnutils.mchirp_eta_to_mass1_mass2( - opts.min_chirp_mass, 0.25) + opts.min_chirp_mass, 0.25 + ) # Are any of these possible? if m1_at_max_m2 <= opts.max_mass1 and m1_at_max_m2 >= opts.min_mass1: min_tot_mass = opts.max_mass2 + m1_at_max_m2 elif m2_at_min_m1 <= opts.max_mass2 and m2_at_min_m1 >= opts.min_mass2: min_tot_mass = opts.min_mass1 + m2_at_min_m1 - elif m1_at_equal_mass <= opts.max_mass1 and \ - m1_at_equal_mass >= opts.min_mass1 and \ - m2_at_equal_mass <= opts.max_mass2 and \ - m2_at_equal_mass >= opts.min_mass2: + elif ( + m1_at_equal_mass <= opts.max_mass1 + and m1_at_equal_mass >= opts.min_mass1 + and m2_at_equal_mass <= opts.max_mass2 + and m2_at_equal_mass >= opts.min_mass2 + ): min_tot_mass = m1_at_equal_mass + m2_at_equal_mass # So either the restriction is low enough to be redundant, or is # removing all the parameter space @@ -673,7 +813,8 @@ def verify_mass_range_options(opts, parser, nonSpin=False): if opts.max_eta: # Get the value of m1,m2 at max_eta, min_chirp_mass max_eta_m1, max_eta_m2 = pnutils.mchirp_eta_to_mass1_mass2( - opts.min_chirp_mass, opts.max_eta) + opts.min_chirp_mass, opts.max_eta + ) max_eta_min_tot_mass = max_eta_m1 + max_eta_m2 if max_eta_min_tot_mass > min_tot_mass: # Okay, eta does restrict this further. Still physical? @@ -695,10 +836,12 @@ def verify_mass_range_options(opts, parser, nonSpin=False): # Either it is found at min_m2, or at max_m1, or it doesn't intersect # at all. # First let's get the masses at both of these possible points - m1_at_min_m2 = pnutils.mchirp_mass1_to_mass2(opts.max_chirp_mass, - opts.min_mass2) - m2_at_max_m1 = pnutils.mchirp_mass1_to_mass2(opts.max_chirp_mass, - opts.max_mass1) + m1_at_min_m2 = pnutils.mchirp_mass1_to_mass2( + opts.max_chirp_mass, opts.min_mass2 + ) + m2_at_max_m1 = pnutils.mchirp_mass1_to_mass2( + opts.max_chirp_mass, opts.max_mass1 + ) # Are either of these possible? if m1_at_min_m2 <= opts.max_mass1 and m1_at_min_m2 >= opts.min_mass1: max_tot_mass = opts.min_mass2 + m1_at_min_m2 @@ -718,7 +861,8 @@ def verify_mass_range_options(opts, parser, nonSpin=False): if opts.min_eta: # Get the value of m1,m2 at max_eta, min_chirp_mass min_eta_m1, min_eta_m2 = pnutils.mchirp_eta_to_mass1_mass2( - opts.max_chirp_mass, opts.min_eta) + opts.max_chirp_mass, opts.min_eta + ) min_eta_max_tot_mass = min_eta_m1 + min_eta_m2 if min_eta_max_tot_mass < max_tot_mass: # Okay, eta does restrict this further. Still physical? @@ -737,14 +881,18 @@ def verify_mass_range_options(opts, parser, nonSpin=False): # Similar to above except this can affect both the minimum and maximum # total mass. Need to identify where the line of max_eta intersects # the parameter space, and if it affects mass restrictions. - m1_at_min_m2 = pnutils.eta_mass1_to_mass2(opts.max_eta, opts.min_mass2, - return_mass_heavier=True) - m2_at_min_m1 = pnutils.eta_mass1_to_mass2(opts.max_eta, opts.min_mass1, - return_mass_heavier=False) - m1_at_max_m2 = pnutils.eta_mass1_to_mass2(opts.max_eta, opts.max_mass2, - return_mass_heavier=True) - m2_at_max_m1 = pnutils.eta_mass1_to_mass2(opts.max_eta, opts.max_mass1, - return_mass_heavier=False) + m1_at_min_m2 = pnutils.eta_mass1_to_mass2( + opts.max_eta, opts.min_mass2, return_mass_heavier=True + ) + m2_at_min_m1 = pnutils.eta_mass1_to_mass2( + opts.max_eta, opts.min_mass1, return_mass_heavier=False + ) + m1_at_max_m2 = pnutils.eta_mass1_to_mass2( + opts.max_eta, opts.max_mass2, return_mass_heavier=True + ) + m2_at_max_m1 = pnutils.eta_mass1_to_mass2( + opts.max_eta, opts.max_mass1, return_mass_heavier=False + ) # Check for restrictions on the minimum total mass # Are either of these possible? if m1_at_min_m2 <= opts.max_mass1 and m1_at_min_m2 >= opts.min_mass1: @@ -757,8 +905,9 @@ def verify_mass_range_options(opts, parser, nonSpin=False): elif m2_at_min_m1 > opts.max_mass2: # This is the redundant case, ignore min_tot_mass = opts.min_total_mass - elif opts.max_eta == 0.25 and (m1_at_min_m2 < opts.min_mass2 or \ - m2_at_min_m1 > opts.min_mass1): + elif opts.max_eta == 0.25 and ( + m1_at_min_m2 < opts.min_mass2 or m2_at_min_m1 > opts.min_mass1 + ): # This just catches potential roundoff issues in the case that # max-eta is not used min_tot_mass = opts.min_total_mass @@ -793,14 +942,18 @@ def verify_mass_range_options(opts, parser, nonSpin=False): # Same as max_eta. # Need to identify where the line of max_eta intersects # the parameter space, and if it affects mass restrictions. - m1_at_min_m2 = pnutils.eta_mass1_to_mass2(opts.min_eta, opts.min_mass2, - return_mass_heavier=True) - m2_at_min_m1 = pnutils.eta_mass1_to_mass2(opts.min_eta, opts.min_mass1, - return_mass_heavier=False) - m1_at_max_m2 = pnutils.eta_mass1_to_mass2(opts.min_eta, opts.max_mass2, - return_mass_heavier=True) - m2_at_max_m1 = pnutils.eta_mass1_to_mass2(opts.min_eta, opts.max_mass1, - return_mass_heavier=False) + m1_at_min_m2 = pnutils.eta_mass1_to_mass2( + opts.min_eta, opts.min_mass2, return_mass_heavier=True + ) + m2_at_min_m1 = pnutils.eta_mass1_to_mass2( + opts.min_eta, opts.min_mass1, return_mass_heavier=False + ) + m1_at_max_m2 = pnutils.eta_mass1_to_mass2( + opts.min_eta, opts.max_mass2, return_mass_heavier=True + ) + m2_at_max_m1 = pnutils.eta_mass1_to_mass2( + opts.min_eta, opts.max_mass1, return_mass_heavier=False + ) # Check for restrictions on the maximum total mass # Are either of these possible? @@ -853,24 +1006,29 @@ def verify_mass_range_options(opts, parser, nonSpin=False): if opts.nsbh_flag: parser.error("Must supply --max_ns_spin_mag with --nsbh-flag") # Can ignore this if no NSs will be generated - elif opts.min_mass2 < (opts.ns_bh_boundary_mass or - massRangeParameters.default_nsbh_boundary_mass): - parser.error("Must supply --max-ns-spin-mag for the chosen" - " value of --min_mass2") + elif opts.min_mass2 < ( + opts.ns_bh_boundary_mass or massRangeParameters.default_nsbh_boundary_mass + ): + parser.error( + "Must supply --max-ns-spin-mag for the chosen value of --min_mass2" + ) else: opts.max_ns_spin_mag = opts.max_bh_spin_mag if opts.max_bh_spin_mag is None: if opts.nsbh_flag: parser.error("Must supply --max_bh_spin_mag with --nsbh-flag") # Can ignore this if no BHs will be generated - if opts.max_mass1 >= (opts.ns_bh_boundary_mass or - massRangeParameters.default_nsbh_boundary_mass): - parser.error("Must supply --max-bh-spin-mag for the chosen" - " value of --max_mass1") + if opts.max_mass1 >= ( + opts.ns_bh_boundary_mass or massRangeParameters.default_nsbh_boundary_mass + ): + parser.error( + "Must supply --max-bh-spin-mag for the chosen value of --max_mass1" + ) else: opts.max_bh_spin_mag = opts.max_ns_spin_mag -class massRangeParameters(object): + +class massRangeParameters: """ This class holds all of the options that are parsed in the function insert_mass_range_option_group @@ -879,58 +1037,71 @@ class massRangeParameters(object): provided on the command line """ - default_nsbh_boundary_mass = 3. - default_ns_eos = '2H' + default_nsbh_boundary_mass = 3.0 + default_ns_eos = "2H" default_delta_bh_spin = 0.1 default_delta_ns_mass = 0.1 - def __init__(self, minMass1, maxMass1, minMass2, maxMass2, - maxNSSpinMag=0, maxBHSpinMag=0, maxTotMass=None, - minTotMass=None, maxEta=None, minEta=0, - max_chirp_mass=None, min_chirp_mass=None, - ns_bh_boundary_mass=None, nsbhFlag=False, - remnant_mass_threshold=None, ns_eos=None, use_eos_max_ns_mass=False, - delta_bh_spin=None, delta_ns_mass=None): + def __init__( + self, + minMass1, + maxMass1, + minMass2, + maxMass2, + maxNSSpinMag=0, + maxBHSpinMag=0, + maxTotMass=None, + minTotMass=None, + maxEta=None, + minEta=0, + max_chirp_mass=None, + min_chirp_mass=None, + ns_bh_boundary_mass=None, + nsbhFlag=False, + remnant_mass_threshold=None, + ns_eos=None, + use_eos_max_ns_mass=False, + delta_bh_spin=None, + delta_ns_mass=None, + ): """ Initialize an instance of the massRangeParameters by providing all options directly. See the help message associated with any code that uses the metric options for more details of how to set each of these. For e.g. pycbc_aligned_stoch_bank --help """ - self.minMass1=minMass1 - self.maxMass1=maxMass1 - self.minMass2=minMass2 - self.maxMass2=maxMass2 - self.maxNSSpinMag=maxNSSpinMag - self.maxBHSpinMag=maxBHSpinMag + self.minMass1 = minMass1 + self.maxMass1 = maxMass1 + self.minMass2 = minMass2 + self.maxMass2 = maxMass2 + self.maxNSSpinMag = maxNSSpinMag + self.maxBHSpinMag = maxBHSpinMag self.minTotMass = minMass1 + minMass2 if minTotMass and (minTotMass > self.minTotMass): self.minTotMass = minTotMass self.maxTotMass = maxMass1 + maxMass2 if maxTotMass and (maxTotMass < self.maxTotMass): self.maxTotMass = maxTotMass - self.maxTotMass=maxTotMass - self.minTotMass=minTotMass + self.maxTotMass = maxTotMass + self.minTotMass = minTotMass if maxEta: - self.maxEta=maxEta + self.maxEta = maxEta else: - self.maxEta=0.25 + self.maxEta = 0.25 self.max_chirp_mass = max_chirp_mass self.min_chirp_mass = min_chirp_mass - self.minEta=minEta + self.minEta = minEta self.ns_bh_boundary_mass = ( - ns_bh_boundary_mass or self.default_nsbh_boundary_mass) - self.nsbhFlag=nsbhFlag + ns_bh_boundary_mass or self.default_nsbh_boundary_mass + ) + self.nsbhFlag = nsbhFlag self.remnant_mass_threshold = remnant_mass_threshold - self.ns_eos = ( - ns_eos or self.default_ns_eos) - self.delta_bh_spin = ( - delta_bh_spin or self.default_delta_bh_spin) - self.delta_ns_mass = ( - delta_ns_mass or self.default_delta_ns_mass) + self.ns_eos = ns_eos or self.default_ns_eos + self.delta_bh_spin = delta_bh_spin or self.default_delta_bh_spin + self.delta_ns_mass = delta_ns_mass or self.default_delta_ns_mass self.use_eos_max_ns_mass = use_eos_max_ns_mass if self.remnant_mass_threshold is not None: - if self.ns_eos != '2H': + if self.ns_eos != "2H": errMsg = """ By setting a value for --remnant-mass-threshold you have asked to filter out EM dim NS-BH templates. @@ -941,17 +1112,15 @@ def __init__(self, minMass1, maxMass1, minMass2, maxMass2, raise ValueError(errMsg) if use_eos_max_ns_mass: _, max_ns_g_mass = load_ns_sequence(self.ns_eos) - if(self.maxMass2 > max_ns_g_mass): - errMsg = """ + if self.maxMass2 > max_ns_g_mass: + errMsg = f""" The maximum NS mass supported by this EOS is - {0}. Please set --max-mass2 to this value or run + {max_ns_g_mass - 0.0000000001}. Please set --max-mass2 to this value or run without the --use-eos-max-ns-mass flag. - """.format(max_ns_g_mass-0.0000000001) + """ raise ValueError(errMsg) - self.delta_bh_spin = ( - delta_bh_spin or self.default_delta_bh_spin) - self.delta_ns_mass = ( - delta_ns_mass or self.default_delta_ns_mass) + self.delta_bh_spin = delta_bh_spin or self.default_delta_bh_spin + self.delta_ns_mass = delta_ns_mass or self.default_delta_ns_mass # FIXME: This may be inaccurate if Eta limits are given # This will not cause any problems, but maybe could be fixed. @@ -969,7 +1138,6 @@ def __init__(self, minMass1, maxMass1, minMass2, maxMass2, errMsg += "Check input options." raise ValueError(errMsg) - @classmethod def from_argparse(cls, opts, nonSpin=False): """ @@ -981,27 +1149,44 @@ def from_argparse(cls, opts, nonSpin=False): have already been called before initializing the class. """ if nonSpin: - return cls(opts.min_mass1, opts.max_mass1, opts.min_mass2, - opts.max_mass2, maxTotMass=opts.max_total_mass, - minTotMass=opts.min_total_mass, maxEta=opts.max_eta, - minEta=opts.min_eta, max_chirp_mass=opts.max_chirp_mass, - min_chirp_mass=opts.min_chirp_mass, - remnant_mass_threshold=opts.remnant_mass_threshold, - ns_eos=opts.ns_eos, use_eos_max_ns_mass=opts.use_eos_max_ns_mass, - delta_bh_spin=opts.delta_bh_spin, delta_ns_mass=opts.delta_ns_mass) - else: - return cls(opts.min_mass1, opts.max_mass1, opts.min_mass2, - opts.max_mass2, maxTotMass=opts.max_total_mass, - minTotMass=opts.min_total_mass, maxEta=opts.max_eta, - minEta=opts.min_eta, maxNSSpinMag=opts.max_ns_spin_mag, - maxBHSpinMag=opts.max_bh_spin_mag, - nsbhFlag=opts.nsbh_flag, - max_chirp_mass=opts.max_chirp_mass, - min_chirp_mass=opts.min_chirp_mass, - ns_bh_boundary_mass=opts.ns_bh_boundary_mass, - remnant_mass_threshold=opts.remnant_mass_threshold, - ns_eos=opts.ns_eos, use_eos_max_ns_mass=opts.use_eos_max_ns_mass, - delta_bh_spin=opts.delta_bh_spin, delta_ns_mass=opts.delta_ns_mass) + return cls( + opts.min_mass1, + opts.max_mass1, + opts.min_mass2, + opts.max_mass2, + maxTotMass=opts.max_total_mass, + minTotMass=opts.min_total_mass, + maxEta=opts.max_eta, + minEta=opts.min_eta, + max_chirp_mass=opts.max_chirp_mass, + min_chirp_mass=opts.min_chirp_mass, + remnant_mass_threshold=opts.remnant_mass_threshold, + ns_eos=opts.ns_eos, + use_eos_max_ns_mass=opts.use_eos_max_ns_mass, + delta_bh_spin=opts.delta_bh_spin, + delta_ns_mass=opts.delta_ns_mass, + ) + return cls( + opts.min_mass1, + opts.max_mass1, + opts.min_mass2, + opts.max_mass2, + maxTotMass=opts.max_total_mass, + minTotMass=opts.min_total_mass, + maxEta=opts.max_eta, + minEta=opts.min_eta, + maxNSSpinMag=opts.max_ns_spin_mag, + maxBHSpinMag=opts.max_bh_spin_mag, + nsbhFlag=opts.nsbh_flag, + max_chirp_mass=opts.max_chirp_mass, + min_chirp_mass=opts.min_chirp_mass, + ns_bh_boundary_mass=opts.ns_bh_boundary_mass, + remnant_mass_threshold=opts.remnant_mass_threshold, + ns_eos=opts.ns_eos, + use_eos_max_ns_mass=opts.use_eos_max_ns_mass, + delta_bh_spin=opts.delta_bh_spin, + delta_ns_mass=opts.delta_ns_mass, + ) def is_outside_range(self, mass1, mass2, spin1z, spin2z): """ @@ -1020,25 +1205,37 @@ def is_outside_range(self, mass1, mass2, spin1z, spin2z): return 1 # Spin1 test if self.nsbhFlag: - if (abs(spin1z) > self.maxBHSpinMag * 1.001): + if abs(spin1z) > self.maxBHSpinMag * 1.001: return 1 else: spin1zM = abs(spin1z) - if not( (mass1 * 1.001 > self.ns_bh_boundary_mass \ - and spin1zM <= self.maxBHSpinMag * 1.001) \ - or (mass1 < self.ns_bh_boundary_mass * 1.001 \ - and spin1zM <= self.maxNSSpinMag * 1.001)): + if not ( + ( + mass1 * 1.001 > self.ns_bh_boundary_mass + and spin1zM <= self.maxBHSpinMag * 1.001 + ) + or ( + mass1 < self.ns_bh_boundary_mass * 1.001 + and spin1zM <= self.maxNSSpinMag * 1.001 + ) + ): return 1 # Spin2 test if self.nsbhFlag: - if (abs(spin2z) > self.maxNSSpinMag * 1.001): + if abs(spin2z) > self.maxNSSpinMag * 1.001: return 1 else: spin2zM = abs(spin2z) - if not( (mass2 * 1.001 > self.ns_bh_boundary_mass \ - and spin2zM <= self.maxBHSpinMag * 1.001) \ - or (mass2 < self.ns_bh_boundary_mass * 1.001 and \ - spin2zM <= self.maxNSSpinMag * 1.001)): + if not ( + ( + mass2 * 1.001 > self.ns_bh_boundary_mass + and spin2zM <= self.maxBHSpinMag * 1.001 + ) + or ( + mass2 < self.ns_bh_boundary_mass * 1.001 + and spin2zM <= self.maxNSSpinMag * 1.001 + ) + ): return 1 # Total mass test mTot = mass1 + mass2 @@ -1055,17 +1252,16 @@ def is_outside_range(self, mass1, mass2, spin1z, spin2z): return 1 # Chirp mass test - chirp_mass = mTot * eta**(3./5.) - if self.min_chirp_mass is not None \ - and chirp_mass * 1.001 < self.min_chirp_mass: + chirp_mass = mTot * eta ** (3.0 / 5.0) + if self.min_chirp_mass is not None and chirp_mass * 1.001 < self.min_chirp_mass: return 1 - if self.max_chirp_mass is not None \ - and chirp_mass > self.max_chirp_mass * 1.001: + if self.max_chirp_mass is not None and chirp_mass > self.max_chirp_mass * 1.001: return 1 return 0 -class ethincaParameters(object): + +class ethincaParameters: """ This class holds all of the options that are parsed in the function insert_ethinca_metric_options @@ -1073,34 +1269,45 @@ class ethincaParameters(object): from the __init__ function, providing directly the options normally provided on the command line """ - def __init__(self, pnOrder, cutoff, freqStep, fLow=None, full_ethinca=False, - time_ethinca=False): + + def __init__( + self, + pnOrder, + cutoff, + freqStep, + fLow=None, + full_ethinca=False, + time_ethinca=False, + ): """ Initialize an instance of ethincaParameters by providing all options directly. See the insert_ethinca_metric_options() function for explanation or e.g. run pycbc_geom_nonspinbank --help """ - self.full_ethinca=full_ethinca - self.time_ethinca=time_ethinca - self.doEthinca= self.full_ethinca or self.time_ethinca - self.pnOrder=pnOrder - self.cutoff=cutoff - self.freqStep=freqStep + self.full_ethinca = full_ethinca + self.time_ethinca = time_ethinca + self.doEthinca = self.full_ethinca or self.time_ethinca + self.pnOrder = pnOrder + self.cutoff = cutoff + self.freqStep = freqStep # independent fLow for ethinca metric is currently not used - self.fLow=fLow + self.fLow = fLow # check that ethinca options make sense if self.full_ethinca and self.time_ethinca: err_msg = "It does not make sense to ask me to do the time " err_msg += "restricted ethinca and also the full ethinca." raise ValueError(err_msg) - if self.doEthinca and not ( - cutoff in pnutils.named_frequency_cutoffs.keys()): - raise ValueError("Need a valid cutoff formula to calculate " - "ethinca! Possible values are "+ - str(tuple(pnutils.named_frequency_cutoffs.keys()))) + if self.doEthinca and cutoff not in pnutils.named_frequency_cutoffs.keys(): + raise ValueError( + "Need a valid cutoff formula to calculate " + "ethinca! Possible values are " + + str(tuple(pnutils.named_frequency_cutoffs.keys())) + ) if self.doEthinca and not freqStep: - raise ValueError("Need to specify a cutoff frequency step to " - "calculate ethinca! (ethincaFreqStep)") + raise ValueError( + "Need to specify a cutoff frequency step to " + "calculate ethinca! (ethincaFreqStep)" + ) @classmethod def from_argparse(cls, opts): @@ -1112,58 +1319,82 @@ def from_argparse(cls, opts): verify_ethinca_metric_options have already been called before initializing the class. """ - return cls(opts.ethinca_pn_order, opts.filter_cutoff, - opts.ethinca_frequency_step, fLow=None, + return cls( + opts.ethinca_pn_order, + opts.filter_cutoff, + opts.ethinca_frequency_step, + fLow=None, full_ethinca=opts.calculate_ethinca_metric, - time_ethinca=opts.calculate_time_metric_components) + time_ethinca=opts.calculate_time_metric_components, + ) + def insert_ethinca_metric_options(parser): """ Adds the options used to calculate the ethinca metric, if required. - + Parameters - ----------- + ---------- parser : object OptionParser instance. + """ - ethincaGroup = parser.add_argument_group("Ethinca metric options", - "Options used in the calculation of Gamma metric " - "components for the ethinca coincidence test and for " - "assigning high-frequency cutoffs to templates.") + ethincaGroup = parser.add_argument_group( + "Ethinca metric options", + "Options used in the calculation of Gamma metric " + "components for the ethinca coincidence test and for " + "assigning high-frequency cutoffs to templates.", + ) ethinca_methods = ethincaGroup.add_mutually_exclusive_group() - ethinca_methods.add_argument("--calculate-time-metric-components", - action="store_true", default=False, - help="If given, the ethinca metric will be calculated " - "for only the time component, and stored in the Gamma0 " - "entry of the sngl_inspiral table. OPTIONAL, default=False") - ethinca_methods.add_argument("--calculate-ethinca-metric", - action="store_true", default=False, - help="If given, the ethinca metric will be calculated " - "and stored in the Gamma entries of the sngl_inspiral " - "table. OPTIONAL, default=False") - ethincaGroup.add_argument("--ethinca-pn-order", - default=None, choices=get_ethinca_orders(), - help="Specify a PN order to be used in calculating the " - "ethinca metric. OPTIONAL: if not specified, the same " - "order will be used as for the bank metric.") - ethincaGroup.add_argument("--filter-cutoff", - default=None, - choices=tuple(pnutils.named_frequency_cutoffs.keys()), - help="Specify an upper frequency cutoff formula for the " - "ethinca metric calculation, and for the values of f_final" - " assigned to the templates. REQUIRED if the " - "calculate-ethinca-metric option is given.") - ethincaGroup.add_argument("--ethinca-frequency-step", action="store", - type=float, default=10., - help="Control the precision of the upper frequency cutoff." - " For speed, the metric is calculated only for discrete " - "f_max values with a spacing given by this option. Each " - "template is assigned the metric for the f_max closest to " - "its analytical cutoff formula. OPTIONAL, default=10. " - "UNITS=Hz") + ethinca_methods.add_argument( + "--calculate-time-metric-components", + action="store_true", + default=False, + help="If given, the ethinca metric will be calculated " + "for only the time component, and stored in the Gamma0 " + "entry of the sngl_inspiral table. OPTIONAL, default=False", + ) + ethinca_methods.add_argument( + "--calculate-ethinca-metric", + action="store_true", + default=False, + help="If given, the ethinca metric will be calculated " + "and stored in the Gamma entries of the sngl_inspiral " + "table. OPTIONAL, default=False", + ) + ethincaGroup.add_argument( + "--ethinca-pn-order", + default=None, + choices=get_ethinca_orders(), + help="Specify a PN order to be used in calculating the " + "ethinca metric. OPTIONAL: if not specified, the same " + "order will be used as for the bank metric.", + ) + ethincaGroup.add_argument( + "--filter-cutoff", + default=None, + choices=tuple(pnutils.named_frequency_cutoffs.keys()), + help="Specify an upper frequency cutoff formula for the " + "ethinca metric calculation, and for the values of f_final" + " assigned to the templates. REQUIRED if the " + "calculate-ethinca-metric option is given.", + ) + ethincaGroup.add_argument( + "--ethinca-frequency-step", + action="store", + type=float, + default=10.0, + help="Control the precision of the upper frequency cutoff." + " For speed, the metric is calculated only for discrete " + "f_max values with a spacing given by this option. Each " + "template is assigned the metric for the f_max closest to " + "its analytical cutoff formula. OPTIONAL, default=10. " + "UNITS=Hz", + ) return ethincaGroup + def verify_ethinca_metric_options(opts, parser): """ Checks that the necessary options are given for the ethinca metric @@ -1175,20 +1406,26 @@ def verify_ethinca_metric_options(opts, parser): Result of parsing the input options with OptionParser parser : object The OptionParser instance. + """ - if opts.filter_cutoff is not None and not (opts.filter_cutoff in - pnutils.named_frequency_cutoffs.keys()): - parser.error("Need a valid cutoff formula to calculate ethinca or " - "assign filter f_final values! Possible values are " - +str(tuple(pnutils.named_frequency_cutoffs.keys()))) - if (opts.calculate_ethinca_metric or opts.calculate_time_metric_components)\ - and not opts.ethinca_frequency_step: - parser.error("Need to specify a cutoff frequency step to calculate " - "ethinca!") - if not (opts.calculate_ethinca_metric or\ - opts.calculate_time_metric_components) and opts.ethinca_pn_order: - parser.error("Can't specify an ethinca PN order if not " - "calculating ethinca metric!") + if opts.filter_cutoff is not None and opts.filter_cutoff not in pnutils.named_frequency_cutoffs.keys(): + parser.error( + "Need a valid cutoff formula to calculate ethinca or " + "assign filter f_final values! Possible values are " + + str(tuple(pnutils.named_frequency_cutoffs.keys())) + ) + if ( + opts.calculate_ethinca_metric or opts.calculate_time_metric_components + ) and not opts.ethinca_frequency_step: + parser.error("Need to specify a cutoff frequency step to calculate ethinca!") + if ( + not (opts.calculate_ethinca_metric or opts.calculate_time_metric_components) + and opts.ethinca_pn_order + ): + parser.error( + "Can't specify an ethinca PN order if not calculating ethinca metric!" + ) + def check_ethinca_against_bank_params(ethincaParams, metricParams): """ @@ -1200,16 +1437,20 @@ def check_ethinca_against_bank_params(ethincaParams, metricParams): ---------- ethincaParams: instance of ethincaParameters metricParams: instance of metricParameters + """ if ethincaParams.doEthinca: if metricParams.f0 != metricParams.fLow: - raise ValueError("If calculating ethinca metric, f0 and f-low " - "must be equal!") - if ethincaParams.fLow is not None and ( - ethincaParams.fLow != metricParams.fLow): - raise ValueError("Ethinca metric calculation does not currently " - "support a f-low value different from the bank " - "metric!") + raise ValueError( + "If calculating ethinca metric, f0 and f-low must be equal!" + ) + if ethincaParams.fLow is not None and (ethincaParams.fLow != metricParams.fLow): + raise ValueError( + "Ethinca metric calculation does not currently " + "support a f-low value different from the bank " + "metric!" + ) if ethincaParams.pnOrder is None: ethincaParams.pnOrder = metricParams.pnOrder - else: pass + else: + pass diff --git a/pycbc/tmpltbank/partitioned_bank.py b/pycbc/tmpltbank/partitioned_bank.py index 7f2c9ce3f72..1f48db9c5dd 100644 --- a/pycbc/tmpltbank/partitioned_bank.py +++ b/pycbc/tmpltbank/partitioned_bank.py @@ -15,29 +15,33 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import copy -import numpy import logging +import numpy + from pycbc.tmpltbank import coord_utils -logger = logging.getLogger('pycbc.tmpltbank.partitioned_bank') +logger = logging.getLogger("pycbc.tmpltbank.partitioned_bank") + -class PartitionedTmpltbank(object): +class PartitionedTmpltbank: """ This class is used to hold a template bank partitioned into numerous bins based on position in the Cartesian parameter space where the axes are the principal components. It can also be used to hold intermediary products used while constructing (e.g.) a stochastic template bank. """ - def __init__(self, mass_range_params, metric_params, ref_freq, - bin_spacing, bin_range_check=1): + + def __init__( + self, mass_range_params, metric_params, ref_freq, bin_spacing, bin_range_check=1 + ): """ Set up the partitioned template bank class. The combination of the reference frequency, the bin spacing and the metric dictates how the parameter space will be partitioned. Parameters - ----------- + ---------- mass_range_params : massRangeParameters object An initialized massRangeParameters object holding the details of the mass and spin ranges being considered. @@ -65,6 +69,7 @@ def __init__(self, mass_range_params, metric_params, ref_freq, When computing matches consider points in the corresponding bin and all bins +/- this value in both chi_1 and chi_2 directions. DEFAULT = 1. + """ # Flags to be used in other methods of this class. Initialized here for # simplicity @@ -77,8 +82,9 @@ def __init__(self, mass_range_params, metric_params, ref_freq, self.bin_spacing = bin_spacing # Get parameter space extent - vals = coord_utils.estimate_mass_range(1000000, mass_range_params, - metric_params, ref_freq, covary=True) + vals = coord_utils.estimate_mass_range( + 1000000, mass_range_params, metric_params, ref_freq, covary=True + ) chi1_max = vals[0].max() chi1_min = vals[0].min() chi1_diff = chi1_max - chi1_min @@ -87,10 +93,10 @@ def __init__(self, mass_range_params, metric_params, ref_freq, chi2_diff = chi2_max - chi2_min # Add a little bit extra as we may not have reached the edges. # FIXME: Maybe better to use the numerical code to find maxima here? - chi1_min = chi1_min - 0.1*chi1_diff - chi1_max = chi1_max + 0.1*chi1_diff - chi2_min = chi2_min - 0.1*chi2_diff - chi2_max = chi2_max + 0.1*chi2_diff + chi1_min = chi1_min - 0.1 * chi1_diff + chi1_max = chi1_max + 0.1 * chi1_diff + chi2_min = chi2_min - 0.1 * chi2_diff + chi2_max = chi2_max + 0.1 * chi2_diff massbank = {} bank = {} @@ -101,7 +107,7 @@ def __init__(self, mass_range_params, metric_params, ref_freq, for j in range(-2, int((chi2_max - chi2_min) // bin_spacing + 2)): bank[i][j] = [] massbank[i][j] = {} - massbank[i][j]['mass1s'] = numpy.array([]) + massbank[i][j]["mass1s"] = numpy.array([]) self.massbank = massbank self.bank = bank @@ -120,13 +126,14 @@ def __init__(self, mass_range_params, metric_params, ref_freq, self.bin_loop_order = coord_utils.outspiral_loop(self.bin_range_check) def get_point_from_bins_and_idx(self, chi1_bin, chi2_bin, idx): - """Find masses and spins given bin numbers and index. + """ + Find masses and spins given bin numbers and index. Given the chi1 bin, chi2 bin and an index, return the masses and spins of the point at that index. Will fail if no point exists there. Parameters - ----------- + ---------- chi1_bin : int The bin number for chi1. chi2_bin : int @@ -135,7 +142,7 @@ def get_point_from_bins_and_idx(self, chi1_bin, chi2_bin, idx): The index within the chi1, chi2 bin. Returns - -------- + ------- mass1 : float Mass of heavier body. mass2 : float @@ -144,15 +151,15 @@ def get_point_from_bins_and_idx(self, chi1_bin, chi2_bin, idx): Spin of heavier body. spin2z : float Spin of lighter body. + """ - mass1 = self.massbank[chi1_bin][chi2_bin]['mass1s'][idx] - mass2 = self.massbank[chi1_bin][chi2_bin]['mass2s'][idx] - spin1z = self.massbank[chi1_bin][chi2_bin]['spin1s'][idx] - spin2z = self.massbank[chi1_bin][chi2_bin]['spin2s'][idx] + mass1 = self.massbank[chi1_bin][chi2_bin]["mass1s"][idx] + mass2 = self.massbank[chi1_bin][chi2_bin]["mass2s"][idx] + spin1z = self.massbank[chi1_bin][chi2_bin]["spin1s"][idx] + spin2z = self.massbank[chi1_bin][chi2_bin]["spin2s"][idx] return mass1, mass2, spin1z, spin2z - def get_freq_map_and_normalizations(self, frequency_list, - upper_freq_formula): + def get_freq_map_and_normalizations(self, frequency_list, upper_freq_formula): """ If using the --vary-fupper capability we need to store the mapping between index and frequencies in the list. We also precalculate the @@ -160,11 +167,12 @@ def get_freq_map_and_normalizations(self, frequency_list, overlaps to account for abrupt changes in termination frequency. Parameters - ----------- + ---------- frequency_list : array of floats The frequencies for which the metric has been computed and lie within the parameter space being considered. upper_freq_formula : string + """ self.frequency_map = {} self.normalization_map = {} @@ -174,8 +182,9 @@ def get_freq_map_and_normalizations(self, frequency_list, for idx, frequency in enumerate(frequency_list): self.frequency_map[frequency] = idx - self.normalization_map[frequency] = \ - (self.metric_params.moments['I7'][frequency])**0.5 + self.normalization_map[frequency] = ( + (self.metric_params.moments["I7"][frequency]) ** 0.5 + ) def find_point_bin(self, chi_coords): """ @@ -184,16 +193,17 @@ def find_point_bin(self, chi_coords): these indices. Parameters - ----------- + ---------- chi_coords : numpy.array The position of the point in the chi coordinates. Returns - -------- + ------- chi1_bin : int Index of the chi_1 bin. chi2_bin : int Index of the chi_2 bin. + """ # Identify bin chi1_bin = int((chi_coords[0] - self.chi1_min) // self.bin_spacing) @@ -201,7 +211,6 @@ def find_point_bin(self, chi_coords): self.check_bin_existence(chi1_bin, chi2_bin) return chi1_bin, chi2_bin - def check_bin_existence(self, chi1_bin, chi2_bin): """ Given indices for bins in chi1 and chi2 space check that the bin @@ -209,29 +218,33 @@ def check_bin_existence(self, chi1_bin, chi2_bin): all bins within +/- self.bin_range_check and add if not present. Parameters - ----------- + ---------- chi1_bin : int The index of the chi1_bin to check chi2_bin : int The index of the chi2_bin to check + """ bin_range_check = self.bin_range_check # Check if this bin actually exists. If not add it - if ( (chi1_bin < self.min_chi1_bin+bin_range_check) or - (chi1_bin > self.max_chi1_bin-bin_range_check) or - (chi2_bin < self.min_chi2_bin+bin_range_check) or - (chi2_bin > self.max_chi2_bin-bin_range_check) ): - for temp_chi1 in range(chi1_bin-bin_range_check, - chi1_bin+bin_range_check+1): + if ( + (chi1_bin < self.min_chi1_bin + bin_range_check) + or (chi1_bin > self.max_chi1_bin - bin_range_check) + or (chi2_bin < self.min_chi2_bin + bin_range_check) + or (chi2_bin > self.max_chi2_bin - bin_range_check) + ): + for temp_chi1 in range( + chi1_bin - bin_range_check, chi1_bin + bin_range_check + 1 + ): if temp_chi1 not in self.massbank: self.massbank[temp_chi1] = {} self.bank[temp_chi1] = {} - for temp_chi2 in range(chi2_bin-bin_range_check, - chi2_bin+bin_range_check+1): + for temp_chi2 in range( + chi2_bin - bin_range_check, chi2_bin + bin_range_check + 1 + ): if temp_chi2 not in self.massbank[temp_chi1]: self.massbank[temp_chi1][temp_chi2] = {} - self.massbank[temp_chi1][temp_chi2]['mass1s'] =\ - numpy.array([]) + self.massbank[temp_chi1][temp_chi2]["mass1s"] = numpy.array([]) self.bank[temp_chi1][temp_chi2] = [] def calc_point_distance(self, chi_coords): @@ -240,17 +253,18 @@ def calc_point_distance(self, chi_coords): distance. Parameters - ----------- + ---------- chi_coords : numpy.array The position of the point in the chi coordinates. Returns - -------- + ------- min_dist : float The smallest **SQUARED** metric distance between the test point and the bank. indexes : The chi1_bin, chi2_bin and position within that bin at which the closest matching point lies. + """ chi1_bin, chi2_bin = self.find_point_bin(chi_coords) min_dist = 1000000000 @@ -258,8 +272,7 @@ def calc_point_distance(self, chi_coords): for chi1_bin_offset, chi2_bin_offset in self.bin_loop_order: curr_chi1_bin = chi1_bin + chi1_bin_offset curr_chi2_bin = chi2_bin + chi2_bin_offset - for idx, bank_chis in \ - enumerate(self.bank[curr_chi1_bin][curr_chi2_bin]): + for idx, bank_chis in enumerate(self.bank[curr_chi1_bin][curr_chi2_bin]): dist = coord_utils.calc_point_dist(chi_coords, bank_chis) if dist < min_dist: min_dist = dist @@ -272,7 +285,7 @@ def test_point_distance(self, chi_coords, distance_threshold): than the supplied distance theshold. Parameters - ----------- + ---------- chi_coords : numpy.array The position of the point in the chi coordinates. distance_threshold : float @@ -281,7 +294,7 @@ def test_point_distance(self, chi_coords, distance_threshold): use 1 - 0.97 = 0.03 for this value. Returns - -------- + ------- Boolean True if point is within the distance threshold. False if not. @@ -294,8 +307,7 @@ def test_point_distance(self, chi_coords, distance_threshold): dist = coord_utils.calc_point_dist(chi_coords, bank_chis) if dist < distance_threshold: return True - else: - return False + return False def calc_point_distance_vary(self, chi_coords, point_fupper, mus): """ @@ -305,7 +317,7 @@ def calc_point_distance_vary(self, chi_coords, point_fupper, mus): change a lot. Parameters - ----------- + ---------- chi_coords : numpy.array The position of the point in the chi coordinates. point_fupper : float @@ -317,12 +329,13 @@ def calc_point_distance_vary(self, chi_coords, point_fupper, mus): each value of the upper frequency cutoff. Returns - -------- + ------- min_dist : float The smallest **SQUARED** metric distance between the test point and the bank. indexes : The chi1_bin, chi2_bin and position within that bin at which the closest matching point lies. + """ chi1_bin, chi2_bin = self.find_point_bin(chi_coords) min_dist = 1000000000 @@ -332,30 +345,28 @@ def calc_point_distance_vary(self, chi_coords, point_fupper, mus): curr_chi2_bin = chi2_bin + chi2_bin_offset # No points = Next iteration curr_bank = self.massbank[curr_chi1_bin][curr_chi2_bin] - if not curr_bank['mass1s'].size: + if not curr_bank["mass1s"].size: continue # *NOT* the same of .min and .max - f_upper = numpy.minimum(point_fupper, curr_bank['freqcuts']) - f_other = numpy.maximum(point_fupper, curr_bank['freqcuts']) + f_upper = numpy.minimum(point_fupper, curr_bank["freqcuts"]) + f_other = numpy.maximum(point_fupper, curr_bank["freqcuts"]) # NOTE: freq_idxes is a vector! freq_idxes = numpy.array([self.frequency_map[f] for f in f_upper]) # vecs1 gives a 2x2 vector: idx0 = stored index, idx1 = mu index vecs1 = mus[freq_idxes, :] # vecs2 gives a 2x2 vector: idx0 = stored index, idx1 = mu index range_idxes = numpy.arange(len(freq_idxes)) - vecs2 = curr_bank['mus'][range_idxes, freq_idxes, :] + vecs2 = curr_bank["mus"][range_idxes, freq_idxes, :] # Now do the sums - dists = (vecs1 - vecs2)*(vecs1 - vecs2) + dists = (vecs1 - vecs2) * (vecs1 - vecs2) # This reduces to 1D: idx = stored index dists = numpy.sum(dists, axis=1) - norm_upper = numpy.array([self.normalization_map[f] \ - for f in f_upper]) - norm_other = numpy.array([self.normalization_map[f] \ - for f in f_other]) + norm_upper = numpy.array([self.normalization_map[f] for f in f_upper]) + norm_other = numpy.array([self.normalization_map[f] for f in f_other]) norm_fac = norm_upper / norm_other - renormed_dists = 1 - (1 - dists)*norm_fac + renormed_dists = 1 - (1 - dists) * norm_fac curr_min_dist = renormed_dists.min() if curr_min_dist < min_dist: min_dist = curr_min_dist @@ -363,8 +374,9 @@ def calc_point_distance_vary(self, chi_coords, point_fupper, mus): return min_dist, indexes - def test_point_distance_vary(self, chi_coords, point_fupper, mus, - distance_threshold): + def test_point_distance_vary( + self, chi_coords, point_fupper, mus, distance_threshold + ): """ Test if distance between point and the bank is greater than distance threshold while allowing the metric to @@ -373,7 +385,7 @@ def test_point_distance_vary(self, chi_coords, point_fupper, mus, change a lot. Parameters - ----------- + ---------- chi_coords : numpy.array The position of the point in the chi coordinates. point_fupper : float @@ -389,9 +401,10 @@ def test_point_distance_vary(self, chi_coords, point_fupper, mus, use 1 - 0.97 = 0.03 for this value. Returns - -------- + ------- Boolean True if point is within the distance threshold. False if not. + """ chi1_bin, chi2_bin = self.find_point_bin(chi_coords) for chi1_bin_offset, chi2_bin_offset in self.bin_loop_order: @@ -399,41 +412,39 @@ def test_point_distance_vary(self, chi_coords, point_fupper, mus, curr_chi2_bin = chi2_bin + chi2_bin_offset # No points = Next iteration curr_bank = self.massbank[curr_chi1_bin][curr_chi2_bin] - if not curr_bank['mass1s'].size: + if not curr_bank["mass1s"].size: continue # *NOT* the same of .min and .max - f_upper = numpy.minimum(point_fupper, curr_bank['freqcuts']) - f_other = numpy.maximum(point_fupper, curr_bank['freqcuts']) + f_upper = numpy.minimum(point_fupper, curr_bank["freqcuts"]) + f_other = numpy.maximum(point_fupper, curr_bank["freqcuts"]) # NOTE: freq_idxes is a vector! freq_idxes = numpy.array([self.frequency_map[f] for f in f_upper]) # vecs1 gives a 2x2 vector: idx0 = stored index, idx1 = mu index vecs1 = mus[freq_idxes, :] # vecs2 gives a 2x2 vector: idx0 = stored index, idx1 = mu index range_idxes = numpy.arange(len(freq_idxes)) - vecs2 = curr_bank['mus'][range_idxes,freq_idxes,:] + vecs2 = curr_bank["mus"][range_idxes, freq_idxes, :] # Now do the sums - dists = (vecs1 - vecs2)*(vecs1 - vecs2) + dists = (vecs1 - vecs2) * (vecs1 - vecs2) # This reduces to 1D: idx = stored index dists = numpy.sum(dists, axis=1) # I wonder if this line actually speeds things up? if (dists > distance_threshold).all(): continue # This is only needed for close templates, should we prune? - norm_upper = numpy.array([self.normalization_map[f] \ - for f in f_upper]) - norm_other = numpy.array([self.normalization_map[f] \ - for f in f_other]) + norm_upper = numpy.array([self.normalization_map[f] for f in f_upper]) + norm_other = numpy.array([self.normalization_map[f] for f in f_other]) norm_fac = norm_upper / norm_other - renormed_dists = 1 - (1 - dists)*norm_fac + renormed_dists = 1 - (1 - dists) * norm_fac if (renormed_dists < distance_threshold).any(): return True - else: - return False + return False - def add_point_by_chi_coords(self, chi_coords, mass1, mass2, spin1z, spin2z, - point_fupper=None, mus=None): + def add_point_by_chi_coords( + self, chi_coords, mass1, mass2, spin1z, spin2z, point_fupper=None, mus=None + ): """ Add a point to the partitioned template bank. The point_fupper and mus kwargs must be provided for all templates if the vary fupper capability @@ -443,7 +454,7 @@ def add_point_by_chi_coords(self, chi_coords, mass1, mass2, spin1z, spin2z, add_point_by_masses, which will do translations and then call this. Parameters - ----------- + ---------- chi_coords : numpy.array The position of the point in the chi coordinates. mass1 : float @@ -460,45 +471,51 @@ def add_point_by_chi_coords(self, chi_coords, mass1, mass2, spin1z, spin2z, A 2D array where idx 0 holds the upper frequency cutoff and idx 1 holds the coordinates in the [not covaried] mu parameter space for each value of the upper frequency cutoff. + """ chi1_bin, chi2_bin = self.find_point_bin(chi_coords) self.bank[chi1_bin][chi2_bin].append(copy.deepcopy(chi_coords)) curr_bank = self.massbank[chi1_bin][chi2_bin] - if curr_bank['mass1s'].size: - curr_bank['mass1s'] = numpy.append(curr_bank['mass1s'], - numpy.array([mass1])) - curr_bank['mass2s'] = numpy.append(curr_bank['mass2s'], - numpy.array([mass2])) - curr_bank['spin1s'] = numpy.append(curr_bank['spin1s'], - numpy.array([spin1z])) - curr_bank['spin2s'] = numpy.append(curr_bank['spin2s'], - numpy.array([spin2z])) + if curr_bank["mass1s"].size: + curr_bank["mass1s"] = numpy.append( + curr_bank["mass1s"], numpy.array([mass1]) + ) + curr_bank["mass2s"] = numpy.append( + curr_bank["mass2s"], numpy.array([mass2]) + ) + curr_bank["spin1s"] = numpy.append( + curr_bank["spin1s"], numpy.array([spin1z]) + ) + curr_bank["spin2s"] = numpy.append( + curr_bank["spin2s"], numpy.array([spin2z]) + ) if point_fupper is not None: - curr_bank['freqcuts'] = numpy.append(curr_bank['freqcuts'], - numpy.array([point_fupper])) + curr_bank["freqcuts"] = numpy.append( + curr_bank["freqcuts"], numpy.array([point_fupper]) + ) # Mus needs to append onto axis 0. See below for contents of # the mus variable if mus is not None: - curr_bank['mus'] = numpy.append(curr_bank['mus'], - numpy.array([mus[:,:]]), axis=0) + curr_bank["mus"] = numpy.append( + curr_bank["mus"], numpy.array([mus[:, :]]), axis=0 + ) else: - curr_bank['mass1s'] = numpy.array([mass1]) - curr_bank['mass2s'] = numpy.array([mass2]) - curr_bank['spin1s'] = numpy.array([spin1z]) - curr_bank['spin2s'] = numpy.array([spin2z]) + curr_bank["mass1s"] = numpy.array([mass1]) + curr_bank["mass2s"] = numpy.array([mass2]) + curr_bank["spin1s"] = numpy.array([spin1z]) + curr_bank["spin2s"] = numpy.array([spin2z]) if point_fupper is not None: - curr_bank['freqcuts'] = numpy.array([point_fupper]) + curr_bank["freqcuts"] = numpy.array([point_fupper]) # curr_bank['mus'] is a 3D array # NOTE: mu relates to the non-covaried Cartesian coordinate system # Axis 0: Template index # Axis 1: Frequency cutoff index # Axis 2: Mu coordinate index if mus is not None: - curr_bank['mus'] = numpy.array([mus[:,:]]) + curr_bank["mus"] = numpy.array([mus[:, :]]) - def add_point_by_masses(self, mass1, mass2, spin1z, spin2z, - vary_fupper=False): + def add_point_by_masses(self, mass1, mass2, spin1z, spin2z, vary_fupper=False): """ Add a point to the template bank. This differs from add point to bank as it assumes that the chi coordinates and the products needed to use @@ -510,7 +527,7 @@ def add_point_by_masses(self, mass1, mass2, spin1z, spin2z, not do for speed concerns. Parameters - ----------- + ---------- mass1 : float Mass of the heavier body mass2 : float @@ -519,6 +536,7 @@ def add_point_by_masses(self, mass1, mass2, spin1z, spin2z, Spin of the heavier body spin2z : float Spin of the lighter body + """ # Test that masses are the expected way around (ie. mass1 > mass2) if mass2 > mass1: @@ -531,46 +549,50 @@ def add_point_by_masses(self, mass1, mass2, spin1z, spin2z, self.spin_warning_given = True # These that masses obey the restrictions of mass_range_params - if self.mass_range_params.is_outside_range(mass1, mass2, spin1z, - spin2z): + if self.mass_range_params.is_outside_range(mass1, mass2, spin1z, spin2z): err_msg = "Point with masses given by " - err_msg += "%f %f %f %f " %(mass1, mass2, spin1z, spin2z) + err_msg += "%f %f %f %f " % (mass1, mass2, spin1z, spin2z) err_msg += "(mass1, mass2, spin1z, spin2z) is not consistent " err_msg += "with the provided command-line restrictions on masses " err_msg += "and spins." raise ValueError(err_msg) # Get chi coordinates - chi_coords = coord_utils.get_cov_params(mass1, mass2, spin1z, spin2z, - self.metric_params, - self.ref_freq) + chi_coords = coord_utils.get_cov_params( + mass1, mass2, spin1z, spin2z, self.metric_params, self.ref_freq + ) # Get mus and best fupper for this point, if needed if vary_fupper: mass_dict = {} - mass_dict['m1'] = numpy.array([mass1]) - mass_dict['m2'] = numpy.array([mass2]) - mass_dict['s1z'] = numpy.array([spin1z]) - mass_dict['s2z'] = numpy.array([spin2z]) + mass_dict["m1"] = numpy.array([mass1]) + mass_dict["m2"] = numpy.array([mass2]) + mass_dict["s1z"] = numpy.array([spin1z]) + mass_dict["s2z"] = numpy.array([spin2z]) freqs = numpy.array(list(self.frequency_map.keys()), dtype=float) - freq_cutoff = coord_utils.return_nearest_cutoff(\ - self.upper_freq_formula, mass_dict, freqs) + freq_cutoff = coord_utils.return_nearest_cutoff( + self.upper_freq_formula, mass_dict, freqs + ) freq_cutoff = freq_cutoff[0] - lambdas = coord_utils.get_chirp_params\ - (mass1, mass2, spin1z, spin2z, self.metric_params.f0, - self.metric_params.pnOrder) + lambdas = coord_utils.get_chirp_params( + mass1, + mass2, + spin1z, + spin2z, + self.metric_params.f0, + self.metric_params.pnOrder, + ) mus = [] for freq in self.frequency_map: - mus.append(coord_utils.get_mu_params(lambdas, - self.metric_params, freq) ) + mus.append(coord_utils.get_mu_params(lambdas, self.metric_params, freq)) mus = numpy.array(mus) else: - freq_cutoff=None - mus=None - - self.add_point_by_chi_coords(chi_coords, mass1, mass2, spin1z, spin2z, - point_fupper=freq_cutoff, mus=mus) + freq_cutoff = None + mus = None + self.add_point_by_chi_coords( + chi_coords, mass1, mass2, spin1z, spin2z, point_fupper=freq_cutoff, mus=mus + ) def add_tmpltbank_from_xml_table(self, sngl_table, vary_fupper=False): """ @@ -578,16 +600,22 @@ def add_tmpltbank_from_xml_table(self, sngl_table, vary_fupper=False): into the partitioned template bank object. Parameters - ----------- + ---------- sngl_table : sngl_inspiral_table List of sngl_inspiral templates. vary_fupper : False If given also include the additional information needed to compute distances with a varying upper frequency cutoff. + """ for sngl in sngl_table: - self.add_point_by_masses(sngl.mass1, sngl.mass2, sngl.spin1z, - sngl.spin2z, vary_fupper=vary_fupper) + self.add_point_by_masses( + sngl.mass1, + sngl.mass2, + sngl.spin1z, + sngl.spin2z, + vary_fupper=vary_fupper, + ) def add_tmpltbank_from_hdf_file(self, hdf_fp, vary_fupper=False): """ @@ -596,23 +624,30 @@ def add_tmpltbank_from_hdf_file(self, hdf_fp, vary_fupper=False): object. Parameters - ----------- + ---------- hdf_fp : h5py.File object The template bank in HDF5 format. vary_fupper : False If given also include the additional information needed to compute distances with a varying upper frequency cutoff. + """ - mass1s = hdf_fp['mass1'][:] - mass2s = hdf_fp['mass2'][:] - spin1zs = hdf_fp['spin1z'][:] - spin2zs = hdf_fp['spin2z'][:] + mass1s = hdf_fp["mass1"][:] + mass2s = hdf_fp["mass2"][:] + spin1zs = hdf_fp["spin1z"][:] + spin2zs = hdf_fp["spin2z"][:] for idx in range(len(mass1s)): - self.add_point_by_masses(mass1s[idx], mass2s[idx], spin1zs[idx], - spin2zs[idx], vary_fupper=vary_fupper) + self.add_point_by_masses( + mass1s[idx], + mass2s[idx], + spin1zs[idx], + spin2zs[idx], + vary_fupper=vary_fupper, + ) def output_all_points(self): - """Return all points in the bank. + """ + Return all points in the bank. Return all points in the bank as lists of m1, m2, spin1z, spin2z. @@ -626,6 +661,7 @@ def output_all_points(self): List of spin1z values. spin2z : list List of spin2z values. + """ mass1 = [] mass2 = [] @@ -633,11 +669,11 @@ def output_all_points(self): spin2z = [] for i in self.massbank.keys(): for j in self.massbank[i].keys(): - for k in range(len(self.massbank[i][j]['mass1s'])): + for k in range(len(self.massbank[i][j]["mass1s"])): curr_bank = self.massbank[i][j] - mass1.append(curr_bank['mass1s'][k]) - mass2.append(curr_bank['mass2s'][k]) - spin1z.append(curr_bank['spin1s'][k]) - spin2z.append(curr_bank['spin2s'][k]) + mass1.append(curr_bank["mass1s"][k]) + mass2.append(curr_bank["mass2s"][k]) + spin1z.append(curr_bank["spin1s"][k]) + spin2z.append(curr_bank["spin2s"][k]) return mass1, mass2, spin1z, spin2z diff --git a/pycbc/tmpltbank/sky_grid.py b/pycbc/tmpltbank/sky_grid.py index 26f72483f2d..63a24d98ffc 100644 --- a/pycbc/tmpltbank/sky_grid.py +++ b/pycbc/tmpltbank/sky_grid.py @@ -1,19 +1,21 @@ -"""Functionality for handling grids of points in the sky for coherent SNR +""" +Functionality for handling grids of points in the sky for coherent SNR calculation via `pycbc_multi_inspiral`. The main operation to be performed on these points is calculating the antenna pattern functions and time delays from the Earth center for a network of detectors. """ -import numpy as np import h5py +import numpy as np -from pycbc.detector import Detector from pycbc.conversions import ensurearray +from pycbc.detector import Detector class SkyGrid: def __init__(self, ra, dec, detectors, ref_gps_time): - """Initialize a sky grid from a list of RA/dec coordinates. + """ + Initialize a sky grid from a list of RA/dec coordinates. Parameters ---------- @@ -31,15 +33,16 @@ def __init__(self, ra, dec, detectors, ref_gps_time): Reference GPS time associated with the sky grid. This will be used when calculating the antenna pattern functions and time delays from Earth center. + """ # We store the points in a 2D array internally, first dimension runs # over the list of points, second dimension is RA/dec. # Question: should we use Astropy sky positions instead? - ra, dec, _ = ensurearray(ra,dec) + ra, dec, _ = ensurearray(ra, dec) if (ra < 0).any() or (ra > 2 * np.pi).any(): - raise ValueError('RA must be in the range [0,2π]') - if (dec < -np.pi/2).any() or (dec > np.pi/2).any(): - raise ValueError('DEC must be in the range [-π/2, π/2]') + raise ValueError("RA must be in the range [0,2π]") + if (dec < -np.pi / 2).any() or (dec > np.pi / 2).any(): + raise ValueError("DEC must be in the range [-π/2, π/2]") self.positions = np.vstack([ra, dec]).T self.detectors = sorted(detectors) self.ref_gps_time = ref_gps_time @@ -59,20 +62,23 @@ def ras(self): @property def decs(self): - """Returns all declinations in radians, where π/2 is the North pole, - -π/2 is the South pole, and 0 is the celestial equator.""" + """ + Returns all declinations in radians, where π/2 is the North pole, + -π/2 is the South pole, and 0 is the celestial equator. + """ return self.positions[:, 1] @classmethod def from_cli(cls, cli_parser, cli_args): - """Initialize a sky grid from command-line interface, via argparse + """ + Initialize a sky grid from command-line interface, via argparse objects. """ if cli_args.sky_grid is not None: if cli_args.ra is not None or cli_args.dec is not None: cli_parser.error( - 'Please provide either a sky grid via --sky-grid or a ' - 'single sky position via --ra and --dec, not both' + "Please provide either a sky grid via --sky-grid or a " + "single sky position via --ra and --dec, not both" ) return cls.read_from_file(cli_args.sky_grid) if cli_args.ra is not None and cli_args.dec is not None: @@ -80,37 +86,38 @@ def from_cli(cls, cli_parser, cli_args): [cli_args.ra], [cli_args.dec], cli_args.instruments, - cli_args.trigger_time + cli_args.trigger_time, ) cli_parser.error( - 'Please specify a sky grid via --sky-grid or a single sky ' - 'position via --ra and --dec' + "Please specify a sky grid via --sky-grid or a single sky " + "position via --ra and --dec" ) @classmethod def read_from_file(cls, path): """Initialize a sky grid from a given HDF5 file.""" - with h5py.File(path, 'r') as hf: - ra = hf['ra'][:] - dec = hf['dec'][:] - detectors = hf.attrs['detectors'] - ref_gps_time = hf.attrs['ref_gps_time'] + with h5py.File(path, "r") as hf: + ra = hf["ra"][:] + dec = hf["dec"][:] + detectors = hf.attrs["detectors"] + ref_gps_time = hf.attrs["ref_gps_time"] return cls(ra, dec, detectors, ref_gps_time) def write_to_file(self, path, extra_attrs=None, extra_datasets=None): """Writes a sky grid to an HDF5 file.""" - with h5py.File(path, 'w') as hf: - hf['ra'] = self.ras - hf['dec'] = self.decs - hf.attrs['detectors'] = self.detectors - hf.attrs['ref_gps_time'] = self.ref_gps_time - for attribute in (extra_attrs or {}): + with h5py.File(path, "w") as hf: + hf["ra"] = self.ras + hf["dec"] = self.decs + hf.attrs["detectors"] = self.detectors + hf.attrs["ref_gps_time"] = self.ref_gps_time + for attribute in extra_attrs or {}: hf.attrs[attribute] = extra_attrs[attribute] - for dataset in (extra_datasets or {}): + for dataset in extra_datasets or {}: hf[dataset] = extra_datasets[dataset] def calculate_antenna_patterns(self): - """Calculate the antenna pattern functions at each point in the grid + """ + Calculate the antenna pattern functions at each point in the grid for the list of GW detectors specified at instantiation. Return a dict, keyed by detector name, whose items are 2-dimensional Numpy arrays. The first dimension of these arrays runs over the sky grid, and the @@ -127,7 +134,8 @@ def calculate_antenna_patterns(self): return result def calculate_time_delays(self): - """Calculate the time delays from the Earth center to each GW detector + """ + Calculate the time delays from the Earth center to each GW detector specified at instantiation, for each point in the grid. Return a dict, keyed by detector name, whose items are 1-dimensional Numpy arrays containing the time delays for each sky point. @@ -143,4 +151,4 @@ def calculate_time_delays(self): return result -__all__ = ['SkyGrid'] +__all__ = ["SkyGrid"] diff --git a/pycbc/transforms.py b/pycbc/transforms.py index 7af5c9a8648..a35ad990ea3 100644 --- a/pycbc/transforms.py +++ b/pycbc/transforms.py @@ -16,23 +16,21 @@ This modules provides classes and functions for transforming parameters. """ -import os import logging +import os + import numpy -from pycbc import conversions -from pycbc import coordinates -from pycbc import cosmology -from pycbc.io import record -from pycbc.waveform import parameters +from pycbc import VARARGS_DELIM, conversions, coordinates, cosmology from pycbc.boundaries import Bounds -from pycbc import VARARGS_DELIM +from pycbc.io import record from pycbc.pnutils import jframe_to_l0frame +from pycbc.waveform import parameters -logger = logging.getLogger('pycbc.transforms') +logger = logging.getLogger("pycbc.transforms") -class BaseTransform(object): +class BaseTransform: """A base class for transforming between two sets of parameters.""" name = None @@ -52,7 +50,8 @@ def transform(self, maps): raise NotImplementedError("Not added.") def inverse_transform(self, maps): - """The inverse conversions of transform. This function transforms from + """ + The inverse conversions of transform. This function transforms from outputs to inputs. """ raise NotImplementedError("Not added.") @@ -67,7 +66,8 @@ def inverse_jacobian(self, maps): @staticmethod def format_output(old_maps, new_maps): - """This function takes the returned dict from `transform` and converts + """ + This function takes the returned dict from `transform` and converts it to the same datatype as the input. Parameters @@ -81,8 +81,8 @@ def format_output(old_maps, new_maps): ------- {FieldArray, dict} The old_maps object with new keys from new_maps. - """ + """ # if input is FieldArray then return FieldArray if isinstance(old_maps, record.FieldArray): keys = new_maps.keys() @@ -95,19 +95,18 @@ def format_output(old_maps, new_maps): return old_maps # if input is dict then return dict - elif isinstance(old_maps, dict): + if isinstance(old_maps, dict): out = old_maps.copy() out.update(new_maps) return out # else error - else: - raise TypeError("Input type must be FieldArray or dict.") + raise TypeError("Input type must be FieldArray or dict.") @classmethod - def from_config(cls, cp, section, outputs, - skip_opts=None, additional_opts=None): - """Initializes a transform from the given section. + def from_config(cls, cp, section, outputs, skip_opts=None, additional_opts=None): + """ + Initializes a transform from the given section. Parameters ---------- @@ -130,6 +129,7 @@ def from_config(cls, cp, section, outputs, ------- cls An instance of the class. + """ tag = outputs if skip_opts is None: @@ -158,13 +158,14 @@ def from_config(cls, cp, section, outputs, # check that the outputs matches if outputs - out.outputs != set() or out.outputs - outputs != set(): raise ValueError( - "outputs of class do not match outputs specified " "in section" + "outputs of class do not match outputs specified in section" ) return out class CustomTransform(BaseTransform): - """Allows for any transform to be defined. + """ + Allows for any transform to be defined. Parameters ---------- @@ -199,8 +200,7 @@ class CustomTransform(BaseTransform): name = "custom" - def __init__(self, input_args, output_args, transform_functions, - jacobian=None): + def __init__(self, input_args, output_args, transform_functions, jacobian=None): if isinstance(input_args, str): input_args = [input_args] if isinstance(output_args, str): @@ -221,7 +221,8 @@ def _createscratch(self, shape=1): ) def _copytoscratch(self, maps): - """Copies the data in maps to the scratch space. + """ + Copies the data in maps to the scratch space. If the maps contain arrays that are not the same shape as the scratch space, a new scratch space will be created. @@ -252,7 +253,8 @@ def _getslice(self, maps): return getslice def transform(self, maps): - """Applies the transform functions to the given maps object. + """ + Applies the transform functions to the given maps object. Parameters ---------- @@ -264,6 +266,7 @@ def transform(self, maps): A map object containing the transformed variables, along with the original variables. The type of the output will be the same as the input. + """ if self.transform_functions is None: raise NotImplementedError("no transform function(s) provided") @@ -290,7 +293,8 @@ def jacobian(self, maps): @classmethod def from_config(cls, cp, section, outputs): - """Loads a CustomTransform from the given config file. + """ + Loads a CustomTransform from the given config file. Example section: @@ -305,8 +309,7 @@ def from_config(cls, cp, section, outputs): """ tag = outputs outputs = set(outputs.split(VARARGS_DELIM)) - inputs = map(str.strip, - cp.get_opt_tag(section, "inputs", tag).split(",")) + inputs = map(str.strip, cp.get_opt_tag(section, "inputs", tag).split(",")) # get the functions for each output transform_functions = {} for var in outputs: @@ -322,7 +325,8 @@ def from_config(cls, cp, section, outputs): class CustomTransformMultiOutputs(CustomTransform): - """Allows for any transform to be defined. Based on CustomTransform, + """ + Allows for any transform to be defined. Based on CustomTransform, but also supports multi-returning value functions. Parameters @@ -337,26 +341,31 @@ class CustomTransformMultiOutputs(CustomTransform): jacobian : str, optional String giving a jacobian function. The function must be in terms of the input arguments. + """ name = "custom_multi" - def __init__(self, input_args, output_args, transform_functions, - jacobian=None): - super(CustomTransformMultiOutputs, self).__init__( - input_args, output_args, transform_functions, jacobian) + def __init__(self, input_args, output_args, transform_functions, jacobian=None): + super().__init__( + input_args, output_args, transform_functions, jacobian + ) def transform(self, maps): - """Applies the transform functions to the given maps object. + """ + Applies the transform functions to the given maps object. + Parameters ---------- maps : dict, or FieldArray + Returns ------- dict or FieldArray A map object containing the transformed variables, along with the original variables. The type of the output will be the same as the input. + """ if self.transform_functions is None: raise NotImplementedError("no transform function(s) provided") @@ -368,16 +377,17 @@ def transform(self, maps): # func[0] is the function itself, func[1] is the index, # this supports multiple returning values function out = { - p: self._scratch[func[0]][func[1]][getslice] if - len(self._scratch[func[0]]) > 1 else - self._scratch[func[0]][getslice] - for p, func in self.transform_functions.items() - } + p: self._scratch[func[0]][func[1]][getslice] + if len(self._scratch[func[0]]) > 1 + else self._scratch[func[0]][getslice] + for p, func in self.transform_functions.items() + } return self.format_output(maps, out) @classmethod def from_config(cls, cp, section, outputs): - """Loads a CustomTransformMultiOutputs from the given config file. + """ + Loads a CustomTransformMultiOutputs from the given config file. Example section: @@ -392,8 +402,7 @@ def from_config(cls, cp, section, outputs): tag = outputs outputs = list(outputs.split(VARARGS_DELIM)) all_vars = ", ".join(outputs) - inputs = map(str.strip, - cp.get_opt_tag(section, "inputs", tag).split(",")) + inputs = map(str.strip, cp.get_opt_tag(section, "inputs", tag).split(",")) # get the functions for each output transform_functions = {} output_index = slice(None, None, None) @@ -403,7 +412,7 @@ def from_config(cls, cp, section, outputs): func = cp.get_opt_tag(section, var, tag) except Exception: func = cp.get_opt_tag(section, all_vars, tag) - output_index = slice(outputs.index(var), outputs.index(var)+1) + output_index = slice(outputs.index(var), outputs.index(var) + 1) transform_functions[var] = [func, output_index] s = "-".join([section, tag]) if cp.has_option(s, "jacobian"): @@ -444,10 +453,11 @@ def __init__( self.q_param = q_param self._inputs = [self.mchirp_param, self.q_param] self._outputs = [self.mass1_param, self.mass2_param] - super(MchirpQToMass1Mass2, self).__init__() + super().__init__() def transform(self, maps): - """This function transforms from chirp mass and mass ratio to component + """ + This function transforms from chirp mass and mass ratio to component masses. Parameters @@ -470,6 +480,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} out[self.mass1_param] = conversions.mass1_from_mchirp_q( @@ -481,7 +492,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms from component masses to chirp mass and + """ + This function transforms from component masses to chirp mass and mass ratio. Parameters @@ -504,6 +516,7 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} m1 = maps[self.mass1_param] @@ -513,20 +526,22 @@ def inverse_transform(self, maps): return self.format_output(maps, out) def jacobian(self, maps): - """Returns the Jacobian for transforming mchirp and q to mass1 and + """ + Returns the Jacobian for transforming mchirp and q to mass1 and mass2. """ mchirp = maps[self.mchirp_param] q = maps[self.q_param] - return mchirp * ((1.0 + q) / q ** 3.0) ** (2.0 / 5) + return mchirp * ((1.0 + q) / q**3.0) ** (2.0 / 5) def inverse_jacobian(self, maps): - """Returns the Jacobian for transforming mass1 and mass2 to + """ + Returns the Jacobian for transforming mass1 and mass2 to mchirp and q. """ m1 = maps[self.mass1_param] m2 = maps[self.mass2_param] - return conversions.mchirp_from_mass1_mass2(m1, m2) / m2 ** 2.0 + return conversions.mchirp_from_mass1_mass2(m1, m2) / m2**2.0 class MchirpEtaToMass1Mass2(BaseTransform): @@ -537,7 +552,8 @@ class MchirpEtaToMass1Mass2(BaseTransform): _outputs = [parameters.mass1, parameters.mass2] def transform(self, maps): - """This function transforms from chirp mass and symmetric mass ratio to + """ + This function transforms from chirp mass and symmetric mass ratio to component masses. Parameters @@ -560,6 +576,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} out[parameters.mass1] = conversions.mass1_from_mchirp_eta( @@ -571,7 +588,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms from component masses to chirp mass and + """ + This function transforms from component masses to chirp mass and symmetric mass ratio. Parameters @@ -594,6 +612,7 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} m1 = maps[parameters.mass1] @@ -603,7 +622,8 @@ def inverse_transform(self, maps): return self.format_output(maps, out) def jacobian(self, maps): - """Returns the Jacobian for transforming mchirp and eta to mass1 and + """ + Returns the Jacobian for transforming mchirp and eta to mass1 and mass2. """ mchirp = maps[parameters.mchirp] @@ -613,7 +633,8 @@ def jacobian(self, maps): return mchirp * (m1 - m2) / (m1 + m2) ** 3 def inverse_jacobian(self, maps): - """Returns the Jacobian for transforming mass1 and mass2 to + """ + Returns the Jacobian for transforming mass1 and mass2 to mchirp and eta. """ m1 = maps[parameters.mass1] @@ -636,7 +657,8 @@ def __init__(self, ref_mass=1.4): self.ref_mass = ref_mass def transform(self, maps): - """This function transforms from chirp distance to luminosity distance, + """ + This function transforms from chirp distance to luminosity distance, given the chirp mass. Parameters @@ -658,6 +680,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} out[parameters.distance] = conversions.distance_from_chirp_distance_mchirp( @@ -668,7 +691,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms from luminosity distance to chirp distance, + """ + This function transforms from luminosity distance to chirp distance, given the chirp mass. Parameters @@ -690,6 +714,7 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} out[parameters.chirp_distance] = conversions.chirp_distance( @@ -698,7 +723,8 @@ def inverse_transform(self, maps): return self.format_output(maps, out) def jacobian(self, maps): - """Returns the Jacobian for transforming chirp distance to + """ + Returns the Jacobian for transforming chirp distance to luminosity distance, given the chirp mass. """ ref_mass = 1.4 @@ -706,7 +732,8 @@ def jacobian(self, maps): return (2.0 ** (-1.0 / 5) * self.ref_mass / mchirp) ** (-5.0 / 6) def inverse_jacobian(self, maps): - """Returns the Jacobian for transforming luminosity distance to + """ + Returns the Jacobian for transforming luminosity distance to chirp distance, given the chirp mass. """ ref_mass = 1.4 @@ -715,22 +742,39 @@ def inverse_jacobian(self, maps): class AlignTotalSpin(BaseTransform): - """Converts angles from total angular momentum J frame to orbital angular - momentum L (waveform) frame""" + """ + Converts angles from total angular momentum J frame to orbital angular + momentum L (waveform) frame + """ name = "align_total_spin" - _inputs = [parameters.thetajn, parameters.spin1x, parameters.spin1y, - parameters.spin1z, parameters.spin2x, parameters.spin2y, - parameters.spin2z, parameters.mass1, parameters.mass2, - parameters.f_ref, "phi_ref"] - _outputs = [parameters.inclination, parameters.spin1x, parameters.spin1y, - parameters.spin1z, parameters.spin2x, parameters.spin2y, - parameters.spin2z] + _inputs = [ + parameters.thetajn, + parameters.spin1x, + parameters.spin1y, + parameters.spin1z, + parameters.spin2x, + parameters.spin2y, + parameters.spin2z, + parameters.mass1, + parameters.mass2, + parameters.f_ref, + "phi_ref", + ] + _outputs = [ + parameters.inclination, + parameters.spin1x, + parameters.spin1y, + parameters.spin1z, + parameters.spin2x, + parameters.spin2y, + parameters.spin2z, + ] def __init__(self): self.inputs = set(self._inputs) self.outputs = set(self._outputs) - super(AlignTotalSpin, self).__init__() + super().__init__() def transform(self, maps): """ @@ -740,24 +784,33 @@ def transform(self, maps): Note: the spins are assumed to be given in the frame defined by the orbital angular momentum. """ - if isinstance(maps, dict): maps = record.FieldArray.from_kwargs(**maps) newfields = [n for n in self._outputs if n not in maps.fieldnames] - newmaps = maps.add_fields([numpy.zeros(len(maps))]*len(newfields), - names=newfields) + newmaps = maps.add_fields( + [numpy.zeros(len(maps))] * len(newfields), names=newfields + ) for item in newmaps: - if not all(s == 0.0 for s in - [item[parameters.spin1x], item[parameters.spin1y], - item[parameters.spin2x], item[parameters.spin2y]]): - + if not all( + s == 0.0 + for s in [ + item[parameters.spin1x], + item[parameters.spin1y], + item[parameters.spin2x], + item[parameters.spin2y], + ] + ): # Calculate the quantities required by jframe_to_l0frame s1_a, s1_az, s1_pol = coordinates.cartesian_to_spherical( - item[parameters.spin1x], item[parameters.spin1y], - item[parameters.spin1z]) + item[parameters.spin1x], + item[parameters.spin1y], + item[parameters.spin1z], + ) s2_a, s2_az, s2_pol = coordinates.cartesian_to_spherical( - item[parameters.spin2x], item[parameters.spin2y], - item[parameters.spin2z]) + item[parameters.spin2x], + item[parameters.spin2y], + item[parameters.spin2z], + ) out = jframe_to_l0frame( item[parameters.mass1], @@ -770,7 +823,7 @@ def transform(self, maps): spin2_a=s2_a, spin1_polar=s1_pol, spin2_polar=s2_pol, - spin12_deltaphi=s1_az-s2_az + spin12_deltaphi=s1_az - s2_az, ) for key in out: @@ -782,7 +835,8 @@ def transform(self, maps): class SphericalToCartesian(BaseTransform): - """Converts spherical coordinates to cartesian. + """ + Converts spherical coordinates to cartesian. Parameters ---------- @@ -798,6 +852,7 @@ class SphericalToCartesian(BaseTransform): The name of the azimuthal angle parameter. polar : str The name of the polar angle parameter. + """ name = "spherical_to_cartesian" @@ -811,10 +866,11 @@ def __init__(self, x, y, z, radial, azimuthal, polar): self.azimuthal = azimuthal self._inputs = [self.radial, self.azimuthal, self.polar] self._outputs = [self.x, self.y, self.z] - super(SphericalToCartesian, self).__init__() + super().__init__() def transform(self, maps): - """This function transforms from spherical to cartesian spins. + """ + This function transforms from spherical to cartesian spins. Parameters ---------- @@ -839,6 +895,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ a = self.radial az = self.azimuthal @@ -848,7 +905,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms from cartesian to spherical spins. + """ + This function transforms from cartesian to spherical spins. Parameters ---------- @@ -859,6 +917,7 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ x = self.x y = self.y @@ -869,7 +928,8 @@ def inverse_transform(self, maps): class SphericalSpin1ToCartesianSpin1(SphericalToCartesian): - """Converts spherical spin parameters (radial and two angles) to + """ + Converts spherical spin parameters (radial and two angles) to catesian spin parameters. This class only transforms spins for the first component mass. @@ -886,16 +946,17 @@ def __init__(self): "removed in a future update. Please use %s instead, " "passing spin1x, spin1y, spin1z, spin1_a, " "spin1_azimuthal, spin1_polar as arguments.", - self.name, SphericalToCartesian.name + self.name, + SphericalToCartesian.name, ) - super(SphericalSpin1ToCartesianSpin1, self).__init__( - "spin1x", "spin1y", "spin1z", "spin1_a", - "spin1_azimuthal", "spin1_polar" + super().__init__( + "spin1x", "spin1y", "spin1z", "spin1_a", "spin1_azimuthal", "spin1_polar" ) class SphericalSpin2ToCartesianSpin2(SphericalToCartesian): - """Converts spherical spin parameters (radial and two angles) to + """ + Converts spherical spin parameters (radial and two angles) to catesian spin parameters. This class only transforms spins for the first component mass. @@ -912,11 +973,11 @@ def __init__(self): "removed in a future update. Please use %s instead, " "passing spin2x, spin2y, spin2z, spin2_a, " "spin2_azimuthal, spin2_polar as arguments.", - self.name, SphericalToCartesian.name + self.name, + SphericalToCartesian.name, ) - super(SphericalSpin2ToCartesianSpin2, self).__init__( - "spin2x", "spin2y", "spin2z", - "spin2_a", "spin2_azimuthal", "spin2_polar" + super().__init__( + "spin2x", "spin2y", "spin2z", "spin2_a", "spin2_azimuthal", "spin2_polar" ) @@ -929,7 +990,8 @@ class DistanceToRedshift(BaseTransform): _outputs = [parameters.redshift] def transform(self, maps): - """This function transforms from distance to redshift. + """ + This function transforms from distance to redshift. Parameters ---------- @@ -950,6 +1012,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {parameters.redshift: cosmology.redshift(maps[parameters.distance])} return self.format_output(maps, out) @@ -968,7 +1031,8 @@ class AlignedMassSpinToCartesianSpin(BaseTransform): ] def transform(self, maps): - """This function transforms from aligned mass-weighted spins to + """ + This function transforms from aligned mass-weighted spins to cartesian spins aligned along the z-axis. Parameters @@ -980,6 +1044,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ mass1 = maps[parameters.mass1] mass2 = maps[parameters.mass2] @@ -993,7 +1058,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms from component masses and cartesian spins + """ + This function transforms from component masses and cartesian spins to mass-weighted spin parameters aligned with the angular momentum. Parameters @@ -1005,14 +1071,14 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ mass1 = maps[parameters.mass1] spin1z = maps[parameters.spin1z] mass2 = maps[parameters.mass2] spin2z = maps[parameters.spin2z] out = { - parameters.chi_eff: - conversions.chi_eff(mass1, mass2, spin1z, spin2z), + parameters.chi_eff: conversions.chi_eff(mass1, mass2, spin1z, spin2z), "chi_a": conversions.chi_a(mass1, mass2, spin1z, spin2z), } return self.format_output(maps, out) @@ -1022,8 +1088,7 @@ class PrecessionMassSpinToCartesianSpin(BaseTransform): """Converts mass-weighted spins to cartesian x-y plane spins.""" name = "precession_mass_spin_to_cartesian_spin" - _inputs = [parameters.mass1, parameters.mass2, - "xi1", "xi2", "phi_a", "phi_s"] + _inputs = [parameters.mass1, parameters.mass2, "xi1", "xi2", "phi_a", "phi_s"] _outputs = [ parameters.mass1, parameters.mass2, @@ -1034,7 +1099,8 @@ class PrecessionMassSpinToCartesianSpin(BaseTransform): ] def transform(self, maps): - """This function transforms from mass-weighted spins to caretsian spins + """ + This function transforms from mass-weighted spins to caretsian spins in the x-y plane. Parameters @@ -1046,8 +1112,8 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. - """ + """ # find primary and secondary masses # since functions in conversions.py map to primary/secondary masses m_p = conversions.primary_mass(maps["mass1"], maps["mass2"]) @@ -1108,7 +1174,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms from component masses and cartesian spins to + """ + This function transforms from component masses and cartesian spins to mass-weighted spin parameters perpendicular with the angular momentum. Parameters @@ -1120,8 +1187,8 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. - """ + """ # convert out = {} xi1 = conversions.primary_xi( @@ -1193,7 +1260,8 @@ class CartesianSpinToChiP(BaseTransform): _outputs = ["chi_p"] def transform(self, maps): - """This function transforms from component masses and caretsian spins + """ + This function transforms from component masses and caretsian spins to chi_p. Parameters @@ -1209,6 +1277,7 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} out["chi_p"] = conversions.chi_p( @@ -1223,7 +1292,8 @@ def transform(self, maps): class LambdaFromTOVFile(BaseTransform): - """Transforms mass values corresponding to Lambda values for a given EOS + """ + Transforms mass values corresponding to Lambda values for a given EOS interpolating from the mass-Lambda data for that EOS read in from an external ASCII file. @@ -1291,6 +1361,7 @@ class LambdaFromTOVFile(BaseTransform): The names and order of columns in the ``mass_lambda_file``. Must contain at least 'mass' and 'lambda'. If not provided, will assume the order is ('mass', 'lambda'). + """ name = "lambda_from_tov_file" @@ -1316,7 +1387,7 @@ def __init__( dtype = [(fname, float) for fname in file_columns] data = numpy.loadtxt(self._mass_lambda_file, dtype=dtype) self._data = data - super(LambdaFromTOVFile, self).__init__() + super().__init__() @property def mass_param(self): @@ -1334,28 +1405,32 @@ def data(self): @property def mass_data(self): - """Returns the mass data read from the mass-Lambda data file for + """ + Returns the mass data read from the mass-Lambda data file for an EOS. """ return self._data["mass"] @property def lambda_data(self): - """Returns the Lambda data read from the mass-Lambda data file for + """ + Returns the Lambda data read from the mass-Lambda data file for an EOS. """ return self._data["lambda"] @property def distance(self): - """Returns the fixed distance to transform mass samples from detector + """ + Returns the fixed distance to transform mass samples from detector to source frame if one is specified. """ return self._distance @staticmethod def lambda_from_tov_data(m_src, mass_data, lambda_data): - """Returns Lambda corresponding to a given mass interpolating from the + """ + Returns Lambda corresponding to a given mass interpolating from the TOV data. Parameters @@ -1371,6 +1446,7 @@ def lambda_from_tov_data(m_src, mass_data, lambda_data): ------- lambdav : float The Lambda corresponding to the mass `m` for the EOS considered. + """ if m_src > mass_data.max(): # assume black hole @@ -1380,7 +1456,8 @@ def lambda_from_tov_data(m_src, mass_data, lambda_data): return lambdav def transform(self, maps): - """Computes the transformation of mass to Lambda. + """ + Computes the transformation of mass to Lambda. Parameters ---------- @@ -1393,6 +1470,7 @@ def transform(self, maps): out : dict or FieldArray A map between the transformed variable name and value(s), along with the original variable name and value(s). + """ m = maps[self._mass_param] if self.redshift_mass: @@ -1428,13 +1506,14 @@ def from_config(cls, cp, section, outputs): else: additional_opts = None skip_opts = None - return super(LambdaFromTOVFile, cls).from_config( + return super().from_config( cp, section, outputs, skip_opts=skip_opts, additional_opts=additional_opts ) class LambdaFromMultipleTOVFiles(BaseTransform): - """Uses multiple equation of states. + """ + Uses multiple equation of states. Parameters ---------- @@ -1453,6 +1532,7 @@ class LambdaFromMultipleTOVFiles(BaseTransform): The names and order of columns in the ``mass_lambda_file``. Must contain at least 'mass' and 'lambda'. If not provided, will assume the order is ('radius', 'mass', 'lambda'). + """ name = "lambda_from_multiple_tov_files" @@ -1475,7 +1555,7 @@ def __init__( self._outputs = [lambda_param] # create a dictionary of the EOS files from the map_file self._eos_files = {} - with open(self._map_file, "r") as fp: + with open(self._map_file) as fp: for line in fp: fname = line.rstrip("\n") eosidx = int(os.path.basename(fname).split(".")[0]) @@ -1485,7 +1565,7 @@ def __init__( if file_columns is None: file_columns = ("radius", "mass", "lambda") self._file_columns = file_columns - super(LambdaFromMultipleTOVFiles, self).__init__() + super().__init__() @property def mass_param(self): @@ -1499,20 +1579,23 @@ def lambda_param(self): @property def map_file(self): - """Returns the mass data read from the mass-Lambda data file for + """ + Returns the mass data read from the mass-Lambda data file for an EOS. """ return self._map_file @property def distance(self): - """Returns the fixed distance to transform mass samples from detector + """ + Returns the fixed distance to transform mass samples from detector to source frame if one is specified. """ return self._distance def get_eos(self, eos_index): - """Gets the EOS for the given index. + """ + Gets the EOS for the given index. If the index is not in range returns None. """ @@ -1542,10 +1625,9 @@ def transform(self, maps): eos = self.get_eos(eos_index) if eos is not None: return eos.transform(maps) - else: - # no eos, just return nan - out = {self._lambda_param: numpy.nan} - return self.format_output(maps, out) + # no eos, just return nan + out = {self._lambda_param: numpy.nan} + return self.format_output(maps, out) @classmethod def from_config(cls, cp, section, outputs): @@ -1556,38 +1638,51 @@ def from_config(cls, cp, section, outputs): else: additional_opts = None skip_opts = None - return super(LambdaFromMultipleTOVFiles, cls).from_config( + return super().from_config( cp, section, outputs, skip_opts=skip_opts, additional_opts=additional_opts ) class GEOToSSB(BaseTransform): - """Converts arrival time, sky localization, and polarization angle in the - geocentric frame to the corresponding values in the SSB frame.""" + """ + Converts arrival time, sky localization, and polarization angle in the + geocentric frame to the corresponding values in the SSB frame. + """ name = "geo_to_ssb" default_params_name = { - 'default_tc_geo': parameters.tc, - 'default_longitude_geo': parameters.ra, - 'default_latitude_geo': parameters.dec, - 'default_polarization_geo': parameters.polarization, - 'default_tc_ssb': parameters.tc, - 'default_longitude_ssb': parameters.eclipticlongitude, - 'default_latitude_ssb': parameters.eclipticlatitude, - 'default_polarization_ssb': parameters.polarization + "default_tc_geo": parameters.tc, + "default_longitude_geo": parameters.ra, + "default_latitude_geo": parameters.dec, + "default_polarization_geo": parameters.polarization, + "default_tc_ssb": parameters.tc, + "default_longitude_ssb": parameters.eclipticlongitude, + "default_latitude_ssb": parameters.eclipticlatitude, + "default_polarization_ssb": parameters.polarization, } def __init__( - self, tc_geo_param=None, longitude_geo_param=None, - latitude_geo_param=None, polarization_geo_param=None, - tc_ssb_param=None, longitude_ssb_param=None, - latitude_ssb_param=None, polarization_ssb_param=None + self, + tc_geo_param=None, + longitude_geo_param=None, + latitude_geo_param=None, + polarization_geo_param=None, + tc_ssb_param=None, + longitude_ssb_param=None, + latitude_ssb_param=None, + polarization_ssb_param=None, ): - params = [tc_geo_param, longitude_geo_param, - latitude_geo_param, polarization_geo_param, - tc_ssb_param, longitude_ssb_param, - latitude_ssb_param, polarization_ssb_param] + params = [ + tc_geo_param, + longitude_geo_param, + latitude_geo_param, + polarization_geo_param, + tc_ssb_param, + longitude_ssb_param, + latitude_ssb_param, + polarization_ssb_param, + ] for index in range(len(params)): if params[index] is None: @@ -1602,15 +1697,24 @@ def __init__( self.longitude_ssb_param = params[5] self.latitude_ssb_param = params[6] self.polarization_ssb_param = params[7] - self._inputs = [self.tc_geo_param, self.longitude_geo_param, - self.latitude_geo_param, self.polarization_geo_param] - self._outputs = [self.tc_ssb_param, self.longitude_ssb_param, - self.latitude_ssb_param, self.polarization_ssb_param] - - super(GEOToSSB, self).__init__() + self._inputs = [ + self.tc_geo_param, + self.longitude_geo_param, + self.latitude_geo_param, + self.polarization_geo_param, + ] + self._outputs = [ + self.tc_ssb_param, + self.longitude_ssb_param, + self.latitude_ssb_param, + self.polarization_ssb_param, + ] + + super().__init__() def transform(self, maps): - """This function transforms arrival time, sky localization, + """ + This function transforms arrival time, sky localization, and polarization angle in the geocentric frame to the corresponding values in the SSB frame. @@ -1623,18 +1727,25 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} - out[self.tc_ssb_param], out[self.longitude_ssb_param], \ - out[self.latitude_ssb_param], out[self.polarization_ssb_param] = \ - coordinates.geo_to_ssb( - maps[self.tc_geo_param], maps[self.longitude_geo_param], - maps[self.latitude_geo_param], maps[self.polarization_geo_param] - ) + ( + out[self.tc_ssb_param], + out[self.longitude_ssb_param], + out[self.latitude_ssb_param], + out[self.polarization_ssb_param], + ) = coordinates.geo_to_ssb( + maps[self.tc_geo_param], + maps[self.longitude_geo_param], + maps[self.latitude_geo_param], + maps[self.polarization_geo_param], + ) return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms arrival time, sky localization, + """ + This function transforms arrival time, sky localization, and polarization angle in the SSB frame to the corresponding values in the geocentric frame. @@ -1647,14 +1758,20 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} - out[self.tc_geo_param], out[self.longitude_geo_param], \ - out[self.latitude_geo_param], out[self.polarization_geo_param] = \ - coordinates.ssb_to_geo( - maps[self.tc_ssb_param], maps[self.longitude_ssb_param], - maps[self.latitude_ssb_param], maps[self.polarization_ssb_param] - ) + ( + out[self.tc_geo_param], + out[self.longitude_geo_param], + out[self.latitude_geo_param], + out[self.polarization_geo_param], + ) = coordinates.ssb_to_geo( + maps[self.tc_ssb_param], + maps[self.longitude_ssb_param], + maps[self.latitude_ssb_param], + maps[self.polarization_ssb_param], + ) return self.format_output(maps, out) @classmethod @@ -1665,61 +1782,76 @@ def from_config(cls, cp, section, outputs): # get custom variable names variables = { - 'tc-geo': cls.default_params_name['default_tc_geo'], - 'longitude-geo': cls.default_params_name['default_longitude_geo'], - 'latitude-geo': cls.default_params_name['default_latitude_geo'], - 'polarization-geo': cls.default_params_name[ - 'default_polarization_geo'], - 'tc-ssb': cls.default_params_name['default_tc_ssb'], - 'longitude-ssb': cls.default_params_name['default_longitude_ssb'], - 'latitude-ssb': cls.default_params_name['default_latitude_ssb'], - 'polarization-ssb': cls.default_params_name[ - 'default_polarization_ssb'] + "tc-geo": cls.default_params_name["default_tc_geo"], + "longitude-geo": cls.default_params_name["default_longitude_geo"], + "latitude-geo": cls.default_params_name["default_latitude_geo"], + "polarization-geo": cls.default_params_name["default_polarization_geo"], + "tc-ssb": cls.default_params_name["default_tc_ssb"], + "longitude-ssb": cls.default_params_name["default_longitude_ssb"], + "latitude-ssb": cls.default_params_name["default_latitude_ssb"], + "polarization-ssb": cls.default_params_name["default_polarization_ssb"], } - for param_name in variables.keys(): - name_underline = param_name.replace('-', '_') + for param_name in variables: + name_underline = param_name.replace("-", "_") if cp.has_option("-".join([section, outputs]), param_name): skip_opts.append(param_name) additional_opts.update( - {name_underline+'_param': cp.get_opt_tag( - section, param_name, tag)}) + { + name_underline + "_param": cp.get_opt_tag( + section, param_name, tag + ) + } + ) else: additional_opts.update( - {name_underline+'_param': variables[param_name]}) + {name_underline + "_param": variables[param_name]} + ) - return super(GEOToSSB, cls).from_config( - cp, section, outputs, skip_opts=skip_opts, - additional_opts=additional_opts + return super().from_config( + cp, section, outputs, skip_opts=skip_opts, additional_opts=additional_opts ) class LISAToSSB(BaseTransform): - """Converts arrival time, sky localization, and polarization angle in the - LISA frame to the corresponding values in the SSB frame.""" + """ + Converts arrival time, sky localization, and polarization angle in the + LISA frame to the corresponding values in the SSB frame. + """ name = "lisa_to_ssb" default_params_name = { - 'default_tc_lisa': parameters.tc, - 'default_longitude_lisa': parameters.eclipticlongitude, - 'default_latitude_lisa': parameters.eclipticlatitude, - 'default_polarization_lisa': parameters.polarization, - 'default_tc_ssb': parameters.tc, - 'default_longitude_ssb': parameters.eclipticlongitude, - 'default_latitude_ssb': parameters.eclipticlatitude, - 'default_polarization_ssb': parameters.polarization + "default_tc_lisa": parameters.tc, + "default_longitude_lisa": parameters.eclipticlongitude, + "default_latitude_lisa": parameters.eclipticlatitude, + "default_polarization_lisa": parameters.polarization, + "default_tc_ssb": parameters.tc, + "default_longitude_ssb": parameters.eclipticlongitude, + "default_latitude_ssb": parameters.eclipticlatitude, + "default_polarization_ssb": parameters.polarization, } def __init__( - self, tc_lisa_param=None, longitude_lisa_param=None, - latitude_lisa_param=None, polarization_lisa_param=None, - tc_ssb_param=None, longitude_ssb_param=None, - latitude_ssb_param=None, polarization_ssb_param=None + self, + tc_lisa_param=None, + longitude_lisa_param=None, + latitude_lisa_param=None, + polarization_lisa_param=None, + tc_ssb_param=None, + longitude_ssb_param=None, + latitude_ssb_param=None, + polarization_ssb_param=None, ): - params = [tc_lisa_param, longitude_lisa_param, - latitude_lisa_param, polarization_lisa_param, - tc_ssb_param, longitude_ssb_param, - latitude_ssb_param, polarization_ssb_param] + params = [ + tc_lisa_param, + longitude_lisa_param, + latitude_lisa_param, + polarization_lisa_param, + tc_ssb_param, + longitude_ssb_param, + latitude_ssb_param, + polarization_ssb_param, + ] for index in range(len(params)): if params[index] is None: key = list(self.default_params_name.keys())[index] @@ -1733,14 +1865,23 @@ def __init__( self.longitude_ssb_param = params[5] self.latitude_ssb_param = params[6] self.polarization_ssb_param = params[7] - self._inputs = [self.tc_lisa_param, self.longitude_lisa_param, - self.latitude_lisa_param, self.polarization_lisa_param] - self._outputs = [self.tc_ssb_param, self.longitude_ssb_param, - self.latitude_ssb_param, self.polarization_ssb_param] - super(LISAToSSB, self).__init__() + self._inputs = [ + self.tc_lisa_param, + self.longitude_lisa_param, + self.latitude_lisa_param, + self.polarization_lisa_param, + ] + self._outputs = [ + self.tc_ssb_param, + self.longitude_ssb_param, + self.latitude_ssb_param, + self.polarization_ssb_param, + ] + super().__init__() def transform(self, maps): - """This function transforms arrival time, sky localization, + """ + This function transforms arrival time, sky localization, and polarization angle in the LISA frame to the corresponding values in the SSB frame. @@ -1753,18 +1894,25 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} - out[self.tc_ssb_param], out[self.longitude_ssb_param], \ - out[self.latitude_ssb_param], out[self.polarization_ssb_param] = \ - coordinates.lisa_to_ssb( - maps[self.tc_lisa_param], maps[self.longitude_lisa_param], - maps[self.latitude_lisa_param], maps[self.polarization_lisa_param] - ) + ( + out[self.tc_ssb_param], + out[self.longitude_ssb_param], + out[self.latitude_ssb_param], + out[self.polarization_ssb_param], + ) = coordinates.lisa_to_ssb( + maps[self.tc_lisa_param], + maps[self.longitude_lisa_param], + maps[self.latitude_lisa_param], + maps[self.polarization_lisa_param], + ) return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms arrival time, sky localization, + """ + This function transforms arrival time, sky localization, and polarization angle in the SSB frame to the corresponding values in the LISA frame. @@ -1777,15 +1925,20 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} - out[self.tc_lisa_param], out[self.longitude_lisa_param], \ - out[self.latitude_lisa_param], \ - out[self.polarization_lisa_param] = \ - coordinates.ssb_to_lisa( - maps[self.tc_ssb_param], maps[self.longitude_ssb_param], - maps[self.latitude_ssb_param], maps[self.polarization_ssb_param] - ) + ( + out[self.tc_lisa_param], + out[self.longitude_lisa_param], + out[self.latitude_lisa_param], + out[self.polarization_lisa_param], + ) = coordinates.ssb_to_lisa( + maps[self.tc_ssb_param], + maps[self.longitude_ssb_param], + maps[self.latitude_ssb_param], + maps[self.polarization_ssb_param], + ) return self.format_output(maps, out) @classmethod @@ -1796,62 +1949,76 @@ def from_config(cls, cp, section, outputs): # get custom variable names variables = { - 'tc-lisa': cls.default_params_name['default_tc_lisa'], - 'longitude-lisa': cls.default_params_name[ - 'default_longitude_lisa'], - 'latitude-lisa': cls.default_params_name['default_latitude_lisa'], - 'polarization-lisa': cls.default_params_name[ - 'default_polarization_lisa'], - 'tc-ssb': cls.default_params_name['default_tc_ssb'], - 'longitude-ssb': cls.default_params_name['default_longitude_ssb'], - 'latitude-ssb': cls.default_params_name['default_latitude_ssb'], - 'polarization-ssb': cls.default_params_name[ - 'default_polarization_ssb'] + "tc-lisa": cls.default_params_name["default_tc_lisa"], + "longitude-lisa": cls.default_params_name["default_longitude_lisa"], + "latitude-lisa": cls.default_params_name["default_latitude_lisa"], + "polarization-lisa": cls.default_params_name["default_polarization_lisa"], + "tc-ssb": cls.default_params_name["default_tc_ssb"], + "longitude-ssb": cls.default_params_name["default_longitude_ssb"], + "latitude-ssb": cls.default_params_name["default_latitude_ssb"], + "polarization-ssb": cls.default_params_name["default_polarization_ssb"], } - for param_name in variables.keys(): - name_underline = param_name.replace('-', '_') + for param_name in variables: + name_underline = param_name.replace("-", "_") if cp.has_option("-".join([section, outputs]), param_name): skip_opts.append(param_name) additional_opts.update( - {name_underline+'_param': cp.get_opt_tag( - section, param_name, tag)}) + { + name_underline + "_param": cp.get_opt_tag( + section, param_name, tag + ) + } + ) else: additional_opts.update( - {name_underline+'_param': variables[param_name]}) + {name_underline + "_param": variables[param_name]} + ) - return super(LISAToSSB, cls).from_config( - cp, section, outputs, skip_opts=skip_opts, - additional_opts=additional_opts + return super().from_config( + cp, section, outputs, skip_opts=skip_opts, additional_opts=additional_opts ) class LISAToGEO(BaseTransform): - """Converts arrival time, sky localization, and polarization angle in the - LISA frame to the corresponding values in the geocentric frame.""" + """ + Converts arrival time, sky localization, and polarization angle in the + LISA frame to the corresponding values in the geocentric frame. + """ name = "lisa_to_geo" default_params_name = { - 'default_tc_lisa': parameters.tc, - 'default_longitude_lisa': parameters.eclipticlongitude, - 'default_latitude_lisa': parameters.eclipticlatitude, - 'default_polarization_lisa': parameters.polarization, - 'default_tc_geo': parameters.tc, - 'default_longitude_geo': parameters.ra, - 'default_latitude_geo': parameters.dec, - 'default_polarization_geo': parameters.polarization + "default_tc_lisa": parameters.tc, + "default_longitude_lisa": parameters.eclipticlongitude, + "default_latitude_lisa": parameters.eclipticlatitude, + "default_polarization_lisa": parameters.polarization, + "default_tc_geo": parameters.tc, + "default_longitude_geo": parameters.ra, + "default_latitude_geo": parameters.dec, + "default_polarization_geo": parameters.polarization, } def __init__( - self, tc_lisa_param=None, longitude_lisa_param=None, - latitude_lisa_param=None, polarization_lisa_param=None, - tc_geo_param=None, longitude_geo_param=None, - latitude_geo_param=None, polarization_geo_param=None + self, + tc_lisa_param=None, + longitude_lisa_param=None, + latitude_lisa_param=None, + polarization_lisa_param=None, + tc_geo_param=None, + longitude_geo_param=None, + latitude_geo_param=None, + polarization_geo_param=None, ): - params = [tc_lisa_param, longitude_lisa_param, - latitude_lisa_param, polarization_lisa_param, - tc_geo_param, longitude_geo_param, - latitude_geo_param, polarization_geo_param] + params = [ + tc_lisa_param, + longitude_lisa_param, + latitude_lisa_param, + polarization_lisa_param, + tc_geo_param, + longitude_geo_param, + latitude_geo_param, + polarization_geo_param, + ] for index in range(len(params)): if params[index] is None: key = list(self.default_params_name.keys())[index] @@ -1865,14 +2032,23 @@ def __init__( self.longitude_geo_param = params[5] self.latitude_geo_param = params[6] self.polarization_geo_param = params[7] - self._inputs = [self.tc_lisa_param, self.longitude_lisa_param, - self.latitude_lisa_param, self.polarization_lisa_param] - self._outputs = [self.tc_geo_param, self.longitude_geo_param, - self.latitude_geo_param, self.polarization_geo_param] - super(LISAToGEO, self).__init__() + self._inputs = [ + self.tc_lisa_param, + self.longitude_lisa_param, + self.latitude_lisa_param, + self.polarization_lisa_param, + ] + self._outputs = [ + self.tc_geo_param, + self.longitude_geo_param, + self.latitude_geo_param, + self.polarization_geo_param, + ] + super().__init__() def transform(self, maps): - """This function transforms arrival time, sky localization, + """ + This function transforms arrival time, sky localization, and polarization angle in the LISA frame to the corresponding values in the geocentric frame. @@ -1885,18 +2061,25 @@ def transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} - out[self.tc_geo_param], out[self.longitude_geo_param], \ - out[self.latitude_geo_param], out[self.polarization_geo_param] = \ - coordinates.lisa_to_geo( - maps[self.tc_lisa_param], maps[self.longitude_lisa_param], - maps[self.latitude_lisa_param], maps[self.polarization_lisa_param] - ) + ( + out[self.tc_geo_param], + out[self.longitude_geo_param], + out[self.latitude_geo_param], + out[self.polarization_geo_param], + ) = coordinates.lisa_to_geo( + maps[self.tc_lisa_param], + maps[self.longitude_lisa_param], + maps[self.latitude_lisa_param], + maps[self.polarization_lisa_param], + ) return self.format_output(maps, out) def inverse_transform(self, maps): - """This function transforms arrival time, sky localization, + """ + This function transforms arrival time, sky localization, and polarization angle in the geocentric frame to the corresponding values in the LISA frame. @@ -1909,15 +2092,20 @@ def inverse_transform(self, maps): out : dict A dict with key as parameter name and value as numpy.array or float of transformed values. + """ out = {} - out[self.tc_lisa_param], out[self.longitude_lisa_param], \ - out[self.latitude_lisa_param], \ - out[self.polarization_lisa_param] = \ - coordinates.geo_to_lisa( - maps[self.tc_geo_param], maps[self.longitude_geo_param], - maps[self.latitude_geo_param], maps[self.polarization_geo_param] - ) + ( + out[self.tc_lisa_param], + out[self.longitude_lisa_param], + out[self.latitude_lisa_param], + out[self.polarization_lisa_param], + ) = coordinates.geo_to_lisa( + maps[self.tc_geo_param], + maps[self.longitude_geo_param], + maps[self.latitude_geo_param], + maps[self.polarization_geo_param], + ) return self.format_output(maps, out) @classmethod @@ -1928,37 +2116,39 @@ def from_config(cls, cp, section, outputs): # get custom variable names variables = { - 'tc-lisa': cls.default_params_name['default_tc_lisa'], - 'longitude-lisa': cls.default_params_name[ - 'default_longitude_lisa'], - 'latitude-lisa': cls.default_params_name['default_latitude_lisa'], - 'polarization-lisa': cls.default_params_name[ - 'default_polarization_lisa'], - 'tc-geo': cls.default_params_name['default_tc_geo'], - 'longitude-geo': cls.default_params_name['default_longitude_geo'], - 'latitude-geo': cls.default_params_name['default_latitude_geo'], - 'polarization-geo': cls.default_params_name[ - 'default_polarization_geo'] + "tc-lisa": cls.default_params_name["default_tc_lisa"], + "longitude-lisa": cls.default_params_name["default_longitude_lisa"], + "latitude-lisa": cls.default_params_name["default_latitude_lisa"], + "polarization-lisa": cls.default_params_name["default_polarization_lisa"], + "tc-geo": cls.default_params_name["default_tc_geo"], + "longitude-geo": cls.default_params_name["default_longitude_geo"], + "latitude-geo": cls.default_params_name["default_latitude_geo"], + "polarization-geo": cls.default_params_name["default_polarization_geo"], } - for param_name in variables.keys(): - name_underline = param_name.replace('-', '_') + for param_name in variables: + name_underline = param_name.replace("-", "_") if cp.has_option("-".join([section, outputs]), param_name): skip_opts.append(param_name) additional_opts.update( - {name_underline+'_param': cp.get_opt_tag( - section, param_name, tag)}) + { + name_underline + "_param": cp.get_opt_tag( + section, param_name, tag + ) + } + ) else: additional_opts.update( - {name_underline+'_param': variables[param_name]}) + {name_underline + "_param": variables[param_name]} + ) - return super(LISAToGEO, cls).from_config( - cp, section, outputs, skip_opts=skip_opts, - additional_opts=additional_opts + return super().from_config( + cp, section, outputs, skip_opts=skip_opts, additional_opts=additional_opts ) class Log(BaseTransform): - """Applies a log transform from an `inputvar` parameter to an `outputvar` + """ + Applies a log transform from an `inputvar` parameter to an `outputvar` parameter. This is the inverse of the exponent transform. Parameters @@ -1967,6 +2157,7 @@ class Log(BaseTransform): The name of the parameter to transform. outputvar : str The name of the transformed parameter. + """ name = "log" @@ -1976,7 +2167,7 @@ def __init__(self, inputvar, outputvar): self._outputvar = outputvar self._inputs = [inputvar] self._outputs = [outputvar] - super(Log, self).__init__() + super().__init__() @property def inputvar(self): @@ -1989,7 +2180,8 @@ def outputvar(self): return self._outputvar def transform(self, maps): - r"""Computes :math:`\log(x)`. + r""" + Computes :math:`\log(x)`. Parameters ---------- @@ -2002,13 +2194,15 @@ def transform(self, maps): out : dict or FieldArray A map between the transformed variable name and value(s), along with the original variable name and value(s). + """ x = maps[self._inputvar] out = {self._outputvar: numpy.log(x)} return self.format_output(maps, out) def inverse_transform(self, maps): - r"""Computes :math:`y = e^{x}`. + r""" + Computes :math:`y = e^{x}`. Parameters ---------- @@ -2021,13 +2215,15 @@ def inverse_transform(self, maps): out : dict or FieldArray A map between the transformed variable name and value(s), along with the original variable name and value(s). + """ y = maps[self._outputvar] out = {self._inputvar: numpy.exp(y)} return self.format_output(maps, out) def jacobian(self, maps): - r"""Computes the Jacobian of :math:`y = \log(x)`. + r""" + Computes the Jacobian of :math:`y = \log(x)`. This is: @@ -2045,12 +2241,14 @@ def jacobian(self, maps): ------- float The value of the jacobian at the given point(s). + """ x = maps[self._inputvar] return 1.0 / x def inverse_jacobian(self, maps): - r"""Computes the Jacobian of :math:`y = e^{x}`. + r""" + Computes the Jacobian of :math:`y = e^{x}`. This is: @@ -2068,13 +2266,15 @@ def inverse_jacobian(self, maps): ------- float The value of the jacobian at the given point(s). + """ x = maps[self._outputvar] return numpy.exp(x) class Logit(BaseTransform): - r"""Applies a logit transform from an `inputvar` parameter to an `outputvar` + r""" + Applies a logit transform from an `inputvar` parameter to an `outputvar` parameter. This is the inverse of the logistic transform. Typically, the input of the logit function is assumed to have domain @@ -2090,6 +2290,7 @@ class Logit(BaseTransform): domain : tuple or distributions.bounds.Bounds, optional The domain of the input parameter. Can be any finite interval. Default is (0., 1.). + """ name = "logit" @@ -2099,12 +2300,11 @@ def __init__(self, inputvar, outputvar, domain=(0.0, 1.0)): self._outputvar = outputvar self._inputs = [inputvar] self._outputs = [outputvar] - self._bounds = Bounds(domain[0], domain[1], - btype_min="open", btype_max="open") + self._bounds = Bounds(domain[0], domain[1], btype_min="open", btype_max="open") # shortcuts for quick access later self._a = domain[0] self._b = domain[1] - super(Logit, self).__init__() + super().__init__() @property def inputvar(self): @@ -2123,7 +2323,8 @@ def bounds(self): @staticmethod def logit(x, a=0.0, b=1.0): - r"""Computes the logit function with domain :math:`x \in (a, b)`. + r""" + Computes the logit function with domain :math:`x \in (a, b)`. This is given by: @@ -2147,12 +2348,14 @@ def logit(x, a=0.0, b=1.0): ------- float The logit of x. + """ return numpy.log(x - a) - numpy.log(b - x) @staticmethod def logistic(x, a=0.0, b=1.0): - r"""Computes the logistic function with range :math:`\in (a, b)`. + r""" + Computes the logistic function with range :math:`\in (a, b)`. This is given by: @@ -2178,12 +2381,14 @@ def logistic(x, a=0.0, b=1.0): ------- float The logistic of x. + """ expx = numpy.exp(x) return (a + b * expx) / (1.0 + expx) def transform(self, maps): - r"""Computes :math:`\mathrm{logit}(x; a, b)`. + r""" + Computes :math:`\mathrm{logit}(x; a, b)`. The domain :math:`a, b` of :math:`x` are given by the class's bounds. @@ -2198,6 +2403,7 @@ def transform(self, maps): out : dict or FieldArray A map between the transformed variable name and value(s), along with the original variable name and value(s). + """ x = maps[self._inputvar] # check that x is in bounds @@ -2210,7 +2416,8 @@ def transform(self, maps): return self.format_output(maps, out) def inverse_transform(self, maps): - r"""Computes :math:`y = \mathrm{logistic}(x; a,b)`. + r""" + Computes :math:`y = \mathrm{logistic}(x; a,b)`. The codomain :math:`a, b` of :math:`y` are given by the class's bounds. @@ -2225,13 +2432,15 @@ def inverse_transform(self, maps): out : dict or FieldArray A map between the transformed variable name and value(s), along with the original variable name and value(s). + """ y = maps[self._outputvar] out = {self._inputvar: self.logistic(y, self._a, self._b)} return self.format_output(maps, out) def jacobian(self, maps): - r"""Computes the Jacobian of :math:`y = \mathrm{logit}(x; a,b)`. + r""" + Computes the Jacobian of :math:`y = \mathrm{logit}(x; a,b)`. This is: @@ -2251,18 +2460,20 @@ def jacobian(self, maps): ------- float The value of the jacobian at the given point(s). + """ x = maps[self._inputvar] # check that x is in bounds isin = self._bounds.__contains__(x) if isinstance(isin, numpy.ndarray) and not isin.all(): raise ValueError("one or more values are not in bounds") - elif not isin: - raise ValueError("{} is not in bounds".format(x)) + if not isin: + raise ValueError(f"{x} is not in bounds") return (self._b - self._a) / ((x - self._a) * (self._b - x)) def inverse_jacobian(self, maps): - r"""Computes the Jacobian of :math:`y = \mathrm{logistic}(x; a,b)`. + r""" + Computes the Jacobian of :math:`y = \mathrm{logistic}(x; a,b)`. This is: @@ -2282,15 +2493,16 @@ def inverse_jacobian(self, maps): ------- float The value of the jacobian at the given point(s). + """ x = maps[self._outputvar] expx = numpy.exp(x) return expx * (self._b - self._a) / (1.0 + expx) ** 2.0 @classmethod - def from_config(cls, cp, section, outputs, - skip_opts=None, additional_opts=None): - """Initializes a Logit transform from the given section. + def from_config(cls, cp, section, outputs, skip_opts=None, additional_opts=None): + """ + Initializes a Logit transform from the given section. The section must specify an input and output variable name. The domain of the input may be specified using `min-{input}`, `max-{input}`. @@ -2326,11 +2538,12 @@ def from_config(cls, cp, section, outputs, ------- cls An instance of the class. + """ # pull out the minimum, maximum values of the input variable inputvar = cp.get_opt_tag(section, "inputvar", outputs) s = "-".join([section, outputs]) - opt = "min-{}".format(inputvar) + opt = f"min-{inputvar}" if skip_opts is None: skip_opts = [] if additional_opts is None: @@ -2342,20 +2555,19 @@ def from_config(cls, cp, section, outputs, skip_opts.append(opt) else: a = None - opt = "max-{}".format(inputvar) + opt = f"max-{inputvar}" if cp.has_option(s, opt): b = cp.get_opt_tag(section, opt, outputs) skip_opts.append(opt) else: b = None - if a is None and b is not None or b is None and a is not None: + if (a is None and b is not None) or (b is None and a is not None): raise ValueError( - "if providing a min(max)-{}, must also provide " - "a max(min)-{}".format(inputvar, inputvar) + f"if providing a min(max)-{inputvar}, must also provide a max(min)-{inputvar}" ) - elif a is not None: + if a is not None: additional_opts.update({"domain": (float(a), float(b))}) - return super(Logit, cls).from_config( + return super().from_config( cp, section, outputs, skip_opts, additional_opts ) @@ -2424,7 +2636,8 @@ class DistanceToChirpDistance(ChirpDistanceToDistance): class CartesianToSpherical(SphericalToCartesian): - """Converts spherical coordinates to cartesian. + """ + Converts spherical coordinates to cartesian. Parameters ---------- @@ -2440,6 +2653,7 @@ class CartesianToSpherical(SphericalToCartesian): The name of the azimuthal angle parameter. polar : str The name of the polar angle parameter. + """ name = "cartesian_to_spherical" @@ -2450,7 +2664,7 @@ class CartesianToSpherical(SphericalToCartesian): inverse_jacobian = inverse.jacobian def __init__(self, *args): - super(CartesianToSpherical, self).__init__(*args) + super().__init__(*args) # swap inputs and outputs outputs = self._inputs inputs = self._outputs @@ -2461,7 +2675,8 @@ def __init__(self, *args): class CartesianSpin1ToSphericalSpin1(CartesianToSpherical): - """The inverse of SphericalSpin1ToCartesianSpin1. + """ + The inverse of SphericalSpin1ToCartesianSpin1. **Deprecation Warning:** This will be removed in a future update. Use :py:class:`CartesianToSpherical` with spin-parameter names passed in @@ -2476,16 +2691,17 @@ def __init__(self): "removed in a future update. Please use %s instead, " "passing spin1x, spin1y, spin1z, spin1_a, " "spin1_azimuthal, spin1_polar as arguments.", - self.name, CartesianToSpherical.name + self.name, + CartesianToSpherical.name, ) - super(CartesianSpin1ToSphericalSpin1, self).__init__( - "spin1x", "spin1y", "spin1z", - "spin1_a", "spin1_azimuthal", "spin1_polar" + super().__init__( + "spin1x", "spin1y", "spin1z", "spin1_a", "spin1_azimuthal", "spin1_polar" ) class CartesianSpin2ToSphericalSpin2(CartesianToSpherical): - """The inverse of SphericalSpin2ToCartesianSpin2. + """ + The inverse of SphericalSpin2ToCartesianSpin2. **Deprecation Warning:** This will be removed in a future update. Use :py:class:`CartesianToSpherical` with spin-parameter names passed in @@ -2500,11 +2716,11 @@ def __init__(self): "removed in a future update. Please use %s instead, " "passing spin2x, spin2y, spin2z, spin2_a, " "spin2_azimuthal, spin2_polar as arguments.", - self.name, CartesianToSpherical.name + self.name, + CartesianToSpherical.name, ) - super(CartesianSpin2ToSphericalSpin2, self).__init__( - "spin2x", "spin2y", "spin2z", - "spin2_a", "spin2_azimuthal", "spin2_polar" + super().__init__( + "spin2x", "spin2y", "spin2z", "spin2_a", "spin2_azimuthal", "spin2_polar" ) @@ -2556,15 +2772,26 @@ class SSBToGEO(GEOToSSB): inverse_transform = inverse.transform def __init__( - self, tc_geo_param=None, longitude_geo_param=None, - latitude_geo_param=None, polarization_geo_param=None, - tc_ssb_param=None, longitude_ssb_param=None, - latitude_ssb_param=None, polarization_ssb_param=None + self, + tc_geo_param=None, + longitude_geo_param=None, + latitude_geo_param=None, + polarization_geo_param=None, + tc_ssb_param=None, + longitude_ssb_param=None, + latitude_ssb_param=None, + polarization_ssb_param=None, ): - params = [tc_geo_param, longitude_geo_param, - latitude_geo_param, polarization_geo_param, - tc_ssb_param, longitude_ssb_param, - latitude_ssb_param, polarization_ssb_param] + params = [ + tc_geo_param, + longitude_geo_param, + latitude_geo_param, + polarization_geo_param, + tc_ssb_param, + longitude_ssb_param, + latitude_ssb_param, + polarization_ssb_param, + ] for index in range(len(params)): if params[index] is None: key = list(self.default_params_name.keys())[index] @@ -2578,10 +2805,18 @@ def __init__( self.longitude_ssb_param = params[5] self.latitude_ssb_param = params[6] self.polarization_ssb_param = params[7] - self._inputs = [self.tc_ssb_param, self.longitude_ssb_param, - self.latitude_ssb_param, self.polarization_ssb_param] - self._outputs = [self.tc_geo_param, self.longitude_geo_param, - self.latitude_geo_param, self.polarization_geo_param] + self._inputs = [ + self.tc_ssb_param, + self.longitude_ssb_param, + self.latitude_ssb_param, + self.polarization_ssb_param, + ] + self._outputs = [ + self.tc_geo_param, + self.longitude_geo_param, + self.latitude_geo_param, + self.polarization_geo_param, + ] class SSBToLISA(LISAToSSB): @@ -2593,15 +2828,26 @@ class SSBToLISA(LISAToSSB): inverse_transform = inverse.transform def __init__( - self, tc_lisa_param=None, longitude_lisa_param=None, - latitude_lisa_param=None, polarization_lisa_param=None, - tc_ssb_param=None, longitude_ssb_param=None, - latitude_ssb_param=None, polarization_ssb_param=None + self, + tc_lisa_param=None, + longitude_lisa_param=None, + latitude_lisa_param=None, + polarization_lisa_param=None, + tc_ssb_param=None, + longitude_ssb_param=None, + latitude_ssb_param=None, + polarization_ssb_param=None, ): - params = [tc_lisa_param, longitude_lisa_param, - latitude_lisa_param, polarization_lisa_param, - tc_ssb_param, longitude_ssb_param, - latitude_ssb_param, polarization_ssb_param] + params = [ + tc_lisa_param, + longitude_lisa_param, + latitude_lisa_param, + polarization_lisa_param, + tc_ssb_param, + longitude_ssb_param, + latitude_ssb_param, + polarization_ssb_param, + ] for index in range(len(params)): if params[index] is None: key = list(self.default_params_name.keys())[index] @@ -2615,10 +2861,18 @@ def __init__( self.longitude_ssb_param = params[5] self.latitude_ssb_param = params[6] self.polarization_ssb_param = params[7] - self._inputs = [self.tc_ssb_param, self.longitude_ssb_param, - self.latitude_ssb_param, self.polarization_ssb_param] - self._outputs = [self.tc_lisa_param, self.longitude_lisa_param, - self.latitude_lisa_param, self.polarization_lisa_param] + self._inputs = [ + self.tc_ssb_param, + self.longitude_ssb_param, + self.latitude_ssb_param, + self.polarization_ssb_param, + ] + self._outputs = [ + self.tc_lisa_param, + self.longitude_lisa_param, + self.latitude_lisa_param, + self.polarization_lisa_param, + ] class GEOToLISA(LISAToGEO): @@ -2630,15 +2884,26 @@ class GEOToLISA(LISAToGEO): inverse_transform = inverse.transform def __init__( - self, tc_lisa_param=None, longitude_lisa_param=None, - latitude_lisa_param=None, polarization_lisa_param=None, - tc_geo_param=None, longitude_geo_param=None, - latitude_geo_param=None, polarization_geo_param=None + self, + tc_lisa_param=None, + longitude_lisa_param=None, + latitude_lisa_param=None, + polarization_lisa_param=None, + tc_geo_param=None, + longitude_geo_param=None, + latitude_geo_param=None, + polarization_geo_param=None, ): - params = [tc_lisa_param, longitude_lisa_param, - latitude_lisa_param, polarization_lisa_param, - tc_geo_param, longitude_geo_param, - latitude_geo_param, polarization_geo_param] + params = [ + tc_lisa_param, + longitude_lisa_param, + latitude_lisa_param, + polarization_lisa_param, + tc_geo_param, + longitude_geo_param, + latitude_geo_param, + polarization_geo_param, + ] for index in range(len(params)): if params[index] is None: key = list(self.default_params_name.keys())[index] @@ -2652,14 +2917,23 @@ def __init__( self.longitude_geo_param = params[5] self.latitude_geo_param = params[6] self.polarization_geo_param = params[7] - self._inputs = [self.tc_geo_param, self.longitude_geo_param, - self.latitude_geo_param, self.polarization_geo_param] - self._outputs = [self.tc_lisa_param, self.longitude_lisa_param, - self.latitude_lisa_param, self.polarization_lisa_param] + self._inputs = [ + self.tc_geo_param, + self.longitude_geo_param, + self.latitude_geo_param, + self.polarization_geo_param, + ] + self._outputs = [ + self.tc_lisa_param, + self.longitude_lisa_param, + self.latitude_lisa_param, + self.polarization_lisa_param, + ] class Exponent(Log): - """Applies an exponent transform to an `inputvar` parameter. + """ + Applies an exponent transform to an `inputvar` parameter. This is the inverse of the log transform. @@ -2669,6 +2943,7 @@ class Exponent(Log): The name of the parameter to transform. outputvar : str The name of the transformed parameter. + """ name = "exponent" @@ -2679,11 +2954,12 @@ class Exponent(Log): inverse_jacobian = inverse.jacobian def __init__(self, inputvar, outputvar): - super(Exponent, self).__init__(outputvar, inputvar) + super().__init__(outputvar, inputvar) class Logistic(Logit): - r"""Applies a logistic transform from an `input` parameter to an `output` + r""" + Applies a logistic transform from an `input` parameter to an `output` parameter. This is the inverse of the logit transform. Typically, the output of the logistic function has range :math:`\in [0,1)`. @@ -2699,6 +2975,7 @@ class Logistic(Logit): frange : tuple or distributions.bounds.Bounds, optional The range of the output parameter. Can be any finite interval. Default is (0., 1.). + """ name = "logistic" @@ -2709,7 +2986,7 @@ class Logistic(Logit): inverse_jacobian = inverse.jacobian def __init__(self, inputvar, outputvar, codomain=(0.0, 1.0)): - super(Logistic, self).__init__(outputvar, inputvar, domain=codomain) + super().__init__(outputvar, inputvar, domain=codomain) @property def bounds(self): @@ -2717,9 +2994,9 @@ def bounds(self): return self._bounds @classmethod - def from_config(cls, cp, section, outputs, - skip_opts=None, additional_opts=None): - """Initializes a Logistic transform from the given section. + def from_config(cls, cp, section, outputs, skip_opts=None, additional_opts=None): + """ + Initializes a Logistic transform from the given section. The section must specify an input and output variable name. The codomain of the output may be specified using `min-{output}`, @@ -2755,6 +3032,7 @@ def from_config(cls, cp, section, outputs, ------- cls An instance of the class. + """ # pull out the minimum, maximum values of the output variable outputvar = cp.get_opt_tag(section, "output", outputs) @@ -2765,26 +3043,25 @@ def from_config(cls, cp, section, outputs, else: additional_opts = additional_opts.copy() s = "-".join([section, outputs]) - opt = "min-{}".format(outputvar) + opt = f"min-{outputvar}" if cp.has_option(s, opt): a = cp.get_opt_tag(section, opt, outputs) skip_opts.append(opt) else: a = None - opt = "max-{}".format(outputvar) + opt = f"max-{outputvar}" if cp.has_option(s, opt): b = cp.get_opt_tag(section, opt, outputs) skip_opts.append(opt) else: b = None - if a is None and b is not None or b is None and a is not None: + if (a is None and b is not None) or (b is None and a is not None): raise ValueError( - "if providing a min(max)-{}, must also provide " - "a max(min)-{}".format(outputvar, outputvar) + f"if providing a min(max)-{outputvar}, must also provide a max(min)-{outputvar}" ) - elif a is not None: + if a is not None: additional_opts.update({"codomain": (float(a), float(b))}) - return super(Logistic, cls).from_config( + return super().from_config( cp, section, outputs, skip_opts, additional_opts ) @@ -2907,12 +3184,12 @@ def from_config(cls, cp, section, outputs, ] ) -common_cbc_transforms = common_cbc_forward_transforms \ - + common_cbc_inverse_transforms +common_cbc_transforms = common_cbc_forward_transforms + common_cbc_inverse_transforms def get_common_cbc_transforms(requested_params, variable_args, valid_params=None): - """Determines if any additional parameters from the InferenceFile are + """ + Determines if any additional parameters from the InferenceFile are needed to get derived parameters that user has asked for. First it will try to add any base parameters that are required to calculate @@ -2934,6 +3211,7 @@ def get_common_cbc_transforms(requested_params, variable_args, valid_params=None Updated list of parameters that user wants. all_c : list List of BaseTransforms to apply. + """ variable_args = ( set(variable_args) if not isinstance(variable_args, set) else variable_args @@ -2959,8 +3237,9 @@ def get_common_cbc_transforms(requested_params, variable_args, valid_params=None # calculated from base parameters from_base_c = [] for converter in common_cbc_inverse_transforms: - if converter.outputs.issubset(variable_args) or \ - converter.outputs.isdisjoint(requested_params): + if converter.outputs.issubset(variable_args) or converter.outputs.isdisjoint( + requested_params + ): continue intersect = converter.outputs.intersection(requested_params) if ( @@ -2992,7 +3271,8 @@ def get_common_cbc_transforms(requested_params, variable_args, valid_params=None def apply_transforms(samples, transforms, inverse=False): - """Applies a list of BaseTransform instances on a mapping object. + """ + Applies a list of BaseTransform instances on a mapping object. Parameters ---------- @@ -3009,6 +3289,7 @@ def apply_transforms(samples, transforms, inverse=False): ------- samples : {FieldArray, dict} Mapping object with transforms applied. Same type as input. + """ if inverse: transforms = transforms[::-1] @@ -3024,7 +3305,8 @@ def apply_transforms(samples, transforms, inverse=False): def compute_jacobian(samples, transforms, inverse=False): - """Computes the jacobian of the list of transforms at the given sample + """ + Computes the jacobian of the list of transforms at the given sample points. Parameters @@ -3041,6 +3323,7 @@ def compute_jacobian(samples, transforms, inverse=False): ------- float : The product of the jacobians of all fo the transforms. + """ j = 1.0 if inverse: @@ -3053,7 +3336,8 @@ def compute_jacobian(samples, transforms, inverse=False): def order_transforms(transforms): - """Orders transforms to ensure proper chaining. + """ + Orders transforms to ensure proper chaining. For example, if `transforms = [B, A, C]`, and `A` produces outputs needed by `B`, the transforms will be re-rorderd to `[A, B, C]`. @@ -3068,9 +3352,10 @@ def order_transforms(transforms): list : List of transformed ordered such that forward transforms can be carried out without error. + """ # get a set of all inputs and all outputs - outputs = set().union(*[set(t.outputs)-set(t.inputs) for t in transforms]) + outputs = set().union(*[set(t.outputs) - set(t.inputs) for t in transforms]) out = [] remaining = [t for t in transforms] while remaining: @@ -3087,7 +3372,8 @@ def order_transforms(transforms): def read_transforms_from_config(cp, section="transforms"): - """Returns a list of PyCBC transform instances for a section in the + """ + Returns a list of PyCBC transform instances for a section in the given configuration file. If the transforms are nested (i.e., the output of one transform is the @@ -3105,6 +3391,7 @@ def read_transforms_from_config(cp, section="transforms"): ------- list A list of the parsed transforms. + """ trans = [] for subsection in cp.get_subsections(section): diff --git a/pycbc/types/__init__.py b/pycbc/types/__init__.py index 891c5bdb8dc..f375fde36ea 100644 --- a/pycbc/types/__init__.py +++ b/pycbc/types/__init__.py @@ -1,5 +1,5 @@ +from .aligned import check_aligned from .array import * -from .timeseries import * from .frequencyseries import * from .optparse import * -from .aligned import check_aligned +from .timeseries import * diff --git a/pycbc/types/aligned.py b/pycbc/types/aligned.py index 6c7bcee0b13..a26faca5106 100644 --- a/pycbc/types/aligned.py +++ b/pycbc/types/aligned.py @@ -27,28 +27,33 @@ whether or not its memory is aligned. It further provides functions for creating zeros and empty (unitialized) arrays with this class. """ + import numpy as _np + from pycbc import PYCBC_ALIGNMENT + def check_aligned(ndarr): - return ((ndarr.ctypes.data % PYCBC_ALIGNMENT) == 0) + return (ndarr.ctypes.data % PYCBC_ALIGNMENT) == 0 + def zeros(n, dtype): d = _np.dtype(dtype) - nbytes = (d.itemsize)*int(n) - tmp = _np.zeros(nbytes+PYCBC_ALIGNMENT, dtype=_np.uint8) - address = tmp.__array_interface__['data'][0] - offset = (PYCBC_ALIGNMENT - address%PYCBC_ALIGNMENT)%PYCBC_ALIGNMENT - ret_ary = tmp[offset:offset+nbytes].view(dtype=d) + nbytes = (d.itemsize) * int(n) + tmp = _np.zeros(nbytes + PYCBC_ALIGNMENT, dtype=_np.uint8) + address = tmp.__array_interface__["data"][0] + offset = (PYCBC_ALIGNMENT - address % PYCBC_ALIGNMENT) % PYCBC_ALIGNMENT + ret_ary = tmp[offset : offset + nbytes].view(dtype=d) del tmp return ret_ary + def empty(n, dtype): d = _np.dtype(dtype) - nbytes = (d.itemsize)*int(n) - tmp = _np.empty(nbytes+PYCBC_ALIGNMENT, dtype=_np.uint8) - address = tmp.__array_interface__['data'][0] - offset = (PYCBC_ALIGNMENT - address%PYCBC_ALIGNMENT)%PYCBC_ALIGNMENT - ret_ary = tmp[offset:offset+nbytes].view(dtype=d) + nbytes = (d.itemsize) * int(n) + tmp = _np.empty(nbytes + PYCBC_ALIGNMENT, dtype=_np.uint8) + address = tmp.__array_interface__["data"][0] + offset = (PYCBC_ALIGNMENT - address % PYCBC_ALIGNMENT) % PYCBC_ALIGNMENT + ret_ary = tmp[offset : offset + nbytes].view(dtype=d) del tmp return ret_ary diff --git a/pycbc/types/array.py b/pycbc/types/array.py index dc646c43f97..64d94d567f5 100644 --- a/pycbc/types/array.py +++ b/pycbc/types/array.py @@ -26,126 +26,141 @@ This modules provides a device independent Array class based on PyCUDA and Numpy. """ -BACKEND_PREFIX="pycbc.types.array_" +BACKEND_PREFIX = "pycbc.types.array_" import os as _os - from functools import wraps import h5py - import numpy as _numpy -from numpy import float32, float64, complex64, complex128, ones +from numpy import complex64, complex128, float32, float64, ones from numpy.linalg import norm import pycbc.scheme as _scheme -from pycbc.scheme import schemed, cpuonly -from pycbc.opt import LimitedSizeDict from pycbc.libutils import import_optional +from pycbc.opt import LimitedSizeDict +from pycbc.scheme import cpuonly, schemed -_lal = import_optional('lal') +_lal = import_optional("lal") #! FIXME: the uint32 datatype has not been fully tested, # we should restrict any functions that do not allow an # array of uint32 integers -_ALLOWED_DTYPES = [_numpy.float32, _numpy.float64, _numpy.complex64, - _numpy.complex128, _numpy.uint32, _numpy.int32, int] +_ALLOWED_DTYPES = [ + _numpy.float32, + _numpy.float64, + _numpy.complex64, + _numpy.complex128, + _numpy.uint32, + _numpy.int32, + int, +] try: _ALLOWED_SCALARS = [int, long, float, complex] + _ALLOWED_DTYPES except NameError: _ALLOWED_SCALARS = [int, float, complex] + _ALLOWED_DTYPES + def _convert_to_scheme(ary): if not isinstance(ary._scheme, _scheme.mgr.state.__class__): converted_array = Array(ary, dtype=ary._data.dtype) ary._data = converted_array._data ary._scheme = _scheme.mgr.state - + + def _convert(func): @wraps(func) def convert(self, *args, **kwargs): _convert_to_scheme(self) return func(self, *args, **kwargs) + return convert - + + def _nocomplex(func): @wraps(func) def nocomplex(self, *args, **kwargs): - if self.kind == 'real': + if self.kind == "real": return func(self, *args, **kwargs) - else: - raise TypeError( func.__name__ + " does not support complex types") + raise TypeError(func.__name__ + " does not support complex types") + return nocomplex + def _noreal(func): @wraps(func) def noreal(self, *args, **kwargs): - if self.kind == 'complex': + if self.kind == "complex": return func(self, *args, **kwargs) - else: - raise TypeError( func.__name__ + " does not support real types") + raise TypeError(func.__name__ + " does not support real types") + return noreal + def force_precision_to_match(scalar, precision): if _numpy.iscomplexobj(scalar): - if precision == 'single': + if precision == "single": return _numpy.complex64(scalar) - else: - return _numpy.complex128(scalar) - else: - if precision == 'single': - return _numpy.float32(scalar) - else: - return _numpy.float64(scalar) + return _numpy.complex128(scalar) + if precision == "single": + return _numpy.float32(scalar) + return _numpy.float64(scalar) + def common_kind(*dtypes): for dtype in dtypes: - if dtype.kind == 'c': + if dtype.kind == "c": return dtype return dtypes[0] - -@schemed(BACKEND_PREFIX) + + +@schemed(BACKEND_PREFIX) def _to_device(array): - """ Move input to device """ + """Move input to device""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + + @schemed(BACKEND_PREFIX) def _copy_base_array(array): - """ Copy a backend array""" + """Copy a backend array""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + @schemed(BACKEND_PREFIX) def _scheme_matches_base_array(array): - """ Check that input matches array type for scheme """ + """Check that input matches array type for scheme""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + def check_same_len_precision(a, b): - """Check that the two arguments have the same length and precision. + """ + Check that the two arguments have the same length and precision. Raises ValueError if they do not. """ if len(a) != len(b): - msg = 'lengths do not match ({} vs {})'.format( - len(a), len(b)) + msg = f"lengths do not match ({len(a)} vs {len(b)})" raise ValueError(msg) if a.precision != b.precision: - msg = 'precisions do not match ({} vs {})'.format( - a.precision, b.precision) + msg = f"precisions do not match ({a.precision} vs {b.precision})" raise TypeError(msg) -class Array(object): - """Array used to do numeric calculations on a various compute + +class Array: + """ + Array used to do numeric calculations on a various compute devices. It is a convience wrapper around numpy, and pycuda. """ def __init__(self, initial_array, dtype=None, copy=True): - """ initial_array: An array-like object as specified by NumPy, this + """ + initial_array: An array-like object as specified by NumPy, this also includes instances of an underlying data type as described in section 3 or an instance of the PYCBC Array class itself. This object is used to populate the data of the array. @@ -158,67 +173,66 @@ def __init__(self, initial_array, dtype=None, copy=True): created, and so all arguments that would force a copy are ignored. The default is to copy the given object. """ - self._scheme=_scheme.mgr.state + self._scheme = _scheme.mgr.state self._saved = LimitedSizeDict(size_limit=2**5) - - #Unwrap initial_array + + # Unwrap initial_array if isinstance(initial_array, Array): initial_array = initial_array._data if not copy: if not _scheme_matches_base_array(initial_array): raise TypeError("Cannot avoid a copy of this array") - else: - self._data = initial_array + self._data = initial_array # Check that the dtype is supported. if self._data.dtype not in _ALLOWED_DTYPES: - raise TypeError(str(self._data.dtype) + ' is not supported') + raise TypeError(str(self._data.dtype) + " is not supported") if dtype and dtype != self._data.dtype: raise TypeError("Can only set dtype when allowed to copy data") - if copy: # First we will check the dtype that we are given - if not hasattr(initial_array, 'dtype'): + if not hasattr(initial_array, "dtype"): initial_array = _numpy.array(initial_array) # Determine the dtype to use - if dtype is not None: + if dtype is not None: dtype = _numpy.dtype(dtype) if dtype not in _ALLOWED_DTYPES: - raise TypeError(str(dtype) + ' is not supported') - if dtype.kind != 'c' and initial_array.dtype.kind == 'c': - raise TypeError(str(initial_array.dtype) + ' cannot be cast as ' + str(dtype)) + raise TypeError(str(dtype) + " is not supported") + if dtype.kind != "c" and initial_array.dtype.kind == "c": + raise TypeError( + str(initial_array.dtype) + " cannot be cast as " + str(dtype) + ) elif initial_array.dtype in _ALLOWED_DTYPES: dtype = initial_array.dtype + elif initial_array.dtype.kind == "c": + dtype = complex128 else: - if initial_array.dtype.kind == 'c': - dtype = complex128 - else: - dtype = float64 - + dtype = float64 + # Cast to the final dtype if needed if initial_array.dtype != dtype: initial_array = initial_array.astype(dtype) - - #Create new instance with initial_array as initialization. + + # Create new instance with initial_array as initialization. if issubclass(type(self._scheme), _scheme.CPUScheme): - if hasattr(initial_array, 'get'): + if hasattr(initial_array, "get"): self._data = _numpy.array(initial_array.get()) else: self._data = _numpy.array(initial_array, dtype=dtype, ndmin=1) elif _scheme_matches_base_array(initial_array): - self._data = _copy_base_array(initial_array) # pylint:disable=assignment-from-no-return + self._data = _copy_base_array(initial_array) # pylint:disable=assignment-from-no-return else: initial_array = _numpy.array(initial_array, dtype=dtype, ndmin=1) - self._data = _to_device(initial_array) # pylint:disable=assignment-from-no-return + self._data = _to_device(initial_array) # pylint:disable=assignment-from-no-return def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): inputs = [i.numpy() if isinstance(i, Array) else i for i in inputs] ret = getattr(ufunc, method)(*inputs, **kwargs) - if hasattr(ret, 'shape') and ret.shape == self.shape: + if hasattr(ret, "shape") and ret.shape == self.shape: ret = self._return(ret) return ret @@ -231,7 +245,7 @@ def __array__(self, dtype=None): @property def shape(self): return self._data.shape - + def _memoize_single(func): @wraps(func) def memoize_single(self, arg): @@ -240,28 +254,31 @@ def memoize_single(self, arg): if badh in self._saved: return self._saved[badh] - res = func(self, arg) # pylint:disable=not-callable + res = func(self, arg) # pylint:disable=not-callable self._saved[badh] = res return res + return memoize_single def _returnarray(func): @wraps(func) def returnarray(self, *args, **kwargs): - return Array(func(self, *args, **kwargs), copy=False) # pylint:disable=not-callable + return Array(func(self, *args, **kwargs), copy=False) # pylint:disable=not-callable + return returnarray def _returntype(func): @wraps(func) def returntype(self, *args, **kwargs): - ary = func(self, *args, **kwargs) # pylint:disable=not-callable + ary = func(self, *args, **kwargs) # pylint:disable=not-callable if ary is NotImplemented: return NotImplemented return self._return(ary) + return returntype - + def _return(self, ary): - """Wrap the ary to return an Array type """ + """Wrap the ary to return an Array type""" if isinstance(ary, Array): return ary return Array(ary, copy=False) @@ -274,7 +291,7 @@ def checkother(self, *args): self._typecheck(other) if type(other) in _ALLOWED_SCALARS: other = force_precision_to_match(other, self.precision) - nargs +=(other,) + nargs += (other,) elif isinstance(other, type(self)) or type(other) is Array: check_same_len_precision(self, other) _convert_to_scheme(other) @@ -282,7 +299,8 @@ def checkother(self, *args): else: return NotImplemented - return func(self, *nargs) # pylint:disable=not-callable + return func(self, *nargs) # pylint:disable=not-callable + return checkother def _vcheckother(func): @@ -296,11 +314,12 @@ def vcheckother(self, *args): _convert_to_scheme(other) nargs += (other._data,) else: - raise TypeError('array argument required') + raise TypeError("array argument required") + + return func(self, *nargs) # pylint:disable=not-callable - return func(self, *nargs) # pylint:disable=not-callable return vcheckother - + def _vrcheckother(func): @wraps(func) def vrcheckother(self, *args): @@ -311,94 +330,96 @@ def vrcheckother(self, *args): _convert_to_scheme(other) nargs += (other._data,) else: - raise TypeError('array argument required') + raise TypeError("array argument required") + + return func(self, *nargs) # pylint:disable=not-callable - return func(self, *nargs) # pylint:disable=not-callable return vrcheckother def _icheckother(func): @wraps(func) def icheckother(self, other): - """ Checks the input to in-place operations """ + """Checks the input to in-place operations""" self._typecheck(other) if type(other) in _ALLOWED_SCALARS: - if self.kind == 'real' and type(other) == complex: - raise TypeError('dtypes are incompatible') + if self.kind == "real" and type(other) == complex: + raise TypeError("dtypes are incompatible") other = force_precision_to_match(other, self.precision) elif isinstance(other, type(self)) or type(other) is Array: check_same_len_precision(self, other) - if self.kind == 'real' and other.kind == 'complex': - raise TypeError('dtypes are incompatible') + if self.kind == "real" and other.kind == "complex": + raise TypeError("dtypes are incompatible") _convert_to_scheme(other) other = other._data else: return NotImplemented - return func(self, other) # pylint:disable=not-callable + return func(self, other) # pylint:disable=not-callable + return icheckother def _typecheck(self, other): - """ Additional typechecking for other. Placeholder for use by derived - types. """ - pass + Additional typechecking for other. Placeholder for use by derived + types. + """ @_returntype @_convert @_checkother - def __mul__(self,other): - """ Multiply by an Array or a scalar and return an Array. """ + def __mul__(self, other): + """Multiply by an Array or a scalar and return an Array.""" return self._data * other __rmul__ = __mul__ @_convert @_icheckother - def __imul__(self,other): - """ Multiply by an Array or a scalar and return an Array. """ + def __imul__(self, other): + """Multiply by an Array or a scalar and return an Array.""" self._data *= other return self @_returntype @_convert @_checkother - def __add__(self,other): - """ Add Array to Array or scalar and return an Array. """ + def __add__(self, other): + """Add Array to Array or scalar and return an Array.""" return self._data + other __radd__ = __add__ - + def fill(self, value): self._data.fill(value) @_convert @_icheckother - def __iadd__(self,other): - """ Add Array to Array or scalar and return an Array. """ + def __iadd__(self, other): + """Add Array to Array or scalar and return an Array.""" self._data += other return self @_convert @_checkother @_returntype - def __truediv__(self,other): - """ Divide Array by Array or scalar and return an Array. """ + def __truediv__(self, other): + """Divide Array by Array or scalar and return an Array.""" return self._data / other @_returntype @_convert @_checkother - def __rtruediv__(self,other): - """ Divide Array by Array or scalar and return an Array. """ + def __rtruediv__(self, other): + """Divide Array by Array or scalar and return an Array.""" return self._data.__rtruediv__(other) @_convert @_icheckother - def __itruediv__(self,other): - """ Divide Array by Array or scalar and return an Array. """ + def __itruediv__(self, other): + """Divide Array by Array or scalar and return an Array.""" self._data /= other return self - + __div__ = __truediv__ __idiv__ = __itruediv__ __rdiv__ = __rtruediv__ @@ -406,55 +427,55 @@ def __itruediv__(self,other): @_returntype @_convert def __neg__(self): - """ Return negation of self """ - return - self._data + """Return negation of self""" + return -self._data @_returntype @_convert @_checkother - def __sub__(self,other): - """ Subtract Array or scalar from Array and return an Array. """ + def __sub__(self, other): + """Subtract Array or scalar from Array and return an Array.""" return self._data - other @_returntype @_convert @_checkother - def __rsub__(self,other): - """ Subtract Array or scalar from Array and return an Array. """ + def __rsub__(self, other): + """Subtract Array or scalar from Array and return an Array.""" return self._data.__rsub__(other) @_convert @_icheckother - def __isub__(self,other): - """ Subtract Array or scalar from Array and return an Array. """ + def __isub__(self, other): + """Subtract Array or scalar from Array and return an Array.""" self._data -= other return self @_returntype @_convert @_checkother - def __pow__(self,other): - """ Exponentiate Array by scalar """ - return self._data ** other + def __pow__(self, other): + """Exponentiate Array by scalar""" + return self._data**other @_returntype @_convert def __abs__(self): - """ Return absolute value of Array """ + """Return absolute value of Array""" return abs(self._data) def __len__(self): - """ Return length of Array """ + """Return length of Array""" return len(self._data) def __str__(self): return str(self._data) - + @property def ndim(self): return self._data.ndim - def __eq__(self,other): + def __eq__(self, other): """ This is the Python special method invoked whenever the '==' comparison is used. It will return true if the data of two @@ -487,8 +508,8 @@ def __eq__(self,other): ------- boolean: 'True' if the types, dtypes, lengths, and data of the two objects are each identical. - """ + """ # Writing the first test as below allows this method to be safely # called from subclasses. if type(self) != type(other): @@ -512,7 +533,7 @@ def __eq__(self,other): # of that array of booleans is True. return (sary == oary).all() - def almost_equal_elem(self,other,tol,relative=True): + def almost_equal_elem(self, other, tol, relative=True): """ Compare whether two array types are almost equal, element by element. @@ -550,14 +571,15 @@ def almost_equal_elem(self,other,tol,relative=True): Returns ------- - boolean + boolean 'True' if the data agree within the tolerance, as interpreted by the 'relative' keyword, and if the types, lengths, and dtypes are exactly the same. + """ # Check that the tolerance is non-negative and raise an # exception otherwise. - if (tol<0): + if tol < 0: raise ValueError("Tolerance cannot be negative") # Check that the meta-data agree; the type check is written in # this way so that this method may be safely called from @@ -572,15 +594,15 @@ def almost_equal_elem(self,other,tol,relative=True): # The numpy() method will move any GPU memory onto the CPU. # Slow, but the user was warned. - diff = abs(self.numpy()-other.numpy()) + diff = abs(self.numpy() - other.numpy()) if relative: - cmpary = tol*abs(self.numpy()) + cmpary = tol * abs(self.numpy()) else: - cmpary = tol*ones(len(self),dtype=self.dtype) + cmpary = tol * ones(len(self), dtype=self.dtype) - return (diff<=cmpary).all() + return (diff <= cmpary).all() - def almost_equal_norm(self,other,tol,relative=True): + def almost_equal_norm(self, other, tol, relative=True): """ Compare whether two array types are almost equal, normwise. @@ -604,7 +626,7 @@ def almost_equal_norm(self,other,tol,relative=True): other another Python object, that should be tested for almost-equality with 'self', based on their norms. - tol + tol a non-negative number, the tolerance, which is interpreted as either a relative tolerance (the default) or an absolute tolerance. @@ -619,10 +641,11 @@ def almost_equal_norm(self,other,tol,relative=True): 'True' if the data agree within the tolerance, as interpreted by the 'relative' keyword, and if the types, lengths, and dtypes are exactly the same. + """ # Check that the tolerance is non-negative and raise an # exception otherwise. - if (tol<0): + if tol < 0: raise ValueError("Tolerance cannot be negative") # Check that the meta-data agree; the type check is written in # this way so that this method may be safely called from @@ -637,36 +660,35 @@ def almost_equal_norm(self,other,tol,relative=True): # The numpy() method will move any GPU memory onto the CPU. # Slow, but the user was warned. - diff = self.numpy()-other.numpy() + diff = self.numpy() - other.numpy() dnorm = norm(diff) if relative: - return (dnorm <= tol*norm(self)) - else: - return (dnorm <= tol) + return dnorm <= tol * norm(self) + return dnorm <= tol @_returntype @_convert def real(self): - """ Return real part of Array """ + """Return real part of Array""" return Array(self._data.real, copy=True) @_returntype @_convert def imag(self): - """ Return imaginary part of Array """ + """Return imaginary part of Array""" return Array(self._data.imag, copy=True) @_returntype @_convert def conj(self): - """ Return complex conjugate of Array. """ + """Return complex conjugate of Array.""" return self._data.conj() - + @_returntype @_convert @schemed(BACKEND_PREFIX) def squared_norm(self): - """ Return the elementwise squared norm of the array """ + """Return the elementwise squared norm of the array""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -676,7 +698,8 @@ def squared_norm(self): @_convert @schemed(BACKEND_PREFIX) def multiply_and_add(self, other, mult_fac): - """ Return other multiplied by mult_fac and with self added. + """ + Return other multiplied by mult_fac and with self added. Self is modified in place and returned as output. Precisions of inputs must match. """ @@ -688,8 +711,7 @@ def multiply_and_add(self, other, mult_fac): @_convert @schemed(BACKEND_PREFIX) def inner(self, other): - """ Return the inner product of the array with complex conjugation. - """ + """Return the inner product of the array with complex conjugation.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -698,16 +720,15 @@ def inner(self, other): @_convert @schemed(BACKEND_PREFIX) def vdot(self, other): - """ Return the inner product of the array with complex conjugation. - """ + """Return the inner product of the array with complex conjugation.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @_convert @schemed(BACKEND_PREFIX) - def clear(self): - """ Clear out the values of the array. """ + def clear(self): + """Clear out the values of the array.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -716,8 +737,7 @@ def clear(self): @_convert @schemed(BACKEND_PREFIX) def weighted_inner(self, other, weight): - """ Return the inner product of the array with complex conjugation. - """ + """Return the inner product of the array with complex conjugation.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -725,7 +745,7 @@ def weighted_inner(self, other, weight): @_convert @schemed(BACKEND_PREFIX) def sum(self): - """ Return the sum of the the array. """ + """Return the sum of the the array.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -734,25 +754,25 @@ def sum(self): @_convert @schemed(BACKEND_PREFIX) def cumsum(self): - """ Return the cumulative sum of the the array. """ + """Return the cumulative sum of the the array.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @_convert @_nocomplex @schemed(BACKEND_PREFIX) def max(self): - """ Return the maximum value in the array. """ + """Return the maximum value in the array.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @_convert @_nocomplex @schemed(BACKEND_PREFIX) def max_loc(self): - """Return the maximum value in the array along with the index location """ + """Return the maximum value in the array along with the index location""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -760,7 +780,7 @@ def max_loc(self): @_convert @schemed(BACKEND_PREFIX) def abs_arg_max(self): - """ Return location of the maximum argument max """ + """Return location of the maximum argument max""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) @@ -777,11 +797,11 @@ def abs_max_loc(self): @_nocomplex @schemed(BACKEND_PREFIX) def min(self): - """ Return the maximum value in the array. """ + """Return the maximum value in the array.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @_returnarray @_convert @schemed(BACKEND_PREFIX) @@ -794,15 +814,16 @@ def take(self, indices): @_vcheckother @schemed(BACKEND_PREFIX) def dot(self, other): - """ Return the dot product""" + """Return the dot product""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @schemed(BACKEND_PREFIX) def _getvalue(self, index): - """Helper function to return a single value from an array. May be very - slow if the memory is on a gpu. + """ + Helper function to return a single value from an array. May be very + slow if the memory is on a gpu. """ err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" @@ -812,50 +833,47 @@ def _getvalue(self, index): @_returntype def _getslice(self, index): return self._return(self._data[index]) - + @_convert def __getitem__(self, index): - """ Return items from the Array. This not guaranteed to be fast for - returning single values. + """ + Return items from the Array. This not guaranteed to be fast for + returning single values. """ if isinstance(index, slice): return self._getslice(index) - else: - return self._getvalue(index) + return self._getvalue(index) @_convert def resize(self, new_size): - """Resize self to new_size - """ + """Resize self to new_size""" if new_size == len(self): return + self._saved = LimitedSizeDict(size_limit=2**5) + new_arr = zeros(new_size, dtype=self.dtype) + if len(self) <= new_size: + new_arr[0 : len(self)] = self else: - self._saved = LimitedSizeDict(size_limit=2**5) - new_arr = zeros(new_size, dtype=self.dtype) - if len(self) <= new_size: - new_arr[0:len(self)] = self - else: - new_arr[:] = self[0:new_size] - - self._data = new_arr._data + new_arr[:] = self[0:new_size] + + self._data = new_arr._data @_convert def roll(self, shift): - """shift vector - """ + """Shift vector""" new_arr = zeros(len(self), dtype=self.dtype) if shift < 0: shift = shift - len(self) * (shift // len(self)) - + if shift == 0: return - - new_arr[0:shift] = self[len(self)-shift: len(self)] - new_arr[shift:len(self)] = self[0:len(self)-shift] - + + new_arr[0:shift] = self[len(self) - shift : len(self)] + new_arr[shift : len(self)] = self[0 : len(self) - shift] + self._saved = LimitedSizeDict(size_limit=2**5) - + self._data = new_arr._data @_returntype @@ -863,31 +881,31 @@ def roll(self, shift): def astype(self, dtype): if _numpy.dtype(self.dtype) == _numpy.dtype(dtype): return self - else: - return self._data.astype(dtype) - + return self._data.astype(dtype) + @schemed(BACKEND_PREFIX) def _copy(self, self_ref, other_ref): - """Helper function to copy between two arrays. The arrays references - should be bare array types and not `Array` class instances. + """ + Helper function to copy between two arrays. The arrays references + should be bare array types and not `Array` class instances. """ err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @_convert def __setitem__(self, index, other): - if isinstance(other,Array): + if isinstance(other, Array): _convert_to_scheme(other) - if self.kind == 'real' and other.kind == 'complex': - raise ValueError('Cannot set real value with complex') + if self.kind == "real" and other.kind == "complex": + raise ValueError("Cannot set real value with complex") - if isinstance(index,slice): + if isinstance(index, slice): self_ref = self._data[index] other_ref = other._data else: - self_ref = self._data[index:index+1] + self_ref = self._data[index : index + 1] other_ref = other._data self._copy(self_ref, other_ref) @@ -896,36 +914,34 @@ def __setitem__(self, index, other): if isinstance(index, slice): self[index].fill(other) else: - self[index:index+1].fill(other) + self[index : index + 1].fill(other) else: - raise TypeError('Can only copy data from another Array') + raise TypeError("Can only copy data from another Array") @property def precision(self): if self.dtype == float32 or self.dtype == complex64: - return 'single' - else: - return 'double' - + return "single" + return "double" + @property def kind(self): if self.dtype == float32 or self.dtype == float64: - return 'real' - elif self.dtype == complex64 or self.dtype == complex128: - return 'complex' - else: - return 'unknown' + return "real" + if self.dtype == complex64 or self.dtype == complex128: + return "complex" + return "unknown" @property @_convert def data(self): - """Returns the internal python array """ + """Returns the internal python array""" return self._data @data.setter - def data(self,other): + def data(self, other): dtype = None - if hasattr(other,'dtype'): + if hasattr(other, "dtype"): dtype = other.dtype temp = Array(other, dtype=dtype) self._data = temp._data @@ -934,15 +950,15 @@ def data(self,other): @_convert @schemed(BACKEND_PREFIX) def ptr(self): - """ Returns a pointer to the memory of this array """ + """Returns a pointer to the memory of this array""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @property def itemsize(self): return self.dtype.itemsize - + @property def nbytes(self): return len(self.data) * self.itemsize @@ -951,23 +967,23 @@ def nbytes(self): @cpuonly @_convert def _swighelper(self): - """ Used internally by SWIG typemaps to ensure @_convert - is called and scheme is correct """ - return self; + Used internally by SWIG typemaps to ensure @_convert + is called and scheme is correct + """ + return self @_convert @schemed(BACKEND_PREFIX) def numpy(self): - """ Returns a Numpy Array that contains this data """ + """Returns a Numpy Array that contains this data""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) - + @_convert def lal(self): - """ Returns a LAL Object that contains this data """ - + """Returns a LAL Object that contains this data""" lal_data = None if self._data.dtype == float32: lal_data = _lal.CreateREAL4Vector(len(self)) @@ -985,7 +1001,7 @@ def lal(self): @property def dtype(self): return self._data.dtype - + def save(self, path, group=None): """ Save array to a Numpy .npy, hdf, or text file. When saving a complex array as @@ -997,8 +1013,8 @@ def save(self, path, group=None): ---------- path: string Destination file path. Must end with either .hdf, .npy or .txt. - - group: string + + group: string Additional name for internal storage use. Ex. hdf storage uses this as the key value. @@ -1006,34 +1022,37 @@ def save(self, path, group=None): ------ ValueError If path does not end in .npy or .txt. - """ + """ ext = _os.path.splitext(path)[1] - if ext == '.npy': + if ext == ".npy": _numpy.save(path, self.numpy()) - elif ext == '.txt': - if self.kind == 'real': + elif ext == ".txt": + if self.kind == "real": _numpy.savetxt(path, self.numpy()) - elif self.kind == 'complex': - output = _numpy.vstack((self.numpy().real, - self.numpy().imag)).T + elif self.kind == "complex": + output = _numpy.vstack((self.numpy().real, self.numpy().imag)).T _numpy.savetxt(path, output) - elif ext == '.hdf': - key = 'data' if group is None else group - with h5py.File(path, 'a') as f: - f.create_dataset(key, data=self.numpy(), compression='gzip', - compression_opts=9, shuffle=True) + elif ext == ".hdf": + key = "data" if group is None else group + with h5py.File(path, "a") as f: + f.create_dataset( + key, + data=self.numpy(), + compression="gzip", + compression_opts=9, + shuffle=True, + ) else: - raise ValueError('Path must end with .npy, .txt, or .hdf') - - @_convert + raise ValueError("Path must end with .npy, .txt, or .hdf") + + @_convert def trim_zeros(self): - """Remove the leading and trailing zeros. - """ + """Remove the leading and trailing zeros.""" tmp = self.numpy() - f = len(self)-len(_numpy.trim_zeros(tmp, trim='f')) - b = len(self)-len(_numpy.trim_zeros(tmp, trim='b')) - return self[f:len(self)-b] + f = len(self) - len(_numpy.trim_zeros(tmp, trim="f")) + b = len(self) - len(_numpy.trim_zeros(tmp, trim="b")) + return self[f : len(self) - b] @_returntype @_convert @@ -1047,67 +1066,74 @@ def view(self, dtype): ---------- dtype : numpy dtype (one of float32, float64, complex64 or complex128) The new dtype that should be used to interpret the bytes of self + """ return self._data.view(dtype) def copy(self): - """ Return copy of this array """ + """Return copy of this array""" return self._return(self.data.copy()) - + def __lt__(self, other): return self.numpy().__lt__(other) - + def __le__(self, other): return self.numpy().__le__(other) - + def __ne__(self, other): return self.numpy().__ne__(other) - + def __gt__(self, other): return self.numpy().__gt__(other) - + def __ge__(self, other): return self.numpy().__ge__(other) - + + # Convenience functions for determining dtypes def real_same_precision_as(data): - if data.precision == 'single': + if data.precision == "single": return float32 - elif data.precision == 'double': + if data.precision == "double": return float64 + def complex_same_precision_as(data): - if data.precision == 'single': + if data.precision == "single": return complex64 - elif data.precision == 'double': + if data.precision == "double": return complex128 + def _return_array(func): @wraps(func) def return_array(*args, **kwds): return Array(func(*args, **kwds), copy=False) + return return_array + @_return_array @schemed(BACKEND_PREFIX) def zeros(length, dtype=float64): - """ Return an Array filled with zeros. - """ + """Return an Array filled with zeros.""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + @_return_array @schemed(BACKEND_PREFIX) def empty(length, dtype=float64): - """ Return an empty Array (no initialization) - """ + """Return an empty Array (no initialization)""" err_msg = "This function is a stub that should be overridden using " err_msg += "the scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + def load_array(path, group=None): - """Load an Array from an HDF5, ASCII or Numpy file. The file type is + """ + Load an Array from an HDF5, ASCII or Numpy file. The file type is inferred from the file extension, which must be `.hdf`, `.txt` or `.npy`. For ASCII and Numpy files with a single column, a real array is returned. @@ -1131,24 +1157,28 @@ def load_array(path, group=None): If path does not end with a supported extension. For Numpy and ASCII input files, this is also raised if the array does not have 1 or 2 dimensions. + """ ext = _os.path.splitext(path)[1] - if ext == '.npy': + if ext == ".npy": data = _numpy.load(path) - elif ext == '.txt': + elif ext == ".txt": data = _numpy.loadtxt(path) - elif ext == '.hdf': - key = 'data' if group is None else group - with h5py.File(path, 'r') as f: + elif ext == ".hdf": + key = "data" if group is None else group + with h5py.File(path, "r") as f: array = Array(f[key]) return array else: - raise ValueError('Path must end with .npy, .hdf, or .txt') + raise ValueError("Path must end with .npy, .hdf, or .txt") if data.ndim == 1: return Array(data) - elif data.ndim == 2: - return Array(data[:,0] + 1j*data[:,1]) - - raise ValueError('File has %s dimensions, cannot convert to Array, \ - must be 1 (real) or 2 (complex)' % data.ndim) + if data.ndim == 2: + return Array(data[:, 0] + 1j * data[:, 1]) + + raise ValueError( + "File has %s dimensions, cannot convert to Array, \ + must be 1 (real) or 2 (complex)" + % data.ndim + ) diff --git a/pycbc/types/array_cuda.py b/pycbc/types/array_cuda.py index 2a9d330db45..aceef9047a0 100644 --- a/pycbc/types/array_cuda.py +++ b/pycbc/types/array_cuda.py @@ -21,45 +21,50 @@ # # ============================================================================= # -"""Pycuda based -""" +"""Pycuda based""" + +import numpy as np import pycuda import pycuda.driver +import pycuda.gpuarray from pycuda.elementwise import ElementwiseKernel +from pycuda.gpuarray import GPUArray, _get_common_dtype, empty from pycuda.reduction import ReductionKernel -from pycuda.tools import get_or_register_dtype -from pycuda.tools import context_dependent_memoize -from pycuda.tools import dtype_to_ctype -from pytools import match_precision -from pycuda.gpuarray import _get_common_dtype, empty, GPUArray -import pycuda.gpuarray from pycuda.scan import InclusiveScanKernel -import numpy as np +from pycuda.tools import ( + context_dependent_memoize, + dtype_to_ctype, + get_or_register_dtype, +) +from pytools import match_precision include_complex = """ #include """ + @context_dependent_memoize def get_cumsum_kernel(dtype): return InclusiveScanKernel(dtype, "a+b", preamble=include_complex) + def icumsum(vec): krnl = get_cumsum_kernel(vec.dtype) return krnl(vec) + @context_dependent_memoize def call_prepare(self, sz, allocator): MAX_BLOCK_COUNT = 1024 - SMALL_SEQ_COUNT = 4 + SMALL_SEQ_COUNT = 4 - if sz <= self.block_size*SMALL_SEQ_COUNT*MAX_BLOCK_COUNT: - total_block_size = SMALL_SEQ_COUNT*self.block_size + if sz <= self.block_size * SMALL_SEQ_COUNT * MAX_BLOCK_COUNT: + total_block_size = SMALL_SEQ_COUNT * self.block_size block_count = (sz + total_block_size - 1) // total_block_size seq_count = SMALL_SEQ_COUNT else: block_count = MAX_BLOCK_COUNT - macroblock_size = block_count*self.block_size + macroblock_size = block_count * self.block_size seq_count = (sz + macroblock_size - 1) // macroblock_size if block_count == 1: @@ -68,126 +73,175 @@ def call_prepare(self, sz, allocator): result = empty((block_count,), self.dtype_out, allocator) grid_size = (block_count, 1) - block_size = (self.block_size, 1, 1) + block_size = (self.block_size, 1, 1) return result, block_count, seq_count, grid_size, block_size -class LowerLatencyReductionKernel(ReductionKernel): - def __init__(self, dtype_out, - neutral, reduce_expr, map_expr=None, arguments=None, - name="reduce_kernel", keep=False, options=None, preamble=""): - ReductionKernel.__init__(self, dtype_out, - neutral, reduce_expr, map_expr, arguments, - name, keep, options, preamble) - - self.shared_size=self.block_size*self.dtype_out.itemsize +class LowerLatencyReductionKernel(ReductionKernel): + def __init__( + self, + dtype_out, + neutral, + reduce_expr, + map_expr=None, + arguments=None, + name="reduce_kernel", + keep=False, + options=None, + preamble="", + ): + ReductionKernel.__init__( + self, + dtype_out, + neutral, + reduce_expr, + map_expr, + arguments, + name, + keep, + options, + preamble, + ) + + self.shared_size = self.block_size * self.dtype_out.itemsize def __call__(self, *args, **kwargs): f = self.stage1_func - s1_invocation_args = [] + s1_invocation_args = [] for arg in args: s1_invocation_args.append(arg.gpudata) sz = args[0].size - result, block_count, seq_count, grid_size, block_size = call_prepare(self, sz, args[0].allocator) + result, block_count, seq_count, grid_size, block_size = call_prepare( + self, sz, args[0].allocator + ) - f(grid_size, block_size, None, - *([result.gpudata]+s1_invocation_args+[seq_count, sz]), - shared_size=self.shared_size) + f( + grid_size, + block_size, + None, + *([result.gpudata] + s1_invocation_args + [seq_count, sz]), + shared_size=self.shared_size, + ) while True: f = self.stage2_func sz = result.size result2 = result - result, block_count, seq_count, grid_size, block_size = call_prepare(self, sz, args[0].allocator) - - f(grid_size, block_size, None, - *([result.gpudata, result2.gpudata]+s1_invocation_args+[seq_count, sz]), - shared_size=self.shared_size) + result, block_count, seq_count, grid_size, block_size = call_prepare( + self, sz, args[0].allocator + ) + + f( + grid_size, + block_size, + None, + *( + [result.gpudata, result2.gpudata] + + s1_invocation_args + + [seq_count, sz] + ), + shared_size=self.shared_size, + ) if block_count == 1: return result - @context_dependent_memoize def get_norm_kernel(dtype_x, dtype_out): return ElementwiseKernel( - "%(tp_x)s *x, %(tp_z)s *z" % { - "tp_x": dtype_to_ctype(dtype_x), - "tp_z": dtype_to_ctype(dtype_out), - }, - "z[i] = norm(x[i])", - "normalize") + "%(tp_x)s *x, %(tp_z)s *z" + % { + "tp_x": dtype_to_ctype(dtype_x), + "tp_z": dtype_to_ctype(dtype_out), + }, + "z[i] = norm(x[i])", + "normalize", + ) + def squared_norm(self): a = self.data - dtype_out = match_precision(np.dtype('float64'), a.dtype) + dtype_out = match_precision(np.dtype("float64"), a.dtype) out = a._new_like_me(dtype=dtype_out) krnl = get_norm_kernel(a.dtype, dtype_out) krnl(a, out) - return out + return out + # FIXME: Write me! -#def multiply_and_add(self, other, mult_fac): +# def multiply_and_add(self, other, mult_fac): # """ # Return other multiplied by mult_fac and with self added. # Self will be modified in place. This requires all inputs to be of the same # precision. # """ - + + @context_dependent_memoize def get_weighted_inner_kernel(dtype_x, dtype_y, dtype_w, dtype_out): if (dtype_x == np.complex64) or (dtype_x == np.complex128): - inner_map="conj(x[i])*y[i]/w[i]" + inner_map = "conj(x[i])*y[i]/w[i]" else: - inner_map="x[i]*y[i]/w[i]" - return LowerLatencyReductionKernel(dtype_out, - neutral="0", - arguments="%(tp_x)s *x, %(tp_y)s *y, %(tp_w)s *w" % { - "tp_x": dtype_to_ctype(dtype_x), - "tp_y": dtype_to_ctype(dtype_y), - "tp_w": dtype_to_ctype(dtype_w), - }, - reduce_expr="a+b", - map_expr=inner_map, - name="weighted_inner") + inner_map = "x[i]*y[i]/w[i]" + return LowerLatencyReductionKernel( + dtype_out, + neutral="0", + arguments="%(tp_x)s *x, %(tp_y)s *y, %(tp_w)s *w" + % { + "tp_x": dtype_to_ctype(dtype_x), + "tp_y": dtype_to_ctype(dtype_y), + "tp_w": dtype_to_ctype(dtype_w), + }, + reduce_expr="a+b", + map_expr=inner_map, + name="weighted_inner", + ) + @context_dependent_memoize def get_inner_kernel(dtype_x, dtype_y, dtype_out): if (dtype_x == np.complex64) or (dtype_x == np.complex128): - inner_map="conj(x[i])*y[i]" + inner_map = "conj(x[i])*y[i]" else: - inner_map="x[i]*y[i]" - return LowerLatencyReductionKernel(dtype_out, - neutral="0", - arguments="%(tp_x)s *x, %(tp_y)s *y" % { - "tp_x": dtype_to_ctype(dtype_x), - "tp_y": dtype_to_ctype(dtype_y), - }, - reduce_expr="a+b", - map_expr=inner_map, - name="inner") + inner_map = "x[i]*y[i]" + return LowerLatencyReductionKernel( + dtype_out, + neutral="0", + arguments="%(tp_x)s *x, %(tp_y)s *y" + % { + "tp_x": dtype_to_ctype(dtype_x), + "tp_y": dtype_to_ctype(dtype_y), + }, + reduce_expr="a+b", + map_expr=inner_map, + name="inner", + ) + def inner(self, b): a = self.data - dtype_out = _get_common_dtype(a,b) + dtype_out = _get_common_dtype(a, b) krnl = get_inner_kernel(a.dtype, b.dtype, dtype_out) return krnl(a, b).get().max() - + + vdot = inner + def weighted_inner(self, b, w): if w is None: - return self.inner(b) + return self.inner(b) a = self.data dtype_out = _get_common_dtype(a, b) krnl = get_weighted_inner_kernel(a.dtype, b.dtype, w.dtype, dtype_out) return krnl(a, b, w).get().max() -# Define PYCUDA MAXLOC for both single and double precission ################## - + +# Define PYCUDA MAXLOC for both single and double precission ################## + maxloc_preamble = """ struct MAXLOCN{ @@ -235,123 +289,182 @@ def weighted_inner(self, b, w): } """ - -maxloc_preamble_single = """ + +maxloc_preamble_single = ( + """ #define MAXLOCN maxlocs #define TTYPE float #define LTYPE int -""" + maxloc_preamble +""" + + maxloc_preamble +) -maxloc_preamble_double = """ +maxloc_preamble_double = ( + """ #define MAXLOCN maxlocd #define TTYPE double #define LTYPE long -""" + maxloc_preamble - +""" + + maxloc_preamble +) + maxloc_dtype_double = np.dtype([("max", np.float64), ("loc", np.int64)]) maxloc_dtype_single = np.dtype([("max", np.float32), ("loc", np.int32)]) -if type(pycuda).__name__ not in ('MagicMock', '_MockModule'): +if type(pycuda).__name__ not in ("MagicMock", "_MockModule"): maxloc_dtype_single = get_or_register_dtype("maxlocs", dtype=maxloc_dtype_single) maxloc_dtype_double = get_or_register_dtype("maxlocd", dtype=maxloc_dtype_double) - mls = LowerLatencyReductionKernel(maxloc_dtype_single, neutral = "maxloc_start()", - reduce_expr="maxloc_red(a, b)", map_expr="maxloc_map(x[i], i)", - arguments="float *x", preamble=maxloc_preamble_single) - - mld = LowerLatencyReductionKernel(maxloc_dtype_double, neutral = "maxloc_start()", - reduce_expr="maxloc_red(a, b)", map_expr="maxloc_map(x[i], i)", - arguments="double *x", preamble=maxloc_preamble_double) - - max_loc_map = {'single':mls,'double':mld} - - - amls = LowerLatencyReductionKernel(maxloc_dtype_single, neutral = "maxloc_start()", - reduce_expr="maxloc_red(a, b)", map_expr="maxloc_map(abs(x[i]), i)", - arguments="float *x", preamble=maxloc_preamble_single) - - amld = LowerLatencyReductionKernel(maxloc_dtype_double, neutral = "maxloc_start()", - reduce_expr="maxloc_red(a, b)", map_expr="maxloc_map(abs(x[i]), i)", - arguments="double *x", preamble=maxloc_preamble_double) - - amlsc = LowerLatencyReductionKernel(maxloc_dtype_single, neutral = "maxloc_start()", - reduce_expr="maxloc_red(a, b)", map_expr="maxloc_map(abs(x[i]), i)", - arguments="pycuda::complex *x", preamble=maxloc_preamble_single) - - amldc = LowerLatencyReductionKernel(maxloc_dtype_double, neutral = "maxloc_start()", - reduce_expr="maxloc_red(a, b)", map_expr="maxloc_map(abs(x[i]), i)", - arguments="pycuda::complex *x", preamble=maxloc_preamble_double) - - abs_max_loc_map = {'single':{ 'real':amls, 'complex':amlsc }, 'double':{ 'real':amld, 'complex':amldc }} + mls = LowerLatencyReductionKernel( + maxloc_dtype_single, + neutral="maxloc_start()", + reduce_expr="maxloc_red(a, b)", + map_expr="maxloc_map(x[i], i)", + arguments="float *x", + preamble=maxloc_preamble_single, + ) + + mld = LowerLatencyReductionKernel( + maxloc_dtype_double, + neutral="maxloc_start()", + reduce_expr="maxloc_red(a, b)", + map_expr="maxloc_map(x[i], i)", + arguments="double *x", + preamble=maxloc_preamble_double, + ) + + max_loc_map = {"single": mls, "double": mld} + + amls = LowerLatencyReductionKernel( + maxloc_dtype_single, + neutral="maxloc_start()", + reduce_expr="maxloc_red(a, b)", + map_expr="maxloc_map(abs(x[i]), i)", + arguments="float *x", + preamble=maxloc_preamble_single, + ) + + amld = LowerLatencyReductionKernel( + maxloc_dtype_double, + neutral="maxloc_start()", + reduce_expr="maxloc_red(a, b)", + map_expr="maxloc_map(abs(x[i]), i)", + arguments="double *x", + preamble=maxloc_preamble_double, + ) + + amlsc = LowerLatencyReductionKernel( + maxloc_dtype_single, + neutral="maxloc_start()", + reduce_expr="maxloc_red(a, b)", + map_expr="maxloc_map(abs(x[i]), i)", + arguments="pycuda::complex *x", + preamble=maxloc_preamble_single, + ) + + amldc = LowerLatencyReductionKernel( + maxloc_dtype_double, + neutral="maxloc_start()", + reduce_expr="maxloc_red(a, b)", + map_expr="maxloc_map(abs(x[i]), i)", + arguments="pycuda::complex *x", + preamble=maxloc_preamble_double, + ) + + abs_max_loc_map = { + "single": {"real": amls, "complex": amlsc}, + "double": {"real": amld, "complex": amldc}, + } else: max_loc_map = {} abs_max_loc_map = {} + def zeros(length, dtype=np.float64): result = GPUArray(length, dtype=dtype) nwords = result.nbytes / 4 pycuda.driver.memset_d32(result.gpudata, 0, nwords) return result + def ptr(self): return self._data.ptr + def dot(self, other): - return pycuda.gpuarray.dot(self._data,other).get().max() + return pycuda.gpuarray.dot(self._data, other).get().max() + def min(self): return pycuda.gpuarray.min(self._data).get().max() + def abs_max_loc(self): maxloc = abs_max_loc_map[self.precision][self.kind](self._data) maxloc = maxloc.get() - return float(maxloc['max']),int(maxloc['loc']) + return float(maxloc["max"]), int(maxloc["loc"]) + def cumsum(self): - tmp = self.data*1 + tmp = self.data * 1 return icumsum(tmp) + def max(self): return pycuda.gpuarray.max(self._data).get().max() + def max_loc(self): maxloc = max_loc_map[self.precision](self._data) maxloc = maxloc.get() - return float(maxloc['max']),int(maxloc['loc']) - + return float(maxloc["max"]), int(maxloc["loc"]) + + def take(self, indices): if not isinstance(indices, pycuda.gpuarray.GPUArray): indices = pycuda.gpuarray.to_gpu(indices) return pycuda.gpuarray.take(self.data, indices) - + + def numpy(self): return self._data.get() - + + def _copy(self, self_ref, other_ref): - if (len(other_ref) <= len(self_ref)) : + if len(other_ref) <= len(self_ref): from pycuda.elementwise import get_copy_kernel + func = get_copy_kernel(self.dtype, other_ref.dtype) - func.prepared_async_call(self_ref._grid, self_ref._block, None, - self_ref.gpudata, other_ref.gpudata, - self_ref.mem_size) + func.prepared_async_call( + self_ref._grid, + self_ref._block, + None, + self_ref.gpudata, + other_ref.gpudata, + self_ref.mem_size, + ) else: raise RuntimeError("The arrays must the same length") + def _getvalue(self, index): return self._data.get()[index] - + + def sum(self): return pycuda.gpuarray.sum(self._data).get().max() - + + def clear(self): n32 = self.data.nbytes / 4 pycuda.driver.memset_d32(self.data.gpudata, 0, n32) - + + def _scheme_matches_base_array(array): if isinstance(array, pycuda.gpuarray.GPUArray): return True - else: - return False + return False + def _copy_base_array(array): data = pycuda.gpuarray.GPUArray((array.size), array.dtype) @@ -359,8 +472,6 @@ def _copy_base_array(array): pycuda.driver.memcpy_dtod(data.gpudata, array.gpudata, array.nbytes) return data + def _to_device(array): return pycuda.gpuarray.to_gpu(array) - - - diff --git a/pycbc/types/array_cupy.py b/pycbc/types/array_cupy.py index bc5b23cf3a6..714ddf18cb0 100644 --- a/pycbc/types/array_cupy.py +++ b/pycbc/types/array_cupy.py @@ -21,118 +21,134 @@ # # ============================================================================= # -"""Cupy based CPU backend for PyCBC Array -""" +"""Cupy based CPU backend for PyCBC Array""" + import cupy as cp + from pycbc.types.array import common_kind, complex128, float64 + def zeros(length, dtype=cp.float64): return cp.zeros(length, dtype=dtype) + def empty(length, dtype=cp.float64): return cp.empty(length, dtype=dtype) + def ptr(self): return self.data.data.mem.ptr + def dot(self, other): - return cp.dot(self._data,other) + return cp.dot(self._data, other) + def min(self): return self.data.min() + def abs_max_loc(self): - if self.kind == 'real': + if self.kind == "real": tmp = abs(self.data) ind = cp.argmax(tmp) return tmp[ind], ind - else: - tmp = self.data.real ** 2.0 - tmp += self.data.imag ** 2.0 - ind = cp.argmax(tmp) - return tmp[ind] ** 0.5, ind + tmp = self.data.real**2.0 + tmp += self.data.imag**2.0 + ind = cp.argmax(tmp) + return tmp[ind] ** 0.5, ind + def cumsum(self): return self.data.cumsum() + def max(self): return self.data.max() + def max_loc(self): ind = cp.argmax(self.data) return self.data[ind], ind + def take(self, indices): return self.data.take(indices) + def weighted_inner(self, other, weight): - """ Return the inner product of the array with complex conjugation. - """ + """Return the inner product of the array with complex conjugation.""" if weight is None: return self.inner(other) cdtype = common_kind(self.dtype, other.dtype) - if cdtype.kind == 'c': + if cdtype.kind == "c": acum_dtype = complex128 else: acum_dtype = float64 return cp.sum(self.data.conj() * other / weight, dtype=acum_dtype) + def abs_arg_max(self): if self.dtype == cp.float32 or self.dtype == cp.float64: return cp.argmax(abs(self.data)) - else: - return abs_arg_max_complex(self._data) + return abs_arg_max_complex(self._data) + def inner(self, other): - """ Return the inner product of the array with complex conjugation. - """ + """Return the inner product of the array with complex conjugation.""" cdtype = common_kind(self.dtype, other.dtype) - if cdtype.kind == 'c': + if cdtype.kind == "c": return cp.sum(self.data.conj() * other, dtype=complex128) - else: - return inner_real(self.data, other) + return inner_real(self.data, other) + def vdot(self, other): - """ Return the inner product of the array with complex conjugation. - """ + """Return the inner product of the array with complex conjugation.""" return cp.vdot(self.data, other) + def squared_norm(self): - """ Return the elementwise squared norm of the array """ - return (self.data.real**2 + self.data.imag**2) + """Return the elementwise squared norm of the array""" + return self.data.real**2 + self.data.imag**2 + def numpy(self): return cp.asnumpy(self.data) + def _copy(self, self_ref, other_ref): self_ref[:] = other_ref[:] + def _getvalue(self, index): return self._data[index] + def sum(self): - if self.kind == 'real': - return cp.sum(self._data,dtype=float64) - else: - return cp.sum(self._data,dtype=complex128) + if self.kind == "real": + return cp.sum(self._data, dtype=float64) + return cp.sum(self._data, dtype=complex128) + def clear(self): self[:] = 0 + def _scheme_matches_base_array(array): if isinstance(array, cp.ndarray): return True - else: - return False + return False + def _to_device(array): return cp.asarray(array) + def numpy(self): return cp.asnumpy(self._data) + def _copy_base_array(array): return array.copy() - diff --git a/pycbc/types/config.py b/pycbc/types/config.py index f2b50eadc71..b45d4cb41a2 100644 --- a/pycbc/types/config.py +++ b/pycbc/types/config.py @@ -25,14 +25,15 @@ This module provides a wrapper to the ConfigParser utilities for pycbc. This module is described in the page here: """ -import re -import os + +import configparser as ConfigParser import itertools import logging +import os +import re from io import StringIO -import configparser as ConfigParser -logger = logging.getLogger('pycbc.types.config') +logger = logging.getLogger("pycbc.types.config") class DeepCopyableConfigParser(ConfigParser.ConfigParser): @@ -67,14 +68,14 @@ def __init__( deleteTuples=None, skip_extended=False, sanitize_newline=True, - delete_sharedoptions_sections=True + delete_sharedoptions_sections=True, ): """ Initialize an InterpolatingConfigParser. This reads the input configuration files, overrides values if necessary and performs the interpolation. - Parameters - ----------- + Parameters + ---------- configFiles : Path to .ini file, or list of paths The file(s) to be read in and parsed. overrideTuples : List of (section, option, value) tuples @@ -89,9 +90,10 @@ def __init__( is provided, the entire section will be deleted. Returns - -------- + ------- InterpolatingConfigParser Initialized InterpolatingConfigParser instance. + """ if configFiles is None: configFiles = [] @@ -113,10 +115,11 @@ def __init__( # that are special to ConfigParser. So any variable containing a % or a # $ is ignored. env_vals = { - key: value for key, value in os.environ.items() - if '%' not in value and '$' not in value + key: value + for key, value in os.environ.items() + if "%" not in value and "$" not in value } - self.read_dict({'environment': env_vals}) + self.read_dict({"environment": env_vals}) self.read_ini_file(configFiles) @@ -124,9 +127,7 @@ def __init__( self.split_multi_sections() # Populate shared options from the [sharedoptions] section - self.populate_shared_sections( - delete_sections=delete_sharedoptions_sections - ) + self.populate_shared_sections(delete_sections=delete_sharedoptions_sections) # Do deletes from command line for delete in deleteTuples: @@ -137,9 +138,7 @@ def __init__( "no such section in configuration." % delete ) - logger.info( - "Deleting section %s from configuration", delete[0] - ) + logger.info("Deleting section %s from configuration", delete[0]) elif len(delete) == 2: if self.remove_option(delete[0], delete[1]) is False: raise ValueError( @@ -148,14 +147,13 @@ def __init__( ) logger.info( - "Deleting option %s from section %s in " "configuration", + "Deleting option %s from section %s in configuration", delete[1], delete[0], ) else: raise ValueError( - "Deletes must be tuples of length 1 or 2. " - "Got %s." % str(delete) + "Deletes must be tuples of length 1 or 2. Got %s." % str(delete) ) # Do overrides from command line @@ -174,8 +172,7 @@ def __init__( self.add_section(section) self.set(section, option, value) logger.info( - "Overriding section %s option %s with value %s " - "in configuration.", + "Overriding section %s option %s with value %s in configuration.", section, option, value, @@ -201,30 +198,29 @@ def __init__( @classmethod def from_cli(cls, opts): - """Initialize the config parser using options parsed from the command + """ + Initialize the config parser using options parsed from the command line. The parsed options ``opts`` must include options provided by :py:func:`add_workflow_command_line_group`. Parameters - ----------- + ---------- opts : argparse.ArgumentParser The command line arguments parsed by argparse + """ # read configuration file logger.info("Reading configuration file") if opts.config_overrides is not None: overrides = [ - tuple(override.split(":", 2)) - for override in opts.config_overrides + tuple(override.split(":", 2)) for override in opts.config_overrides ] else: overrides = None if opts.config_delete is not None: - deletes = [ - tuple(delete.split(":")) for delete in opts.config_delete - ] + deletes = [tuple(delete.split(":")) for delete in opts.config_delete] else: deletes = None return cls(opts.config_files, overrides, deleteTuples=deletes) @@ -247,6 +243,7 @@ def read_ini_file(self, fpath): ------- cp : ConfigParser The ConfigParser class containing the read in .ini file + """ # Read the file @@ -254,7 +251,7 @@ def read_ini_file(self, fpath): for filename in fpath: parser = ConfigParser.ConfigParser() - parser.optionxform=str + parser.optionxform = str parser.read(filename) for section in parser.sections(): @@ -263,10 +260,14 @@ def read_ini_file(self, fpath): section_options = parser.options(section) - option_intersection = options_seen[section].intersection(section_options) + option_intersection = options_seen[section].intersection( + section_options + ) if option_intersection: - raise ValueError(f"Duplicate option(s) {', '.join(option_intersection)} found in section '{section}' in file '{filename}'") + raise ValueError( + f"Duplicate option(s) {', '.join(option_intersection)} found in section '{section}' in file '{filename}'" + ) options_seen[section].update(section_options) @@ -276,10 +277,9 @@ def get_subsections(self, section_name): """Return a list of subsections for the given section name""" # Keep only subsection names subsections = [ - sec[len(section_name) + 1:] + sec[len(section_name) + 1 :] for sec in self.sections() - if sec.startswith(section_name + "-") - and not sec.endswith('defaultvalues') + if sec.startswith(section_name + "-") and not sec.endswith("defaultvalues") ] for sec in subsections: @@ -287,9 +287,7 @@ def get_subsections(self, section_name): # The format [section-subsection-tag] is okay. Just # check that [section-subsection] section exists. If not it is possible # the user is trying to use an subsection name with '-' in it - if (len(sp) > 1) and not self.has_section( - "%s-%s" % (section_name, sp[0]) - ): + if (len(sp) > 1) and not self.has_section("%s-%s" % (section_name, sp[0])): raise ValueError( "Workflow uses the '-' as a delimiter so " "this is interpreted as section-subsection-tag. " @@ -304,10 +302,9 @@ def get_subsections(self, section_name): if len(subsections) > 0: return [sec.split("-")[0] for sec in subsections] - elif self.has_section(section_name): + if self.has_section(section_name): return [""] - else: - return [] + return [] def perform_extended_interpolation(self): """ @@ -323,7 +320,6 @@ def perform_extended_interpolation(self): Nested interpolation is not supported here. """ - # Do not allow any interpolation of the section names for section in self.sections(): for option, value in self.items(section): @@ -343,11 +339,10 @@ def sanitize_newline(self): newlines with spaces. This is useful for command line conversion and allow multiline configparser inputs without added backslashes """ - # Do not allow any interpolation of the section names for section in self.sections(): for option, value in self.items(section): - new_value = value.replace('\n', ' ').replace('\r', ' ') + new_value = value.replace("\n", " ").replace("\r", " ") self.set(section, option, new_value) def interpolate_string(self, test_string, section): @@ -373,11 +368,11 @@ def interpolate_string(self, test_string, section): The current section of the ConfigParser object Returns - ---------- + ------- test_string : String Interpolated string - """ + """ # First check if any interpolation is needed and abort if not re_obj = re.search(r"\$\{.*?\}", test_string) while re_obj: @@ -431,7 +426,8 @@ def split_multi_sections(self): self.remove_section(section) def populate_shared_sections(self, delete_sections=True): - """Parse the [sharedoptions] section of the ini file. + """ + Parse the [sharedoptions] section of the ini file. That section should contain entries according to: @@ -508,12 +504,11 @@ def add_options_to_section(self, section, items, overwrite_options=False): This will override so that the options+values given in items will replace the original values if the value is set to True. Default = False + """ # Sanity checking if not self.has_section(section): - raise ValueError( - "Section %s not present in ConfigParser." % (section,) - ) + raise ValueError("Section %s not present in ConfigParser." % (section,)) # Check for duplicate options first for option, value in items: @@ -540,8 +535,7 @@ def sanity_check_subsections(self): if section == "pegasus_profile": continue - if section.endswith('-defaultvalues') and \ - not len(section.split('-')) == 2: + if section.endswith("-defaultvalues") and not len(section.split("-")) == 2: # Only allow defaultvalues for top-level sections raise NotImplementedError( "-defaultvalues subsections are only allowed for " @@ -557,9 +551,7 @@ def sanity_check_subsections(self): # be over-written by anything in the sections-proper continue # Check for duplicate options whenever this exists - self.check_duplicate_options( - section, section2, raise_error=True - ) + self.check_duplicate_options(section, section2, raise_error=True) def check_duplicate_options(self, section1, section2, raise_error=False): """ @@ -576,29 +568,27 @@ def check_duplicate_options(self, section1, section2, raise_error=False): If True, raise an error if duplicates are present. Returns - ---------- + ------- duplicates : List List of duplicate options + """ # Sanity checking if not self.has_section(section1): - raise ValueError( - "Section %s not present in ConfigParser." % (section1,) - ) + raise ValueError("Section %s not present in ConfigParser." % (section1,)) if not self.has_section(section2): - raise ValueError( - "Section %s not present in ConfigParser." % (section2,) - ) + raise ValueError("Section %s not present in ConfigParser." % (section2,)) # Are section1 and section2 a section-and-defaultvalues pair? - section_and_default = (section1 == f"{section2}-defaultvalues" or - section2 == f"{section1}-defaultvalues") + section_and_default = ( + section1 == f"{section2}-defaultvalues" + or section2 == f"{section1}-defaultvalues" + ) # Is one the sections defaultvalues, but the other is not the # top-level section? This is to catch the case where we are # comparing section-defaultvalues with section-subsection - if section1.endswith("-defaultvalues") or \ - section2.endswith("-defaultvalues"): + if section1.endswith("-defaultvalues") or section2.endswith("-defaultvalues"): if not section_and_default: # Override the raise_error variable not to error when # defaultvalues are given and the sections are not @@ -612,8 +602,10 @@ def check_duplicate_options(self, section1, section2, raise_error=False): duplicates = [x for x in items1 if x in items2] if duplicates and raise_error: - err_msg = ("The following options appear in both section " - f"{section1} and {section2}: " + ", ".join(duplicates)) + err_msg = ( + "The following options appear in both section " + f"{section1} and {section2}: " + ", ".join(duplicates) + ) if section_and_default: err_msg += ". Default values are unused in this case." raise ValueError(err_msg) @@ -627,7 +619,7 @@ def get_opt_tag(self, section, option, tag): NB calling get_opt_tags() directly is preferred for simplicity. Parameters - ----------- + ---------- self : ConfigParser object The ConfigParser object (automatically passed when this is appended to the ConfigParser class) @@ -639,9 +631,10 @@ def get_opt_tag(self, section, option, tag): The name of the subsection to look in, if not found in [section] Returns - -------- + ------- string The value of the options being searched for + """ return self.get_opt_tags(section, option, [tag]) @@ -654,7 +647,7 @@ def get_opt_tags(self, section, option, tags): values. Will raise a ConfigParser.Error if it cannot find a value. Parameters - ----------- + ---------- self : ConfigParser object The ConfigParser object (automatically passed when this is appended to the ConfigParser class) @@ -666,9 +659,10 @@ def get_opt_tags(self, section, option, tags): The name of subsections to look in, if not found in [section] Returns - -------- + ------- string The value of the options being searched for + """ # Need lower case tag name; also exclude cases with tag=None if tags: @@ -684,9 +678,7 @@ def get_opt_tags(self, section, option, tags): # First, check if there are any default values set: has_defaultvalue = False if self.has_section(f"{section}-defaultvalues"): - return_vals.append( - self.get(f"{section}-defaultvalues", option) - ) + return_vals.append(self.get(f"{section}-defaultvalues", option)) has_defaultvalue = True sub_section_list = [] @@ -700,9 +692,7 @@ def get_opt_tags(self, section, option, tags): if self.has_section("%s-%s" % (section, sub)): if self.has_option("%s-%s" % (section, sub), option): err_section_list.append("%s-%s" % (section, sub)) - return_vals.append( - self.get("%s-%s" % (section, sub), option) - ) + return_vals.append(self.get("%s-%s" % (section, sub), option)) if has_defaultvalue and len(return_vals) > 1: # option supplied which should overwrite the default; @@ -711,14 +701,11 @@ def get_opt_tags(self, section, option, tags): # We also want to recursively go into sections if not return_vals: - err_string += "or in sections [%s]." % ( - "] [".join(section_list) - ) + err_string += "or in sections [%s]." % ("] [".join(section_list)) raise ConfigParser.Error(err_string) if len(return_vals) > 1: - err_string += ( - "and multiple entries found in sections [%s]." - % ("] [".join(err_section_list)) + err_string += "and multiple entries found in sections [%s]." % ( + "] [".join(err_section_list) ) raise ConfigParser.Error(err_string) return return_vals[0] @@ -730,7 +717,7 @@ def has_option_tag(self, section, option, tag): NB calling has_option_tags() directly is preferred for simplicity. Parameters - ----------- + ---------- self : ConfigParser object The ConfigParser object (automatically passed when this is appended to the ConfigParser class) @@ -742,9 +729,10 @@ def has_option_tag(self, section, option, tag): The name of the subsection to look in, if not found in [section] Returns - -------- + ------- Boolean Is the option in the section or [section-tag] + """ return self.has_option_tags(section, option, [tag]) @@ -756,7 +744,7 @@ def has_option_tags(self, section, option, tags): Returns True if the option is found and false if not. Parameters - ----------- + ---------- self : ConfigParser object The ConfigParser object (automatically passed when this is appended to the ConfigParser class) @@ -768,9 +756,10 @@ def has_option_tags(self, section, option, tags): The names of the subsection to look in, if not found in [section] Returns - -------- + ------- Boolean Is the option in the section or [section-tag] (for tag in tags) + """ try: self.get_opt_tags(section, option, tags) diff --git a/pycbc/types/frequencyseries.py b/pycbc/types/frequencyseries.py index b4cb25afdd7..2b5ebbced6d 100644 --- a/pycbc/types/frequencyseries.py +++ b/pycbc/types/frequencyseries.py @@ -17,19 +17,23 @@ """ Provides a class representing a frequency series. """ + import os as _os + import h5py import numpy as _numpy -from pycbc.types.array import Array, _convert, zeros, _noreal -from pycbc.types.utils import determine_epoch -from pycbc.types import float64 from pycbc.libutils import import_optional +from pycbc.types import float64 +from pycbc.types.array import Array, _convert, _noreal, zeros +from pycbc.types.utils import determine_epoch + +_lal = import_optional("lal") -_lal = import_optional('lal') class FrequencySeries(Array): - """Models a frequency series consisting of uniformly sampled scalar values. + """ + Models a frequency series consisting of uniformly sampled scalar values. Parameters ---------- @@ -43,18 +47,21 @@ class FrequencySeries(Array): Sample data type. copy : boolean, optional If True, samples are copied to a new array. + """ def __init__(self, initial_array, delta_f=None, epoch="", dtype=None, copy=True): if len(initial_array) < 1: - raise ValueError('initial_array must contain at least one sample.') + raise ValueError("initial_array must contain at least one sample.") if delta_f is None: try: delta_f = initial_array.delta_f except AttributeError: - raise TypeError('must provide either an initial_array with a delta_f attribute, or a value for delta_f') + raise TypeError( + "must provide either an initial_array with a delta_f attribute, or a value for delta_f" + ) if not delta_f > 0: - raise ValueError('delta_f must be a positive number') + raise ValueError("delta_f must be a positive number") Array.__init__(self, initial_array, dtype=dtype, copy=copy) self._delta_f = delta_f @@ -66,90 +73,88 @@ def _return(self, ary): def _typecheck(self, other): if isinstance(other, FrequencySeries): try: - _numpy.testing.assert_almost_equal(other._delta_f, - self._delta_f) + _numpy.testing.assert_almost_equal(other._delta_f, self._delta_f) except: - raise ValueError('different delta_f') + raise ValueError("different delta_f") # consistency of _epoch is not required because we may want # to combine frequency series estimated at different times # (e.g. PSD estimation) def get_delta_f(self): - """Return frequency between consecutive samples in Hertz. - """ + """Return frequency between consecutive samples in Hertz.""" return self._delta_f - delta_f = property(get_delta_f, - doc="Frequency between consecutive samples in Hertz.") + + delta_f = property( + get_delta_f, doc="Frequency between consecutive samples in Hertz." + ) def get_epoch(self): - """Return frequency series epoch - """ + """Return frequency series epoch""" return self._epoch - - epoch = property(get_epoch, - doc="Frequency series epoch.") + + epoch = property(get_epoch, doc="Frequency series epoch.") def get_sample_frequencies(self): - """Return an Array containing the sample frequencies. - """ + """Return an Array containing the sample frequencies.""" return Array(range(len(self))) * self._delta_f - sample_frequencies = property(get_sample_frequencies, - doc="Array of the sample frequencies.") + + sample_frequencies = property( + get_sample_frequencies, doc="Array of the sample frequencies." + ) def _getslice(self, index): if index.step is not None: new_delta_f = self._delta_f * index.step else: new_delta_f = self._delta_f - return FrequencySeries(Array._getslice(self, index), - delta_f=new_delta_f, - epoch=self._epoch, - copy=False) + return FrequencySeries( + Array._getslice(self, index), + delta_f=new_delta_f, + epoch=self._epoch, + copy=False, + ) def at_frequency(self, freq): - """ Return the value at the specified frequency - """ + """Return the value at the specified frequency""" return self[int(freq / self.delta_f)] @property def start_time(self): - """Return the start time of this vector - """ + """Return the start time of this vector""" return self.epoch @start_time.setter def start_time(self, time): - """ Set the start time - """ + """Set the start time""" self._epoch = float64(time) @property def end_time(self): - """Return the end time of this vector - """ + """Return the end time of this vector""" return self.start_time + self.duration @property def duration(self): - """Return the time duration of this vector - """ + """Return the time duration of this vector""" return 1.0 / self.delta_f @property def delta_t(self): - """Return the time between samples if this were a time series. + """ + Return the time between samples if this were a time series. This assume the time series is even in length! """ return 1.0 / self.sample_rate @property def sample_rate(self): - """Return the sample rate this would have in the time domain. This + """ + Return the sample rate this would have in the time domain. This assumes even length time series! """ return (len(self) - 1) * self.delta_f * 2.0 - def __eq__(self,other): + def __eq__(self, other): """ This is the Python special method invoked whenever the '==' comparison is used. It will return true if the data of two @@ -182,13 +187,13 @@ def __eq__(self,other): ------- boolean: 'True' if the types, dtypes, lengths, epochs, delta_fs and data of the two objects are each identical. + """ - if super(FrequencySeries,self).__eq__(other): - return (self._epoch == other._epoch and self._delta_f == other._delta_f) - else: - return False + if super().__eq__(other): + return self._epoch == other._epoch and self._delta_f == other._delta_f + return False - def almost_equal_elem(self,other,tol,relative=True,dtol=0.0): + def almost_equal_elem(self, other, tol, relative=True, dtol=0.0): """ Compare whether two frequency series are almost equal, element by element. @@ -234,22 +239,27 @@ def almost_equal_elem(self,other,tol,relative=True,dtol=0.0): boolean: 'True' if the data and delta_fs agree within the tolerance, as interpreted by the 'relative' keyword, and if the types, lengths, dtypes, and epochs are exactly the same. + """ # Check that the delta_f tolerance is non-negative; raise an exception # if needed. - if (dtol < 0.0): + if dtol < 0.0: raise ValueError("Tolerance in delta_f cannot be negative") - if super(FrequencySeries,self).almost_equal_elem(other,tol=tol,relative=relative): + if super().almost_equal_elem( + other, tol=tol, relative=relative + ): if relative: - return (self._epoch == other._epoch and - abs(self._delta_f-other._delta_f) <= dtol*self._delta_f) - else: - return (self._epoch == other._epoch and - abs(self._delta_f-other._delta_f) <= dtol) - else: - return False + return ( + self._epoch == other._epoch + and abs(self._delta_f - other._delta_f) <= dtol * self._delta_f + ) + return ( + self._epoch == other._epoch + and abs(self._delta_f - other._delta_f) <= dtol + ) + return False - def almost_equal_norm(self,other,tol,relative=True,dtol=0.0): + def almost_equal_norm(self, other, tol, relative=True, dtol=0.0): """ Compare whether two frequency series are almost equal, normwise. @@ -292,24 +302,30 @@ def almost_equal_norm(self,other,tol,relative=True,dtol=0.0): boolean: 'True' if the data and delta_fs agree within the tolerance, as interpreted by the 'relative' keyword, and if the types, lengths, dtypes, and epochs are exactly the same. + """ # Check that the delta_f tolerance is non-negative; raise an exception # if needed. - if (dtol < 0.0): + if dtol < 0.0: raise ValueError("Tolerance in delta_f cannot be negative") - if super(FrequencySeries,self).almost_equal_norm(other,tol=tol,relative=relative): + if super().almost_equal_norm( + other, tol=tol, relative=relative + ): if relative: - return (self._epoch == other._epoch and - abs(self._delta_f-other._delta_f) <= dtol*self._delta_f) - else: - return (self._epoch == other._epoch and - abs(self._delta_f-other._delta_f) <= dtol) - else: - return False + return ( + self._epoch == other._epoch + and abs(self._delta_f - other._delta_f) <= dtol * self._delta_f + ) + return ( + self._epoch == other._epoch + and abs(self._delta_f - other._delta_f) <= dtol + ) + return False @_convert def lal(self): - """Produces a LAL frequency series object equivalent to self. + """ + Produces a LAL frequency series object equivalent to self. Returns ------- @@ -323,28 +339,36 @@ def lal(self): ------ TypeError If frequency series is stored in GPU memory. - """ + """ lal_data = None if self._epoch is None: - ep = _lal.LIGOTimeGPS(0,0) + ep = _lal.LIGOTimeGPS(0, 0) else: ep = _lal.LIGOTimeGPS(self._epoch) if self._data.dtype == _numpy.float32: - lal_data = _lal.CreateREAL4FrequencySeries("",ep,0,self.delta_f,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateREAL4FrequencySeries( + "", ep, 0, self.delta_f, _lal.SecondUnit, len(self) + ) elif self._data.dtype == _numpy.float64: - lal_data = _lal.CreateREAL8FrequencySeries("",ep,0,self.delta_f,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateREAL8FrequencySeries( + "", ep, 0, self.delta_f, _lal.SecondUnit, len(self) + ) elif self._data.dtype == _numpy.complex64: - lal_data = _lal.CreateCOMPLEX8FrequencySeries("",ep,0,self.delta_f,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateCOMPLEX8FrequencySeries( + "", ep, 0, self.delta_f, _lal.SecondUnit, len(self) + ) elif self._data.dtype == _numpy.complex128: - lal_data = _lal.CreateCOMPLEX16FrequencySeries("",ep,0,self.delta_f,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateCOMPLEX16FrequencySeries( + "", ep, 0, self.delta_f, _lal.SecondUnit, len(self) + ) lal_data.data.data[:] = self.numpy() return lal_data - def save(self, path, group=None, ifo='P1'): + def save(self, path, group=None, ifo="P1"): """ Save frequency series to a Numpy .npy, hdf, or text file. The first column contains the sample frequencies, the second contains the values. @@ -365,63 +389,68 @@ def save(self, path, group=None, ifo='P1'): ------ ValueError If path does not end in .npy or .txt. - """ + """ ext = _os.path.splitext(path)[1] - if ext == '.npy': - output = _numpy.vstack((self.sample_frequencies.numpy(), - self.numpy())).T + if ext == ".npy": + output = _numpy.vstack((self.sample_frequencies.numpy(), self.numpy())).T _numpy.save(path, output) - elif ext == '.txt': - if self.kind == 'real': - output = _numpy.vstack((self.sample_frequencies.numpy(), - self.numpy())).T - elif self.kind == 'complex': - output = _numpy.vstack((self.sample_frequencies.numpy(), - self.numpy().real, - self.numpy().imag)).T + elif ext == ".txt": + if self.kind == "real": + output = _numpy.vstack( + (self.sample_frequencies.numpy(), self.numpy()) + ).T + elif self.kind == "complex": + output = _numpy.vstack( + ( + self.sample_frequencies.numpy(), + self.numpy().real, + self.numpy().imag, + ) + ).T _numpy.savetxt(path, output) - elif ext == '.xml' or path.endswith('.xml.gz'): - from pycbc.io.ligolw import make_psd_xmldoc + elif ext == ".xml" or path.endswith(".xml.gz"): from igwn_ligolw import utils - if self.kind != 'real': - raise ValueError('XML only supports real frequency series') + from pycbc.io.ligolw import make_psd_xmldoc + + if self.kind != "real": + raise ValueError("XML only supports real frequency series") output = self.lal() - output.name = 'psd' + output.name = "psd" # When writing in this format we must *not* have the 0 values at # frequencies less than flow. To resolve this we set the first # non-zero value < flow. data_lal = output.data.data - first_idx = _numpy.argmax(data_lal>0) + first_idx = _numpy.argmax(data_lal > 0) if not first_idx == 0: data_lal[:first_idx] = data_lal[first_idx] psddict = {ifo: output} - utils.write_filename( - make_psd_xmldoc(psddict), - path, - compress='auto' - ) - elif ext == '.hdf': - key = 'data' if group is None else group - with h5py.File(path, 'a') as f: - ds = f.create_dataset(key, data=self.numpy(), - compression='gzip', - compression_opts=9, shuffle=True) + utils.write_filename(make_psd_xmldoc(psddict), path, compress="auto") + elif ext == ".hdf": + key = "data" if group is None else group + with h5py.File(path, "a") as f: + ds = f.create_dataset( + key, + data=self.numpy(), + compression="gzip", + compression_opts=9, + shuffle=True, + ) if self.epoch is not None: - ds.attrs['epoch'] = float(self.epoch) - ds.attrs['delta_f'] = float(self.delta_f) + ds.attrs["epoch"] = float(self.epoch) + ds.attrs["delta_f"] = float(self.delta_f) else: - raise ValueError('Path must end with .npy, .txt, .xml, .xml.gz ' - 'or .hdf') + raise ValueError("Path must end with .npy, .txt, .xml, .xml.gz or .hdf") def to_frequencyseries(self): - """ Return frequency series """ + """Return frequency series""" return self @_noreal def to_timeseries(self, delta_t=None): - """ Return the Fourier transform of this time series. + """ + Return the Fourier transform of this time series. Note that this assumes even length time series! @@ -436,39 +465,47 @@ def to_timeseries(self, delta_t=None): ------- TimeSeries: The inverse fourier transform of this frequency series. + """ from pycbc.fft import ifft from pycbc.types import TimeSeries, real_same_precision_as - nat_delta_t = 1.0 / ((len(self)-1)*2) / self.delta_f + + nat_delta_t = 1.0 / ((len(self) - 1) * 2) / self.delta_f if not delta_t: delta_t = nat_delta_t # add 0.5 to round integer - tlen = int(1.0 / self.delta_f / delta_t + 0.5) + tlen = int(1.0 / self.delta_f / delta_t + 0.5) flen = int(tlen / 2 + 1) if flen < len(self): - raise ValueError("The value of delta_t (%s) would be " - "undersampled. Maximum delta_t " - "is %s." % (delta_t, nat_delta_t)) + raise ValueError( + "The value of delta_t (%s) would be " + "undersampled. Maximum delta_t " + "is %s." % (delta_t, nat_delta_t) + ) if not delta_t: tmp = self else: - tmp = FrequencySeries(zeros(flen, dtype=self.dtype), - delta_f=self.delta_f, epoch=self.epoch, - copy=False) - tmp[:len(self)] = self[:] - - f = TimeSeries(zeros(tlen, - dtype=real_same_precision_as(self)), - delta_t=delta_t, copy=False) + tmp = FrequencySeries( + zeros(flen, dtype=self.dtype), + delta_f=self.delta_f, + epoch=self.epoch, + copy=False, + ) + tmp[: len(self)] = self[:] + + f = TimeSeries( + zeros(tlen, dtype=real_same_precision_as(self)), delta_t=delta_t, copy=False + ) ifft(tmp, f) f._delta_t = delta_t return f @_noreal def cyclic_time_shift(self, dt): - """Shift the data and timestamps by a given number of seconds + """ + Shift the data and timestamps by a given number of seconds Shift the data and timestamps in the time domain a given number of seconds. To just change the time stamps, do ts.start_time += dt. @@ -486,15 +523,19 @@ def cyclic_time_shift(self, dt): ------- data : pycbc.types.FrequencySeries The time shifted frequency series. + """ from pycbc.waveform import apply_fseries_time_shift + data = apply_fseries_time_shift(self, dt) data.start_time = self.start_time - dt return data - def match(self, other, psd=None, - low_frequency_cutoff=None, high_frequency_cutoff=None): - """ Return the match between the two TimeSeries or FrequencySeries. + def match( + self, other, psd=None, low_frequency_cutoff=None, high_frequency_cutoff=None + ): + """ + Return the match between the two TimeSeries or FrequencySeries. Return the match between two waveforms. This is equivalent to the overlap maximized over time and phase. By default, the other vector will be @@ -519,9 +560,10 @@ def match(self, other, psd=None, match: float index: int The number of samples to shift to get the match. + """ - from pycbc.types import TimeSeries from pycbc.filter import match + from pycbc.types import TimeSeries if isinstance(other, TimeSeries): if other.duration != self.duration: @@ -538,25 +580,30 @@ def match(self, other, psd=None, psd = psd.copy() psd.resize(len(self)) - return match(self, other, psd=psd, - low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff) + return match( + self, + other, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + ) def plot(self, **kwds): - """ Basic plot of this frequency series - """ + """Basic plot of this frequency series""" from matplotlib import pyplot - if self.kind == 'real': + if self.kind == "real": plot = pyplot.plot(self.sample_frequencies, self, **kwds) return plot - elif self.kind == 'complex': + if self.kind == "complex": plot1 = pyplot.plot(self.sample_frequencies, self.real(), **kwds) plot2 = pyplot.plot(self.sample_frequencies, self.imag(), **kwds) return plot1, plot2 + def load_frequencyseries(path, group=None): - """Load a FrequencySeries from an HDF5, ASCII or Numpy file. The file type + """ + Load a FrequencySeries from an HDF5, ASCII or Numpy file. The file type is inferred from the file extension, which must be `.hdf`, `.txt` or `.npy`. @@ -587,29 +634,34 @@ def load_frequencyseries(path, group=None): If the path does not end in a supported extension. For Numpy and ASCII input files, this is also raised if the array does not have 2 or 3 dimensions. + """ ext = _os.path.splitext(path)[1] - if ext == '.npy': + if ext == ".npy": data = _numpy.load(path) - elif ext == '.txt': + elif ext == ".txt": data = _numpy.loadtxt(path) - elif ext == '.hdf': - key = 'data' if group is None else group - with h5py.File(path, 'r') as f: + elif ext == ".hdf": + key = "data" if group is None else group + with h5py.File(path, "r") as f: data = f[key][:] - delta_f = f[key].attrs['delta_f'] - epoch = f[key].attrs['epoch'] if 'epoch' in f[key].attrs else None + delta_f = f[key].attrs["delta_f"] + epoch = f[key].attrs["epoch"] if "epoch" in f[key].attrs else None series = FrequencySeries(data, delta_f=delta_f, epoch=epoch) return series else: - raise ValueError('Path must end with .npy, .hdf, or .txt') + raise ValueError("Path must end with .npy, .hdf, or .txt") delta_f = (data[-1][0] - data[0][0]) / (len(data) - 1) if data.ndim == 2: - return FrequencySeries(data[:,1], delta_f=delta_f, epoch=None) - elif data.ndim == 3: - return FrequencySeries(data[:,1] + 1j*data[:,2], delta_f=delta_f, - epoch=None) - - raise ValueError('File has %s dimensions, cannot convert to FrequencySeries, \ - must be 2 (real) or 3 (complex)' % data.ndim) + return FrequencySeries(data[:, 1], delta_f=delta_f, epoch=None) + if data.ndim == 3: + return FrequencySeries( + data[:, 1] + 1j * data[:, 2], delta_f=delta_f, epoch=None + ) + + raise ValueError( + "File has %s dimensions, cannot convert to FrequencySeries, \ + must be 2 (real) or 3 (complex)" + % data.ndim + ) diff --git a/pycbc/types/optparse.py b/pycbc/types/optparse.py index e9382e16a59..b1faa0a4e30 100644 --- a/pycbc/types/optparse.py +++ b/pycbc/types/optparse.py @@ -19,59 +19,64 @@ This modules contains extensions for use with argparse """ -import copy -import warnings import argparse -import re +import copy import math +import re +import warnings from collections import defaultdict class DictWithDefaultReturn(defaultdict): default_set = False ifo_set = False + def __bool__(self): if self.items() and not all(entry is None for entry in self.values()): # True if any values are explictly set. return True - elif self['RANDOM_STRING_314324'] is not None: + if self["RANDOM_STRING_314324"] is not None: # Or true if the default value was set # NOTE: This stores the string RANDOM_STRING_314324 in the dict # so subsequent calls will be caught in the first test here. return True - else: - # Else false - return False + # Else false + return False + # Python 2 and 3 have different conventions for boolean method __nonzero__ = __bool__ + class MultiDetOptionAction(argparse.Action): # Initialise the same as the standard 'append' action - def __init__(self, - option_strings, - dest, - nargs='+', - const=None, - default=None, - type=None, - choices=None, - required=False, - help=None, - metavar=None): + def __init__( + self, + option_strings, + dest, + nargs="+", + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None, + ): if type is not None: self.internal_type = type else: self.internal_type = str new_default = DictWithDefaultReturn(lambda: default) - #new_default.default_value=default + # new_default.default_value=default if nargs == 0: - raise ValueError('nargs for append actions must be > 0; if arg ' - 'strings are not supplying the value to append, ' - 'the append const action may be more appropriate') + raise ValueError( + "nargs for append actions must be > 0; if arg " + "strings are not supplying the value to append, " + "the append const action may be more appropriate" + ) if const is not None and nargs != argparse.OPTIONAL: - raise ValueError('nargs must be %r to supply const' - % argparse.OPTIONAL) - super(MultiDetOptionAction, self).__init__( + raise ValueError("nargs must be %r to supply const" % argparse.OPTIONAL) + super().__init__( option_strings=option_strings, dest=dest, nargs=nargs, @@ -81,18 +86,19 @@ def __init__(self, choices=choices, required=required, help=help, - metavar=metavar) + metavar=metavar, + ) def __call__(self, parser, namespace, values, option_string=None): # Again this is modified from the standard argparse 'append' action - err_msg = "Issue with option: %s \n" %(self.dest,) - err_msg += "Received value: %s \n" %(' '.join(values),) + err_msg = "Issue with option: %s \n" % (self.dest,) + err_msg += "Received value: %s \n" % (" ".join(values),) if getattr(namespace, self.dest, None) is None: setattr(namespace, self.dest, DictWithDefaultReturn()) items = getattr(namespace, self.dest) items = copy.copy(items) for value in values: - value = value.split(':') + value = value.split(":") if len(value) == 2: # "Normal" case, all ifos supplied independently as "H1:VALUE" if items.default_set: @@ -112,7 +118,7 @@ def __call__(self, parser, namespace, values, option_string=None): err_msg += "If you are supplying a value for all ifos, you " err_msg += "cannot also supply values for specific ifos." raise ValueError(err_msg) - #items.default_value = self.internal_type(value[0]) + # items.default_value = self.internal_type(value[0]) new_default = self.internal_type(value[0]) items.default_factory = lambda: new_default items.default_set = True @@ -123,6 +129,7 @@ def __call__(self, parser, namespace, values, option_string=None): raise ValueError(err_msg) setattr(namespace, self.dest, items) + class MultiDetOptionActionSpecial(MultiDetOptionAction): """ This class in an extension of the MultiDetOptionAction class to handle @@ -131,37 +138,38 @@ class MultiDetOptionActionSpecial(MultiDetOptionAction): be provided uniquely for each ifo. The dictionary key is set to H1 and the value to H1:CHANNEL_NAME for this example. """ + def __call__(self, parser, namespace, values, option_string=None): # Again this is modified from the standard argparse 'append' action - err_msg = "Issue with option: %s \n" %(self.dest,) - err_msg += "Received value: %s \n" %(' '.join(values),) + err_msg = "Issue with option: %s \n" % (self.dest,) + err_msg += "Received value: %s \n" % (" ".join(values),) if getattr(namespace, self.dest, None) is None: setattr(namespace, self.dest, {}) items = getattr(namespace, self.dest) items = copy.copy(items) for value in values: - value_split = value.split(':') + value_split = value.split(":") if len(value_split) == 2: # "Normal" case, all ifos supplied independently as "H1:VALUE" if value_split[0] in items: - err_msg += "Multiple values supplied for ifo %s.\n" \ - %(value_split[0],) - err_msg += "Already have %s." %(items[value_split[0]]) + err_msg += "Multiple values supplied for ifo %s.\n" % ( + value_split[0], + ) + err_msg += "Already have %s." % (items[value_split[0]]) raise ValueError(err_msg) - else: - items[value_split[0]] = value + items[value_split[0]] = value elif len(value_split) == 3: # This is an unadvertised feature. It is used for cases where I # want to pretend H1 data is actually L1 (or similar). So if I # supply --channel-name H1:L1:LDAS-STRAIN I can use L1 data and # pretend it is H1 internally. if value_split[0] in items: - err_msg += "Multiple values supplied for ifo %s.\n" \ - %(value_split[0],) - err_msg += "Already have %s." %(items[value_split[0]]) + err_msg += "Multiple values supplied for ifo %s.\n" % ( + value_split[0], + ) + err_msg += "Already have %s." % (items[value_split[0]]) raise ValueError(err_msg) - else: - items[value_split[0]] = ':'.join(value_split[1:3]) + items[value_split[0]] = ":".join(value_split[1:3]) else: err_msg += "The character ':' is used to deliminate the " err_msg += "ifo and the value. It must appear exactly " @@ -169,34 +177,37 @@ def __call__(self, parser, namespace, values, option_string=None): raise ValueError(err_msg) setattr(namespace, self.dest, items) + class MultiDetMultiColonOptionAction(MultiDetOptionAction): - """A special case of `MultiDetOptionAction` which allows one to use + """ + A special case of `MultiDetOptionAction` which allows one to use arguments containing colons, such as `V1:FOOBAR:1`. The first colon is assumed to be the separator between the detector and the argument. All subsequent colons are kept as part of the argument. Unlike `MultiDetOptionAction`, all arguments must be prefixed by the corresponding detector. """ + def __call__(self, parser, namespace, values, option_string=None): - err_msg = ('Issue with option: {}\n' - 'Received value: {}\n').format(self.dest, ' '.join(values)) + err_msg = ("Issue with option: {}\nReceived value: {}\n").format( + self.dest, " ".join(values) + ) if getattr(namespace, self.dest, None) is None: setattr(namespace, self.dest, {}) items = copy.copy(getattr(namespace, self.dest)) for value in values: - if ':' not in value: - err_msg += ("Each argument must contain at least one ':' " - "character") + if ":" not in value: + err_msg += "Each argument must contain at least one ':' character" raise ValueError(err_msg) - detector, argument = value.split(':', 1) + detector, argument = value.split(":", 1) if detector in items: - err_msg += ('Multiple values supplied for detector {},\n' - 'already have {}.') + err_msg += "Multiple values supplied for detector {},\nalready have {}." err_msg = err_msg.format(detector, items[detector]) raise ValueError(err_msg) items[detector] = self.internal_type(argument) setattr(namespace, self.dest, items) + class MultiDetOptionAppendAction(MultiDetOptionAction): def __call__(self, parser, namespace, values, option_string=None): # Again this is modified from the standard argparse 'append' action @@ -205,7 +216,7 @@ def __call__(self, parser, namespace, values, option_string=None): items = getattr(namespace, self.dest) items = copy.copy(items) for value in values: - value = value.split(':') + value = value.split(":") if len(value) == 2: # "Normal" case, all ifos supplied independetly as "H1:VALUE" if value[0] in items: @@ -213,40 +224,44 @@ def __call__(self, parser, namespace, values, option_string=None): else: items[value[0]] = [self.internal_type(value[1])] else: - err_msg = "Issue with option: %s \n" %(self.dest,) - err_msg += "Received value: %s \n" %(' '.join(values),) + err_msg = "Issue with option: %s \n" % (self.dest,) + err_msg += "Received value: %s \n" % (" ".join(values),) err_msg += "The character ':' is used to distinguish the " err_msg += "ifo and the value. It must be given exactly once " err_msg += "for all entries" raise ValueError(err_msg) setattr(namespace, self.dest, items) + class DictOptionAction(argparse.Action): # Initialise the same as the standard 'append' action - def __init__(self, - option_strings, - dest, - nargs='+', - const=None, - default=None, - type=None, - choices=None, - required=False, - help=None, - metavar=None): + def __init__( + self, + option_strings, + dest, + nargs="+", + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None, + ): if type is not None: self.internal_type = type else: self.internal_type = str new_default = DictWithDefaultReturn(lambda: default) if nargs == 0: - raise ValueError('nargs for append actions must be > 0; if arg ' - 'strings are not supplying the value to append, ' - 'the append const action may be more appropriate') + raise ValueError( + "nargs for append actions must be > 0; if arg " + "strings are not supplying the value to append, " + "the append const action may be more appropriate" + ) if const is not None and nargs != argparse.OPTIONAL: - raise ValueError('nargs must be %r to supply const' - % argparse.OPTIONAL) - super(DictOptionAction, self).__init__( + raise ValueError("nargs must be %r to supply const" % argparse.OPTIONAL) + super().__init__( option_strings=option_strings, dest=dest, nargs=nargs, @@ -256,20 +271,21 @@ def __init__(self, choices=choices, required=required, help=help, - metavar=metavar) + metavar=metavar, + ) def __call__(self, parser, namespace, values, option_string=None): # Again this is modified from the standard argparse 'append' action - err_msg = "Issue with option: %s \n" %(self.dest,) - err_msg += "Received value: %s \n" %(' '.join(values),) + err_msg = "Issue with option: %s \n" % (self.dest,) + err_msg += "Received value: %s \n" % (" ".join(values),) if getattr(namespace, self.dest, None) is None: setattr(namespace, self.dest, {}) items = getattr(namespace, self.dest) items = copy.copy(items) for value in values: - if values == ['{}']: + if values == ["{}"]: break - value = value.split(':') + value = value.split(":") if len(value) == 2: # "Normal" case, all extra arguments supplied independently # as "param:VALUE" @@ -281,8 +297,10 @@ def __call__(self, parser, namespace, values, option_string=None): raise ValueError(err_msg) setattr(namespace, self.dest, items) + class MultiDetDictOptionAction(DictOptionAction): - """A special case of `DictOptionAction` which allows one to use + """ + A special case of `DictOptionAction` which allows one to use argument containing the detector (channel) name, such as `DETECTOR:PARAM:VALUE`. The first colon is the name of detector, the second colon is the name of parameter, the third colon is the value. @@ -290,55 +308,64 @@ class MultiDetDictOptionAction(DictOptionAction): detector, such as `PARAM:VALUE`, this will assume each detector has same values of those parameters. """ + def __call__(self, parser, namespace, values, option_string=None): # Again this is modified from the standard argparse 'append' action - err_msg = ('Issue with option: {}\n' - 'Received value: {}\n').format(self.dest, ' '.join(values)) + err_msg = ("Issue with option: {}\nReceived value: {}\n").format( + self.dest, " ".join(values) + ) if getattr(namespace, self.dest, None) is None: setattr(namespace, self.dest, {}) items = copy.copy(getattr(namespace, self.dest)) detector_args = {} for value in values: - if values == ['{}']: + if values == ["{}"]: break - if value.count(':') == 2: - detector, param_value = value.split(':', 1) - param, val = param_value.split(':') + if value.count(":") == 2: + detector, param_value = value.split(":", 1) + param, val = param_value.split(":") if detector not in detector_args: detector_args[detector] = {param: self.internal_type(val)} if param in detector_args[detector]: - err_msg += ("Multiple values supplied for the same " - "parameter {} under detector {},\n" - "already have {}.") - err_msg = err_msg.format(param, detector, - detector_args[detector][param]) + err_msg += ( + "Multiple values supplied for the same " + "parameter {} under detector {},\n" + "already have {}." + ) + err_msg = err_msg.format( + param, detector, detector_args[detector][param] + ) else: detector_args[detector][param] = self.internal_type(val) - elif value.count(':') == 1: - param, val = value.split(':') - for detector in getattr(namespace, 'instruments'): + elif value.count(":") == 1: + param, val = value.split(":") + for detector in namespace.instruments: if detector not in detector_args: - detector_args[detector] = \ - {param: self.internal_type(val)} + detector_args[detector] = {param: self.internal_type(val)} if param in detector_args[detector]: - err_msg += ("Multiple values supplied for the same " - "parameter {} under detector {},\n" - "already have {}.") + err_msg += ( + "Multiple values supplied for the same " + "parameter {} under detector {},\n" + "already have {}." + ) err_msg = err_msg.format( - param, detector, - detector_args[detector][param]) + param, detector, detector_args[detector][param] + ) else: - detector_args[detector][param] = \ - self.internal_type(val) + detector_args[detector][param] = self.internal_type(val) else: - err_msg += ("Use format `DETECTOR:PARAM:VALUE` for each " - "detector, or use `PARAM:VALUE` for all.") + err_msg += ( + "Use format `DETECTOR:PARAM:VALUE` for each " + "detector, or use `PARAM:VALUE` for all." + ) raise ValueError(err_msg) items = detector_args setattr(namespace, self.dest, items) + def required_opts(opt, parser, opt_list, required_by=None): - """Check that all the opts are defined + """ + Check that all the opts are defined Parameters ---------- @@ -349,17 +376,20 @@ def required_opts(opt, parser, opt_list, required_by=None): opt_list : list of strings required_by : string, optional the option that requires these options (if applicable) + """ for name in opt_list: - attr = name[2:].replace('-', '_') + attr = name[2:].replace("-", "_") if not hasattr(opt, attr) or (getattr(opt, attr) is None): err_str = "%s is missing " % name if required_by is not None: err_str += ", required by %s" % required_by parser.error(err_str) + def required_opts_multi_ifo(opt, parser, ifo, opt_list, required_by=None): - """Check that all the opts are defined + """ + Check that all the opts are defined Parameters ---------- @@ -371,9 +401,10 @@ def required_opts_multi_ifo(opt, parser, ifo, opt_list, required_by=None): opt_list : list of strings required_by : string, optional the option that requires these options (if applicable) + """ for name in opt_list: - attr = name[2:].replace('-', '_') + attr = name[2:].replace("-", "_") try: if getattr(opt, attr)[ifo] is None: raise KeyError @@ -383,8 +414,10 @@ def required_opts_multi_ifo(opt, parser, ifo, opt_list, required_by=None): err_str += ", required by %s" % required_by parser.error(err_str) + def ensure_one_opt(opt, parser, opt_list): - """ Check that one and only one in the opt_list is defined in opt + """ + Check that one and only one in the opt_list is defined in opt Parameters ---------- @@ -393,24 +426,24 @@ def ensure_one_opt(opt, parser, opt_list): parser : object OptionParser instance. opt_list : list of strings - """ + """ the_one = None for name in opt_list: - attr = name[2:].replace('-', '_') + attr = name[2:].replace("-", "_") if hasattr(opt, attr) and (getattr(opt, attr) is not None): if the_one is None: the_one = name else: - parser.error("%s and %s are mutually exculsive" \ - % (the_one, name)) + parser.error("%s and %s are mutually exculsive" % (the_one, name)) if the_one is None: - parser.error("you must supply one of the following %s" \ - % (', '.join(opt_list))) + parser.error("you must supply one of the following %s" % (", ".join(opt_list))) + def ensure_one_opt_multi_ifo(opt, parser, ifo, opt_list): - """ Check that one and only one in the opt_list is defined in opt + """ + Check that one and only one in the opt_list is defined in opt Parameters ---------- @@ -419,11 +452,11 @@ def ensure_one_opt_multi_ifo(opt, parser, ifo, opt_list): parser : object OptionParser instance. opt_list : list of strings - """ + """ the_one = None for name in opt_list: - attr = name[2:].replace('-', '_') + attr = name[2:].replace("-", "_") try: if getattr(opt, attr)[ifo] is None: raise KeyError @@ -433,12 +466,11 @@ def ensure_one_opt_multi_ifo(opt, parser, ifo, opt_list): if the_one is None: the_one = name else: - parser.error("%s and %s are mutually exculsive" \ - % (the_one, name)) + parser.error("%s and %s are mutually exculsive" % (the_one, name)) if the_one is None: - parser.error("you must supply one of the following %s" \ - % (', '.join(opt_list))) + parser.error("you must supply one of the following %s" % (", ".join(opt_list))) + def copy_opts_for_single_ifo(opt, ifo): """ @@ -448,11 +480,13 @@ def copy_opts_for_single_ifo(opt, ifo): """ opt = copy.deepcopy(opt) for arg, val in vars(opt).items(): - if isinstance(val, DictWithDefaultReturn) or \ - (isinstance(val, dict) and ifo in val): + if isinstance(val, DictWithDefaultReturn) or ( + isinstance(val, dict) and ifo in val + ): setattr(opt, arg, getattr(opt, arg)[ifo]) return opt + def convert_to_process_params_dict(opt): """ Takes the namespace object (opt) from the multi-detector interface and @@ -467,13 +501,13 @@ def convert_to_process_params_dict(opt): if isinstance(val[key], list): for item in val[key]: if item is not None: - new_val.append(':'.join([key, str(item)])) - else: - if val[key] is not None: - new_val.append(':'.join([key, str(val[key])])) + new_val.append(":".join([key, str(item)])) + elif val[key] is not None: + new_val.append(":".join([key, str(val[key])])) setattr(opt, arg, new_val) return vars(opt) + def _positive_type(s, dtype=None): """ Ensure argument is positive and convert type to dtype @@ -490,6 +524,7 @@ def _positive_type(s, dtype=None): raise argparse.ArgumentTypeError(err_msg) return value + def _nonnegative_type(s, dtype=None): """ Ensure argument is positive or zero and convert type to dtype @@ -506,6 +541,7 @@ def _nonnegative_type(s, dtype=None): raise argparse.ArgumentTypeError(err_msg) return value + def positive_float(s): """ Ensure argument is a positive real number and return it as float. @@ -514,6 +550,7 @@ def positive_float(s): """ return _positive_type(s, dtype=float) + def nonnegative_float(s): """ Ensure argument is a positive real number or zero and return it as float. @@ -522,6 +559,7 @@ def nonnegative_float(s): """ return _nonnegative_type(s, dtype=float) + def positive_int(s): """ Ensure argument is a positive integer and return it as int. @@ -530,6 +568,7 @@ def positive_int(s): """ return _positive_type(s, dtype=int) + def nonnegative_int(s): """ Ensure argument is a positive integer or zero and return it as int. @@ -538,6 +577,7 @@ def nonnegative_int(s): """ return _nonnegative_type(s, dtype=int) + def angle_as_radians(s): """ Interpret argument as a string defining an angle, which will be converted @@ -559,18 +599,16 @@ def angle_as_radians(s): # if `s` converts to a float then there is no unit, so assume radians try: value = float(s) - warnings.warn( - f'Angle units not specified for {value}, assuming radians' - ) + warnings.warn(f"Angle units not specified for {value}, assuming radians") return value except: pass # looks like we have units, so do some parsing - rematch = re.match('([0-9.e+-]+) *(deg|rad)', s) + rematch = re.match("([0-9.e+-]+) *(deg|rad)", s) value = float(rematch.group(1)) unit = rematch.group(2) - if unit == 'deg': + if unit == "deg": return math.radians(value) - if unit == 'rad': + if unit == "rad": return value - raise argparse.ArgumentTypeError(f'Unknown unit {unit}') + raise argparse.ArgumentTypeError(f"Unknown unit {unit}") diff --git a/pycbc/types/timeseries.py b/pycbc/types/timeseries.py index d90044be154..69c17c5c0ba 100644 --- a/pycbc/types/timeseries.py +++ b/pycbc/types/timeseries.py @@ -17,24 +17,32 @@ """ Provides a class representing a time series. """ + import os as _os -import h5py +import h5py import numpy as _numpy +from igwn_segments import segment, segmentlist from scipy.io.wavfile import write as write_wav -from igwn_segments import segmentlist, segment -from pycbc.types.array import Array, _convert, complex_same_precision_as, zeros -from pycbc.types.utils import determine_epoch -from pycbc.types.array import _nocomplex -from pycbc.types.frequencyseries import FrequencySeries -from pycbc.types import float32, float64 from pycbc.libutils import import_optional +from pycbc.types import float32, float64 +from pycbc.types.array import ( + Array, + _convert, + _nocomplex, + complex_same_precision_as, + zeros, +) +from pycbc.types.frequencyseries import FrequencySeries +from pycbc.types.utils import determine_epoch + +_lal = import_optional("lal") -_lal = import_optional('lal') class TimeSeries(Array): - """Models a time series consisting of uniformly sampled scalar values. + """ + Models a time series consisting of uniformly sampled scalar values. Parameters ---------- @@ -48,54 +56,56 @@ class TimeSeries(Array): Sample data type. copy : boolean, optional If True, samples are copied to a new array. + """ - def __init__(self, initial_array, delta_t=None, - epoch="", dtype=None, copy=True): + def __init__(self, initial_array, delta_t=None, epoch="", dtype=None, copy=True): if len(initial_array) < 1: - raise ValueError('initial_array must contain at least one sample.') + raise ValueError("initial_array must contain at least one sample.") if delta_t is None: try: delta_t = initial_array.delta_t except AttributeError: - raise TypeError('must provide either an initial_array with a delta_t attribute, or a value for delta_t') + raise TypeError( + "must provide either an initial_array with a delta_t attribute, or a value for delta_t" + ) if not delta_t > 0: - raise ValueError('delta_t must be a positive number') + raise ValueError("delta_t must be a positive number") self._epoch = determine_epoch(epoch, initial_array) Array.__init__(self, initial_array, dtype=dtype, copy=copy) self._delta_t = delta_t - def to_astropy(self, name='pycbc'): - """ Return an astropy.timeseries.TimeSeries instance - """ - from astropy.timeseries import TimeSeries as ATimeSeries + def to_astropy(self, name="pycbc"): + """Return an astropy.timeseries.TimeSeries instance""" from astropy.time import Time + from astropy.timeseries import TimeSeries as ATimeSeries from astropy.units import s - start = Time(float(self.start_time), format='gps', scale='utc') + start = Time(float(self.start_time), format="gps", scale="utc") delta = self.delta_t * s - return ATimeSeries({name: self.numpy()}, - time_start=start, - time_delta=delta, - n_samples=len(self)) + return ATimeSeries( + {name: self.numpy()}, + time_start=start, + time_delta=delta, + n_samples=len(self), + ) def epoch_close(self, other): - """ Check if the epoch is close enough to allow operations """ + """Check if the epoch is close enough to allow operations""" if self._epoch is None or other._epoch is None: return False dt = abs(float(self.start_time - other.start_time)) return dt <= 1e-7 def sample_rate_close(self, other): - """ Check if the sample rate is close enough to allow operations """ - + """Check if the sample rate is close enough to allow operations""" # compare our delta_t either to a another time series' or # to a given sample rate (float) if isinstance(other, TimeSeries): odelta_t = other.delta_t else: - odelta_t = 1.0/other + odelta_t = 1.0 / other if (odelta_t - self.delta_t) / self.delta_t > 1e-4: return False @@ -111,11 +121,13 @@ def _return(self, ary): def _typecheck(self, other): if isinstance(other, TimeSeries): if not self.sample_rate_close(other): - raise ValueError('different delta_t, {} vs {}'.format( - self.delta_t, other.delta_t)) + raise ValueError( + f"different delta_t, {self.delta_t} vs {other.delta_t}" + ) if not self.epoch_close(other): - raise ValueError('different epoch, {} vs {}'.format( - self.start_time, other.start_time)) + raise ValueError( + f"different epoch, {self.start_time} vs {other.start_time}" + ) def _getslice(self, index): # Set the new epoch - index.start or self._epoch may be None @@ -123,8 +135,9 @@ def _getslice(self, index): new_epoch = self._epoch else: if index.start < 0: - raise ValueError(('Negative start index ({})' - ' not supported').format(index.start)) + raise ValueError( + f"Negative start index ({index.start}) not supported" + ) new_epoch = self._epoch + index.start * self._delta_t if index.step is not None: @@ -132,12 +145,13 @@ def _getslice(self, index): else: new_delta_t = self._delta_t - return TimeSeries(Array._getslice(self, index), new_delta_t, - new_epoch, copy=False) - + return TimeSeries( + Array._getslice(self, index), new_delta_t, new_epoch, copy=False + ) def prepend_zeros(self, num): - """Prepend num zeros onto the beginning of this TimeSeries. Update also + """ + Prepend num zeros onto the beginning of this TimeSeries. Update also epoch to include this prepending. """ self.resize(len(self) + num) @@ -145,99 +159,92 @@ def prepend_zeros(self, num): self._epoch = self._epoch - num * self._delta_t def append_zeros(self, num): - """Append num zeros onto the end of this TimeSeries. - """ + """Append num zeros onto the end of this TimeSeries.""" self.resize(len(self) + num) def get_delta_t(self): - """Return time between consecutive samples in seconds. - """ + """Return time between consecutive samples in seconds.""" return self._delta_t - delta_t = property(get_delta_t, - doc="Time between consecutive samples in seconds.") + + delta_t = property(get_delta_t, doc="Time between consecutive samples in seconds.") def get_duration(self): - """Return duration of time series in seconds. - """ + """Return duration of time series in seconds.""" return len(self) * self._delta_t - duration = property(get_duration, - doc="Duration of time series in seconds.") + + duration = property(get_duration, doc="Duration of time series in seconds.") def get_sample_rate(self): - """Return the sample rate of the time series. - """ - return 1.0/self.delta_t - sample_rate = property(get_sample_rate, - doc="The sample rate of the time series.") + """Return the sample rate of the time series.""" + return 1.0 / self.delta_t - def time_slice(self, start, end, mode='floor'): - """Return the slice of the time series that contains the time range + sample_rate = property(get_sample_rate, doc="The sample rate of the time series.") + + def time_slice(self, start, end, mode="floor"): + """ + Return the slice of the time series that contains the time range in GPS seconds. """ if start < self.start_time: - raise ValueError('Time series does not contain a time as early as %s' % start) + raise ValueError( + "Time series does not contain a time as early as %s" % start + ) if end > self.end_time: - raise ValueError('Time series does not contain a time as late as %s' % end) + raise ValueError("Time series does not contain a time as late as %s" % end) start_idx = float(start - self.start_time) * self.sample_rate end_idx = float(end - self.start_time) * self.sample_rate - if _numpy.isclose(start_idx, round(start_idx), rtol=0, atol=1E-3): + if _numpy.isclose(start_idx, round(start_idx), rtol=0, atol=1e-3): start_idx = round(start_idx) - if _numpy.isclose(end_idx, round(end_idx), rtol=0, atol=1E-3): + if _numpy.isclose(end_idx, round(end_idx), rtol=0, atol=1e-3): end_idx = round(end_idx) - if mode == 'floor': + if mode == "floor": start_idx = int(start_idx) end_idx = int(end_idx) - elif mode == 'nearest': + elif mode == "nearest": start_idx = int(round(start_idx)) end_idx = int(round(end_idx)) else: - raise ValueError("Invalid mode: {}".format(mode)) + raise ValueError(f"Invalid mode: {mode}") return self[start_idx:end_idx] @property def delta_f(self): - """Return the delta_f this ts would have in the frequency domain - """ + """Return the delta_f this ts would have in the frequency domain""" return 1.0 / self.duration @property def start_time(self): - """Return time series start time. - """ + """Return time series start time.""" return self._epoch @start_time.setter def start_time(self, time): - """ Set the start time - """ + """Set the start time""" self._epoch = float64(time) def get_end_time(self): - """Return time series end time. - """ + """Return time series end time.""" return self._epoch + self.get_duration() - end_time = property(get_end_time, - doc="Time series end time.") + + end_time = property(get_end_time, doc="Time series end time.") def get_sample_times(self): - """Return an Array containing the sample times. - """ + """Return an Array containing the sample times.""" if self._epoch is None: return Array(range(len(self))) * self._delta_t - else: - return Array(range(len(self))) * self._delta_t + float(self._epoch) - sample_times = property(get_sample_times, - doc="Array containing the sample times.") + return Array(range(len(self))) * self._delta_t + float(self._epoch) - def at_time(self, time, nearest_sample=False, - interpolate=None, extrapolate=None): - """Return the value of the TimeSeries at the specified GPS time. + sample_times = property(get_sample_times, doc="Array containing the sample times.") + + def at_time(self, time, nearest_sample=False, interpolate=None, extrapolate=None): + """ + Return the value of the TimeSeries at the specified GPS time. Parameters ---------- @@ -253,6 +260,7 @@ def at_time(self, time, nearest_sample=False, extrapolate: str or float, None Value to return if time is outside the range of the vector or method of extrapolating the value. + """ if nearest_sample: time = time + self.delta_t / 2.0 @@ -265,13 +273,13 @@ def at_time(self, time, nearest_sample=False, if _numpy.isscalar(extrapolate) and _numpy.isreal(extrapolate): fill_value = extrapolate facl = facr = 0 - if interpolate == 'quadratic': + if interpolate == "quadratic": facl = facr = 1.1 - elif interpolate == 'linear': + elif interpolate == "linear": facl, facr = 0.1, 1.1 - left = (vtime >= self.start_time + self.delta_t * facl) - right = (vtime < self.end_time - self.delta_t * facr) + left = vtime >= self.start_time + self.delta_t * facl + right = vtime < self.end_time - self.delta_t * facr keep_idx = _numpy.where(left & right)[0] vtime = vtime[keep_idx] else: @@ -281,11 +289,11 @@ def at_time(self, time, nearest_sample=False, i = _numpy.asarray(_numpy.floor(fi)).astype(int) di = fi - i - if interpolate == 'linear': + if interpolate == "linear": a = self[i] - b = self[i+1] + b = self[i + 1] ans = a + (b - a) * di - elif interpolate == 'quadratic': + elif interpolate == "quadratic": c = self.data[i] xr = self.data[i + 1] - c xl = self.data[i - 1] - c @@ -308,7 +316,7 @@ def at_time(self, time, nearest_sample=False, at_times = at_time - def __eq__(self,other): + def __eq__(self, other): """ This is the Python special method invoked whenever the '==' comparison is used. It will return true if the data of two @@ -341,13 +349,13 @@ def __eq__(self,other): ------- boolean: 'True' if the types, dtypes, lengths, epochs, delta_ts and data of the two objects are each identical. + """ - if super(TimeSeries,self).__eq__(other): - return (self._epoch == other._epoch and self._delta_t == other._delta_t) - else: - return False + if super().__eq__(other): + return self._epoch == other._epoch and self._delta_t == other._delta_t + return False - def almost_equal_elem(self,other,tol,relative=True,dtol=0.0): + def almost_equal_elem(self, other, tol, relative=True, dtol=0.0): """ Compare whether two time series are almost equal, element by element. @@ -393,22 +401,25 @@ def almost_equal_elem(self,other,tol,relative=True,dtol=0.0): boolean: 'True' if the data and delta_ts agree within the tolerance, as interpreted by the 'relative' keyword, and if the types, lengths, dtypes, and epochs are exactly the same. + """ # Check that the delta_t tolerance is non-negative; raise an exception # if needed. - if (dtol < 0.0): + if dtol < 0.0: raise ValueError("Tolerance in delta_t cannot be negative") - if super(TimeSeries,self).almost_equal_elem(other,tol=tol,relative=relative): + if super().almost_equal_elem(other, tol=tol, relative=relative): if relative: - return (self._epoch == other._epoch and - abs(self._delta_t-other._delta_t) <= dtol*self._delta_t) - else: - return (self._epoch == other._epoch and - abs(self._delta_t-other._delta_t) <= dtol) - else: - return False + return ( + self._epoch == other._epoch + and abs(self._delta_t - other._delta_t) <= dtol * self._delta_t + ) + return ( + self._epoch == other._epoch + and abs(self._delta_t - other._delta_t) <= dtol + ) + return False - def almost_equal_norm(self,other,tol,relative=True,dtol=0.0): + def almost_equal_norm(self, other, tol, relative=True, dtol=0.0): """ Compare whether two time series are almost equal, normwise. @@ -451,24 +462,28 @@ def almost_equal_norm(self,other,tol,relative=True,dtol=0.0): boolean: 'True' if the data and delta_ts agree within the tolerance, as interpreted by the 'relative' keyword, and if the types, lengths, dtypes, and epochs are exactly the same. + """ # Check that the delta_t tolerance is non-negative; raise an exception # if needed. - if (dtol < 0.0): + if dtol < 0.0: raise ValueError("Tolerance in delta_t cannot be negative") - if super(TimeSeries,self).almost_equal_norm(other,tol=tol,relative=relative): + if super().almost_equal_norm(other, tol=tol, relative=relative): if relative: - return (self._epoch == other._epoch and - abs(self._delta_t-other._delta_t) <= dtol*self._delta_t) - else: - return (self._epoch == other._epoch and - abs(self._delta_t-other._delta_t) <= dtol) - else: - return False + return ( + self._epoch == other._epoch + and abs(self._delta_t - other._delta_t) <= dtol * self._delta_t + ) + return ( + self._epoch == other._epoch + and abs(self._delta_t - other._delta_t) <= dtol + ) + return False @_convert def lal(self): - """Produces a LAL time series object equivalent to self. + """ + Produces a LAL time series object equivalent to self. Returns ------- @@ -482,25 +497,35 @@ def lal(self): ------ TypeError If time series is stored in GPU memory. + """ lal_data = None ep = _lal.LIGOTimeGPS(self._epoch) if self._data.dtype == _numpy.float32: - lal_data = _lal.CreateREAL4TimeSeries("",ep,0,self.delta_t,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateREAL4TimeSeries( + "", ep, 0, self.delta_t, _lal.SecondUnit, len(self) + ) elif self._data.dtype == _numpy.float64: - lal_data = _lal.CreateREAL8TimeSeries("",ep,0,self.delta_t,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateREAL8TimeSeries( + "", ep, 0, self.delta_t, _lal.SecondUnit, len(self) + ) elif self._data.dtype == _numpy.complex64: - lal_data = _lal.CreateCOMPLEX8TimeSeries("",ep,0,self.delta_t,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateCOMPLEX8TimeSeries( + "", ep, 0, self.delta_t, _lal.SecondUnit, len(self) + ) elif self._data.dtype == _numpy.complex128: - lal_data = _lal.CreateCOMPLEX16TimeSeries("",ep,0,self.delta_t,_lal.SecondUnit,len(self)) + lal_data = _lal.CreateCOMPLEX16TimeSeries( + "", ep, 0, self.delta_t, _lal.SecondUnit, len(self) + ) lal_data.data.data[:] = self.numpy() return lal_data def crop(self, left, right): - """ Remove given seconds from either end of time series + """ + Remove given seconds from either end of time series Parameters ---------- @@ -513,27 +538,31 @@ def crop(self, left, right): ------- cropped : pycbc.types.TimeSeries The reduced time series + """ if left + right > self.duration: - raise ValueError('Cannot crop more data than we have') + raise ValueError("Cannot crop more data than we have") s = int(left * self.sample_rate) e = len(self) - int(right * self.sample_rate) return self[s:e] def save_to_wav(self, file_name): - """ Save this time series to a wav format audio file. + """ + Save this time series to a wav format audio file. Parameters ---------- file_name : string The output file name + """ - scaled = _numpy.int16(self.numpy()/max(abs(self)) * 32767) + scaled = _numpy.int16(self.numpy() / max(abs(self)) * 32767) write_wav(file_name, int(self.sample_rate), scaled) def psd(self, segment_duration, **kwds): - """ Calculate the power spectral density of this time series. + """ + Calculate the power spectral density of this time series. Use the `pycbc.psd.welch` method to estimate the psd of this time segment. For more complete options, please see that function. @@ -549,18 +578,20 @@ def psd(self, segment_duration, **kwds): ------- psd : FrequencySeries Frequency series containing the estimated PSD. + """ from pycbc.psd import welch + seg_len = int(round(segment_duration * self.sample_rate)) seg_stride = int(seg_len / 2) - return welch(self, seg_len=seg_len, - seg_stride=seg_stride, - **kwds) - + return welch(self, seg_len=seg_len, seg_stride=seg_stride, **kwds) + # map between tapering string in sim_inspiral table or inspiral # code option and lalsimulation constants - def taper_timeseries(self, location=None, tapermethod='lal', return_lal=False, taper_window=None): + def taper_timeseries( + self, location=None, tapermethod="lal", return_lal=False, taper_window=None + ): """ Taper either or both ends of a time series using wrapped LALSimulation functions or a constant window taper. @@ -582,40 +613,47 @@ def taper_timeseries(self, location=None, tapermethod='lal', return_lal=False, t return_lal : Boolean If True, return a wrapped LAL time series object, else return a PyCBC time series. + """ import lalsimulation as sim - - if hasattr(location, 'decode'): + + if hasattr(location, "decode"): location = location.decode() - - if hasattr(tapermethod, 'decode'): + + if hasattr(tapermethod, "decode"): tapermethod = tapermethod.decode() taper_map = { - 'TAPER_NONE' : None, - 'TAPER_START' : sim.SIM_INSPIRAL_TAPER_START, - 'start' : sim.SIM_INSPIRAL_TAPER_START, - 'TAPER_END' : sim.SIM_INSPIRAL_TAPER_END, - 'end' : sim.SIM_INSPIRAL_TAPER_END, - 'TAPER_STARTEND': sim.SIM_INSPIRAL_TAPER_STARTEND, - 'startend' : sim.SIM_INSPIRAL_TAPER_STARTEND} + "TAPER_NONE": None, + "TAPER_START": sim.SIM_INSPIRAL_TAPER_START, + "start": sim.SIM_INSPIRAL_TAPER_START, + "TAPER_END": sim.SIM_INSPIRAL_TAPER_END, + "end": sim.SIM_INSPIRAL_TAPER_END, + "TAPER_STARTEND": sim.SIM_INSPIRAL_TAPER_STARTEND, + "startend": sim.SIM_INSPIRAL_TAPER_STARTEND, + } taper_func_map = { _numpy.dtype(float32): sim.SimInspiralREAL4WaveTaper, - _numpy.dtype(float64): sim.SimInspiralREAL8WaveTaper} + _numpy.dtype(float64): sim.SimInspiralREAL8WaveTaper, + } tsdata = self if location is None: - raise ValueError("Must specify a tapering method (function was called" - "with location=None)") - if location not in taper_map.keys(): - raise ValueError("Unknown location %s, valid locations are %s" % \ - (location, ", ".join(taper_map.keys()))) + raise ValueError( + "Must specify a tapering method (function was calledwith location=None)" + ) + if location not in taper_map: + raise ValueError( + "Unknown location %s, valid locations are %s" + % (location, ", ".join(taper_map.keys())) + ) if tsdata.dtype not in (float32, float64): - raise TypeError("Strain dtype must be float32 or float64, not " - + str(tsdata.dtype)) - if tapermethod == 'lal': + raise TypeError( + "Strain dtype must be float32 or float64, not " + str(tsdata.dtype) + ) + if tapermethod == "lal": taper_func = taper_func_map[tsdata.dtype] # make a LAL TimeSeries to pass to the LALSim function ts_lal = tsdata.astype(tsdata.dtype).lal() @@ -623,32 +661,36 @@ def taper_timeseries(self, location=None, tapermethod='lal', return_lal=False, t taper_func(ts_lal.data, taper_map[location]) if return_lal: return ts_lal - else: - return TimeSeries(ts_lal.data.data[:], delta_t=ts_lal.deltaT, - epoch=ts_lal.epoch) - elif tapermethod == 'constant': + return TimeSeries( + ts_lal.data.data[:], delta_t=ts_lal.deltaT, epoch=ts_lal.epoch + ) + if tapermethod == "constant": # constant window tapering if taper_window is None: - raise ValueError("If taper_method is 'constant', taper_window must be set") - + raise ValueError( + "If taper_method is 'constant', taper_window must be set" + ) + gate_params = [] - if location in ('TAPER_START', 'start', 'TAPER_STARTEND'): + if location in ("TAPER_START", "start", "TAPER_STARTEND"): first_nonzero = _numpy.nonzero(tsdata)[0][0] nonzero_starttime = tsdata.start_time + first_nonzero * tsdata.delta_t gate_params.append((nonzero_starttime, 0, taper_window)) - if location in ('TAPER_END', 'end', 'TAPER_STARTEND'): + if location in ("TAPER_END", "end", "TAPER_STARTEND"): last_nonzero = _numpy.nonzero(tsdata)[0][-1] nonzero_endtime = tsdata.end_time - last_nonzero * tsdata.delta_t gate_params.append((nonzero_endtime - taper_window, 0, taper_window)) from pycbc.strain import gate_data + return gate_data(tsdata, gate_params) - else: - raise ValueError("Unknown tapering method %s, valid methods are lal and constant" % \ - (tapermethod)) + raise ValueError( + "Unknown tapering method %s, valid methods are lal and constant" + % (tapermethod) + ) - def get_gate_indices(self, time, window): - """Calculates the indices at which a gate should be applied. + """ + Calculates the indices at which a gate should be applied. Parameters ---------- @@ -663,19 +705,29 @@ def get_gate_indices(self, time, window): The left index of the gate rindex: int The right index of the gate + """ st = float(self.start_time) dt = float(self.delta_t) lindex = int((time - window - st) / dt) rindex = int((time + window - st) / dt) - lindex = lindex if lindex >= 0 else 0 - rindex = rindex if rindex <= len(self) else len(self) + lindex = max(lindex, 0) + rindex = min(rindex, len(self)) return lindex, rindex - def gate(self, time, window=0.25, method='taper', copy=True, - taper_width=0.25, invpsd=None, paint_method='toeplitz', - paint_invmat=None): - """ Gate out portion of time series + def gate( + self, + time, + window=0.25, + method="taper", + copy=True, + taper_width=0.25, + invpsd=None, + paint_method="toeplitz", + paint_invmat=None, + ): + """ + Gate out portion of time series Parameters ---------- @@ -707,59 +759,61 @@ def gate(self, time, window=0.25, method='taper', copy=True, ------- data: pycbc.types.TimeSeries Gated time series + """ data = self.copy() if copy else self - if method == 'taper': + if method == "taper": from pycbc.strain import gate_data + return gate_data(data, [(time, window, taper_width)]) - elif method == 'paint': + if method == "paint": # Uses the hole-filling method of # https://arxiv.org/pdf/1908.05644.pdf - from pycbc.strain.gate import (gate_and_paint, - gate_and_paint_matmul) + from pycbc.strain.gate import gate_and_paint, gate_and_paint_matmul from pycbc.waveform.utils import apply_fd_time_shift + if invpsd is None: # These are some bare minimum settings, normally you # should probably provide a psd - invpsd = 1. / self.filter_psd(self.duration/32, self.delta_f, 0) + invpsd = 1.0 / self.filter_psd(self.duration / 32, self.delta_f, 0) lindex, rindex = self.get_gate_indices(time, window) rindex_time = float(self.start_time + rindex * self.delta_t) offset = rindex_time - (time + window) if offset == 0: - if paint_method == 'toeplitz': + if paint_method == "toeplitz": return gate_and_paint(data, lindex, rindex, invpsd, copy=False) - elif paint_method == 'matmul': - return gate_and_paint_matmul(data, lindex, rindex, invpsd, - invmat=paint_invmat, copy=False) - else: - raise ValueError(f'Unrecognized paint_method input {paint_method}') + if paint_method == "matmul": + return gate_and_paint_matmul( + data, lindex, rindex, invpsd, invmat=paint_invmat, copy=False + ) + raise ValueError(f"Unrecognized paint_method input {paint_method}") + # time shift such that gate end time lands on a specific data sample + fdata = data.to_frequencyseries() + fdata = apply_fd_time_shift(fdata, offset + fdata.epoch, copy=False) + # gate and paint in time domain + data = fdata.to_timeseries() + if paint_method == "toeplitz": + data = gate_and_paint(data, lindex, rindex, invpsd, copy=False) + elif paint_method == "matmul": + data = gate_and_paint_matmul( + data, lindex, rindex, invpsd, invmat=paint_invmat, copy=False + ) else: - # time shift such that gate end time lands on a specific data sample - fdata = data.to_frequencyseries() - fdata = apply_fd_time_shift(fdata, offset + fdata.epoch, copy=False) - # gate and paint in time domain - data = fdata.to_timeseries() - if paint_method == 'toeplitz': - data = gate_and_paint(data, lindex, rindex, invpsd, copy=False) - elif paint_method == 'matmul': - data = gate_and_paint_matmul(data, lindex, rindex, invpsd, - invmat=paint_invmat, copy=False) - else: - raise ValueError(f'Unrecognized paint_method input {paint_method}') - # shift back to the original time - fdata = data.to_frequencyseries() - fdata = apply_fd_time_shift(fdata, -offset + fdata.epoch, copy=False) - tdata = fdata.to_timeseries() - return tdata - elif method == 'hard': + raise ValueError(f"Unrecognized paint_method input {paint_method}") + # shift back to the original time + fdata = data.to_frequencyseries() + fdata = apply_fd_time_shift(fdata, -offset + fdata.epoch, copy=False) + tdata = fdata.to_timeseries() + return tdata + if method == "hard": tslice = data.time_slice(time - window, time + window) tslice[:] = 0 return data - else: - raise ValueError('Invalid method name: {}'.format(method)) + raise ValueError(f"Invalid method name: {method}") def filter_psd(self, segment_duration, delta_f, flow): - """ Calculate the power spectral density of this time series. + """ + Calculate the power spectral density of this time series. Use the `pycbc.psd.welch` method to estimate the psd of this time segment. The psd is then truncated in the time domain to the segment duration @@ -779,19 +833,29 @@ def filter_psd(self, segment_duration, delta_f, flow): ------- psd : FrequencySeries Frequency series containing the estimated PSD. + """ from pycbc.psd import interpolate, inverse_spectrum_truncation + p = self.psd(segment_duration) samples = int(round(p.sample_rate * segment_duration)) p = interpolate(p, delta_f) - return inverse_spectrum_truncation(p, samples, - low_frequency_cutoff=flow, - trunc_method='hann') - - def whiten(self, segment_duration, max_filter_duration, trunc_method='hann', - remove_corrupted=True, low_frequency_cutoff=None, - return_psd=False, **kwds): - """ Return a whitened time series + return inverse_spectrum_truncation( + p, samples, low_frequency_cutoff=flow, trunc_method="hann" + ) + + def whiten( + self, + segment_duration, + max_filter_duration, + trunc_method="hann", + remove_corrupted=True, + low_frequency_cutoff=None, + return_psd=False, + **kwds, + ): + """ + Return a whitened time series Parameters ---------- @@ -820,33 +884,46 @@ def whiten(self, segment_duration, max_filter_duration, trunc_method='hann', ------- whitened_data : TimeSeries The whitened time series + """ - from pycbc.psd import inverse_spectrum_truncation, interpolate + from pycbc.psd import interpolate, inverse_spectrum_truncation + # Estimate the noise spectrum psd = self.psd(segment_duration, **kwds) psd = interpolate(psd, self.delta_f) max_filter_len = int(round(max_filter_duration * self.sample_rate)) # Interpolate and smooth to the desired corruption length - psd = inverse_spectrum_truncation(psd, - max_filter_len=max_filter_len, - low_frequency_cutoff=low_frequency_cutoff, - trunc_method=trunc_method) + psd = inverse_spectrum_truncation( + psd, + max_filter_len=max_filter_len, + low_frequency_cutoff=low_frequency_cutoff, + trunc_method=trunc_method, + ) # Whiten the data by the asd white = (self.to_frequencyseries() / psd**0.5).to_timeseries() if remove_corrupted: - white = white[int(max_filter_len/2):int(len(self)-max_filter_len/2)] + white = white[int(max_filter_len / 2) : int(len(self) - max_filter_len / 2)] if return_psd: return white, psd return white - def qtransform(self, delta_t=None, delta_f=None, logfsteps=None, - frange=None, qrange=(4,64), mismatch=0.2, return_complex=False): - """ Return the interpolated 2d qtransform of this data + def qtransform( + self, + delta_t=None, + delta_f=None, + logfsteps=None, + frange=None, + qrange=(4, 64), + mismatch=0.2, + return_complex=False, + ): + """ + Return the interpolated 2d qtransform of this data Parameters ---------- @@ -874,16 +951,19 @@ def qtransform(self, delta_t=None, delta_f=None, logfsteps=None, The frequencies that the qtransform is sampled. qplane : numpy.ndarray (2d) The two dimensional interpolated qtransform of this time series. + """ - from pycbc.filter.qtransform import qtiling, qplane from scipy.interpolate import RectBivariateSpline as interp2d + from pycbc.filter.qtransform import qplane, qtiling + if frange is None: frange = (30, int(self.sample_rate / 2 * 8)) q_base = qtiling(self, qrange, frange, mismatch) - _, times, freqs, q_plane = qplane(q_base, self.to_frequencyseries(), - return_complex=return_complex) + _, times, freqs, q_plane = qplane( + q_base, self.to_frequencyseries(), return_complex=return_complex + ) if logfsteps and delta_f: raise ValueError("Provide only one (or none) of delta_f and logfsteps") @@ -891,20 +971,18 @@ def qtransform(self, delta_t=None, delta_f=None, logfsteps=None, if delta_f or delta_t or logfsteps: if return_complex: interp_amp = interp2d(freqs, times, abs(q_plane), kx=1, ky=1) - interp_phase = interp2d(freqs, times, _numpy.angle(q_plane), - kx=1, ky=1) + interp_phase = interp2d(freqs, times, _numpy.angle(q_plane), kx=1, ky=1) else: interp = interp2d(freqs, times, q_plane, kx=1, ky=1) if delta_t: - times = _numpy.arange(float(self.start_time), - float(self.end_time), delta_t) + times = _numpy.arange(float(self.start_time), float(self.end_time), delta_t) if delta_f: freqs = _numpy.arange(int(frange[0]), int(frange[1]), delta_f) if logfsteps: - freqs = _numpy.logspace(_numpy.log10(frange[0]), - _numpy.log10(frange[1]), - logfsteps) + freqs = _numpy.logspace( + _numpy.log10(frange[0]), _numpy.log10(frange[1]), logfsteps + ) if delta_f or delta_t or logfsteps: if return_complex: @@ -916,7 +994,8 @@ def qtransform(self, delta_t=None, delta_f=None, logfsteps=None, return times, freqs, q_plane def notch_fir(self, f1, f2, order, beta=5.0, remove_corrupted=True): - """ notch filter the time series using an FIR filtered generated from + """ + Notch filter the time series using an FIR filtered generated from the ideal response passed through a time-domain kaiser window (beta = 5.0) @@ -937,15 +1016,18 @@ def notch_fir(self, f1, f2, order, beta=5.0, remove_corrupted=True): Number of corrupted samples on each side of the time series beta: float Beta parameter of the kaiser window that sets the side lobe attenuation. + """ from pycbc.filter import notch_fir + ts = notch_fir(self, f1, f2, order, beta=beta) if remove_corrupted: - ts = ts[order:len(ts)-order] + ts = ts[order : len(ts) - order] return ts def lowpass_fir(self, frequency, order, beta=5.0, remove_corrupted=True): - """ Lowpass filter the time series using an FIR filtered generated from + """ + Lowpass filter the time series using an FIR filtered generated from the ideal response passed through a kaiser window (beta = 5.0) Parameters @@ -962,15 +1044,18 @@ def lowpass_fir(self, frequency, order, beta=5.0, remove_corrupted=True): If True, the region of the time series corrupted by the filtering is excised before returning. If false, the corrupted regions are not excised and the full time series is returned. + """ from pycbc.filter import lowpass_fir + ts = lowpass_fir(self, frequency, order, beta=beta) if remove_corrupted: - ts = ts[order:len(ts)-order] + ts = ts[order : len(ts) - order] return ts def highpass_fir(self, frequency, order, beta=5.0, remove_corrupted=True): - """ Highpass filter the time series using an FIR filtered generated from + """ + Highpass filter the time series using an FIR filtered generated from the ideal response passed through a kaiser window (beta = 5.0) Parameters @@ -987,15 +1072,18 @@ def highpass_fir(self, frequency, order, beta=5.0, remove_corrupted=True): If True, the region of the time series corrupted by the filtering is excised before returning. If false, the corrupted regions are not excised and the full time series is returned. + """ from pycbc.filter import highpass_fir + ts = highpass_fir(self, frequency, order, beta=beta) if remove_corrupted: - ts = ts[order:len(ts)-order] + ts = ts[order : len(ts) - order] return ts def fir_zero_filter(self, coeff): - """Filter the timeseries with a set of FIR coefficients + """ + Filter the timeseries with a set of FIR coefficients Parameters ---------- @@ -1007,15 +1095,18 @@ def fir_zero_filter(self, coeff): filtered_series: pycbc.types.TimeSeries Return the filtered timeseries, which has been properly shifted to account for the FIR filter delay and the corrupted regions zeroed out. + """ from pycbc.filter import fir_zero_filter + return self._return(fir_zero_filter(coeff, self)) def resample(self, delta_t): - """ Resample this time series to the new delta_t + """ + Resample this time series to the new delta_t Parameters - ----------- + ---------- delta_t: float The time step to resample the times series to. @@ -1023,11 +1114,13 @@ def resample(self, delta_t): ------- resampled_ts: pycbc.types.TimeSeries The resample timeseries at the new time interval delta_t. + """ from pycbc.filter import resample_to_delta_t + return resample_to_delta_t(self, delta_t) - def save(self, path, group = None): + def save(self, path, group=None): """ Save time series to a Numpy .npy, hdf, or text file. The first column contains the sample times, the second contains the values. @@ -1048,39 +1141,43 @@ def save(self, path, group = None): ------ ValueError If path does not end in .npy or .txt. - """ + """ ext = _os.path.splitext(path)[1] - if ext == '.npy': + if ext == ".npy": output = _numpy.vstack((self.sample_times.numpy(), self.numpy())).T _numpy.save(path, output) - elif ext == '.txt': - if self.kind == 'real': - output = _numpy.vstack((self.sample_times.numpy(), - self.numpy())).T - elif self.kind == 'complex': - output = _numpy.vstack((self.sample_times.numpy(), - self.numpy().real, - self.numpy().imag)).T + elif ext == ".txt": + if self.kind == "real": + output = _numpy.vstack((self.sample_times.numpy(), self.numpy())).T + elif self.kind == "complex": + output = _numpy.vstack( + (self.sample_times.numpy(), self.numpy().real, self.numpy().imag) + ).T _numpy.savetxt(path, output) - elif ext =='.hdf': - key = 'data' if group is None else group - with h5py.File(path, 'a') as f: - ds = f.create_dataset(key, data=self.numpy(), - compression='gzip', - compression_opts=9, shuffle=True) - ds.attrs['start_time'] = float(self.start_time) - ds.attrs['delta_t'] = float(self.delta_t) + elif ext == ".hdf": + key = "data" if group is None else group + with h5py.File(path, "a") as f: + ds = f.create_dataset( + key, + data=self.numpy(), + compression="gzip", + compression_opts=9, + shuffle=True, + ) + ds.attrs["start_time"] = float(self.start_time) + ds.attrs["delta_t"] = float(self.delta_t) else: - raise ValueError('Path must end with .npy, .txt or .hdf') + raise ValueError("Path must end with .npy, .txt or .hdf") def to_timeseries(self): - """ Return time series""" + """Return time series""" return self @_nocomplex def to_frequencyseries(self, delta_f=None): - """ Return the Fourier transform of this time series + """ + Return the Fourier transform of this time series Parameters ---------- @@ -1092,35 +1189,43 @@ def to_frequencyseries(self, delta_f=None): ------- FrequencySeries: The fourier transform of this time series. + """ from pycbc.fft import fft + if not delta_f: delta_f = 1.0 / self.duration # add 0.5 to round integer - tlen = int(1.0 / delta_f / self.delta_t + 0.5) + tlen = int(1.0 / delta_f / self.delta_t + 0.5) flen = int(tlen / 2 + 1) if tlen < len(self): - raise ValueError("The value of delta_f (%s) would be " - "undersampled. Maximum delta_f " - "is %s." % (delta_f, 1.0 / self.duration)) + raise ValueError( + "The value of delta_f (%s) would be " + "undersampled. Maximum delta_f " + "is %s." % (delta_f, 1.0 / self.duration) + ) if not delta_f: tmp = self else: - tmp = TimeSeries(zeros(tlen, dtype=self.dtype), - delta_t=self.delta_t, epoch=self.start_time) - tmp[:len(self)] = self[:] + tmp = TimeSeries( + zeros(tlen, dtype=self.dtype), + delta_t=self.delta_t, + epoch=self.start_time, + ) + tmp[: len(self)] = self[:] - f = FrequencySeries(zeros(flen, - dtype=complex_same_precision_as(self)), - delta_f=delta_f) + f = FrequencySeries( + zeros(flen, dtype=complex_same_precision_as(self)), delta_f=delta_f + ) fft(tmp, f) f._delta_f = delta_f return f def inject(self, other, copy=True): - """Return copy of self with other injected into it. + """ + Return copy of self with other injected into it. The other vector will be resized and time shifted with sub-sample precision before adding. This assumes that one can assume zeros @@ -1128,15 +1233,14 @@ def inject(self, other, copy=True): """ # only handle equal sample rate for now. if not self.sample_rate_close(other): - raise ValueError('Sample rate must be the same') + raise ValueError("Sample rate must be the same") # determine if we want to inject in place or not if copy: ts = self.copy() else: ts = self # Other is disjoint - if ((other.start_time >= ts.end_time) or - (ts.start_time > other.end_time)): + if (other.start_time >= ts.end_time) or (ts.start_time > other.end_time): return ts other = other.copy() @@ -1178,7 +1282,8 @@ def inject(self, other, copy=True): @_nocomplex def cyclic_time_shift(self, dt): - """Shift the data and timestamps by a given number of seconds + """ + Shift the data and timestamps by a given number of seconds Shift the data and timestamps in the time domain a given number of seconds. To just change the time stamps, do ts.start_time += dt. @@ -1196,6 +1301,7 @@ def cyclic_time_shift(self, dt): ------- data : pycbc.types.TimeSeries The time shifted time series. + """ # We do this in the frequency domain to allow us to do sub-sample # time shifts. This also results in the shift being circular. It @@ -1203,9 +1309,11 @@ def cyclic_time_shift(self, dt): # where the time shift can be done with an exact number of samples. return self.to_frequencyseries().cyclic_time_shift(dt).to_timeseries() - def match(self, other, psd=None, - low_frequency_cutoff=None, high_frequency_cutoff=None): - """ Return the match between the two TimeSeries or FrequencySeries. + def match( + self, other, psd=None, low_frequency_cutoff=None, high_frequency_cutoff=None + ): + """ + Return the match between the two TimeSeries or FrequencySeries. Return the match between two waveforms. This is equivalent to the overlap maximized over time and phase. By default, the other vector will be @@ -1228,13 +1336,18 @@ def match(self, other, psd=None, match: float index: int The number of samples to shift to get the match. - """ - return self.to_frequencyseries().match(other, psd=psd, - low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff) - def detrend(self, type='linear'): - """ Remove linear trend from the data + """ + return self.to_frequencyseries().match( + other, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + ) + + def detrend(self, type="linear"): + """ + Remove linear trend from the data Remove a linear trend from the data to improve the approximation that the data is circularly convolved, this helps reduce the size of filter @@ -1245,19 +1358,20 @@ def detrend(self, type='linear'): type: str The choice of detrending. The default ('linear') removes a linear least squares fit. 'constant' removes only the mean of the data. + """ from scipy.signal import detrend + return self._return(detrend(self.numpy(), type=type)) def plot(self, **kwds): - """ Basic plot of this time series - """ + """Basic plot of this time series""" from matplotlib import pyplot - if self.kind == 'real': + if self.kind == "real": plot = pyplot.plot(self.sample_times, self, **kwds) return plot - elif self.kind == 'complex': + if self.kind == "complex": plot1 = pyplot.plot(self.sample_times, self.real(), **kwds) plot2 = pyplot.plot(self.sample_times, self.imag(), **kwds) return plot1, plot2 @@ -1267,14 +1381,13 @@ def bool_to_segmentlist(self): Convert a boolean pycbc TimeSeries (this must be bool or integer) to an igwn_segments.segmentlist of (start, end) in GPS seconds. """ - # Is the data truthlike? # bools or numbers are OK, but we require finite values arr = self.numpy() - if arr.dtype.kind not in ['b', 'i']: + if arr.dtype.kind not in ["b", "i"]: raise TypeError( - 'To use bool_to_segmentlist, we require that the timeseries ' - 'is boolean or integer' + "To use bool_to_segmentlist, we require that the timeseries " + "is boolean or integer" ) segs = segmentlist([]) @@ -1292,7 +1405,7 @@ def bool_to_segmentlist(self): # starts = False to True transitions starts = _numpy.flatnonzero((~b[:-1]) & b[1:]) # ends = True to False Transitions - ends = _numpy.flatnonzero(b[:-1] & (~b[1:])) + ends = _numpy.flatnonzero(b[:-1] & (~b[1:])) # Convert indices to GPS times starts_time = self.start_time + starts * self.delta_t @@ -1304,8 +1417,10 @@ def bool_to_segmentlist(self): return segs.coalesce() + def load_timeseries(path, group=None): - """Load a TimeSeries from an HDF5, ASCII or Numpy file. The file type is + """ + Load a TimeSeries from an HDF5, ASCII or Numpy file. The file type is inferred from the file extension, which must be `.hdf`, `.txt` or `.npy`. For ASCII and Numpy files, the first column of the array is assumed to @@ -1335,29 +1450,33 @@ def load_timeseries(path, group=None): If path does not end in a supported extension. For Numpy and ASCII input files, this is also raised if the array does not have 2 or 3 dimensions. + """ ext = _os.path.splitext(path)[1] - if ext == '.npy': + if ext == ".npy": data = _numpy.load(path) - elif ext == '.txt': + elif ext == ".txt": data = _numpy.loadtxt(path) - elif ext == '.hdf': - key = 'data' if group is None else group - with h5py.File(path, 'r') as f: + elif ext == ".hdf": + key = "data" if group is None else group + with h5py.File(path, "r") as f: data = f[key][:] - series = TimeSeries(data, delta_t=f[key].attrs['delta_t'], - epoch=f[key].attrs['start_time']) + series = TimeSeries( + data, delta_t=f[key].attrs["delta_t"], epoch=f[key].attrs["start_time"] + ) return series else: - raise ValueError('Path must end with .npy, .hdf, or .txt') + raise ValueError("Path must end with .npy, .hdf, or .txt") delta_t = (data[-1][0] - data[0][0]) / (len(data) - 1) epoch = _lal.LIGOTimeGPS(data[0][0]) if data.ndim == 2: - return TimeSeries(data[:,1], delta_t=delta_t, epoch=epoch) - elif data.ndim == 3: - return TimeSeries(data[:,1] + 1j*data[:,2], - delta_t=delta_t, epoch=epoch) - - raise ValueError('File has %s dimensions, cannot convert to TimeSeries, \ - must be 2 (real) or 3 (complex)' % data.ndim) + return TimeSeries(data[:, 1], delta_t=delta_t, epoch=epoch) + if data.ndim == 3: + return TimeSeries(data[:, 1] + 1j * data[:, 2], delta_t=delta_t, epoch=epoch) + + raise ValueError( + "File has %s dimensions, cannot convert to TimeSeries, \ + must be 2 (real) or 3 (complex)" + % data.ndim + ) diff --git a/pycbc/types/utils.py b/pycbc/types/utils.py index 37b001911c6..e4e9ea4b48b 100644 --- a/pycbc/types/utils.py +++ b/pycbc/types/utils.py @@ -1,12 +1,14 @@ import logging + import numpy as _numpy -from numpy import float64 +from numpy import float64 from pycbc.libutils import import_optional -logger = logging.getLogger('pycbc.type.utils') +logger = logging.getLogger("pycbc.type.utils") + +_lal = import_optional("lal") -_lal = import_optional('lal') def determine_epoch(epoch, initial_array): """ @@ -28,7 +30,7 @@ def determine_epoch(epoch, initial_array): Parameters ---------- - epoch: + epoch: float64/number-type, LIGOTimeGPS, None initial_array: Array - only really matters if this has an _epoch set already @@ -36,12 +38,11 @@ def determine_epoch(epoch, initial_array): Returns ------- epoch: float64 or None - see logic above - """ - + """ if isinstance(epoch, float64) or epoch is None: return epoch - + if epoch == "": # The default has been given, try these: try: @@ -63,16 +64,15 @@ def determine_epoch(epoch, initial_array): # It looks like this is an array/list/tuple, so float conversion could # succeed, but we shouldn't be trying it if not is_ltg and not _numpy.isscalar(epoch): - # Its not a + # Its not a raise TypeError("epoch must be a number, not array-like") - + try: # Okay we have gone through the special cases now, just try it and see return float64(epoch) except TypeError as e: # Give something helpful before failing. logger.warning( - "epoch cannot be determined: " - f"type: {type(epoch)}, value: {epoch}" + f"epoch cannot be determined: type: {type(epoch)}, value: {epoch}" ) - raise e \ No newline at end of file + raise e diff --git a/pycbc/vetoes/__init__.py b/pycbc/vetoes/__init__.py index 7cf1014bd76..6058e3e7509 100644 --- a/pycbc/vetoes/__init__.py +++ b/pycbc/vetoes/__init__.py @@ -1,3 +1,3 @@ -from .chisq import * -from .bank_chisq import * from .autochisq import * +from .bank_chisq import * +from .chisq import * diff --git a/pycbc/vetoes/autochisq.py b/pycbc/vetoes/autochisq.py index 3d12902e31f..e76f11e20a4 100644 --- a/pycbc/vetoes/autochisq.py +++ b/pycbc/vetoes/autochisq.py @@ -13,18 +13,27 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -from pycbc.filter import make_frequency_series -from pycbc.filter import matched_filter_core -from pycbc.types import Array -import numpy as np import logging -BACKEND_PREFIX="pycbc.vetoes.autochisq_" +import numpy as np + +from pycbc.filter import make_frequency_series, matched_filter_core +from pycbc.types import Array +BACKEND_PREFIX = "pycbc.vetoes.autochisq_" -def autochisq_from_precomputed(sn, corr_sn, hautocorr, indices, - stride=1, num_points=None, oneside=None, - twophase=True, maxvalued=False): + +def autochisq_from_precomputed( + sn, + corr_sn, + hautocorr, + indices, + stride=1, + num_points=None, + oneside=None, + twophase=True, + maxvalued=False, +): """ Compute correlation (two sided) between template and data and compares with autocorrelation of the template: C(t) = IFFT(A*A/S(f)) @@ -65,63 +74,66 @@ def autochisq_from_precomputed(sn, corr_sn, hautocorr, indices, number of degrees of freedom autochisq: Array[float] autochisq values corresponding to the time instances defined by indices + """ Nsnr = len(sn) achisq = np.zeros(len(indices)) - num_points_all = int(Nsnr/stride) + num_points_all = int(Nsnr / stride) if num_points is None: num_points = num_points_all - if (num_points > num_points_all): - num_points = num_points_all + num_points = min(num_points, num_points_all) snrabs = np.abs(sn[indices]) cphi_array = (sn[indices]).real / snrabs sphi_array = (sn[indices]).imag / snrabs - start_point = - stride*num_points - end_point = stride*num_points+1 - if oneside == 'left': + start_point = -stride * num_points + end_point = stride * num_points + 1 + if oneside == "left": achisq_idx_list = np.arange(start_point, 0, stride) - elif oneside == 'right': + elif oneside == "right": achisq_idx_list = np.arange(stride, end_point, stride) else: achisq_idx_list_pt1 = np.arange(start_point, 0, stride) achisq_idx_list_pt2 = np.arange(stride, end_point, stride) - achisq_idx_list = np.append(achisq_idx_list_pt1, - achisq_idx_list_pt2) + achisq_idx_list = np.append(achisq_idx_list_pt1, achisq_idx_list_pt2) hauto_corr_vec = hautocorr[achisq_idx_list] - hauto_norm = hauto_corr_vec.real*hauto_corr_vec.real + hauto_norm = hauto_corr_vec.real * hauto_corr_vec.real # REMOVE THIS LINE TO REPRODUCE OLD RESULTS - hauto_norm += hauto_corr_vec.imag*hauto_corr_vec.imag + hauto_norm += hauto_corr_vec.imag * hauto_corr_vec.imag chisq_norm = 1.0 - hauto_norm - for ip,ind in enumerate(indices): + for ip, ind in enumerate(indices): curr_achisq_idx_list = achisq_idx_list + ind cphi = cphi_array[ip] sphi = sphi_array[ip] # By construction, the other "phase" of the SNR is 0 - snr_ind = sn[ind].real*cphi + sn[ind].imag*sphi + snr_ind = sn[ind].real * cphi + sn[ind].imag * sphi # Wrap index if needed (maybe should fail in this case?) if curr_achisq_idx_list[0] < 0: curr_achisq_idx_list[curr_achisq_idx_list < 0] += Nsnr if curr_achisq_idx_list[-1] > (Nsnr - 1): - curr_achisq_idx_list[curr_achisq_idx_list > (Nsnr-1)] -= Nsnr + curr_achisq_idx_list[curr_achisq_idx_list > (Nsnr - 1)] -= Nsnr - z = corr_sn[curr_achisq_idx_list].real*cphi + \ - corr_sn[curr_achisq_idx_list].imag*sphi - dz = z - hauto_corr_vec.real*snr_ind - curr_achisq_list = dz*dz/chisq_norm + z = ( + corr_sn[curr_achisq_idx_list].real * cphi + + corr_sn[curr_achisq_idx_list].imag * sphi + ) + dz = z - hauto_corr_vec.real * snr_ind + curr_achisq_list = dz * dz / chisq_norm if twophase: chisq_norm = 1.0 - hauto_norm - z = -corr_sn[curr_achisq_idx_list].real*sphi + \ - corr_sn[curr_achisq_idx_list].imag*cphi - dz = z - hauto_corr_vec.imag*snr_ind - curr_achisq_list += dz*dz/chisq_norm + z = ( + -corr_sn[curr_achisq_idx_list].real * sphi + + corr_sn[curr_achisq_idx_list].imag * cphi + ) + dz = z - hauto_corr_vec.imag * snr_ind + curr_achisq_list += dz * dz / chisq_norm if maxvalued: achisq[ip] = curr_achisq_list.max() @@ -136,18 +148,28 @@ def autochisq_from_precomputed(sn, corr_sn, hautocorr, indices, return dof, achisq -class SingleDetAutoChisq(object): - """Class that handles precomputation and memory management for efficiently + +class SingleDetAutoChisq: + """ + Class that handles precomputation and memory management for efficiently running the auto chisq in a single detector inspiral analysis. """ - def __init__(self, stride, num_points, onesided=None, twophase=False, - reverse_template=False, take_maximum_value=False, - maximal_value_dof=None): + + def __init__( + self, + stride, + num_points, + onesided=None, + twophase=False, + reverse_template=False, + take_maximum_value=False, + maximal_value_dof=None, + ): """ Initialize autochisq calculation instance Parameters - ----------- + ---------- stride : int Number of sample points between points at which auto-chisq is calculated. @@ -170,6 +192,7 @@ def __init__(self, stride, num_points, onesided=None, twophase=False, maximal_value_dof : int, required if using take_maximum_value If using take_maximum_value the expected value is not known. This value specifies what to store in the cont_chisq_dof output. + """ if stride > 0: self.do = True @@ -179,13 +202,13 @@ def __init__(self, stride, num_points, onesided=None, twophase=False, self.num_points = num_points self.stride = stride self.one_sided = onesided - if (onesided is not None): + if onesided is not None: self.dof = self.dof * 2 self.two_phase = twophase if self.two_phase: self.dof = self.dof * 2 self.reverse_template = reverse_template - self.take_maximum_value=take_maximum_value + self.take_maximum_value = take_maximum_value if self.take_maximum_value: if maximal_value_dof is None: err_msg = "Must provide the maximal_value_dof keyword " @@ -199,13 +222,22 @@ def __init__(self, stride, num_points, onesided=None, twophase=False, else: self.do = False - def values(self, sn, indices, template, psd, norm, stilde=None, - low_frequency_cutoff=None, high_frequency_cutoff=None): + def values( + self, + sn, + indices, + template, + psd, + norm, + stilde=None, + low_frequency_cutoff=None, + high_frequency_cutoff=None, + ): """ Calculate the auto-chisq at the specified indices. Parameters - ----------- + ---------- sn : Array[complex] SNR time series of the template for which auto-chisq is being computed. Provided unnormalized. @@ -232,6 +264,7 @@ def values(self, sn, indices, template, psd, norm, stilde=None, requested sample indices dof: int, approx number of statistical degrees of freedom + """ if not (self.do and len(indices) > 0): return None, None @@ -249,9 +282,9 @@ def values(self, sn, indices, template, psd, norm, stilde=None, htilde, psd=psd, low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff + high_frequency_cutoff=high_frequency_cutoff, ) - Pt = Pt * (1./ Pt[0]) + Pt = Pt * (1.0 / Pt[0]) self._autocor = Array(Pt, copy=True) else: Pt, _, P_norm = matched_filter_core( @@ -259,14 +292,14 @@ def values(self, sn, indices, template, psd, norm, stilde=None, htilde, psd=psd, low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff + high_frequency_cutoff=high_frequency_cutoff, ) # T-reversed template has same norm as forward template # so we can normalize using that # FIXME: Here sigmasq has to be cast to a float or the # code is really slow ... why?? - norm_fac = P_norm / float(((template.sigmasq(psd))**0.5)) + norm_fac = P_norm / float((template.sigmasq(psd)) ** 0.5) Pt *= norm_fac self._autocor = Array(Pt, copy=True) self._autocor_id = key @@ -274,13 +307,13 @@ def values(self, sn, indices, template, psd, norm, stilde=None, logging.debug("...Calculating autochisquare") sn = sn * norm if self.reverse_template: - assert(stilde is not None) + assert stilde is not None asn, _, ahnrm = matched_filter_core( htilde.conj(), stilde, low_frequency_cutoff=low_frequency_cutoff, high_frequency_cutoff=high_frequency_cutoff, - h_norm=template.sigmasq(psd) + h_norm=template.sigmasq(psd), ) correlation_snr = asn * ahnrm else: @@ -289,24 +322,25 @@ def values(self, sn, indices, template, psd, norm, stilde=None, achi_list = np.array([]) index_list = np.array(indices) dof, achi_list = autochisq_from_precomputed( - sn, correlation_snr, + sn, + correlation_snr, self._autocor, index_list, stride=self.stride, num_points=self.num_points, oneside=self.one_sided, twophase=self.two_phase, - maxvalued=self.take_maximum_value + maxvalued=self.take_maximum_value, ) self.dof = dof return achi_list, dof class SingleDetSkyMaxAutoChisq(SingleDetAutoChisq): - """Stub for precessing auto chisq if anyone ever wants to code it up. - """ + """Stub for precessing auto chisq if anyone ever wants to code it up.""" + def __init__(self, *args, **kwds): - super(SingleDetSkyMaxAutoChisq, self).__init__(*args, **kwds) + super().__init__(*args, **kwds) def values(self, *args, **kwargs): if self.do: @@ -314,5 +348,3 @@ def values(self, *args, **kwargs): err_msg += "been written. If you want to use it, why not help " err_msg += "write it?" raise NotImplementedError(err_msg) - else: - return None diff --git a/pycbc/vetoes/bank_chisq.py b/pycbc/vetoes/bank_chisq.py index 80b4f3cca82..07ab1e10a7e 100644 --- a/pycbc/vetoes/bank_chisq.py +++ b/pycbc/vetoes/bank_chisq.py @@ -21,14 +21,19 @@ # # ============================================================================= # -import logging, numpy -from pycbc.types import Array, zeros, real_same_precision_as, TimeSeries -from pycbc.filter import overlap_cplx, matched_filter_core -from pycbc.waveform import FilterBank +import logging from math import sqrt +import numpy + +from pycbc.filter import matched_filter_core, overlap_cplx +from pycbc.types import Array, TimeSeries, real_same_precision_as, zeros +from pycbc.waveform import FilterBank + + def segment_snrs(filters, stilde, psd, low_frequency_cutoff): - """ This functions calculates the snr of each bank veto template against + """ + This functions calculates the snr of each bank veto template against the segment Parameters @@ -44,6 +49,7 @@ def segment_snrs(filters, stilde, psd, low_frequency_cutoff): ------- snr (list): List of snr time series. norm (list): List of normalizations factors for the snr time series. + """ snrs = [] norms = [] @@ -51,8 +57,12 @@ def segment_snrs(filters, stilde, psd, low_frequency_cutoff): for bank_template in filters: # For every template compute the snr against the stilde segment snr, _, norm = matched_filter_core( - bank_template, stilde, h_norm=bank_template.sigmasq(psd), - psd=None, low_frequency_cutoff=low_frequency_cutoff) + bank_template, + stilde, + h_norm=bank_template.sigmasq(psd), + psd=None, + low_frequency_cutoff=low_frequency_cutoff, + ) # SNR time series stored here snrs.append(snr) # Template normalization factor stored here @@ -60,8 +70,10 @@ def segment_snrs(filters, stilde, psd, low_frequency_cutoff): return snrs, norms + def template_overlaps(bank_filters, template, psd, low_frequency_cutoff): - """ This functions calculates the overlaps between the template and the + """ + This functions calculates the overlaps between the template and the bank veto templates. Parameters @@ -74,31 +86,43 @@ def template_overlaps(bank_filters, template, psd, low_frequency_cutoff): Returns ------- overlaps: List of complex overlap values. + """ overlaps = [] template_ow = template / psd for bank_template in bank_filters: - overlap = overlap_cplx(template_ow, bank_template, - low_frequency_cutoff=low_frequency_cutoff, normalized=False) + overlap = overlap_cplx( + template_ow, + bank_template, + low_frequency_cutoff=low_frequency_cutoff, + normalized=False, + ) norm = sqrt(1 / template.sigmasq(psd) / bank_template.sigmasq(psd)) overlaps.append(overlap * norm) - if (abs(overlaps[-1]) > 0.99): + if abs(overlaps[-1]) > 0.99: errMsg = "Overlap > 0.99 between bank template and filter. " errMsg += "This bank template will not be used to calculate " errMsg += "bank chisq for this filter template. The expected " errMsg += "value will be added to the chisq to account for " errMsg += "the removal of this template.\n" - errMsg += "Masses of filter template: %e %e\n" \ - %(template.params.mass1, template.params.mass2) - errMsg += "Masses of bank filter template: %e %e\n" \ - %(bank_template.params.mass1, bank_template.params.mass2) - errMsg += "Overlap: %e" %(abs(overlaps[-1])) + errMsg += "Masses of filter template: %e %e\n" % ( + template.params.mass1, + template.params.mass2, + ) + errMsg += "Masses of bank filter template: %e %e\n" % ( + bank_template.params.mass1, + bank_template.params.mass2, + ) + errMsg += "Overlap: %e" % (abs(overlaps[-1])) logging.info(errMsg) return overlaps -def bank_chisq_from_filters(tmplt_snr, tmplt_norm, bank_snrs, bank_norms, - tmplt_bank_matches, indices=None): - """ This function calculates and returns a TimeSeries object containing the + +def bank_chisq_from_filters( + tmplt_snr, tmplt_norm, bank_snrs, bank_norms, tmplt_bank_matches, indices=None +): + """ + This function calculates and returns a TimeSeries object containing the bank veto calculated over a segment. Parameters @@ -124,13 +148,14 @@ def bank_chisq_from_filters(tmplt_snr, tmplt_norm, bank_snrs, bank_norms, Returns ------- bank_chisq: TimeSeries of the bank vetos + """ if indices is not None: tmplt_snr = Array(tmplt_snr, copy=False) bank_snrs_tmp = [] for bank_snr in bank_snrs: bank_snrs_tmp.append(bank_snr.take(indices)) - bank_snrs=bank_snrs_tmp + bank_snrs = bank_snrs_tmp # Initialise bank_chisq as 0s everywhere bank_chisq = zeros(len(tmplt_snr), dtype=real_same_precision_as(tmplt_snr)) @@ -138,15 +163,15 @@ def bank_chisq_from_filters(tmplt_snr, tmplt_norm, bank_snrs, bank_norms, # Loop over all the bank templates for i in range(len(bank_snrs)): bank_match = tmplt_bank_matches[i] - if (abs(bank_match) > 0.99): + if abs(bank_match) > 0.99: # Not much point calculating bank_chisquared if the bank template # is very close to the filter template. Can also hit numerical # error due to approximations made in this calculation. # The value of 2 is the expected addition to the chisq for this # template - bank_chisq += 2. + bank_chisq += 2.0 continue - bank_norm = sqrt((1 - bank_match*bank_match.conj()).real) + bank_norm = sqrt((1 - bank_match * bank_match.conj()).real) bank_SNR = bank_snrs[i] * (bank_norms[i] / bank_norm) tmplt_SNR = tmplt_snr * (bank_match.conj() * tmplt_norm / bank_norm) @@ -158,16 +183,24 @@ def bank_chisq_from_filters(tmplt_snr, tmplt_norm, bank_snrs, bank_norms, if indices is not None: return bank_chisq - else: - return TimeSeries(bank_chisq, delta_t=tmplt_snr.delta_t, - epoch=tmplt_snr.start_time, copy=False) - -class SingleDetBankVeto(object): - """This class reads in a template bank file for a bank veto, handles the - memory management of its filters internally, and calculates the bank - veto TimeSeries. + return TimeSeries( + bank_chisq, + delta_t=tmplt_snr.delta_t, + epoch=tmplt_snr.start_time, + copy=False, + ) + + +class SingleDetBankVeto: """ - def __init__(self, bank_file, flen, delta_f, f_low, cdtype, approximant=None, **kwds): + This class reads in a template bank file for a bank veto, handles the + memory management of its filters internally, and calculates the bank + veto TimeSeries. + """ + + def __init__( + self, bank_file, flen, delta_f, f_low, cdtype, approximant=None, **kwds + ): if bank_file is not None: self.do = True @@ -178,14 +211,18 @@ def __init__(self, bank_file, flen, delta_f, f_low, cdtype, approximant=None, ** self.delta_f = delta_f self.f_low = f_low self.seg_len_freq = flen - self.seg_len_time = (self.seg_len_freq-1)*2 + self.seg_len_time = (self.seg_len_freq - 1) * 2 logging.info("Read in bank veto template bank") - bank_veto_bank = FilterBank(bank_file, - self.seg_len_freq, - self.delta_f, self.cdtype, - low_frequency_cutoff=f_low, - approximant=approximant, **kwds) + bank_veto_bank = FilterBank( + bank_file, + self.seg_len_freq, + self.delta_f, + self.cdtype, + low_frequency_cutoff=f_low, + approximant=approximant, + **kwds, + ) self.filters = list(bank_veto_bank) self.dof = len(bank_veto_bank) * 2 @@ -220,22 +257,25 @@ def values(self, template, psd, stilde, snrv, norm, indices): requested sample indices bank_chisq_dof: int, approx number of statistical degrees of freedom + """ if not self.do: return None, None logging.debug("...Doing bank veto") overlaps = self.cache_overlaps(template, psd) bank_veto_snrs, bank_veto_norms = self.cache_segment_snrs(stilde, psd) - chisq = bank_chisq_from_filters(snrv, norm, bank_veto_snrs, - bank_veto_norms, overlaps, indices) + chisq = bank_chisq_from_filters( + snrv, norm, bank_veto_snrs, bank_veto_norms, overlaps, indices + ) dof = numpy.repeat(self.dof, len(chisq)) return chisq, dof + class SingleDetSkyMaxBankVeto(SingleDetBankVeto): - """Stub for precessing bank veto if anyone ever wants to code it up. - """ + """Stub for precessing bank veto if anyone ever wants to code it up.""" + def __init__(self, *args, **kwds): - super(SingleDetSkyMaxBankVeto, self).__init__(*args, **kwds) + super().__init__(*args, **kwds) def values(self, *args, **kwargs): if self.do: @@ -243,5 +283,4 @@ def values(self, *args, **kwargs): err_msg += "been written. If you want to use it, why not help " err_msg += "write it?" raise NotImplementedError(err_msg) - else: - return None, None + return None, None diff --git a/pycbc/vetoes/chisq.py b/pycbc/vetoes/chisq.py index abce8c74647..3de25ee7254 100644 --- a/pycbc/vetoes/chisq.py +++ b/pycbc/vetoes/chisq.py @@ -21,21 +21,36 @@ # # ============================================================================= # -import numpy, logging, math, pycbc.fft +import logging +import math -from pycbc.types import zeros, real_same_precision_as, TimeSeries, complex_same_precision_as -from pycbc.filter import sigmasq_series, make_frequency_series, matched_filter_core, get_cutoff_indices -from pycbc.scheme import schemed +import numpy + +import pycbc.fft import pycbc.pnutils +from pycbc.filter import ( + get_cutoff_indices, + make_frequency_series, + matched_filter_core, + sigmasq_series, +) +from pycbc.scheme import schemed +from pycbc.types import ( + TimeSeries, + complex_same_precision_as, + real_same_precision_as, + zeros, +) + +BACKEND_PREFIX = "pycbc.vetoes.chisq_" -BACKEND_PREFIX="pycbc.vetoes.chisq_" def power_chisq_bins_from_sigmasq_series(sigmasq_series, num_bins, kmin, kmax): - """Returns bins of equal power for use with the chisq functions + """ + Returns bins of equal power for use with the chisq functions Parameters ---------- - sigmasq_series: FrequencySeries A frequency series containing the cumulative power of a filter template preweighted by a psd. @@ -48,24 +63,25 @@ def power_chisq_bins_from_sigmasq_series(sigmasq_series, num_bins, kmin, kmax): Returns ------- - bins: List of ints A list of the edges of the chisq bins is returned. + """ sigmasq = sigmasq_series[kmax - 1] edge_vec = numpy.arange(0, num_bins) * sigmasq / num_bins - bins = numpy.searchsorted(sigmasq_series[kmin:kmax], edge_vec, side='right') + bins = numpy.searchsorted(sigmasq_series[kmin:kmax], edge_vec, side="right") bins += kmin return numpy.append(bins, kmax) -def power_chisq_bins(htilde, num_bins, psd, low_frequency_cutoff=None, - high_frequency_cutoff=None): - """Returns bins of equal power for use with the chisq functions +def power_chisq_bins( + htilde, num_bins, psd, low_frequency_cutoff=None, high_frequency_cutoff=None +): + """ + Returns bins of equal power for use with the chisq functions Parameters ---------- - htilde: FrequencySeries A frequency series containing the template waveform num_bins: int @@ -80,16 +96,19 @@ def power_chisq_bins(htilde, num_bins, psd, low_frequency_cutoff=None, Returns ------- - bins: List of ints A list of the edges of the chisq bins is returned. + """ - sigma_vec = sigmasq_series(htilde, psd, low_frequency_cutoff, - high_frequency_cutoff).numpy() - kmin, kmax = get_cutoff_indices(low_frequency_cutoff, - high_frequency_cutoff, - htilde.delta_f, - (len(htilde)-1)*2) + sigma_vec = sigmasq_series( + htilde, psd, low_frequency_cutoff, high_frequency_cutoff + ).numpy() + kmin, kmax = get_cutoff_indices( + low_frequency_cutoff, + high_frequency_cutoff, + htilde.delta_f, + (len(htilde) - 1) * 2, + ) return power_chisq_bins_from_sigmasq_series(sigma_vec, num_bins, kmin, kmax) @@ -99,17 +118,18 @@ def chisq_accum_bin(chisq, q): err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) + @schemed(BACKEND_PREFIX) def shift_sum(v1, shifts, bins): - """ Calculate the time shifted sum of the FrequencySeries - """ + """Calculate the time shifted sum of the FrequencySeries""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) def power_chisq_at_points_from_precomputed(corr, snr, snr_norm, bins, indices): - """Calculate the chisq timeseries from precomputed values for only select points. + """ + Calculate the chisq timeseries from precomputed values for only select points. This function calculates the chisq at each point by explicitly time shifting and summing each bin. No FFT is involved. @@ -132,23 +152,29 @@ def power_chisq_at_points_from_precomputed(corr, snr, snr_norm, bins, indices): ------- chisq: Array An array containing only the chisq at the selected points. + """ num_bins = len(bins) - 1 - chisq = shift_sum(corr, indices, bins) # pylint:disable=assignment-from-no-return - return (chisq * num_bins - (snr.conj() * snr).real) * (snr_norm ** 2.0) + chisq = shift_sum(corr, indices, bins) # pylint:disable=assignment-from-no-return + return (chisq * num_bins - (snr.conj() * snr).real) * (snr_norm**2.0) + _q_l = None _qtilde_l = None _chisq_l = None -def power_chisq_from_precomputed(corr, snr, snr_norm, bins, indices=None, return_bins=False): - """Calculate the chisq timeseries from precomputed values. + + +def power_chisq_from_precomputed( + corr, snr, snr_norm, bins, indices=None, return_bins=False +): + """ + Calculate the chisq timeseries from precomputed values. This function calculates the chisq at all times by performing an inverse FFT of each bin. Parameters ---------- - corr: FrequencySeries The produce of the template and data in the frequency domain. snr: TimeSeries @@ -166,6 +192,7 @@ def power_chisq_from_precomputed(corr, snr, snr_norm, bins, indices=None, return Returns ------- chisq: TimeSeries + """ # Get workspace memory global _q_l, _qtilde_l, _chisq_l @@ -188,42 +215,46 @@ def power_chisq_from_precomputed(corr, snr, snr_norm, bins, indices=None, return chisq = zeros(len(snr), dtype=real_same_precision_as(snr)) _chisq_l = chisq else: - chisq = _chisq_l[0:len(snr)] + chisq = _chisq_l[0 : len(snr)] chisq.clear() num_bins = len(bins) - 1 for j in range(num_bins): k_min = int(bins[j]) - k_max = int(bins[j+1]) + k_max = int(bins[j + 1]) qtilde[k_min:k_max] = corr[k_min:k_max] pycbc.fft.ifft(qtilde, q) qtilde[k_min:k_max].clear() if return_bins: - bin_snrs.append(TimeSeries(q * snr_norm * num_bins ** 0.5, - delta_t=snr.delta_t, - epoch=snr.start_time)) + bin_snrs.append( + TimeSeries( + q * snr_norm * num_bins**0.5, + delta_t=snr.delta_t, + epoch=snr.start_time, + ) + ) if indices is not None: chisq_accum_bin(chisq, q.take(indices)) else: chisq_accum_bin(chisq, q) - chisq = (chisq * num_bins - snr.squared_norm()) * (snr_norm ** 2.0) + chisq = (chisq * num_bins - snr.squared_norm()) * (snr_norm**2.0) if indices is None: chisq = TimeSeries(chisq, delta_t=snr.delta_t, epoch=snr.start_time, copy=False) if return_bins: return chisq, bin_snrs - else: - return chisq + return chisq def fastest_power_chisq_at_points(corr, snr, snrv, snr_norm, bins, indices): - """Calculate the chisq values for only selected points. + """ + Calculate the chisq values for only selected points. This function looks at the number of points to be evaluated and selects the fastest method (FFT, or direct time shift and sum). In either case, @@ -247,23 +278,30 @@ def fastest_power_chisq_at_points(corr, snr, snrv, snr_norm, bins, indices): ------- chisq: Array An array containing only the chisq at the selected points. + """ import pycbc.scheme + if isinstance(pycbc.scheme.mgr.state, pycbc.scheme.CPUScheme): # We don't have that many points so do the direct time shift. - return power_chisq_at_points_from_precomputed(corr, snrv, - snr_norm, bins, indices) - else: - # We have a lot of points so it is faster to use the fourier transform - return power_chisq_from_precomputed(corr, snr, snr_norm, bins, - indices=indices) - - -def power_chisq(template, data, num_bins, psd, - low_frequency_cutoff=None, - high_frequency_cutoff=None, - return_bins=False): - """Calculate the chisq timeseries + return power_chisq_at_points_from_precomputed( + corr, snrv, snr_norm, bins, indices + ) + # We have a lot of points so it is faster to use the fourier transform + return power_chisq_from_precomputed(corr, snr, snr_norm, bins, indices=indices) + + +def power_chisq( + template, + data, + num_bins, + psd, + low_frequency_cutoff=None, + high_frequency_cutoff=None, + return_bins=False, +): + """ + Calculate the chisq timeseries Parameters ---------- @@ -289,24 +327,30 @@ def power_chisq(template, data, num_bins, psd, ------- chisq: TimeSeries TimeSeries containing the chisq values for all times. + """ htilde = make_frequency_series(template) stilde = make_frequency_series(data) - bins = power_chisq_bins(htilde, num_bins, psd, low_frequency_cutoff, - high_frequency_cutoff) + bins = power_chisq_bins( + htilde, num_bins, psd, low_frequency_cutoff, high_frequency_cutoff + ) corra = zeros((len(htilde) - 1) * 2, dtype=htilde.dtype) - total_snr, corr, tnorm = matched_filter_core(htilde, stilde, psd, - low_frequency_cutoff, high_frequency_cutoff, - corr_out=corra) + total_snr, corr, tnorm = matched_filter_core( + htilde, stilde, psd, low_frequency_cutoff, high_frequency_cutoff, corr_out=corra + ) - return power_chisq_from_precomputed(corr, total_snr, tnorm, bins, return_bins=return_bins) + return power_chisq_from_precomputed( + corr, total_snr, tnorm, bins, return_bins=return_bins + ) -class SingleDetPowerChisq(object): - """Class that handles precomputation and memory management for efficiently +class SingleDetPowerChisq: + """ + Class that handles precomputation and memory management for efficiently running the power chisq in a single detector inspiral analysis. """ + def __init__(self, num_bins=0, snr_threshold=None): if not (num_bins == "0" or num_bins == 0): self.do = True @@ -319,35 +363,34 @@ def __init__(self, num_bins=0, snr_threshold=None): @staticmethod def parse_option(row, arg): - safe_dict = {'max': max, 'min': min} + safe_dict = {"max": max, "min": min} safe_dict.update(row.__dict__) safe_dict.update(math.__dict__) safe_dict.update(pycbc.pnutils.__dict__) - return eval(arg, {"__builtins__":None}, safe_dict) + return eval(arg, {"__builtins__": None}, safe_dict) def cached_chisq_bins(self, template, psd): from pycbc.opt import LimitedSizeDict key = id(psd) - if not hasattr(psd, '_chisq_cached_key'): + if not hasattr(psd, "_chisq_cached_key"): psd._chisq_cached_key = {} - if not hasattr(template, '_bin_cache'): + if not hasattr(template, "_bin_cache"): template._bin_cache = LimitedSizeDict(size_limit=2**2) - if key not in template._bin_cache or id(template.params) not in psd._chisq_cached_key: + if ( + key not in template._bin_cache + or id(template.params) not in psd._chisq_cached_key + ): psd._chisq_cached_key[id(template.params)] = True num_bins = int(self.parse_option(template, self.num_bins)) - if hasattr(psd, 'sigmasq_vec') and \ - template.approximant in psd.sigmasq_vec: + if hasattr(psd, "sigmasq_vec") and template.approximant in psd.sigmasq_vec: kmin = int(template.f_lower / psd.delta_f) kmax = template.end_idx bins = power_chisq_bins_from_sigmasq_series( - psd.sigmasq_vec[template.approximant], - num_bins, - kmin, - kmax + psd.sigmasq_vec[template.approximant], num_bins, kmin, kmax ) else: bins = power_chisq_bins(template, num_bins, psd, template.f_lower) @@ -356,7 +399,8 @@ def cached_chisq_bins(self, template, psd): return template._bin_cache[key] def values(self, corr, snrv, snr_norm, psd, indices, template): - """ Calculate the chisq at points given by indices. + """ + Calculate the chisq at points given by indices. Returns ------- @@ -367,6 +411,7 @@ def values(self, corr, snrv, snr_norm, psd, indices, template): chisq_dof: Array Number of statistical degrees of freedom for the chisq test in the given template, equal to 2 * num_bins - 2 + """ if self.do: num_above = len(indices) @@ -374,7 +419,7 @@ def values(self, corr, snrv, snr_norm, psd, indices, template): if self.snr_threshold: above = abs(snrv * snr_norm) > self.snr_threshold num_above = above.sum() - logging.info('%s above chisq activation threshold' % num_above) + logging.info("%s above chisq activation threshold" % num_above) above_indices = indices[above] above_snrv = snrv[above] chisq_out = numpy.zeros(len(indices), dtype=numpy.float32) @@ -386,51 +431,65 @@ def values(self, corr, snrv, snr_norm, psd, indices, template): bins = self.cached_chisq_bins(template, psd) # len(bins) is number of bin edges, num_bins = len(bins) - 1 dof = (len(bins) - 1) * 2 - 2 - _chisq = power_chisq_at_points_from_precomputed(corr, - above_snrv, snr_norm, bins, above_indices) + _chisq = power_chisq_at_points_from_precomputed( + corr, above_snrv, snr_norm, bins, above_indices + ) if self.snr_threshold: if num_above > 0: chisq_out[above] = _chisq + elif num_above == 0: + chisq_out = numpy.zeros(0, dtype=numpy.float32) else: - if num_above == 0: - chisq_out = numpy.zeros(0, dtype=numpy.float32) - else: - chisq_out = _chisq + chisq_out = _chisq - return chisq_out, numpy.repeat(dof, len(indices))# dof * numpy.ones_like(indices) - else: - return None, None + return chisq_out, numpy.repeat( + dof, len(indices) + ) # dof * numpy.ones_like(indices) + return None, None class SingleDetSkyMaxPowerChisq(SingleDetPowerChisq): - """Class that handles precomputation and memory management for efficiently + """ + Class that handles precomputation and memory management for efficiently running the power chisq in a single detector inspiral analysis when maximizing analytically over sky location. """ + def __init__(self, **kwds): - super(SingleDetSkyMaxPowerChisq, self).__init__(**kwds) + super().__init__(**kwds) self.template_mem = None self.corr_mem = None def calculate_chisq_bins(self, template, psd): - """ Obtain the chisq bins for this template and PSD. - """ + """Obtain the chisq bins for this template and PSD.""" num_bins = int(self.parse_option(template, self.num_bins)) - if hasattr(psd, 'sigmasq_vec') and \ - template.approximant in psd.sigmasq_vec: + if hasattr(psd, "sigmasq_vec") and template.approximant in psd.sigmasq_vec: kmin = int(template.f_lower / psd.delta_f) kmax = template.end_idx bins = power_chisq_bins_from_sigmasq_series( - psd.sigmasq_vec[template.approximant], num_bins, kmin, kmax) + psd.sigmasq_vec[template.approximant], num_bins, kmin, kmax + ) else: bins = power_chisq_bins(template, num_bins, psd, template.f_lower) return bins - def values(self, corr_plus, corr_cross, snrv, psd, - indices, template_plus, template_cross, u_vals, - hplus_cross_corr, hpnorm, hcnorm): - """ Calculate the chisq at points given by indices. + def values( + self, + corr_plus, + corr_cross, + snrv, + psd, + indices, + template_plus, + template_cross, + u_vals, + hplus_cross_corr, + hpnorm, + hcnorm, + ): + """ + Calculate the chisq at points given by indices. Returns ------- @@ -440,13 +499,14 @@ def values(self, corr_plus, corr_cross, snrv, psd, chisq_dof: Array Number of statistical degrees of freedom for the chisq test in the given template + """ if self.do: num_above = len(indices) if self.snr_threshold: above = abs(snrv) > self.snr_threshold num_above = above.sum() - logging.info('%s above chisq activation threshold' % num_above) + logging.info("%s above chisq activation threshold" % num_above) above_indices = indices[above] above_snrv = snrv[above] u_vals = u_vals[above] @@ -458,16 +518,18 @@ def values(self, corr_plus, corr_cross, snrv, psd, if num_above > 0: chisq = [] - curr_tmplt_mult_fac = 0. - curr_corr_mult_fac = 0. - if self.template_mem is None or \ - (not len(self.template_mem) == len(template_plus)): - self.template_mem = zeros(len(template_plus), - dtype=complex_same_precision_as(corr_plus)) - if self.corr_mem is None or \ - (not len(self.corr_mem) == len(corr_plus)): - self.corr_mem = zeros(len(corr_plus), - dtype=complex_same_precision_as(corr_plus)) + curr_tmplt_mult_fac = 0.0 + curr_corr_mult_fac = 0.0 + if self.template_mem is None or ( + not len(self.template_mem) == len(template_plus) + ): + self.template_mem = zeros( + len(template_plus), dtype=complex_same_precision_as(corr_plus) + ) + if self.corr_mem is None or (not len(self.corr_mem) == len(corr_plus)): + self.corr_mem = zeros( + len(corr_plus), dtype=complex_same_precision_as(corr_plus) + ) tmplt_data = template_cross.data corr_data = corr_cross.data @@ -483,26 +545,32 @@ def values(self, corr_plus, corr_cross, snrv, psd, # Construct template from _plus and _cross # Note that this modifies in place, so we store that and # revert on the next pass. - template = template_cross.multiply_and_add(template_plus, - local_u_val-curr_tmplt_mult_fac) + template = template_cross.multiply_and_add( + template_plus, local_u_val - curr_tmplt_mult_fac + ) curr_tmplt_mult_fac = local_u_val template.f_lower = template_plus.f_lower template.params = template_plus.params # Construct the corr vector - norm_fac = local_u_val*local_u_val + 1 + norm_fac = local_u_val * local_u_val + 1 norm_fac += 2 * local_u_val * hplus_cross_corr norm_fac = hcnorm / (norm_fac**0.5) hp_fac = local_u_val * hpnorm / hcnorm - corr = corr_cross.multiply_and_add(corr_plus, - hp_fac - curr_corr_mult_fac) + corr = corr_cross.multiply_and_add( + corr_plus, hp_fac - curr_corr_mult_fac + ) curr_corr_mult_fac = hp_fac bins = self.calculate_chisq_bins(template, psd) dof = (len(bins) - 1) * 2 - 2 - curr_chisq = power_chisq_at_points_from_precomputed(corr, - above_local_snr/ norm_fac, norm_fac, - bins, above_local_indices) + curr_chisq = power_chisq_at_points_from_precomputed( + corr, + above_local_snr / norm_fac, + norm_fac, + bins, + above_local_indices, + ) chisq.append(curr_chisq[0]) chisq = numpy.array(chisq) # Must reset corr and template to original values! @@ -515,6 +583,7 @@ def values(self, corr_plus, corr_cross, snrv, psd, else: rchisq = chisq - return rchisq, numpy.repeat(dof, len(indices))# dof * numpy.ones_like(indices) - else: - return None, None + return rchisq, numpy.repeat( + dof, len(indices) + ) # dof * numpy.ones_like(indices) + return None, None diff --git a/pycbc/vetoes/chisq_cuda.py b/pycbc/vetoes/chisq_cuda.py index 5b7dd094b64..edd50179bec 100644 --- a/pycbc/vetoes/chisq_cuda.py +++ b/pycbc/vetoes/chisq_cuda.py @@ -21,22 +21,27 @@ # # ============================================================================= # -import pycuda.driver, numpy -from pycuda.elementwise import ElementwiseKernel -from pycuda.tools import context_dependent_memoize, dtype_to_ctype +import numpy +import pycuda.driver import pycuda.gpuarray from mako.template import Template from pycuda.compiler import SourceModule +from pycuda.elementwise import ElementwiseKernel +from pycuda.tools import context_dependent_memoize, dtype_to_ctype + @context_dependent_memoize def get_accum_diff_sq_kernel(dtype_x, dtype_z): return ElementwiseKernel( - "%(tp_a)s *x, %(tp_c)s *z" % { - "tp_a": dtype_to_ctype(dtype_x), - "tp_c": dtype_to_ctype(dtype_z), - }, - "x[i] += norm(z[i]) ", - "chisq_accum") + "%(tp_a)s *x, %(tp_c)s *z" + % { + "tp_a": dtype_to_ctype(dtype_x), + "tp_c": dtype_to_ctype(dtype_z), + }, + "x[i] += norm(z[i]) ", + "chisq_accum", + ) + def chisq_accum_bin(chisq, q): krnl = get_accum_diff_sq_kernel(chisq.dtype, q.dtype) @@ -216,6 +221,8 @@ def chisq_accum_bin(chisq, q): """) _pchisq_cache = {} + + def get_pchisq_fn(np, fuse_correlate=False): if np not in _pchisq_cache: nt = 256 @@ -228,12 +235,14 @@ def get_pchisq_fn(np, fuse_correlate=False): _pchisq_cache[np] = (fn, nt) return _pchisq_cache[np] + _pchisq_cache_pow2 = {} + + def get_pchisq_fn_pow2(np, fuse_correlate=False): if np not in _pchisq_cache_pow2: nt = 256 - mod = SourceModule(chisqkernel_pow2.render(NT=nt, NP=np, - fuse=fuse_correlate)) + mod = SourceModule(chisqkernel_pow2.render(NT=nt, NP=np, fuse=fuse_correlate)) fn = mod.get_function("power_chisq_at_points_%s_pow2" % (np)) if fuse_correlate: fn.prepare("PPPI" + "I" * np + "PPPI") @@ -242,48 +251,55 @@ def get_pchisq_fn_pow2(np, fuse_correlate=False): _pchisq_cache_pow2[np] = (fn, nt) return _pchisq_cache_pow2[np] + def get_cached_bin_layout(bins): bv, kmin, kmax = [], [], [] - for i in range(len(bins)-1): - s, e = bins[i], bins[i+1] + for i in range(len(bins) - 1): + s, e = bins[i], bins[i + 1] BS = 4096 if (e - s) < BS: bv.append(i) kmin.append(s) kmax.append(e) else: - k = list(numpy.arange(s, e, BS/2)) + k = list(numpy.arange(s, e, BS / 2)) kmin += k kmax += k[1:] + [e] - bv += [i]*len(k) + bv += [i] * len(k) bv = pycuda.gpuarray.to_gpu_async(numpy.array(bv, dtype=numpy.uint32)) kmin = pycuda.gpuarray.to_gpu_async(numpy.array(kmin, dtype=numpy.uint32)) kmax = pycuda.gpuarray.to_gpu_async(numpy.array(kmax, dtype=numpy.uint32)) return kmin, kmax, bv + def shift_sum_points(num, N, arg_tuple): - #fuse = 'fuse' in corr.gpu_callback_method + # fuse = 'fuse' in corr.gpu_callback_method fuse = False - fn, nt = get_pchisq_fn(num, fuse_correlate = fuse) + fn, nt = get_pchisq_fn(num, fuse_correlate=fuse) corr, outp, phase, np, nb, N, kmin, kmax, bv, nbins = arg_tuple args = [(nb, 1), (nt, 1, 1)] if fuse: args += [corr.htilde.data.gpudata, corr.stilde.data.gpudata] else: args += [corr.data.gpudata] - args +=[outp.gpudata, N] + phase[0:num] + [kmin.gpudata, kmax.gpudata, bv.gpudata, nbins] + args += ( + [outp.gpudata, N] + + phase[0:num] + + [kmin.gpudata, kmax.gpudata, bv.gpudata, nbins] + ) fn.prepared_call(*args) - outp = outp[num*nbins:] + outp = outp[num * nbins :] phase = phase[num:] np -= num return outp, phase, np + def shift_sum_points_pow2(num, arg_tuple): - #fuse = 'fuse' in corr.gpu_callback_method + # fuse = 'fuse' in corr.gpu_callback_method fuse = False - fn, nt = get_pchisq_fn_pow2(num, fuse_correlate = fuse) + fn, nt = get_pchisq_fn_pow2(num, fuse_correlate=fuse) corr, outp, points, np, nb, N, kmin, kmax, bv, nbins = arg_tuple args = [(nb, 1), (nt, 1, 1)] @@ -291,20 +307,27 @@ def shift_sum_points_pow2(num, arg_tuple): args += [corr.htilde.data.gpudata, corr.stilde.data.gpudata] else: args += [corr.data.gpudata] - args += [outp.gpudata, N] + points[0:num] + [kmin.gpudata, - kmax.gpudata, bv.gpudata, nbins] + args += ( + [outp.gpudata, N] + + points[0:num] + + [kmin.gpudata, kmax.gpudata, bv.gpudata, nbins] + ) fn.prepared_call(*args) - outp = outp[num*nbins:] + outp = outp[num * nbins :] points = points[num:] np -= num return outp, points, np + _pow2_cache = {} + + def get_cached_pow2(N): if N not in _pow2_cache: - _pow2_cache[N] = not(N & (N-1)) + _pow2_cache[N] = not (N & (N - 1)) return _pow2_cache[N] + def shift_sum(corr, points, bins): kmin, kmax, bv = get_cached_bin_layout(bins) nb = len(kmin) @@ -334,16 +357,13 @@ def shift_sum(corr, points, bins): cargs = (corr, outp, phase, np, nb, N, kmin, kmax, bv, nbins) if np >= 4: - outp, phase, np = shift_sum_points(4, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(4, cargs) # pylint:disable=no-value-for-parameter elif np >= 3: - outp, phase, np = shift_sum_points(3, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(3, cargs) # pylint:disable=no-value-for-parameter elif np >= 2: - outp, phase, np = shift_sum_points(2, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(2, cargs) # pylint:disable=no-value-for-parameter elif np == 1: - outp, phase, np = shift_sum_points(1, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(1, cargs) # pylint:disable=no-value-for-parameter o = outc.get() return (o.conj() * o).sum(axis=1).real - - - diff --git a/pycbc/vetoes/chisq_cupy.py b/pycbc/vetoes/chisq_cupy.py index 83dc2ba991d..0539610c326 100644 --- a/pycbc/vetoes/chisq_cupy.py +++ b/pycbc/vetoes/chisq_cupy.py @@ -23,23 +23,22 @@ # import functools -import numpy + import cupy as cp +import numpy from mako.template import Template from pycbc.constants import TWOPI LALARGS = { - 'TWOPI': TWOPI, + "TWOPI": TWOPI, } accum_diff_sq_kernel = cp.ElementwiseKernel( - "X input", - "raw Y output", - "output[i] += norm(input)", - "accum_diff_sq_kernel" + "X input", "raw Y output", "output[i] += norm(input)", "accum_diff_sq_kernel" ) + def chisq_accum_bin(chisq, q): accum_diff_sq_kernel(q.data, chisq.data) @@ -215,51 +214,54 @@ def chisq_accum_bin(chisq, q): } """) -@functools.lru_cache(maxsize=None) + +@functools.cache def get_pchisq_fn(np, fuse_correlate=False): nt = 256 fn = cp.RawKernel( chisqkernel.render(NT=nt, NP=np, fuse=fuse_correlate, **LALARGS), - f'power_chisq_at_points_{np}', - backend='nvcc' + f"power_chisq_at_points_{np}", + backend="nvcc", ) return fn, nt -@functools.lru_cache(maxsize=None) +@functools.cache def get_pchisq_fn_pow2(np, fuse_correlate=False): nt = 256 fn = cp.RawKernel( chisqkernel_pow2.render(NT=nt, NP=np, fuse=fuse_correlate, **LALARGS), - f'power_chisq_at_points_{np}_pow2', - backend='nvcc' + f"power_chisq_at_points_{np}_pow2", + backend="nvcc", ) return fn, nt + def get_cached_bin_layout(bins): bv, kmin, kmax = [], [], [] - for i in range(len(bins)-1): - s, e = bins[i], bins[i+1] + for i in range(len(bins) - 1): + s, e = bins[i], bins[i + 1] BS = 4096 if (e - s) < BS: bv.append(i) kmin.append(s) kmax.append(e) else: - k = list(numpy.arange(s, e, BS/2)) + k = list(numpy.arange(s, e, BS / 2)) kmin += k kmax += k[1:] + [e] - bv += [i]*len(k) + bv += [i] * len(k) bv = cp.array(bv, dtype=cp.uint32) kmin = cp.array(kmin, dtype=cp.uint32) kmax = cp.array(kmax, dtype=cp.uint32) return kmin, kmax, bv + def shift_sum_points(num, N, arg_tuple): - #fuse = 'fuse' in corr.gpu_callback_method + # fuse = 'fuse' in corr.gpu_callback_method fuse = False - fn, nt = get_pchisq_fn(num, fuse_correlate = fuse) + fn, nt = get_pchisq_fn(num, fuse_correlate=fuse) corr, outp, phase, np, nb, N, kmin, kmax, bv, nbins = arg_tuple if fuse: args = [corr.htilde.data, corr.stilde.data] @@ -272,17 +274,18 @@ def shift_sum_points(num, N, arg_tuple): (nt,), *args, ) - - outp = outp[num*nbins:] + + outp = outp[num * nbins :] phase = phase[num:] np -= num return outp, phase, np + def shift_sum_points_pow2(num, arg_tuple): - #fuse = 'fuse' in corr.gpu_callback_method + # fuse = 'fuse' in corr.gpu_callback_method fuse = False - fn, nt = get_pchisq_fn_pow2(num, fuse_correlate = fuse) + fn, nt = get_pchisq_fn_pow2(num, fuse_correlate=fuse) corr, outp, points, np, nb, N, kmin, kmax, bv, nbins = arg_tuple if fuse: @@ -290,20 +293,18 @@ def shift_sum_points_pow2(num, arg_tuple): else: args = [corr.data] args += [outp, N] + points[0:num] + [kmin, kmax, bv, nbins] - fn( - (nb,), - (nt,), - tuple(args) - ) - - outp = outp[num*nbins:] + fn((nb,), (nt,), tuple(args)) + + outp = outp[num * nbins :] points = points[num:] np -= num return outp, points, np -@functools.lru_cache(maxsize=None) + +@functools.cache def get_cached_pow2(N): - return not(N & (N-1)) + return not (N & (N - 1)) + def shift_sum(corr, points, bins): kmin, kmax, bv = get_cached_bin_layout(bins) @@ -334,13 +335,12 @@ def shift_sum(corr, points, bins): cargs = (corr, outp, phase, np, nb, N, kmin, kmax, bv, nbins) if np >= 4: - outp, phase, np = shift_sum_points(4, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(4, cargs) # pylint:disable=no-value-for-parameter elif np >= 3: - outp, phase, np = shift_sum_points(3, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(3, cargs) # pylint:disable=no-value-for-parameter elif np >= 2: - outp, phase, np = shift_sum_points(2, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(2, cargs) # pylint:disable=no-value-for-parameter elif np == 1: - outp, phase, np = shift_sum_points(1, cargs) # pylint:disable=no-value-for-parameter + outp, phase, np = shift_sum_points(1, cargs) # pylint:disable=no-value-for-parameter return cp.asnumpy((outc.conj() * outc).sum(axis=1).real) - diff --git a/pycbc/vetoes/sgchisq.py b/pycbc/vetoes/sgchisq.py index bede804b48d..d4ad41ba186 100644 --- a/pycbc/vetoes/sgchisq.py +++ b/pycbc/vetoes/sgchisq.py @@ -1,25 +1,28 @@ -"""Chisq based on sine-gaussian tiles. +""" +Chisq based on sine-gaussian tiles. See https://arxiv.org/abs/1709.08974 for a discussion. """ import numpy -from pycbc.waveform.utils import apply_fseries_time_shift +from pycbc.events import ranking from pycbc.filter import sigma -from pycbc.waveform import sinegauss from pycbc.vetoes.chisq import SingleDetPowerChisq -from pycbc.events import ranking +from pycbc.waveform import sinegauss +from pycbc.waveform.utils import apply_fseries_time_shift + class SingleDetSGChisq(SingleDetPowerChisq): - """Class that handles precomputation and memory management for efficiently + """ + Class that handles precomputation and memory management for efficiently running the sine-Gaussian chisq """ - returns = {'sg_chisq': numpy.float32} - def __init__(self, bank, num_bins=0, - snr_threshold=None, - chisq_locations=None): - """ Create sine-Gaussian Chisq Calculator + returns = {"sg_chisq": numpy.float32} + + def __init__(self, bank, num_bins=0, snr_threshold=None, chisq_locations=None): + """ + Create sine-Gaussian Chisq Calculator Parameters ---------- @@ -35,6 +38,7 @@ def __init__(self, bank, num_bins=0, The offset is relative to the end frequency of the approximant. The region is a boolean expression such as 'mtotal>40' indicating which templates to apply this set of sine-Gaussians to. + """ if snr_threshold is not None: self.do = True @@ -43,8 +47,8 @@ def __init__(self, bank, num_bins=0, self.params = {} for descr in chisq_locations: region, values = descr.split(":") - mask = bank.table.parse_boolargs([(1, region), (0, 'else')])[0] - hashes = bank.table['template_hash'][mask.astype(bool)] + mask = bank.table.parse_boolargs([(1, region), (0, "else")])[0] + hashes = bank.table["template_hash"][mask.astype(bool)] for h in hashes: self.params[h] = values else: @@ -53,24 +57,31 @@ def __init__(self, bank, num_bins=0, @staticmethod def insert_option_group(parser): group = parser.add_argument_group("Sine-Gaussian Chisq") - group.add_argument("--sgchisq-snr-threshold", type=float, - help="Minimum SNR threshold to use SG chisq") - group.add_argument("--sgchisq-locations", type=str, nargs='+', + group.add_argument( + "--sgchisq-snr-threshold", + type=float, + help="Minimum SNR threshold to use SG chisq", + ) + group.add_argument( + "--sgchisq-locations", + type=str, + nargs="+", help="Frequency offsets and quality factors of the sine-Gaussians" - " to use, format 'region-boolean:q1-offset1,q2-offset2'. " - "Offset is relative to the end frequency of the approximant." - " Region is a boolean expression selecting templates to " - "apply the sine-Gaussians to, ex. 'mtotal>40'") + " to use, format 'region-boolean:q1-offset1,q2-offset2'. " + "Offset is relative to the end frequency of the approximant." + " Region is a boolean expression selecting templates to " + "apply the sine-Gaussians to, ex. 'mtotal>40'", + ) @classmethod def from_cli(cls, args, bank, chisq_bins): - return cls(bank, chisq_bins, - args.sgchisq_snr_threshold, - args.sgchisq_locations) + return cls(bank, chisq_bins, args.sgchisq_snr_threshold, args.sgchisq_locations) - def values(self, stilde, template, psd, snrv, snr_norm, - bchisq, bchisq_dof, indices): - """ Calculate sine-Gaussian chisq + def values( + self, stilde, template, psd, snrv, snr_norm, bchisq, bchisq_dof, indices + ): + """ + Calculate sine-Gaussian chisq Parameters ---------- @@ -95,13 +106,14 @@ def values(self, stilde, template, psd, snrv, snr_norm, ------- chisq: Array Chisq values, one for each sample index + """ if not self.do: return None if template.params.template_hash not in self.params: return numpy.ones(len(snrv)) - values = self.params[template.params.template_hash].split(',') + values = self.params[template.params.template_hash].split(",") # Get the chisq bins to use as the frequency reference point bins = self.cached_chisq_bins(template, psd) @@ -110,7 +122,7 @@ def values(self, stilde, template, psd, snrv, snr_norm, chisq = numpy.ones(len(snrv)) gtem = [None for _ in values] for i, snrvi in enumerate(snrv): - #Skip if newsnr too low + # Skip if newsnr too low snr = abs(snrvi * snr_norm) nsnr = ranking.newsnr(snr, bchisq[i] / bchisq_dof[i]) if nsnr < self.snr_threshold: @@ -133,7 +145,7 @@ def values(self, stilde, template, psd, snrv, snr_norm, # as constant over the last 2 chisq bins. We cannot use the final # chisq bin edge as it does not have to be where the waveform # terminates. - fstep = (bins[-2] - bins[-3]) + fstep = bins[-2] - bins[-3] fpeak = (bins[-2] + fstep) * template.delta_f # This is 90% of the Nyquist frequency of the data @@ -145,7 +157,7 @@ def values(self, stilde, template, psd, snrv, snr_norm, # Calculate the sum of SNR^2 for the sine-Gaussians specified for idxx, descr in enumerate(values): # Get the q and frequency offset from the descriptor - q, offset = descr.split('-') + q, offset = descr.split("-") q, offset = float(q), float(offset) fcen = fpeak + offset flow = max(kmin * template.delta_f, fcen - qwindow) @@ -159,25 +171,32 @@ def values(self, stilde, template, psd, snrv, snr_norm, kmin = int(flow / template.delta_f) kmax = int(fhigh / template.delta_f) - #Calculate sine-gaussian tile + # Calculate sine-gaussian tile if gtem[idxx] is None: # These are always the same values for a template, so # if computing 10 sgchisq points, don't want to call # this 10 times (for each SG template) - gtem[idxx] = sinegauss.fd_sine_gaussian(1.0, q, fcen, flow, - len(template) * template.delta_f, - template.delta_f).astype(numpy.complex64) - gsigma = sigma(gtem[idxx], psd=psd, - low_frequency_cutoff=flow, - high_frequency_cutoff=fhigh) - #Calculate the SNR of the tile + gtem[idxx] = sinegauss.fd_sine_gaussian( + 1.0, + q, + fcen, + flow, + len(template) * template.delta_f, + template.delta_f, + ).astype(numpy.complex64) + gsigma = sigma( + gtem[idxx], + psd=psd, + low_frequency_cutoff=flow, + high_frequency_cutoff=fhigh, + ) + # Calculate the SNR of the tile gsnr = (gtem[idxx][kmin:kmax] * stilde_shift[kmin:kmax]).sum() gsnr *= 4.0 * gtem[idxx].delta_f / gsigma - chisq[i] += abs(gsnr)**2.0 + chisq[i] += abs(gsnr) ** 2.0 dof += 2 if dof == 0: chisq[i] = 1 else: chisq[i] /= dof return chisq - diff --git a/pycbc/waveform/SpinTaylorF2.py b/pycbc/waveform/SpinTaylorF2.py index f9bf13ad9be..cf1276c985e 100644 --- a/pycbc/waveform/SpinTaylorF2.py +++ b/pycbc/waveform/SpinTaylorF2.py @@ -16,18 +16,18 @@ # Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, # MA 02111-1307 USA -import numpy -from numpy import sqrt, double, complex128 -from math import pow, log, cos, sin, acos, atan2 +from math import acos, atan2, cos, log, pow, sin +import numpy +from numpy import complex128, double, sqrt from pycuda.elementwise import ElementwiseKernel +from pycbc.constants import GAMMA, MRSUN_SI, MTSUN_SI, PC_SI, PI from pycbc.libutils import pkg_config_header_strings from pycbc.types import FrequencySeries, zeros from pycbc.waveform.utils import ceilpow2 -from pycbc.constants import MTSUN_SI, PC_SI, PI, MRSUN_SI, GAMMA -preamble = f""" +preamble = """ #include #include """ @@ -228,7 +228,8 @@ """ -spintaylorf2_kernel = ElementwiseKernel("""pycuda::complex *htildeP, +spintaylorf2_kernel = ElementwiseKernel( + """pycuda::complex *htildeP, pycuda::complex *htildeC, int kmin, int phase_order, int amplitude_order, double delta_f, double PI, @@ -255,161 +256,339 @@ double IM_SBfac3, double IM_SBfac4, double psiJ_P, double psiJ_C, double gamma0""", - spintaylorf2_text, "spintaylorf2_kernel", - preamble=preamble, options=pkg_config_header_strings([])) + spintaylorf2_text, + "spintaylorf2_kernel", + preamble=preamble, + options=pkg_config_header_strings([]), +) + def spintaylorf2(**kwds): - """ Return a SpinTaylorF2 waveform using CUDA to generate the phase and amplitude - """ + """Return a SpinTaylorF2 waveform using CUDA to generate the phase and amplitude""" #####Pull out the input arguments##### - f_lower = double(kwds['f_lower']) - delta_f = double(kwds['delta_f']) - distance = double(kwds['distance']) - mass1 = double(kwds['mass1']) - mass2 = double(kwds['mass2']) - spin1x = double(kwds['spin1x']) - spin1y = double(kwds['spin1y']) - spin1z = double(kwds['spin1z']) - phi0 = double(kwds['coa_phase']) #Orbital Phase at coalescence - phase_order = int(kwds['phase_order']) - amplitude_order = int(kwds['amplitude_order']) - inclination = double(kwds['inclination']) + f_lower = double(kwds["f_lower"]) + delta_f = double(kwds["delta_f"]) + distance = double(kwds["distance"]) + mass1 = double(kwds["mass1"]) + mass2 = double(kwds["mass2"]) + spin1x = double(kwds["spin1x"]) + spin1y = double(kwds["spin1y"]) + spin1z = double(kwds["spin1z"]) + phi0 = double(kwds["coa_phase"]) # Orbital Phase at coalescence + phase_order = int(kwds["phase_order"]) + amplitude_order = int(kwds["amplitude_order"]) + inclination = double(kwds["inclination"]) lnhatx = sin(inclination) - lnhaty = 0. + lnhaty = 0.0 lnhatz = cos(inclination) - psi = 0. + psi = 0.0 - tC= -1.0 / delta_f + tC = -1.0 / delta_f M = mass1 + mass2 eta = mass1 * mass2 / (M * M) m_sec = M * MTSUN_SI piM = PI * m_sec - vISCO = 1. / sqrt(6.) + vISCO = 1.0 / sqrt(6.0) fISCO = vISCO * vISCO * vISCO / piM f_max = ceilpow2(fISCO) n = int(f_max / delta_f + 1) kmax = int(fISCO / delta_f) kmin = int(numpy.ceil(f_lower / delta_f)) - kmax = kmax if (kmax 0.) else 1. - Jx0 = mass1*mass2*lnhatx/v0 + mass1*mass1*spin1x - Jy0 = mass1*mass2*lnhaty/v0 + mass1*mass1*spin1y - Jz0 = mass1*mass2*lnhatz/v0 + mass1*mass1*spin1z - thetaJ = acos(Jz0 / sqrt(Jx0**2+Jy0**2+Jz0**2)) - psiJ = atan2(Jy0, -Jx0) # FIXME: check that Jy0 and Jx0 are not both 0 + v0 = pow(piM * kmin * delta_f, 1.0 / 3) + chi = sqrt(spin1x**2 + spin1y**2 + spin1z**2) + kappa = ( + (lnhatx * spin1x + lnhaty * spin1y + lnhatz * spin1z) / chi + if (chi > 0.0) + else 1.0 + ) + Jx0 = mass1 * mass2 * lnhatx / v0 + mass1 * mass1 * spin1x + Jy0 = mass1 * mass2 * lnhaty / v0 + mass1 * mass1 * spin1y + Jz0 = mass1 * mass2 * lnhatz / v0 + mass1 * mass1 * spin1z + thetaJ = acos(Jz0 / sqrt(Jx0**2 + Jy0**2 + Jz0**2)) + psiJ = atan2(Jy0, -Jx0) # FIXME: check that Jy0 and Jx0 are not both 0 # Rotate Lnhat back to frame where J is along z, to figure out initial alpha - rotLx = lnhatx*cos(thetaJ)*cos(psiJ) - lnhaty*cos(thetaJ)*sin(psiJ) + lnhatz*sin(thetaJ) - rotLy = lnhatx*sin(psiJ) + lnhaty*cos(psiJ) - alpha0 = atan2(rotLy, rotLx) # FIXME: check that rotLy and rotLx are not both 0 - psiJ_P =psiJ + psi - psiJ_C =psiJ + psi + PI/4. + rotLx = ( + lnhatx * cos(thetaJ) * cos(psiJ) + - lnhaty * cos(thetaJ) * sin(psiJ) + + lnhatz * sin(thetaJ) + ) + rotLy = lnhatx * sin(psiJ) + lnhaty * cos(psiJ) + alpha0 = atan2(rotLy, rotLx) # FIXME: check that rotLy and rotLx are not both 0 + psiJ_P = psiJ + psi + psiJ_C = psiJ + psi + PI / 4.0 #####Calculate the Coefficients##### - #quadparam = 1. - gamma0 = mass1*chi/mass2 - #Calculate the spin corrections + # quadparam = 1. + gamma0 = mass1 * chi / mass2 + # Calculate the spin corrections # FIXME should use pycbc's function, but sigma has different expression # in Andy's code, double check # pn_beta, pn_sigma, pn_gamma = pycbc.pnutils.mass1_mass2_spin1z_spin2z_to_beta_sigma_gamma( # mass1, mass2, chi*kappa, 0) # FIXME: spin2 is taken to be 0 - pn_beta = (113.*mass1/(12.*M) - 19.*eta/6.)*chi*kappa - pn_sigma = ( (5.*(3.*kappa*kappa-1.)/2.) + (7. - kappa*kappa)/96. ) * (mass1*mass1*chi*chi/M/M) - pn_gamma = (5.*(146597. + 7056.*eta)*mass1/(2268.*M) - 10.*eta*(1276. + 153.*eta)/81.)*chi*kappa - prec_fac0 = 5.*(4. + 3.*mass2/mass1)/64. - dtdv2 = 743./336. + 11.*eta/4. - dtdv3 = -4.*PI + pn_beta - dtdv4 = 3058673./1016064. + 5429.*eta/1008. + 617.*eta*eta/144. - pn_sigma - dtdv5 = (-7729./672.+13.*eta/8.)*PI + 9.*pn_gamma/40. + pn_beta = (113.0 * mass1 / (12.0 * M) - 19.0 * eta / 6.0) * chi * kappa + pn_sigma = ( + (5.0 * (3.0 * kappa * kappa - 1.0) / 2.0) + (7.0 - kappa * kappa) / 96.0 + ) * (mass1 * mass1 * chi * chi / M / M) + pn_gamma = ( + ( + 5.0 * (146597.0 + 7056.0 * eta) * mass1 / (2268.0 * M) + - 10.0 * eta * (1276.0 + 153.0 * eta) / 81.0 + ) + * chi + * kappa + ) + prec_fac0 = 5.0 * (4.0 + 3.0 * mass2 / mass1) / 64.0 + dtdv2 = 743.0 / 336.0 + 11.0 * eta / 4.0 + dtdv3 = -4.0 * PI + pn_beta + dtdv4 = ( + 3058673.0 / 1016064.0 + + 5429.0 * eta / 1008.0 + + 617.0 * eta * eta / 144.0 + - pn_sigma + ) + dtdv5 = (-7729.0 / 672.0 + 13.0 * eta / 8.0) * PI + 9.0 * pn_gamma / 40.0 #####Calculate the Initial Euler Angles alpha_ref, beta_ref=0 and zeta_ref##### - gam = gamma0*v0 - sqrtfac = sqrt(1. + 2.*kappa*gam + gam*gam) + gam = gamma0 * v0 + sqrtfac = sqrt(1.0 + 2.0 * kappa * gam + gam * gam) logv0 = log(v0) - logfac1 = log(1. + kappa*gam + sqrtfac) + logfac1 = log(1.0 + kappa * gam + sqrtfac) logfac2 = log(kappa + gam + sqrtfac) v02 = v0 * v0 v03 = v0 * v02 kappa2 = kappa * kappa kappa3 = kappa2 * kappa gamma02 = gamma0 * gamma0 - gamma03 = gamma02 *gamma0 - - alpha_ref = prec_fac0*( logfac2 *( dtdv2*gamma0 + dtdv3*kappa - dtdv5*kappa/(2.*gamma02) + dtdv4/(2.*gamma0) - dtdv4*kappa2/(2.*gamma0) + (dtdv5*kappa3)/(2.*gamma02) ) + logfac1*( - dtdv2*gamma0*kappa - dtdv3 + kappa*gamma03/2. - gamma03*kappa3/2. ) + logv0 *( dtdv2*gamma0*kappa + dtdv3 - kappa*gamma03/2. + gamma03*kappa3/2. ) + sqrtfac *( dtdv3 + dtdv4*v0/2. + dtdv5/gamma02/3. + dtdv4*kappa/(2.*gamma0) + dtdv5*kappa*v0/(6.*gamma0) - dtdv5*kappa2/(2.*gamma02) - 1/(3.*v03) - gamma0*kappa/(6.*v02) - dtdv2/v0 - gamma02/(3.*v0) + gamma02*kappa2/(2.*v0) + dtdv5*v02/3. )) - alpha0 - - zeta_ref = prec_fac0*( dtdv3*gamma0*kappa*v0 + dtdv4*v0 + logfac2 *(-dtdv2*gamma0 - dtdv3*kappa + dtdv5*kappa/(2.*gamma02) - dtdv4/(2.*gamma0) + dtdv4*kappa2/(2.*gamma0) - dtdv5*kappa3/(2.*gamma02) ) + logv0 *( kappa*gamma03/2. - gamma03*kappa3/2. ) + logfac1 *( dtdv2*gamma0*kappa + dtdv3 - kappa*gamma03/2. + gamma03*kappa3/2. ) - 1/(3.*v03) - gamma0*kappa/(2.*v02) - dtdv2/v0 + dtdv4*gamma0*kappa*v02/2. + dtdv5*v02/2. + sqrtfac *( -dtdv3 - dtdv4*v0/2. - dtdv5/(3.*gamma02) - dtdv4*kappa/(2.*gamma0) - dtdv5*kappa*v0/(6.*gamma0) + dtdv5*kappa2/(2.*gamma02) + 1/(3.*v03) + gamma0*kappa/(6.*v02) + dtdv2/v0 + gamma02/(3.*v0) - gamma02*kappa2/(2.*v0) - dtdv5*v02/3. ) + dtdv5*gamma0*kappa*v03/3. ) + gamma03 = gamma02 * gamma0 + + alpha_ref = ( + prec_fac0 + * ( + logfac2 + * ( + dtdv2 * gamma0 + + dtdv3 * kappa + - dtdv5 * kappa / (2.0 * gamma02) + + dtdv4 / (2.0 * gamma0) + - dtdv4 * kappa2 / (2.0 * gamma0) + + (dtdv5 * kappa3) / (2.0 * gamma02) + ) + + logfac1 + * ( + -dtdv2 * gamma0 * kappa + - dtdv3 + + kappa * gamma03 / 2.0 + - gamma03 * kappa3 / 2.0 + ) + + logv0 + * ( + dtdv2 * gamma0 * kappa + + dtdv3 + - kappa * gamma03 / 2.0 + + gamma03 * kappa3 / 2.0 + ) + + sqrtfac + * ( + dtdv3 + + dtdv4 * v0 / 2.0 + + dtdv5 / gamma02 / 3.0 + + dtdv4 * kappa / (2.0 * gamma0) + + dtdv5 * kappa * v0 / (6.0 * gamma0) + - dtdv5 * kappa2 / (2.0 * gamma02) + - 1 / (3.0 * v03) + - gamma0 * kappa / (6.0 * v02) + - dtdv2 / v0 + - gamma02 / (3.0 * v0) + + gamma02 * kappa2 / (2.0 * v0) + + dtdv5 * v02 / 3.0 + ) + ) + - alpha0 + ) + + zeta_ref = prec_fac0 * ( + dtdv3 * gamma0 * kappa * v0 + + dtdv4 * v0 + + logfac2 + * ( + -dtdv2 * gamma0 + - dtdv3 * kappa + + dtdv5 * kappa / (2.0 * gamma02) + - dtdv4 / (2.0 * gamma0) + + dtdv4 * kappa2 / (2.0 * gamma0) + - dtdv5 * kappa3 / (2.0 * gamma02) + ) + + logv0 * (kappa * gamma03 / 2.0 - gamma03 * kappa3 / 2.0) + + logfac1 + * ( + dtdv2 * gamma0 * kappa + + dtdv3 + - kappa * gamma03 / 2.0 + + gamma03 * kappa3 / 2.0 + ) + - 1 / (3.0 * v03) + - gamma0 * kappa / (2.0 * v02) + - dtdv2 / v0 + + dtdv4 * gamma0 * kappa * v02 / 2.0 + + dtdv5 * v02 / 2.0 + + sqrtfac + * ( + -dtdv3 + - dtdv4 * v0 / 2.0 + - dtdv5 / (3.0 * gamma02) + - dtdv4 * kappa / (2.0 * gamma0) + - dtdv5 * kappa * v0 / (6.0 * gamma0) + + dtdv5 * kappa2 / (2.0 * gamma02) + + 1 / (3.0 * v03) + + gamma0 * kappa / (6.0 * v02) + + dtdv2 / v0 + + gamma02 / (3.0 * v0) + - gamma02 * kappa2 / (2.0 * v0) + - dtdv5 * v02 / 3.0 + ) + + dtdv5 * gamma0 * kappa * v03 / 3.0 + ) #####Calculate the Complex sideband factors, mm=2 is first entry##### - RE_SBfac0= (1.+cos(thetaJ)**2)/2. - RE_SBfac1= sin(2.*thetaJ) - RE_SBfac2= 3.*sin(thetaJ)**2 - RE_SBfac3= -sin(2.*thetaJ) - RE_SBfac4= (1.+cos(thetaJ)**2)/2. - IM_SBfac0= -cos(thetaJ) - IM_SBfac1= -2.*sin(thetaJ) - IM_SBfac2= 0. - IM_SBfac3= -2.*sin(thetaJ) - IM_SBfac4= cos(thetaJ) + RE_SBfac0 = (1.0 + cos(thetaJ) ** 2) / 2.0 + RE_SBfac1 = sin(2.0 * thetaJ) + RE_SBfac2 = 3.0 * sin(thetaJ) ** 2 + RE_SBfac3 = -sin(2.0 * thetaJ) + RE_SBfac4 = (1.0 + cos(thetaJ) ** 2) / 2.0 + IM_SBfac0 = -cos(thetaJ) + IM_SBfac1 = -2.0 * sin(thetaJ) + IM_SBfac2 = 0.0 + IM_SBfac3 = -2.0 * sin(thetaJ) + IM_SBfac4 = cos(thetaJ) #####Calculate the PN terms##### - theta = -11831./9240. - lambdaa = -1987./3080.0 - pfaN = 3.0/(128.0 * eta) - pfa2 = 5.0*(743.0/84 + 11.0 * eta)/9.0 - pfa3 = -16.0*PI + 4.0*pn_beta - pfa4 = 5.0*(3058.673/7.056 + 5429.0/7.0 * eta + 617.0 * eta*eta)/72.0 - \ - 10.0*pn_sigma - pfa5 = 5.0/9.0 * (7729.0/84.0 - 13.0 * eta) * PI - pn_gamma - pfl5 = 5.0/3.0 * (7729.0/84.0 - 13.0 * eta) * PI - pn_gamma * 3 - pfa6 = (11583.231236531/4.694215680 - 640.0/3.0 * PI * PI- \ - 6848.0/21.0*GAMMA) + \ - eta * (-15335.597827/3.048192 + 2255./12. * PI * \ - PI - 1760./3.*theta +12320./9.*lambdaa) + \ - eta*eta * 76055.0/1728.0 - \ - eta*eta*eta* 127825.0/1296.0 - pfl6 = -6848.0/21.0 - pfa7 = PI * 5.0/756.0 * ( 15419335.0/336.0 + 75703.0/2.0 * eta - \ - 14809.0 * eta*eta) - - FTaN = 32.0 * eta*eta / 5.0 - FTa2 = -(12.47/3.36 + 3.5/1.2 * eta) + theta = -11831.0 / 9240.0 + lambdaa = -1987.0 / 3080.0 + pfaN = 3.0 / (128.0 * eta) + pfa2 = 5.0 * (743.0 / 84 + 11.0 * eta) / 9.0 + pfa3 = -16.0 * PI + 4.0 * pn_beta + pfa4 = ( + 5.0 * (3058.673 / 7.056 + 5429.0 / 7.0 * eta + 617.0 * eta * eta) / 72.0 + - 10.0 * pn_sigma + ) + pfa5 = 5.0 / 9.0 * (7729.0 / 84.0 - 13.0 * eta) * PI - pn_gamma + pfl5 = 5.0 / 3.0 * (7729.0 / 84.0 - 13.0 * eta) * PI - pn_gamma * 3 + pfa6 = ( + (11583.231236531 / 4.694215680 - 640.0 / 3.0 * PI * PI - 6848.0 / 21.0 * GAMMA) + + eta + * ( + -15335.597827 / 3.048192 + + 2255.0 / 12.0 * PI * PI + - 1760.0 / 3.0 * theta + + 12320.0 / 9.0 * lambdaa + ) + + eta * eta * 76055.0 / 1728.0 + - eta * eta * eta * 127825.0 / 1296.0 + ) + pfl6 = -6848.0 / 21.0 + pfa7 = ( + PI + * 5.0 + / 756.0 + * (15419335.0 / 336.0 + 75703.0 / 2.0 * eta - 14809.0 * eta * eta) + ) + + FTaN = 32.0 * eta * eta / 5.0 + FTa2 = -(12.47 / 3.36 + 3.5 / 1.2 * eta) FTa3 = 4.0 * PI - FTa4 = -(44.711/9.072 - 92.71/5.04 * eta - 6.5/1.8 * eta*eta) - FTa5 = -(81.91/6.72 + 58.3/2.4 * eta) * PI - FTa6 = (664.3739519/6.9854400 + 16.0/3.0 * PI*PI - - 17.12/1.05 * GAMMA + - (4.1/4.8 * PI*PI - 134.543/7.776) * eta - - 94.403/3.024 * eta*eta - 7.75/3.24 * eta*eta*eta) - FTl6 = -8.56/1.05 - FTa7 = -(162.85/5.04 - 214.745/1.728 * eta - 193.385/3.024 * eta*eta) \ - * PI - - dETaN = 2 * -eta/2.0 - dETa1 = 2 * -(3.0/4.0 + 1.0/12.0 * eta) - dETa2 = 3 * -(27.0/8.0 - 19.0/8.0 * eta + 1./24.0 * eta*eta) - dETa3 = 4 * -(67.5/6.4 - (344.45/5.76 - 20.5/9.6 * PI*PI) * - eta + 15.5/9.6 * eta*eta + 3.5/518.4 * eta*eta*eta) - - amp0 = -4. * mass1 * mass2 / (1.0e+06 * distance * PC_SI ) * \ - MRSUN_SI * MTSUN_SI * sqrt(PI/12.0) - - htildeP = FrequencySeries(zeros(n,dtype=complex128), delta_f=delta_f, copy=False) - htildeC = FrequencySeries(zeros(n,dtype=complex128), delta_f=delta_f, copy=False) - spintaylorf2_kernel(htildeP.data[kmin:kmax], htildeC.data[kmin:kmax], - kmin, phase_order, amplitude_order, delta_f, PI, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, FTaN, FTa2, - FTa3, FTa4, FTa5, FTa6, - FTl6, FTa7, dETaN, dETa1, dETa2, dETa3, - amp0, tC, phi0, - kappa, prec_fac0, alpha_ref, zeta_ref, - dtdv2, dtdv3, dtdv4, dtdv5, - RE_SBfac0, RE_SBfac1, RE_SBfac2, RE_SBfac3, RE_SBfac4, - IM_SBfac0, IM_SBfac1, IM_SBfac2, IM_SBfac3, IM_SBfac4, - psiJ_P, psiJ_C, gamma0) + FTa4 = -(44.711 / 9.072 - 92.71 / 5.04 * eta - 6.5 / 1.8 * eta * eta) + FTa5 = -(81.91 / 6.72 + 58.3 / 2.4 * eta) * PI + FTa6 = ( + 664.3739519 / 6.9854400 + + 16.0 / 3.0 * PI * PI + - 17.12 / 1.05 * GAMMA + + (4.1 / 4.8 * PI * PI - 134.543 / 7.776) * eta + - 94.403 / 3.024 * eta * eta + - 7.75 / 3.24 * eta * eta * eta + ) + FTl6 = -8.56 / 1.05 + FTa7 = -(162.85 / 5.04 - 214.745 / 1.728 * eta - 193.385 / 3.024 * eta * eta) * PI + + dETaN = 2 * -eta / 2.0 + dETa1 = 2 * -(3.0 / 4.0 + 1.0 / 12.0 * eta) + dETa2 = 3 * -(27.0 / 8.0 - 19.0 / 8.0 * eta + 1.0 / 24.0 * eta * eta) + dETa3 = 4 * -( + 67.5 / 6.4 + - (344.45 / 5.76 - 20.5 / 9.6 * PI * PI) * eta + + 15.5 / 9.6 * eta * eta + + 3.5 / 518.4 * eta * eta * eta + ) + + amp0 = ( + -4.0 + * mass1 + * mass2 + / (1.0e06 * distance * PC_SI) + * MRSUN_SI + * MTSUN_SI + * sqrt(PI / 12.0) + ) + + htildeP = FrequencySeries(zeros(n, dtype=complex128), delta_f=delta_f, copy=False) + htildeC = FrequencySeries(zeros(n, dtype=complex128), delta_f=delta_f, copy=False) + spintaylorf2_kernel( + htildeP.data[kmin:kmax], + htildeC.data[kmin:kmax], + kmin, + phase_order, + amplitude_order, + delta_f, + PI, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + FTaN, + FTa2, + FTa3, + FTa4, + FTa5, + FTa6, + FTl6, + FTa7, + dETaN, + dETa1, + dETa2, + dETa3, + amp0, + tC, + phi0, + kappa, + prec_fac0, + alpha_ref, + zeta_ref, + dtdv2, + dtdv3, + dtdv4, + dtdv5, + RE_SBfac0, + RE_SBfac1, + RE_SBfac2, + RE_SBfac3, + RE_SBfac4, + IM_SBfac0, + IM_SBfac1, + IM_SBfac2, + IM_SBfac3, + IM_SBfac4, + psiJ_P, + psiJ_C, + gamma0, + ) return htildeP, htildeC diff --git a/pycbc/waveform/__init__.py b/pycbc/waveform/__init__.py index 3e5a46755e2..07fe083b1dc 100644 --- a/pycbc/waveform/__init__.py +++ b/pycbc/waveform/__init__.py @@ -1,11 +1,13 @@ -from pycbc.waveform.waveform import * -from pycbc.waveform.utils import * from pycbc.waveform.bank import * -from pycbc.waveform.ringdown import * from pycbc.waveform.parameters import * -from pycbc.waveform.waveform_modes import (get_td_waveform_modes, - get_fd_waveform_modes) -from pycbc.waveform.plugin import (retrieve_waveform_plugins, - add_custom_waveform, - add_length_estimator) +from pycbc.waveform.plugin import ( + add_custom_waveform, + add_length_estimator, + retrieve_waveform_plugins, +) +from pycbc.waveform.ringdown import * +from pycbc.waveform.utils import * +from pycbc.waveform.waveform import * +from pycbc.waveform.waveform_modes import get_fd_waveform_modes, get_td_waveform_modes + retrieve_waveform_plugins() diff --git a/pycbc/waveform/bank.py b/pycbc/waveform/bank.py index 65388fbe2e8..ced381bbe09 100644 --- a/pycbc/waveform/bank.py +++ b/pycbc/waveform/bank.py @@ -25,76 +25,83 @@ """ This module provides classes that describe banks of waveforms """ -import types + +import hashlib import logging import os.path -import h5py +import types +import warnings from copy import copy + +import h5py import numpy as np -from igwn_ligolw import lsctables, utils as ligolw_utils -import pycbc.waveform +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils + +import pycbc.io import pycbc.pnutils +import pycbc.waveform import pycbc.waveform.compress from pycbc import DYN_RANGE_FAC -from pycbc.types import FrequencySeries, zeros -import pycbc.io from pycbc.io.ligolw import LIGOLWContentHandler -import hashlib -import warnings +from pycbc.types import FrequencySeries, zeros def sigma_cached(self, psd): - """ Cache sigma calculate for use in tandem with the FilterBank class - """ - if not hasattr(self, '_sigmasq'): + """Cache sigma calculate for use in tandem with the FilterBank class""" + if not hasattr(self, "_sigmasq"): from pycbc.opt import LimitedSizeDict + self._sigmasq = LimitedSizeDict(size_limit=2**5) key = id(psd) - if not hasattr(psd, '_sigma_cached_key'): + if not hasattr(psd, "_sigma_cached_key"): psd._sigma_cached_key = {} if key not in self._sigmasq or id(self) not in psd._sigma_cached_key: psd._sigma_cached_key[id(self)] = True # If possible, we precalculate the sigmasq vector for all possible waveforms if pycbc.waveform.waveform_norm_exists(self.approximant): - if not hasattr(psd, 'sigmasq_vec'): + if not hasattr(psd, "sigmasq_vec"): psd.sigmasq_vec = {} if self.approximant not in psd.sigmasq_vec: - psd.sigmasq_vec[self.approximant] = \ + psd.sigmasq_vec[self.approximant] = ( pycbc.waveform.get_waveform_filter_norm( - self.approximant, - psd, - len(psd), - psd.delta_f, - self.min_f_lower + self.approximant, psd, len(psd), psd.delta_f, self.min_f_lower ) + ) - if not hasattr(self, 'sigma_scale'): + if not hasattr(self, "sigma_scale"): # Get an amplitude normalization (mass dependant constant norm) amp_norm = pycbc.waveform.get_template_amplitude_norm( - self.params, approximant=self.approximant) + self.params, approximant=self.approximant + ) amp_norm = 1 if amp_norm is None else amp_norm self.sigma_scale = (DYN_RANGE_FAC * amp_norm) ** 2.0 curr_sigmasq = psd.sigmasq_vec[self.approximant] kmin = int(self.f_lower / psd.delta_f) - self._sigmasq[key] = self.sigma_scale * \ - (curr_sigmasq[self.end_idx-1] - curr_sigmasq[kmin]) + self._sigmasq[key] = self.sigma_scale * ( + curr_sigmasq[self.end_idx - 1] - curr_sigmasq[kmin] + ) else: - if not hasattr(self, 'sigma_view'): + if not hasattr(self, "sigma_view"): from pycbc.filter.matchedfilter import get_cutoff_indices - N = (len(self) -1) * 2 + + N = (len(self) - 1) * 2 kmin, kmax = get_cutoff_indices( - self.min_f_lower or self.f_lower, self.end_frequency, - self.delta_f, N) + self.min_f_lower or self.f_lower, + self.end_frequency, + self.delta_f, + N, + ) self.sslice = slice(kmin, kmax) self.sigma_view = self[self.sslice].squared_norm() * 4.0 * self.delta_f - if not hasattr(psd, 'invsqrt'): + if not hasattr(psd, "invsqrt"): psd.invsqrt = 1.0 / psd self._sigmasq[key] = self.sigma_view.inner(psd.invsqrt[self.sslice]) @@ -103,7 +110,8 @@ def sigma_cached(self, psd): # helper function for parsing approximant strings def boolargs_from_apprxstr(approximant_strs): - """Parses a list of strings specifying an approximant and where that + """ + Parses a list of strings specifying an approximant and where that approximant should be used into a list that can be understood by FieldArray.parse_boolargs. @@ -121,14 +129,16 @@ def boolargs_from_apprxstr(approximant_strs): boolargs : list A list of tuples giving the approximant and where to apply them. This can be passed directly to `FieldArray.parse_boolargs`. + """ if not isinstance(approximant_strs, list): approximant_strs = [approximant_strs] - return [tuple(arg.split(':')) for arg in approximant_strs] + return [tuple(arg.split(":")) for arg in approximant_strs] def add_approximant_arg(parser, default=None, help=None): - """Adds an approximant argument to the given parser. + """ + Adds an approximant argument to the given parser. Parameters ---------- @@ -139,35 +149,42 @@ def add_approximant_arg(parser, default=None, help=None): help : {None, str} Provide a custom help message. If None, will use a descriptive message on how to specify the approximant. + """ if help is None: - help=str("The approximant(s) to use. Multiple approximants to use " - "in different regions may be provided. If multiple " - "approximants are provided, every one but the last must be " - "be followed by a conditional statement defining where that " - "approximant should be used. Conditionals can be any boolean " - "test understood by numpy. For example, 'Apprx:(mtotal > 4) & " - "(mchirp <= 5)' would use approximant 'Apprx' where total mass " - "is > 4 and chirp mass is <= 5. " - "Conditionals are applied in order, with each successive one " - "only applied to regions not covered by previous arguments. " - "For example, `'TaylorF2:mtotal < 4' 'IMRPhenomD:mchirp < 3'` " - "would result in IMRPhenomD being used where chirp mass is < 3 " - "and total mass is >= 4. The last approximant given may use " - "'else' as the conditional or include no conditional. In either " - "case, this will cause the last approximant to be used in any " - "remaning regions after all the previous conditionals have been " - "applied. For the full list of possible parameters to apply " - "conditionals to, see WaveformArray.default_fields(). Math " - "operations may also be used on parameters; syntax is python, " - "with any operation recognized by numpy.") - parser.add_argument("--approximant", nargs='+', type=str, default=default, - metavar='APPRX[:COND]', - help=help) + help = ("The approximant(s) to use. Multiple approximants to use " + "in different regions may be provided. If multiple " + "approximants are provided, every one but the last must be " + "be followed by a conditional statement defining where that " + "approximant should be used. Conditionals can be any boolean " + "test understood by numpy. For example, 'Apprx:(mtotal > 4) & " + "(mchirp <= 5)' would use approximant 'Apprx' where total mass " + "is > 4 and chirp mass is <= 5. " + "Conditionals are applied in order, with each successive one " + "only applied to regions not covered by previous arguments. " + "For example, `'TaylorF2:mtotal < 4' 'IMRPhenomD:mchirp < 3'` " + "would result in IMRPhenomD being used where chirp mass is < 3 " + "and total mass is >= 4. The last approximant given may use " + "'else' as the conditional or include no conditional. In either " + "case, this will cause the last approximant to be used in any " + "remaning regions after all the previous conditionals have been " + "applied. For the full list of possible parameters to apply " + "conditionals to, see WaveformArray.default_fields(). Math " + "operations may also be used on parameters; syntax is python, " + "with any operation recognized by numpy.") + parser.add_argument( + "--approximant", + nargs="+", + type=str, + default=default, + metavar="APPRX[:COND]", + help=help, + ) def parse_approximant_arg(approximant_arg, warray): - """Given an approximant arg (see add_approximant_arg) and a field + """ + Given an approximant arg (see add_approximant_arg) and a field array, figures out what approximant to use for each template in the array. Parameters @@ -184,9 +201,11 @@ def parse_approximant_arg(approximant_arg, warray): array A numpy array listing the approximants to use for each element in the warray. + """ return warray.parse_boolargs(boolargs_from_apprxstr(approximant_arg))[0] + def tuple_to_hash(tuple_to_be_hashed): """ Return a hash for a numpy array, avoids native (unsafe) python3 hash function @@ -201,14 +220,15 @@ def tuple_to_hash(tuple_to_be_hashed): ------- int an integer representation of the hashed array + """ - h = hashlib.blake2b(np.array(tuple_to_be_hashed).tobytes('C'), - digest_size=8) + h = hashlib.blake2b(np.array(tuple_to_be_hashed).tobytes("C"), digest_size=8) return np.frombuffer(h.digest(), dtype=int)[0] -class TemplateBank(object): - r"""Class to provide some basic helper functions and information +class TemplateBank: + r""" + Class to provide some basic helper functions and information about elements of a template bank. Parameters @@ -260,39 +280,45 @@ class TemplateBank(object): (left open after initialization). Otherwise, None. extra_args : {None, dict} Any extra keyword arguments that were provided on initialization. + """ - def __init__(self, filename, approximant=None, parameters=None, - **kwds): + + def __init__(self, filename, approximant=None, parameters=None, **kwds): self.has_compressed_waveforms = False ext = os.path.basename(filename) - if ext.endswith(('.xml', '.xml.gz', '.xmlgz')): + if ext.endswith((".xml", ".xml.gz", ".xmlgz")): self.filehandler = None self.indoc = ligolw_utils.load_filename( - filename, False, contenthandler=LIGOLWContentHandler) + filename, False, contenthandler=LIGOLWContentHandler + ) self.table = lsctables.SnglInspiralTable.get_table(self.indoc) - self.table = pycbc.io.WaveformArray.from_ligolw_table(self.table, - columns=parameters) + self.table = pycbc.io.WaveformArray.from_ligolw_table( + self.table, columns=parameters + ) # inclination stored in xml alpha3 column names = list(self.table.dtype.names) - names = tuple([n if n != 'alpha3' else 'inclination' for n in names]) + names = tuple([n if n != "alpha3" else "inclination" for n in names]) # low frequency cutoff in xml alpha6 column - names = tuple([n if n!= 'alpha6' else 'f_lower' for n in names]) + names = tuple([n if n != "alpha6" else "f_lower" for n in names]) self.table.dtype.names = names - elif ext.endswith(('hdf', '.h5', '.hdf5')): + elif ext.endswith(("hdf", ".h5", ".hdf5")): self.indoc = None - f = pycbc.io.HFile(filename, 'r') + f = pycbc.io.HFile(filename, "r") self.filehandler = f try: - fileparams = list(f.attrs['parameters']) + fileparams = list(f.attrs["parameters"]) except KeyError: # just assume all of the top-level groups are the parameters fileparams = list(f.keys()) - logging.info("WARNING: no parameters attribute found. " - "Assuming that %s " %(', '.join(fileparams)) + - "are the parameters.") + logging.info( + "WARNING: no parameters attribute found. " + "Assuming that %s " + % (", ".join(fileparams)) + + "are the parameters." + ) tmp_params = [] # At this point fileparams might be bytes. Fix if it is for param in fileparams: @@ -307,14 +333,12 @@ def __init__(self, filename, approximant=None, parameters=None, # need to be loaded if parameters is None: parameters = fileparams - common_fields = list(pycbc.io.WaveformArray(1, - names=parameters).fieldnames) - add_fields = list(set(parameters) & - (set(fileparams) - set(common_fields))) + common_fields = list(pycbc.io.WaveformArray(1, names=parameters).fieldnames) + add_fields = list(set(parameters) & (set(fileparams) - set(common_fields))) # load dtype = [] data = {} - for key in common_fields+add_fields: + for key in common_fields + add_fields: data[key] = f[key][:] dtype.append((key, data[key].dtype)) num = f[fileparams[0]].size @@ -322,59 +346,74 @@ def __init__(self, filename, approximant=None, parameters=None, for key in data: self.table[key] = data[key] # add the compressed waveforms, if they exist - self.has_compressed_waveforms = 'compressed_waveforms' in f + self.has_compressed_waveforms = "compressed_waveforms" in f else: - raise ValueError("Unsupported template bank file extension %s" %( - ext)) + raise ValueError("Unsupported template bank file extension %s" % (ext)) # if approximant is specified, override whatever was in the file # (if anything was in the file) if approximant is not None: # get the approximant for each template - dtype = h5py.string_dtype(encoding='utf-8') - apprxs = np.array(self.parse_approximant(approximant), - dtype=dtype) - if 'approximant' not in self.table.fieldnames: - self.table = self.table.add_fields(apprxs, 'approximant') + dtype = h5py.string_dtype(encoding="utf-8") + apprxs = np.array(self.parse_approximant(approximant), dtype=dtype) + if "approximant" not in self.table.fieldnames: + self.table = self.table.add_fields(apprxs, "approximant") else: - self.table['approximant'] = apprxs + self.table["approximant"] = apprxs self.extra_args = kwds self.ensure_hash() @property def parameters(self): - """tuple: The parameters loaded from the input file. + """ + tuple: The parameters loaded from the input file. Same as `table.fieldnames`. """ return self.table.fieldnames def ensure_hash(self): - """Ensure that there is a correctly populated template_hash. + """ + Ensure that there is a correctly populated template_hash. Check for a correctly populated template_hash and create if it doesn't already exist. """ fields = self.table.fieldnames - if 'template_hash' in fields: + if "template_hash" in fields: return # The fields to use in making a template hash - hash_fields = ['mass1', 'mass2', 'inclination', - 'spin1x', 'spin1y', 'spin1z', - 'spin2x', 'spin2y', 'spin2z',] + hash_fields = [ + "mass1", + "mass2", + "inclination", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + ] fields = [f for f in hash_fields if f in fields] - template_hash = np.array([tuple_to_hash(v) for v in zip(*[self.table[p] - for p in fields])]) + template_hash = np.array( + [tuple_to_hash(v) for v in zip(*[self.table[p] for p in fields])] + ) if not np.unique(template_hash).size == template_hash.size: - raise RuntimeError("Some template hashes clash. This should not " - "happen.") - self.table = self.table.add_fields(template_hash, 'template_hash') - - def write_to_hdf(self, filename, start_index=None, stop_index=None, - force=False, skip_fields=None, - write_compressed_waveforms=True): - """Writes self to the given hdf file. + raise RuntimeError("Some template hashes clash. This should not happen.") + self.table = self.table.add_fields(template_hash, "template_hash") + + def write_to_hdf( + self, + filename, + start_index=None, + stop_index=None, + force=False, + skip_fields=None, + write_compressed_waveforms=True, + ): + """ + Writes self to the given hdf file. Parameters ---------- @@ -403,56 +442,60 @@ def write_to_hdf(self, filename, start_index=None, stop_index=None, ------- pycbc.io.HFile The file handler to the output hdf file (left open). + """ - if not filename.endswith(('.hdf', '.h5', '.hdf5')): + if not filename.endswith((".hdf", ".h5", ".hdf5")): raise ValueError("Unrecoginized file extension") if os.path.exists(filename) and not force: - raise IOError("File %s already exists" %(filename)) - f = pycbc.io.HFile(filename, 'w') + raise OSError("File %s already exists" % (filename)) + f = pycbc.io.HFile(filename, "w") parameters = self.parameters if skip_fields is not None: if not isinstance(skip_fields, list): skip_fields = [skip_fields] parameters = [p for p in parameters if p not in skip_fields] # save the parameters - f.attrs['parameters'] = parameters + f.attrs["parameters"] = parameters write_tbl = self.table[start_index:stop_index] for p in parameters: f[p] = write_tbl[p] if write_compressed_waveforms and self.has_compressed_waveforms: for tmplt_hash in write_tbl.template_hash: - compressed_waveform = pycbc.waveform.compress.CompressedWaveform.from_hdf( - self.filehandler, tmplt_hash, - load_now=True) + compressed_waveform = ( + pycbc.waveform.compress.CompressedWaveform.from_hdf( + self.filehandler, tmplt_hash, load_now=True + ) + ) compressed_waveform.write_to_hdf(f, tmplt_hash) return f def end_frequency(self, index): - """ Return the end frequency of the waveform at the given index value - """ - if hasattr(self.table[index], 'f_final'): + """Return the end frequency of the waveform at the given index value""" + if hasattr(self.table[index], "f_final"): return self.table[index].f_final return pycbc.waveform.get_waveform_end_frequency( - self.table[index], - approximant=self.approximant(index), - **self.extra_args) + self.table[index], approximant=self.approximant(index), **self.extra_args + ) def parse_approximant(self, approximant): - """Parses the given approximant argument, returning the approximant to + """ + Parses the given approximant argument, returning the approximant to use for each template in self. This is done by calling `parse_approximant_arg` using self's table as the array; see that - function for more details.""" + function for more details. + """ return parse_approximant_arg(approximant, self.table) def approximant(self, index): - """ Return the name of the approximant ot use at the given index - """ - if 'approximant' not in self.table.fieldnames: - raise ValueError("approximant not found in input file and no " - "approximant was specified on initialization") + """Return the name of the approximant ot use at the given index""" + if "approximant" not in self.table.fieldnames: + raise ValueError( + "approximant not found in input file and no " + "approximant was specified on initialization" + ) apx = self.table["approximant"][index] - if hasattr(apx, 'decode'): + if hasattr(apx, "decode"): apx = apx.decode() return apx @@ -461,16 +504,18 @@ def __len__(self): def template_thinning(self, inj_filter_rejector): """Remove templates from bank that are far from all injections.""" - if not inj_filter_rejector.enabled or \ - inj_filter_rejector.chirp_time_window is None: + if ( + not inj_filter_rejector.enabled + or inj_filter_rejector.chirp_time_window is None + ): # Do nothing! return injection_parameters = inj_filter_rejector.injection_params.table fref = inj_filter_rejector.f_lower threshold = inj_filter_rejector.chirp_time_window - m1= self.table['mass1'] - m2= self.table['mass2'] + m1 = self.table["mass1"] + m2 = self.table["mass2"] tau0_temp, _ = pycbc.pnutils.mass1_mass2_to_tau0_tau3(m1, m2, fref) indices = [] @@ -478,50 +523,59 @@ def template_thinning(self, inj_filter_rejector): tau0_temp = tau0_temp[sort] for inj in injection_parameters: - tau0_inj, _ = \ - pycbc.pnutils.mass1_mass2_to_tau0_tau3(inj.mass1, inj.mass2, - fref) + tau0_inj, _ = pycbc.pnutils.mass1_mass2_to_tau0_tau3( + inj.mass1, inj.mass2, fref + ) lid = np.searchsorted(tau0_temp, tau0_inj - threshold) rid = np.searchsorted(tau0_temp, tau0_inj + threshold) inj_indices = sort[lid:rid] indices.append(inj_indices) indices_combined = np.concatenate(indices) - indices_unique= np.unique(indices_combined) + indices_unique = np.unique(indices_combined) self.table = self.table[indices_unique] def ensure_standard_filter_columns(self, low_frequency_cutoff=None): - """ Initialize FilterBank common fields + """ + Initialize FilterBank common fields Parameters ---------- low_frequency_cutoff: {float, None}, Optional A low frequency cutoff which overrides any given within the template bank file. - """ + """ # Make sure we have a template duration field - if not hasattr(self.table, 'template_duration'): - self.table = self.table.add_fields(np.zeros(len(self.table), - dtype=np.float32), 'template_duration') + if not hasattr(self.table, "template_duration"): + self.table = self.table.add_fields( + np.zeros(len(self.table), dtype=np.float32), "template_duration" + ) # Make sure we have a f_lower field if low_frequency_cutoff is not None: - if not hasattr(self.table, 'f_lower'): + if not hasattr(self.table, "f_lower"): vec = np.zeros(len(self.table), dtype=np.float32) - self.table = self.table.add_fields(vec, 'f_lower') - self.table['f_lower'][:] = low_frequency_cutoff + self.table = self.table.add_fields(vec, "f_lower") + self.table["f_lower"][:] = low_frequency_cutoff - self.min_f_lower = min(self.table['f_lower']) - if self.f_lower is None and self.min_f_lower == 0.: - raise ValueError('Invalid low-frequency cutoff settings') + self.min_f_lower = min(self.table["f_lower"]) + if self.f_lower is None and self.min_f_lower == 0.0: + raise ValueError("Invalid low-frequency cutoff settings") class LiveFilterBank(TemplateBank): - def __init__(self, filename, sample_rate, minimum_buffer, - approximant=None, increment=8, parameters=None, - low_frequency_cutoff=None, - **kwds): + def __init__( + self, + filename, + sample_rate, + minimum_buffer, + approximant=None, + increment=8, + parameters=None, + low_frequency_cutoff=None, + **kwds, + ): self.increment = increment self.filename = filename @@ -529,17 +583,19 @@ def __init__(self, filename, sample_rate, minimum_buffer, self.minimum_buffer = minimum_buffer self.f_lower = low_frequency_cutoff - super(LiveFilterBank, self).__init__(filename, approximant=approximant, - parameters=parameters, **kwds) + super().__init__( + filename, approximant=approximant, parameters=parameters, **kwds + ) self.ensure_standard_filter_columns(low_frequency_cutoff=low_frequency_cutoff) self.param_lookup = {} for i, p in enumerate(self.table): - key = (p.mass1, p.mass2, p.spin1z, p.spin2z) - assert(key not in self.param_lookup) # Uh, oh, template confusion! + key = (p.mass1, p.mass2, p.spin1z, p.spin2z) + assert key not in self.param_lookup # Uh, oh, template confusion! self.param_lookup[key] = i def round_up(self, num): - """Determine the length to use for this waveform by rounding. + """ + Determine the length to use for this waveform by rounding. Parameters ---------- @@ -552,6 +608,7 @@ def round_up(self, num): The rounded size to use for the waveform buffer in samples. This is calculated using an internal `increment` attribute, which determines the discreteness of the rounding. + """ inc = self.increment size = np.ceil(num / self.sample_rate / inc) * self.sample_rate * inc @@ -563,7 +620,8 @@ def getslice(self, sindex): return instance def id_from_param(self, param_tuple): - """Get the index of this template based on its param tuple + """ + Get the index of this template based on its param tuple Parameters ---------- @@ -571,9 +629,10 @@ def id_from_param(self, param_tuple): Tuple of the parameters which uniquely identify this template Returns - -------- + ------- index : int The ordered index that this template has in the template bank. + """ return self.param_lookup[param_tuple] @@ -584,7 +643,8 @@ def __getitem__(self, index): return self.get_template(index) def freq_resolution_for_template(self, index): - """Compute the correct resolution for a frequency series that contains + """ + Compute the correct resolution for a frequency series that contains a given template in the bank. """ from pycbc.waveform.waveform import props @@ -592,21 +652,20 @@ def freq_resolution_for_template(self, index): time_duration = self.minimum_buffer time_duration += 0.5 params = props(self.table[index]) - params.pop('approximant') + params.pop("approximant") approximant = self.approximant(index) waveform_duration = pycbc.waveform.get_waveform_filter_length_in_time( approximant, **params ) if waveform_duration is None: - raise RuntimeError( - 'Template waveform {approximant} not recognized!' - ) + raise RuntimeError("Template waveform {approximant} not recognized!") time_duration += waveform_duration td_samples = self.round_up(time_duration * self.sample_rate) return self.sample_rate / float(td_samples) def get_template(self, index, delta_f=None): - """Calculate and return the frequency-domain waveform for the template + """ + Calculate and return the frequency-domain waveform for the template with the given index. The frequency resolution can optionally be given. Parameters @@ -621,6 +680,7 @@ def get_template(self, index, delta_f=None): ------- htilde: FrequencySeries Template waveform in the frequency domain. + """ approximant = self.approximant(index) f_end = self.end_frequency(index) @@ -630,7 +690,7 @@ def get_template(self, index, delta_f=None): delta_f = self.freq_resolution_for_template(index) flen = round(self.sample_rate / (2 * delta_f) + 1) - assert flen & 1, f'flen should be odd, but got {flen}' + assert flen & 1, f"flen should be odd, but got {flen}" if f_end is None or f_end >= (flen * delta_f): f_end = (flen - 1) * delta_f @@ -640,27 +700,33 @@ def get_template(self, index, delta_f=None): approximant, 1.0 / delta_f, index, - flow + flow, ) # Get the waveform filter distance = 1.0 / DYN_RANGE_FAC htilde = pycbc.waveform.get_waveform_filter( - zeros(flen, dtype=np.complex64), self.table[index], - approximant=approximant, f_lower=flow, f_final=f_end, - delta_f=delta_f, delta_t=1.0 / self.sample_rate, distance=distance, - **self.extra_args) + zeros(flen, dtype=np.complex64), + self.table[index], + approximant=approximant, + f_lower=flow, + f_final=f_end, + delta_f=delta_f, + delta_t=1.0 / self.sample_rate, + distance=distance, + **self.extra_args, + ) # If available, record the total duration (which may # include ringdown) and the duration up to merger since they will be # erased by the type conversion below. ttotal = template_duration = -1 time_offset = None - if hasattr(htilde, 'length_in_time'): + if hasattr(htilde, "length_in_time"): ttotal = htilde.length_in_time - if hasattr(htilde, 'chirp_length'): + if hasattr(htilde, "chirp_length"): template_duration = htilde.chirp_length - if hasattr(htilde, 'time_offset'): + if hasattr(htilde, "time_offset"): time_offset = htilde.time_offset self.table[index].template_duration = template_duration @@ -681,111 +747,130 @@ def get_template(self, index, delta_f=None): # Add sigmasq as a method of this instance htilde.sigmasq = types.MethodType(sigma_cached, htilde) - htilde.id = self.id_from_param((htilde.params.mass1, - htilde.params.mass2, - htilde.params.spin1z, - htilde.params.spin2z)) + htilde.id = self.id_from_param( + ( + htilde.params.mass1, + htilde.params.mass2, + htilde.params.spin1z, + htilde.params.spin2z, + ) + ) return htilde class FilterBank(TemplateBank): - def __init__(self, filename, filter_length, delta_f, dtype, - out=None, max_template_length=None, - approximant=None, parameters=None, - enable_compressed_waveforms=True, - low_frequency_cutoff=None, - waveform_decompression_method=None, - **kwds): + def __init__( + self, + filename, + filter_length, + delta_f, + dtype, + out=None, + max_template_length=None, + approximant=None, + parameters=None, + enable_compressed_waveforms=True, + low_frequency_cutoff=None, + waveform_decompression_method=None, + **kwds, + ): self.out = out self.dtype = dtype self.f_lower = low_frequency_cutoff self.filename = filename self.delta_f = delta_f - self.N = (filter_length - 1 ) * 2 + self.N = (filter_length - 1) * 2 self.delta_t = 1.0 / (self.N * self.delta_f) self.filter_length = filter_length self.max_template_length = max_template_length self.enable_compressed_waveforms = enable_compressed_waveforms self.waveform_decompression_method = waveform_decompression_method - super(FilterBank, self).__init__(filename, approximant=approximant, - parameters=parameters, **kwds) + super().__init__( + filename, approximant=approximant, parameters=parameters, **kwds + ) self.ensure_standard_filter_columns(low_frequency_cutoff=low_frequency_cutoff) - def get_decompressed_waveform(self, tempout, index, f_lower=None, - approximant=None, df=None): - """Returns a frequency domain decompressed waveform for the template + def get_decompressed_waveform( + self, tempout, index, f_lower=None, approximant=None, df=None + ): + """ + Returns a frequency domain decompressed waveform for the template in the bank corresponding to the index taken in as an argument. The decompressed waveform is obtained by interpolating in frequency space, the amplitude and phase points for the compressed template that are - read in from the bank.""" - - from pycbc.waveform.waveform import props + read in from the bank. + """ from pycbc.waveform import get_waveform_filter_length_in_time + from pycbc.waveform.waveform import props # Get the template hash corresponding to the template index taken in as argument tmplt_hash = self.table.template_hash[index] # Read the compressed waveform from the bank file compressed_waveform = pycbc.waveform.compress.CompressedWaveform.from_hdf( - self.filehandler, tmplt_hash, - load_now=True) + self.filehandler, tmplt_hash, load_now=True + ) # Get the interpolation method to be used to decompress the waveform - if self.waveform_decompression_method is not None : + if self.waveform_decompression_method is not None: decompression_method = self.waveform_decompression_method - else : + else: decompression_method = compressed_waveform.interpolation logging.info("Decompressing waveform using %s", decompression_method) - if df is not None : + if df is not None: delta_f = df - else : + else: delta_f = self.delta_f # Create memory space for writing the decompressed waveform - decomp_scratch = FrequencySeries(tempout[0:self.filter_length], delta_f=delta_f, copy=False) + decomp_scratch = FrequencySeries( + tempout[0 : self.filter_length], delta_f=delta_f, copy=False + ) # Get the decompressed waveform - hdecomp = compressed_waveform.decompress(out=decomp_scratch, f_lower=f_lower, interpolation=decompression_method) + hdecomp = compressed_waveform.decompress( + out=decomp_scratch, f_lower=f_lower, interpolation=decompression_method + ) p = props(self.table[index]) - p.pop('approximant') + p.pop("approximant") try: tmpltdur = self.table[index].template_duration except AttributeError: tmpltdur = None - if tmpltdur is None or tmpltdur==0.0 : + if tmpltdur is None or tmpltdur == 0.0: tmpltdur = get_waveform_filter_length_in_time(approximant, **p) hdecomp.chirp_length = tmpltdur hdecomp.length_in_time = hdecomp.chirp_length return hdecomp - def generate_with_delta_f_and_max_freq(self, t_num, max_freq, delta_f, - low_frequency_cutoff=None, - cached_mem=None): + def generate_with_delta_f_and_max_freq( + self, t_num, max_freq, delta_f, low_frequency_cutoff=None, cached_mem=None + ): """Generate the template with index t_num using custom length.""" approximant = self.approximant(t_num) # Don't want to use INTERP waveforms in here - if approximant.endswith('_INTERP'): - approximant = approximant.replace('_INTERP', '') + if approximant.endswith("_INTERP"): + approximant = approximant.replace("_INTERP", "") # Using SPAtmplt here is bad as the stored cbrt and logv get # recalculated as we change delta_f values. Fall back to TaylorF2 # in lalsimulation. - if approximant == 'SPAtmplt': - approximant = 'TaylorF2' + if approximant == "SPAtmplt": + approximant = "TaylorF2" if cached_mem is None: wav_len = int(max_freq / delta_f) + 1 cached_mem = zeros(wav_len, dtype=np.complex64) full_calculate_waveform = True - if (self.has_compressed_waveforms and self.enable_compressed_waveforms): + if self.has_compressed_waveforms and self.enable_compressed_waveforms: try: htilde = self.get_decompressed_waveform( tempout, index, f_lower=low_frequency_cutoff, approximant=approximant, - df=None + df=None, ) full_calculate_waveform = False except KeyError: @@ -800,9 +885,14 @@ def generate_with_delta_f_and_max_freq(self, t_num, max_freq, delta_f, if full_calculate_waveform: htilde = pycbc.waveform.get_waveform_filter( - cached_mem, self.table[t_num], approximant=approximant, - f_lower=low_frequency_cutoff, f_final=max_freq, delta_f=delta_f, - distance=1./DYN_RANGE_FAC, delta_t=1./(2.*max_freq) + cached_mem, + self.table[t_num], + approximant=approximant, + f_lower=low_frequency_cutoff, + f_final=max_freq, + delta_f=delta_f, + distance=1.0 / DYN_RANGE_FAC, + delta_t=1.0 / (2.0 * max_freq), ) return htilde @@ -817,30 +907,25 @@ def __getitem__(self, index): approximant = self.approximant(index) f_end = self.end_frequency(index) if f_end is None or f_end >= (self.filter_length * self.delta_f): - f_end = (self.filter_length-1) * self.delta_f + f_end = (self.filter_length - 1) * self.delta_f # Find the start frequency, if variable - f_low = find_variable_start_frequency(approximant, - self.table[index], - self.f_lower, - self.max_template_length) - logging.info('%s: generating %s from %s Hz' % (index, approximant, f_low)) + f_low = find_variable_start_frequency( + approximant, self.table[index], self.f_lower, self.max_template_length + ) + logging.info("%s: generating %s from %s Hz" % (index, approximant, f_low)) # Clear the storage memory - poke = tempout.data # pylint:disable=unused-variable + poke = tempout.data # pylint:disable=unused-variable tempout.clear() # Get the waveform filter distance = 1.0 / DYN_RANGE_FAC full_calculate_waveform = True - if (self.has_compressed_waveforms and self.enable_compressed_waveforms): + if self.has_compressed_waveforms and self.enable_compressed_waveforms: try: htilde = self.get_decompressed_waveform( - tempout, - index, - f_lower=f_low, - approximant=approximant, - df=None + tempout, index, f_lower=f_low, approximant=approximant, df=None ) full_calculate_waveform = False except KeyError: @@ -855,9 +940,14 @@ def __getitem__(self, index): if full_calculate_waveform: htilde = pycbc.waveform.get_waveform_filter( - tempout[0:self.filter_length], self.table[index], - approximant=approximant, f_lower=f_low, f_final=f_end, - delta_f=self.delta_f, delta_t=self.delta_t, distance=distance, + tempout[0 : self.filter_length], + self.table[index], + approximant=approximant, + f_lower=f_low, + f_final=f_end, + delta_f=self.delta_f, + delta_t=self.delta_t, + distance=distance, **self.extra_args, ) @@ -865,9 +955,9 @@ def __getitem__(self, index): # include ringdown) and the duration up to merger since they will be # erased by the type conversion below. ttotal = template_duration = None - if hasattr(htilde, 'length_in_time'): + if hasattr(htilde, "length_in_time"): ttotal = htilde.length_in_time - if hasattr(htilde, 'chirp_length'): + if hasattr(htilde, "chirp_length"): template_duration = htilde.chirp_length self.table[index].template_duration = template_duration @@ -888,43 +978,54 @@ def __getitem__(self, index): return htilde -def find_variable_start_frequency(approximant, parameters, f_start, max_length, - delta_f = 1): - """ Find a frequency value above the starting frequency that results in a +def find_variable_start_frequency( + approximant, parameters, f_start, max_length, delta_f=1 +): + """ + Find a frequency value above the starting frequency that results in a waveform shorter than max_length. """ - if (f_start is None): + if f_start is None: f = parameters.f_lower - elif (max_length is not None): + elif max_length is not None: l = max_length + 1 f = f_start - delta_f while l > max_length: f += delta_f - l = pycbc.waveform.get_waveform_filter_length_in_time(approximant, - parameters, f_lower=f) - else : + l = pycbc.waveform.get_waveform_filter_length_in_time( + approximant, parameters, f_lower=f + ) + else: f = f_start return f class FilterBankSkyMax(TemplateBank): - def __init__(self, filename, filter_length, delta_f, - dtype, out_plus=None, out_cross=None, - max_template_length=None, parameters=None, - low_frequency_cutoff=None, **kwds): + def __init__( + self, + filename, + filter_length, + delta_f, + dtype, + out_plus=None, + out_cross=None, + max_template_length=None, + parameters=None, + low_frequency_cutoff=None, + **kwds, + ): self.out_plus = out_plus self.out_cross = out_cross self.dtype = dtype self.f_lower = low_frequency_cutoff self.filename = filename self.delta_f = delta_f - self.N = (filter_length - 1 ) * 2 + self.N = (filter_length - 1) * 2 self.delta_t = 1.0 / (self.N * self.delta_f) self.filter_length = filter_length self.max_template_length = max_template_length - super(FilterBankSkyMax, self).__init__(filename, parameters=parameters, - **kwds) + super().__init__(filename, parameters=parameters, **kwds) self.ensure_standard_filter_columns(low_frequency_cutoff=low_frequency_cutoff) @@ -944,18 +1045,17 @@ def __getitem__(self, index): # Get the end of the waveform if applicable (only for SPAtmplt atm) f_end = self.end_frequency(index) if f_end is None or f_end >= (self.filter_length * self.delta_f): - f_end = (self.filter_length-1) * self.delta_f + f_end = (self.filter_length - 1) * self.delta_f # Find the start frequency, if variable - f_low = find_variable_start_frequency(approximant, - self.table[index], - self.f_lower, - self.max_template_length) - logging.info('%s: generating %s from %s Hz', index, approximant, f_low) + f_low = find_variable_start_frequency( + approximant, self.table[index], self.f_lower, self.max_template_length + ) + logging.info("%s: generating %s from %s Hz", index, approximant, f_low) # What does this do??? - poke1 = tempoutplus.data # pylint:disable=unused-variable - poke2 = tempoutcross.data # pylint:disable=unused-variable + poke1 = tempoutplus.data # pylint:disable=unused-variable + poke2 = tempoutcross.data # pylint:disable=unused-variable # Clear the storage memory tempoutplus.clear() @@ -964,13 +1064,19 @@ def __getitem__(self, index): # Get the waveform filter distance = 1.0 / DYN_RANGE_FAC hplus, hcross = pycbc.waveform.get_two_pol_waveform_filter( - tempoutplus[0:self.filter_length], - tempoutcross[0:self.filter_length], self.table[index], - approximant=approximant, f_lower=f_low, - f_final=f_end, delta_f=self.delta_f, delta_t=self.delta_t, - distance=distance, **self.extra_args) + tempoutplus[0 : self.filter_length], + tempoutcross[0 : self.filter_length], + self.table[index], + approximant=approximant, + f_lower=f_low, + f_final=f_end, + delta_f=self.delta_f, + delta_t=self.delta_t, + distance=distance, + **self.extra_args, + ) - if hasattr(hplus, 'chirp_length') and hplus.chirp_length is not None: + if hasattr(hplus, "chirp_length") and hplus.chirp_length is not None: self.table[index].template_duration = hplus.chirp_length hplus = hplus.astype(self.dtype) @@ -997,7 +1103,15 @@ def __getitem__(self, index): return hplus, hcross -__all__ = ('sigma_cached', 'boolargs_from_apprxstr', 'add_approximant_arg', - 'parse_approximant_arg', 'tuple_to_hash', 'TemplateBank', - 'LiveFilterBank', 'FilterBank', 'find_variable_start_frequency', - 'FilterBankSkyMax') +__all__ = ( + "FilterBank", + "FilterBankSkyMax", + "LiveFilterBank", + "TemplateBank", + "add_approximant_arg", + "boolargs_from_apprxstr", + "find_variable_start_frequency", + "parse_approximant_arg", + "sigma_cached", + "tuple_to_hash", +) diff --git a/pycbc/waveform/compress.py b/pycbc/waveform/compress.py index 448a6863b71..3ca3451b4a5 100644 --- a/pycbc/waveform/compress.py +++ b/pycbc/waveform/compress.py @@ -20,21 +20,29 @@ # # ============================================================================= # -""" Utilities for handling frequency compressed an unequally spaced frequency +""" +Utilities for handling frequency compressed an unequally spaced frequency domain waveforms. """ -import numpy, logging, h5py, time + +import logging +import time + +import h5py +import numpy from scipy import interpolate from pycbc import filter +from pycbc.constants import MTSUN_SI +from pycbc.io.hdf import HFile +from pycbc.scheme import schemed from pycbc.types import FrequencySeries, real_same_precision_as from pycbc.waveform import utils -from pycbc.scheme import schemed -from pycbc.io.hdf import HFile -from pycbc.constants import MTSUN_SI + def rough_time_estimate(m1, m2, flow, fudge_length=1.1, fudge_min=0.02): - """ A very rough estimate of the duration of the waveform. + """ + A very rough estimate of the duration of the waveform. An estimate of the waveform duration starting from flow. This is intended to be fast but not necessarily accurate. It should be an overestimate of @@ -60,20 +68,19 @@ def rough_time_estimate(m1, m2, flow, fudge_length=1.1, fudge_min=0.02): ------- time: float Time from flow untill the end of the waveform + """ m = m1 + m2 msun = m * MTSUN_SI - t = 5.0 / 256.0 * m * m * msun / (m1 * m2) / \ - (numpy.pi * msun * flow) ** (8.0 / 3.0) + t = 5.0 / 256.0 * m * m * msun / (m1 * m2) / (numpy.pi * msun * flow) ** (8.0 / 3.0) # fudge factoriness - return .022 if t < 0 else (t + fudge_min) * fudge_length + return 0.022 if t < 0 else (t + fudge_min) * fudge_length + -def mchirp_compression(m1, m2, fmin, fmax, - min_seglen=0.02, - df_multiple=None, - scale=1): - """Return the frequencies needed to compress a waveform with the given +def mchirp_compression(m1, m2, fmin, fmax, min_seglen=0.02, df_multiple=None, scale=1): + """ + Return the frequencies needed to compress a waveform with the given chirp mass. This is based on the estimate in rough_time_estimate. Parameters @@ -99,12 +106,13 @@ def mchirp_compression(m1, m2, fmin, fmax, ------- array The frequencies at which to evaluate the compressed waveform. + """ sample_points = [] f = fmin while f < fmax: if df_multiple is not None: - f = int(f/df_multiple)*df_multiple + f = int(f / df_multiple) * df_multiple sample_points.append(f) f += 1.0 / rough_time_estimate(m1, m2, f, fudge_min=min_seglen) * scale # add the last point @@ -112,9 +120,12 @@ def mchirp_compression(m1, m2, fmin, fmax, sample_points.append(fmax) return numpy.array(sample_points) -def spa_compression(htilde, fmin, fmax, min_seglen=0.02, - sample_frequencies=None, scale=1): - """Returns the frequencies needed to compress the given frequency domain + +def spa_compression( + htilde, fmin, fmax, min_seglen=0.02, sample_frequencies=None, scale=1 +): + """ + Returns the frequencies needed to compress the given frequency domain waveform. This is done by estimating t(f) of the waveform using the stationary phase approximation. @@ -139,50 +150,69 @@ def spa_compression(htilde, fmin, fmax, min_seglen=0.02, ------- array The frequencies at which to evaluate the compressed waveform. + """ if sample_frequencies is None: sample_frequencies = htilde.sample_frequencies.numpy() - kmin = int(fmin/htilde.delta_f) - kmax = int(fmax/htilde.delta_f) - tf = abs(utils.time_from_frequencyseries(htilde, - sample_frequencies=sample_frequencies).data[kmin:kmax]) + kmin = int(fmin / htilde.delta_f) + kmax = int(fmax / htilde.delta_f) + tf = abs( + utils.time_from_frequencyseries( + htilde, sample_frequencies=sample_frequencies + ).data[kmin:kmax] + ) sample_frequencies = sample_frequencies[kmin:kmax] sample_points = [] f = fmin while f < fmax: - f = int(f/htilde.delta_f)*htilde.delta_f + f = int(f / htilde.delta_f) * htilde.delta_f sample_points.append(f) jj = numpy.searchsorted(sample_frequencies, f) - f += 1./(tf[jj:].max()+min_seglen) * scale + f += 1.0 / (tf[jj:].max() + min_seglen) * scale # add the last point if sample_points[-1] < fmax: sample_points.append(fmax) return numpy.array(sample_points) -compression_algorithms = { - 'mchirp': mchirp_compression, - 'spa': spa_compression - } + +compression_algorithms = {"mchirp": mchirp_compression, "spa": spa_compression} + def _vecdiff(htilde, hinterp, fmin, fmax, psd=None): - return 1 - abs(filter.overlap_cplx(htilde, hinterp, - low_frequency_cutoff=fmin, - high_frequency_cutoff=fmax, - psd=psd)) + return 1 - abs( + filter.overlap_cplx( + htilde, + hinterp, + low_frequency_cutoff=fmin, + high_frequency_cutoff=fmax, + psd=psd, + ) + ) + def vecdiff(htilde, hinterp, sample_points, psd=None): - """Computes a statistic indicating between which sample points a waveform + """ + Computes a statistic indicating between which sample points a waveform and the interpolated waveform differ the most. """ - vecdiffs = numpy.zeros(sample_points.size-1, dtype=float) - for kk,thisf in enumerate(sample_points[:-1]): - nextf = sample_points[kk+1] + vecdiffs = numpy.zeros(sample_points.size - 1, dtype=float) + for kk, thisf in enumerate(sample_points[:-1]): + nextf = sample_points[kk + 1] vecdiffs[kk] = abs(_vecdiff(htilde, hinterp, thisf, nextf, psd=psd)) return vecdiffs -def compress_waveform(htilde, sample_points, tolerance, interpolation, - precision, decomp_scratch=None, psd=None): - """Retrieves the amplitude and phase at the desired sample points, and adds + +def compress_waveform( + htilde, + sample_points, + tolerance, + interpolation, + precision, + decomp_scratch=None, + psd=None, +): + """ + Retrieves the amplitude and phase at the desired sample points, and adds frequency points in order to ensure that the interpolated waveform has a mismatch with the full waveform that is <= the desired tolerance. The mismatch is computed by finding 1-overlap between `htilde` and the @@ -226,6 +256,7 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, ------- CompressedWaveform The compressed waveform data; see `CompressedWaveform` for details. + """ fmin = sample_points.min() df = htilde.delta_f @@ -238,14 +269,20 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, comp_amp = amp.take(sample_index) comp_phase = phase.take(sample_index) outdf = df if decomp_scratch is None else None - + t1 = time.time() - hdecomp = fd_decompress(comp_amp, comp_phase, sample_points, - out=decomp_scratch, df=outdf, f_lower=fmin, - interpolation=interpolation) + hdecomp = fd_decompress( + comp_amp, + comp_phase, + sample_points, + out=decomp_scratch, + df=outdf, + f_lower=fmin, + interpolation=interpolation, + ) # This will be overwritten in the loop to store the final pass time htime = time.time() - t1 - + kmax = min(len(htilde), len(hdecomp)) htilde = htilde[:kmax] hdecomp = hdecomp[:kmax] @@ -254,23 +291,25 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, s2 = filter.sigma(htilde, psd=psd, low_frequency_cutoff=fmin) if psd is not None: - htilde2 = htilde / (psd[:len(htilde)] * s2) + htilde2 = htilde / (psd[: len(htilde)] * s2) else: htilde2 = htilde / s2 # Do a first test to see if we are done - mismatch = 1. - abs(filter.overlap_cplx(hdecomp / s1, htilde2, - low_frequency_cutoff=fmin, normalized=False)) + mismatch = 1.0 - abs( + filter.overlap_cplx( + hdecomp / s1, htilde2, low_frequency_cutoff=fmin, normalized=False + ) + ) if mismatch > tolerance: # Calculate the overlap errors within each frequency bins. # We use this to determine where to add more interpolation points vecdiffs = vecdiff(htilde, hdecomp, sample_points, psd=psd) - # We will find where in the frequency series the interpolated waveform - # has the smallest overlap with the full waveform, + # has the smallest overlap with the full waveform, # We try to add a new interpolation point in every frequency bin - # that fails this check. Continue untill the overall reconstruction + # that fails this check. Continue untill the overall reconstruction # waveform meets our mismatch target with the origianl waveform added_points = [] iteration_count = 0 @@ -279,12 +318,12 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, while mismatch > tolerance: iteration_start_time = time.time() iteration_count += 1 - + # Pick the worst bins num_bad = (vecdiffs > tolerance).sum() vsort = vecdiffs.argsort()[::-1] - - # This add fraction of the bad segments, up to a maximum + + # This add fraction of the bad segments, up to a maximum # If there are no bad segments, we still try to add the first single # one (can be large numerical error in the veddiff calculation, so # rounding cause all to be below the tolerance yet thte full fails). @@ -297,14 +336,16 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, new_addidxs = [] for minpt in selected_segments: # Calculate midpoint using indices to avoid float drift issues - add_freq = (sample_points[minpt] + sample_points[minpt+1]) / 2.0 + add_freq = (sample_points[minpt] + sample_points[minpt + 1]) / 2.0 addidx = int(add_freq / df) if addidx not in sample_index and addidx not in new_addidxs: new_addidxs.append(addidx) # Don't propose points within a sample of existing ones new_addidxs = numpy.array(new_addidxs) - valid = ~numpy.any(abs(new_addidxs[:, None] - numpy.array(added_points)) <= 2, axis=1) + valid = ~numpy.any( + abs(new_addidxs[:, None] - numpy.array(added_points)) <= 2, axis=1 + ) new_addidxs = list(new_addidxs[valid]) # --- 3. Update and Sort --- @@ -314,29 +355,39 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, sample_points = (sample_index * df).astype(real_same_precision_as(htilde)) comp_amp = amp.take(sample_index) comp_phase = phase.take(sample_index) - + # --- 4. Decompress --- t1 = time.time() - hdecomp = fd_decompress(comp_amp, comp_phase, sample_points, - out=decomp_scratch, df=outdf, - f_lower=fmin, interpolation=interpolation) + hdecomp = fd_decompress( + comp_amp, + comp_phase, + sample_points, + out=decomp_scratch, + df=outdf, + f_lower=fmin, + interpolation=interpolation, + ) htime = time.time() - t1 s1 = filter.sigma(hdecomp, psd=psd, low_frequency_cutoff=fmin) # --- 5. Re-evaluate global mismatch --- - mismatch = 1. - abs(filter.overlap_cplx(hdecomp / s1, htilde2, - low_frequency_cutoff=fmin, - normalized=False)) + mismatch = 1.0 - abs( + filter.overlap_cplx( + hdecomp / s1, htilde2, low_frequency_cutoff=fmin, normalized=False + ) + ) - o = filter.overlap_cplx(hdecomp / s1, htilde2, - low_frequency_cutoff=fmin, - normalized=False) + o = filter.overlap_cplx( + hdecomp / s1, htilde2, low_frequency_cutoff=fmin, normalized=False + ) if mismatch <= tolerance: - mismatch = 1. - abs(filter.overlap_cplx(hdecomp / s1, htilde2, - low_frequency_cutoff=fmin, - normalized=False)) + mismatch = 1.0 - abs( + filter.overlap_cplx( + hdecomp / s1, htilde2, low_frequency_cutoff=fmin, normalized=False + ) + ) else: # Calculate the overlap errors within each frequency bins. # We use this to determine where to add more interpolation points @@ -347,8 +398,11 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, # --- 6. Iteration Logging --- logging.debug( "Iter %i: mismatch %.6f, added %i points (total %i), iter time %.2f ms", - iteration_count, mismatch, len(new_addidxs), len(sample_points), - (time.time() - iteration_start_time) * 1000 + iteration_count, + mismatch, + len(new_addidxs), + len(sample_points), + (time.time() - iteration_start_time) * 1000, ) if len(new_addidxs) == 0: @@ -357,38 +411,46 @@ def compress_waveform(htilde, sample_points, tolerance, interpolation, # Cast compression_factor to float to avoid HDF5 TypeErrors compression_factor = float(len(htilde)) / float(len(sample_points)) - + logging.info( "mismatch: %f, N points: %i (%i added), compression:%.3e, final decomp time %.2f ms", - mismatch, len(comp_amp), len(added_points), compression_factor, htime * 1000, + mismatch, + len(comp_amp), + len(added_points), + compression_factor, + htime * 1000, ) - return CompressedWaveform(sample_points, comp_amp, comp_phase, - interpolation=interpolation, - tolerance=tolerance, mismatch=mismatch, - precision=precision, - compression_factor=compression_factor) + return CompressedWaveform( + sample_points, + comp_amp, + comp_phase, + interpolation=interpolation, + tolerance=tolerance, + mismatch=mismatch, + precision=precision, + compression_factor=compression_factor, + ) + + _precision_map = { - 'float32': 'single', - 'float64': 'double', - 'complex64': 'single', - 'complex128': 'double' + "float32": "single", + "float64": "double", + "complex64": "single", + "complex128": "double", } -_complex_dtypes = { - 'single': numpy.complex64, - 'double': numpy.complex128 -} +_complex_dtypes = {"single": numpy.complex64, "double": numpy.complex128} + +_real_dtypes = {"single": numpy.float32, "double": numpy.float64} -_real_dtypes = { - 'single': numpy.float32, - 'double': numpy.float64 -} @schemed("pycbc.waveform.decompress_") -def inline_cubic_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): - """Generate a frequency-domain waveform via cubic interpolation +def inline_cubic_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): + """ + Generate a frequency-domain waveform via cubic interpolation from sampled amplitude and phase. The sample frequency locations for the amplitude and phase must be the same. This function may be less accurate than scipy's linear interpolation, but should be @@ -432,11 +494,14 @@ def inline_cubic_interp(amp, phase, sample_frequencies, output, """ return - + + @schemed("pycbc.waveform.decompress_") -def inline_quadratic_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): - """Generate a frequency-domain waveform via quadratic interpolation +def inline_quadratic_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): + """ + Generate a frequency-domain waveform via quadratic interpolation from sampled amplitude and phase. The sample frequency locations for the amplitude and phase must be the same. This function may be less accurate than scipy's linear interpolation, but should be @@ -481,10 +546,13 @@ def inline_quadratic_interp(amp, phase, sample_frequencies, output, """ return + @schemed("pycbc.waveform.decompress_") -def inline_linear_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): - """Generate a frequency-domain waveform via linear interpolation +def inline_linear_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): + """ + Generate a frequency-domain waveform via linear interpolation from sampled amplitude and phase. The sample frequency locations for the amplitude and phase must be the same. This function may be less accurate than scipy's linear interpolation, but should be @@ -529,10 +597,13 @@ def inline_linear_interp(amp, phase, sample_frequencies, output, """ return + @schemed("pycbc.waveform.decompress_") -def inline_quartic_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): - """Generate a frequency-domain waveform via quartic interpolation +def inline_quartic_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): + """ + Generate a frequency-domain waveform via quartic interpolation from sampled amplitude and phase. The sample frequency locations for the amplitude and phase must be the same. This function may be less accurate than scipy's quadratic interpolation, but should be @@ -580,9 +651,18 @@ def inline_quartic_interp(amp, phase, sample_frequencies, output, """ return -def fd_decompress(amp, phase, sample_frequencies, out=None, df=None, - f_lower=None, interpolation='inline_linear'): - """Decompresses an FD waveform using the given amplitude, phase, and the + +def fd_decompress( + amp, + phase, + sample_frequencies, + out=None, + df=None, + f_lower=None, + interpolation="inline_linear", +): + """ + Decompresses an FD waveform using the given amplitude, phase, and the frequencies at which they are sampled at. Parameters @@ -616,30 +696,34 @@ def fd_decompress(amp, phase, sample_frequencies, out=None, df=None, out : FrequencySeries If out was provided, writes to that array. Otherwise, a new FrequencySeries with the decompressed waveform. + """ precision = _precision_map[sample_frequencies.dtype.name] - if _precision_map[amp.dtype.name] != precision or \ - _precision_map[phase.dtype.name] != precision: - raise ValueError("amp, phase, and sample_points must all have the " - "same precision") + if ( + _precision_map[amp.dtype.name] != precision + or _precision_map[phase.dtype.name] != precision + ): + raise ValueError( + "amp, phase, and sample_points must all have the same precision" + ) if out is None: if df is None: raise ValueError("Either provide output memory or a df") - hlen = int(numpy.ceil(sample_frequencies.max()/df+1)) - out = FrequencySeries(numpy.zeros(hlen, - dtype=_complex_dtypes[precision]), copy=False, - delta_f=df) + hlen = int(numpy.ceil(sample_frequencies.max() / df + 1)) + out = FrequencySeries( + numpy.zeros(hlen, dtype=_complex_dtypes[precision]), copy=False, delta_f=df + ) else: # check for precision compatibility - if out.precision == 'double' and precision == 'single': + if out.precision == "double" and precision == "single": amp = amp.astype(numpy.float64) phase = phase.astype(numpy.float64) sample_frequencies = sample_frequencies.astype(numpy.float64) df = out.delta_f hlen = len(out) if f_lower is None: - imin = 0 # pylint:disable=unused-variable + imin = 0 # pylint:disable=unused-variable f_lower = sample_frequencies[0] start_index = 0 else: @@ -647,52 +731,62 @@ def fd_decompress(amp, phase, sample_frequencies, out=None, df=None, raise ValueError("f_lower is > than the maximum sample frequency") if f_lower < sample_frequencies.min(): raise ValueError("f_lower is < than the minimum sample frequency") - imin = int(numpy.searchsorted(sample_frequencies, f_lower, - side='right')) - 1 # pylint:disable=unused-variable - start_index = int(numpy.ceil(f_lower/df)) + imin = int(numpy.searchsorted(sample_frequencies, f_lower, side="right")) - 1 # pylint:disable=unused-variable + start_index = int(numpy.ceil(f_lower / df)) if start_index >= hlen: - raise ValueError('requested f_lower >= largest frequency in out') + raise ValueError("requested f_lower >= largest frequency in out") # interpolate the amplitude and the phase if interpolation == "inline_linear": # Call the scheme-dependent function - inline_linear_interp(amp, phase, sample_frequencies, out, - df, f_lower, imin, start_index) + inline_linear_interp( + amp, phase, sample_frequencies, out, df, f_lower, imin, start_index + ) elif interpolation == "inline_quadratic": # Call the scheme-dependent function - inline_quadratic_interp(amp, phase, sample_frequencies, out, - df, f_lower, imin, start_index) + inline_quadratic_interp( + amp, phase, sample_frequencies, out, df, f_lower, imin, start_index + ) elif interpolation == "inline_cubic": # Call the scheme-dependent function - inline_cubic_interp(amp, phase, sample_frequencies, out, - df, f_lower, imin, start_index) + inline_cubic_interp( + amp, phase, sample_frequencies, out, df, f_lower, imin, start_index + ) elif interpolation == "inline_quartic": # Call the scheme-dependent function - inline_quartic_interp(amp, phase, sample_frequencies, out, - df, f_lower, imin, start_index) + inline_quartic_interp( + amp, phase, sample_frequencies, out, df, f_lower, imin, start_index + ) else: # use scipy for fancier interpolation sample_frequencies = numpy.array(sample_frequencies) amp = numpy.array(amp) phase = numpy.array(phase) outfreq = out.sample_frequencies.numpy() - amp_interp = interpolate.interp1d(sample_frequencies, amp, - kind=interpolation, - bounds_error=False, - fill_value=0., - assume_sorted=True) - phase_interp = interpolate.interp1d(sample_frequencies, phase, - kind=interpolation, - bounds_error=False, - fill_value=0., - assume_sorted=True) + amp_interp = interpolate.interp1d( + sample_frequencies, + amp, + kind=interpolation, + bounds_error=False, + fill_value=0.0, + assume_sorted=True, + ) + phase_interp = interpolate.interp1d( + sample_frequencies, + phase, + kind=interpolation, + bounds_error=False, + fill_value=0.0, + assume_sorted=True, + ) A = amp_interp(outfreq) phi = phase_interp(outfreq) - out.data[:] = A*numpy.cos(phi) + (1j)*A*numpy.sin(phi) + out.data[:] = A * numpy.cos(phi) + (1j) * A * numpy.sin(phi) return out -class CompressedWaveform(object): - """Class that stores information about a compressed waveform. +class CompressedWaveform: + """ + Class that stores information about a compressed waveform. Parameters ---------- @@ -737,12 +831,21 @@ class CompressedWaveform(object): precision : {'double', str} The precision used to generate and store the compressed waveform points. Options are 'double' or 'single'; default is 'double + """ - def __init__(self, sample_points, amplitude, phase, - interpolation=None, tolerance=None, mismatch=None, - precision='double', load_to_memory=True, - compression_factor=None): + def __init__( + self, + sample_points, + amplitude, + phase, + interpolation=None, + tolerance=None, + mismatch=None, + precision="double", + load_to_memory=True, + compression_factor=None, + ): self._sample_points = sample_points self._amplitude = amplitude self._phase = phase @@ -753,10 +856,10 @@ def __init__(self, sample_points, amplitude, phase, # save their filenames self._filenames = {} self._groupnames = {} - for arrname in ['sample_points', 'amplitude', 'phase']: + for arrname in ["sample_points", "amplitude", "phase"]: try: - fname = getattr(self, '_{}'.format(arrname)).file.filename - gname = getattr(self, '_{}'.format(arrname)).name + fname = getattr(self, f"_{arrname}").file.filename + gname = getattr(self, f"_{arrname}").name except AttributeError: fname = None gname = None @@ -769,7 +872,7 @@ def __init__(self, sample_points, amplitude, phase, self.precision = precision def _get(self, param): - val = getattr(self, '_%s' %param) + val = getattr(self, "_%s" % param) if isinstance(val, h5py.Dataset): try: val = self._cache[param] @@ -779,7 +882,7 @@ def _get(self, param): except ValueError: # this can happen if the file is closed; if so, open it # and get the data - fp = HFile(self._filenames[param], 'r') + fp = HFile(self._filenames[param], "r") val = fp[self._groupnames[param]][:] fp.close() if self.load_to_memory: @@ -788,7 +891,8 @@ def _get(self, param): @property def amplitude(self): - """The amplitude of the waveform at the `sample_points`. + """ + The amplitude of the waveform at the `sample_points`. This is always returned as an array; the same logic as for `sample_points` is used to determine whether or not to cache in @@ -797,12 +901,14 @@ def amplitude(self): Returns ------- amplitude : Array + """ - return self._get('amplitude') + return self._get("amplitude") @property def phase(self): - """The phase of the waveform as the `sample_points`. + """ + The phase of the waveform as the `sample_points`. This is always returned as an array; the same logic as for `sample_points` returned as an array; the same logic as for @@ -812,12 +918,14 @@ def phase(self): Returns ------- phase : Array + """ - return self._get('phase') + return self._get("phase") @property def sample_points(self): - """The frequencies at which the compressed waveform is sampled. + """ + The frequencies at which the compressed waveform is sampled. This is always returned as an array, even if the stored `sample_points` is an @@ -828,15 +936,17 @@ def sample_points(self): Returns ------- sample_points : Array + """ - return self._get('sample_points') + return self._get("sample_points") def clear_cache(self): """Clear self's cache of amplitude, phase, and sample_points.""" self._cache.clear() def decompress(self, out=None, df=None, f_lower=None, interpolation=None): - """Decompress self. + """ + Decompress self. Parameters ---------- @@ -859,18 +969,26 @@ def decompress(self, out=None, df=None, f_lower=None, interpolation=None): ------- FrequencySeries The decompressed waveform. + """ if f_lower is None: # use the minimum of the samlpe points f_lower = self.sample_points.min() if interpolation is None: interpolation = self.interpolation - return fd_decompress(self.amplitude, self.phase, self.sample_points, - out=out, df=df, f_lower=f_lower, - interpolation=interpolation) + return fd_decompress( + self.amplitude, + self.phase, + self.sample_points, + out=out, + df=df, + f_lower=f_lower, + interpolation=interpolation, + ) def write_to_hdf(self, fp, template_hash, root=None, precision=None): - """Write the compressed waveform to the given hdf file handler. + """ + Write the compressed waveform to the given hdf file handler. The waveform is written to: `fp['[{root}/]compressed_waveforms/{template_hash}/{param}']`, @@ -893,34 +1011,39 @@ def write_to_hdf(self, fp, template_hash, root=None, precision=None): None provided, will use whatever their current precision is. This will raise an error if the parameters have single precision but the requested precision is double. + """ if root is None: - root = '' + root = "" else: - root = '%s/'%(root) + root = "%s/" % (root) if precision is None: precision = self.precision - elif precision == 'double' and self.precision == 'single': + elif precision == "double" and self.precision == "single": raise ValueError("cannot cast single precision to double") outdtype = _real_dtypes[precision] - group = '%scompressed_waveforms/%s' %(root, str(template_hash)) - for param in ['amplitude', 'phase', 'sample_points']: - fp.create_dataset('%s/%s' %(group, param), - data=self._get(param).astype(outdtype), - compression='gzip', - shuffle=True, - compression_opts=9) + group = "%scompressed_waveforms/%s" % (root, str(template_hash)) + for param in ["amplitude", "phase", "sample_points"]: + fp.create_dataset( + "%s/%s" % (group, param), + data=self._get(param).astype(outdtype), + compression="gzip", + shuffle=True, + compression_opts=9, + ) fp_group = fp[group] - fp_group.attrs['mismatch'] = self.mismatch - fp_group.attrs['interpolation'] = self.interpolation - fp_group.attrs['tolerance'] = self.tolerance - fp_group.attrs['precision'] = precision - fp_group.attrs['compression_factor'] = self.compression_factor + fp_group.attrs["mismatch"] = self.mismatch + fp_group.attrs["interpolation"] = self.interpolation + fp_group.attrs["tolerance"] = self.tolerance + fp_group.attrs["precision"] = precision + fp_group.attrs["compression_factor"] = self.compression_factor @classmethod - def from_hdf(cls, fp, template_hash, root=None, load_to_memory=True, - load_now=False): - """Load a compressed waveform from the given hdf file handler. + def from_hdf( + cls, fp, template_hash, root=None, load_to_memory=True, load_now=False + ): + """ + Load a compressed waveform from the given hdf file handler. The waveform is retrieved from: `fp['[{root}/]compressed_waveforms/{template_hash}/{param}']`, @@ -947,24 +1070,29 @@ def from_hdf(cls, fp, template_hash, root=None, load_to_memory=True, ------- CompressedWaveform An instance of this class with parameters loaded from the hdf file. + """ if root is None: - root = '' + root = "" else: - root = '%s/'%(root) - group = '%scompressed_waveforms/%s' %(root, str(template_hash)) + root = "%s/" % (root) + group = "%scompressed_waveforms/%s" % (root, str(template_hash)) fp_group = fp[group] - sample_points = fp_group['sample_points'] - amp = fp_group['amplitude'] - phase = fp_group['phase'] + sample_points = fp_group["sample_points"] + amp = fp_group["amplitude"] + phase = fp_group["phase"] if load_now: sample_points = sample_points[:] amp = amp[:] phase = phase[:] - return cls(sample_points, amp, phase, - interpolation=fp_group.attrs['interpolation'], - tolerance=fp_group.attrs['tolerance'], - mismatch=fp_group.attrs['mismatch'], - precision=fp_group.attrs['precision'], - compression_factor=fp_group.attrs['compression_factor'], - load_to_memory=load_to_memory) + return cls( + sample_points, + amp, + phase, + interpolation=fp_group.attrs["interpolation"], + tolerance=fp_group.attrs["tolerance"], + mismatch=fp_group.attrs["mismatch"], + precision=fp_group.attrs["precision"], + compression_factor=fp_group.attrs["compression_factor"], + load_to_memory=load_to_memory, + ) diff --git a/pycbc/waveform/decompress_cpu.py b/pycbc/waveform/decompress_cpu.py index ed49e88a202..e7954d5d4a5 100644 --- a/pycbc/waveform/decompress_cpu.py +++ b/pycbc/waveform/decompress_cpu.py @@ -20,19 +20,29 @@ # # ============================================================================= # -""" Utilities for handling frequency compressed an unequally spaced frequency +""" +Utilities for handling frequency compressed an unequally spaced frequency domain waveforms. """ + import numpy -from ..types import real_same_precision_as -from ..types import complex_same_precision_as -from .decompress_cpu_cython import (decomp_ccode_double, decomp_ccode_float, - decomp_qcode_double, decomp_qcode_float, - decomp_tcode_double, decomp_tcode_float, - decomp_Qcode_double, decomp_Qcode_float) -def inline_linear_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): +from ..types import complex_same_precision_as, real_same_precision_as +from .decompress_cpu_cython import ( + decomp_ccode_double, + decomp_ccode_float, + decomp_Qcode_double, + decomp_qcode_double, + decomp_Qcode_float, + decomp_qcode_float, + decomp_tcode_double, + decomp_tcode_float, +) + + +def inline_linear_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): rprec = real_same_precision_as(output) cprec = complex_same_precision_as(output) @@ -43,17 +53,21 @@ def inline_linear_interp(amp, phase, sample_frequencies, output, h = numpy.array(output.data, copy=False, dtype=cprec) hlen = len(output) delta_f = float(df) - if output.precision == 'single': - decomp_ccode_float(h, delta_f, hlen, start_index, sample_frequencies, - amp, phase, sflen, imin) + if output.precision == "single": + decomp_ccode_float( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) else: - decomp_ccode_double(h, delta_f, hlen, start_index, sample_frequencies, - amp, phase, sflen, imin) + decomp_ccode_double( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) return output -def inline_quadratic_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): + +def inline_quadratic_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): rprec = real_same_precision_as(output) cprec = complex_same_precision_as(output) @@ -64,17 +78,21 @@ def inline_quadratic_interp(amp, phase, sample_frequencies, output, h = numpy.array(output.data, copy=False, dtype=cprec) hlen = len(output) delta_f = float(df) - if output.precision == 'single': - decomp_qcode_float(h, delta_f, hlen, start_index, - sample_frequencies, amp, phase, sflen, imin) + if output.precision == "single": + decomp_qcode_float( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) else: - decomp_qcode_double(h, delta_f, hlen, start_index, - sample_frequencies, amp, phase, sflen, imin) + decomp_qcode_double( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) return output -def inline_cubic_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): + +def inline_cubic_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): rprec = real_same_precision_as(output) cprec = complex_same_precision_as(output) @@ -85,17 +103,21 @@ def inline_cubic_interp(amp, phase, sample_frequencies, output, h = numpy.array(output.data, copy=False, dtype=cprec) hlen = len(output) delta_f = float(df) - if output.precision == 'single': - decomp_tcode_float(h, delta_f, hlen, start_index, - sample_frequencies, amp, phase, sflen, imin) + if output.precision == "single": + decomp_tcode_float( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) else: - decomp_tcode_double(h, delta_f, hlen, start_index, - sample_frequencies, amp, phase, sflen, imin) + decomp_tcode_double( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) return output -def inline_quartic_interp(amp, phase, sample_frequencies, output, - df, f_lower, imin, start_index): + +def inline_quartic_interp( + amp, phase, sample_frequencies, output, df, f_lower, imin, start_index +): rprec = real_same_precision_as(output) cprec = complex_same_precision_as(output) @@ -106,11 +128,13 @@ def inline_quartic_interp(amp, phase, sample_frequencies, output, h = numpy.array(output.data, copy=False, dtype=cprec) hlen = len(output) delta_f = float(df) - if output.precision == 'single': - decomp_Qcode_float(h, delta_f, hlen, start_index, - sample_frequencies, amp, phase, sflen, imin) + if output.precision == "single": + decomp_Qcode_float( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) else: - decomp_Qcode_double(h, delta_f, hlen, start_index, - sample_frequencies, amp, phase, sflen, imin) + decomp_Qcode_double( + h, delta_f, hlen, start_index, sample_frequencies, amp, phase, sflen, imin + ) return output diff --git a/pycbc/waveform/decompress_cuda.py b/pycbc/waveform/decompress_cuda.py index 8dc832bd4b4..aa5a09495c7 100644 --- a/pycbc/waveform/decompress_cuda.py +++ b/pycbc/waveform/decompress_cuda.py @@ -21,10 +21,11 @@ # # ============================================================================= # -import numpy import mako.template +import numpy from pycuda import gpuarray from pycuda.compiler import SourceModule + import pycbc.scheme from pycbc.types import zeros @@ -251,6 +252,8 @@ """) dckernel_cache = {} + + def get_dckernel(slen): # Right now, hardcoding the number of threads per block nt = 1024 @@ -273,7 +276,8 @@ def get_dckernel(slen): dckernel_cache[nb] = (fn1, fn2, freq_tex, amp_tex, phase_tex, nt, nb) return dckernel_cache[nb] -class CUDALinearInterpolate(object): + +class CUDALinearInterpolate: def __init__(self, output): self.output = output.data.gpudata self.df = numpy.float32(output.delta_f) @@ -292,7 +296,7 @@ def __init__(self, output): def interpolate(self, flow, freqs, amps, phases): flow = numpy.float32(flow) texlen = numpy.int32(len(freqs)) - fmax = numpy.float32(freqs[texlen-1]) + fmax = numpy.float32(freqs[texlen - 1]) freqs_gpu = gpuarray.to_gpu(freqs) freqs_gpu.bind_to_texref_ext(self.freq_tex, allow_offset=False) amps_gpu = gpuarray.to_gpu(amps) @@ -301,19 +305,34 @@ def interpolate(self, flow, freqs, amps, phases): phases_gpu.bind_to_texref_ext(self.phase_tex, allow_offset=False) fn1 = self.fn1.prepared_call fn2 = self.fn2.prepared_call - fn1((1, 1), (self.nb, 1, 1), self.lower, self.upper, texlen, self.df, flow, fmax) - fn2((self.nb, 1), (self.nt, 1, 1), self.output, self.df, self.hlen, flow, fmax, texlen, self.lower, self.upper) + fn1( + (1, 1), (self.nb, 1, 1), self.lower, self.upper, texlen, self.df, flow, fmax + ) + fn2( + (self.nb, 1), + (self.nt, 1, 1), + self.output, + self.df, + self.hlen, + flow, + fmax, + texlen, + self.lower, + self.upper, + ) pycbc.scheme.mgr.state.context.synchronize() - return + def inline_linear_interp(amps, phases, freqs, output, df, flow, imin, start_index): # Note that imin and start_index are ignored in the GPU code; they are only # needed for CPU. - if output.precision == 'double': - raise NotImplementedError("Double precision linear interpolation not currently supported on CUDA scheme") + if output.precision == "double": + raise NotImplementedError( + "Double precision linear interpolation not currently supported on CUDA scheme" + ) flow = numpy.float32(flow) texlen = numpy.int32(len(freqs)) - fmax = numpy.float32(freqs[texlen-1]) + fmax = numpy.float32(freqs[texlen - 1]) hlen = numpy.int32(len(output)) (fn1, fn2, ftex, atex, ptex, nt, nb) = get_dckernel(hlen) freqs_gpu = gpuarray.to_gpu(freqs) diff --git a/pycbc/waveform/decompress_cupy.py b/pycbc/waveform/decompress_cupy.py index 6a53ba356b4..d63741e3ffc 100644 --- a/pycbc/waveform/decompress_cupy.py +++ b/pycbc/waveform/decompress_cupy.py @@ -236,6 +236,8 @@ """) dckernel_cache = {} + + def get_dckernel(slen): # Right now, hardcoding the number of threads per block nt = 1024 @@ -252,7 +254,8 @@ def get_dckernel(slen): return dckernel_cache[nb] -class CUPYLinearInterpolate(object): + +class CUPYLinearInterpolate: def __init__(self, output): self.output = output.data self.df = np.float32(output.delta_f) @@ -268,27 +271,42 @@ def __init__(self, output): def interpolate(self, flow, freqs, amps, phases): flow = np.float32(flow) texlen = np.int32(len(freqs)) - fmax = np.float32(freqs[texlen-1]) + fmax = np.float32(freqs[texlen - 1]) freqs_gpu = cp.asarray(freqs) amps_gpu = cp.asarray(amps) phases_gpu = cp.asarray(phases) self.fn1( - (1,) , (self.nb,), - (self.lower, self.upper, texlen, self.df, flow, freqs_gpu)) + (1,), (self.nb,), (self.lower, self.upper, texlen, self.df, flow, freqs_gpu) + ) self.fn2( - (self.nb,), (self.nt,), - (self.output, self.df, self.hlen, flow, fmax, texlen, freqs_gpu, amps_gpu, phases_gpu, self.lower, self.upper) + (self.nb,), + (self.nt,), + ( + self.output, + self.df, + self.hlen, + flow, + fmax, + texlen, + freqs_gpu, + amps_gpu, + phases_gpu, + self.lower, + self.upper, + ), ) - return + def inline_linear_interp(amps, phases, freqs, output, df, flow, imin, start_index): # Note that imin and start_index are ignored in the GPU code; they are only # needed for CPU. - if output.precision == 'double': - raise NotImplementedError("Double precision linear interpolation not currently supported on CUDA scheme") + if output.precision == "double": + raise NotImplementedError( + "Double precision linear interpolation not currently supported on CUDA scheme" + ) flow = np.float32(flow) texlen = np.int32(len(freqs)) - fmax = np.float32(freqs[texlen-1]) + fmax = np.float32(freqs[texlen - 1]) hlen = np.int32(len(output)) (fn1, fn2, nt, nb) = get_dckernel(hlen) @@ -300,12 +318,22 @@ def inline_linear_interp(amps, phases, freqs, output, df, flow, imin, start_inde g_out = output.data lower = cp.zeros(nb, dtype=np.int32) upper = cp.zeros(nb, dtype=np.int32) - fn1( - (1,), (nb,), - (lower, upper, texlen, df, flow, freqs_gpu) - ) + fn1((1,), (nb,), (lower, upper, texlen, df, flow, freqs_gpu)) fn2( - (nb,), (nt,), - (g_out, df, hlen, flow, fmax, texlen, freqs_gpu, amps_gpu, phases_gpu, lower, upper) + (nb,), + (nt,), + ( + g_out, + df, + hlen, + flow, + fmax, + texlen, + freqs_gpu, + amps_gpu, + phases_gpu, + lower, + upper, + ), ) return output diff --git a/pycbc/waveform/generator.py b/pycbc/waveform/generator.py index 4534e9e53b4..d4501b4fe48 100644 --- a/pycbc/waveform/generator.py +++ b/pycbc/waveform/generator.py @@ -24,31 +24,30 @@ """ This modules provides classes for generating waveforms. """ -import os + import logging +import os +from abc import ABCMeta, abstractmethod -from abc import (ABCMeta, abstractmethod) +from numpy import pi -from . import waveform -from .waveform import (FailedWaveformError) -from . import ringdown -from . import supernovae -from . import waveform_modes -from pycbc.types import TimeSeries -from pycbc.waveform import parameters -from pycbc.waveform.utils import apply_fseries_time_shift, \ - ceilpow2, apply_fd_time_shift +from pycbc import strain from pycbc.detector import Detector from pycbc.pool import use_mpi -from pycbc import strain -from numpy import pi +from pycbc.types import TimeSeries +from pycbc.waveform import parameters +from pycbc.waveform.utils import apply_fd_time_shift, apply_fseries_time_shift, ceilpow2 +from . import ringdown, supernovae, waveform, waveform_modes +from .waveform import FailedWaveformError # utility functions/class failed_counter = 0 -class BaseGenerator(object): - r"""A wrapper class to call a waveform generator with a set of frozen + +class BaseGenerator: + r""" + A wrapper class to call a waveform generator with a set of frozen parameters and a set of variable parameters. The frozen parameters and values, along with a list of variable parameter names, are set at initialization. This way, repeated calls can be made to the underlying @@ -83,9 +82,12 @@ class BaseGenerator(object): current_params : dict A dictionary of the frozen keyword arguments and variable arguments that were last passed to the waveform generator. + """ - def __init__(self, generator, variable_args=(), record_failures=False, - **frozen_params): + + def __init__( + self, generator, variable_args=(), record_failures=False, **frozen_params + ): self.generator = generator self.variable_args = tuple(variable_args) self.frozen_params = frozen_params @@ -98,8 +100,9 @@ def __init__(self, generator, variable_args=(), record_failures=False, # If we are under mpi, then failed waveform will be stored by # mpi rank to avoid file writing conflicts. We'll check for this # upfront - self.record_failures = (record_failures or - ('PYCBC_RECORD_FAILED_WAVEFORMS' in os.environ)) + self.record_failures = record_failures or ( + "PYCBC_RECORD_FAILED_WAVEFORMS" in os.environ + ) self.mpi_enabled, _, self.mpi_rank = use_mpi() @property @@ -108,106 +111,116 @@ def static_args(self): return self.frozen_params def generate(self, **kwargs): - """Generates a waveform from the keyword args. The current params + """ + Generates a waveform from the keyword args. The current params are updated with the given kwargs, then the generator is called. """ self.current_params.update(kwargs) return self._generate_from_current() def _add_pregenerate(self, func): - """ Adds a function that will be called by the generator function + """ + Adds a function that will be called by the generator function before waveform generation. """ self._pregenerate_functions.append(func) def _postgenerate(self, res): - """Allows the waveform returned by the generator function to be + """ + Allows the waveform returned by the generator function to be manipulated before returning. """ return res def _gdecorator(generate_func): - """A decorator that allows for seemless pre/post manipulation of + """ + A decorator that allows for seemless pre/post manipulation of the waveform generator function. """ + def dostuff(self): for func in self._pregenerate_functions: self.current_params = func(self.current_params) - res = generate_func(self) # pylint:disable=not-callable + res = generate_func(self) # pylint:disable=not-callable return self._postgenerate(res) + return dostuff @_gdecorator def _generate_from_current(self): - """Generates a waveform from the current parameters. - """ + """Generates a waveform from the current parameters.""" try: new_waveform = self.generator(**self.current_params) return new_waveform except RuntimeError as e: if self.record_failures: - from pycbc.io.hdf import dump_state, HFile + from pycbc.io.hdf import HFile, dump_state global failed_counter if self.mpi_enabled: - outname = 'failed/params_%s.hdf' % self.mpi_rank + outname = "failed/params_%s.hdf" % self.mpi_rank else: - outname = 'failed/params.hdf' + outname = "failed/params.hdf" - if not os.path.exists('failed'): - os.makedirs('failed') + if not os.path.exists("failed"): + os.makedirs("failed") with HFile(outname) as f: - dump_state(self.current_params, f, - dsetname=str(failed_counter)) + dump_state(self.current_params, f, dsetname=str(failed_counter)) failed_counter += 1 # we'll get a RuntimeError if lalsimulation failed to generate # the waveform for whatever reason - strparams = ' | '.join(['{}: {}'.format( - p, str(val)) for p, val in self.current_params.items()]) - raise FailedWaveformError("Failed to generate waveform with " - "parameters:\n{}\nError was: {}" - .format(strparams, e)) + strparams = " | ".join( + [f"{p}: {val!s}" for p, val in self.current_params.items()] + ) + raise FailedWaveformError( + "Failed to generate waveform with " + f"parameters:\n{strparams}\nError was: {e}" + ) class BaseCBCGenerator(BaseGenerator): - """Adds ability to convert from various derived parameters to parameters + """ + Adds ability to convert from various derived parameters to parameters needed by the waveform generators. """ - possible_args = set(parameters.td_waveform_params + - parameters.fd_waveform_params + - ['taper']) + possible_args = set( + parameters.td_waveform_params + parameters.fd_waveform_params + ["taper"] + ) """set: The set of names of arguments that may be used in the `variable_args` or `frozen_params`. """ def __init__(self, generator, variable_args=(), **frozen_params): - super(BaseCBCGenerator, self).__init__(generator, - variable_args=variable_args, **frozen_params) + super().__init__( + generator, variable_args=variable_args, **frozen_params + ) # decorate the generator function with a list of functions that convert # parameters to those used by the waveform generation interface - all_args = set(list(self.frozen_params.keys()) + - list(self.variable_args)) + all_args = set(list(self.frozen_params.keys()) + list(self.variable_args)) # check that there are no unused (non-calibration) parameters - calib_args = set([a for a in self.variable_args if - a.startswith('calib_')]) + calib_args = set([a for a in self.variable_args if a.startswith("calib_")]) all_args = all_args - calib_args unused_args = all_args - self.possible_args if len(unused_args): - logging.warning("WARNING: The following parameters are generally " - "not used by CBC waveform generators: %s. If you " - "have provided a transform that converted these " - "into known parameters (e.g., mchirp, q to " - "mass1, mass2) or you are using a custom model " - "that uses these parameters, you can safely " - "ignore this message.", ', '.join(unused_args)) + logging.warning( + "WARNING: The following parameters are generally " + "not used by CBC waveform generators: %s. If you " + "have provided a transform that converted these " + "into known parameters (e.g., mchirp, q to " + "mass1, mass2) or you are using a custom model " + "that uses these parameters, you can safely " + "ignore this message.", + ", ".join(unused_args), + ) class FDomainCBCGenerator(BaseCBCGenerator): - """Generates frequency-domain CBC waveforms in the radiation frame. + """ + Generates frequency-domain CBC waveforms in the radiation frame. Uses `waveform.get_fd_waveform` as a generator function to create frequency- domain CBC waveforms in the radiation frame; i.e., with no @@ -227,13 +240,16 @@ class FDomainCBCGenerator(BaseCBCGenerator): ) """ + def __init__(self, variable_args=(), **frozen_params): - super(FDomainCBCGenerator, self).__init__(waveform.get_fd_waveform, - variable_args=variable_args, **frozen_params) + super().__init__( + waveform.get_fd_waveform, variable_args=variable_args, **frozen_params + ) class FDomainCBCModesGenerator(BaseCBCGenerator): - """Generates frequency-domain CBC waveform modes. + """ + Generates frequency-domain CBC waveform modes. Uses :py:func:`waveform_modes.get_fd_waveform_modes` as a generator function to create frequency-domain CBC waveforms mode-by-mode, without @@ -241,14 +257,18 @@ class FDomainCBCModesGenerator(BaseCBCGenerator): For details, on methods and arguments, see :py:class:`BaseGenerator`. """ + def __init__(self, variable_args=(), **frozen_params): - super(FDomainCBCModesGenerator, self).__init__( + super().__init__( waveform_modes.get_fd_waveform_modes, - variable_args=variable_args, **frozen_params) + variable_args=variable_args, + **frozen_params, + ) class TDomainCBCGenerator(BaseCBCGenerator): - """Create time domain CBC waveforms in the radiation frame. + """ + Create time domain CBC waveforms in the radiation frame. Uses waveform.get_td_waveform as a generator function to create time- domain CBC waveforms in the radiation frame; i.e., with no detector @@ -268,28 +288,34 @@ class TDomainCBCGenerator(BaseCBCGenerator): ) """ + def __init__(self, variable_args=(), **frozen_params): - super(TDomainCBCGenerator, self).__init__(waveform.get_td_waveform, - variable_args=variable_args, **frozen_params) + super().__init__( + waveform.get_td_waveform, variable_args=variable_args, **frozen_params + ) def _postgenerate(self, res): - """Applies a taper if it is in current params. - """ + """Applies a taper if it is in current params.""" hp, hc = res - if 'taper' in self.current_params: - location = self.current_params['taper'] - hp = hp.taper_timeseries(location=location, - tapermethod=self.current_params.get('taper_method', 'lal'), - taper_window=self.current_params.get('taper_window')) - hc = hc.taper_timeseries(location=location, - tapermethod=self.current_params.get('taper_method', 'lal'), - taper_window=self.current_params.get('taper_window')) - + if "taper" in self.current_params: + location = self.current_params["taper"] + hp = hp.taper_timeseries( + location=location, + tapermethod=self.current_params.get("taper_method", "lal"), + taper_window=self.current_params.get("taper_window"), + ) + hc = hc.taper_timeseries( + location=location, + tapermethod=self.current_params.get("taper_method", "lal"), + taper_window=self.current_params.get("taper_window"), + ) + return hp, hc class TDomainCBCModesGenerator(BaseCBCGenerator): - """Generates time domain CBC waveform modes. + """ + Generates time domain CBC waveform modes. Uses :py:func:`waveform_modes.get_td_waveform_modes` as a generator function to create time-domain CBC waveforms mode-by-mode, without applying @@ -298,30 +324,37 @@ class TDomainCBCModesGenerator(BaseCBCGenerator): For details, on methods and arguments, see :py:class:`BaseGenerator`. """ + def __init__(self, variable_args=(), **frozen_params): - super(TDomainCBCModesGenerator, self).__init__( + super().__init__( waveform_modes.get_td_waveform_modes, - variable_args=variable_args, **frozen_params) + variable_args=variable_args, + **frozen_params, + ) def _postgenerate(self, res): - """Applies a taper if it is in current params. - """ - if 'taper' in self.current_params: - location = self.current_params['taper'] + """Applies a taper if it is in current params.""" + if "taper" in self.current_params: + location = self.current_params["taper"] for mode in res: ulm, vlm = res[mode] - ulm = ulm.taper_timeseries(location=location, - tapermethod=self.current_params.get('taper_method', 'lal'), - taper_window=self.current_params.get('taper_window')) - vlm = vlm.taper_timeseries(location=location, - tapermethod=self.current_params.get('taper_method', 'lal'), - taper_window=self.current_params.get('taper_window')) + ulm = ulm.taper_timeseries( + location=location, + tapermethod=self.current_params.get("taper_method", "lal"), + taper_window=self.current_params.get("taper_window"), + ) + vlm = vlm.taper_timeseries( + location=location, + tapermethod=self.current_params.get("taper_method", "lal"), + taper_window=self.current_params.get("taper_window"), + ) res[mode] = (ulm, vlm) return res class FDomainMassSpinRingdownGenerator(BaseGenerator): - """Uses ringdown.get_fd_from_final_mass_spin as a generator function to + """ + Uses ringdown.get_fd_from_final_mass_spin as a generator function to create frequency-domain ringdown waveforms with higher modes in the radiation frame; i.e., with no detector response function applied. For more details, see BaseGenerator. @@ -343,13 +376,18 @@ class FDomainMassSpinRingdownGenerator(BaseGenerator): ) """ + def __init__(self, variable_args=(), **frozen_params): - super(FDomainMassSpinRingdownGenerator, self).__init__(ringdown.get_fd_from_final_mass_spin, - variable_args=variable_args, **frozen_params) + super().__init__( + ringdown.get_fd_from_final_mass_spin, + variable_args=variable_args, + **frozen_params, + ) class FDomainFreqTauRingdownGenerator(BaseGenerator): - """Uses ringdown.get_fd_from_freqtau as a generator function to + """ + Uses ringdown.get_fd_from_freqtau as a generator function to create frequency-domain ringdown waveforms with higher modes in the radiation frame; i.e., with no detector response function applied. For more details, see BaseGenerator. @@ -371,13 +409,16 @@ class FDomainFreqTauRingdownGenerator(BaseGenerator): ) """ + def __init__(self, variable_args=(), **frozen_params): - super(FDomainFreqTauRingdownGenerator, self).__init__(ringdown.get_fd_from_freqtau, - variable_args=variable_args, **frozen_params) + super().__init__( + ringdown.get_fd_from_freqtau, variable_args=variable_args, **frozen_params + ) class TDomainMassSpinRingdownGenerator(BaseGenerator): - """Uses ringdown.get_td_from_final_mass_spin as a generator function to + """ + Uses ringdown.get_td_from_final_mass_spin as a generator function to create time-domain ringdown waveforms with higher modes in the radiation frame; i.e., with no detector response function applied. For more details, see BaseGenerator. @@ -399,13 +440,18 @@ class TDomainMassSpinRingdownGenerator(BaseGenerator): ) """ + def __init__(self, variable_args=(), **frozen_params): - super(TDomainMassSpinRingdownGenerator, self).__init__(ringdown.get_td_from_final_mass_spin, - variable_args=variable_args, **frozen_params) + super().__init__( + ringdown.get_td_from_final_mass_spin, + variable_args=variable_args, + **frozen_params, + ) class TDomainFreqTauRingdownGenerator(BaseGenerator): - """Uses ringdown.get_td_from_freqtau as a generator function to + """ + Uses ringdown.get_td_from_freqtau as a generator function to create time-domain ringdown waveforms with higher modes in the radiation frame; i.e., with no detector response function applied. For more details, see BaseGenerator. @@ -427,19 +473,25 @@ class TDomainFreqTauRingdownGenerator(BaseGenerator): ) """ + def __init__(self, variable_args=(), **frozen_params): - super(TDomainFreqTauRingdownGenerator, self).__init__(ringdown.get_td_from_freqtau, - variable_args=variable_args, **frozen_params) + super().__init__( + ringdown.get_td_from_freqtau, variable_args=variable_args, **frozen_params + ) class TDomainSupernovaeGenerator(BaseGenerator): - """Uses supernovae.py to create time domain core-collapse supernovae waveforms + """ + Uses supernovae.py to create time domain core-collapse supernovae waveforms using a set of Principal Components provided in a .hdf file. """ + def __init__(self, variable_args=(), **frozen_params): - super(TDomainSupernovaeGenerator, - self).__init__(supernovae.get_corecollapse_bounce, - variable_args=variable_args, **frozen_params) + super().__init__( + supernovae.get_corecollapse_bounce, + variable_args=variable_args, + **frozen_params, + ) # @@ -452,7 +504,8 @@ def __init__(self, variable_args=(), **frozen_params): class BaseFDomainDetFrameGenerator(metaclass=ABCMeta): - r"""Base generator for frquency-domain waveforms in a detector frame. + r""" + Base generator for frquency-domain waveforms in a detector frame. Parameters ---------- @@ -505,8 +558,16 @@ class BaseFDomainDetFrameGenerator(metaclass=ABCMeta): that set the binary's location. """ - def __init__(self, rFrameGeneratorClass, epoch, detectors=None, - variable_args=(), recalib=None, gates=None, **frozen_params): + def __init__( + self, + rFrameGeneratorClass, + epoch, + detectors=None, + variable_args=(), + recalib=None, + gates=None, + **frozen_params, + ): # initialize frozen & current parameters: self.current_params = frozen_params.copy() self._static_args = frozen_params.copy() @@ -522,7 +583,8 @@ def __init__(self, rFrameGeneratorClass, epoch, detectors=None, rframe_variables = list(set(self.variable_args) - self.location_args) # initialize the radiation frame generator self.rframe_generator = rFrameGeneratorClass( - variable_args=rframe_variables, **frozen_params) + variable_args=rframe_variables, **frozen_params + ) self.set_epoch(epoch) # set calibration model self.recalib = recalib @@ -530,15 +592,21 @@ def __init__(self, rFrameGeneratorClass, epoch, detectors=None, # location variables are specified if detectors is not None: self.detectors = {det: Detector(det) for det in detectors} - missing_args = [arg for arg in self.location_args if not - (arg in self.current_params or arg in self.variable_args)] + missing_args = [ + arg + for arg in self.location_args + if not (arg in self.current_params or arg in self.variable_args) + ] if any(missing_args): - raise ValueError("detectors provided, but missing location " - "parameters %s. " %(', '.join(missing_args)) + - "These must be either in the frozen params or the " - "variable args.") + raise ValueError( + "detectors provided, but missing location " + "parameters %s. " + % (", ".join(missing_args)) + + "These must be either in the frozen params or the " + "variable args." + ) else: - self.detectors = {'RF': None} + self.detectors = {"RF": None} self.detector_names = sorted(self.detectors.keys()) self.gates = gates @@ -553,7 +621,8 @@ def static_args(self): @property def epoch(self): - """The GPS start time of the frequency series returned by the generate + """ + The GPS start time of the frequency series returned by the generate function. A time shift is applied to the waveform equal to tc-epoch. Update by using ``set_epoch`` """ @@ -561,18 +630,16 @@ def epoch(self): @abstractmethod def generate(self, **kwargs): - """The function that generates the waveforms. - """ - pass + """The function that generates the waveforms.""" @abstractmethod def select_rframe_generator(self, approximant): """Method to select waveform generator based on an approximant.""" - pass class FDomainDetFrameGenerator(BaseFDomainDetFrameGenerator): - r"""Generates frequency-domain waveform in a specific frame. + r""" + Generates frequency-domain waveform in a specific frame. Generates a waveform using the given radiation frame generator class, and applies the detector response function and appropriate time offset. @@ -640,7 +707,7 @@ class FDomainDetFrameGenerator(BaseFDomainDetFrameGenerator): """ - location_args = set(['tc', 'ra', 'dec', 'polarization']) + location_args = set(["tc", "ra", "dec", "polarization"]) """set(['tc', 'ra', 'dec', 'polarization']): The set of location parameters. These are not passed to the rFrame generator class; instead, they are used to apply the detector response @@ -660,66 +727,73 @@ class FDomainDetFrameGenerator(BaseFDomainDetFrameGenerator): """ def generate(self, **kwargs): - """Generates a waveform, applies a time shift and the detector response + """ + Generates a waveform, applies a time shift and the detector response function from the given kwargs. """ self.current_params.update(kwargs) - rfparams = {param: self.current_params[param] - for param in kwargs if param not in self.location_args} + rfparams = { + param: self.current_params[param] + for param in kwargs + if param not in self.location_args + } hp, hc = self.rframe_generator.generate(**rfparams) if isinstance(hp, TimeSeries): - df = self.current_params['delta_f'] + df = self.current_params["delta_f"] hp = hp.to_frequencyseries(delta_f=df) hc = hc.to_frequencyseries(delta_f=df) # time-domain waveforms will not be shifted so that the peak amp # happens at the end of the time series (as they are for f-domain), # so we add an additional shift to account for it - tshift = 1./df - abs(hp._epoch) + tshift = 1.0 / df - abs(hp._epoch) else: - tshift = 0. + tshift = 0.0 hp._epoch = hc._epoch = self._epoch h = {} - if self.detector_names != ['RF']: - ra = self.current_params['ra'] - dec = self.current_params['dec'] - ref_tc = self.current_params['tc'] - pol = self.current_params['polarization'] - refframe = self.current_params.get('tc_ref_frame', 'geocentric') + if self.detector_names != ["RF"]: + ra = self.current_params["ra"] + dec = self.current_params["dec"] + ref_tc = self.current_params["tc"] + pol = self.current_params["polarization"] + refframe = self.current_params.get("tc_ref_frame", "geocentric") for detname, det in self.detectors.items(): tc = det.arrival_time(ref_tc, ra, dec, refframe) # apply response function fp, fc = det.antenna_pattern(ra, dec, pol, tc) - thish = fp*hp + fc*hc + thish = fp * hp + fc * hc # apply time shift - h[detname] = apply_fd_time_shift(thish, tc+tshift, copy=False) + h[detname] = apply_fd_time_shift(thish, tc + tshift, copy=False) if self.recalib: # recalibrate with given calibration model - h[detname] = \ - self.recalib[detname].map_to_adjust(h[detname], - **self.current_params) + h[detname] = self.recalib[detname].map_to_adjust( + h[detname], **self.current_params + ) else: # no detector response, just use the + polarization - if 'tc' in self.current_params: - hp = apply_fd_time_shift(hp, self.current_params['tc']+tshift, - copy=False) - h['RF'] = hp + if "tc" in self.current_params: + hp = apply_fd_time_shift( + hp, self.current_params["tc"] + tshift, copy=False + ) + h["RF"] = hp if self.gates is not None: # resize all to nearest power of 2 for d in h.values(): - d.resize(ceilpow2(len(d)-1) + 1) + d.resize(ceilpow2(len(d) - 1) + 1) h = strain.apply_gates_to_fd(h, self.gates) return h @staticmethod def select_rframe_generator(approximant, domain): - """Returns a radiation frame generator class based on the approximant + """ + Returns a radiation frame generator class based on the approximant string. """ return select_waveform_generator(approximant, domain) class FDomainDetFrameTwoPolGenerator(BaseFDomainDetFrameGenerator): - r"""Generates frequency-domain waveform in a specific frame. + r""" + Generates frequency-domain waveform in a specific frame. Generates both polarizations of a waveform using the given radiation frame generator class, and applies the time shift. Detector response functions @@ -774,7 +848,8 @@ class FDomainDetFrameTwoPolGenerator(BaseFDomainDetFrameGenerator): function. """ - location_args = set(['tc', 'ra', 'dec']) + + location_args = set(["tc", "ra", "dec"]) """ set(['tc', 'ra', 'dec']): The set of location parameters. These are not passed to the rFrame generator class; instead, they are used to apply the detector response @@ -793,55 +868,64 @@ class FDomainDetFrameTwoPolGenerator(BaseFDomainDetFrameGenerator): """ def generate(self, **kwargs): - """Generates a waveform polarizations and applies a time shift. + """ + Generates a waveform polarizations and applies a time shift. Returns ------- dict : Dictionary of ``detector names -> (hp, hc)``, where ``hp, hc`` are the plus and cross polarization, respectively. + """ self.current_params.update(kwargs) - rfparams = {param: self.current_params[param] - for param in kwargs if param not in self.location_args} + rfparams = { + param: self.current_params[param] + for param in kwargs + if param not in self.location_args + } hp, hc = self.rframe_generator.generate(**rfparams) if isinstance(hp, TimeSeries): - df = self.current_params['delta_f'] + df = self.current_params["delta_f"] hp = hp.to_frequencyseries(delta_f=df) hc = hc.to_frequencyseries(delta_f=df) # time-domain waveforms will not be shifted so that the peak amp # happens at the end of the time series (as they are for f-domain), # so we add an additional shift to account for it - tshift = 1./df - abs(hp._epoch) + tshift = 1.0 / df - abs(hp._epoch) else: - tshift = 0. + tshift = 0.0 hp._epoch = hc._epoch = self._epoch h = {} - if self.detector_names != ['RF']: + if self.detector_names != ["RF"]: for detname, det in self.detectors.items(): - refframe = self.current_params.get('tc_ref_frame', 'geocentric') - ra = self.current_params['ra'] - dec = self.current_params['dec'] - ref_tc = self.current_params['tc'] + refframe = self.current_params.get("tc_ref_frame", "geocentric") + ra = self.current_params["ra"] + dec = self.current_params["dec"] + ref_tc = self.current_params["tc"] tc = det.arrival_time(ref_tc, ra, dec, refframe) # apply time shift - dethp = apply_fd_time_shift(hp, tc+tshift, copy=True) - dethc = apply_fd_time_shift(hc, tc+tshift, copy=True) + dethp = apply_fd_time_shift(hp, tc + tshift, copy=True) + dethc = apply_fd_time_shift(hc, tc + tshift, copy=True) if self.recalib: # recalibrate with given calibration model dethp = self.recalib[detname].map_to_adjust( - dethp, **self.current_params) + dethp, **self.current_params + ) dethc = self.recalib[detname].map_to_adjust( - dethc, **self.current_params) + dethc, **self.current_params + ) h[detname] = (dethp, dethc) else: # no detector response, just use the + polarization - if 'tc' in self.current_params: - hp = apply_fd_time_shift(hp, self.current_params['tc']+tshift, - copy=False) - hc = apply_fd_time_shift(hc, self.current_params['tc']+tshift, - copy=False) - h['RF'] = (hp, hc) + if "tc" in self.current_params: + hp = apply_fd_time_shift( + hp, self.current_params["tc"] + tshift, copy=False + ) + hc = apply_fd_time_shift( + hc, self.current_params["tc"] + tshift, copy=False + ) + h["RF"] = (hp, hc) if self.gates is not None: # resize all to nearest power of 2 hps = {} @@ -849,8 +933,8 @@ def generate(self, **kwargs): for det in h: hp = h[det] hc = h[det] - hp.resize(ceilpow2(len(hp)-1) + 1) - hc.resize(ceilpow2(len(hc)-1) + 1) + hp.resize(ceilpow2(len(hp) - 1) + 1) + hc.resize(ceilpow2(len(hc) - 1) + 1) hps[det] = hp hcs[det] = hc hps = strain.apply_gates_to_fd(hps, self.gates) @@ -860,13 +944,16 @@ def generate(self, **kwargs): @staticmethod def select_rframe_generator(approximant, domain): - """Returns a radiation frame generator class based on the approximant + """ + Returns a radiation frame generator class based on the approximant string. """ return select_waveform_generator(approximant, domain) + class FDomainDetFrameTwoPolNoRespGenerator(BaseFDomainDetFrameGenerator): - r"""Generates frequency-domain waveform in a specific frame. + r""" + Generates frequency-domain waveform in a specific frame. Generates both polarizations of a waveform using the given radiation frame generator class, and applies the time shift. Detector response functions @@ -923,24 +1010,26 @@ class FDomainDetFrameTwoPolNoRespGenerator(BaseFDomainDetFrameGenerator): """ def generate(self, **kwargs): - """Generates a waveform polarizations + """ + Generates a waveform polarizations Returns ------- dict : Dictionary of ``detector names -> (hp, hc)``, where ``hp, hc`` are the plus and cross polarization, respectively. + """ self.current_params.update(kwargs) hp, hc = self.rframe_generator.generate(**self.current_params) if isinstance(hp, TimeSeries): - df = self.current_params['delta_f'] + df = self.current_params["delta_f"] hp = hp.to_frequencyseries(delta_f=df) hc = hc.to_frequencyseries(delta_f=df) # time-domain waveforms will not be shifted so that the peak amp # happens at the end of the time series (as they are for f-domain), # so we add an additional shift to account for it - tshift = 1./df - abs(hp._epoch) + tshift = 1.0 / df - abs(hp._epoch) hp = apply_fseries_time_shift(hp, tshift, copy=True) hc = apply_fseries_time_shift(hc, tshift, copy=True) @@ -950,27 +1039,27 @@ def generate(self, **kwargs): for detname in self.detectors: if self.recalib: # recalibrate with given calibration model - hp = self.recalib[detname].map_to_adjust( - hp, **self.current_params) - hc = self.recalib[detname].map_to_adjust( - hc, **self.current_params) + hp = self.recalib[detname].map_to_adjust(hp, **self.current_params) + hc = self.recalib[detname].map_to_adjust(hc, **self.current_params) h[detname] = (hp.copy(), hc.copy()) return h @staticmethod def select_rframe_generator(approximant, domain): - """Returns a radiation frame generator class based on the approximant + """ + Returns a radiation frame generator class based on the approximant string. """ return select_waveform_generator(approximant, domain) class FDomainDetFrameTwoPhaseGenerator(BaseFDomainDetFrameGenerator): - r"""Generates frequency-domain waveform in a specific frame. - + r""" + Generates frequency-domain waveform in a specific frame. + This class assumes that the radiation-frame waveform can be decomposed in terms of a phase phi such that - + h = h_c * cos(phi) + h_s * sin(phi), where h_c and h_s are the waveform evaluated at phi = 0 and phi = pi/2 @@ -1034,9 +1123,10 @@ class FDomainDetFrameTwoPhaseGenerator(BaseFDomainDetFrameGenerator): variable_args : tuple The list of names of arguments that are passed to the generate function. + """ - location_args = set(['tc', 'ra', 'dec', 'polarization']) + location_args = set(["tc", "ra", "dec", "polarization"]) """set(['tc', 'ra', 'dec', 'polarization']): The set of location parameters. These are not passed to the rFrame generator class; instead, they are used to apply the detector response @@ -1056,24 +1146,29 @@ class FDomainDetFrameTwoPhaseGenerator(BaseFDomainDetFrameGenerator): """ def generate(self, phases=None, ref_phase=None, **kwargs): - """Generates a waveform, applies a time shift and the detector response + """ + Generates a waveform, applies a time shift and the detector response function from the given kwargs. """ self.current_params.update(kwargs) - rfparams = {param: self.current_params[param] - for param in kwargs if param not in self.location_args} + rfparams = { + param: self.current_params[param] + for param in kwargs + if param not in self.location_args + } # generate the cosine term: ref_phase = 0 - if rfparams[ref_phase] != 0.: - raise ValueError(f'Reference phase {ref_phase}={rfparams[ref_phase]} is ' - 'not zero') + if rfparams[ref_phase] != 0.0: + raise ValueError( + f"Reference phase {ref_phase}={rfparams[ref_phase]} is not zero" + ) hpc, hcc = self.rframe_generator.generate(**rfparams) # generate the sine term: shift all phases by pi/2 sin_params = rfparams.copy() for i in phases: - sin_params[i] = rfparams[i] + pi/2 + sin_params[i] = rfparams[i] + pi / 2 hps, hcs = self.rframe_generator.generate(**sin_params) if isinstance(hpc, TimeSeries): - df = self.current_params['delta_f'] + df = self.current_params["delta_f"] hpc = hpc.to_frequencyseries(delta_f=df) hcc = hcc.to_frequencyseries(delta_f=df) hps = hps.to_frequencyseries(delta_f=df) @@ -1081,61 +1176,65 @@ def generate(self, phases=None, ref_phase=None, **kwargs): # time-domain waveforms will not be shifted so that the peak amp # happens at the end of the time series (as they are for f-domain), # so we add an additional shift to account for it - tshift = 1./df - abs(hpc._epoch) + tshift = 1.0 / df - abs(hpc._epoch) else: - tshift = 0. + tshift = 0.0 hpc._epoch = hcc._epoch = hps._epoch = hcs._epoch = self._epoch h = {} - if self.detector_names != ['RF']: - ra = self.current_params['ra'] - dec = self.current_params['dec'] - ref_tc = self.current_params['tc'] - pol = self.current_params['polarization'] - refframe = self.current_params.get('tc_ref_frame', 'geocentric') + if self.detector_names != ["RF"]: + ra = self.current_params["ra"] + dec = self.current_params["dec"] + ref_tc = self.current_params["tc"] + pol = self.current_params["polarization"] + refframe = self.current_params.get("tc_ref_frame", "geocentric") for detname, det in self.detectors.items(): tc = det.arrival_time(ref_tc, ra, dec, refframe) # apply response function fp, fc = det.antenna_pattern(ra, dec, pol, tc) - thishc = fp*hpc + fc*hcc - thishs = fp*hps + fc*hcs + thishc = fp * hpc + fc * hcc + thishs = fp * hps + fc * hcs # apply time shift - hc = apply_fd_time_shift(thishc, tc+tshift, copy=False) - hs = apply_fd_time_shift(thishs, tc+tshift, copy=False) + hc = apply_fd_time_shift(thishc, tc + tshift, copy=False) + hs = apply_fd_time_shift(thishs, tc + tshift, copy=False) if self.recalib: # recalibrate with given calibration model - hc = self.recalib[detname].map_to_adjust(hc, - **self.current_params) - hs = self.recalib[detname].map_to_adjust(hs, - **self.current_params) + hc = self.recalib[detname].map_to_adjust(hc, **self.current_params) + hs = self.recalib[detname].map_to_adjust(hs, **self.current_params) h[detname] = (hc, hs) else: # no detector response, just use the + polarization - if 'tc' in self.current_params: - hpc = apply_fd_time_shift(hpc, self.current_params['tc']+tshift, - copy=False) - hps = apply_fd_time_shift(hps, self.current_params['tc']+tshift, - copy=False) - h['RF'] = (hpc, hps) + if "tc" in self.current_params: + hpc = apply_fd_time_shift( + hpc, self.current_params["tc"] + tshift, copy=False + ) + hps = apply_fd_time_shift( + hps, self.current_params["tc"] + tshift, copy=False + ) + h["RF"] = (hpc, hps) if self.gates is not None: # resize all to nearest power of 2 for ifo, (hc, hs) in h.items(): - hc.resize(ceilpow2(len(hc)-1) + 1) - hs.resize(ceilpow2(len(hs)-1) + 1) + hc.resize(ceilpow2(len(hc) - 1) + 1) + hs.resize(ceilpow2(len(hs) - 1) + 1) # apply gates to wfs - h[ifo] = (strain.gate_data(hc, self.gates[ifo]), - strain.gate_data(hs, self.gates[ifo])) + h[ifo] = ( + strain.gate_data(hc, self.gates[ifo]), + strain.gate_data(hs, self.gates[ifo]), + ) return h @staticmethod def select_rframe_generator(approximant, domain): - """Returns a radiation frame generator class based on the approximant + """ + Returns a radiation frame generator class based on the approximant string. """ return select_waveform_generator(approximant, domain) class FDomainDetFrameModesGenerator(BaseFDomainDetFrameGenerator): - r"""Generates frequency-domain waveform modes in a specific frame. + r""" + Generates frequency-domain waveform modes in a specific frame. Generates both polarizations of every waveform mode using the given radiation frame generator class, and applies the time shift. Detector @@ -1191,7 +1290,8 @@ class FDomainDetFrameModesGenerator(BaseFDomainDetFrameGenerator): function. """ - location_args = set(['tc', 'ra', 'dec']) + + location_args = set(["tc", "ra", "dec"]) """ set(['tc', 'ra', 'dec']): The set of location parameters. These are not passed to the rFrame generator class; instead, they are used to apply the detector response @@ -1210,7 +1310,8 @@ class FDomainDetFrameModesGenerator(BaseFDomainDetFrameGenerator): """ def generate(self, **kwargs): - """Generates and returns a waveform decompsed into separate modes. + """ + Generates and returns a waveform decompsed into separate modes. Returns ------- @@ -1219,61 +1320,67 @@ def generate(self, **kwargs): ``ulm, vlm`` are the frequency-domain representations of the real and imaginary parts, respectively, of the complex time series representation of the ``hlm``. + """ self.current_params.update(kwargs) - rfparams = {param: self.current_params[param] - for param in kwargs if param not in self.location_args} + rfparams = { + param: self.current_params[param] + for param in kwargs + if param not in self.location_args + } hlms = self.rframe_generator.generate(**rfparams) h = {det: {} for det in self.detectors} for mode in hlms: ulm, vlm = hlms[mode] if isinstance(ulm, TimeSeries): - df = self.current_params['delta_f'] + df = self.current_params["delta_f"] ulm = ulm.to_frequencyseries(delta_f=df) vlm = vlm.to_frequencyseries(delta_f=df) # time-domain waveforms will not be shifted so that the peak # amplitude happens at the end of the time series (as they are # for f-domain), so we add an additional shift to account for # it - tshift = 1./df - abs(ulm._epoch) + tshift = 1.0 / df - abs(ulm._epoch) else: - tshift = 0. + tshift = 0.0 ulm._epoch = vlm._epoch = self._epoch - if self.detector_names != ['RF']: + if self.detector_names != ["RF"]: for detname, det in self.detectors.items(): - refframe = self.current_params.get('tc_ref_frame', 'geocentric') - ra = self.current_params['ra'] - dec = self.current_params['dec'] - ref_tc = self.current_params['tc'] + refframe = self.current_params.get("tc_ref_frame", "geocentric") + ra = self.current_params["ra"] + dec = self.current_params["dec"] + ref_tc = self.current_params["tc"] tc = det.arrival_time(ref_tc, ra, dec, refframe) # apply time shift - detulm = apply_fd_time_shift(ulm, tc+tshift, copy=True) - detvlm = apply_fd_time_shift(vlm, tc+tshift, copy=True) + detulm = apply_fd_time_shift(ulm, tc + tshift, copy=True) + detvlm = apply_fd_time_shift(vlm, tc + tshift, copy=True) if self.recalib: # recalibrate with given calibration model detulm = self.recalib[detname].map_to_adjust( - detulm, **self.current_params) + detulm, **self.current_params + ) detvlm = self.recalib[detname].map_to_adjust( - detvlm, **self.current_params) + detvlm, **self.current_params + ) h[detname][mode] = (detulm, detvlm) else: # no detector response, just apply time shift - if 'tc' in self.current_params: - ulm = apply_fd_time_shift(ulm, - self.current_params['tc']+tshift, - copy=False) - vlm = apply_fd_time_shift(vlm, - self.current_params['tc']+tshift, - copy=False) - h['RF'][mode] = (ulm, vlm) + if "tc" in self.current_params: + ulm = apply_fd_time_shift( + ulm, self.current_params["tc"] + tshift, copy=False + ) + vlm = apply_fd_time_shift( + vlm, self.current_params["tc"] + tshift, copy=False + ) + h["RF"][mode] = (ulm, vlm) if self.gates is not None: # resize all to nearest power of 2 ulms = {} vlms = {} for det in h: ulm, vlm = h[det][mode] - ulm.resize(ceilpow2(len(ulm)-1) + 1) - vlm.resize(ceilpow2(len(vlm)-1) + 1) + ulm.resize(ceilpow2(len(ulm) - 1) + 1) + vlm.resize(ceilpow2(len(vlm) - 1) + 1) ulms[det] = ulm vlms[det] = vlm ulms = strain.apply_gates_to_fd(ulms, self.gates) @@ -1284,14 +1391,16 @@ def generate(self, **kwargs): @staticmethod def select_rframe_generator(approximant, domain): - """Returns a radiation frame generator class based on the approximant + """ + Returns a radiation frame generator class based on the approximant string. """ return select_waveform_modes_generator(approximant, domain) class FDomainDirectDetFrameGenerator(BaseCBCGenerator): - """Generates frequency-domain waveforms directly in the detector frame. + """ + Generates frequency-domain waveforms directly in the detector frame. Uses :py:func:`waveform.get_fd_det_waveform` as a generator function to create frequency-domain CBC waveforms that include the detector @@ -1299,6 +1408,7 @@ class FDomainDirectDetFrameGenerator(BaseCBCGenerator): For details, on methods and arguments, see :py:class:`BaseCBCGenerator`. """ + def __init__( self, rFrameGeneratorClass=None, @@ -1307,7 +1417,7 @@ def __init__( variable_args=(), gates=None, recalib=None, - **frozen_params + **frozen_params, ): if rFrameGeneratorClass is not None: @@ -1320,13 +1430,9 @@ def __init__( self.detectors = detectors if gates is not None: - raise RuntimeError( - f"{self.__class__.__name__} does not support `gates`" - ) + raise RuntimeError(f"{self.__class__.__name__} does not support `gates`") if recalib is not None: - raise RuntimeError( - f"{self.__class__.__name__} does not support `recalib`" - ) + raise RuntimeError(f"{self.__class__.__name__} does not support `recalib`") if detectors is None: raise ValueError( @@ -1334,9 +1440,7 @@ def __init__( ) super().__init__( - waveform.get_fd_det_waveform, - variable_args=variable_args, - **frozen_params + waveform.get_fd_det_waveform, variable_args=variable_args, **frozen_params ) def set_epoch(self, epoch): @@ -1345,7 +1449,8 @@ def set_epoch(self, epoch): @property def epoch(self): - """The GPS start time of the frequency series returned by the generate + """ + The GPS start time of the frequency series returned by the generate function. A time shift is applied to the waveform equal to tc-epoch. Update by using ``set_epoch`` """ @@ -1353,21 +1458,24 @@ def epoch(self): @staticmethod def select_rframe_generator(approximant): - """Returns the radiation frame generator. + """ + Returns the radiation frame generator. Returns ``None`` since this class does not support generating waveforms in the radiation frame. """ - return None + return def generate(self, **kwargs): - """Generates and returns a waveform in the detector frame. + """ + Generates and returns a waveform in the detector frame. Returns ------- dict : Dictionary of ``detector names -> h``, where is the waveform in the specified detector. + """ wfs = super().generate(ifos=self.detectors, **kwargs) for det in self.detectors: @@ -1384,6 +1492,7 @@ def generate(self, **kwargs): # ============================================================================= # + def get_td_generator(approximant, modes=False): """Returns the time-domain generator for the given approximant.""" if approximant in waveform.td_approximants(): @@ -1392,15 +1501,15 @@ def get_td_generator(approximant, modes=False): return TDomainCBCGenerator if approximant in ringdown.ringdown_td_approximants: - if approximant == 'TdQNMfromFinalMassSpin': + if approximant == "TdQNMfromFinalMassSpin": return TDomainMassSpinRingdownGenerator return TDomainFreqTauRingdownGenerator if approximant in supernovae.supernovae_td_approximants: return TDomainSupernovaeGenerator - raise ValueError(f"No time-domain generator found for " - "approximant: {approximant}") + raise ValueError("No time-domain generator found for approximant: {approximant}") + def get_fd_generator(approximant, modes=False): """Returns the frequency-domain generator for the given approximant.""" @@ -1410,15 +1519,18 @@ def get_fd_generator(approximant, modes=False): return FDomainCBCGenerator if approximant in ringdown.ringdown_fd_approximants: - if approximant == 'FdQNMfromFinalMassSpin': + if approximant == "FdQNMfromFinalMassSpin": return FDomainMassSpinRingdownGenerator return FDomainFreqTauRingdownGenerator - raise ValueError(f"No frequency-domain generator found for " - "approximant: {approximant}") + raise ValueError( + "No frequency-domain generator found for approximant: {approximant}" + ) + def select_waveform_generator(approximant, domain=None): - """Returns the single-IFO generator for the approximant. + """ + Returns the single-IFO generator for the approximant. Parameters ---------- @@ -1446,24 +1558,27 @@ def select_waveform_generator(approximant, domain=None): Get generator object: >>> from pycbc.waveform.generator import select_waveform_generator >>> select_waveform_generator(waveform.fd_approximants()[0]) - """ - if domain not in {None, 'td', 'fd'}: - raise ValueError(f"Invalid domain '{domain}'. " - "Must be one of: None, 'td', or 'fd'.") + """ + if domain not in {None, "td", "fd"}: + raise ValueError( + f"Invalid domain '{domain}'. Must be one of: None, 'td', or 'fd'." + ) - if domain == 'td': + if domain == "td": return get_td_generator(approximant) - elif domain == 'fd': + if domain == "fd": return get_fd_generator(approximant) - elif domain is None: + if domain is None: try: return get_fd_generator(approximant) except ValueError: return get_td_generator(approximant) + def select_waveform_modes_generator(approximant, domain=None): - """Returns the single-IFO modes generator for the approximant. + """ + Returns the single-IFO modes generator for the approximant. Parameters ---------- @@ -1478,17 +1593,18 @@ def select_waveform_modes_generator(approximant, domain=None): ------- generator : (PyCBC generator instance) A waveform modes generator object. - """ - if domain not in {None, 'td', 'fd'}: - raise ValueError(f"Invalid domain '{domain}'. " - "Must be one of: None, 'td', or 'fd'.") + """ + if domain not in {None, "td", "fd"}: + raise ValueError( + f"Invalid domain '{domain}'. Must be one of: None, 'td', or 'fd'." + ) - if domain == 'td': + if domain == "td": return get_td_generator(approximant, modes=True) - elif domain == 'fd': + if domain == "fd": return get_fd_generator(approximant, modes=True) - elif domain is None: + if domain is None: try: return get_fd_generator(approximant, modes=True) except ValueError: diff --git a/pycbc/waveform/multiband.py b/pycbc/waveform/multiband.py index 14c85e19693..f622453bd33 100644 --- a/pycbc/waveform/multiband.py +++ b/pycbc/waveform/multiband.py @@ -1,12 +1,13 @@ -""" Tools and functions to calculate interpolate waveforms using multi-banding -""" +"""Tools and functions to calculate interpolate waveforms using multi-banding""" + import numpy from pycbc.types import TimeSeries, zeros def multiband_fd_waveform(bands=None, lengths=None, overlap=0, **p): - """ Generate a fourier domain waveform using multibanding + """ + Generate a fourier domain waveform using multibanding Speed up generation of a fouerier domain waveform using multibanding. This allows for multi-rate sampling of the frequeny space. Each band is @@ -36,19 +37,20 @@ def multiband_fd_waveform(bands=None, lengths=None, overlap=0, **p): Plus polarization hc: pycbc.type.FrequencySeries Cross polarization + """ from pycbc.waveform import get_fd_waveform if isinstance(bands, str): - bands = [float(s) for s in bands.split(' ')] + bands = [float(s) for s in bands.split(" ")] if isinstance(lengths, str): - lengths = [float(s) for s in lengths.split(' ')] + lengths = [float(s) for s in lengths.split(" ")] - p['approximant'] = p['base_approximant'] - df = p['delta_f'] - fmax = p['f_final'] - flow = p['f_lower'] + p["approximant"] = p["base_approximant"] + df = p["delta_f"] + fmax = p["f_final"] + flow = p["f_lower"] bands = [flow] + bands + [fmax] dfs = [df] + [1.0 / l for l in lengths] @@ -56,13 +58,15 @@ def multiband_fd_waveform(bands=None, lengths=None, overlap=0, **p): dt = 1.0 / (2.0 * fmax) tlen = int(1.0 / dt / df) flen = tlen / 2 + 1 - wf_plus = TimeSeries(zeros(tlen, dtype=numpy.float32), - copy=False, delta_t=dt, epoch=-1.0/df) - wf_cross = TimeSeries(zeros(tlen, dtype=numpy.float32), - copy=False, delta_t=dt, epoch=-1.0/df) + wf_plus = TimeSeries( + zeros(tlen, dtype=numpy.float32), copy=False, delta_t=dt, epoch=-1.0 / df + ) + wf_cross = TimeSeries( + zeros(tlen, dtype=numpy.float32), copy=False, delta_t=dt, epoch=-1.0 / df + ) # Iterate over the sub-bands - for i in range(len(lengths)+1): + for i in range(len(lengths) + 1): taper_start = taper_end = False if i != 0: taper_start = True @@ -71,17 +75,17 @@ def multiband_fd_waveform(bands=None, lengths=None, overlap=0, **p): # Generate waveform for sub-band of full waveform start = bands[i] - stop = bands[i+1] + stop = bands[i + 1] p2 = p.copy() - p2['delta_f'] = dfs[i] - p2['f_lower'] = start - p2['f_final'] = stop + p2["delta_f"] = dfs[i] + p2["f_lower"] = start + p2["f_final"] = stop if taper_start: - p2['f_lower'] -= overlap / 2.0 + p2["f_lower"] -= overlap / 2.0 if taper_end: - p2['f_final'] += overlap / 2.0 + p2["f_final"] += overlap / 2.0 tlen = int(1.0 / dt / dfs[i]) flen = tlen / 2 + 1 @@ -89,24 +93,26 @@ def multiband_fd_waveform(bands=None, lengths=None, overlap=0, **p): hp, hc = get_fd_waveform(**p2) # apply window function to smooth over transition regions - kmin = int(p2['f_lower'] / dfs[i]) - kmax = int(p2['f_final'] / dfs[i]) + kmin = int(p2["f_lower"] / dfs[i]) + kmax = int(p2["f_final"] / dfs[i]) taper = numpy.hanning(int(overlap * 2 / dfs[i])) for wf, h in zip([wf_plus, wf_cross], [hp, hc]): h = h.astype(numpy.complex64) if taper_start: - h[kmin:kmin + len(taper) // 2] *= taper[:len(taper)//2] + h[kmin : kmin + len(taper) // 2] *= taper[: len(taper) // 2] if taper_end: l, r = kmax - (len(taper) - len(taper) // 2), kmax - h[l:r] *= taper[len(taper)//2:] + h[l:r] *= taper[len(taper) // 2 :] # add frequency band to total and use fft to interpolate h.resize(flen) h = h.to_timeseries() - wf[len(wf)-len(h):] += h + wf[len(wf) - len(h) :] += h - return (wf_plus.to_frequencyseries().astype(hp.dtype), - wf_cross.to_frequencyseries().astype(hp.dtype)) + return ( + wf_plus.to_frequencyseries().astype(hp.dtype), + wf_cross.to_frequencyseries().astype(hp.dtype), + ) diff --git a/pycbc/waveform/nltides.py b/pycbc/waveform/nltides.py index cb83bb9f057..7964061ed9e 100644 --- a/pycbc/waveform/nltides.py +++ b/pycbc/waveform/nltides.py @@ -1,12 +1,14 @@ -""" Utilities for introducing nonlinear tidal effects into waveform approximants -""" +"""Utilities for introducing nonlinear tidal effects into waveform approximants""" + import numpy import pycbc.conversions from pycbc.constants import PI + def nltides_fourier_phase_difference(f, delta_f, f0, amplitude, n, m1, m2): - r"""Calculate the change to the Fourier phase change due + r""" + Calculate the change to the Fourier phase change due to non-linear tides. Note that the Fourier phase Psi(f) is not the same as the gravitational-wave phase phi(f) and is computed by @@ -33,52 +35,65 @@ def nltides_fourier_phase_difference(f, delta_f, f0, amplitude, n, m1, m2): ------- delta_psi: numpy.array Fourier phase as a function of frequency - """ - kmin = int(f0/delta_f) + """ + kmin = int(f0 / delta_f) kmax = len(f) - f_ref, t_of_f_factor, phi_of_f_factor = \ - pycbc.conversions.nltides_coefs(amplitude, n, m1, m2) + f_ref, t_of_f_factor, phi_of_f_factor = pycbc.conversions.nltides_coefs( + amplitude, n, m1, m2 + ) # Fourier phase shift below f0 from \Delta \phi(f) delta_psi_f_le_f0 = numpy.ones(kmin) - delta_psi_f_le_f0 *= - phi_of_f_factor * (f0/f_ref)**(n-3.) + delta_psi_f_le_f0 *= -phi_of_f_factor * (f0 / f_ref) ** (n - 3.0) # Fourier phase shift above f0 from \Delta \phi(f) - delta_psi_f_gt_f0 = - phi_of_f_factor * (f[kmin:kmax]/f_ref)**(n-3.) + delta_psi_f_gt_f0 = -phi_of_f_factor * (f[kmin:kmax] / f_ref) ** (n - 3.0) # Fourier phase shift below f0 from 2 pi f \Delta t(f) - delta_psi_f_le_f0 += 2.0 * PI * f[0:kmin] * t_of_f_factor * \ - (f0/f_ref)**(n-4.) + delta_psi_f_le_f0 += ( + 2.0 * PI * f[0:kmin] * t_of_f_factor * (f0 / f_ref) ** (n - 4.0) + ) # Fourier phase shift above f0 from 2 pi f \Delta t(f) - delta_psi_f_gt_f0 += 2.0 * PI * f[kmin:kmax] * t_of_f_factor * \ - (f[kmin:kmax]/f_ref)**(n-4.) + delta_psi_f_gt_f0 += ( + 2.0 * PI * f[kmin:kmax] * t_of_f_factor * (f[kmin:kmax] / f_ref) ** (n - 4.0) + ) # Return the shift to the Fourier phase return numpy.concatenate((delta_psi_f_le_f0, delta_psi_f_gt_f0), axis=0) def nonlinear_tidal_spa(**kwds): - """Generates a frequency-domain waveform that implements the + """ + Generates a frequency-domain waveform that implements the TaylorF2+NL tide model described in https://arxiv.org/abs/1808.07013 """ - from pycbc import waveform from pycbc.types import Array # We start with the standard TaylorF2 based waveform - kwds.pop('approximant') + kwds.pop("approximant") hp, hc = waveform.get_fd_waveform(approximant="TaylorF2", **kwds) # Add the phasing difference from the nonlinear tides f = numpy.arange(len(hp)) * hp.delta_f - pd = Array(numpy.exp(-1.0j * nltides_fourier_phase_difference(f, - hp.delta_f, - kwds['f0'], kwds['amplitude'], kwds['n'], - kwds['mass1'], kwds['mass2'])), - dtype=hp.dtype) + pd = Array( + numpy.exp( + -1.0j + * nltides_fourier_phase_difference( + f, + hp.delta_f, + kwds["f0"], + kwds["amplitude"], + kwds["n"], + kwds["mass1"], + kwds["mass2"], + ) + ), + dtype=hp.dtype, + ) hp *= pd hc *= pd return hp, hc diff --git a/pycbc/waveform/parameters.py b/pycbc/waveform/parameters.py index 428bbcdedd9..de501ac4e49 100644 --- a/pycbc/waveform/parameters.py +++ b/pycbc/waveform/parameters.py @@ -22,10 +22,10 @@ # # ============================================================================= # -"""Classes to define common parameters used for waveform generation. -""" +"""Classes to define common parameters used for waveform generation.""" from collections import OrderedDict + try: from collections import UserList except ImportError: @@ -39,13 +39,16 @@ # ============================================================================= # + class Parameter(str): - """A class that stores information about a parameter. This is done by + """ + A class that stores information about a parameter. This is done by sub-classing string, adding additional attributes. """ - def __new__(cls, name, dtype=None, default=None, label=None, - description="No description."): + def __new__( + cls, name, dtype=None, default=None, label=None, description="No description." + ): obj = str.__new__(cls, name) obj.name = name obj.dtype = dtype @@ -54,28 +57,35 @@ def __new__(cls, name, dtype=None, default=None, label=None, obj.description = description return obj - def docstr(self, prefix='', include_label=True): - """Returns a string summarizing the parameter. Format is: + def docstr(self, prefix="", include_label=True): + """ + Returns a string summarizing the parameter. Format is: ``name`` : {``default``, ``dtype``} ``description`` Label: ``label``. """ - dtype_str = str(self.dtype).replace("", '') - dtype_str = dtype_str.replace("", "") + dtype_str = dtype_str.replace(" *htilde, int kmin, double delta_f, double PI, +phenomC_kernel = ElementwiseKernel( + """pycuda::complex *htilde, int kmin, double delta_f, double PI, double eta, double Xi, double distance, double m_sec, double piM, double Mfrd, double pfaN, double pfa2, double pfa3, double pfa4, @@ -125,11 +124,13 @@ double A2, double A3, double A4, double A5, double A5imag, double A6, double A6log, double A6imag, double g1, double del1, double del2, double Q""", - phenomC_text, "phenomC_kernel", - preamble=preamble) + phenomC_text, + "phenomC_kernel", + preamble=preamble, +) -def FinalSpin( Xi, eta ): +def FinalSpin(Xi, eta): """Computes the spin of the final BH that gets formed after merger. This is done usingn Eq 5-6 of arXiv:0710.3345""" s4 = -0.129 s5 = -0.384 @@ -137,39 +138,52 @@ def FinalSpin( Xi, eta ): t2 = -3.454 t3 = 2.353 etaXi = eta * Xi - eta2 = eta*eta - finspin = (Xi + s4*Xi*etaXi + s5*etaXi*eta + t0*etaXi + 2.*(3.**0.5)*eta + t2*eta2 + t3*eta2*eta) + eta2 = eta * eta + finspin = ( + Xi + + s4 * Xi * etaXi + + s5 * etaXi * eta + + t0 * etaXi + + 2.0 * (3.0**0.5) * eta + + t2 * eta2 + + t3 * eta2 * eta + ) if finspin > 1.0: raise ValueError("Value of final spin > 1.0. Aborting") - else: - return finspin + return finspin + -def fRD( a, M): +def fRD(a, M): """Calculate the ring-down frequency for the final Kerr BH. Using Eq. 5.5 of Main paper""" - f = (C_SI**3.0 / (2.0*PI*G_SI*M*MSUN_SI)) * (1.5251 - 1.1568*(1.0-a)**0.1292) + f = (C_SI**3.0 / (2.0 * PI * G_SI * M * MSUN_SI)) * ( + 1.5251 - 1.1568 * (1.0 - a) ** 0.1292 + ) return f -def Qa( a ): + +def Qa(a): """Calculate the quality factor of ring-down, using Eq 5.6 of Main paper""" - return (0.7 + 1.4187*(1.0-a)**-0.4990) + return 0.7 + 1.4187 * (1.0 - a) ** -0.4990 -#Functions to calculate the Tanh window, defined in Eq 5.8 of the main paper + +# Functions to calculate the Tanh window, defined in Eq 5.8 of the main paper def imrphenomc_tmplt(**kwds): - """ Return an IMRPhenomC waveform using CUDA to generate the phase and amplitude - Main Paper: arXiv:1005.3306 + """ + Return an IMRPhenomC waveform using CUDA to generate the phase and amplitude + Main Paper: arXiv:1005.3306 """ # Pull out the input arguments - f_min = float128(kwds['f_lower']) - f_max = float128(kwds['f_final']) - delta_f = float128(kwds['delta_f']) - distance = float128(kwds['distance']) - mass1 = float128(kwds['mass1']) - mass2 = float128(kwds['mass2']) - spin1z = float128(kwds['spin1z']) - spin2z = float128(kwds['spin2z']) - - if 'out' in kwds: - out = kwds['out'] + f_min = float128(kwds["f_lower"]) + f_max = float128(kwds["f_final"]) + delta_f = float128(kwds["delta_f"]) + distance = float128(kwds["distance"]) + mass1 = float128(kwds["mass1"]) + mass2 = float128(kwds["mass2"]) + spin1z = float128(kwds["spin1z"]) + spin2z = float128(kwds["spin2z"]) + + if "out" in kwds: + out = kwds["out"] else: out = None @@ -177,15 +191,17 @@ def imrphenomc_tmplt(**kwds): M = mass1 + mass2 eta = mass1 * mass2 / (M * M) Xi = (mass1 * spin1z / M) + (mass2 * spin2z / M) - Xisum = 2.*Xi - Xiprod = Xi*Xi - Xi2 = Xi*Xi + Xisum = 2.0 * Xi + Xiprod = Xi * Xi + Xi2 = Xi * Xi m_sec = M * MTSUN_SI piM = PI * m_sec ## The units of distance given as input is taken to pe Mpc. Converting to SI - distance *= (1.0e6 * PC_SI / (2. * sqrt(5. / (64.*PI)) * M * MRSUN_SI * M * MTSUN_SI)) + distance *= ( + 1.0e6 * PC_SI / (2.0 * sqrt(5.0 / (64.0 * PI)) * M * MRSUN_SI * M * MTSUN_SI) + ) # Check if the value of f_max is correctly given, else replace with the fCut # used in the PhenomB code in lalsimulation. The various coefficients come @@ -205,38 +221,38 @@ def imrphenomc_tmplt(**kwds): z201 = 5.962e-01 z202 = -5.600e-02 z211 = 1.520e-01 - z210 = -2.970e+00 - z220 = 1.312e+01 - - z301 = -3.283e+01 - z302 = 8.859e+00 - z311 = 2.931e+01 - z310 = 7.954e+01 - z320 = -4.349e+02 - - z401 = 1.619e+02 - z402 = -4.702e+01 - z411 = -1.751e+02 - z410 = -3.225e+02 - z420 = 1.587e+03 - - z501 = -6.320e+02 - z502 = 2.463e+02 - z511 = 1.048e+03 - z510 = 3.355e+02 - z520 = -5.115e+03 - - z601 = -4.809e+01 - z602 = -3.643e+02 - z611 = -5.215e+02 - z610 = 1.870e+03 - z620 = 7.354e+02 - - z701 = 4.149e+00 - z702 = -4.070e+00 - z711 = -8.752e+01 - z710 = -4.897e+01 - z720 = 6.665e+02 + z210 = -2.970e00 + z220 = 1.312e01 + + z301 = -3.283e01 + z302 = 8.859e00 + z311 = 2.931e01 + z310 = 7.954e01 + z320 = -4.349e02 + + z401 = 1.619e02 + z402 = -4.702e01 + z411 = -1.751e02 + z410 = -3.225e02 + z420 = 1.587e03 + + z501 = -6.320e02 + z502 = 2.463e02 + z511 = 1.048e03 + z510 = 3.355e02 + z520 = -5.115e03 + + z601 = -4.809e01 + z602 = -3.643e02 + z611 = -5.215e02 + z610 = 1.870e03 + z620 = 7.354e02 + + z701 = 4.149e00 + z702 = -4.070e00 + z711 = -8.752e01 + z710 = -4.897e01 + z720 = 6.665e02 z801 = -5.472e-02 z802 = 2.094e-02 @@ -244,13 +260,13 @@ def imrphenomc_tmplt(**kwds): z810 = 1.151e-01 z820 = 9.640e-01 - z901 = -1.235e+00 + z901 = -1.235e00 z902 = 3.423e-01 - z911 = 6.062e+00 - z910 = 5.949e+00 - z920 = -1.069e+01 + z911 = 6.062e00 + z910 = 5.949e00 + z920 = -1.069e01 - eta2 = eta*eta + eta2 = eta * eta Xi2 = Xiprod # Calculate alphas, gamma, deltas from Table II and Eq 5.14 of Main paper @@ -267,11 +283,11 @@ def imrphenomc_tmplt(**kwds): del2 = z901 * Xi + z902 * Xi2 + z911 * eta * Xi + z910 * eta + z920 * eta2 # Get the spin of the final BH - afin = FinalSpin( Xi, eta ) - Q = Qa( abs(afin) ) + afin = FinalSpin(Xi, eta) + Q = Qa(abs(afin)) # Get the fRD - frd = fRD( abs(afin), M) + frd = fRD(abs(afin), M) Mfrd = frd * m_sec # Define the frequencies where SPA->PM->RD @@ -287,88 +303,152 @@ def imrphenomc_tmplt(**kwds): # Now use this frequency for calculation of betas # calculate beta1 and beta2, that appear in Eq 5.7 in the main paper. - b2 = ((-5./3.)* a1 * pow(Mfrd,(-8./3.)) - a2/(Mfrd*Mfrd) - \ - (a3/3.)*pow(Mfrd,(-4./3.)) + (2./3.)* a5 * pow(Mfrd,(-1./3.)) + a6)/eta - - psiPMrd = (a1 * pow(Mfrd,(-5./3.)) + a2/Mfrd + a3 * pow(Mfrd,(-1./3.)) + \ - a4 + a5 * pow(Mfrd,(2./3.)) + a6 * Mfrd)/eta + b2 = ( + (-5.0 / 3.0) * a1 * pow(Mfrd, (-8.0 / 3.0)) + - a2 / (Mfrd * Mfrd) + - (a3 / 3.0) * pow(Mfrd, (-4.0 / 3.0)) + + (2.0 / 3.0) * a5 * pow(Mfrd, (-1.0 / 3.0)) + + a6 + ) / eta + + psiPMrd = ( + a1 * pow(Mfrd, (-5.0 / 3.0)) + + a2 / Mfrd + + a3 * pow(Mfrd, (-1.0 / 3.0)) + + a4 + + a5 * pow(Mfrd, (2.0 / 3.0)) + + a6 * Mfrd + ) / eta b1 = psiPMrd - (b2 * Mfrd) ### Calculate the PN coefficients, Eq A3 - A5 of main paper ### - pfaN = 3.0/(128.0 * eta) - pfa2 = (3715./756.) + (55.*eta/9.0) - pfa3 = -16.0*PI + (113./3.)*Xi - 38.*eta*Xisum/3. - pfa4 = (152.93365/5.08032) - 50.*Xi2 + eta*(271.45/5.04 + 1.25*Xiprod) + \ - 3085.*eta2/72. - pfa5 = PI*(386.45/7.56 - 65.*eta/9.) - \ - Xi*(735.505/2.268 + 130.*eta/9.) + Xisum*(1285.0*eta/8.1 + 170.*eta2/9.) - \ - 10.*Xi2*Xi/3. + 10.*eta*Xi*Xiprod - pfa6 = 11583.231236531/4.694215680 - 640.0*PI*PI/3. - \ - 6848.0*GAMMA/21. - 684.8*log(64.)/6.3 + \ - eta*(2255.*PI*PI/12. - 15737.765635/3.048192) + \ - 76.055*eta2/1.728 - (127.825*eta2*eta/1.296) + \ - 2920.*PI*Xi/3. - (175. - 1490.*eta)*Xi2/3. - \ - (1120.*PI/3. - 1085.*Xi/3.)*eta*Xisum + \ - (269.45*eta/3.36 - 2365.*eta2/6.)*Xiprod - - pfa6log = -6848./63. - - pfa7 = PI*(770.96675/2.54016 + 378.515*eta/1.512 - 740.45*eta2/7.56) - \ - Xi*(20373.952415/3.048192 + 1509.35*eta/2.24 - 5786.95*eta2/4.32) + \ - Xisum*(4862.041225*eta/1.524096 + 1189.775*eta2/1.008 - 717.05*eta2*eta/2.16 - 830.*eta*Xi2/3. + 35.*eta2*Xiprod/3.) - \ - 560.*PI*Xi2 + 20.*PI*eta*Xiprod + \ - Xi2*Xi*(945.55/1.68 - 85.*eta) + Xi*Xiprod*(396.65*eta/1.68 + 255.*eta2) - - - xdotaN = 64.*eta/5. - xdota2 = -7.43/3.36 - 11.*eta/4. - xdota3 = 4.*PI - 11.3*Xi/1.2 + 19.*eta*Xisum/6. - xdota4 = 3.4103/1.8144 + 5*Xi2 + eta*(13.661/2.016 - Xiprod/8.) + 5.9*eta2/1.8 - xdota5 = -PI*(41.59/6.72 + 189.*eta/8.) - Xi*(31.571/1.008 - 116.5*eta/2.4) + \ - Xisum*(21.863*eta/1.008 - 79.*eta2/6.) - 3*Xi*Xi2/4. + \ - 9.*eta*Xi*Xiprod/4. - xdota6 = 164.47322263/1.39708800 - 17.12*GAMMA/1.05 + \ - 16.*PI*PI/3 - 8.56*log(16.)/1.05 + \ - eta*(45.1*PI*PI/4.8 - 561.98689/2.17728) + \ - 5.41*eta2/8.96 - 5.605*eta*eta2/2.592 - 80.*PI*Xi/3. + \ - eta*Xisum*(20.*PI/3. - 113.5*Xi/3.6) + \ - Xi2*(64.153/1.008 - 45.7*eta/3.6) - \ - Xiprod*(7.87*eta/1.44 - 30.37*eta2/1.44) - - xdota6log = -856./105. - - xdota7 = -PI*(4.415/4.032 - 358.675*eta/6.048 - 91.495*eta2/1.512) - \ - Xi*(252.9407/2.7216 - 845.827*eta/6.048 + 415.51*eta2/8.64) + \ - Xisum*(158.0239*eta/5.4432 - 451.597*eta2/6.048 + 20.45*eta2*eta/4.32 + 107.*eta*Xi2/6. - 5.*eta2*Xiprod/24.) + \ - 12.*PI*Xi2 - Xi2*Xi*(150.5/2.4 + eta/8.) + \ - Xi*Xiprod*(10.1*eta/2.4 + 3.*eta2/8.) - - - AN = 8.*eta*sqrt(PI/5.) - A2 = (-107. + 55.*eta)/42. - A3 = 2.*PI - 4.*Xi/3. + 2.*eta*Xisum/3. - A4 = -2.173/1.512 - eta*(10.69/2.16 - 2.*Xiprod) + 2.047*eta2/1.512 - A5 = -10.7*PI/2.1 + eta*(3.4*PI/2.1) - - A5imag = -24.*eta - - A6 = 270.27409/6.46800 - 8.56*GAMMA/1.05 + \ - 2.*PI*PI/3. + \ - eta*(4.1*PI*PI/9.6 - 27.8185/3.3264) - \ - 20.261*eta2/2.772 + 11.4635*eta*eta2/9.9792 - \ - 4.28*log(16.)/1.05 - - A6log = -428./105. - - A6imag = 4.28*PI/1.05 + pfaN = 3.0 / (128.0 * eta) + pfa2 = (3715.0 / 756.0) + (55.0 * eta / 9.0) + pfa3 = -16.0 * PI + (113.0 / 3.0) * Xi - 38.0 * eta * Xisum / 3.0 + pfa4 = ( + (152.93365 / 5.08032) + - 50.0 * Xi2 + + eta * (271.45 / 5.04 + 1.25 * Xiprod) + + 3085.0 * eta2 / 72.0 + ) + pfa5 = ( + PI * (386.45 / 7.56 - 65.0 * eta / 9.0) + - Xi * (735.505 / 2.268 + 130.0 * eta / 9.0) + + Xisum * (1285.0 * eta / 8.1 + 170.0 * eta2 / 9.0) + - 10.0 * Xi2 * Xi / 3.0 + + 10.0 * eta * Xi * Xiprod + ) + pfa6 = ( + 11583.231236531 / 4.694215680 + - 640.0 * PI * PI / 3.0 + - 6848.0 * GAMMA / 21.0 + - 684.8 * log(64.0) / 6.3 + + eta * (2255.0 * PI * PI / 12.0 - 15737.765635 / 3.048192) + + 76.055 * eta2 / 1.728 + - (127.825 * eta2 * eta / 1.296) + + 2920.0 * PI * Xi / 3.0 + - (175.0 - 1490.0 * eta) * Xi2 / 3.0 + - (1120.0 * PI / 3.0 - 1085.0 * Xi / 3.0) * eta * Xisum + + (269.45 * eta / 3.36 - 2365.0 * eta2 / 6.0) * Xiprod + ) + + pfa6log = -6848.0 / 63.0 + + pfa7 = ( + PI * (770.96675 / 2.54016 + 378.515 * eta / 1.512 - 740.45 * eta2 / 7.56) + - Xi * (20373.952415 / 3.048192 + 1509.35 * eta / 2.24 - 5786.95 * eta2 / 4.32) + + Xisum + * ( + 4862.041225 * eta / 1.524096 + + 1189.775 * eta2 / 1.008 + - 717.05 * eta2 * eta / 2.16 + - 830.0 * eta * Xi2 / 3.0 + + 35.0 * eta2 * Xiprod / 3.0 + ) + - 560.0 * PI * Xi2 + + 20.0 * PI * eta * Xiprod + + Xi2 * Xi * (945.55 / 1.68 - 85.0 * eta) + + Xi * Xiprod * (396.65 * eta / 1.68 + 255.0 * eta2) + ) + + xdotaN = 64.0 * eta / 5.0 + xdota2 = -7.43 / 3.36 - 11.0 * eta / 4.0 + xdota3 = 4.0 * PI - 11.3 * Xi / 1.2 + 19.0 * eta * Xisum / 6.0 + xdota4 = ( + 3.4103 / 1.8144 + + 5 * Xi2 + + eta * (13.661 / 2.016 - Xiprod / 8.0) + + 5.9 * eta2 / 1.8 + ) + xdota5 = ( + -PI * (41.59 / 6.72 + 189.0 * eta / 8.0) + - Xi * (31.571 / 1.008 - 116.5 * eta / 2.4) + + Xisum * (21.863 * eta / 1.008 - 79.0 * eta2 / 6.0) + - 3 * Xi * Xi2 / 4.0 + + 9.0 * eta * Xi * Xiprod / 4.0 + ) + xdota6 = ( + 164.47322263 / 1.39708800 + - 17.12 * GAMMA / 1.05 + + 16.0 * PI * PI / 3 + - 8.56 * log(16.0) / 1.05 + + eta * (45.1 * PI * PI / 4.8 - 561.98689 / 2.17728) + + 5.41 * eta2 / 8.96 + - 5.605 * eta * eta2 / 2.592 + - 80.0 * PI * Xi / 3.0 + + eta * Xisum * (20.0 * PI / 3.0 - 113.5 * Xi / 3.6) + + Xi2 * (64.153 / 1.008 - 45.7 * eta / 3.6) + - Xiprod * (7.87 * eta / 1.44 - 30.37 * eta2 / 1.44) + ) + + xdota6log = -856.0 / 105.0 + + xdota7 = ( + -PI * (4.415 / 4.032 - 358.675 * eta / 6.048 - 91.495 * eta2 / 1.512) + - Xi * (252.9407 / 2.7216 - 845.827 * eta / 6.048 + 415.51 * eta2 / 8.64) + + Xisum + * ( + 158.0239 * eta / 5.4432 + - 451.597 * eta2 / 6.048 + + 20.45 * eta2 * eta / 4.32 + + 107.0 * eta * Xi2 / 6.0 + - 5.0 * eta2 * Xiprod / 24.0 + ) + + 12.0 * PI * Xi2 + - Xi2 * Xi * (150.5 / 2.4 + eta / 8.0) + + Xi * Xiprod * (10.1 * eta / 2.4 + 3.0 * eta2 / 8.0) + ) + + AN = 8.0 * eta * sqrt(PI / 5.0) + A2 = (-107.0 + 55.0 * eta) / 42.0 + A3 = 2.0 * PI - 4.0 * Xi / 3.0 + 2.0 * eta * Xisum / 3.0 + A4 = -2.173 / 1.512 - eta * (10.69 / 2.16 - 2.0 * Xiprod) + 2.047 * eta2 / 1.512 + A5 = -10.7 * PI / 2.1 + eta * (3.4 * PI / 2.1) + + A5imag = -24.0 * eta + + A6 = ( + 270.27409 / 6.46800 + - 8.56 * GAMMA / 1.05 + + 2.0 * PI * PI / 3.0 + + eta * (4.1 * PI * PI / 9.6 - 27.8185 / 3.3264) + - 20.261 * eta2 / 2.772 + + 11.4635 * eta * eta2 / 9.9792 + - 4.28 * log(16.0) / 1.05 + ) + + A6log = -428.0 / 105.0 + + A6imag = 4.28 * PI / 1.05 ### Define other parameters needed by waveform generation ### kmin = int(f_min / delta_f) kmax = int(f_max / delta_f) - n = kmax + 1; - + n = kmax + 1 if not out: - htilde = FrequencySeries(zeros(n,dtype=numpy.complex128), delta_f=delta_f, copy=False) + htilde = FrequencySeries( + zeros(n, dtype=numpy.complex128), delta_f=delta_f, copy=False + ) else: if type(out) is not Array: raise TypeError("Output must be an instance of Array") @@ -378,16 +458,61 @@ def imrphenomc_tmplt(**kwds): raise TypeError("Output array is the wrong dtype") htilde = FrequencySeries(out, delta_f=delta_f, copy=False) - phenomC_kernel(htilde.data[kmin:kmax], kmin, delta_f, PI, eta, Xi, distance, - m_sec, piM, Mfrd, - pfaN, pfa2, pfa3, pfa4, pfa5, pfa6, pfa6log, pfa7, - a1, a2, a3, a4, a5, a6, b1, b2, - Mf1, Mf2, Mf0, d1, d2, d0, - xdota2, xdota3, xdota4, xdota5, xdota6, xdota6log, - xdota7, xdotaN, AN, A2, A3, A4, A5, - A5imag, A6, A6log, A6imag, - g1, del1, del2, Q ) + phenomC_kernel( + htilde.data[kmin:kmax], + kmin, + delta_f, + PI, + eta, + Xi, + distance, + m_sec, + piM, + Mfrd, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfa6, + pfa6log, + pfa7, + a1, + a2, + a3, + a4, + a5, + a6, + b1, + b2, + Mf1, + Mf2, + Mf0, + d1, + d2, + d0, + xdota2, + xdota3, + xdota4, + xdota5, + xdota6, + xdota6log, + xdota7, + xdotaN, + AN, + A2, + A3, + A4, + A5, + A5imag, + A6, + A6log, + A6imag, + g1, + del1, + del2, + Q, + ) hp = htilde hc = htilde * 1j return hp, hc - diff --git a/pycbc/waveform/ringdown.py b/pycbc/waveform/ringdown.py index b9018f190ef..0391e6f8485 100644 --- a/pycbc/waveform/ringdown.py +++ b/pycbc/waveform/ringdown.py @@ -22,28 +22,27 @@ # # ============================================================================= # -"""Generate ringdown templates in the time and frequency domain. -""" +"""Generate ringdown templates in the time and frequency domain.""" import numpy from pycbc.libutils import import_optional -pykerr = import_optional('pykerr') -lal = import_optional('lal') -from pycbc.types import (TimeSeries, FrequencySeries, float64, complex128, - zeros) -from pycbc.waveform.waveform import get_obj_attrs + +pykerr = import_optional("pykerr") +lal = import_optional("lal") +from pycbc.constants import C_SI, G_SI, MSUN_SI, PC_SI from pycbc.conversions import get_lm_f0tau_allmodes -from pycbc.constants import MSUN_SI, G_SI, C_SI, PC_SI +from pycbc.types import FrequencySeries, TimeSeries, complex128, float64, zeros +from pycbc.waveform.waveform import get_obj_attrs -qnm_required_args = ['f_0', 'tau', 'amp', 'phi'] -mass_spin_required_args = ['final_mass','final_spin', 'lmns', 'inclination'] -freqtau_required_args = ['lmns'] -td_args = {'delta_t': None, 't_final': None, 'taper': False} -fd_args = {'t_0': 0, 'delta_f': None, 'f_lower': 0, 'f_final': None} +qnm_required_args = ["f_0", "tau", "amp", "phi"] +mass_spin_required_args = ["final_mass", "final_spin", "lmns", "inclination"] +freqtau_required_args = ["lmns"] +td_args = {"delta_t": None, "t_final": None, "taper": False} +fd_args = {"t_0": 0, "delta_f": None, "f_lower": 0, "f_final": None} -max_freq = 16384/2. -min_dt = 1. / (2 * max_freq) +max_freq = 16384 / 2.0 +min_dt = 1.0 / (2 * max_freq) pi = numpy.pi two_pi = 2 * numpy.pi pi_sq = numpy.pi * numpy.pi @@ -51,8 +50,10 @@ # Input parameters ############################################################ + def props(obj, required, domain_args, **kwargs): - """ Return a dictionary built from the combination of defaults, kwargs, + """ + Return a dictionary built from the combination of defaults, kwargs, and the attributes of the given object. """ # Get the attributes of the template object @@ -66,12 +67,14 @@ def props(obj, required, domain_args, **kwargs): # Check if the required arguments are given for arg in required: if arg not in input_params: - raise ValueError('Please provide ' + str(arg)) + raise ValueError("Please provide " + str(arg)) return input_params + def format_lmns(lmns): - """Checks if the format of the parameter lmns is correct, returning the + """ + Checks if the format of the parameter lmns is correct, returning the appropriate format if not, and raise an error if nmodes=0. The required format for the ringdown approximants is a list of lmn modes @@ -85,20 +88,21 @@ def format_lmns(lmns): will return the appropriate list of strings. If a different format is given, raise an error. """ - # Catch case of lmns given as float (as int injection values are cast # to float by pycbc_create_injections), cast to int, then string if isinstance(lmns, float): lmns = str(int(lmns)) # Case 1: the lmns are given as a string, e.g. '221 331' if isinstance(lmns, str): - lmns = lmns.split(' ') + lmns = lmns.split(" ") # Case 2: the lmns are given as strings in a list, e.g. ['221', '331'] elif isinstance(lmns, list): pass else: - raise ValueError('Format of parameter lmns not recognized. See ' - 'approximant documentation for more info.') + raise ValueError( + "Format of parameter lmns not recognized. See " + "approximant documentation for more info." + ) out = [] # Cycle over the lmns to ensure that we get back a list of strings that @@ -111,55 +115,61 @@ def format_lmns(lmns): # Try to convert to int and then str, to ensure the right format lmn = str(int(lmn)) if len(lmn) != 3: - raise ValueError('Format of parameter lmns not recognized. See ' - 'approximant documentation for more info.') - elif int(lmn[2]) == 0: - raise ValueError('Number of overtones (nmodes) must be greater ' - 'than zero in lmn={}.'.format(lmn)) + raise ValueError( + "Format of parameter lmns not recognized. See " + "approximant documentation for more info." + ) + if int(lmn[2]) == 0: + raise ValueError( + "Number of overtones (nmodes) must be greater " + f"than zero in lmn={lmn}." + ) out.append(lmn) return out + def parse_mode(lmn): - """Extracts overtones from an lmn. - """ + """Extracts overtones from an lmn.""" lm, nmodes = lmn[0:2], int(lmn[2]) overtones = [] for n in range(nmodes): - mode = lm + '{}'.format(n) + mode = lm + f"{n}" overtones.append(mode) return overtones def lm_amps_phases(**kwargs): - r"""Takes input_params and return dictionaries with amplitudes and phases + r""" + Takes input_params and return dictionaries with amplitudes and phases of each overtone of a specific lm mode, checking that all of them are given. Will also look for dbetas and dphis. If ``(dphi|dbeta)`` (i.e., without a mode suffix) are provided, they will be used for all modes that don't explicitly set a ``(dphi|dbeta){lmn}``. """ - lmns = format_lmns(kwargs['lmns']) + lmns = format_lmns(kwargs["lmns"]) amps = {} phis = {} dbetas = {} dphis = {} # reference mode - ref_amp = kwargs.pop('ref_amp', None) + ref_amp = kwargs.pop("ref_amp", None) if ref_amp is None: # default to the 220 mode - ref_amp = 'amp220' + ref_amp = "amp220" # check for reference dphi and dbeta - ref_dbeta = kwargs.pop('dbeta', 0.) - ref_dphi = kwargs.pop('dphi', 0.) - if isinstance(ref_amp, str) and ref_amp.startswith('amp'): + ref_dbeta = kwargs.pop("dbeta", 0.0) + ref_dphi = kwargs.pop("dphi", 0.0) + if isinstance(ref_amp, str) and ref_amp.startswith("amp"): # assume a mode was provided; check if the mode exists - ref_mode = ref_amp.replace('amp', '') + ref_mode = ref_amp.replace("amp", "") try: ref_amp = kwargs.pop(ref_amp) amps[ref_mode] = ref_amp except KeyError: - raise ValueError("Must provide an amplitude for the reference " - "mode {}".format(ref_amp)) + raise ValueError( + f"Must provide an amplitude for the reference mode {ref_amp}" + ) else: ref_mode = None # Get amplitudes and phases of the modes @@ -169,58 +179,62 @@ def lm_amps_phases(**kwargs): # skip the reference mode if mode != ref_mode: try: - amps[mode] = kwargs['amp' + mode] * ref_amp + amps[mode] = kwargs["amp" + mode] * ref_amp except KeyError: - raise ValueError('amp{} is required'.format(mode)) + raise ValueError(f"amp{mode} is required") try: - phis[mode] = kwargs['phi' + mode] + phis[mode] = kwargs["phi" + mode] except KeyError: - raise ValueError('phi{} is required'.format(mode)) - dphis[mode] = kwargs.pop('dphi'+mode, ref_dphi) - dbetas[mode] = kwargs.pop('dbeta'+mode, ref_dbeta) + raise ValueError(f"phi{mode} is required") + dphis[mode] = kwargs.pop("dphi" + mode, ref_dphi) + dbetas[mode] = kwargs.pop("dbeta" + mode, ref_dbeta) return amps, phis, dbetas, dphis def lm_freqs_taus(**kwargs): - """Take input_params and return dictionaries with frequencies and damping + """ + Take input_params and return dictionaries with frequencies and damping times of each overtone of a specific lm mode, checking that all of them are given. """ - lmns = format_lmns(kwargs['lmns']) + lmns = format_lmns(kwargs["lmns"]) freqs, taus = {}, {} for lmn in lmns: overtones = parse_mode(lmn) for mode in overtones: try: - freqs[mode] = kwargs['f_' + mode] + freqs[mode] = kwargs["f_" + mode] except KeyError: - raise ValueError('f_{} is required'.format(mode)) + raise ValueError(f"f_{mode} is required") try: - taus[mode] = kwargs['tau_' + mode] + taus[mode] = kwargs["tau_" + mode] except KeyError: - raise ValueError('tau_{} is required'.format(mode)) + raise ValueError(f"tau_{mode} is required") return freqs, taus def lm_arbitrary_harmonics(**kwargs): - """Take input_params and return dictionaries with arbitrary harmonics + """ + Take input_params and return dictionaries with arbitrary harmonics for each mode. """ - lmns = format_lmns(kwargs['lmns']) + lmns = format_lmns(kwargs["lmns"]) pols = {} polnms = {} for lmn in lmns: overtones = parse_mode(lmn) for mode in overtones: - pols[mode] = kwargs.pop('pol{}'.format(mode), None) - polnms[mode] = kwargs.pop('polnm{}'.format(mode), None) + pols[mode] = kwargs.pop(f"pol{mode}", None) + polnms[mode] = kwargs.pop(f"polnm{mode}", None) return pols, polnms # Functions to obtain t_final, f_final and output vector ###################### + def qnm_time_decay(tau, decay): - """Return the time at which the amplitude of the + """ + Return the time at which the amplitude of the ringdown falls to decay of the peak amplitude. Parameters @@ -235,12 +249,14 @@ def qnm_time_decay(tau, decay): t_decay : float The time at which the amplitude of the time-domain ringdown falls to decay of the peak amplitude. + """ return -tau * numpy.log(decay) def qnm_freq_decay(f_0, tau, decay): - """Return the frequency at which the amplitude of the + """ + Return the frequency at which the amplitude of the ringdown falls to decay of the peak amplitude. Parameters @@ -257,93 +273,96 @@ def qnm_freq_decay(f_0, tau, decay): f_decay : float The frequency at which the amplitude of the frequency-domain ringdown falls to decay of the peak amplitude. + """ q_0 = pi * f_0 * tau - alpha = 1. / decay - alpha_sq = 1. / decay / decay + alpha = 1.0 / decay + alpha_sq = 1.0 / decay / decay # Expression obtained analytically under the assumption # that 1./alpha_sq, q_0^2 >> 1 - q_sq = (alpha_sq + 4*q_0*q_0 + alpha*numpy.sqrt(alpha_sq + 16*q_0*q_0))/4. + q_sq = ( + alpha_sq + 4 * q_0 * q_0 + alpha * numpy.sqrt(alpha_sq + 16 * q_0 * q_0) + ) / 4.0 return numpy.sqrt(q_sq) / pi / tau def lm_tfinal(damping_times): - """Return the maximum t_final of the modes given, with t_final the time + """ + Return the maximum t_final of the modes given, with t_final the time at which the amplitude falls to 1/1000 of the peak amplitude """ if isinstance(damping_times, dict): t_max = {} for lmn in damping_times.keys(): - t_max[lmn] = qnm_time_decay(damping_times[lmn], 1./1000) + t_max[lmn] = qnm_time_decay(damping_times[lmn], 1.0 / 1000) t_final = max(t_max.values()) else: - t_final = qnm_time_decay(damping_times, 1./1000) + t_final = qnm_time_decay(damping_times, 1.0 / 1000) return t_final def lm_deltat(freqs, damping_times): - """Return the minimum delta_t of all the modes given, with delta_t given by + """ + Return the minimum delta_t of all the modes given, with delta_t given by the inverse of the frequency at which the amplitude of the ringdown falls to 1/1000 of the peak amplitude. """ if isinstance(freqs, dict) and isinstance(damping_times, dict): dt = {} for lmn in freqs.keys(): - dt[lmn] = 1. / qnm_freq_decay(freqs[lmn], - damping_times[lmn], 1./1000) + dt[lmn] = 1.0 / qnm_freq_decay(freqs[lmn], damping_times[lmn], 1.0 / 1000) delta_t = min(dt.values()) elif isinstance(freqs, dict) and not isinstance(damping_times, dict): - raise ValueError('Missing damping times.') + raise ValueError("Missing damping times.") elif isinstance(damping_times, dict) and not isinstance(freqs, dict): - raise ValueError('Missing frequencies.') + raise ValueError("Missing frequencies.") else: - delta_t = 1. / qnm_freq_decay(freqs, damping_times, 1./1000) + delta_t = 1.0 / qnm_freq_decay(freqs, damping_times, 1.0 / 1000) - if delta_t < min_dt: - delta_t = min_dt + delta_t = max(delta_t, min_dt) return delta_t def lm_ffinal(freqs, damping_times): - """Return the maximum f_final of the modes given, with f_final the + """ + Return the maximum f_final of the modes given, with f_final the frequency at which the amplitude falls to 1/1000 of the peak amplitude """ if isinstance(freqs, dict) and isinstance(damping_times, dict): f_max = {} for lmn in freqs.keys(): - f_max[lmn] = qnm_freq_decay(freqs[lmn], - damping_times[lmn], 1./1000) + f_max[lmn] = qnm_freq_decay(freqs[lmn], damping_times[lmn], 1.0 / 1000) f_final = max(f_max.values()) elif isinstance(freqs, dict) and not isinstance(damping_times, dict): - raise ValueError('Missing damping times.') + raise ValueError("Missing damping times.") elif isinstance(damping_times, dict) and not isinstance(freqs, dict): - raise ValueError('Missing frequencies.') + raise ValueError("Missing frequencies.") else: - f_final = qnm_freq_decay(freqs, damping_times, 1./1000) - if f_final > max_freq: - f_final = max_freq + f_final = qnm_freq_decay(freqs, damping_times, 1.0 / 1000) + f_final = min(f_final, max_freq) return f_final def lm_deltaf(damping_times): - """Return the minimum delta_f of all the modes given, with delta_f given by + """ + Return the minimum delta_f of all the modes given, with delta_f given by the inverse of the time at which the amplitude of the ringdown falls to 1/1000 of the peak amplitude. """ if isinstance(damping_times, dict): df = {} for lmn in damping_times.keys(): - df[lmn] = 1. / qnm_time_decay(damping_times[lmn], 1./1000) + df[lmn] = 1.0 / qnm_time_decay(damping_times[lmn], 1.0 / 1000) delta_f = min(df.values()) else: - delta_f = 1. / qnm_time_decay(damping_times, 1./1000) + delta_f = 1.0 / qnm_time_decay(damping_times, 1.0 / 1000) return delta_f -def td_output_vector(freqs, damping_times, taper=False, - delta_t=None, t_final=None): - """Return an empty TimeSeries with the appropriate size to fit all +def td_output_vector(freqs, damping_times, taper=False, delta_t=None, t_final=None): + """ + Return an empty TimeSeries with the appropriate size to fit all the quasi-normal modes present in freqs, damping_times """ if not delta_t: @@ -354,14 +373,17 @@ def td_output_vector(freqs, damping_times, taper=False, # Different modes will have different tapering window-size # Find maximum window size to create long enough output vector if taper: - max_tau = max(damping_times.values()) if \ - isinstance(damping_times, dict) else damping_times - kmax += int(max_tau/delta_t) + max_tau = ( + max(damping_times.values()) + if isinstance(damping_times, dict) + else damping_times + ) + kmax += int(max_tau / delta_t) outplus = TimeSeries(zeros(kmax, dtype=float64), delta_t=delta_t) outcross = TimeSeries(zeros(kmax, dtype=float64), delta_t=delta_t) if taper: # Change epoch of output vector if tapering will be applied - start = - max_tau + start = -max_tau # To ensure that t=0 is still in output vector start -= start % delta_t outplus._epoch, outcross._epoch = start, start @@ -369,7 +391,8 @@ def td_output_vector(freqs, damping_times, taper=False, def fd_output_vector(freqs, damping_times, delta_f=None, f_final=None): - """Return an empty FrequencySeries with the appropriate size to fit all + """ + Return an empty FrequencySeries with the appropriate size to fit all the quasi-normal modes present in freqs, damping_times """ if not delta_f: @@ -384,10 +407,20 @@ def fd_output_vector(freqs, damping_times, delta_f=None, f_final=None): # Spherical harmonics and Kerr factor ######################################### -def spher_harms(harmonics='spherical', l=None, m=None, n=0, - inclination=0., azimuthal=0., - spin=None, pol=None, polnm=None): - r"""Return the +/-m harmonic polarizations. + +def spher_harms( + harmonics="spherical", + l=None, + m=None, + n=0, + inclination=0.0, + azimuthal=0.0, + spin=None, + pol=None, + polnm=None, +): + r""" + Return the +/-m harmonic polarizations. This will return either spherical, spheroidal, or an arbitrary complex number depending on what ``harmonics`` is set to. If harmonics is set to @@ -430,39 +463,33 @@ def spher_harms(harmonics='spherical', l=None, m=None, n=0, The harmonic of the +m mode. xlnm : complex The harmonic of the -m mode. + """ - if harmonics == 'spherical': + if harmonics == "spherical": if lal is None: - raise ImportError( - "lal must be installed for spherical " - "harmonics" - ) - xlm = lal.SpinWeightedSphericalHarmonic(inclination, azimuthal, -2, - l, m) - xlnm = lal.SpinWeightedSphericalHarmonic(inclination, azimuthal, -2, - l, -m) - elif harmonics == 'spheroidal': + raise ImportError("lal must be installed for spherical harmonics") + xlm = lal.SpinWeightedSphericalHarmonic(inclination, azimuthal, -2, l, m) + xlnm = lal.SpinWeightedSphericalHarmonic(inclination, azimuthal, -2, l, -m) + elif harmonics == "spheroidal": if spin is None: raise ValueError("must provide a spin for spheroidal harmonics") if pykerr is None: - raise ImportError("pykerr must be installed for spheroidal " - "harmonics") + raise ImportError("pykerr must be installed for spheroidal harmonics") xlm = pykerr.spheroidal(inclination, spin, l, m, n, phi=azimuthal) xlnm = pykerr.spheroidal(inclination, spin, l, -m, n, phi=azimuthal) - elif harmonics == 'arbitrary': + elif harmonics == "arbitrary": if pol is None or polnm is None: - raise ValueError('must provide a pol and a polnm for arbitrary ' - 'harmonics') - xlm = numpy.exp(1j*pol) - xlnm = numpy.exp(1j*polnm) + raise ValueError("must provide a pol and a polnm for arbitrary harmonics") + xlm = numpy.exp(1j * pol) + xlnm = numpy.exp(1j * polnm) else: - raise ValueError("harmonics must be either spherical, spheroidal, " - "or arbitrary") + raise ValueError("harmonics must be either spherical, spheroidal, or arbitrary") return xlm, xlnm def Kerr_factor(final_mass, distance): - """Return the factor final_mass/distance (in dimensionless units) for Kerr + """ + Return the factor final_mass/distance (in dimensionless units) for Kerr ringdowns """ # Convert solar masses to meters @@ -476,12 +503,27 @@ def Kerr_factor(final_mass, distance): #### Basic functions to generate damped sinusoid ###################################################### -def td_damped_sinusoid(f_0, tau, amp, phi, times, - l=2, m=2, n=0, inclination=0., azimuthal=0., - dphi=0., dbeta=0., - harmonics='spherical', final_spin=None, - pol=None, polnm=None): - r"""Return a time domain damped sinusoid (plus and cross polarizations) + +def td_damped_sinusoid( + f_0, + tau, + amp, + phi, + times, + l=2, + m=2, + n=0, + inclination=0.0, + azimuthal=0.0, + dphi=0.0, + dbeta=0.0, + harmonics="spherical", + final_spin=None, + pol=None, + polnm=None, +): + r""" + Return a time domain damped sinusoid (plus and cross polarizations) with central frequency f_0, damping time tau, amplitude amp and phase phi. This returns the plus and cross polarization of the QNM, defined as @@ -578,44 +620,68 @@ def td_damped_sinusoid(f_0, tau, amp, phi, times, The plus polarization. hcross : numpy.ndarray The cross polarization. + """ # evaluate the harmonics - xlm, xlnm = spher_harms(harmonics=harmonics, l=l, m=m, n=n, - inclination=inclination, azimuthal=azimuthal, - spin=final_spin, pol=pol, polnm=polnm) + xlm, xlnm = spher_harms( + harmonics=harmonics, + l=l, + m=m, + n=n, + inclination=inclination, + azimuthal=azimuthal, + spin=final_spin, + pol=pol, + polnm=polnm, + ) # generate the +/-m modes # we measure things as deviations from circular polarization, which occurs # when h_{l-m} = (-1)^l h_{lm}^*; that implies that # phi_{l-m} = - phi_{lm} and A_{l-m} = (-1)^l A_{lm} omegalm = two_pi * f_0 * times - damping = -times/tau + damping = -times / tau # check for negative times mask = times < 0 if mask.any(): - damping[mask] = 10*times[mask]/tau + damping[mask] = 10 * times[mask] / tau if m == 0: # no -m, just calculate - hlm = xlm * amp * numpy.exp(damping + 1j*(omegalm + phi)) + hlm = xlm * amp * numpy.exp(damping + 1j * (omegalm + phi)) else: # amplitude if dbeta == 0: alm = alnm = amp else: - beta = pi/4 + dbeta + beta = pi / 4 + dbeta alm = 2**0.5 * amp * numpy.cos(beta) alnm = 2**0.5 * amp * numpy.sin(beta) # phase - phinm = l*pi + dphi - phi - hlm = xlm * alm * numpy.exp(damping + 1j*(omegalm + phi)) \ - + xlnm * alnm * numpy.exp(damping - 1j*(omegalm - phinm)) + phinm = l * pi + dphi - phi + hlm = xlm * alm * numpy.exp( + damping + 1j * (omegalm + phi) + ) + xlnm * alnm * numpy.exp(damping - 1j * (omegalm - phinm)) return hlm.real, hlm.imag -def fd_damped_sinusoid(f_0, tau, amp, phi, freqs, t_0=0., - l=2, m=2, n=0, inclination=0., azimuthal=0., - harmonics='spherical', final_spin=None, - pol=None, polnm=None): - r"""Return the frequency domain version of a damped sinusoid. +def fd_damped_sinusoid( + f_0, + tau, + amp, + phi, + freqs, + t_0=0.0, + l=2, + m=2, + n=0, + inclination=0.0, + azimuthal=0.0, + harmonics="spherical", + final_spin=None, + pol=None, + polnm=None, +): + r""" + Return the frequency domain version of a damped sinusoid. This is the frequency domain version :py:func:`td_damped_sinusoid` without a taper if an infinite sample rate were used to resolve the step function @@ -676,25 +742,37 @@ def fd_damped_sinusoid(f_0, tau, amp, phi, freqs, t_0=0., The plus polarization. hctilde : numpy.ndarray The cross polarization. + """ # evaluate the harmonics if inclination is None: - inclination = 0. + inclination = 0.0 if azimuthal is None: - azimuthal = 0. - xlm, xlnm = spher_harms(harmonics=harmonics, l=l, m=m, n=n, - inclination=inclination, azimuthal=azimuthal, - spin=final_spin, pol=pol, polnm=polnm) + azimuthal = 0.0 + xlm, xlnm = spher_harms( + harmonics=harmonics, + l=l, + m=m, + n=n, + inclination=inclination, + azimuthal=azimuthal, + spin=final_spin, + pol=pol, + polnm=polnm, + ) # we'll assume circular polarization - xp = xlm + (-1)**l * xlnm - xc = xlm - (-1)**l * xlnm - denominator = 1 + (4j * pi * freqs * tau) - \ - (4 * pi_sq * (freqs*freqs - f_0*f_0) * tau*tau) + xp = xlm + (-1) ** l * xlnm + xc = xlm - (-1) ** l * xlnm + denominator = ( + 1 + + (4j * pi * freqs * tau) + - (4 * pi_sq * (freqs * freqs - f_0 * f_0) * tau * tau) + ) norm = amp * tau / denominator if t_0 != 0: time_shift = numpy.exp(-1j * two_pi * freqs * t_0) norm *= time_shift - A1 = (1 + 2j * pi * freqs * tau) + A1 = 1 + 2j * pi * freqs * tau A2 = two_pi * f_0 * tau # Analytical expression for the Fourier transform of the ringdown hptilde = norm * xp * (A1 * numpy.cos(phi) - A2 * numpy.sin(phi)) @@ -706,8 +784,10 @@ def fd_damped_sinusoid(f_0, tau, amp, phi, freqs, t_0=0., #### Base multi-mode for all approximants ###################################################### + def multimode_base(input_params, domain, freq_tau_approximant=False): - """Return a superposition of damped sinusoids in either time or frequency + """ + Return a superposition of damped sinusoids in either time or frequency domains with parameters set by input_params. Parameters @@ -734,80 +814,108 @@ def multimode_base(input_params, domain, freq_tau_approximant=False): hcross : TimeSeries The cross phase of a ringdown with the lm modes specified and n overtones in the chosen domain (time or frequency). + """ - input_params['lmns'] = format_lmns(input_params['lmns']) + input_params["lmns"] = format_lmns(input_params["lmns"]) amps, phis, dbetas, dphis = lm_amps_phases(**input_params) pols, polnms = lm_arbitrary_harmonics(**input_params) # get harmonics argument try: - harmonics = input_params['harmonics'] + harmonics = input_params["harmonics"] except KeyError: - harmonics = 'spherical' + harmonics = "spherical" # we'll need the final spin for spheroidal harmonics - if harmonics == 'spheroidal': - final_spin = input_params['final_spin'] + if harmonics == "spheroidal": + final_spin = input_params["final_spin"] else: final_spin = None # add inclination and azimuthal if they aren't provided - if 'inclination' not in input_params: - input_params['inclination'] = 0. - if 'azimuthal' not in input_params: - input_params['azimuthal'] = 0. + if "inclination" not in input_params: + input_params["inclination"] = 0.0 + if "azimuthal" not in input_params: + input_params["azimuthal"] = 0.0 # figure out the frequencies and damping times if freq_tau_approximant: freqs, taus = lm_freqs_taus(**input_params) - norm = 1. + norm = 1.0 else: - freqs, taus = get_lm_f0tau_allmodes(input_params['final_mass'], - input_params['final_spin'], input_params['lmns']) - norm = Kerr_factor(input_params['final_mass'], - input_params['distance']) if 'distance' in input_params.keys() \ - else 1. + freqs, taus = get_lm_f0tau_allmodes( + input_params["final_mass"], input_params["final_spin"], input_params["lmns"] + ) + norm = ( + Kerr_factor(input_params["final_mass"], input_params["distance"]) + if "distance" in input_params.keys() + else 1.0 + ) for mode, freq in freqs.items(): - if 'delta_f{}'.format(mode) in input_params: - freqs[mode] += input_params['delta_f{}'.format(mode)]*freq + if f"delta_f{mode}" in input_params: + freqs[mode] += input_params[f"delta_f{mode}"] * freq for mode, tau in taus.items(): - if 'delta_tau{}'.format(mode) in input_params: - taus[mode] += input_params['delta_tau{}'.format(mode)]*tau + if f"delta_tau{mode}" in input_params: + taus[mode] += input_params[f"delta_tau{mode}"] * tau # setup the output - if domain == 'td': - outplus, outcross = td_output_vector(freqs, taus, - input_params['taper'], input_params['delta_t'], - input_params['t_final']) + if domain == "td": + outplus, outcross = td_output_vector( + freqs, + taus, + input_params["taper"], + input_params["delta_t"], + input_params["t_final"], + ) sample_times = outplus.sample_times.numpy() - elif domain == 'fd': - kmin = int(input_params['f_lower'] / input_params['delta_f']) - outplus, outcross = fd_output_vector(freqs, taus, - input_params['delta_f'], - input_params['f_final']) + elif domain == "fd": + kmin = int(input_params["f_lower"] / input_params["delta_f"]) + outplus, outcross = fd_output_vector( + freqs, taus, input_params["delta_f"], input_params["f_final"] + ) sample_freqs = outplus.sample_frequencies.numpy()[kmin:] else: - raise ValueError('unrecognised domain argument {}; ' - 'must be either fd or td'.format(domain)) + raise ValueError( + f"unrecognised domain argument {domain}; must be either fd or td" + ) # cyclce over the modes, generating the waveforms for lmn in freqs: - if amps[lmn] == 0.: + if amps[lmn] == 0.0: # skip continue - if domain == 'td': + if domain == "td": hplus, hcross = td_damped_sinusoid( - freqs[lmn], taus[lmn], amps[lmn], phis[lmn], sample_times, - l=int(lmn[0]), m=int(lmn[1]), n=int(lmn[2]), - inclination=input_params['inclination'], - azimuthal=input_params['azimuthal'], - dphi=dphis[lmn], dbeta=dbetas[lmn], - harmonics=harmonics, final_spin=final_spin, - pol=pols[lmn], polnm=polnms[lmn]) + freqs[lmn], + taus[lmn], + amps[lmn], + phis[lmn], + sample_times, + l=int(lmn[0]), + m=int(lmn[1]), + n=int(lmn[2]), + inclination=input_params["inclination"], + azimuthal=input_params["azimuthal"], + dphi=dphis[lmn], + dbeta=dbetas[lmn], + harmonics=harmonics, + final_spin=final_spin, + pol=pols[lmn], + polnm=polnms[lmn], + ) outplus += hplus outcross += hcross - elif domain == 'fd': + elif domain == "fd": hplus, hcross = fd_damped_sinusoid( - freqs[lmn], taus[lmn], amps[lmn], phis[lmn], sample_freqs, - l=int(lmn[0]), m=int(lmn[1]), n=int(lmn[2]), - inclination=input_params['inclination'], - azimuthal=input_params['azimuthal'], - harmonics=harmonics, final_spin=final_spin, - pol=pols[lmn], polnm=polnms[lmn]) + freqs[lmn], + taus[lmn], + amps[lmn], + phis[lmn], + sample_freqs, + l=int(lmn[0]), + m=int(lmn[1]), + n=int(lmn[2]), + inclination=input_params["inclination"], + azimuthal=input_params["azimuthal"], + harmonics=harmonics, + final_spin=final_spin, + pol=pols[lmn], + polnm=polnms[lmn], + ) outplus[kmin:] += hplus outcross[kmin:] += hcross return norm * outplus, norm * outcross @@ -817,8 +925,10 @@ def multimode_base(input_params, domain, freq_tau_approximant=False): #### Approximants ###################################################### + def get_td_from_final_mass_spin(template=None, **kwargs): - """Return time domain ringdown with all the modes specified. + """ + Return time domain ringdown with all the modes specified. Parameters ---------- @@ -924,12 +1034,15 @@ def get_td_from_final_mass_spin(template=None, **kwargs): hcross : TimeSeries The cross phase of a ringdown with the lm modes specified and n overtones in time domain. + """ input_params = props(template, mass_spin_required_args, td_args, **kwargs) - return multimode_base(input_params, domain='td') + return multimode_base(input_params, domain="td") + def get_fd_from_final_mass_spin(template=None, **kwargs): - """Return frequency domain ringdown with all the modes specified. + """ + Return frequency domain ringdown with all the modes specified. Parameters ---------- @@ -1019,12 +1132,15 @@ def get_fd_from_final_mass_spin(template=None, **kwargs): hcrosstilde : FrequencySeries The cross phase of a ringdown with the lm modes specified and n overtones in frequency domain. + """ input_params = props(template, mass_spin_required_args, fd_args, **kwargs) - return multimode_base(input_params, domain='fd') + return multimode_base(input_params, domain="fd") + def get_td_from_freqtau(template=None, **kwargs): - """Return time domain ringdown with all the modes specified. + """ + Return time domain ringdown with all the modes specified. Parameters ---------- @@ -1122,12 +1238,15 @@ def get_td_from_freqtau(template=None, **kwargs): hcross : TimeSeries The cross phase of a ringdown with the lm modes specified and n overtones in time domain. + """ input_params = props(template, freqtau_required_args, td_args, **kwargs) - return multimode_base(input_params, domain='td', freq_tau_approximant=True) + return multimode_base(input_params, domain="td", freq_tau_approximant=True) + def get_fd_from_freqtau(template=None, **kwargs): - """Return frequency domain ringdown with all the modes specified. + """ + Return frequency domain ringdown with all the modes specified. Parameters ---------- @@ -1222,16 +1341,20 @@ def get_fd_from_freqtau(template=None, **kwargs): hcrosstilde : FrequencySeries The cross phase of a ringdown with the lm modes specified and n overtones in frequency domain. + """ input_params = props(template, freqtau_required_args, fd_args, **kwargs) - return multimode_base(input_params, domain='fd', freq_tau_approximant=True) + return multimode_base(input_params, domain="fd", freq_tau_approximant=True) + # Approximant names ########################################################### ringdown_fd_approximants = { - 'FdQNMfromFinalMassSpin': get_fd_from_final_mass_spin, - 'FdQNMfromFreqTau': get_fd_from_freqtau} + "FdQNMfromFinalMassSpin": get_fd_from_final_mass_spin, + "FdQNMfromFreqTau": get_fd_from_freqtau, +} ringdown_td_approximants = { - 'TdQNMfromFinalMassSpin': get_td_from_final_mass_spin, - 'TdQNMfromFreqTau': get_td_from_freqtau} + "TdQNMfromFinalMassSpin": get_td_from_final_mass_spin, + "TdQNMfromFreqTau": get_td_from_freqtau, +} diff --git a/pycbc/waveform/sinegauss.py b/pycbc/waveform/sinegauss.py index bac4ad967e2..7836246b809 100644 --- a/pycbc/waveform/sinegauss.py +++ b/pycbc/waveform/sinegauss.py @@ -1,17 +1,20 @@ -""" Generation of sine-Gaussian bursty type things -""" +"""Generation of sine-Gaussian bursty type things""" -import pycbc.types -import numpy import functools +import numpy + +import pycbc.types + + @functools.lru_cache(maxsize=128) def cached_arange(kmax, delta_f): return numpy.arange(0, kmax) * delta_f def fd_sine_gaussian(amp, quality, central_frequency, fmin, fmax, delta_f): - """ Generate a Fourier domain sine-Gaussian + """ + Generate a Fourier domain sine-Gaussian Parameters ---------- @@ -33,6 +36,7 @@ def fd_sine_gaussian(amp, quality, central_frequency, fmin, fmax, delta_f): ------- sg: pycbc.types.Frequencyseries A Fourier domain sine-Gaussian + """ # Optimization note: Ian has profiled and done optimization on this # function. If further speed up is needed caching the v vector and @@ -42,7 +46,7 @@ def fd_sine_gaussian(amp, quality, central_frequency, fmin, fmax, delta_f): kmax = int(round(fmax / delta_f)) pi = numpy.pi - tau = (quality / (2 * pi * central_frequency)) + tau = quality / (2 * pi * central_frequency) quality_sq = quality**2 f = cached_arange(kmax, delta_f) @@ -62,36 +66,29 @@ def fd_sine_gaussian(amp, quality, central_frequency, fmin, fmax, delta_f): # Only consider points larger than kmin indices[kmin:] = 1 # Find frequencies at which first term is equal to exp_term_cutoff - low_freq_first_term = ( - central_frequency - (-exp_term_cutoff)**0.5 / (tau * pi) - ) - high_freq_first_term = ( - central_frequency + (-exp_term_cutoff)**0.5 / (tau * pi) - ) - low_freq_first_idx = max(kmin, int(low_freq_first_term//delta_f)) - high_freq_first_idx = min(kmax, int(high_freq_first_term//delta_f)) + low_freq_first_term = central_frequency - (-exp_term_cutoff) ** 0.5 / (tau * pi) + high_freq_first_term = central_frequency + (-exp_term_cutoff) ** 0.5 / (tau * pi) + low_freq_first_idx = max(kmin, int(low_freq_first_term // delta_f)) + high_freq_first_idx = min(kmax, int(high_freq_first_term // delta_f)) # Find frequency at which second term drops to exp_term_cutoff - high_freq_second_idx = ( - int(-exp_term_cutoff / quality_sq * central_frequency // delta_f) + high_freq_second_idx = int( + -exp_term_cutoff / quality_sq * central_frequency // delta_f ) - exp_term_1 = -( - tau * pi * - (f[low_freq_first_idx:high_freq_first_idx] - central_frequency) - )**2.0 + exp_term_1 = ( + -( + (tau * pi * (f[low_freq_first_idx:high_freq_first_idx] - central_frequency)) + ** 2.0 + ) + ) A_term = amp * (pi**0.5) / 2 * tau - v[low_freq_first_idx:high_freq_first_idx] = ( - A_term * numpy.exp(exp_term_1) - ) + v[low_freq_first_idx:high_freq_first_idx] = A_term * numpy.exp(exp_term_1) # If the first term is already less than e**50 don't need the second # term at all ... It's often the case that the second term is not needed. if high_freq_second_idx > kmin: - exp_term_2 = ( - -quality_sq * f[kmin:high_freq_second_idx] / central_frequency - ) - v[kmin:high_freq_second_idx] *= (1 + numpy.exp(exp_term_2)) + exp_term_2 = -quality_sq * f[kmin:high_freq_second_idx] / central_frequency + v[kmin:high_freq_second_idx] *= 1 + numpy.exp(exp_term_2) return pycbc.types.FrequencySeries(v, delta_f=delta_f, copy=False) - diff --git a/pycbc/waveform/spa_tmplt.py b/pycbc/waveform/spa_tmplt.py index d56479239d0..2562cebb367 100644 --- a/pycbc/waveform/spa_tmplt.py +++ b/pycbc/waveform/spa_tmplt.py @@ -18,22 +18,26 @@ # Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, # MA 02111-1307 USA -"""This module contains functions for generating common SPA template precalculated - vectors. """ -from math import sqrt, log +This module contains functions for generating common SPA template precalculated +vectors. +""" + import warnings +from math import log, sqrt + import numpy import pycbc.pnutils +from pycbc.constants import GAMMA, MRSUN_SI, MTSUN_SI, PC_SI, PI +from pycbc.libutils import import_optional from pycbc.scheme import schemed -from pycbc.types import FrequencySeries, Array, complex64, float32, zeros +from pycbc.types import Array, FrequencySeries, complex64, float32, zeros from pycbc.waveform.utils import ceilpow2 -from pycbc.constants import PI, GAMMA, MTSUN_SI, PC_SI, MRSUN_SI -from pycbc.libutils import import_optional -lal = import_optional('lal') -lalsimulation = import_optional('lalsimulation') +lal = import_optional("lal") +lalsimulation = import_optional("lalsimulation") + def findchirp_chirptime(m1, m2, fLower, porder): # variables used to compute chirp time @@ -41,21 +45,28 @@ def findchirp_chirptime(m1, m2, fLower, porder): m2 = float(m2) m = m1 + m2 eta = m1 * m2 / m / m - c0T = c2T = c3T = c4T = c5T = c6T = c6LogT = c7T = 0. + c0T = c2T = c3T = c4T = c5T = c6T = c6LogT = c7T = 0.0 # All implemented option if porder == -1: porder = 7 if porder >= 7: - c7T = PI * (14809.0 * eta * eta / 378.0 - 75703.0 * eta / 756.0 - 15419335.0 / 127008.0) + c7T = PI * ( + 14809.0 * eta * eta / 378.0 - 75703.0 * eta / 756.0 - 15419335.0 / 127008.0 + ) if porder >= 6: - c6T = GAMMA * 6848.0 / 105.0 - 10052469856691.0 / 23471078400.0 +\ - PI * PI * 128.0 / 3.0 + \ - eta * (3147553127.0 / 3048192.0 - PI * PI * 451.0 / 12.0) -\ - eta * eta * 15211.0 / 1728.0 + eta * eta * eta * 25565.0 / 1296.0 +\ - eta * eta * eta * 25565.0 / 1296.0 + numpy.log(4.0) * 6848.0 / 105.0 + c6T = ( + GAMMA * 6848.0 / 105.0 + - 10052469856691.0 / 23471078400.0 + + PI * PI * 128.0 / 3.0 + + eta * (3147553127.0 / 3048192.0 - PI * PI * 451.0 / 12.0) + - eta * eta * 15211.0 / 1728.0 + + eta * eta * eta * 25565.0 / 1296.0 + + eta * eta * eta * 25565.0 / 1296.0 + + numpy.log(4.0) * 6848.0 / 105.0 + ) c6LogT = 6848.0 / 105.0 if porder >= 5: @@ -68,7 +79,7 @@ def findchirp_chirptime(m1, m2, fLower, porder): c0T = 5.0 * m * MTSUN_SI / (256.0 * eta) # This is the PN parameter v evaluated at the lower freq. cutoff - xT = pow (PI * m * MTSUN_SI * fLower, 1.0 / 3.0) + xT = pow(PI * m * MTSUN_SI * fLower, 1.0 / 3.0) x2T = xT * xT x3T = xT * x2T x4T = x2T * x2T @@ -83,8 +94,19 @@ def findchirp_chirptime(m1, m2, fLower, porder): # This formula works for any PN order, because # higher order coeffs will be set to zero. - return c0T * (1 + c2T * x2T + c3T * x3T + c4T * x4T + c5T * x5T + - (c6T + c6LogT * numpy.log(xT)) * x6T + c7T * x7T) / x8T + return ( + c0T + * ( + 1 + + c2T * x2T + + c3T * x3T + + c4T * x4T + + c5T * x5T + + (c6T + c6LogT * numpy.log(xT)) * x6T + + c7T * x7T + ) + / x8T + ) def spa_length_in_time(**kwds): @@ -93,10 +115,10 @@ def spa_length_in_time(**kwds): based on the masses, PN order, and low-frequency cut-off. """ - m1 = kwds['mass1'] - m2 = kwds['mass2'] - flow = kwds['f_lower'] - porder = int(kwds['phase_order']) + m1 = kwds["mass1"] + m2 = kwds["mass2"] + flow = kwds["f_lower"] + porder = int(kwds["phase_order"]) # For now, we call the swig-wrapped function below in # lalinspiral. Eventually would be nice to replace this @@ -105,53 +127,57 @@ def spa_length_in_time(**kwds): def spa_amplitude_factor(**kwds): - m1 = kwds['mass1'] - m2 = kwds['mass2'] + m1 = kwds["mass1"] + m2 = kwds["mass2"] _, eta = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) - FTaN = 32. * eta * eta / 5. - dETaN = 2. * -eta / 2. + FTaN = 32.0 * eta * eta / 5.0 + dETaN = 2.0 * -eta / 2.0 M = m1 + m2 m_sec = M * MTSUN_SI piM = PI * m_sec - amp0 = 4. * m1 * m2 / (1e6 * PC_SI) * MRSUN_SI * MTSUN_SI * sqrt(PI / 12.) + amp0 = 4.0 * m1 * m2 / (1e6 * PC_SI) * MRSUN_SI * MTSUN_SI * sqrt(PI / 12.0) - fac = numpy.sqrt(-dETaN / FTaN) * amp0 * (piM ** (-7./6.)) + fac = numpy.sqrt(-dETaN / FTaN) * amp0 * (piM ** (-7.0 / 6.0)) return -fac _prec = None + + def spa_tmplt_precondition(length, delta_f, kmin=0): - """Return the amplitude portion of the TaylorF2 approximant, used to precondition + """ + Return the amplitude portion of the TaylorF2 approximant, used to precondition the strain data. The result is cached, and so should not be modified, only read. """ global _prec if _prec is None or _prec.delta_f != delta_f or len(_prec) < length: - v = numpy.arange(0, (kmin + length*2), 1.) * delta_f - v = numpy.power(v[1:len(v)], -7./6.) + v = numpy.arange(0, (kmin + length * 2), 1.0) * delta_f + v = numpy.power(v[1 : len(v)], -7.0 / 6.0) _prec = FrequencySeries(v, delta_f=delta_f, dtype=float32) - return _prec[kmin:kmin + length] + return _prec[kmin : kmin + length] def spa_tmplt_norm(psd, length, delta_f, f_lower): amp = spa_tmplt_precondition(length, delta_f) k_min = int(f_lower / delta_f) - sigma = (amp[k_min:length].numpy() ** 2. / psd[k_min:length].numpy()) + sigma = amp[k_min:length].numpy() ** 2.0 / psd[k_min:length].numpy() norm_vec = numpy.zeros(length) - norm_vec[k_min:length] = sigma.cumsum() * 4. * delta_f + norm_vec[k_min:length] = sigma.cumsum() * 4.0 * delta_f return norm_vec def spa_tmplt_end(**kwds): - return pycbc.pnutils.f_SchwarzISCO(kwds['mass1'] + kwds['mass2']) + return pycbc.pnutils.f_SchwarzISCO(kwds["mass1"] + kwds["mass2"]) def spa_distance(psd, mass1, mass2, lower_frequency_cutoff, snr=8): - """ Return the distance at a given snr (default=8) of the SPA TaylorF2 + """ + Return the distance at a given snr (default=8) of the SPA TaylorF2 template. """ kend = int(spa_tmplt_end(mass1=mass1, mass2=mass2) / psd.delta_f) @@ -164,30 +190,42 @@ def spa_distance(psd, mass1, mass2, lower_frequency_cutoff, snr=8): @schemed("pycbc.waveform.spa_tmplt_") -def spa_tmplt_engine(htilde, kmin, phase_order, delta_f, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, amp_factor): - """ Calculate the spa tmplt phase - """ +def spa_tmplt_engine( + htilde, + kmin, + phase_order, + delta_f, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, +): + """Calculate the spa tmplt phase""" err_msg = "This function is a stub that should be overridden using the " err_msg += "scheme. You shouldn't be seeing this error!" raise ValueError(err_msg) def spa_tmplt(**kwds): - """ Generate a minimal TaylorF2 approximant with optimizations for the sin/cos - """ - distance = kwds['distance'] - mass1 = kwds['mass1'] - mass2 = kwds['mass2'] - s1z = kwds['spin1z'] - s2z = kwds['spin2z'] - phase_order = int(kwds['phase_order']) - #amplitude_order = int(kwds['amplitude_order']) - spin_order = int(kwds['spin_order']) - - if 'out' in kwds: - out = kwds['out'] + """Generate a minimal TaylorF2 approximant with optimizations for the sin/cos""" + distance = kwds["distance"] + mass1 = kwds["mass1"] + mass2 = kwds["mass2"] + s1z = kwds["spin1z"] + s2z = kwds["spin2z"] + phase_order = int(kwds["phase_order"]) + # amplitude_order = int(kwds['amplitude_order']) + spin_order = int(kwds["spin_order"]) + + if "out" in kwds: + out = kwds["out"] else: out = None @@ -195,18 +233,15 @@ def spa_tmplt(**kwds): lal_pars = lal.CreateDict() if phase_order != -1: - lalsimulation.SimInspiralWaveformParamsInsertPNPhaseOrder( - lal_pars, phase_order) + lalsimulation.SimInspiralWaveformParamsInsertPNPhaseOrder(lal_pars, phase_order) if spin_order != -1: - lalsimulation.SimInspiralWaveformParamsInsertPNSpinOrder( - lal_pars, spin_order) + lalsimulation.SimInspiralWaveformParamsInsertPNSpinOrder(lal_pars, spin_order) # Calculate the PN terms phasing = lalsimulation.SimInspiralTaylorF2AlignedPhasing( - float(mass1), float(mass2), - float(s1z), float(s2z), - lal_pars) + float(mass1), float(mass2), float(s1z), float(s2z), lal_pars + ) pfaN = phasing.v[0] pfa2 = phasing.v[2] / pfaN @@ -221,51 +256,80 @@ def spa_tmplt(**kwds): piM = PI * (mass1 + mass2) * MTSUN_SI - if 'sample_points' not in kwds: - f_lower = kwds['f_lower'] - delta_f = kwds['delta_f'] + if "sample_points" not in kwds: + f_lower = kwds["f_lower"] + delta_f = kwds["delta_f"] kmin = int(f_lower / float(delta_f)) # Get max frequency one way or another # f_final is assigned default value 0 in parameters.py - if 'f_final' in kwds and kwds['f_final'] > 0.: - fstop = kwds['f_final'] - elif 'f_upper' in kwds: - fstop = kwds['f_upper'] - warnings.warn('f_upper is deprecated in favour of f_final!', - DeprecationWarning) + if "f_final" in kwds and kwds["f_final"] > 0.0: + fstop = kwds["f_final"] + elif "f_upper" in kwds: + fstop = kwds["f_upper"] + warnings.warn( + "f_upper is deprecated in favour of f_final!", DeprecationWarning + ) else: # Schwarzschild ISCO frequency - vISCO = 1. / sqrt(6.) + vISCO = 1.0 / sqrt(6.0) fstop = vISCO * vISCO * vISCO / piM if fstop <= f_lower: - raise ValueError("cannot generate waveform! f_lower >= f_final" - f" ({f_lower}, {fstop})") + raise ValueError( + f"cannot generate waveform! f_lower >= f_final ({f_lower}, {fstop})" + ) kmax = int(fstop / delta_f) f_max = ceilpow2(fstop) n = int(f_max / delta_f) + 1 if not out: - htilde = FrequencySeries(zeros(n, dtype=numpy.complex64), delta_f=delta_f, copy=False) + htilde = FrequencySeries( + zeros(n, dtype=numpy.complex64), delta_f=delta_f, copy=False + ) else: if type(out) is not Array: raise TypeError("Output must be an instance of Array") - if len(out) < kmax: - kmax = len(out) + kmax = min(kmax, len(out)) if out.dtype != complex64: raise TypeError("Output array is the wrong dtype") htilde = FrequencySeries(out, delta_f=delta_f, copy=False) - spa_tmplt_engine(htilde[kmin:kmax], kmin, phase_order, - delta_f, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, amp_factor) + spa_tmplt_engine( + htilde[kmin:kmax], + kmin, + phase_order, + delta_f, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, + ) else: from .spa_tmplt_cpu import spa_tmplt_inline_sequence - htilde = numpy.empty(len(kwds['sample_points']), dtype=numpy.complex64) + + htilde = numpy.empty(len(kwds["sample_points"]), dtype=numpy.complex64) spa_tmplt_inline_sequence( - piM, pfaN, pfa2, pfa3, pfa4, pfa5, pfl5, pfa6, pfl6, pfa7, - amp_factor, kwds['sample_points'], htilde) + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, + kwds["sample_points"], + htilde, + ) return htilde diff --git a/pycbc/waveform/spa_tmplt_cuda.py b/pycbc/waveform/spa_tmplt_cuda.py index d67ea978e11..2c912be0bbd 100644 --- a/pycbc/waveform/spa_tmplt_cuda.py +++ b/pycbc/waveform/spa_tmplt_cuda.py @@ -23,7 +23,7 @@ from pycuda.elementwise import ElementwiseKernel -from pycbc.constants import TWOPI, PI_4 +from pycbc.constants import PI_4, TWOPI preamble = "" @@ -75,20 +75,51 @@ htilde[i]._M_im = - psin * amp2; """ -taylorf2_kernel = ElementwiseKernel("""pycuda::complex *htilde, int kmin, int phase_order, +taylorf2_kernel = ElementwiseKernel( + """pycuda::complex *htilde, int kmin, int phase_order, float delta_f, float TWOPI, float PI_4, float piM, float pfaN, float pfa2, float pfa3, float pfa4, float pfa5, float pfl5, float pfa6, float pfl6, float pfa7, float amp""", - taylorf2_text, "SPAtmplt", - preamble=preamble) + taylorf2_text, + "SPAtmplt", + preamble=preamble, +) -def spa_tmplt_engine(htilde, kmin, phase_order, - delta_f, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, amp_factor): - """ Calculate the spa tmplt phase - """ - taylorf2_kernel(htilde.data, kmin, phase_order, - delta_f, TWOPI, PI_4, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, amp_factor) + +def spa_tmplt_engine( + htilde, + kmin, + phase_order, + delta_f, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, +): + """Calculate the spa tmplt phase""" + taylorf2_kernel( + htilde.data, + kmin, + phase_order, + delta_f, + TWOPI, + PI_4, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, + ) diff --git a/pycbc/waveform/spa_tmplt_cupy.py b/pycbc/waveform/spa_tmplt_cupy.py index fb98e17f094..e9299792cf7 100644 --- a/pycbc/waveform/spa_tmplt_cupy.py +++ b/pycbc/waveform/spa_tmplt_cupy.py @@ -24,7 +24,7 @@ import cupy as cp import mako.template -from pycbc.constants import PI_4, TWOPI, LN2 +from pycbc.constants import LN2, PI_4, TWOPI taylorf2_text = mako.template.Template(""" const float f = (i + kmin ) * delta_f; @@ -74,7 +74,7 @@ htilde.real(pcos * amp2); htilde.imag(-psin * amp2); -""").render(TWOPI=TWOPI, PI_4=PI_4, LN4=2*LN2) +""").render(TWOPI=TWOPI, PI_4=PI_4, LN4=2 * LN2) taylorf2_kernel = cp.ElementwiseKernel( @@ -89,13 +89,38 @@ ) -def spa_tmplt_engine(htilde, kmin, phase_order, - delta_f, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, amp_factor): - """ Calculate the spa tmplt phase - """ - taylorf2_kernel(kmin, phase_order, - delta_f, piM, pfaN, - pfa2, pfa3, pfa4, pfa5, pfl5, - pfa6, pfl6, pfa7, amp_factor, htilde.data) +def spa_tmplt_engine( + htilde, + kmin, + phase_order, + delta_f, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, +): + """Calculate the spa tmplt phase""" + taylorf2_kernel( + kmin, + phase_order, + delta_f, + piM, + pfaN, + pfa2, + pfa3, + pfa4, + pfa5, + pfl5, + pfa6, + pfl6, + pfa7, + amp_factor, + htilde.data, + ) diff --git a/pycbc/waveform/supernovae.py b/pycbc/waveform/supernovae.py index c36ec1b4afd..50795ef1fd2 100644 --- a/pycbc/waveform/supernovae.py +++ b/pycbc/waveform/supernovae.py @@ -1,52 +1,54 @@ -"""Generate core-collapse supernovae waveform for core bounce and +""" +Generate core-collapse supernovae waveform for core bounce and subsequent postbounce oscillations. """ import numpy -from pycbc.types import TimeSeries + from pycbc.io.hdf import HFile +from pycbc.types import TimeSeries _pc_dict = {} def get_corecollapse_bounce(**kwargs): - """ Generates core bounce and postbounce waveform by using principal + """ + Generates core bounce and postbounce waveform by using principal component basis vectors from a .hdf file. The waveform parameters are the coefficients of the principal components and the distance. The number of principal components used can also be varied. """ - try: - principal_components = _pc_dict['principal_components'] + principal_components = _pc_dict["principal_components"] except KeyError: - with HFile(kwargs['principal_components_file'], 'r') as pc_file: - principal_components = numpy.array(pc_file['principal_components']) - _pc_dict['principal_components'] = principal_components + with HFile(kwargs["principal_components_file"], "r") as pc_file: + principal_components = numpy.array(pc_file["principal_components"]) + _pc_dict["principal_components"] = principal_components - if 'coefficients_array' in kwargs: - coefficients_array = kwargs['coefficients_array'] + if "coefficients_array" in kwargs: + coefficients_array = kwargs["coefficients_array"] else: - coeffs_keys = [x for x in kwargs if x.startswith('coeff_')] + coeffs_keys = [x for x in kwargs if x.startswith("coeff_")] coeffs_keys = numpy.sort(numpy.array(coeffs_keys)) coefficients_array = numpy.array([kwargs[x] for x in coeffs_keys]) - no_of_pcs = int(kwargs['no_of_pcs']) + no_of_pcs = int(kwargs["no_of_pcs"]) coefficients_array = coefficients_array[:no_of_pcs] principal_components = principal_components[:no_of_pcs] pc_len = len(principal_components) assert len(coefficients_array) == pc_len - distance = kwargs['distance'] - mpc_conversion = 3.08567758128e+22 + distance = kwargs["distance"] + mpc_conversion = 3.08567758128e22 distance *= mpc_conversion strain = numpy.dot(coefficients_array, principal_components) / distance - delta_t = kwargs['delta_t'] + delta_t = kwargs["delta_t"] outhp = TimeSeries(strain, delta_t=delta_t) outhc = TimeSeries(numpy.zeros(len(strain)), delta_t=delta_t) return outhp, outhc # Approximant names ########################################################### -supernovae_td_approximants = {'CoreCollapseBounce': get_corecollapse_bounce} +supernovae_td_approximants = {"CoreCollapseBounce": get_corecollapse_bounce} diff --git a/pycbc/waveform/utils.py b/pycbc/waveform/utils.py index 0e55223c583..feedbc92d5e 100644 --- a/pycbc/waveform/utils.py +++ b/pycbc/waveform/utils.py @@ -22,36 +22,39 @@ # # ============================================================================= # -"""This module contains convenience utilities for manipulating waveforms -""" +"""This module contains convenience utilities for manipulating waveforms""" from math import frexp -import numpy +import numpy from scipy import signal +from pycbc.constants import PI from pycbc.scheme import schemed from pycbc.types import ( - TimeSeries, FrequencySeries, Array, - complex_same_precision_as, real_same_precision_as + Array, + FrequencySeries, + TimeSeries, + complex_same_precision_as, + real_same_precision_as, ) -from pycbc.constants import PI - -def ceilpow2(n): - """convenience function to determine a power-of-2 upper frequency limit""" - signif,exponent = frexp(n) - if (signif < 0): - return 1; - if (signif == 0.5): - exponent -= 1; - return (1) << exponent; -def coalign_waveforms(h1, h2, psd=None, - low_frequency_cutoff=None, - high_frequency_cutoff=None, - resize=True): - """ Return two time series which are aligned in time and phase. +def ceilpow2(n): + """Convenience function to determine a power-of-2 upper frequency limit""" + signif, exponent = frexp(n) + if signif < 0: + return 1 + if signif == 0.5: + exponent -= 1 + return (1) << exponent + + +def coalign_waveforms( + h1, h2, psd=None, low_frequency_cutoff=None, high_frequency_cutoff=None, resize=True +): + """ + Return two time series which are aligned in time and phase. The alignment is only to the nearest sample point and all changes to the phase are made to the first input waveform. Waveforms should not be split @@ -80,8 +83,10 @@ def coalign_waveforms(h1, h2, psd=None, The shifted waveform to align with h2 h2: pycbc.type.TimeSeries The resized (if necessary) waveform to align with h1. + """ from pycbc.filter import matched_filter + mlen = ceilpow2(max(len(h1), len(h2))) h1 = h1.copy() @@ -91,23 +96,30 @@ def coalign_waveforms(h1, h2, psd=None, h1.resize(mlen) h2.resize(mlen) elif len(h1) != len(h2) or len(h2) % 2 != 0: - raise ValueError("Time series must be the same size and even if you do " - "not allow resizing") - - snr = matched_filter(h1, h2, psd=psd, - low_frequency_cutoff=low_frequency_cutoff, - high_frequency_cutoff=high_frequency_cutoff) - - _, l = snr.abs_max_loc() - rotation = snr[l] / abs(snr[l]) + raise ValueError( + "Time series must be the same size and even if you do not allow resizing" + ) + + snr = matched_filter( + h1, + h2, + psd=psd, + low_frequency_cutoff=low_frequency_cutoff, + high_frequency_cutoff=high_frequency_cutoff, + ) + + _, l = snr.abs_max_loc() + rotation = snr[l] / abs(snr[l]) h1 = (h1.to_frequencyseries() * rotation).to_timeseries() h1.roll(l) h1 = TimeSeries(h1, delta_t=h2.delta_t, epoch=h2.start_time) return h1, h2 + def phase_from_frequencyseries(htilde, remove_start_phase=True): - """Returns the phase from the given frequency-domain waveform. This assumes + """ + Returns the phase from the given frequency-domain waveform. This assumes that the waveform has been sampled finely enough that the phase cannot change by more than pi radians between each step. @@ -122,16 +134,17 @@ def phase_from_frequencyseries(htilde, remove_start_phase=True): ------- FrequencySeries The phase of the waveform as a function of frequency. + """ - p = numpy.unwrap(numpy.angle(htilde.data)).astype( - real_same_precision_as(htilde)) + p = numpy.unwrap(numpy.angle(htilde.data)).astype(real_same_precision_as(htilde)) if remove_start_phase: p += -p[0] - return FrequencySeries(p, delta_f=htilde.delta_f, epoch=htilde.epoch, - copy=False) + return FrequencySeries(p, delta_f=htilde.delta_f, epoch=htilde.epoch, copy=False) + def amplitude_from_frequencyseries(htilde): - """Returns the amplitude of the given frequency-domain waveform as a + """ + Returns the amplitude of the given frequency-domain waveform as a FrequencySeries. Parameters @@ -143,14 +156,17 @@ def amplitude_from_frequencyseries(htilde): ------- FrequencySeries The amplitude of the waveform as a function of frequency. + """ amp = abs(htilde.data).astype(real_same_precision_as(htilde)) - return FrequencySeries(amp, delta_f=htilde.delta_f, epoch=htilde.epoch, - copy=False) + return FrequencySeries(amp, delta_f=htilde.delta_f, epoch=htilde.epoch, copy=False) + -def time_from_frequencyseries(htilde, sample_frequencies=None, - discont_threshold=0.99*numpy.pi): - """Computes time as a function of frequency from the given +def time_from_frequencyseries( + htilde, sample_frequencies=None, discont_threshold=0.99 * numpy.pi +): + """ + Computes time as a function of frequency from the given frequency-domain waveform. This assumes the stationary phase approximation. Any frequencies lower than the first non-zero value in htilde are assigned the time at the first non-zero value. Times for any @@ -182,26 +198,32 @@ def time_from_frequencyseries(htilde, sample_frequencies=None, ------- FrequencySeries The time evolution of the waveform as a function of frequency. + """ if sample_frequencies is None: sample_frequencies = htilde.sample_frequencies.numpy() phase = phase_from_frequencyseries(htilde).data dphi = numpy.diff(phase) - time = -dphi / (2.*numpy.pi*numpy.diff(sample_frequencies)) + time = -dphi / (2.0 * numpy.pi * numpy.diff(sample_frequencies)) nzidx = numpy.nonzero(abs(htilde.data))[0] kmin, kmax = nzidx[0], nzidx[-2] # exclude everything after a discontinuity discont_idx = numpy.where(abs(dphi[kmin:]) >= discont_threshold)[0] if discont_idx.size != 0: - kmax = min(kmax, kmin + discont_idx[0]-1) + kmax = min(kmax, kmin + discont_idx[0] - 1) time[:kmin] = time[kmin] time[kmax:] = time[kmax] - return FrequencySeries(time.astype(real_same_precision_as(htilde)), - delta_f=htilde.delta_f, epoch=htilde.epoch, - copy=False) + return FrequencySeries( + time.astype(real_same_precision_as(htilde)), + delta_f=htilde.delta_f, + epoch=htilde.epoch, + copy=False, + ) + def phase_from_polarizations(h_plus, h_cross, remove_start_phase=True): - """Return gravitational wave phase + """ + Return gravitational wave phase Return the gravitation-wave phase from the h_plus and h_cross polarizations of the waveform. The returned phase is always @@ -222,6 +244,7 @@ def phase_from_polarizations(h_plus, h_cross, remove_start_phase=True): A TimeSeries containing the gravitational wave phase. Examples + -------- --------s >>> from pycbc.waveform import get_td_waveform, phase_from_polarizations >>> hp, hc = get_td_waveform(approximant="TaylorT4", mass1=10, mass2=10, @@ -230,14 +253,16 @@ def phase_from_polarizations(h_plus, h_cross, remove_start_phase=True): """ p = numpy.unwrap(numpy.arctan2(h_cross.data, h_plus.data)).astype( - real_same_precision_as(h_plus)) + real_same_precision_as(h_plus) + ) if remove_start_phase: p += -p[0] - return TimeSeries(p, delta_t=h_plus.delta_t, epoch=h_plus.start_time, - copy=False) + return TimeSeries(p, delta_t=h_plus.delta_t, epoch=h_plus.start_time, copy=False) + def amplitude_from_polarizations(h_plus, h_cross): - """Return gravitational wave amplitude + """ + Return gravitational wave amplitude Return the gravitation-wave amplitude from the h_plus and h_cross polarizations of the waveform. @@ -267,8 +292,10 @@ def amplitude_from_polarizations(h_plus, h_cross): amp = (h_plus.squared_norm() + h_cross.squared_norm()) ** (0.5) return TimeSeries(amp, delta_t=h_plus.delta_t, epoch=h_plus.start_time) + def frequency_from_polarizations(h_plus, h_cross): - """Return gravitational wave frequency + """ + Return gravitational wave frequency Return the gravitation-wave frequency as a function of time from the h_plus and h_cross polarizations of the waveform. @@ -299,20 +326,26 @@ def frequency_from_polarizations(h_plus, h_cross): """ phase = phase_from_polarizations(h_plus, h_cross) - freq = numpy.diff(phase) / ( 2 * PI * phase.delta_t ) + freq = numpy.diff(phase) / (2 * PI * phase.delta_t) start_time = phase.start_time + phase.delta_t / 2 - return TimeSeries(freq.astype(real_same_precision_as(h_plus)), - delta_t=phase.delta_t, epoch=start_time) + return TimeSeries( + freq.astype(real_same_precision_as(h_plus)), + delta_t=phase.delta_t, + epoch=start_time, + ) @schemed("pycbc.waveform.utils_") def apply_fseries_time_shift(htilde, dt, kmin=0, copy=True): - """Shifts a frequency domain waveform in time. The waveform is assumed to + """ + Shifts a frequency domain waveform in time. The waveform is assumed to be sampled at equal frequency intervals. """ + def apply_fd_time_shift(htilde, shifttime, kmin=0, fseries=None, copy=True): - """Shifts a frequency domain waveform in time. The shift applied is + """ + Shifts a frequency domain waveform in time. The shift applied is shiftime - htilde.epoch. Parameters @@ -336,27 +369,32 @@ def apply_fd_time_shift(htilde, shifttime, kmin=0, fseries=None, copy=True): A frequency series with the waveform shifted to the new time. If makecopy is True, will be a new frequency series; if makecopy is False, will be the same as htilde. + """ dt = float(shifttime - htilde.epoch) - if dt == 0.: + if dt == 0.0: # no shift to apply, just copy if desired if copy: - htilde = 1. * htilde + htilde = 1.0 * htilde elif isinstance(htilde, FrequencySeries): # FrequencySeries means equally sampled in frequency, use faster shifting htilde = apply_fseries_time_shift(htilde, dt, kmin=kmin, copy=copy) else: if fseries is None: fseries = htilde.sample_frequencies.numpy() - shift = Array(numpy.exp(-2j*numpy.pi*dt*fseries), - dtype=complex_same_precision_as(htilde)) + shift = Array( + numpy.exp(-2j * numpy.pi * dt * fseries), + dtype=complex_same_precision_as(htilde), + ) if copy: - htilde = 1. * htilde + htilde = 1.0 * htilde htilde *= shift return htilde -def td_taper(out, start, end, beta=8, side='left'): - """Applies a taper to the given TimeSeries. + +def td_taper(out, start, end, beta=8, side="left"): + """ + Applies a taper to the given TimeSeries. A half-kaiser window is used for the roll-off. @@ -381,27 +419,30 @@ def td_taper(out, start, end, beta=8, side='left'): ------- TimeSeries The tapered time series. + """ out = out.copy() width = end - start winlen = 2 * int(width / out.delta_t) - window = Array(signal.get_window(('kaiser', beta), winlen)) + window = Array(signal.get_window(("kaiser", beta), winlen)) xmin = int((start - out.start_time) / out.delta_t) - xmax = xmin + winlen//2 - if side == 'left': - out[xmin:xmax] *= window[:winlen//2] + xmax = xmin + winlen // 2 + if side == "left": + out[xmin:xmax] *= window[: winlen // 2] if xmin > 0: out[:xmin].clear() - elif side == 'right': - out[xmin:xmax] *= window[winlen//2:] + elif side == "right": + out[xmin:xmax] *= window[winlen // 2 :] if xmax < len(out): out[xmax:].clear() else: - raise ValueError("unrecognized side argument {}".format(side)) + raise ValueError(f"unrecognized side argument {side}") return out -def fd_taper(out, start, end, beta=8, side='left'): - """Applies a taper to the given FrequencySeries. + +def fd_taper(out, start, end, beta=8, side="left"): + """ + Applies a taper to the given FrequencySeries. A half-kaiser window is used for the roll-off. @@ -425,27 +466,30 @@ def fd_taper(out, start, end, beta=8, side='left'): ------- FrequencySeries The tapered frequency series. + """ out = out.copy() width = end - start winlen = 2 * int(width / out.delta_f) - window = Array(signal.get_window(('kaiser', beta), winlen)) + window = Array(signal.get_window(("kaiser", beta), winlen)) kmin = int(start / out.delta_f) - kmax = kmin + winlen//2 - if side == 'left': - out[kmin:kmax] *= window[:winlen//2] - out[:kmin] *= 0. - elif side == 'right': - out[kmin:kmax] *= window[winlen//2:] - out[kmax:] *= 0. + kmax = kmin + winlen // 2 + if side == "left": + out[kmin:kmax] *= window[: winlen // 2] + out[:kmin] *= 0.0 + elif side == "right": + out[kmin:kmax] *= window[winlen // 2 :] + out[kmax:] *= 0.0 else: - raise ValueError("unrecognized side argument {}".format(side)) + raise ValueError(f"unrecognized side argument {side}") return out -def fd_to_td(htilde, delta_t=None, left_window=None, right_window=None, - left_beta=8, right_beta=8): - """Converts a FD waveform to TD. +def fd_to_td( + htilde, delta_t=None, left_window=None, right_window=None, left_beta=8, right_beta=8 +): + """ + Converts a FD waveform to TD. A window can optionally be applied using ``fd_taper`` to the left or right side of the waveform before being converted to the time domain. @@ -472,18 +516,20 @@ def fd_to_td(htilde, delta_t=None, left_window=None, right_window=None, ------- TimeSeries The time-series representation of ``htilde``. + """ if left_window is not None: start, end = left_window - htilde = fd_taper(htilde, start, end, side='left', beta=left_beta) + htilde = fd_taper(htilde, start, end, side="left", beta=left_beta) if right_window is not None: start, end = right_window - htilde = fd_taper(htilde, start, end, side='right', beta=right_beta) + htilde = fd_taper(htilde, start, end, side="right", beta=right_beta) return htilde.to_timeseries(delta_t=delta_t) def redshift_waveform(srch, z, tref=0): - """Redshifts a time-domain or frequency-domain waveform. + """ + Redshifts a time-domain or frequency-domain waveform. The waveform is stretched in time by :math:`(1+z)` and its (time-domain) amplitude increased by :math:`(1+z)`. A time shift is also applied to the @@ -503,16 +549,17 @@ def redshift_waveform(srch, z, tref=0): ------- TimeSeries or FrequencySeries The red-shifted waveform. The return type will be the same as `srch`. + """ isfs = isinstance(srch, FrequencySeries) if isfs: redshifted = srch.to_timeseries() - redshifted *= 1+z + redshifted *= 1 + z else: - redshifted = (1+z) * srch - redshifted._delta_t *= 1+z + redshifted = (1 + z) * srch + redshifted._delta_t *= 1 + z # find the location of tref in the original time series - tindex = (tref - srch.start_time)/srch.delta_t + tindex = (tref - srch.start_time) / srch.delta_t # find what it's been stretched to and subtract that off from the start # time so as to keep the reference time in the same spot tnew = tindex * redshifted.delta_t + redshifted.start_time diff --git a/pycbc/waveform/utils_cuda.py b/pycbc/waveform/utils_cuda.py index bc295b770e3..70279a26796 100644 --- a/pycbc/waveform/utils_cuda.py +++ b/pycbc/waveform/utils_cuda.py @@ -22,13 +22,16 @@ # # ============================================================================= # -"""This module contains the CUDA-specific code for - convenience utilities for manipulating waveforms """ -from pycbc.types import FrequencySeries +This module contains the CUDA-specific code for +convenience utilities for manipulating waveforms +""" + +import numpy from mako.template import Template from pycuda.compiler import SourceModule -import numpy + +from pycbc.types import FrequencySeries time_shift_kernel = Template(""" __global__ void fseries_ts(float2 *out, float phi, @@ -77,12 +80,16 @@ fseries_ts_fn = mod.get_function("fseries_ts") fseries_ts_fn.prepare("Pfii") + def apply_fseries_time_shift(htilde, dt, kmin=0, copy=True): - """Shifts a frequency domain waveform in time. The waveform is assumed to + """ + Shifts a frequency domain waveform in time. The waveform is assumed to be sampled at equal frequency intervals. """ - if htilde.precision != 'single': - raise NotImplementedError("CUDA version of apply_fseries_time_shift only supports single precision") + if htilde.precision != "single": + raise NotImplementedError( + "CUDA version of apply_fseries_time_shift only supports single precision" + ) if copy: out = htilde.copy() @@ -98,6 +105,7 @@ def apply_fseries_time_shift(htilde, dt, kmin=0, copy=True): phi = numpy.float32(-2 * numpy.pi * dt * htilde.delta_f) fseries_ts_fn.prepared_call((nb, 1), (nt, 1, 1), out.data.gpudata, phi, kmin, kmax) if copy: - htilde = FrequencySeries(out, delta_f=htilde.delta_f, epoch=htilde.epoch, - copy=False) + htilde = FrequencySeries( + out, delta_f=htilde.delta_f, epoch=htilde.epoch, copy=False + ) return htilde diff --git a/pycbc/waveform/utils_cupy.py b/pycbc/waveform/utils_cupy.py index 22b1fdc5a76..5789a7e306b 100644 --- a/pycbc/waveform/utils_cupy.py +++ b/pycbc/waveform/utils_cupy.py @@ -22,16 +22,19 @@ # # ============================================================================= # -"""This module contains the CuPy-specific code for - convenience utilities for manipulating waveforms """ -from pycbc.types import FrequencySeries +This module contains the CuPy-specific code for +convenience utilities for manipulating waveforms +""" + import cupy as xp +from pycbc.types import FrequencySeries def apply_fseries_time_shift(htilde, dt, kmin=0, copy=True): - """Shifts a frequency domain waveform in time. The waveform is assumed to + """ + Shifts a frequency domain waveform in time. The waveform is assumed to be sampled at equal frequency intervals. """ out = xp.array(htilde.data, copy=copy) @@ -39,8 +42,9 @@ def apply_fseries_time_shift(htilde, dt, kmin=0, copy=True): kmax = len(htilde) fstimeshift(out, phi, kmin, kmax) if copy: - htilde = FrequencySeries(out, delta_f=htilde.delta_f, - epoch=htilde.epoch, copy=False) + htilde = FrequencySeries( + out, delta_f=htilde.delta_f, epoch=htilde.epoch, copy=False + ) return htilde @@ -50,5 +54,4 @@ def fstimeshift(freqseries, phi, kmin, kmax): # FIXME: Convert to ElementwiseKernel and use kmin and max in that. idx = xp.arange(len(freqseries)) phase_shift = xp.exp(phi * idx) - freqseries[:] = freqseries[:] * phase_shift - + freqseries[:] = freqseries[:] * phase_shift diff --git a/pycbc/waveform/waveform.py b/pycbc/waveform/waveform.py index 7d5ae66e672..cefb7a958c8 100644 --- a/pycbc/waveform/waveform.py +++ b/pycbc/waveform/waveform.py @@ -22,38 +22,56 @@ # # ============================================================================= # -"""Convenience functions to genenerate gravitational wave templates and +""" +Convenience functions to genenerate gravitational wave templates and waveforms. """ +import inspect import os -import lal, numpy -from pycbc.types import TimeSeries, FrequencySeries, zeros, Array -from pycbc.types import real_same_precision_as, complex_same_precision_as + +import lal +import numpy + +import pycbc import pycbc.scheme as _scheme -import inspect -from pycbc.fft import fft -from pycbc import pnutils, libutils -from pycbc.waveform import utils as wfutils -from pycbc.waveform import parameters +from pycbc import libutils, pnutils from pycbc.conversions import get_final_from_initial, tau_from_final_mass_spin +from pycbc.fft import fft from pycbc.filter import interpolate_complex_frequency, resample_to_delta_t -import pycbc -from .spa_tmplt import spa_tmplt, spa_tmplt_norm, spa_tmplt_end, \ - spa_tmplt_precondition, spa_amplitude_factor, \ - spa_length_in_time +from pycbc.types import ( + Array, + FrequencySeries, + TimeSeries, + complex_same_precision_as, + real_same_precision_as, + zeros, +) +from pycbc.waveform import parameters +from pycbc.waveform import utils as wfutils + +from .spa_tmplt import ( + spa_amplitude_factor, + spa_length_in_time, + spa_tmplt, + spa_tmplt_end, + spa_tmplt_norm, + spa_tmplt_precondition, +) + class NoWaveformError(Exception): - """This should be raised if generating a waveform would just result in all + """ + This should be raised if generating a waveform would just result in all zeros being returned, e.g., if a requested `f_final` is <= `f_lower`. """ - pass + class FailedWaveformError(Exception): - """This should be raised if a waveform fails to generate. - """ - pass + """This should be raised if a waveform fails to generate.""" + + # If this is set to True, waveform generation codes will try to regenerate # waveforms with known failure conditions to try to avoid the failure. For @@ -61,12 +79,12 @@ class FailedWaveformError(Exception): # If this is set to False waveform failures will always raise exceptions fail_tolerant_waveform_generation = True -default_args = \ - (parameters.fd_waveform_params.default_dict() + - parameters.td_waveform_params).default_dict() +default_args = ( + parameters.fd_waveform_params.default_dict() + parameters.td_waveform_params +).default_dict() -default_sgburst_args = {'eccentricity':0, 'polarization':0} -sgburst_required_args = ['q','frequency','hrss'] +default_sgburst_args = {"eccentricity": 0, "polarization": 0} +sgburst_required_args = ["q", "frequency", "hrss"] # td, fd, filter waveforms generated on the CPU _lalsim_td_approximants = {} @@ -74,8 +92,10 @@ class FailedWaveformError(Exception): _lalsim_enum = {} _lalsim_sgburst_approximants = {} + def _check_lal_pars(p): - """ Create a laldict object from the dictionary of waveform parameters + """ + Create a laldict object from the dictionary of waveform parameters Parameters ---------- @@ -86,137 +106,197 @@ def _check_lal_pars(p): ------- laldict: LalDict The lal type dictionary to pass to the lalsimulation waveform functions. + """ lal_pars = lal.CreateDict() - #nonGRparams can be straightforwardly added if needed, however they have to + # nonGRparams can be straightforwardly added if needed, however they have to # be invoked one by one - if p['phase_order']!=-1: - lalsimulation.SimInspiralWaveformParamsInsertPNPhaseOrder(lal_pars,int(p['phase_order'])) - if p['amplitude_order']!=-1: - lalsimulation.SimInspiralWaveformParamsInsertPNAmplitudeOrder(lal_pars,int(p['amplitude_order'])) - if p['spin_order']!=-1: - lalsimulation.SimInspiralWaveformParamsInsertPNSpinOrder(lal_pars,int(p['spin_order'])) - if p['tidal_order']!=-1: - lalsimulation.SimInspiralWaveformParamsInsertPNTidalOrder(lal_pars, p['tidal_order']) - if p['eccentricity_order']!=-1: - lalsimulation.SimInspiralWaveformParamsInsertPNEccentricityOrder(lal_pars, p['eccentricity_order']) - if p['lambda1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalLambda1(lal_pars, p['lambda1']) - if p['lambda2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalLambda2(lal_pars, p['lambda2']) - if p['lambda_octu1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarLambda1(lal_pars, p['lambda_octu1']) - if p['lambda_octu2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarLambda2(lal_pars, p['lambda_octu2']) - if p['quadfmode1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalQuadrupolarFMode1(lal_pars, p['quadfmode1']) - if p['quadfmode2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalQuadrupolarFMode2(lal_pars, p['quadfmode2']) - if p['octufmode1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarFMode1(lal_pars, p['octufmode1']) - if p['octufmode2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarFMode2(lal_pars, p['octufmode2']) - if p['dquad_mon1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertdQuadMon1(lal_pars, p['dquad_mon1']) - if p['dquad_mon2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertdQuadMon2(lal_pars, p['dquad_mon2']) - if p['numrel_data']: - lalsimulation.SimInspiralWaveformParamsInsertNumRelData(lal_pars, str(p['numrel_data'])) - if p['modes_choice']: - lalsimulation.SimInspiralWaveformParamsInsertModesChoice(lal_pars, p['modes_choice']) - if p['frame_axis']: - lalsimulation.SimInspiralWaveformParamsInsertFrameAxis(lal_pars, p['frame_axis']) - if p['side_bands']: - lalsimulation.SimInspiralWaveformParamsInsertSideband(lal_pars, p['side_bands']) - if p['mode_array'] is not None: + if p["phase_order"] != -1: + lalsimulation.SimInspiralWaveformParamsInsertPNPhaseOrder( + lal_pars, int(p["phase_order"]) + ) + if p["amplitude_order"] != -1: + lalsimulation.SimInspiralWaveformParamsInsertPNAmplitudeOrder( + lal_pars, int(p["amplitude_order"]) + ) + if p["spin_order"] != -1: + lalsimulation.SimInspiralWaveformParamsInsertPNSpinOrder( + lal_pars, int(p["spin_order"]) + ) + if p["tidal_order"] != -1: + lalsimulation.SimInspiralWaveformParamsInsertPNTidalOrder( + lal_pars, p["tidal_order"] + ) + if p["eccentricity_order"] != -1: + lalsimulation.SimInspiralWaveformParamsInsertPNEccentricityOrder( + lal_pars, p["eccentricity_order"] + ) + if p["lambda1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalLambda1( + lal_pars, p["lambda1"] + ) + if p["lambda2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalLambda2( + lal_pars, p["lambda2"] + ) + if p["lambda_octu1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarLambda1( + lal_pars, p["lambda_octu1"] + ) + if p["lambda_octu2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarLambda2( + lal_pars, p["lambda_octu2"] + ) + if p["quadfmode1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalQuadrupolarFMode1( + lal_pars, p["quadfmode1"] + ) + if p["quadfmode2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalQuadrupolarFMode2( + lal_pars, p["quadfmode2"] + ) + if p["octufmode1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarFMode1( + lal_pars, p["octufmode1"] + ) + if p["octufmode2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertTidalOctupolarFMode2( + lal_pars, p["octufmode2"] + ) + if p["dquad_mon1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertdQuadMon1( + lal_pars, p["dquad_mon1"] + ) + if p["dquad_mon2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertdQuadMon2( + lal_pars, p["dquad_mon2"] + ) + if p["numrel_data"]: + lalsimulation.SimInspiralWaveformParamsInsertNumRelData( + lal_pars, str(p["numrel_data"]) + ) + if p["modes_choice"]: + lalsimulation.SimInspiralWaveformParamsInsertModesChoice( + lal_pars, p["modes_choice"] + ) + if p["frame_axis"]: + lalsimulation.SimInspiralWaveformParamsInsertFrameAxis( + lal_pars, p["frame_axis"] + ) + if p["side_bands"]: + lalsimulation.SimInspiralWaveformParamsInsertSideband(lal_pars, p["side_bands"]) + if p["mode_array"] is not None: ma = lalsimulation.SimInspiralCreateModeArray() - for l,m in p['mode_array']: + for l, m in p["mode_array"]: lalsimulation.SimInspiralModeArrayActivateMode(ma, l, m) lalsimulation.SimInspiralWaveformParamsInsertModeArray(lal_pars, ma) - #TestingGR parameters: - if p['dchi0'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi0(lal_pars,p['dchi0']) - if p['dchi1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi1(lal_pars,p['dchi1']) - if p['dchi2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi2(lal_pars,p['dchi2']) - if p['dchi3'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi3(lal_pars,p['dchi3']) - if p['dchi4'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi4(lal_pars,p['dchi4']) - if p['dchi5'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi5(lal_pars,p['dchi5']) - if p['dchi5l'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi5L(lal_pars,p['dchi5l']) - if p['dchi6'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi6(lal_pars,p['dchi6']) - if p['dchi6l'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi6L(lal_pars,p['dchi6l']) - if p['dchi7'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi7(lal_pars,p['dchi7']) - if p['dalpha1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha1(lal_pars,p['dalpha1']) - if p['dalpha2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha2(lal_pars,p['dalpha2']) - if p['dalpha3'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha3(lal_pars,p['dalpha3']) - if p['dalpha4'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha4(lal_pars,p['dalpha4']) - if p['dalpha5'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha5(lal_pars,p['dalpha5']) - if p['dbeta1'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDBeta1(lal_pars,p['dbeta1']) - if p['dbeta2'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDBeta2(lal_pars,p['dbeta2']) - if p['dbeta3'] is not None: - lalsimulation.SimInspiralWaveformParamsInsertNonGRDBeta3(lal_pars,p['dbeta3']) + # TestingGR parameters: + if p["dchi0"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi0(lal_pars, p["dchi0"]) + if p["dchi1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi1(lal_pars, p["dchi1"]) + if p["dchi2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi2(lal_pars, p["dchi2"]) + if p["dchi3"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi3(lal_pars, p["dchi3"]) + if p["dchi4"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi4(lal_pars, p["dchi4"]) + if p["dchi5"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi5(lal_pars, p["dchi5"]) + if p["dchi5l"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi5L(lal_pars, p["dchi5l"]) + if p["dchi6"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi6(lal_pars, p["dchi6"]) + if p["dchi6l"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi6L(lal_pars, p["dchi6l"]) + if p["dchi7"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi7(lal_pars, p["dchi7"]) + if p["dalpha1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha1( + lal_pars, p["dalpha1"] + ) + if p["dalpha2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha2( + lal_pars, p["dalpha2"] + ) + if p["dalpha3"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha3( + lal_pars, p["dalpha3"] + ) + if p["dalpha4"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha4( + lal_pars, p["dalpha4"] + ) + if p["dalpha5"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDAlpha5( + lal_pars, p["dalpha5"] + ) + if p["dbeta1"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDBeta1(lal_pars, p["dbeta1"]) + if p["dbeta2"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDBeta2(lal_pars, p["dbeta2"]) + if p["dbeta3"] is not None: + lalsimulation.SimInspiralWaveformParamsInsertNonGRDBeta3(lal_pars, p["dbeta3"]) return lal_pars + def _lalsim_td_waveform(**p): lal_pars = _check_lal_pars(p) - #nonGRparams can be straightforwardly added if needed, however they have to + # nonGRparams can be straightforwardly added if needed, however they have to # be invoked one by one try: hp1, hc1 = lalsimulation.SimInspiralChooseTDWaveform( - float(pnutils.solar_mass_to_kg(p['mass1'])), - float(pnutils.solar_mass_to_kg(p['mass2'])), - float(p['spin1x']), float(p['spin1y']), float(p['spin1z']), - float(p['spin2x']), float(p['spin2y']), float(p['spin2z']), - pnutils.megaparsecs_to_meters(float(p['distance'])), - float(p['inclination']), float(p['coa_phase']), - float(p['long_asc_nodes']), float(p['eccentricity']), float(p['mean_per_ano']), - float(p['delta_t']), float(p['f_lower']), float(p['f_ref']), - lal_pars, - _lalsim_enum[p['approximant']]) + float(pnutils.solar_mass_to_kg(p["mass1"])), + float(pnutils.solar_mass_to_kg(p["mass2"])), + float(p["spin1x"]), + float(p["spin1y"]), + float(p["spin1z"]), + float(p["spin2x"]), + float(p["spin2y"]), + float(p["spin2z"]), + pnutils.megaparsecs_to_meters(float(p["distance"])), + float(p["inclination"]), + float(p["coa_phase"]), + float(p["long_asc_nodes"]), + float(p["eccentricity"]), + float(p["mean_per_ano"]), + float(p["delta_t"]), + float(p["f_lower"]), + float(p["f_ref"]), + lal_pars, + _lalsim_enum[p["approximant"]], + ) except RuntimeError: if not fail_tolerant_waveform_generation: raise # For some cases failure modes can occur. Here we add waveform-specific # instructions to try to work with waveforms that are known to fail. - if 'SEOBNRv3' in p['approximant']: + if "SEOBNRv3" in p["approximant"]: # Try doubling the sample time and redoing. # Don't want to get stuck in a loop though! - if 'delta_t_orig' not in p: - p['delta_t_orig'] = p['delta_t'] - p['delta_t'] = p['delta_t'] / 2. - if p['delta_t_orig'] / p['delta_t'] > 9: + if "delta_t_orig" not in p: + p["delta_t_orig"] = p["delta_t"] + p["delta_t"] = p["delta_t"] / 2.0 + if p["delta_t_orig"] / p["delta_t"] > 9: raise hp, hc = _lalsim_td_waveform(**p) - p['delta_t'] = p['delta_t_orig'] - hp = resample_to_delta_t(hp, hp.delta_t*2) - hc = resample_to_delta_t(hc, hc.delta_t*2) + p["delta_t"] = p["delta_t_orig"] + hp = resample_to_delta_t(hp, hp.delta_t * 2) + hc = resample_to_delta_t(hc, hc.delta_t * 2) return hp, hc raise - #lal.DestroyDict(lal_pars) + # lal.DestroyDict(lal_pars) hp = TimeSeries(hp1.data.data[:], delta_t=hp1.deltaT, epoch=hp1.epoch) hc = TimeSeries(hc1.data.data[:], delta_t=hc1.deltaT, epoch=hc1.epoch) return hp, hc + _lalsim_td_waveform.required = parameters.cbc_td_required + def _spintaylor_aligned_prec_swapper(**p): """ SpinTaylorF2 is only single spin, it also struggles with anti-aligned spin @@ -226,77 +306,92 @@ def _spintaylor_aligned_prec_swapper(**p): the case of nonaligned doublespin systems the code will fail at the waveform generator level. """ - orig_approximant = p['approximant'] - if p['spin2x'] == 0 and p['spin2y'] == 0 and p['spin1x'] == 0 and \ - p['spin1y'] == 0: - p['approximant'] = 'TaylorF2' + orig_approximant = p["approximant"] + if p["spin2x"] == 0 and p["spin2y"] == 0 and p["spin1x"] == 0 and p["spin1y"] == 0: + p["approximant"] = "TaylorF2" else: - p['approximant'] = 'SpinTaylorF2' + p["approximant"] = "SpinTaylorF2" hp, hc = _lalsim_fd_waveform(**p) - p['approximant'] = orig_approximant + p["approximant"] = orig_approximant return hp, hc + def _lalsim_fd_waveform(**p): lal_pars = _check_lal_pars(p) hp1, hc1 = lalsimulation.SimInspiralChooseFDWaveform( - float(pnutils.solar_mass_to_kg(p['mass1'])), - float(pnutils.solar_mass_to_kg(p['mass2'])), - float(p['spin1x']), float(p['spin1y']), float(p['spin1z']), - float(p['spin2x']), float(p['spin2y']), float(p['spin2z']), - pnutils.megaparsecs_to_meters(float(p['distance'])), - float(p['inclination']), float(p['coa_phase']), - float(p['long_asc_nodes']), float(p['eccentricity']), float(p['mean_per_ano']), - p['delta_f'], float(p['f_lower']), float(p['f_final']), float(p['f_ref']), - lal_pars, - _lalsim_enum[p['approximant']]) - - hp = FrequencySeries(hp1.data.data[:], delta_f=hp1.deltaF, - epoch=hp1.epoch) - - hc = FrequencySeries(hc1.data.data[:], delta_f=hc1.deltaF, - epoch=hc1.epoch) - #lal.DestroyDict(lal_pars) + float(pnutils.solar_mass_to_kg(p["mass1"])), + float(pnutils.solar_mass_to_kg(p["mass2"])), + float(p["spin1x"]), + float(p["spin1y"]), + float(p["spin1z"]), + float(p["spin2x"]), + float(p["spin2y"]), + float(p["spin2z"]), + pnutils.megaparsecs_to_meters(float(p["distance"])), + float(p["inclination"]), + float(p["coa_phase"]), + float(p["long_asc_nodes"]), + float(p["eccentricity"]), + float(p["mean_per_ano"]), + p["delta_f"], + float(p["f_lower"]), + float(p["f_final"]), + float(p["f_ref"]), + lal_pars, + _lalsim_enum[p["approximant"]], + ) + + hp = FrequencySeries(hp1.data.data[:], delta_f=hp1.deltaF, epoch=hp1.epoch) + + hc = FrequencySeries(hc1.data.data[:], delta_f=hc1.deltaF, epoch=hc1.epoch) + # lal.DestroyDict(lal_pars) return hp, hc + _lalsim_fd_waveform.required = parameters.cbc_fd_required + def _lalsim_sgburst_waveform(**p): - hp, hc = lalsimulation.SimBurstSineGaussian(float(p['q']), - float(p['frequency']), - float(p['hrss']), - float(p['eccentricity']), - float(p['polarization']), - float(p['delta_t'])) + hp, hc = lalsimulation.SimBurstSineGaussian( + float(p["q"]), + float(p["frequency"]), + float(p["hrss"]), + float(p["eccentricity"]), + float(p["polarization"]), + float(p["delta_t"]), + ) hp = TimeSeries(hp.data.data[:], delta_t=hp.deltaT, epoch=hp.epoch) hc = TimeSeries(hc.data.data[:], delta_t=hc.deltaT, epoch=hc.epoch) return hp, hc + # Populate waveform approximants from lalsimulation if the library is # available try: import lalsimulation - for approx_enum in range(0, lalsimulation.NumApproximants): + + for approx_enum in range(lalsimulation.NumApproximants): if lalsimulation.SimInspiralImplementedTDApproximants(approx_enum): approx_name = lalsimulation.GetStringFromApproximant(approx_enum) _lalsim_enum[approx_name] = approx_enum _lalsim_td_approximants[approx_name] = _lalsim_td_waveform - for approx_enum in range(0, lalsimulation.NumApproximants): + for approx_enum in range(lalsimulation.NumApproximants): if lalsimulation.SimInspiralImplementedFDApproximants(approx_enum): approx_name = lalsimulation.GetStringFromApproximant(approx_enum) _lalsim_enum[approx_name] = approx_enum _lalsim_fd_approximants[approx_name] = _lalsim_fd_waveform # sine-Gaussian burst - for approx_enum in range(0, lalsimulation.NumApproximants): + for approx_enum in range(lalsimulation.NumApproximants): if lalsimulation.SimInspiralImplementedFDApproximants(approx_enum): approx_name = lalsimulation.GetStringFromApproximant(approx_enum) _lalsim_enum[approx_name] = approx_enum _lalsim_sgburst_approximants[approx_name] = _lalsim_sgburst_waveform except ImportError: - lalsimulation = libutils.import_optional('lalsimulation') + lalsimulation = libutils.import_optional("lalsimulation") cpu_sgburst = _lalsim_sgburst_approximants cpu_td = dict(_lalsim_td_approximants.items()) @@ -309,72 +404,89 @@ def _lalsim_sgburst_waveform(**p): if pycbc.HAVE_CUDA: from pycbc.waveform.pycbc_phenomC_tmplt import imrphenomc_tmplt from pycbc.waveform.SpinTaylorF2 import spintaylorf2 as cuda_spintaylorf2 + _cuda_fd_approximants["IMRPhenomC"] = imrphenomc_tmplt _cuda_fd_approximants["SpinTaylorF2"] = cuda_spintaylorf2 -cuda_td = dict(list(_lalsim_td_approximants.items()) + list(_cuda_td_approximants.items())) -cuda_fd = dict(list(_lalsim_fd_approximants.items()) + list(_cuda_fd_approximants.items())) +cuda_td = dict( + list(_lalsim_td_approximants.items()) + list(_cuda_td_approximants.items()) +) +cuda_fd = dict( + list(_lalsim_fd_approximants.items()) + list(_cuda_fd_approximants.items()) +) # List the various available approximants #################################### + def print_td_approximants(): print("LalSimulation Approximants") - for approx in _lalsim_td_approximants.keys(): + for approx in _lalsim_td_approximants: print(" " + approx) print("CUDA Approximants") - for approx in _cuda_td_approximants.keys(): + for approx in _cuda_td_approximants: print(" " + approx) + def print_fd_approximants(): print("LalSimulation Approximants") - for approx in _lalsim_fd_approximants.keys(): + for approx in _lalsim_fd_approximants: print(" " + approx) print("CUDA Approximants") - for approx in _cuda_fd_approximants.keys(): + for approx in _cuda_fd_approximants: print(" " + approx) + def print_sgburst_approximants(): print("LalSimulation Approximants") - for approx in _lalsim_sgburst_approximants.keys(): + for approx in _lalsim_sgburst_approximants: print(" " + approx) + def td_approximants(scheme=_scheme.mgr.state): - """Return a list containing the available time domain approximants for - the given processing scheme. + """ + Return a list containing the available time domain approximants for + the given processing scheme. """ return list(td_wav[type(scheme)].keys()) + def fd_approximants(scheme=_scheme.mgr.state): - """Return a list containing the available fourier domain approximants for - the given processing scheme. + """ + Return a list containing the available fourier domain approximants for + the given processing scheme. """ return list(fd_wav[type(scheme)].keys()) + def sgburst_approximants(scheme=_scheme.mgr.state): - """Return a list containing the available time domain sgbursts for - the given processing scheme. + """ + Return a list containing the available time domain sgbursts for + the given processing scheme. """ return list(sgburst_wav[type(scheme)].keys()) + def filter_approximants(scheme=_scheme.mgr.state): - """Return a list of fourier domain approximants including those - written specifically as templates. + """ + Return a list of fourier domain approximants including those + written specifically as templates. """ return list(filter_wav[type(scheme)].keys()) + # Input parameter handling ################################################### + def get_obj_attrs(obj): - """ Return a dictionary built from the attributes of the given object. - """ + """Return a dictionary built from the attributes of the given object.""" pr = {} if obj is not None: if isinstance(obj, numpy.record): for name in obj.dtype.names: pr[name] = getattr(obj, name) - elif hasattr(obj, '__dict__') and obj.__dict__: + elif hasattr(obj, "__dict__") and obj.__dict__: pr = obj.__dict__ - elif hasattr(obj, '__slots__'): + elif hasattr(obj, "__slots__"): for slot in obj.__slots__: if hasattr(obj, slot): pr[slot] = getattr(obj, slot) @@ -384,7 +496,7 @@ def get_obj_attrs(obj): for name in dir(obj): try: value = getattr(obj, name) - if not name.startswith('__') and not inspect.ismethod(value): + if not name.startswith("__") and not inspect.ismethod(value): pr[name] = value except: continue @@ -393,7 +505,8 @@ def get_obj_attrs(obj): def parse_mode_array(input_params): - """Ensures mode_array argument in a dictionary of input parameters is + """ + Ensures mode_array argument in a dictionary of input parameters is a list of tuples of (l, m), where l and m are ints. Accepted formats for the ``mode_array`` argument is a list of tuples of @@ -401,8 +514,8 @@ def parse_mode_array(input_params): the modes (e.g., ``22 33 44``), or an array of ints or floats (e.g., ``[22., 33., 44.]``. """ - if 'mode_array' in input_params and input_params['mode_array'] is not None: - mode_array = input_params['mode_array'] + if "mode_array" in input_params and input_params["mode_array"] is not None: + mode_array = input_params["mode_array"] if isinstance(mode_array, str): mode_array = mode_array.split() if not isinstance(mode_array, (numpy.ndarray, list)): @@ -414,9 +527,9 @@ def parse_mode_array(input_params): ma = str(int(ma)) # if ma is a str convert to (int, int) (e.g., '22' -> (2, 2)) if isinstance(ma, str): - if len(ma) == 2: # format is "22", presumed m is positive + if len(ma) == 2: # format is "22", presumed m is positive l, m = ma - elif len(ma) == 3: # format is "2+2", "2-2", signed m + elif len(ma) == 3: # format is "2+2", "2-2", signed m l = ma[0] m = ma[1:] else: @@ -424,12 +537,13 @@ def parse_mode_array(input_params): ma = (int(l), int(m)) mode_array[ii] = ma - input_params['mode_array'] = mode_array + input_params["mode_array"] = mode_array return input_params def props(obj, **kwargs): - """ Return a dictionary built from the combination of defaults, kwargs, + """ + Return a dictionary built from the combination of defaults, kwargs, and the attributes of the given object. """ pr = get_obj_attrs(obj) @@ -444,24 +558,26 @@ def props(obj, **kwargs): def check_args(args, required_args): - """ check that required args are given """ + """Check that required args are given""" missing = [] for arg in required_args: if (arg not in args) or (args[arg] is None): missing.append(arg) if len(missing) != 0: - raise ValueError("Please provide {}".format(', '.join(missing))) + raise ValueError("Please provide {}".format(", ".join(missing))) + # Input parameter handling for bursts ######################################## + def props_sgburst(obj, **kwargs): pr = {} if obj is not None: for name in dir(obj): try: value = getattr(obj, name) - if not name.startswith('__') and not inspect.ismethod(value): + if not name.startswith("__") and not inspect.ismethod(value): pr[name] = value except: continue @@ -474,28 +590,36 @@ def props_sgburst(obj, **kwargs): return input_params + # Waveform generation ######################################################## fd_sequence = {} fd_det_sequence = {} fd_det = {} + def _lalsim_fd_sequence(**p): - """ Shim to interface to lalsimulation SimInspiralChooseFDWaveformSequence - """ + """Shim to interface to lalsimulation SimInspiralChooseFDWaveformSequence""" lal_pars = _check_lal_pars(p) hp, hc = lalsimulation.SimInspiralChooseFDWaveformSequence( - float(p['coa_phase']), - float(pnutils.solar_mass_to_kg(p['mass1'])), - float(pnutils.solar_mass_to_kg(p['mass2'])), - float(p['spin1x']), float(p['spin1y']), float(p['spin1z']), - float(p['spin2x']), float(p['spin2y']), float(p['spin2z']), - float(p['f_ref']), - pnutils.megaparsecs_to_meters(float(p['distance'])), - float(p['inclination']), - lal_pars, - _lalsim_enum[p['approximant']], - p['sample_points'].lal()) + float(p["coa_phase"]), + float(pnutils.solar_mass_to_kg(p["mass1"])), + float(pnutils.solar_mass_to_kg(p["mass2"])), + float(p["spin1x"]), + float(p["spin1y"]), + float(p["spin1z"]), + float(p["spin2x"]), + float(p["spin2y"]), + float(p["spin2z"]), + float(p["f_ref"]), + pnutils.megaparsecs_to_meters(float(p["distance"])), + float(p["inclination"]), + lal_pars, + _lalsim_enum[p["approximant"]], + p["sample_points"].lal(), + ) return Array(hp.data.data), Array(hc.data.data) + + _lalsim_fd_sequence.required = parameters.cbc_fd_required for apx in _lalsim_enum: @@ -503,7 +627,8 @@ def _lalsim_fd_sequence(**p): def get_fd_waveform_sequence(template=None, **kwds): - """Return values of the waveform evaluated at the sequence of frequency + """ + Return values of the waveform evaluated at the sequence of frequency points. The waveform generator doesn't include detector response. Parameters @@ -522,25 +647,27 @@ def get_fd_waveform_sequence(template=None, **kwds): hcrosstilde: Array The cross phase of the waveform in frequency domain evaluated at the frequency points. + """ input_params = props(template, **kwds) - input_params['delta_f'] = -1 - input_params['f_lower'] = -1 - if input_params['approximant'] not in fd_sequence: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) - wav_gen = fd_sequence[input_params['approximant']] - if hasattr(wav_gen, 'required'): + input_params["delta_f"] = -1 + input_params["f_lower"] = -1 + if input_params["approximant"] not in fd_sequence: + raise ValueError("Approximant %s not available" % (input_params["approximant"])) + wav_gen = fd_sequence[input_params["approximant"]] + if hasattr(wav_gen, "required"): required = wav_gen.required else: required = parameters.fd_required - if not isinstance(input_params['sample_points'], Array): - input_params['sample_points'] = Array(input_params['sample_points']) + if not isinstance(input_params["sample_points"], Array): + input_params["sample_points"] = Array(input_params["sample_points"]) check_args(input_params, required) return wav_gen(**input_params) + def get_fd_det_waveform_sequence(template=None, **kwds): - """Return values of the waveform evaluated at the sequence of frequency + """ + Return values of the waveform evaluated at the sequence of frequency points. The waveform generator includes detector response. Parameters @@ -557,28 +684,35 @@ def get_fd_det_waveform_sequence(template=None, **kwds): The detector-frame waveform (with detector response) in frequency domain evaluated at the frequency points. Keys are requested data channels, values are FrequencySeries. + """ input_params = props(template, **kwds) - if input_params['approximant'] not in fd_det_sequence: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) - wav_gen = fd_det_sequence[input_params['approximant']] - if hasattr(wav_gen, 'required'): + if input_params["approximant"] not in fd_det_sequence: + raise ValueError("Approximant %s not available" % (input_params["approximant"])) + wav_gen = fd_det_sequence[input_params["approximant"]] + if hasattr(wav_gen, "required"): required = wav_gen.required else: required = parameters.fd_det_sequence_required check_args(input_params, required) return wav_gen(**input_params) + get_fd_waveform_sequence.__doc__ = get_fd_waveform_sequence.__doc__.format( - params=parameters.fd_waveform_sequence_params.docstr(prefix=" ", - include_label=False)) + params=parameters.fd_waveform_sequence_params.docstr( + prefix=" ", include_label=False + ) +) get_fd_det_waveform_sequence.__doc__ = get_fd_det_waveform_sequence.__doc__.format( - params=parameters.fd_waveform_sequence_params.docstr(prefix=" ", - include_label=False)) + params=parameters.fd_waveform_sequence_params.docstr( + prefix=" ", include_label=False + ) +) + def get_td_waveform(template=None, **kwargs): - """Return the plus and cross polarizations of a time domain waveform. + """ + Return the plus and cross polarizations of a time domain waveform. Parameters ---------- @@ -594,26 +728,29 @@ def get_td_waveform(template=None, **kwargs): The plus polarization of the waveform. hcross: TimeSeries The cross polarization of the waveform. + """ input_params = props(template, **kwargs) wav_gen = td_wav[type(_scheme.mgr.state)] - if input_params['approximant'] not in wav_gen: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) - wav_gen = wav_gen[input_params['approximant']] - if hasattr(wav_gen, 'required'): + if input_params["approximant"] not in wav_gen: + raise ValueError("Approximant %s not available" % (input_params["approximant"])) + wav_gen = wav_gen[input_params["approximant"]] + if hasattr(wav_gen, "required"): required = wav_gen.required else: required = parameters.td_required check_args(input_params, required) return wav_gen(**input_params) + get_td_waveform.__doc__ = get_td_waveform.__doc__.format( - params=parameters.td_waveform_params.docstr(prefix=" ", - include_label=False)) + params=parameters.td_waveform_params.docstr(prefix=" ", include_label=False) +) + def get_fd_waveform(template=None, **kwargs): - """Return a frequency domain gravitational waveform. + """ + Return a frequency domain gravitational waveform. Parameters ---------- @@ -629,28 +766,29 @@ def get_fd_waveform(template=None, **kwargs): The plus phase of the waveform in frequency domain. hcrosstilde: FrequencySeries The cross phase of the waveform in frequency domain. + """ input_params = props(template, **kwargs) wav_gen = fd_wav[type(_scheme.mgr.state)] - if input_params['approximant'] not in wav_gen: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) + if input_params["approximant"] not in wav_gen: + raise ValueError("Approximant %s not available" % (input_params["approximant"])) try: - ffunc = input_params.pop('f_final_func') - if ffunc != '': + ffunc = input_params.pop("f_final_func") + if ffunc != "": # convert the frequency function to a value - input_params['f_final'] = pnutils.named_frequency_cutoffs[ffunc]( - input_params) + input_params["f_final"] = pnutils.named_frequency_cutoffs[ffunc]( + input_params + ) # if the f_final is < f_lower, raise a NoWaveformError - if 'f_final' in input_params and \ - (input_params['f_lower']+input_params['delta_f'] >= - input_params['f_final']): - raise NoWaveformError("cannot generate waveform: f_lower >= " - "f_final") + if "f_final" in input_params and ( + input_params["f_lower"] + input_params["delta_f"] + >= input_params["f_final"] + ): + raise NoWaveformError("cannot generate waveform: f_lower >= f_final") except KeyError: pass - wav_gen = wav_gen[input_params['approximant']] - if hasattr(wav_gen, 'required'): + wav_gen = wav_gen[input_params["approximant"]] + if hasattr(wav_gen, "required"): required = wav_gen.required else: required = parameters.fd_required @@ -659,11 +797,13 @@ def get_fd_waveform(template=None, **kwargs): get_fd_waveform.__doc__ = get_fd_waveform.__doc__.format( - params=parameters.fd_waveform_params.docstr(prefix=" ", - include_label=False)) + params=parameters.fd_waveform_params.docstr(prefix=" ", include_label=False) +) + def get_fd_waveform_from_td(**params): - """ Return time domain version of fourier domain approximant. + """ + Return time domain version of fourier domain approximant. This returns a frequency domain version of a fourier domain approximant, with padding and tapering at the start of the waveform. @@ -680,9 +820,10 @@ def get_fd_waveform_from_td(**params): Plus polarization time series hc: pycbc.types.FrequencySeries Cross polarization time series + """ nparams = params.copy() - if not 'taper_method' in params: + if "taper_method" not in params: # determine the duration to use for an automatic tapering choice. # If taper method specified, assume they have set f_lower as they # want exactly. @@ -690,10 +831,10 @@ def get_fd_waveform_from_td(**params): while full_duration < duration * 1.5: full_duration = get_waveform_filter_length_in_time(**nparams) - nparams['f_lower'] -= 1 + nparams["f_lower"] -= 1 - if 'f_ref' not in nparams and 'f_lower' in nparams: - nparams['f_ref'] = nparams['f_lower'] + if "f_ref" not in nparams and "f_lower" in nparams: + nparams["f_ref"] = nparams["f_lower"] # We'll try to do the right thing and figure out what the frequency # end is. Otherwise, we'll just assume 2048 Hz. @@ -701,45 +842,53 @@ def get_fd_waveform_from_td(**params): # approximants try: f_end = get_waveform_end_frequency(**params) - delta_t = (0.5 / pnutils.nearest_larger_binary_number(f_end)) + delta_t = 0.5 / pnutils.nearest_larger_binary_number(f_end) except: delta_t = 1.0 / 2048 - nparams['delta_t'] = delta_t + nparams["delta_t"] = delta_t hp, hc = get_td_waveform(**nparams) # Resize to the right duration - tsamples = int(1.0 / params['delta_f'] / delta_t) + tsamples = int(1.0 / params["delta_f"] / delta_t) if tsamples < len(hp): - raise ValueError("The frequency spacing (df = {}) is too low to " - "generate the {} approximant from the time " - "domain".format(params['delta_f'], params['approximant'])) + raise ValueError( + "The frequency spacing (df = {}) is too low to " + "generate the {} approximant from the time " + "domain".format(params["delta_f"], params["approximant"]) + ) hp.resize(tsamples) hc.resize(tsamples) - if not 'taper_method' in params: + if "taper_method" not in params: # apply the tapering, we will use a safety factor here to allow for # somewhat inaccurate duration difference estimation. window = (full_duration - duration) * 0.8 hp = wfutils.td_taper(hp, hp.start_time, hp.start_time + window) hc = wfutils.td_taper(hc, hc.start_time, hc.start_time + window) else: - hp = hp.taper_timeseries(location=params['taper'], - tapermethod=params['taper_method'], - taper_window=params['taper_window']) - hc = hc.taper_timeseries(location=params['taper'], - tapermethod=params['taper_method'], - taper_window=params['taper_window']) + hp = hp.taper_timeseries( + location=params["taper"], + tapermethod=params["taper_method"], + taper_window=params["taper_window"], + ) + hc = hc.taper_timeseries( + location=params["taper"], + tapermethod=params["taper_method"], + taper_window=params["taper_window"], + ) # avoid wraparound hp = hp.to_frequencyseries().cyclic_time_shift(hp.start_time) hc = hc.to_frequencyseries().cyclic_time_shift(hc.start_time) return hp, hc + def get_fd_det_waveform(template=None, **kwargs): - """Return a frequency domain gravitational waveform. + """ + Return a frequency domain gravitational waveform. The waveform generator includes detector response. Parameters @@ -755,25 +904,28 @@ def get_fd_det_waveform(template=None, **kwargs): dict The detector-frame waveform (with detector response) in frequency domain. Keys are requested data channels, values are FrequencySeries. + """ input_params = props(template, **kwargs) - if input_params['approximant'] not in fd_det: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) - wav_gen = fd_det[input_params['approximant']] - if hasattr(wav_gen, 'required'): + if input_params["approximant"] not in fd_det: + raise ValueError("Approximant %s not available" % (input_params["approximant"])) + wav_gen = fd_det[input_params["approximant"]] + if hasattr(wav_gen, "required"): required = wav_gen.required else: required = parameters.fd_required check_args(input_params, required) return wav_gen(**input_params) + get_fd_det_waveform.__doc__ = get_fd_det_waveform.__doc__.format( - params=parameters.fd_waveform_params.docstr(prefix=" ", - include_label=False)) + params=parameters.fd_waveform_params.docstr(prefix=" ", include_label=False) +) + def _base_get_td_waveform_from_fd(template=None, rwrap=None, **params): - """ The base function to calculate time domain version of fourier + """ + The base function to calculate time domain version of fourier domain approximant which not include or includes detector response. Called by `get_td_waveform_from_fd` and `get_td_det_waveform_from_fd_det`. """ @@ -784,11 +936,14 @@ def _base_get_td_waveform_from_fd(template=None, rwrap=None, **params): # In the `pycbc.waveform.parameters` module, spin1z and # spin2z have the default value 0. Users must have input # masses, so no else is needed. - mass_spin_params = set(['mass1', 'mass2', 'spin1z', 'spin2z']) + mass_spin_params = set(["mass1", "mass2", "spin1z", "spin2z"]) if mass_spin_params.issubset(set(nparams.keys())): m_final, spin_final = get_final_from_initial( - mass1=nparams['mass1'], mass2=nparams['mass2'], - spin1z=nparams['spin1z'], spin2z=nparams['spin2z']) + mass1=nparams["mass1"], + mass2=nparams["mass2"], + spin1z=nparams["spin1z"], + spin2z=nparams["spin2z"], + ) rwrap = tau_from_final_mass_spin(m_final, spin_final) * 10 if rwrap < 5: # Long enough for very massive BBHs in XG detectors, @@ -796,35 +951,37 @@ def _base_get_td_waveform_from_fd(template=None, rwrap=None, **params): # computational burden for 2G cases. rwrap = 5 - if nparams['approximant'] not in _filter_time_lengths: - raise ValueError("Approximant %s _filter_time_lengths function \ - not available" % (nparams['approximant'])) + if nparams["approximant"] not in _filter_time_lengths: + raise ValueError( + "Approximant %s _filter_time_lengths function \ + not available" + % (nparams["approximant"]) + ) # determine the duration to use full_duration = duration = get_waveform_filter_length_in_time(**nparams) while full_duration < duration * 1.5: full_duration = get_waveform_filter_length_in_time(**nparams) - nparams['f_lower'] *= 0.99 - if 't_obs_start' in nparams and \ - full_duration >= nparams['t_obs_start']: + nparams["f_lower"] *= 0.99 + if "t_obs_start" in nparams and full_duration >= nparams["t_obs_start"]: break - if 'f_ref' not in nparams: - nparams['f_ref'] = params['f_lower'] + if "f_ref" not in nparams: + nparams["f_ref"] = params["f_lower"] # factor to ensure the vectors are all large enough. We don't need to # completely trust our duration estimator in this case, at a small # increase in computational cost - fudge_duration = (max(0, full_duration) + .1 + rwrap) * 1.5 - fsamples = int(fudge_duration / nparams['delta_t']) + fudge_duration = (max(0, full_duration) + 0.1 + rwrap) * 1.5 + fsamples = int(fudge_duration / nparams["delta_t"]) N = pnutils.nearest_larger_binary_number(fsamples) - fudge_duration = N * nparams['delta_t'] + fudge_duration = N * nparams["delta_t"] - nparams['delta_f'] = 1.0 / fudge_duration - tsize = int(1.0 / nparams['delta_t'] / nparams['delta_f']) + nparams["delta_f"] = 1.0 / fudge_duration + tsize = int(1.0 / nparams["delta_t"] / nparams["delta_f"]) fsize = tsize // 2 + 1 - if nparams['approximant'] not in fd_det: + if nparams["approximant"] not in fd_det: hp, hc = get_fd_waveform(**nparams) # Resize to the right sample rate hp.resize(fsize) @@ -834,26 +991,33 @@ def _base_get_td_waveform_from_fd(template=None, rwrap=None, **params): hp = hp.cyclic_time_shift(-rwrap) hc = hc.cyclic_time_shift(-rwrap) - hp = wfutils.fd_to_td(hp, delta_t=params['delta_t'], - left_window=(nparams['f_lower'], - params['f_lower'])) - hc = wfutils.fd_to_td(hc, delta_t=params['delta_t'], - left_window=(nparams['f_lower'], - params['f_lower'])) + hp = wfutils.fd_to_td( + hp, + delta_t=params["delta_t"], + left_window=(nparams["f_lower"], params["f_lower"]), + ) + hc = wfutils.fd_to_td( + hc, + delta_t=params["delta_t"], + left_window=(nparams["f_lower"], params["f_lower"]), + ) return hp, hc - else: - wfs = get_fd_det_waveform(**nparams) - for ifo in wfs.keys(): - wfs[ifo].resize(fsize) - # avoid wraparound - wfs[ifo] = wfs[ifo].cyclic_time_shift(-rwrap) - wfs[ifo] = wfutils.fd_to_td(wfs[ifo], delta_t=kwds['delta_t'], - left_window=(nparams['f_lower'], - kwds['f_lower'])) - return wfs + wfs = get_fd_det_waveform(**nparams) + for ifo in wfs.keys(): + wfs[ifo].resize(fsize) + # avoid wraparound + wfs[ifo] = wfs[ifo].cyclic_time_shift(-rwrap) + wfs[ifo] = wfutils.fd_to_td( + wfs[ifo], + delta_t=kwds["delta_t"], + left_window=(nparams["f_lower"], kwds["f_lower"]), + ) + return wfs + def get_td_waveform_from_fd(rwrap=None, **params): - """ Return time domain version of fourier domain approximant. + """ + Return time domain version of fourier domain approximant. This returns a time domain version of a fourier domain approximant, with padding and tapering at the start of the waveform. @@ -874,11 +1038,14 @@ def get_td_waveform_from_fd(rwrap=None, **params): Plus polarization time series hc: pycbc.types.TimeSeries Cross polarization time series + """ return _base_get_td_waveform_from_fd(None, rwrap, **params) + def get_td_det_waveform_from_fd_det(template=None, rwrap=None, **params): - """ Return time domain version of fourier domain approximant which + """ + Return time domain version of fourier domain approximant which includes detector response, with padding and tapering at the start of the waveform. @@ -897,35 +1064,37 @@ def get_td_det_waveform_from_fd_det(template=None, rwrap=None, **params): dict The detector-frame waveform (with detector response) in time domain. Keys are requested data channels. + """ return _base_get_td_waveform_from_fd(template, rwrap, **params) -get_td_det_waveform_from_fd_det.__doc__ = \ + +get_td_det_waveform_from_fd_det.__doc__ = ( get_td_det_waveform_from_fd_det.__doc__.format( - params=parameters.td_waveform_params.docstr(prefix=" ", - include_label=False)) + params=parameters.td_waveform_params.docstr(prefix=" ", include_label=False) + ) +) -def get_interpolated_fd_waveform(dtype=numpy.complex64, return_hc=True, - **params): - """ Return a fourier domain waveform approximant, using interpolation - """ + +def get_interpolated_fd_waveform(dtype=numpy.complex64, return_hc=True, **params): + """Return a fourier domain waveform approximant, using interpolation""" def rulog2(val): return 2.0 ** numpy.ceil(numpy.log2(float(val))) - orig_approx = params['approximant'] - params['approximant'] = params['approximant'].replace('_INTERP', '') - df = params['delta_f'] + orig_approx = params["approximant"] + params["approximant"] = params["approximant"].replace("_INTERP", "") + df = params["delta_f"] - if 'duration' not in params: + if "duration" not in params: duration = get_waveform_filter_length_in_time(**params) - elif params['duration'] > 0: - duration = params['duration'] + elif params["duration"] > 0: + duration = params["duration"] else: err_msg = "Waveform duration must be greater than 0." raise ValueError(err_msg) - #FIXME We should try to get this length directly somehow + # FIXME We should try to get this length directly somehow # I think this number should be conservative ringdown_padding = 0.5 @@ -934,9 +1103,8 @@ def rulog2(val): # off the inspiral when using ringdown_padding - 0.5. # Also, if ringdown_padding is set to a very small # value we can see cases where the ringdown is chopped. - if df_min > 0.5: - df_min = 0.5 - params['delta_f'] = df_min + df_min = min(df_min, 0.5) + params["delta_f"] = df_min hp, hc = get_fd_waveform(**params) hp = hp.astype(dtype) if return_hc: @@ -947,8 +1115,8 @@ def rulog2(val): f_end = get_waveform_end_frequency(**params) if f_end is None: f_end = (len(hp) - 1) * hp.delta_f - if 'f_final' in params and params['f_final'] > 0: - f_end_params = params['f_final'] + if "f_final" in params and params["f_final"] > 0: + f_end_params = params["f_final"] if f_end is not None: f_end = min(f_end_params, f_end) @@ -958,17 +1126,18 @@ def rulog2(val): if hc is not None: hc = hc[:n_min] - offset = int(ringdown_padding * (len(hp)-1)*2 * hp.delta_f) + offset = int(ringdown_padding * (len(hp) - 1) * 2 * hp.delta_f) - hp = interpolate_complex_frequency(hp, df, zeros_offset=offset, side='left') + hp = interpolate_complex_frequency(hp, df, zeros_offset=offset, side="left") if hc is not None: - hc = interpolate_complex_frequency(hc, df, zeros_offset=offset, - side='left') - params['approximant'] = orig_approx + hc = interpolate_complex_frequency(hc, df, zeros_offset=offset, side="left") + params["approximant"] = orig_approx return hp, hc + def get_sgburst_waveform(template=None, **kwargs): - """Return the plus and cross polarizations of a time domain + """ + Return the plus and cross polarizations of a time domain sine-Gaussian burst waveform. Parameters @@ -996,8 +1165,9 @@ def get_sgburst_waveform(template=None, **kwargs): The plus polarization of the waveform. hcross: TimeSeries The cross polarization of the waveform. + """ - input_params = props_sgburst(template,**kwargs) + input_params = props_sgburst(template, **kwargs) for arg in sgburst_required_args: if arg not in input_params: @@ -1005,6 +1175,7 @@ def get_sgburst_waveform(template=None, **kwargs): return _lalsim_sgburst_waveform(**input_params) + # Waveform filter routines ################################################### # Organize Filter Generators @@ -1012,15 +1183,18 @@ def get_sgburst_waveform(template=None, **kwargs): _cuda_fd_filters = {} _cupy_fd_filters = {} -_cuda_fd_filters['SPAtmplt'] = spa_tmplt -_cupy_fd_filters['SPAtmplt'] = spa_tmplt -_inspiral_fd_filters['SPAtmplt'] = spa_tmplt +_cuda_fd_filters["SPAtmplt"] = spa_tmplt +_cupy_fd_filters["SPAtmplt"] = spa_tmplt +_inspiral_fd_filters["SPAtmplt"] = spa_tmplt filter_wav = _scheme.ChooseBySchemeDict() -filter_wav.update( {_scheme.CPUScheme:_inspiral_fd_filters, - _scheme.CUDAScheme:_cuda_fd_filters, - _scheme.CUPYScheme:_cupy_fd_filters, - } ) +filter_wav.update( + { + _scheme.CPUScheme: _inspiral_fd_filters, + _scheme.CUDAScheme: _cuda_fd_filters, + _scheme.CUPYScheme: _cupy_fd_filters, + } +) # Organize functions for function conditioning/precalculated values _filter_norms = {} @@ -1029,81 +1203,85 @@ def get_sgburst_waveform(template=None, **kwargs): _template_amplitude_norms = {} _filter_time_lengths = {} + def seobnrv2_final_frequency(**kwds): - return pnutils.get_final_freq("SEOBNRv2", kwds['mass1'], kwds['mass2'], - kwds['spin1z'], kwds['spin2z']) + return pnutils.get_final_freq( + "SEOBNRv2", kwds["mass1"], kwds["mass2"], kwds["spin1z"], kwds["spin2z"] + ) + def get_imr_length(approx, **kwds): - """Call through to pnutils to obtain IMR waveform durations - """ - m1 = float(kwds['mass1']) - m2 = float(kwds['mass2']) - s1z = float(kwds['spin1z']) - s2z = float(kwds['spin2z']) - f_low = float(kwds['f_lower']) + """Call through to pnutils to obtain IMR waveform durations""" + m1 = float(kwds["mass1"]) + m2 = float(kwds["mass2"]) + s1z = float(kwds["spin1z"]) + s2z = float(kwds["spin2z"]) + f_low = float(kwds["f_lower"]) # 10% margin of error is incorporated in the pnutils function return pnutils.get_imr_duration(m1, m2, s1z, s2z, f_low, approximant=approx) + def seobnrv2_length_in_time(**kwds): - """Stub for holding the calculation of SEOBNRv2* waveform duration. - """ + """Stub for holding the calculation of SEOBNRv2* waveform duration.""" return get_imr_length("SEOBNRv2", **kwds) + def seobnrv4_length_in_time(**kwds): - """Stub for holding the calculation of SEOBNRv4* waveform duration. - """ + """Stub for holding the calculation of SEOBNRv4* waveform duration.""" return get_imr_length("SEOBNRv4", **kwds) + def seobnrv5_length_in_time(**kwds): - """Stub for holding the calculation of SEOBNRv5_ROM waveform duration. - """ + """Stub for holding the calculation of SEOBNRv5_ROM waveform duration.""" return get_imr_length("SEOBNRv5_ROM", **kwds) + def imrphenomd_length_in_time(**kwds): - """Stub for holding the calculation of IMRPhenomD waveform duration. - """ + """Stub for holding the calculation of IMRPhenomD waveform duration.""" return get_imr_length("IMRPhenomD", **kwds) + def imrphenomhm_length_in_time(**kwargs): - """Estimates the duration of IMRPhenom waveforms that include higher modes. - """ + """Estimates the duration of IMRPhenom waveforms that include higher modes.""" # Default maximum node number for IMRPhenomHM is 4 # The relevant lower order approximant here is IMRPhenomD return get_hm_length_in_time("IMRPhenomD", 4, **kwargs) + def seobnrv4hm_length_in_time(**kwargs): - """ Estimates the duration of SEOBNRv4HM waveforms that include higher modes. - """ + """Estimates the duration of SEOBNRv4HM waveforms that include higher modes.""" # Default maximum node number for SEOBNRv4HM is 5 # The relevant lower order approximant here is SEOBNRv4 - return get_hm_length_in_time('SEOBNRv4', 5, **kwargs) + return get_hm_length_in_time("SEOBNRv4", 5, **kwargs) + def get_hm_length_in_time(lor_approx, maxm_default, **kwargs): kwargs = parse_mode_array(kwargs) - if 'mode_array' in kwargs and kwargs['mode_array'] is not None: - maxm = max(m for _, m in kwargs['mode_array']) + if "mode_array" in kwargs and kwargs["mode_array"] is not None: + maxm = max(m for _, m in kwargs["mode_array"]) else: maxm = maxm_default try: - flow = kwargs['f_lower'] + flow = kwargs["f_lower"] except KeyError: raise ValueError("must provide a f_lower") - kwargs['f_lower'] = flow * 2./maxm + kwargs["f_lower"] = flow * 2.0 / maxm return get_imr_length(lor_approx, **kwargs) + _filter_norms["SPAtmplt"] = spa_tmplt_norm _filter_preconditions["SPAtmplt"] = spa_tmplt_precondition _filter_ends["SPAtmplt"] = spa_tmplt_end _filter_ends["TaylorF2"] = spa_tmplt_end -#_filter_ends["SEOBNRv1_ROM_EffectiveSpin"] = seobnrv2_final_frequency -#_filter_ends["SEOBNRv1_ROM_DoubleSpin"] = seobnrv2_final_frequency -#_filter_ends["SEOBNRv2_ROM_EffectiveSpin"] = seobnrv2_final_frequency -#_filter_ends["SEOBNRv2_ROM_DoubleSpin"] = seobnrv2_final_frequency -#_filter_ends["SEOBNRv2_ROM_DoubleSpin_HI"] = seobnrv2_final_frequency +# _filter_ends["SEOBNRv1_ROM_EffectiveSpin"] = seobnrv2_final_frequency +# _filter_ends["SEOBNRv1_ROM_DoubleSpin"] = seobnrv2_final_frequency +# _filter_ends["SEOBNRv2_ROM_EffectiveSpin"] = seobnrv2_final_frequency +# _filter_ends["SEOBNRv2_ROM_DoubleSpin"] = seobnrv2_final_frequency +# _filter_ends["SEOBNRv2_ROM_DoubleSpin_HI"] = seobnrv2_final_frequency # PhenomD returns higher frequencies than this, so commenting this out for now -#_filter_ends["IMRPhenomC"] = seobnrv2_final_frequency -#_filter_ends["IMRPhenomD"] = seobnrv2_final_frequency +# _filter_ends["IMRPhenomC"] = seobnrv2_final_frequency +# _filter_ends["IMRPhenomD"] = seobnrv2_final_frequency _template_amplitude_norms["SPAtmplt"] = spa_amplitude_factor _filter_time_lengths["SPAtmplt"] = spa_length_in_time @@ -1136,29 +1314,34 @@ def get_hm_length_in_time(lor_approx, maxm_default, **kwargs): # Also add generators for switching between approximants apx_name = "SpinTaylorF2_SWAPPER" -cpu_fd[apx_name] = _spintaylor_aligned_prec_swapper +cpu_fd[apx_name] = _spintaylor_aligned_prec_swapper _filter_time_lengths[apx_name] = _filter_time_lengths["SpinTaylorF2"] -from . nltides import nonlinear_tidal_spa +from .nltides import nonlinear_tidal_spa + cpu_fd["TaylorF2NL"] = nonlinear_tidal_spa from .premerger import premerger_taylorf2 -cpu_fd['PreTaylorF2'] = premerger_taylorf2 + +cpu_fd["PreTaylorF2"] = premerger_taylorf2 from .multiband import multiband_fd_waveform -cpu_fd['multiband'] = multiband_fd_waveform + +cpu_fd["multiband"] = multiband_fd_waveform # Load external waveforms ##################################################### -if 'PYCBC_WAVEFORM' in os.environ: - mods = os.environ['PYCBC_WAVEFORM'].split(':') +if "PYCBC_WAVEFORM" in os.environ: + mods = os.environ["PYCBC_WAVEFORM"].split(":") for mod in mods: - mhandle = __import__(mod, fromlist=['']) - mhandle.add_me(cpu_fd=cpu_fd, - cpu_td=cpu_td, - filter_time_lengths=_filter_time_lengths) + mhandle = __import__(mod, fromlist=[""]) + mhandle.add_me( + cpu_fd=cpu_fd, cpu_td=cpu_td, filter_time_lengths=_filter_time_lengths + ) + def td_fd_waveform_transform(approximant): - '''If the waveform approximant is in time domain, make a frequency domain + """ + If the waveform approximant is in time domain, make a frequency domain version using 'get_fd_waveform_from_td'; If the waveform approximant is in frequency domain, do interpolation for waveforms with a time length estimator, and make a time domain version using 'get_td_waveform_from_fd' @@ -1167,7 +1350,8 @@ def td_fd_waveform_transform(approximant): ---------- approximant: string The name of a waveform approximant. - ''' + + """ fd_apx = list(cpu_fd.keys()) td_apx = list(cpu_td.keys()) @@ -1177,7 +1361,7 @@ def td_fd_waveform_transform(approximant): if approximant in fd_apx and (approximant in _filter_time_lengths): # We can do interpolation for waveforms that have a time length - apx_int = approximant + '_INTERP' + apx_int = approximant + "_INTERP" cpu_fd[apx_int] = get_interpolated_fd_waveform _filter_time_lengths[apx_int] = _filter_time_lengths[approximant] @@ -1186,90 +1370,88 @@ def td_fd_waveform_transform(approximant): # (ex. IMRPhenomXX) cpu_td[approximant] = get_td_waveform_from_fd + for apx in list(_filter_time_lengths.keys()) + list(cpu_fd.keys()): td_fd_waveform_transform(apx) td_wav = _scheme.ChooseBySchemeDict() fd_wav = _scheme.ChooseBySchemeDict() -td_wav.update({_scheme.CPUScheme:cpu_td,_scheme.CUDAScheme:cuda_td}) -fd_wav.update({_scheme.CPUScheme:cpu_fd,_scheme.CUDAScheme:cuda_fd}) -sgburst_wav = {_scheme.CPUScheme:cpu_sgburst} +td_wav.update({_scheme.CPUScheme: cpu_td, _scheme.CUDAScheme: cuda_td}) +fd_wav.update({_scheme.CPUScheme: cpu_fd, _scheme.CUDAScheme: cuda_fd}) +sgburst_wav = {_scheme.CPUScheme: cpu_sgburst} + def get_waveform_filter(out, template=None, **kwargs): - """Return a frequency domain waveform filter for the specified approximant - """ + """Return a frequency domain waveform filter for the specified approximant""" n = len(out) input_params = props(template, **kwargs) - if input_params['approximant'] in filter_approximants(_scheme.mgr.state): + if input_params["approximant"] in filter_approximants(_scheme.mgr.state): wav_gen = filter_wav[type(_scheme.mgr.state)] - htilde = wav_gen[input_params['approximant']](out=out, **input_params) + htilde = wav_gen[input_params["approximant"]](out=out, **input_params) htilde.resize(n) htilde.chirp_length = get_waveform_filter_length_in_time(**input_params) htilde.length_in_time = htilde.chirp_length return htilde - if input_params['approximant'] in fd_approximants(_scheme.mgr.state): + if input_params["approximant"] in fd_approximants(_scheme.mgr.state): wav_gen = fd_wav[type(_scheme.mgr.state)] duration = get_waveform_filter_length_in_time(**input_params) - hp, _ = wav_gen[input_params['approximant']](duration=duration, - return_hc=False, **input_params) + hp, _ = wav_gen[input_params["approximant"]]( + duration=duration, return_hc=False, **input_params + ) hp.resize(n) - out[0:len(hp)] = hp[:] + out[0 : len(hp)] = hp[:] hp.data = out hp.length_in_time = hp.chirp_length = duration return hp - elif input_params['approximant'] in td_approximants(_scheme.mgr.state): + if input_params["approximant"] in td_approximants(_scheme.mgr.state): wav_gen = td_wav[type(_scheme.mgr.state)] - hp, _ = wav_gen[input_params['approximant']](**input_params) + hp, _ = wav_gen[input_params["approximant"]](**input_params) # taper the time series hp if required - if 'taper' in input_params.keys() and \ - input_params['taper'] is not None: - hp = wfutils.taper_timeseries(hp, input_params['taper'], - return_lal=False) + if "taper" in input_params.keys() and input_params["taper"] is not None: + hp = wfutils.taper_timeseries(hp, input_params["taper"], return_lal=False) return td_waveform_to_fd_waveform(hp, out=out) - else: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) - -def td_waveform_to_fd_waveform(waveform, out=None, length=None, - buffer_length=100): - """ Convert a time domain into a frequency domain waveform by FFT. - As a waveform is assumed to "wrap" in the time domain one must be - careful to ensure the waveform goes to 0 at both "boundaries". To - ensure this is done correctly the waveform must have the epoch set such - the merger time is at t=0 and the length of the waveform should be - shorter than the desired length of the FrequencySeries (times 2 - 1) - so that zeroes can be suitably pre- and post-pended before FFTing. - If given, out is a memory array to be used as the output of the FFT. - If not given memory is allocated internally. - If present the length of the returned FrequencySeries is determined - from the length out. If out is not given the length can be provided - expicitly, or it will be chosen as the nearest power of 2. If choosing - length explicitly the waveform length + buffer_length is used when - choosing the nearest binary number so that some zero padding is always - added. + raise ValueError("Approximant %s not available" % (input_params["approximant"])) + + +def td_waveform_to_fd_waveform(waveform, out=None, length=None, buffer_length=100): + """ + Convert a time domain into a frequency domain waveform by FFT. + As a waveform is assumed to "wrap" in the time domain one must be + careful to ensure the waveform goes to 0 at both "boundaries". To + ensure this is done correctly the waveform must have the epoch set such + the merger time is at t=0 and the length of the waveform should be + shorter than the desired length of the FrequencySeries (times 2 - 1) + so that zeroes can be suitably pre- and post-pended before FFTing. + If given, out is a memory array to be used as the output of the FFT. + If not given memory is allocated internally. + If present the length of the returned FrequencySeries is determined + from the length out. If out is not given the length can be provided + expicitly, or it will be chosen as the nearest power of 2. If choosing + length explicitly the waveform length + buffer_length is used when + choosing the nearest binary number so that some zero padding is always + added. """ # Figure out lengths and set out if needed if out is None: if length is None: - N = pnutils.nearest_larger_binary_number(len(waveform) + \ - buffer_length) - n = int(N//2) + 1 + N = pnutils.nearest_larger_binary_number(len(waveform) + buffer_length) + n = int(N // 2) + 1 else: n = length - N = (n-1)*2 + N = (n - 1) * 2 out = zeros(n, dtype=complex_same_precision_as(waveform)) else: n = len(out) - N = (n-1)*2 - delta_f = 1. / (N * waveform.delta_t) + N = (n - 1) * 2 + delta_f = 1.0 / (N * waveform.delta_t) # total duration of the waveform tmplt_length = len(waveform) * waveform.delta_t @@ -1279,13 +1461,13 @@ def td_waveform_to_fd_waveform(waveform, out=None, length=None, err_msg += "not supported in this function. Please shorten the " err_msg += "waveform appropriately before calling this function or " err_msg += "increase the allowed waveform length. " - err_msg += "Waveform length (in samples): {}".format(len(waveform)) - err_msg += ". Intended length: {}.".format(N) + err_msg += f"Waveform length (in samples): {len(waveform)}" + err_msg += f". Intended length: {N}." raise ValueError(err_msg) # for IMR templates the zero of time is at max amplitude (merger) # thus the start time is minus the duration of the template from # lower frequency cutoff to merger, i.e. minus the 'chirp time' - tChirp = - float( waveform.start_time ) # conversion from LIGOTimeGPS + tChirp = -float(waveform.start_time) # conversion from LIGOTimeGPS waveform.resize(N) k_zero = int(waveform.start_time / waveform.delta_t) waveform.roll(k_zero) @@ -1295,8 +1477,10 @@ def td_waveform_to_fd_waveform(waveform, out=None, length=None, htilde.chirp_length = tChirp return htilde + def get_two_pol_waveform_filter(outplus, outcross, template, **kwargs): - """Return a frequency domain waveform filter for the specified approximant. + """ + Return a frequency domain waveform filter for the specified approximant. Unlike get_waveform_filter this function returns both h_plus and h_cross components of the waveform, which are needed for searches where h_plus and h_cross are not related by a simple phase shift. @@ -1304,47 +1488,52 @@ def get_two_pol_waveform_filter(outplus, outcross, template, **kwargs): n = len(outplus) # If we don't have an inclination column alpha3 might be used - if not hasattr(template, 'inclination') and 'inclination' not in kwargs: - if hasattr(template, 'alpha3'): - kwargs['inclination'] = template.alpha3 + if not hasattr(template, "inclination") and "inclination" not in kwargs: + if hasattr(template, "alpha3"): + kwargs["inclination"] = template.alpha3 input_params = props(template, **kwargs) - if input_params['approximant'] in fd_approximants(_scheme.mgr.state): + if input_params["approximant"] in fd_approximants(_scheme.mgr.state): wav_gen = fd_wav[type(_scheme.mgr.state)] - hp, hc = wav_gen[input_params['approximant']](**input_params) + hp, hc = wav_gen[input_params["approximant"]](**input_params) hp.resize(n) hc.resize(n) - outplus[0:len(hp)] = hp[:] + outplus[0 : len(hp)] = hp[:] hp = FrequencySeries(outplus, delta_f=hp.delta_f, copy=False) - outcross[0:len(hc)] = hc[:] + outcross[0 : len(hc)] = hc[:] hc = FrequencySeries(outcross, delta_f=hc.delta_f, copy=False) hp.chirp_length = get_waveform_filter_length_in_time(**input_params) hp.length_in_time = hp.chirp_length hc.chirp_length = hp.chirp_length hc.length_in_time = hp.length_in_time return hp, hc - elif input_params['approximant'] in td_approximants(_scheme.mgr.state): + if input_params["approximant"] in td_approximants(_scheme.mgr.state): # N: number of time samples required - N = (n-1)*2 - delta_f = 1.0 / (N * input_params['delta_t']) + N = (n - 1) * 2 + delta_f = 1.0 / (N * input_params["delta_t"]) wav_gen = td_wav[type(_scheme.mgr.state)] - hp, hc = wav_gen[input_params['approximant']](**input_params) + hp, hc = wav_gen[input_params["approximant"]](**input_params) # taper the time series hp if required - if 'taper' in input_params.keys() and \ - input_params['taper'] is not None: - hp = hp.taper_timeseries(location=input_params['taper'], - tapermethod=input_params.get('taper_method', 'lal'), - taper_window=input_params.get('taper_window'), return_lal=False) - hc = hc.taper_timeseries(location=input_params['taper'], - tapermethod=input_params.get('taper_method', 'lal'), - taper_window=input_params.get('taper_window'), return_lal=False) + if "taper" in input_params.keys() and input_params["taper"] is not None: + hp = hp.taper_timeseries( + location=input_params["taper"], + tapermethod=input_params.get("taper_method", "lal"), + taper_window=input_params.get("taper_window"), + return_lal=False, + ) + hc = hc.taper_timeseries( + location=input_params["taper"], + tapermethod=input_params.get("taper_method", "lal"), + taper_window=input_params.get("taper_window"), + return_lal=False, + ) # total duration of the waveform tmplt_length = len(hp) * hp.delta_t # for IMR templates the zero of time is at max amplitude (merger) # thus the start time is minus the duration of the template from # lower frequency cutoff to merger, i.e. minus the 'chirp time' - tChirp = - float( hp.start_time ) # conversion from LIGOTimeGPS + tChirp = -float(hp.start_time) # conversion from LIGOTimeGPS hp.resize(N) hc.resize(N) k_zero = int(hp.start_time / hp.delta_t) @@ -1359,78 +1548,93 @@ def get_two_pol_waveform_filter(outplus, outcross, template, **kwargs): hc_tilde.length_in_time = tmplt_length hc_tilde.chirp_length = tChirp return hp_tilde, hc_tilde - else: - raise ValueError("Approximant %s not available" % - (input_params['approximant'])) + raise ValueError("Approximant %s not available" % (input_params["approximant"])) + def waveform_norm_exists(approximant): if approximant in _filter_norms: return True - else: - return False + return False + def get_template_amplitude_norm(template=None, **kwargs): - """ Return additional constant template normalization. This only affects - the effective distance calculation. Returns None for all templates with a - physically meaningful amplitude. """ - input_params = props(template,**kwargs) - approximant = kwargs['approximant'] + Return additional constant template normalization. This only affects + the effective distance calculation. Returns None for all templates with a + physically meaningful amplitude. + """ + input_params = props(template, **kwargs) + approximant = kwargs["approximant"] if approximant in _template_amplitude_norms: return _template_amplitude_norms[approximant](**input_params) - else: - return None + return None + def get_waveform_filter_precondition(approximant, length, delta_f): - """Return the data preconditioning factor for this approximant. - """ + """Return the data preconditioning factor for this approximant.""" if approximant in _filter_preconditions: return _filter_preconditions[approximant](length, delta_f) - else: - return None + return None + def get_waveform_filter_norm(approximant, psd, length, delta_f, f_lower): - """ Return the normalization vector for the approximant - """ + """Return the normalization vector for the approximant""" if approximant in _filter_norms: return _filter_norms[approximant](psd, length, delta_f, f_lower) - else: - return None + return None + def get_waveform_end_frequency(template=None, **kwargs): - """Return the stop frequency of a template - """ - input_params = props(template,**kwargs) - approximant = kwargs['approximant'] + """Return the stop frequency of a template""" + input_params = props(template, **kwargs) + approximant = kwargs["approximant"] if approximant in _filter_ends: return _filter_ends[approximant](**input_params) - else: - return None + return None + def get_waveform_filter_length_in_time(approximant, template=None, **kwargs): - """For filter templates, return the length in time of the template. - """ + """For filter templates, return the length in time of the template.""" kwargs = props(template, **kwargs) if approximant in _filter_time_lengths: return _filter_time_lengths[approximant](**kwargs) - else: - return None - -__all__ = ["get_td_waveform", "get_td_det_waveform_from_fd_det", - "get_fd_waveform", "get_fd_waveform_sequence", - "get_fd_det_waveform", "get_fd_det_waveform_sequence", - "get_fd_waveform_from_td", - "print_td_approximants", "print_fd_approximants", - "td_approximants", "fd_approximants", - "get_waveform_filter", "filter_approximants", - "get_waveform_filter_norm", "get_waveform_end_frequency", - "waveform_norm_exists", "get_template_amplitude_norm", - "get_waveform_filter_length_in_time", "get_sgburst_waveform", - "print_sgburst_approximants", "sgburst_approximants", - "td_waveform_to_fd_waveform", "get_two_pol_waveform_filter", - "NoWaveformError", "FailedWaveformError", "get_td_waveform_from_fd", - 'cpu_fd', 'cpu_td', 'fd_sequence', 'fd_det_sequence', 'fd_det', - '_filter_time_lengths'] + return None + + +__all__ = [ + "FailedWaveformError", + "NoWaveformError", + "_filter_time_lengths", + "cpu_fd", + "cpu_td", + "fd_approximants", + "fd_det", + "fd_det_sequence", + "fd_sequence", + "filter_approximants", + "get_fd_det_waveform", + "get_fd_det_waveform_sequence", + "get_fd_waveform", + "get_fd_waveform_from_td", + "get_fd_waveform_sequence", + "get_sgburst_waveform", + "get_td_det_waveform_from_fd_det", + "get_td_waveform", + "get_td_waveform_from_fd", + "get_template_amplitude_norm", + "get_two_pol_waveform_filter", + "get_waveform_end_frequency", + "get_waveform_filter", + "get_waveform_filter_length_in_time", + "get_waveform_filter_norm", + "print_fd_approximants", + "print_sgburst_approximants", + "print_td_approximants", + "sgburst_approximants", + "td_approximants", + "td_waveform_to_fd_waveform", + "waveform_norm_exists", +] diff --git a/pycbc/waveform/waveform_modes.py b/pycbc/waveform/waveform_modes.py index d0873f0be84..bb92322156c 100644 --- a/pycbc/waveform/waveform_modes.py +++ b/pycbc/waveform/waveform_modes.py @@ -13,42 +13,52 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Provides functions and utilities for generating waveforms mode-by-mode. -""" +"""Provides functions and utilities for generating waveforms mode-by-mode.""" from string import Formatter + import lal from pycbc import libutils, pnutils -from pycbc.types import (TimeSeries, FrequencySeries) from pycbc.constants import MSUN_SI, PC_SI -from .waveform import (props, _check_lal_pars, check_args) +from pycbc.types import FrequencySeries, TimeSeries + from . import parameters +from .waveform import _check_lal_pars, check_args, props + +lalsimulation = libutils.import_optional("lalsimulation") -lalsimulation = libutils.import_optional('lalsimulation') def _formatdocstr(docstr): - """Utility for formatting docstrings with parameter information. - """ + """Utility for formatting docstrings with parameter information.""" return docstr.format( - **{_p[1]: getattr(parameters, _p[1]).docstr( - prefix=" ", include_label=False).lstrip(' ') - for _p in Formatter().parse(docstr) if _p[1] is not None - }) + **{ + _p[1]: getattr(parameters, _p[1]) + .docstr(prefix=" ", include_label=False) + .lstrip(" ") + for _p in Formatter().parse(docstr) + if _p[1] is not None + } + ) def _formatdocstrlist(docstr, paramlist, skip_params=None): - """Utility for formatting docstrings with parameter information. - """ + """Utility for formatting docstrings with parameter information.""" if skip_params is None: skip_params = [] - pl = '\n'.join([_p.docstr(prefix=" ", include_label=False) - for _p in paramlist if _p not in skip_params]) + pl = "\n".join( + [ + _p.docstr(prefix=" ", include_label=False) + for _p in paramlist + if _p not in skip_params + ] + ) return docstr.format(params=pl) def sum_modes(hlms, inclination, phi): - """Applies spherical harmonics and sums modes to produce a plus and cross + """ + Applies spherical harmonics and sums modes to produce a plus and cross polarization. Parameters @@ -67,6 +77,7 @@ def sum_modes(hlms, inclination, phi): complex float or array The plus and cross polarization as a complex number. The real part gives the plus, the negative imaginary part the cross. + """ out = None for mode in hlms: @@ -81,35 +92,47 @@ def sum_modes(hlms, inclination, phi): def default_modes(approximant): - """Returns the default modes for the given approximant. - """ + """Returns the default modes for the given approximant.""" # FIXME: this should be replaced to a call to a lalsimulation function, # whenever that's added - if approximant in ['IMRPhenomXPHM', 'IMRPhenomXHM']: + if approximant in ["IMRPhenomXPHM", "IMRPhenomXHM"]: # according to arXiv:2004.06503 ma = [(2, 2), (2, 1), (3, 3), (3, 2), (4, 4)] # add the -m modes ma += [(l, -m) for l, m in ma] - elif approximant in ['IMRPhenomPv3HM', 'IMRPhenomHM']: + elif approximant in ["IMRPhenomPv3HM", "IMRPhenomHM"]: # according to arXiv:1911.06050 ma = [(2, 2), (2, 1), (3, 3), (3, 2), (4, 4), (4, 3)] # add the -m modes ma += [(l, -m) for l, m in ma] - elif approximant.startswith('NRSur7dq4'): + elif approximant.startswith("NRSur7dq4"): # according to arXiv:1905.09300 - ma = [(l, m) for l in [2, 3, 4] for m in range(-l, l+1)] - elif approximant.startswith('NRHybSur3dq8'): + ma = [(l, m) for l in [2, 3, 4] for m in range(-l, l + 1)] + elif approximant.startswith("NRHybSur3dq8"): # according to arXiv:1812.07865 - ma = [(2, 0), (2, 1), (2, 2), (3, 0), (3, 1), (3, 2), - (3, 3), (4, 2), (4, 3), (4, 4), (5, 5)] + ma = [ + (2, 0), + (2, 1), + (2, 2), + (3, 0), + (3, 1), + (3, 2), + (3, 3), + (4, 2), + (4, 3), + (4, 4), + (5, 5), + ] else: - raise ValueError("I don't know what the default modes are for " - "approximant {}, sorry!".format(approximant)) + raise ValueError( + f"I don't know what the default modes are for approximant {approximant}, sorry!" + ) return ma def get_glm(l, m, theta): - r"""The maginitude of the :math:`{}_{-2}Y_{\ell m}`. + r""" + The maginitude of the :math:`{}_{-2}Y_{\ell m}`. The spin-weighted spherical harmonics can be written as :math:`{}_{-2}Y_{\ell m}(\theta, \phi) = g_{\ell m}(\theta)e^{i m \phi}`. @@ -128,12 +151,14 @@ def get_glm(l, m, theta): ------- float : The amplitude of the harmonic at the given polar angle. + """ - return lal.SpinWeightedSphericalHarmonic(theta, 0., -2, l, m).real + return lal.SpinWeightedSphericalHarmonic(theta, 0.0, -2, l, m).real def get_nrsur_modes(**params): - """Generates NRSurrogate waveform mode-by-mode. + """ + Generates NRSurrogate waveform mode-by-mode. All waveform parameters should be provided as keyword arguments. Recognized parameters are listed below. Unrecognized arguments are ignored. @@ -163,28 +188,38 @@ def get_nrsur_modes(**params): ------- dict : Dictionary of ``(l, m)`` -> ``(h_+, -h_x)`` ``TimeSeries``. + """ laldict = _check_lal_pars(params) ret = lalsimulation.SimInspiralPrecessingNRSurModes( - params['delta_t'], - params['mass1']*MSUN_SI, - params['mass2']*MSUN_SI, - params['spin1x'], params['spin1y'], params['spin1z'], - params['spin2x'], params['spin2y'], params['spin2z'], - params['f_lower'], params['f_ref'], - params['distance']*1e6*PC_SI, laldict, - getattr(lalsimulation, params['approximant']) + params["delta_t"], + params["mass1"] * MSUN_SI, + params["mass2"] * MSUN_SI, + params["spin1x"], + params["spin1y"], + params["spin1z"], + params["spin2x"], + params["spin2y"], + params["spin2z"], + params["f_lower"], + params["f_ref"], + params["distance"] * 1e6 * PC_SI, + laldict, + getattr(lalsimulation, params["approximant"]), ) hlms = {} while ret: - hlm = TimeSeries(ret.mode.data.data, delta_t=ret.mode.deltaT, - epoch=ret.mode.epoch) + hlm = TimeSeries( + ret.mode.data.data, delta_t=ret.mode.deltaT, epoch=ret.mode.epoch + ) hlms[ret.l, ret.m] = (hlm.real(), hlm.imag()) ret = ret.next return hlms + def get_nrhybsur_modes(**params): - """Generates NRHybSur3dq8 waveform mode-by-mode. + """ + Generates NRHybSur3dq8 waveform mode-by-mode. All waveform parameters should be provided as keyword arguments. Recognized parameters are listed below. Unrecognized arguments are ignored. @@ -210,21 +245,25 @@ def get_nrhybsur_modes(**params): ------- dict : Dictionary of ``(l, m)`` -> ``(h_+, -h_x)`` ``TimeSeries``. + """ laldict = _check_lal_pars(params) ret = lalsimulation.SimIMRNRHybSur3dq8Modes( - params['delta_t'], - params['mass1']*MSUN_SI, - params['mass2']*MSUN_SI, - params['spin1z'], - params['spin2z'], - params['f_lower'], params['f_ref'], - params['distance']*1e6*PC_SI, laldict + params["delta_t"], + params["mass1"] * MSUN_SI, + params["mass2"] * MSUN_SI, + params["spin1z"], + params["spin2z"], + params["f_lower"], + params["f_ref"], + params["distance"] * 1e6 * PC_SI, + laldict, ) hlms = {} while ret: - hlm = TimeSeries(ret.mode.data.data, delta_t=ret.mode.deltaT, - epoch=ret.mode.epoch) + hlm = TimeSeries( + ret.mode.data.data, delta_t=ret.mode.deltaT, epoch=ret.mode.epoch + ) hlms[ret.l, ret.m] = (hlm.real(), hlm.imag()) ret = ret.next return hlms @@ -233,30 +272,34 @@ def get_nrhybsur_modes(**params): get_nrsur_modes.__doc__ = _formatdocstr(get_nrsur_modes.__doc__) get_nrhybsur_modes.__doc__ = _formatdocstr(get_nrhybsur_modes.__doc__) + def get_lalsimulation_approximant(approximant): import lalsimulation as ls + return { - 'EOBNRv2': ls.EOBNRv2, - 'EOBNRv2HM': ls.EOBNRv2HM, - 'IMRPhenomTPHM': ls.IMRPhenomTPHM, - 'NRSur7dq2': ls.NRSur7dq2, - 'NRSur7dq4': ls.NRSur7dq4, - 'NRHybSur3dq8': ls.NRHybSur3dq8, - 'pSEOBNRv4HM_PA': ls.pSEOBNRv4HM_PA, - 'SEOBNRv4HM_PA': ls.SEOBNRv4HM_PA, - 'SEOBNRv4P': ls.SEOBNRv4P, - 'SEOBNRv4PHM': ls.SEOBNRv4PHM, - 'SpinTaylorT1': ls.SpinTaylorT1, - 'SpinTaylorT4': ls.SpinTaylorT4, - 'SpinTaylorT5': ls.SpinTaylorT5, - 'TaylorT1': ls.TaylorT1, - 'TaylorT2': ls.TaylorT2, - 'TaylorT3': ls.TaylorT3, - 'TaylorT4': ls.TaylorT4, - }[approximant] + "EOBNRv2": ls.EOBNRv2, + "EOBNRv2HM": ls.EOBNRv2HM, + "IMRPhenomTPHM": ls.IMRPhenomTPHM, + "NRSur7dq2": ls.NRSur7dq2, + "NRSur7dq4": ls.NRSur7dq4, + "NRHybSur3dq8": ls.NRHybSur3dq8, + "pSEOBNRv4HM_PA": ls.pSEOBNRv4HM_PA, + "SEOBNRv4HM_PA": ls.SEOBNRv4HM_PA, + "SEOBNRv4P": ls.SEOBNRv4P, + "SEOBNRv4PHM": ls.SEOBNRv4PHM, + "SpinTaylorT1": ls.SpinTaylorT1, + "SpinTaylorT4": ls.SpinTaylorT4, + "SpinTaylorT5": ls.SpinTaylorT5, + "TaylorT1": ls.TaylorT1, + "TaylorT2": ls.TaylorT2, + "TaylorT3": ls.TaylorT3, + "TaylorT4": ls.TaylorT4, + }[approximant] + def get_lalsimulation_modes(**params): - """Generates approximant waveform mode-by-mode. + """ + Generates approximant waveform mode-by-mode. All waveform parameters should be provided as keyword arguments. Recognized parameters are listed below. Unrecognized arguments are ignored. @@ -288,61 +331,71 @@ def get_lalsimulation_modes(**params): ------- dict : Dictionary of ``(l, m)`` -> ``(h_+, -h_x)`` ``TimeSeries``. + """ ell_max = 5 - if 'ell_max' in params: - ell_max = params['ell_max'] + if "ell_max" in params: + ell_max = params["ell_max"] laldict = _check_lal_pars(params) ret = lalsimulation.SimInspiralChooseTDModes( - params['coa_phase'], - params['delta_t'], - params['mass1']*MSUN_SI, - params['mass2']*MSUN_SI, - params['spin1x'], - params['spin1y'], - params['spin1z'], - params['spin2x'], - params['spin2y'], - params['spin2z'], - params['f_lower'], params['f_ref'], - params['distance']*1e6*PC_SI, laldict, + params["coa_phase"], + params["delta_t"], + params["mass1"] * MSUN_SI, + params["mass2"] * MSUN_SI, + params["spin1x"], + params["spin1y"], + params["spin1z"], + params["spin2x"], + params["spin2y"], + params["spin2z"], + params["f_lower"], + params["f_ref"], + params["distance"] * 1e6 * PC_SI, + laldict, ell_max, - get_lalsimulation_approximant(params['approximant']) + get_lalsimulation_approximant(params["approximant"]), ) hlms = {} while ret: - hlm = TimeSeries(ret.mode.data.data, delta_t=ret.mode.deltaT, - epoch=ret.mode.epoch) + hlm = TimeSeries( + ret.mode.data.data, delta_t=ret.mode.deltaT, epoch=ret.mode.epoch + ) hlms[(ret.l, ret.m)] = (hlm.real(), hlm.imag()) ret = ret.next return hlms + def get_imrphenomxh_modes(**params): - """Generates ``IMRPhenomXHM`` waveforms mode-by-mode. """ - approx = params['approximant'] - if not approx.startswith('IMRPhenomX'): + """Generates ``IMRPhenomXHM`` waveforms mode-by-mode.""" + approx = params["approximant"] + if not approx.startswith("IMRPhenomX"): raise ValueError("unsupported approximant") - mode_array = params.pop('mode_array', None) + mode_array = params.pop("mode_array", None) if mode_array is None: mode_array = default_modes(approx) - if 'f_final' not in params: + if "f_final" not in params: # setting to 0 will default to ringdown frequency - params['f_final'] = 0. + params["f_final"] = 0.0 hlms = {} - for (l, m) in mode_array: - params['mode_array'] = [(l, m)] + for l, m in mode_array: + params["mode_array"] = [(l, m)] laldict = _check_lal_pars(params) hlm = lalsimulation.SimIMRPhenomXHMGenerateFDOneMode( - float(pnutils.solar_mass_to_kg(params['mass1'])), - float(pnutils.solar_mass_to_kg(params['mass2'])), - float(params['spin1z']), - float(params['spin2z']), l, m, - pnutils.megaparsecs_to_meters(float(params['distance'])), - params['f_lower'], params['f_final'], params['delta_f'], - params['coa_phase'], params['f_ref'], - laldict) - hlm = FrequencySeries(hlm.data.data, delta_f=hlm.deltaF, - epoch=hlm.epoch) + float(pnutils.solar_mass_to_kg(params["mass1"])), + float(pnutils.solar_mass_to_kg(params["mass2"])), + float(params["spin1z"]), + float(params["spin2z"]), + l, + m, + pnutils.megaparsecs_to_meters(float(params["distance"])), + params["f_lower"], + params["f_final"], + params["delta_f"], + params["coa_phase"], + params["f_ref"], + laldict, + ) + hlm = FrequencySeries(hlm.data.data, delta_f=hlm.deltaF, epoch=hlm.epoch) # Plus, cross strains without Y_lm. # (-1)**(l) factor ALREADY included in FDOneMode hplm = 0.5 * hlm # Plus strain @@ -353,29 +406,32 @@ def get_imrphenomxh_modes(**params): return hlms -_mode_waveform_td = {'EOBNRv2': get_lalsimulation_modes, - 'EOBNRv2HM': get_lalsimulation_modes, - 'IMRPhenomTPHM': get_lalsimulation_modes, - 'NRSur7dq2': get_lalsimulation_modes, - 'NRSur7dq4': get_nrsur_modes, - 'NRHybSur3dq8': get_nrhybsur_modes, - 'pSEOBNRv4HM_PA': get_lalsimulation_modes, - 'SEOBNRv4HM_PA': get_lalsimulation_modes, - 'SEOBNRv4P': get_lalsimulation_modes, - 'SEOBNRv4PHM': get_lalsimulation_modes, - 'SpinTaylorT1': get_lalsimulation_modes, - 'SpinTaylorT4': get_lalsimulation_modes, - 'SpinTaylorT5': get_lalsimulation_modes, - 'TaylorT1': get_lalsimulation_modes, - 'TaylorT2': get_lalsimulation_modes, - 'TaylorT3': get_lalsimulation_modes, - 'TaylorT4': get_lalsimulation_modes, - } -_mode_waveform_fd = {'IMRPhenomXHM': get_imrphenomxh_modes, - } +_mode_waveform_td = { + "EOBNRv2": get_lalsimulation_modes, + "EOBNRv2HM": get_lalsimulation_modes, + "IMRPhenomTPHM": get_lalsimulation_modes, + "NRSur7dq2": get_lalsimulation_modes, + "NRSur7dq4": get_nrsur_modes, + "NRHybSur3dq8": get_nrhybsur_modes, + "pSEOBNRv4HM_PA": get_lalsimulation_modes, + "SEOBNRv4HM_PA": get_lalsimulation_modes, + "SEOBNRv4P": get_lalsimulation_modes, + "SEOBNRv4PHM": get_lalsimulation_modes, + "SpinTaylorT1": get_lalsimulation_modes, + "SpinTaylorT4": get_lalsimulation_modes, + "SpinTaylorT5": get_lalsimulation_modes, + "TaylorT1": get_lalsimulation_modes, + "TaylorT2": get_lalsimulation_modes, + "TaylorT3": get_lalsimulation_modes, + "TaylorT4": get_lalsimulation_modes, +} +_mode_waveform_fd = { + "IMRPhenomXHM": get_imrphenomxh_modes, +} # 'IMRPhenomXPHM':get_imrphenomhm_modes needs to be implemented # LAL function do not split strain mode by mode + def fd_waveform_mode_approximants(): """Frequency domain approximants that will return separate modes.""" return sorted(_mode_waveform_fd.keys()) @@ -387,7 +443,8 @@ def td_waveform_mode_approximants(): def get_fd_waveform_modes(template=None, **kwargs): - r"""Generates frequency domain waveforms, but does not sum over the modes. + r""" + Generates frequency domain waveforms, but does not sum over the modes. The returned values are the frequency-domain equivalents of the real and imaginary parts of the complex :math:`\mathfrak{{h}}_{{\ell m}}(t)` time @@ -411,24 +468,27 @@ def get_fd_waveform_modes(template=None, **kwargs): vlm : dict Dictionary of mode tuples -> fourier transform of the imaginary part of the hlm time series, as a :py:class:`pycbc.types.FrequencySeries`. + """ params = props(template, **kwargs) required = parameters.fd_required check_args(params, required) - apprx = params['approximant'] + apprx = params["approximant"] if apprx not in _mode_waveform_fd: - raise ValueError("I don't support approximant {}, sorry" - .format(apprx)) + raise ValueError(f"I don't support approximant {apprx}, sorry") return _mode_waveform_fd[apprx](**params) get_fd_waveform_modes.__doc__ = _formatdocstrlist( - get_fd_waveform_modes.__doc__, parameters.fd_waveform_params, - skip_params=['inclination', 'coa_phase']) + get_fd_waveform_modes.__doc__, + parameters.fd_waveform_params, + skip_params=["inclination", "coa_phase"], +) def get_td_waveform_modes(template=None, **kwargs): - r"""Generates time domain waveforms, but does not sum over the modes. + r""" + Generates time domain waveforms, but does not sum over the modes. The returned values are the real and imaginary parts of the complex :math:`\mathfrak{{h}}_{{\ell m}}(t)`. These are defined such that the plus @@ -456,17 +516,19 @@ def get_td_waveform_modes(template=None, **kwargs): vlm : dict Dictionary of mode tuples -> imaginary part of the hlm, as a :py:class:`pycbc.types.TimeSeries`. + """ params = props(template, **kwargs) required = parameters.td_required check_args(params, required) - apprx = params['approximant'] + apprx = params["approximant"] if apprx not in _mode_waveform_td: - raise ValueError("I don't support approximant {}, sorry" - .format(apprx)) + raise ValueError(f"I don't support approximant {apprx}, sorry") return _mode_waveform_td[apprx](**params) get_td_waveform_modes.__doc__ = _formatdocstrlist( - get_td_waveform_modes.__doc__, parameters.td_waveform_params, - skip_params=['inclination', 'coa_phase']) + get_td_waveform_modes.__doc__, + parameters.td_waveform_params, + skip_params=["inclination", "coa_phase"], +) diff --git a/pycbc/workflow/__init__.py b/pycbc/workflow/__init__.py index be7ebdf0b53..69e55ec7643 100644 --- a/pycbc/workflow/__init__.py +++ b/pycbc/workflow/__init__.py @@ -26,28 +26,29 @@ performing a coincident CBC matched-filter analysis on gravitational-wave interferometer data """ + import os.path +from pycbc.workflow.coincidence import * from pycbc.workflow.configuration import * from pycbc.workflow.core import * +from pycbc.workflow.datafind import * +from pycbc.workflow.dq import * from pycbc.workflow.grb_utils import * +from pycbc.workflow.injection import * from pycbc.workflow.jobsetup import * -from pycbc.workflow.psd import * from pycbc.workflow.matched_filter import * -from pycbc.workflow.datafind import * -from pycbc.workflow.segment import * -from pycbc.workflow.tmpltbank import * -from pycbc.workflow.psdfiles import * -from pycbc.workflow.splittable import * -from pycbc.workflow.coincidence import * -from pycbc.workflow.injection import * -from pycbc.workflow.plotting import * from pycbc.workflow.minifollowups import * -from pycbc.workflow.dq import * -from pycbc.workflow.versioning import * # Set the pycbc workflow specific pegasus configuration and planning files from pycbc.workflow.pegasus_workflow import PEGASUS_FILE_DIRECTORY +from pycbc.workflow.plotting import * +from pycbc.workflow.psd import * +from pycbc.workflow.psdfiles import * +from pycbc.workflow.segment import * +from pycbc.workflow.splittable import * +from pycbc.workflow.tmpltbank import * +from pycbc.workflow.versioning import * # Set the configuration file base directory -INI_FILE_DIRECTORY = os.path.join(os.path.dirname(__file__), 'ini_files') +INI_FILE_DIRECTORY = os.path.join(os.path.dirname(__file__), "ini_files") diff --git a/pycbc/workflow/coincidence.py b/pycbc/workflow/coincidence.py index 1e5f1072ef3..72f4579d73c 100644 --- a/pycbc/workflow/coincidence.py +++ b/pycbc/workflow/coincidence.py @@ -27,23 +27,25 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/coincidence.html """ -import os import logging +import os import igwn_segments as segments -from pycbc.workflow.core import FileList, make_analysis_dir, Executable, Node, File +from pycbc.workflow.core import Executable, File, FileList, Node, make_analysis_dir + +logger = logging.getLogger("pycbc.workflow.coincidence") -logger = logging.getLogger('pycbc.workflow.coincidence') class PyCBCBank2HDFExecutable(Executable): """Converts xml tmpltbank to hdf format""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, bank_file): node = Node(self) - node.add_input_opt('--bank-file', bank_file) - node.new_output_file_opt(bank_file.segment, '.hdf', '--output-file') + node.add_input_opt("--bank-file", bank_file) + node.new_output_file_opt(bank_file.segment, ".hdf", "--output-file") return node @@ -51,12 +53,14 @@ class PyCBCTrig2HDFExecutable(Executable): """Converts xml triggers to hdf format, grouped by template hash""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, trig_files, bank_file): node = Node(self) - node.add_input_opt('--bank-file', bank_file) - node.add_input_list_opt('--trigger-files', trig_files) - node.new_output_file_opt(trig_files[0].segment, '.hdf', - '--output-file', use_tmp_subdirs=True) + node.add_input_opt("--bank-file", bank_file) + node.add_input_list_opt("--trigger-files", trig_files) + node.new_output_file_opt( + trig_files[0].segment, ".hdf", "--output-file", use_tmp_subdirs=True + ) return node @@ -64,15 +68,16 @@ class PyCBCFitByTemplateExecutable(Executable): """Calculates values that describe the background distribution template by template""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, trig_file, bank_file, veto_file, veto_name): node = Node(self) # Executable objects are initialized with ifo information - node.add_opt('--ifo', self.ifo_string) - node.add_input_opt('--trigger-file', trig_file) - node.add_input_opt('--bank-file', bank_file) - node.add_input_opt('--veto-file', veto_file) - node.add_opt('--veto-segment-name', veto_name) - node.new_output_file_opt(trig_file.segment, '.hdf', '--output') + node.add_opt("--ifo", self.ifo_string) + node.add_input_opt("--trigger-file", trig_file) + node.add_input_opt("--bank-file", bank_file) + node.add_input_opt("--veto-file", veto_file) + node.add_opt("--veto-segment-name", veto_name) + node.new_output_file_opt(trig_file.segment, ".hdf", "--output") return node @@ -80,11 +85,12 @@ class PyCBCFitOverParamExecutable(Executable): """Smooths the background distribution parameters over a continuous parameter""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, raw_fit_file, bank_file): node = Node(self) - node.add_input_opt('--template-fit-file', raw_fit_file) - node.add_input_opt('--bank-file', bank_file) - node.new_output_file_opt(raw_fit_file.segment, '.hdf', '--output') + node.add_input_opt("--template-fit-file", raw_fit_file) + node.add_input_opt("--bank-file", bank_file) + node.new_output_file_opt(raw_fit_file.segment, ".hdf", "--output") return node @@ -92,61 +98,80 @@ class PyCBCFindCoincExecutable(Executable): """Find coinc triggers using a folded interval method""" current_retention_level = Executable.ALL_TRIGGERS - def create_node(self, trig_files, bank_file, stat_files, veto_file, - veto_name, template_str, pivot_ifo, fixed_ifo, tags=None): + + def create_node( + self, + trig_files, + bank_file, + stat_files, + veto_file, + veto_name, + template_str, + pivot_ifo, + fixed_ifo, + tags=None, + ): if tags is None: tags = [] segs = trig_files.get_times_covered_by_files() seg = segments.segment(segs[0][0], segs[-1][1]) node = Node(self) - node.add_input_opt('--template-bank', bank_file) - node.add_input_list_opt('--trigger-files', trig_files) + node.add_input_opt("--template-bank", bank_file) + node.add_input_list_opt("--trigger-files", trig_files) if len(stat_files) > 0: node.add_input_list_opt( - '--statistic-files', - stat_files, - check_existing_options=False + "--statistic-files", stat_files, check_existing_options=False ) if veto_file is not None: - node.add_input_opt('--veto-files', veto_file) - node.add_opt('--segment-name', veto_name) - node.add_opt('--pivot-ifo', pivot_ifo) - node.add_opt('--fixed-ifo', fixed_ifo) - node.add_opt('--template-fraction-range', template_str) - node.new_output_file_opt(seg, '.hdf', '--output-file', tags=tags) + node.add_input_opt("--veto-files", veto_file) + node.add_opt("--segment-name", veto_name) + node.add_opt("--pivot-ifo", pivot_ifo) + node.add_opt("--fixed-ifo", fixed_ifo) + node.add_opt("--template-fraction-range", template_str) + node.new_output_file_opt(seg, ".hdf", "--output-file", tags=tags) return node + class PyCBCFindSnglsExecutable(Executable): """Calculate single-detector ranking statistic for triggers""" current_retention_level = Executable.ALL_TRIGGERS - file_input_options = ['--statistic-files'] - def create_node(self, trig_files, bank_file, stat_files, veto_file, - veto_name, template_str, tags=None): + file_input_options = ["--statistic-files"] + + def create_node( + self, + trig_files, + bank_file, + stat_files, + veto_file, + veto_name, + template_str, + tags=None, + ): if tags is None: tags = [] segs = trig_files.get_times_covered_by_files() seg = segments.segment(segs[0][0], segs[-1][1]) node = Node(self) - node.add_input_opt('--template-bank', bank_file) - node.add_input_list_opt('--trigger-files', trig_files) + node.add_input_opt("--template-bank", bank_file) + node.add_input_list_opt("--trigger-files", trig_files) if len(stat_files) > 0: node.add_input_list_opt( - '--statistic-files', - stat_files, - check_existing_options=False + "--statistic-files", stat_files, check_existing_options=False ) if veto_file is not None: - node.add_input_opt('--veto-files', veto_file) - node.add_opt('--segment-name', veto_name) - node.add_opt('--template-fraction-range', template_str) - node.new_output_file_opt(seg, '.hdf', '--output-file', tags=tags) + node.add_input_opt("--veto-files", veto_file) + node.add_opt("--segment-name", veto_name) + node.add_opt("--template-fraction-range", template_str) + node.new_output_file_opt(seg, ".hdf", "--output-file", tags=tags) return node + class PyCBCStatMapExecutable(Executable): """Calculate FAP, IFAR, etc for coincs""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, coinc_files, ifos, tags=None): if tags is None: tags = [] @@ -154,15 +179,17 @@ def create_node(self, coinc_files, ifos, tags=None): seg = segments.segment(segs[0][0], segs[-1][1]) node = Node(self) - node.add_input_list_opt('--coinc-files', coinc_files) - node.add_opt('--ifos', ifos) - node.new_output_file_opt(seg, '.hdf', '--output-file', tags=tags) + node.add_input_list_opt("--coinc-files", coinc_files) + node.add_opt("--ifos", ifos) + node.new_output_file_opt(seg, ".hdf", "--output-file", tags=tags) return node + class PyCBCSnglsStatMapExecutable(Executable): """Calculate FAP, IFAR, etc for singles""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, sngls_files, ifo, tags=None): if tags is None: tags = [] @@ -170,9 +197,9 @@ def create_node(self, sngls_files, ifo, tags=None): seg = segments.segment(segs[0][0], segs[-1][1]) node = Node(self) - node.add_input_list_opt('--sngls-files', sngls_files) - node.add_opt('--ifos', ifo) - node.new_output_file_opt(seg, '.hdf', '--output-file', tags=tags) + node.add_input_list_opt("--sngls-files", sngls_files) + node.add_opt("--ifos", ifo) + node.new_output_file_opt(seg, ".hdf", "--output-file", tags=tags) return node @@ -180,42 +207,43 @@ class PyCBCStatMapInjExecutable(Executable): """Calculate FAP, IFAR, etc for coincs for injections""" current_retention_level = Executable.MERGED_TRIGGERS - def create_node(self, coinc_files, full_data, - ifos, tags=None): + + def create_node(self, coinc_files, full_data, ifos, tags=None): if tags is None: tags = [] segs = coinc_files.get_times_covered_by_files() seg = segments.segment(segs[0][0], segs[-1][1]) node = Node(self) - node.add_input_list_opt('--zero-lag-coincs', coinc_files) + node.add_input_list_opt("--zero-lag-coincs", coinc_files) if isinstance(full_data, list): - node.add_input_list_opt('--full-data-background', full_data) + node.add_input_list_opt("--full-data-background", full_data) else: - node.add_input_opt('--full-data-background', full_data) + node.add_input_opt("--full-data-background", full_data) - node.add_opt('--ifos', ifos) - node.new_output_file_opt(seg, '.hdf', '--output-file', tags=tags) + node.add_opt("--ifos", ifos) + node.new_output_file_opt(seg, ".hdf", "--output-file", tags=tags) return node + class PyCBCSnglsStatMapInjExecutable(Executable): """Calculate FAP, IFAR, etc for singles for injections""" current_retention_level = Executable.MERGED_TRIGGERS - def create_node(self, sngls_files, background_file, - ifos, tags=None): + + def create_node(self, sngls_files, background_file, ifos, tags=None): if tags is None: tags = [] segs = sngls_files.get_times_covered_by_files() seg = segments.segment(segs[0][0], segs[-1][1]) node = Node(self) - node.add_input_list_opt('--sngls-files', sngls_files) - node.add_input_opt('--full-data-background', background_file) + node.add_input_list_opt("--sngls-files", sngls_files) + node.add_input_opt("--full-data-background", background_file) - node.add_opt('--ifos', ifos) - node.new_output_file_opt(seg, '.hdf', '--output-file', tags=tags) + node.add_opt("--ifos", ifos) + node.new_output_file_opt(seg, ".hdf", "--output-file", tags=tags) return node @@ -223,17 +251,21 @@ class PyCBCHDFInjFindExecutable(Executable): """Find injections in the hdf files output""" current_retention_level = Executable.MERGED_TRIGGERS - def create_node(self, inj_coinc_file, inj_xml_file, veto_file, veto_name, tags=None): + + def create_node( + self, inj_coinc_file, inj_xml_file, veto_file, veto_name, tags=None + ): if tags is None: tags = [] node = Node(self) - node.add_input_list_opt('--trigger-file', inj_coinc_file) - node.add_input_list_opt('--injection-file', inj_xml_file) + node.add_input_list_opt("--trigger-file", inj_coinc_file) + node.add_input_list_opt("--injection-file", inj_xml_file) if veto_name is not None: - node.add_input_opt('--veto-file', veto_file) - node.add_opt('--segment-name', veto_name) - node.new_output_file_opt(inj_xml_file[0].segment, '.hdf', - '--output-file', tags=tags) + node.add_input_opt("--veto-file", veto_file) + node.add_opt("--segment-name", veto_name) + node.new_output_file_opt( + inj_xml_file[0].segment, ".hdf", "--output-file", tags=tags + ) return node @@ -241,23 +273,29 @@ class PyCBCDistributeBackgroundBins(Executable): """Distribute coinc files among different background bins""" current_retention_level = Executable.ALL_TRIGGERS + def create_node(self, coinc_files, bank_file, background_bins, tags=None): if tags is None: tags = [] node = Node(self) - node.add_input_list_opt('--coinc-files', coinc_files) - node.add_input_opt('--bank-file', bank_file) - node.add_opt('--background-bins', ' '.join(background_bins)) - - names = [b.split(':')[0] for b in background_bins] - - output_files = [File(coinc_files[0].ifo_list, - self.name, - coinc_files[0].segment, - directory=self.out_dir, - tags = tags + ['mbin-%s' % i], - extension='.hdf') for i in range(len(background_bins))] - node.add_output_list_opt('--output-files', output_files) + node.add_input_list_opt("--coinc-files", coinc_files) + node.add_input_opt("--bank-file", bank_file) + node.add_opt("--background-bins", " ".join(background_bins)) + + names = [b.split(":")[0] for b in background_bins] + + output_files = [ + File( + coinc_files[0].ifo_list, + self.name, + coinc_files[0].segment, + directory=self.out_dir, + tags=tags + ["mbin-%s" % i], + extension=".hdf", + ) + for i in range(len(background_bins)) + ] + node.add_output_list_opt("--output-files", output_files) node.names = names return node @@ -266,13 +304,15 @@ class PyCBCCombineStatmap(Executable): """Combine coincs over different bins and apply trials factor""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, statmap_files, tags=None): if tags is None: tags = [] node = Node(self) - node.add_input_list_opt('--statmap-files', statmap_files) - node.new_output_file_opt(statmap_files[0].segment, '.hdf', - '--output-file', tags=tags) + node.add_input_list_opt("--statmap-files", statmap_files) + node.new_output_file_opt( + statmap_files[0].segment, ".hdf", "--output-file", tags=tags + ) return node @@ -280,32 +320,36 @@ class PyCBCAddStatmap(PyCBCCombineStatmap): """Combine statmap files and add FARs over different coinc types""" current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, statmap_files, background_files, tags=None): if tags is None: tags = [] - node = super(PyCBCAddStatmap, self).create_node(statmap_files, - tags=tags) + node = super().create_node(statmap_files, tags=tags) # Enforce upper case ctags = [t.upper() for t in (tags + self.tags)] - if 'INJECTIONS' in ctags: - node.add_input_list_opt('--background-files', background_files) + if "INJECTIONS" in ctags: + node.add_input_list_opt("--background-files", background_files) return node class PyCBCExcludeZerolag(Executable): - """ Remove times of zerolag coincidences of all types from exclusive - background """ + """ + Remove times of zerolag coincidences of all types from exclusive + background + """ + current_retention_level = Executable.MERGED_TRIGGERS + def create_node(self, statmap_file, other_statmap_files, tags=None): if tags is None: tags = [] node = Node(self) - node.add_input_opt('--statmap-file', statmap_file) - node.add_input_list_opt('--other-statmap-files', - other_statmap_files) - node.new_output_file_opt(statmap_file.segment, '.hdf', - '--output-file', tags=None) + node.add_input_opt("--statmap-file", statmap_file) + node.add_input_list_opt("--other-statmap-files", other_statmap_files) + node.new_output_file_opt( + statmap_file.segment, ".hdf", "--output-file", tags=None + ) return node @@ -318,72 +362,76 @@ class CensorForeground(Executable): current_retention_level = Executable.MERGED_TRIGGERS -def make_foreground_censored_veto(workflow, bg_file, veto_file, veto_name, - censored_name, out_dir, tags=None): +def make_foreground_censored_veto( + workflow, bg_file, veto_file, veto_name, censored_name, out_dir, tags=None +): tags = [] if tags is None else tags - node = CensorForeground(workflow.cp, 'foreground_censor', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--foreground-triggers', bg_file) - node.add_input_opt('--veto-file', veto_file) - node.add_opt('--segment-name', veto_name) - node.add_opt('--output-segment-name', censored_name) - node.new_output_file_opt(workflow.analysis_time, '.xml', '--output-file') + node = CensorForeground( + workflow.cp, "foreground_censor", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--foreground-triggers", bg_file) + node.add_input_opt("--veto-file", veto_file) + node.add_opt("--segment-name", veto_name) + node.add_opt("--output-segment-name", censored_name) + node.new_output_file_opt(workflow.analysis_time, ".xml", "--output-file") workflow += node return node.output_files[0] -def merge_single_detector_hdf_files(workflow, bank_file, trigger_files, out_dir, tags=None): +def merge_single_detector_hdf_files( + workflow, bank_file, trigger_files, out_dir, tags=None +): if tags is None: tags = [] make_analysis_dir(out_dir) out = FileList() for ifo in workflow.ifos: - node = MergeExecutable(workflow.cp, 'hdf_trigger_merge', - ifos=ifo, out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--bank-file', bank_file) - node.add_input_list_opt('--trigger-files', trigger_files.find_output_with_ifo(ifo)) - node.new_output_file_opt(workflow.analysis_time, '.hdf', '--output-file') + node = MergeExecutable( + workflow.cp, "hdf_trigger_merge", ifos=ifo, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--bank-file", bank_file) + node.add_input_list_opt( + "--trigger-files", trigger_files.find_output_with_ifo(ifo) + ) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node out += node.output_files return out -def setup_trigger_fitting(workflow, insps, hdfbank, veto_file, veto_name, - output_dir=None, tags=None): - if not workflow.cp.has_option('workflow-coincidence', 'do-trigger-fitting'): +def setup_trigger_fitting( + workflow, insps, hdfbank, veto_file, veto_name, output_dir=None, tags=None +): + if not workflow.cp.has_option("workflow-coincidence", "do-trigger-fitting"): return FileList() - else: - smoothed_fit_files = FileList() - for i in workflow.ifos: - ifo_insp = [insp for insp in insps if (insp.ifo == i)] - assert len(ifo_insp)==1 - ifo_insp = ifo_insp[0] - raw_exe = PyCBCFitByTemplateExecutable(workflow.cp, - 'fit_by_template', ifos=i, - out_dir=output_dir, - tags=tags) - raw_node = raw_exe.create_node(ifo_insp, hdfbank, - veto_file, veto_name) - workflow += raw_node - smooth_exe = PyCBCFitOverParamExecutable(workflow.cp, - 'fit_over_param', ifos=i, - out_dir=output_dir, - tags=tags) - smooth_node = smooth_exe.create_node(raw_node.output_file, - hdfbank) - workflow += smooth_node - smoothed_fit_files += smooth_node.output_files - return smoothed_fit_files - - -def find_injections_in_hdf_coinc(workflow, inj_coinc_file, inj_xml_file, - veto_file, veto_name, out_dir, tags=None): + smoothed_fit_files = FileList() + for i in workflow.ifos: + ifo_insp = [insp for insp in insps if (insp.ifo == i)] + assert len(ifo_insp) == 1 + ifo_insp = ifo_insp[0] + raw_exe = PyCBCFitByTemplateExecutable( + workflow.cp, "fit_by_template", ifos=i, out_dir=output_dir, tags=tags + ) + raw_node = raw_exe.create_node(ifo_insp, hdfbank, veto_file, veto_name) + workflow += raw_node + smooth_exe = PyCBCFitOverParamExecutable( + workflow.cp, "fit_over_param", ifos=i, out_dir=output_dir, tags=tags + ) + smooth_node = smooth_exe.create_node(raw_node.output_file, hdfbank) + workflow += smooth_node + smoothed_fit_files += smooth_node.output_files + return smoothed_fit_files + + +def find_injections_in_hdf_coinc( + workflow, inj_coinc_file, inj_xml_file, veto_file, veto_name, out_dir, tags=None +): if tags is None: tags = [] make_analysis_dir(out_dir) - exe = PyCBCHDFInjFindExecutable(workflow.cp, 'hdfinjfind', - ifos=workflow.ifos, - out_dir=out_dir, tags=tags) + exe = PyCBCHDFInjFindExecutable( + workflow.cp, "hdfinjfind", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) node = exe.create_node(inj_coinc_file, inj_xml_file, veto_file, veto_name) workflow += node return node.output_files[0] @@ -393,15 +441,15 @@ def convert_bank_to_hdf(workflow, xmlbank, out_dir, tags=None): """Return the template bank in hdf format""" if tags is None: tags = [] - #FIXME, make me not needed + # FIXME, make me not needed if len(xmlbank) > 1: - raise ValueError('Can only convert a single template bank') + raise ValueError("Can only convert a single template bank") - logger.info('convert template bank to HDF') + logger.info("convert template bank to HDF") make_analysis_dir(out_dir) - bank2hdf_exe = PyCBCBank2HDFExecutable(workflow.cp, 'bank2hdf', - ifos=workflow.ifos, - out_dir=out_dir, tags=tags) + bank2hdf_exe = PyCBCBank2HDFExecutable( + workflow.cp, "bank2hdf", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) bank2hdf_node = bank2hdf_exe.create_node(xmlbank[0]) workflow.add_node(bank2hdf_node) return bank2hdf_node.output_files @@ -411,17 +459,18 @@ def convert_trig_to_hdf(workflow, hdfbank, xml_trigger_files, out_dir, tags=None """Return the list of hdf5 trigger files outputs""" if tags is None: tags = [] - #FIXME, make me not needed - logger.info('convert single inspiral trigger files to hdf5') + # FIXME, make me not needed + logger.info("convert single inspiral trigger files to hdf5") make_analysis_dir(out_dir) trig_files = FileList() - for ifo, insp_group in zip(*xml_trigger_files.categorize_by_attr('ifo')): - trig2hdf_exe = PyCBCTrig2HDFExecutable(workflow.cp, 'trig2hdf', - ifos=ifo, out_dir=out_dir, tags=tags) - _, insp_bundles = insp_group.categorize_by_attr('segment') + for ifo, insp_group in zip(*xml_trigger_files.categorize_by_attr("ifo")): + trig2hdf_exe = PyCBCTrig2HDFExecutable( + workflow.cp, "trig2hdf", ifos=ifo, out_dir=out_dir, tags=tags + ) + _, insp_bundles = insp_group.categorize_by_attr("segment") for insps in insp_bundles: - trig2hdf_node = trig2hdf_exe.create_node(insps, hdfbank[0]) + trig2hdf_node = trig2hdf_exe.create_node(insps, hdfbank[0]) workflow.add_node(trig2hdf_node) trig_files += trig2hdf_node.output_files return trig_files @@ -430,11 +479,11 @@ def convert_trig_to_hdf(workflow, hdfbank, xml_trigger_files, out_dir, tags=None def setup_statmap(workflow, ifos, coinc_files, out_dir, tags=None): tags = [] if tags is None else tags - statmap_exe = PyCBCStatMapExecutable(workflow.cp, 'statmap', - ifos=ifos, - tags=tags, out_dir=out_dir) + statmap_exe = PyCBCStatMapExecutable( + workflow.cp, "statmap", ifos=ifos, tags=tags, out_dir=out_dir + ) - ifolist = ' '.join(ifos) + ifolist = " ".join(ifos) stat_node = statmap_exe.create_node(coinc_files, ifolist) workflow.add_node(stat_node) return stat_node.output_file @@ -443,206 +492,260 @@ def setup_statmap(workflow, ifos, coinc_files, out_dir, tags=None): def setup_sngls_statmap(workflow, ifo, sngls_files, out_dir, tags=None): tags = [] if tags is None else tags - statmap_exe = PyCBCSnglsStatMapExecutable(workflow.cp, 'sngls_statmap', - ifos=ifo, - tags=tags, out_dir=out_dir) + statmap_exe = PyCBCSnglsStatMapExecutable( + workflow.cp, "sngls_statmap", ifos=ifo, tags=tags, out_dir=out_dir + ) stat_node = statmap_exe.create_node(sngls_files, ifo) workflow.add_node(stat_node) return stat_node.output_file -def setup_statmap_inj(workflow, ifos, coinc_files, background_file, - out_dir, tags=None): +def setup_statmap_inj(workflow, ifos, coinc_files, background_file, out_dir, tags=None): tags = [] if tags is None else tags - statmap_exe = PyCBCStatMapInjExecutable(workflow.cp, - 'statmap_inj', - ifos=ifos, - tags=tags, out_dir=out_dir) + statmap_exe = PyCBCStatMapInjExecutable( + workflow.cp, "statmap_inj", ifos=ifos, tags=tags, out_dir=out_dir + ) - ifolist = ' '.join(ifos) - stat_node = statmap_exe.create_node(FileList(coinc_files), - background_file, - ifolist) + ifolist = " ".join(ifos) + stat_node = statmap_exe.create_node(FileList(coinc_files), background_file, ifolist) workflow.add_node(stat_node) return stat_node.output_files[0] -def setup_sngls_statmap_inj(workflow, ifo, sngls_inj_files, background_file, - out_dir, tags=None): +def setup_sngls_statmap_inj( + workflow, ifo, sngls_inj_files, background_file, out_dir, tags=None +): tags = [] if tags is None else tags - statmap_exe = PyCBCSnglsStatMapInjExecutable(workflow.cp, - 'sngls_statmap_inj', - ifos=ifo, - tags=tags, - out_dir=out_dir) + statmap_exe = PyCBCSnglsStatMapInjExecutable( + workflow.cp, "sngls_statmap_inj", ifos=ifo, tags=tags, out_dir=out_dir + ) - stat_node = statmap_exe.create_node(sngls_inj_files, - background_file, - ifo) + stat_node = statmap_exe.create_node(sngls_inj_files, background_file, ifo) workflow.add_node(stat_node) return stat_node.output_files[0] -def setup_interval_coinc_inj(workflow, hdfbank, - inj_trig_files, stat_files, - background_file, veto_file, veto_name, - out_dir, pivot_ifo, fixed_ifo, tags=None): +def setup_interval_coinc_inj( + workflow, + hdfbank, + inj_trig_files, + stat_files, + background_file, + veto_file, + veto_name, + out_dir, + pivot_ifo, + fixed_ifo, + tags=None, +): """ This function sets up exact match coincidence for injections """ if tags is None: tags = [] make_analysis_dir(out_dir) - logger.info('Setting up coincidence for injections') + logger.info("Setting up coincidence for injections") # Wall time knob and memory knob - factor = int(workflow.cp.get_opt_tags('workflow-coincidence', - 'parallelization-factor', tags)) + factor = int( + workflow.cp.get_opt_tags("workflow-coincidence", "parallelization-factor", tags) + ) ifiles = {} - for ifo, ifi in zip(*inj_trig_files.categorize_by_attr('ifo')): + for ifo, ifi in zip(*inj_trig_files.categorize_by_attr("ifo")): ifiles[ifo] = ifi[0] injinj_files = FileList() for ifo in ifiles: # ifiles is keyed on ifo injinj_files.append(ifiles[ifo]) - findcoinc_exe = PyCBCFindCoincExecutable(workflow.cp, - 'coinc', - ifos=ifiles.keys(), - tags=tags + ['injinj'], - out_dir=out_dir) + findcoinc_exe = PyCBCFindCoincExecutable( + workflow.cp, + "coinc", + ifos=ifiles.keys(), + tags=tags + ["injinj"], + out_dir=out_dir, + ) bg_files = [] for i in range(factor): - group_str = '%s/%s' % (i, factor) - coinc_node = findcoinc_exe.create_node(injinj_files, hdfbank, - stat_files, - veto_file, veto_name, - group_str, - pivot_ifo, - fixed_ifo, - tags=['JOB'+str(i)]) + group_str = "%s/%s" % (i, factor) + coinc_node = findcoinc_exe.create_node( + injinj_files, + hdfbank, + stat_files, + veto_file, + veto_name, + group_str, + pivot_ifo, + fixed_ifo, + tags=["JOB" + str(i)], + ) bg_files += coinc_node.output_files workflow.add_node(coinc_node) - logger.info('...leaving coincidence for injections') - - return setup_statmap_inj(workflow, ifiles.keys(), bg_files, - background_file, out_dir, - tags=tags + [veto_name]) - - -def setup_interval_coinc(workflow, hdfbank, trig_files, stat_files, - veto_file, veto_name, out_dir, pivot_ifo, - fixed_ifo, tags=None): + logger.info("...leaving coincidence for injections") + + return setup_statmap_inj( + workflow, + ifiles.keys(), + bg_files, + background_file, + out_dir, + tags=tags + [veto_name], + ) + + +def setup_interval_coinc( + workflow, + hdfbank, + trig_files, + stat_files, + veto_file, + veto_name, + out_dir, + pivot_ifo, + fixed_ifo, + tags=None, +): """ This function sets up exact match coincidence """ if tags is None: tags = [] make_analysis_dir(out_dir) - logger.info('Setting up coincidence') + logger.info("Setting up coincidence") - ifos, _ = trig_files.categorize_by_attr('ifo') - findcoinc_exe = PyCBCFindCoincExecutable(workflow.cp, 'coinc', - ifos=ifos, - tags=tags, out_dir=out_dir) + ifos, _ = trig_files.categorize_by_attr("ifo") + findcoinc_exe = PyCBCFindCoincExecutable( + workflow.cp, "coinc", ifos=ifos, tags=tags, out_dir=out_dir + ) # Wall time knob and memory knob - factor = int(workflow.cp.get_opt_tags('workflow-coincidence', - 'parallelization-factor', - [findcoinc_exe.ifo_string] + tags)) + factor = int( + workflow.cp.get_opt_tags( + "workflow-coincidence", + "parallelization-factor", + [findcoinc_exe.ifo_string] + tags, + ) + ) statmap_files = [] bg_files = FileList() for i in range(factor): - group_str = '%s/%s' % (i, factor) - coinc_node = findcoinc_exe.create_node(trig_files, hdfbank, - stat_files, - veto_file, veto_name, - group_str, - pivot_ifo, - fixed_ifo, - tags=['JOB'+str(i)]) + group_str = "%s/%s" % (i, factor) + coinc_node = findcoinc_exe.create_node( + trig_files, + hdfbank, + stat_files, + veto_file, + veto_name, + group_str, + pivot_ifo, + fixed_ifo, + tags=["JOB" + str(i)], + ) bg_files += coinc_node.output_files workflow.add_node(coinc_node) - statmap_files = setup_statmap(workflow, ifos, bg_files, - out_dir, tags=tags) + statmap_files = setup_statmap(workflow, ifos, bg_files, out_dir, tags=tags) - logger.info('...leaving coincidence ') + logger.info("...leaving coincidence ") return statmap_files -def setup_sngls(workflow, hdfbank, trig_files, stat_files, - veto_file, veto_name, out_dir, tags=None): +def setup_sngls( + workflow, hdfbank, trig_files, stat_files, veto_file, veto_name, out_dir, tags=None +): """ This function sets up getting statistic values for single-detector triggers """ - ifos, _ = trig_files.categorize_by_attr('ifo') - findsngls_exe = PyCBCFindSnglsExecutable(workflow.cp, 'sngls', ifos=ifos, - tags=tags, out_dir=out_dir) + ifos, _ = trig_files.categorize_by_attr("ifo") + findsngls_exe = PyCBCFindSnglsExecutable( + workflow.cp, "sngls", ifos=ifos, tags=tags, out_dir=out_dir + ) # Wall time knob and memory knob - factor = int(workflow.cp.get_opt_tags('workflow-coincidence', - 'parallelization-factor', - [findsngls_exe.ifo_string] + tags)) + factor = int( + workflow.cp.get_opt_tags( + "workflow-coincidence", + "parallelization-factor", + [findsngls_exe.ifo_string] + tags, + ) + ) statmap_files = [] bg_files = FileList() for i in range(factor): - group_str = '%s/%s' % (i, factor) - sngls_node = findsngls_exe.create_node(trig_files, hdfbank, - stat_files, - veto_file, veto_name, - group_str, - tags=['JOB'+str(i)]) + group_str = "%s/%s" % (i, factor) + sngls_node = findsngls_exe.create_node( + trig_files, + hdfbank, + stat_files, + veto_file, + veto_name, + group_str, + tags=["JOB" + str(i)], + ) bg_files += sngls_node.output_files workflow.add_node(sngls_node) - statmap_files = setup_sngls_statmap(workflow, ifos[0], bg_files, - out_dir, tags=tags) + statmap_files = setup_sngls_statmap(workflow, ifos[0], bg_files, out_dir, tags=tags) - logger.info('...leaving coincidence ') + logger.info("...leaving coincidence ") return statmap_files -def setup_sngls_inj(workflow, hdfbank, inj_trig_files, - stat_files, background_file, veto_file, veto_name, - out_dir, tags=None): +def setup_sngls_inj( + workflow, + hdfbank, + inj_trig_files, + stat_files, + background_file, + veto_file, + veto_name, + out_dir, + tags=None, +): """ This function sets up getting statistic values for single-detector triggers from injections """ - ifos, _ = inj_trig_files.categorize_by_attr('ifo') - findsnglsinj_exe = PyCBCFindSnglsExecutable(workflow.cp, 'sngls', ifos=ifos, - tags=tags, out_dir=out_dir) + ifos, _ = inj_trig_files.categorize_by_attr("ifo") + findsnglsinj_exe = PyCBCFindSnglsExecutable( + workflow.cp, "sngls", ifos=ifos, tags=tags, out_dir=out_dir + ) # Wall time knob and memory knob exe_str_tags = [findsnglsinj_exe.ifo_string] + tags - factor = int(workflow.cp.get_opt_tags('workflow-coincidence', - 'parallelization-factor', - exe_str_tags)) + factor = int( + workflow.cp.get_opt_tags( + "workflow-coincidence", "parallelization-factor", exe_str_tags + ) + ) statmap_files = [] bg_files = FileList() for i in range(factor): - group_str = '%s/%s' % (i, factor) - sngls_node = findsnglsinj_exe.create_node(inj_trig_files, hdfbank, - stat_files, - veto_file, veto_name, - group_str, - tags=['JOB'+str(i)]) + group_str = "%s/%s" % (i, factor) + sngls_node = findsnglsinj_exe.create_node( + inj_trig_files, + hdfbank, + stat_files, + veto_file, + veto_name, + group_str, + tags=["JOB" + str(i)], + ) bg_files += sngls_node.output_files workflow.add_node(sngls_node) - statmap_files = setup_sngls_statmap_inj(workflow, ifos[0], bg_files, - background_file, - out_dir, tags=tags) + statmap_files = setup_sngls_statmap_inj( + workflow, ifos[0], bg_files, background_file, out_dir, tags=tags + ) - logger.info('...leaving coincidence ') + logger.info("...leaving coincidence ") return statmap_files @@ -651,7 +754,7 @@ def select_files_by_ifo_combination(ifocomb, insps): This function selects single-detector files ('insps') for a given ifo combination """ inspcomb = FileList() - for ifo, ifile in zip(*insps.categorize_by_attr('ifo')): + for ifo, ifile in zip(*insps.categorize_by_attr("ifo")): if ifo in ifocomb: inspcomb += ifile @@ -670,48 +773,50 @@ def get_ordered_ifo_list(ifocomb, ifo_ids): # combination_prec stores precedence info for the detectors in the combination combination_prec = {ifo: ifo_ids[ifo] for ifo in ifocomb} - ordered_ifo_list = sorted(combination_prec, key = combination_prec.get) + ordered_ifo_list = sorted(combination_prec, key=combination_prec.get) pivot_ifo = ordered_ifo_list[0] fixed_ifo = ordered_ifo_list[1] - return pivot_ifo, fixed_ifo, ''.join(ordered_ifo_list) + return pivot_ifo, fixed_ifo, "".join(ordered_ifo_list) -def setup_combine_statmap(workflow, final_bg_file_list, bg_file_list, - out_dir, tags=None): +def setup_combine_statmap( + workflow, final_bg_file_list, bg_file_list, out_dir, tags=None +): """ Combine the statmap files into one background file """ if tags is None: tags = [] make_analysis_dir(out_dir) - logger.info('Setting up combine statmap') + logger.info("Setting up combine statmap") - cstat_exe_name = os.path.basename(workflow.cp.get("executables", - "combine_statmap")) - if cstat_exe_name == 'pycbc_combine_statmap': + cstat_exe_name = os.path.basename(workflow.cp.get("executables", "combine_statmap")) + if cstat_exe_name == "pycbc_combine_statmap": cstat_class = PyCBCCombineStatmap - elif cstat_exe_name == 'pycbc_add_statmap': + elif cstat_exe_name == "pycbc_add_statmap": cstat_class = PyCBCAddStatmap else: - raise NotImplementedError('executable should be ' - 'pycbc_combine_statmap or pycbc_add_statmap') + raise NotImplementedError( + "executable should be pycbc_combine_statmap or pycbc_add_statmap" + ) - cstat_exe = cstat_class(workflow.cp, 'combine_statmap', ifos=workflow.ifos, - tags=tags, out_dir=out_dir) + cstat_exe = cstat_class( + workflow.cp, "combine_statmap", ifos=workflow.ifos, tags=tags, out_dir=out_dir + ) - if cstat_exe_name == 'pycbc_combine_statmap': + if cstat_exe_name == "pycbc_combine_statmap": combine_statmap_node = cstat_exe.create_node(final_bg_file_list) - elif cstat_exe_name == 'pycbc_add_statmap': - combine_statmap_node = cstat_exe.create_node(final_bg_file_list, - bg_file_list) + elif cstat_exe_name == "pycbc_add_statmap": + combine_statmap_node = cstat_exe.create_node(final_bg_file_list, bg_file_list) workflow.add_node(combine_statmap_node) return combine_statmap_node.output_file -def setup_exclude_zerolag(workflow, statmap_file, other_statmap_files, - out_dir, ifos, tags=None): +def setup_exclude_zerolag( + workflow, statmap_file, other_statmap_files, out_dir, ifos, tags=None +): """ Exclude single triggers close to zerolag triggers from forming any background events @@ -719,22 +824,27 @@ def setup_exclude_zerolag(workflow, statmap_file, other_statmap_files, if tags is None: tags = [] make_analysis_dir(out_dir) - logger.info('Setting up exclude zerolag') - - exc_zerolag_exe = PyCBCExcludeZerolag(workflow.cp, 'exclude_zerolag', - ifos=ifos, tags=tags, - out_dir=out_dir) - exc_zerolag_node = exc_zerolag_exe.create_node(statmap_file, - other_statmap_files, - tags=None) + logger.info("Setting up exclude zerolag") + + exc_zerolag_exe = PyCBCExcludeZerolag( + workflow.cp, "exclude_zerolag", ifos=ifos, tags=tags, out_dir=out_dir + ) + exc_zerolag_node = exc_zerolag_exe.create_node( + statmap_file, other_statmap_files, tags=None + ) workflow.add_node(exc_zerolag_node) return exc_zerolag_node.output_file -def rerank_coinc_followup(workflow, statmap_file, bank_file, out_dir, - tags=None, - injection_file=None, - ranking_file=None): +def rerank_coinc_followup( + workflow, + statmap_file, + bank_file, + out_dir, + tags=None, + injection_file=None, + ranking_file=None, +): if tags is None: tags = [] @@ -743,63 +853,73 @@ def rerank_coinc_followup(workflow, statmap_file, bank_file, out_dir, if not workflow.cp.has_section("workflow-rerank"): logger.info("No reranking done in this workflow") return statmap_file - else: - logger.info("Setting up reranking of candidates") + logger.info("Setting up reranking of candidates") # Generate reduced data files (maybe this could also be used elsewhere?) stores = FileList([]) for ifo in workflow.ifos: - make_analysis_dir('strain_files') - node = Executable(workflow.cp, 'strain_data_reduce', ifos=[ifo], - out_dir='strain_files', tags=tags).create_node() - node.add_opt('--gps-start-time', workflow.analysis_time[0]) - node.add_opt('--gps-end-time', workflow.analysis_time[1]) + make_analysis_dir("strain_files") + node = Executable( + workflow.cp, + "strain_data_reduce", + ifos=[ifo], + out_dir="strain_files", + tags=tags, + ).create_node() + node.add_opt("--gps-start-time", workflow.analysis_time[0]) + node.add_opt("--gps-end-time", workflow.analysis_time[1]) if injection_file: - node.add_input_opt('--injection-file', injection_file) + node.add_input_opt("--injection-file", injection_file) - fil = node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file') + fil = node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") stores.append(fil) workflow += node # Generate trigger input file - node = Executable(workflow.cp, 'rerank_trigger_input', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--statmap-file', statmap_file) - node.add_input_opt('--bank-file', bank_file) - trigfil = node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file') + node = Executable( + workflow.cp, + "rerank_trigger_input", + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags, + ).create_node() + node.add_input_opt("--statmap-file", statmap_file) + node.add_input_opt("--bank-file", bank_file) + trigfil = node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node # Parallelize coinc trigger followup - factor = int(workflow.cp.get_opt_tags("workflow-rerank", - "parallelization-factor", tags)) - exe = Executable(workflow.cp, 'coinc_followup', ifos=workflow.ifos, - out_dir=out_dir, tags=tags) + factor = int( + workflow.cp.get_opt_tags("workflow-rerank", "parallelization-factor", tags) + ) + exe = Executable( + workflow.cp, "coinc_followup", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) stat_files = FileList([]) for i in range(factor): node = exe.create_node() - node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file', tags=[str(i)]) - node.add_multiifo_input_list_opt('--hdf-store', stores) - node.add_input_opt('--input-file', trigfil) - node.add_opt('--start-index', str(i)) - node.add_opt('--stride', factor) + node.new_output_file_opt( + workflow.analysis_time, ".hdf", "--output-file", tags=[str(i)] + ) + node.add_multiifo_input_list_opt("--hdf-store", stores) + node.add_input_opt("--input-file", trigfil) + node.add_opt("--start-index", str(i)) + node.add_opt("--stride", factor) workflow += node stat_files += node.output_files - exe = Executable(workflow.cp, 'rerank_coincs', ifos=workflow.ifos, - out_dir=out_dir, tags=tags) + exe = Executable( + workflow.cp, "rerank_coincs", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) node = exe.create_node() - node.add_input_list_opt('--stat-files', stat_files) - node.add_input_opt('--statmap-file', statmap_file) - node.add_input_opt('--followup-file', trigfil) + node.add_input_list_opt("--stat-files", stat_files) + node.add_input_opt("--statmap-file", statmap_file) + node.add_input_opt("--followup-file", trigfil) if ranking_file: - node.add_input_opt('--ranking-file', ranking_file) + node.add_input_opt("--ranking-file", ranking_file) - node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file') + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node return node.output_file diff --git a/pycbc/workflow/configparser_test.py b/pycbc/workflow/configparser_test.py index bf6bf344b83..ba6c8af4dac 100644 --- a/pycbc/workflow/configparser_test.py +++ b/pycbc/workflow/configparser_test.py @@ -1,12 +1,15 @@ -import re import copy +import re + try: import configparser as ConfigParser except ImportError: import ConfigParser -def parse_workflow_ini_file(cpFile,parsed_filepath=None): - """Read a .ini file in, parse it as described in the documentation linked + +def parse_workflow_ini_file(cpFile, parsed_filepath=None): + """ + Read a .ini file in, parse it as described in the documentation linked to above, and return the parsed ini file. Parameters @@ -19,6 +22,7 @@ def parse_workflow_ini_file(cpFile,parsed_filepath=None): Returns ------- cp: The parsed ConfigParser class containing the read in .ini file + """ # First read the .ini file cp = read_ini_file(cpFile) @@ -30,7 +34,7 @@ def parse_workflow_ini_file(cpFile,parsed_filepath=None): # We use the same formatting as the new configparser module when doing # ExtendedInterpolation # This is described at http://docs.python.org/3.4/library/configparser.html - #cp = perform_extended_interpolation(cp) + # cp = perform_extended_interpolation(cp) # Split sections like [inspiral&tmplt] into [inspiral] and [tmplt] cp = split_multi_sections(cp) @@ -42,7 +46,7 @@ def parse_workflow_ini_file(cpFile,parsed_filepath=None): # Dump parsed .ini file if needed if parsed_filepath: - fp = open(parsed_filepath,'w') + fp = open(parsed_filepath, "w") cp.write(fp) fp.close() @@ -50,7 +54,8 @@ def parse_workflow_ini_file(cpFile,parsed_filepath=None): def read_ini_file(cpFile): - """Read a .ini file and return it as a ConfigParser class. + """ + Read a .ini file and return it as a ConfigParser class. This function does none of the parsing/combining of sections. It simply reads the file and returns it unedited @@ -61,19 +66,20 @@ def read_ini_file(cpFile): Returns ------- cp: The ConfigParser class containing the read in .ini file - """ + """ # Initialise ConfigParser class - cp = ConfigParser.ConfigParser(\ - interpolation=ConfigParser.ExtendedInterpolation()) + cp = ConfigParser.ConfigParser(interpolation=ConfigParser.ExtendedInterpolation()) # Read the file - fp = open(cpFile,'r') + fp = open(cpFile) cp.read_file(fp) fp.close() return cp -def perform_extended_interpolation(cp,preserve_orig_file=False): - """Filter through an ini file and replace all examples of + +def perform_extended_interpolation(cp, preserve_orig_file=False): + """ + Filter through an ini file and replace all examples of ExtendedInterpolation formatting with the exact value. For values like ${example} this is replaced with the value that corresponds to the option called example ***in the same section*** @@ -96,6 +102,7 @@ def perform_extended_interpolation(cp,preserve_orig_file=False): Returns ------- cp: parsed ConfigParser object + """ # Deepcopy the cp object if needed if preserve_orig_file: @@ -103,21 +110,23 @@ def perform_extended_interpolation(cp,preserve_orig_file=False): # Do not allow any interpolation of the section names for section in cp.sections(): - for option,value in cp.items(section): - # Check the option name - newStr = interpolate_string(option,cp,section) - if newStr != option: - cp.set(section,newStr,value) - cp.remove_option(section,option) - # Check the value - newStr = interpolate_string(value,cp,section) - if newStr != value: - cp.set(section,option,newStr) + for option, value in cp.items(section): + # Check the option name + newStr = interpolate_string(option, cp, section) + if newStr != option: + cp.set(section, newStr, value) + cp.remove_option(section, option) + # Check the value + newStr = interpolate_string(value, cp, section) + if newStr != value: + cp.set(section, option, newStr) return cp -def interpolate_string(testString,cp,section): - """Take a string and replace all example of ExtendedInterpolation formatting + +def interpolate_string(testString, cp, section): + """ + Take a string and replace all example of ExtendedInterpolation formatting within the string with the exact value. For values like ${example} this is replaced with the value that corresponds @@ -140,11 +149,11 @@ def interpolate_string(testString,cp,section): The current section of the ConfigParser object Returns - ---------- + ------- testString: String Interpolated string - """ + """ # First check if any interpolation is needed and abort if not reObj = re.search(r"\$\{.*?\}", testString) while reObj: @@ -152,19 +161,23 @@ def interpolate_string(testString,cp,section): # instance of a string contained within ${....} repString = (reObj).group(0)[2:-1] # Need to test which of the two formats we have - splitString = repString.split('|') + splitString = repString.split("|") if len(splitString) == 1: - testString = testString.replace('${'+repString+'}',\ - cp.get(section,splitString[0])) + testString = testString.replace( + "${" + repString + "}", cp.get(section, splitString[0]) + ) if len(splitString) == 2: - testString = testString.replace('${'+repString+'}',\ - cp.get(splitString[0],splitString[1])) + testString = testString.replace( + "${" + repString + "}", cp.get(splitString[0], splitString[1]) + ) reObj = re.search(r"\$\{.*?\}", testString) return testString -def split_multi_sections(cp,preserve_orig_file=False): - """Parse through a supplied ConfigParser object and splits any sections + +def split_multi_sections(cp, preserve_orig_file=False): + """ + Parse through a supplied ConfigParser object and splits any sections labelled with an "&" sign (for e.g. [inspiral&tmpltbank]) into [inspiral] and [tmpltbank] sections. If these individual sections already exist they will be appended to. If an option exists in both the [inspiral] and @@ -179,8 +192,9 @@ def split_multi_sections(cp,preserve_orig_file=False): Default = False Returns - ---------- + ------- cp: The ConfigParser class + """ # Deepcopy the cp object if needed if preserve_orig_file: @@ -189,20 +203,22 @@ def split_multi_sections(cp,preserve_orig_file=False): # Begin by looping over all sections for section in cp.sections(): # Only continue if section needs splitting - if '&' not in section: + if "&" not in section: continue # Get list of section names to add these options to - splitSections = section.split('&') + splitSections = section.split("&") for newSec in splitSections: # Add sections if they don't already exist if not cp.has_section(newSec): cp.add_section(newSec) - add_options_to_section(cp,newSec,cp.items(section)) + add_options_to_section(cp, newSec, cp.items(section)) cp.remove_section(section) return cp + def sanity_check_subsections(cp): - """This function goes through the ConfigParset and checks that any options + """ + This function goes through the ConfigParset and checks that any options given in the [SECTION_NAME] section are not also given in any [SECTION_NAME-SUBSECTION] sections. @@ -211,21 +227,25 @@ def sanity_check_subsections(cp): cp: The ConfigParser class Returns - ---------- + ------- None + """ # Loop over the sections in the ini file for section in cp.sections(): # Loop over the sections again for section2 in cp.sections(): # Check if any are subsections of section - if section2.startswith(section + '-'): + if section2.startswith(section + "-"): # Check for duplicate options whenever this exists - check_duplicate_options(cp,section,section2,raise_error=True) + check_duplicate_options(cp, section, section2, raise_error=True) -def add_options_to_section(cp,section,items,preserve_orig_file=False,\ - overwrite_options=False): - """Add a set of options and values to a section of a ConfigParser object. + +def add_options_to_section( + cp, section, items, preserve_orig_file=False, overwrite_options=False +): + """ + Add a set of options and values to a section of a ConfigParser object. Will throw an error if any of the options being added already exist, this behaviour can be overridden if desired @@ -247,30 +267,36 @@ def add_options_to_section(cp,section,items,preserve_orig_file=False,\ items. This will override so that the options+values given in items will replace the original values if the value is set to True. Default = True + Returns - ---------- + ------- cp: The ConfigParser class + """ # Sanity checking if not cp.has_section(section): - raise ValueError('Section %s not present in ConfigParser.' %(section,)) + raise ValueError("Section %s not present in ConfigParser." % (section,)) # Deepcopy the cp object if needed if preserve_orig_file: cp = copy.deepcopy(cp) # Check for duplicate options first - for option,value in items: + for option, value in items: if not overwrite_options: if option in cp.options(section): - raise ValueError('Option %s exists in both original' + \ - 'ConfigParser and input list' %(option,)) - cp.set(section,option,value) + raise ValueError( + "Option %s exists in both original" + + "ConfigParser and input list" % (option,) + ) + cp.set(section, option, value) return cp -def check_duplicate_options(cp,section1,section2,raise_error=False): - """Check for duplicate options in two sections, section1 and section2. + +def check_duplicate_options(cp, section1, section2, raise_error=False): + """ + Check for duplicate options in two sections, section1 and section2. Will return True if there are duplicate options and False if not @@ -286,15 +312,16 @@ def check_duplicate_options(cp,section1,section2,raise_error=False): Default = False Returns - ---------- + ------- duplicate: List List of duplicate options + """ # Sanity checking if not cp.has_section(section1): - raise ValueError('Section %s not present in ConfigParser.'%(section1,)) + raise ValueError("Section %s not present in ConfigParser." % (section1,)) if not cp.has_section(section2): - raise ValueError('Section %s not present in ConfigParser.'%(section2,)) + raise ValueError("Section %s not present in ConfigParser." % (section2,)) items1 = cp.items(section1) items2 = cp.items(section2) @@ -303,8 +330,9 @@ def check_duplicate_options(cp,section1,section2,raise_error=False): duplicates = [x for x in items1 if x in items2] if duplicates and raise_error: - raise ValueError('The following options appear in both section ' +\ - '%s and %s: %s' \ - %(section1,section2,duplicates.join(' '))) + raise ValueError( + "The following options appear in both section " + + "%s and %s: %s" % (section1, section2, duplicates.join(" ")) + ) return duplicates diff --git a/pycbc/workflow/configuration.py b/pycbc/workflow/configuration.py index 4ccc09190c0..2b952557a31 100644 --- a/pycbc/workflow/configuration.py +++ b/pycbc/workflow/configuration.py @@ -27,25 +27,25 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/ahope/initialization_inifile.html """ -import re -import os +import hashlib import logging -import stat +import os +import re import shutil +import stat import subprocess -from shutil import which import urllib.parse -import hashlib +from shutil import which from pycbc.types.config import InterpolatingConfigParser -logger = logging.getLogger('pycbc.workflow.configuration') +logger = logging.getLogger("pycbc.workflow.configuration") # NOTE urllib is weird. For some reason it only allows known schemes and will # give *wrong* results, rather then failing, if you use something like gsiftp # We can add schemes explicitly, as below, but be careful with this! -urllib.parse.uses_relative.append('osdf') -urllib.parse.uses_netloc.append('osdf') +urllib.parse.uses_relative.append("osdf") +urllib.parse.uses_netloc.append("osdf") def hash_compare(filename_1, filename_2, chunk_size=None, max_chunks=None): @@ -69,27 +69,24 @@ def hash_compare(filename_1, filename_2, chunk_size=None, max_chunks=None): ------- hash : string The hexdigest() after a sha1 hash of (part of) the file - """ + """ if max_chunks is None and chunk_size is not None: max_chunks = 10 elif chunk_size is None: max_chunks = 1 - with open(filename_1, 'rb') as f1: - with open(filename_2, 'rb') as f2: - for _ in range(max_chunks): - h1 = hashlib.sha1(f1.read(chunk_size)).hexdigest() - h2 = hashlib.sha1(f2.read(chunk_size)).hexdigest() - if h1 != h2: - return False + with open(filename_1, "rb") as f1, open(filename_2, "rb") as f2: + for _ in range(max_chunks): + h1 = hashlib.sha1(f1.read(chunk_size)).hexdigest() + h2 = hashlib.sha1(f2.read(chunk_size)).hexdigest() + if h1 != h2: + return False return True def resolve_url_http(url, u, filename): - """Helper function used by `resolve_url()` to handle HTTP and HTTPS URLs. - """ - + """Helper function used by `resolve_url()` to handle HTTP and HTTPS URLs.""" # Would like to move ciecplib import to top using import_optional, but # it needs to be available when documentation runs in the CI, and I # can't get it to install in the GitHub CI @@ -97,48 +94,46 @@ def resolve_url_http(url, u, filename): headers = None - if u.netloc == 'git.ligo.org': + if u.netloc == "git.ligo.org": # We need to do two ugly special things to download a file from # git.ligo.org. First, we need to pass a per-user GitLab Personal Access # Token via the headers. Second, we need to translate the raw-download # URL scheme to a different scheme that uses GitLab's REST API. re_match = re.match( - 'https://git.ligo.org/([^ ]+(? 6: warn_msg = "This job has way too many tags. " - warn_msg += "Current tags are {}. ".format(' '.join(tags)) - warn_msg += "Current executable {}.".format(self.name) + warn_msg += "Current tags are {}. ".format(" ".join(tags)) + warn_msg += f"Current executable {self.name}." logger.warning(warn_msg) if len(tags) != 0: - self.tagged_name = "{0}-{1}".format(self.name, '_'.join(tags)) + self.tagged_name = "{0}-{1}".format(self.name, "_".join(tags)) else: self.tagged_name = self.name if self.ifo_string is not None: - self.tagged_name = "{0}-{1}".format(self.tagged_name, - self.ifo_string) + self.tagged_name = f"{self.tagged_name}-{self.ifo_string}" # Determine the sections from the ini file that will configure # this executable @@ -590,10 +621,10 @@ def update_current_tags(self, tags): sec_tags = tags + self.ifo_list else: sec_tags = tags - for sec_len in range(1, len(sec_tags)+1): + for sec_len in range(1, len(sec_tags) + 1): for tag_permutation in permutations(sec_tags, sec_len): - joined_name = '-'.join(tag_permutation) - section = '{0}-{1}'.format(self.name, joined_name.lower()) + joined_name = "-".join(tag_permutation) + section = f"{self.name}-{joined_name.lower()}" if self.cp.has_section(section): sections.append(section) @@ -615,22 +646,25 @@ def update_current_tags(self, tags): self._add_ini_opts(self.cp, sec) else: warn_string = "warning: config file is missing section " - warn_string += "[{0}]".format(sec) + warn_string += f"[{sec}]" logger.warning(warn_string) # get uppermost section - if self.cp.has_section(f'{self.name}-defaultvalues'): - self._add_ini_opts(self.cp, f'{self.name}-defaultvalues', - ignore_existing=True) + if self.cp.has_section(f"{self.name}-defaultvalues"): + self._add_ini_opts( + self.cp, f"{self.name}-defaultvalues", ignore_existing=True + ) def update_output_directory(self, out_dir=None): - """Update the default output directory for output files. + """ + Update the default output directory for output files. Parameters - ----------- + ---------- out_dir : string (optional, default=None) If provided use this as the output directory. Else choose this automatically from the tags. + """ # Determine the output directory if out_dir is not None: @@ -648,7 +682,8 @@ def update_output_directory(self, out_dir=None): make_analysis_dir(self.out_dir) def _set_pegasus_profile_options(self): - """Set the pegasus-profile settings for this Executable. + """ + Set the pegasus-profile settings for this Executable. These are a property of the Executable and not of nodes that it will spawn. Therefore it *cannot* be updated without also changing values @@ -657,9 +692,8 @@ def _set_pegasus_profile_options(self): """ # Executable- and tag-specific profile information for sec in self.sections: - if self.cp.has_section('pegasus_profile-{0}'.format(sec)): - self.add_ini_profile(self.cp, - 'pegasus_profile-{0}'.format(sec)) + if self.cp.has_section(f"pegasus_profile-{sec}"): + self.add_ini_profile(self.cp, f"pegasus_profile-{sec}") class Workflow(pegasus_workflow.Workflow): @@ -668,6 +702,7 @@ class Workflow(pegasus_workflow.Workflow): functions for finding input files using time and keywords. It can also generate cache files from the inputs. """ + def __init__(self, args, name=None): """ Create a pycbc workflow @@ -676,17 +711,18 @@ def __init__(self, args, name=None): ---------- args : argparse.ArgumentParser The command line options to initialize a CBC workflow. + """ # Parse ini file self.cp = WorkflowConfigParser.from_cli(args) self.args = args - if hasattr(args, 'dax_file'): + if hasattr(args, "dax_file"): dax_file = args.dax_file or None else: dax_file = None - if hasattr(args, 'dax_file_directory'): + if hasattr(args, "dax_file_directory"): output_dir = args.dax_file_directory or args.output_dir or None else: output_dir = args.output_dir or None @@ -697,7 +733,7 @@ def __init__(self, args, name=None): else: cache_file = None - super(Workflow, self).__init__( + super().__init__( name=name if name is not None else args.workflow_name, directory=output_dir, cache_file=cache_file, @@ -706,23 +742,23 @@ def __init__(self, args, name=None): # Set global values start_time = end_time = 0 - if self.cp.has_option('workflow', 'start-time'): + if self.cp.has_option("workflow", "start-time"): start_time = int(self.cp.get("workflow", "start-time")) - if self.cp.has_option('workflow', 'end-time'): + if self.cp.has_option("workflow", "end-time"): end_time = int(self.cp.get("workflow", "end-time")) self.analysis_time = segments.segment([start_time, end_time]) # Set the ifos to analyse ifos = [] - if self.cp.has_section('workflow-ifos'): - for ifo in self.cp.options('workflow-ifos'): + if self.cp.has_section("workflow-ifos"): + for ifo in self.cp.options("workflow-ifos"): ifos.append(ifo.upper()) self.ifos = ifos self.ifos.sort(key=str.lower) self.get_ifo_combinations() - self.ifo_string = ''.join(self.ifos) + self.ifo_string = "".join(self.ifos) # Set up input and output file lists for workflow self._inputs = FileList([]) @@ -733,33 +769,34 @@ def __init__(self, args, name=None): def output_map(self): args = self.args - if hasattr(args, 'output_map') and args.output_map is not None: + if hasattr(args, "output_map") and args.output_map is not None: return args.output_map if self.in_workflow is not False: - name = self.name + '.map' + name = self.name + ".map" else: - name = 'output.map' + name = "output.map" - path = os.path.join(self.out_dir, name) + path = os.path.join(self.out_dir, name) return path @property def sites(self): """List of all possible exucution sites for jobs in this workflow""" sites = set() - sites.add('local') - if self.cp.has_option('pegasus_profile', 'pycbc|primary_site'): - site = self.cp.get('pegasus_profile', 'pycbc|primary_site') + sites.add("local") + if self.cp.has_option("pegasus_profile", "pycbc|primary_site"): + site = self.cp.get("pegasus_profile", "pycbc|primary_site") else: # The default if not chosen - site = 'condorpool_symlink' + site = "condorpool_symlink" sites.add(site) - subsections = [sec for sec in self.cp.sections() - if sec.startswith('pegasus_profile-')] + subsections = [ + sec for sec in self.cp.sections() if sec.startswith("pegasus_profile-") + ] for subsec in subsections: - if self.cp.has_option(subsec, 'pycbc|site'): - site = self.cp.get(subsec, 'pycbc|site') + if self.cp.has_option(subsec, "pycbc|site"): + site = self.cp.get(subsec, "pycbc|site") sites.add(site) return list(sites) @@ -768,23 +805,22 @@ def staging_site(self): """Site to use for staging to/from each site""" staging_site = {} for site in self.sites: - if site in ['condorpool_shared']: + if site in ["condorpool_shared"]: staging_site[site] = site else: - staging_site[site] = 'local' + staging_site[site] = "local" return staging_site @property def staging_site_str(self): - return ','.join(['='.join(x) for x in self.staging_site.items()]) + return ",".join(["=".join(x) for x in self.staging_site.items()]) @property def exec_sites_str(self): - return ','.join(self.sites) + return ",".join(self.sites) - def execute_node(self, node, verbatim_exe = False): - """ Execute this node immediately on the local machine - """ + def execute_node(self, node, verbatim_exe=False): + """Execute this node immediately on the local machine""" node.executed = True # Check that the PFN is for a file or path @@ -796,15 +832,15 @@ def execute_node(self, node, verbatim_exe = False): # or it may have been marked nonlocal. That's # fine, we'll resolve the URL and make a local # entry. - pfn = node.executable.get_pfn('nonlocal') + pfn = node.executable.get_pfn("nonlocal") resolved = resolve_url( - pfn, - permissions=stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR + pfn, permissions=stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR ) node.executable.clear_pfns() - node.executable.add_pfn(urljoin('file:', pathname2url(resolved)), - site='local') + node.executable.add_pfn( + urljoin("file:", pathname2url(resolved)), site="local" + ) cmd_list = node.get_command_line() @@ -814,15 +850,17 @@ def execute_node(self, node, verbatim_exe = False): os.chdir(out_dir) # Make call - make_external_call(cmd_list, out_dir=os.path.join(out_dir, 'logs'), - out_basename=node.executable.name) + make_external_call( + cmd_list, + out_dir=os.path.join(out_dir, "logs"), + out_basename=node.executable.name, + ) # Change back os.chdir(curr_dir) for fil in node._outputs: fil.node = None - fil.add_pfn(urljoin('file:', pathname2url(fil.storage_path)), - site='local') + fil.add_pfn(urljoin("file:", pathname2url(fil.storage_path)), site="local") def save(self, filename=None, output_map_path=None, root=True): # FIXME: Too close to pegasus to live here and not in pegasus_workflow @@ -831,14 +869,14 @@ def save(self, filename=None, output_map_path=None, root=True): output_map_path = self.output_map output_map_file = pegasus_workflow.File(os.path.basename(output_map_path)) - output_map_file.add_pfn(output_map_path, site='local') + output_map_file.add_pfn(output_map_path, site="local") self.output_map_file = output_map_file if self.in_workflow: self._as_job.set_subworkflow_properties( output_map_file, staging_site=self.staging_site, - cache_file=self.cache_file + cache_file=self.cache_file, ) self._as_job.add_planner_args(**self._as_job.pycbc_planner_args) @@ -850,7 +888,7 @@ def save(self, filename=None, output_map_path=None, root=True): self.add_container(container) # save the configuration file - ini_file = os.path.join(self.out_dir, self.name + '.ini') + ini_file = os.path.join(self.out_dir, self.name + ".ini") # This shouldn't already exist, but just in case if os.path.isfile(ini_file): @@ -859,29 +897,32 @@ def save(self, filename=None, output_map_path=None, root=True): err_msg += ini_file raise ValueError(err_msg) - with open(ini_file, 'w') as fp: + with open(ini_file, "w") as fp: self.cp.write(fp) # save the sites file # checking for in_workflow prevents sub-workflows from making # extra unused sites.yml if not self.in_workflow: - catalog_path = os.path.join(self.out_dir, 'sites.yml') + catalog_path = os.path.join(self.out_dir, "sites.yml") make_catalog(self.cp, self.out_dir).write(catalog_path) # save the dax file - super(Workflow, self).save(filename=filename, - output_map_path=output_map_path, - submit_now=self.args.submit_now, - plan_now=self.args.plan_now, - root=root) + super().save( + filename=filename, + output_map_path=output_map_path, + submit_now=self.args.submit_now, + plan_now=self.args.plan_now, + root=root, + ) def save_config(self, fname, output_dir, cp=None): - """ Writes configuration file to disk and returns a pycbc.workflow.File + """ + Writes configuration file to disk and returns a pycbc.workflow.File instance for the configuration file. Parameters - ----------- + ---------- fname : string The filename of the configuration file written to disk. output_dir : string @@ -893,13 +934,15 @@ def save_config(self, fname, output_dir, cp=None): ------- FileList The FileList object with the configuration file. + """ cp = self.cp if cp is None else cp ini_file_path = os.path.abspath(os.path.join(output_dir, fname)) with open(ini_file_path, "w") as fp: cp.write(fp) - ini_file = File(self.ifos, "", self.analysis_time, - file_url="file://" + ini_file_path) + ini_file = File( + self.ifos, "", self.analysis_time, file_url="file://" + ini_file_path + ) # set the physical file name ini_file.add_pfn(ini_file_path, "local") # set the storage path to be the same @@ -913,13 +956,14 @@ def get_ifo_combinations(self): """ self.ifo_combinations = [] for n in range(len(self.ifos)): - self.ifo_combinations += [''.join(ifos).lower() for ifos in - combinations(self.ifos, n + 1)] + self.ifo_combinations += [ + "".join(ifos).lower() for ifos in combinations(self.ifos, n + 1) + ] class Node(pegasus_workflow.Node): def __init__(self, executable, valid_seg=None): - super(Node, self).__init__(executable.get_transformation()) + super().__init__(executable.get_transformation()) self.executable = executable self.executed = False self.set_category(executable.name) @@ -947,12 +991,12 @@ def get_command_line(self): tmpargs = [] for a in arglist: if not isinstance(a, File): - tmpargs += a.split(' ') + tmpargs += a.split(" ") else: tmpargs.append(a) arglist = tmpargs - arglist = [a for a in arglist if a != ''] + arglist = [a for a in arglist if a != ""] arglist = [a.storage_path if isinstance(a, File) else a for a in arglist] @@ -962,14 +1006,21 @@ def get_command_line(self): return [exe_path] + arglist - def new_output_file_opt(self, valid_seg, extension, option_name, tags=None, - store_file=None, use_tmp_subdirs=False): + def new_output_file_opt( + self, + valid_seg, + extension, + option_name, + tags=None, + store_file=None, + use_tmp_subdirs=False, + ): """ This function will create a workflow.File object corresponding to the given information and then add that file as output of this node. Parameters - ----------- + ---------- valid_seg : igwn_segments.segment The time span over which the job is valid for. extension : string @@ -986,6 +1037,7 @@ def new_output_file_opt(self, valid_seg, extension, option_name, tags=None, This file is to be added to the output mapper and will be stored in the specified output location if True. If false file will be removed when no longer needed in the workflow. + """ if tags is None: tags = [] @@ -997,55 +1049,73 @@ def new_output_file_opt(self, valid_seg, extension, option_name, tags=None, if tag not in all_tags: all_tags.append(tag) - store_file = store_file if store_file is not None else self.executable.retain_files + store_file = ( + store_file if store_file is not None else self.executable.retain_files + ) - fil = File(self.executable.ifo_list, self.executable.name, - valid_seg, extension=extension, store_file=store_file, - directory=self.executable.out_dir, tags=all_tags, - use_tmp_subdirs=use_tmp_subdirs) + fil = File( + self.executable.ifo_list, + self.executable.name, + valid_seg, + extension=extension, + store_file=store_file, + directory=self.executable.out_dir, + tags=all_tags, + use_tmp_subdirs=use_tmp_subdirs, + ) self.add_output_opt(option_name, fil) return fil def add_multiifo_input_list_opt(self, opt, inputs): - """ Add an option that determines a list of inputs from multiple - detectors. Files will be supplied as --opt ifo1:input1 ifo2:input2 - ..... + """ + Add an option that determines a list of inputs from multiple + detectors. Files will be supplied as --opt ifo1:input1 ifo2:input2 + ..... """ # NOTE: Here we have to use the raw arguments functionality as the # file and ifo are not space separated. self.add_raw_arg(opt) - self.add_raw_arg(' ') + self.add_raw_arg(" ") for infile in inputs: self.add_raw_arg(infile.ifo) - self.add_raw_arg(':') + self.add_raw_arg(":") self.add_raw_arg(infile.name) - self.add_raw_arg(' ') + self.add_raw_arg(" ") self.add_input(infile) def add_multiifo_output_list_opt(self, opt, outputs): - """ Add an option that determines a list of outputs from multiple - detectors. Files will be supplied as --opt ifo1:input1 ifo2:input2 - ..... + """ + Add an option that determines a list of outputs from multiple + detectors. Files will be supplied as --opt ifo1:input1 ifo2:input2 + ..... """ # NOTE: Here we have to use the raw arguments functionality as the # file and ifo are not space separated. self.add_raw_arg(opt) - self.add_raw_arg(' ') + self.add_raw_arg(" ") for outfile in outputs: self.add_raw_arg(outfile.ifo) - self.add_raw_arg(':') + self.add_raw_arg(":") self.add_raw_arg(outfile.name) - self.add_raw_arg(' ') + self.add_raw_arg(" ") self.add_output(outfile) - def new_multiifo_output_list_opt(self, opt, ifos, analysis_time, extension, - tags=None, store_file=None, - use_tmp_subdirs=False): - """ Add an option that determines a list of outputs from multiple - detectors. Files will be supplied as --opt ifo1:input1 ifo2:input2 - ..... - File names are created internally from the provided extension and - analysis time. + def new_multiifo_output_list_opt( + self, + opt, + ifos, + analysis_time, + extension, + tags=None, + store_file=None, + use_tmp_subdirs=False, + ): + """ + Add an option that determines a list of outputs from multiple + detectors. Files will be supplied as --opt ifo1:input1 ifo2:input2 + ..... + File names are created internally from the provided extension and + analysis time. """ if tags is None: tags = [] @@ -1055,14 +1125,21 @@ def new_multiifo_output_list_opt(self, opt, ifos, analysis_time, extension, all_tags.append(tag) output_files = FileList([]) - store_file = store_file if store_file is not None \ - else self.executable.retain_files + store_file = ( + store_file if store_file is not None else self.executable.retain_files + ) for ifo in ifos: - curr_file = File(ifo, self.executable.name, analysis_time, - extension=extension, store_file=store_file, - directory=self.executable.out_dir, tags=all_tags, - use_tmp_subdirs=use_tmp_subdirs) + curr_file = File( + ifo, + self.executable.name, + analysis_time, + extension=extension, + store_file=store_file, + directory=self.executable.out_dir, + tags=all_tags, + use_tmp_subdirs=use_tmp_subdirs, + ) output_files.append(curr_file) self.add_multiifo_output_list_opt(opt, output_files) @@ -1084,13 +1161,13 @@ def output_file(self): if len(out_files) != 1: err_msg = "output_file property is only valid if there is a single" err_msg += " output file. Here there are " - err_msg += "%d output files." %(len(out_files)) + err_msg += "%d output files." % (len(out_files)) raise ValueError(err_msg) return out_files[0] class File(pegasus_workflow.File): - ''' + """ This class holds the details of an individual output file This file(s) may be pre-supplied, generated from within the workflow command line script, or generated within the workflow. The important stuff @@ -1110,10 +1187,20 @@ class File(pegasus_workflow.File): >> c = File("H1", "INSPIRAL_S6LOWMASS", segments.segment(815901601, 815902001), directory="/home/spxiwh", extension="xml.gz" ) - ''' - def __init__(self, ifos, exe_name, segs, file_url=None, - extension=None, directory=None, tags=None, - store_file=True, use_tmp_subdirs=False): + """ + + def __init__( + self, + ifos, + exe_name, + segs, + file_url=None, + extension=None, + directory=None, + tags=None, + store_file=True, + use_tmp_subdirs=False, + ): """ Create a File instance @@ -1152,6 +1239,7 @@ def __init__(self, ifos, exe_name, segs, file_url=None, This is a list of descriptors describing what this file is. For e.g. this might be ["BNSINJECTIONS" ,"LOWMASS","CAT_2_VETO"]. These are used in file naming. + """ self.metadata = {} @@ -1162,9 +1250,9 @@ def __init__(self, ifos, exe_name, segs, file_url=None, self.ifo_list = ifos if self.ifo_list is not None: - self.ifo_string = ''.join(self.ifo_list) + self.ifo_string = "".join(self.ifo_list) else: - self.ifo_string = 'file' + self.ifo_string = "file" self.description = exe_name @@ -1186,14 +1274,14 @@ def __init__(self, ifos, exe_name, segs, file_url=None, if tags is None: tags = [] - if '' in tags: - logger.warning('DO NOT GIVE EMPTY TAGS (from %s)', exe_name) - tags.remove('') + if "" in tags: + logger.warning("DO NOT GIVE EMPTY TAGS (from %s)", exe_name) + tags.remove("") self.tags = tags if len(self.tags): - self.tag_str = '_'.join(tags) - tagged_description = '_'.join([self.description] + tags) + self.tag_str = "_".join(tags) + tagged_description = "_".join([self.description] + tags) else: tagged_description = self.description @@ -1203,26 +1291,31 @@ def __init__(self, ifos, exe_name, segs, file_url=None, if not file_url: if not extension: - raise TypeError("a file extension required if a file_url " - "is not provided") + raise TypeError( + "a file extension required if a file_url is not provided" + ) if not directory: - raise TypeError("a directory is required if a file_url is " - "not provided") + raise TypeError("a directory is required if a file_url is not provided") - filename = self._filename(self.ifo_string, self.tagged_description, - extension, self.segment_list.extent()) + filename = self._filename( + self.ifo_string, + self.tagged_description, + extension, + self.segment_list.extent(), + ) path = os.path.join(directory, filename) if not os.path.isabs(path): path = os.path.join(os.getcwd(), path) - file_url = urllib.parse.urlunparse(['file', 'localhost', path, - None, None, None]) + file_url = urllib.parse.urlunparse( + ["file", "localhost", path, None, None, None] + ) if use_tmp_subdirs and len(self.segment_list): pegasus_lfn = str(int(self.segment_list.extent()[0]))[:-4] - pegasus_lfn = pegasus_lfn + '/' + os.path.basename(file_url) + pegasus_lfn = pegasus_lfn + "/" + os.path.basename(file_url) else: pegasus_lfn = os.path.basename(file_url) - super(File, self).__init__(pegasus_lfn) + super().__init__(pegasus_lfn) if store_file: self.storage_path = urllib.parse.urlsplit(file_url).path @@ -1230,20 +1323,21 @@ def __init__(self, ifos, exe_name, segs, file_url=None, self.storage_path = None def __getstate__(self): - """ Allow the workflow.File to be picklable. This disables the usage of + """ + Allow the workflow.File to be picklable. This disables the usage of the internal cache entry. """ for i, seg in enumerate(self.segment_list): self.segment_list[i] = segments.segment(float(seg[0]), float(seg[1])) self.cache_entry = None safe_dict = copy.copy(self.__dict__) - safe_dict['cache_entry'] = None + safe_dict["cache_entry"] = None return safe_dict # FIXME: This is a pegasus_workflow thing (don't think it's needed at all!) # use the pegasus function directly (maybe not). def add_metadata(self, key, value): - """ Add arbitrary metadata to this file """ + """Add arbitrary metadata to this file""" self.metadata[key] = value @property @@ -1254,10 +1348,9 @@ def ifo(self): """ if len(self.ifo_list) == 1: return self.ifo_list[0] - else: - err = "self.ifo_list must contain only one ifo to access the " - err += "ifo property. %s." %(str(self.ifo_list),) - raise TypeError(err) + err = "self.ifo_list must contain only one ifo to access the " + err += "ifo property. %s." % (str(self.ifo_list),) + raise TypeError(err) @property def segment(self): @@ -1267,10 +1360,9 @@ def segment(self): """ if len(self.segment_list) == 1: return self.segment_list[0] - else: - err = "self.segment_list must only contain one segment to access" - err += " the segment property. %s." %(str(self.segment_list),) - raise TypeError(err) + err = "self.segment_list must only contain one segment to access" + err += " the segment property. %s." % (str(self.segment_list),) + raise TypeError(err) @property def cache_entry(self): @@ -1278,14 +1370,19 @@ def cache_entry(self): Returns a CacheEntry instance for File. """ if self.storage_path is None: - raise ValueError('This file is temporary and so a lal ' - 'cache entry cannot be made') - - file_url = urllib.parse.urlunparse(['file', 'localhost', - self.storage_path, None, - None, None]) - cache_entry = lal.utils.CacheEntry(self.ifo_string, - self.tagged_description, self.segment_list.extent(), file_url) + raise ValueError( + "This file is temporary and so a lal cache entry cannot be made" + ) + + file_url = urllib.parse.urlunparse( + ["file", "localhost", self.storage_path, None, None, None] + ) + cache_entry = lal.utils.CacheEntry( + self.ifo_string, + self.tagged_description, + self.segment_list.extent(), + file_url, + ) cache_entry.workflow_file = self return cache_entry @@ -1294,18 +1391,16 @@ def _filename(self, ifo, description, extension, segment): Construct the standard output filename. Should only be used internally of the File class. """ - if extension.startswith('.'): - extension = extension[1:] + extension = extension.removeprefix(".") # Follow the frame convention of using integer filenames, # but stretching to cover partially covered seconds. start = int(segment[0]) end = int(math.ceil(segment[1])) - duration = str(end-start) + duration = str(end - start) start = str(start) - return "%s-%s-%s-%s.%s" % (ifo, description.upper(), start, - duration, extension) + return "%s-%s-%s-%s.%s" % (ifo, description.upper(), start, duration, extension) @classmethod def from_path(cls, path, attrs=None, **kwargs): @@ -1314,52 +1409,54 @@ def from_path(cls, path, attrs=None, **kwargs): """ if attrs is None: attrs = {} - if attrs and 'ifos' in attrs: - ifos = attrs['ifos'] + if attrs and "ifos" in attrs: + ifos = attrs["ifos"] else: - ifos = ['H1', 'K1', 'L1', 'V1'] - if attrs and 'exe_name' in attrs: - exe_name = attrs['exe_name'] + ifos = ["H1", "K1", "L1", "V1"] + if attrs and "exe_name" in attrs: + exe_name = attrs["exe_name"] else: - exe_name = 'INPUT' - if attrs and 'segs' in attrs: - segs = attrs['segs'] + exe_name = "INPUT" + if attrs and "segs" in attrs: + segs = attrs["segs"] else: segs = segments.segment([1, 2000000000]) - if attrs and 'tags' in attrs: - tags = attrs['tags'] + if attrs and "tags" in attrs: + tags = attrs["tags"] else: tags = [] curr_file = cls(ifos, exe_name, segs, path, tags=tags, **kwargs) return curr_file + class FileList(list): - ''' + """ This class holds a list of File objects. It inherits from the built-in list class, but also allows a number of features. ONLY pycbc.workflow.File instances should be within a FileList instance. - ''' + """ + entry_class = File def categorize_by_attr(self, attribute): - ''' + """ Function to categorize a FileList by a File object attribute (eg. 'segment', 'ifo', 'description'). Parameters - ----------- + ---------- attribute : string File object attribute to categorize FileList Returns - -------- + ------- keys : list A list of values for an attribute groups : list A list of FileLists - ''' + """ # need to sort FileList otherwise using groupby without sorting does # 'AAABBBCCDDAABB' -> ['AAA','BBB','CC','DD','AA','BB'] # and using groupby with sorting does @@ -1376,13 +1473,14 @@ def categorize_by_attr(self, attribute): return keys, groups def find_output(self, ifo, time): - '''Returns one File most appropriate at the given time/time range. + """ + Returns one File most appropriate at the given time/time range. Return one File that covers the given time, or is most appropriate for the supplied time range. Parameters - ----------- + ---------- ifo : string Name of the ifo (or ifos) that the file should be valid for. time : int/float/LIGOGPStime or tuple containing two values @@ -1394,31 +1492,32 @@ def find_output(self, ifo, time): self.find_output_in_range Returns - -------- + ------- pycbc_file : pycbc.workflow.File instance The File that corresponds to the time or time range - ''' + + """ # Determine whether I have a specific time, or a range of times try: lenTime = len(time) except TypeError: # This is if I have a single time - outFile = self.find_output_at_time(ifo,time) + outFile = self.find_output_at_time(ifo, time) else: # This is if I have a range of times if lenTime == 2: - outFile = self.find_output_in_range(ifo,time[0],time[1]) + outFile = self.find_output_in_range(ifo, time[0], time[1]) # This is if I got a list that had more (or less) than 2 entries if len(time) != 2: raise TypeError("I do not understand the input variable time") return outFile def find_output_at_time(self, ifo, time): - ''' + """ Return File that covers the given time. Parameters - ----------- + ---------- ifo : string Name of the ifo (or ifos) that the File should correspond to time : int/float/LIGOGPStime @@ -1426,23 +1525,23 @@ def find_output_at_time(self, ifo, time): File covers the time this will return None. Returns - -------- + ------- list of File classes The Files that corresponds to the time. - ''' + + """ # Get list of Files that overlap time, for given ifo outFiles = [i for i in self if ifo in i.ifo_list and time in i.segment_list] if len(outFiles) == 0: # No OutFile at this time return None - elif len(outFiles) == 1: + if len(outFiles) == 1: # 1 OutFile at this time (good!) return outFiles - else: - # Multiple output files. Currently this is valid, but we may want - # to demand exclusivity later, or in certain cases. Hence the - # separation. - return outFiles + # Multiple output files. Currently this is valid, but we may want + # to demand exclusivity later, or in certain cases. Hence the + # separation. + return outFiles def find_outputs_in_range(self, ifo, current_segment, useSplitLists=False): """ @@ -1451,25 +1550,29 @@ def find_outputs_in_range(self, ifo, current_segment, useSplitLists=False): largest overlap with the supplied time range. Parameters - ----------- + ---------- ifo : string Name of the ifo (or ifos) that the File should correspond to current_segment : igwn_segments.segment The segment of time that files must intersect. Returns - -------- + ------- FileList class The list of Files that are most appropriate for the time range + """ currsegment_list = segments.segmentlist([current_segment]) # Get all files overlapping the window - overlap_files = self.find_all_output_in_range(ifo, current_segment, - useSplitLists=useSplitLists) + overlap_files = self.find_all_output_in_range( + ifo, current_segment, useSplitLists=useSplitLists + ) # By how much do they overlap? - overlap_windows = [abs(i.segment_list & currsegment_list) for i in overlap_files] + overlap_windows = [ + abs(i.segment_list & currsegment_list) for i in overlap_files + ] if not overlap_windows: return [] @@ -1477,22 +1580,22 @@ def find_outputs_in_range(self, ifo, current_segment, useSplitLists=False): # Return the File with the biggest overlap # Note if two File have identical overlap, the first is used # to define the valid segment - overlap_windows = numpy.array(overlap_windows, dtype = int) + overlap_windows = numpy.array(overlap_windows, dtype=int) segmentLst = overlap_files[overlap_windows.argmax()].segment_list # Get all output files with the exact same segment definition - output_files = [f for f in overlap_files if f.segment_list==segmentLst] + output_files = [f for f in overlap_files if f.segment_list == segmentLst] return output_files def find_output_in_range(self, ifo, start, end): - ''' + """ Return the File that is most appropriate for the supplied time range. That is, the File whose coverage time has the largest overlap with the supplied time range. If no Files overlap the supplied time window, will return None. Parameters - ----------- + ---------- ifo : string Name of the ifo (or ifos) that the File should correspond to start : int/float/LIGOGPStime @@ -1501,10 +1604,11 @@ def find_output_in_range(self, ifo, start, end): The end of the time range of interest Returns - -------- + ------- File class The File that is most appropriate for the time range - ''' + + """ currsegment_list = segments.segmentlist([segments.segment(start, end)]) # First filter Files corresponding to ifo @@ -1514,24 +1618,21 @@ def find_output_in_range(self, ifo, start, end): # No OutFiles correspond to that ifo return None # Filter OutFiles to those overlapping the given window - currSeg = segments.segment([start,end]) - outFiles = [i for i in outFiles \ - if i.segment_list.intersects_segment(currSeg)] + currSeg = segments.segment([start, end]) + outFiles = [i for i in outFiles if i.segment_list.intersects_segment(currSeg)] if len(outFiles) == 0: # No OutFile overlap that time period return None - elif len(outFiles) == 1: + if len(outFiles) == 1: # One OutFile overlaps that period return outFiles[0] - else: - overlap_windows = [abs(i.segment_list & currsegment_list) \ - for i in outFiles] - # Return the File with the biggest overlap - # Note if two File have identical overlap, this will return - # the first File in the list - overlap_windows = numpy.array(overlap_windows, dtype = int) - return outFiles[overlap_windows.argmax()] + overlap_windows = [abs(i.segment_list & currsegment_list) for i in outFiles] + # Return the File with the biggest overlap + # Note if two File have identical overlap, this will return + # the first File in the list + overlap_windows = numpy.array(overlap_windows, dtype=int) + return outFiles[overlap_windows.argmax()] def find_all_output_in_range(self, ifo, currSeg, useSplitLists=False): """ @@ -1540,8 +1641,9 @@ def find_all_output_in_range(self, ifo, currSeg, useSplitLists=False): if not useSplitLists: # Slower, but simpler method outFiles = [i for i in self if ifo in i.ifo_list] - outFiles = [i for i in outFiles - if i.segment_list.intersects_segment(currSeg)] + outFiles = [ + i for i in outFiles if i.segment_list.intersects_segment(currSeg) + ] else: # Faster, but more complicated # Basically only check if a subset of files intersects_segment by @@ -1549,8 +1651,7 @@ def find_all_output_in_range(self, ifo, currSeg, useSplitLists=False): if not self._check_split_list_validity(): # FIXME: DO NOT hard code this. self._temporal_split_list(100) - startIdx = int((currSeg[0] - self._splitListsStart) / - self._splitListsStep) + startIdx = int((currSeg[0] - self._splitListsStart) / self._splitListsStep) # Add some small rounding here endIdx = (currSeg[1] - self._splitListsStart) / self._splitListsStep endIdx = int(endIdx - 0.000001) @@ -1559,10 +1660,14 @@ def find_all_output_in_range(self, ifo, currSeg, useSplitLists=False): for idx in range(startIdx, endIdx + 1): if idx < 0 or idx >= self._splitListsNum: continue - outFilesTemp = [i for i in self._splitLists[idx] - if ifo in i.ifo_list] - outFiles.extend([i for i in outFilesTemp - if i.segment_list.intersects_segment(currSeg)]) + outFilesTemp = [i for i in self._splitLists[idx] if ifo in i.ifo_list] + outFiles.extend( + [ + i + for i in outFilesTemp + if i.segment_list.intersects_segment(currSeg) + ] + ) # Remove duplicates outFiles = list(set(outFiles)) @@ -1573,7 +1678,7 @@ def find_output_with_tag(self, tag, fail_if_not_single_file=False): Find all files who have tag in self.tags Parameters - ----------- + ---------- tag : string Tag used to seive the file names fail_if_not_single_file : boolean @@ -1581,13 +1686,14 @@ def find_output_with_tag(self, tag, fail_if_not_single_file=False): user expects to find a single file with the desired tag in its name Returns - -------- + ------- FileList/File class If fail_if_not_single_file is False the FileList Containing File instances with tag in self.tags is returned, otherwise the single File with tag in self.tags is returned (if the sanity check requested with fail_if_not_single_file=True is passed) + """ # Enforce upper case tag = tag.upper() @@ -1640,7 +1746,7 @@ def convert_to_lal_cache(self): pass return lal_cache - def _temporal_split_list(self,numSubLists): + def _temporal_split_list(self, numSubLists): """ This internal function is used to speed the code up in cases where a number of operations are being made to determine if files overlap a @@ -1657,8 +1763,8 @@ def _temporal_split_list(self,numSubLists): invalid. Currently the testing for this is pretty basic """ # Assume segment lists are coalesced! - startTime = float( min([i.segment_list[0][0] for i in self])) - endTime = float( max([i.segment_list[-1][-1] for i in self])) + startTime = float(min([i.segment_list[0][0] for i in self])) + endTime = float(max([i.segment_list[-1][-1] for i in self])) step = (endTime - startTime) / float(numSubLists) # Set up storage @@ -1676,8 +1782,7 @@ def _temporal_split_list(self,numSubLists): startIdx = int(startIdx - 0.001) endIdx = int(endIdx + 0.001) - if startIdx < 0: - startIdx = 0 + startIdx = max(startIdx, 0) if endIdx >= numSubLists: endIdx = numSubLists - 1 @@ -1698,30 +1803,28 @@ def _check_split_list_validity(self): split lists are still valid. """ # FIXME: Currently very primitive, but needs to be fast - if not (hasattr(self,"_splitListsSet") and (self._splitListsSet)): - return False - elif len(self) != self._splitListsLength: + if not (hasattr(self, "_splitListsSet") and (self._splitListsSet)) or len(self) != self._splitListsLength: return False - else: - return True + return True @classmethod def load(cls, filename): """ Load a FileList from a pickle file """ - f = open(filename, 'r') + f = open(filename) return pickle.load(f) def dump(self, filename): """ Output this FileList to a pickle file """ - f = open(filename, 'w') + f = open(filename, "w") pickle.dump(self, f) def to_file_object(self, name, out_dir): - """Dump to a pickle file and return an File object reference + """ + Dump to a pickle file and return an File object reference Parameters ---------- @@ -1733,24 +1836,38 @@ def to_file_object(self, name, out_dir): Returns ------- file : AhopeFile + """ make_analysis_dir(out_dir) - file_ref = File('ALL', name, self.get_times_covered_by_files(), - extension='.pkl', directory=out_dir) + file_ref = File( + "ALL", + name, + self.get_times_covered_by_files(), + extension=".pkl", + directory=out_dir, + ) self.dump(file_ref.storage_path) return file_ref class SegFile(File): - ''' + """ This class inherits from the File class, and is designed to store workflow output files containing a segment dict. This is identical in usage to File except for an additional kwarg for holding the segment dictionary, if it is known at workflow run time. - ''' - def __init__(self, ifo_list, description, valid_segment, - segment_dict=None, seg_summ_dict=None, **kwargs): + """ + + def __init__( + self, + ifo_list, + description, + valid_segment, + segment_dict=None, + seg_summ_dict=None, + **kwargs, + ): """ See File.__init__ for a full set of documentation for how to call this class. The only thing unique and added to this class is @@ -1759,7 +1876,7 @@ def __init__(self, ifo_list, description, valid_segment, we key by dict[ifo:name]. Parameters - ------------ + ---------- ifo_list : string or list (required) See File.__init__ description : string (required) @@ -1773,8 +1890,7 @@ def __init__(self, ifo_list, description, valid_segment, instance of the class. """ - super(SegFile, self).__init__(ifo_list, description, valid_segment, - **kwargs) + super().__init__(ifo_list, description, valid_segment, **kwargs) # To avoid confusion with the segment_list property of the parent class # we refer to this as valid_segments here self.valid_segments = self.segment_list @@ -1782,12 +1898,14 @@ def __init__(self, ifo_list, description, valid_segment, self.seg_summ_dict = seg_summ_dict @classmethod - def from_segment_list(cls, description, segmentlist, name, ifo, - seg_summ_list=None, **kwargs): - """ Initialize a SegFile object from a segmentlist. + def from_segment_list( + cls, description, segmentlist, name, ifo, seg_summ_list=None, **kwargs + ): + """ + Initialize a SegFile object from a segmentlist. Parameters - ------------ + ---------- description : string (required) See File.__init__ segmentlist : igwn_segments.segmentslist @@ -1800,23 +1918,27 @@ def from_segment_list(cls, description, segmentlist, name, ifo, Specify the segment_summary segmentlist that goes along with the segmentlist. Default=None, in this case segment_summary is taken from the valid_segment of the SegFile class. + """ seglistdict = segments.segmentlistdict() - seglistdict[ifo + ':' + name] = segmentlist + seglistdict[ifo + ":" + name] = segmentlist seg_summ_dict = None if seg_summ_list is not None: seg_summ_dict = segments.segmentlistdict() - seg_summ_dict[ifo + ':' + name] = seg_summ_list - return cls.from_segment_list_dict(description, seglistdict, - seg_summ_dict=seg_summ_dict, **kwargs) + seg_summ_dict[ifo + ":" + name] = seg_summ_list + return cls.from_segment_list_dict( + description, seglistdict, seg_summ_dict=seg_summ_dict, **kwargs + ) @classmethod - def from_multi_segment_list(cls, description, segmentlists, names, ifos, - seg_summ_lists=None, **kwargs): - """ Initialize a SegFile object from a list of segmentlists. + def from_multi_segment_list( + cls, description, segmentlists, names, ifos, seg_summ_lists=None, **kwargs + ): + """ + Initialize a SegFile object from a list of segmentlists. Parameters - ------------ + ---------- description : string (required) See File.__init__ segmentlists : List of igwn_segments.segmentslist @@ -1829,29 +1951,38 @@ def from_multi_segment_list(cls, description, segmentlists, names, ifos, Specify the segment_summary segmentlists that go along with the segmentlists. Default=None, in this case segment_summary is taken from the valid_segment of the SegFile class. + """ seglistdict = segments.segmentlistdict() for name, ifo, segmentlist in zip(names, ifos, segmentlists): - seglistdict[ifo + ':' + name] = segmentlist + seglistdict[ifo + ":" + name] = segmentlist if seg_summ_lists is not None: seg_summ_dict = segments.segmentlistdict() for name, ifo, seg_summ_list in zip(names, ifos, seg_summ_lists): - seg_summ_dict[ifo + ':' + name] = seg_summ_list + seg_summ_dict[ifo + ":" + name] = seg_summ_list else: seg_summ_dict = None - return cls.from_segment_list_dict(description, seglistdict, - seg_summ_dict=seg_summ_dict, **kwargs) + return cls.from_segment_list_dict( + description, seglistdict, seg_summ_dict=seg_summ_dict, **kwargs + ) @classmethod - def from_segment_list_dict(cls, description, segmentlistdict, - ifo_list=None, valid_segment=None, - file_exists=False, seg_summ_dict=None, - **kwargs): - """ Initialize a SegFile object from a segmentlistdict. + def from_segment_list_dict( + cls, + description, + segmentlistdict, + ifo_list=None, + valid_segment=None, + file_exists=False, + seg_summ_dict=None, + **kwargs, + ): + """ + Initialize a SegFile object from a segmentlistdict. Parameters - ------------ + ---------- description : string (required) See File.__init__ segmentlistdict : igwn_segments.segmentslistdict @@ -1867,14 +1998,14 @@ def from_segment_list_dict(cls, description, segmentlistdict, exists on disk and so there is no need to write again. seg_summ_dict : igwn_segments.segmentslistdict Optional. See SegFile.__init__. + """ if ifo_list is None: - ifo_set = set([i.split(':')[0] for i in segmentlistdict.keys()]) + ifo_set = set([i.split(":")[0] for i in segmentlistdict.keys()]) ifo_list = list(ifo_set) ifo_list.sort() if valid_segment is None: - if seg_summ_dict and \ - numpy.any([len(v) for _, v in seg_summ_dict.items()]): + if seg_summ_dict and numpy.any([len(v) for _, v in seg_summ_dict.items()]): # Only come here if seg_summ_dict is supplied and it is # not empty. valid_segment = seg_summ_dict.extent_all() @@ -1884,7 +2015,7 @@ def from_segment_list_dict(cls, description, segmentlistdict, except: # Numpty probably didn't supply a # igwn_segments.segmentlistdict - segmentlistdict=segments.segmentlistdict(segmentlistdict) + segmentlistdict = segments.segmentlistdict(segmentlistdict) try: valid_segment = segmentlistdict.extent_all() except ValueError: @@ -1893,15 +2024,21 @@ def from_segment_list_dict(cls, description, segmentlistdict, warn_msg = "No information with which to set valid " warn_msg += "segment." logger.warning(warn_msg) - valid_segment = segments.segment([0,1]) - instnc = cls(ifo_list, description, valid_segment, - segment_dict=segmentlistdict, seg_summ_dict=seg_summ_dict, - **kwargs) + valid_segment = segments.segment([0, 1]) + instnc = cls( + ifo_list, + description, + valid_segment, + segment_dict=segmentlistdict, + seg_summ_dict=seg_summ_dict, + **kwargs, + ) if not file_exists: instnc.to_segment_xml() else: - instnc.add_pfn(urljoin('file:', pathname2url(instnc.storage_path)), - site='local') + instnc.add_pfn( + urljoin("file:", pathname2url(instnc.storage_path)), site="local" + ) return instnc @classmethod @@ -1911,14 +2048,16 @@ def from_segment_xml(cls, xml_file, **kwargs): xml segment table. Parameters - ----------- + ---------- xml_file : file object file object for segment xml file + """ # load xmldocument and SegmentDefTable and SegmentTables - fp = open(xml_file, 'rb') + fp = open(xml_file, "rb") xmldoc = ligolw_utils.load_fileobj( - fp, compress='auto', contenthandler=LIGOLWContentHandler) + fp, compress="auto", contenthandler=LIGOLWContentHandler + ) seg_def_table = lsctables.SegmentDefTable.get_table(xmldoc) seg_table = lsctables.SegmentTable.get_table(xmldoc) @@ -1930,22 +2069,23 @@ def from_segment_xml(cls, xml_file, **kwargs): seg_id = {} for seg_def in seg_def_table: # Here we want to encode ifo and segment name - full_channel_name = ':'.join([str(seg_def.ifos), - str(seg_def.name)]) + full_channel_name = ":".join([str(seg_def.ifos), str(seg_def.name)]) seg_id[int(seg_def.segment_def_id)] = full_channel_name segs[full_channel_name] = segments.segmentlist() seg_summ[full_channel_name] = segments.segmentlist() for seg in seg_table: seg_obj = segments.segment( - lal.LIGOTimeGPS(seg.start_time, seg.start_time_ns), - lal.LIGOTimeGPS(seg.end_time, seg.end_time_ns)) + lal.LIGOTimeGPS(seg.start_time, seg.start_time_ns), + lal.LIGOTimeGPS(seg.end_time, seg.end_time_ns), + ) segs[seg_id[int(seg.segment_def_id)]].append(seg_obj) for seg in seg_sum_table: seg_obj = segments.segment( - lal.LIGOTimeGPS(seg.start_time, seg.start_time_ns), - lal.LIGOTimeGPS(seg.end_time, seg.end_time_ns)) + lal.LIGOTimeGPS(seg.start_time, seg.start_time_ns), + lal.LIGOTimeGPS(seg.end_time, seg.end_time_ns), + ) seg_summ[seg_id[int(seg.segment_def_id)]].append(seg_obj) for seg_name in seg_id.values(): @@ -1953,12 +2093,18 @@ def from_segment_xml(cls, xml_file, **kwargs): xmldoc.unlink() fp.close() - curr_url = urllib.parse.urlunparse(['file', 'localhost', xml_file, - None, None, None]) + curr_url = urllib.parse.urlunparse( + ["file", "localhost", xml_file, None, None, None] + ) - return cls.from_segment_list_dict('SEGMENTS', segs, file_url=curr_url, - file_exists=True, - seg_summ_dict=seg_summ, **kwargs) + return cls.from_segment_list_dict( + "SEGMENTS", + segs, + file_url=curr_url, + file_exists=True, + seg_summ_dict=seg_summ, + **kwargs, + ) def remove_short_sci_segs(self, minSegLength): """ @@ -1967,10 +2113,11 @@ def remove_short_sci_segs(self, minSegLength): these segments. Parameters - ----------- + ---------- minSegLength : int Maximum length of science segments. Segments shorter than this will be removed. + """ newsegment_list = segments.segmentlist() for key, seglist in self.segment_dict.items(): @@ -1989,12 +2136,11 @@ def parse_segdict_key(self, key): """ Return ifo and name from the segdict key. """ - splt = key.split(':') + splt = key.split(":") if len(splt) == 2: return splt[0], splt[1] - else: - err_msg = "Key should be of the format 'ifo:name', got %s." %(key,) - raise ValueError(err_msg) + err_msg = "Key should be of the format 'ifo:name', got %s." % (key,) + raise ValueError(err_msg) def to_segment_xml(self, override_file_if_exists=False): """ @@ -2008,40 +2154,50 @@ def to_segment_xml(self, override_file_if_exists=False): for key, seglist in self.segment_dict.items(): ifo, name = self.parse_segdict_key(key) # Ensure we have LIGOTimeGPS - fsegs = [(lal.LIGOTimeGPS(seg[0]), - lal.LIGOTimeGPS(seg[1])) for seg in seglist] + fsegs = [ + (lal.LIGOTimeGPS(seg[0]), lal.LIGOTimeGPS(seg[1])) for seg in seglist + ] if self.seg_summ_dict is None: - vsegs = [(lal.LIGOTimeGPS(seg[0]), - lal.LIGOTimeGPS(seg[1])) \ - for seg in self.valid_segments] + vsegs = [ + (lal.LIGOTimeGPS(seg[0]), lal.LIGOTimeGPS(seg[1])) + for seg in self.valid_segments + ] else: - vsegs = [(lal.LIGOTimeGPS(seg[0]), - lal.LIGOTimeGPS(seg[1])) \ - for seg in self.seg_summ_dict[key]] + vsegs = [ + (lal.LIGOTimeGPS(seg[0]), lal.LIGOTimeGPS(seg[1])) + for seg in self.seg_summ_dict[key] + ] # Add using glue library to set all segment tables with ligolw_segments.LigolwSegments(outdoc, process) as x: - x.add(ligolw_segments.LigolwSegmentList(active=fsegs, - instruments=set([ifo]), name=name, - version=1, valid=vsegs)) + x.add( + ligolw_segments.LigolwSegmentList( + active=fsegs, + instruments=set([ifo]), + name=name, + version=1, + valid=vsegs, + ) + ) # write file - url = urljoin('file:', pathname2url(self.storage_path)) - if not override_file_if_exists or not self.has_pfn(url, site='local'): - self.add_pfn(url, site='local') + url = urljoin("file:", pathname2url(self.storage_path)) + if not override_file_if_exists or not self.has_pfn(url, site="local"): + self.add_pfn(url, site="local") ligolw_utils.write_filename(outdoc, self.storage_path) -def make_external_call(cmdList, out_dir=None, out_basename='external_call', - shell=False, fail_on_error=True): +def make_external_call( + cmdList, out_dir=None, out_basename="external_call", shell=False, fail_on_error=True +): """ Use this to make an external call using the python subprocess module. See the subprocess documentation for more details of how this works. http://docs.python.org/2/library/subprocess.html Parameters - ----------- + ---------- cmdList : list of strings This list of strings contains the command to be run. See the subprocess documentation for more details. @@ -2064,19 +2220,20 @@ def make_external_call(cmdList, out_dir=None, out_basename='external_call', and out_basename options. Returns - -------- + ------- exitCode : int The code returned by the process. + """ if out_dir: - outBase = os.path.join(out_dir,out_basename) - errFile = outBase + '.err' - errFP = open(errFile, 'w') - outFile = outBase + '.out' - outFP = open(outFile, 'w') - cmdFile = outBase + '.sh' - cmdFP = open(cmdFile, 'w') - cmdFP.write(' '.join(cmdList)) + outBase = os.path.join(out_dir, out_basename) + errFile = outBase + ".err" + errFP = open(errFile, "w") + outFile = outBase + ".out" + outFP = open(outFile, "w") + cmdFile = outBase + ".sh" + cmdFP = open(cmdFile, "w") + cmdFP.write(" ".join(cmdList)) cmdFP.close() else: errFile = None @@ -2085,18 +2242,22 @@ def make_external_call(cmdList, out_dir=None, out_basename='external_call', errFP = None outFP = None - msg = "Making external call %s" %(' '.join(cmdList)) + msg = "Making external call %s" % (" ".join(cmdList)) logger.info(msg) - errCode = subprocess.call(cmdList, stderr=errFP, stdout=outFP,\ - shell=shell) + errCode = subprocess.call(cmdList, stderr=errFP, stdout=outFP, shell=shell) if errFP: errFP.close() if outFP: outFP.close() if errCode and fail_on_error: - raise CalledProcessErrorMod(errCode, ' '.join(cmdList), - errFile=errFile, outFile=outFile, cmdFile=cmdFile) + raise CalledProcessErrorMod( + errCode, + " ".join(cmdList), + errFile=errFile, + outFile=outFile, + cmdFile=cmdFile, + ) logger.info("Call successful, or error checking disabled.") @@ -2106,30 +2267,30 @@ class CalledProcessErrorMod(Exception): and checking has been requested. This should not be accessed by the user it is used only within make_external_call. """ - def __init__(self, returncode, cmd, errFile=None, outFile=None, - cmdFile=None): + + def __init__(self, returncode, cmd, errFile=None, outFile=None, cmdFile=None): self.returncode = returncode self.cmd = cmd self.errFile = errFile self.outFile = outFile self.cmdFile = cmdFile + def __str__(self): - msg = "Command '%s' returned non-zero exit status %d.\n" \ - %(self.cmd, self.returncode) + msg = "Command '%s' returned non-zero exit status %d.\n" % ( + self.cmd, + self.returncode, + ) if self.errFile: - msg += "Stderr can be found in %s .\n" %(self.errFile) + msg += "Stderr can be found in %s .\n" % (self.errFile) if self.outFile: - msg += "Stdout can be found in %s .\n" %(self.outFile) + msg += "Stdout can be found in %s .\n" % (self.outFile) if self.cmdFile: - msg += "The failed command has been printed in %s ." %(self.cmdFile) + msg += "The failed command has been printed in %s ." % (self.cmdFile) return msg def resolve_url_to_file( - curr_pfn, - attrs=None, - hash_max_chunks=10, - hash_chunk_size=int(1e6) + curr_pfn, attrs=None, hash_max_chunks=10, hash_chunk_size=int(1e6) ): """ Resolves a PFN into a workflow.File object. @@ -2158,9 +2319,9 @@ def resolve_url_to_file( hash_max_chunks and hash_chunk_size are used to decide how much of the files to check before they are considered the same, and not copied. """ - cvmfsstr1 = 'file:///cvmfs/' - cvmfsstr2 = 'file://localhost/cvmfs/' - osdfstr1 = 'osdf:///' # Technically this isn't CVMFS, but same handling! + cvmfsstr1 = "file:///cvmfs/" + cvmfsstr2 = "file://localhost/cvmfs/" + osdfstr1 = "osdf:///" # Technically this isn't CVMFS, but same handling! cvmfsstrs = (cvmfsstr1, cvmfsstr2, osdfstr1) # Get LFN @@ -2168,10 +2329,10 @@ def resolve_url_to_file( curr_lfn = os.path.basename(urlp.path) # Does this already exist as a File? - if curr_lfn in file_input_from_config_dict.keys(): + if curr_lfn in file_input_from_config_dict: file_pfn = file_input_from_config_dict[curr_lfn][2] # If the PFNs are different, but LFNs are the same then fail. - assert(file_pfn == curr_pfn) + assert file_pfn == curr_pfn curr_file = file_input_from_config_dict[curr_lfn][1] else: # Use resolve_url to download file/symlink as appropriate @@ -2186,10 +2347,10 @@ def resolve_url_to_file( if curr_pfn.startswith(cvmfsstrs): # Add PFNs for nonlocal sites for special cases (e.g. CVMFS). # This block could be extended as needed - curr_file.add_pfn(curr_pfn, site='all') + curr_file.add_pfn(curr_pfn, site="all") else: - pfn_local = urljoin('file:', pathname2url(local_file_path)) - curr_file.add_pfn(pfn_local, 'local') + pfn_local = urljoin("file:", pathname2url(local_file_path)) + curr_file.add_pfn(pfn_local, "local") # Store the file to avoid later duplication tuple_val = (local_file_path, curr_file, curr_pfn) file_input_from_config_dict[curr_lfn] = tuple_val @@ -2202,7 +2363,7 @@ def configparser_value_to_file(cp, sec, opt, attrs=None): and option in the workflow configuration parser. Parameters - ----------- + ---------- cp : ConfigParser object The ConfigParser object holding the workflow configuration settings sec : string @@ -2212,9 +2373,10 @@ def configparser_value_to_file(cp, sec, opt, attrs=None): attrs : list to specify the 4 attributes of the file. Returns - -------- + ------- fileobj_from_path : workflow.File object obtained from the path specified by opt, within sec, in cp. + """ path = cp.get(sec, opt) fileobj_from_path = resolve_url_to_file(path, attrs=attrs) @@ -2227,14 +2389,15 @@ def get_full_analysis_chunk(science_segs): and return a single segment spanning that full time. Parameters - ----------- + ---------- science_segs : ifo-keyed dictionary of igwn_segments.segmentlist instances The list of times that are being analysed in this workflow. Returns - -------- + ------- fullSegment : igwn_segments.segment The segment spanning the first and last time point contained in science_segs. + """ extents = [science_segs[ifo].extent() for ifo in science_segs.keys()] min, max = extents[0] @@ -2251,8 +2414,9 @@ def get_random_label(): """ Get a random label string to use when clustering jobs. """ - return ''.join(random.choice(string.ascii_uppercase + string.digits) \ - for _ in range(15)) + return "".join( + random.choice(string.ascii_uppercase + string.digits) for _ in range(15) + ) def resolve_td_option(val_str, valid_seg): @@ -2275,12 +2439,12 @@ def resolve_td_option(val_str, valid_seg): The function will just return value_a. """ # Track if we've already found a matching option - output = '' + output = "" # Strip any whitespace, and split on comma - curr_vals = val_str.replace(' ', '').strip().split(',') + curr_vals = val_str.replace(" ", "").strip().split(",") # Resolving the simple case is trivial and can be done immediately. - if len(curr_vals) == 1 and '[' not in curr_vals[0]: + if len(curr_vals) == 1 and "[" not in curr_vals[0]: return curr_vals[0] # Loop over all possible values @@ -2288,28 +2452,28 @@ def resolve_td_option(val_str, valid_seg): start = int(valid_seg[0]) end = int(valid_seg[1]) # Extract limits for each case, and check overlap with valid_seg - if '[' in cval: - bopt = cval.split('[')[1].split(']')[0] - start, end = bopt.split(':') - cval = cval.replace('[' + bopt + ']', '') + if "[" in cval: + bopt = cval.split("[")[1].split("]")[0] + start, end = bopt.split(":") + cval = cval.replace("[" + bopt + "]", "") curr_seg = segments.segment(int(start), int(end)) # The segments module is a bit weird so we need to check if the two # overlap using the following code. If valid_seg is fully within # curr_seg this will be true. - if curr_seg.intersects(valid_seg) and \ - (curr_seg & valid_seg == valid_seg): + if curr_seg.intersects(valid_seg) and (curr_seg & valid_seg == valid_seg): if output: err_msg = "Time-dependent options must be disjoint." raise ValueError(err_msg) output = cval if not output: - err_msg = "Could not resolve option {}".format(val_str) + err_msg = f"Could not resolve option {val_str}" raise ValueError return output def add_workflow_settings_cli(parser, include_subdax_opts=False): - """Adds workflow options to an argument parser. + """ + Adds workflow options to an argument parser. Parameters ---------- @@ -2319,39 +2483,61 @@ def add_workflow_settings_cli(parser, include_subdax_opts=False): If True, will add output-map and dax-file-directory options to the parser. These can be used for workflows that are generated as a subdax of another workflow. Default is False. + """ wfgrp = parser.add_argument_group("Options for setting workflow files") - wfgrp.add_argument("--workflow-name", required=True, - help="Name of the workflow.") - wfgrp.add_argument("--tags", nargs="+", default=[], - help="Append the given tags to file names.") - wfgrp.add_argument("--output-dir", default=None, - help="Path to directory where the workflow will be " - "written. Default is to use " - "{workflow-name}_output.") - wfgrp.add_argument("--cache-file", default=None, - help="Path to input file containing list of files to " - "be reused (the 'input_map' file)") - wfgrp.add_argument("--plan-now", default=False, action='store_true', - help="If given, workflow will immediately be planned " - "on completion of workflow generation but not " - "submitted to the condor pool. A start script " - "will be created to submit to condor.") - wfgrp.add_argument("--submit-now", default=False, action='store_true', - help="If given, workflow will immediately be submitted " - "on completion of workflow generation") - wfgrp.add_argument("--dax-file", default=None, - help="Path to DAX file. Default is to write to the " - "output directory with name " - "{workflow-name}.dax.") + wfgrp.add_argument("--workflow-name", required=True, help="Name of the workflow.") + wfgrp.add_argument( + "--tags", nargs="+", default=[], help="Append the given tags to file names." + ) + wfgrp.add_argument( + "--output-dir", + default=None, + help="Path to directory where the workflow will be " + "written. Default is to use " + "{workflow-name}_output.", + ) + wfgrp.add_argument( + "--cache-file", + default=None, + help="Path to input file containing list of files to " + "be reused (the 'input_map' file)", + ) + wfgrp.add_argument( + "--plan-now", + default=False, + action="store_true", + help="If given, workflow will immediately be planned " + "on completion of workflow generation but not " + "submitted to the condor pool. A start script " + "will be created to submit to condor.", + ) + wfgrp.add_argument( + "--submit-now", + default=False, + action="store_true", + help="If given, workflow will immediately be submitted " + "on completion of workflow generation", + ) + wfgrp.add_argument( + "--dax-file", + default=None, + help="Path to DAX file. Default is to write to the " + "output directory with name " + "{workflow-name}.dax.", + ) if include_subdax_opts: - wfgrp.add_argument("--output-map", default=None, - help="Path to an output map file.") - wfgrp.add_argument("--dax-file-directory", default=None, - help="Put dax files (including output map, " - "sites.yml etc. in this directory. The use " - "case for this is when running a sub-workflow " - "under pegasus the outputs need to be copied " - "back to the appropriate directory, and " - "using this as --dax-file-directory . allows " - "that to be done.") + wfgrp.add_argument( + "--output-map", default=None, help="Path to an output map file." + ) + wfgrp.add_argument( + "--dax-file-directory", + default=None, + help="Put dax files (including output map, " + "sites.yml etc. in this directory. The use " + "case for this is when running a sub-workflow " + "under pegasus the outputs need to be copied " + "back to the appropriate directory, and " + "using this as --dax-file-directory . allows " + "that to be done.", + ) diff --git a/pycbc/workflow/datafind.py b/pycbc/workflow/datafind.py index 3690e4b1500..1cf4871e9dd 100644 --- a/pycbc/workflow/datafind.py +++ b/pycbc/workflow/datafind.py @@ -29,28 +29,29 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/ahope/datafind.html """ -import os, copy +import copy import logging +import os import urllib.parse import igwn_segments as segments -from igwn_ligolw import utils, ligolw from gwdatafind import find_urls as find_frame_urls +from igwn_ligolw import ligolw, utils -from pycbc.workflow.core import SegFile, File, FileList, make_analysis_dir from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.workflow.core import File, FileList, SegFile, make_analysis_dir # NOTE urllib is weird. For some reason it only allows known schemes and will # give *wrong* results, rather then failing, if you use something like gsiftp # We can add schemes explicitly, as below, but be careful with this! # (urllib is used indirectly through lal.Cache objects) -urllib.parse.uses_relative.append('osdf') -urllib.parse.uses_netloc.append('osdf') +urllib.parse.uses_relative.append("osdf") +urllib.parse.uses_netloc.append("osdf") -logger = logging.getLogger('pycbc.workflow.datafind') +logger = logging.getLogger("pycbc.workflow.datafind") -def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, - tags=None): + +def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, tags=None): """ Setup datafind section of the workflow. This section is responsible for generating, or setting up the workflow to generate, a list of files that @@ -87,7 +88,7 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, FIXME: Filenames may not be unique with current codes! Returns - -------- + ------- datafindOuts : OutGroupList List of all the datafind output files for use later in the pipeline. sci_avlble_file : SegFile @@ -101,6 +102,7 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, sci_avlble_name : string The name with which the analysable time is stored in the sci_avlble_file. + """ if tags is None: tags = [] @@ -109,52 +111,51 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, cp = workflow.cp # Parse for options in ini file - datafind_method = cp.get_opt_tags("workflow-datafind", - "datafind-method", tags) + datafind_method = cp.get_opt_tags("workflow-datafind", "datafind-method", tags) - if cp.has_option_tags("workflow-datafind", - "datafind-check-segment-gaps", tags): - checkSegmentGaps = cp.get_opt_tags("workflow-datafind", - "datafind-check-segment-gaps", tags) + if cp.has_option_tags("workflow-datafind", "datafind-check-segment-gaps", tags): + checkSegmentGaps = cp.get_opt_tags( + "workflow-datafind", "datafind-check-segment-gaps", tags + ) else: checkSegmentGaps = "no_test" - if cp.has_option_tags("workflow-datafind", - "datafind-check-frames-exist", tags): - checkFramesExist = cp.get_opt_tags("workflow-datafind", - "datafind-check-frames-exist", tags) + if cp.has_option_tags("workflow-datafind", "datafind-check-frames-exist", tags): + checkFramesExist = cp.get_opt_tags( + "workflow-datafind", "datafind-check-frames-exist", tags + ) else: checkFramesExist = "no_test" - if cp.has_option_tags("workflow-datafind", - "datafind-check-segment-summary", tags): - checkSegmentSummary = cp.get_opt_tags("workflow-datafind", - "datafind-check-segment-summary", tags) + if cp.has_option_tags("workflow-datafind", "datafind-check-segment-summary", tags): + checkSegmentSummary = cp.get_opt_tags( + "workflow-datafind", "datafind-check-segment-summary", tags + ) else: checkSegmentSummary = "no_test" logger.info("Starting datafind with setup_datafind_runtime_generated") if datafind_method == "AT_RUNTIME_MULTIPLE_CACHES": - datafindcaches, datafindouts = \ - setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, - outputDir, tags=tags) + datafindcaches, datafindouts = setup_datafind_runtime_cache_multi_calls_perifo( + cp, scienceSegs, outputDir, tags=tags + ) elif datafind_method == "AT_RUNTIME_SINGLE_CACHES": - datafindcaches, datafindouts = \ - setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, - outputDir, tags=tags) + datafindcaches, datafindouts = setup_datafind_runtime_cache_single_call_perifo( + cp, scienceSegs, outputDir, tags=tags + ) elif datafind_method == "AT_RUNTIME_MULTIPLE_FRAMES": - datafindcaches, datafindouts = \ - setup_datafind_runtime_frames_multi_calls_perifo(cp, scienceSegs, - outputDir, tags=tags) + datafindcaches, datafindouts = setup_datafind_runtime_frames_multi_calls_perifo( + cp, scienceSegs, outputDir, tags=tags + ) elif datafind_method == "AT_RUNTIME_SINGLE_FRAMES": - datafindcaches, datafindouts = \ - setup_datafind_runtime_frames_single_call_perifo(cp, scienceSegs, - outputDir, tags=tags) + datafindcaches, datafindouts = setup_datafind_runtime_frames_single_call_perifo( + cp, scienceSegs, outputDir, tags=tags + ) elif datafind_method == "AT_RUNTIME_FAKE_DATA": pass elif datafind_method == "FROM_PREGENERATED_LCF_FILES": ifos = scienceSegs.keys() - datafindcaches, datafindouts = \ - setup_datafind_from_pregenerated_lcf_files(cp, ifos, - outputDir, tags=tags) + datafindcaches, datafindouts = setup_datafind_from_pregenerated_lcf_files( + cp, ifos, outputDir, tags=tags + ) else: msg = """Entry datafind-method in [workflow-datafind] does not have " expected value. Valid values are @@ -165,29 +166,37 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, raise ValueError(msg) using_backup_server = False - if datafind_method == "AT_RUNTIME_MULTIPLE_FRAMES" or \ - datafind_method == "AT_RUNTIME_SINGLE_FRAMES": - if cp.has_option_tags("workflow-datafind", - "datafind-backup-datafind-server", tags): + if ( + datafind_method == "AT_RUNTIME_MULTIPLE_FRAMES" + or datafind_method == "AT_RUNTIME_SINGLE_FRAMES" + ): + if cp.has_option_tags( + "workflow-datafind", "datafind-backup-datafind-server", tags + ): using_backup_server = True - backup_server = cp.get_opt_tags("workflow-datafind", - "datafind-backup-datafind-server", tags) + backup_server = cp.get_opt_tags( + "workflow-datafind", "datafind-backup-datafind-server", tags + ) cp_new = copy.deepcopy(cp) - cp_new.set("workflow-datafind", - "datafind-ligo-datafind-server", backup_server) - cp_new.set('datafind', 'urltype', 'gsiftp') - backup_datafindcaches, backup_datafindouts =\ - setup_datafind_runtime_frames_single_call_perifo(cp_new, - scienceSegs, outputDir, tags=tags) - backup_datafindouts = datafind_keep_unique_backups(\ - backup_datafindouts, datafindouts) + cp_new.set( + "workflow-datafind", "datafind-ligo-datafind-server", backup_server + ) + cp_new.set("datafind", "urltype", "gsiftp") + backup_datafindcaches, backup_datafindouts = ( + setup_datafind_runtime_frames_single_call_perifo( + cp_new, scienceSegs, outputDir, tags=tags + ) + ) + backup_datafindouts = datafind_keep_unique_backups( + backup_datafindouts, datafindouts + ) datafindcaches.extend(backup_datafindcaches) datafindouts.extend(backup_datafindouts) logger.info("setup_datafind_runtime_generated completed") # If we don't have frame files covering all times we can update the science # segments. - if checkSegmentGaps in ['warn','update_times','raise_error']: + if checkSegmentGaps in ["warn", "update_times", "raise_error"]: logger.info("Checking science segments against datafind output....") newScienceSegs = get_science_segs_from_datafind_outs(datafindcaches) logger.info("New segments calculated from data find output.....") @@ -195,7 +204,7 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, for ifo in scienceSegs.keys(): # If no science segments in input then do nothing if not scienceSegs[ifo]: - msg = "No science segments are present for ifo %s, " %(ifo) + msg = "No science segments are present for ifo %s, " % (ifo) msg += "the segment metadata indicates there is no analyzable" msg += " strain data between the selected GPS start and end " msg += "times." @@ -203,27 +212,27 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, continue if ifo not in newScienceSegs: msg = "No data frames were found corresponding to the science " - msg += "segments for ifo %s" %(ifo) + msg += "segments for ifo %s" % (ifo) logger.error(msg) missingData = True - if checkSegmentGaps == 'update_times': + if checkSegmentGaps == "update_times": scienceSegs[ifo] = segments.segmentlist() continue missing = scienceSegs[ifo] - newScienceSegs[ifo] if abs(missing): - msg = "From ifo %s we are missing frames covering:" %(ifo) + msg = "From ifo %s we are missing frames covering:" % (ifo) msg += "\n%s" % "\n".join(map(str, missing)) missingData = True logger.error(msg) - if checkSegmentGaps == 'update_times': + if checkSegmentGaps == "update_times": # Remove missing time, so that we can carry on if desired logger.info("Updating science segments for ifo %s.", ifo) scienceSegs[ifo] = scienceSegs[ifo] - missing - if checkSegmentGaps == 'raise_error' and missingData: + if checkSegmentGaps == "raise_error" and missingData: raise ValueError("Workflow cannot find needed data, exiting.") logger.info("Done checking, any discrepancies are reported above.") - elif checkSegmentGaps == 'no_test': + elif checkSegmentGaps == "no_test": pass else: errMsg = "checkSegmentGaps kwarg must take a value from 'no_test', " @@ -231,10 +240,11 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, raise ValueError(errMsg) # Do all of the frame files that were returned actually exist? - if checkFramesExist in ['warn','update_times','raise_error']: + if checkFramesExist in ["warn", "update_times", "raise_error"]: logger.info("Verifying that all frames exist on disk.") - missingFrSegs, missingFrames = \ - get_missing_segs_from_frame_file_cache(datafindcaches) + missingFrSegs, missingFrames = get_missing_segs_from_frame_file_cache( + datafindcaches + ) missingFlag = False for ifo in missingFrames.keys(): # If no data in the input then do nothing @@ -256,29 +266,29 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, if len(dfout.pfns) == 0: new_list.append(frame) else: - msg = "Frame %s not found locally. "\ - %(frame.url,) - msg += "Replacing with remote url(s) %s." \ - %(str([a.url for a in dfout.pfns]),) + msg = "Frame %s not found locally. " % (frame.url,) + msg += "Replacing with remote url(s) %s." % ( + str([a.url for a in dfout.pfns]), + ) logger.info(msg) break else: new_list.append(frame) missingFrames[ifo] = new_list if missingFrames[ifo]: - msg = "From ifo %s we are missing the following frames:" %(ifo) - msg +='\n'.join([a.url for a in missingFrames[ifo]]) + msg = "From ifo %s we are missing the following frames:" % (ifo) + msg += "\n".join([a.url for a in missingFrames[ifo]]) missingFlag = True logger.error(msg) - if checkFramesExist == 'update_times': + if checkFramesExist == "update_times": # Remove missing times, so that we can carry on if desired logger.info("Updating science times for ifo %s.", ifo) scienceSegs[ifo] = scienceSegs[ifo] - missingFrSegs[ifo] - if checkFramesExist == 'raise_error' and missingFlag: + if checkFramesExist == "raise_error" and missingFlag: raise ValueError("Workflow cannot find all frames, exiting.") logger.info("Finished checking frames.") - elif checkFramesExist == 'no_test': + elif checkFramesExist == "no_test": pass else: errMsg = "checkFramesExist kwarg must take a value from 'no_test', " @@ -287,7 +297,7 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, # Check if there are cases where frames exist, but no entry in the segment # summary table are present. - if checkSegmentSummary in ['warn', 'raise_error']: + if checkSegmentSummary in ["warn", "raise_error"]: logger.info("Checking the segment summary table against frames.") dfScienceSegs = get_science_segs_from_datafind_outs(datafindcaches) missingFlag = False @@ -305,7 +315,7 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, seg_summary_times = seg_file.seg_summ_dict for ifo in dfScienceSegs.keys(): curr_seg_summ_times = seg_summary_times[ifo + ":" + sci_seg_name] - missing = (dfScienceSegs[ifo] & seg_file.valid_segments) + missing = dfScienceSegs[ifo] & seg_file.valid_segments missing.coalesce() missing = missing - curr_seg_summ_times missing.coalesce() @@ -316,22 +326,22 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, missing2 = missing2 - curr_seg_summ_times missing2.coalesce() if abs(missing): - msg = "From ifo %s the following times have frames, " %(ifo) + msg = "From ifo %s the following times have frames, " % (ifo) msg += "but are not covered in the segment summary table." msg += "\n%s" % "\n".join(map(str, missing)) logger.error(msg) missingFlag = True if abs(missing2): - msg = "From ifo %s the following times have frames, " %(ifo) + msg = "From ifo %s the following times have frames, " % (ifo) msg += "are science, and are not covered in the segment " msg += "summary table." msg += "\n%s" % "\n".join(map(str, missing2)) logger.error(msg) missingFlag = True - if checkSegmentSummary == 'raise_error' and missingFlag: + if checkSegmentSummary == "raise_error" and missingFlag: errMsg = "Segment_summary discrepancy detected, exiting." raise ValueError(errMsg) - elif checkSegmentSummary == 'no_test': + elif checkSegmentSummary == "no_test": pass else: errMsg = "checkSegmentSummary kwarg must take a value from 'no_test', " @@ -343,25 +353,30 @@ def setup_datafind_workflow(workflow, scienceSegs, outputDir, seg_file=None, # NOTE: Should this be overrideable in the config file? sci_avlble_name = "SCIENCE_AVAILABLE" for ifo in scienceSegs.keys(): - sci_avlble_dict[ifo + ':' + sci_avlble_name] = scienceSegs[ifo] - - sci_avlble_file = SegFile.from_segment_list_dict('SCIENCE_AVAILABLE', - sci_avlble_dict, ifo_list = scienceSegs.keys(), - valid_segment=workflow.analysis_time, - extension='.xml', tags=tags, directory=outputDir) + sci_avlble_dict[ifo + ":" + sci_avlble_name] = scienceSegs[ifo] + + sci_avlble_file = SegFile.from_segment_list_dict( + "SCIENCE_AVAILABLE", + sci_avlble_dict, + ifo_list=scienceSegs.keys(), + valid_segment=workflow.analysis_time, + extension=".xml", + tags=tags, + directory=outputDir, + ) logger.info("Leaving datafind module") if datafind_method == "AT_RUNTIME_FAKE_DATA": datafindouts = None else: - datafindouts = FileList(datafindouts) - + datafindouts = FileList(datafindouts) return datafindouts, sci_avlble_file, scienceSegs, sci_avlble_name -def setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, - outputDir, tags=None): +def setup_datafind_runtime_cache_multi_calls_perifo( + cp, scienceSegs, outputDir, tags=None +): """ This function uses the `gwdatafind` library to obtain the location of all the frame files that will be needed to cover the analysis of the data @@ -376,7 +391,7 @@ def setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, themselves.) Parameters - ----------- + ---------- cp : ConfigParser.ConfigParser instance This contains a representation of the information stored within the workflow configuration files @@ -394,7 +409,7 @@ def setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, FIXME: Filenames may not be unique with current codes! Returns - -------- + ------- datafindcaches : list of glue.lal.Cache instances The glue.lal.Cache representations of the various calls to the datafind server and the returned frame files. @@ -411,11 +426,12 @@ def setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, logger.info("Querying datafind server for all science segments.") for ifo, scienceSegsIfo in scienceSegs.items(): observatory = ifo[0].upper() - frameType = cp.get_opt_tags("workflow-datafind", - "datafind-%s-frame-type" % (ifo.lower()), tags) + frameType = cp.get_opt_tags( + "workflow-datafind", "datafind-%s-frame-type" % (ifo.lower()), tags + ) for seg in scienceSegsIfo: - msg = "Finding data between %d and %d " %(seg[0],seg[1]) - msg += "for ifo %s" %(ifo) + msg = "Finding data between %d and %d " % (seg[0], seg[1]) + msg += "for ifo %s" % (ifo) logger.info(msg) # WARNING: For now the workflow will expect times to be in integer seconds startTime = int(seg[0]) @@ -431,7 +447,7 @@ def setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, startTime, endTime, ifo, - tags=tags + tags=tags, ) except: cache, cache_file = run_datafind_instance( @@ -442,14 +458,16 @@ def setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, startTime, endTime, ifo, - tags=tags + tags=tags, ) datafindouts.append(cache_file) datafindcaches.append(cache) return datafindcaches, datafindouts -def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, - tags=None): + +def setup_datafind_runtime_cache_single_call_perifo( + cp, scienceSegs, outputDir, tags=None +): """ This function uses the `gwdatafind` library to obtain the location of all the frame files that will be needed to cover the analysis of the data @@ -464,7 +482,7 @@ def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, themselves.) Parameters - ----------- + ---------- cp : ConfigParser.ConfigParser instance This contains a representation of the information stored within the workflow configuration files @@ -482,7 +500,7 @@ def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, FIXME: Filenames may not be unique with current codes! Returns - -------- + ------- datafindcaches : list of glue.lal.Cache instances The glue.lal.Cache representations of the various calls to the datafind server and the returned frame files. @@ -496,7 +514,7 @@ def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, # We want to ignore gaps as the detectors go up and down and calling this # way will give gaps. See the setup_datafind_runtime_generated function # for datafind calls that only query for data that will exist - cp.set("datafind","on_gaps","ignore") + cp.set("datafind", "on_gaps", "ignore") # Now ready to loop over the input segments datafindouts = [] @@ -506,11 +524,10 @@ def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, observatory = ifo[0].upper() checked_times = segments.segmentlist([]) frame_types = cp.get_opt_tags( - "workflow-datafind", - "datafind-%s-frame-type" % (ifo.lower()), tags + "workflow-datafind", "datafind-%s-frame-type" % (ifo.lower()), tags ) # Check if this is one type, or time varying - frame_types = frame_types.replace(' ', '').strip().split(',') + frame_types = frame_types.replace(" ", "").strip().split(",") for ftype in frame_types: # Check the times, default to full time initially # This REQUIRES a coalesced segment list to work @@ -518,17 +535,17 @@ def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, end = int(scienceSegsIfo[-1][1]) # Then check for limits. We're expecting something like: # value[start:end], so need to extract value, start and end - if '[' in ftype: + if "[" in ftype: # This gets start and end out - bopt = ftype.split('[')[1].split(']')[0] - newstart, newend = bopt.split(':') + bopt = ftype.split("[")[1].split("]")[0] + newstart, newend = bopt.split(":") # Then check if the times are within science time start = max(int(newstart), start) end = min(int(newend), end) if end <= start: continue # This extracts value - ftype = ftype.split('[')[0] + ftype = ftype.split("[")[0] curr_times = segments.segment(start, end) # The times here must be distinct. We cannot have two different # frame files at the same time from the same ifo. @@ -540,33 +557,21 @@ def setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, outputDir, # Ask datafind where the frames are try: cache, cache_file = run_datafind_instance( - cp, - outputDir, - observatory, - ftype, - start, - end, - ifo, - tags=tags + cp, outputDir, observatory, ftype, start, end, ifo, tags=tags ) except: cache, cache_file = run_datafind_instance( - cp, - outputDir, - observatory, - ftype, - start, - end, - ifo, - tags=tags + cp, outputDir, observatory, ftype, start, end, ifo, tags=tags ) datafindouts.append(cache_file) datafindcaches.append(cache) return datafindcaches, datafindouts -def setup_datafind_runtime_frames_single_call_perifo(cp, scienceSegs, - outputDir, tags=None): + +def setup_datafind_runtime_frames_single_call_perifo( + cp, scienceSegs, outputDir, tags=None +): """ This function uses the `gwdatafind` library to obtain the location of all the frame files that will be needed to cover the analysis of the data @@ -581,7 +586,7 @@ def setup_datafind_runtime_frames_single_call_perifo(cp, scienceSegs, files as input. Parameters - ----------- + ---------- cp : ConfigParser.ConfigParser instance This contains a representation of the information stored within the workflow configuration files @@ -599,7 +604,7 @@ def setup_datafind_runtime_frames_single_call_perifo(cp, scienceSegs, FIXME: Filenames may not be unique with current codes! Returns - -------- + ------- datafindcaches : list of glue.lal.Cache instances The glue.lal.Cache representations of the various calls to the datafind server and the returned frame files. @@ -607,16 +612,18 @@ def setup_datafind_runtime_frames_single_call_perifo(cp, scienceSegs, List of all the datafind output files for use later in the pipeline. """ - datafindcaches, _ = \ - setup_datafind_runtime_cache_single_call_perifo(cp, scienceSegs, - outputDir, tags=tags) + datafindcaches, _ = setup_datafind_runtime_cache_single_call_perifo( + cp, scienceSegs, outputDir, tags=tags + ) datafindouts = convert_cachelist_to_filelist(datafindcaches) return datafindcaches, datafindouts -def setup_datafind_runtime_frames_multi_calls_perifo(cp, scienceSegs, - outputDir, tags=None): + +def setup_datafind_runtime_frames_multi_calls_perifo( + cp, scienceSegs, outputDir, tags=None +): """ This function uses the `gwdatafind` library to obtain the location of all the frame files that will be needed to cover the analysis of the data @@ -631,7 +638,7 @@ def setup_datafind_runtime_frames_multi_calls_perifo(cp, scienceSegs, files as input. Parameters - ----------- + ---------- cp : ConfigParser.ConfigParser instance This contains a representation of the information stored within the workflow configuration files @@ -649,7 +656,7 @@ def setup_datafind_runtime_frames_multi_calls_perifo(cp, scienceSegs, FIXME: Filenames may not be unique with current codes! Returns - -------- + ------- datafindcaches : list of glue.lal.Cache instances The glue.lal.Cache representations of the various calls to the datafind server and the returned frame files. @@ -657,21 +664,22 @@ def setup_datafind_runtime_frames_multi_calls_perifo(cp, scienceSegs, List of all the datafind output files for use later in the pipeline. """ - datafindcaches, _ = \ - setup_datafind_runtime_cache_multi_calls_perifo(cp, scienceSegs, - outputDir, tags=tags) + datafindcaches, _ = setup_datafind_runtime_cache_multi_calls_perifo( + cp, scienceSegs, outputDir, tags=tags + ) datafindouts = convert_cachelist_to_filelist(datafindcaches) return datafindcaches, datafindouts + def setup_datafind_from_pregenerated_lcf_files(cp, ifos, outputDir, tags=None): """ This function is used if you want to run with pregenerated lcf frame cache files. Parameters - ----------- + ---------- cp : ConfigParser.ConfigParser instance This contains a representation of the information stored within the workflow configuration files @@ -688,12 +696,13 @@ def setup_datafind_from_pregenerated_lcf_files(cp, ifos, outputDir, tags=None): the Files and uniqueify the actual filename. Returns - -------- + ------- datafindcaches : list of glue.lal.Cache instances The glue.lal.Cache representations of the various calls to the datafind server and the returned frame files. datafindOuts : pycbc.workflow.core.FileList List of all the datafind output files for use later in the pipeline. + """ from glue import lal @@ -702,31 +711,35 @@ def setup_datafind_from_pregenerated_lcf_files(cp, ifos, outputDir, tags=None): datafindcaches = [] for ifo in ifos: - search_string = "datafind-pregenerated-cache-file-%s" %(ifo.lower(),) - frame_cache_file_name = cp.get_opt_tags("workflow-datafind", - search_string, tags=tags) - curr_cache = lal.Cache.fromfilenames([frame_cache_file_name], - coltype=lal.LIGOTimeGPS) + search_string = "datafind-pregenerated-cache-file-%s" % (ifo.lower(),) + frame_cache_file_name = cp.get_opt_tags( + "workflow-datafind", search_string, tags=tags + ) + curr_cache = lal.Cache.fromfilenames( + [frame_cache_file_name], coltype=lal.LIGOTimeGPS + ) curr_cache.ifo = ifo datafindcaches.append(curr_cache) datafindouts = convert_cachelist_to_filelist(datafindcaches) return datafindcaches, datafindouts + def convert_cachelist_to_filelist(datafindcache_list): """ Take as input a list of glue.lal.Cache objects and return a pycbc FileList containing all frames within those caches. Parameters - ----------- + ---------- datafindcache_list : list of glue.lal.Cache objects The list of cache files to convert. Returns - -------- + ------- datafind_filelist : FileList of frame File objects The list of frame files. + """ prev_file = None prev_name = None @@ -740,10 +753,10 @@ def convert_cachelist_to_filelist(datafindcache_list): curr_ifo = cache.ifo for frame in cache: # Pegasus doesn't like "localhost" in URLs. - frame.url = frame.url.replace('file://localhost', 'file://') + frame.url = frame.url.replace("file://localhost", "file://") # Not sure why it happens in OSDF URLs!! # May need to remove use of Cache objects - frame.url = frame.url.replace('osdf://localhost/', 'osdf:///') + frame.url = frame.url.replace("osdf://localhost/", "osdf:///") # Create one File() object for each unique frame file that we # get back in the cache. @@ -752,28 +765,33 @@ def convert_cachelist_to_filelist(datafindcache_list): this_name = os.path.basename(frame.url) if (prev_file is None) or (prev_name != this_name): - currFile = File(curr_ifo, frame.description, - frame.segment, file_url=frame.url, use_tmp_subdirs=True) + currFile = File( + curr_ifo, + frame.description, + frame.segment, + file_url=frame.url, + use_tmp_subdirs=True, + ) datafind_filelist.append(currFile) prev_file = currFile # Populate the PFNs for the File() we just created - cvmfs_urls = ('file:///cvmfs/', 'osdf://') + cvmfs_urls = ("file:///cvmfs/", "osdf://") if frame.url.startswith(cvmfs_urls): # Frame is on CVMFS/OSDF, so let all sites read it directly. - currFile.add_pfn(frame.url, site='all') - elif frame.url.startswith('file://'): + currFile.add_pfn(frame.url, site="all") + elif frame.url.startswith("file://"): # Frame not on CVMFS, so may need transferring. # Be careful here! If all your frames files are on site # = local and you try to run on OSG, it will likely # overwhelm the condor file transfer process! - for site in ['local', 'condorpool_shared']: + for site in ["local", "condorpool_shared"]: currFile.add_pfn(frame.url, site=site) else: # Frame is at some unknown URL. Pegasus will decide how to deal # with this, but will likely transfer to local site first, and # from there transfer to remote sites as needed. - currFile.add_pfn(frame.url, site='notlocal') + currFile.add_pfn(frame.url, site="notlocal") return datafind_filelist @@ -791,9 +809,10 @@ def get_science_segs_from_datafind_outs(datafindcaches): List of all the datafind output files. Returns - -------- + ------- newScienceSegs : Dictionary of ifo keyed igwn_segments.segmentlist instances The times covered by the frames found in datafindOuts. + """ newScienceSegs = {} for cache in datafindcaches: @@ -807,6 +826,7 @@ def get_science_segs_from_datafind_outs(datafindcaches): newScienceSegs[ifo].coalesce() return newScienceSegs + def get_missing_segs_from_frame_file_cache(datafindcaches): """ This function will use os.path.isfile to determine if all the frame files @@ -814,16 +834,17 @@ def get_missing_segs_from_frame_file_cache(datafindcaches): then be used to update the science times if needed. Parameters - ----------- + ---------- datafindcaches : OutGroupList List of all the datafind output files. Returns - -------- + ------- missingFrameSegs : Dict. of ifo keyed igwn_segments.segmentlist instances The times corresponding to missing frames found in datafindOuts. missingFrames: Dict. of ifo keyed lal.Cache instances The list of missing frames + """ from glue import lal @@ -833,14 +854,15 @@ def get_missing_segs_from_frame_file_cache(datafindcaches): if len(cache) > 0: # Don't bother if these are not file:// urls, assume all urls in # one cache file must be the same type - if not cache[0].scheme == 'file': - warn_msg = "We have %s entries in the " %(cache[0].scheme,) + if not cache[0].scheme == "file": + warn_msg = "We have %s entries in the " % (cache[0].scheme,) warn_msg += "cache file. I do not check if these exist." logger.warning(warn_msg) continue _, currMissingFrames = cache.checkfilesexist(on_missing="warn") - missingSegs = segments.segmentlist(e.segment \ - for e in currMissingFrames).coalesce() + missingSegs = segments.segmentlist( + e.segment for e in currMissingFrames + ).coalesce() ifo = cache.ifo if ifo not in missingFrameSegs: missingFrameSegs[ifo] = missingSegs @@ -860,44 +882,46 @@ def get_segment_summary_times(scienceFile, segmentName): for the flag given by segmentName. Parameters - ----------- + ---------- scienceFile : SegFile The segment file that we want to use to determine this. segmentName : string The DQ flag to search for times in the segment_summary table. Returns - --------- + ------- summSegList : igwn_segments.segmentlist The times that are covered in the segment summary table. + """ # Parse the segmentName - segmentName = segmentName.split(':') - if not len(segmentName) in [2, 3]: + segmentName = segmentName.split(":") + if len(segmentName) not in [2, 3]: raise ValueError(f"Invalid channel name {segmentName}.") ifo = segmentName[0] channel = segmentName[1] - version = '' + version = "" if len(segmentName) == 3: version = int(segmentName[2]) # Load the filename xmldoc = utils.load_filename( scienceFile.cache_entry.path, - compress='auto', - contenthandler=LIGOLWContentHandler + compress="auto", + contenthandler=LIGOLWContentHandler, ) # Get the segment_def_id for the segmentName segmentDefTable = ligolw.Table.get_table(xmldoc, "segment_definer") for entry in segmentDefTable: if (entry.ifos == ifo) and (entry.name == channel): - if len(segmentName) == 2 or (entry.version==version): + if len(segmentName) == 2 or (entry.version == version): segDefID = entry.segment_def_id break else: - raise ValueError("Cannot find channel %s in segment_definer table."\ - %(segmentName)) + raise ValueError( + "Cannot find channel %s in segment_definer table." % (segmentName) + ) # Get the segmentlist corresponding to this segmentName in segment_summary segmentSummTable = ligolw.Table.get_table(xmldoc, "segment_summary") @@ -910,8 +934,10 @@ def get_segment_summary_times(scienceFile, segmentName): return summSegList -def run_datafind_instance(cp, outputDir, observatory, frameType, - startTime, endTime, ifo, tags=None): + +def run_datafind_instance( + cp, outputDir, observatory, frameType, startTime, endTime, ifo, tags=None +): """ This function will query the datafind server once to find frames between the specified times for the specified frame type and observatory. @@ -945,7 +971,7 @@ def run_datafind_instance(cp, outputDir, observatory, frameType, FIXME: Filenames may not be unique with current codes! Returns - -------- + ------- dfCache : glue.lal.Cache instance The glue.lal.Cache representation of the call to the datafind server and the returned frame files. @@ -959,12 +985,9 @@ def run_datafind_instance(cp, outputDir, observatory, frameType, tags = [] # Determine if we should override the default datafind server - if cp.has_option_tags("workflow-datafind", - "datafind-ligo-datafind-server", tags): + if cp.has_option_tags("workflow-datafind", "datafind-ligo-datafind-server", tags): datafind_server = cp.get_opt_tags( - "workflow-datafind", - "datafind-ligo-datafind-server", - tags + "workflow-datafind", "datafind-ligo-datafind-server", tags ) else: datafind_server = None @@ -974,35 +997,37 @@ def run_datafind_instance(cp, outputDir, observatory, frameType, # given). dfKwargs = {} # By default ignore missing frames, this case is dealt with outside of here - dfKwargs['on_gaps'] = 'ignore' + dfKwargs["on_gaps"] = "ignore" if cp.has_section("datafind"): for item, value in cp.items("datafind"): dfKwargs[item] = value for tag in tags: - if cp.has_section('datafind-%s' %(tag)): - for item, value in cp.items("datafind-%s" %(tag)): + if cp.has_section("datafind-%s" % (tag)): + for item, value in cp.items("datafind-%s" % (tag)): dfKwargs[item] = value # It is useful to print the corresponding command to the logs # directory to check if this was expected. - log_datafind_command(observatory, frameType, startTime, endTime, - os.path.join(outputDir,'logs'), **dfKwargs) + log_datafind_command( + observatory, + frameType, + startTime, + endTime, + os.path.join(outputDir, "logs"), + **dfKwargs, + ) logger.debug("Asking datafind server for frames.") dfCache = lal.Cache.from_urls( find_frame_urls( - observatory, - frameType, - startTime, - endTime, - host=datafind_server, - **dfKwargs + observatory, frameType, startTime, endTime, host=datafind_server, **dfKwargs ), ) logger.debug("Frames returned") # workflow format output file - cache_file = File(ifo, 'DATAFIND', seg, extension='lcf', - directory=outputDir, tags=tags) - cache_file.add_pfn(cache_file.cache_entry.path, site='local') + cache_file = File( + ifo, "DATAFIND", seg, extension="lcf", directory=outputDir, tags=tags + ) + cache_file.add_pfn(cache_file.cache_entry.path, site="local") dfCache.ifo = ifo # Dump output to file @@ -1010,20 +1035,24 @@ def run_datafind_instance(cp, outputDir, observatory, frameType, # FIXME: CANNOT use dfCache.tofile because it will print 815901601.00000 # as a gps time which is incompatible with the lal cache format # (and the C codes) which demand an integer. - #dfCache.tofile(fP) + # dfCache.tofile(fP) for entry in dfCache: start = str(int(entry.segment[0])) duration = str(int(abs(entry.segment))) - print("%s %s %s %s %s" \ - % (entry.observatory, entry.description, start, duration, entry.url), file=fP) + print( + "%s %s %s %s %s" + % (entry.observatory, entry.description, start, duration, entry.url), + file=fP, + ) entry.segment = segments.segment(int(entry.segment[0]), int(entry.segment[1])) fP.close() return dfCache, cache_file -def log_datafind_command(observatory, frameType, startTime, endTime, - outputDir, **dfKwargs): +def log_datafind_command( + observatory, frameType, startTime, endTime, outputDir, **dfKwargs +): """ This command will print an equivalent gw_data_find command to disk that can be used to debug why the internal datafind module is not working. @@ -1031,34 +1060,47 @@ def log_datafind_command(observatory, frameType, startTime, endTime, # FIXME: This does not accurately reproduce the call as assuming the # kwargs will be the same is wrong, so some things need to be converted # "properly" to the command line equivalent. - gw_command = ['gw_data_find', '--observatory', observatory, - '--type', frameType, - '--gps-start-time', str(startTime), - '--gps-end-time', str(endTime)] + gw_command = [ + "gw_data_find", + "--observatory", + observatory, + "--type", + frameType, + "--gps-start-time", + str(startTime), + "--gps-end-time", + str(endTime), + ] for name, value in dfKwargs.items(): - if name == 'match': + if name == "match": gw_command.append("--match") gw_command.append(str(value)) - elif name == 'urltype': + elif name == "urltype": gw_command.append("--url-type") gw_command.append(str(value)) - elif name == 'on_gaps': + elif name == "on_gaps": pass else: - errMsg = "Unknown datafind kwarg given: %s. " %(name) - errMsg+= "This argument is stripped in the logged .sh command." + errMsg = "Unknown datafind kwarg given: %s. " % (name) + errMsg += "This argument is stripped in the logged .sh command." logger.warning(errMsg) - fileName = "%s-%s-%d-%d.sh" \ - %(observatory, frameType, startTime, endTime-startTime) + fileName = "%s-%s-%d-%d.sh" % ( + observatory, + frameType, + startTime, + endTime - startTime, + ) filePath = os.path.join(outputDir, fileName) - fP = open(filePath, 'w') - fP.write(' '.join(gw_command)) + fP = open(filePath, "w") + fP.write(" ".join(gw_command)) fP.close() + def datafind_keep_unique_backups(backup_outs, orig_outs): - """This function will take a list of backup datafind files, presumably + """ + This function will take a list of backup datafind files, presumably obtained by querying a remote datafind server, e.g. CIT, and compares these against a list of original datafind files, presumably obtained by querying the local datafind server. Only the datafind files in the backup @@ -1066,16 +1108,17 @@ def datafind_keep_unique_backups(backup_outs, orig_outs): to use only files that are missing from the local cluster. Parameters - ----------- + ---------- backup_outs : FileList List of datafind files from the remote datafind server. orig_outs : FileList List of datafind files from the local datafind server. Returns - -------- + ------- FileList List of datafind files in backup_outs and not in orig_outs. + """ # NOTE: This function is not optimized and could be made considerably # quicker if speed becomes in issue. With 4s frame files this might @@ -1093,7 +1136,7 @@ def datafind_keep_unique_backups(backup_outs, orig_outs): orig_out = orig_outs[index_num] pfns = list(file.pfns) # This shouldn't happen, but catch if it does - assert(len(pfns) == 1) - orig_out.add_pfn(pfns[0].url, site='notlocal') + assert len(pfns) == 1 + orig_out.add_pfn(pfns[0].url, site="notlocal") return return_list diff --git a/pycbc/workflow/dq.py b/pycbc/workflow/dq.py index 361a448e5e6..ec293dd12c0 100644 --- a/pycbc/workflow/dq.py +++ b/pycbc/workflow/dq.py @@ -23,9 +23,10 @@ # import logging -from pycbc.workflow.core import (FileList, Executable, Node, make_analysis_dir) -logger = logging.getLogger('pycbc.workflow.dq') +from pycbc.workflow.core import Executable, FileList, Node, make_analysis_dir + +logger = logging.getLogger("pycbc.workflow.dq") class PyCBCBinTemplatesDQExecutable(Executable): @@ -33,36 +34,46 @@ class PyCBCBinTemplatesDQExecutable(Executable): def create_node(self, workflow, ifo, template_bank_file): node = Node(self) - node.add_opt('--ifo', ifo) - node.add_input_opt('--bank-file', template_bank_file) - node.new_output_file_opt( - workflow.analysis_time, '.hdf', '--output-file') + node.add_opt("--ifo", ifo) + node.add_input_opt("--bank-file", template_bank_file) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") return node class PyCBCBinTriggerRatesDQExecutable(Executable): current_retention_level = Executable.MERGED_TRIGGERS - def create_node(self, workflow, flag_file, flag_name, - analysis_segment_file, analysis_segment_name, - trig_file, template_bins_file): + def create_node( + self, + workflow, + flag_file, + flag_name, + analysis_segment_file, + analysis_segment_name, + trig_file, + template_bins_file, + ): node = Node(self) - node.add_input_opt('--template-bins-file', template_bins_file) - node.add_input_opt('--trig-file', trig_file) - node.add_input_opt('--flag-file', flag_file) - node.add_opt('--flag-name', flag_name) - node.add_input_opt('--analysis-segment-file', analysis_segment_file) - node.add_opt('--analysis-segment-name', analysis_segment_name) - node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file') + node.add_input_opt("--template-bins-file", template_bins_file) + node.add_input_opt("--trig-file", trig_file) + node.add_input_opt("--flag-file", flag_file) + node.add_opt("--flag-name", flag_name) + node.add_input_opt("--analysis-segment-file", analysis_segment_file) + node.add_opt("--analysis-segment-name", analysis_segment_name) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") return node -def setup_dq_reranking(workflow, insps, bank, - analyzable_seg_file, - analyzable_name, - dq_seg_file, - output_dir=None, tags=None): +def setup_dq_reranking( + workflow, + insps, + bank, + analyzable_seg_file, + analyzable_name, + dq_seg_file, + output_dir=None, + tags=None, +): logger.info("Setting up dq reranking") make_analysis_dir(output_dir) output_files = FileList() @@ -70,18 +81,21 @@ def setup_dq_reranking(workflow, insps, bank, if tags is None: tags = [] - dq_labels = workflow.cp.get_subsections('workflow-data_quality') + dq_labels = workflow.cp.get_subsections("workflow-data_quality") dq_ifos = {} dq_names = {} dq_types = {} for dql in dq_labels: dq_ifos[dql] = workflow.cp.get_opt_tags( - 'workflow-data_quality', 'dq-ifo', [dql]) + "workflow-data_quality", "dq-ifo", [dql] + ) dq_names[dql] = workflow.cp.get_opt_tags( - 'workflow-data_quality', 'dq-name', [dql]) + "workflow-data_quality", "dq-name", [dql] + ) dq_types[dql] = workflow.cp.get_opt_tags( - 'workflow-data_quality', 'dq-type', [dql]) + "workflow-data_quality", "dq-type", [dql] + ) ifos = set(dq_ifos.values()) & set(workflow.ifos) @@ -97,22 +111,18 @@ def setup_dq_reranking(workflow, insps, bank, # get triggers for this ifo ifo_insp = [insp for insp in insps if (insp.ifo == ifo)] - assert len(ifo_insp) == 1, \ - f"Received more than one inspiral file for {ifo}" + assert len(ifo_insp) == 1, f"Received more than one inspiral file for {ifo}" ifo_insp = ifo_insp[0] # calculate template bins for this ifo bin_templates_exe = PyCBCBinTemplatesDQExecutable( - workflow.cp, - 'bin_templates', - ifos=ifo, - out_dir=output_dir, - tags=tags) + workflow.cp, "bin_templates", ifos=ifo, out_dir=output_dir, tags=tags + ) bin_templates_node = bin_templates_exe.create_node(workflow, ifo, bank) workflow += bin_templates_node template_bins_file = bin_templates_node.output_file - if dq_type == 'flag': + if dq_type == "flag": flag_file = dq_seg_file flag_name = dq_name else: @@ -121,10 +131,11 @@ def setup_dq_reranking(workflow, insps, bank, # calculate trigger rates during dq flags bin_triggers_exe = PyCBCBinTriggerRatesDQExecutable( workflow.cp, - 'bin_trigger_rates_dq', + "bin_trigger_rates_dq", ifos=ifo, out_dir=output_dir, - tags=dq_tags) + tags=dq_tags, + ) bin_triggers_node = bin_triggers_exe.create_node( workflow, flag_file, @@ -132,7 +143,8 @@ def setup_dq_reranking(workflow, insps, bank, analyzable_seg_file, analyzable_name, ifo_insp, - template_bins_file) + template_bins_file, + ) workflow += bin_triggers_node output_files += bin_triggers_node.output_files output_labels += [dq_label] diff --git a/pycbc/workflow/grb_utils.py b/pycbc/workflow/grb_utils.py index f478ac828aa..54d9d079a30 100644 --- a/pycbc/workflow/grb_utils.py +++ b/pycbc/workflow/grb_utils.py @@ -28,21 +28,20 @@ http://pycbc.org/pycbc/latest/html/workflow.html """ -import os import logging +import os + import numpy as np -from scipy.stats import rayleigh from gwdatafind.utils import filename_metadata +from scipy.stats import rayleigh from pycbc import makedir -from pycbc.workflow.core import \ - File, FileList, resolve_url_to_file, \ - Executable, Node +from pycbc.workflow.core import Executable, File, FileList, Node, resolve_url_to_file from pycbc.workflow.jobsetup import select_generic_executable from pycbc.workflow.pegasus_workflow import SubWorkflow from pycbc.workflow.plotting import PlotExecutable -logger = logging.getLogger('pycbc.workflow.grb_utils') +logger = logging.getLogger("pycbc.workflow.grb_utils") def _select_grb_pp_class(wflow, curr_exe): @@ -62,13 +61,14 @@ def _select_grb_pp_class(wflow, curr_exe): * job.create_node() and * job.get_valid_times(ifo, ) + """ - exe_path = wflow.cp.get('executables', curr_exe) + exe_path = wflow.cp.get("executables", curr_exe) exe_name = os.path.basename(exe_path) exe_to_class_map = { - 'pycbc_grb_trig_combiner': PycbcGrbTrigCombinerExecutable, - 'pycbc_grb_trig_cluster': PycbcGrbTrigClusterExecutable, - 'pycbc_grb_inj_finder': PycbcGrbInjFinderExecutable + "pycbc_grb_trig_combiner": PycbcGrbTrigCombinerExecutable, + "pycbc_grb_trig_cluster": PycbcGrbTrigClusterExecutable, + "pycbc_grb_inj_finder": PycbcGrbInjFinderExecutable, } if exe_name not in exe_to_class_map: raise ValueError(f"No job class exists for executable {curr_exe}") @@ -92,7 +92,7 @@ def set_grb_start_end(cp, start, end): The end of the workflow analysis time. Returns - -------- + ------- cp : pycbc.workflow.configuration.WorkflowConfigParser object The modified WorkflowConfigParser object. @@ -104,7 +104,7 @@ def set_grb_start_end(cp, start, end): def make_gating_node(workflow, datafind_files, outdir=None, tags=None): - ''' + """ Generate jobs for autogating the data for PyGRB runs. Parameters @@ -120,32 +120,36 @@ def make_gating_node(workflow, datafind_files, outdir=None, tags=None): that would be produced in multiple calls to this function. Returns - -------- + ------- condition_strain_nodes : list List containing the pycbc.workflow.core.Node objects representing the autogating jobs. condition_strain_outs : pycbc.workflow.core.FileList FileList containing the pycbc.workflow.core.File objects representing the gated frame files. - ''' + """ cp = workflow.cp if tags is None: tags = [] - condition_strain_class = select_generic_executable(workflow, - "condition_strain") + condition_strain_class = select_generic_executable(workflow, "condition_strain") condition_strain_nodes = [] condition_strain_outs = FileList([]) for ifo in workflow.ifos: - input_files = FileList([datafind_file for datafind_file in - datafind_files if datafind_file.ifo == ifo]) - condition_strain_jobs = condition_strain_class(cp, "condition_strain", - ifos=ifo, - out_dir=outdir, - tags=tags) - condition_strain_node, condition_strain_out = \ - condition_strain_jobs.create_node(input_files, tags=tags) + input_files = FileList( + [ + datafind_file + for datafind_file in datafind_files + if datafind_file.ifo == ifo + ] + ) + condition_strain_jobs = condition_strain_class( + cp, "condition_strain", ifos=ifo, out_dir=outdir, tags=tags + ) + condition_strain_node, condition_strain_out = condition_strain_jobs.create_node( + input_files, tags=tags + ) condition_strain_nodes.append(condition_strain_node) condition_strain_outs.extend(FileList([condition_strain_out])) @@ -153,8 +157,10 @@ def make_gating_node(workflow, datafind_files, outdir=None, tags=None): def fermi_core_tail_model( - sky_err, rad, core_frac=0.98, core_sigma=3.6, tail_sigma=29.6): - """Fermi systematic error model following + sky_err, rad, core_frac=0.98, core_sigma=3.6, tail_sigma=29.6 +): + """ + Fermi systematic error model following https://arxiv.org/abs/1909.03006, with default values valid before 11 September 2019. @@ -169,21 +175,28 @@ def fermi_core_tail_model( Size of the GBM systematic tail component. Returns + ------- _______ tuple Tuple containing the core and tail probability distributions as a function of radius. + """ scaledsq = sky_err**2 / -2 / np.log(0.32) return ( - frac * (1 - np.exp(-0.5 * (rad / np.sqrt(scaledsq + sigma**2))**2)) - for frac, sigma - in zip([core_frac, 1 - core_frac], [core_sigma, tail_sigma])) + frac * (1 - np.exp(-0.5 * (rad / np.sqrt(scaledsq + sigma**2)) ** 2)) + for frac, sigma in zip([core_frac, 1 - core_frac], [core_sigma, tail_sigma]) + ) def get_sky_grid_scale( - sky_error=0.0, containment=0.9, upscale=False, fermi_sys=False, - precision=1e-3, **kwargs): + sky_error=0.0, + containment=0.9, + upscale=False, + fermi_sys=False, + precision=1e-3, + **kwargs, +): """ Calculate the angular radius corresponding to a desired localization uncertainty level. This is used to generate the search @@ -214,16 +227,20 @@ def get_sky_grid_scale( Additional keyword arguments passed to `fermi_core_tail_model`. Returns + ------- _______ float Sky error radius in degrees. + """ if fermi_sys: lims = (0.5, 4) radii = np.linspace( - lims[0] * sky_error, lims[1] * sky_error, - int((lims[1] - lims[0]) * sky_error / precision) + 1) + lims[0] * sky_error, + lims[1] * sky_error, + int((lims[1] - lims[0]) * sky_error / precision) + 1, + ) core, tail = fermi_core_tail_model(sky_error, radii, **kwargs) out = radii[(abs(core + tail - containment)).argmin()] else: @@ -240,18 +257,22 @@ def get_sky_grid_scale( def make_skygrid_node(workflow, out_dir, tags=None): """ Adds a job to the workflow to produce the PyGRB search skygrid.""" - tags = [] if tags is None else tags # Initialize job node - grb_name = workflow.cp.get('workflow', 'trigger-name') - extra_tags = ['GRB'+grb_name] - node = Executable(workflow.cp, 'make_sky_grid', - ifos=workflow.ifos, out_dir=out_dir, - tags=tags+extra_tags).create_node() - node.add_opt('--instruments', ' '.join(workflow.ifos)) - node.new_output_file_opt(workflow.analysis_time, '.h5', '--output', - tags=extra_tags, store_file=True) + grb_name = workflow.cp.get("workflow", "trigger-name") + extra_tags = ["GRB" + grb_name] + node = Executable( + workflow.cp, + "make_sky_grid", + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags + extra_tags, + ).create_node() + node.add_opt("--instruments", " ".join(workflow.ifos)) + node.new_output_file_opt( + workflow.analysis_time, ".h5", "--output", tags=extra_tags, store_file=True + ) # Add job node to the workflow workflow += node @@ -271,8 +292,8 @@ def generate_tc_prior(wflow, tc_path, buffer_seg): Path where the configuration file for the prior needs to be written. buffer_seg : segmentlist Start and end times of the buffer segment encapsulating the onsource. - """ + """ # Write the tc-prior configuration file if it does not exist if os.path.exists(tc_path): raise ValueError("Refusing to overwrite %s." % tc_path) @@ -283,26 +304,33 @@ def generate_tc_prior(wflow, tc_path, buffer_seg): tc_file.write("max-tc = %s\n\n" % wflow.analysis_time[1]) tc_file.write("[constraint-tc]\n") tc_file.write("name = custom\n") - tc_file.write("constraint_arg = (tc < %s) | (tc > %s)\n" % - (buffer_seg[0], buffer_seg[1])) + tc_file.write( + "constraint_arg = (tc < %s) | (tc > %s)\n" % (buffer_seg[0], buffer_seg[1]) + ) tc_file.close() # Add the tc-prior configuration file url to wflow.cp if necessary - tc_file_path = "file://"+tc_path + tc_file_path = "file://" + tc_path for inj_sec in wflow.cp.get_subsections("injections"): - config_urls = wflow.cp.get("workflow-injections", - inj_sec+"-config-files") + config_urls = wflow.cp.get("workflow-injections", inj_sec + "-config-files") config_urls = [url.strip() for url in config_urls.split(",")] if tc_file_path not in config_urls: config_urls += [tc_file_path] - config_urls = ', '.join([str(item) for item in config_urls]) - wflow.cp.set("workflow-injections", - inj_sec+"-config-files", - config_urls) - - -def setup_pygrb_pp_workflow(wf, pp_dir, seg_dir, segment, bank_file, - insp_files, inj_files, inj_insp_files, inj_tags): + config_urls = ", ".join([str(item) for item in config_urls]) + wflow.cp.set("workflow-injections", inj_sec + "-config-files", config_urls) + + +def setup_pygrb_pp_workflow( + wf, + pp_dir, + seg_dir, + segment, + bank_file, + insp_files, + inj_files, + inj_insp_files, + inj_tags, +): """ Generate post-processing section of PyGRB offline workflow @@ -329,18 +357,16 @@ def setup_pygrb_pp_workflow(wf, pp_dir, seg_dir, segment, bank_file, Contains triggers after clustering inj_find_files : FileList FOUNDMISSED FileList covering all injection sets + """ # Begin setting up trig combiner job(s) # Select executable class and initialize exe_class = _select_grb_pp_class(wf, "trig_combiner") job_instance = exe_class(wf.cp, "trig_combiner") # Create node for coherent no injections jobs - node, trig_files = job_instance.create_node(wf.ifo_string, - seg_dir, - segment, - insp_files, - pp_dir, - bank_file) + node, trig_files = job_instance.create_node( + wf.ifo_string, seg_dir, segment, insp_files, pp_dir, bank_file + ) wf.add_node(node) # Trig clustering for each trig file @@ -358,16 +384,14 @@ def setup_pygrb_pp_workflow(wf, pp_dir, seg_dir, segment, bank_file, job_instance = exe_class(wf.cp, "inj_finder") inj_find_files = FileList([]) for inj_tag in inj_tags: - tag_inj_files = FileList([f for f in inj_files - if inj_tag in f.tags]) + tag_inj_files = FileList([f for f in inj_files if inj_tag in f.tags]) # The here stems from the injection group information # being stored in the second tag. This could be improved # depending on the final implementation of injections - tag_insp_files = FileList([f for f in inj_insp_files - if inj_tag in f.tags[1]]) + tag_insp_files = FileList([f for f in inj_insp_files if inj_tag in f.tags[1]]) node, inj_find_file = job_instance.create_node( - tag_inj_files, tag_insp_files, - bank_file, pp_dir) + tag_inj_files, tag_insp_files, bank_file, pp_dir + ) wf.add_node(node) inj_find_files.append(inj_find_file) @@ -375,7 +399,8 @@ def setup_pygrb_pp_workflow(wf, pp_dir, seg_dir, segment, bank_file, class PycbcGrbTrigCombinerExecutable(Executable): - """ The class responsible for creating jobs + """ + The class responsible for creating jobs for ''pycbc_grb_trig_combiner''. """ @@ -383,14 +408,15 @@ class PycbcGrbTrigCombinerExecutable(Executable): def __init__(self, cp, name): super().__init__(cp=cp, name=name) - self.trigger_name = cp.get('workflow', 'trigger-name') - self.trig_start_time = cp.get('workflow', 'start-time') - self.num_trials = int(cp.get('trig_combiner', 'num-trials')) + self.trigger_name = cp.get("workflow", "trigger-name") + self.trig_start_time = cp.get("workflow", "start-time") + self.num_trials = int(cp.get("trig_combiner", "num-trials")) - def create_node(self, ifo_tag, seg_dir, segment, insp_files, - out_dir, bank_file, tags=None): + def create_node( + self, ifo_tag, seg_dir, segment, insp_files, out_dir, bank_file, tags=None + ): node = Node(self) - node.add_opt('--verbose') + node.add_opt("--verbose") node.add_opt("--ifo-tag", ifo_tag) node.add_opt("--grb-name", self.trigger_name) node.add_opt("--trig-start-time", self.trig_start_time) @@ -401,18 +427,20 @@ def create_node(self, ifo_tag, seg_dir, segment, insp_files, # Prepare output file tag user_tag = f"PYGRB_GRB{self.trigger_name}" if tags: - user_tag += "_{}".format(tags) + user_tag += f"_{tags}" # Add on/off source and off trial outputs output_files = FileList([]) - outfile_types = ['ALL_TIMES', 'ONSOURCE', 'OFFSOURCE'] + outfile_types = ["ALL_TIMES", "ONSOURCE", "OFFSOURCE"] for i in range(self.num_trials): - outfile_types.append("OFFTRIAL_{}".format(i+1)) + outfile_types.append(f"OFFTRIAL_{i + 1}") for out_type in outfile_types: - out_name = "{}-{}_{}-{}-{}.h5".format( - ifo_tag, user_tag, out_type, - segment[0], segment[1]-segment[0]) - out_file = File(ifo_tag, 'trig_combiner', segment, - file_url=os.path.join(out_dir, out_name)) + out_name = f"{ifo_tag}-{user_tag}_{out_type}-{segment[0]}-{segment[1] - segment[0]}.h5" + out_file = File( + ifo_tag, + "trig_combiner", + segment, + file_url=os.path.join(out_dir, out_name), + ) node.add_output(out_file) output_files.append(out_file) @@ -420,7 +448,8 @@ def create_node(self, ifo_tag, seg_dir, segment, insp_files, class PycbcGrbTrigClusterExecutable(Executable): - """ The class responsible for creating jobs + """ + The class responsible for creating jobs for ''pycbc_grb_trig_cluster''. """ @@ -435,44 +464,42 @@ def create_node(self, in_file, out_dir): # Determine output file name ifotag, filetag, segment = filename_metadata(in_file.name) start, end = segment - out_name = "{}-{}_CLUSTERED-{}-{}.h5".format(ifotag, filetag, - start, end-start) - out_file = File(ifotag, 'trig_cluster', segment, - file_url=os.path.join(out_dir, out_name)) + out_name = f"{ifotag}-{filetag}_CLUSTERED-{start}-{end - start}.h5" + out_file = File( + ifotag, "trig_cluster", segment, file_url=os.path.join(out_dir, out_name) + ) node.add_output(out_file) return node, out_file class PycbcGrbInjFinderExecutable(Executable): - """The class responsible for creating jobs for ``pycbc_grb_inj_finder`` - """ + """The class responsible for creating jobs for ``pycbc_grb_inj_finder``""" + current_retention_level = Executable.ALL_TRIGGERS def __init__(self, cp, exe_name): super().__init__(cp=cp, name=exe_name) - def create_node(self, inj_files, inj_insp_files, bank_file, - out_dir, tags=None): + def create_node(self, inj_files, inj_insp_files, bank_file, out_dir, tags=None): if tags is None: tags = [] node = Node(self) - node.add_input_list_opt('--input-files', inj_insp_files) - node.add_input_list_opt('--inj-files', inj_files) - node.add_input_opt('--bank-file', bank_file) + node.add_input_list_opt("--input-files", inj_insp_files) + node.add_input_list_opt("--inj-files", inj_files) + node.add_input_opt("--bank-file", bank_file) ifo_tag, desc, segment = filename_metadata(inj_files[0].name) - desc = '_'.join(desc.split('_')[:-1]) - out_name = "{}-{}_FOUNDMISSED-{}-{}.h5".format( - ifo_tag, desc, segment[0], abs(segment)) - out_file = File(ifo_tag, 'inj_finder', segment, - os.path.join(out_dir, out_name), tags=tags) + desc = "_".join(desc.split("_")[:-1]) + out_name = f"{ifo_tag}-{desc}_FOUNDMISSED-{segment[0]}-{abs(segment)}.h5" + out_file = File( + ifo_tag, "inj_finder", segment, os.path.join(out_dir, out_name), tags=tags + ) node.add_output(out_file) return node, out_file def build_segment_filelist(seg_dir): """Construct a FileList instance containing all segments txt files""" - # Needs to be in this order for consistency with _read_seg_files file_names = ["bufferSeg.txt", "offSourceSeg.txt", "onSourceSeg.txt"] seg_files = [os.path.join(seg_dir, fn) for fn in file_names] @@ -482,111 +509,139 @@ def build_segment_filelist(seg_dir): return seg_files -def make_pygrb_plot(workflow, exec_name, out_dir, - ifo=None, inj_file=None, trig_file=None, - onsource_file=None, bank_file=None, - seg_files=None, sky_grid_file=None, - veto_file=None, tags=None, **kwargs): +def make_pygrb_plot( + workflow, + exec_name, + out_dir, + ifo=None, + inj_file=None, + trig_file=None, + onsource_file=None, + bank_file=None, + seg_files=None, + sky_grid_file=None, + veto_file=None, + tags=None, + **kwargs, +): """Adds a node for a plot of PyGRB results to the workflow""" - tags = [] if tags is None else tags # Initialize job node with its tags - grb_name = workflow.cp.get('workflow', 'trigger-name') - extra_tags = ['GRB'+grb_name] + grb_name = workflow.cp.get("workflow", "trigger-name") + extra_tags = ["GRB" + grb_name] if ifo: extra_tags.append(ifo) - node = PlotExecutable(workflow.cp, exec_name, ifos=workflow.ifos, - out_dir=out_dir, - tags=tags+extra_tags).create_node() + node = PlotExecutable( + workflow.cp, + exec_name, + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags + extra_tags, + ).create_node() if trig_file: - node.add_input_opt('--trig-file', trig_file) + node.add_input_opt("--trig-file", trig_file) # Pass the veto and segment files and options if seg_files: - node.add_input_list_opt('--seg-files', seg_files) + node.add_input_list_opt("--seg-files", seg_files) if sky_grid_file: - node.add_input_opt('--sky-grid', sky_grid_file) + node.add_input_opt("--sky-grid", sky_grid_file) if veto_file: - node.add_input_opt('--veto-file', veto_file) + node.add_input_opt("--veto-file", veto_file) # Option to show the onsource trial if this is a plot of all data - if exec_name == 'pygrb_plot_snr_timeseries' and 'alltimes' in tags: - node.add_opt('--onsource') - if exec_name in ['pygrb_plot_injs_results', - 'pygrb_plot_snr_timeseries']: - trig_time = workflow.cp.get('workflow', 'trigger-time') - node.add_opt('--trigger-time', trig_time) + if exec_name == "pygrb_plot_snr_timeseries" and "alltimes" in tags: + node.add_opt("--onsource") + if exec_name in ["pygrb_plot_injs_results", "pygrb_plot_snr_timeseries"]: + trig_time = workflow.cp.get("workflow", "trigger-time") + node.add_opt("--trigger-time", trig_time) # Pass the injection file as an input File instance - if inj_file is not None and exec_name not in \ - ['pygrb_plot_skygrid', 'pygrb_plot_stats_distribution', - 'pycbc_plot_bank_corner']: - node.add_input_opt('--found-missed-file', inj_file) - if exec_name == 'pycbc_plot_bank_corner': + if inj_file is not None and exec_name not in [ + "pygrb_plot_skygrid", + "pygrb_plot_stats_distribution", + "pycbc_plot_bank_corner", + ]: + node.add_input_opt("--found-missed-file", inj_file) + if exec_name == "pycbc_plot_bank_corner": if inj_file is not None: - node.add_input_opt('--bank-file', inj_file) + node.add_input_opt("--bank-file", inj_file) else: - node.add_input_opt('--bank-file', bank_file) + node.add_input_opt("--bank-file", bank_file) # IFO option if ifo: - node.add_opt('--ifo', ifo) + node.add_opt("--ifo", ifo) # Output files and final input file (passed as a File instance) - if exec_name == 'pygrb_efficiency': + if exec_name == "pygrb_efficiency": # In this case tags[0] is the offtrial number - node.add_input_opt('--bank-file', bank_file) - node.add_opt('--trial-name', tags[0]) - node.add_opt('--injection-set-name', tags[1]) + node.add_input_opt("--bank-file", bank_file) + node.add_opt("--trial-name", tags[0]) + node.add_opt("--injection-set-name", tags[1]) # Output the sensitivity plot - if kwargs['plot_bkgd']: - node.new_output_file_opt(workflow.analysis_time, '.png', - '--background-output-file', - tags=extra_tags+['max_background']) + if kwargs["plot_bkgd"]: + node.new_output_file_opt( + workflow.analysis_time, + ".png", + "--background-output-file", + tags=extra_tags + ["max_background"], + ) # Output the exclusion distance plot and table else: - node.add_input_opt('--onsource-file', - onsource_file) - node.new_output_file_opt(workflow.analysis_time, '.png', - '--onsource-output-file', - tags=['onsource']+extra_tags) - node.new_output_file_opt(workflow.analysis_time, '.json', - '--exclusion-dist-output-file', - tags=extra_tags) - elif exec_name == 'pycbc_plot_bank_corner': - node.new_output_file_opt(workflow.analysis_time, '.png', - '--output-plot-file', tags=extra_tags+tags) + node.add_input_opt("--onsource-file", onsource_file) + node.new_output_file_opt( + workflow.analysis_time, + ".png", + "--onsource-output-file", + tags=["onsource"] + extra_tags, + ) + node.new_output_file_opt( + workflow.analysis_time, + ".json", + "--exclusion-dist-output-file", + tags=extra_tags, + ) + elif exec_name == "pycbc_plot_bank_corner": + node.new_output_file_opt( + workflow.analysis_time, ".png", "--output-plot-file", tags=extra_tags + tags + ) else: - node.new_output_file_opt(workflow.analysis_time, '.png', - '--output-file', tags=extra_tags) - if exec_name in ['pygrb_plot_coh_ifosnr', 'pygrb_plot_null_stats'] \ - and 'zoomin' in tags: - node.add_opt('--zoom-in') + node.new_output_file_opt( + workflow.analysis_time, ".png", "--output-file", tags=extra_tags + ) + if ( + exec_name in ["pygrb_plot_coh_ifosnr", "pygrb_plot_null_stats"] + and "zoomin" in tags + ): + node.add_opt("--zoom-in") # Quantity to be displayed on the y-axis of the plot - if exec_name in ['pygrb_plot_chisq_veto', 'pygrb_plot_null_stats', - 'pygrb_plot_snr_timeseries']: - node.add_opt('--y-variable', tags[0]) + if exec_name in [ + "pygrb_plot_chisq_veto", + "pygrb_plot_null_stats", + "pygrb_plot_snr_timeseries", + ]: + node.add_opt("--y-variable", tags[0]) # Quantity to be displayed on the x-axis of the plot - elif exec_name == 'pygrb_plot_stats_distribution': - node.add_opt('--x-variable', tags[0]) - elif exec_name == 'pygrb_plot_injs_results': + elif exec_name == "pygrb_plot_stats_distribution": + node.add_opt("--x-variable", tags[0]) + elif exec_name == "pygrb_plot_injs_results": # Variables to plot on x and y axes - node.add_opt('--y-variable', tags[0]) - node.add_opt('--x-variable', tags[1]) + node.add_opt("--y-variable", tags[0]) + node.add_opt("--x-variable", tags[1]) # Flag to plot found over missed or missed over found - if tags[2] == 'missed-on-top': - node.add_opt('--'+tags[2]) + if tags[2] == "missed-on-top": + node.add_opt("--" + tags[2]) # Enable log axes - subsection = '_'.join(tags[0:2]) - for log_flag in ['x-log', 'y-log']: - if workflow.cp.has_option_tags(exec_name, log_flag, - tags=[subsection]): - node.add_opt('--'+log_flag) - elif exec_name == 'pycbc_plot_bank_corner': - node.add_opt('--no-suptitle') - #if inj_file: + subsection = "_".join(tags[0:2]) + for log_flag in ["x-log", "y-log"]: + if workflow.cp.has_option_tags(exec_name, log_flag, tags=[subsection]): + node.add_opt("--" + log_flag) + elif exec_name == "pycbc_plot_bank_corner": + node.add_opt("--no-suptitle") + # if inj_file: # node.add_opt('--title', f'\"{tags[0]} injections\"') # params = workflow.cp.get_opt_tags(exec_name, 'parameters', ['injs']) - #else: + # else: # node.add_opt('--title', '\"Template bank\"') # params = workflow.cp.get_opt_tags(exec_name, 'parameters', ['bank']) - #node.add_opt('--parameters', params) + # node.add_opt('--parameters', params) # Add job node to workflow workflow += node @@ -594,32 +649,35 @@ def make_pygrb_plot(workflow, exec_name, out_dir, return node, node.output_files -def make_pygrb_info_table(workflow, exec_name, out_dir, in_files=None, - tags=None): +def make_pygrb_info_table(workflow, exec_name, out_dir, in_files=None, tags=None): """ Setup a job to create an html snippet with the GRB trigger information or exlusion distances information. """ - # Organize tags tags = [] if tags is None else tags - grb_name = workflow.cp.get('workflow', 'trigger-name') - extra_tags = ['GRB'+grb_name] + grb_name = workflow.cp.get("workflow", "trigger-name") + extra_tags = ["GRB" + grb_name] # Initialize job node - node = PlotExecutable(workflow.cp, exec_name, - ifos=workflow.ifos, out_dir=out_dir, - tags=tags+extra_tags).create_node() + node = PlotExecutable( + workflow.cp, + exec_name, + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags + extra_tags, + ).create_node() # Options - if exec_name == 'pygrb_grb_info_table': - node.add_opt('--ifos', ' '.join(workflow.ifos)) - elif exec_name == 'pygrb_exclusion_dist_table': - node.add_input_opt('--input-files', in_files) + if exec_name == "pygrb_grb_info_table": + node.add_opt("--ifos", " ".join(workflow.ifos)) + elif exec_name == "pygrb_exclusion_dist_table": + node.add_input_opt("--input-files", in_files) # Output - node.new_output_file_opt(workflow.analysis_time, '.html', - '--output-file', tags=extra_tags) + node.new_output_file_opt( + workflow.analysis_time, ".html", "--output-file", tags=extra_tags + ) # Add job node to workflow workflow += node @@ -627,61 +685,85 @@ def make_pygrb_info_table(workflow, exec_name, out_dir, in_files=None, return node, node.output_files -def make_pygrb_injs_tables(workflow, out_dir, bank_file, off_file, seg_files, - inj_file=None, on_file=None, veto_file=None, - tags=None): +def make_pygrb_injs_tables( + workflow, + out_dir, + bank_file, + off_file, + seg_files, + inj_file=None, + on_file=None, + veto_file=None, + tags=None, +): """ Adds a job to make quiet-found and missed-found injection tables, - or loudest trigger(s) table.""" - + or loudest trigger(s) table. + """ tags = [] if tags is None else tags # Executable - exec_name = 'pygrb_page_tables' + exec_name = "pygrb_page_tables" # Initialize job node - grb_name = workflow.cp.get('workflow', 'trigger-name') - extra_tags = ['GRB'+grb_name] - node = PlotExecutable(workflow.cp, exec_name, - ifos=workflow.ifos, out_dir=out_dir, - tags=tags+extra_tags).create_node() + grb_name = workflow.cp.get("workflow", "trigger-name") + extra_tags = ["GRB" + grb_name] + node = PlotExecutable( + workflow.cp, + exec_name, + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags + extra_tags, + ).create_node() # Pass the bank-file - node.add_input_opt('--bank-file', bank_file) + node.add_input_opt("--bank-file", bank_file) # Offsource input file (or equivalently trigger file for injections) offsource_file = off_file - node.add_input_opt('--offsource-file', offsource_file) + node.add_input_opt("--offsource-file", offsource_file) # Pass the veto and segment files (as File instances) if veto_file: - node.add_input_opt('--veto-file', veto_file) - node.add_input_list_opt('--seg-files', seg_files) + node.add_input_opt("--veto-file", veto_file) + node.add_input_list_opt("--seg-files", seg_files) # Handle input/output for injections if inj_file: # Found-missed injection file (passed as File instance) - node.add_input_opt('--found-missed-file', inj_file) + node.add_input_opt("--found-missed-file", inj_file) # Missed-found and quiet-found injections html output files - for mf_or_qf in ['missed-found', 'quiet-found']: - mf_or_qf_tags = [mf_or_qf.upper().replace('-', '_')] - node.new_output_file_opt(workflow.analysis_time, '.html', - '--'+mf_or_qf+'-injs-output-file', - tags=extra_tags+mf_or_qf_tags) + for mf_or_qf in ["missed-found", "quiet-found"]: + mf_or_qf_tags = [mf_or_qf.upper().replace("-", "_")] + node.new_output_file_opt( + workflow.analysis_time, + ".html", + "--" + mf_or_qf + "-injs-output-file", + tags=extra_tags + mf_or_qf_tags, + ) # Quiet-found injections h5 output file - node.new_output_file_opt(workflow.analysis_time, '.h5', - '--quiet-found-injs-h5-output-file', - tags=extra_tags+['QUIET_FOUND']) + node.new_output_file_opt( + workflow.analysis_time, + ".h5", + "--quiet-found-injs-h5-output-file", + tags=extra_tags + ["QUIET_FOUND"], + ) # Handle input/output for onsource/offsource else: - src_type = 'offsource-trigs' + src_type = "offsource-trigs" if on_file: - src_type = 'onsource-trig' + src_type = "onsource-trig" # Pass onsource input File instance - node.add_input_opt('--onsource-file', on_file) + node.add_input_opt("--onsource-file", on_file) # Loudest offsource/onsource triggers html and h5 output files - src_type_tags = [src_type.upper().replace('-', '_')] - node.new_output_file_opt(workflow.analysis_time, '.html', - '--loudest-'+src_type+'-output-file', - tags=extra_tags+src_type_tags) - node.new_output_file_opt(workflow.analysis_time, '.h5', - '--loudest-'+src_type+'-h5-output-file', - tags=extra_tags+src_type_tags) + src_type_tags = [src_type.upper().replace("-", "_")] + node.new_output_file_opt( + workflow.analysis_time, + ".html", + "--loudest-" + src_type + "-output-file", + tags=extra_tags + src_type_tags, + ) + node.new_output_file_opt( + workflow.analysis_time, + ".h5", + "--loudest-" + src_type + "-h5-output-file", + tags=extra_tags + src_type_tags, + ) # Add job node to the workflow workflow += node @@ -690,10 +772,18 @@ def make_pygrb_injs_tables(workflow, out_dir, bank_file, off_file, seg_files, # Based on setup_single_det_minifollowups -def setup_pygrb_minifollowups(workflow, followups_file, trigger_file, - dax_output, out_dir, seg_files=None, - veto_file=None, tags=None): - """ Create plots that followup the the loudest PyGRB triggers or +def setup_pygrb_minifollowups( + workflow, + followups_file, + trigger_file, + dax_output, + out_dir, + seg_files=None, + veto_file=None, + tags=None, +): + """ + Create plots that followup the the loudest PyGRB triggers or missed injections from an HDF file. Parameters @@ -713,79 +803,93 @@ def setup_pygrb_minifollowups(workflow, followups_file, trigger_file, The veto definer file tags: {None, optional} Tags to add to the minifollowups executables - """ - logging.info('Entering minifollowups module') + """ + logging.info("Entering minifollowups module") - if not workflow.cp.has_section('workflow-minifollowups'): - msg = 'There is no [workflow-minifollowups] section in ' - msg += 'the configuration file' + if not workflow.cp.has_section("workflow-minifollowups"): + msg = "There is no [workflow-minifollowups] section in " + msg += "the configuration file" logging.info(msg) - logging.info('Leaving minifollowups') - return + logging.info("Leaving minifollowups") + return None tags = [] if tags is None else tags makedir(dax_output) # Turn the config file into a File instance - config_path = os.path.abspath(dax_output + '/' + - '_'.join(tags) + '_minifollowup.ini') - workflow.cp.write(open(config_path, 'w')) + config_path = os.path.abspath( + dax_output + "/" + "_".join(tags) + "_minifollowup.ini" + ) + workflow.cp.write(open(config_path, "w")) config_file = resolve_url_to_file(config_path) # wikifile = curr_ifo + '_'.join(tags) + 'loudest_table.txt' - wikifile = '_'.join(tags) + 'loudest_table.txt' + wikifile = "_".join(tags) + "loudest_table.txt" # Create the node - exe = Executable(workflow.cp, 'pygrb_minifollowups', - ifos=workflow.ifos, out_dir=dax_output, - tags=tags) + exe = Executable( + workflow.cp, + "pygrb_minifollowups", + ifos=workflow.ifos, + out_dir=dax_output, + tags=tags, + ) node = exe.create_node() - node.add_input_opt('--trig-file', trigger_file) + node.add_input_opt("--trig-file", trigger_file) # Grab and pass all necessary files as File instances if seg_files: - node.add_input_list_opt('--seg-files', seg_files) + node.add_input_list_opt("--seg-files", seg_files) if veto_file: - node.add_input_opt('--veto-file', veto_file) - node.add_input_opt('--config-files', config_file) - node.add_input_opt('--followups-file', followups_file) - node.add_opt('--wiki-file', wikifile) + node.add_input_opt("--veto-file", veto_file) + node.add_input_opt("--config-files", config_file) + node.add_input_opt("--followups-file", followups_file) + node.add_opt("--wiki-file", wikifile) if tags: - node.add_list_opt('--tags', tags) - node.new_output_file_opt(workflow.analysis_time, '.dax', '--dax-file') - node.new_output_file_opt(workflow.analysis_time, '.dax.map', - '--output-map') + node.add_list_opt("--tags", tags) + node.new_output_file_opt(workflow.analysis_time, ".dax", "--dax-file") + node.new_output_file_opt(workflow.analysis_time, ".dax.map", "--output-map") name = node.output_files[0].name - assert name.endswith('.dax') + assert name.endswith(".dax") map_file = node.output_files[1] - assert map_file.name.endswith('.map') + assert map_file.name.endswith(".map") - node.add_opt('--workflow-name', name) - node.add_opt('--output-dir', out_dir) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", name) + node.add_opt("--output-dir", out_dir) + node.add_opt("--dax-file-directory", ".") workflow += node # Execute this in a sub-workflow fil = node.output_files[0] job = SubWorkflow(fil.name, is_planned=False) - job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) job.add_into_workflow(workflow) - logging.info('Leaving minifollowups module') + logging.info("Leaving minifollowups module") return job def setup_pygrb_results_workflow( - workflow, res_dir, trig_files, full_injs_files, inj_files, bank_file, - seg_dir, sky_grid_file, veto_file=None, tags=None, - explicit_dependencies=None): - """Create subworkflow to produce plots, tables, + workflow, + res_dir, + trig_files, + full_injs_files, + inj_files, + bank_file, + seg_dir, + sky_grid_file, + veto_file=None, + tags=None, + explicit_dependencies=None, +): + """ + Create subworkflow to produce plots, tables, and results webpage for a PyGRB analysis. Parameters @@ -805,58 +909,61 @@ def setup_pygrb_results_workflow( Tags to add to the executables explicit_dependencies: {None, optional} nodes that must precede this - """ + """ tags = [] if tags is None else tags - dax_output = res_dir+'/webpage_daxes' + dax_output = res_dir + "/webpage_daxes" # _workflow.makedir(dax_output) makedir(dax_output) # Create the node - exe = Executable(workflow.cp, 'pygrb_results_workflow', - ifos=workflow.ifo_string, out_dir=dax_output, - tags=tags) + exe = Executable( + workflow.cp, + "pygrb_results_workflow", + ifos=workflow.ifo_string, + out_dir=dax_output, + tags=tags, + ) node = exe.create_node() # Grab and pass all necessary files - node.add_input_list_opt('--trig-files', trig_files) + node.add_input_list_opt("--trig-files", trig_files) # node.add_input_opt('--config-files', config_file) - node.add_input_list_opt('--full-inj-files', full_injs_files) - node.add_input_list_opt('--inj-files', inj_files) - node.add_input_opt('--bank-file', bank_file) - node.add_input_opt('--sky-grid', sky_grid_file) - node.add_opt('--segment-dir', seg_dir) + node.add_input_list_opt("--full-inj-files", full_injs_files) + node.add_input_list_opt("--inj-files", inj_files) + node.add_input_opt("--bank-file", bank_file) + node.add_input_opt("--sky-grid", sky_grid_file) + node.add_opt("--segment-dir", seg_dir) if veto_file: - node.add_input_opt('--veto-file', veto_file) + node.add_input_opt("--veto-file", veto_file) if tags: - node.add_list_opt('--tags', tags) + node.add_list_opt("--tags", tags) - node.new_output_file_opt(workflow.analysis_time, '.dax', - '--dax-file', tags=tags) - node.new_output_file_opt(workflow.analysis_time, '.map', - '--output-map', tags=tags) + node.new_output_file_opt(workflow.analysis_time, ".dax", "--dax-file", tags=tags) + node.new_output_file_opt(workflow.analysis_time, ".map", "--output-map", tags=tags) # + ['MAP'], use_tmp_subdirs=True) name = node.output_files[0].name - assert name.endswith('.dax') + assert name.endswith(".dax") map_file = node.output_files[1] - assert map_file.name.endswith('.map') - node.add_opt('--workflow-name', name) + assert map_file.name.endswith(".map") + node.add_opt("--workflow-name", name) # This is the output dir for the products of this node, namely dax and map - node.add_opt('--output-dir', res_dir) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--output-dir", res_dir) + node.add_opt("--dax-file-directory", ".") # Turn the config file into a File instance - config_path = os.path.abspath(dax_output + '/' + - '_'.join(tags) + 'webpage.ini') - workflow.cp.write(open(config_path, 'w')) + config_path = os.path.abspath(dax_output + "/" + "_".join(tags) + "webpage.ini") + workflow.cp.write(open(config_path, "w")) config_file = resolve_url_to_file(config_path) - node.add_input_opt('--config-files', config_file) + node.add_input_opt("--config-files", config_file) # Track additional ini file produced by pycbc_pygrb_results_workflow - out_file = File(workflow.ifos, - 'pygrb_results_workflow', - workflow.analysis_time, - file_url=os.path.join(dax_output, name+'.ini')) + out_file = File( + workflow.ifos, + "pygrb_results_workflow", + workflow.analysis_time, + file_url=os.path.join(dax_output, name + ".ini"), + ) node.add_output(out_file) # Add node to the workflow @@ -867,9 +974,9 @@ def setup_pygrb_results_workflow( # Execute this in a sub-workflow job = SubWorkflow(name, is_planned=False) # , _id='results') - job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) job.add_into_workflow(workflow) return node.output_files diff --git a/pycbc/workflow/inference_followups.py b/pycbc/workflow/inference_followups.py index 1af2655b35b..90d6656c31b 100644 --- a/pycbc/workflow/inference_followups.py +++ b/pycbc/workflow/inference_followups.py @@ -16,21 +16,29 @@ """ Module that contains functions for setting up the inference workflow. """ + import logging -from pycbc.workflow.core import (Executable, makedir) -from pycbc.workflow.plotting import PlotExecutable from pycbc.results import layout +from pycbc.workflow.core import Executable, makedir +from pycbc.workflow.plotting import PlotExecutable -logger = logging.getLogger('pycbc.workflow.inference_followups') +logger = logging.getLogger("pycbc.workflow.inference_followups") -def make_inference_plot(workflow, input_file, output_dir, - name, analysis_seg=None, - tags=None, input_file_opt='input-file', - output_file_extension='.png', - add_to_workflow=False): - """Boiler-plate function for creating a standard plotting job. +def make_inference_plot( + workflow, + input_file, + output_dir, + name, + analysis_seg=None, + tags=None, + input_file_opt="input-file", + output_file_extension=".png", + add_to_workflow=False, +): + """ + Boiler-plate function for creating a standard plotting job. Parameters ---------- @@ -63,6 +71,7 @@ def make_inference_plot(workflow, input_file, output_dir, ------- pycbc.workflow.plotting.PlotExecutable The job node for creating the plot. + """ # default values if tags is None: @@ -77,37 +86,37 @@ def make_inference_plot(workflow, input_file, output_dir, # appropriate escapes to the parameters option so pegasus will render it # properly (see _params_for_pegasus for details). parameters = None - if workflow.cp.has_option(name, 'parameters'): - parameters = workflow.cp.get(name, 'parameters') - workflow.cp.remove_option(name, 'parameters') + if workflow.cp.has_option(name, "parameters"): + parameters = workflow.cp.get(name, "parameters") + workflow.cp.remove_option(name, "parameters") # make a node for plotting the posterior as a corner plot - node = PlotExecutable(workflow.cp, name, ifos=workflow.ifos, - out_dir=output_dir, - tags=tags).create_node() + node = PlotExecutable( + workflow.cp, name, ifos=workflow.ifos, out_dir=output_dir, tags=tags + ).create_node() # add back the parameters option if it was specified if parameters is not None: node.add_opt("--parameters", _params_for_pegasus(parameters)) # and put the opt back in the config file in memory - workflow.cp.set(name, 'parameters', parameters) + workflow.cp.set(name, "parameters", parameters) # add input and output options if isinstance(input_file, list): # list of input files are given, use input_list_opt - node.add_input_list_opt("--{}".format(input_file_opt), input_file) + node.add_input_list_opt(f"--{input_file_opt}", input_file) else: # assume just a single file - node.add_input_opt("--{}".format(input_file_opt), input_file) - node.new_output_file_opt(analysis_seg, output_file_extension, - "--output-file") + node.add_input_opt(f"--{input_file_opt}", input_file) + node.new_output_file_opt(analysis_seg, output_file_extension, "--output-file") # add node to workflow if add_to_workflow: workflow += node return node -def make_inference_prior_plot(workflow, config_file, output_dir, - name="plot_prior", - analysis_seg=None, tags=None): - """Sets up the corner plot of the priors in the workflow. +def make_inference_prior_plot( + workflow, config_file, output_dir, name="plot_prior", analysis_seg=None, tags=None +): + """ + Sets up the corner plot of the priors in the workflow. Parameters ---------- @@ -131,18 +140,32 @@ def make_inference_prior_plot(workflow, config_file, output_dir, ------- pycbc.workflow.FileList A list of the output files. + """ - node = make_inference_plot(workflow, config_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - input_file_opt='config-file', - add_to_workflow=True) + node = make_inference_plot( + workflow, + config_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + input_file_opt="config-file", + add_to_workflow=True, + ) return node.output_files -def create_posterior_files(workflow, samples_files, output_dir, - parameters=None, name="extract_posterior", - analysis_seg=None, tags=None): - """Sets up job to create posterior files from some given samples files. +def create_posterior_files( + workflow, + samples_files, + output_dir, + parameters=None, + name="extract_posterior", + analysis_seg=None, + tags=None, +): + """ + Sets up job to create posterior files from some given samples files. Parameters ---------- @@ -166,6 +189,7 @@ def create_posterior_files(workflow, samples_files, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ if analysis_seg is None: analysis_seg = workflow.analysis_time @@ -177,18 +201,18 @@ def create_posterior_files(workflow, samples_files, output_dir, # appropriate escapes to the parameters option so pegasus will render it # properly (see _params_for_pegasus for details). parameters = None - if workflow.cp.has_option(name, 'parameters'): - parameters = workflow.cp.get(name, 'parameters') - workflow.cp.remove_option(name, 'parameters') - extract_posterior_exe = Executable(workflow.cp, name, - ifos=workflow.ifos, - out_dir=output_dir) + if workflow.cp.has_option(name, "parameters"): + parameters = workflow.cp.get(name, "parameters") + workflow.cp.remove_option(name, "parameters") + extract_posterior_exe = Executable( + workflow.cp, name, ifos=workflow.ifos, out_dir=output_dir + ) node = extract_posterior_exe.create_node() # add back the parameters option if it was specified if parameters is not None: node.add_opt("--parameters", _params_for_pegasus(parameters)) # and put the opt back in the config file in memory - workflow.cp.set(name, 'parameters', parameters) + workflow.cp.set(name, "parameters", parameters) if not isinstance(samples_files, list): samples_files = [samples_files] node.add_input_list_opt("--input-file", samples_files) @@ -198,10 +222,16 @@ def create_posterior_files(workflow, samples_files, output_dir, return node.output_files -def create_fits_file(workflow, inference_file, output_dir, - name="create_fits_file", - analysis_seg=None, tags=None): - """Sets up job to create fits files from some given samples files. +def create_fits_file( + workflow, + inference_file, + output_dir, + name="create_fits_file", + analysis_seg=None, + tags=None, +): + """ + Sets up job to create fits files from some given samples files. Parameters ---------- @@ -225,14 +255,15 @@ def create_fits_file(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ if analysis_seg is None: analysis_seg = workflow.analysis_time if tags is None: tags = [] - create_fits_exe = Executable(workflow.cp, name, - ifos=workflow.ifos, - out_dir=output_dir) + create_fits_exe = Executable( + workflow.cp, name, ifos=workflow.ifos, out_dir=output_dir + ) node = create_fits_exe.create_node() node.add_input_opt("--input-file", inference_file) node.new_output_file_opt(analysis_seg, ".fits", "--output-file", tags=tags) @@ -241,10 +272,11 @@ def create_fits_file(workflow, inference_file, output_dir, return node.output_files -def make_inference_skymap(workflow, fits_file, output_dir, - name="plot_skymap", analysis_seg=None, - tags=None): - """Sets up the skymap plot. +def make_inference_skymap( + workflow, fits_file, output_dir, name="plot_skymap", analysis_seg=None, tags=None +): + """ + Sets up the skymap plot. Parameters ---------- @@ -268,18 +300,32 @@ def make_inference_skymap(workflow, fits_file, output_dir, ------- pycbc.workflow.FileList A list of result and output files. + """ - node = make_inference_plot(workflow, fits_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=True) + node = make_inference_plot( + workflow, + fits_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=True, + ) return node.output_files -def make_inference_summary_table(workflow, inference_file, output_dir, - parameters=None, print_metadata=None, - name="table_summary", - analysis_seg=None, tags=None): - """Sets up the html table summarizing parameter estimates. +def make_inference_summary_table( + workflow, + inference_file, + output_dir, + parameters=None, + print_metadata=None, + name="table_summary", + analysis_seg=None, + tags=None, +): + """ + Sets up the html table summarizing parameter estimates. Parameters ---------- @@ -309,13 +355,20 @@ def make_inference_summary_table(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ # we'll use make_inference_plot even though this isn't a plot; the # setup is the same, we just change the file extension - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - output_file_extension='.html', - add_to_workflow=False) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + output_file_extension=".html", + add_to_workflow=False, + ) # now add the parameters and print metadata options; these are pulled # from separate sections in the workflow config file, which is why we # add them separately here @@ -327,11 +380,18 @@ def make_inference_summary_table(workflow, inference_file, output_dir, return node.output_files -def make_inference_posterior_plot(workflow, inference_file, output_dir, - parameters=None, plot_prior_from_file=None, - name="plot_posterior", - analysis_seg=None, tags=None): - """Sets up the corner plot of the posteriors in the workflow. +def make_inference_posterior_plot( + workflow, + inference_file, + output_dir, + parameters=None, + plot_prior_from_file=None, + name="plot_posterior", + analysis_seg=None, + tags=None, +): + """ + Sets up the corner plot of the posteriors in the workflow. Parameters ---------- @@ -359,25 +419,38 @@ def make_inference_posterior_plot(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ # create the node, but delay adding it to the workflow so we can add # the prior file if it is requested - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=False) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=False, + ) if parameters is not None: node.add_opt("--parameters", _params_for_pegasus(parameters)) if plot_prior_from_file is not None: - node.add_input_opt('--plot-prior', plot_prior_from_file) + node.add_input_opt("--plot-prior", plot_prior_from_file) # now add the node to workflow workflow += node return node.output_files -def make_inference_samples_plot(workflow, inference_file, output_dir, - name="plot_samples", - analysis_seg=None, tags=None): - """Sets up a plot of the samples versus iteration (for MCMC samplers). +def make_inference_samples_plot( + workflow, + inference_file, + output_dir, + name="plot_samples", + analysis_seg=None, + tags=None, +): + """ + Sets up a plot of the samples versus iteration (for MCMC samplers). Parameters ---------- @@ -401,17 +474,30 @@ def make_inference_samples_plot(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=True) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=True, + ) return node.output_files -def make_inference_acceptance_rate_plot(workflow, inference_file, output_dir, - name="plot_acceptance_rate", - analysis_seg=None, tags=None): - """Sets up a plot of the acceptance rate (for MCMC samplers). +def make_inference_acceptance_rate_plot( + workflow, + inference_file, + output_dir, + name="plot_acceptance_rate", + analysis_seg=None, + tags=None, +): + """ + Sets up a plot of the acceptance rate (for MCMC samplers). Parameters ---------- @@ -435,17 +521,30 @@ def make_inference_acceptance_rate_plot(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=True) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=True, + ) return node.output_files -def make_inference_plot_mcmc_history(workflow, inference_file, output_dir, - name="plot_mcmc_history", - analysis_seg=None, tags=None): - """Sets up a plot showing the checkpoint history of an MCMC sampler. +def make_inference_plot_mcmc_history( + workflow, + inference_file, + output_dir, + name="plot_mcmc_history", + analysis_seg=None, + tags=None, +): + """ + Sets up a plot showing the checkpoint history of an MCMC sampler. Parameters ---------- @@ -469,17 +568,30 @@ def make_inference_plot_mcmc_history(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=True) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=True, + ) return node.output_files -def make_inference_dynesty_run_plot(workflow, inference_file, output_dir, - name="plot_dynesty_run", - analysis_seg=None, tags=None): - """Sets up a debugging plot for the dynesty run (for Dynesty sampler). +def make_inference_dynesty_run_plot( + workflow, + inference_file, + output_dir, + name="plot_dynesty_run", + analysis_seg=None, + tags=None, +): + """ + Sets up a debugging plot for the dynesty run (for Dynesty sampler). Parameters ---------- @@ -503,17 +615,30 @@ def make_inference_dynesty_run_plot(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=True) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=True, + ) return node.output_files -def make_inference_dynesty_trace_plot(workflow, inference_file, output_dir, - name="plot_dynesty_traceplot", - analysis_seg=None, tags=None): - """Sets up a trace plot for the dynesty run (for Dynesty sampler). +def make_inference_dynesty_trace_plot( + workflow, + inference_file, + output_dir, + name="plot_dynesty_traceplot", + analysis_seg=None, + tags=None, +): + """ + Sets up a trace plot for the dynesty run (for Dynesty sampler). Parameters ---------- @@ -537,18 +662,32 @@ def make_inference_dynesty_trace_plot(workflow, inference_file, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ - node = make_inference_plot(workflow, inference_file, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=True) + node = make_inference_plot( + workflow, + inference_file, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=True, + ) return node.output_files -def make_inference_pp_table(workflow, posterior_files, output_dir, - parameters=None, injection_samples_map=None, - name="pp_table_summary", - analysis_seg=None, tags=None): - """Performs a PP, writing results to an html table. +def make_inference_pp_table( + workflow, + posterior_files, + output_dir, + parameters=None, + injection_samples_map=None, + name="pp_table_summary", + analysis_seg=None, + tags=None, +): + """ + Performs a PP, writing results to an html table. Parameters ---------- @@ -578,28 +717,43 @@ def make_inference_pp_table(workflow, posterior_files, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ # we'll use make_inference_plot even though this isn't a plot; the # setup is the same, we just change the file extension - node = make_inference_plot(workflow, posterior_files, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - output_file_extension='.html', - add_to_workflow=False) + node = make_inference_plot( + workflow, + posterior_files, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + output_file_extension=".html", + add_to_workflow=False, + ) # add the parameters and inj/samples map if parameters is not None: node.add_opt("--parameters", _params_for_pegasus(parameters)) if injection_samples_map is not None: - node.add_opt("--injection-samples-map", - _params_for_pegasus(injection_samples_map)) + node.add_opt( + "--injection-samples-map", _params_for_pegasus(injection_samples_map) + ) workflow += node return node.output_files -def make_inference_pp_plot(workflow, posterior_files, output_dir, - parameters=None, injection_samples_map=None, - name="plot_pp", - analysis_seg=None, tags=None): - """Sets up a pp plot in the workflow. +def make_inference_pp_plot( + workflow, + posterior_files, + output_dir, + parameters=None, + injection_samples_map=None, + name="plot_pp", + analysis_seg=None, + tags=None, +): + """ + Sets up a pp plot in the workflow. Parameters ---------- @@ -628,26 +782,41 @@ def make_inference_pp_plot(workflow, posterior_files, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ - node = make_inference_plot(workflow, posterior_files, output_dir, - name, analysis_seg=analysis_seg, tags=tags, - add_to_workflow=False) + node = make_inference_plot( + workflow, + posterior_files, + output_dir, + name, + analysis_seg=analysis_seg, + tags=tags, + add_to_workflow=False, + ) # add the parameters and inj/samples map if parameters is not None: node.add_opt("--parameters", _params_for_pegasus(parameters)) if injection_samples_map is not None: - node.add_opt("--injection-samples-map", - _params_for_pegasus(injection_samples_map)) + node.add_opt( + "--injection-samples-map", _params_for_pegasus(injection_samples_map) + ) # now add the node to workflow workflow += node return node.output_files -def make_inference_inj_recovery_plot(workflow, posterior_files, output_dir, - parameter, injection_samples_map=None, - name="inj_recovery", - analysis_seg=None, tags=None): - """Sets up the recovered versus injected parameter plot in the workflow. +def make_inference_inj_recovery_plot( + workflow, + posterior_files, + output_dir, + parameter, + injection_samples_map=None, + name="inj_recovery", + analysis_seg=None, + tags=None, +): + """ + Sets up the recovered versus injected parameter plot in the workflow. Parameters ---------- @@ -676,13 +845,20 @@ def make_inference_inj_recovery_plot(workflow, posterior_files, output_dir, ------- pycbc.workflow.FileList A list of output files. + """ # arguments are the same as plot_pp, so just call that with the # different executable name return make_inference_pp_plot( - workflow, posterior_files, output_dir, parameters=parameter, + workflow, + posterior_files, + output_dir, + parameters=parameter, injection_samples_map=injection_samples_map, - name=name, analysis_seg=analysis_seg, tags=tags) + name=name, + analysis_seg=analysis_seg, + tags=tags, + ) def get_plot_group(cp, section_tag): @@ -690,8 +866,11 @@ def get_plot_group(cp, section_tag): group_prefix = "plot-group-" # parameters for the summary plots plot_groups = {} - opts = [opt for opt in cp.options("workflow-{}".format(section_tag)) - if opt.startswith(group_prefix)] + opts = [ + opt + for opt in cp.options(f"workflow-{section_tag}") + if opt.startswith(group_prefix) + ] for opt in opts: group = opt.replace(group_prefix, "").replace("-", "_") plot_groups[group] = cp.get_opt_tag("workflow", opt, section_tag) @@ -699,7 +878,8 @@ def get_plot_group(cp, section_tag): def get_diagnostic_plots(workflow): - """Determines what diagnostic plots to create based on workflow. + """ + Determines what diagnostic plots to create based on workflow. The plots to create are based on what executable's are specified in the workflow's config file. A list of strings is returned giving the diagnostic @@ -716,24 +896,25 @@ def get_diagnostic_plots(workflow): ------- list : List of names of diagnostic plots. + """ diagnostics = [] if "plot_samples" in workflow.cp.options("executables"): - diagnostics.append('samples') + diagnostics.append("samples") if "plot_acceptance_rate" in workflow.cp.options("executables"): - diagnostics.append('acceptance_rate') + diagnostics.append("acceptance_rate") if "plot_mcmc_history" in workflow.cp.options("executables"): - diagnostics.append('mcmc_history') + diagnostics.append("mcmc_history") if "plot_dynesty_run" in workflow.cp.options("executables"): - diagnostics.append('dynesty_run') + diagnostics.append("dynesty_run") if "plot_dynesty_traceplot" in workflow.cp.options("executables"): - diagnostics.append('dynesty_traceplot') + diagnostics.append("dynesty_traceplot") return diagnostics -def make_diagnostic_plots(workflow, diagnostics, samples_file, label, rdir, - tags=None): - """Makes diagnostic plots. +def make_diagnostic_plots(workflow, diagnostics, samples_file, label, rdir, tags=None): + """ + Makes diagnostic plots. Diagnostic plots are sampler-specific plots the provide information on how the sampler performed. All diagnostic plots use the output file @@ -764,79 +945,102 @@ def make_diagnostic_plots(workflow, diagnostics, samples_file, label, rdir, dict : Dictionary of diagnostic name -> list of files giving the plots that will be created. + """ if tags is None: tags = [] out = {} if not isinstance(samples_file, list): samples_file = [samples_file] - if 'samples' in diagnostics: + if "samples" in diagnostics: # files for samples summary subsection - base = "samples/{}".format(label) + base = f"samples/{label}" samples_plots = [] for kk, sf in enumerate(samples_file): samples_plots += make_inference_samples_plot( - workflow, sf, rdir[base], + workflow, + sf, + rdir[base], analysis_seg=workflow.analysis_time, - tags=tags+[label, str(kk)]) - out['samples'] = samples_plots + tags=tags + [label, str(kk)], + ) + out["samples"] = samples_plots layout.group_layout(rdir[base], samples_plots) - if 'acceptance_rate' in diagnostics: + if "acceptance_rate" in diagnostics: # files for samples acceptance_rate subsection - base = "acceptance_rate/{}".format(label) + base = f"acceptance_rate/{label}" acceptance_plots = [] for kk, sf in enumerate(samples_file): acceptance_plots += make_inference_acceptance_rate_plot( - workflow, sf, rdir[base], + workflow, + sf, + rdir[base], analysis_seg=workflow.analysis_time, - tags=tags+[label, str(kk)]) - out['acceptance_rate'] = acceptance_plots + tags=tags + [label, str(kk)], + ) + out["acceptance_rate"] = acceptance_plots layout.single_layout(rdir[base], acceptance_plots) - if 'mcmc_history' in diagnostics: + if "mcmc_history" in diagnostics: # files for samples mcmc history subsection - base = "mcmc_history/{}".format(label) + base = f"mcmc_history/{label}" history_plots = [] for kk, sf in enumerate(samples_file): history_plots += make_inference_plot_mcmc_history( - workflow, sf, rdir[base], + workflow, + sf, + rdir[base], analysis_seg=workflow.analysis_time, - tags=tags+[label, str(kk)]) - out['mcmc_history'] = history_plots + tags=tags + [label, str(kk)], + ) + out["mcmc_history"] = history_plots layout.single_layout(rdir[base], history_plots) - if 'dynesty_run' in diagnostics: + if "dynesty_run" in diagnostics: # files for dynesty run subsection - base = "dynesty_run/{}".format(label) + base = f"dynesty_run/{label}" dynesty_run_plots = [] for kk, sf in enumerate(samples_file): dynesty_run_plots += make_inference_dynesty_run_plot( - workflow, sf, rdir[base], + workflow, + sf, + rdir[base], analysis_seg=workflow.analysis_time, - tags=tags+[label, str(kk)]) - out['dynesty_run'] = dynesty_run_plots + tags=tags + [label, str(kk)], + ) + out["dynesty_run"] = dynesty_run_plots layout.single_layout(rdir[base], dynesty_run_plots) - if 'dynesty_traceplot' in diagnostics: + if "dynesty_traceplot" in diagnostics: # files for samples dynesty tyrace plots subsection - base = "dynesty_traceplot/{}".format(label) + base = f"dynesty_traceplot/{label}" dynesty_trace_plots = [] for kk, sf in enumerate(samples_file): dynesty_trace_plots += make_inference_dynesty_trace_plot( - workflow, sf, rdir[base], + workflow, + sf, + rdir[base], analysis_seg=workflow.analysis_time, - tags=tags+[label, str(kk)]) - out['dynesty_traceplot'] = dynesty_trace_plots + tags=tags + [label, str(kk)], + ) + out["dynesty_traceplot"] = dynesty_trace_plots layout.single_layout(rdir[base], dynesty_trace_plots) return out -def make_posterior_workflow(workflow, samples_files, config_file, label, - rdir, posterior_file_dir='posterior_files', - tags=None): - """Adds jobs to a workflow that make a posterior file and subsequent plots. +def make_posterior_workflow( + workflow, + samples_files, + config_file, + label, + rdir, + posterior_file_dir="posterior_files", + tags=None, +): + """ + Adds jobs to a workflow that make a posterior file and subsequent plots. A posterior file is first created from the given samples file(s). The settings for extracting the posterior are set by the @@ -918,6 +1122,7 @@ def make_posterior_workflow(workflow, samples_files, config_file, label, posterior_plots : list List of posterior plots that will be created. These will be saved to ``posteriors/LABEL/`` in the results directory. + """ # the list of plots to go in the summary summary_files = [] @@ -929,88 +1134,120 @@ def make_posterior_workflow(workflow, samples_files, config_file, label, # figure out what parameters user wants to plot from workflow configuration # parameters for the summary plots - summary_plot_params = get_plot_group(workflow.cp, 'summary_plots') + summary_plot_params = get_plot_group(workflow.cp, "summary_plots") # parameters to plot in large corner plots - plot_params = get_plot_group(workflow.cp, 'plot_params') + plot_params = get_plot_group(workflow.cp, "plot_params") # get parameters for the summary tables - table_params = workflow.cp.get_opt_tag('workflow', 'table-params', - 'summary_table') + table_params = workflow.cp.get_opt_tag("workflow", "table-params", "summary_table") # get any metadata that should be printed - if workflow.cp.has_option('workflow-summary_table', 'print-metadata'): - table_metadata = workflow.cp.get_opt_tag('workflow', 'print-metadata', - 'summary_table') + if workflow.cp.has_option("workflow-summary_table", "print-metadata"): + table_metadata = workflow.cp.get_opt_tag( + "workflow", "print-metadata", "summary_table" + ) else: table_metadata = None # figure out if we are making a skymap - make_skymap = ("create_fits_file" in workflow.cp.options("executables") and - "plot_skymap" in workflow.cp.options("executables")) + make_skymap = "create_fits_file" in workflow.cp.options( + "executables" + ) and "plot_skymap" in workflow.cp.options("executables") - make_prior = ("plot_prior" in workflow.cp.options("executables")) + make_prior = "plot_prior" in workflow.cp.options("executables") _config = None if make_prior: _config = config_file # make node for running extract samples posterior_file = create_posterior_files( - workflow, samples_files, posterior_file_dir, - analysis_seg=analysis_seg, tags=tags+[label])[0] + workflow, + samples_files, + posterior_file_dir, + analysis_seg=analysis_seg, + tags=tags + [label], + )[0] # summary table - summary_files += (make_inference_summary_table( - workflow, posterior_file, rdir.base, - parameters=table_params, print_metadata=table_metadata, - analysis_seg=analysis_seg, - tags=tags+[label]),) + summary_files += ( + make_inference_summary_table( + workflow, + posterior_file, + rdir.base, + parameters=table_params, + print_metadata=table_metadata, + analysis_seg=analysis_seg, + tags=tags + [label], + ), + ) # summary posteriors summary_plots = [] for group, params in summary_plot_params.items(): summary_plots += make_inference_posterior_plot( - workflow, posterior_file, rdir.base, - name='plot_posterior_summary', - parameters=params, plot_prior_from_file=_config, + workflow, + posterior_file, + rdir.base, + name="plot_posterior_summary", + parameters=params, + plot_prior_from_file=_config, analysis_seg=analysis_seg, - tags=tags+[label, group]) + tags=tags + [label, group], + ) # sky map if make_skymap: # create the fits file fits_file = create_fits_file( - workflow, posterior_file, rdir.base, analysis_seg=analysis_seg, - tags=tags+[label])[0] + workflow, + posterior_file, + rdir.base, + analysis_seg=analysis_seg, + tags=tags + [label], + )[0] # now plot the skymap skymap_plot = make_inference_skymap( - workflow, fits_file, rdir.base, analysis_seg=analysis_seg, - tags=tags+[label]) + workflow, + fits_file, + rdir.base, + analysis_seg=analysis_seg, + tags=tags + [label], + ) summary_plots += skymap_plot summary_files += list(layout.grouper(summary_plots, 2)) # files for posteriors summary subsection - base = "posteriors/{}".format(label) + base = f"posteriors/{label}" posterior_plots = [] for group, params in plot_params.items(): posterior_plots += make_inference_posterior_plot( - workflow, posterior_file, rdir[base], - parameters=params, plot_prior_from_file=_config, + workflow, + posterior_file, + rdir[base], + parameters=params, + plot_prior_from_file=_config, analysis_seg=analysis_seg, - tags=tags+[label, group]) + tags=tags + [label, group], + ) layout.single_layout(rdir[base], posterior_plots) prior_plots = [] # files for priors summary section if make_prior: - base = "priors/{}".format(label) + base = f"priors/{label}" prior_plots += make_inference_prior_plot( - workflow, config_file, rdir[base], - analysis_seg=workflow.analysis_time, tags=tags+[label]) + workflow, + config_file, + rdir[base], + analysis_seg=workflow.analysis_time, + tags=tags + [label], + ) layout.single_layout(rdir[base], prior_plots) return posterior_file, summary_files, prior_plots, posterior_plots def _params_for_pegasus(parameters): - """Escapes $ and escapes in parameters string for pegasus. + """ + Escapes $ and escapes in parameters string for pegasus. Pegaus kickstart tries to do variable substitution if it sees a ``$``, and it will strip away back slashes. This can be problematic when trying to use @@ -1023,7 +1260,8 @@ def _params_for_pegasus(parameters): parameters : list or str The parameters argument to modify. If a list, the output will be converted to a space-separated string. + """ if isinstance(parameters, list): parameters = " ".join(parameters) - return parameters.replace('\\', '\\\\').replace('$', '\\$') + return parameters.replace("\\", "\\\\").replace("$", "\\$") diff --git a/pycbc/workflow/injection.py b/pycbc/workflow/injection.py index d7900dcc719..00c2ea0a494 100644 --- a/pycbc/workflow/injection.py +++ b/pycbc/workflow/injection.py @@ -32,140 +32,150 @@ import logging import os.path -from pycbc.workflow.core import FileList, make_analysis_dir, Node -from pycbc.workflow.core import Executable, resolve_url_to_file +from pycbc.workflow.core import ( + Executable, + FileList, + Node, + make_analysis_dir, + resolve_url_to_file, +) from pycbc.workflow.jobsetup import ( - PycbcCreateInjectionsExecutable, select_generic_executable) + PycbcCreateInjectionsExecutable, + select_generic_executable, +) + +logger = logging.getLogger("pycbc.workflow.injection") -logger = logging.getLogger('pycbc.workflow.injection') def veto_injections(workflow, inj_file, veto_file, veto_name, out_dir, tags=None): tags = [] if tags is None else tags make_analysis_dir(out_dir) - node = Executable(workflow.cp, 'strip_injections', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_opt('--segment-name', veto_name) - node.add_input_opt('--veto-file', veto_file) - node.add_input_opt('--injection-file', inj_file) - node.add_opt('--ifos', ' '.join(workflow.ifos)) - node.new_output_file_opt(workflow.analysis_time, '.xml', '--output-file') + node = Executable( + workflow.cp, "strip_injections", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_opt("--segment-name", veto_name) + node.add_input_opt("--veto-file", veto_file) + node.add_input_opt("--injection-file", inj_file) + node.add_opt("--ifos", " ".join(workflow.ifos)) + node.new_output_file_opt(workflow.analysis_time, ".xml", "--output-file") workflow += node return node.output_files[0] class PyCBCOptimalSNRExecutable(Executable): """Compute optimal SNR for injections""" + current_retention_level = Executable.ALL_TRIGGERS def create_node(self, workflow, inj_file, precalc_psd_files, group_str): node = Node(self) _, ext = os.path.splitext(inj_file.name) - node.add_input_opt('--input-file', inj_file) - node.add_opt('--injection-fraction-range', group_str) - node.add_opt('--ifos', ' '.join(workflow.ifos)) - node.add_input_list_opt('--time-varying-psds', precalc_psd_files) - node.new_output_file_opt(workflow.analysis_time, ext, - '--output-file') + node.add_input_opt("--input-file", inj_file) + node.add_opt("--injection-fraction-range", group_str) + node.add_opt("--ifos", " ".join(workflow.ifos)) + node.add_input_list_opt("--time-varying-psds", precalc_psd_files) + node.new_output_file_opt(workflow.analysis_time, ext, "--output-file") return node class PyCBCMergeHDFExecutable(Executable): """Merge HDF injection files executable class""" + current_retention_level = Executable.MERGED_TRIGGERS def create_node(self, workflow, input_files): node = Node(self) - node.add_input_list_opt('--injection-files', input_files) - node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file') + node.add_input_list_opt("--injection-files", input_files) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") return node -def compute_inj_optimal_snr(workflow, inj_file, precalc_psd_files, out_dir, - tags=None): - "Set up a job for computing optimal SNRs of a sim_inspiral file." +def compute_inj_optimal_snr(workflow, inj_file, precalc_psd_files, out_dir, tags=None): + """Set up a job for computing optimal SNRs of a sim_inspiral file.""" if tags is None: tags = [] try: - factor = int(workflow.cp.get_opt_tags('workflow-optimal-snr', - 'parallelization-factor', - tags)) + factor = int( + workflow.cp.get_opt_tags( + "workflow-optimal-snr", "parallelization-factor", tags + ) + ) except Exception as e: logger.warning(e) factor = 1 if factor == 1: # parallelization factor not given - default to single optimal snr job - opt_snr_exe = PyCBCOptimalSNRExecutable(workflow.cp, 'optimal_snr', - ifos=workflow.ifos, - out_dir=out_dir, tags=tags) - node = opt_snr_exe.create_node(workflow, inj_file, - precalc_psd_files, '0/1') + opt_snr_exe = PyCBCOptimalSNRExecutable( + workflow.cp, "optimal_snr", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) + node = opt_snr_exe.create_node(workflow, inj_file, precalc_psd_files, "0/1") workflow += node return node.output_files[0] opt_snr_split_files = [] for i in range(factor): - group_str = '%s/%s' % (i, factor) - opt_snr_exe = PyCBCOptimalSNRExecutable(workflow.cp, 'optimal_snr', - ifos=workflow.ifos, - out_dir=out_dir, - tags=tags + [str(i)]) - opt_snr_exe.update_current_retention_level( - Executable.INTERMEDIATE_PRODUCT) - node = opt_snr_exe.create_node(workflow, inj_file, precalc_psd_files, - group_str) + group_str = "%s/%s" % (i, factor) + opt_snr_exe = PyCBCOptimalSNRExecutable( + workflow.cp, + "optimal_snr", + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags + [str(i)], + ) + opt_snr_exe.update_current_retention_level(Executable.INTERMEDIATE_PRODUCT) + node = opt_snr_exe.create_node(workflow, inj_file, precalc_psd_files, group_str) opt_snr_split_files += [node.output_files[0]] workflow += node hdfcombine_exe = PyCBCMergeHDFExecutable( - workflow.cp, - 'optimal_snr_merge', - ifos=workflow.ifos, - out_dir=out_dir, - tags=tags + workflow.cp, "optimal_snr_merge", ifos=workflow.ifos, out_dir=out_dir, tags=tags ) - hdfcombine_node = hdfcombine_exe.create_node( - workflow, - opt_snr_split_files - ) + hdfcombine_node = hdfcombine_exe.create_node(workflow, opt_snr_split_files) workflow += hdfcombine_node return hdfcombine_node.output_files[0] + def cut_distant_injections(workflow, inj_file, out_dir, tags=None): - "Set up a job for removing injections that are too distant to be seen" + """Set up a job for removing injections that are too distant to be seen""" if tags is None: tags = [] - node = Executable(workflow.cp, 'inj_cut', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--input', inj_file) - node.new_output_file_opt(workflow.analysis_time, '.xml', '--output-file') + node = Executable( + workflow.cp, "inj_cut", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--input", inj_file) + node.new_output_file_opt(workflow.analysis_time, ".xml", "--output-file") workflow += node return node.output_files[0] + def inj_to_hdf(workflow, inj_file, out_dir, tags=None): - """ Convert injection file to hdf format. + """ + Convert injection file to hdf format. If the file is already PyCBC HDF format, this will just make a copy. """ if tags is None: tags = [] - node = Executable(workflow.cp, 'inj2hdf', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--injection-file', inj_file) - node.new_output_file_opt(workflow.analysis_time, '.hdf', '--output-file') + node = Executable( + workflow.cp, "inj2hdf", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--injection-file", inj_file) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node return node.output_file -def setup_injection_workflow(workflow, output_dir=None, - inj_section_name='injections', tags=None): + +def setup_injection_workflow( + workflow, output_dir=None, inj_section_name="injections", tags=None +): """ This function is the gateway for setting up injection-generation jobs in a workflow. It should be possible for this function to support a number @@ -174,7 +184,7 @@ def setup_injection_workflow(workflow, output_dir=None, inspinj) there are currently no subfunctions in this moudle. Parameters - ----------- + ---------- workflow : pycbc.workflow.core.Workflow The Workflow instance that the coincidence jobs will be added to. output_dir : path @@ -189,13 +199,14 @@ def setup_injection_workflow(workflow, output_dir=None, by this call to the workflow. This will be used in output names. Returns - -------- + ------- inj_files : pycbc.workflow.core.FileList The list of injection files created by this call. inj_tags : list of strings The tag corresponding to each injection file and used to uniquely identify them. The FileList class contains functions to search based on tags. + """ if tags is None: tags = [] @@ -210,26 +221,32 @@ def setup_injection_workflow(workflow, output_dir=None, inj_tags = [] inj_files = FileList([]) - for section in workflow.cp.get_subsections(inj_section_name): + for section in workflow.cp.get_subsections(inj_section_name): inj_tag = section.upper() curr_tags = tags + [inj_tag] # Parse for options in ini file - injection_method = workflow.cp.get_opt_tags("workflow-injections", - "injections-method", - curr_tags) + injection_method = workflow.cp.get_opt_tags( + "workflow-injections", "injections-method", curr_tags + ) if injection_method in ["IN_WORKFLOW", "AT_RUNTIME"]: - exe = select_generic_executable(workflow, 'injections') - inj_job = exe(workflow.cp, inj_section_name, - out_dir=output_dir, ifos='HL', - tags=curr_tags) + exe = select_generic_executable(workflow, "injections") + inj_job = exe( + workflow.cp, + inj_section_name, + out_dir=output_dir, + ifos="HL", + tags=curr_tags, + ) if exe is PycbcCreateInjectionsExecutable: - config_urls = workflow.cp.get('workflow-injections', - section+'-config-files') - config_urls = config_urls.split(',') - config_files = FileList([resolve_url_to_file(cf.strip()) - for cf in config_urls]) + config_urls = workflow.cp.get( + "workflow-injections", section + "-config-files" + ) + config_urls = config_urls.split(",") + config_files = FileList( + [resolve_url_to_file(cf.strip()) for cf in config_urls] + ) node, inj_file = inj_job.create_node(config_files) else: node = inj_job.create_node(full_segment) @@ -240,15 +257,9 @@ def setup_injection_workflow(workflow, output_dir=None, inj_file = node.output_files[0] inj_files.append(inj_file) elif injection_method == "PREGENERATED": - file_attrs = { - 'ifos': ['HL'], - 'segs': full_segment, - 'tags': curr_tags - } + file_attrs = {"ifos": ["HL"], "segs": full_segment, "tags": curr_tags} injection_path = workflow.cp.get_opt_tags( - "workflow-injections", - "injections-pregenerated-file", - curr_tags + "workflow-injections", "injections-pregenerated-file", curr_tags ) curr_file = resolve_url_to_file(injection_path, attrs=file_attrs) inj_files.append(curr_file) diff --git a/pycbc/workflow/jobsetup.py b/pycbc/workflow/jobsetup.py index 68dd9c47d64..7c25e5818a9 100644 --- a/pycbc/workflow/jobsetup.py +++ b/pycbc/workflow/jobsetup.py @@ -28,21 +28,29 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/ahope.html """ -import math, os +import math +import os + import igwn_segments as segments + from pycbc.workflow.core import Executable, File, FileList, Node + def int_gps_time_to_str(t): - """Takes an integer GPS time, either given as int or lal.LIGOTimeGPS, and + """ + Takes an integer GPS time, either given as int or lal.LIGOTimeGPS, and converts it to a string. If a LIGOTimeGPS with nonzero decimal part is - given, raises a ValueError.""" + given, raises a ValueError. + """ int_t = int(t) - if abs(float(t - int_t)) > 0.: - raise ValueError('Need an integer GPS time, got %s' % str(t)) + if abs(float(t - int_t)) > 0.0: + raise ValueError("Need an integer GPS time, got %s" % str(t)) return str(int_t) + def select_tmpltbank_class(curr_exe): - """ This function returns a class that is appropriate for setting up + """ + This function returns a class that is appropriate for setting up template bank jobs within workflow. Parameters @@ -51,26 +59,30 @@ def select_tmpltbank_class(curr_exe): The name of the executable to be used for generating template banks. Returns - -------- + ------- exe_class : Sub-class of pycbc.workflow.core.Executable that holds utility functions appropriate for the given executable. Instances of the class ('jobs') **must** have methods * job.create_node() and * job.get_valid_times(ifo, ) + """ exe_to_class_map = { - 'pycbc_geom_nonspinbank' : PyCBCTmpltbankExecutable, - 'pycbc_aligned_stoch_bank': PyCBCTmpltbankExecutable + "pycbc_geom_nonspinbank": PyCBCTmpltbankExecutable, + "pycbc_aligned_stoch_bank": PyCBCTmpltbankExecutable, } try: return exe_to_class_map[curr_exe] except KeyError: raise NotImplementedError( - "No job class exists for executable %s, exiting" % curr_exe) + "No job class exists for executable %s, exiting" % curr_exe + ) + def select_matchedfilter_class(curr_exe): - """ This function returns a class that is appropriate for setting up + """ + This function returns a class that is appropriate for setting up matched-filtering jobs within workflow. Parameters @@ -79,28 +91,32 @@ def select_matchedfilter_class(curr_exe): The name of the matched filter executable to be used. Returns - -------- + ------- exe_class : Sub-class of pycbc.workflow.core.Executable that holds utility functions appropriate for the given executable. Instances of the class ('jobs') **must** have methods * job.create_node() and * job.get_valid_times(ifo, ) + """ exe_to_class_map = { - 'pycbc_inspiral' : PyCBCInspiralExecutable, - 'pycbc_inspiral_skymax' : PyCBCInspiralExecutable, - 'pycbc_multi_inspiral' : PyCBCMultiInspiralExecutable, + "pycbc_inspiral": PyCBCInspiralExecutable, + "pycbc_inspiral_skymax": PyCBCInspiralExecutable, + "pycbc_multi_inspiral": PyCBCMultiInspiralExecutable, } try: return exe_to_class_map[curr_exe] except KeyError: # also conceivable to introduce a default class?? raise NotImplementedError( - "No job class exists for executable %s, exiting" % curr_exe) + "No job class exists for executable %s, exiting" % curr_exe + ) + def select_generic_executable(workflow, exe_tag): - """ Returns a class that is appropriate for setting up jobs to run executables + """ + Returns a class that is appropriate for setting up jobs to run executables having specific tags in the workflow config. Executables should not be "specialized" jobs fitting into one of the select_XXX_class functions above, i.e. not a matched filter or template @@ -116,30 +132,41 @@ def select_generic_executable(workflow, exe_tag): the option giving the executable path in the [executables] section. Returns - -------- + ------- exe_class : Sub-class of pycbc.workflow.core.Executable that holds utility functions appropriate for the given executable. Instances of the class ('jobs') **must** have a method job.create_node() + """ exe_path = workflow.cp.get("executables", exe_tag) exe_name = os.path.basename(exe_path) exe_to_class_map = { - 'ligolw_add' : LigolwAddExecutable, - 'lalapps_inspinj' : LalappsInspinjExecutable, - 'pycbc_create_injections' : PycbcCreateInjectionsExecutable, - 'pycbc_condition_strain' : PycbcConditionStrainExecutable + "ligolw_add": LigolwAddExecutable, + "lalapps_inspinj": LalappsInspinjExecutable, + "pycbc_create_injections": PycbcCreateInjectionsExecutable, + "pycbc_condition_strain": PycbcConditionStrainExecutable, } try: return exe_to_class_map[exe_name] except KeyError: # Should we try some sort of default class?? raise NotImplementedError( - "No job class exists for executable %s, exiting" % exe_name) + "No job class exists for executable %s, exiting" % exe_name + ) -def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, - datafind_outs, parents=None, - allow_overlap=True): - """ This function sets up a set of single ifo jobs. A basic overview of how this + +def sngl_ifo_job_setup( + workflow, + ifo, + out_files, + curr_exe_job, + science_segs, + datafind_outs, + parents=None, + allow_overlap=True, +): + """ + This function sets up a set of single ifo jobs. A basic overview of how this works is as follows: * (1) Identify the length of data that each job needs to read in, and what @@ -158,7 +185,7 @@ def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, * END LOOPING OVER SCIENCE SEGMENTS Parameters - ----------- + ---------- workflow: pycbc.workflow.core.Workflow An instance of the Workflow class that manages the constructed workflow. ifo : string @@ -183,12 +210,12 @@ def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, at all. Returns - -------- + ------- out_files : pycbc.workflow.core.FileList A list of the files that will be generated by this step in the workflow. - """ + """ ########### (1) ############ # Get the times that can be analysed and needed data lengths data_length, valid_chunk, valid_length = identify_needed_data(curr_exe_job) @@ -199,14 +226,16 @@ def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, ########### (2) ############ # Initialize the class that identifies how many jobs are needed and the # shift between them. - segmenter = JobSegmenter(data_length, valid_chunk, valid_length, - curr_seg, curr_exe_job) + segmenter = JobSegmenter( + data_length, valid_chunk, valid_length, curr_seg, curr_exe_job + ) for job_num in range(segmenter.num_jobs): ############## (3) ############# # Figure out over what times this job will be valid for - job_valid_seg = segmenter.get_valid_times_for_job(job_num, - allow_overlap=allow_overlap) + job_valid_seg = segmenter.get_valid_times_for_job( + job_num, allow_overlap=allow_overlap + ) ############## (4) ############# # Get the data that this job should read in @@ -216,11 +245,14 @@ def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, # Identify parents/inputs to the job if parents: # Find the set of files with the best overlap - curr_parent = parents.find_outputs_in_range(ifo, job_valid_seg, - useSplitLists=True) + curr_parent = parents.find_outputs_in_range( + ifo, job_valid_seg, useSplitLists=True + ) if not curr_parent: - err_string = ("No parent jobs found overlapping %d to %d." - %(job_valid_seg[0], job_valid_seg[1])) + err_string = "No parent jobs found overlapping %d to %d." % ( + job_valid_seg[0], + job_valid_seg[1], + ) err_string += "\nThis is a bad error! Contact a developer." raise ValueError(err_string) else: @@ -228,15 +260,17 @@ def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, curr_dfouts = None if datafind_outs: - curr_dfouts = datafind_outs.find_all_output_in_range(ifo, - job_data_seg, useSplitLists=True) + curr_dfouts = datafind_outs.find_all_output_in_range( + ifo, job_data_seg, useSplitLists=True + ) if not curr_dfouts: - err_str = ("No datafind jobs found overlapping %d to %d." - %(job_data_seg[0],job_data_seg[1])) + err_str = "No datafind jobs found overlapping %d to %d." % ( + job_data_seg[0], + job_data_seg[1], + ) err_str += "\nThis shouldn't happen. Contact a developer." raise ValueError(err_str) - ############## (6) ############# # Make node and add to workflow @@ -247,26 +281,34 @@ def sngl_ifo_job_setup(workflow, ifo, out_files, curr_exe_job, science_segs, for parent in curr_parent: if len(curr_parent) != 1: - bank_tag = [t for t in parent.tags if 'bank' in t.lower()] + bank_tag = [t for t in parent.tags if "bank" in t.lower()] curr_exe_job.update_current_tags(bank_tag + exe_tags) # We should generate unique names automatically, but it is a # pain until we can set the output names for all Executables - node = curr_exe_job.create_node(job_data_seg, job_valid_seg, - parent=parent, - df_parents=curr_dfouts) + node = curr_exe_job.create_node( + job_data_seg, job_valid_seg, parent=parent, df_parents=curr_dfouts + ) workflow.add_node(node) curr_out_files = node.output_files # FIXME: Here we remove PSD files if they are coming through. # This should be done in a better way. On to-do list. - curr_out_files = [i for i in curr_out_files if 'PSD_FILE'\ - not in i.tags] + curr_out_files = [i for i in curr_out_files if "PSD_FILE" not in i.tags] out_files += curr_out_files return out_files -def multi_ifo_coherent_job_setup(workflow, out_files, curr_exe_job, - science_segs, datafind_outs, output_dir, - parents=None, slide_dict=None, tags=None): + +def multi_ifo_coherent_job_setup( + workflow, + out_files, + curr_exe_job, + science_segs, + datafind_outs, + output_dir, + parents=None, + slide_dict=None, + tags=None, +): """ Method for setting up coherent inspiral jobs. """ @@ -279,15 +321,15 @@ def multi_ifo_coherent_job_setup(workflow, out_files, curr_exe_job, skygrid_file = None input_files = FileList(datafind_outs) for f in datafind_outs: - if 'IPN_SKY_POINTS' in f.description: + if "IPN_SKY_POINTS" in f.description: ipn_sky_points = f input_files.remove(f) - elif 'vetoes' in f.description: + elif "vetoes" in f.description: input_files.remove(f) - elif 'INPUT_BANK_VETO_BANK' in f.description: + elif "INPUT_BANK_VETO_BANK" in f.description: bank_veto = f input_files.remove(f) - elif 'make_sky_grid' in f.description: + elif "make_sky_grid" in f.description: skygrid_file = f input_files.remove(f) @@ -297,11 +339,17 @@ def multi_ifo_coherent_job_setup(workflow, out_files, curr_exe_job, for split_bank in parents: tag = list(tags) tag.append(split_bank.tag_str) - node = curr_exe_job.create_node(data_seg, job_valid_seg, - parent=split_bank, dfParents=input_files, - bankVetoBank=bank_veto, - skygrid_file=skygrid_file, ipn_file=ipn_sky_points, - slide=slide_dict, tags=tag) + node = curr_exe_job.create_node( + data_seg, + job_valid_seg, + parent=split_bank, + dfParents=input_files, + bankVetoBank=bank_veto, + skygrid_file=skygrid_file, + ipn_file=ipn_sky_points, + slide=slide_dict, + tags=tag, + ) workflow.add_node(node) split_bank_counter += 1 curr_out_files.extend(node.output_files) @@ -311,10 +359,17 @@ def multi_ifo_coherent_job_setup(workflow, out_files, curr_exe_job, tag = list(tags) tag.append(inj_file.tag_str) tag.append(split_bank.tag_str) - node = curr_exe_job.create_node(data_seg, job_valid_seg, - parent=split_bank, inj_file=inj_file, tags=tag, - dfParents=input_files, bankVetoBank=bank_veto, - skygrid_file=skygrid_file, ipn_file=ipn_sky_points) + node = curr_exe_job.create_node( + data_seg, + job_valid_seg, + parent=split_bank, + inj_file=inj_file, + tags=tag, + dfParents=input_files, + bankVetoBank=bank_veto, + skygrid_file=skygrid_file, + ipn_file=ipn_sky_points, + ) workflow.add_node(node) split_bank_counter += 1 curr_out_files.extend(node.output_files) @@ -325,24 +380,25 @@ def multi_ifo_coherent_job_setup(workflow, out_files, curr_exe_job, # IWHNOTE: This will not be needed when coh_PTF is retired, but it is # okay to do this. It just means you can't access these files # later. - curr_out_files = [i for i in curr_out_files if 'PSD_FILE'\ - not in i.tags] + curr_out_files = [i for i in curr_out_files if "PSD_FILE" not in i.tags] out_files += curr_out_files return out_files + def identify_needed_data(curr_exe_job): - """ This function will identify the length of data that a specific + """ + This function will identify the length of data that a specific executable needs to analyse and what part of that data is valid (ie. inspiral doesn't analyse the first or last 8s of data it reads in). Parameters - ----------- + ---------- curr_exe_job : Job An instance of the Job class that has a get_valid times method. Returns - -------- + ------- dataLength : float The amount of data (in seconds) that each instance of the job must read in. @@ -354,6 +410,7 @@ def identify_needed_data(curr_exe_job): valid_length : float The maximum length of data each job can be valid for. This is abs(valid_segment). + """ # Set up the condorJob class for the current executable data_lengths, valid_chunks = curr_exe_job.get_valid_times() @@ -365,35 +422,42 @@ def identify_needed_data(curr_exe_job): return data_lengths, valid_chunks, valid_lengths -class JobSegmenter(object): - """ This class is used when running sngl_ifo_job_setup to determine what times +class JobSegmenter: + """ + This class is used when running sngl_ifo_job_setup to determine what times should be analysed be each job and what data is needed. """ - def __init__(self, data_lengths, valid_chunks, valid_lengths, curr_seg, - curr_exe_class): - """ Initialize class. """ + def __init__( + self, data_lengths, valid_chunks, valid_lengths, curr_seg, curr_exe_class + ): + """Initialize class.""" self.exe_class = curr_exe_class self.curr_seg = curr_seg self.curr_seg_length = float(abs(curr_seg)) - self.data_length, self.valid_chunk, self.valid_length = \ - self.pick_tile_size(self.curr_seg_length, data_lengths, - valid_chunks, valid_lengths) + self.data_length, self.valid_chunk, self.valid_length = self.pick_tile_size( + self.curr_seg_length, data_lengths, valid_chunks, valid_lengths + ) self.data_chunk = segments.segment([0, self.data_length]) self.data_loss = self.data_length - abs(self.valid_chunk) if self.data_loss < 0: - raise ValueError("pycbc.workflow.jobsetup needs fixing! Please contact a developer") + raise ValueError( + "pycbc.workflow.jobsetup needs fixing! Please contact a developer" + ) if self.curr_seg_length < self.data_length: self.num_jobs = 0 return # How many jobs do we need - self.num_jobs = int( math.ceil( (self.curr_seg_length \ - - self.data_loss) / float(self.valid_length) )) + self.num_jobs = int( + math.ceil( + (self.curr_seg_length - self.data_loss) / float(self.valid_length) + ) + ) if self.curr_seg_length == self.data_length: # If the segment length is identical to the data length then I @@ -401,62 +465,61 @@ def __init__(self, data_lengths, valid_chunks, valid_lengths, curr_seg, self.job_time_shift = 0 else: # What is the incremental shift between jobs - self.job_time_shift = (self.curr_seg_length - self.data_length) / \ - float(self.num_jobs - 1) + self.job_time_shift = (self.curr_seg_length - self.data_length) / float( + self.num_jobs - 1 + ) def pick_tile_size(self, seg_size, data_lengths, valid_chunks, valid_lengths): - """ Choose job tiles size based on science segment length """ - + """Choose job tiles size based on science segment length""" if len(valid_lengths) == 1: return data_lengths[0], valid_chunks[0], valid_lengths[0] - else: - # Pick the tile size that is closest to 1/3 of the science segment - target_size = seg_size / 3 - pick, pick_diff = 0, abs(valid_lengths[0] - target_size) - for i, size in enumerate(valid_lengths): - if abs(size - target_size) < pick_diff: - pick, pick_diff = i, abs(size - target_size) - return data_lengths[pick], valid_chunks[pick], valid_lengths[pick] + # Pick the tile size that is closest to 1/3 of the science segment + target_size = seg_size / 3 + pick, pick_diff = 0, abs(valid_lengths[0] - target_size) + for i, size in enumerate(valid_lengths): + if abs(size - target_size) < pick_diff: + pick, pick_diff = i, abs(size - target_size) + return data_lengths[pick], valid_chunks[pick], valid_lengths[pick] def get_valid_times_for_job(self, num_job, allow_overlap=True): - """ Get the times for which this job is valid. """ + """Get the times for which this job is valid.""" # small factor of 0.0001 to avoid float round offs causing us to # miss a second at end of segments. - shift_dur = self.curr_seg[0] + int(self.job_time_shift * num_job\ - + 0.0001) + shift_dur = self.curr_seg[0] + int(self.job_time_shift * num_job + 0.0001) job_valid_seg = self.valid_chunk.shift(shift_dur) # If we need to recalculate the valid times to avoid overlap if not allow_overlap: - data_per_job = (self.curr_seg_length - self.data_loss) / \ - float(self.num_jobs) - lower_boundary = num_job*data_per_job + \ - self.valid_chunk[0] + self.curr_seg[0] + data_per_job = (self.curr_seg_length - self.data_loss) / float( + self.num_jobs + ) + lower_boundary = ( + num_job * data_per_job + self.valid_chunk[0] + self.curr_seg[0] + ) upper_boundary = data_per_job + lower_boundary # NOTE: Convert to int after calculating both boundaries # small factor of 0.0001 to avoid float round offs causing us to # miss a second at end of segments. lower_boundary = int(lower_boundary) upper_boundary = int(upper_boundary + 0.0001) - if lower_boundary < job_valid_seg[0] or \ - upper_boundary > job_valid_seg[1]: - err_msg = ("Workflow is attempting to generate output " - "from a job at times where it is not valid.") + if lower_boundary < job_valid_seg[0] or upper_boundary > job_valid_seg[1]: + err_msg = ( + "Workflow is attempting to generate output " + "from a job at times where it is not valid." + ) raise ValueError(err_msg) - job_valid_seg = segments.segment([lower_boundary, - upper_boundary]) + job_valid_seg = segments.segment([lower_boundary, upper_boundary]) return job_valid_seg def get_data_times_for_job(self, num_job): - """ Get the data that this job will read in. """ + """Get the data that this job will read in.""" # small factor of 0.0001 to avoid float round offs causing us to # miss a second at end of segments. - shift_dur = self.curr_seg[0] + int(self.job_time_shift * num_job\ - + 0.0001) + shift_dur = self.curr_seg[0] + int(self.job_time_shift * num_job + 0.0001) job_data_seg = self.data_chunk.shift(shift_dur) # Sanity check that all data is used if num_job == 0: if job_data_seg[0] != self.curr_seg[0]: - err= "Job is not using data from the start of the " + err = "Job is not using data from the start of the " err += "science segment. It should be using all data." raise ValueError(err) if num_job == (self.num_jobs - 1): @@ -465,118 +528,133 @@ def get_data_times_for_job(self, num_job): err += "science segment. It should be using all data." raise ValueError(err) - if hasattr(self.exe_class, 'zero_pad_data_extend'): - job_data_seg = self.exe_class.zero_pad_data_extend(job_data_seg, - self.curr_seg) + if hasattr(self.exe_class, "zero_pad_data_extend"): + job_data_seg = self.exe_class.zero_pad_data_extend( + job_data_seg, self.curr_seg + ) return job_data_seg class PyCBCInspiralExecutable(Executable): - """ The class used to create jobs for pycbc_inspiral Executable. """ + """The class used to create jobs for pycbc_inspiral Executable.""" current_retention_level = Executable.ALL_TRIGGERS - time_dependent_options = ['--channel-name'] - - def __init__(self, cp, exe_name, ifo=None, out_dir=None, - injection_file=None, tags=None, reuse_executable=False): + time_dependent_options = ["--channel-name"] + + def __init__( + self, + cp, + exe_name, + ifo=None, + out_dir=None, + injection_file=None, + tags=None, + reuse_executable=False, + ): if tags is None: tags = [] - super().__init__(cp, exe_name, ifo, out_dir, tags=tags, - reuse_executable=reuse_executable, - set_submit_subdir=False) + super().__init__( + cp, + exe_name, + ifo, + out_dir, + tags=tags, + reuse_executable=reuse_executable, + set_submit_subdir=False, + ) self.cp = cp self.injection_file = injection_file - self.ext = '.hdf' + self.ext = ".hdf" self.num_threads = 1 - if self.get_opt('processing-scheme') is not None: - stxt = self.get_opt('processing-scheme') - if len(stxt.split(':')) > 1: - self.num_threads = stxt.split(':')[1] + if self.get_opt("processing-scheme") is not None: + stxt = self.get_opt("processing-scheme") + if len(stxt.split(":")) > 1: + self.num_threads = stxt.split(":")[1] - def create_node(self, data_seg, valid_seg, parent=None, df_parents=None, - tags=None): + def create_node(self, data_seg, valid_seg, parent=None, df_parents=None, tags=None): if tags is None: tags = [] node = Node(self, valid_seg=valid_seg) - if not self.has_opt('pad-data'): - raise ValueError("The option pad-data is a required option of " - "%s. Please check the ini file." % self.name) - pad_data = int(self.get_opt('pad-data')) + if not self.has_opt("pad-data"): + raise ValueError( + "The option pad-data is a required option of " + "%s. Please check the ini file." % self.name + ) + pad_data = int(self.get_opt("pad-data")) # set remaining options flags - node.add_opt('--gps-start-time', - int_gps_time_to_str(data_seg[0] + pad_data)) - node.add_opt('--gps-end-time', - int_gps_time_to_str(data_seg[1] - pad_data)) - node.add_opt('--trig-start-time', int_gps_time_to_str(valid_seg[0])) - node.add_opt('--trig-end-time', int_gps_time_to_str(valid_seg[1])) + node.add_opt("--gps-start-time", int_gps_time_to_str(data_seg[0] + pad_data)) + node.add_opt("--gps-end-time", int_gps_time_to_str(data_seg[1] - pad_data)) + node.add_opt("--trig-start-time", int_gps_time_to_str(valid_seg[0])) + node.add_opt("--trig-end-time", int_gps_time_to_str(valid_seg[1])) if self.injection_file is not None: - node.add_input_opt('--injection-file', self.injection_file) + node.add_input_opt("--injection-file", self.injection_file) # set the input and output files fil = node.new_output_file_opt( valid_seg, self.ext, - '--output', + "--output", tags=tags, store_file=self.retain_files, - use_tmp_subdirs=True + use_tmp_subdirs=True, ) # For inspiral jobs we overrwrite the "relative.submit.dir" # attribute to avoid too many files in one sub-directory - curr_rel_dir = fil.name.split('/')[0] - node.add_profile('pegasus', 'relative.submit.dir', - self.pegasus_name + '_' + curr_rel_dir) + curr_rel_dir = fil.name.split("/")[0] + node.add_profile( + "pegasus", "relative.submit.dir", self.pegasus_name + "_" + curr_rel_dir + ) # Must ensure this is not a LIGOGPS as JSON won't understand it data_seg = segments.segment([int(data_seg[0]), int(data_seg[1])]) - fil.add_metadata('data_seg', data_seg) - node.add_input_opt('--bank-file', parent) + fil.add_metadata("data_seg", data_seg) + node.add_input_opt("--bank-file", parent) if df_parents is not None: - node.add_input_list_opt('--frame-files', df_parents) + node.add_input_list_opt("--frame-files", df_parents) return node def get_valid_times(self): - """ Determine possible dimensions of needed input and valid output - """ - - if self.cp.has_option('workflow-matchedfilter', - 'min-analysis-segments'): - min_analysis_segs = int(self.cp.get('workflow-matchedfilter', - 'min-analysis-segments')) + """Determine possible dimensions of needed input and valid output""" + if self.cp.has_option("workflow-matchedfilter", "min-analysis-segments"): + min_analysis_segs = int( + self.cp.get("workflow-matchedfilter", "min-analysis-segments") + ) else: min_analysis_segs = 0 - if self.cp.has_option('workflow-matchedfilter', - 'max-analysis-segments'): - max_analysis_segs = int(self.cp.get('workflow-matchedfilter', - 'max-analysis-segments')) + if self.cp.has_option("workflow-matchedfilter", "max-analysis-segments"): + max_analysis_segs = int( + self.cp.get("workflow-matchedfilter", "max-analysis-segments") + ) else: # Choose ridiculously large default value max_analysis_segs = 1000 - if self.cp.has_option('workflow-matchedfilter', 'min-analysis-length'): - min_analysis_length = int(self.cp.get('workflow-matchedfilter', - 'min-analysis-length')) + if self.cp.has_option("workflow-matchedfilter", "min-analysis-length"): + min_analysis_length = int( + self.cp.get("workflow-matchedfilter", "min-analysis-length") + ) else: min_analysis_length = 0 - if self.cp.has_option('workflow-matchedfilter', 'max-analysis-length'): - max_analysis_length = int(self.cp.get('workflow-matchedfilter', - 'max-analysis-length')) + if self.cp.has_option("workflow-matchedfilter", "max-analysis-length"): + max_analysis_length = int( + self.cp.get("workflow-matchedfilter", "max-analysis-length") + ) else: # Choose a ridiculously large default value max_analysis_length = 100000 - segment_length = int(self.get_opt('segment-length')) + segment_length = int(self.get_opt("segment-length")) pad_data = 0 - if self.has_opt('pad-data'): - pad_data += int(self.get_opt('pad-data')) + if self.has_opt("pad-data"): + pad_data += int(self.get_opt("pad-data")) # NOTE: Currently the tapered data is ignored as it is short and # will lie within the segment start/end pad. This means that @@ -588,15 +666,15 @@ def get_valid_times(self): # pad take normal values. When using zero-padding this data will # be used for SNR generation. - #if self.has_opt('taper-data'): + # if self.has_opt('taper-data'): # pad_data += int(self.get_opt( 'taper-data' )) - if self.has_opt('allow-zero-padding'): - self.zero_padding=True + if self.has_opt("allow-zero-padding"): + self.zero_padding = True else: - self.zero_padding=False + self.zero_padding = False - start_pad = int(self.get_opt( 'segment-start-pad')) - end_pad = int(self.get_opt('segment-end-pad')) + start_pad = int(self.get_opt("segment-start-pad")) + end_pad = int(self.get_opt("segment-end-pad")) seg_ranges = range(min_analysis_segs, max_analysis_segs + 1) data_lengths = [] @@ -604,16 +682,17 @@ def get_valid_times(self): for nsegs in seg_ranges: analysis_length = (segment_length - start_pad - end_pad) * nsegs if not self.zero_padding: - data_length = analysis_length + pad_data * 2 \ - + start_pad + end_pad + data_length = analysis_length + pad_data * 2 + start_pad + end_pad start = pad_data + start_pad end = data_length - pad_data - end_pad else: data_length = analysis_length + pad_data * 2 start = pad_data end = data_length - pad_data - if data_length > max_analysis_length: continue - if data_length < min_analysis_length: continue + if data_length > max_analysis_length: + continue + if data_length < min_analysis_length: + continue data_lengths += [data_length] valid_regions += [segments.segment(start, end)] # If min_analysis_length is given, ensure that it is added as an option @@ -632,7 +711,8 @@ def get_valid_times(self): return data_lengths, valid_regions def zero_pad_data_extend(self, job_data_seg, curr_seg): - """When using zero padding, *all* data is analysable, but the setup + """ + When using zero padding, *all* data is analysable, but the setup functions must include the padding data where it is available so that we are not zero-padding in the middle of science segments. This function takes a job_data_seg, that is chosen for a particular node @@ -641,13 +721,13 @@ def zero_pad_data_extend(self, job_data_seg, curr_seg): """ if self.zero_padding is False: return job_data_seg - else: - start_pad = int(self.get_opt( 'segment-start-pad')) - end_pad = int(self.get_opt('segment-end-pad')) - new_data_start = max(curr_seg[0], job_data_seg[0] - start_pad) - new_data_end = min(curr_seg[1], job_data_seg[1] + end_pad) - new_data_seg = segments.segment([new_data_start, new_data_end]) - return new_data_seg + start_pad = int(self.get_opt("segment-start-pad")) + end_pad = int(self.get_opt("segment-end-pad")) + new_data_start = max(curr_seg[0], job_data_seg[0] - start_pad) + new_data_end = min(curr_seg[1], job_data_seg[1] + end_pad) + new_data_seg = segments.segment([new_data_start, new_data_end]) + return new_data_seg + # FIXME: This is probably misnamed, this is really GRBInspiralExectuable. # There's nothing coherent here, it's just that data segment stuff is @@ -657,118 +737,149 @@ class PyCBCMultiInspiralExecutable(Executable): The class responsible for setting up jobs for the pycbc_multi_inspiral executable. """ + current_retention_level = Executable.ALL_TRIGGERS # bank-veto-bank-file is a file input option for pycbc_multi_inspiral - file_input_options = Executable.file_input_options + \ - ['--bank-veto-bank-file'] - - def __init__(self, cp, name, ifo=None, injection_file=None, - gate_files=None, out_dir=None, tags=None): + file_input_options = Executable.file_input_options + ["--bank-veto-bank-file"] + + def __init__( + self, + cp, + name, + ifo=None, + injection_file=None, + gate_files=None, + out_dir=None, + tags=None, + ): if tags is None: tags = [] super().__init__(cp, name, ifo, out_dir=out_dir, tags=tags) self.injection_file = injection_file - self.data_seg = segments.segment(int(cp.get('workflow', 'start-time')), - int(cp.get('workflow', 'end-time'))) + self.data_seg = segments.segment( + int(cp.get("workflow", "start-time")), int(cp.get("workflow", "end-time")) + ) self.num_threads = 1 - def create_node(self, data_seg, valid_seg, parent=None, inj_file=None, - dfParents=None, bankVetoBank=None, - skygrid_file=None, ipn_file=None, - slide=None, tags=None): + def create_node( + self, + data_seg, + valid_seg, + parent=None, + inj_file=None, + dfParents=None, + bankVetoBank=None, + skygrid_file=None, + ipn_file=None, + slide=None, + tags=None, + ): if tags is None: tags = [] node = Node(self) if not dfParents: - raise ValueError("%s must be supplied with frame or cache files" - % self.name) + raise ValueError( + "%s must be supplied with frame or cache files" % self.name + ) # If doing single IFO search, make sure slides are disabled - if len(self.ifo_list) < 2 and \ - (self.get_opt('--do-short-slides') is not None or \ - self.get_opt('--short-slide-offset') is not None): - raise ValueError("Cannot run with time slides in a single IFO " - "configuration! Please edit your configuration " - "file accordingly.") + if len(self.ifo_list) < 2 and ( + self.get_opt("--do-short-slides") is not None + or self.get_opt("--short-slide-offset") is not None + ): + raise ValueError( + "Cannot run with time slides in a single IFO " + "configuration! Please edit your configuration " + "file accordingly." + ) # Set instuments node.add_opt("--instruments", " ".join(self.ifo_list)) - pad_data = self.get_opt('pad-data') + pad_data = self.get_opt("pad-data") if pad_data is None: - raise ValueError("The option pad-data is a required option of " - "%s. Please check the ini file." % self.name) + raise ValueError( + "The option pad-data is a required option of " + "%s. Please check the ini file." % self.name + ) # Feed in bank_veto_bank.xml, if given - if self.cp.has_option('workflow-inspiral', 'bank-veto-bank-file'): - node.add_input_opt('--bank-veto-bank-file', bankVetoBank) + if self.cp.has_option("workflow-inspiral", "bank-veto-bank-file"): + node.add_input_opt("--bank-veto-bank-file", bankVetoBank) # Set time options - node.add_opt('--gps-start-time', data_seg[0] + int(pad_data)) - node.add_opt('--gps-end-time', data_seg[1] - int(pad_data)) - node.add_opt('--trig-start-time', valid_seg[0]) - node.add_opt('--trig-end-time', valid_seg[1]) - node.add_opt('--trigger-time', self.cp.get('workflow', 'trigger-time')) + node.add_opt("--gps-start-time", data_seg[0] + int(pad_data)) + node.add_opt("--gps-end-time", data_seg[1] - int(pad_data)) + node.add_opt("--trig-start-time", valid_seg[0]) + node.add_opt("--trig-end-time", valid_seg[1]) + node.add_opt("--trigger-time", self.cp.get("workflow", "trigger-time")) # Set the input and output files - node.new_output_file_opt(data_seg, '.hdf', '--output', - tags=tags, store_file=self.retain_files) - node.add_input_opt('--bank-file', parent, ) + node.new_output_file_opt( + data_seg, ".hdf", "--output", tags=tags, store_file=self.retain_files + ) + node.add_input_opt( + "--bank-file", + parent, + ) if dfParents is not None: - frame_arg = '--frame-files' + frame_arg = "--frame-files" for frame_file in dfParents: frame_arg += f" {frame_file.ifo}:{frame_file.name}" node.add_input(frame_file) node.add_arg(frame_arg) if skygrid_file is not None: - node.add_input_opt('--sky-grid', skygrid_file) + node.add_input_opt("--sky-grid", skygrid_file) if ipn_file is not None: - node.add_input_opt('--sky-positions-file', ipn_file) + node.add_input_opt("--sky-positions-file", ipn_file) if inj_file is not None: - if self.get_opt('--do-short-slides') is not None or \ - self.get_opt('--short-slide-offset') is not None: - raise ValueError("Cannot run with short slides in an " - "injection job. Please edit your " - "configuration file accordingly.") - node.add_input_opt('--injection-file', inj_file) + if ( + self.get_opt("--do-short-slides") is not None + or self.get_opt("--short-slide-offset") is not None + ): + raise ValueError( + "Cannot run with short slides in an " + "injection job. Please edit your " + "configuration file accordingly." + ) + node.add_input_opt("--injection-file", inj_file) if slide is not None: for ifo in self.ifo_list: - node.add_opt('--%s-slide-segment' % ifo.lower(), slide[ifo]) + node.add_opt("--%s-slide-segment" % ifo.lower(), slide[ifo]) # Channels channel_names = {} for ifo in self.ifo_list: channel_names[ifo] = self.cp.get_opt_tags( - "workflow", "%s-channel-name" % ifo.lower(), "") - channel_names_str = \ - " ".join([val for key, val in channel_names.items()]) + "workflow", "%s-channel-name" % ifo.lower(), "" + ) + channel_names_str = " ".join([val for key, val in channel_names.items()]) node.add_opt("--channel-name", channel_names_str) return node def get_valid_times(self): - pad_data = int(self.get_opt('pad-data')) + pad_data = int(self.get_opt("pad-data")) if self.has_opt("segment-start-pad"): pad_data = int(self.get_opt("pad-data")) start_pad = int(self.get_opt("segment-start-pad")) end_pad = int(self.get_opt("segment-end-pad")) valid_start = self.data_seg[0] + pad_data + start_pad valid_end = self.data_seg[1] - pad_data - end_pad - elif self.has_opt('analyse-segment-end'): + elif self.has_opt("analyse-segment-end"): safety = 1 - deadtime = int(self.get_opt('segment-length')) / 2 - spec_len = int(self.get_opt('inverse-spec-length')) / 2 - valid_start = (self.data_seg[0] + deadtime - spec_len + pad_data - - safety) + deadtime = int(self.get_opt("segment-length")) / 2 + spec_len = int(self.get_opt("inverse-spec-length")) / 2 + valid_start = self.data_seg[0] + deadtime - spec_len + pad_data - safety valid_end = self.data_seg[1] - spec_len - pad_data - safety else: - overlap = int(self.get_opt('segment-length')) / 4 + overlap = int(self.get_opt("segment-length")) / 4 valid_start = self.data_seg[0] + overlap + pad_data valid_end = self.data_seg[1] - overlap - pad_data @@ -776,13 +887,23 @@ def get_valid_times(self): class PyCBCTmpltbankExecutable(Executable): - """ The class used to create jobs for pycbc_geom_nonspin_bank Executable and + """ + The class used to create jobs for pycbc_geom_nonspin_bank Executable and any other Executables using the same command line option groups. """ current_retention_level = Executable.MERGED_TRIGGERS - def __init__(self, cp, exe_name, ifo=None, out_dir=None, - tags=None, write_psd=False, psd_files=None): + + def __init__( + self, + cp, + exe_name, + ifo=None, + out_dir=None, + tags=None, + write_psd=False, + psd_files=None, + ): if tags is None: tags = [] super().__init__(cp, exe_name, ifo, out_dir, tags=tags) @@ -796,44 +917,55 @@ def create_node(self, data_seg, valid_seg, parent=None, df_parents=None, tags=No node = Node(self) if not df_parents: - raise ValueError("%s must be supplied with data file(s)" - % self.name) + raise ValueError("%s must be supplied with data file(s)" % self.name) - pad_data = int(self.get_opt('pad-data')) + pad_data = int(self.get_opt("pad-data")) if pad_data is None: - raise ValueError("The option pad-data is a required option of " - "%s. Please check the ini file." % self.name) + raise ValueError( + "The option pad-data is a required option of " + "%s. Please check the ini file." % self.name + ) # set the remaining option flags - node.add_opt('--gps-start-time', - int_gps_time_to_str(data_seg[0] + pad_data)) - node.add_opt('--gps-end-time', - int_gps_time_to_str(data_seg[1] - pad_data)) + node.add_opt("--gps-start-time", int_gps_time_to_str(data_seg[0] + pad_data)) + node.add_opt("--gps-end-time", int_gps_time_to_str(data_seg[1] - pad_data)) # set the input and output files # Add the PSD file if needed if self.write_psd: - node.new_output_file_opt(valid_seg, '.txt', '--psd-output', - tags=tags+['PSD_FILE'], store_file=self.retain_files) - node.new_output_file_opt(valid_seg, '.xml.gz', '--output-file', - tags=tags, store_file=self.retain_files) - node.add_input_list_opt('--frame-files', df_parents) + node.new_output_file_opt( + valid_seg, + ".txt", + "--psd-output", + tags=tags + ["PSD_FILE"], + store_file=self.retain_files, + ) + node.new_output_file_opt( + valid_seg, + ".xml.gz", + "--output-file", + tags=tags, + store_file=self.retain_files, + ) + node.add_input_list_opt("--frame-files", df_parents) return node def create_nodata_node(self, valid_seg, tags=None): - """ A simplified version of create_node that creates a node that does + """ + A simplified version of create_node that creates a node that does not need to read in data. Parameters - ----------- + ---------- valid_seg : igwn_segments.segment The segment over which to declare the node valid. Usually this would be the duration of the analysis. Returns - -------- + ------- node : pycbc.workflow.core.Node The instance corresponding to the created node. + """ if tags is None: tags = [] @@ -842,12 +974,17 @@ def create_nodata_node(self, valid_seg, tags=None): # Set the output file # Add the PSD file if needed if self.write_psd: - node.new_output_file_opt(valid_seg, '.txt', '--psd-output', - tags=tags+['PSD_FILE'], - store_file=self.retain_files) - - node.new_output_file_opt(valid_seg, '.xml.gz', '--output-file', - store_file=self.retain_files) + node.new_output_file_opt( + valid_seg, + ".txt", + "--psd-output", + tags=tags + ["PSD_FILE"], + store_file=self.retain_files, + ) + + node.new_output_file_opt( + valid_seg, ".xml.gz", "--output-file", store_file=self.retain_files + ) if self.psd_files is not None: should_add = False @@ -860,14 +997,13 @@ def create_nodata_node(self, valid_seg, tags=None): should_add = True if should_add: - node.add_input_opt('--psd-file', psd_file) + node.add_input_opt("--psd-file", psd_file) return node def get_valid_times(self): - pad_data = int(self.get_opt( 'pad-data')) - analysis_length = int(self.cp.get('workflow-tmpltbank', - 'analysis-length')) + pad_data = int(self.get_opt("pad-data")) + analysis_length = int(self.cp.get("workflow-tmpltbank", "analysis-length")) data_length = analysis_length + pad_data * 2 start = pad_data end = data_length - pad_data @@ -875,12 +1011,13 @@ def get_valid_times(self): class LigolwAddExecutable(Executable): - """ The class used to create nodes for the ligolw_add Executable. """ + """The class used to create nodes for the ligolw_add Executable.""" current_retention_level = Executable.INTERMEDIATE_PRODUCT - def create_node(self, jobSegment, input_files, output=None, - use_tmp_subdirs=True, tags=None): + def create_node( + self, jobSegment, input_files, output=None, use_tmp_subdirs=True, tags=None + ): if tags is None: tags = [] node = Node(self) @@ -893,11 +1030,16 @@ def create_node(self, jobSegment, input_files, output=None, node.add_input_arg(fil) if output: - node.add_output_opt('--output', output) + node.add_output_opt("--output", output) else: - node.new_output_file_opt(jobSegment, '.xml.gz', '--output', - tags=tags, store_file=self.retain_files, - use_tmp_subdirs=use_tmp_subdirs) + node.new_output_file_opt( + jobSegment, + ".xml.gz", + "--output", + tags=tags, + store_file=self.retain_files, + use_tmp_subdirs=use_tmp_subdirs, + ) return node @@ -905,6 +1047,7 @@ class PycbcSplitInspinjExecutable(Executable): """ The class responsible for running the pycbc_split_inspinj executable """ + current_retention_level = Executable.INTERMEDIATE_PRODUCT def __init__(self, cp, exe_name, num_splits, ifo=None, out_dir=None): @@ -916,7 +1059,7 @@ def create_node(self, parent, tags=None): tags = [] node = Node(self) - node.add_input_opt('--input-file', parent) + node.add_input_opt("--input-file", parent) if parent.name.endswith("gz"): ext = ".xml.gz" @@ -925,15 +1068,21 @@ def create_node(self, parent, tags=None): out_files = FileList([]) for i in range(self.num_splits): - curr_tag = 'split%d' % i + curr_tag = "split%d" % i curr_tags = parent.tags + [curr_tag] job_tag = parent.description + "_" + self.name.upper() - out_file = File(parent.ifo_list, job_tag, parent.segment, - extension=ext, directory=self.out_dir, - tags=curr_tags, store_file=self.retain_files) + out_file = File( + parent.ifo_list, + job_tag, + parent.segment, + extension=ext, + directory=self.out_dir, + tags=curr_tags, + store_file=self.retain_files, + ) out_files.append(out_file) - node.add_output_list_opt('--output-files', out_files) + node.add_output_list_opt("--output-files", out_files) return node @@ -941,8 +1090,10 @@ class LalappsInspinjExecutable(Executable): """ The class used to create jobs for the lalapps_inspinj Executable. """ + current_retention_level = Executable.FINAL_RESULT - extension = '.xml' + extension = ".xml" + def create_node(self, segment, exttrig_file=None, tags=None): if tags is None: tags = [] @@ -952,56 +1103,64 @@ def create_node(self, segment, exttrig_file=None, tags=None): # This allows the desired number of injections to be given explicitly # in the config file. Used for coh_PTF as segment length is unknown # before run time. - if self.get_opt('write-compress') is not None: - self.extension = '.xml.gz' + if self.get_opt("write-compress") is not None: + self.extension = ".xml.gz" # Check if these injections are using trigger information to choose # sky positions for the simulated signals - if (self.get_opt('l-distr') == 'exttrig' and exttrig_file is not None \ - and 'trigger' in exttrig_file.description): + if ( + self.get_opt("l-distr") == "exttrig" + and exttrig_file is not None + and "trigger" in exttrig_file.description + ): # Use an XML file containing trigger information triggered = True - node.add_input_opt('--exttrig-file', exttrig_file) - elif (self.get_opt('l-distr') == 'ipn' and exttrig_file is not None \ - and 'IPN' in exttrig_file.description): + node.add_input_opt("--exttrig-file", exttrig_file) + elif ( + self.get_opt("l-distr") == "ipn" + and exttrig_file is not None + and "IPN" in exttrig_file.description + ): # Use an IPN sky points file triggered = True - node.add_input_opt('--ipn-file', exttrig_file) - elif (self.get_opt('l-distr') != 'exttrig') \ - and (self.get_opt('l-distr') != 'ipn' and not \ - self.has_opt('ipn-file')): + node.add_input_opt("--ipn-file", exttrig_file) + elif (self.get_opt("l-distr") != "exttrig") and ( + self.get_opt("l-distr") != "ipn" and not self.has_opt("ipn-file") + ): # Use no trigger information for generating injections triggered = False else: err_msg = "The argument 'l-distr' passed to the " err_msg += "%s job has the value " % self.tagged_name - err_msg += "'%s' but you have not " % self.get_opt('l-distr') + err_msg += "'%s' but you have not " % self.get_opt("l-distr") err_msg += "provided the corresponding ExtTrig or IPN file. " err_msg += "Please check your configuration files and try again." raise ValueError(err_msg) if triggered: - num_injs = int(self.cp.get_opt_tags('workflow-injections', - 'num-injs', curr_tags)) + num_injs = int( + self.cp.get_opt_tags("workflow-injections", "num-injs", curr_tags) + ) inj_tspace = float(segment[1] - segment[0]) / num_injs - node.add_opt('--time-interval', inj_tspace) - node.add_opt('--time-step', inj_tspace) + node.add_opt("--time-interval", inj_tspace) + node.add_opt("--time-step", inj_tspace) - node.new_output_file_opt(segment, self.extension, '--output', - store_file=self.retain_files) + node.new_output_file_opt( + segment, self.extension, "--output", store_file=self.retain_files + ) - node.add_opt('--gps-start-time', int_gps_time_to_str(segment[0])) - node.add_opt('--gps-end-time', int_gps_time_to_str(segment[1])) + node.add_opt("--gps-start-time", int_gps_time_to_str(segment[0])) + node.add_opt("--gps-end-time", int_gps_time_to_str(segment[1])) return node class PycbcSplitBankExecutable(Executable): - """ The class responsible for creating jobs for pycbc_hdf5_splitbank. """ + """The class responsible for creating jobs for pycbc_hdf5_splitbank.""" - extension = '.hdf' + extension = ".hdf" current_retention_level = Executable.ALL_TRIGGERS - def __init__(self, cp, exe_name, num_banks, - ifo=None, out_dir=None): + + def __init__(self, cp, exe_name, num_banks, ifo=None, out_dir=None): super().__init__(cp, exe_name, ifo, out_dir, tags=[]) self.num_banks = int(num_banks) @@ -1015,40 +1174,47 @@ def create_node(self, bank, tags=None): The File containing the template bank to be split Returns - -------- + ------- node : pycbc.workflow.core.Node The node to run the job + """ if tags is None: tags = [] node = Node(self) - node.add_input_opt('--bank-file', bank) + node.add_input_opt("--bank-file", bank) # Get the output (taken from inspiral.py) out_files = FileList([]) n_dp = math.ceil(math.log10(self.num_banks)) - for i in range( 0, self.num_banks): - curr_tag = (f'bank%0{n_dp}d') % (i) + for i in range(self.num_banks): + curr_tag = (f"bank%0{n_dp}d") % (i) # FIXME: What should the tags actually be? The job.tags values are # currently ignored. curr_tags = bank.tags + [curr_tag] + tags job_tag = bank.description + "_" + self.name.upper() - out_file = File(bank.ifo_list, job_tag, bank.segment, - extension=self.extension, directory=self.out_dir, - tags=curr_tags, store_file=self.retain_files) + out_file = File( + bank.ifo_list, + job_tag, + bank.segment, + extension=self.extension, + directory=self.out_dir, + tags=curr_tags, + store_file=self.retain_files, + ) out_files.append(out_file) - node.add_output_list_opt('--output-filenames', out_files) + node.add_output_list_opt("--output-filenames", out_files) return node class PycbcSplitBankXmlExecutable(PycbcSplitBankExecutable): - """ Subclass resonsible for creating jobs for pycbc_splitbank. """ + """Subclass resonsible for creating jobs for pycbc_splitbank.""" - extension='.xml.gz' + extension = ".xml.gz" class PycbcConditionStrainExecutable(Executable): - """ The class responsible for creating jobs for pycbc_condition_strain. """ + """The class responsible for creating jobs for pycbc_condition_strain.""" current_retention_level = Executable.ALL_TRIGGERS @@ -1058,36 +1224,47 @@ def create_node(self, input_files, tags=None): node = Node(self) start_time = self.cp.get("workflow", "start-time") end_time = self.cp.get("workflow", "end-time") - node.add_opt('--gps-start-time', start_time) - node.add_opt('--gps-end-time', end_time) - node.add_input_list_opt('--frame-files', input_files) - - out_file = File(self.ifo, "gated", - segments.segment(int(start_time), int(end_time)), - directory=self.out_dir, store_file=self.retain_files, - extension=input_files[0].name.split('.', 1)[-1], - tags=tags) - node.add_output_opt('--output-strain-file', out_file) - - out_gates_file = File(self.ifo, "output_gates", - segments.segment(int(start_time), int(end_time)), - directory=self.out_dir, extension='txt', - store_file=self.retain_files, tags=tags) - node.add_output_opt('--output-gates-file', out_gates_file) + node.add_opt("--gps-start-time", start_time) + node.add_opt("--gps-end-time", end_time) + node.add_input_list_opt("--frame-files", input_files) + + out_file = File( + self.ifo, + "gated", + segments.segment(int(start_time), int(end_time)), + directory=self.out_dir, + store_file=self.retain_files, + extension=input_files[0].name.split(".", 1)[-1], + tags=tags, + ) + node.add_output_opt("--output-strain-file", out_file) + + out_gates_file = File( + self.ifo, + "output_gates", + segments.segment(int(start_time), int(end_time)), + directory=self.out_dir, + extension="txt", + store_file=self.retain_files, + tags=tags, + ) + node.add_output_opt("--output-gates-file", out_gates_file) return node, out_file class PycbcCreateInjectionsExecutable(Executable): - """ The class responsible for creating jobs + """ + The class responsible for creating jobs for ``pycbc_create_injections``. """ current_retention_level = Executable.ALL_TRIGGERS - extension = '.hdf' + extension = ".hdf" def create_node(self, config_files=None, seed=None, tags=None): - """ Set up a CondorDagmanNode class to run ``pycbc_create_injections``. + """ + Set up a CondorDagmanNode class to run ``pycbc_create_injections``. Parameters ---------- @@ -1100,11 +1277,11 @@ def create_node(self, config_files=None, seed=None, tags=None): A list of tags to include in filenames. Returns - -------- + ------- node : pycbc.workflow.core.Node The node to run the job. - """ + """ # default for tags is empty list tags = [] if tags is None else tags @@ -1119,23 +1296,21 @@ def create_node(self, config_files=None, seed=None, tags=None): node.add_input_list_opt("--config-files", config_files) if seed: node.add_opt("--seed", seed) - injection_file = node.new_output_file_opt(analysis_time, - self.extension, - "--output-file", - tags=tags) + injection_file = node.new_output_file_opt( + analysis_time, self.extension, "--output-file", tags=tags + ) return node, injection_file class PycbcInferenceExecutable(Executable): - """ The class responsible for creating jobs for ``pycbc_inference``. - """ + """The class responsible for creating jobs for ``pycbc_inference``.""" current_retention_level = Executable.ALL_TRIGGERS - def create_node(self, config_file, seed=None, tags=None, - analysis_time=None): - """ Set up a pegasus.Node instance to run ``pycbc_inference``. + def create_node(self, config_file, seed=None, tags=None, analysis_time=None): + """ + Set up a pegasus.Node instance to run ``pycbc_inference``. Parameters ---------- @@ -1148,9 +1323,10 @@ def create_node(self, config_file, seed=None, tags=None, A list of tags to include in filenames. Returns - -------- + ------- node : pycbc.workflow.core.Node The node to run the job. + """ # default for tags is empty list tags = [] if tags is None else tags @@ -1164,27 +1340,26 @@ def create_node(self, config_file, seed=None, tags=None, node.add_input_opt("--config-file", config_file) if seed is not None: node.add_opt("--seed", seed) - inference_file = node.new_output_file_opt(analysis_time, - ".hdf", "--output-file", - tags=tags) - if self.cp.has_option("pegasus_profile-inference", - "condor|+CheckpointSig"): + inference_file = node.new_output_file_opt( + analysis_time, ".hdf", "--output-file", tags=tags + ) + if self.cp.has_option("pegasus_profile-inference", "condor|+CheckpointSig"): err_msg = "This is not yet supported/tested with pegasus 5. " err_msg += "Please reimplement this (with unittest :-) )." raise ValueError(err_msg) - #ckpt_file_name = "{}.checkpoint".format(inference_file.name) - #ckpt_file = dax.File(ckpt_file_name) + # ckpt_file_name = "{}.checkpoint".format(inference_file.name) + # ckpt_file = dax.File(ckpt_file_name) # DO NOT call pegasus API stuff outside of # pegasus_workflow.py. - #node._dax_node.uses(ckpt_file, link=dax.Link.OUTPUT, + # node._dax_node.uses(ckpt_file, link=dax.Link.OUTPUT, # register=False, transfer=False) return node, inference_file class PycbcHDFSplitInjExecutable(Executable): - """ The class responsible for creating jobs for ``pycbc_hdf_splitinj``. - """ + """The class responsible for creating jobs for ``pycbc_hdf_splitinj``.""" + current_retention_level = Executable.ALL_TRIGGERS def __init__(self, cp, exe_name, num_splits, ifo=None, out_dir=None): @@ -1195,15 +1370,21 @@ def create_node(self, parent, tags=None): if tags is None: tags = [] node = Node(self) - node.add_input_opt('--input-file', parent) + node.add_input_opt("--input-file", parent) out_files = FileList([]) for i in range(self.num_splits): - curr_tag = 'split%d' % i + curr_tag = "split%d" % i curr_tags = parent.tags + [curr_tag] job_tag = parent.description + "_" + self.name.upper() - out_file = File(parent.ifo_list, job_tag, parent.segment, - extension='.hdf', directory=self.out_dir, - tags=curr_tags, store_file=self.retain_files) + out_file = File( + parent.ifo_list, + job_tag, + parent.segment, + extension=".hdf", + directory=self.out_dir, + tags=curr_tags, + store_file=self.retain_files, + ) out_files.append(out_file) - node.add_output_list_opt('--output-files', out_files) + node.add_output_list_opt("--output-files", out_files) return node diff --git a/pycbc/workflow/matched_filter.py b/pycbc/workflow/matched_filter.py index f6ed4693900..252b4b40542 100644 --- a/pycbc/workflow/matched_filter.py +++ b/pycbc/workflow/matched_filter.py @@ -28,22 +28,29 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/NOTYETCREATED.html """ - -import os import logging +import os from pycbc.workflow.core import FileList, make_analysis_dir -from pycbc.workflow.jobsetup import (select_matchedfilter_class, - sngl_ifo_job_setup, - multi_ifo_coherent_job_setup) +from pycbc.workflow.jobsetup import ( + multi_ifo_coherent_job_setup, + select_matchedfilter_class, + sngl_ifo_job_setup, +) -logger = logging.getLogger('pycbc.workflow.matched_filter') +logger = logging.getLogger("pycbc.workflow.matched_filter") -def setup_matchedfltr_workflow(workflow, science_segs, datafind_outs, - tmplt_banks, output_dir=None, - injection_file=None, tags=None): - ''' +def setup_matchedfltr_workflow( + workflow, + science_segs, + datafind_outs, + tmplt_banks, + output_dir=None, + injection_file=None, + tags=None, +): + """ This function aims to be the gateway for setting up a set of matched-filter jobs in a workflow. This function is intended to support multiple different ways/codes that could be used for doing this. For now the only @@ -53,7 +60,7 @@ def setup_matchedfltr_workflow(workflow, science_segs, datafind_outs, there is data and a template bank file. Parameters - ----------- + ---------- Workflow : pycbc.workflow.core.Workflow The workflow instance that the coincidence jobs will be added to. science_segs : ifo-keyed dictionary of igwn_segments.segmentlist instances @@ -81,7 +88,8 @@ def setup_matchedfltr_workflow(workflow, science_segs, datafind_outs, any intermediate products produced within this stage of the workflow. If you require access to any intermediate products produced at this stage you can call the various sub-functions directly. - ''' + + """ if tags is None: tags = [] logger.info("Entering matched-filtering setup module.") @@ -89,22 +97,33 @@ def setup_matchedfltr_workflow(workflow, science_segs, datafind_outs, cp = workflow.cp # Parse for options in .ini file - mfltrMethod = cp.get_opt_tags("workflow-matchedfilter", "matchedfilter-method", - tags) + mfltrMethod = cp.get_opt_tags( + "workflow-matchedfilter", "matchedfilter-method", tags + ) # Could have a number of choices here if mfltrMethod == "WORKFLOW_INDEPENDENT_IFOS": logger.info("Adding matched-filter jobs to workflow.") - inspiral_outs = setup_matchedfltr_dax_generated(workflow, science_segs, - datafind_outs, tmplt_banks, output_dir, - injection_file=injection_file, - tags=tags) + inspiral_outs = setup_matchedfltr_dax_generated( + workflow, + science_segs, + datafind_outs, + tmplt_banks, + output_dir, + injection_file=injection_file, + tags=tags, + ) elif mfltrMethod == "WORKFLOW_MULTIPLE_IFOS": logger.info("Adding matched-filter jobs to workflow.") - inspiral_outs = setup_matchedfltr_dax_generated_multi(workflow, - science_segs, datafind_outs, tmplt_banks, - output_dir, injection_file=injection_file, - tags=tags) + inspiral_outs = setup_matchedfltr_dax_generated_multi( + workflow, + science_segs, + datafind_outs, + tmplt_banks, + output_dir, + injection_file=injection_file, + tags=tags, + ) else: errMsg = "Matched filter method not recognized. Must be one of " errMsg += "WORKFLOW_INDEPENDENT_IFOS or WORKFLOW_MULTIPLE_IFOS." @@ -113,11 +132,17 @@ def setup_matchedfltr_workflow(workflow, science_segs, datafind_outs, logger.info("Leaving matched-filtering setup module.") return inspiral_outs -def setup_matchedfltr_dax_generated(workflow, science_segs, datafind_outs, - tmplt_banks, output_dir, - injection_file=None, - tags=None): - ''' + +def setup_matchedfltr_dax_generated( + workflow, + science_segs, + datafind_outs, + tmplt_banks, + output_dir, + injection_file=None, + tags=None, +): + """ Setup matched-filter jobs that are generated as part of the workflow. This module can support any matched-filter code that is similar in principle to @@ -125,7 +150,7 @@ def setup_matchedfltr_dax_generated(workflow, science_segs, datafind_outs, Executable and Job sub-classes (see jobutils.py). Parameters - ----------- + ---------- workflow : pycbc.workflow.core.Workflow The Workflow instance that the coincidence jobs will be added to. science_segs : ifo-keyed dictionary of igwn_segments.segmentlist instances @@ -153,7 +178,8 @@ def setup_matchedfltr_dax_generated(workflow, science_segs, datafind_outs, any intermediate products produced within this stage of the workflow. If you require access to any intermediate products produced at this stage you can call the various sub-functions directly. - ''' + + """ if tags is None: tags = [] # Need to get the exe to figure out what sections are analysed, what is @@ -162,7 +188,7 @@ def setup_matchedfltr_dax_generated(workflow, science_segs, datafind_outs, cp = workflow.cp ifos = science_segs.keys() - match_fltr_exe = os.path.basename(cp.get('executables','inspiral')) + match_fltr_exe = os.path.basename(cp.get("executables", "inspiral")) # Select the appropriate class exe_class = select_matchedfilter_class(match_fltr_exe) @@ -174,21 +200,38 @@ def setup_matchedfltr_dax_generated(workflow, science_segs, datafind_outs, # it would probably require a new module for ifo in ifos: logger.info("Setting up matched-filtering for %s.", ifo) - job_instance = exe_class(workflow.cp, 'inspiral', ifo=ifo, - out_dir=output_dir, - injection_file=injection_file, - tags=tags) - - sngl_ifo_job_setup(workflow, ifo, inspiral_outs, job_instance, - science_segs[ifo], datafind_outs, - parents=tmplt_banks, allow_overlap=False) + job_instance = exe_class( + workflow.cp, + "inspiral", + ifo=ifo, + out_dir=output_dir, + injection_file=injection_file, + tags=tags, + ) + + sngl_ifo_job_setup( + workflow, + ifo, + inspiral_outs, + job_instance, + science_segs[ifo], + datafind_outs, + parents=tmplt_banks, + allow_overlap=False, + ) return inspiral_outs -def setup_matchedfltr_dax_generated_multi(workflow, science_segs, datafind_outs, - tmplt_banks, output_dir, - injection_file=None, - tags=None): - ''' + +def setup_matchedfltr_dax_generated_multi( + workflow, + science_segs, + datafind_outs, + tmplt_banks, + output_dir, + injection_file=None, + tags=None, +): + """ Setup matched-filter jobs that are generated as part of the workflow in which a single job reads in and generates triggers over multiple ifos. This module can support any matched-filter code that is similar in @@ -196,7 +239,7 @@ def setup_matchedfltr_dax_generated_multi(workflow, science_segs, datafind_outs, needed to define Executable and Job sub-classes (see jobutils.py). Parameters - ----------- + ---------- workflow : pycbc.workflow.core.Workflow The Workflow instance that the coincidence jobs will be added to. science_segs : ifo-keyed dictionary of igwn_segments.segmentlist instances @@ -225,7 +268,8 @@ def setup_matchedfltr_dax_generated_multi(workflow, science_segs, datafind_outs, any intermediate products produced within this stage of the workflow. If you require access to any intermediate products produced at this stage you can call the various sub-functions directly. - ''' + + """ if tags is None: tags = [] # Need to get the exe to figure out what sections are analysed, what is @@ -234,25 +278,23 @@ def setup_matchedfltr_dax_generated_multi(workflow, science_segs, datafind_outs, cp = workflow.cp ifos = sorted(science_segs.keys()) - match_fltr_exe = os.path.basename(cp.get('executables','inspiral')) + match_fltr_exe = os.path.basename(cp.get("executables", "inspiral")) # List for holding the output inspiral_outs = FileList([]) - logger.info("Setting up matched-filtering for %s.", ' '.join(ifos)) + logger.info("Setting up matched-filtering for %s.", " ".join(ifos)) - if match_fltr_exe == 'pycbc_multi_inspiral': + if match_fltr_exe == "pycbc_multi_inspiral": exe_class = select_matchedfilter_class(match_fltr_exe) - bool_sg = ['make_sky_grid' in f.description for f in datafind_outs] + bool_sg = ["make_sky_grid" in f.description for f in datafind_outs] n_sg = sum(bool_sg) if n_sg == 0: - cp.set('inspiral', 'ra', - cp.get('workflow', 'ra')) - cp.set('inspiral', 'dec', - cp.get('workflow', 'dec')) + cp.set("inspiral", "ra", cp.get("workflow", "ra")) + cp.set("inspiral", "dec", cp.get("workflow", "dec")) elif n_sg > 1: - msg = f'{datafind_outs} has {n_sg} sky-grid files, ' - msg += 'instead of only one.' + msg = f"{datafind_outs} has {n_sg} sky-grid files, " + msg += "instead of only one." raise RuntimeError(msg) # Code lines for Fermi GBM are commented out for the time being # from pycbc.workflow.grb_utils import get_sky_grid_scale @@ -271,27 +313,41 @@ def setup_matchedfltr_dax_generated_multi(workflow, science_segs, datafind_outs, # str(abs(science_segs[ifos[0]][0]) - \ # 2 * int(cp.get('inspiral', 'pad-data')))) - job_instance = exe_class(workflow.cp, 'inspiral', ifo=ifos, - out_dir=output_dir, - injection_file=injection_file, - tags=tags) + job_instance = exe_class( + workflow.cp, + "inspiral", + ifo=ifos, + out_dir=output_dir, + injection_file=injection_file, + tags=tags, + ) if cp.has_option("workflow", "do-long-slides") and "slide" in tags[-1]: slide_num = int(tags[-1].replace("slide", "")) - logger.info( - "Setting up matched-filtering for slide %d", - slide_num - ) + logger.info("Setting up matched-filtering for slide %d", slide_num) slide_shift = int(cp.get("inspiral", "segment-length")) - time_slide_dict = {ifo: (slide_num + 1) * ix * slide_shift - for ix, ifo in enumerate(ifos)} - multi_ifo_coherent_job_setup(workflow, inspiral_outs, job_instance, - science_segs, datafind_outs, - output_dir, parents=tmplt_banks, - slide_dict=time_slide_dict) + time_slide_dict = { + ifo: (slide_num + 1) * ix * slide_shift for ix, ifo in enumerate(ifos) + } + multi_ifo_coherent_job_setup( + workflow, + inspiral_outs, + job_instance, + science_segs, + datafind_outs, + output_dir, + parents=tmplt_banks, + slide_dict=time_slide_dict, + ) else: - multi_ifo_coherent_job_setup(workflow, inspiral_outs, job_instance, - science_segs, datafind_outs, - output_dir, parents=tmplt_banks) + multi_ifo_coherent_job_setup( + workflow, + inspiral_outs, + job_instance, + science_segs, + datafind_outs, + output_dir, + parents=tmplt_banks, + ) else: # Select the appropriate class raise ValueError("Not currently supported.") diff --git a/pycbc/workflow/minifollowups.py b/pycbc/workflow/minifollowups.py index a847c2dcae8..579446b4948 100644 --- a/pycbc/workflow/minifollowups.py +++ b/pycbc/workflow/minifollowups.py @@ -20,9 +20,9 @@ import igwn_segments as segments from pycbc.events import coinc -from pycbc.workflow.core import Executable, FileList -from pycbc.workflow.core import makedir, resolve_url_to_file -from pycbc.workflow.plotting import PlotExecutable, requirestr, excludestr +from pycbc.workflow.core import Executable, FileList, makedir, resolve_url_to_file +from pycbc.workflow.plotting import PlotExecutable, excludestr, requirestr + try: # Python 3 from itertools import zip_longest @@ -31,19 +31,29 @@ from itertools import izip_longest as zip_longest from pycbc.workflow.pegasus_workflow import SubWorkflow -logger = logging.getLogger('pycbc.workflow.minifollowups') +logger = logging.getLogger("pycbc.workflow.minifollowups") + def grouper(iterable, n, fillvalue=None): - """ Create a list of n length tuples - """ + """Create a list of n length tuples""" args = [iter(iterable)] * n return zip_longest(*args, fillvalue=fillvalue) -def setup_foreground_minifollowups(workflow, coinc_file, single_triggers, - tmpltbank_file, insp_segs, insp_data_name, - insp_anal_name, dax_output, out_dir, - tags=None): - """ Create plots that followup the Nth loudest coincident injection + +def setup_foreground_minifollowups( + workflow, + coinc_file, + single_triggers, + tmpltbank_file, + insp_segs, + insp_data_name, + insp_anal_name, + dax_output, + out_dir, + tags=None, +): + """ + Create plots that followup the Nth loudest coincident injection from a statmap produced HDF file. Parameters @@ -75,48 +85,55 @@ def setup_foreground_minifollowups(workflow, coinc_file, single_triggers, layout: list A list of tuples which specify the displayed file layout for the minifollops plots. + """ - logger.info('Entering minifollowups module') + logger.info("Entering minifollowups module") - if not workflow.cp.has_section('workflow-minifollowups'): - msg = 'There is no [workflow-minifollowups] section in ' - msg += 'configuration file' + if not workflow.cp.has_section("workflow-minifollowups"): + msg = "There is no [workflow-minifollowups] section in " + msg += "configuration file" logger.info(msg) - logger.info('Leaving minifollowups') + logger.info("Leaving minifollowups") return tags = [] if tags is None else tags makedir(dax_output) # turn the config file into a File class - config_path = os.path.abspath(dax_output + '/' + '_'.join(tags) + 'foreground_minifollowup.ini') - workflow.cp.write(open(config_path, 'w')) + config_path = os.path.abspath( + dax_output + "/" + "_".join(tags) + "foreground_minifollowup.ini" + ) + workflow.cp.write(open(config_path, "w")) config_file = resolve_url_to_file(config_path) - exe = Executable(workflow.cp, 'foreground_minifollowup', - ifos=workflow.ifos, out_dir=dax_output, tags=tags) + exe = Executable( + workflow.cp, + "foreground_minifollowup", + ifos=workflow.ifos, + out_dir=dax_output, + tags=tags, + ) node = exe.create_node() - node.add_input_opt('--config-files', config_file) - node.add_input_opt('--bank-file', tmpltbank_file) - node.add_input_opt('--statmap-file', coinc_file) - node.add_multiifo_input_list_opt('--single-detector-triggers', - single_triggers) - node.add_input_opt('--inspiral-segments', insp_segs) - node.add_opt('--inspiral-data-read-name', insp_data_name) - node.add_opt('--inspiral-data-analyzed-name', insp_anal_name) + node.add_input_opt("--config-files", config_file) + node.add_input_opt("--bank-file", tmpltbank_file) + node.add_input_opt("--statmap-file", coinc_file) + node.add_multiifo_input_list_opt("--single-detector-triggers", single_triggers) + node.add_input_opt("--inspiral-segments", insp_segs) + node.add_opt("--inspiral-data-read-name", insp_data_name) + node.add_opt("--inspiral-data-analyzed-name", insp_anal_name) if tags: - node.add_list_opt('--tags', tags) - node.new_output_file_opt(workflow.analysis_time, '.dax', '--dax-file') - node.new_output_file_opt(workflow.analysis_time, '.dax.map', '--output-map') + node.add_list_opt("--tags", tags) + node.new_output_file_opt(workflow.analysis_time, ".dax", "--dax-file") + node.new_output_file_opt(workflow.analysis_time, ".dax.map", "--output-map") name = node.output_files[0].name map_file = node.output_files[1] - node.add_opt('--workflow-name', name) - node.add_opt('--output-dir', out_dir) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", name) + node.add_opt("--output-dir", out_dir) + node.add_opt("--dax-file-directory", ".") workflow += node @@ -125,22 +142,33 @@ def setup_foreground_minifollowups(workflow, coinc_file, single_triggers, # determine if a staging site has been specified job = SubWorkflow(fil.name, is_planned=False) - input_files = [tmpltbank_file, coinc_file, insp_segs] + \ - single_triggers + input_files = [tmpltbank_file, coinc_file, insp_segs] + single_triggers job.add_inputs(*input_files) - job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) job.add_into_workflow(workflow) - logger.info('Leaving minifollowups module') - -def setup_single_det_minifollowups(workflow, single_trig_file, tmpltbank_file, - insp_segs, insp_data_name, insp_anal_name, - dax_output, out_dir, veto_file=None, - veto_segment_name=None, fg_file=None, - fg_name=None, statfiles=None, - tags=None): - """ Create plots that followup the Nth loudest clustered single detector + logger.info("Leaving minifollowups module") + + +def setup_single_det_minifollowups( + workflow, + single_trig_file, + tmpltbank_file, + insp_segs, + insp_data_name, + insp_anal_name, + dax_output, + out_dir, + veto_file=None, + veto_segment_name=None, + fg_file=None, + fg_name=None, + statfiles=None, + tags=None, +): + """ + Create plots that followup the Nth loudest clustered single detector triggers from a merged single detector trigger HDF file. Parameters @@ -164,19 +192,21 @@ def setup_single_det_minifollowups(workflow, single_trig_file, tmpltbank_file, statistic. tags: {None, optional} Tags to add to the minifollowups executables + Returns ------- layout: list A list of tuples which specify the displayed file layout for the minifollops plots. + """ - logger.info('Entering minifollowups module') + logger.info("Entering minifollowups module") - if not workflow.cp.has_section('workflow-sngl_minifollowups'): - msg = 'There is no [workflow-sngl_minifollowups] section in ' - msg += 'configuration file' + if not workflow.cp.has_section("workflow-sngl_minifollowups"): + msg = "There is no [workflow-sngl_minifollowups] section in " + msg += "configuration file" logger.info(msg) - logger.info('Leaving minifollowups') + logger.info("Leaving minifollowups") return tags = [] if tags is None else tags @@ -184,49 +214,54 @@ def setup_single_det_minifollowups(workflow, single_trig_file, tmpltbank_file, # turn the config file into a File class curr_ifo = single_trig_file.ifo - config_path = os.path.abspath(dax_output + '/' + curr_ifo + \ - '_'.join(tags) + 'singles_minifollowup.ini') - workflow.cp.write(open(config_path, 'w')) + config_path = os.path.abspath( + dax_output + "/" + curr_ifo + "_".join(tags) + "singles_minifollowup.ini" + ) + workflow.cp.write(open(config_path, "w")) config_file = resolve_url_to_file(config_path) - exe = Executable(workflow.cp, 'singles_minifollowup', - ifos=curr_ifo, out_dir=dax_output, tags=tags) + exe = Executable( + workflow.cp, + "singles_minifollowup", + ifos=curr_ifo, + out_dir=dax_output, + tags=tags, + ) node = exe.create_node() - node.add_input_opt('--config-files', config_file) - node.add_input_opt('--bank-file', tmpltbank_file) - node.add_input_opt('--single-detector-file', single_trig_file) - node.add_input_opt('--inspiral-segments', insp_segs) - node.add_opt('--inspiral-data-read-name', insp_data_name) - node.add_opt('--inspiral-data-analyzed-name', insp_anal_name) - node.add_opt('--instrument', curr_ifo) + node.add_input_opt("--config-files", config_file) + node.add_input_opt("--bank-file", tmpltbank_file) + node.add_input_opt("--single-detector-file", single_trig_file) + node.add_input_opt("--inspiral-segments", insp_segs) + node.add_opt("--inspiral-data-read-name", insp_data_name) + node.add_opt("--inspiral-data-analyzed-name", insp_anal_name) + node.add_opt("--instrument", curr_ifo) if veto_file is not None: - assert(veto_segment_name is not None) - node.add_input_opt('--veto-file', veto_file) - node.add_opt('--veto-segment-name', veto_segment_name) + assert veto_segment_name is not None + node.add_input_opt("--veto-file", veto_file) + node.add_opt("--veto-segment-name", veto_segment_name) if fg_file is not None: - assert(fg_name is not None) - node.add_input_opt('--foreground-censor-file', fg_file) - node.add_opt('--foreground-segment-name', fg_name) + assert fg_name is not None + node.add_input_opt("--foreground-censor-file", fg_file) + node.add_opt("--foreground-segment-name", fg_name) if statfiles: node.add_input_list_opt( - '--statistic-files', + "--statistic-files", statfiles, check_existing_options=False, ) if tags: - node.add_list_opt('--tags', tags) - node.new_output_file_opt(workflow.analysis_time, '.dax', '--dax-file') - node.new_output_file_opt(workflow.analysis_time, '.dax.map', - '--output-map') + node.add_list_opt("--tags", tags) + node.new_output_file_opt(workflow.analysis_time, ".dax", "--dax-file") + node.new_output_file_opt(workflow.analysis_time, ".dax.map", "--output-map") name = node.output_files[0].name map_file = node.output_files[1] - node.add_opt('--workflow-name', name) - node.add_opt('--output-dir', out_dir) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", name) + node.add_opt("--output-dir", out_dir) + node.add_opt("--dax-file-directory", ".") workflow += node @@ -240,18 +275,28 @@ def setup_single_det_minifollowups(workflow, single_trig_file, tmpltbank_file, if statfiles: input_files += statfiles job.add_inputs(*input_files) - job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) job.add_into_workflow(workflow) - logger.info('Leaving minifollowups module') - - -def setup_injection_minifollowups(workflow, injection_file, inj_xml_file, - single_triggers, tmpltbank_file, - insp_segs, insp_data_name, insp_anal_name, - dax_output, out_dir, tags=None): - """ Create plots that followup the closest missed injections + logger.info("Leaving minifollowups module") + + +def setup_injection_minifollowups( + workflow, + injection_file, + inj_xml_file, + single_triggers, + tmpltbank_file, + insp_segs, + insp_data_name, + insp_anal_name, + dax_output, + out_dir, + tags=None, +): + """ + Create plots that followup the closest missed injections Parameters ---------- @@ -279,47 +324,54 @@ def setup_injection_minifollowups(workflow, injection_file, inj_xml_file, layout: list A list of tuples which specify the displayed file layout for the minifollops plots. + """ - logger.info('Entering injection minifollowups module') + logger.info("Entering injection minifollowups module") - if not workflow.cp.has_section('workflow-injection_minifollowups'): - msg = 'There is no [workflow-injection_minifollowups] section in ' - msg += 'configuration file' + if not workflow.cp.has_section("workflow-injection_minifollowups"): + msg = "There is no [workflow-injection_minifollowups] section in " + msg += "configuration file" logger.info(msg) - logger.info('Leaving minifollowups') + logger.info("Leaving minifollowups") return tags = [] if tags is None else tags makedir(dax_output) # turn the config file into a File class - config_path = os.path.abspath(dax_output + '/' + '_'.join(tags) + 'injection_minifollowup.ini') - workflow.cp.write(open(config_path, 'w')) + config_path = os.path.abspath( + dax_output + "/" + "_".join(tags) + "injection_minifollowup.ini" + ) + workflow.cp.write(open(config_path, "w")) config_file = resolve_url_to_file(config_path) - exe = Executable(workflow.cp, 'injection_minifollowup', ifos=workflow.ifos, out_dir=dax_output) + exe = Executable( + workflow.cp, "injection_minifollowup", ifos=workflow.ifos, out_dir=dax_output + ) node = exe.create_node() - node.add_input_opt('--config-files', config_file) - node.add_input_opt('--bank-file', tmpltbank_file) - node.add_input_opt('--injection-file', injection_file) - node.add_input_opt('--injection-xml-file', inj_xml_file) - node.add_multiifo_input_list_opt('--single-detector-triggers', single_triggers) - node.add_input_opt('--inspiral-segments', insp_segs) - node.add_opt('--inspiral-data-read-name', insp_data_name) - node.add_opt('--inspiral-data-analyzed-name', insp_anal_name) + node.add_input_opt("--config-files", config_file) + node.add_input_opt("--bank-file", tmpltbank_file) + node.add_input_opt("--injection-file", injection_file) + node.add_input_opt("--injection-xml-file", inj_xml_file) + node.add_multiifo_input_list_opt("--single-detector-triggers", single_triggers) + node.add_input_opt("--inspiral-segments", insp_segs) + node.add_opt("--inspiral-data-read-name", insp_data_name) + node.add_opt("--inspiral-data-analyzed-name", insp_anal_name) if tags: - node.add_list_opt('--tags', tags) - node.new_output_file_opt(workflow.analysis_time, '.dax', '--dax-file', tags=tags) - node.new_output_file_opt(workflow.analysis_time, '.dax.map', '--output-map', tags=tags) + node.add_list_opt("--tags", tags) + node.new_output_file_opt(workflow.analysis_time, ".dax", "--dax-file", tags=tags) + node.new_output_file_opt( + workflow.analysis_time, ".dax.map", "--output-map", tags=tags + ) name = node.output_files[0].name map_file = node.output_files[1] - node.add_opt('--workflow-name', name) - node.add_opt('--output-dir', out_dir) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", name) + node.add_opt("--output-dir", out_dir) + node.add_opt("--dax-file-directory", ".") workflow += node @@ -330,40 +382,45 @@ def setup_injection_minifollowups(workflow, injection_file, inj_xml_file, input_files = [tmpltbank_file, injection_file, inj_xml_file, insp_segs] input_files += single_triggers job.add_inputs(*input_files) - job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) job.add_into_workflow(workflow) - logger.info('Leaving injection minifollowups module') + logger.info("Leaving injection minifollowups module") class SingleTemplateExecutable(PlotExecutable): - """Class to be used for to create workflow.Executable instances for the + """ + Class to be used for to create workflow.Executable instances for the pycbc_single_template executable. Basically inherits directly from PlotExecutable. """ - time_dependent_options = ['--channel-name', '--frame-type'] + + time_dependent_options = ["--channel-name", "--frame-type"] class SingleTimeFreqExecutable(PlotExecutable): - """Class to be used for to create workflow.Executable instances for the + """ + Class to be used for to create workflow.Executable instances for the pycbc_plot_singles_timefreq executable. Basically inherits directly from PlotExecutable. """ - time_dependent_options = ['--channel-name', '--frame-type'] + + time_dependent_options = ["--channel-name", "--frame-type"] class PlotQScanExecutable(PlotExecutable): - """Class to be used for to create workflow.Executable instances for the + """ + Class to be used for to create workflow.Executable instances for the pycbc_plot_qscan executable. Basically inherits directly from PlotExecutable. """ - time_dependent_options = ['--channel-name', '--frame-type'] + time_dependent_options = ["--channel-name", "--frame-type"] -def get_single_template_params(curr_idx, times, bank_data, - bank_id, fsdt, tids): + +def get_single_template_params(curr_idx, times, bank_data, bank_id, fsdt, tids): """ A function to get the parameters needed for the make_single_template_files function. @@ -391,53 +448,64 @@ def get_single_template_params(curr_idx, times, bank_data, """ params = {} for ifo in times: - params['%s_end_time' % ifo] = times[ifo][curr_idx] + params["%s_end_time" % ifo] = times[ifo][curr_idx] try: # Only present for precessing, so may not exist - params['u_vals_%s' % ifo] = \ - fsdt[ifo][ifo]['u_vals'][tids[ifo][curr_idx]] + params["u_vals_%s" % ifo] = fsdt[ifo][ifo]["u_vals"][tids[ifo][curr_idx]] except: pass - params['mean_time'] = coinc.mean_if_greater_than_zero( + params["mean_time"] = coinc.mean_if_greater_than_zero( [times[ifo][curr_idx] for ifo in times] )[0] - params['mass1'] = bank_data['mass1'][bank_id] - params['mass2'] = bank_data['mass2'][bank_id] - params['spin1z'] = bank_data['spin1z'][bank_id] - params['spin2z'] = bank_data['spin2z'][bank_id] - params['f_lower'] = bank_data['f_lower'][bank_id] - if 'approximant' in bank_data: - params['approximant'] = bank_data['approximant'][bank_id] + params["mass1"] = bank_data["mass1"][bank_id] + params["mass2"] = bank_data["mass2"][bank_id] + params["spin1z"] = bank_data["spin1z"][bank_id] + params["spin2z"] = bank_data["spin2z"][bank_id] + params["f_lower"] = bank_data["f_lower"][bank_id] + if "approximant" in bank_data: + params["approximant"] = bank_data["approximant"][bank_id] # don't require precessing template info if not present try: - params['spin1x'] = bank_data['spin1x'][bank_id] - params['spin1y'] = bank_data['spin1y'][bank_id] - params['spin2x'] = bank_data['spin2x'][bank_id] - params['spin2y'] = bank_data['spin2y'][bank_id] - params['inclination'] = bank_data['inclination'][bank_id] + params["spin1x"] = bank_data["spin1x"][bank_id] + params["spin1y"] = bank_data["spin1y"][bank_id] + params["spin2x"] = bank_data["spin2x"][bank_id] + params["spin2y"] = bank_data["spin2y"][bank_id] + params["inclination"] = bank_data["inclination"][bank_id] except KeyError: pass # optional eccentric parameters, present when using eccentric waveform models try: - params['eccentricity'] = bank_data['eccentricity'][bank_id] - params['rel_anomaly'] = bank_data['rel_anomaly'][bank_id] + params["eccentricity"] = bank_data["eccentricity"][bank_id] + params["rel_anomaly"] = bank_data["rel_anomaly"][bank_id] except KeyError: pass return params -def make_single_template_files(workflow, segs, ifo, data_read_name, - analyzed_name, params, out_dir, inj_file=None, - exclude=None, require=None, tags=None, - store_file=False, use_mean_time=False, - use_exact_inj_params=False): - """Function for creating jobs to run the pycbc_single_template code and +def make_single_template_files( + workflow, + segs, + ifo, + data_read_name, + analyzed_name, + params, + out_dir, + inj_file=None, + exclude=None, + require=None, + tags=None, + store_file=False, + use_mean_time=False, + use_exact_inj_params=False, +): + """ + Function for creating jobs to run the pycbc_single_template code and add these jobs to the workflow. Parameters - ----------- + ---------- workflow : workflow.Workflow instance The pycbc.workflow.Workflow instance to add these jobs to. segs : workflow.File instance @@ -484,79 +552,90 @@ def make_single_template_files(workflow, segs, ifo, data_read_name, but instead use the injection closest to the filter time as a template. Returns - -------- + ------- output_files : workflow.FileList The list of workflow.Files created in this function. + """ tags = [] if tags is None else tags makedir(out_dir) - name = 'single_template' + name = "single_template" secs = requirestr(workflow.cp.get_subsections(name), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) # Reanalyze the time around the trigger in each detector - curr_exe = SingleTemplateExecutable(workflow.cp, 'single_template', - ifos=[ifo], out_dir=out_dir, - tags=tags) - start = int(params[ifo + '_end_time']) + curr_exe = SingleTemplateExecutable( + workflow.cp, "single_template", ifos=[ifo], out_dir=out_dir, tags=tags + ) + start = int(params[ifo + "_end_time"]) end = start + 1 cseg = segments.segment([start, end]) node = curr_exe.create_node(valid_seg=cseg) if use_exact_inj_params: - node.add_opt('--use-params-of-closest-injection') + node.add_opt("--use-params-of-closest-injection") else: - node.add_opt('--mass1', "%.6f" % params['mass1']) - node.add_opt('--mass2', "%.6f" % params['mass2']) - node.add_opt('--spin1z',"%.6f" % params['spin1z']) - node.add_opt('--spin2z',"%.6f" % params['spin2z']) - node.add_opt('--template-start-frequency', - "%.6f" % params['f_lower']) + node.add_opt("--mass1", "%.6f" % params["mass1"]) + node.add_opt("--mass2", "%.6f" % params["mass2"]) + node.add_opt("--spin1z", "%.6f" % params["spin1z"]) + node.add_opt("--spin2z", "%.6f" % params["spin2z"]) + node.add_opt("--template-start-frequency", "%.6f" % params["f_lower"]) # Is this precessing? - if 'u_vals' in params or 'u_vals_%s' % ifo in params: - node.add_opt('--spin1x',"%.6f" % params['spin1x']) - node.add_opt('--spin1y',"%.6f" % params['spin1y']) - node.add_opt('--spin2x',"%.6f" % params['spin2x']) - node.add_opt('--spin2y',"%.6f" % params['spin2y']) - node.add_opt('--inclination',"%.6f" % params['inclination']) + if "u_vals" in params or "u_vals_%s" % ifo in params: + node.add_opt("--spin1x", "%.6f" % params["spin1x"]) + node.add_opt("--spin1y", "%.6f" % params["spin1y"]) + node.add_opt("--spin2x", "%.6f" % params["spin2x"]) + node.add_opt("--spin2y", "%.6f" % params["spin2y"]) + node.add_opt("--inclination", "%.6f" % params["inclination"]) try: - node.add_opt('--u-val',"%.6f" % params['u_vals']) + node.add_opt("--u-val", "%.6f" % params["u_vals"]) except: - node.add_opt('--u-val', - "%.6f" % params['u_vals_%s' % ifo]) + node.add_opt("--u-val", "%.6f" % params["u_vals_%s" % ifo]) # If this is an eccentricity search - if 'eccentricity' in params: - node.add_opt('--eccentricity', "%.6f" % params['eccentricity']) - node.add_opt('--rel-anomaly', "%.6f" % params['rel_anomaly']) + if "eccentricity" in params: + node.add_opt("--eccentricity", "%.6f" % params["eccentricity"]) + node.add_opt("--rel-anomaly", "%.6f" % params["rel_anomaly"]) - if params[ifo + '_end_time'] > 0 and not use_mean_time: - trig_time = params[ifo + '_end_time'] + if params[ifo + "_end_time"] > 0 and not use_mean_time: + trig_time = params[ifo + "_end_time"] else: - trig_time = params['mean_time'] + trig_time = params["mean_time"] - node.add_opt('--trigger-time', f"{trig_time:.6f}") - node.add_input_opt('--inspiral-segments', segs) + node.add_opt("--trigger-time", f"{trig_time:.6f}") + node.add_input_opt("--inspiral-segments", segs) if inj_file is not None: - node.add_input_opt('--injection-file', inj_file) - node.add_opt('--data-read-name', data_read_name) - node.add_opt('--data-analyzed-name', analyzed_name) - node.new_output_file_opt(workflow.analysis_time, '.hdf', - '--output-file', store_file=store_file) + node.add_input_opt("--injection-file", inj_file) + node.add_opt("--data-read-name", data_read_name) + node.add_opt("--data-analyzed-name", analyzed_name) + node.new_output_file_opt( + workflow.analysis_time, ".hdf", "--output-file", store_file=store_file + ) workflow += node return node.output_files -def make_single_template_plots(workflow, segs, data_read_name, analyzed_name, - params, out_dir, inj_file=None, exclude=None, - data_segments=None, - require=None, tags=None, params_str=None, - use_exact_inj_params=False): - """Function for creating jobs to run the pycbc_single_template code and +def make_single_template_plots( + workflow, + segs, + data_read_name, + analyzed_name, + params, + out_dir, + inj_file=None, + exclude=None, + data_segments=None, + require=None, + tags=None, + params_str=None, + use_exact_inj_params=False, +): + """ + Function for creating jobs to run the pycbc_single_template code and to run the associated plotting code pycbc_single_template_plots and add these jobs to the workflow. Parameters - ----------- + ---------- workflow : workflow.Workflow instance The pycbc.workflow.Workflow instance to add these jobs to. segs : workflow.File instance @@ -604,17 +683,18 @@ def make_single_template_plots(workflow, segs, data_read_name, analyzed_name, but instead use the injection closest to the filter time as a template. Returns - -------- + ------- hdf_files : workflow.FileList The list of workflow.Files created by single_template jobs in this function. plot_files : workflow.FileList The list of workflow.Files created by single_template_plot jobs in this function. + """ tags = [] if tags is None else tags makedir(out_dir) - name = 'single_template_plot' + name = "single_template_plot" secs = requirestr(workflow.cp.get_subsections(name), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) @@ -622,8 +702,11 @@ def make_single_template_plots(workflow, segs, data_read_name, analyzed_name, plot_files = FileList([]) valid = {} for ifo in workflow.ifos: - valid[ifo] = params['mean_time'] in data_segments[ifo] if data_segments \ - else params['%s_end_time' % ifo] > 0 + valid[ifo] = ( + params["mean_time"] in data_segments[ifo] + if data_segments + else params["%s_end_time" % ifo] > 0 + ) for tag in secs: for ifo in workflow.ifos: if not valid[ifo]: @@ -642,176 +725,223 @@ def make_single_template_plots(workflow, segs, data_read_name, analyzed_name, require=require, tags=tags + [tag], store_file=False, - use_exact_inj_params=use_exact_inj_params + use_exact_inj_params=use_exact_inj_params, ) hdf_files += data # Make the plot for this trigger and detector - node = PlotExecutable(workflow.cp, name, ifos=[ifo], - out_dir=out_dir, tags=[tag] + tags).create_node() - node.add_input_opt('--single-template-file', data[0]) - node.new_output_file_opt(workflow.analysis_time, '.png', - '--output-file') - title="'%s SNR and chi^2 timeseries" %(ifo) + node = PlotExecutable( + workflow.cp, name, ifos=[ifo], out_dir=out_dir, tags=[tag] + tags + ).create_node() + node.add_input_opt("--single-template-file", data[0]) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") + title = "'%s SNR and chi^2 timeseries" % (ifo) if params_str is not None: - title+= " using %s" %(params_str) - title+="'" - node.add_opt('--plot-title', title) + title += " using %s" % (params_str) + title += "'" + node.add_opt("--plot-title", title) caption = "'The SNR and chi^2 timeseries around the injection" if params_str is not None: - caption += " using %s" %(params_str) + caption += " using %s" % (params_str) if use_exact_inj_params: caption += ". The injection itself was used as the template.'" else: caption += ". The template used has the following parameters: " - caption += "mass1=%s, mass2=%s, spin1z=%s, spin2z=%s'"\ - %(params['mass1'], params['mass2'], params['spin1z'], - params['spin2z']) - node.add_opt('--plot-caption', caption) + caption += "mass1=%s, mass2=%s, spin1z=%s, spin2z=%s'" % ( + params["mass1"], + params["mass2"], + params["spin1z"], + params["spin2z"], + ) + node.add_opt("--plot-caption", caption) workflow += node plot_files += node.output_files return hdf_files, plot_files -def make_plot_waveform_plot(workflow, params, out_dir, ifos, exclude=None, - require=None, tags=None): - """ Add plot_waveform jobs to the workflow. - """ + +def make_plot_waveform_plot( + workflow, params, out_dir, ifos, exclude=None, require=None, tags=None +): + """Add plot_waveform jobs to the workflow.""" tags = [] if tags is None else tags makedir(out_dir) - name = 'single_template_plot' + name = "single_template_plot" secs = requirestr(workflow.cp.get_subsections(name), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - node = PlotExecutable(workflow.cp, 'plot_waveform', ifos=ifos, - out_dir=out_dir, tags=[tag] + tags).create_node() - node.add_opt('--mass1', "%.6f" % params['mass1']) - node.add_opt('--mass2', "%.6f" % params['mass2']) - node.add_opt('--spin1z',"%.6f" % params['spin1z']) - node.add_opt('--spin2z',"%.6f" % params['spin2z']) - if 'u_vals' in params: + node = PlotExecutable( + workflow.cp, "plot_waveform", ifos=ifos, out_dir=out_dir, tags=[tag] + tags + ).create_node() + node.add_opt("--mass1", "%.6f" % params["mass1"]) + node.add_opt("--mass2", "%.6f" % params["mass2"]) + node.add_opt("--spin1z", "%.6f" % params["spin1z"]) + node.add_opt("--spin2z", "%.6f" % params["spin2z"]) + if "u_vals" in params: # Precessing options - node.add_opt('--spin1x',"%.6f" % params['spin1x']) - node.add_opt('--spin2x',"%.6f" % params['spin2x']) - node.add_opt('--spin1y',"%.6f" % params['spin1y']) - node.add_opt('--spin2y',"%.6f" % params['spin2y']) - node.add_opt('--inclination',"%.6f" % params['inclination']) - node.add_opt('--u-val', "%.6f" % params['u_vals']) - if 'eccentricity' in params: - node.add_opt('--eccentricity', "%.6f" % params['eccentricity']) - node.add_opt('--rel-anomaly', "%.6f" % params['rel_anomaly']) - node.new_output_file_opt(workflow.analysis_time, '.png', - '--output-file') + node.add_opt("--spin1x", "%.6f" % params["spin1x"]) + node.add_opt("--spin2x", "%.6f" % params["spin2x"]) + node.add_opt("--spin1y", "%.6f" % params["spin1y"]) + node.add_opt("--spin2y", "%.6f" % params["spin2y"]) + node.add_opt("--inclination", "%.6f" % params["inclination"]) + node.add_opt("--u-val", "%.6f" % params["u_vals"]) + if "eccentricity" in params: + node.add_opt("--eccentricity", "%.6f" % params["eccentricity"]) + node.add_opt("--rel-anomaly", "%.6f" % params["rel_anomaly"]) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node files += node.output_files return files -def make_inj_info(workflow, injection_file, injection_index, num, out_dir, - tags=None): + +def make_inj_info(workflow, injection_file, injection_index, num, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - name = 'page_injinfo' + name = "page_injinfo" files = FileList([]) - node = PlotExecutable(workflow.cp, name, ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--injection-file', injection_file) - node.add_opt('--injection-index', str(injection_index)) - node.add_opt('--n-nearest', str(num)) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node = PlotExecutable( + workflow.cp, name, ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--injection-file", injection_file) + node.add_opt("--injection-index", str(injection_index)) + node.add_opt("--n-nearest", str(num)) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node files += node.output_files return files -def make_coinc_info(workflow, singles, bank, coinc_file, out_dir, - n_loudest=None, trig_id=None, file_substring=None, - sort_order=None, sort_var=None, title=None, tags=None): + +def make_coinc_info( + workflow, + singles, + bank, + coinc_file, + out_dir, + n_loudest=None, + trig_id=None, + file_substring=None, + sort_order=None, + sort_var=None, + title=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - name = 'page_coincinfo' + name = "page_coincinfo" files = FileList([]) - node = PlotExecutable(workflow.cp, name, ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--single-trigger-files', singles) - node.add_input_opt('--statmap-file', coinc_file) - node.add_input_opt('--bank-file', bank) + node = PlotExecutable( + workflow.cp, name, ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--single-trigger-files", singles) + node.add_input_opt("--statmap-file", coinc_file) + node.add_input_opt("--bank-file", bank) if sort_order: - node.add_opt('--sort-order', sort_order) + node.add_opt("--sort-order", sort_order) if sort_var: - node.add_opt('--sort-variable', sort_var) + node.add_opt("--sort-variable", sort_var) if n_loudest is not None: - node.add_opt('--n-loudest', str(n_loudest)) + node.add_opt("--n-loudest", str(n_loudest)) if trig_id is not None: - node.add_opt('--trigger-id', str(trig_id)) + node.add_opt("--trigger-id", str(trig_id)) if title is not None: - node.add_opt('--title', f'"{title}"') + node.add_opt("--title", f'"{title}"') if file_substring is not None: - node.add_opt('--statmap-file-subspace-name', file_substring) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node.add_opt("--statmap-file-subspace-name", file_substring) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node files += node.output_files return files -def make_sngl_ifo(workflow, sngl_file, bank_file, trigger_id, out_dir, ifo, - statfiles=None, title=None, tags=None): - """Setup a job to create sngl detector sngl ifo html summary snippet. - """ + +def make_sngl_ifo( + workflow, + sngl_file, + bank_file, + trigger_id, + out_dir, + ifo, + statfiles=None, + title=None, + tags=None, +): + """Setup a job to create sngl detector sngl ifo html summary snippet.""" tags = [] if tags is None else tags makedir(out_dir) - name = 'page_snglinfo' + name = "page_snglinfo" files = FileList([]) - node = PlotExecutable(workflow.cp, name, ifos=[ifo], - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--single-trigger-file', sngl_file) - node.add_input_opt('--bank-file', bank_file) - node.add_opt('--trigger-id', str(trigger_id)) - node.add_opt('--instrument', ifo) + node = PlotExecutable( + workflow.cp, name, ifos=[ifo], out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--single-trigger-file", sngl_file) + node.add_input_opt("--bank-file", bank_file) + node.add_opt("--trigger-id", str(trigger_id)) + node.add_opt("--instrument", ifo) if statfiles is not None: node.add_input_list_opt( - '--statistic-files', + "--statistic-files", statfiles, check_existing_options=False, ) if title is not None: - node.add_opt('--title', f'"{title}"') - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node.add_opt("--title", f'"{title}"') + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node files += node.output_files return files -def make_trigger_timeseries(workflow, singles, ifo_times, out_dir, special_tids=None, - exclude=None, require=None, tags=None): +def make_trigger_timeseries( + workflow, + singles, + ifo_times, + out_dir, + special_tids=None, + exclude=None, + require=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - name = 'plot_trigger_timeseries' + name = "plot_trigger_timeseries" secs = requirestr(workflow.cp.get_subsections(name), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - node = PlotExecutable(workflow.cp, name, ifos=workflow.ifos, - out_dir=out_dir, tags=[tag] + tags).create_node() - node.add_multiifo_input_list_opt('--single-trigger-files', singles) - node.add_opt('--times', ifo_times) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node = PlotExecutable( + workflow.cp, name, ifos=workflow.ifos, out_dir=out_dir, tags=[tag] + tags + ).create_node() + node.add_multiifo_input_list_opt("--single-trigger-files", singles) + node.add_opt("--times", ifo_times) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") if special_tids is not None: - node.add_opt('--special-trigger-ids', special_tids) + node.add_opt("--special-trigger-ids", special_tids) workflow += node files += node.output_files return files -def make_qscan_plot(workflow, ifo, trig_time, out_dir, injection_file=None, - data_segments=None, time_window=100, tags=None): - """ Generate a make_qscan node and add it to workflow. + +def make_qscan_plot( + workflow, + ifo, + trig_time, + out_dir, + injection_file=None, + data_segments=None, + time_window=100, + tags=None, +): + """ + Generate a make_qscan node and add it to workflow. This function generates a single node of the plot_qscan executable and adds it to the current workflow. Parent/child relationships are set by the input/output files automatically. Parameters - ----------- + ---------- workflow: pycbc.workflow.core.Workflow The workflow class that stores the jobs that will be run. ifo: str @@ -837,13 +967,15 @@ def make_qscan_plot(workflow, ifo, trig_time, out_dir, injection_file=None, cases. tags: list (optional, default=None) List of tags to add to the created nodes, which determine file naming. + """ tags = [] if tags is None else tags makedir(out_dir) - name = 'plot_qscan' + name = "plot_qscan" - curr_exe = PlotQScanExecutable(workflow.cp, name, ifos=[ifo], - out_dir=out_dir, tags=tags) + curr_exe = PlotQScanExecutable( + workflow.cp, name, ifos=[ifo], out_dir=out_dir, tags=tags + ) # Determine start/end times, using data segments if needed. # Begin by choosing "optimal" times @@ -859,54 +991,64 @@ def make_qscan_plot(workflow, ifo, trig_time, out_dir, injection_file=None, if trig_time in seg: data_seg = seg break - elif trig_time == -1.0: - node.add_opt('--gps-start-time', int(trig_time)) - node.add_opt('--gps-end-time', int(trig_time)) - node.add_opt('--center-time', trig_time) + if trig_time == -1.0: + node.add_opt("--gps-start-time", int(trig_time)) + node.add_opt("--gps-end-time", int(trig_time)) + node.add_opt("--center-time", trig_time) caption_string = "'No trigger in %s'" % ifo - node.add_opt('--plot-caption', caption_string) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.add_opt("--plot-caption", caption_string) + node.new_output_file_opt( + workflow.analysis_time, ".png", "--output-file" + ) workflow += node return node.output_files else: - err_msg = "Trig time {} ".format(trig_time) + err_msg = f"Trig time {trig_time} " err_msg += "does not seem to lie within any data segments. " err_msg += "This shouldn't be possible, please ask for help!" raise ValueError(err_msg) # Check for pad-data - if curr_exe.has_opt('pad-data'): - pad_data = int(curr_exe.get_opt('pad-data')) + if curr_exe.has_opt("pad-data"): + pad_data = int(curr_exe.get_opt("pad-data")) else: pad_data = 0 # We only read data that's available. The code must handle the case # of not much data being available. - if end > (data_seg[1] - pad_data): - end = data_seg[1] - pad_data - if start < (data_seg[0] + pad_data): - start = data_seg[0] + pad_data + end = min(end, data_seg[1] - pad_data) + start = max(start, data_seg[0] + pad_data) - node.add_opt('--gps-start-time', int(start)) - node.add_opt('--gps-end-time', int(end)) - node.add_opt('--center-time', trig_time) + node.add_opt("--gps-start-time", int(start)) + node.add_opt("--gps-end-time", int(end)) + node.add_opt("--center-time", trig_time) if injection_file is not None: - node.add_input_opt('--injection-file', injection_file) + node.add_input_opt("--injection-file", injection_file) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node return node.output_files -def make_singles_timefreq(workflow, single, bank_file, trig_time, out_dir, - veto_file=None, time_window=10, data_segments=None, - tags=None): - """ Generate a singles_timefreq node and add it to workflow. + +def make_singles_timefreq( + workflow, + single, + bank_file, + trig_time, + out_dir, + veto_file=None, + time_window=10, + data_segments=None, + tags=None, +): + """ + Generate a singles_timefreq node and add it to workflow. This function generates a single node of the singles_timefreq executable and adds it to the current workflow. Parent/child relationships are set by the input/output files automatically. Parameters - ----------- + ---------- workflow: pycbc.workflow.core.Workflow The workflow class that stores the jobs that will be run. single: pycbc.workflow.core.File instance @@ -935,13 +1077,15 @@ def make_singles_timefreq(workflow, single, bank_file, trig_time, out_dir, trigger. This **must** be coalesced. tags: list (optional, default=None) List of tags to add to the created nodes, which determine file naming. + """ tags = [] if tags is None else tags makedir(out_dir) - name = 'plot_singles_timefreq' + name = "plot_singles_timefreq" - curr_exe = SingleTimeFreqExecutable(workflow.cp, name, ifos=[single.ifo], - out_dir=out_dir, tags=tags) + curr_exe = SingleTimeFreqExecutable( + workflow.cp, name, ifos=[single.ifo], out_dir=out_dir, tags=tags + ) # Determine start/end times, using data segments if needed. # Begin by choosing "optimal" times @@ -949,8 +1093,8 @@ def make_singles_timefreq(workflow, single, bank_file, trig_time, out_dir, end = trig_time + time_window node = curr_exe.create_node(valid_seg=segments.segment([start, end])) - node.add_input_opt('--trig-file', single) - node.add_input_opt('--bank-file', bank_file) + node.add_input_opt("--trig-file", single) + node.add_input_opt("--bank-file", bank_file) # Then if data_segments is available, check against that, and move if # needed @@ -960,33 +1104,35 @@ def make_singles_timefreq(workflow, single, bank_file, trig_time, out_dir, if trig_time in seg: data_seg = seg break - elif trig_time == -1.0: - node.add_opt('--gps-start-time', int(trig_time)) - node.add_opt('--gps-end-time', int(trig_time)) - node.add_opt('--center-time', trig_time) + if trig_time == -1.0: + node.add_opt("--gps-start-time", int(trig_time)) + node.add_opt("--gps-end-time", int(trig_time)) + node.add_opt("--center-time", trig_time) if veto_file: - node.add_input_opt('--veto-file', veto_file) + node.add_input_opt("--veto-file", veto_file) - node.add_opt('--detector', single.ifo) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.add_opt("--detector", single.ifo) + node.new_output_file_opt( + workflow.analysis_time, ".png", "--output-file" + ) workflow += node return node.output_files else: - err_msg = "Trig time {} ".format(trig_time) + err_msg = f"Trig time {trig_time} " err_msg += "does not seem to lie within any data segments. " err_msg += "This shouldn't be possible, please ask for help!" raise ValueError(err_msg) # Check for pad-data - if curr_exe.has_opt('pad-data'): - pad_data = int(curr_exe.get_opt('pad-data')) + if curr_exe.has_opt("pad-data"): + pad_data = int(curr_exe.get_opt("pad-data")) else: pad_data = 0 if abs(data_seg) < (2 * time_window + 2 * pad_data): tl = 2 * time_window + 2 * pad_data - err_msg = "I was asked to use {} seconds of data ".format(tl) + err_msg = f"I was asked to use {tl} seconds of data " err_msg += "to run a plot_singles_timefreq job. However, I have " - err_msg += "only {} seconds available.".format(abs(data_seg)) + err_msg += f"only {abs(data_seg)} seconds available." raise ValueError(err_msg) if data_seg[0] > (start - pad_data): start = data_seg[0] + pad_data @@ -999,24 +1145,26 @@ def make_singles_timefreq(workflow, single, bank_file, trig_time, out_dir, err_msg = "I shouldn't be here! Go ask Ian what he broke." raise ValueError(err_msg) - node.add_opt('--gps-start-time', int(start)) - node.add_opt('--gps-end-time', int(end)) - node.add_opt('--center-time', trig_time) + node.add_opt("--gps-start-time", int(start)) + node.add_opt("--gps-end-time", int(end)) + node.add_opt("--center-time", trig_time) if veto_file: - node.add_input_opt('--veto-file', veto_file) + node.add_input_opt("--veto-file", veto_file) - node.add_opt('--detector', single.ifo) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.add_opt("--detector", single.ifo) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node return node.output_files + def make_skipped_html(workflow, skipped_data, out_dir, tags): """ Make a html snippet from the list of skipped background coincidences """ - exe = Executable(workflow.cp, 'html_snippet', - ifos=workflow.ifos, out_dir=out_dir, tags=tags) + exe = Executable( + workflow.cp, "html_snippet", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) node = exe.create_node() @@ -1030,12 +1178,12 @@ def make_skipped_html(workflow, skipped_data, out_dir, tags): parsed_data[ifo][time] = parsed_data[ifo][time] + 1 n_events = len(skipped_data) - html_string = '"{} background events have been skipped '.format(n_events) - html_string += 'because one of their single triggers already appears ' - html_string += 'in the events followed up above. ' - html_string += 'Specifically, the following single detector triggers ' - html_string += 'were found in these coincidences. ' - html_template = '{} event at time {} appeared {} times. ' + html_string = f'"{n_events} background events have been skipped ' + html_string += "because one of their single triggers already appears " + html_string += "in the events followed up above. " + html_string += "Specifically, the following single detector triggers " + html_string += "were found in these coincidences. " + html_template = "{} event at time {} appeared {} times. " for ifo in parsed_data: for time in parsed_data[ifo]: n_occurances = parsed_data[ifo][time] @@ -1043,17 +1191,25 @@ def make_skipped_html(workflow, skipped_data, out_dir, tags): html_string += '"' - node.add_opt('--html-text', html_string) - node.add_opt('--title', '"Events were skipped"') - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node.add_opt("--html-text", html_string) + node.add_opt("--title", '"Events were skipped"') + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node files = node.output_files return files -def make_upload_files(workflow, psd_files, snr_timeseries, xml_all, - event_id, approximant, out_dir, channel_name, - tags=None): +def make_upload_files( + workflow, + psd_files, + snr_timeseries, + xml_all, + event_id, + approximant, + out_dir, + channel_name, + tags=None, +): """ Make files including xml, skymap fits and plots for uploading to gracedb for a given event @@ -1084,97 +1240,92 @@ def make_upload_files(workflow, psd_files, snr_timeseries, xml_all, ------- all_output_files: FileList List of all output files from this process + """ logging.info("Setting up upload files") indiv_xml_exe = Executable( - workflow.cp, - 'generate_xml', - ifos=workflow.ifos, out_dir=out_dir, - tags=tags + workflow.cp, "generate_xml", ifos=workflow.ifos, out_dir=out_dir, tags=tags ) logging.info("Setting up XML generation") xml_node = indiv_xml_exe.create_node() - xml_node.add_input_opt('--input-file', xml_all) - xml_node.add_opt('--event-id', event_id) - xml_node.add_input_list_opt('--psd-files', psd_files) - xml_node.add_input_list_opt('--snr-timeseries', snr_timeseries) - xml_node.add_opt('--channel-name', channel_name) + xml_node.add_input_opt("--input-file", xml_all) + xml_node.add_opt("--event-id", event_id) + xml_node.add_input_list_opt("--psd-files", psd_files) + xml_node.add_input_list_opt("--snr-timeseries", snr_timeseries) + xml_node.add_opt("--channel-name", channel_name) xml_node.new_output_file_opt( - workflow.analysis_time, - '.png', - '--snr-timeseries-plot', - tags=['snr'] + workflow.analysis_time, ".png", "--snr-timeseries-plot", tags=["snr"] ) xml_node.new_output_file_opt( - workflow.analysis_time, - '.png', - '--psd-plot', - tags=['psd'] + workflow.analysis_time, ".png", "--psd-plot", tags=["psd"] ) xml_out = xml_node.new_output_file_opt( - workflow.analysis_time, - '.xml', - '--output-file' + workflow.analysis_time, ".xml", "--output-file" ) workflow += xml_node logging.info("Setting up bayestar generation") bayestar_exe = Executable( - workflow.cp, - 'bayestar', - ifos=workflow.ifos, - out_dir=out_dir, - tags=tags + workflow.cp, "bayestar", ifos=workflow.ifos, out_dir=out_dir, tags=tags ) bayestar_node = bayestar_exe.create_node() - bayestar_node.add_input_opt('--event-xml', xml_out) + bayestar_node.add_input_opt("--event-xml", xml_out) fits_out = bayestar_node.new_output_file_opt( workflow.analysis_time, - '.fits', - '--output-file', + ".fits", + "--output-file", ) # This will be called if the approximant is within the bank - if approximant == b'SPAtmplt': + if approximant == b"SPAtmplt": # Bayestar doesn't use the SPAtmplt approximant - approximant = b'TaylorF2' + approximant = b"TaylorF2" if approximant is not None: - bayestar_node.add_opt('--approximant', approximant.decode()) + bayestar_node.add_opt("--approximant", approximant.decode()) workflow += bayestar_node logging.info("Setting up skymap plot generation") skymap_plot_exe = PlotExecutable( - workflow.cp, - 'skymap_plot', - ifos=workflow.ifos, - out_dir=out_dir, - tags=tags + workflow.cp, "skymap_plot", ifos=workflow.ifos, out_dir=out_dir, tags=tags ) skymap_plot_node = skymap_plot_exe.create_node() - skymap_plot_node.add_input_opt('', fits_out) + skymap_plot_node.add_input_opt("", fits_out) skymap_plot_node.new_output_file_opt( workflow.analysis_time, - '.png', - '-o', + ".png", + "-o", ) workflow += skymap_plot_node - all_output_files = xml_node.output_files + bayestar_node.output_files + \ - skymap_plot_node.output_files + all_output_files = ( + xml_node.output_files + + bayestar_node.output_files + + skymap_plot_node.output_files + ) return all_output_files -def setup_upload_prep_minifollowups(workflow, coinc_file, xml_all_file, - single_triggers, psd_files, - tmpltbank_file, insp_segs, insp_data_name, - insp_anal_name, dax_output, out_dir, - tags=None): - """ Create plots that followup the Nth loudest coincident injection +def setup_upload_prep_minifollowups( + workflow, + coinc_file, + xml_all_file, + single_triggers, + psd_files, + tmpltbank_file, + insp_segs, + insp_data_name, + insp_anal_name, + dax_output, + out_dir, + tags=None, +): + """ + Create plots that followup the Nth loudest coincident injection from a statmap produced HDF file. Parameters @@ -1213,14 +1364,15 @@ def setup_upload_prep_minifollowups(workflow, coinc_file, xml_all_file, layout: list A list of tuples which specify the displayed file layout for the minifollowups plots. + """ - logger.info('Entering minifollowups module') + logger.info("Entering minifollowups module") - if not workflow.cp.has_section('workflow-minifollowups'): - msg = 'There is no [workflow-minifollowups] section in ' - msg += 'configuration file' + if not workflow.cp.has_section("workflow-minifollowups"): + msg = "There is no [workflow-minifollowups] section in " + msg += "configuration file" logger.info(msg) - logger.info('Leaving minifollowups') + logger.info("Leaving minifollowups") return tags = [] if tags is None else tags @@ -1228,37 +1380,42 @@ def setup_upload_prep_minifollowups(workflow, coinc_file, xml_all_file, makedir(out_dir) # turn the config file into a File class - config_path = os.path.abspath(dax_output + '/' + '_'.join(tags) + \ - 'upload_prep_minifollowup.ini') - workflow.cp.write(open(config_path, 'w')) + config_path = os.path.abspath( + dax_output + "/" + "_".join(tags) + "upload_prep_minifollowup.ini" + ) + workflow.cp.write(open(config_path, "w")) config_file = resolve_url_to_file(config_path) - exe = Executable(workflow.cp, 'upload_prep_minifollowup', - ifos=workflow.ifos, out_dir=dax_output, tags=tags) + exe = Executable( + workflow.cp, + "upload_prep_minifollowup", + ifos=workflow.ifos, + out_dir=dax_output, + tags=tags, + ) node = exe.create_node() - node.add_input_opt('--config-files', config_file) - node.add_input_opt('--xml-all-file', xml_all_file) - node.add_input_opt('--bank-file', tmpltbank_file) - node.add_input_opt('--statmap-file', coinc_file) - node.add_multiifo_input_list_opt('--single-detector-triggers', - single_triggers) - node.add_multiifo_input_list_opt('--psd-files', psd_files) - node.add_input_opt('--inspiral-segments', insp_segs) - node.add_opt('--inspiral-data-read-name', insp_data_name) - node.add_opt('--inspiral-data-analyzed-name', insp_anal_name) + node.add_input_opt("--config-files", config_file) + node.add_input_opt("--xml-all-file", xml_all_file) + node.add_input_opt("--bank-file", tmpltbank_file) + node.add_input_opt("--statmap-file", coinc_file) + node.add_multiifo_input_list_opt("--single-detector-triggers", single_triggers) + node.add_multiifo_input_list_opt("--psd-files", psd_files) + node.add_input_opt("--inspiral-segments", insp_segs) + node.add_opt("--inspiral-data-read-name", insp_data_name) + node.add_opt("--inspiral-data-analyzed-name", insp_anal_name) if tags: - node.add_list_opt('--tags', tags) - node.new_output_file_opt(workflow.analysis_time, '.dax', '--dax-file') - node.new_output_file_opt(workflow.analysis_time, '.dax.map', '--output-map') + node.add_list_opt("--tags", tags) + node.new_output_file_opt(workflow.analysis_time, ".dax", "--dax-file") + node.new_output_file_opt(workflow.analysis_time, ".dax.map", "--output-map") name = node.output_files[0].name map_file = node.output_files[1] - node.add_opt('--workflow-name', name) - node.add_opt('--output-dir', out_dir) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", name) + node.add_opt("--output-dir", out_dir) + node.add_opt("--dax-file-directory", ".") workflow += node @@ -1267,11 +1424,14 @@ def setup_upload_prep_minifollowups(workflow, coinc_file, xml_all_file, # determine if a staging site has been specified job = SubWorkflow(fil.name, is_planned=False) - input_files = [xml_all_file, tmpltbank_file, coinc_file, insp_segs] + \ - single_triggers + psd_files + input_files = ( + [xml_all_file, tmpltbank_file, coinc_file, insp_segs] + + single_triggers + + psd_files + ) job.add_inputs(*input_files) - job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) job.add_into_workflow(workflow) - logger.info('Leaving minifollowups module') + logger.info("Leaving minifollowups module") diff --git a/pycbc/workflow/pegasus_sites.py b/pycbc/workflow/pegasus_sites.py index 5224e487a43..f8a77dcb9af 100644 --- a/pycbc/workflow/pegasus_sites.py +++ b/pycbc/workflow/pegasus_sites.py @@ -7,7 +7,8 @@ # # ============================================================================= # -""" This module provides default site catalogs, which should be suitable for +""" +This module provides default site catalogs, which should be suitable for most use cases. You can override individual details here. It should also be possible to implement a new site, but not sure how that would work in practice. """ @@ -20,13 +21,20 @@ from urllib.parse import urljoin from urllib.request import pathname2url -from Pegasus.api import Directory, FileServer, Site, Operation, Namespace -from Pegasus.api import Arch, OS, SiteCatalog - -from pycbc.version import last_release, version, release # noqa +from Pegasus.api import ( + OS, + Arch, + Directory, + FileServer, + Namespace, + Operation, + Site, + SiteCatalog, +) +from pycbc.version import last_release, release, version -logger = logging.getLogger('pycbc.workflow.pegasus_sites') +logger = logging.getLogger("pycbc.workflow.pegasus_sites") if release: sing_version = version @@ -36,48 +44,55 @@ # NOTE urllib is weird. For some reason it only allows known schemes and will # give *wrong* results, rather then failing, if you use something like gsiftp # We can add schemes explicitly, as below, but be careful with this! -urllib.parse.uses_relative.append('gsiftp') -urllib.parse.uses_netloc.append('gsiftp') +urllib.parse.uses_relative.append("gsiftp") +urllib.parse.uses_netloc.append("gsiftp") -KNOWN_SITES = ['local', 'condorpool_symlink', - 'condorpool_copy', 'condorpool_shared', 'osg'] +KNOWN_SITES = [ + "local", + "condorpool_symlink", + "condorpool_copy", + "condorpool_shared", + "osg", +] def add_site_pegasus_profile(site, cp): """Add options from [pegasus_profile] in configparser to site""" # Add global profile information - if cp.has_section('pegasus_profile'): - add_ini_site_profile(site, cp, 'pegasus_profile') + if cp.has_section("pegasus_profile"): + add_ini_site_profile(site, cp, "pegasus_profile") # Add site-specific profile information - if cp.has_section('pegasus_profile-{}'.format(site.name)): - add_ini_site_profile(site, cp, 'pegasus_profile-{}'.format(site.name)) + if cp.has_section(f"pegasus_profile-{site.name}"): + add_ini_site_profile(site, cp, f"pegasus_profile-{site.name}") def add_ini_site_profile(site, cp, sec): """Add options from sec in configparser to site""" for opt in cp.options(sec): - namespace = opt.split('|')[0] - if namespace in ('pycbc', 'container'): + namespace = opt.split("|")[0] + if namespace in ("pycbc", "container"): continue value = cp.get(sec, opt).strip() - key = opt.split('|')[1] + key = opt.split("|")[1] site.add_profiles(Namespace(namespace), key=key, value=value) def add_local_site(sitecat, cp, local_path, local_url): """Add the local site to site catalog""" # local_url must end with a '/' - if not local_url.endswith('/'): - local_url = local_url + '/' + if not local_url.endswith("/"): + local_url = local_url + "/" local = Site("local", arch=Arch.X86_64, os_type=OS.LINUX) add_site_pegasus_profile(local, cp) - local_dir = Directory(Directory.SHARED_SCRATCH, - path=os.path.join(local_path, 'local-site-scratch')) - local_file_serv = FileServer(urljoin(local_url, 'local-site-scratch'), - Operation.ALL) + local_dir = Directory( + Directory.SHARED_SCRATCH, path=os.path.join(local_path, "local-site-scratch") + ) + local_file_serv = FileServer( + urljoin(local_url, "local-site-scratch"), Operation.ALL + ) local_dir.add_file_servers(local_file_serv) local.add_directories(local_dir) @@ -91,24 +106,19 @@ def add_condorpool_symlink_site(sitecat, cp): add_site_pegasus_profile(site, cp) site.add_profiles(Namespace.PEGASUS, key="style", value="condor") - site.add_profiles(Namespace.PEGASUS, key="data.configuration", - value="nonsharedfs") - site.add_profiles(Namespace.PEGASUS, key='transfer.bypass.input.staging', - value="true") - site.add_profiles(Namespace.PEGASUS, key='auxillary.local', - value="true") - site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", - value="False") - site.add_profiles(Namespace.CONDOR, key="should_transfer_files", - value="Yes") - site.add_profiles(Namespace.CONDOR, key="when_to_transfer_output", - value="ON_EXIT_OR_EVICT") - site.add_profiles(Namespace.CONDOR, key="My.DESIRED_Sites", - value='"nogrid"') - site.add_profiles(Namespace.CONDOR, key="My.IS_GLIDEIN", - value='"False"') - site.add_profiles(Namespace.CONDOR, key="My.flock_local", - value="True") + site.add_profiles(Namespace.PEGASUS, key="data.configuration", value="nonsharedfs") + site.add_profiles( + Namespace.PEGASUS, key="transfer.bypass.input.staging", value="true" + ) + site.add_profiles(Namespace.PEGASUS, key="auxillary.local", value="true") + site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", value="False") + site.add_profiles(Namespace.CONDOR, key="should_transfer_files", value="Yes") + site.add_profiles( + Namespace.CONDOR, key="when_to_transfer_output", value="ON_EXIT_OR_EVICT" + ) + site.add_profiles(Namespace.CONDOR, key="My.DESIRED_Sites", value='"nogrid"') + site.add_profiles(Namespace.CONDOR, key="My.IS_GLIDEIN", value='"False"') + site.add_profiles(Namespace.CONDOR, key="My.flock_local", value="True") site.add_profiles(Namespace.DAGMAN, key="retry", value="2") sitecat.add_sites(site) @@ -119,27 +129,21 @@ def add_condorpool_copy_site(sitecat, cp): add_site_pegasus_profile(site, cp) site.add_profiles(Namespace.PEGASUS, key="style", value="condor") - site.add_profiles(Namespace.PEGASUS, key="data.configuration", - value="condorio") - site.add_profiles(Namespace.PEGASUS, key='transfer.bypass.input.staging', - value="true") + site.add_profiles(Namespace.PEGASUS, key="data.configuration", value="condorio") + site.add_profiles( + Namespace.PEGASUS, key="transfer.bypass.input.staging", value="true" + ) # This explicitly disables symlinking - site.add_profiles(Namespace.PEGASUS, key='nosymlink', - value=True) - site.add_profiles(Namespace.PEGASUS, key='auxillary.local', - value="true") - site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", - value="False") - site.add_profiles(Namespace.CONDOR, key="should_transfer_files", - value="Yes") - site.add_profiles(Namespace.CONDOR, key="when_to_transfer_output", - value="ON_EXIT_OR_EVICT") - site.add_profiles(Namespace.CONDOR, key="My.DESIRED_Sites", - value='"nogrid"') - site.add_profiles(Namespace.CONDOR, key="My.IS_GLIDEIN", - value='"False"') - site.add_profiles(Namespace.CONDOR, key="My.flock_local", - value="True") + site.add_profiles(Namespace.PEGASUS, key="nosymlink", value=True) + site.add_profiles(Namespace.PEGASUS, key="auxillary.local", value="true") + site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", value="False") + site.add_profiles(Namespace.CONDOR, key="should_transfer_files", value="Yes") + site.add_profiles( + Namespace.CONDOR, key="when_to_transfer_output", value="ON_EXIT_OR_EVICT" + ) + site.add_profiles(Namespace.CONDOR, key="My.DESIRED_Sites", value='"nogrid"') + site.add_profiles(Namespace.CONDOR, key="My.IS_GLIDEIN", value='"False"') + site.add_profiles(Namespace.CONDOR, key="My.flock_local", value="True") site.add_profiles(Namespace.DAGMAN, key="retry", value="2") sitecat.add_sites(site) @@ -147,53 +151,49 @@ def add_condorpool_copy_site(sitecat, cp): def add_condorpool_shared_site(sitecat, cp, local_path, local_url): """Add condorpool_shared site to site catalog""" # local_url must end with a '/' - if not local_url.endswith('/'): - local_url = local_url + '/' + if not local_url.endswith("/"): + local_url = local_url + "/" site = Site("condorpool_shared", arch=Arch.X86_64, os_type=OS.LINUX) add_site_pegasus_profile(site, cp) # It's annoying that this is needed! - local_dir = Directory(Directory.SHARED_SCRATCH, - path=os.path.join(local_path, 'cpool-site-scratch')) - local_file_serv = FileServer(urljoin(local_url, 'cpool-site-scratch'), - Operation.ALL) + local_dir = Directory( + Directory.SHARED_SCRATCH, path=os.path.join(local_path, "cpool-site-scratch") + ) + local_file_serv = FileServer( + urljoin(local_url, "cpool-site-scratch"), Operation.ALL + ) local_dir.add_file_servers(local_file_serv) site.add_directories(local_dir) site.add_profiles(Namespace.PEGASUS, key="style", value="condor") - site.add_profiles(Namespace.PEGASUS, key="data.configuration", - value="sharedfs") - site.add_profiles(Namespace.PEGASUS, key='transfer.bypass.input.staging', - value="true") - site.add_profiles(Namespace.PEGASUS, key='auxillary.local', - value="true") - site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", - value="False") - site.add_profiles(Namespace.CONDOR, key="should_transfer_files", - value="Yes") - site.add_profiles(Namespace.CONDOR, key="when_to_transfer_output", - value="ON_EXIT_OR_EVICT") - site.add_profiles(Namespace.CONDOR, key="My.DESIRED_Sites", - value='"nogrid"') - site.add_profiles(Namespace.CONDOR, key="My.IS_GLIDEIN", - value='"False"') - site.add_profiles(Namespace.CONDOR, key="My.flock_local", - value="True") + site.add_profiles(Namespace.PEGASUS, key="data.configuration", value="sharedfs") + site.add_profiles( + Namespace.PEGASUS, key="transfer.bypass.input.staging", value="true" + ) + site.add_profiles(Namespace.PEGASUS, key="auxillary.local", value="true") + site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", value="False") + site.add_profiles(Namespace.CONDOR, key="should_transfer_files", value="Yes") + site.add_profiles( + Namespace.CONDOR, key="when_to_transfer_output", value="ON_EXIT_OR_EVICT" + ) + site.add_profiles(Namespace.CONDOR, key="My.DESIRED_Sites", value='"nogrid"') + site.add_profiles(Namespace.CONDOR, key="My.IS_GLIDEIN", value='"False"') + site.add_profiles(Namespace.CONDOR, key="My.flock_local", value="True") site.add_profiles(Namespace.DAGMAN, key="retry", value="2") # Need to set PEGASUS_HOME - peg_home = which('pegasus-plan') + peg_home = which("pegasus-plan") if peg_home is None: raise RuntimeError( - 'pegasus-plan command not found. ' - 'Make sure Pegasus is correctly installed.' + "pegasus-plan command not found. Make sure Pegasus is correctly installed." ) - if not peg_home.endswith('bin/pegasus-plan'): + if not peg_home.endswith("bin/pegasus-plan"): raise RuntimeError( - f'path to pegasus-plan is weird: {peg_home}. ' - 'Make sure Pegasus is correctly installed.' + f"path to pegasus-plan is weird: {peg_home}. " + "Make sure Pegasus is correctly installed." ) - peg_home = peg_home.replace('bin/pegasus-plan', '') + peg_home = peg_home.replace("bin/pegasus-plan", "") site.add_profiles(Namespace.ENV, key="PEGASUS_HOME", value=peg_home) sitecat.add_sites(site) @@ -210,77 +210,83 @@ def add_osg_site(sitecat, cp): site = Site("osg", arch=Arch.X86_64, os_type=OS.LINUX) add_site_pegasus_profile(site, cp) site.add_profiles(Namespace.PEGASUS, key="style", value="condor") - site.add_profiles(Namespace.PEGASUS, key="data.configuration", - value="condorio") - site.add_profiles(Namespace.PEGASUS, key='transfer.bypass.input.staging', - value="true") - site.add_profiles(Namespace.CONDOR, key="should_transfer_files", - value="Yes") - site.add_profiles(Namespace.CONDOR, key="when_to_transfer_output", - value="ON_SUCCESS") - site.add_profiles(Namespace.CONDOR, key="success_exit_code", - value="0") - site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", - value="True") - site.add_profiles(Namespace.CONDOR, key="ulog_execute_attrs", - value="GLIDEIN_Site") - site.add_profiles(Namespace.CONDOR, key="My.InitializeModulesEnv", - value="False") - site.add_profiles(Namespace.CONDOR, key="My.SingularityCleanEnv", - value="True") + site.add_profiles(Namespace.PEGASUS, key="data.configuration", value="condorio") + site.add_profiles( + Namespace.PEGASUS, key="transfer.bypass.input.staging", value="true" + ) + site.add_profiles(Namespace.CONDOR, key="should_transfer_files", value="Yes") + site.add_profiles( + Namespace.CONDOR, key="when_to_transfer_output", value="ON_SUCCESS" + ) + site.add_profiles(Namespace.CONDOR, key="success_exit_code", value="0") + site.add_profiles(Namespace.CONDOR, key="My.OpenScienceGrid", value="True") + site.add_profiles(Namespace.CONDOR, key="ulog_execute_attrs", value="GLIDEIN_Site") + site.add_profiles(Namespace.CONDOR, key="My.InitializeModulesEnv", value="False") + site.add_profiles(Namespace.CONDOR, key="My.SingularityCleanEnv", value="True") # These numbers below correspond to the codes in table B.2 here: # https://htcondor.readthedocs.io/en/24.0/codes-other-values/job-event-log-codes.html # Values recommended by a condor expert - site.add_profiles(Namespace.CONDOR, key="My.DAGManNodesMask", - value=r"\"0,1,2,4,5,7,8,9,10,11,12,13,16,17,24,27,35,36,40\"") - site.add_profiles(Namespace.CONDOR, key="Requirements", - value="(HAS_SINGULARITY =?= TRUE) && " - "(IS_GLIDEIN =?= True) && " - "(HAS_CVMFS_singularity_opensciencegrid_org =?= True)") + site.add_profiles( + Namespace.CONDOR, + key="My.DAGManNodesMask", + value=r"\"0,1,2,4,5,7,8,9,10,11,12,13,16,17,24,27,35,36,40\"", + ) + site.add_profiles( + Namespace.CONDOR, + key="Requirements", + value="(HAS_SINGULARITY =?= TRUE) && " + "(IS_GLIDEIN =?= True) && " + "(HAS_CVMFS_singularity_opensciencegrid_org =?= True)", + ) cvmfs_loc = '"/cvmfs/singularity.opensciencegrid.org/pycbc/pycbc-el8:v' cvmfs_loc += sing_version + '"' - site.add_profiles(Namespace.CONDOR, key="My.SingularityImage", - value=cvmfs_loc) + site.add_profiles(Namespace.CONDOR, key="My.SingularityImage", value=cvmfs_loc) # On OSG failure rate is high site.add_profiles(Namespace.DAGMAN, key="retry", value="4") - site.add_profiles(Namespace.ENV, key="LAL_DATA_PATH", - value="/cvmfs/software.igwn.org/pycbc/lalsuite-extra/current/share/lalsimulation") + site.add_profiles( + Namespace.ENV, + key="LAL_DATA_PATH", + value="/cvmfs/software.igwn.org/pycbc/lalsuite-extra/current/share/lalsimulation", + ) # Add MKL location to LD_LIBRARY_PATH for OSG - site.add_profiles(Namespace.ENV, key="LD_LIBRARY_PATH", - value="/usr/local/lib:/.singularity.d/libs") + site.add_profiles( + Namespace.ENV, + key="LD_LIBRARY_PATH", + value="/usr/local/lib:/.singularity.d/libs", + ) sitecat.add_sites(site) def add_site(sitecat, sitename, cp, out_dir=None): """Add site sitename to site catalog""" # Allow local site scratch to be overriden for any site which uses it - sec = 'pegasus_profile-{}'.format(sitename) - opt = 'pycbc|site-scratch' + sec = f"pegasus_profile-{sitename}" + opt = "pycbc|site-scratch" if cp.has_option(sec, opt): out_dir = os.path.abspath(cp.get(sec, opt)) - if cp.has_option(sec, 'pycbc|unique-scratch'): - scratchdir = tempfile.mkdtemp(prefix='pycbc-tmp_', dir=out_dir) + if cp.has_option(sec, "pycbc|unique-scratch"): + scratchdir = tempfile.mkdtemp(prefix="pycbc-tmp_", dir=out_dir) os.chmod(scratchdir, 0o755) try: - os.symlink(scratchdir, '{}-site-scratch'.format(sitename)) + os.symlink(scratchdir, f"{sitename}-site-scratch") except OSError: pass out_dir = scratchdir elif out_dir is None: out_dir = os.getcwd() - local_url = urljoin('file://', pathname2url(out_dir)) - if sitename == 'local': + local_url = urljoin("file://", pathname2url(out_dir)) + if sitename == "local": add_local_site(sitecat, cp, out_dir, local_url) - elif sitename == 'condorpool_symlink': + elif sitename == "condorpool_symlink": add_condorpool_symlink_site(sitecat, cp) - elif sitename == 'condorpool_copy': + elif sitename == "condorpool_copy": add_condorpool_copy_site(sitecat, cp) - elif sitename == 'condorpool_shared': + elif sitename == "condorpool_shared": add_condorpool_shared_site(sitecat, cp, out_dir, local_url) - elif sitename == 'osg': + elif sitename == "osg": add_osg_site(sitecat, cp) else: - raise ValueError("Do not recognize site {}".format(sitename)) + raise ValueError(f"Do not recognize site {sitename}") def make_catalog(cp, out_dir): diff --git a/pycbc/workflow/pegasus_workflow.py b/pycbc/workflow/pegasus_workflow.py index 52457268795..2294b274a8f 100644 --- a/pycbc/workflow/pegasus_workflow.py +++ b/pycbc/workflow/pegasus_workflow.py @@ -23,56 +23,57 @@ # # ============================================================================= # -""" This module provides thin wrappers around Pegasus.DAX3 functionality that +""" +This module provides thin wrappers around Pegasus.DAX3 functionality that provides additional abstraction and argument handling. """ + +import logging import os import shutil -import logging -import tempfile import subprocess +import tempfile import warnings -from packaging import version -from urllib.request import pathname2url from urllib.parse import urljoin, urlsplit +from urllib.request import pathname2url import Pegasus.api as dax +from packaging import version -logger = logging.getLogger('pycbc.workflow.pegasus_workflow') +logger = logging.getLogger("pycbc.workflow.pegasus_workflow") -PEGASUS_FILE_DIRECTORY = os.path.join(os.path.dirname(__file__), - 'pegasus_files') +PEGASUS_FILE_DIRECTORY = os.path.join(os.path.dirname(__file__), "pegasus_files") -class ProfileShortcuts(object): - """ Container of common methods for setting pegasus profile information +class ProfileShortcuts: + """ + Container of common methods for setting pegasus profile information on Executables and nodes. This class expects to be inherited from and for a add_profile method to be implemented. """ + def set_memory(self, size): - """ Set the amount of memory that is required in megabytes - """ - self.add_profile('condor', 'request_memory', '%sM' % size) + """Set the amount of memory that is required in megabytes""" + self.add_profile("condor", "request_memory", "%sM" % size) def set_storage(self, size): - """ Set the amount of storage required in megabytes - """ - self.add_profile('condor', 'request_disk', '%sM' % size) + """Set the amount of storage required in megabytes""" + self.add_profile("condor", "request_disk", "%sM" % size) def set_num_cpus(self, number): - self.add_profile('condor', 'request_cpus', number) + self.add_profile("condor", "request_cpus", number) def set_universe(self, universe): - if universe == 'standard': + if universe == "standard": self.add_profile("pegasus", "gridstart", "none") self.add_profile("condor", "universe", universe) def set_category(self, category): - self.add_profile('dagman', 'category', category) + self.add_profile("dagman", "category", category) def set_priority(self, priority): - self.add_profile('dagman', 'priority', priority) + self.add_profile("dagman", "priority", priority) def set_num_retries(self, number): self.add_profile("dagman", "retry", number) @@ -82,12 +83,13 @@ def set_execution_site(self, site): class Executable(ProfileShortcuts): - """ The workflow representation of an Executable - """ + """The workflow representation of an Executable""" + id = 0 - def __init__(self, name, os='linux', - arch='x86_64', installed=False, - container=None): + + def __init__( + self, name, os="linux", arch="x86_64", installed=False, container=None + ): self.logical_name = name + "_ID%s" % str(Executable.id) self.pegasus_name = name Executable.id += 1 @@ -107,20 +109,15 @@ def create_transformation(self, site, url): is_stageable=self.installed, arch=self.arch, os_type=self.os, - container=self.container + container=self.container, ) transform.pycbc_name = self.pegasus_name for (namespace, key), value in self.profiles.items(): - transform.add_profiles( - dax.Namespace(namespace), - key=key, - value=value - ) + transform.add_profiles(dax.Namespace(namespace), key=key, value=value) self.transformations[site] = transform def add_profile(self, namespace, key, value): - """ Add profile information to this executable - """ + """Add profile information to this executable""" if self.transformations: err_msg = "Need code changes to be able to add profiles " err_msg += "after transformations are created." @@ -129,11 +126,16 @@ def add_profile(self, namespace, key, value): class Transformation(dax.Transformation): - def is_same_as(self, other): - test_vals = ['namespace', 'version'] - test_site_vals = ['arch', 'os_type', 'os_release', - 'os_version', 'bypass', 'container'] + test_vals = ["namespace", "version"] + test_site_vals = [ + "arch", + "os_type", + "os_release", + "os_version", + "bypass", + "container", + ] # Check for logical name first if not self.pycbc_name == other.pycbc_name: return False @@ -172,7 +174,7 @@ def is_same_as(self, other): class Node(ProfileShortcuts): def __init__(self, transformation): self.in_workflow = False - self.transformation=transformation + self.transformation = transformation self._inputs = [] self._outputs = [] self._dax_node = dax.Job(transformation) @@ -189,17 +191,17 @@ def __init__(self, transformation): self._raw_options = [] def add_arg(self, arg): - """ Add an argument - """ + """Add an argument""" if not isinstance(arg, File): arg = str(arg) self._args += [arg] def add_raw_arg(self, arg): - """ Add an argument to the command line of this job, but do *NOT* add - white space between arguments. This can be added manually by adding - ' ' if needed + """ + Add an argument to the command line of this job, but do *NOT* add + white space between arguments. This can be added manually by adding + ' ' if needed """ if not isinstance(arg, File): arg = str(arg) @@ -207,10 +209,10 @@ def add_raw_arg(self, arg): self._raw_options += [arg] def add_opt(self, opt, value=None, check_existing_options=True, **kwargs): # pylint:disable=unused-argument - """ Add an option - """ - if check_existing_options and (opt in self._options - or opt in self._raw_options): + """Add an option""" + if check_existing_options and ( + opt in self._options or opt in self._raw_options + ): err_msg = ( "Trying to set option %s with value %s, but it " "has already been provided by the configuration file. " @@ -225,16 +227,14 @@ def add_opt(self, opt, value=None, check_existing_options=True, **kwargs): # py else: self._options += [opt] - #private functions to add input and output data sources/sinks + # private functions to add input and output data sources/sinks def _add_input(self, inp): - """ Add as source of input data - """ + """Add as source of input data""" self._inputs += [inp] self._dax_node.add_inputs(inp) def _add_output(self, out): - """ Add as destination of output data - """ + """Add as destination of output data""" self._outputs += [out] out.node = self stage_out = out.storage_path is not None @@ -242,117 +242,93 @@ def _add_output(self, out): # public functions to add options, arguments with or without data sources def add_input(self, inp): - """Declares an input file without adding it as a command-line option. - """ + """Declares an input file without adding it as a command-line option.""" self._add_input(inp) def add_output(self, inp): - """Declares an output file without adding it as a command-line option. - """ + """Declares an output file without adding it as a command-line option.""" self._add_output(inp) def add_input_opt(self, opt, inp, **kwargs): - """ Add an option that determines an input - """ + """Add an option that determines an input""" self.add_opt(opt, inp._dax_repr(), **kwargs) self._add_input(inp) def add_output_opt(self, opt, out, **kwargs): - """ Add an option that determines an output - """ + """Add an option that determines an output""" self.add_opt(opt, out._dax_repr(), **kwargs) self._add_output(out) def add_output_list_opt(self, opt, outputs, **kwargs): - """ Add an option that determines a list of outputs - """ + """Add an option that determines a list of outputs""" self.add_opt(opt, **kwargs) # Never check existing options for list option values - if 'check_existing_options' in kwargs: - kwargs['check_existing_options'] = False + if "check_existing_options" in kwargs: + kwargs["check_existing_options"] = False for out in outputs: self.add_opt(out, **kwargs) self._add_output(out) def add_input_list_opt(self, opt, inputs, **kwargs): - """ Add an option that determines a list of inputs - """ + """Add an option that determines a list of inputs""" self.add_opt(opt, **kwargs) # Never check existing options for list option values - if 'check_existing_options' in kwargs: - del kwargs['check_existing_options'] + kwargs.pop("check_existing_options", None) for inp in inputs: - self.add_opt( - inp, - check_existing_options=False, - **kwargs - ) + self.add_opt(inp, check_existing_options=False, **kwargs) self._add_input(inp) def add_list_opt(self, opt, values, **kwargs): - """ Add an option with a list of non-file parameters. - """ + """Add an option with a list of non-file parameters.""" self.add_opt(opt, **kwargs) # Never check existing options for list option values - if 'check_existing_options' in kwargs: - del kwargs['check_existing_options'] + kwargs.pop("check_existing_options", None) for val in values: - self.add_opt( - val, - check_existing_options=False, - **kwargs - ) + self.add_opt(val, check_existing_options=False, **kwargs) def add_input_arg(self, inp): - """ Add an input as an argument - """ + """Add an input as an argument""" self.add_arg(inp._dax_repr()) self._add_input(inp) def add_output_arg(self, out): - """ Add an output as an argument - """ + """Add an output as an argument""" self.add_arg(out._dax_repr()) self._add_output(out) def new_output_file_opt(self, opt, name): - """ Add an option and return a new file handle - """ + """Add an option and return a new file handle""" fil = File(name) self.add_output_opt(opt, fil) return fil # functions to describe properties of this node def add_profile(self, namespace, key, value): - """ Add profile information to this node at the DAX level - """ - self._dax_node.add_profiles( - dax.Namespace(namespace), - key=key, - value=value - ) + """Add profile information to this node at the DAX level""" + self._dax_node.add_profiles(dax.Namespace(namespace), key=key, value=value) def _finalize(self): if len(self._raw_options): - raw_args = [''.join([str(a) for a in self._raw_options])] + raw_args = ["".join([str(a) for a in self._raw_options])] else: raw_args = [] args = self._args + raw_args + self._options self._dax_node.add_args(*args) -class Workflow(object): - """ - """ - def __init__(self, name='my_workflow', directory=None, cache_file=None, - dax_file_name=None): +class Workflow: + """ """ + + def __init__( + self, name="my_workflow", directory=None, cache_file=None, dax_file_name=None + ): # Pegasus logging is fairly verbose, quieten it down a bit # This sets the logger to one level less verbose than the root # (pycbc) logger curr_level = logging.getLogger().level # Get the logger associated with the Pegasus workflow import - pegasus_logger = logging.getLogger('Pegasus') + pegasus_logger = logging.getLogger("Pegasus") pegasus_logger.setLevel(curr_level + 10) self.name = name self._rc = dax.ReplicaCatalog() @@ -374,19 +350,19 @@ def __init__(self, name='my_workflow', directory=None, cache_file=None, self.in_workflow = False self.sub_workflows = [] if dax_file_name is None: - self.filename = self.name + '.dax' + self.filename = self.name + ".dax" else: self.filename = dax_file_name self._adag = dax.Workflow(self.filename) # A pegasus job version of this workflow for use if it is included # within a larger workflow - self._as_job = SubWorkflow(self.filename, is_planned=False, - _id=self.name) + self._as_job = SubWorkflow(self.filename, is_planned=False, _id=self.name) self._swinputs = [] def add_workflow(self, workflow): - """ Add a sub-workflow to this workflow + """ + Add a sub-workflow to this workflow This function adds a sub-workflow of Workflow class to this workflow. Parent child relationships are determined by data dependencies @@ -395,6 +371,7 @@ def add_workflow(self, workflow): ---------- workflow : Workflow instance The sub-workflow to add to this one + """ workflow.in_workflow = self self.sub_workflows += [workflow] @@ -414,7 +391,9 @@ def add_explicit_dependancy(self, parent, child): ---------- parent : Node, Workflow or SubWorkflow instance child : Node, Workflow or SubWorkflow instance + """ + def convert(thing): if isinstance(thing, Workflow): return thing._as_job @@ -422,7 +401,7 @@ def convert(thing): return thing._dax_node if isinstance(thing, SubWorkflow): return thing - raise TypeError('ayee, cannot handle this dependancy!') + raise TypeError("ayee, cannot handle this dependancy!") self._adag.add_dependency(convert(parent), children=[convert(child)]) @@ -444,11 +423,13 @@ def add_subworkflow_dependancy(self, parent_workflow, child_workflow): child_workflow : Workflow instance The sub-workflow to add as the child dependence. Must be a sub-workflow of this workflow. + """ self.add_explicit_dependancy(parent_workflow, child_workflow) def add_transformation(self, tranformation): - """ Add a transformation to this workflow + """ + Add a transformation to this workflow Adds the input transformation to this workflow. @@ -456,11 +437,13 @@ def add_transformation(self, tranformation): ---------- transformation : Pegasus.api.Transformation The transformation to be added. + """ self._tc.add_transformations(tranformation) def add_container(self, container): - """ Add a container to this workflow + """ + Add a container to this workflow Adds the input container to this workflow. @@ -468,11 +451,13 @@ def add_container(self, container): ---------- container : Pegasus.api.Container The container to be added. + """ self._tc.add_containers(container) def add_node(self, node): - """ Add a node to this workflow + """ + Add a node to this workflow This function adds nodes to the workflow. It also determines parent/child relations from the inputs to this job. @@ -481,6 +466,7 @@ def add_node(self, node): ---------- node : pycbc.workflow.pegasus_workflow.Node A node that should be executed as part of this workflow. + """ node._finalize() node.in_workflow = self @@ -496,9 +482,11 @@ def add_node(self, node): break else: self._transformations += [node.transformation] - lgc = (hasattr(node, 'executable') - and node.executable.container is not None - and node.executable.container not in self._containers) + lgc = ( + hasattr(node, "executable") + and node.executable.container is not None + and node.executable.container not in self._containers + ) if lgc: self._containers.append(node.executable.container) @@ -516,14 +504,15 @@ def add_node(self, node): # Don't need to do anything here. continue - elif inp.node is not None and not inp.node.in_workflow: + if inp.node is not None and not inp.node.in_workflow: # This error should be rare, but can happen. If a Node hasn't # yet been added to a workflow, this logic breaks. Always add # nodes in order that files will be produced. - raise ValueError('Parents of this node must be added to the ' - 'workflow first.') + raise ValueError( + "Parents of this node must be added to the workflow first." + ) - elif inp.node is None: + if inp.node is None: # File is external to the workflow (e.g. a pregenerated # template bank). (if inp.node is None) if inp not in self._inputs: @@ -536,10 +525,12 @@ def add_node(self, node): self._inputs += [inp] self._swinputs += [inp] else: - err_msg = ("I don't understand how to deal with an input file " - "here. Ian doesn't think this message should be " - "possible, but if you get here something has gone " - "wrong and will need debugging!") + err_msg = ( + "I don't understand how to deal with an input file " + "here. Ian doesn't think this message should be " + "possible, but if you get here something has gone " + "wrong and will need debugging!" + ) raise ValueError(err_msg) # Record the outputs that this node generates @@ -550,17 +541,18 @@ def add_node(self, node): def __add__(self, other): if isinstance(other, Node): return self.add_node(other) - elif isinstance(other, Workflow): + if isinstance(other, Workflow): return self.add_workflow(other) - else: - raise TypeError('Cannot add type %s to this workflow' % type(other)) + raise TypeError("Cannot add type %s to this workflow" % type(other)) def traverse_workflow_io(self): - """ If input is needed from another workflow within a larger + """ + If input is needed from another workflow within a larger hierarchical workflow, determine the path for the file to reach the destination and add the file to workflows input / output as needed. """ + def root_path(v): path = [v] while v.in_workflow: @@ -580,29 +572,34 @@ def root_path(v): # to a workflow that contains the job which needs it. for idx in range(input_root.index(common)): child_wflow = input_root[idx] - parent_wflow = input_root[idx+1] + parent_wflow = input_root[idx + 1] if inp not in child_wflow._as_job.get_outputs(): child_wflow._as_job.add_outputs(inp, stage_out=True) parent_wflow._outputs += [inp] # Set out needed file so it gets staged downwards towards the # job that needs it. - for wf in workflow_root[:workflow_root.index(common)]: + for wf in workflow_root[: workflow_root.index(common)]: if inp not in wf._as_job.get_inputs(): wf._as_job.add_inputs(inp) for wf in self.sub_workflows: wf.traverse_workflow_io() - def save(self, filename=None, submit_now=False, plan_now=False, - output_map_path=None, root=True): - """ Write this workflow to DAX file and plan/submit it if necessary - """ + def save( + self, + filename=None, + submit_now=False, + plan_now=False, + output_map_path=None, + root=True, + ): + """Write this workflow to DAX file and plan/submit it if necessary""" if filename is None: filename = self.filename if output_map_path is None: - output_map_path = 'output.map' + output_map_path = "output.map" # Handle setting up io for inter-workflow file use ahead of time # so that when daxes are saved the metadata is complete @@ -616,7 +613,7 @@ def save(self, filename=None, submit_now=False, plan_now=False, sub.output_map_file.insert_into_dax(self._rc, self.sites) sub_workflow_file = File(sub.filename) pfn = os.path.join(os.getcwd(), sub.filename) - sub_workflow_file.add_pfn(pfn, site='local') + sub_workflow_file.add_pfn(pfn, site="local") sub_workflow_file.insert_into_dax(self._rc, self.sites) # add workflow input files pfns for local site to dax @@ -628,10 +625,10 @@ def save(self, filename=None, submit_now=False, plan_now=False, # Add TC into workflow self._adag.add_transformation_catalog(self._tc) - with open(output_map_path, 'w') as f: + with open(output_map_path, "w") as f: for out in self._outputs: try: - f.write(out.output_map_str() + '\n') + f.write(out.output_map_str() + "\n") except ValueError: # There was no storage path pass @@ -646,53 +643,50 @@ def save(self, filename=None, submit_now=False, plan_now=False, os.chdir(olddir) def plan_and_submit(self, submit_now=True): - """ Plan and optionally submit the workflow now. - """ - + """Plan and optionally submit the workflow now.""" planner_args = {} - planner_args['submit'] = submit_now + planner_args["submit"] = submit_now # Get properties file - would be nice to add extra properties here. - prop_file = os.path.join(PEGASUS_FILE_DIRECTORY, - 'pegasus-properties.conf') - planner_args['conf'] = prop_file + prop_file = os.path.join(PEGASUS_FILE_DIRECTORY, "pegasus-properties.conf") + planner_args["conf"] = prop_file # Cache file, if there is one if self.cache_file is not None: - planner_args['cache'] = [self.cache_file] + planner_args["cache"] = [self.cache_file] # Not really sure what this does, but Karan said to use it. Seems to # matter for subworkflows - planner_args['output_sites'] = ['local'] + planner_args["output_sites"] = ["local"] # Site options - planner_args['sites'] = self.sites - planner_args['staging_sites'] = self.staging_site + planner_args["sites"] = self.sites + planner_args["staging_sites"] = self.staging_site # Make tmpdir for submitfiles # default directory is the system default, but is overrideable # This should probably be moved to core.py? - submit_opts = 'pegasus_profile', 'pycbc|submit-directory' + submit_opts = "pegasus_profile", "pycbc|submit-directory" submit_dir = None if self.cp.has_option(*submit_opts): submit_dir = self.cp.get(*submit_opts) - submitdir = tempfile.mkdtemp(prefix='pycbc-tmp_', dir=submit_dir) + submitdir = tempfile.mkdtemp(prefix="pycbc-tmp_", dir=submit_dir) os.chmod(submitdir, 0o755) try: - os.remove('submitdir') + os.remove("submitdir") except FileNotFoundError: pass - os.symlink(submitdir, 'submitdir') - planner_args['dir'] = submitdir + os.symlink(submitdir, "submitdir") + planner_args["dir"] = submitdir # Other options - planner_args['cluster'] = ['label,horizontal'] - planner_args['relative_dir'] = 'work' - planner_args['cleanup'] = 'inplace' + planner_args["cluster"] = ["label,horizontal"] + planner_args["relative_dir"] = "work" + planner_args["cleanup"] = "inplace" # This quietens the planner a bit. We cannot set the verbosity # directly, which would be better. So be careful, if changing the # pegasus.mode property, it will change the verbosity (a lot). - planner_args['quiet'] = 1 + planner_args["quiet"] = 1 # FIXME: The location of output.map is hardcoded in the properties # file. This is overridden for subworkflows, but is not for @@ -700,44 +694,46 @@ def plan_and_submit(self, submit_now=True): # we should include the location explicitly here. # Need to set this to avoid pegasus pulling in other environment - os.environ['PEGASUS_UPDATE_PYTHONPATH'] = '0' + os.environ["PEGASUS_UPDATE_PYTHONPATH"] = "0" self._adag.plan(**planner_args) # Set up convenience scripts - with open('status', 'w') as fp: - fp.write('export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-status ') - fp.write(f'--long {submitdir}/work $@') + with open("status", "w") as fp: + fp.write("export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-status ") + fp.write(f"--long {submitdir}/work $@") - with open('debug', 'w') as fp: - fp.write('export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-analyzer -r ') - fp.write(f'-v {submitdir}/work $@') + with open("debug", "w") as fp: + fp.write("export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-analyzer -r ") + fp.write(f"-v {submitdir}/work $@") - with open('stop', 'w') as fp: - fp.write('export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-remove ') - fp.write(f'{submitdir}/work $@') + with open("stop", "w") as fp: + fp.write("export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-remove ") + fp.write(f"{submitdir}/work $@") - with open('start', 'w') as fp: - fp.write('export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-run ') - fp.write(f'{submitdir}/work $@') + with open("start", "w") as fp: + fp.write("export PEGASUS_UPDATE_PYTHONPATH=0; pegasus-run ") + fp.write(f"{submitdir}/work $@") - os.chmod('status', 0o755) - os.chmod('debug', 0o755) - os.chmod('stop', 0o755) - os.chmod('start', 0o755) + os.chmod("status", 0o755) + os.chmod("debug", 0o755) + os.chmod("stop", 0o755) + os.chmod("start", 0o755) - os.makedirs('workflow/planning', exist_ok=True) + os.makedirs("workflow/planning", exist_ok=True) - shutil.copy2(prop_file, 'workflow/planning') - shutil.copy2(os.path.join(submitdir, 'work', 'braindump.yml'), - 'workflow/planning') + shutil.copy2(prop_file, "workflow/planning") + shutil.copy2( + os.path.join(submitdir, "work", "braindump.yml"), "workflow/planning" + ) if self.cache_file is not None: - shutil.copy2(self.cache_file, 'workflow/planning') + shutil.copy2(self.cache_file, "workflow/planning") class SubWorkflow(dax.SubWorkflow): - """Workflow job representation of a SubWorkflow. + """ + Workflow job representation of a SubWorkflow. This follows the Pegasus nomenclature where there are Workflows, Jobs and SubWorkflows. Be careful though! A SubWorkflow is actually a Job, not a @@ -753,8 +749,7 @@ def __init__(self, *args, **kwargs): self.pycbc_planner_args = {} def add_into_workflow(self, container_wflow): - """Add this Job into a container Workflow - """ + """Add this Job into a container Workflow""" self.add_planner_args(**self.pycbc_planner_args) # Set this to None so code will fail if more planner args are added @@ -763,51 +758,54 @@ def add_into_workflow(self, container_wflow): def add_planner_arg(self, value, option): if self.pycbc_planner_args is None: - err_msg = ("We cannot add arguments to the SubWorkflow planning " - "stage after this is added to the parent workflow.") + err_msg = ( + "We cannot add arguments to the SubWorkflow planning " + "stage after this is added to the parent workflow." + ) raise ValueError(err_msg) self.pycbc_planner_args[value] = option - def set_subworkflow_properties(self, output_map_file, - staging_site, - cache_file): + def set_subworkflow_properties(self, output_map_file, staging_site, cache_file): - self.add_planner_arg('pegasus.dir.storage.mapper.replica.file', - os.path.basename(output_map_file.name)) + self.add_planner_arg( + "pegasus.dir.storage.mapper.replica.file", + os.path.basename(output_map_file.name), + ) # Ensure output_map_file has the for_planning flag set. There's no # API way to set this after the File is initialized, so we have to # change the attribute here. # WORSE, we only want to set this if the pegasus *planner* is version # 5.0.4 or larger try: - sproc_out = subprocess.check_output(['pegasus-version']).strip() + sproc_out = subprocess.check_output(["pegasus-version"]).strip() sproc_out = sproc_out.decode() - if version.parse(sproc_out) >= version.parse('5.0.4'): - output_map_file.for_planning=True + if version.parse(sproc_out) >= version.parse("5.0.4"): + output_map_file.for_planning = True except: logging.warning("Could not execute pegasus-version, assuming >= 5.0.4") - output_map_file.for_planning=True + output_map_file.for_planning = True self.add_inputs(output_map_file) # I think this is needed to deal with cases where the subworkflow file # does not exist at submission time. bname = os.path.splitext(os.path.basename(self.file))[0] - self.add_planner_arg('basename', bname) - self.add_planner_arg('output_sites', ['local']) - self.add_planner_arg('cleanup', 'inplace') - self.add_planner_arg('cluster', ['label', 'horizontal']) - self.add_planner_arg('verbose', 3) + self.add_planner_arg("basename", bname) + self.add_planner_arg("output_sites", ["local"]) + self.add_planner_arg("cleanup", "inplace") + self.add_planner_arg("cluster", ["label", "horizontal"]) + self.add_planner_arg("verbose", 3) if cache_file: - self.add_planner_arg('cache', [cache_file]) + self.add_planner_arg("cache", [cache_file]) if staging_site: - self.add_planner_arg('staging_sites', staging_site) + self.add_planner_arg("staging_sites", staging_site) class File(dax.File): - """ The workflow representation of a physical file + """ + The workflow representation of a physical file An object that represents a file from the perspective of setting up a workflow. The file may or may not exist at the time of workflow generation. @@ -815,6 +813,7 @@ class File(dax.File): A storage path is also available to indicate the desired final destination of this file. """ + def __init__(self, name): self.name = name self.node = None @@ -837,9 +836,8 @@ def dax_repr(self): def output_map_str(self): if self.storage_path: - return '%s %s pool="%s"' % (self.name, self.storage_path, 'local') - else: - raise ValueError('This file does not have a storage path') + return '%s %s pool="%s"' % (self.name, self.storage_path, "local") + raise ValueError("This file does not have a storage path") def add_pfn(self, url, site): """ @@ -847,17 +845,16 @@ def add_pfn(self, url, site): """ self.input_pfns.append((url, site)) - def has_pfn(self, url, site='local'): + def has_pfn(self, url, site="local"): """ Check if the url, site is already associated to this File. If site is not provided, we will assume it is 'local'. """ - return (((url, site) in self.input_pfns) - or ((url, 'all') in self.input_pfns)) + return ((url, site) in self.input_pfns) or ((url, "all") in self.input_pfns) def insert_into_dax(self, rep_cat, sites): - for (url, site) in self.input_pfns: - if site == 'all': + for url, site in self.input_pfns: + if site == "all": for curr_site in sites: rep_cat.add_replica(curr_site, self, url) else: @@ -866,18 +863,20 @@ def insert_into_dax(self, rep_cat, sites): @classmethod def from_path(cls, path): """Takes a path and returns a File object with the path as the PFN.""" - warnings.warn("The from_path method in pegasus_workflow is " - "deprecated. Please use File.from_path (for " - "output files) in core.py or resolve_url_to_file " - "in core.py (for input files) instead.", - DeprecationWarning) + warnings.warn( + "The from_path method in pegasus_workflow is " + "deprecated. Please use File.from_path (for " + "output files) in core.py or resolve_url_to_file " + "in core.py (for input files) instead.", + DeprecationWarning, + ) urlparts = urlsplit(path) - site = 'nonlocal' - if (urlparts.scheme == '' or urlparts.scheme == 'file'): + site = "nonlocal" + if urlparts.scheme == "" or urlparts.scheme == "file": if os.path.isfile(urlparts.path): path = os.path.abspath(urlparts.path) - path = urljoin('file:', pathname2url(path)) - site = 'local' + path = urljoin("file:", pathname2url(path)) + site = "local" fil = cls(os.path.basename(path)) fil.add_pfn(path, site=site) diff --git a/pycbc/workflow/plotting.py b/pycbc/workflow/plotting.py index dbf65ed0f30..462e68fa11b 100644 --- a/pycbc/workflow/plotting.py +++ b/pycbc/workflow/plotting.py @@ -27,12 +27,12 @@ """ import logging -from urllib.request import pathname2url from urllib.parse import urljoin +from urllib.request import pathname2url -from pycbc.workflow.core import File, FileList, makedir, Executable +from pycbc.workflow.core import Executable, File, FileList, makedir -logger = logging.getLogger('pycbc.workflow.plotting') +logger = logging.getLogger("pycbc.workflow.plotting") def excludestr(tags, substr): @@ -52,8 +52,8 @@ def requirestr(tags, substr): class PlotExecutable(Executable): - """ plot executable - """ + """plot executable""" + current_retention_level = Executable.FINAL_RESULT # plots and final results should get the highest priority @@ -64,55 +64,67 @@ def create_node(self, **kwargs): return node -def make_template_plot(workflow, bank_file, out_dir, bins=None, - tags=None): +def make_template_plot(workflow, bank_file, out_dir, bins=None, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'plot_bank', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--bank-file', bank_file) - - if workflow.cp.has_option_tags('workflow-coincidence', 'background-bins', tags=tags): - bins = workflow.cp.get_opt_tags('workflow-coincidence', 'background-bins', tags=tags) + node = PlotExecutable( + workflow.cp, "plot_bank", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--bank-file", bank_file) + + if workflow.cp.has_option_tags( + "workflow-coincidence", "background-bins", tags=tags + ): + bins = workflow.cp.get_opt_tags( + "workflow-coincidence", "background-bins", tags=tags + ) if bins: - node.add_opt('--background-bins', bins) + node.add_opt("--background-bins", bins) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node return node.output_files[0] -def make_range_plot(workflow, psd_files, out_dir, exclude=None, require=None, - tags=None): +def make_range_plot( + workflow, psd_files, out_dir, exclude=None, require=None, tags=None +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_range'), require) + secs = requirestr(workflow.cp.get_subsections("plot_range"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - node = PlotExecutable(workflow.cp, 'plot_range', ifos=workflow.ifos, - out_dir=out_dir, tags=[tag] + tags).create_node() - node.add_input_list_opt('--psd-files', psd_files) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node = PlotExecutable( + workflow.cp, + "plot_range", + ifos=workflow.ifos, + out_dir=out_dir, + tags=[tag] + tags, + ).create_node() + node.add_input_list_opt("--psd-files", psd_files) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node files += node.output_files return files -def make_spectrum_plot(workflow, psd_files, out_dir, tags=None, - hdf_group=None, precalc_psd_files=None): +def make_spectrum_plot( + workflow, psd_files, out_dir, tags=None, hdf_group=None, precalc_psd_files=None +): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'plot_spectrum', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--psd-files', psd_files) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node = PlotExecutable( + workflow.cp, "plot_spectrum", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--psd-files", psd_files) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") if hdf_group is not None: - node.add_opt('--hdf-group', hdf_group) + node.add_opt("--hdf-group", hdf_group) if precalc_psd_files is not None and len(precalc_psd_files) == 1: - node.add_input_list_opt('--psd-file', precalc_psd_files) + node.add_input_list_opt("--psd-file", precalc_psd_files) workflow += node return node.output_files[0] @@ -121,435 +133,564 @@ def make_spectrum_plot(workflow, psd_files, out_dir, tags=None, def make_segments_plot(workflow, seg_files, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'plot_segments', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--segment-files', seg_files) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node = PlotExecutable( + workflow.cp, "plot_segments", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--segment-files", seg_files) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node def make_gating_plot(workflow, insp_files, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'plot_gating', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--input-file', insp_files) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node = PlotExecutable( + workflow.cp, "plot_gating", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--input-file", insp_files) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node def make_throughput_plot(workflow, insp_files, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'plot_throughput', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--input-file', insp_files) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node = PlotExecutable( + workflow.cp, "plot_throughput", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--input-file", insp_files) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node return node.output_files[0] -def make_foreground_table(workflow, trig_file, bank_file, out_dir, - singles=None, extension='.html', tags=None, - hierarchical_level=None): +def make_foreground_table( + workflow, + trig_file, + bank_file, + out_dir, + singles=None, + extension=".html", + tags=None, + hierarchical_level=None, +): if hierarchical_level is not None and tags: - tags = [("HIERARCHICAL_LEVEL_{:02d}".format( - hierarchical_level))] + tags + tags = [(f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}")] + tags elif hierarchical_level is not None and not tags: - tags = ["HIERARCHICAL_LEVEL_{:02d}".format(hierarchical_level)] + tags = [f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}"] elif hierarchical_level is None and not tags: tags = [] makedir(out_dir) - exe = PlotExecutable(workflow.cp, 'page_foreground', - ifos=trig_file.ifo_list, - out_dir=out_dir, tags=tags) + exe = PlotExecutable( + workflow.cp, + "page_foreground", + ifos=trig_file.ifo_list, + out_dir=out_dir, + tags=tags, + ) node = exe.create_node() - node.add_input_opt('--bank-file', bank_file) - node.add_input_opt('--trigger-file', trig_file) + node.add_input_opt("--bank-file", bank_file) + node.add_input_opt("--trigger-file", trig_file) if hierarchical_level is not None: - node.add_opt('--use-hierarchical-level', hierarchical_level) + node.add_opt("--use-hierarchical-level", hierarchical_level) if singles is not None: - node.add_input_list_opt('--single-detector-triggers', singles) - node.new_output_file_opt(bank_file.segment, extension, '--output-file') + node.add_input_list_opt("--single-detector-triggers", singles) + node.new_output_file_opt(bank_file.segment, extension, "--output-file") workflow += node return node.output_files[0] -def make_sensitivity_plot(workflow, inj_file, out_dir, exclude=None, - require=None, tags=None): +def make_sensitivity_plot( + workflow, inj_file, out_dir, exclude=None, require=None, tags=None +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_sensitivity'), require) + secs = requirestr(workflow.cp.get_subsections("plot_sensitivity"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - node = PlotExecutable(workflow.cp, 'plot_sensitivity', ifos=workflow.ifos, - out_dir=out_dir, tags=[tag] + tags).create_node() - node.add_input_opt('--injection-file', inj_file) - node.new_output_file_opt(inj_file.segment, '.png', '--output-file') + node = PlotExecutable( + workflow.cp, + "plot_sensitivity", + ifos=workflow.ifos, + out_dir=out_dir, + tags=[tag] + tags, + ).create_node() + node.add_input_opt("--injection-file", inj_file) + node.new_output_file_opt(inj_file.segment, ".png", "--output-file") workflow += node files += node.output_files return files -def make_coinc_snrchi_plot(workflow, inj_file, inj_trig, stat_file, trig_file, - out_dir, exclude=None, require=None, tags=None): +def make_coinc_snrchi_plot( + workflow, + inj_file, + inj_trig, + stat_file, + trig_file, + out_dir, + exclude=None, + require=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_coinc_snrchi'), require) + secs = requirestr(workflow.cp.get_subsections("plot_coinc_snrchi"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - exe = PlotExecutable(workflow.cp, 'plot_coinc_snrchi', - ifos=inj_trig.ifo_list, - out_dir=out_dir, tags=[tag] + tags) + exe = PlotExecutable( + workflow.cp, + "plot_coinc_snrchi", + ifos=inj_trig.ifo_list, + out_dir=out_dir, + tags=[tag] + tags, + ) node = exe.create_node() - node.add_input_opt('--found-injection-file', inj_file) - node.add_input_opt('--single-injection-file', inj_trig) - node.add_input_opt('--coinc-statistic-file', stat_file) - node.add_input_opt('--single-trigger-file', trig_file) - node.new_output_file_opt(inj_file.segment, '.png', '--output-file') + node.add_input_opt("--found-injection-file", inj_file) + node.add_input_opt("--single-injection-file", inj_trig) + node.add_input_opt("--coinc-statistic-file", stat_file) + node.add_input_opt("--single-trigger-file", trig_file) + node.new_output_file_opt(inj_file.segment, ".png", "--output-file") workflow += node files += node.output_files return files -def make_inj_table(workflow, inj_file, out_dir, missed=False, singles=None, - tags=None): +def make_inj_table(workflow, inj_file, out_dir, missed=False, singles=None, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'page_injections', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() + node = PlotExecutable( + workflow.cp, "page_injections", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() - node.add_input_opt('--injection-file', inj_file) + node.add_input_opt("--injection-file", inj_file) if missed: - node.add_opt('--show-missed') + node.add_opt("--show-missed") if singles is not None: - node.add_multiifo_input_list_opt('--single-trigger-files', singles) + node.add_multiifo_input_list_opt("--single-trigger-files", singles) - node.new_output_file_opt(inj_file.segment, '.html', '--output-file') + node.new_output_file_opt(inj_file.segment, ".html", "--output-file") workflow += node return node.output_files[0] -def make_seg_table(workflow, seg_files, seg_names, out_dir, tags=None, - title_text=None, description=None): - """ Creates a node in the workflow for writing the segment summary +def make_seg_table( + workflow, + seg_files, + seg_names, + out_dir, + tags=None, + title_text=None, + description=None, +): + """ + Creates a node in the workflow for writing the segment summary table. Returns a File instances for the output file. """ seg_files = list(seg_files) seg_names = list(seg_names) - if tags is None: tags = [] + if tags is None: + tags = [] makedir(out_dir) - node = PlotExecutable(workflow.cp, 'page_segtable', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--segment-files', seg_files) + node = PlotExecutable( + workflow.cp, "page_segtable", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--segment-files", seg_files) quoted_seg_names = [] for s in seg_names: quoted_seg_names.append("'" + s + "'") - node.add_opt('--segment-names', ' '.join(quoted_seg_names)) - node.add_opt('--ifos', ' '.join(workflow.ifos)) + node.add_opt("--segment-names", " ".join(quoted_seg_names)) + node.add_opt("--ifos", " ".join(workflow.ifos)) if description: - node.add_opt('--description', "'" + description + "'") + node.add_opt("--description", "'" + description + "'") if title_text: - node.add_opt('--title-text', "'" + title_text + "'") - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node.add_opt("--title-text", "'" + title_text + "'") + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node return node.output_files[0] def make_veto_table(workflow, out_dir, vetodef_file=None, tags=None): - """ Creates a node in the workflow for writing the veto_definer + """ + Creates a node in the workflow for writing the veto_definer table. Returns a File instances for the output file. """ if vetodef_file is None: - if not workflow.cp.has_option_tags("workflow-segments", - "segments-veto-definer-file", []): + if not workflow.cp.has_option_tags( + "workflow-segments", "segments-veto-definer-file", [] + ): return None - vetodef_file = workflow.cp.get_opt_tags("workflow-segments", - "segments-veto-definer-file", []) - file_url = urljoin('file:', pathname2url(vetodef_file)) - vdf_file = File(workflow.ifos, 'VETO_DEFINER', - workflow.analysis_time, file_url=file_url) - vdf_file.add_pfn(file_url, site='local') + vetodef_file = workflow.cp.get_opt_tags( + "workflow-segments", "segments-veto-definer-file", [] + ) + file_url = urljoin("file:", pathname2url(vetodef_file)) + vdf_file = File( + workflow.ifos, "VETO_DEFINER", workflow.analysis_time, file_url=file_url + ) + vdf_file.add_pfn(file_url, site="local") else: vdf_file = vetodef_file - if tags is None: tags = [] + if tags is None: + tags = [] makedir(out_dir) - node = PlotExecutable(workflow.cp, 'page_vetotable', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--veto-definer-file', vdf_file) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node = PlotExecutable( + workflow.cp, "page_vetotable", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--veto-definer-file", vdf_file) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node return node.output_files[0] def make_seg_plot(workflow, seg_files, out_dir, seg_names=None, tags=None): - """ Creates a node in the workflow for plotting science, and veto segments. - """ + """Creates a node in the workflow for plotting science, and veto segments.""" seg_files = list(seg_files) - if tags is None: tags = [] + if tags is None: + tags = [] makedir(out_dir) - node = PlotExecutable(workflow.cp, 'page_segplot', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--segment-files', seg_files) + node = PlotExecutable( + workflow.cp, "page_segplot", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--segment-files", seg_files) quoted_seg_names = [] for s in seg_names: quoted_seg_names.append("'" + s + "'") - node.add_opt('--segment-names', ' '.join(quoted_seg_names)) - node.add_opt('--ifos', ' '.join(workflow.ifos)) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node.add_opt("--segment-names", " ".join(quoted_seg_names)) + node.add_opt("--ifos", " ".join(workflow.ifos)) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow += node return node.output_files[0] -def make_ifar_plot(workflow, trigger_file, out_dir, tags=None, - hierarchical_level=None, executable='page_ifar'): - """ Creates a node in the workflow for plotting cumulative histogram +def make_ifar_plot( + workflow, + trigger_file, + out_dir, + tags=None, + hierarchical_level=None, + executable="page_ifar", +): + """ + Creates a node in the workflow for plotting cumulative histogram of IFAR values. """ - if hierarchical_level is not None and tags: - tags = [("HIERARCHICAL_LEVEL_{:02d}".format( - hierarchical_level))] + tags + tags = [(f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}")] + tags elif hierarchical_level is not None and not tags: - tags = ["HIERARCHICAL_LEVEL_{:02d}".format(hierarchical_level)] + tags = [f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}"] elif hierarchical_level is None and not tags: tags = [] makedir(out_dir) - exe = PlotExecutable(workflow.cp, executable, ifos=trigger_file.ifo_list, - out_dir=out_dir, tags=tags) + exe = PlotExecutable( + workflow.cp, executable, ifos=trigger_file.ifo_list, out_dir=out_dir, tags=tags + ) node = exe.create_node() - node.add_input_opt('--trigger-file', trigger_file) + node.add_input_opt("--trigger-file", trigger_file) if hierarchical_level is not None: - node.add_opt('--use-hierarchical-level', hierarchical_level) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.add_opt("--use-hierarchical-level", hierarchical_level) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node return node.output_files[0] -def make_farstat_plot(workflow, trigger_files, out_dir, tags=None, - hierarchical_level=None, require=None, executable='page_farstat'): - """ Creates a node in the workflow for plotting cumulative histogram + +def make_farstat_plot( + workflow, + trigger_files, + out_dir, + tags=None, + hierarchical_level=None, + require=None, + executable="page_farstat", +): + """ + Creates a node in the workflow for plotting cumulative histogram of IFAR vs stat values. """ - makedir(out_dir) opt = list(trigger_files.values()) - ifo_combos = ' '.join(sorted(trigger_files.keys(), key=lambda x: (len(x), x))) + ifo_combos = " ".join(sorted(trigger_files.keys(), key=lambda x: (len(x), x))) - exe = PlotExecutable(workflow.cp, executable, ifos=workflow.ifos, - out_dir=out_dir, tags=tags) + exe = PlotExecutable( + workflow.cp, executable, ifos=workflow.ifos, out_dir=out_dir, tags=tags + ) node = exe.create_node() - node.add_input_list_opt('--trigger-files', opt) - node.add_opt('--ifo-combos', ifo_combos) - node.new_output_file_opt(workflow.analysis_time, '.png', '--output-file') + node.add_input_list_opt("--trigger-files", opt) + node.add_opt("--ifo-combos", ifo_combos) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output-file") workflow += node return node.output_files[0] -def make_snrchi_plot(workflow, trig_files, veto_file, veto_name, - out_dir, exclude=None, require=None, tags=None): + +def make_snrchi_plot( + workflow, + trig_files, + veto_file, + veto_name, + out_dir, + exclude=None, + require=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_snrchi'), require) + secs = requirestr(workflow.cp.get_subsections("plot_snrchi"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: for trig_file in trig_files: - exe = PlotExecutable(workflow.cp, 'plot_snrchi', - ifos=trig_file.ifo_list, - out_dir=out_dir, - tags=[tag] + tags) + exe = PlotExecutable( + workflow.cp, + "plot_snrchi", + ifos=trig_file.ifo_list, + out_dir=out_dir, + tags=[tag] + tags, + ) node = exe.create_node() node.set_memory(15000) - node.add_input_opt('--trigger-file', trig_file) + node.add_input_opt("--trigger-file", trig_file) if veto_file is not None: - node.add_input_opt('--veto-file', veto_file) - node.add_opt('--segment-name', veto_name) - node.new_output_file_opt(trig_file.segment, '.png', '--output-file') + node.add_input_opt("--veto-file", veto_file) + node.add_opt("--segment-name", veto_name) + node.new_output_file_opt(trig_file.segment, ".png", "--output-file") workflow += node files += node.output_files return files -def make_foundmissed_plot(workflow, inj_file, out_dir, exclude=None, - require=None, tags=None): +def make_foundmissed_plot( + workflow, inj_file, out_dir, exclude=None, require=None, tags=None +): if tags is None: tags = [] makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_foundmissed'), require) + secs = requirestr(workflow.cp.get_subsections("plot_foundmissed"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - exe = PlotExecutable(workflow.cp, 'plot_foundmissed', ifos=workflow.ifos, - out_dir=out_dir, tags=[tag] + tags) + exe = PlotExecutable( + workflow.cp, + "plot_foundmissed", + ifos=workflow.ifos, + out_dir=out_dir, + tags=[tag] + tags, + ) node = exe.create_node() - ext = '.html' if exe.has_opt('dynamic') else '.png' - node.add_input_opt('--injection-file', inj_file) - node.new_output_file_opt(inj_file.segment, ext, '--output-file') + ext = ".html" if exe.has_opt("dynamic") else ".png" + node.add_input_opt("--injection-file", inj_file) + node.new_output_file_opt(inj_file.segment, ext, "--output-file") workflow += node files += node.output_files return files -def make_snrratehist_plot(workflow, bg_file, out_dir, closed_box=False, - tags=None, hierarchical_level=None): +def make_snrratehist_plot( + workflow, bg_file, out_dir, closed_box=False, tags=None, hierarchical_level=None +): if hierarchical_level is not None and tags: - tags = [("HIERARCHICAL_LEVEL_{:02d}".format( - hierarchical_level))] + tags + tags = [(f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}")] + tags elif hierarchical_level is not None and not tags: - tags = ["HIERARCHICAL_LEVEL_{:02d}".format(hierarchical_level)] + tags = [f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}"] elif hierarchical_level is None and not tags: tags = [] makedir(out_dir) - exe = PlotExecutable(workflow.cp, 'plot_snrratehist', - ifos=bg_file.ifo_list, - out_dir=out_dir, tags=tags) + exe = PlotExecutable( + workflow.cp, + "plot_snrratehist", + ifos=bg_file.ifo_list, + out_dir=out_dir, + tags=tags, + ) node = exe.create_node() - node.add_input_opt('--trigger-file', bg_file) + node.add_input_opt("--trigger-file", bg_file) if hierarchical_level is not None: - node.add_opt('--use-hierarchical-level', hierarchical_level) + node.add_opt("--use-hierarchical-level", hierarchical_level) if closed_box: - node.add_opt('--closed-box') + node.add_opt("--closed-box") - node.new_output_file_opt(bg_file.segment, '.png', '--output-file') + node.new_output_file_opt(bg_file.segment, ".png", "--output-file") workflow += node return node.output_files[0] -def make_snrifar_plot(workflow, bg_file, out_dir, closed_box=False, - cumulative=True, tags=None, hierarchical_level=None): +def make_snrifar_plot( + workflow, + bg_file, + out_dir, + closed_box=False, + cumulative=True, + tags=None, + hierarchical_level=None, +): if hierarchical_level is not None and tags: - tags = [("HIERARCHICAL_LEVEL_{:02d}".format( - hierarchical_level))] + tags + tags = [(f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}")] + tags elif hierarchical_level is not None and not tags: - tags = ["HIERARCHICAL_LEVEL_{:02d}".format(hierarchical_level)] + tags = [f"HIERARCHICAL_LEVEL_{hierarchical_level:02d}"] elif hierarchical_level is None and not tags: tags = [] makedir(out_dir) - exe = PlotExecutable(workflow.cp, 'plot_snrifar', ifos=bg_file.ifo_list, - out_dir=out_dir, tags=tags) + exe = PlotExecutable( + workflow.cp, "plot_snrifar", ifos=bg_file.ifo_list, out_dir=out_dir, tags=tags + ) node = exe.create_node() - node.add_input_opt('--trigger-file', bg_file) + node.add_input_opt("--trigger-file", bg_file) if hierarchical_level is not None: - node.add_opt('--use-hierarchical-level', hierarchical_level) + node.add_opt("--use-hierarchical-level", hierarchical_level) if closed_box: - node.add_opt('--closed-box') + node.add_opt("--closed-box") if not cumulative: - node.add_opt('--not-cumulative') + node.add_opt("--not-cumulative") - node.new_output_file_opt(bg_file.segment, '.png', '--output-file') + node.new_output_file_opt(bg_file.segment, ".png", "--output-file") workflow += node return node.output_files[0] -def make_results_web_page(workflow, results_dir, template='orange', - explicit_dependencies=None): - template_path = 'templates/'+template+'.html' +def make_results_web_page( + workflow, results_dir, template="orange", explicit_dependencies=None +): + template_path = "templates/" + template + ".html" - out_dir = workflow.cp.get('results_page', 'output-path') + out_dir = workflow.cp.get("results_page", "output-path") makedir(out_dir) - node = PlotExecutable(workflow.cp, 'results_page', ifos=workflow.ifos, - out_dir=out_dir).create_node() - node.add_opt('--plots-dir', results_dir) - node.add_opt('--template-file', template_path) + node = PlotExecutable( + workflow.cp, "results_page", ifos=workflow.ifos, out_dir=out_dir + ).create_node() + node.add_opt("--plots-dir", results_dir) + node.add_opt("--template-file", template_path) workflow += node if explicit_dependencies is not None: for dep in explicit_dependencies: workflow.add_explicit_dependancy(dep, node) -def make_single_hist(workflow, trig_file, veto_file, veto_name, - out_dir, bank_file=None, exclude=None, - require=None, tags=None): +def make_single_hist( + workflow, + trig_file, + veto_file, + veto_name, + out_dir, + bank_file=None, + exclude=None, + require=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_hist'), require) + secs = requirestr(workflow.cp.get_subsections("plot_hist"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - node = PlotExecutable(workflow.cp, 'plot_hist', - ifos=trig_file.ifo, - out_dir=out_dir, - tags=[tag] + tags).create_node() + node = PlotExecutable( + workflow.cp, + "plot_hist", + ifos=trig_file.ifo, + out_dir=out_dir, + tags=[tag] + tags, + ).create_node() if veto_file is not None: - node.add_opt('--segment-name', veto_name) - node.add_input_opt('--veto-file', veto_file) - node.add_input_opt('--trigger-file', trig_file) + node.add_opt("--segment-name", veto_name) + node.add_input_opt("--veto-file", veto_file) + node.add_input_opt("--trigger-file", trig_file) if bank_file: - node.add_input_opt('--bank-file', bank_file) - node.new_output_file_opt(trig_file.segment, '.png', '--output-file') + node.add_input_opt("--bank-file", bank_file) + node.new_output_file_opt(trig_file.segment, ".png", "--output-file") workflow += node files += node.output_files return files -def make_binned_hist(workflow, trig_file, veto_file, veto_name, - out_dir, bank_file, exclude=None, - require=None, tags=None): +def make_binned_hist( + workflow, + trig_file, + veto_file, + veto_name, + out_dir, + bank_file, + exclude=None, + require=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_binnedhist'), require) + secs = requirestr(workflow.cp.get_subsections("plot_binnedhist"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: - node = PlotExecutable(workflow.cp, 'plot_binnedhist', - ifos=trig_file.ifo, - out_dir=out_dir, - tags=[tag] + tags).create_node() - node.add_opt('--ifo', trig_file.ifo) + node = PlotExecutable( + workflow.cp, + "plot_binnedhist", + ifos=trig_file.ifo, + out_dir=out_dir, + tags=[tag] + tags, + ).create_node() + node.add_opt("--ifo", trig_file.ifo) if veto_file is not None: - node.add_opt('--veto-segment-name', veto_name) - node.add_input_opt('--veto-file', veto_file) - node.add_input_opt('--trigger-file', trig_file) - node.add_input_opt('--bank-file', bank_file) - node.new_output_file_opt(trig_file.segment, '.png', '--output-file') + node.add_opt("--veto-segment-name", veto_name) + node.add_input_opt("--veto-file", veto_file) + node.add_input_opt("--trigger-file", trig_file) + node.add_input_opt("--bank-file", bank_file) + node.new_output_file_opt(trig_file.segment, ".png", "--output-file") workflow += node files += node.output_files return files -def make_singles_plot(workflow, trig_files, bank_file, veto_file, veto_name, - out_dir, exclude=None, require=None, tags=None): +def make_singles_plot( + workflow, + trig_files, + bank_file, + veto_file, + veto_name, + out_dir, + exclude=None, + require=None, + tags=None, +): tags = [] if tags is None else tags makedir(out_dir) - secs = requirestr(workflow.cp.get_subsections('plot_singles'), require) + secs = requirestr(workflow.cp.get_subsections("plot_singles"), require) secs = excludestr(secs, exclude) secs = excludestr(secs, workflow.ifo_combinations) files = FileList([]) for tag in secs: for trig_file in trig_files: - node = PlotExecutable(workflow.cp, 'plot_singles', - ifos=trig_file.ifo, - out_dir=out_dir, - tags=[tag] + tags).create_node() + node = PlotExecutable( + workflow.cp, + "plot_singles", + ifos=trig_file.ifo, + out_dir=out_dir, + tags=[tag] + tags, + ).create_node() node.set_memory(15000) - node.add_input_opt('--bank-file', bank_file) + node.add_input_opt("--bank-file", bank_file) if veto_file is not None: - node.add_input_opt('--veto-file', veto_file) - node.add_opt('--segment-name', veto_name) - node.add_opt('--detector', trig_file.ifo) - node.add_input_opt('--single-trig-file', trig_file) - node.new_output_file_opt(trig_file.segment, '.png', '--output-file') + node.add_input_opt("--veto-file", veto_file) + node.add_opt("--segment-name", veto_name) + node.add_opt("--detector", trig_file.ifo) + node.add_input_opt("--single-trig-file", trig_file) + node.new_output_file_opt(trig_file.segment, ".png", "--output-file") workflow += node files += node.output_files return files @@ -558,13 +699,17 @@ def make_singles_plot(workflow, trig_files, bank_file, veto_file, veto_name, def make_dq_flag_trigger_rate_plot(workflow, dq_file, dq_label, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'plot_dq_flag_likelihood', - ifos=dq_file.ifo, out_dir=out_dir, - tags=tags).create_node() - node.add_input_opt('--dq-file', dq_file) - node.add_opt('--dq-label', dq_label) - node.add_opt('--ifo', dq_file.ifo) - node.new_output_file_opt(dq_file.segment, '.png', '--output-file') + node = PlotExecutable( + workflow.cp, + "plot_dq_flag_likelihood", + ifos=dq_file.ifo, + out_dir=out_dir, + tags=tags, + ).create_node() + node.add_input_opt("--dq-file", dq_file) + node.add_opt("--dq-label", dq_label) + node.add_opt("--ifo", dq_file.ifo) + node.new_output_file_opt(dq_file.segment, ".png", "--output-file") workflow += node return node.output_files[0] @@ -572,11 +717,12 @@ def make_dq_flag_trigger_rate_plot(workflow, dq_file, dq_label, out_dir, tags=No def make_dq_segment_table(workflow, dq_file, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'page_dq_table', ifos=dq_file.ifo, - out_dir=out_dir, tags=tags).create_node() - node.add_input_opt('--dq-file', dq_file) - node.add_opt('--ifo', dq_file.ifo) - node.new_output_file_opt(dq_file.segment, '.html', '--output-file') + node = PlotExecutable( + workflow.cp, "page_dq_table", ifos=dq_file.ifo, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_opt("--dq-file", dq_file) + node.add_opt("--ifo", dq_file.ifo) + node.new_output_file_opt(dq_file.segment, ".html", "--output-file") workflow += node return node.output_files[0] @@ -584,12 +730,16 @@ def make_dq_segment_table(workflow, dq_file, out_dir, tags=None): def make_template_bin_table(workflow, dq_file, out_dir, tags=None): tags = [] if tags is None else tags makedir(out_dir) - node = PlotExecutable(workflow.cp, 'page_template_bin_table', - ifos=dq_file.ifo, out_dir=out_dir, - tags=tags).create_node() - node.add_input_opt('--dq-file', dq_file) - node.add_opt('--ifo', dq_file.ifo) - node.new_output_file_opt(dq_file.segment, '.html', '--output-file') + node = PlotExecutable( + workflow.cp, + "page_template_bin_table", + ifos=dq_file.ifo, + out_dir=out_dir, + tags=tags, + ).create_node() + node.add_input_opt("--dq-file", dq_file) + node.add_opt("--ifo", dq_file.ifo) + node.new_output_file_opt(dq_file.segment, ".html", "--output-file") workflow += node return node.output_files[0] @@ -605,18 +755,11 @@ def make_bank_compression_plots(workflow, bank_files, out_dir, tags=None): "plot_bank_compression", ifos=workflow.ifos, out_dir=out_dir, - tags=[tag] + tags + tags=[tag] + tags, ).create_node() - node.add_input_list_opt( - "--bank-files", - bank_files - ) + node.add_input_list_opt("--bank-files", bank_files) - node.new_output_file_opt( - workflow.analysis_time, - '.png', - '--output' - ) + node.new_output_file_opt(workflow.analysis_time, ".png", "--output") workflow += node files += node.output_files return files diff --git a/pycbc/workflow/psd.py b/pycbc/workflow/psd.py index 5a0e1dfa198..68c0947593a 100644 --- a/pycbc/workflow/psd.py +++ b/pycbc/workflow/psd.py @@ -14,52 +14,56 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This module is responsible for setting up PSD-related jobs in workflows. -""" +"""This module is responsible for setting up PSD-related jobs in workflows.""" import logging from igwn_segments import segmentlist -from pycbc.workflow.core import FileList, make_analysis_dir, Executable -from pycbc.workflow.core import SegFile +from pycbc.workflow.core import Executable, FileList, SegFile, make_analysis_dir -logger = logging.getLogger('pycbc.workflow.psd') +logger = logging.getLogger("pycbc.workflow.psd") class CalcPSDExecutable(Executable): current_retention_level = Executable.ALL_TRIGGERS + class MergePSDFiles(Executable): current_retention_level = Executable.MERGED_TRIGGERS + def chunks(l, n): - """ Yield n successive chunks from l. - """ + """Yield n successive chunks from l.""" newn = int(len(l) / n) - for i in range(0, n-1): - yield l[i*newn:i*newn+newn] - yield l[n*newn-newn:] + for i in range(n - 1): + yield l[i * newn : i * newn + newn] + yield l[n * newn - newn :] + def merge_psds(workflow, files, ifo, out_dir, tags=None): make_analysis_dir(out_dir) tags = [] if not tags else tags - node = MergePSDFiles(workflow.cp, 'merge_psds', - ifos=ifo, out_dir=out_dir, - tags=tags).create_node() - node.add_input_list_opt('--psd-files', files) - node.new_output_file_opt(workflow.analysis_time, '.hdf', '--output-file') + node = MergePSDFiles( + workflow.cp, "merge_psds", ifos=ifo, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--psd-files", files) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node return node.output_files[0] -def setup_psd_calculate(workflow, frame_files, ifo, segments, - segment_name, out_dir, tags=None): + +def setup_psd_calculate( + workflow, frame_files, ifo, segments, segment_name, out_dir, tags=None +): make_analysis_dir(out_dir) tags = [] if not tags else tags - if workflow.cp.has_option_tags('workflow-psd', 'parallelization-factor', tags=tags): - num_parts = int(workflow.cp.get_opt_tags('workflow-psd', - 'parallelization-factor', - tags=tags)) + if workflow.cp.has_option_tags("workflow-psd", "parallelization-factor", tags=tags): + num_parts = int( + workflow.cp.get_opt_tags( + "workflow-psd", "parallelization-factor", tags=tags + ) + ) else: num_parts = 1 @@ -70,55 +74,74 @@ def setup_psd_calculate(workflow, frame_files, ifo, segments, psd_files = FileList([]) for i, segs in enumerate(segment_lists): - seg_file = SegFile.from_segment_list('%s_%s' %(segment_name, i), - segmentlist(segs), segment_name, ifo, - valid_segment=workflow.analysis_time, - extension='xml', directory=out_dir) - - psd_files += [make_psd_file(workflow, frame_files, seg_file, - segment_name, out_dir, - tags=tags + ['PART%s' % i])] + seg_file = SegFile.from_segment_list( + "%s_%s" % (segment_name, i), + segmentlist(segs), + segment_name, + ifo, + valid_segment=workflow.analysis_time, + extension="xml", + directory=out_dir, + ) + + psd_files += [ + make_psd_file( + workflow, + frame_files, + seg_file, + segment_name, + out_dir, + tags=tags + ["PART%s" % i], + ) + ] return merge_psds(workflow, psd_files, ifo, out_dir, tags=tags) -def make_psd_file(workflow, frame_files, segment_file, segment_name, out_dir, - tags=None): + +def make_psd_file( + workflow, frame_files, segment_file, segment_name, out_dir, tags=None +): make_analysis_dir(out_dir) tags = [] if not tags else tags - exe = CalcPSDExecutable(workflow.cp, 'calculate_psd', - ifos=segment_file.ifo, out_dir=out_dir, - tags=tags) + exe = CalcPSDExecutable( + workflow.cp, "calculate_psd", ifos=segment_file.ifo, out_dir=out_dir, tags=tags + ) node = exe.create_node() - node.add_input_opt('--analysis-segment-file', segment_file) - node.add_opt('--segment-name', segment_name) + node.add_input_opt("--analysis-segment-file", segment_file) + node.add_opt("--segment-name", segment_name) - if frame_files and not exe.has_opt('frame-type'): - node.add_input_list_opt('--frame-files', frame_files) + if frame_files and not exe.has_opt("frame-type"): + node.add_input_list_opt("--frame-files", frame_files) - node.new_output_file_opt(workflow.analysis_time, '.hdf', '--output-file') + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node return node.output_files[0] + class AvgPSDExecutable(Executable): current_retention_level = Executable.FINAL_RESULT -def make_average_psd(workflow, psd_files, out_dir, tags=None, - output_fmt='.txt'): + +def make_average_psd(workflow, psd_files, out_dir, tags=None, output_fmt=".txt"): make_analysis_dir(out_dir) tags = [] if tags is None else tags - node = AvgPSDExecutable(workflow.cp, 'average_psd', ifos=workflow.ifos, - out_dir=out_dir, tags=tags).create_node() - node.add_input_list_opt('--input-files', psd_files) + node = AvgPSDExecutable( + workflow.cp, "average_psd", ifos=workflow.ifos, out_dir=out_dir, tags=tags + ).create_node() + node.add_input_list_opt("--input-files", psd_files) if len(workflow.ifos) > 1: - node.new_output_file_opt(workflow.analysis_time, output_fmt, - '--detector-avg-file') + node.new_output_file_opt( + workflow.analysis_time, output_fmt, "--detector-avg-file" + ) - node.new_multiifo_output_list_opt('--time-avg-file', workflow.ifos, - workflow.analysis_time, output_fmt, tags=tags) + node.new_multiifo_output_list_opt( + "--time-avg-file", workflow.ifos, workflow.analysis_time, output_fmt, tags=tags + ) workflow += node return node.output_files + # keep namespace clean -__all__ = ['make_psd_file', 'make_average_psd', 'setup_psd_calculate', 'merge_psds'] +__all__ = ["make_average_psd", "make_psd_file", "merge_psds", "setup_psd_calculate"] diff --git a/pycbc/workflow/psdfiles.py b/pycbc/workflow/psdfiles.py index aacb1f0c1be..33c5f56b125 100644 --- a/pycbc/workflow/psdfiles.py +++ b/pycbc/workflow/psdfiles.py @@ -29,17 +29,18 @@ # FIXME: Is this module still relevant for any code? Can it be removed? -import logging import configparser as ConfigParser +import logging + +from pycbc.workflow.core import FileList, make_analysis_dir, resolve_url_to_file -from pycbc.workflow.core import FileList -from pycbc.workflow.core import make_analysis_dir, resolve_url_to_file +logger = logging.getLogger("pycbc.workflow.psdfiles") -logger = logging.getLogger('pycbc.workflow.psdfiles') -def setup_psd_workflow(workflow, science_segs, datafind_outs, - output_dir=None, tags=None): - ''' +def setup_psd_workflow( + workflow, science_segs, datafind_outs, output_dir=None, tags=None +): + """ Setup static psd section of CBC workflow. At present this only supports pregenerated psd files, in the future these could be created within the workflow. @@ -59,10 +60,11 @@ def setup_psd_workflow(workflow, science_segs, datafind_outs, that would be produced in multiple calls to this function. Returns - -------- + ------- psd_files : pycbc.workflow.core.FileList The FileList holding the psd files, 0 or 1 per ifo - ''' + + """ if tags is None: tags = [] logger.info("Entering static psd module.") @@ -71,8 +73,7 @@ def setup_psd_workflow(workflow, science_segs, datafind_outs, # Parse for options in ini file. try: - psdMethod = cp.get_opt_tags("workflow-psd", "psd-method", - tags) + psdMethod = cp.get_opt_tags("workflow-psd", "psd-method", tags) except: # Predefined PSD sare optional, just return an empty list if not # provided. @@ -91,7 +92,7 @@ def setup_psd_workflow(workflow, science_segs, datafind_outs, def setup_psd_pregenerated(workflow, tags=None): - ''' + """ Setup CBC workflow to use pregenerated psd files. The file given in cp.get('workflow','pregenerated-psd-file-(ifo)') will be used as the --psd-file argument to geom_nonspinbank, geom_aligned_bank @@ -106,33 +107,33 @@ def setup_psd_pregenerated(workflow, tags=None): that would be produced in multiple calls to this function. Returns - -------- + ------- psd_files : pycbc.workflow.core.FileList The FileList holding the gating files - ''' + + """ if tags is None: tags = [] psd_files = FileList([]) cp = workflow.cp global_seg = workflow.analysis_time - file_attrs = {'segs': global_seg, 'tags': tags} + file_attrs = {"segs": global_seg, "tags": tags} # Check for one psd for all ifos try: - pre_gen_file = cp.get_opt_tags('workflow-psd', - 'psd-pregenerated-file', tags) - file_attrs['ifos'] = workflow.ifos + pre_gen_file = cp.get_opt_tags("workflow-psd", "psd-pregenerated-file", tags) + file_attrs["ifos"] = workflow.ifos curr_file = resolve_url_to_file(pre_gen_file, attrs=file_attrs) psd_files.append(curr_file) except ConfigParser.Error: # Check for one psd per ifo for ifo in workflow.ifos: try: - pre_gen_file = cp.get_opt_tags('workflow-psd', - 'psd-pregenerated-file-%s' % ifo.lower(), - tags) - file_attrs['ifos'] = [ifo] + pre_gen_file = cp.get_opt_tags( + "workflow-psd", "psd-pregenerated-file-%s" % ifo.lower(), tags + ) + file_attrs["ifos"] = [ifo] curr_file = resolve_url_to_file(pre_gen_file, attrs=file_attrs) psd_files.append(curr_file) @@ -140,7 +141,5 @@ def setup_psd_pregenerated(workflow, tags=None): # It's unlikely, but not impossible, that only some ifos # will have pregenerated PSDs logger.warning("No psd file specified for IFO %s.", ifo) - pass return psd_files - diff --git a/pycbc/workflow/segment.py b/pycbc/workflow/segment.py index cbb3bd9a35f..b1deb23cf06 100644 --- a/pycbc/workflow/segment.py +++ b/pycbc/workflow/segment.py @@ -27,40 +27,42 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/ahope/segments.html """ -import os -import shutil import itertools import logging +import os +import shutil import igwn_segments as segments from igwn_segments import utils as segmentsUtils -from pycbc.workflow.core import SegFile, make_analysis_dir -from pycbc.workflow.core import resolve_url +from pycbc.workflow.core import SegFile, make_analysis_dir, resolve_url + +logger = logging.getLogger("pycbc.workflow.segment") -logger = logging.getLogger('pycbc.workflow.segment') def save_veto_definer(cp, out_dir, tags=None): - """ Retrieve the veto definer file and save it locally + """ + Retrieve the veto definer file and save it locally Parameters - ----------- + ---------- cp : ConfigParser instance out_dir : path tags : list of strings Used to retrieve subsections of the ini file for configuration options. + """ if tags is None: tags = [] make_analysis_dir(out_dir) - veto_def_url = cp.get_opt_tags("workflow-segments", - "segments-veto-definer-url", tags) + veto_def_url = cp.get_opt_tags( + "workflow-segments", "segments-veto-definer-url", tags + ) veto_def_base_name = os.path.basename(veto_def_url) - veto_def_new_path = os.path.abspath(os.path.join(out_dir, - veto_def_base_name)) + veto_def_new_path = os.path.abspath(os.path.join(out_dir, veto_def_base_name)) # Don't need to do this if already done - resolve_url(veto_def_url,out_dir) + resolve_url(veto_def_url, out_dir) # and update location cp.set("workflow-segments", "segments-veto-definer-file", veto_def_new_path) @@ -68,7 +70,8 @@ def save_veto_definer(cp, out_dir, tags=None): def get_segments_file(workflow, name, option_name, out_dir, tags=None): - """Get cumulative segments from option name syntax for each ifo. + """ + Get cumulative segments from option name syntax for each ifo. Use syntax of configparser string to define the resulting segment_file e.x. option_name = +up_flag1,+up_flag2,+up_flag3,-down_flag1,-down_flag2 @@ -86,12 +89,14 @@ def get_segments_file(workflow, name, option_name, out_dir, tags=None): Used to retrieve subsections of the ini file for configuration options. - returns - -------- + Returns + ------- seg_file: pycbc.workflow.SegFile SegFile intance that points to the segment xml file on disk. + """ from pycbc.dq import query_str + make_analysis_dir(out_dir) cp = workflow.cp start = workflow.analysis_time[0] @@ -105,15 +110,14 @@ def get_segments_file(workflow, name, option_name, out_dir, tags=None): if cp.has_option("workflow-segments", "segments-veto-definer-url"): veto_definer = save_veto_definer(workflow.cp, out_dir) - cache = cp.has_option_tags("workflow-segments", 'enable-query-caching', tags) + cache = cp.has_option_tags("workflow-segments", "enable-query-caching", tags) if cache: - logger.info('Caching queries enabled') + logger.info("Caching queries enabled") # Check for provided server server = "https://segments.ligo.org" if cp.has_option_tags("workflow-segments", "segments-database-url", tags): - server = cp.get_opt_tags("workflow-segments", - "segments-database-url", tags) + server = cp.get_opt_tags("workflow-segments", "segments-database-url", tags) if cp.has_option_tags("workflow-segments", "segments-source", tags): source = cp.get_opt_tags("workflow-segments", "segments-source", tags) @@ -121,9 +125,9 @@ def get_segments_file(workflow, name, option_name, out_dir, tags=None): source = "any" if source == "file": - local_file_path = \ - resolve_url(cp.get_opt_tags("workflow-segments", - option_name+"-file", tags)) + local_file_path = resolve_url( + cp.get_opt_tags("workflow-segments", option_name + "-file", tags) + ) pfn = os.path.join(out_dir, os.path.basename(local_file_path)) shutil.move(local_file_path, pfn) return SegFile.from_segment_xml(pfn) @@ -131,7 +135,7 @@ def get_segments_file(workflow, name, option_name, out_dir, tags=None): segs = {} for ifo in workflow.ifos: flag_str = cp.get_opt_tags("workflow-segments", option_name, [ifo]) - key = ifo + ':' + name + key = ifo + ":" + name if flag_str.upper() == "OFF": segs[key] = segments.segmentlist([]) @@ -139,15 +143,25 @@ def get_segments_file(workflow, name, option_name, out_dir, tags=None): all_seg = segments.segment([start, end]) segs[key] = segments.segmentlist([all_seg]) else: - segs[key] = query_str(ifo, flag_str, start, end, - source=source, server=server, - veto_definer=veto_definer, cache=cache) + segs[key] = query_str( + ifo, + flag_str, + start, + end, + source=source, + server=server, + veto_definer=veto_definer, + cache=cache, + ) logger.info("%s: got %s flags", ifo, option_name) - return SegFile.from_segment_list_dict(name, segs, - extension='.xml', - valid_segment=workflow.analysis_time, - directory=out_dir) + return SegFile.from_segment_list_dict( + name, + segs, + extension=".xml", + valid_segment=workflow.analysis_time, + directory=out_dir, + ) def get_triggered_coherent_segment(workflow, sciencesegs): @@ -157,47 +171,43 @@ def get_triggered_coherent_segment(workflow, sciencesegs): are insufficient for a search. Parameters - ----------- + ---------- workflow : pycbc.workflow.core.Workflow The workflow instance that the calculated segments belong to. sciencesegs : dict Dictionary of all science segments within analysis time. Returns - -------- + ------- onsource : igwn_segments.segmentlistdict or None A dictionary containing the on source segments for network IFOs, or None if no segments are available that meet the requirements. offsource : igwn_segments.segmentlistdict A dictionary containing the off source segments for network IFOs - """ + """ # Load parsed workflow config options cp = workflow.cp - triggertime = int(os.path.basename(cp.get('workflow', 'trigger-time'))) - minduration = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'min-duration'))) - maxduration = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'max-duration'))) - onbefore = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'on-before'))) - onafter = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'on-after'))) - padding = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'pad-data'))) + triggertime = int(os.path.basename(cp.get("workflow", "trigger-time"))) + minduration = int( + os.path.basename(cp.get("workflow-exttrig_segments", "min-duration")) + ) + maxduration = int( + os.path.basename(cp.get("workflow-exttrig_segments", "max-duration")) + ) + onbefore = int(os.path.basename(cp.get("workflow-exttrig_segments", "on-before"))) + onafter = int(os.path.basename(cp.get("workflow-exttrig_segments", "on-after"))) + padding = int(os.path.basename(cp.get("workflow-exttrig_segments", "pad-data"))) if cp.has_option("workflow-condition_strain", "do-gating"): - padding += int(os.path.basename(cp.get("condition_strain", - "pad-data"))) - quanta = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'quanta'))) + padding += int(os.path.basename(cp.get("condition_strain", "pad-data"))) + quanta = int(os.path.basename(cp.get("workflow-exttrig_segments", "quanta"))) # Check available data segments meet criteria specified in arguments commonsegs = sciencesegs.extract_common(sciencesegs.keys()) offsrclist = commonsegs[tuple(commonsegs.keys())[0]] if len(offsrclist) > 1: - logger.info("Removing network segments that do not contain trigger " - "time") + logger.info("Removing network segments that do not contain trigger time") for seg in offsrclist: if triggertime in seg: offsrc = seg @@ -208,10 +218,15 @@ def get_triggered_coherent_segment(workflow, sciencesegs): offsrc = segments.segment(int(offsrc[0]), int(offsrc[1])) if abs(offsrc) < minduration + 2 * padding: - fail = segments.segment([triggertime - minduration / 2. - padding, - triggertime + minduration / 2. + padding]) - logger.warning("Available network segment shorter than minimum " - "allowed duration.") + fail = segments.segment( + [ + triggertime - minduration / 2.0 - padding, + triggertime + minduration / 2.0 + padding, + ] + ) + logger.warning( + "Available network segment shorter than minimum allowed duration." + ) offsource = segments.segmentlistdict() for iifo in sciencesegs: offsource[iifo] = segments.segmentlist([fail]) @@ -219,71 +234,80 @@ def get_triggered_coherent_segment(workflow, sciencesegs): # Will segment duration be the maximum desired length or not? if abs(offsrc) >= maxduration + 2 * padding: - logger.info("Available network science segment duration (%ds) is " - "greater than the maximum allowed segment length (%ds). " - "Truncating...", abs(offsrc), maxduration) + logger.info( + "Available network science segment duration (%ds) is " + "greater than the maximum allowed segment length (%ds). " + "Truncating...", + abs(offsrc), + maxduration, + ) else: - logger.info("Available network science segment duration (%ds) is " - "less than the maximum allowed segment length (%ds).", - abs(offsrc), maxduration) + logger.info( + "Available network science segment duration (%ds) is " + "less than the maximum allowed segment length (%ds).", + abs(offsrc), + maxduration, + ) - logger.info("%ds of padding applied at beginning and end of segment.", - padding) + logger.info("%ds of padding applied at beginning and end of segment.", padding) # Construct on-source onstart = triggertime - onbefore onend = triggertime + onafter oncentre = onstart + ((onbefore + onafter) / 2) onsrc = segments.segment(onstart, onend) - logger.info("Constructed ON-SOURCE: duration %ds (%ds before to %ds after" - " trigger).", abs(onsrc), triggertime - onsrc[0], - onsrc[1] - triggertime) + logger.info( + "Constructed ON-SOURCE: duration %ds (%ds before to %ds after trigger).", + abs(onsrc), + triggertime - onsrc[0], + onsrc[1] - triggertime, + ) onsrc = segments.segmentlist([onsrc]) # Maximal, centred coherent network segment - idealsegment = segments.segment(int(oncentre - padding - - 0.5 * maxduration), - int(oncentre + padding + - 0.5 * maxduration)) + idealsegment = segments.segment( + int(oncentre - padding - 0.5 * maxduration), + int(oncentre + padding + 0.5 * maxduration), + ) # Construct off-source if idealsegment in offsrc: offsrc = idealsegment elif idealsegment[1] not in offsrc: - offsrc &= segments.segment(offsrc[1] - maxduration - 2 * padding, - offsrc[1]) + offsrc &= segments.segment(offsrc[1] - maxduration - 2 * padding, offsrc[1]) elif idealsegment[0] not in offsrc: - offsrc &= segments.segment(offsrc[0], - offsrc[0] + maxduration + 2 * padding) + offsrc &= segments.segment(offsrc[0], offsrc[0] + maxduration + 2 * padding) # Trimming off-source excess = (abs(offsrc) - 2 * padding) % quanta if excess != 0: - logger.info("Trimming %ds excess time to make OFF-SOURCE duration a " - "multiple of %ds", excess, quanta) - offset = (offsrc[0] + abs(offsrc) / 2.) - oncentre + logger.info( + "Trimming %ds excess time to make OFF-SOURCE duration a multiple of %ds", + excess, + quanta, + ) + offset = (offsrc[0] + abs(offsrc) / 2.0) - oncentre if 2 * abs(offset) > excess: if offset < 0: - offsrc &= segments.segment(offsrc[0] + excess, - offsrc[1]) + offsrc &= segments.segment(offsrc[0] + excess, offsrc[1]) elif offset > 0: - offsrc &= segments.segment(offsrc[0], - offsrc[1] - excess) + offsrc &= segments.segment(offsrc[0], offsrc[1] - excess) assert abs(offsrc) % quanta == 2 * padding else: - logger.info("This will make OFF-SOURCE symmetrical about trigger " - "time.") + logger.info("This will make OFF-SOURCE symmetrical about trigger time.") start = int(offsrc[0] - offset + excess / 2) end = int(offsrc[1] - offset - round(float(excess) / 2)) offsrc = segments.segment(start, end) assert abs(offsrc) % quanta == 2 * padding - logger.info("Constructed OFF-SOURCE: duration %ds (%ds before to %ds " - "after trigger).", abs(offsrc) - 2 * padding, - triggertime - offsrc[0] - padding, - offsrc[1] - triggertime - padding) + logger.info( + "Constructed OFF-SOURCE: duration %ds (%ds before to %ds after trigger).", + abs(offsrc) - 2 * padding, + triggertime - offsrc[0] - padding, + offsrc[1] - triggertime - padding, + ) offsrc = segments.segmentlist([offsrc]) # Put segments into segmentlistdicts @@ -304,29 +328,30 @@ def generate_triggered_segment(workflow, out_dir, sciencesegs): else: min_ifos = 2 - triggertime = int(os.path.basename(cp.get('workflow', 'trigger-time'))) - minbefore = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'min-before'))) - minafter = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'min-after'))) - minduration = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'min-duration'))) - onbefore = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'on-before'))) - onafter = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'on-after'))) - padding = int(os.path.basename(cp.get('workflow-exttrig_segments', - 'pad-data'))) + triggertime = int(os.path.basename(cp.get("workflow", "trigger-time"))) + minbefore = int(os.path.basename(cp.get("workflow-exttrig_segments", "min-before"))) + minafter = int(os.path.basename(cp.get("workflow-exttrig_segments", "min-after"))) + minduration = int( + os.path.basename(cp.get("workflow-exttrig_segments", "min-duration")) + ) + onbefore = int(os.path.basename(cp.get("workflow-exttrig_segments", "on-before"))) + onafter = int(os.path.basename(cp.get("workflow-exttrig_segments", "on-after"))) + padding = int(os.path.basename(cp.get("workflow-exttrig_segments", "pad-data"))) if cp.has_option("workflow-condition_strain", "do-gating"): - padding += int(os.path.basename(cp.get("condition_strain", - "pad-data"))) + padding += int(os.path.basename(cp.get("condition_strain", "pad-data"))) # How many IFOs meet minimum data requirements? - min_seg = segments.segment(triggertime - onbefore - minbefore - padding, - triggertime + onafter + minafter + padding) - scisegs = segments.segmentlistdict({ifo: sciencesegs[ifo] - for ifo in sciencesegs if min_seg in sciencesegs[ifo] - and abs(sciencesegs[ifo]) >= minduration}) + min_seg = segments.segment( + triggertime - onbefore - minbefore - padding, + triggertime + onafter + minafter + padding, + ) + scisegs = segments.segmentlistdict( + { + ifo: sciencesegs[ifo] + for ifo in sciencesegs + if min_seg in sciencesegs[ifo] and abs(sciencesegs[ifo]) >= minduration + } + ) # Find highest number of IFOs that give an acceptable coherent segment num_ifos = len(scisegs) while num_ifos >= min_ifos: @@ -337,25 +362,24 @@ def generate_triggered_segment(workflow, out_dir, sciencesegs): for ifo_combo in ifo_combos: ifos = "".join(ifo_combo) logger.info("Calculating optimal segment for %s.", ifos) - segs = segments.segmentlistdict({ifo: scisegs[ifo] - for ifo in ifo_combo}) + segs = segments.segmentlistdict({ifo: scisegs[ifo] for ifo in ifo_combo}) onsource[ifos], offsource[ifos] = get_triggered_coherent_segment( - workflow, segs) + workflow, segs + ) # Which combination gives the longest coherent segment? - valid_combs = [iifos for iifos in onsource - if onsource[iifos] is not None] + valid_combs = [iifos for iifos in onsource if onsource[iifos] is not None] if len(valid_combs) == 0: # If none, offsource dict will contain segments showing criteria # that have not been met, for use in plotting seg_lens = { - ifos: abs(next(iter(offsource[ifos].values()))[0]) - for ifos in offsource + ifos: abs(next(iter(offsource[ifos].values()))[0]) for ifos in offsource } best_comb = max(seg_lens, key=seg_lens.get) - logger.info("No combination of %d IFOs with suitable science " - "segment.", num_ifos) + logger.info( + "No combination of %d IFOs with suitable science segment.", num_ifos + ) else: # Identify best analysis segment seg_lens = { @@ -364,29 +388,30 @@ def generate_triggered_segment(workflow, out_dir, sciencesegs): } best_comb = max(seg_lens, key=seg_lens.get) logger.info( - "Calculated science segments, best combination is %s", - best_comb + "Calculated science segments, best combination is %s", best_comb ) offsourceSegfile = os.path.join(out_dir, "offSourceSeg.txt") - segmentsUtils.tosegwizard(open(offsourceSegfile, "w"), - list(offsource[best_comb].values())[0]) + segmentsUtils.tosegwizard( + open(offsourceSegfile, "w"), list(offsource[best_comb].values())[0] + ) onsourceSegfile = os.path.join(out_dir, "onSourceSeg.txt") - segmentsUtils.tosegwizard(open(onsourceSegfile, "w"), - list(onsource[best_comb].values())[0]) + segmentsUtils.tosegwizard( + open(onsourceSegfile, "w"), list(onsource[best_comb].values())[0] + ) - bufferleft = int(cp.get('workflow-exttrig_segments', - 'num-buffer-before')) - bufferright = int(cp.get('workflow-exttrig_segments', - 'num-buffer-after')) + bufferleft = int(cp.get("workflow-exttrig_segments", "num-buffer-before")) + bufferright = int(cp.get("workflow-exttrig_segments", "num-buffer-after")) onlen = onbefore + onafter bufferSegment = segments.segment( - triggertime - onbefore - bufferleft * onlen, - triggertime + onafter + bufferright * onlen) + triggertime - onbefore - bufferleft * onlen, + triggertime + onafter + bufferright * onlen, + ) bufferSegfile = os.path.join(out_dir, "bufferSeg.txt") - segmentsUtils.tosegwizard(open(bufferSegfile, "w"), - segments.segmentlist([bufferSegment])) + segmentsUtils.tosegwizard( + open(bufferSegfile, "w"), segments.segmentlist([bufferSegment]) + ) return onsource[best_comb], offsource[best_comb], bufferSegment @@ -397,14 +422,15 @@ def generate_triggered_segment(workflow, out_dir, sciencesegs): return None, offsource[best_comb], None except UnboundLocalError: # Catches the case where the while loop above is never taken. - min_seg = segments.segmentlistdict({ - ifo: segments.segmentlist([min_seg]) - for ifo in sciencesegs - }) + min_seg = segments.segmentlistdict( + {ifo: segments.segmentlist([min_seg]) for ifo in sciencesegs} + ) return None, min_seg, None + def get_flag_segments_file(workflow, name, option_name, out_dir, tags=None): - """Get segments from option name syntax for each ifo for indivudal flags. + """ + Get segments from option name syntax for each ifo for indivudal flags. Use syntax of configparser string to define the resulting segment_file e.x. option_name = +up_flag1,+up_flag2,+up_flag3,-down_flag1,-down_flag2 @@ -423,12 +449,14 @@ def get_flag_segments_file(workflow, name, option_name, out_dir, tags=None): Used to retrieve subsections of the ini file for configuration options. - returns - -------- + Returns + ------- seg_file: pycbc.workflow.SegFile SegFile intance that points to the segment xml file on disk. + """ from pycbc.dq import query_str + make_analysis_dir(out_dir) cp = workflow.cp start = workflow.analysis_time[0] @@ -442,23 +470,22 @@ def get_flag_segments_file(workflow, name, option_name, out_dir, tags=None): if cp.has_option("workflow-segments", "segments-veto-definer-url"): veto_definer = save_veto_definer(workflow.cp, out_dir) - cache = cp.has_option_tags("workflow-segments", 'enable-query-caching', tags) + cache = cp.has_option_tags("workflow-segments", "enable-query-caching", tags) if cache: - logger.info('Caching query') + logger.info("Caching query") # Check for provided server server = "https://segments.ligo.org" if cp.has_option_tags("workflow-segments", "segments-database-url", tags): - server = cp.get_opt_tags("workflow-segments", - "segments-database-url", tags) + server = cp.get_opt_tags("workflow-segments", "segments-database-url", tags) source = "any" if cp.has_option_tags("workflow-segments", "segments-source", tags): source = cp.get_opt_tags("workflow-segments", "segments-source", tags) if source == "file": - local_file_path = \ - resolve_url(cp.get_opt_tags("workflow-segments", - option_name+"-file", tags)) + local_file_path = resolve_url( + cp.get_opt_tags("workflow-segments", option_name + "-file", tags) + ) pfn = os.path.join(out_dir, os.path.basename(local_file_path)) shutil.move(local_file_path, pfn) return SegFile.from_segment_xml(pfn) @@ -467,20 +494,30 @@ def get_flag_segments_file(workflow, name, option_name, out_dir, tags=None): for ifo in workflow.ifos: if cp.has_option_tags("workflow-segments", option_name, [ifo]): flag_str = cp.get_opt_tags("workflow-segments", option_name, [ifo]) - flag_list = flag_str.split(',') + flag_list = flag_str.split(",") for flag in flag_list: flag_name = flag[1:] - if len(flag_name.split(':')) > 1: - flag_name = name.split(':')[1] - key = ifo + ':' + flag_name - segs[key] = query_str(ifo, flag, start, end, - source=source, server=server, - veto_definer=veto_definer, cache=cache) + if len(flag_name.split(":")) > 1: + flag_name = name.split(":")[1] + key = ifo + ":" + flag_name + segs[key] = query_str( + ifo, + flag, + start, + end, + source=source, + server=server, + veto_definer=veto_definer, + cache=cache, + ) logger.info("%s: got %s segments", ifo, flag_name) else: logger.info("%s: no segments requested", ifo) - return SegFile.from_segment_list_dict(name, segs, - extension='.xml', - valid_segment=workflow.analysis_time, - directory=out_dir) + return SegFile.from_segment_list_dict( + name, + segs, + extension=".xml", + valid_segment=workflow.analysis_time, + directory=out_dir, + ) diff --git a/pycbc/workflow/splittable.py b/pycbc/workflow/splittable.py index 7473fcac091..309975119c3 100644 --- a/pycbc/workflow/splittable.py +++ b/pycbc/workflow/splittable.py @@ -28,20 +28,23 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/NOTYETCREATED.html """ -import os -import logging import glob +import logging import math - -from pycbc.workflow.core import FileList, make_analysis_dir, File -from pycbc.workflow.jobsetup import (PycbcSplitBankExecutable, - PycbcSplitBankXmlExecutable, PycbcSplitInspinjExecutable, - PycbcHDFSplitInjExecutable) - +import os from urllib.parse import urljoin from urllib.request import pathname2url -logger = logging.getLogger('pycbc.workflow.splittable') +from pycbc.workflow.core import File, FileList, make_analysis_dir +from pycbc.workflow.jobsetup import ( + PycbcHDFSplitInjExecutable, + PycbcSplitBankExecutable, + PycbcSplitBankXmlExecutable, + PycbcSplitInspinjExecutable, +) + +logger = logging.getLogger("pycbc.workflow.splittable") + def select_splitfilejob_instance(curr_exe): """ @@ -56,31 +59,33 @@ def select_splitfilejob_instance(curr_exe): The name of the section storing options for this executble Returns - -------- + ------- exe class : sub-class of pycbc.workflow.core.Executable The class that holds the utility functions appropriate for the given Executable. This class **must** contain * exe_class.create_job() and the job returned by this **must** contain * job.create_node() + """ - if curr_exe == 'pycbc_hdf5_splitbank': + if curr_exe == "pycbc_hdf5_splitbank": exe_class = PycbcSplitBankExecutable - elif curr_exe == 'pycbc_splitbank': + elif curr_exe == "pycbc_splitbank": exe_class = PycbcSplitBankXmlExecutable - elif curr_exe == 'pycbc_split_inspinj': + elif curr_exe == "pycbc_split_inspinj": exe_class = PycbcSplitInspinjExecutable - elif curr_exe == 'pycbc_hdf_splitinj': + elif curr_exe == "pycbc_hdf_splitinj": exe_class = PycbcHDFSplitInjExecutable else: # Should we try some sort of default class?? - err_string = "No class exists for Executable %s" %(curr_exe,) + err_string = "No class exists for Executable %s" % (curr_exe,) raise NotImplementedError(err_string) return exe_class + def setup_splittable_workflow(workflow, input_tables, out_dir=None, tags=None): - ''' + """ This function aims to be the gateway for code that is responsible for taking some input file containing some table, and splitting into multiple files containing different parts of that table. For now the only supported operation @@ -88,7 +93,7 @@ def setup_splittable_workflow(workflow, input_tables, out_dir=None, tags=None): template bank xml files. Parameters - ----------- + ---------- workflow : pycbc.workflow.core.Workflow The Workflow instance that the jobs will be added to. input_tables : pycbc.workflow.core.FileList @@ -97,23 +102,26 @@ def setup_splittable_workflow(workflow, input_tables, out_dir=None, tags=None): The directory in which output will be written. Returns - -------- + ------- split_table_outs : pycbc.workflow.core.FileList The list of split up files as output from this job. - ''' + + """ if tags is None: tags = [] logger.info("Entering split output files module.") make_analysis_dir(out_dir) # Parse for options in .ini file - splitMethod = workflow.cp.get_opt_tags("workflow-splittable", - "splittable-method", tags) + splitMethod = workflow.cp.get_opt_tags( + "workflow-splittable", "splittable-method", tags + ) if splitMethod == "IN_WORKFLOW": # Scope here for choosing different options logger.info("Adding split output file jobs to workflow.") - split_table_outs = setup_splittable_dax_generated(workflow, - input_tables, out_dir, tags) + split_table_outs = setup_splittable_dax_generated( + workflow, input_tables, out_dir, tags + ) elif splitMethod == "MANUAL_DIRECTORY": logger.info("Registering pre-existing split files from directory.") split_table_outs = setup_splittable_manual_directory(workflow, tags) @@ -129,23 +137,24 @@ def setup_splittable_workflow(workflow, input_tables, out_dir=None, tags=None): logger.info("Leaving split output files module.") return split_table_outs + def setup_splittable_manual_directory(workflow, tags=None): """ - New function to glob a directory and register existing files as + New function to glob a directory and register existing files as workflow products. """ if tags is None: tags = [] cp = workflow.cp - + # Get directory from config bank_dir = cp.get_opt_tags("workflow-splittable", "tmpltbank-directory", tags) - + if not os.path.isabs(bank_dir): bank_dir = os.path.abspath(bank_dir) # Glob all HDF files - bank_paths = sorted(glob.glob(os.path.join(bank_dir, '*.hdf'))) + bank_paths = sorted(glob.glob(os.path.join(bank_dir, "*.hdf"))) if not bank_paths: raise ValueError(f"No .hdf files found in {bank_dir}") @@ -153,29 +162,30 @@ def setup_splittable_manual_directory(workflow, tags=None): tmplt_banks = FileList([]) for i, path in enumerate(bank_paths): - bank_tag = ('bank%0{}d'.format(n_dp)) % i + bank_tag = (f"bank%0{n_dp}d") % i abs_path = os.path.abspath(path) - pfn_local = urljoin('file:', pathname2url(abs_path)) - + pfn_local = urljoin("file:", pathname2url(abs_path)) + # Create a File object that Pegasus recognizes as an existing input curr_file = File( workflow.ifos, - 'TMPLTBANK', + "TMPLTBANK", workflow.analysis_time, file_url=pfn_local, - tags=tags + [bank_tag] + tags=tags + [bank_tag], ) - curr_file.add_pfn(pfn_local, site='local') + curr_file.add_pfn(pfn_local, site="local") tmplt_banks.append(curr_file) - + return tmplt_banks + def setup_splittable_dax_generated(workflow, input_tables, out_dir, tags): - ''' + """ Function for setting up the splitting jobs as part of the workflow. Parameters - ----------- + ---------- workflow : pycbc.workflow.core.Workflow The Workflow instance that the jobs will be added to. input_tables : pycbc.workflow.core.FileList @@ -184,33 +194,34 @@ def setup_splittable_dax_generated(workflow, input_tables, out_dir, tags): The directory in which output will be written. Returns - -------- + ------- split_table_outs : pycbc.workflow.core.FileList The list of split up files as output from this job. - ''' + + """ cp = workflow.cp # Get values from ini file try: - num_splits = cp.get_opt_tags("workflow-splittable", - "splittable-num-banks", tags) + num_splits = cp.get_opt_tags( + "workflow-splittable", "splittable-num-banks", tags + ) except BaseException: - inj_interval = int(cp.get_opt_tags("workflow-splittable", - "splitinjtable-interval", tags)) - if cp.has_option_tags("em_bright_filter", "max-keep", tags) and \ - cp.has_option("workflow-injections", "em-bright-only"): - num_injs = int(cp.get_opt_tags("em_bright_filter", "max-keep", - tags)) + inj_interval = int( + cp.get_opt_tags("workflow-splittable", "splitinjtable-interval", tags) + ) + if cp.has_option_tags("em_bright_filter", "max-keep", tags) and cp.has_option( + "workflow-injections", "em-bright-only" + ): + num_injs = int(cp.get_opt_tags("em_bright_filter", "max-keep", tags)) else: # This needed to be changed from num-injs to ninjections in order # to work properly with pycbc_create_injections - num_injs = int(cp.get_opt_tags("workflow-injections", - "ninjections", tags)) + num_injs = int(cp.get_opt_tags("workflow-injections", "ninjections", tags)) inj_tspace = float(abs(workflow.analysis_time)) / num_injs num_splits = int(inj_interval // inj_tspace) + 1 - split_exe_tag = cp.get_opt_tags("workflow-splittable", - "splittable-exe-tag", tags) + split_exe_tag = cp.get_opt_tags("workflow-splittable", "splittable-exe-tag", tags) split_exe = os.path.basename(cp.get("executables", split_exe_tag)) # Select the appropriate class exe_class = select_splitfilejob_instance(split_exe) @@ -219,12 +230,10 @@ def setup_splittable_dax_generated(workflow, input_tables, out_dir, tags): out_file_groups = FileList([]) # Set up the condorJob class for the current executable - curr_exe_job = exe_class(workflow.cp, split_exe_tag, num_splits, - out_dir=out_dir) + curr_exe_job = exe_class(workflow.cp, split_exe_tag, num_splits, out_dir=out_dir) for input in input_tables: node = curr_exe_job.create_node(input, tags=tags) workflow.add_node(node) out_file_groups += node.output_files return out_file_groups - diff --git a/pycbc/workflow/tmpltbank.py b/pycbc/workflow/tmpltbank.py index cb61a38337d..79cdd8c8045 100644 --- a/pycbc/workflow/tmpltbank.py +++ b/pycbc/workflow/tmpltbank.py @@ -28,24 +28,33 @@ https://ldas-jobs.ligo.caltech.edu/~cbc/docs/pycbc/ahope/template_bank.html """ - -import os +import configparser as ConfigParser import logging import math -import configparser as ConfigParser +import os import pycbc -from pycbc.workflow.core import FileList, Executable -from pycbc.workflow.core import make_analysis_dir, resolve_url_to_file +from pycbc.workflow.core import ( + Executable, + FileList, + make_analysis_dir, + resolve_url_to_file, +) from pycbc.workflow.jobsetup import select_tmpltbank_class, sngl_ifo_job_setup -logger = logging.getLogger('pycbc.workflow.tmpltbank') +logger = logging.getLogger("pycbc.workflow.tmpltbank") -def setup_tmpltbank_workflow(workflow, science_segs, datafind_outs, - output_dir=None, psd_files=None, tags=None, - return_format=None): - ''' +def setup_tmpltbank_workflow( + workflow, + science_segs, + datafind_outs, + output_dir=None, + psd_files=None, + tags=None, + return_format=None, +): + """ Setup template bank section of CBC workflow. This function is responsible for deciding which of the various template bank workflow generation utilities should be used. @@ -68,10 +77,11 @@ def setup_tmpltbank_workflow(workflow, science_segs, datafind_outs, that would be produced in multiple calls to this function. Returns - -------- + ------- tmplt_banks : pycbc.workflow.core.FileList The FileList holding the details of all the template bank jobs. - ''' + + """ if tags is None: tags = [] logger.info("Entering template bank generation module.") @@ -79,8 +89,7 @@ def setup_tmpltbank_workflow(workflow, science_segs, datafind_outs, cp = workflow.cp # Parse for options in ini file - tmpltbankMethod = cp.get_opt_tags("workflow-tmpltbank", "tmpltbank-method", - tags) + tmpltbankMethod = cp.get_opt_tags("workflow-tmpltbank", "tmpltbank-method", tags) # There can be a large number of different options here, for e.g. to set # up fixed bank, or maybe something else @@ -90,19 +99,24 @@ def setup_tmpltbank_workflow(workflow, science_segs, datafind_outs, # Else we assume template banks will be generated in the workflow elif tmpltbankMethod == "WORKFLOW_INDEPENDENT_IFOS": logger.info("Adding template bank jobs to workflow.") - tmplt_banks = setup_tmpltbank_dax_generated(workflow, science_segs, - datafind_outs, output_dir, tags=tags, - psd_files=psd_files) + tmplt_banks = setup_tmpltbank_dax_generated( + workflow, + science_segs, + datafind_outs, + output_dir, + tags=tags, + psd_files=psd_files, + ) elif tmpltbankMethod == "WORKFLOW_INDEPENDENT_IFOS_NODATA": logger.info("Adding template bank jobs to workflow.") - tmplt_banks = setup_tmpltbank_without_frames(workflow, output_dir, - tags=tags, independent_ifos=True, - psd_files=psd_files) + tmplt_banks = setup_tmpltbank_without_frames( + workflow, output_dir, tags=tags, independent_ifos=True, psd_files=psd_files + ) elif tmpltbankMethod == "WORKFLOW_NO_IFO_VARIATION_NODATA": logger.info("Adding template bank jobs to workflow.") - tmplt_banks = setup_tmpltbank_without_frames(workflow, output_dir, - tags=tags, independent_ifos=False, - psd_files=psd_files) + tmplt_banks = setup_tmpltbank_without_frames( + workflow, output_dir, tags=tags, independent_ifos=False, psd_files=psd_files + ) else: errMsg = "Template bank method not recognized. Must be either " errMsg += "PREGENERATED_BANK, WORKFLOW_INDEPENDENT_IFOS " @@ -115,28 +129,30 @@ def setup_tmpltbank_workflow(workflow, science_segs, datafind_outs, # a conversion from xml.gz or xml to hdf is supported, but not vice # versa. If a return_format is not specified the function returns # the bank in the format as it was inputted. - tmplt_bank_filename=tmplt_banks[0].name - ext = tmplt_bank_filename.split('.', 1)[1] + tmplt_bank_filename = tmplt_banks[0].name + ext = tmplt_bank_filename.split(".", 1)[1] logger.info("Input bank is a %s file", ext) - if return_format is None : + if return_format is None: tmplt_banks_return = tmplt_banks - elif return_format in ('hdf', 'h5', 'hdf5'): - if ext in ('hdf', 'h5', 'hdf5') or ext in ('xml.gz' , 'xml'): - tmplt_banks_return = pycbc.workflow.convert_bank_to_hdf(workflow, - tmplt_banks, "bank") - else : - if ext == return_format: - tmplt_banks_return = tmplt_banks - else: - raise NotImplementedError("{0} to {1} conversion is not " - "supported.".format(ext, return_format)) + elif return_format in ("hdf", "h5", "hdf5"): + if ext in ("hdf", "h5", "hdf5") or ext in ("xml.gz", "xml"): + tmplt_banks_return = pycbc.workflow.convert_bank_to_hdf( + workflow, tmplt_banks, "bank" + ) + elif ext == return_format: + tmplt_banks_return = tmplt_banks + else: + raise NotImplementedError( + f"{ext} to {return_format} conversion is not supported." + ) logger.info("Leaving template bank generation module.") return tmplt_banks_return -def setup_tmpltbank_dax_generated(workflow, science_segs, datafind_outs, - output_dir, tags=None, - psd_files=None): - ''' + +def setup_tmpltbank_dax_generated( + workflow, science_segs, datafind_outs, output_dir, tags=None, psd_files=None +): + """ Setup template bank jobs that are generated as part of the CBC workflow. This function will add numerous jobs to the CBC workflow using configuration options from the .ini file. The following executables are @@ -163,10 +179,11 @@ def setup_tmpltbank_dax_generated(workflow, science_segs, datafind_outs, The file list containing predefined PSDs, if provided. Returns - -------- + ------- tmplt_banks : pycbc.workflow.core.FileList The FileList holding the details of all the template bank jobs. - ''' + + """ if tags is None: tags = [] cp = workflow.cp @@ -175,7 +192,7 @@ def setup_tmpltbank_dax_generated(workflow, science_segs, datafind_outs, # will require a bit of effort here .... ifos = science_segs.keys() - tmplt_bank_exe = os.path.basename(cp.get('executables', 'tmpltbank')) + tmplt_bank_exe = os.path.basename(cp.get("executables", "tmpltbank")) # Select the appropriate class exe_class = select_tmpltbank_class(tmplt_bank_exe) @@ -183,24 +200,31 @@ def setup_tmpltbank_dax_generated(workflow, science_segs, datafind_outs, tmplt_banks = FileList([]) for ifo in ifos: - job_instance = exe_class(workflow.cp, 'tmpltbank', ifo=ifo, - out_dir=output_dir, - tags=tags) + job_instance = exe_class( + workflow.cp, "tmpltbank", ifo=ifo, out_dir=output_dir, tags=tags + ) # Check for the write_psd flag if cp.has_option_tags("workflow-tmpltbank", "tmpltbank-write-psd-file", tags): job_instance.write_psd = True else: job_instance.write_psd = False - sngl_ifo_job_setup(workflow, ifo, tmplt_banks, job_instance, - science_segs[ifo], datafind_outs, - allow_overlap=True) + sngl_ifo_job_setup( + workflow, + ifo, + tmplt_banks, + job_instance, + science_segs[ifo], + datafind_outs, + allow_overlap=True, + ) return tmplt_banks -def setup_tmpltbank_without_frames(workflow, output_dir, - tags=None, independent_ifos=False, - psd_files=None): - ''' + +def setup_tmpltbank_without_frames( + workflow, output_dir, tags=None, independent_ifos=False, psd_files=None +): + """ Setup CBC workflow to use a template bank (or banks) that are generated in the workflow, but do not use the data to estimate a PSD, and therefore do not vary over the duration of the workflow. This can either generate one @@ -223,10 +247,11 @@ def setup_tmpltbank_without_frames(workflow, output_dir, The file list containing predefined PSDs, if provided. Returns - -------- + ------- tmplt_banks : pycbc.workflow.core.FileList The FileList holding the details of the template bank(s). - ''' + + """ if tags is None: tags = [] cp = workflow.cp @@ -237,9 +262,9 @@ def setup_tmpltbank_without_frames(workflow, output_dir, ifos = workflow.ifos fullSegment = workflow.analysis_time - tmplt_bank_exe = os.path.basename(cp.get('executables','tmpltbank')) + tmplt_bank_exe = os.path.basename(cp.get("executables", "tmpltbank")) # Can not use lalapps_template bank with this - if tmplt_bank_exe == 'lalapps_tmpltbank': + if tmplt_bank_exe == "lalapps_tmpltbank": errMsg = "Lalapps_tmpltbank cannot be used to generate template banks " errMsg += "without using frames. Try another code." raise ValueError(errMsg) @@ -262,18 +287,23 @@ def setup_tmpltbank_without_frames(workflow, output_dir, exe_instance.write_psd = False for ifo in ifoList: - job_instance = exe_instance(workflow.cp, 'tmpltbank', ifo=ifo, - out_dir=output_dir, - tags=tags, - psd_files=psd_files) + job_instance = exe_instance( + workflow.cp, + "tmpltbank", + ifo=ifo, + out_dir=output_dir, + tags=tags, + psd_files=psd_files, + ) node = job_instance.create_nodata_node(fullSegment) workflow.add_node(node) tmplt_banks += node.output_files return tmplt_banks + def setup_tmpltbank_pregenerated(workflow, tags=None): - ''' + """ Setup CBC workflow to use a pregenerated template bank. The bank given in cp.get('workflow','pregenerated-template-bank') will be used as the input file for all matched-filtering jobs. If this option is @@ -289,10 +319,11 @@ def setup_tmpltbank_pregenerated(workflow, tags=None): that would be produced in multiple calls to this function. Returns - -------- + ------- tmplt_banks : pycbc.workflow.core.FileList The FileList holding the details of the template bank. - ''' + + """ if tags is None: tags = [] # Currently this uses the *same* fixed bank for all ifos. @@ -303,23 +334,26 @@ def setup_tmpltbank_pregenerated(workflow, tags=None): cp = workflow.cp global_seg = workflow.analysis_time - file_attrs = {'segs' : global_seg, 'tags' : tags} + file_attrs = {"segs": global_seg, "tags": tags} try: # First check if we have a bank for all ifos - pre_gen_bank = cp.get_opt_tags('workflow-tmpltbank', - 'tmpltbank-pregenerated-bank', tags) - file_attrs['ifos'] = workflow.ifos + pre_gen_bank = cp.get_opt_tags( + "workflow-tmpltbank", "tmpltbank-pregenerated-bank", tags + ) + file_attrs["ifos"] = workflow.ifos curr_file = resolve_url_to_file(pre_gen_bank, attrs=file_attrs) tmplt_banks.append(curr_file) except ConfigParser.Error: # Okay then I must have banks for each ifo for ifo in workflow.ifos: try: - pre_gen_bank = cp.get_opt_tags('workflow-tmpltbank', - 'tmpltbank-pregenerated-bank-%s' % ifo.lower(), - tags) - file_attrs['ifos'] = [ifo] + pre_gen_bank = cp.get_opt_tags( + "workflow-tmpltbank", + "tmpltbank-pregenerated-bank-%s" % ifo.lower(), + tags, + ) + file_attrs["ifos"] = [ifo] curr_file = resolve_url_to_file(pre_gen_bank, attrs=file_attrs) tmplt_banks.append(curr_file) @@ -327,52 +361,42 @@ def setup_tmpltbank_pregenerated(workflow, tags=None): err_msg = "Cannot find pregerated template bank in section " err_msg += "[workflow-tmpltbank] or any tagged sections. " if tags: - tagged_secs = " ".join("[workflow-tmpltbank-%s]" \ - %(ifo,) for ifo in workflow.ifos) - err_msg += "Tagged sections are %s. " %(tagged_secs,) + tagged_secs = " ".join( + "[workflow-tmpltbank-%s]" % (ifo,) for ifo in workflow.ifos + ) + err_msg += "Tagged sections are %s. " % (tagged_secs,) err_msg += "I looked for 'tmpltbank-pregenerated-bank' option " - err_msg += "and 'tmpltbank-pregenerated-bank-%s'." %(ifo,) + err_msg += "and 'tmpltbank-pregenerated-bank-%s'." % (ifo,) raise ConfigParser.Error(err_msg) return tmplt_banks -def make_compress_split_banks(workflow, bank_files, out_dir, - tags=None): +def make_compress_split_banks(workflow, bank_files, out_dir, tags=None): tags = [] if tags is None else tags compressed_banks = FileList([]) n_dp = math.ceil(math.log10(len(bank_files))) for i, bank_file in enumerate(bank_files): node = Executable( workflow.cp, - 'compress', + "compress", ifos=workflow.ifos, out_dir=out_dir, - tags=tags + [f'bank%0{n_dp}d' % i] + tags=tags + [f"bank%0{n_dp}d" % i], ).create_node() - node.add_input_opt('--bank-file', bank_file) - node.new_output_file_opt( - workflow.analysis_time, - '.hdf', - '--output' - ) + node.add_input_opt("--bank-file", bank_file) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output") workflow += node compressed_banks += node.output_files return compressed_banks -def make_combine_split_banks(workflow, bank_files, out_dir, - tags=None): + +def make_combine_split_banks(workflow, bank_files, out_dir, tags=None): tags = [] if tags is None else tags - if workflow.cp.has_option_tags( - "workflow-splittable", - "recombine-num-banks", - tags - ): - n_banks_combined = int(workflow.cp.get_opt_tags( - "workflow-splittable", - "recombine-num-banks", - tags - )) + if workflow.cp.has_option_tags("workflow-splittable", "recombine-num-banks", tags): + n_banks_combined = int( + workflow.cp.get_opt_tags("workflow-splittable", "recombine-num-banks", tags) + ) else: n_banks_combined = 1 @@ -381,20 +405,16 @@ def make_combine_split_banks(workflow, bank_files, out_dir, for i in range(n_banks_combined): node = Executable( workflow.cp, - 'combine_banks', + "combine_banks", ifos=workflow.ifos, out_dir=out_dir, - tags=tags + [f'%0{n_dp}d' % i] + tags=tags + [f"%0{n_dp}d" % i], ).create_node() start = int(i / n_banks_combined * len(bank_files)) end = int((i + 1) / n_banks_combined * len(bank_files)) bank_files_subset = bank_files[start:end] - node.add_input_list_opt('--input-filenames', bank_files_subset) - node.new_output_file_opt( - workflow.analysis_time, - '.hdf', - '--output-file' - ) + node.add_input_list_opt("--input-filenames", bank_files_subset) + node.new_output_file_opt(workflow.analysis_time, ".hdf", "--output-file") workflow += node out_files += node.output_files diff --git a/pycbc/workflow/versioning.py b/pycbc/workflow/versioning.py index c0af9b88c85..7120f659bbf 100644 --- a/pycbc/workflow/versioning.py +++ b/pycbc/workflow/versioning.py @@ -25,18 +25,20 @@ Module to generate/manage the executable used for version information in workflows """ -import os + import logging +import os from pycbc.workflow.core import Executable -logger = logging.getLogger('pycbc.workflow.versioning') +logger = logging.getLogger("pycbc.workflow.versioning") class VersioningExecutable(Executable): """ Executable for getting version information """ + current_retention_level = Executable.FINAL_RESULT @@ -46,7 +48,7 @@ def make_versioning_page(workflow, config_parser, out_dir, tags=None): """ vers_exe = VersioningExecutable( workflow.cp, - 'page_versioning', + "page_versioning", out_dir=out_dir, ifos=workflow.ifos, tags=tags, @@ -54,7 +56,7 @@ def make_versioning_page(workflow, config_parser, out_dir, tags=None): node = vers_exe.create_node() config_names = [] exes = [] - for name, path in config_parser.items('executables'): + for name, path in config_parser.items("executables"): exe_to_test = os.path.basename(path) if exe_to_test in exes: # executable is already part of the list, @@ -66,9 +68,9 @@ def make_versioning_page(workflow, config_parser, out_dir, tags=None): else: config_names.append(name) exes.append(exe_to_test) - node.add_list_opt('--executables', exes) - node.add_list_opt('--executables-names', config_names) - node.new_output_file_opt(workflow.analysis_time, '.html', '--output-file') + node.add_list_opt("--executables", exes) + node.add_list_opt("--executables-names", config_names) + node.new_output_file_opt(workflow.analysis_time, ".html", "--output-file") workflow.add_node(node) return node, node.output_files From cde2d8680fb6ca3187a3e653214414eee603dadf Mon Sep 17 00:00:00 2001 From: Ian Harry Date: Mon, 27 Jul 2026 15:57:05 +0100 Subject: [PATCH 2/5] Check the bin files too --- .qlty/qlty.toml | 30 + bin/all_sky_search/pycbc_add_statmap | 792 +++++---- bin/all_sky_search/pycbc_apply_rerank | 134 +- bin/all_sky_search/pycbc_average_psd | 93 +- bin/all_sky_search/pycbc_bin_templates | 62 +- bin/all_sky_search/pycbc_bin_trigger_rates_dq | 110 +- bin/all_sky_search/pycbc_calculate_psd | 102 +- bin/all_sky_search/pycbc_coinc_findtrigs | 446 +++-- bin/all_sky_search/pycbc_coinc_hdfinjfind | 369 ++-- bin/all_sky_search/pycbc_coinc_mergetrigs | 167 +- bin/all_sky_search/pycbc_coinc_statmap | 465 ++--- bin/all_sky_search/pycbc_coinc_statmap_inj | 119 +- .../pycbc_combine_coincident_events | 90 +- bin/all_sky_search/pycbc_combine_statmap | 196 ++- .../pycbc_cut_merge_triggers_to_tmpltbank | 154 +- .../pycbc_distribute_background_bins | 66 +- bin/all_sky_search/pycbc_dtphase | 245 +-- bin/all_sky_search/pycbc_exclude_zerolag | 132 +- bin/all_sky_search/pycbc_fit_sngls_binned | 501 ++++-- .../pycbc_fit_sngls_by_template | 348 ++-- .../pycbc_fit_sngls_over_multiparam | 369 ++-- bin/all_sky_search/pycbc_fit_sngls_over_param | 240 +-- .../pycbc_fit_sngls_split_binned | 570 +++--- bin/all_sky_search/pycbc_followup_file | 111 +- bin/all_sky_search/pycbc_foreground_censor | 67 +- bin/all_sky_search/pycbc_get_loudest_params | 138 +- bin/all_sky_search/pycbc_make_bayestar_skymap | 81 +- bin/all_sky_search/pycbc_merge_psds | 58 +- bin/all_sky_search/pycbc_plot_kde_vals | 116 +- .../pycbc_prepare_xml_for_gracedb | 142 +- bin/all_sky_search/pycbc_reduce_template_bank | 47 +- bin/all_sky_search/pycbc_rerank_passthrough | 50 +- bin/all_sky_search/pycbc_sngls_findtrigs | 230 ++- bin/all_sky_search/pycbc_sngls_pastro | 469 ++--- bin/all_sky_search/pycbc_sngls_statmap | 380 ++-- bin/all_sky_search/pycbc_sngls_statmap_inj | 154 +- bin/all_sky_search/pycbc_strip_injections | 69 +- bin/all_sky_search/pycbc_template_kde_calc | 378 ++-- bin/all_sky_search/pycbc_template_kde_max | 39 +- .../pycbc_template_recovery_hist | 127 +- .../pycbc_upload_single_event_to_gracedb | 144 +- bin/bank/pycbc_aligned_bank_cat | 100 +- bin/bank/pycbc_aligned_stoch_bank | 222 ++- bin/bank/pycbc_bank_verification | 305 ++-- bin/bank/pycbc_brute_bank | 500 ++++-- bin/bank/pycbc_coinc_bank2hdf | 106 +- bin/bank/pycbc_geom_aligned_2dstack | 422 +++-- bin/bank/pycbc_geom_aligned_bank | 238 ++- bin/bank/pycbc_geom_nonspinbank | 176 +- bin/bank/pycbc_tmpltbank_to_chi_params | 125 +- bin/hwinj/pycbc_generate_hwinj | 579 ++++--- bin/hwinj/pycbc_generate_hwinj_from_xml | 76 +- bin/hwinj/pycbc_insert_frame_hwinj | 71 +- bin/hwinj/pycbc_plot_hwinj | 48 +- bin/inference/pycbc_inference | 122 +- .../pycbc_inference_create_calibration_config | 368 ++-- bin/inference/pycbc_inference_create_fits | 98 +- bin/inference/pycbc_inference_extract_samples | 231 +-- bin/inference/pycbc_inference_model_stats | 88 +- bin/inference/pycbc_inference_monitor | 71 +- .../pycbc_inference_plot_acceptance_rate | 82 +- bin/inference/pycbc_inference_plot_acf | 99 +- bin/inference/pycbc_inference_plot_acl | 43 +- .../pycbc_inference_plot_dynesty_run | 44 +- .../pycbc_inference_plot_dynesty_traceplot | 50 +- .../pycbc_inference_plot_gelman_rubin | 90 +- bin/inference/pycbc_inference_plot_geweke | 107 +- .../pycbc_inference_plot_inj_recovery | 84 +- .../pycbc_inference_plot_mcmc_history | 156 +- bin/inference/pycbc_inference_plot_movie | 315 ++-- bin/inference/pycbc_inference_plot_posterior | 287 +-- bin/inference/pycbc_inference_plot_pp | 108 +- bin/inference/pycbc_inference_plot_prior | 102 +- bin/inference/pycbc_inference_plot_samples | 82 +- bin/inference/pycbc_inference_plot_skymap | 29 +- ...cbc_inference_plot_thermodynamic_integrand | 70 +- .../pycbc_inference_pp_table_summary | 99 +- bin/inference/pycbc_inference_savage_dickey | 245 +-- .../pycbc_inference_start_from_samples | 33 +- bin/inference/pycbc_inference_table_summary | 153 +- bin/inference/pycbc_validate_test_posterior | 46 +- bin/live/pycbc_live_collate_triggers | 142 +- bin/live/pycbc_live_collated_dq_trigger_rates | 122 +- bin/live/pycbc_live_combine_dq_trigger_rates | 79 +- ...ycbc_live_combine_single_significance_fits | 170 +- ...ive_plot_combined_single_significance_fits | 203 ++- .../pycbc_live_plot_single_significance_fits | 170 +- bin/live/pycbc_live_single_significance_fits | 226 +-- ...pycbc_live_supervise_collated_trigger_fits | 606 +++---- .../pycbc_foreground_minifollowup | 241 +-- .../pycbc_injection_minifollowup | 444 +++-- bin/minifollowups/pycbc_page_coincinfo | 247 +-- bin/minifollowups/pycbc_page_injinfo | 188 +- bin/minifollowups/pycbc_page_snglinfo | 180 +- bin/minifollowups/pycbc_plot_chigram | 90 +- .../pycbc_plot_trigger_timeseries | 135 +- bin/minifollowups/pycbc_single_template_plot | 100 +- bin/minifollowups/pycbc_sngl_minifollowup | 352 ++-- .../pycbc_upload_prep_minifollowup | 144 +- .../pycbc_banksim_plot_eff_fitting_factor | 150 +- .../pycbc_banksim_plot_fitting_factors | 180 +- bin/plotting/pycbc_banksim_table_point_injs | 128 +- bin/plotting/pycbc_create_html_snippet | 27 +- bin/plotting/pycbc_faithsim_plots | 24 +- bin/plotting/pycbc_ifar_catalog | 264 +-- bin/plotting/pycbc_mass_area_plot | 26 +- bin/plotting/pycbc_mchirp_plots | 14 +- bin/plotting/pycbc_page_banktriggerrate | 62 +- bin/plotting/pycbc_page_coinc_snrchi | 277 +-- bin/plotting/pycbc_page_dq_table | 49 +- bin/plotting/pycbc_page_fars_vs_stat | 181 +- bin/plotting/pycbc_page_foreground | 245 +-- bin/plotting/pycbc_page_foundmissed | 339 ++-- bin/plotting/pycbc_page_ifar | 359 ++-- bin/plotting/pycbc_page_injtable | 202 ++- bin/plotting/pycbc_page_recovery | 268 +-- bin/plotting/pycbc_page_segments | 53 +- bin/plotting/pycbc_page_segplot | 117 +- bin/plotting/pycbc_page_segtable | 122 +- bin/plotting/pycbc_page_sensitivity | 500 +++--- bin/plotting/pycbc_page_snrchi | 138 +- bin/plotting/pycbc_page_snrifar | 427 +++-- bin/plotting/pycbc_page_snrratehist | 282 +-- bin/plotting/pycbc_page_template_bin_table | 80 +- bin/plotting/pycbc_page_versioning | 67 +- bin/plotting/pycbc_page_vetotable | 68 +- bin/plotting/pycbc_plot_background_coincs | 69 +- bin/plotting/pycbc_plot_bank_bins | 175 +- bin/plotting/pycbc_plot_bank_compression | 209 +-- bin/plotting/pycbc_plot_bank_corner | 227 +-- bin/plotting/pycbc_plot_dq_flag_likelihood | 75 +- bin/plotting/pycbc_plot_dq_likelihood_vs_time | 66 +- bin/plotting/pycbc_plot_dq_percentiles | 87 +- bin/plotting/pycbc_plot_gate_triggers | 193 ++- bin/plotting/pycbc_plot_gating | 85 +- bin/plotting/pycbc_plot_hist | 141 +- bin/plotting/pycbc_plot_multiifo_dtphase | 262 +-- bin/plotting/pycbc_plot_psd_file | 160 +- bin/plotting/pycbc_plot_psd_timefreq | 167 +- bin/plotting/pycbc_plot_qscan | 317 ++-- bin/plotting/pycbc_plot_range | 97 +- bin/plotting/pycbc_plot_range_vs_mtot | 110 +- bin/plotting/pycbc_plot_singles_timefreq | 281 +-- bin/plotting/pycbc_plot_singles_vs_params | 182 +- bin/plotting/pycbc_plot_throughput | 75 +- bin/plotting/pycbc_plot_trigrate | 298 ++-- bin/plotting/pycbc_plot_vt_ratio | 190 +- bin/plotting/pycbc_plot_waveform | 313 ++-- bin/population/pycbc_multiifo_pastro | 530 +++--- bin/population/pycbc_population_plots | 138 +- bin/population/pycbc_population_rates | 222 ++- bin/pycbc_banksim | 490 ++++-- bin/pycbc_banksim_combine_banks | 40 +- bin/pycbc_banksim_match_combine | 197 ++- bin/pycbc_banksim_skymax | 589 ++++--- bin/pycbc_compress_bank | 301 ++-- bin/pycbc_condition_strain | 113 +- bin/pycbc_convertinjfiletohdf | 125 +- bin/pycbc_create_injections | 202 ++- bin/pycbc_data_store | 53 +- bin/pycbc_faithsim | 345 ++-- bin/pycbc_faithsim_collect_results | 8 +- bin/pycbc_fit_sngl_trigs | 442 +++-- bin/pycbc_generate_mock_data | 314 ++-- bin/pycbc_get_ffinal | 279 +-- bin/pycbc_gwosc_segment_query | 74 +- bin/pycbc_hdf5_splitbank | 90 +- bin/pycbc_hdf_splitinj | 31 +- bin/pycbc_inj_cut | 92 +- bin/pycbc_inspiral | 723 +++++--- bin/pycbc_inspiral_skymax | 612 ++++--- bin/pycbc_live | 1533 ++++++++++------- bin/pycbc_live_nagios_monitor | 79 +- bin/pycbc_make_banksim | 234 ++- bin/pycbc_make_faithsim | 144 +- bin/pycbc_make_html_page | 234 +-- bin/pycbc_make_sky_grid | 120 +- bin/pycbc_make_skymap | 807 +++++---- bin/pycbc_merge_inj_hdf | 49 +- bin/pycbc_multi_inspiral | 423 +++-- bin/pycbc_optimal_snr | 309 ++-- bin/pycbc_optimize_snr | 376 ++-- bin/pycbc_process_sngls | 142 +- bin/pycbc_randomize_inj_dist_by_optsnr | 243 +-- bin/pycbc_single_template | 507 +++--- bin/pycbc_source_probability_offline | 105 +- bin/pycbc_split_inspinj | 35 +- bin/pycbc_splitbank | 118 +- bin/pycbc_upload_xml_to_gracedb | 237 ++- bin/pygrb/pycbc_grb_inj_finder | 132 +- bin/pygrb/pycbc_grb_trig_cluster | 97 +- bin/pygrb/pycbc_grb_trig_combiner | 185 +- bin/pygrb/pycbc_make_offline_grb_workflow | 410 +++-- bin/pygrb/pycbc_pygrb_efficiency | 561 +++--- bin/pygrb/pycbc_pygrb_exclusion_dist_table | 42 +- bin/pygrb/pycbc_pygrb_grb_info_table | 139 +- bin/pygrb/pycbc_pygrb_minifollowups | 203 ++- bin/pygrb/pycbc_pygrb_page_tables | 873 ++++++---- bin/pygrb/pycbc_pygrb_plot_chisq_veto | 242 +-- bin/pygrb/pycbc_pygrb_plot_coh_ifosnr | 150 +- bin/pygrb/pycbc_pygrb_plot_injs_results | 558 +++--- bin/pygrb/pycbc_pygrb_plot_null_stats | 158 +- bin/pygrb/pycbc_pygrb_plot_skygrid | 194 ++- bin/pygrb/pycbc_pygrb_plot_snr_timeseries | 163 +- bin/pygrb/pycbc_pygrb_plot_stats_distribution | 134 +- bin/pygrb/pycbc_pygrb_results_workflow | 343 ++-- .../pycbc_combine_injection_comparisons | 324 ++-- .../pycbc_injection_set_comparison | 469 ++--- ...pycbc_plot_injections_found_both_workflows | 62 +- .../pycbc_plot_injections_missed_one_workflow | 81 +- .../pycbc_plot_vt_ratio_vs_ifar | 130 +- .../pycbc_make_bank_compression_workflow | 135 +- .../pycbc_make_bank_verifier_workflow | 355 ++-- bin/workflows/pycbc_make_faithsim_workflow | 14 +- .../pycbc_make_geom_aligned_bank_workflow | 29 +- .../pycbc_make_inference_inj_workflow | 296 ++-- .../pycbc_make_inference_plots_workflow | 149 +- bin/workflows/pycbc_make_inference_workflow | 169 +- .../pycbc_make_offline_search_workflow | 1072 +++++++----- .../pycbc_make_psd_estimation_workflow | 171 +- bin/workflows/pycbc_make_sbank_workflow | 251 +-- bin/workflows/pycbc_make_uberbank_workflow | 226 +-- 222 files changed, 27716 insertions(+), 19643 deletions(-) diff --git a/.qlty/qlty.toml b/.qlty/qlty.toml index 67a92eea48d..13ff5e8b93b 100644 --- a/.qlty/qlty.toml +++ b/.qlty/qlty.toml @@ -20,6 +20,25 @@ exclude_patterns = [ "test/**", ] +# bin/ holds ~220 Python executables with no .py extension (just a +# `#!/usr/bin/env python` shebang). Qlty's default Python file-type +# detection doesn't pick these up (unlike ruff itself, which sniffs +# shebangs fine when run directly), so they were silently skipped by +# every qlty_check run. Extending the glob list to include bin/** fixes +# that; the rest of the list matches Qlty's own documented default. +[file_types.python] +globs = [ + "*.py", + "*.pyw", + "**/DEPS", + "**/Snakefile", + "**/SConscript", + "**/wscript", + "**/SConstruct", + "**/.gclient", + "bin/**", +] + [[source]] name = "default" default = true @@ -56,3 +75,14 @@ name = "ripgrep" [plugins.definitions.ripgrep.drivers.lint-fixme] script = "rg \"FIXME|XXX\" ${target} --json --word-regexp --only-matching" + +# bin/** (above) is a path-only glob, so it also swept in a handful of +# extensionless files under bin/ that are actually bash, not Python -- +# ruff fails to parse these as Python. Skip ruff specifically for them. +[[exclude]] +plugins = ["ruff"] +file_patterns = [ + "bin/pycbc_copy_output_map", + "bin/pycbc_stageout_failed_workflow", + "bin/inference/run_pycbc_inference", +] diff --git a/bin/all_sky_search/pycbc_add_statmap b/bin/all_sky_search/pycbc_add_statmap index c67753b44a5..eb4afe1b5b2 100755 --- a/bin/all_sky_search/pycbc_add_statmap +++ b/bin/all_sky_search/pycbc_add_statmap @@ -1,68 +1,123 @@ #!/bin/env python -""" Calculate total FAR based on statistic ranking for coincidences in times +""" +Calculate total FAR based on statistic ranking for coincidences in times with more than one ifo combination available. Cluster to keep foreground coincs with the highest stat value. """ -import numpy as np, argparse, logging, pycbc, pycbc.events, pycbc.io + +import argparse +import logging + import igwn_segments as segments +import numpy as np import pycbc.version + +import pycbc import pycbc.conversions as conv +import pycbc.events +import pycbc.io from pycbc.events import significance + def get_ifo_string(fi): # Returns a string of a space-separated list of ifos from input file. # Can be deprecated soon, needed for older coinc_statmap files which # do not have 'ifos' attribute try: # input file has ifos stored as an attribute - istring = fi.attrs['ifos'] + istring = fi.attrs["ifos"] except KeyError: # Foreground group contains the time information for each ifo so # the ifos list can be reconstructed - istring = ' '.join(sorted([k for k in fi['foreground'].keys() - if 'time' in fi['foreground/%s' % k]])) + istring = " ".join( + sorted( + [ + k + for k in fi["foreground"].keys() + if "time" in fi["foreground/%s" % k] + ] + ) + ) return istring + parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--statmap-files', nargs='+', - help="List of coinc files to be combined") -parser.add_argument('--background-files', nargs='+', default=None, +parser.add_argument( + "--statmap-files", nargs="+", help="List of coinc files to be combined" +) +parser.add_argument( + "--background-files", + nargs="+", + default=None, help="full data coinc_statmap files for use in background" - " calculation when used for injections") -parser.add_argument('--censor-ifar-threshold', type=float, default=0.003, + " calculation when used for injections", +) +parser.add_argument( + "--censor-ifar-threshold", + type=float, + default=0.003, help="If provided, only window out foreground triggers with IFAR (years)" - " above the threshold [default=0.003yr]") -parser.add_argument('--veto-window', type=float, default=0.1, - help="Time around each zerolag trigger to window out [default=.1s]") -parser.add_argument('--cluster-window', type=float, - help="Time interval to cluster coincident events over") -parser.add_argument('--cluster-stat', default='stat', - help="What ranking to use when clustering zerolag [default='stat']") -parser.add_argument('--output-coinc-types', action='store_true', + " above the threshold [default=0.003yr]", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) +parser.add_argument( + "--cluster-window", + type=float, + help="Time interval to cluster coincident events over", +) +parser.add_argument( + "--cluster-stat", + default="stat", + help="What ranking to use when clustering zerolag [default='stat']", +) +parser.add_argument( + "--output-coinc-types", + action="store_true", help="Create additional foreground dataset recording coinc type for each" - " event. Mainly for debugging") -parser.add_argument('--max-hierarchical-removal', type=int, default=0, + " event. Mainly for debugging", +) +parser.add_argument( + "--max-hierarchical-removal", + type=int, + default=0, help="Maximum amount of hierarchical removals to carry out. Choose -1 " - "for continuous hierarchical removal until no foreground triggers " - "are louder than the chosen threshold. Choose 0 to not do any " - "hierarchical removals. Choose 1 to do at most 1 hierarchical " - "removal, etc. [default=0]") -parser.add_argument('--hierarchical-removal-window', type=float, default=1., + "for continuous hierarchical removal until no foreground triggers " + "are louder than the chosen threshold. Choose 0 to not do any " + "hierarchical removals. Choose 1 to do at most 1 hierarchical " + "removal, etc. [default=0]", +) +parser.add_argument( + "--hierarchical-removal-window", + type=float, + default=1.0, help="Time around each trigger to window out for a loud trigger in " - "hierarchical removal. [default=1s]") -parser.add_argument('--hierarchical-removal-ifar-thresh', type=float, - default=100., + "hierarchical removal. [default=1s]", +) +parser.add_argument( + "--hierarchical-removal-ifar-thresh", + type=float, + default=100.0, help="Minimum IFAR for a foreground event to be hierarchically removed " - "from background of quieter events (years) [default=100yr]") -parser.add_argument('--hierarchical-removal-against', type=str, - default='none', choices=['none', 'inclusive', 'exclusive'], - help='If doing hierarchical removal, remove foreground ' - 'triggers that are louder than either the "inclusive"' - ' (little-dogs-in) background, or the "exclusive" ' - '(little-dogs-out) background. [default="none"]') + "from background of quieter events (years) [default=100yr]", +) +parser.add_argument( + "--hierarchical-removal-against", + type=str, + default="none", + choices=["none", "inclusive", "exclusive"], + help="If doing hierarchical removal, remove foreground " + 'triggers that are louder than either the "inclusive"' + ' (little-dogs-in) background, or the "exclusive" ' + '(little-dogs-out) background. [default="none"]', +) significance.insert_significance_option_group(parser) -parser.add_argument('--output-file', help="name of output file") +parser.add_argument("--output-file", help="name of output file") args = parser.parse_args() @@ -70,107 +125,111 @@ injection_style = args.background_files != None significance.check_significance_options(args, parser) if args.max_hierarchical_removal and injection_style: - raise NotImplementedError("Hierarchical background removal doesn't make " - "sense for injections.") + raise NotImplementedError( + "Hierarchical background removal doesn't make sense for injections." + ) # Check that the user chose inclusive or exclusive background to perform # hierarchical removals of foreground triggers against. if args.max_hierarchical_removal == 0: - if args.hierarchical_removal_against != 'none': - parser.error("User Error: 0 maximum hierarchical removals chosen but " - "option for --hierarchical-removal-against was given. " - "These are conflicting options. Use with --help for more " - "information.") -else : - if args.hierarchical_removal_against == 'none': - parser.error("--max-hierarchical-removal requires a choice of which " - "background to remove foreground triggers against, " - "inclusive or exclusive. Use with --help for more " - "information.") + if args.hierarchical_removal_against != "none": + parser.error( + "User Error: 0 maximum hierarchical removals chosen but " + "option for --hierarchical-removal-against was given. " + "These are conflicting options. Use with --help for more " + "information." + ) +elif args.hierarchical_removal_against == "none": + parser.error( + "--max-hierarchical-removal requires a choice of which " + "background to remove foreground triggers against, " + "inclusive or exclusive. Use with --help for more " + "information." + ) pycbc.init_logging(args.verbose) -files = [pycbc.io.HFile(n, 'r') for n in args.statmap_files] +files = [pycbc.io.HFile(n, "r") for n in args.statmap_files] f = pycbc.io.HFile(args.output_file, "w") # Work out the combinations of detectors used by each input file all_ifo_combos = [] for fi in files: - ifo_list = get_ifo_string(fi).split(' ') - all_ifo_combos.append(''.join(ifo_list)) + ifo_list = get_ifo_string(fi).split(" ") + all_ifo_combos.append("".join(ifo_list)) -significance_dict = significance.digest_significance_options(all_ifo_combos, - args) +significance_dict = significance.digest_significance_options(all_ifo_combos, args) -logging.info('Copying segments and attributes to %s' % args.output_file) +logging.info("Copying segments and attributes to %s" % args.output_file) # Move segments information into the final file - remove some duplication # in earlier files. Also set up dictionaries to contain segments from the # individual statmap files indiv_segs = segments.segmentlistdict({}) for fi in files: - key = get_ifo_string(fi).replace(' ','') - starts = fi['segments/{}/start'.format(key)][:] - ends = fi['segments/{}/end'.format(key)][:] + key = get_ifo_string(fi).replace(" ", "") + starts = fi[f"segments/{key}/start"][:] + ends = fi[f"segments/{key}/end"][:] indiv_segs[key] = pycbc.events.veto.start_end_to_segments(starts, ends) - f['segments/{}/start'.format(key)] = starts - f['segments/{}/end'.format(key)] = ends - if 'segments/foreground_veto' in fi: - f['segments/%s/foreground_veto/end' % key] = \ - fi['segments/foreground_veto/end'][:] - f['segments/%s/foreground_veto/start' % key] = \ - fi['segments/foreground_veto/start'][:] + f[f"segments/{key}/start"] = starts + f[f"segments/{key}/end"] = ends + if "segments/foreground_veto" in fi: + f["segments/%s/foreground_veto/end" % key] = fi["segments/foreground_veto/end"][ + : + ] + f["segments/%s/foreground_veto/start" % key] = fi[ + "segments/foreground_veto/start" + ][:] for attr_name in fi.attrs: if key not in f: f.create_group(key) f[key].attrs[attr_name] = fi.attrs[attr_name] -logging.info('Combining foreground segments') +logging.info("Combining foreground segments") # combine the segment list from each ifo foreground_segs = segments.segmentlist([]) for k in all_ifo_combos: foreground_segs += indiv_segs[k] -f.attrs['foreground_time'] = abs(foreground_segs) +f.attrs["foreground_time"] = abs(foreground_segs) # Output the segments which are in *any* type of coincidence -f['segments/coinc/start'], f['segments/coinc/end'] = \ +f["segments/coinc/start"], f["segments/coinc/end"] = ( pycbc.events.veto.segments_to_start_end(foreground_segs) +) # obtain list of all ifos involved in the coinc_statmap files -all_ifos = np.unique([ifo for fi in files - for ifo in get_ifo_string(fi).split(' ')]) +all_ifos = np.unique([ifo for fi in files for ifo in get_ifo_string(fi).split(" ")]) # output inherits ifo list -f.attrs['ifos'] = ' '.join(sorted(all_ifos)) -all_ifo_combos = [get_ifo_string(fi).replace(' ','') for fi in files] +f.attrs["ifos"] = " ".join(sorted(all_ifos)) +all_ifo_combos = [get_ifo_string(fi).replace(" ", "") for fi in files] -logging.info('Copying foreground datasets') -for k in files[0]['foreground']: - if not k.startswith('fap') and k not in all_ifos: - pycbc.io.combine_and_copy(f, files, 'foreground/' + k) +logging.info("Copying foreground datasets") +for k in files[0]["foreground"]: + if not k.startswith("fap") and k not in all_ifos: + pycbc.io.combine_and_copy(f, files, "foreground/" + k) if not injection_style: - logging.info('Copying background datasets') - for k in files[0]['background']: + logging.info("Copying background datasets") + for k in files[0]["background"]: if k not in all_ifos: - pycbc.io.combine_and_copy(f, files, 'background/' + k) - for k in files[0]['background_exc']: + pycbc.io.combine_and_copy(f, files, "background/" + k) + for k in files[0]["background_exc"]: if k not in all_ifos: - pycbc.io.combine_and_copy(f, files, 'background_exc/' + k) + pycbc.io.combine_and_copy(f, files, "background_exc/" + k) # create dataset of ifo combination strings fg_coinc_type = np.array([]) for f_in in files: - key = get_ifo_string(f_in).replace(' ','') - combo_repeat = np.array(np.repeat(key.encode('utf8'), - f_in['foreground/stat'].size)) + key = get_ifo_string(f_in).replace(" ", "") + combo_repeat = np.array(np.repeat(key.encode("utf8"), f_in["foreground/stat"].size)) fg_coinc_type = np.concatenate([fg_coinc_type, combo_repeat]) if args.output_coinc_types: - f['foreground/ifo_combination'] = fg_coinc_type + f["foreground/ifo_combination"] = fg_coinc_type -logging.info('Collating triggers into single structure') +logging.info("Collating triggers into single structure") # Initialise arrays for filling with time and trigger ids fg_trig_times = {} @@ -192,93 +251,126 @@ for ifo in all_ifos: # If an ifo does not participate in any given coinc then fill with -1 values for f_in in files: for ifo in all_ifos: - if ifo in f_in['foreground']: - fg_trig_times[ifo] = np.concatenate([fg_trig_times[ifo], - f_in['foreground/{}/time'.format(ifo)][:]]) - fg_trig_ids[ifo] = np.concatenate([fg_trig_ids[ifo], - f_in['foreground/{}/trigger_id'.format(ifo)][:]]) + if ifo in f_in["foreground"]: + fg_trig_times[ifo] = np.concatenate( + [fg_trig_times[ifo], f_in[f"foreground/{ifo}/time"][:]] + ) + fg_trig_ids[ifo] = np.concatenate( + [fg_trig_ids[ifo], f_in[f"foreground/{ifo}/trigger_id"][:]] + ) if not injection_style: - bg_trig_times[ifo] = np.concatenate([bg_trig_times[ifo], - f_in['background/{}/time'.format(ifo)][:]]) - bg_trig_ids[ifo] = np.concatenate([bg_trig_ids[ifo], - f_in['background/{}/trigger_id'.format(ifo)][:]]) - bg_exc_trig_times[ifo] = np.concatenate([bg_exc_trig_times[ifo], - f_in['background_exc/{}/time'.format(ifo)][:]]) - bg_exc_trig_ids[ifo] = np.concatenate([bg_exc_trig_ids[ifo], - f_in['background_exc/{}/trigger_id'.format(ifo)][:]]) + bg_trig_times[ifo] = np.concatenate( + [bg_trig_times[ifo], f_in[f"background/{ifo}/time"][:]] + ) + bg_trig_ids[ifo] = np.concatenate( + [bg_trig_ids[ifo], f_in[f"background/{ifo}/trigger_id"][:]] + ) + bg_exc_trig_times[ifo] = np.concatenate( + [ + bg_exc_trig_times[ifo], + f_in[f"background_exc/{ifo}/time"][:], + ] + ) + bg_exc_trig_ids[ifo] = np.concatenate( + [ + bg_exc_trig_ids[ifo], + f_in[f"background_exc/{ifo}/trigger_id"][:], + ] + ) else: - fg_trig_times[ifo] = np.concatenate([fg_trig_times[ifo], - -1 * np.ones_like(f_in['foreground/stat'][:], - dtype=float)]) - fg_trig_ids[ifo] = np.concatenate([fg_trig_ids[ifo], - -1 * np.ones_like(f_in['foreground/stat'][:], - dtype=int)]) + fg_trig_times[ifo] = np.concatenate( + [ + fg_trig_times[ifo], + -1 * np.ones_like(f_in["foreground/stat"][:], dtype=float), + ] + ) + fg_trig_ids[ifo] = np.concatenate( + [ + fg_trig_ids[ifo], + -1 * np.ones_like(f_in["foreground/stat"][:], dtype=int), + ] + ) if not injection_style: - bg_trig_times[ifo] = np.concatenate([bg_trig_times[ifo], - -1 * np.ones_like(f_in['background/stat'][:], - dtype=float)]) - bg_trig_ids[ifo] = np.concatenate([bg_trig_ids[ifo], - -1 * np.ones_like(f_in['background/stat'][:], - dtype=int)]) - bg_exc_trig_times[ifo] = np.concatenate([bg_exc_trig_times[ifo], - -1 * np.ones_like(f_in['background_exc/stat'][:], - dtype=float)]) - bg_exc_trig_ids[ifo] = np.concatenate([bg_exc_trig_ids[ifo], - -1 * np.ones_like(f_in['background_exc/stat'][:], - dtype=int)]) -n_triggers = f['foreground/stat'].size -logging.info('{} foreground events before clustering'.format(n_triggers)) + bg_trig_times[ifo] = np.concatenate( + [ + bg_trig_times[ifo], + -1 * np.ones_like(f_in["background/stat"][:], dtype=float), + ] + ) + bg_trig_ids[ifo] = np.concatenate( + [ + bg_trig_ids[ifo], + -1 * np.ones_like(f_in["background/stat"][:], dtype=int), + ] + ) + bg_exc_trig_times[ifo] = np.concatenate( + [ + bg_exc_trig_times[ifo], + -1 * np.ones_like(f_in["background_exc/stat"][:], dtype=float), + ] + ) + bg_exc_trig_ids[ifo] = np.concatenate( + [ + bg_exc_trig_ids[ifo], + -1 * np.ones_like(f_in["background_exc/stat"][:], dtype=int), + ] + ) +n_triggers = f["foreground/stat"].size +logging.info(f"{n_triggers} foreground events before clustering") for ifo in all_ifos: - f.create_dataset('foreground/{}/time'.format(ifo), - data=fg_trig_times[ifo]) - f.create_dataset('foreground/{}/trigger_id'.format(ifo), - data=fg_trig_ids[ifo]) + f.create_dataset(f"foreground/{ifo}/time", data=fg_trig_times[ifo]) + f.create_dataset(f"foreground/{ifo}/trigger_id", data=fg_trig_ids[ifo]) if not injection_style: - f.create_dataset('background/{}/time'.format(ifo), - data=bg_trig_times[ifo]) - f.create_dataset('background/{}/trigger_id'.format(ifo), - data=bg_trig_ids[ifo]) - f.create_dataset('background_exc/{}/time'.format(ifo), - data=bg_exc_trig_times[ifo]) - f.create_dataset('background_exc/{}/trigger_id'.format(ifo), - data=bg_exc_trig_ids[ifo]) + f.create_dataset(f"background/{ifo}/time", data=bg_trig_times[ifo]) + f.create_dataset(f"background/{ifo}/trigger_id", data=bg_trig_ids[ifo]) + f.create_dataset( + f"background_exc/{ifo}/time", data=bg_exc_trig_times[ifo] + ) + f.create_dataset( + f"background_exc/{ifo}/trigger_id", data=bg_exc_trig_ids[ifo] + ) # fg_times is a tuple of trigger time arrays -fg_times = (f['foreground/%s/time' % ifo][:] for ifo in all_ifos) +fg_times = (f["foreground/%s/time" % ifo][:] for ifo in all_ifos) # Cluster using the chosen method. Currently only clustering zerolag, # i.e. foreground, so set all timeslide_ids to zero -cidx = pycbc.events.cluster_coincs_multiifo(f[f'foreground/{args.cluster_stat}'][:], - fg_times, - np.zeros(n_triggers), 0, - args.cluster_window) +cidx = pycbc.events.cluster_coincs_multiifo( + f[f"foreground/{args.cluster_stat}"][:], + fg_times, + np.zeros(n_triggers), + 0, + args.cluster_window, +) del fg_times + def filter_dataset(h5file, name, idx): # Dataset needs to be deleted and remade as it is a different size filtered_dset = h5file[name][:][idx] del h5file[name] h5file[name] = filtered_dset + # Downsample the foreground columns to only the loudest ifar between the # multiple files -for key in f['foreground'].keys(): +for key in f["foreground"].keys(): if key not in all_ifos: - filter_dataset(f, 'foreground/%s' % key, cidx) + filter_dataset(f, "foreground/%s" % key, cidx) else: # key is an ifo - for k in f['foreground/%s' % key].keys(): - filter_dataset(f, 'foreground/{}/{}'.format(key, k), cidx) + for k in f["foreground/%s" % key].keys(): + filter_dataset(f, f"foreground/{key}/{k}", cidx) fg_coinc_type = fg_coinc_type[cidx] -n_triggers = f['foreground/stat'].size +n_triggers = f["foreground/stat"].size -logging.info('Calculating event times to determine which types of coinc ' - 'are available') -times_tuple = (f['foreground/{}/time'.format(ifo)][:] for ifo in all_ifos) -test_times = np.array([pycbc.events.mean_if_greater_than_zero(tc)[0] - for tc in zip(*times_tuple)]) +logging.info("Calculating event times to determine which types of coinc are available") +times_tuple = (f[f"foreground/{ifo}/time"][:] for ifo in all_ifos) +test_times = np.array( + [pycbc.events.mean_if_greater_than_zero(tc)[0] for tc in zip(*times_tuple)] +) del times_tuple @@ -288,15 +380,15 @@ is_in_combo_time = {} for key in all_ifo_combos: logging.info("Checking if events are in %s time", key) is_in_combo_time[key] = np.zeros(n_triggers) - end_times = np.array(f['segments/%s/end' % key][:]) - start_times = np.array(f['segments/%s/start' % key][:]) - idx_within_segment = pycbc.events.indices_within_times(test_times, - start_times, - end_times) + end_times = np.array(f["segments/%s/end" % key][:]) + start_times = np.array(f["segments/%s/start" % key][:]) + idx_within_segment = pycbc.events.indices_within_times( + test_times, start_times, end_times + ) is_in_combo_time[key][idx_within_segment] = np.ones_like(idx_within_segment) del idx_within_segment -logging.info('Calculating FAR over all coinc types for foreground events') +logging.info("Calculating FAR over all coinc types for foreground events") far = {} far_exc = {} @@ -307,47 +399,48 @@ if injection_style: # if background files are provided, this is being used for injections # use provided background files to calculate the FARs for bg_fname in args.background_files: - bg_f = pycbc.io.HFile(bg_fname, 'r') - ifo_combo_key = bg_f.attrs['ifos'].replace(' ','') + bg_f = pycbc.io.HFile(bg_fname, "r") + ifo_combo_key = bg_f.attrs["ifos"].replace(" ", "") _, far[ifo_combo_key], _ = significance.get_far( - bg_f['background/stat'][:], - f['foreground/stat'][:], - bg_f['background/decimation_factor'][:], - bg_f.attrs['background_time'], - **significance_dict[ifo_combo_key]) - - _, far_exc[ifo_combo_key], _ = \ - significance.get_far( - bg_f['background_exc/stat'][:], - f['foreground/stat'][:], - bg_f['background_exc/decimation_factor'][:], - bg_f.attrs['background_time_exc'], - **significance_dict[ifo_combo_key]) + bg_f["background/stat"][:], + f["foreground/stat"][:], + bg_f["background/decimation_factor"][:], + bg_f.attrs["background_time"], + **significance_dict[ifo_combo_key], + ) + + _, far_exc[ifo_combo_key], _ = significance.get_far( + bg_f["background_exc/stat"][:], + f["foreground/stat"][:], + bg_f["background_exc/decimation_factor"][:], + bg_f.attrs["background_time_exc"], + **significance_dict[ifo_combo_key], + ) bg_f.close() else: # if not injection style input files, then the input files will have the # background included for f_in in files: - ifo_combo_key = get_ifo_string(f_in).replace(' ','') - _, far[ifo_combo_key], _ = \ - significance.get_far( - f_in['background/stat'][:], - f['foreground/stat'][:], - f_in['background/decimation_factor'][:], - f_in.attrs['background_time'], - **significance_dict[ifo_combo_key]) - - _, far_exc[ifo_combo_key], _ = \ - significance.get_far( - f_in['background_exc/stat'][:], - f['foreground/stat'][:], - f_in['background_exc/decimation_factor'][:], - f_in.attrs['background_time_exc'], - **significance_dict[ifo_combo_key]) + ifo_combo_key = get_ifo_string(f_in).replace(" ", "") + _, far[ifo_combo_key], _ = significance.get_far( + f_in["background/stat"][:], + f["foreground/stat"][:], + f_in["background/decimation_factor"][:], + f_in.attrs["background_time"], + **significance_dict[ifo_combo_key], + ) + + _, far_exc[ifo_combo_key], _ = significance.get_far( + f_in["background_exc/stat"][:], + f["foreground/stat"][:], + f_in["background_exc/decimation_factor"][:], + f_in.attrs["background_time_exc"], + **significance_dict[ifo_combo_key], + ) del ifo_combo_key -logging.info('Combining false alarm rates from all available backgrounds') +logging.info("Combining false alarm rates from all available backgrounds") # Convert dictionary of whether the ifo combination is available at trigger # time into a 2D mask @@ -364,76 +457,74 @@ fg_fars_exc_out = np.sum(isincombo_mask * fg_fars_exc, axis=0) # Apply any limits as appropriate fg_fars_out = significance.apply_far_limit( - fg_fars_out, - significance_dict, - combo=fg_coinc_type) + fg_fars_out, significance_dict, combo=fg_coinc_type +) fg_fars_exc_out = significance.apply_far_limit( - fg_fars_exc_out, - significance_dict, - combo=fg_coinc_type) + fg_fars_exc_out, significance_dict, combo=fg_coinc_type +) -fg_ifar = conv.sec_to_year(1. / fg_fars_out) -fg_ifar_exc = conv.sec_to_year(1. / fg_fars_exc_out) -fg_time = f.attrs['foreground_time'] +fg_ifar = conv.sec_to_year(1.0 / fg_fars_out) +fg_ifar_exc = conv.sec_to_year(1.0 / fg_fars_exc_out) +fg_time = f.attrs["foreground_time"] del isincombo_mask, fg_fars, fg_fars_exc, _ -f.attrs['foreground_time_exc'] = f.attrs['foreground_time'] +f.attrs["foreground_time_exc"] = f.attrs["foreground_time"] if not injection_style: # Construct the foreground censor veto from the clustered candidate times # above the ifar threshold thr = test_times[fg_ifar > args.censor_ifar_threshold] vstart = thr - args.veto_window vend = thr + args.veto_window - vtime = segments.segmentlist([segments.segment(s, e) - for s, e in zip(vstart, vend)]) - logging.info('Censoring %.2f seconds', abs(vtime)) - f.attrs['foreground_time_exc'] -= abs(vtime) - f['segments/foreground_veto/start'] = vstart - f['segments/foreground_veto/end'] = vend + vtime = segments.segmentlist([segments.segment(s, e) for s, e in zip(vstart, vend)]) + logging.info("Censoring %.2f seconds", abs(vtime)) + f.attrs["foreground_time_exc"] -= abs(vtime) + f["segments/foreground_veto/start"] = vstart + f["segments/foreground_veto/end"] = vend # Only output non-exclusive ifar/fap if it is _not_ an injection case - f['foreground/ifar'][:] = fg_ifar - f['foreground/fap'] = 1 - np.exp(-conv.sec_to_year(fg_time) / fg_ifar) + f["foreground/ifar"][:] = fg_ifar + f["foreground/fap"] = 1 - np.exp(-conv.sec_to_year(fg_time) / fg_ifar) del test_times -f['foreground/ifar_exc'][:] = fg_ifar_exc -fg_time_exc = f.attrs['foreground_time_exc'] -f['foreground/fap_exc'] = 1 - np.exp(-conv.sec_to_year(fg_time_exc) / - fg_ifar_exc) +f["foreground/ifar_exc"][:] = fg_ifar_exc +fg_time_exc = f.attrs["foreground_time_exc"] +f["foreground/fap_exc"] = 1 - np.exp(-conv.sec_to_year(fg_time_exc) / fg_ifar_exc) del fg_ifar_exc, fg_ifar # Hierarchical removal stage if args.max_hierarchical_removal == 0: f.close() - logging.info('Not performing hierarchical removal. Done!') + logging.info("Not performing hierarchical removal. Done!") exit() -logging.info('Performing hierarchical removal') +logging.info("Performing hierarchical removal") # Datasets required for hier removal -grps = ['decimation_factor', 'stat', 'template_id', 'timeslide_id', 'ifar'] +grps = ["decimation_factor", "stat", "template_id", "timeslide_id", "ifar"] for ifo in all_ifos: - grps += ['%s/time' % ifo] - grps += ['%s/trigger_id' % ifo] + grps += ["%s/time" % ifo] + grps += ["%s/trigger_id" % ifo] -fg_grps = grps + ['fap', 'fap_exc', 'ifar_exc'] +fg_grps = grps + ["fap", "fap_exc", "ifar_exc"] -combined_bg_data = pycbc.io.DictArray(data={g: f['background/%s' % g][:] - for g in grps}) -combined_fg_data = pycbc.io.DictArray(data={g: f['foreground/%s' % g][:] - for g in fg_grps}) +combined_bg_data = pycbc.io.DictArray(data={g: f["background/%s" % g][:] for g in grps}) +combined_fg_data = pycbc.io.DictArray( + data={g: f["foreground/%s" % g][:] for g in fg_grps} +) # Get coinc type for all coincidences fg_coinc_type = np.array([]) bg_coinc_type = np.array([]) for f_in in files: - key = get_ifo_string(f_in).replace(' ','') - combo_repeat_fg = np.array(np.repeat(key.encode('utf8'), - f_in['foreground/stat'].size)) + key = get_ifo_string(f_in).replace(" ", "") + combo_repeat_fg = np.array( + np.repeat(key.encode("utf8"), f_in["foreground/stat"].size) + ) fg_coinc_type = np.concatenate([fg_coinc_type, combo_repeat_fg]) - combo_repeat_bg = np.array(np.repeat(key.encode('utf8'), - f_in['background/stat'].size)) + combo_repeat_bg = np.array( + np.repeat(key.encode("utf8"), f_in["background/stat"].size) + ) bg_coinc_type = np.concatenate([bg_coinc_type, combo_repeat_bg]) # Apply previously used clustering fg_coinc_type = fg_coinc_type[cidx] @@ -444,29 +535,36 @@ sep_fg_data = {} sep_bg_data = {} final_fg_data = {} for combo in all_ifo_combos: - idx_fg_ct = np.nonzero(fg_coinc_type == combo.encode('utf8')) + idx_fg_ct = np.nonzero(fg_coinc_type == combo.encode("utf8")) sep_fg_data[combo] = combined_fg_data.select(idx_fg_ct) - idx_bg_ct = np.nonzero(bg_coinc_type == combo.encode('utf8')) + idx_bg_ct = np.nonzero(bg_coinc_type == combo.encode("utf8")) sep_bg_data[combo] = combined_bg_data.select(idx_bg_ct) final_fg_data[combo] = pycbc.io.DictArray( - data={k: np.array([], sep_fg_data[combo].data[k].dtype) - for k in sep_fg_data[combo].data}) + data={ + k: np.array([], sep_fg_data[combo].data[k].dtype) + for k in sep_fg_data[combo].data + } + ) final_combined_fg = pycbc.io.DictArray( - data={k: np.array([], combined_fg_data.data[k].dtype) - for k in combined_fg_data.data}) - -fg_time_ct = {f_in.attrs['ifos'].replace(' ',''): f_in.attrs['foreground_time'] - for f_in in files} -bg_time_ct = {f_in.attrs['ifos'].replace(' ',''): f_in.attrs['background_time'] - for f_in in files} + data={ + k: np.array([], combined_fg_data.data[k].dtype) for k in combined_fg_data.data + } +) + +fg_time_ct = { + f_in.attrs["ifos"].replace(" ", ""): f_in.attrs["foreground_time"] for f_in in files +} +bg_time_ct = { + f_in.attrs["ifos"].replace(" ", ""): f_in.attrs["background_time"] for f_in in files +} # Counter for number of removals h_iterations = 0 -if args.hierarchical_removal_against == 'inclusive': - ifar_key = 'ifar' +if args.hierarchical_removal_against == "inclusive": + ifar_key = "ifar" else: - ifar_key = 'ifar_exc' + ifar_key = "ifar_exc" # Break out of loop if max number of removals is reached # or no more triggers above specified IFAR threshold @@ -476,33 +574,37 @@ while True: if h_iterations == 0: # copy over existing data as h0 for combo in all_ifo_combos: - bg_grps[combo] = ['ifar', 'stat', 'timeslide_id'] + bg_grps[combo] = ["ifar", "stat", "timeslide_id"] for key in sep_fg_data[combo].data: - full_fg_key = 'foreground_h0/%s/%s' % (combo, key) + full_fg_key = "foreground_h0/%s/%s" % (combo, key) f[full_fg_key] = sep_fg_data[combo].data[key][:] for key in bg_grps[combo]: - f['background_h0/%s/%s' % (combo, key)] = \ - sep_bg_data[combo].data[key][:] - for key in ['stat', 'timeslide_id']: - f['background_h0/%s' % (key)] = \ - combined_bg_data.data[key][:] + f["background_h0/%s/%s" % (combo, key)] = sep_bg_data[combo].data[key][ + : + ] + for key in ["stat", "timeslide_id"]: + f["background_h0/%s" % (key)] = combined_bg_data.data[key][:] for key in combined_fg_data.data: - f['foreground_h0/%s' % (key)] = \ - combined_fg_data.data[key][:] + f["foreground_h0/%s" % (key)] = combined_fg_data.data[key][:] else: for combo in all_ifo_combos: - bg_grps[combo] = ['ifar', 'stat', 'timeslide_id', - 'decimation_factor', 'template_id'] + bg_grps[combo] = [ + "ifar", + "stat", + "timeslide_id", + "decimation_factor", + "template_id", + ] for key in sep_fg_data[combo].data: - if key.split('/')[0] in all_ifos: + if key.split("/")[0] in all_ifos: bg_grps[combo] += [key] - full_fg_key = 'foreground_h%d/%s/%s' % (h_iterations, combo, - key) - comp_fg_key = 'foreground_h%d/%s/%s' % (h_iterations - 1, - combo, key) - if comp_fg_key in f \ - and sep_fg_data[combo].data[key].size == f[comp_fg_key].size \ - and all(sep_fg_data[combo].data[key][:] == f[comp_fg_key][:]): + full_fg_key = "foreground_h%d/%s/%s" % (h_iterations, combo, key) + comp_fg_key = "foreground_h%d/%s/%s" % (h_iterations - 1, combo, key) + if ( + comp_fg_key in f + and sep_fg_data[combo].data[key].size == f[comp_fg_key].size + and all(sep_fg_data[combo].data[key][:] == f[comp_fg_key][:]) + ): # if the group has not changed create a hard link f[full_fg_key] = f[comp_fg_key] else: @@ -510,98 +612,115 @@ while True: for key in bg_grps[combo]: # The background will (almost certainly) be affected, # so copy as normal - f['background_h%s/%s/%s' % (h_iterations, combo, key)] = \ - sep_bg_data[combo].data[key][:] - f[combo].attrs['foreground_time_h%s' % h_iterations] = \ - fg_time_ct[combo] - f.attrs['foreground_time_h%s' % h_iterations] = fg_time + f["background_h%s/%s/%s" % (h_iterations, combo, key)] = sep_bg_data[ + combo + ].data[key][:] + f[combo].attrs["foreground_time_h%s" % h_iterations] = fg_time_ct[combo] + f.attrs["foreground_time_h%s" % h_iterations] = fg_time for key in combined_bg_data.data: - f['background_h%s/%s' % (h_iterations, key)] = \ - combined_bg_data.data[key][:] + f["background_h%s/%s" % (h_iterations, key)] = combined_bg_data.data[key][:] for key in combined_fg_data.data: - f['foreground_h%s/%s' % (h_iterations, key)] = \ - combined_fg_data.data[key][:] + f["foreground_h%s/%s" % (h_iterations, key)] = combined_fg_data.data[key][:] - if (h_iterations == args.max_hierarchical_removal): - logging.info("Reached hierarchical removal limit of %d" % - args.max_hierarchical_removal) + if h_iterations == args.max_hierarchical_removal: + logging.info( + "Reached hierarchical removal limit of %d" % args.max_hierarchical_removal + ) break - max_each_combo = {combo: sep_fg_data[combo].data[ifar_key][:].argmax() - for combo in all_ifo_combos - if len(sep_fg_data[combo].data[ifar_key][:]) > 0} - max_ifars = {combo: sep_fg_data[combo].data[ifar_key][:][hidx] - for combo, hidx in max_each_combo.items()} - - logging.info('Maximum IFAR values per combination:') + max_each_combo = { + combo: sep_fg_data[combo].data[ifar_key][:].argmax() + for combo in all_ifo_combos + if len(sep_fg_data[combo].data[ifar_key][:]) > 0 + } + max_ifars = { + combo: sep_fg_data[combo].data[ifar_key][:][hidx] + for combo, hidx in max_each_combo.items() + } + + logging.info("Maximum IFAR values per combination:") for k in max_ifars: - logging.info('{}: {:.3g}'.format(k, max_ifars[k])) + logging.info(f"{k}: {max_ifars[k]:.3g}") if args.verbose: # Debug statements max_combd = combined_fg_data.data[ifar_key].argmax() max_combd_ifar = combined_fg_data.data[ifar_key][max_combd] - logging.info('combined: {:.3g}'.format(max_combd_ifar)) + logging.info(f"combined: {max_combd_ifar:.3g}") maxcombo = max(max_ifars, key=lambda k: max_ifars[k]) max_ifar_idx = max_each_combo[maxcombo] max_ifar = max_ifars[maxcombo] if not max_ifar > args.hierarchical_removal_ifar_thresh: - logging.info("Loudest event IFAR of %.3f in %s is less than threshold" - " %f, stopping hierarchical removal" - % (max_ifar, maxcombo, - args.hierarchical_removal_ifar_thresh)) + logging.info( + "Loudest event IFAR of %.3f in %s is less than threshold" + " %f, stopping hierarchical removal" + % (max_ifar, maxcombo, args.hierarchical_removal_ifar_thresh) + ) break h_iterations += 1 # Add the highest ifar (yet to be removed) to the final output - final_fg_data[maxcombo] = final_fg_data[maxcombo] + \ - sep_fg_data[maxcombo].select([max_ifar_idx]) + final_fg_data[maxcombo] = final_fg_data[maxcombo] + sep_fg_data[maxcombo].select( + [max_ifar_idx] + ) maxtime = pycbc.events.mean_if_greater_than_zero( - [sep_fg_data[maxcombo].data['%s/time' % ifo][:][max_ifar_idx] for - ifo in all_ifos if ifo in maxcombo])[0] - logging.info('Removing trigger at time {:.2f} with ifar {:.3g} from {} ' - '& combined foreground '.format(maxtime, max_ifar, maxcombo)) - where_combined = np.flatnonzero(combined_fg_data.data['stat'] == - sep_fg_data[maxcombo].data['stat'][:][max_ifar_idx]) + [ + sep_fg_data[maxcombo].data["%s/time" % ifo][:][max_ifar_idx] + for ifo in all_ifos + if ifo in maxcombo + ] + )[0] + logging.info( + f"Removing trigger at time {maxtime:.2f} with ifar {max_ifar:.3g} from {maxcombo} " + "& combined foreground " + ) + where_combined = np.flatnonzero( + combined_fg_data.data["stat"] + == sep_fg_data[maxcombo].data["stat"][:][max_ifar_idx] + ) sep_fg_data[maxcombo] = sep_fg_data[maxcombo].remove(max_ifar_idx) # Add to final dataset and remove from continuing dataset - final_combined_fg = final_combined_fg + \ - combined_fg_data.select(where_combined) + final_combined_fg = final_combined_fg + combined_fg_data.select(where_combined) combined_fg_data = combined_fg_data.remove(where_combined) fg_coinc_type = np.delete(fg_coinc_type, where_combined) n_triggers -= 1 - logging.info('Removing background triggers at time {} within window ' - '{}s'.format(maxtime, args.hierarchical_removal_window)) + logging.info( + f"Removing background triggers at time {maxtime} within window {args.hierarchical_removal_window}s" + ) for combo in all_ifo_combos: all_hred_idx = [] for ifo in all_ifos: if ifo in combo: - times = sep_bg_data[combo].data['%s/time' % ifo] - hred_ids = np.nonzero(abs(times - maxtime) < - args.hierarchical_removal_window)[0] + times = sep_bg_data[combo].data["%s/time" % ifo] + hred_ids = np.nonzero( + abs(times - maxtime) < args.hierarchical_removal_window + )[0] all_hred_idx += list(hred_ids) - logging.info('Removing {} background triggers from {}'.format( - len(all_hred_idx), combo)) + logging.info( + f"Removing {len(all_hred_idx)} background triggers from {combo}" + ) sep_bg_data[combo] = sep_bg_data[combo].remove(all_hred_idx) hred_ids = [] for ifo in all_ifos: - times = combined_bg_data.data['%s/time' % ifo][:] - within_window = np.flatnonzero(abs(times - maxtime) < - args.hierarchical_removal_window) + times = combined_bg_data.data["%s/time" % ifo][:] + within_window = np.flatnonzero( + abs(times - maxtime) < args.hierarchical_removal_window + ) hred_ids += list(within_window) - logging.info('Removing {} background triggers from combined' - ' background'.format(len(hred_ids))) + logging.info( + f"Removing {len(hred_ids)} background triggers from combined background" + ) combined_bg_data = combined_bg_data.remove(hred_ids) logging.info("Recalculating IFARs") - times_tuple = tuple(combined_fg_data.data[ifo + '/time'][:] - for ifo in all_ifos) - test_times = np.array([pycbc.events.mean_if_greater_than_zero(tc)[0] - for tc in zip(*times_tuple)]) + times_tuple = tuple(combined_fg_data.data[ifo + "/time"][:] for ifo in all_ifos) + test_times = np.array( + [pycbc.events.mean_if_greater_than_zero(tc)[0] for tc in zip(*times_tuple)] + ) for key in all_ifo_combos: # In principle, bg time should be adjusted, but is expected to be a # negligible correction @@ -609,11 +728,12 @@ while True: bg_t_y = conv.sec_to_year(bg_time_ct[key]) fg_t_y = conv.sec_to_year(fg_time_ct[key]) bg_far, fg_far, _ = significance.get_far( - sep_bg_data[key].data['stat'], - sep_fg_data[key].data['stat'], - sep_bg_data[key].data['decimation_factor'], + sep_bg_data[key].data["stat"], + sep_fg_data[key].data["stat"], + sep_bg_data[key].data["decimation_factor"], bg_t_y, - **significance_dict[key]) + **significance_dict[key], + ) fg_far = significance.apply_far_limit( fg_far, significance_dict, @@ -625,31 +745,29 @@ while True: combo=key, ) - sep_bg_data[key].data['ifar'] = 1. / bg_far - sep_fg_data[key].data['ifar'] = 1. / fg_far - sep_fg_data[key].data['fap'] = 1 - \ - np.exp(-fg_t_y * fg_far) + sep_bg_data[key].data["ifar"] = 1.0 / bg_far + sep_fg_data[key].data["ifar"] = 1.0 / fg_far + sep_fg_data[key].data["fap"] = 1 - np.exp(-fg_t_y * fg_far) logging.info("Recalculating combined IFARs") for key in all_ifo_combos: _, far[key], _ = significance.get_far( - sep_bg_data[key].data['stat'], - combined_fg_data.data['stat'], - sep_bg_data[key].data['decimation_factor'], + sep_bg_data[key].data["stat"], + combined_fg_data.data["stat"], + sep_bg_data[key].data["decimation_factor"], bg_time_ct[key], - **significance_dict[key]) + **significance_dict[key], + ) # Set up variable for whether each coincidence is available in each coincidence time is_in_combo_time[key] = np.zeros(n_triggers) - end_times = np.array(f['segments/%s/end' % key][:]) - start_times = np.array(f['segments/%s/start' % key][:]) - idx_within_segment = pycbc.events.indices_within_times(test_times, - start_times, - end_times) - is_in_combo_time[key][idx_within_segment] = \ - np.ones_like(idx_within_segment) - - isincombo_mask = np.array([list(is_in_combo_time[ct]) - for ct in all_ifo_combos]) + end_times = np.array(f["segments/%s/end" % key][:]) + start_times = np.array(f["segments/%s/start" % key][:]) + idx_within_segment = pycbc.events.indices_within_times( + test_times, start_times, end_times + ) + is_in_combo_time[key][idx_within_segment] = np.ones_like(idx_within_segment) + + isincombo_mask = np.array([list(is_in_combo_time[ct]) for ct in all_ifo_combos]) fg_fars = np.array([list(far[ct]) for ct in all_ifo_combos]) fg_fars_out = np.sum(isincombo_mask * fg_fars, axis=0) fg_fars_out = significance.apply_far_limit( @@ -658,31 +776,31 @@ while True: combo=fg_coinc_type, ) # Combine the FARs with the mask to obtain the new ifars - combined_fg_data.data['ifar'] = conv.sec_to_year( - 1. / fg_fars_out) + combined_fg_data.data["ifar"] = conv.sec_to_year(1.0 / fg_fars_out) fg_time -= args.cluster_window - combined_fg_data.data['fap'] = 1 - \ - np.exp(-conv.sec_to_year(fg_time) / combined_fg_data.data['ifar']) + combined_fg_data.data["fap"] = 1 - np.exp( + -conv.sec_to_year(fg_time) / combined_fg_data.data["ifar"] + ) for combo in all_ifo_combos: final_fg_data[combo] = final_fg_data[combo] + sep_fg_data[combo] for key in final_fg_data[combo].data: - full_key = 'foreground/%s/%s' % (combo, key) + full_key = "foreground/%s/%s" % (combo, key) if full_key in f: del f[full_key] f[full_key] = final_fg_data[combo].data[key] final_combined_fg = final_combined_fg + combined_fg_data for key in final_combined_fg.data: - full_key = 'foreground/%s' % (key) + full_key = "foreground/%s" % (key) if full_key in f: del f[full_key] f[full_key] = final_combined_fg.data[key] for key in f: - if 'background' in key and (key + '/ifar') in f: - del f[key + '/ifar'] + if "background" in key and (key + "/ifar") in f: + del f[key + "/ifar"] -f.attrs['hierarchical_removal_iterations'] = h_iterations +f.attrs["hierarchical_removal_iterations"] = h_iterations f.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_apply_rerank b/bin/all_sky_search/pycbc_apply_rerank index 24b6d1ad6c8..1e3a669562e 100644 --- a/bin/all_sky_search/pycbc_apply_rerank +++ b/bin/all_sky_search/pycbc_apply_rerank @@ -1,26 +1,39 @@ #!/bin/env python -"""Rewrite statmap file and rerank candidates using the statistic values +""" +Rewrite statmap file and rerank candidates using the statistic values generated from the followup of candidates. """ -import numpy, argparse, pycbc -from pycbc.io import HFile + +import argparse +from shutil import copyfile + +import numpy + +import pycbc from pycbc.conversions import sec_to_year from pycbc.events import significance -from shutil import copyfile +from pycbc.io import HFile parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--stat-files', nargs='+', - help="Statistic files produced by candidate followup codes") -parser.add_argument('--followup-file', - help="File containing the candidate times which were analyzed") -parser.add_argument('--statmap-file', - help="The statmap file containing the candidates to rerank") +parser.add_argument( + "--stat-files", + nargs="+", + help="Statistic files produced by candidate followup codes", +) +parser.add_argument( + "--followup-file", help="File containing the candidate times which were analyzed" +) +parser.add_argument( + "--statmap-file", help="The statmap file containing the candidates to rerank" +) significance.insert_significance_option_group(parser) -parser.add_argument('--ranking-file', +parser.add_argument( + "--ranking-file", help="Provided only for injection sets, use this file to provide the " - "background to rank candidate significance") -parser.add_argument('--output-file') + "background to rank candidate significance", +) +parser.add_argument("--output-file") args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -28,23 +41,23 @@ pycbc.init_logging(args.verbose) significance.check_significance_options(args, parser) # Reconstruct the full set of statistic values for our candidates -f = HFile(args.followup_file, 'r') -num = len(f['offsets']) # Number of followups done +f = HFile(args.followup_file, "r") +num = len(f["offsets"]) # Number of followups done # Mapping between the followups done and the original candidate list # May be shorter due to duplicates in the original set (which combines # background / background_exc, etc -inv = f['inverse'][:] +inv = f["inverse"][:] stats = numpy.zeros(num) -sections = f.attrs['sections'] +sections = f.attrs["sections"] values = [] starts = [] for fname in args.stat_files: - f = HFile(fname, 'r') - s = f.attrs['start_index'] - v = f['stat'][:] - stride = f.attrs['stride'] + f = HFile(fname, "r") + s = f.attrs["start_index"] + v = f["stat"][:] + stride = f.attrs["stride"] stats[s::stride] = v stats = stats[inv] @@ -53,74 +66,67 @@ stats = stats[inv] copyfile(args.statmap_file, args.output_file) o = HFile(args.output_file) -ifo_combo = o.attrs['ifos'].replace(' ','') +ifo_combo = o.attrs["ifos"].replace(" ", "") -significance_dict = significance.digest_significance_options([ifo_combo], - args) +significance_dict = significance.digest_significance_options([ifo_combo], args) # Update the statistic values for sec in sections: # New stats for this section - nsize = len(o[sec]['stat']) - o[sec]['stat'][...] = stats[:nsize] + nsize = len(o[sec]["stat"]) + o[sec]["stat"][...] = stats[:nsize] # use for next section stats = stats[nsize:] -background_time = o.attrs['background_time'] -coinc_time = o.attrs['foreground_time'] -coinc_time_exc = o.attrs['foreground_time_exc'] -background_time_exc = o.attrs['background_time_exc'] +background_time = o.attrs["background_time"] +coinc_time = o.attrs["foreground_time"] +coinc_time_exc = o.attrs["foreground_time_exc"] +background_time_exc = o.attrs["background_time_exc"] # Injection run if args.ranking_file: - f = HFile(args.ranking_file, 'r') - fstat = o['foreground/stat'][:] - backstat = f['background_exc/stat'][:] - dec = f['background_exc/decimation_factor'][:] + f = HFile(args.ranking_file, "r") + fstat = o["foreground/stat"][:] + backstat = f["background_exc/stat"][:] + dec = f["background_exc/decimation_factor"][:] bnum, fnum = significance.get_n_louder( - backstat, - fstat, - dec, - **significance_dict[ifo_combo]) + backstat, fstat, dec, **significance_dict[ifo_combo] + ) ifar = background_time / (fnum + 1) - fap = 1 - numpy.exp(- coinc_time / ifar) - o['foreground/ifar'][...] = sec_to_year(ifar) - o['foreground/fap'][...] = fap + fap = 1 - numpy.exp(-coinc_time / ifar) + o["foreground/ifar"][...] = sec_to_year(ifar) + o["foreground/fap"][...] = fap - o['foreground/ifar_exc'][...] = o['foreground/ifar'][:] - o['foreground/fap_exc'][...] = o['foreground/fap'][:] + o["foreground/ifar_exc"][...] = o["foreground/ifar"][:] + o["foreground/fap_exc"][...] = o["foreground/fap"][:] # full data run else: - fstat = o['foreground/stat'][:] - backstat = o['background/stat'][:] - dec = o['background/decimation_factor'][:] - dec_exc = o['background_exc/decimation_factor'][:] - backstat_exc = o['background_exc/stat'][:] + fstat = o["foreground/stat"][:] + backstat = o["background/stat"][:] + dec = o["background/decimation_factor"][:] + dec_exc = o["background_exc/decimation_factor"][:] + backstat_exc = o["background_exc/stat"][:] bnum, fnum = significance.get_n_louder( - backstat, - fstat, - dec, - **significance_dict[ifo_combo]) + backstat, fstat, dec, **significance_dict[ifo_combo] + ) bnum_exc, fnum_exc = significance.get_n_louder( - backstat_exc, - fstat, - dec_exc, - **significance_dict[ifo_combo]) + backstat_exc, fstat, dec_exc, **significance_dict[ifo_combo] + ) - o['background/ifar'][...] = sec_to_year(background_time / (bnum + 1)) - o['background_exc/ifar'][...] = sec_to_year(background_time_exc / (bnum_exc + 1)) + o["background/ifar"][...] = sec_to_year(background_time / (bnum + 1)) + o["background_exc/ifar"][...] = sec_to_year(background_time_exc / (bnum_exc + 1)) ifar = background_time / (fnum + 1) - fap = 1 - numpy.exp(- coinc_time / ifar) - o['foreground/ifar'][...] = sec_to_year(ifar) - o['foreground/fap'][...] = fap + fap = 1 - numpy.exp(-coinc_time / ifar) + o["foreground/ifar"][...] = sec_to_year(ifar) + o["foreground/fap"][...] = fap ifar_exc = background_time_exc / (fnum_exc + 1) - fap_exc = 1 - numpy.exp(- coinc_time_exc / ifar_exc) - o['foreground/ifar_exc'][...] = sec_to_year(ifar_exc) - o['foreground/fap_exc'][...] = fap_exc + fap_exc = 1 - numpy.exp(-coinc_time_exc / ifar_exc) + o["foreground/ifar_exc"][...] = sec_to_year(ifar_exc) + o["foreground/fap_exc"][...] = fap_exc diff --git a/bin/all_sky_search/pycbc_average_psd b/bin/all_sky_search/pycbc_average_psd index 0f34ae2d456..9dab991c887 100644 --- a/bin/all_sky_search/pycbc_average_psd +++ b/bin/all_sky_search/pycbc_average_psd @@ -16,88 +16,101 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""This program reads the noise PSDs estimated by pycbc_calculate_psd and +""" +This program reads the noise PSDs estimated by pycbc_calculate_psd and calculates the average PSD over time for each detector, as well as the average PSD across time and detectors. The currently implemented averaging method is -the harmonic mean.""" +the harmonic mean. +""" -import logging import argparse +import logging + import numpy as np + import pycbc from pycbc.io import HFile -from pycbc.types import MultiDetOptionAction, FrequencySeries - +from pycbc.types import FrequencySeries, MultiDetOptionAction parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input-files', nargs='+', required=True, metavar='PATH', - help='HDF5 files from pycbc_calculate_psd (one per ' - 'detector) containing the input PSDs to average.') -parser.add_argument('--time-avg-file', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:PATH', - help='Output file names for single-detector PSDs averaged ' - 'over time.') -parser.add_argument('--detector-avg-file', metavar='PATH', - help='Output file name for the average PSD over time and ' - 'detectors.') +parser.add_argument( + "--input-files", + nargs="+", + required=True, + metavar="PATH", + help="HDF5 files from pycbc_calculate_psd (one per " + "detector) containing the input PSDs to average.", +) +parser.add_argument( + "--time-avg-file", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:PATH", + help="Output file names for single-detector PSDs averaged over time.", +) +parser.add_argument( + "--detector-avg-file", + metavar="PATH", + help="Output file name for the average PSD over time and detectors.", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -dynamic_range_factor = pycbc.DYN_RANGE_FAC ** (-2.) +dynamic_range_factor = pycbc.DYN_RANGE_FAC ** (-2.0) time_avg_psds = {} delta_f = None for input_file in args.input_files: - logging.info('Reading %s', input_file) - f = HFile(input_file, 'r') + logging.info("Reading %s", input_file) + f = HFile(input_file, "r") ifo = tuple(f.keys())[0] - df = f[ifo + '/psds/0'].attrs['delta_f'] + df = f[ifo + "/psds/0"].attrs["delta_f"] if delta_f is None: delta_f = df elif delta_f != df: - raise ValueError('Inconsistent frequency resolution in input PSDs ' - '(%f vs %f)' % (df, delta_f)) - keys = f[ifo + '/psds'].keys() + raise ValueError( + "Inconsistent frequency resolution in input PSDs (%f vs %f)" % (df, delta_f) + ) + keys = f[ifo + "/psds"].keys() - logging.info('Averaging %s over time', ifo) + logging.info("Averaging %s over time", ifo) sum_inv_psds = None count = 0 - for key in f[ifo + '/psds'].keys(): - psd = f[ifo + '/psds/' + key][:] + for key in f[ifo + "/psds"].keys(): + psd = f[ifo + "/psds/" + key][:] if sum_inv_psds is None: - sum_inv_psds = 1. / psd + sum_inv_psds = 1.0 / psd else: - sum_inv_psds += 1. / psd + sum_inv_psds += 1.0 / psd count += 1 avg_psd = count / sum_inv_psds time_avg_psds[ifo] = avg_psd - if ifo in args.time_avg_file and args.time_avg_file[ifo]: - logging.info('Writing %s average over time', ifo) + if args.time_avg_file.get(ifo): + logging.info("Writing %s average over time", ifo) fs = FrequencySeries( - avg_psd.astype(np.float64) * dynamic_range_factor, - delta_f=delta_f) + avg_psd.astype(np.float64) * dynamic_range_factor, delta_f=delta_f + ) fs.save(args.time_avg_file[ifo], ifo=ifo) if args.detector_avg_file: - logging.info('Averaging over detectors') + logging.info("Averaging over detectors") sum_inv_psds = None for ifo, psd in time_avg_psds.items(): if sum_inv_psds is None: - sum_inv_psds = 1. / psd + sum_inv_psds = 1.0 / psd else: - sum_inv_psds += 1. / psd + sum_inv_psds += 1.0 / psd network_psd = len(time_avg_psds) / sum_inv_psds - logging.info('Writing average over detectors') + logging.info("Writing average over detectors") fs = FrequencySeries( - network_psd.astype(np.float64) * dynamic_range_factor, - delta_f=delta_f) - ifo_str = ''.join(sorted(time_avg_psds.keys())) + network_psd.astype(np.float64) * dynamic_range_factor, delta_f=delta_f + ) + ifo_str = "".join(sorted(time_avg_psds.keys())) fs.save(args.detector_avg_file, ifo=ifo_str) -logging.info('Done') - +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_bin_templates b/bin/all_sky_search/pycbc_bin_templates index 196040f2b38..e42622b2762 100755 --- a/bin/all_sky_search/pycbc_bin_templates +++ b/bin/all_sky_search/pycbc_bin_templates @@ -1,8 +1,9 @@ #!/usr/bin/env python -""" Bin templates by their duration -""" -import logging +"""Bin templates by their duration""" + import argparse +import logging + import h5py as h5 import numpy as np @@ -13,43 +14,48 @@ from pycbc.events import background_bin_from_string parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) parser.add_argument("--ifo", type=str, required=True) -parser.add_argument("--f-lower", type=float, default=15., - help='Enforce a uniform low frequency cutoff to ' - 'calculate template duration over the bank') -parser.add_argument('--bank-file', help='hdf format template bank file', - required=True) -parser.add_argument('--background-bins', nargs='+', - help='Used to provide a list of ' - 'precomputed background bins') +parser.add_argument( + "--f-lower", + type=float, + default=15.0, + help="Enforce a uniform low frequency cutoff to " + "calculate template duration over the bank", +) +parser.add_argument("--bank-file", help="hdf format template bank file", required=True) +parser.add_argument( + "--background-bins", + nargs="+", + help="Used to provide a list of precomputed background bins", +) parser.add_argument("--output-file", required=True) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Starting template binning') +logging.info("Starting template binning") -with h5.File(args.bank_file, 'r') as bank: - logging.info('Sorting bank into bins') +with h5.File(args.bank_file, "r") as bank: + logging.info("Sorting bank into bins") data = { - 'mass1': bank['mass1'][:], - 'mass2': bank['mass2'][:], - 'spin1z': bank['spin1z'][:], - 'spin2z': bank['spin2z'][:], - 'f_lower': np.ones_like(bank['mass1'][:]) * args.f_lower - } + "mass1": bank["mass1"][:], + "mass2": bank["mass2"][:], + "spin1z": bank["spin1z"][:], + "spin2z": bank["spin2z"][:], + "f_lower": np.ones_like(bank["mass1"][:]) * args.f_lower, + } bin_dict = background_bin_from_string(args.background_bins, data) - bin_names = [b.split(':')[0] for b in args.background_bins] + bin_names = [b.split(":")[0] for b in args.background_bins] -logging.info('Writing bin template ids to file') -with h5.File(args.output_file, 'w') as f: +logging.info("Writing bin template ids to file") +with h5.File(args.output_file, "w") as f: ifo_grp = f.create_group(args.ifo) for bin_name in bin_names: bin_tids = bin_dict[bin_name] grp = ifo_grp.create_group(bin_name) - grp['tids'] = bin_tids - f.attrs['bank_file'] = args.bank_file - f.attrs['f_lower'] = args.f_lower - f.attrs['background_bins'] = ' '.join(args.background_bins) + grp["tids"] = bin_tids + f.attrs["bank_file"] = args.bank_file + f.attrs["f_lower"] = args.f_lower + f.attrs["background_bins"] = " ".join(args.background_bins) -logging.info('Finished') +logging.info("Finished") diff --git a/bin/all_sky_search/pycbc_bin_trigger_rates_dq b/bin/all_sky_search/pycbc_bin_trigger_rates_dq index db2c22add95..0419a27a07d 100644 --- a/bin/all_sky_search/pycbc_bin_trigger_rates_dq +++ b/bin/all_sky_search/pycbc_bin_trigger_rates_dq @@ -1,21 +1,22 @@ #!/usr/bin/env python -""" Bin triggers by their dq value and calculate trigger rates in each bin -""" -import logging +"""Bin triggers by their dq value and calculate trigger rates in each bin""" + import argparse +import logging -import numpy as np import h5py as h5 - +import numpy as np from igwn_segments import segmentlist import pycbc from pycbc.events import stat as pystat -from pycbc.events.veto import (select_segments_by_definer, - start_end_to_segments, - segments_to_start_end) -from pycbc.types.optparse import MultiDetOptionAction +from pycbc.events.veto import ( + segments_to_start_end, + select_segments_by_definer, + start_end_to_segments, +) from pycbc.io.hdf import SingleDetTriggers +from pycbc.types.optparse import MultiDetOptionAction parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) @@ -25,36 +26,43 @@ parser.add_argument("--flag-file", required=True) parser.add_argument("--flag-name", required=True) parser.add_argument("--analysis-segment-file", required=True) parser.add_argument("--analysis-segment-name", required=True) -parser.add_argument("--gating-windows", nargs='+', - action=MultiDetOptionAction, - help="Seconds to reweight before and after the central" - "time of each gate. Given as detector-values pairs, " - "e.g. H1:-1,2.5 L1:-1,2.5 V1:0,0") -parser.add_argument("--stat-threshold", type=float, default=1., - help="Only consider triggers with --sngl-ranking value " - "above this threshold") +parser.add_argument( + "--gating-windows", + nargs="+", + action=MultiDetOptionAction, + help="Seconds to reweight before and after the central" + "time of each gate. Given as detector-values pairs, " + "e.g. H1:-1,2.5 L1:-1,2.5 V1:0,0", +) +parser.add_argument( + "--stat-threshold", + type=float, + default=1.0, + help="Only consider triggers with --sngl-ranking value above this threshold", +) parser.add_argument("--output-file", required=True) pystat.insert_statistic_option_group( - parser, default_ranking_statistic='single_ranking_only') + parser, default_ranking_statistic="single_ranking_only" +) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Start') +logging.info("Start") -ifo, flag_name = args.flag_name.split(':') +ifo, flag_name = args.flag_name.split(":") if args.gating_windows: gate_times = [] - with h5.File(args.trig_file, 'r') as trig_file: - logging.info('Getting gated times') + with h5.File(args.trig_file, "r") as trig_file: + logging.info("Getting gated times") try: - gating_types = trig_file[f'{ifo}/gating'].keys() + gating_types = trig_file[f"{ifo}/gating"].keys() for gt in gating_types: - gate_times += list(trig_file[f'{ifo}/gating/{gt}/time'][:]) + gate_times += list(trig_file[f"{ifo}/gating/{gt}/time"][:]) gate_times = np.unique(gate_times) except KeyError: - logging.warning('No gating found in trigger file') + logging.warning("No gating found in trigger file") trigs = SingleDetTriggers( args.trig_file, @@ -70,37 +78,35 @@ stat = trigs.get_ranking(args.sngl_ranking) # Get the template bins bin_tids_dict = {} -with h5.File(args.template_bins_file, 'r') as f: +with h5.File(args.template_bins_file, "r") as f: ifo_grp = f[ifo] for bin_name in ifo_grp.keys(): - bin_tids_dict[bin_name] = ifo_grp[bin_name]['tids'][:] + bin_tids_dict[bin_name] = ifo_grp[bin_name]["tids"][:] # get analysis segments analysis_segs = select_segments_by_definer( - args.analysis_segment_file, - segment_name=args.analysis_segment_name, - ifo=ifo) + args.analysis_segment_file, segment_name=args.analysis_segment_name, ifo=ifo +) livetime = abs(analysis_segs) # get flag segments -flag_segs = select_segments_by_definer(args.flag_file, - segment_name=flag_name, - ifo=ifo) +flag_segs = select_segments_by_definer(args.flag_file, segment_name=flag_name, ifo=ifo) # construct gate segments gating_segs = segmentlist([]) if args.gating_windows: - gating_windows = args.gating_windows[ifo].split(',') + gating_windows = args.gating_windows[ifo].split(",") gate_before = float(gating_windows[0]) gate_after = float(gating_windows[1]) if gate_before > 0 or gate_after < 0: - raise ValueError("Gating window values must be negative " - "before gates and positive after gates.") + raise ValueError( + "Gating window values must be negative " + "before gates and positive after gates." + ) if not (gate_before == 0 and gate_after == 0): gating_segs = start_end_to_segments( - gate_times + gate_before, - gate_times + gate_after + gate_times + gate_before, gate_times + gate_after ).coalesce() # make segments into mutually exclusive dq states @@ -122,15 +128,15 @@ def dq_state_at_time(t): # compute and save results -with h5.File(args.output_file, 'w') as f: +with h5.File(args.output_file, "w") as f: ifo_grp = f.create_group(ifo) - all_bin_grp = ifo_grp.create_group('bins') - all_dq_grp = ifo_grp.create_group('dq_segments') + all_bin_grp = ifo_grp.create_group("bins") + all_dq_grp = ifo_grp.create_group("dq_segments") # setup data for each template bin for bin_name, bin_tids in bin_tids_dict.items(): bin_grp = all_bin_grp.create_group(bin_name) - bin_grp['tids'] = bin_tids + bin_grp["tids"] = bin_tids # get the dq states of the triggers in this bin inbin = np.isin(tmplt_ids, bin_tids) @@ -143,20 +149,20 @@ with h5.File(args.output_file, 'w') as f: frac_eff = np.mean(trig_states == state) frac_dt = abs(segs) / livetime dq_rates[state] = frac_eff / frac_dt - bin_grp['dq_rates'] = dq_rates - bin_grp['num_triggers'] = len(trig_times_bin) + bin_grp["dq_rates"] = dq_rates + bin_grp["num_triggers"] = len(trig_times_bin) # save dq state segments for dq_state, segs in dq_state_segs_dict.items(): - name = f'dq_state_{dq_state}' + name = f"dq_state_{dq_state}" dq_grp = all_dq_grp.create_group(name) starts, ends = segments_to_start_end(segs) - dq_grp['segment_starts'] = starts - dq_grp['segment_ends'] = ends - dq_grp['livetime'] = abs(segs) + dq_grp["segment_starts"] = starts + dq_grp["segment_ends"] = ends + dq_grp["livetime"] = abs(segs) - f.attrs['stat'] = f'{ifo}-dq_stat_info' - f.attrs['sngl_ranking'] = args.sngl_ranking - f.attrs['sngl_ranking_threshold'] = args.stat_threshold + f.attrs["stat"] = f"{ifo}-dq_stat_info" + f.attrs["sngl_ranking"] = args.sngl_ranking + f.attrs["sngl_ranking_threshold"] = args.stat_threshold -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_calculate_psd b/bin/all_sky_search/pycbc_calculate_psd index 74d95c98663..431681a08d8 100755 --- a/bin/all_sky_search/pycbc_calculate_psd +++ b/bin/all_sky_search/pycbc_calculate_psd @@ -1,22 +1,39 @@ #!/usr/bin/env python -""" Calculate psd estimates for analysis segments -""" -import logging, argparse, numpy, time, copy +"""Calculate psd estimates for analysis segments""" + +import argparse +import copy +import logging +import time + +import numpy +from igwn_segments import segment, segmentlist from six.moves import zip_longest -import pycbc, pycbc.psd, pycbc.strain, pycbc.events + +import pycbc +import pycbc.events +import pycbc.psd +import pycbc.strain +from pycbc.fft.fftw import set_measure_level from pycbc.io import HFile from pycbc.pool import BroadcastPool as Pool -from pycbc.fft.fftw import set_measure_level from pycbc.workflow import resolve_td_option -from igwn_segments import segmentlist, segment + set_measure_level(0) parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--low-frequency-cutoff", type=float, required=True, - help="The low frequency cutoff to use for filtering (Hz)") -parser.add_argument("--analysis-segment-file", required=True, - help="File defining the segments to estimate PSDs over") +parser.add_argument( + "--low-frequency-cutoff", + type=float, + required=True, + help="The low frequency cutoff to use for filtering (Hz)", +) +parser.add_argument( + "--analysis-segment-file", + required=True, + help="File defining the segments to estimate PSDs over", +) parser.add_argument("--segment-name", help="Name of segment list to use") parser.add_argument("--cores", default=1, type=int) parser.add_argument("--output-file", required=True) @@ -31,20 +48,24 @@ pycbc.init_logging(args.verbose) pycbc.psd.verify_psd_options(args, parser) pycbc.strain.StrainSegments.verify_segment_options(args, parser) - + + def grouper(n, iterable): args = [iter(iterable)] * n return list([e for e in t if e != None] for t in zip_longest(*args)) + def get_psd(input_tuple): - """ Get the PSDs for the given data chunck. This follows the same rules + """ + Get the PSDs for the given data chunck. This follows the same rules as pycbc_inspiral for determining where to calculate PSDs """ seg = input_tuple[0] i = input_tuple[1] - - logging.info('%d: getting strain for %.1f-%.1f (%.1f s)', i, seg[0], - seg[1], abs(seg)) + + logging.info( + "%d: getting strain for %.1f-%.1f (%.1f s)", i, seg[0], seg[1], abs(seg) + ) argstmp = copy.deepcopy(args) argstmp.gps_start_time = int(seg[0]) + args.pad_data argstmp.gps_end_time = int(seg[1]) - args.pad_data @@ -67,10 +88,10 @@ def get_psd(input_tuple): break except RuntimeError: time.sleep(10) - if i == (args.num_data_retries-1): + if i == (args.num_data_retries - 1): raise - logging.info('%d: determining strain segmentation', i) + logging.info("%d: determining strain segmentation", i) strain_segments = pycbc.strain.StrainSegments.from_cli(args, gwstrain) flow = args.low_frequency_cutoff @@ -78,28 +99,31 @@ def get_psd(input_tuple): tlen = strain_segments.time_len delta_f = strain_segments.delta_f - logging.info('%d: calculating psd', i) - psds_and_times = pycbc.psd.generate_overlapping_psds(args, gwstrain, - flen, delta_f, flow, dyn_range_factor=pycbc.DYN_RANGE_FAC) + logging.info("%d: calculating psd", i) + psds_and_times = pycbc.psd.generate_overlapping_psds( + args, gwstrain, flen, delta_f, flow, dyn_range_factor=pycbc.DYN_RANGE_FAC + ) lpsd = [] for start_idx, end_idx, psd in psds_and_times: - start_time = gwstrain.start_time + start_idx/gwstrain.sample_rate - end_time = gwstrain.start_time + end_idx/gwstrain.sample_rate + start_time = gwstrain.start_time + start_idx / gwstrain.sample_rate + end_time = gwstrain.start_time + end_idx / gwstrain.sample_rate lpsd.append((psd.numpy(), psd.delta_f, int(start_time), int(end_time))) return lpsd + # Determine what times to calculate PSDs for ifo = args.channel_name[0:2] -segments = pycbc.events.select_segments_by_definer(args.analysis_segment_file, - args.segment_name, ifo=ifo) +segments = pycbc.events.select_segments_by_definer( + args.analysis_segment_file, args.segment_name, ifo=ifo +) # get rid of duplicate segments which happen when splitting the bank segments = segmentlist(frozenset(segments)) -# Calculate the PSDs -logging.info('%d psds to calculate', len(segments)) +# Calculate the PSDs +logging.info("%d psds to calculate", len(segments)) if len(segments) > 0: pool = Pool(args.cores) @@ -109,25 +133,25 @@ else: psds = [] # Store the PSDs in an hdf file, include some basic metadata -f = HFile(args.output_file, 'w') -psd_group = f.create_group(ifo + '/psds') +f = HFile(args.output_file, "w") +psd_group = f.create_group(ifo + "/psds") inc, start, end = 0, [], [] for gpsd in psds: for psd_numpy, psd_delta_f, s, e in gpsd: - logging.info('writing psd %d', inc) + logging.info("writing psd %d", inc) key = str(inc) start.append(int(s)) end.append(int(e)) - psd_group.create_dataset(key, data=psd_numpy, compression='gzip', - compression_opts=9, shuffle=True) - psd_group[key].attrs['epoch'] = int(s) - psd_group[key].attrs['delta_f'] = float(psd_delta_f) + psd_group.create_dataset( + key, data=psd_numpy, compression="gzip", compression_opts=9, shuffle=True + ) + psd_group[key].attrs["epoch"] = int(s) + psd_group[key].attrs["delta_f"] = float(psd_delta_f) inc += 1 -f[ifo + '/start_time'] = numpy.array(start, dtype=numpy.uint32) -f[ifo + '/end_time'] = numpy.array(end, dtype=numpy.uint32) -f.attrs['low_frequency_cutoff'] = args.low_frequency_cutoff -f.attrs['dynamic_range_factor'] = pycbc.DYN_RANGE_FAC - -logging.info('Done!') +f[ifo + "/start_time"] = numpy.array(start, dtype=numpy.uint32) +f[ifo + "/end_time"] = numpy.array(end, dtype=numpy.uint32) +f.attrs["low_frequency_cutoff"] = args.low_frequency_cutoff +f.attrs["dynamic_range_factor"] = pycbc.DYN_RANGE_FAC +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_coinc_findtrigs b/bin/all_sky_search/pycbc_coinc_findtrigs index 3fa74ec33ca..df9194ba41a 100644 --- a/bin/all_sky_search/pycbc_coinc_findtrigs +++ b/bin/all_sky_search/pycbc_coinc_findtrigs @@ -1,74 +1,131 @@ #!/usr/bin/env python -import copy, argparse, logging, numpy, numpy.random -import shutil, uuid, os.path, atexit +import argparse +import atexit +import copy +import logging +import os.path +import shutil +import uuid + +import numpy +import numpy.random from igwn_segments import infinity +from numpy.random import seed, shuffle + import pycbc -from pycbc.events import veto, coinc, stat, cuts +from pycbc import init_logging, pool +from pycbc.events import coinc, cuts, stat, veto from pycbc.io import HFile -from pycbc import pool, init_logging -from numpy.random import seed, shuffle from pycbc.io.hdf import ReadByTemplate from pycbc.types.optparse import MultiDetOptionAction parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument("--veto-files", nargs='*', action='append', default=[], - help="Optional veto file. Triggers within veto segments " - "contained in the file are ignored") -parser.add_argument("--segment-name", nargs='*', action='append', default=[], - help="Optional, name of veto segment in veto file") -parser.add_argument("--gating-veto-windows", nargs='+', - action=MultiDetOptionAction, - help="Seconds to be vetoed before and after the central time " - "of each gate. Given as detector-values pairs, e.g. " - "H1:-1,2.5 L1:-1,2.5 V1:0,0") -parser.add_argument("--trigger-files", nargs='*', action='append', default=[], - help="Files containing single-detector triggers") -parser.add_argument("--template-bank", required=True, - help="Template bank file in HDF format") -parser.add_argument("--pivot-ifo", required=True, - help="Add the ifo to use as the pivot for multi " - "detector coincidence") -parser.add_argument("--fixed-ifo", required=True, - help="Add the ifo to use as the fixed ifo for " - "multi detector coincidence") +parser.add_argument( + "--veto-files", + nargs="*", + action="append", + default=[], + help="Optional veto file. Triggers within veto segments " + "contained in the file are ignored", +) +parser.add_argument( + "--segment-name", + nargs="*", + action="append", + default=[], + help="Optional, name of veto segment in veto file", +) +parser.add_argument( + "--gating-veto-windows", + nargs="+", + action=MultiDetOptionAction, + help="Seconds to be vetoed before and after the central time " + "of each gate. Given as detector-values pairs, e.g. " + "H1:-1,2.5 L1:-1,2.5 V1:0,0", +) +parser.add_argument( + "--trigger-files", + nargs="*", + action="append", + default=[], + help="Files containing single-detector triggers", +) +parser.add_argument( + "--template-bank", required=True, help="Template bank file in HDF format" +) +parser.add_argument( + "--pivot-ifo", + required=True, + help="Add the ifo to use as the pivot for multi detector coincidence", +) +parser.add_argument( + "--fixed-ifo", + required=True, + help="Add the ifo to use as the fixed ifo for multi detector coincidence", +) # produces a list of lists to allow multiple invocations and multiple args parser.add_argument("--use-maxalpha", action="store_true") -parser.add_argument("--coinc-threshold", type=float, default=0.0, - help="Seconds to add to time-of-flight coincidence window") -parser.add_argument("--timeslide-interval", type=float, - help="Interval between timeslides in seconds. Timeslides are" - " disabled if the option is omitted.") -parser.add_argument("--loudest-keep-values", - default='[6:1]', - help="Apply successive multiplicative levels of" - " decimation to coincs with stat value below the" - " given thresholds. Supply as a comma-separated list" - " of threshold:decimation value pairs surrounded by" - " square brackets (no spaces!). Decimation values must" - " be positive integers." - " Ex. [15:5,10:30,5:30,0:30]." - " Default: no decimation") -parser.add_argument("--template-fraction-range", default="0/1", - help="Optional, analyze only part of template bank. Format" - " PART/NUM_PARTS") -parser.add_argument("--randomize-template-order", action="store_true", - help="Random shuffle templates with fixed seed " - "before selecting range to analyze") -parser.add_argument("--cluster-window", type=float, - help="Optional, window size in seconds to cluster " - "coincidences over the bank") -parser.add_argument("--output-file", - help="File to store the coincident triggers") -parser.add_argument("--batch-singles", default=5000, type=int, - help="Number of single triggers to process at once") -parser.add_argument('--nprocesses', type=int, default=1, - help="Number of processes to use") -parser.add_argument('--stage-input', action='store_true', - help="Stage input files through to speed up" - "access by multiple processes") -parser.add_argument('--stage-input-dir', type=str, default='/dev/shm', - help="Directory to stage input files") +parser.add_argument( + "--coinc-threshold", + type=float, + default=0.0, + help="Seconds to add to time-of-flight coincidence window", +) +parser.add_argument( + "--timeslide-interval", + type=float, + help="Interval between timeslides in seconds. Timeslides are" + " disabled if the option is omitted.", +) +parser.add_argument( + "--loudest-keep-values", + default="[6:1]", + help="Apply successive multiplicative levels of" + " decimation to coincs with stat value below the" + " given thresholds. Supply as a comma-separated list" + " of threshold:decimation value pairs surrounded by" + " square brackets (no spaces!). Decimation values must" + " be positive integers." + " Ex. [15:5,10:30,5:30,0:30]." + " Default: no decimation", +) +parser.add_argument( + "--template-fraction-range", + default="0/1", + help="Optional, analyze only part of template bank. Format PART/NUM_PARTS", +) +parser.add_argument( + "--randomize-template-order", + action="store_true", + help="Random shuffle templates with fixed seed before selecting range to analyze", +) +parser.add_argument( + "--cluster-window", + type=float, + help="Optional, window size in seconds to cluster coincidences over the bank", +) +parser.add_argument("--output-file", help="File to store the coincident triggers") +parser.add_argument( + "--batch-singles", + default=5000, + type=int, + help="Number of single triggers to process at once", +) +parser.add_argument( + "--nprocesses", type=int, default=1, help="Number of processes to use" +) +parser.add_argument( + "--stage-input", + action="store_true", + help="Stage input files through to speed upaccess by multiple processes", +) +parser.add_argument( + "--stage-input-dir", + type=str, + default="/dev/shm", + help="Directory to stage input files", +) stat.insert_statistic_option_group(parser) cuts.insert_cuts_option_group(parser) args = parser.parse_args() @@ -80,39 +137,46 @@ args.trigger_files = sum(args.trigger_files, []) init_logging(args.verbose) + def parse_template_range(num_templates, rangestr): - part = int(rangestr.split('/')[0]) - pieces = int(rangestr.split('/')[1]) + part = int(rangestr.split("/")[0]) + pieces = int(rangestr.split("/")[1]) tmin = int(num_templates / float(pieces) * part) - tmax = int(num_templates / float(pieces) * (part+1)) + tmax = int(num_templates / float(pieces) * (part + 1)) return tmin, tmax -logging.info('Starting...') + +logging.info("Starting...") trigger_cut_dict, template_cut_dict = cuts.ingest_cuts_option_group(args) -num_templates = len(HFile(args.template_bank, "r")['template_hash']) +num_templates = len(HFile(args.template_bank, "r")["template_hash"]) tmin, tmax = parse_template_range(num_templates, args.template_fraction_range) -logging.info('Analyzing template %s - %s' % (tmin, tmax-1)) +logging.info("Analyzing template %s - %s" % (tmin, tmax - 1)) + -class MultiifoTrigs(object): +class MultiifoTrigs: """store trigger info in parallel with ifo name and shift vector""" + def __init__(self): self.ifos = [] self.to_shift = [] self.singles = [] + trigs = MultiifoTrigs() cleanup_files = [] + def exit_cleaning(): for fname in cleanup_files: - logging.info('cleaning up %s', fname) + logging.info("cleaning up %s", fname) try: os.remove(fname) except OSError as e: print(e) - pass + + atexit.register(exit_cleaning) # If some templates have no triggers one or more ifos, there can be no @@ -121,40 +185,43 @@ tids_with_trigs = None for i in range(len(args.trigger_files)): if args.stage_input: - dest = os.path.join(args.stage_input_dir, str(uuid.uuid4()) + '.hdf') - logging.info("Moving %s to shared memory as %s", - args.trigger_files[i], dest) + dest = os.path.join(args.stage_input_dir, str(uuid.uuid4()) + ".hdf") + logging.info("Moving %s to shared memory as %s", args.trigger_files[i], dest) cleanup_files.append(dest) shutil.copyfile(args.trigger_files[i], dest) else: dest = args.trigger_files[i] - logging.info('Opening trigger file %s: %s' % (i, dest)) - reader = ReadByTemplate(dest, - args.template_bank, - args.segment_name, - args.veto_files, - args.gating_veto_windows) + logging.info("Opening trigger file %s: %s" % (i, dest)) + reader = ReadByTemplate( + dest, + args.template_bank, + args.segment_name, + args.veto_files, + args.gating_veto_windows, + ) ifo = reader.ifo trigs.ifos.append(ifo) # We don't have that many triggers, see if we can skip some templates - if len(reader.file[ifo]['template_id']) < 2**27: - uniq = numpy.unique(reader.file[ifo]['template_id'][:]).astype(numpy.int64) + if len(reader.file[ifo]["template_id"]) < 2**27: + uniq = numpy.unique(reader.file[ifo]["template_id"][:]).astype(numpy.int64) if tids_with_trigs is None: tids_with_trigs = numpy.arange(0, num_templates, dtype=numpy.int64) tids_with_trigs = numpy.intersect1d(tids_with_trigs, uniq) # time shift is subtracted from pivot ifo time trigs.to_shift.append(-1 if ifo == args.pivot_ifo else 0) - logging.info('Applying time shift multiple %i to ifo %s' % - (trigs.to_shift[-1], trigs.ifos[-1])) + logging.info( + "Applying time shift multiple %i to ifo %s" + % (trigs.to_shift[-1], trigs.ifos[-1]) + ) trigs.singles.append(reader) # Coinc_segs contains only segments where all ifos are analyzed coinc_segs = veto.start_end_to_segments([-infinity()], [infinity()]) for i, sngl in zip(trigs.ifos, trigs.singles): - coinc_segs = (coinc_segs & sngl.segs) + coinc_segs = coinc_segs & sngl.segs for sngl in trigs.singles: sngl.segs = coinc_segs sngl.valid = veto.segments_to_start_end(sngl.segs) @@ -166,8 +233,9 @@ rank_method = stat.get_statistic_from_opts(args, trigs.ifos) # Earth crossing time, which is approximately 0.085 seconds. TWOEARTH = 0.085 if args.timeslide_interval is not None and args.timeslide_interval <= TWOEARTH: - raise parser.error("The time slide interval should be larger " - "than twice the Earth crossing time.") + raise parser.error( + "The time slide interval should be larger than twice the Earth crossing time." + ) # slide = 0 means don't do timeslides if args.timeslide_interval is None: @@ -193,27 +261,31 @@ template_ids = cuts.apply_template_cuts( template_cut_dict, statistic=rank_method, ifos=trigs.ifos, - template_ids=template_ids) + template_ids=template_ids, +) -logging.info("%d out of %d templates kept after applying template cuts", - len(template_ids), original_bank_len) +logging.info( + "%d out of %d templates kept after applying template cuts", + len(template_ids), + original_bank_len, +) # 'data' will store output of coinc finding # in addition to these lists of coinc info, will also store trigger times and # ids in each ifo -data = {'stat': [], 'decimation_factor': [], 'timeslide_id': [], 'template_id': []} +data = {"stat": [], "decimation_factor": [], "timeslide_id": [], "template_id": []} for ifo in trigs.ifos: - data['%s/time' % ifo] = [] - data['%s/trigger_id' % ifo] = [] + data["%s/time" % ifo] = [] + data["%s/trigger_id" % ifo] = [] factors = [1] threshes = [numpy.inf] -loudest_keep_vals = args.loudest_keep_values.strip('[]').split(',') +loudest_keep_vals = args.loudest_keep_values.strip("[]").split(",") for decstr in loudest_keep_vals: - thresh, factor = decstr.split(':') + thresh, factor = decstr.split(":") if float(factor) % 1: - raise RuntimeError("Non-integer decimation is not supported") + raise RuntimeError("Non-integer decimation is not supported") if int(factor) < 1: raise RuntimeError("Negative or zero decimation does not make sense") if int(factor) == 1: @@ -230,45 +302,43 @@ factors = numpy.array(factors)[threshorder] # Decimation factors are applied successively in descending order of threshold total_factors = numpy.cumprod(factors) + # Gather the coincs from a single template def process_template(tnum): local_data = copy.deepcopy(data) times_full = {} sds_full = {} tids_full = {} - logging.debug('Obtaining trigs for template %i ..' % (tnum)) + logging.debug("Obtaining trigs for template %i .." % (tnum)) for i, sngl in zip(trigs.ifos, trigs.singles): # Apply cuts to triggers tids_uncut = sngl.set_template(tnum) - trigger_keep_ids = cuts.apply_trigger_cuts(sngl, trigger_cut_dict, - statistic=rank_method) + trigger_keep_ids = cuts.apply_trigger_cuts( + sngl, trigger_cut_dict, statistic=rank_method + ) tids_full[i] = tids_uncut[trigger_keep_ids] - times_full[i] = sngl['end_time'][trigger_keep_ids] - logging.debug('%s:%s', i, len(tids_uncut)) + times_full[i] = sngl["end_time"][trigger_keep_ids] + logging.debug("%s:%s", i, len(tids_uncut)) if len(tids_full[i]) < len(tids_uncut): - logging.info("%s triggers cut", - len(tids_uncut) - len(tids_full[i])) - + logging.info("%s triggers cut", len(tids_uncut) - len(tids_full[i])) # get single-detector statistic sds_full[i] = rank_method.single(sngl)[trigger_keep_ids] - mintrigs = min([len(ti) for ti in tids_full.values()]) if mintrigs == 0: - logging.info('No triggers in at least one ifo for template %i, ' - 'skipping' % tnum) + logging.info("No triggers in at least one ifo for template %i, skipping" % tnum) return local_data if type(rank_method.single_dtype) == list: entries = [e[0] for e in rank_method.single_dtype] - if 'snr' in entries: - min_snr = min([numpy.min(sds_full[i]['snr']) for i in trigs.ifos]) + if "snr" in entries: + min_snr = min([numpy.min(sds_full[i]["snr"]) for i in trigs.ifos]) else: min_snr = None - if 'sigmasq'in entries: - max_sigmasq = min([numpy.max(sds_full[i]['sigmasq']) for i in trigs.ifos]) + if "sigmasq" in entries: + max_sigmasq = min([numpy.max(sds_full[i]["sigmasq"]) for i in trigs.ifos]) else: max_sigmasq = None else: @@ -278,8 +348,8 @@ def process_template(tnum): # Test whether sds_full contains single arrays or record arrays # this depends on the stat being used try: - pivot_stat = sds_full[args.pivot_ifo]['snglstat'].copy() - fixed_stat = sds_full[args.fixed_ifo]['snglstat'].copy() + pivot_stat = sds_full[args.pivot_ifo]["snglstat"].copy() + fixed_stat = sds_full[args.fixed_ifo]["snglstat"].copy() except IndexError: pivot_stat = sds_full[args.pivot_ifo].copy() fixed_stat = sds_full[args.fixed_ifo].copy() @@ -297,8 +367,7 @@ def process_template(tnum): start1 = 0 end0 = start0 + args.batch_singles - if end0 > len(sds_full[args.fixed_ifo]): - end0 = len(sds_full[args.fixed_ifo]) + end0 = min(end0, len(sds_full[args.fixed_ifo])) fixed_idxs = fixed_sort[start0:end0] @@ -317,10 +386,9 @@ def process_template(tnum): fixed_times = {i: times[i] for i in fixed_ifos} if len(fixed_ifos) > 1: - fixed_ids, fixed_slide = coinc.time_multi_coincidence(fixed_times, 0., - args.coinc_threshold, - fixed_ifos[0], - fixed_ifos[1]) + fixed_ids, fixed_slide = coinc.time_multi_coincidence( + fixed_times, 0.0, args.coinc_threshold, fixed_ifos[0], fixed_ifos[1] + ) if len(fixed_slide) == 0: start0 += args.batch_singles continue @@ -341,14 +409,15 @@ def process_template(tnum): # For each trigger in the fixed network calculate the limit int # the pivot detector to pass the current decimation threshold pivot_lims[kidx] = rank_method.coinc_lim_for_thresh( - fixed_single_info, threshes[kidx], + fixed_single_info, + threshes[kidx], args.pivot_ifo, time_addition=args.coinc_threshold, min_snr=min_snr, - max_sigmasq=max_sigmasq + max_sigmasq=max_sigmasq, ) if not rank_method.single_increasing: - pivot_lims[kidx] *= -1. + pivot_lims[kidx] *= -1.0 # subtract small amount to account for errors due to rounding pivot_lims[kidx] -= 1e-6 # Get the minimum statistic required for all triggers at the @@ -357,8 +426,7 @@ def process_template(tnum): while start1 < len(sds_full[args.pivot_ifo]): end1 = start1 + args.batch_singles - if end1 > len(sds_full[args.pivot_ifo]): - end1 = len(sds_full[args.pivot_ifo]) + end1 = min(end1, len(sds_full[args.pivot_ifo])) pivot_idxs = pivot_sort[start1:end1] @@ -369,24 +437,28 @@ def process_template(tnum): tids[args.pivot_ifo] = tids_full[args.pivot_ifo][pivot_idxs] # Do time coincidence for slides that will be kept after the last decimation - ids, slide = coinc.time_multi_coincidence(times, - args.timeslide_interval*total_factors[-1], - args.coinc_threshold, - args.pivot_ifo, - args.fixed_ifo) + ids, slide = coinc.time_multi_coincidence( + times, + args.timeslide_interval * total_factors[-1], + args.coinc_threshold, + args.pivot_ifo, + args.fixed_ifo, + ) slide *= total_factors[-1] single_info = [(i, sds[i][ids[i]]) for i in trigs.ifos] cstat = rank_method.rank_stat_coinc( - single_info, slide, args.timeslide_interval, + single_info, + slide, + args.timeslide_interval, to_shift=trigs.to_shift, - time_addition=args.coinc_threshold + time_addition=args.coinc_threshold, ) - #index values of the zerolag triggers + # index values of the zerolag triggers fi = numpy.where(slide == 0)[0] - #index values of the background triggers + # index values of the background triggers bi = numpy.where(slide != 0)[0] bl = bi[cstat[bi] < threshes[-1]] @@ -396,15 +468,17 @@ def process_template(tnum): cstat = cstat[ti] slide = slide[ti] - dec = numpy.concatenate([numpy.ones(len(fi)), numpy.repeat(total_factors[-1], len(bl))]) + dec = numpy.concatenate( + [numpy.ones(len(fi)), numpy.repeat(total_factors[-1], len(bl))] + ) try: - pivot_stat = sds[args.pivot_ifo]['snglstat'].copy() + pivot_stat = sds[args.pivot_ifo]["snglstat"].copy() except IndexError: pivot_stat = sds[args.pivot_ifo].copy() if not rank_method.single_increasing: - pivot_stat *= -1. + pivot_stat *= -1.0 # Starting from the largest decimation threshold, find the first decimation step # where the loudest single detector trigger in pivot can pass the decimation threshold @@ -418,7 +492,6 @@ def process_template(tnum): # loop through decimation steps starting from the first step where passing # the threshold is possible for kidx in range(tidx, len(threshes)): - # Remove triggers in pivot that cannot form coincidences above # the current decimation threshold pivot_cut = numpy.searchsorted(pivot_stat, pivot_lower[kidx]) @@ -429,11 +502,13 @@ def process_template(tnum): test_times[args.pivot_ifo] = times[args.pivot_ifo][pivot_cut:] # Do time coincidence for the current decimation factor - set_ids, set_slide = coinc.time_multi_coincidence(test_times, - args.timeslide_interval*total_factors[kidx - 1], - args.coinc_threshold, - args.pivot_ifo, - args.fixed_ifo) + set_ids, set_slide = coinc.time_multi_coincidence( + test_times, + args.timeslide_interval * total_factors[kidx - 1], + args.coinc_threshold, + args.pivot_ifo, + args.fixed_ifo, + ) set_slide *= total_factors[kidx - 1] # Remove foreground triggers @@ -444,14 +519,17 @@ def process_template(tnum): # where all triggers from the fixed network coincidence are still # together for i in range(1, len(fixed_ifos)): - sets *= set_ids[fixed_ifos[i-1]] == set_ids[fixed_ifos[i]] + sets *= set_ids[fixed_ifos[i - 1]] == set_ids[fixed_ifos[i]] set_ids = {i: set_ids[i][sets] for i in trigs.ifos} set_slide = set_slide[sets] # Only keep coincidences where pivot has a single stat above the threshold # calculated earlier - above = pivot_s[set_ids[args.pivot_ifo]] >= pivot_lims[kidx][set_ids[fixed_ifos[0]]] + above = ( + pivot_s[set_ids[args.pivot_ifo]] + >= pivot_lims[kidx][set_ids[fixed_ifos[0]]] + ) test_ids = {i: fixed_ids[i][set_ids[i][above]] for i in fixed_ifos} test_ids[args.pivot_ifo] = set_ids[args.pivot_ifo][above] + pivot_cut @@ -459,9 +537,12 @@ def process_template(tnum): test_single_info = [(i, sds[i][test_ids[i]]) for i in trigs.ifos] test_cstat = rank_method.rank_stat_coinc( - test_single_info, set_slide, args.timeslide_interval, + test_single_info, + set_slide, + args.timeslide_interval, to_shift=trigs.to_shift, - time_addition=args.coinc_threshold) + time_addition=args.coinc_threshold, + ) # Keep triggers in the current decimation range test_bi = numpy.where(test_cstat >= threshes[kidx])[0] @@ -471,66 +552,73 @@ def process_template(tnum): ids[i] = numpy.concatenate([ids[i], test_ids[i][test_bi]]) cstat = numpy.concatenate([cstat, test_cstat[test_bi]]) slide = numpy.concatenate([slide, set_slide[test_bi]]) - dec = numpy.concatenate([dec, numpy.repeat(total_factors[kidx - 1], len(test_bi))]) + dec = numpy.concatenate( + [dec, numpy.repeat(total_factors[kidx - 1], len(test_bi))] + ) # temporary storage for decimated trigger ids for ifo in trigs.ifos: addtime = times[ifo][ids[ifo]] addtriggerid = tids[ifo][ids[ifo]] - local_data['%s/time' % ifo] += [addtime] - local_data['%s/trigger_id' % ifo] += [addtriggerid] - local_data['stat'] += [cstat] - local_data['decimation_factor'] += [dec] - local_data['timeslide_id'] += [slide] - local_data['template_id'] += [numpy.repeat(tnum, len(slide))] + local_data["%s/time" % ifo] += [addtime] + local_data["%s/trigger_id" % ifo] += [addtriggerid] + local_data["stat"] += [cstat] + local_data["decimation_factor"] += [dec] + local_data["timeslide_id"] += [slide] + local_data["template_id"] += [numpy.repeat(tnum, len(slide))] start1 += args.batch_singles start0 += args.batch_singles return local_data + if args.nprocesses == 1: ldatas = list(map(process_template, list(template_ids))) else: p = pool.BroadcastPool(args.nprocesses) ldatas = p.map(process_template, list(template_ids)) -logging.info('merging data from the templates') +logging.info("merging data from the templates") for ldata in ldatas: for key in data: data[key] += ldata[key] -if len(data['stat']) > 0: +if len(data["stat"]) > 0: for key in data: data[key] = numpy.concatenate(data[key]) -if args.cluster_window and len(data['stat']) > 0: - timestring0 = '%s/time' % args.pivot_ifo - timestring1 = '%s/time' % args.fixed_ifo - cid = coinc.cluster_coincs(data['stat'], data[timestring0], data[timestring1], - data['timeslide_id'], args.timeslide_interval, - args.cluster_window) - -logging.info('saving coincident triggers') -f = HFile(args.output_file, 'w') -if len(data['stat']) > 0: +if args.cluster_window and len(data["stat"]) > 0: + timestring0 = "%s/time" % args.pivot_ifo + timestring1 = "%s/time" % args.fixed_ifo + cid = coinc.cluster_coincs( + data["stat"], + data[timestring0], + data[timestring1], + data["timeslide_id"], + args.timeslide_interval, + args.cluster_window, + ) + +logging.info("saving coincident triggers") +f = HFile(args.output_file, "w") +if len(data["stat"]) > 0: for key in data: var = data[key][cid] if args.cluster_window else data[key] - f.create_dataset(key, data=var, - compression='gzip', - compression_opts=9, - shuffle=True) + f.create_dataset( + key, data=var, compression="gzip", compression_opts=9, shuffle=True + ) # Store coinc segments keyed by detector combination -key = ''.join(sorted(trigs.ifos)) -f['segments/%s/start' % key], f['segments/%s/end' % key] = trigs.singles[0].valid - -f.attrs['timeslide_interval'] = args.timeslide_interval -f.attrs['coinc_time'] = abs(coinc_segs) -f.attrs['num_of_ifos'] = len(args.trigger_files) -f.attrs['pivot'] = args.pivot_ifo -f.attrs['fixed'] = args.fixed_ifo +key = "".join(sorted(trigs.ifos)) +f["segments/%s/start" % key], f["segments/%s/end" % key] = trigs.singles[0].valid + +f.attrs["timeslide_interval"] = args.timeslide_interval +f.attrs["coinc_time"] = abs(coinc_segs) +f.attrs["num_of_ifos"] = len(args.trigger_files) +f.attrs["pivot"] = args.pivot_ifo +f.attrs["fixed"] = args.fixed_ifo for i, sngl in zip(trigs.ifos, trigs.singles): - f.attrs['%s_foreground_time' % i] = abs(sngl.segs) -f.attrs['ifos'] = ' '.join(sorted(trigs.ifos)) + f.attrs["%s_foreground_time" % i] = abs(sngl.segs) +f.attrs["ifos"] = " ".join(sorted(trigs.ifos)) # What does this code actually calculate? if args.timeslide_interval: @@ -538,7 +626,7 @@ if args.timeslide_interval: nslides = int(maxtrigs / args.timeslide_interval) else: nslides = 0 -f.attrs['num_slides'] = nslides +f.attrs["num_slides"] = nslides -logging.info('Done') +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_coinc_hdfinjfind b/bin/all_sky_search/pycbc_coinc_hdfinjfind index b7c74d70ac0..f2047894540 100755 --- a/bin/all_sky_search/pycbc_coinc_hdfinjfind +++ b/bin/all_sky_search/pycbc_coinc_hdfinjfind @@ -1,18 +1,26 @@ #!/usr/bin/python -"""Associate coincident triggers with injections listed in one or more LIGOLW +""" +Associate coincident triggers with injections listed in one or more LIGOLW files. """ -import argparse, logging, types, numpy, os.path -from igwn_ligolw import lsctables, utils as ligolw_utils +import argparse +import logging +import os.path +import types + import igwn_segments as segments +import numpy +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils + import pycbc from pycbc import events, init_logging from pycbc.events import indices_within_segments -from pycbc.types import MultiDetOptionAction from pycbc.inject import CBCHDFInjectionSet from pycbc.io import HFile from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.types import MultiDetOptionAction def hdf_append(f, key, value): @@ -23,90 +31,104 @@ def hdf_append(f, key, value): else: f[key] = value + HFile.append = types.MethodType(hdf_append, HFile) + def keep_ind(times, start, end): - """ Return the list of indices within the list of start and end times - """ + """Return the list of indices within the list of start and end times""" time_sorting = times.argsort() times = times[time_sorting] indices = numpy.array([], dtype=numpy.uint32) - leftidx = numpy.searchsorted(times, start, side='left') - rightidx = numpy.searchsorted(times, end, side='right') + leftidx = numpy.searchsorted(times, start, side="left") + rightidx = numpy.searchsorted(times, end, side="right") for li, ri in zip(leftidx, rightidx): seg_indices = numpy.arange(li, ri, 1).astype(numpy.uint32) - indices=numpy.union1d(seg_indices, indices) + indices = numpy.union1d(seg_indices, indices) return time_sorting[indices] + def xml_to_hdf(table, hdf_file, hdf_key, columns): - """ Save xml columns as hdf columns, only float32 supported atm. - """ + """Save xml columns as hdf columns, only float32 supported atm.""" for col in columns: - # Look for key in the form 'a:b' where we want to + # Look for key in the form 'a:b' where we want to # read data from column a but store it as named output 'b' # This is used to keep the output consistent where the input may not # depending on the injection file format (xml vs hdf) - if ':' in col: - col, new_col = col.split(':') + if ":" in col: + col, new_col = col.split(":") else: new_col = col key = os.path.join(hdf_key, new_col) - hdf_append(hdf_file, key, numpy.array(table.getColumnByName(col), - dtype=numpy.float32)) + hdf_append( + hdf_file, key, numpy.array(table.getColumnByName(col), dtype=numpy.float32) + ) + parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-files', nargs='+', required=True) -parser.add_argument('--injection-files', nargs='+', required=True) -parser.add_argument('--veto-file') -parser.add_argument('--segment-name', default=None, - help='Name of segment list to use for vetoes. Optional') -parser.add_argument('--injection-window', type=float, required=True) -parser.add_argument('--min-required-ifos', type=int, default=2, - help='Minimum number of IFOs required to be observing and' - ' not vetoed for an injection to be counted. Default 2') -parser.add_argument('--optimal-snr-column', nargs='+', - action=MultiDetOptionAction, metavar='DETECTOR:COLUMN', - help='Names of the sim_inspiral columns containing the' - ' optimal SNRs.') -parser.add_argument('--redshift-column', default=None, - help='Name of sim_inspiral column containing redshift. ' - 'Optional') -parser.add_argument('--output-file', required=True) +parser.add_argument("--trigger-files", nargs="+", required=True) +parser.add_argument("--injection-files", nargs="+", required=True) +parser.add_argument("--veto-file") +parser.add_argument( + "--segment-name", + default=None, + help="Name of segment list to use for vetoes. Optional", +) +parser.add_argument("--injection-window", type=float, required=True) +parser.add_argument( + "--min-required-ifos", + type=int, + default=2, + help="Minimum number of IFOs required to be observing and" + " not vetoed for an injection to be counted. Default 2", +) +parser.add_argument( + "--optimal-snr-column", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + help="Names of the sim_inspiral columns containing the optimal SNRs.", +) +parser.add_argument( + "--redshift-column", + default=None, + help="Name of sim_inspiral column containing redshift. Optional", +) +parser.add_argument("--output-file", required=True) args = parser.parse_args() init_logging(args.verbose) -fo = HFile(args.output_file, 'w') +fo = HFile(args.output_file, "w") injection_index = 0 -for trigger_file, injection_file in zip(args.trigger_files, - args.injection_files): - logging.info('Read in the coinc data: %s' % trigger_file) - f = HFile(trigger_file, 'r') +for trigger_file, injection_file in zip(args.trigger_files, args.injection_files): + logging.info("Read in the coinc data: %s" % trigger_file) + f = HFile(trigger_file, "r") # Get list of groups which contain subgroup 'time' # - these will be the IFOs - ifo_list = [key for key in f['foreground'] - if 'time' in f['foreground/%s/' % key]] + ifo_list = [key for key in f["foreground"] if "time" in f["foreground/%s/" % key]] assert len(ifo_list) > 1 # Check required ifos option if len(ifo_list) < args.min_required_ifos: - raise RuntimeError('min-required-ifos (%s) must be <= number of ifos' - ' being searched (%s)' % - (args.min_required_ifos, len(ifo_list))) - fo.attrs['ifos'] = ' '.join(sorted(ifo_list)) - - template_id = f['foreground/template_id'][:] - stat = f['foreground/stat'][:] - ifar_exc = f['foreground/ifar_exc'][:] - fap_exc = f['foreground/fap_exc'][:] + raise RuntimeError( + "min-required-ifos (%s) must be <= number of ifos" + " being searched (%s)" % (args.min_required_ifos, len(ifo_list)) + ) + fo.attrs["ifos"] = " ".join(sorted(ifo_list)) + + template_id = f["foreground/template_id"][:] + stat = f["foreground/stat"][:] + ifar_exc = f["foreground/ifar_exc"][:] + fap_exc = f["foreground/fap_exc"][:] try: - ifar = f['foreground/ifar'][:] - fap = f['foreground/fap'][:] + ifar = f["foreground/ifar"][:] + fap = f["foreground/fap"][:] except KeyError: - logging.info('No inclusive ifar/fap. Proceeding anyway') + logging.info("No inclusive ifar/fap. Proceeding anyway") ifar = None fap = None # using multi-ifo-style trigger file input @@ -114,74 +136,83 @@ for trigger_file, injection_file in zip(args.trigger_files, time_dict = {} trig_dict = {} for ifo in ifo_list: - ifo_times += (f['foreground/%s/time' % ifo][:],) - time_dict[ifo] = f['foreground/%s/time' % ifo][:] - trig_dict[ifo] = f['foreground/%s/trigger_id' % ifo][:] - time = numpy.array([events.mean_if_greater_than_zero(vals)[0] - for vals in zip(*ifo_times)]) + ifo_times += (f["foreground/%s/time" % ifo][:],) + time_dict[ifo] = f["foreground/%s/time" % ifo][:] + trig_dict[ifo] = f["foreground/%s/trigger_id" % ifo][:] + time = numpy.array( + [events.mean_if_greater_than_zero(vals)[0] for vals in zip(*ifo_times)] + ) # We will discard injections which cannot be associated with a # coincident event, thus combine segments over all combinations # of coincident detectors to determine which times to keep any_seg = segments.segmentlist([]) - for key in f['segments']: - if key == 'foreground': + for key in f["segments"]: + if key == "foreground": continue else: - starts = f['/segments/%s/start' % key][:] - ends = f['/segments/%s/end' % key][:] + starts = f["/segments/%s/start" % key][:] + ends = f["/segments/%s/end" % key][:] any_seg += events.start_end_to_segments(starts, ends) ana_start, ana_end = events.segments_to_start_end(any_seg) time_sorting = time.argsort() - logging.info('Read in the injection file') - if '.xml' in injection_file or '.xml.gz' in injection_file: - indoc = ligolw_utils.load_filename(injection_file, False, - contenthandler=LIGOLWContentHandler) + logging.info("Read in the injection file") + if ".xml" in injection_file or ".xml.gz" in injection_file: + indoc = ligolw_utils.load_filename( + injection_file, False, contenthandler=LIGOLWContentHandler + ) sim_table = lsctables.SimInspiralTable.get_table(indoc) - inj_time = numpy.array(sim_table.getColumnByName('geocent_end_time').asarray() + - 1e-9 * sim_table.getColumnByName('geocent_end_time_ns').asarray(), - dtype=numpy.float64) + inj_time = numpy.array( + sim_table.getColumnByName("geocent_end_time").asarray() + + 1e-9 * sim_table.getColumnByName("geocent_end_time_ns").asarray(), + dtype=numpy.float64, + ) else: inj_file = CBCHDFInjectionSet(injection_file) inj_data = inj_file.table - inj_time = inj_data['tc'][:] - - logging.info('Determined the found injections by time') - left = numpy.searchsorted(time[time_sorting], - inj_time - args.injection_window, side='left') - right = numpy.searchsorted(time[time_sorting], - inj_time + args.injection_window, side='right') - found = numpy.where((right-left) == 1)[0] - missed = numpy.where((right-left) == 0)[0] - ambiguous = numpy.where((right-left) > 1)[0] + inj_time = inj_data["tc"][:] + + logging.info("Determined the found injections by time") + left = numpy.searchsorted( + time[time_sorting], inj_time - args.injection_window, side="left" + ) + right = numpy.searchsorted( + time[time_sorting], inj_time + args.injection_window, side="right" + ) + found = numpy.where((right - left) == 1)[0] + missed = numpy.where((right - left) == 0)[0] + ambiguous = numpy.where((right - left) > 1)[0] missed = numpy.concatenate([missed, ambiguous]) - logging.info('Found: %s, Missed: %s Ambiguous: %s' - % (len(found), len(missed), len(ambiguous))) + logging.info( + "Found: %s, Missed: %s Ambiguous: %s" + % (len(found), len(missed), len(ambiguous)) + ) if len(ambiguous) > 0: - logging.warning('More than one coinc trigger found associated ' - 'with injection') + logging.warning("More than one coinc trigger found associated with injection") am = numpy.arange(0, len(inj_time), 1)[left[ambiguous]] bm = numpy.arange(0, len(inj_time), 1)[right[ambiguous]] - logging.info('Removing injections outside of analyzed time') + logging.info("Removing injections outside of analyzed time") ki = keep_ind(inj_time, ana_start, ana_end) found_within_time = numpy.intersect1d(ki, found) missed_within_time = numpy.intersect1d(ki, missed) - logging.info('Found: %s, Missed: %s' % - (len(found_within_time), len(missed_within_time))) + logging.info( + "Found: %s, Missed: %s" % (len(found_within_time), len(missed_within_time)) + ) if args.veto_file: - logging.info('Removing injections in vetoed time') + logging.info("Removing injections in vetoed time") # Put individual detector vetoes into list of all vetoed indices, if an # injection is vetoed in N ifos then its index will appear N times vetoid = numpy.array([]) veto_dict = {} for ifo in ifo_list: - vi, _ = indices_within_segments(inj_time, [args.veto_file], ifo=ifo, - segment_name=args.segment_name) + vi, _ = indices_within_segments( + inj_time, [args.veto_file], ifo=ifo, segment_name=args.segment_name + ) vetoid = numpy.append(vetoid, vi) veto_dict[ifo] = vi @@ -191,32 +222,48 @@ for trigger_file, injection_file in zip(args.trigger_files, # remove injections where the number of unvetoed ifos is less than the # minimum specified by the user - vetoed_all = vetoid_unique[(len(ifo_list) - count_vetoid) - < args.min_required_ifos] - found_after_vetoes = numpy.array([i for i in found_within_time - if i not in vetoed_all]) - missed_after_vetoes = numpy.array([i for i in missed_within_time - if i not in vetoed_all]).astype(int) - logging.info('Found: %s, Missed: %s' % - (len(found_after_vetoes), len(missed_after_vetoes))) + vetoed_all = vetoid_unique[ + (len(ifo_list) - count_vetoid) < args.min_required_ifos + ] + found_after_vetoes = numpy.array( + [i for i in found_within_time if i not in vetoed_all] + ) + missed_after_vetoes = numpy.array( + [i for i in missed_within_time if i not in vetoed_all] + ).astype(int) + logging.info( + "Found: %s, Missed: %s" + % (len(found_after_vetoes), len(missed_after_vetoes)) + ) else: veto_dict = {ifo: [] for ifo in ifo_list} found_after_vetoes = found_within_time missed_after_vetoes = missed_within_time.astype(int) - found_fore = numpy.arange(0, len(stat), 1)[left[found]] found_fore_v = numpy.arange(0, len(stat), 1)[left[found_after_vetoes]] - logging.info('Saving injection information') - if '.xml' in injection_file or '.xml.gz' in injection_file: + logging.info("Saving injection information") + if ".xml" in injection_file or ".xml.gz" in injection_file: ninj = len(sim_table) - columns = ['mass1', 'mass2', 'spin1x', 'spin1y', - 'spin1z', 'spin2x', 'spin2y', 'spin2z', - 'inclination', 'polarization', 'coa_phase', - 'latitude:dec', 'longitude:ra', 'distance'] - xml_to_hdf(sim_table, fo, 'injections', columns) - hdf_append(fo, 'injections/tc', inj_time) + columns = [ + "mass1", + "mass2", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + "inclination", + "polarization", + "coa_phase", + "latitude:dec", + "longitude:ra", + "distance", + ] + xml_to_hdf(sim_table, fo, "injections", columns) + hdf_append(fo, "injections/tc", inj_time) # pick up optimal SNRs for ifo, column in args.optimal_snr_column.items(): @@ -225,20 +272,22 @@ for trigger_file, injection_file in zip(args.trigger_files, # to later calculate decisive optimal snr optimal_snr_all = numpy.array(sim_table.getColumnByName(column)) optimal_snr_all[veto_dict[ifo]] = 0 - hdf_append(fo, 'injections/optimal_snr_%s' % ifo, - optimal_snr_all) + hdf_append(fo, "injections/optimal_snr_%s" % ifo, optimal_snr_all) # pick up redshift if args.redshift_column: - hdf_append(fo, 'injections/redshift', - sim_table.getColumnByName(args.redshift_column)) + hdf_append( + fo, + "injections/redshift", + sim_table.getColumnByName(args.redshift_column), + ) else: # hdf injection format ninj = len(inj_data) # fill these columns if not provided (rest should be there # if we go this far) as they are commonly used for plots - for k in ['spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y', 'spin2z']: + for k in ["spin1x", "spin1y", "spin1z", "spin2x", "spin2y", "spin2z"]: if k not in inj_data.dtype.names: inj_data = inj_data.add_fields(numpy.zeros(ninj), k) @@ -246,23 +295,23 @@ for trigger_file, injection_file in zip(args.trigger_files, data = inj_data[k][:] # set optimal snr to zero for detectors that were vetoed - if 'optimal_snr' in k: - ifo = k.split('_')[-1] + if "optimal_snr" in k: + ifo = k.split("_")[-1] data[veto_dict[ifo]] = 0 - if data.dtype.char == 'U': - data = data.astype('S') + if data.dtype.char == "U": + data = data.astype("S") - hdf_append(fo, 'injections/{}'.format(k), data) + hdf_append(fo, f"injections/{k}", data) # copy over common search info - if 'foreground_time' in f.attrs.keys(): - fo.attrs['foreground_time'] = f.attrs['foreground_time'] - if 'foreground_time_exc' in f.attrs.keys(): - fo.attrs['foreground_time_exc'] = f.attrs['foreground_time_exc'] + if "foreground_time" in f.attrs.keys(): + fo.attrs["foreground_time"] = f.attrs["foreground_time"] + if "foreground_time_exc" in f.attrs.keys(): + fo.attrs["foreground_time_exc"] = f.attrs["foreground_time_exc"] - for key in f['segments'].keys(): - if 'foreground' in key or 'coinc' in key: + for key in f["segments"].keys(): + if "foreground" in key or "coinc" in key: continue if key not in fo: fo.create_group(key) @@ -270,44 +319,52 @@ for trigger_file, injection_file in zip(args.trigger_files, fkey = f[key] else: fkey = f - if 'pivot' in fo[key].attrs: + if "pivot" in fo[key].attrs: # This is a coincident statmap file - fo[key].attrs['pivot'] = fkey.attrs['pivot'] - fo[key].attrs['fixed'] = fkey.attrs['fixed'] - if 'foreground_time' in fkey.attrs.keys(): - fo[key].attrs['foreground_time'] = fkey.attrs['foreground_time'] - if 'foreground_time_exc' in fkey.attrs.keys(): - fo[key].attrs['foreground_time_exc'] = fkey.attrs['foreground_time_exc'] - - hdf_append(fo, 'missed/all', missed + injection_index) - hdf_append(fo, 'missed/within_analysis', missed_within_time + injection_index) - hdf_append(fo, 'missed/after_vetoes', missed_after_vetoes + injection_index) - hdf_append(fo, 'found/template_id', template_id[time_sorting][found_fore]) - hdf_append(fo, 'found/injection_index', found + injection_index) - hdf_append(fo, 'found/stat', stat[time_sorting][found_fore]) - hdf_append(fo, 'found/ifar_exc', ifar_exc[time_sorting][found_fore]) - hdf_append(fo, 'found/fap_exc', ifar_exc[time_sorting][found_fore]) + fo[key].attrs["pivot"] = fkey.attrs["pivot"] + fo[key].attrs["fixed"] = fkey.attrs["fixed"] + if "foreground_time" in fkey.attrs.keys(): + fo[key].attrs["foreground_time"] = fkey.attrs["foreground_time"] + if "foreground_time_exc" in fkey.attrs.keys(): + fo[key].attrs["foreground_time_exc"] = fkey.attrs["foreground_time_exc"] + + hdf_append(fo, "missed/all", missed + injection_index) + hdf_append(fo, "missed/within_analysis", missed_within_time + injection_index) + hdf_append(fo, "missed/after_vetoes", missed_after_vetoes + injection_index) + hdf_append(fo, "found/template_id", template_id[time_sorting][found_fore]) + hdf_append(fo, "found/injection_index", found + injection_index) + hdf_append(fo, "found/stat", stat[time_sorting][found_fore]) + hdf_append(fo, "found/ifar_exc", ifar_exc[time_sorting][found_fore]) + hdf_append(fo, "found/fap_exc", ifar_exc[time_sorting][found_fore]) if ifar is not None: - hdf_append(fo, 'found/ifar', ifar[time_sorting][found_fore]) - hdf_append(fo, 'found/fap', fap[time_sorting][found_fore]) - hdf_append(fo, 'found_after_vetoes/template_id', - template_id[time_sorting][found_fore_v]) - hdf_append(fo, 'found_after_vetoes/injection_index', - found_after_vetoes + injection_index) - hdf_append(fo, 'found_after_vetoes/stat', stat[time_sorting][found_fore_v]) - hdf_append(fo, 'found_after_vetoes/ifar_exc', ifar_exc[time_sorting][found_fore_v]) - hdf_append(fo, 'found_after_vetoes/fap_exc', fap_exc[time_sorting][found_fore_v]) + hdf_append(fo, "found/ifar", ifar[time_sorting][found_fore]) + hdf_append(fo, "found/fap", fap[time_sorting][found_fore]) + hdf_append( + fo, "found_after_vetoes/template_id", template_id[time_sorting][found_fore_v] + ) + hdf_append( + fo, "found_after_vetoes/injection_index", found_after_vetoes + injection_index + ) + hdf_append(fo, "found_after_vetoes/stat", stat[time_sorting][found_fore_v]) + hdf_append(fo, "found_after_vetoes/ifar_exc", ifar_exc[time_sorting][found_fore_v]) + hdf_append(fo, "found_after_vetoes/fap_exc", fap_exc[time_sorting][found_fore_v]) if ifar is not None: - hdf_append(fo, 'found_after_vetoes/ifar', ifar[time_sorting][found_fore_v]) - hdf_append(fo, 'found_after_vetoes/fap', fap[time_sorting][found_fore_v]) + hdf_append(fo, "found_after_vetoes/ifar", ifar[time_sorting][found_fore_v]) + hdf_append(fo, "found_after_vetoes/fap", fap[time_sorting][found_fore_v]) for ifo in ifo_list: - hdf_append(fo, 'found/%s/time' % ifo, - time_dict[ifo][time_sorting][found_fore]) - hdf_append(fo, 'found/%s/trigger_id' % ifo, - trig_dict[ifo][time_sorting][found_fore]) - hdf_append(fo, 'found_after_vetoes/%s/time' % ifo, - time_dict[ifo][time_sorting][found_fore_v]) - hdf_append(fo, 'found_after_vetoes/%s/trigger_id' % ifo, - trig_dict[ifo][time_sorting][found_fore_v]) + hdf_append(fo, "found/%s/time" % ifo, time_dict[ifo][time_sorting][found_fore]) + hdf_append( + fo, "found/%s/trigger_id" % ifo, trig_dict[ifo][time_sorting][found_fore] + ) + hdf_append( + fo, + "found_after_vetoes/%s/time" % ifo, + time_dict[ifo][time_sorting][found_fore_v], + ) + hdf_append( + fo, + "found_after_vetoes/%s/trigger_id" % ifo, + trig_dict[ifo][time_sorting][found_fore_v], + ) injection_index += ninj diff --git a/bin/all_sky_search/pycbc_coinc_mergetrigs b/bin/all_sky_search/pycbc_coinc_mergetrigs index 77637b85b6d..14fe48f2ebc 100755 --- a/bin/all_sky_search/pycbc_coinc_mergetrigs +++ b/bin/all_sky_search/pycbc_coinc_mergetrigs @@ -1,67 +1,79 @@ #!/usr/bin/env python -""" This program adds single detector hdf trigger files together. -""" +"""This program adds single detector hdf trigger files together.""" + +import argparse +import logging + +import h5py +import numpy -import numpy, argparse, h5py, logging -from pycbc.io import HFile from pycbc import add_common_pycbc_options, init_logging +from pycbc.io import HFile + def changes(arr): l = numpy.where(arr[:-1] != arr[1:])[0] - l = numpy.concatenate(([0], l+1, [len(arr)])) + l = numpy.concatenate(([0], l + 1, [len(arr)])) return numpy.unique(l) + def collect(key, files): data = [] for fname in files: - with HFile(fname, 'r') as fin: + with HFile(fname, "r") as fin: if key in fin: data += [fin[key][:]] return numpy.concatenate(data) + def region(f, key, boundaries, ids): dset = f[key] refs = [] for j in range(len(boundaries) - 1): - l, r = boundaries[ids[j]], boundaries[ids[j]+1] - refs.append(dset.regionref[l:r]) - f.create_dataset(key+'_template', data=refs, - dtype=h5py.special_dtype(ref=h5py.RegionReference)) + l, r = boundaries[ids[j]], boundaries[ids[j] + 1] + refs.append(dset.regionref[l:r]) + f.create_dataset( + key + "_template", data=refs, dtype=h5py.special_dtype(ref=h5py.RegionReference) + ) + parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--trigger-files', nargs='+') -parser.add_argument('--output-file', required=True) -parser.add_argument('--bank-file', required=True) -parser.add_argument('--compression-level', type=int, default=6, - help='Set HDF compression level in the output file ' - '(default 6)') +parser.add_argument("--trigger-files", nargs="+") +parser.add_argument("--output-file", required=True) +parser.add_argument("--bank-file", required=True) +parser.add_argument( + "--compression-level", + type=int, + default=6, + help="Set HDF compression level in the output file (default 6)", +) args = parser.parse_args() init_logging(args.verbose) -f = HFile(args.output_file, 'w') +f = HFile(args.output_file, "w") logging.info("getting the list of columns from a representative file") trigger_columns = [] for fname in args.trigger_files: try: - f2 = HFile(fname, 'r') - except IOError as e: - logging.error("Cannot open %s" % fname) + f2 = HFile(fname, "r") + except OSError as e: + logging.exception("Cannot open %s" % fname) raise e ifo = tuple(f2.keys())[0] if len(f2[ifo].keys()) > 0: k = f2[ifo].keys() trigger_columns = list(f2[ifo].keys()) - if not 'template_hash' in trigger_columns: + if "template_hash" not in trigger_columns: f2.close() continue - trigger_columns.remove('search') - trigger_columns.remove('template_hash') - if 'gating' in trigger_columns: - trigger_columns.remove('gating') + trigger_columns.remove("search") + trigger_columns.remove("template_hash") + if "gating" in trigger_columns: + trigger_columns.remove("gating") f2.close() break f2.close() @@ -69,7 +81,7 @@ for fname in args.trigger_files: for col in trigger_columns: logging.info("trigger column: %s", col) -logging.info('reading the metadata from the files') +logging.info("reading the metadata from the files") tpc = numpy.array([], dtype=numpy.float64) frpc = numpy.array([], dtype=numpy.float64) @@ -81,31 +93,31 @@ ends = numpy.array([], dtype=numpy.float64) gating = {} for filename in args.trigger_files: try: - data = HFile(filename, 'r') - except IOError as e: - logging.error('Cannot open %s', filename) + data = HFile(filename, "r") + except OSError as e: + logging.exception("Cannot open %s", filename) raise e ifo_data = data[ifo] - starts = numpy.append(starts, ifo_data['search/start_time'][:]) - ends = numpy.append(ends, ifo_data['search/end_time'][:]) - - if 'templates_per_core' in ifo_data['search'].keys(): - tpc = numpy.append(tpc, ifo_data['search/templates_per_core'][:]) - if 'filter_rate_per_core' in ifo_data['search'].keys(): - frpc = numpy.append(frpc, ifo_data['search/filter_rate_per_core'][:]) - if 'setup_time_fraction' in ifo_data['search'].keys(): - stf = numpy.append(stf, ifo_data['search/setup_time_fraction'][:]) - if 'run_time' in ifo_data['search'].keys(): - rtime = numpy.append(rtime, ifo_data['search/run_time'][:]) - - if 'gating' in ifo_data: + starts = numpy.append(starts, ifo_data["search/start_time"][:]) + ends = numpy.append(ends, ifo_data["search/end_time"][:]) + + if "templates_per_core" in ifo_data["search"].keys(): + tpc = numpy.append(tpc, ifo_data["search/templates_per_core"][:]) + if "filter_rate_per_core" in ifo_data["search"].keys(): + frpc = numpy.append(frpc, ifo_data["search/filter_rate_per_core"][:]) + if "setup_time_fraction" in ifo_data["search"].keys(): + stf = numpy.append(stf, ifo_data["search/setup_time_fraction"][:]) + if "run_time" in ifo_data["search"].keys(): + rtime = numpy.append(rtime, ifo_data["search/run_time"][:]) + + if "gating" in ifo_data: gating_keys = [] - ifo_data['gating'].visit(gating_keys.append) + ifo_data["gating"].visit(gating_keys.append) for gk in gating_keys: - gk_data = ifo_data['gating/' + gk] + gk_data = ifo_data["gating/" + gk] if isinstance(gk_data, h5py.Dataset): - if not gk in gating: + if gk not in gating: gating[gk] = numpy.array([], dtype=numpy.float64) gating[gk] = numpy.append(gating[gk], gk_data[:]) data.close() @@ -115,57 +127,66 @@ starts, uindex = numpy.unique(starts, return_index=True) ends = ends[uindex] sort = starts.argsort() -f['%s/search/start_time' % ifo] = starts[sort] -f['%s/search/end_time' % ifo] = ends[sort] +f["%s/search/start_time" % ifo] = starts[sort] +f["%s/search/end_time" % ifo] = ends[sort] if len(tpc) > 0: - f['%s/search/templates_per_core' % ifo] = tpc + f["%s/search/templates_per_core" % ifo] = tpc if len(frpc) > 0: - f['%s/search/filter_rate_per_core' % ifo] = frpc + f["%s/search/filter_rate_per_core" % ifo] = frpc if len(stf) > 0: - f['%s/search/setup_time_fraction' % ifo] = stf + f["%s/search/setup_time_fraction" % ifo] = stf if len(rtime) > 0: - f['%s/search/run_time' % ifo] = rtime + f["%s/search/run_time" % ifo] = rtime for gk, gv in gating.items(): - f[ifo + '/gating/' + gk] = gv + f[ifo + "/gating/" + gk] = gv -logging.info('set up sorting of triggers and template ids') +logging.info("set up sorting of triggers and template ids") # For fast lookup we need the templates in hash order -hashes = HFile(args.bank_file, 'r')['template_hash'][:] +hashes = HFile(args.bank_file, "r")["template_hash"][:] bank_tids = hashes.argsort() unsort = bank_tids.argsort() hashes = hashes[bank_tids] -trigger_hashes = collect('%s/template_hash' % ifo, args.trigger_files) +trigger_hashes = collect("%s/template_hash" % ifo, args.trigger_files) trigger_sort = trigger_hashes.argsort() trigger_hashes = trigger_hashes[trigger_sort] template_boundaries = changes(trigger_hashes) -template_ids = bank_tids[numpy.searchsorted(hashes, trigger_hashes[template_boundaries[:-1]])] +template_ids = bank_tids[ + numpy.searchsorted(hashes, trigger_hashes[template_boundaries[:-1]]) +] full_boundaries = numpy.searchsorted(trigger_hashes, hashes) full_boundaries = numpy.concatenate([full_boundaries, [len(trigger_hashes)]]) -# get the full boundaries in hash order +# get the full boundaries in hash order del trigger_hashes -idlen = (template_boundaries[1:] - template_boundaries[:-1]) -f.create_dataset('%s/template_id' % ifo, - data=numpy.repeat(template_ids, idlen), - compression='gzip', shuffle=True, - compression_opts=args.compression_level) -f['%s/template_boundaries' % ifo] = full_boundaries[unsort] - -logging.info('reading the trigger columns from the input files') +idlen = template_boundaries[1:] - template_boundaries[:-1] +f.create_dataset( + "%s/template_id" % ifo, + data=numpy.repeat(template_ids, idlen), + compression="gzip", + shuffle=True, + compression_opts=args.compression_level, +) +f["%s/template_boundaries" % ifo] = full_boundaries[unsort] + +logging.info("reading the trigger columns from the input files") for col in trigger_columns: - key = '%s/%s' % (ifo, col) - logging.info('reading %s', col) + key = "%s/%s" % (ifo, col) + logging.info("reading %s", col) data = collect(key, args.trigger_files)[trigger_sort] - logging.info('writing %s to file', col) - dset = f.create_dataset(key, data=data, compression='gzip', - compression_opts=args.compression_level, - shuffle=True) + logging.info("writing %s to file", col) + dset = f.create_dataset( + key, + data=data, + compression="gzip", + compression_opts=args.compression_level, + shuffle=True, + ) del data - region(f, key, full_boundaries, unsort) + region(f, key, full_boundaries, unsort) f.close() -logging.info('done') +logging.info("done") diff --git a/bin/all_sky_search/pycbc_coinc_statmap b/bin/all_sky_search/pycbc_coinc_statmap index 83f5061ab49..aea6d316e4b 100755 --- a/bin/all_sky_search/pycbc_coinc_statmap +++ b/bin/all_sky_search/pycbc_coinc_statmap @@ -7,24 +7,35 @@ the capability of doing hierarchical removal of foreground triggers that are louder than all of the background triggers. We use this to properly assess the FANs of any other gravitational waves in the dataset. """ + import argparse -import logging, numpy -from pycbc.events import veto, coinc, significance -import pycbc.pnutils, pycbc.io +import logging import sys + +import numpy + import pycbc.conversions as conv +import pycbc.io +import pycbc.pnutils +from pycbc.events import coinc, significance, veto -class fw(object): + +class fw: def __init__(self, name): - self.f = pycbc.io.HFile(name, 'w') + self.f = pycbc.io.HFile(name, "w") self.attrs = self.f.attrs def __setitem__(self, name, data): # Make a new item if isn't in the hdf file - if not name in self.f: - self.f.create_dataset(name, data=data, compression="gzip", - compression_opts=9, shuffle=True, - maxshape=data.shape) + if name not in self.f: + self.f.create_dataset( + name, + data=data, + compression="gzip", + compression_opts=9, + shuffle=True, + maxshape=data.shape, + ) # Else reassign values else: self.f[name][:] = data @@ -32,64 +43,92 @@ class fw(object): def __getitem__(self, *args): return self.f.__getitem__(*args) + parser = argparse.ArgumentParser() # General required options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--coinc-files', nargs='+', - help='List of coincidence files used to calculate the ' - 'FAP, FAR, etc.') -parser.add_argument('--ifos', nargs='+', - help='List of ifos used in these coincidence files') -parser.add_argument('--cluster-window', type=float, default=10, - help='Length of time window in seconds to cluster coinc ' - 'events [default=10s]') -parser.add_argument('--veto-window', type=float, default=.1, - help='Time around each zerolag trigger to window out ' - '[default=.1s]') -parser.add_argument('--hierarchical-removal-window', type=float, default=1.0, - help='Time around each trigger to window out for a very ' - 'louder trigger in the hierarchical removal ' - 'procedure [default=1.0s]') -parser.add_argument('--max-hierarchical-removal', type=int, default=0, - help='Maximum number of hierarchical removals to carry ' - 'out. Choose -1 for unlimited hierarchical removal ' - 'until no foreground triggers are louder than the ' - '(inclusive/exclusive) background. Choose 0 to do ' - 'no hierarchical removals, choose 1 to do at most ' - '1 hierarchical removal and so on. If given, must ' - 'also provide --hierarchical-removal-against to ' - 'indicate which background to remove triggers ' - 'against. [default=0]') -parser.add_argument('--hierarchical-removal-against', type=str, - default='none', choices=['none', 'inclusive', 'exclusive'], - help='If doing hierarchical removal, remove foreground ' - 'triggers that are louder than either the "inclusive"' - ' (little-dogs-in) background, or the "exclusive" ' - '(little-dogs-out) background. [default="none"]') -parser.add_argument('--additional-event-times', type=float, nargs="+", - help="Additional event times which will be removed from " - "both inclusive and exclusive background " - "(gps seconds)") +parser.add_argument( + "--coinc-files", + nargs="+", + help="List of coincidence files used to calculate the FAP, FAR, etc.", +) +parser.add_argument( + "--ifos", nargs="+", help="List of ifos used in these coincidence files" +) +parser.add_argument( + "--cluster-window", + type=float, + default=10, + help="Length of time window in seconds to cluster coinc events [default=10s]", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) +parser.add_argument( + "--hierarchical-removal-window", + type=float, + default=1.0, + help="Time around each trigger to window out for a very " + "louder trigger in the hierarchical removal " + "procedure [default=1.0s]", +) +parser.add_argument( + "--max-hierarchical-removal", + type=int, + default=0, + help="Maximum number of hierarchical removals to carry " + "out. Choose -1 for unlimited hierarchical removal " + "until no foreground triggers are louder than the " + "(inclusive/exclusive) background. Choose 0 to do " + "no hierarchical removals, choose 1 to do at most " + "1 hierarchical removal and so on. If given, must " + "also provide --hierarchical-removal-against to " + "indicate which background to remove triggers " + "against. [default=0]", +) +parser.add_argument( + "--hierarchical-removal-against", + type=str, + default="none", + choices=["none", "inclusive", "exclusive"], + help="If doing hierarchical removal, remove foreground " + 'triggers that are louder than either the "inclusive"' + ' (little-dogs-in) background, or the "exclusive" ' + '(little-dogs-out) background. [default="none"]', +) +parser.add_argument( + "--additional-event-times", + type=float, + nargs="+", + help="Additional event times which will be removed from " + "both inclusive and exclusive background " + "(gps seconds)", +) significance.insert_significance_option_group(parser) -parser.add_argument('--output-file') +parser.add_argument("--output-file") args = parser.parse_args() significance.check_significance_options(args, parser) # Check that the user chose inclusive or exclusive background to perform # hierarchical removals of foreground triggers against. if args.max_hierarchical_removal == 0: - if args.hierarchical_removal_against != 'none': - parser.error("User Error: 0 maximum hierarchical removals chosen but " - "option for --hierarchical-removal-against was given. " - "These are conflicting options. Use with --help for more " - "information.") -else : - if args.hierarchical_removal_against == 'none': - parser.error("--max-hierarchical-removal requires a choice of which " - "background to remove foreground triggers against, " - "inclusive or exclusive. Use with --help for more " - "information.") - + if args.hierarchical_removal_against != "none": + parser.error( + "User Error: 0 maximum hierarchical removals chosen but " + "option for --hierarchical-removal-against was given. " + "These are conflicting options. Use with --help for more " + "information." + ) +elif args.hierarchical_removal_against == "none": + parser.error( + "--max-hierarchical-removal requires a choice of which " + "background to remove foreground triggers against, " + "inclusive or exclusive. Use with --help for more " + "information." + ) pycbc.init_logging(args.verbose) @@ -97,29 +136,32 @@ pycbc.init_logging(args.verbose) logging.info("Loading coinc triggers") logging.info("IFO input: %s" % args.ifos) all_trigs = pycbc.io.MultiifoStatmapData(files=args.coinc_files, ifos=args.ifos) -if 'ifos' in all_trigs.attrs: - ifos = all_trigs.attrs['ifos'].split(' ') - logging.info('using ifos from file {}'.format(args.coinc_files[0])) +if "ifos" in all_trigs.attrs: + ifos = all_trigs.attrs["ifos"].split(" ") + logging.info(f"using ifos from file {args.coinc_files[0]}") else: ifos = args.ifos - logging.info('using ifos from command line input') + logging.info("using ifos from command line input") -ifo_combo = ''.join(ifos) +ifo_combo = "".join(ifos) significance_dict = significance.digest_significance_options([ifo_combo], args) logging.info("We have %s triggers" % len(all_trigs.stat)) # Remove triggers from manually flagged times: if args.additional_event_times: - logging.info("Removing triggers around times %s", - " ".join([str(et) for et in args.additional_event_times])) + logging.info( + "Removing triggers around times %s", + " ".join([str(et) for et in args.additional_event_times]), + ) rm_cent_times = numpy.array(args.additional_event_times) rm_start_times = rm_cent_times - args.veto_window rm_end_times = rm_cent_times + args.veto_window manual_rm_idx = [] for ifo in ifos: - rm_idx = veto.indices_within_times(all_trigs.data['%s/time' % ifo], - rm_start_times, rm_end_times) + rm_idx = veto.indices_within_times( + all_trigs.data["%s/time" % ifo], rm_start_times, rm_end_times + ) manual_rm_idx += list(rm_idx) logging.info("Removing %d triggers", len(manual_rm_idx)) all_trigs = all_trigs.remove(manual_rm_idx) @@ -130,7 +172,7 @@ fore_locs = all_trigs.timeslide_id == 0 # Foreground trigger times for ifos fore_time = {} for ifo in ifos: - fore_time[ifo] = all_trigs.data['%s/time' % ifo][fore_locs] + fore_time[ifo] = all_trigs.data["%s/time" % ifo][fore_locs] # Average times of triggers (note that coincs where not all ifos have triggers # will contain -1 sentinel values) fore_time_zip = zip(*fore_time.values()) @@ -144,16 +186,17 @@ remove_start_time = ave_fore_time - args.veto_window remove_end_time = ave_fore_time + args.veto_window # Total amount of time removed around foreground triggers -veto_time = abs(veto.start_end_to_segments(remove_start_time, - remove_end_time).coalesce()) +veto_time = abs( + veto.start_end_to_segments(remove_start_time, remove_end_time).coalesce() +) # Veto indices from list of triggers in the windowed times around fg triggers # This gives exclusive background triggers exc_zero_trigs = all_trigs.remove([]) # Start by copying existing triggers for ifo in ifos: fg_veto_ids = veto.indices_within_times( - exc_zero_trigs.data['%s/time' % ifo], - remove_start_time, remove_end_time) + exc_zero_trigs.data["%s/time" % ifo], remove_start_time, remove_end_time + ) exc_zero_trigs = exc_zero_trigs.remove(fg_veto_ids) logging.info("Clustering coinc triggers (inclusive of zerolag)") @@ -168,68 +211,68 @@ exc_zero_trigs = exc_zero_trigs.cluster(args.cluster_window) logging.info("Dumping foreground triggers") f = fw(args.output_file) -f.attrs['num_of_ifos'] = all_trigs.attrs['num_of_ifos'] -f.attrs['pivot'] = all_trigs.attrs['pivot'] -f.attrs['fixed'] = all_trigs.attrs['fixed'] -if 'ifos' in all_trigs.attrs: - f.attrs['ifos'] = all_trigs.attrs['ifos'] +f.attrs["num_of_ifos"] = all_trigs.attrs["num_of_ifos"] +f.attrs["pivot"] = all_trigs.attrs["pivot"] +f.attrs["fixed"] = all_trigs.attrs["fixed"] +if "ifos" in all_trigs.attrs: + f.attrs["ifos"] = all_trigs.attrs["ifos"] else: - f.attrs['ifos'] = ' '.join(sorted(args.ifos)) + f.attrs["ifos"] = " ".join(sorted(args.ifos)) -f.attrs['timeslide_interval'] = all_trigs.attrs['timeslide_interval'] +f.attrs["timeslide_interval"] = all_trigs.attrs["timeslide_interval"] # Copy over the segment for coincs and singles for key in all_trigs.seg.keys(): - f['segments/%s/start' % key] = all_trigs.seg[key]['start'][:] - f['segments/%s/end' % key] = all_trigs.seg[key]['end'][:] + f["segments/%s/start" % key] = all_trigs.seg[key]["start"][:] + f["segments/%s/end" % key] = all_trigs.seg[key]["end"][:] if fore_locs.sum() > 0: - f['segments/foreground_veto/start'] = remove_start_time - f['segments/foreground_veto/end'] = remove_end_time + f["segments/foreground_veto/start"] = remove_start_time + f["segments/foreground_veto/end"] = remove_end_time for k in all_trigs.data: - f['foreground/' + k] = all_trigs.data[k][fore_locs] + f["foreground/" + k] = all_trigs.data[k][fore_locs] else: # Put SOMETHING in here to avoid failures later - f['segments/foreground_veto/start'] = numpy.array([0]) - f['segments/foreground_veto/end'] = numpy.array([0]) + f["segments/foreground_veto/start"] = numpy.array([0]) + f["segments/foreground_veto/end"] = numpy.array([0]) for k in all_trigs.data: - f['foreground/' + k] = numpy.array([], dtype=all_trigs.data[k].dtype) + f["foreground/" + k] = numpy.array([], dtype=all_trigs.data[k].dtype) # If a particular index of all_trigs.timeslide_id isn't 0, evaluate true. # List of locations that is background. back_locs = all_trigs.timeslide_id != 0 if (back_locs.sum()) == 0: - logging.warning("There were no background events, so we could not " - "assign any statistic values") + logging.warning( + "There were no background events, so we could not assign any statistic values" + ) sys.exit() logging.info("Dumping background triggers (inclusive of zerolag)") for k in all_trigs.data: - f['background/' + k] = all_trigs.data[k][back_locs] + f["background/" + k] = all_trigs.data[k][back_locs] logging.info("Dumping background triggers (exclusive of zerolag)") for k in exc_zero_trigs.data: - f['background_exc/' + k] = exc_zero_trigs.data[k] + f["background_exc/" + k] = exc_zero_trigs.data[k] -maxtime = all_trigs.attrs['%s_foreground_time' % f.attrs['pivot']] +maxtime = all_trigs.attrs["%s_foreground_time" % f.attrs["pivot"]] for ifo in ifos: - if all_trigs.attrs['%s_foreground_time' % ifo] > maxtime: - maxtime = all_trigs.attrs['%s_foreground_time' % ifo] + maxtime = max(maxtime, all_trigs.attrs["%s_foreground_time" % ifo]) -mintime = all_trigs.attrs['%s_foreground_time' % f.attrs['pivot']] +mintime = all_trigs.attrs["%s_foreground_time" % f.attrs["pivot"]] for ifo in ifos: - if all_trigs.attrs['%s_foreground_time' % ifo] < mintime: - mintime = all_trigs.attrs['%s_foreground_time' % ifo] + mintime = min(mintime, all_trigs.attrs["%s_foreground_time" % ifo]) maxtime_exc = maxtime - veto_time mintime_exc = mintime - veto_time -background_time = int(maxtime / all_trigs.attrs['timeslide_interval']) * mintime -coinc_time = float(all_trigs.attrs['coinc_time']) +background_time = int(maxtime / all_trigs.attrs["timeslide_interval"]) * mintime +coinc_time = float(all_trigs.attrs["coinc_time"]) -background_time_exc = int(maxtime_exc / all_trigs.attrs['timeslide_interval']) \ - * mintime_exc +background_time_exc = ( + int(maxtime_exc / all_trigs.attrs["timeslide_interval"]) * mintime_exc +) coinc_time_exc = coinc_time - veto_time logging.info("Calculating FAN from background statistic values") @@ -244,7 +287,8 @@ bg_far, fg_far, sig_info = significance.get_far( fore_stat, all_trigs.decimation_factor[back_locs], background_time, - **significance_dict[ifo_combo]) + **significance_dict[ifo_combo], +) # Cumulative array of exclusive background triggers and the number # of exclusive background triggers louder than each foreground trigger @@ -253,55 +297,48 @@ bg_far_exc, fg_far_exc, exc_sig_info = significance.get_far( fore_stat, exc_zero_trigs.decimation_factor, background_time_exc, - **significance_dict[ifo_combo]) - -fg_far = significance.apply_far_limit( - fg_far, - significance_dict, - combo=ifo_combo) -bg_far = significance.apply_far_limit( - bg_far, - significance_dict, - combo=ifo_combo) + **significance_dict[ifo_combo], +) + +fg_far = significance.apply_far_limit(fg_far, significance_dict, combo=ifo_combo) +bg_far = significance.apply_far_limit(bg_far, significance_dict, combo=ifo_combo) fg_far_exc = significance.apply_far_limit( - fg_far_exc, - significance_dict, - combo=ifo_combo) + fg_far_exc, significance_dict, combo=ifo_combo +) bg_far_exc = significance.apply_far_limit( - bg_far_exc, - significance_dict, - combo=ifo_combo) + bg_far_exc, significance_dict, combo=ifo_combo +) -f['background/ifar'] = conv.sec_to_year(1. / bg_far) -f['background_exc/ifar'] = conv.sec_to_year(1. / bg_far_exc) -f.attrs['background_time'] = background_time -f.attrs['foreground_time'] = coinc_time -f.attrs['background_time_exc'] = background_time_exc -f.attrs['foreground_time_exc'] = coinc_time_exc +f["background/ifar"] = conv.sec_to_year(1.0 / bg_far) +f["background_exc/ifar"] = conv.sec_to_year(1.0 / bg_far_exc) +f.attrs["background_time"] = background_time +f.attrs["foreground_time"] = coinc_time +f.attrs["background_time_exc"] = background_time_exc +f.attrs["foreground_time_exc"] = coinc_time_exc logging.info("calculating ifar/fap values") if fore_locs.sum() > 0: - ifar = 1. / fg_far - fap = 1 - numpy.exp(- coinc_time / ifar) - f['foreground/ifar'] = conv.sec_to_year(ifar) - f['foreground/fap'] = fap + ifar = 1.0 / fg_far + fap = 1 - numpy.exp(-coinc_time / ifar) + f["foreground/ifar"] = conv.sec_to_year(ifar) + f["foreground/fap"] = fap for key, value in sig_info.items(): - f['foreground'].attrs[key] = value - ifar_exc = 1. / fg_far_exc - fap_exc = 1 - numpy.exp(- coinc_time_exc / ifar_exc) - f['foreground/ifar_exc'] = conv.sec_to_year(ifar_exc) - f['foreground/fap_exc'] = fap_exc + f["foreground"].attrs[key] = value + ifar_exc = 1.0 / fg_far_exc + fap_exc = 1 - numpy.exp(-coinc_time_exc / ifar_exc) + f["foreground/ifar_exc"] = conv.sec_to_year(ifar_exc) + f["foreground/fap_exc"] = fap_exc for key, value in exc_sig_info.items(): - f['foreground'].attrs[key + '_exc'] = value + f["foreground"].attrs[key + "_exc"] = value else: - f['foreground/ifar'] = ifar = numpy.array([]) - f['foreground/fap'] = numpy.array([]) - f['foreground/ifar_exc'] = ifar_exc = numpy.array([]) - f['foreground/fap_exc'] = numpy.array([]) + f["foreground/ifar"] = ifar = numpy.array([]) + f["foreground/fap"] = numpy.array([]) + f["foreground/ifar_exc"] = ifar_exc = numpy.array([]) + f["foreground/fap_exc"] = numpy.array([]) -if 'name' in all_trigs.attrs: - f.attrs['name'] = all_trigs.attrs['name'] +if "name" in all_trigs.attrs: + f.attrs["name"] = all_trigs.attrs["name"] # Incorporate hierarchical removal for any other loud triggers logging.info("Beginning hierarchical removal of foreground triggers") @@ -320,12 +357,12 @@ orig_fore_stat = fore_stat # or exclusive background. if args.max_hierarchical_removal != 0: # If user wants to remove against inclusive background. - if args.hierarchical_removal_against == 'inclusive': + if args.hierarchical_removal_against == "inclusive": ifar_foreground = ifar # Otherwise user wants to remove against exclusive background - else : + else: ifar_foreground = ifar_exc -else : +else: # It doesn't matter if you choose ifar_foreground = ifar # or ifar_exc, the while loop below will break # straight away. But this avoids a NameError @@ -342,24 +379,28 @@ ifar_foreground = numpy.append(ifar_foreground, 0) while numpy.any(ifar_foreground >= background_time): # If the user wants to stop doing hierarchical removals after a set # number of iterations then break when that happens. - if (h_iterations == args.max_hierarchical_removal): + if h_iterations == args.max_hierarchical_removal: break # Write foreground trigger info before hierarchical removals for # downstream codes. if h_iterations == 0: - f['background_h%s/stat' % h_iterations] = back_stat - f['background_h%s/ifar' % h_iterations] = conv.sec_to_year(1. / bg_far) - f['background_h%s/timeslide_id' % h_iterations] = all_trigs.data['timeslide_id'][back_locs] - f['foreground_h%s/stat' % h_iterations] = fore_stat - f['foreground_h%s/ifar' % h_iterations] = conv.sec_to_year(ifar) - f['foreground_h%s/fap' % h_iterations] = fap - f['foreground_h%s/template_id' % h_iterations] = all_trigs.data['template_id'][fore_locs] + f["background_h%s/stat" % h_iterations] = back_stat + f["background_h%s/ifar" % h_iterations] = conv.sec_to_year(1.0 / bg_far) + f["background_h%s/timeslide_id" % h_iterations] = all_trigs.data[ + "timeslide_id" + ][back_locs] + f["foreground_h%s/stat" % h_iterations] = fore_stat + f["foreground_h%s/ifar" % h_iterations] = conv.sec_to_year(ifar) + f["foreground_h%s/fap" % h_iterations] = fap + f["foreground_h%s/template_id" % h_iterations] = all_trigs.data["template_id"][ + fore_locs + ] for ifo in args.ifos: - trig_id = all_trigs.data['%s/trigger_id' % ifo][fore_locs] - trig_time = all_trigs.data['%s/time' % ifo][fore_locs] - f['foreground_h%s/%s/time' % (h_iterations,ifo)] = trig_time - f['foreground_h%s/%s/trigger_id' % (h_iterations,ifo)] = trig_id + trig_id = all_trigs.data["%s/trigger_id" % ifo][fore_locs] + trig_time = all_trigs.data["%s/time" % ifo][fore_locs] + f["foreground_h%s/%s/time" % (h_iterations, ifo)] = trig_time + f["foreground_h%s/%s/trigger_id" % (h_iterations, ifo)] = trig_id # Add the iteration number of hierarchical removals done. h_iterations += 1 @@ -377,28 +418,34 @@ while numpy.any(ifar_foreground >= background_time): # Store any foreground trigger's information that we want to # hierarchically remove. - f['foreground/ifar'][orig_fore_idx] = conv.sec_to_year(ifar[max_stat_idx]) - f['foreground/fap'][orig_fore_idx] = fap[max_stat_idx] + f["foreground/ifar"][orig_fore_idx] = conv.sec_to_year(ifar[max_stat_idx]) + f["foreground/fap"][orig_fore_idx] = fap[max_stat_idx] - logging.info("Removing foreground trigger that is louder than the inclusive background.") + logging.info( + "Removing foreground trigger that is louder than the inclusive background." + ) # Remove the foreground trigger and all of the background triggers that # are associated with it. ave_rm_time = 0 for ifo in args.ifos: - ave_rm_time += all_trigs.data['%s/time' % ifo][rm_trig_idx] / len(args.ifos) + ave_rm_time += all_trigs.data["%s/time" % ifo][rm_trig_idx] / len(args.ifos) ind_to_rm = {} for ifo in args.ifos: - ind_to_rm[ifo] = veto.indices_within_times(all_trigs.data['%s/time' % ifo], - [ave_rm_time - args.hierarchical_removal_window], - [ave_rm_time + args.hierarchical_removal_window]) + ind_to_rm[ifo] = veto.indices_within_times( + all_trigs.data["%s/time" % ifo], + [ave_rm_time - args.hierarchical_removal_window], + [ave_rm_time + args.hierarchical_removal_window], + ) indices_to_rm = [] for ifo in args.ifos: indices_to_rm = numpy.concatenate([indices_to_rm, ind_to_rm[ifo]]) all_trigs = all_trigs.remove(indices_to_rm.astype(int)) - logging.info("We have %s triggers after hierarchical removal." % len(all_trigs.stat)) + logging.info( + "We have %s triggers after hierarchical removal." % len(all_trigs.stat) + ) # Step 4: Re-cluster the triggers and calculate the inclusive ifar/fap logging.info("Clustering coinc triggers (inclusive of zerolag)") @@ -413,20 +460,18 @@ while numpy.any(ifar_foreground >= background_time): logging.info("Dumping foreground triggers") logging.info("Dumping background triggers (inclusive of zerolag)") for k in all_trigs.data: - f['background_h%s/' % h_iterations + k] = all_trigs.data[k][back_locs] + f["background_h%s/" % h_iterations + k] = all_trigs.data[k][back_locs] - maxtime = all_trigs.attrs['%s_foreground_time' % f.attrs['pivot']] + maxtime = all_trigs.attrs["%s_foreground_time" % f.attrs["pivot"]] for ifo in ifos: - if all_trigs.attrs['%s_foreground_time' % ifo] > maxtime: - maxtime = all_trigs.attrs['%s_foreground_time' % ifo] + maxtime = max(maxtime, all_trigs.attrs["%s_foreground_time" % ifo]) - mintime = all_trigs.attrs['%s_foreground_time' % f.attrs['pivot']] + mintime = all_trigs.attrs["%s_foreground_time" % f.attrs["pivot"]] for ifo in ifos: - if all_trigs.attrs['%s_foreground_time' % ifo] < mintime: - mintime = all_trigs.attrs['%s_foreground_time' % ifo] + mintime = min(mintime, all_trigs.attrs["%s_foreground_time" % ifo]) - background_time = int(maxtime / all_trigs.attrs['timeslide_interval']) * mintime - coinc_time = float(all_trigs.attrs['coinc_time']) + background_time = int(maxtime / all_trigs.attrs["timeslide_interval"]) * mintime + coinc_time = float(all_trigs.attrs["coinc_time"]) logging.info("Calculating FAN from background statistic values") back_stat = all_trigs.stat[back_locs] @@ -437,13 +482,10 @@ while numpy.any(ifar_foreground >= background_time): all_trigs.decimation_factor[back_locs], background_time, return_counts=True, - **significance_dict[ifo_combo]) - - fg_far = significance.apply_far_limit( - fg_far, - significance_dict, - combo=ifo_combo + **significance_dict[ifo_combo], ) + + fg_far = significance.apply_far_limit(fg_far, significance_dict, combo=ifo_combo) bg_far = significance.apply_far_limit( bg_far, significance_dict, @@ -451,76 +493,77 @@ while numpy.any(ifar_foreground >= background_time): ) # Update the ifar_foreground criteria depending on whether foreground - # triggers are being removed via inclusive or exclusive background. - if args.hierarchical_removal_against == 'inclusive': - ifar_foreground = 1. / fg_far + # triggers are being removed via inclusive or exclusive background. + if args.hierarchical_removal_against == "inclusive": + ifar_foreground = 1.0 / fg_far # Exclusive background doesn't change when removing foreground triggers. # So we don't have to take background ifar, just repopulate ifar_foreground - else : + else: _, fg_far_exc, _ = significance.get_far( exc_zero_trigs.stat, fore_stat, exc_zero_trigs.decimation_factor, background_time_exc, - **significance_dict[ifo_combo]) + **significance_dict[ifo_combo], + ) fg_far_exc = significance.apply_far_limit( - fg_far_exc, - significance_dict, - combo=ifo_combo + fg_far_exc, significance_dict, combo=ifo_combo ) - ifar_foreground = 1. / fg_far_exc + ifar_foreground = 1.0 / fg_far_exc # ifar_foreground has been updated and the code can continue. logging.info("Calculating ifar/fap values") - f['background_h%s/ifar' % h_iterations] = conv.sec_to_year(1. / bg_far) - f.attrs['background_time_h%s' % h_iterations] = background_time - f.attrs['foreground_time_h%s' % h_iterations] = coinc_time + f["background_h%s/ifar" % h_iterations] = conv.sec_to_year(1.0 / bg_far) + f.attrs["background_time_h%s" % h_iterations] = background_time + f.attrs["foreground_time_h%s" % h_iterations] = coinc_time if fore_locs.sum() > 0: # Write ranking statistic to file just for downstream plotting code - f['foreground_h%s/stat' % h_iterations] = fore_stat + f["foreground_h%s/stat" % h_iterations] = fore_stat - ifar = 1. / fg_far - fap = 1 - numpy.exp(- coinc_time / ifar) - f['foreground_h%s/ifar' % h_iterations] = conv.sec_to_year(ifar) - f['foreground_h%s/fap' % h_iterations] = fap + ifar = 1.0 / fg_far + fap = 1 - numpy.exp(-coinc_time / ifar) + f["foreground_h%s/ifar" % h_iterations] = conv.sec_to_year(ifar) + f["foreground_h%s/fap" % h_iterations] = fap for key, value in sig_info.items(): - f['foreground_h%' % h_iterations].attrs[key] = value + f["foreground_h%" % h_iterations].attrs[key] = value # Update ifar and fap for other foreground triggers for i in range(len(ifar)): orig_fore_idx = numpy.where(orig_fore_stat == fore_stat[i])[0][0] - f['foreground/ifar'][orig_fore_idx] = conv.sec_to_year(ifar[i]) - f['foreground/fap'][orig_fore_idx] = fap[i] + f["foreground/ifar"][orig_fore_idx] = conv.sec_to_year(ifar[i]) + f["foreground/fap"][orig_fore_idx] = fap[i] # Save trigger ids for foreground triggers for downstream plotting code. # These don't change with the iterations but should be written at every # level. - f['foreground_h%s/template_id' % h_iterations] = all_trigs.data['template_id'][fore_locs] + f["foreground_h%s/template_id" % h_iterations] = all_trigs.data["template_id"][ + fore_locs + ] for ifo in args.ifos: - trig_id = all_trigs.data['%s/trigger_id' % ifo][fore_locs] - trig_time = all_trigs.data['%s/time' % ifo][fore_locs] - f['foreground_h%s/%s/time' % (h_iterations,ifo)] = trig_time - f['foreground_h%s/%s/trigger_id' % (h_iterations,ifo)] = trig_id - else : - f['foreground_h%s/stat' % h_iterations] = numpy.array([]) - f['foreground_h%s/ifar' % h_iterations] = numpy.array([]) - f['foreground_h%s/fap' % h_iterations] = numpy.array([]) - f['foreground_h%s/template_id' % h_iterations] = numpy.array([]) + trig_id = all_trigs.data["%s/trigger_id" % ifo][fore_locs] + trig_time = all_trigs.data["%s/time" % ifo][fore_locs] + f["foreground_h%s/%s/time" % (h_iterations, ifo)] = trig_time + f["foreground_h%s/%s/trigger_id" % (h_iterations, ifo)] = trig_id + else: + f["foreground_h%s/stat" % h_iterations] = numpy.array([]) + f["foreground_h%s/ifar" % h_iterations] = numpy.array([]) + f["foreground_h%s/fap" % h_iterations] = numpy.array([]) + f["foreground_h%s/template_id" % h_iterations] = numpy.array([]) for ifo in args.ifos: - f['foreground_h%s/%s/time' % (h_iterations,ifo)] = numpy.array([]) - f['foreground_h%s/%s/trigger_id' % (h_iterations,ifo)] = numpy.array([]) + f["foreground_h%s/%s/time" % (h_iterations, ifo)] = numpy.array([]) + f["foreground_h%s/%s/trigger_id" % (h_iterations, ifo)] = numpy.array([]) # Write to file how many hierarchical removals were implemented. -f.attrs['hierarchical_removal_iterations'] = h_iterations +f.attrs["hierarchical_removal_iterations"] = h_iterations # Write whether hierarchical removals were removed against the # inclusive background or the exclusive background. Have to use # numpy.bytes_ datatype. if h_iterations != 0: hrm_method = args.hierarchical_removal_against - f.attrs['hierarchical_removal_method'] = numpy.bytes_(hrm_method) + f.attrs["hierarchical_removal_method"] = numpy.bytes_(hrm_method) logging.info("Done") diff --git a/bin/all_sky_search/pycbc_coinc_statmap_inj b/bin/all_sky_search/pycbc_coinc_statmap_inj index a0486593946..8fd3047ec19 100644 --- a/bin/all_sky_search/pycbc_coinc_statmap_inj +++ b/bin/all_sky_search/pycbc_coinc_statmap_inj @@ -4,29 +4,46 @@ The program combines coincident output files generated by pycbc_coinc_findtrigs to generated a mapping between SNR and FAP, along with producing the combined foreground and background triggers """ -import argparse, logging, pycbc.io, numpy -from pycbc.events import significance + +import argparse +import logging + +import numpy + import pycbc.conversions as conv +import pycbc.io from pycbc import init_logging +from pycbc.events import significance parser = argparse.ArgumentParser() # General required options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--cluster-window', type=float, default=10, - help='Length of time window in seconds to cluster coinc ' - 'events [default=10s]') -parser.add_argument('--zero-lag-coincs', nargs='+', - help='Files containing the injection zerolag coincidences') -parser.add_argument('--full-data-background', - help='background file from full data for use in analyzing ' - 'injection coincs') -parser.add_argument('--veto-window', type=float, default=.1, - help='Time around each zerolag trigger to window out ' - '[default=.1s]') -parser.add_argument('--ifos', nargs='+', - help='List of ifos used in these coincidence files') +parser.add_argument( + "--cluster-window", + type=float, + default=10, + help="Length of time window in seconds to cluster coinc events [default=10s]", +) +parser.add_argument( + "--zero-lag-coincs", + nargs="+", + help="Files containing the injection zerolag coincidences", +) +parser.add_argument( + "--full-data-background", + help="background file from full data for use in analyzing injection coincs", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) +parser.add_argument( + "--ifos", nargs="+", help="List of ifos used in these coincidence files" +) significance.insert_significance_option_group(parser) -parser.add_argument('--output-file') +parser.add_argument("--output-file") args = parser.parse_args() init_logging(args.verbose) @@ -38,62 +55,58 @@ window = args.cluster_window logging.info("Loading coinc zerolag triggers") zdata = pycbc.io.MultiifoStatmapData(files=args.zero_lag_coincs, ifos=args.ifos) -if 'ifos' in zdata.attrs: - ifos = zdata.attrs['ifos'].split(' ') - logging.info('using ifos from file {}'.format(args.zero_lag_coincs[0])) +if "ifos" in zdata.attrs: + ifos = zdata.attrs["ifos"].split(" ") + logging.info(f"using ifos from file {args.zero_lag_coincs[0]}") else: ifos = args.ifos - logging.info('using ifos from command line input') + logging.info("using ifos from command line input") -ifo_key = ''.join(ifos) +ifo_key = "".join(ifos) significance_dict = significance.digest_significance_options([ifo_key], args) zdata = zdata.cluster(window) f = pycbc.io.HFile(args.output_file, "w") -f.attrs['num_of_ifos'] = zdata.attrs['num_of_ifos'] -f.attrs['pivot'] = zdata.attrs['pivot'] -f.attrs['fixed'] = zdata.attrs['fixed'] -f.attrs['timeslide_interval'] = zdata.attrs['timeslide_interval'] -f.attrs['ifos'] = ' '.join(sorted(ifos)) +f.attrs["num_of_ifos"] = zdata.attrs["num_of_ifos"] +f.attrs["pivot"] = zdata.attrs["pivot"] +f.attrs["fixed"] = zdata.attrs["fixed"] +f.attrs["timeslide_interval"] = zdata.attrs["timeslide_interval"] +f.attrs["ifos"] = " ".join(sorted(ifos)) # Copy over the segment for coincs and singles for key in zdata.seg.keys(): - f['segments/%s/start' % key] = zdata.seg[key]['start'][:] - f['segments/%s/end' % key] = zdata.seg[key]['end'][:] + f["segments/%s/start" % key] = zdata.seg[key]["start"][:] + f["segments/%s/end" % key] = zdata.seg[key]["end"][:] -logging.info('writing zero lag triggers') +logging.info("writing zero lag triggers") if len(zdata) > 0: for key in zdata.data: - f['foreground/%s' % key] = zdata.data[key] + f["foreground/%s" % key] = zdata.data[key] else: for key in zdata.data: - f['foreground/%s' % key] = numpy.array([], dtype=zdata.data[key].dtype) + f["foreground/%s" % key] = numpy.array([], dtype=zdata.data[key].dtype) -logging.info('calculating statistics excluding zerolag') +logging.info("calculating statistics excluding zerolag") fb = pycbc.io.HFile(args.full_data_background, "r") # we expect the injfull file to contain injection data as pivot # and fullinj to contain full data as pivot -background_time = float(fb.attrs['background_time']) -coinc_time = float(fb.attrs['foreground_time']) -back_stat = fb['background_exc/stat'][:] -dec_fac = fb['background_exc/decimation_factor'][:] +background_time = float(fb.attrs["background_time"]) +coinc_time = float(fb.attrs["foreground_time"]) +back_stat = fb["background_exc/stat"][:] +dec_fac = fb["background_exc/decimation_factor"][:] -f.attrs['background_time_exc'] = background_time -f.attrs['foreground_time_exc'] = coinc_time -f.attrs['background_time'] = background_time -f.attrs['foreground_time'] = coinc_time +f.attrs["background_time_exc"] = background_time +f.attrs["foreground_time_exc"] = coinc_time +f.attrs["background_time"] = background_time +f.attrs["foreground_time"] = coinc_time if len(zdata) > 0: - _, fg_far_exc, exc_sig_info = significance.get_far( - back_stat, - zdata.stat, - dec_fac, - background_time, - **significance_dict[ifo_key]) + back_stat, zdata.stat, dec_fac, background_time, **significance_dict[ifo_key] + ) fg_far_exc = significance.apply_far_limit( fg_far_exc, @@ -101,15 +114,15 @@ if len(zdata) > 0: combo=ifo_key, ) - ifar_exc = 1. / fg_far_exc - fap_exc = 1 - numpy.exp(- coinc_time / ifar_exc) - f['foreground/ifar_exc'] = conv.sec_to_year(ifar_exc) - f['foreground/fap_exc'] = fap_exc + ifar_exc = 1.0 / fg_far_exc + fap_exc = 1 - numpy.exp(-coinc_time / ifar_exc) + f["foreground/ifar_exc"] = conv.sec_to_year(ifar_exc) + f["foreground/fap_exc"] = fap_exc for key, value in exc_sig_info.items(): - f['foreground'].attrs[key + '_exc'] = value + f["foreground"].attrs[key + "_exc"] = value else: - f['foreground/ifar_exc'] = numpy.array([]) - f['foreground/fap_exc'] = numpy.array([]) + f["foreground/ifar_exc"] = numpy.array([]) + f["foreground/fap_exc"] = numpy.array([]) logging.info("Done") diff --git a/bin/all_sky_search/pycbc_combine_coincident_events b/bin/all_sky_search/pycbc_combine_coincident_events index f5ac146d476..062104cf817 100644 --- a/bin/all_sky_search/pycbc_combine_coincident_events +++ b/bin/all_sky_search/pycbc_combine_coincident_events @@ -3,17 +3,24 @@ This program combines together a set of STATMAP files from disjoint times. The resulting file would contain triggers from the full set of input files """ -import numpy + import argparse +import numpy + import pycbc from pycbc.io import HFile + def com(f, files, group): - """ Combine the same column from multiple files into another file f""" - f[group] = numpy.concatenate\ - ([fi[group][:] if group in fi else numpy.array([], dtype=numpy.uint32)\ - for fi in files]) + """Combine the same column from multiple files into another file f""" + f[group] = numpy.concatenate( + [ + fi[group][:] if group in fi else numpy.array([], dtype=numpy.uint32) + for fi in files + ] + ) + def com_with_detector_key(f, files, group): """ @@ -28,10 +35,10 @@ def com_with_detector_key(f, files, group): # It may not be the *same* detector for each file. # What detector are we dealing with - if group.endswith('1'): - ifo_name = f.attrs['detector_1'] - elif group.endswith('2'): - ifo_name = f.attrs['detector_2'] + if group.endswith("1"): + ifo_name = f.attrs["detector_1"] + elif group.endswith("2"): + ifo_name = f.attrs["detector_2"] else: raise ValueError("Group name must end in 1 or 2, got %s" % group) @@ -40,23 +47,24 @@ def com_with_detector_key(f, files, group): # For the remaining files we must check the detector name for nfp in files[1:]: # What detector do we need in this file - if nfp.attrs['detector_1'] == ifo_name: - new_det_num = '1' - elif nfp.attrs['detector_2'] == ifo_name: - new_det_num = '2' + if nfp.attrs["detector_1"] == ifo_name: + new_det_num = "1" + elif nfp.attrs["detector_2"] == ifo_name: + new_det_num = "2" else: raise ValueError("Cannot find detector %s in input file" % ifo_name) new_group = group[:-1] + new_det_num data_for_catting.append(nfp[new_group][:]) f[group] = numpy.concatenate(data_for_catting) - + parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--statmap-files', nargs='+', - help="List of coinc files to be redistributed") -parser.add_argument('--output-file', help="name of output file") +parser.add_argument( + "--statmap-files", nargs="+", help="List of coinc files to be redistributed" +) +parser.add_argument("--output-file", help="name of output file") args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -66,34 +74,40 @@ files = [HFile(n) for n in args.statmap_files] # Start setting some of the attributes f = HFile(args.output_file, "w") # It's not guaranteed that all files will follow this, so be careful later! -f.attrs['detector_1'] = files[0].attrs['detector_1'] -f.attrs['detector_2'] = files[0].attrs['detector_2'] - -f.attrs['background_time'] = \ - sum([cfp.attrs['background_time'] for cfp in files]) -f.attrs['foreground_time'] = \ - sum([cfp.attrs['foreground_time'] for cfp in files]) -f.attrs['background_time_exc'] = \ - sum([cfp.attrs['background_time_exc'] for cfp in files]) -f.attrs['foreground_time_exc'] = \ - sum([cfp.attrs['foreground_time_exc'] for cfp in files]) +f.attrs["detector_1"] = files[0].attrs["detector_1"] +f.attrs["detector_2"] = files[0].attrs["detector_2"] + +f.attrs["background_time"] = sum([cfp.attrs["background_time"] for cfp in files]) +f.attrs["foreground_time"] = sum([cfp.attrs["foreground_time"] for cfp in files]) +f.attrs["background_time_exc"] = sum( + [cfp.attrs["background_time_exc"] for cfp in files] +) +f.attrs["foreground_time_exc"] = sum( + [cfp.attrs["foreground_time_exc"] for cfp in files] +) # Combine segments -for key in files[0]['segments'].keys(): - com(f, files, 'segments/%s/start' % key) - com(f, files, 'segments/%s/end' % key) +for key in files[0]["segments"].keys(): + com(f, files, "segments/%s/start" % key) + com(f, files, "segments/%s/end" % key) # copy over all the columns in the foreground group. A few special cases here -for fg_bg_key in ['foreground', 'background', 'background_exc']: +for fg_bg_key in ["foreground", "background", "background_exc"]: for key in files[0][fg_bg_key].keys(): - if key not in ['time1', 'time2', 'trigger_id1', 'trigger_id2', - 'fap', 'fap_exc']: - com(f, files, '%s/%s' % (fg_bg_key,key)) - elif key in ['time1', 'time2', 'trigger_id1', 'trigger_id2']: + if key not in [ + "time1", + "time2", + "trigger_id1", + "trigger_id2", + "fap", + "fap_exc", + ]: + com(f, files, "%s/%s" % (fg_bg_key, key)) + elif key in ["time1", "time2", "trigger_id1", "trigger_id2"]: # Check if all files use the same detector convention - com_with_detector_key(f, files, '%s/%s' % (fg_bg_key,key)) + com_with_detector_key(f, files, "%s/%s" % (fg_bg_key, key)) else: # Do not store FAP numbers ... Could be recalculated. continue - + f.close() diff --git a/bin/all_sky_search/pycbc_combine_statmap b/bin/all_sky_search/pycbc_combine_statmap index 1c8928c313e..7a638920ec2 100755 --- a/bin/all_sky_search/pycbc_combine_statmap +++ b/bin/all_sky_search/pycbc_combine_statmap @@ -1,58 +1,77 @@ #!/bin/env python -""" Apply a trials factor based on the number of available detector +""" +Apply a trials factor based on the number of available detector combinations at the time of coincidence. This clusters to find the most significant foreground, but leaves the background triggers alone. """ -import numpy, argparse, logging, pycbc, pycbc.events, pycbc.io +import argparse +import logging + import igwn_segments as segments +import numpy + +import pycbc +import pycbc.events +import pycbc.io parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--statmap-files', nargs='+', - help="List of coinc files to be redistributed") -parser.add_argument('--cluster-window', type=float) -parser.add_argument('--censor-ifar-threshold', type=float, default=0.003, +parser.add_argument( + "--statmap-files", nargs="+", help="List of coinc files to be redistributed" +) +parser.add_argument("--cluster-window", type=float) +parser.add_argument( + "--censor-ifar-threshold", + type=float, + default=0.003, help="If provided, only window out foreground triggers with IFAR (years)" - "above the threshold [default=0.003yr]") -parser.add_argument('--veto-window', type=float, default=0.1, - help="Time around each zerolag trigger to window out [default=.1s]") -parser.add_argument('--output-file', help="name of output file") + "above the threshold [default=0.003yr]", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) +parser.add_argument("--output-file", help="name of output file") args = parser.parse_args() pycbc.init_logging(args.verbose) -files = [pycbc.io.HFile(n, 'r') for n in args.statmap_files] +files = [pycbc.io.HFile(n, "r") for n in args.statmap_files] f = pycbc.io.HFile(args.output_file, "w") -logging.info('Copying segments and attributes to %s' % args.output_file) +logging.info("Copying segments and attributes to %s" % args.output_file) # Move segments information into the final file - remove some duplication # in earlier files for fi in files: - for key in fi['segments']: - if key.startswith('foreground') or key.startswith('background'): + for key in fi["segments"]: + if key.startswith("foreground") or key.startswith("background"): continue - f['segments/%s/end' % key] = fi['segments/%s/end' % key][:] - f['segments/%s/start' % key] = fi['segments/%s/start' % key][:] - if 'segments/foreground_veto' in fi: - f['segments/%s/foreground_veto/end' % key] = \ - fi['segments/foreground_veto/end'][:] - f['segments/%s/foreground_veto/start' % key] = \ - fi['segments/foreground_veto/start'][:] + f["segments/%s/end" % key] = fi["segments/%s/end" % key][:] + f["segments/%s/start" % key] = fi["segments/%s/start" % key][:] + if "segments/foreground_veto" in fi: + f["segments/%s/foreground_veto/end" % key] = fi[ + "segments/foreground_veto/end" + ][:] + f["segments/%s/foreground_veto/start" % key] = fi[ + "segments/foreground_veto/start" + ][:] for attr_name in fi.attrs: if key not in f: - f.create_group(key) + f.create_group(key) f[key].attrs[attr_name] = fi.attrs[attr_name] # Set up dictionaries to contain segments from the individual statmap files indiv_segs = segments.segmentlistdict({}) # loop through statmap files and put segments into segmentlistdicts for fi in files: - key = fi.attrs['ifos'].replace(' ','') + key = fi.attrs["ifos"].replace(" ", "") # get analysed segments from individual statmap files - starts = fi['segments/{}/start'.format(key)][:] - ends = fi['segments/{}/end'.format(key)][:] + starts = fi[f"segments/{key}/start"][:] + ends = fi[f"segments/{key}/end"][:] indiv_segs[key] = pycbc.events.veto.start_end_to_segments(starts, ends) if len(indiv_segs.values()) == 1: @@ -63,24 +82,23 @@ else: foreground_segs = numpy.sum(list(indiv_segs.values()), axis=0) # Total zerolag analysis time -f.attrs['foreground_time'] = abs(foreground_segs) +f.attrs["foreground_time"] = abs(foreground_segs) # obtain list of all ifos involved in the coinc_statmap files -all_ifos = numpy.unique([ifo for fi in files for ifo in \ - fi.attrs['ifos'].split(' ')]) -f.attrs['ifos'] = ' '.join(sorted(all_ifos)) +all_ifos = numpy.unique([ifo for fi in files for ifo in fi.attrs["ifos"].split(" ")]) +f.attrs["ifos"] = " ".join(sorted(all_ifos)) -logging.info('Generating list of datasets in input files') +logging.info("Generating list of datasets in input files") key_set = pycbc.io.name_all_datasets(files) -logging.info('Copying foreground non-ifo-specific data') +logging.info("Copying foreground non-ifo-specific data") # copy and concatenate all the columns in the foreground group # from all files /except/ the IFO groups for key in key_set: - if key.startswith('foreground') and not any([ifo in key for ifo in all_ifos]): + if key.startswith("foreground") and not any([ifo in key for ifo in all_ifos]): pycbc.io.combine_and_copy(f, files, key) -logging.info('Collating triggers into single structure') +logging.info("Collating triggers into single structure") all_trig_times = {} all_trig_ids = {} @@ -92,41 +110,52 @@ for ifo in all_ifos: # If an ifo does not participate in any given coinc then fill with -1 values for f_in in files: for ifo in all_ifos: - if ifo in f_in['foreground']: - all_trig_times[ifo] = numpy.concatenate( [all_trig_times[ifo], - f_in['foreground/{}/time'.format(ifo)][:]] ) - all_trig_ids[ifo] = numpy.concatenate( [all_trig_ids[ifo], - f_in['foreground/{}/trigger_id'.format(ifo)][:]] ) + if ifo in f_in["foreground"]: + all_trig_times[ifo] = numpy.concatenate( + [all_trig_times[ifo], f_in[f"foreground/{ifo}/time"][:]] + ) + all_trig_ids[ifo] = numpy.concatenate( + [all_trig_ids[ifo], f_in[f"foreground/{ifo}/trigger_id"][:]] + ) else: - all_trig_times[ifo] = numpy.concatenate( [all_trig_times[ifo], - -1 * numpy.ones_like(f_in['foreground/fap'][:], - dtype=numpy.uint32)] ) - all_trig_ids[ifo] = numpy.concatenate( [all_trig_ids[ifo], - -1 * numpy.ones_like(f_in['foreground/fap'][:], - dtype=numpy.uint32)] ) + all_trig_times[ifo] = numpy.concatenate( + [ + all_trig_times[ifo], + -1 * numpy.ones_like(f_in["foreground/fap"][:], dtype=numpy.uint32), + ] + ) + all_trig_ids[ifo] = numpy.concatenate( + [ + all_trig_ids[ifo], + -1 * numpy.ones_like(f_in["foreground/fap"][:], dtype=numpy.uint32), + ] + ) f_in.close() -logging.info('Clustering triggers for loudest ifar value') +logging.info("Clustering triggers for loudest ifar value") for ifo in all_ifos: - f['foreground/{}/time'.format(ifo)] = all_trig_times[ifo] - f['foreground/{}/trigger_id'.format(ifo)] = all_trig_ids[ifo] + f[f"foreground/{ifo}/time"] = all_trig_times[ifo] + f[f"foreground/{ifo}/trigger_id"] = all_trig_ids[ifo] -ifar_stat = numpy.core.records.fromarrays([f['foreground/ifar'][:], - f['foreground/stat'][:]], - names='ifar,stat') +ifar_stat = numpy.core.records.fromarrays( + [f["foreground/ifar"][:], f["foreground/stat"][:]], names="ifar,stat" +) # all_times is a tuple of trigger time arrays -all_times = (f['foreground/%s/time' % ifo][:] for ifo in all_ifos) +all_times = (f["foreground/%s/time" % ifo][:] for ifo in all_ifos) + def argmax(v): return numpy.argsort(v)[-1] + # Currently only clustering zerolag, i.e. foreground, so set all timeslide_ids # to zero -cidx = pycbc.events.cluster_coincs_multiifo(ifar_stat, all_times, - numpy.zeros(len(ifar_stat)), 0, - args.cluster_window, argmax) +cidx = pycbc.events.cluster_coincs_multiifo( + ifar_stat, all_times, numpy.zeros(len(ifar_stat)), 0, args.cluster_window, argmax +) + def filter_dataset(h5file, name, idx): # Dataset needs to be deleted and remade as it is a different size @@ -135,57 +164,58 @@ def filter_dataset(h5file, name, idx): h5file[name] = filtered_dset return idx + # Downsample the foreground columns to only the loudest ifar between the # multiple files -for key in f['foreground'].keys(): +for key in f["foreground"].keys(): if key not in all_ifos: - id = filter_dataset(f, 'foreground/%s' % key, cidx) + id = filter_dataset(f, "foreground/%s" % key, cidx) else: # key is an ifo - for k in f['foreground/%s' % key].keys(): - id = filter_dataset(f, 'foreground/{}/{}'.format(key, k), cidx) + for k in f["foreground/%s" % key].keys(): + id = filter_dataset(f, f"foreground/{key}/{k}", cidx) -logging.info('Applying trials factor') +logging.info("Applying trials factor") # Recalculate event times after clustering for trials factor calculation -clustered_times = (f['foreground/%s/time' % ifo][:] for ifo in all_ifos) +clustered_times = (f["foreground/%s/time" % ifo][:] for ifo in all_ifos) # Trials factor is how many possible 2+IFO combinations are 'on' # at the time of the coincidence -trials_factors = numpy.zeros_like(f['foreground/ifar'][:]) -test_times = numpy.array([pycbc.events.mean_if_greater_than_zero(tc)[0] - for tc in zip(*clustered_times)]) +trials_factors = numpy.zeros_like(f["foreground/ifar"][:]) +test_times = numpy.array( + [pycbc.events.mean_if_greater_than_zero(tc)[0] for tc in zip(*clustered_times)] +) # Iterate over different ifo combinations -for key in f['segments']: - if key.startswith('foreground') or key.startswith('background'): +for key in f["segments"]: + if key.startswith("foreground") or key.startswith("background"): continue - end_times = numpy.array(f['segments/%s/end' % key][:]) - start_times = numpy.array(f['segments/%s/start' % key][:]) - idx_within_segment = pycbc.events.indices_within_times(test_times, - start_times, - end_times) + end_times = numpy.array(f["segments/%s/end" % key][:]) + start_times = numpy.array(f["segments/%s/start" % key][:]) + idx_within_segment = pycbc.events.indices_within_times( + test_times, start_times, end_times + ) trials_factors[idx_within_segment] += numpy.ones_like(idx_within_segment) -f['foreground/ifar'][:] = f['foreground/ifar'][:] / trials_factors -f['foreground/ifar_exc'][:] = f['foreground/ifar_exc'][:] / trials_factors +f["foreground/ifar"][:] = f["foreground/ifar"][:] / trials_factors +f["foreground/ifar_exc"][:] = f["foreground/ifar_exc"][:] / trials_factors -f.attrs['foreground_time_exc'] = f.attrs['foreground_time'] +f.attrs["foreground_time_exc"] = f.attrs["foreground_time"] # Construct the foreground censor veto from the clustered candidate times # above the ifar threshold -thr = test_times[f['foreground/ifar'][:] > args.censor_ifar_threshold] +thr = test_times[f["foreground/ifar"][:] > args.censor_ifar_threshold] vstart = thr - args.veto_window vend = thr + args.veto_window -vtime = segments.segmentlist([segments.segment(s, e) - for s, e in zip(vstart, vend)]) -logging.info('Censoring %.2f seconds', abs(vtime)) -f.attrs['foreground_time_exc'] -= abs(vtime) -f['segments/foreground_veto/start'] = vstart -f['segments/foreground_veto/end'] = vend - -#TODO: add in background combinations +vtime = segments.segmentlist([segments.segment(s, e) for s, e in zip(vstart, vend)]) +logging.info("Censoring %.2f seconds", abs(vtime)) +f.attrs["foreground_time_exc"] -= abs(vtime) +f["segments/foreground_veto/start"] = vstart +f["segments/foreground_veto/end"] = vend + +# TODO: add in background combinations # If there is a background set (full_data as opposed to injection run), then # recalculate the values for its triggers as well f.close() -logging.info('done') +logging.info("done") diff --git a/bin/all_sky_search/pycbc_cut_merge_triggers_to_tmpltbank b/bin/all_sky_search/pycbc_cut_merge_triggers_to_tmpltbank index fafbb03bf12..05b1d8c94a7 100644 --- a/bin/all_sky_search/pycbc_cut_merge_triggers_to_tmpltbank +++ b/bin/all_sky_search/pycbc_cut_merge_triggers_to_tmpltbank @@ -20,58 +20,83 @@ Reduce a MERGE triggers file to a reduced template bank """ -import logging import argparse -import numpy +import logging + import h5py +import numpy + import pycbc -from pycbc.io import HFile from pycbc import load_source +from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-file", required=True, - help="Input merge triggers HDF file.") -parser.add_argument("--output-file", required=True, - help="Output merge triggers HDF file.") -parser.add_argument("--full-template-bank", required=True, - help="The original full template bank HDF file.") -parser.add_argument("--filter-func-file", required=True, - help="This can be provided to give a function to define " - "which points are covered by the template bank " - "bounds, and which are not. The file should contain " - "a function called filter_tmpltbank, which should " - "take as call profile the template bank HDF object " - "and return a boolean (accept=1/reject=0) array.") +parser.add_argument( + "--input-file", required=True, help="Input merge triggers HDF file." +) +parser.add_argument( + "--output-file", required=True, help="Output merge triggers HDF file." +) +parser.add_argument( + "--full-template-bank", + required=True, + help="The original full template bank HDF file.", +) +parser.add_argument( + "--filter-func-file", + required=True, + help="This can be provided to give a function to define " + "which points are covered by the template bank " + "bounds, and which are not. The file should contain " + "a function called filter_tmpltbank, which should " + "take as call profile the template bank HDF object " + "and return a boolean (accept=1/reject=0) array.", +) opt = parser.parse_args() pycbc.init_logging(opt.verbose) -bank_fd = HFile(opt.full_template_bank, 'r') +bank_fd = HFile(opt.full_template_bank, "r") -modl = load_source('filter_func', opt.filter_func_file) +modl = load_source("filter_func", opt.filter_func_file) func = modl.filter_tmpltbank bool_arr = func(bank_fd) -logging.info("Downselecting templates. Started with %d templates, now have " - "%d after downselecting.", len(bool_arr), bool_arr.sum()) +logging.info( + "Downselecting templates. Started with %d templates, now have " + "%d after downselecting.", + len(bool_arr), + bool_arr.sum(), +) tids = numpy.arange(len(bool_arr))[bool_arr] -hashes = bank_fd['template_hash'][:] +hashes = bank_fd["template_hash"][:] bank_tids = hashes.argsort() unsort = bank_tids.argsort() -copy_params = ['bank_chisq', 'bank_chisq_dof', 'chisq', 'chisq_dof', - 'coa_phase', 'cont_chisq', 'cont_chisq_dof','end_time', - 'sg_chisq', 'sigmasq', 'snr', 'template_duration'] - -ifd = HFile(opt.input_file, 'r') +copy_params = [ + "bank_chisq", + "bank_chisq_dof", + "chisq", + "chisq_dof", + "coa_phase", + "cont_chisq", + "cont_chisq_dof", + "end_time", + "sg_chisq", + "sigmasq", + "snr", + "template_duration", +] + +ifd = HFile(opt.input_file, "r") ifos = list(ifd.keys()) -assert(len(ifos) == 1) +assert len(ifos) == 1 ifo = ifos[0] -ofd = HFile(opt.output_file, 'w') +ofd = HFile(opt.output_file, "w") ofd.create_group(ifo) new_boundaries = [] old_boundaries = [] @@ -82,80 +107,81 @@ for tid in tids: if not tid_count % 1000: logging.info("Processing template %d of %d", tid_count, len(tids)) # Where is it's lower boundary - boundary1 = ifd[ifo+'/template_boundaries'][tid] + boundary1 = ifd[ifo + "/template_boundaries"][tid] # Upper boundary is harder # Position in sorted hashed list pos = unsort[tid] if pos == len(bool_arr) - 1: # If it's the last one, then go to the end - boundary2 = len(ifd[ifo+'/template_duration']) + boundary2 = len(ifd[ifo + "/template_duration"]) else: # Otherwise find the next template boundary, which is tricksy - boundary2 = ifd[ifo+'/template_boundaries'][bank_tids[pos+1]] + boundary2 = ifd[ifo + "/template_boundaries"][bank_tids[pos + 1]] # Check this is sane - test_tids = ifd[ifo+'/template_id'][boundary1:boundary2] + test_tids = ifd[ifo + "/template_id"][boundary1:boundary2] if (test_tids - tid).any(): - raise ValueError() - old_boundaries.append((boundary1,boundary2)) + raise ValueError + old_boundaries.append((boundary1, boundary2)) if new_boundaries: - new_boundaries.append((new_boundaries[-1][1], - new_boundaries[-1][1]+boundary2-boundary1)) + new_boundaries.append( + (new_boundaries[-1][1], new_boundaries[-1][1] + boundary2 - boundary1) + ) else: - new_boundaries.append((0,boundary2-boundary1)) + new_boundaries.append((0, boundary2 - boundary1)) template_boundaries = [tmpx[0] for tmpx in new_boundaries] -ofd[ifo]['template_boundaries'] = template_boundaries +ofd[ifo]["template_boundaries"] = template_boundaries for c in copy_params: logging.info("Copying parameter " + c) - currdtype=ifd[ifo][c][:2].dtype - temp_array=numpy.zeros([new_boundaries[-1][1]], dtype=currdtype) + currdtype = ifd[ifo][c][:2].dtype + temp_array = numpy.zeros([new_boundaries[-1][1]], dtype=currdtype) for i in range(len(old_boundaries)): old_bound = old_boundaries[i] new_bound = new_boundaries[i] - curr_data = ifd[ifo][c][old_bound[0]:old_bound[1]] - temp_array[new_bound[0]:new_bound[1]] = curr_data + curr_data = ifd[ifo][c][old_bound[0] : old_bound[1]] + temp_array[new_bound[0] : new_bound[1]] = curr_data ofd[ifo][c] = temp_array refs = [] for i in range(len(new_boundaries)): new_bound = new_boundaries[i] - refs.append(ofd[ifo][c].regionref[new_bound[0]:new_bound[1]]) - ofd[ifo].create_dataset\ - (c + '_template', data=refs, - dtype=h5py.special_dtype(ref=h5py.RegionReference)) + refs.append(ofd[ifo][c].regionref[new_bound[0] : new_bound[1]]) + ofd[ifo].create_dataset( + c + "_template", data=refs, dtype=h5py.special_dtype(ref=h5py.RegionReference) + ) logging.info("Updating template IDs") -c = 'template_id' -currdtype=ifd[ifo][c][:2].dtype -temp_array=numpy.zeros([new_boundaries[-1][1]], dtype=currdtype) -temp_array2=numpy.zeros([new_boundaries[-1][1]], dtype=currdtype) +c = "template_id" +currdtype = ifd[ifo][c][:2].dtype +temp_array = numpy.zeros([new_boundaries[-1][1]], dtype=currdtype) +temp_array2 = numpy.zeros([new_boundaries[-1][1]], dtype=currdtype) for i in range(len(old_boundaries)): old_bound = old_boundaries[i] new_bound = new_boundaries[i] - curr_data = ifd[ifo][c][old_bound[0]:old_bound[1]] - temp_array2[new_bound[0]:new_bound[1]] = curr_data - temp_array[new_bound[0]:new_bound[1]] = i + curr_data = ifd[ifo][c][old_bound[0] : old_bound[1]] + temp_array2[new_bound[0] : new_bound[1]] = curr_data + temp_array[new_bound[0] : new_bound[1]] = i ofd[ifo][c] = temp_array -ofd[ifo][c+'_orig'] = temp_array2 +ofd[ifo][c + "_orig"] = temp_array2 refs = [] refs2 = [] for i in range(len(new_boundaries)): new_bound = new_boundaries[i] - refs.append(ofd[ifo][c].regionref[new_bound[0]:new_bound[1]]) - refs2.append(ofd[ifo][c+'_orig'].regionref[new_bound[0]:new_bound[1]]) + refs.append(ofd[ifo][c].regionref[new_bound[0] : new_bound[1]]) + refs2.append(ofd[ifo][c + "_orig"].regionref[new_bound[0] : new_bound[1]]) -ofd[ifo].create_dataset\ - (c + '_template', data=refs, - dtype=h5py.special_dtype(ref=h5py.RegionReference)) -ofd[ifo].create_dataset\ - (c + '_orig_template', data=refs2, - dtype=h5py.special_dtype(ref=h5py.RegionReference)) +ofd[ifo].create_dataset( + c + "_template", data=refs, dtype=h5py.special_dtype(ref=h5py.RegionReference) +) +ofd[ifo].create_dataset( + c + "_orig_template", data=refs2, dtype=h5py.special_dtype(ref=h5py.RegionReference) +) # Copy some of the unchanged groups -ifd.copy(ifo+'/gating', ofd[ifo]) -ifd.copy(ifo+'/search', ofd[ifo]) +ifd.copy(ifo + "/gating", ofd[ifo]) +ifd.copy(ifo + "/search", ofd[ifo]) # Copy attributes logging.info("Copying attributes") diff --git a/bin/all_sky_search/pycbc_distribute_background_bins b/bin/all_sky_search/pycbc_distribute_background_bins index e8b746fcc58..d1f83bedb9a 100644 --- a/bin/all_sky_search/pycbc_distribute_background_bins +++ b/bin/all_sky_search/pycbc_distribute_background_bins @@ -1,27 +1,39 @@ #!/bin/env python -import argparse, numpy, pycbc.events, logging, pycbc.events, pycbc.io +import argparse +import logging + +import numpy + +import pycbc.events +import pycbc.io parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--coinc-files', nargs='+', - help="List of coinc files to be redistributed") -parser.add_argument('--background-bins', nargs='+', - help="Ordered list of mass bin upper boundaries. " - "An ordered list of type-boundary pairs, applied sequentially." - "Must provide a name (can be any unique string for tagging " - "purposes), the parameter to bin on, and the membership " - "condition via 'lt' / 'gt' operators. " - "Ex. name1:component:lt2 name2:total:lt15 name3:SEOBNRv2Peak:gt1000") -parser.add_argument('--f-lower', - help="Lower frequency cutoff for evaluating template duration. Should" - " be equal to the lower cutoff used in inspiral jobs") -parser.add_argument('--bank-file', - help="hdf format template bank file") -parser.add_argument('--output-files', nargs='+', - help="list of output file names, one for each mass bin") +parser.add_argument( + "--coinc-files", nargs="+", help="List of coinc files to be redistributed" +) +parser.add_argument( + "--background-bins", + nargs="+", + help="Ordered list of mass bin upper boundaries. " + "An ordered list of type-boundary pairs, applied sequentially." + "Must provide a name (can be any unique string for tagging " + "purposes), the parameter to bin on, and the membership " + "condition via 'lt' / 'gt' operators. " + "Ex. name1:component:lt2 name2:total:lt15 name3:SEOBNRv2Peak:gt1000", +) +parser.add_argument( + "--f-lower", + help="Lower frequency cutoff for evaluating template duration. Should" + " be equal to the lower cutoff used in inspiral jobs", +) +parser.add_argument("--bank-file", help="hdf format template bank file") +parser.add_argument( + "--output-files", nargs="+", help="list of output file names, one for each mass bin" +) args = parser.parse_args() -if 'duration' in args.background_bins and not args.f_lower: +if "duration" in args.background_bins and not args.f_lower: raise RuntimeError("Can't bin on template duration without f-lower!") pycbc.init_logging(args.verbose) @@ -30,21 +42,25 @@ if len(args.output_files) != len(args.background_bins): raise ValueError("Number of mass bins and output files does not match") f = pycbc.io.HFile(args.bank_file) -data = {'mass1':f['mass1'][:], 'mass2':f['mass2'][:], - 'spin1z':f['spin1z'][:], 'spin2z':f['spin2z'][:]} +data = { + "mass1": f["mass1"][:], + "mass2": f["mass2"][:], + "spin1z": f["spin1z"][:], + "spin2z": f["spin2z"][:], +} if args.f_lower: - data['f_lower'] = float(args.f_lower) + data["f_lower"] = float(args.f_lower) locs_dict = pycbc.events.background_bin_from_string(args.background_bins, data) -names = [b.split(':')[0] for b in args.background_bins] +names = [b.split(":")[0] for b in args.background_bins] d = pycbc.io.StatmapData(files=args.coinc_files) -logging.info('%s coinc triggers' % len(d)) +logging.info("%s coinc triggers" % len(d)) for name, outname in zip(names, args.output_files): # select the coincs from only this bin and save to a single combined file locs = locs_dict[name] e = d.select(numpy.isin(d.template_id, locs)) - logging.info('%s coincs in mass bin: %s' % (len(e), name)) + logging.info("%s coincs in mass bin: %s" % (len(e), name)) e.save(outname) f = pycbc.io.HFile(outname) - f.attrs['name'] = name + f.attrs["name"] = name diff --git a/bin/all_sky_search/pycbc_dtphase b/bin/all_sky_search/pycbc_dtphase index 77f35e99bd1..d19a3549671 100644 --- a/bin/all_sky_search/pycbc_dtphase +++ b/bin/all_sky_search/pycbc_dtphase @@ -13,66 +13,110 @@ To get the signal rate this weight should be scaled by the local sensitivity value and by the SNR of the event in the reference detector. """ -import argparse, numpy as np, pycbc.detector, logging +import argparse +import logging +from copy import deepcopy + +import numpy as np from numpy.random import uniform from scipy.stats import norm -from copy import deepcopy -from pycbc.io import HFile +import pycbc.detector +from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--ifos', nargs='+', - help="The ifos to generate a histogram for") -parser.add_argument('--sample-size', type=int, required=True, - help="Approximate number of independent samples to draw " - "for the distribution") -parser.add_argument('--snr-ratio', type=float, required=True, - help="The SNR ratio permitted between reference ifo and " - "all others. Ex. giving 4 permits a ratio of " - "0.25 -> 4") -parser.add_argument('--relative-sensitivities', nargs='+', type=float, - help="Numbers proportional to horizon distance or " - "expected SNR at fixed distance, one for each ifo") -parser.add_argument('--seed', type=int, default=124) -parser.add_argument('--output-file', required=True) -parser.add_argument('--bin-density', type=int, default=1, - help="Number of bins per 1 sigma uncertainty in a " - "parameter. Higher values increase the resolution of " - "the histogram at the expense of storage.") -parser.add_argument('--smoothing-sigma', type=int, default=2, - help="Width of the smoothing kernel in sigmas") -parser.add_argument('--timing-uncertainty', type=float, default=.001, - help="Timing uncertainty to set bin size and smoothing " - "interval [default=.001s]") -parser.add_argument('--phase-uncertainty', type=float, default=0.25, - help="Phase uncertainty used to set bin size and " - "smoothing") -parser.add_argument('--snr-reference', type=float, default=5, - help="Reference SNR to scale SNR uncertainty") -parser.add_argument('--snr-uncertainty', type=float, default=1.0, - help="SNR uncertainty to set bin size and smoothing") -parser.add_argument('--weight-threshold', type=float, default=1e-10, - help="Minimum histogram weight to store as a proportion " - "of the histogram maximum: bins with less weight " - "will not be stored.") -parser.add_argument('--batch-size', type=int, default=1000000) -parser.add_argument('--param-bin-dtype', default='int32', - help="Type to use to store param_bin information. " - "Affects the maximum values that can be taken by " - "bin parameters. Default int32.", - choices=['int8', 'int16', 'int32']) +parser.add_argument("--ifos", nargs="+", help="The ifos to generate a histogram for") +parser.add_argument( + "--sample-size", + type=int, + required=True, + help="Approximate number of independent samples to draw for the distribution", +) +parser.add_argument( + "--snr-ratio", + type=float, + required=True, + help="The SNR ratio permitted between reference ifo and " + "all others. Ex. giving 4 permits a ratio of " + "0.25 -> 4", +) +parser.add_argument( + "--relative-sensitivities", + nargs="+", + type=float, + help="Numbers proportional to horizon distance or " + "expected SNR at fixed distance, one for each ifo", +) +parser.add_argument("--seed", type=int, default=124) +parser.add_argument("--output-file", required=True) +parser.add_argument( + "--bin-density", + type=int, + default=1, + help="Number of bins per 1 sigma uncertainty in a " + "parameter. Higher values increase the resolution of " + "the histogram at the expense of storage.", +) +parser.add_argument( + "--smoothing-sigma", + type=int, + default=2, + help="Width of the smoothing kernel in sigmas", +) +parser.add_argument( + "--timing-uncertainty", + type=float, + default=0.001, + help="Timing uncertainty to set bin size and smoothing interval [default=.001s]", +) +parser.add_argument( + "--phase-uncertainty", + type=float, + default=0.25, + help="Phase uncertainty used to set bin size and smoothing", +) +parser.add_argument( + "--snr-reference", + type=float, + default=5, + help="Reference SNR to scale SNR uncertainty", +) +parser.add_argument( + "--snr-uncertainty", + type=float, + default=1.0, + help="SNR uncertainty to set bin size and smoothing", +) +parser.add_argument( + "--weight-threshold", + type=float, + default=1e-10, + help="Minimum histogram weight to store as a proportion " + "of the histogram maximum: bins with less weight " + "will not be stored.", +) +parser.add_argument("--batch-size", type=int, default=1000000) +parser.add_argument( + "--param-bin-dtype", + default="int32", + help="Type to use to store param_bin information. " + "Affects the maximum values that can be taken by " + "bin parameters. Default int32.", + choices=["int8", "int16", "int32"], +) args = parser.parse_args() if len(args.relative_sensitivities) != len(args.ifos): - parser.error('--relative-sensitivities requires one numerical argument ' - 'for each detector') + parser.error( + "--relative-sensitivities requires one numerical argument for each detector" + ) # Approximate timing error at lower SNRs twidth = args.timing_uncertainty / args.bin_density # Factor of sqrt(2) as we'll be combining two SNRs with independent # uncertainties -serr = args.snr_uncertainty * 2 ** 0.5 +serr = args.snr_uncertainty * 2**0.5 # Reference SNR for bin smoothing sref = args.snr_reference swidth = serr / sref / args.bin_density @@ -82,12 +126,15 @@ pwidth = np.arctan(serr / sref) / args.bin_density srbmax = int(args.snr_ratio / swidth) srbmin = int((1.0 / args.snr_ratio) / swidth) + # Apply a simple smoothing to help account for measurement errors for # weak signals def smooth_param(data, index, wrapped=None): mref = max(data.values()) - bins = np.arange(-args.smoothing_sigma * args.bin_density,\ - args.smoothing_sigma * args.bin_density + 1) + bins = np.arange( + -args.smoothing_sigma * args.bin_density, + args.smoothing_sigma * args.bin_density + 1, + ) kernel = norm.pdf(bins, scale=args.bin_density) nweights = {} @@ -128,6 +175,7 @@ def smooth_param(data, index, wrapped=None): nweights[tnkey] = weight return nweights + d = {ifo: pycbc.detector.Detector(ifo) for ifo in args.ifos} pycbc.init_logging(args.verbose) @@ -140,17 +188,19 @@ bin_dtype = args.param_bin_dtype max_td = twidth * (np.iinfo(bin_dtype).max - np.iinfo(bin_dtype).min + 1) if max_td < 0.0425: - raise RuntimeError('Max allowed time difference is less than the earth ' - 'travel time. Some observatories may be further apart ' - 'than this.') + raise RuntimeError( + "Max allowed time difference is less than the earth " + "travel time. Some observatories may be further apart " + "than this." + ) # Store results for each ifo as a reference. The reference ifo is the # ifo which gets the smallest amplitude. This allows us to get the correct # symmetries handled under ifo switch and apply more consistent treatment # of error uncertainties. -f = HFile(args.output_file, 'w') +f = HFile(args.output_file, "w") for ifo0 in args.ifos: - logging.info('Storing results using %s as a reference', ifo0) + logging.info("Storing results using %s as a reference", ifo0) other_ifos = deepcopy(args.ifos) other_ifos.remove(ifo0) @@ -159,13 +209,13 @@ for ifo0 in args.ifos: weights = {} for k in range(chunks): nsamples += size - logging.info('generating %s samples', size) + logging.info("generating %s samples", size) # Choose random sky location and polarizations from # an isotropic population ra = uniform(0, 2 * np.pi, size=size) - dec = np.arccos(uniform(-1., 1., size=size)) - np.pi/2 - inc = np.arccos(uniform(-1., 1., size=size)) + dec = np.arccos(uniform(-1.0, 1.0, size=size)) - np.pi / 2 + inc = np.arccos(uniform(-1.0, 1.0, size=size)) pol = uniform(0, 2 * np.pi, size=size) ic = np.cos(inc) ip = 0.5 * (1.0 + ic * ic) @@ -176,17 +226,17 @@ for ifo0 in args.ifos: data[ifo] = {} fp, fc = d[ifo].antenna_pattern(ra, dec, pol, 0) sp, sc = fp * ip, fc * ic - data[ifo]['s'] = (sp ** 2. + sc ** 2.) ** 0.5 * rs - data[ifo]['t'] = d[ifo].time_delay_from_earth_center(ra, dec, 0) - data[ifo]['p'] = np.arctan2(sc, sp) + data[ifo]["s"] = (sp**2.0 + sc**2.0) ** 0.5 * rs + data[ifo]["t"] = d[ifo].time_delay_from_earth_center(ra, dec, 0) + data[ifo]["p"] = np.arctan2(sc, sp) # Bin the data bind = [] keep = None for ifo1 in other_ifos: - dt = (data[ifo0]['t'] - data[ifo1]['t']) - dp = (data[ifo0]['p'] - data[ifo1]['p']) % (2. * np.pi) - sr = (data[ifo1]['s'] / data[ifo0]['s']) + dt = data[ifo0]["t"] - data[ifo1]["t"] + dp = (data[ifo0]["p"] - data[ifo1]["p"]) % (2.0 * np.pi) + sr = data[ifo1]["s"] / data[ifo0]["s"] dtbin = (dt / twidth).astype(int) dpbin = (dp / pwidth).astype(int) srbin = (sr / swidth).astype(int) @@ -202,7 +252,7 @@ for ifo0 in args.ifos: # use first ifo as reference for weights bind = [a[keep] for a in bind] - w = data[ifo0]['s'][keep] ** 3. + w = data[ifo0]["s"][keep] ** 3.0 for i, key in enumerate(zip(*bind)): if key not in weights: weights[key] = 0 @@ -210,30 +260,31 @@ for ifo0 in args.ifos: ol = l l = len(weights.values()) - logging.info('%s, %s, %s, %s', l, l - ol, (l - ol) / float(size), - l / float(nsamples)) + logging.info( + "%s, %s, %s, %s", l, l - ol, (l - ol) / float(size), l / float(nsamples) + ) - logging.info('applying smoothing') + logging.info("applying smoothing") # apply smoothing iteratively pwrap = (2 * np.pi) / pwidth - for i in range(len(args.ifos)-1): - logging.info('%s-phase', len(weights)) + for i in range(len(args.ifos) - 1): + logging.info("%s-phase", len(weights)) weights = smooth_param(weights, i * 3 + 1, wrapped=pwrap) - logging.info('%s-time', len(weights)) + logging.info("%s-time", len(weights)) weights = smooth_param(weights, i * 3 + 0) - logging.info('%s-amp', len(weights)) + logging.info("%s-amp", len(weights)) weights = smooth_param(weights, i * 3 + 2) - logging.info('smoothing done: %s', len(weights)) + logging.info("smoothing done: %s", len(weights)) - logging.info('converting to numpy arrays and normalizing') + logging.info("converting to numpy arrays and normalizing") keys = np.array(list(weights.keys())) values = np.array(list(weights.values())) values /= values.max() - logging.info('Removing bins outside of SNR ratio limits') + logging.info("Removing bins outside of SNR ratio limits") n_precut = len(keys) keep = None - for i in range(len(args.ifos)-1): + for i in range(len(args.ifos) - 1): srbin = np.array(list(zip(*keys))[i * 3 + 2]) if keep is None: keep = (srbin <= srbmax) & (srbin >= srbmin) @@ -241,37 +292,39 @@ for ifo0 in args.ifos: keep = keep & (srbin <= srbmax) & (srbin >= srbmin) keys = keys[keep] values = values[keep] - logging.info('Removed %s bins', n_precut - len(keys)) + logging.info("Removed %s bins", n_precut - len(keys)) - logging.info('discarding bins below threshold to limit storage') + logging.info("discarding bins below threshold to limit storage") l = values > args.weight_threshold keys = keys[l].astype(bin_dtype) values = values[l].astype(np.float32) - logging.info('Final length: %s', len(keys)) + logging.info("Final length: %s", len(keys)) - logging.info('Presorting by keys for downstream use') + logging.info("Presorting by keys for downstream use") ncol = keys.shape[1] - pdtype = [('c%s' % i, bin_dtype) for i in range(ncol)] + pdtype = [("c%s" % i, bin_dtype) for i in range(ncol)] keys_bin = np.zeros(len(values), dtype=pdtype) for i in range(ncol): - keys_bin['c%s' % i] = keys[:, i] + keys_bin["c%s" % i] = keys[:, i] lsort = keys_bin.argsort() keys_bin = keys_bin[lsort] values = values[lsort] - logging.info('Writing results to file') - f.create_dataset('%s/param_bin' % ifo0, data=keys_bin, compression='gzip', - compression_opts=7) - f.create_dataset('%s/weights' % ifo0, data=values, compression='gzip', - compression_opts=7) - -f.attrs['sensitivity_ratios'] = args.relative_sensitivities -f.attrs['srbmin'] = srbmin -f.attrs['srbmax'] = srbmax -f.attrs['twidth'] = twidth -f.attrs['pwidth'] = pwidth -f.attrs['swidth'] = swidth -f.attrs['ifos'] = args.ifos -f.attrs['stat'] = 'phasetd_newsnr_%s' % ''.join(args.ifos) - -logging.info('Done') + logging.info("Writing results to file") + f.create_dataset( + "%s/param_bin" % ifo0, data=keys_bin, compression="gzip", compression_opts=7 + ) + f.create_dataset( + "%s/weights" % ifo0, data=values, compression="gzip", compression_opts=7 + ) + +f.attrs["sensitivity_ratios"] = args.relative_sensitivities +f.attrs["srbmin"] = srbmin +f.attrs["srbmax"] = srbmax +f.attrs["twidth"] = twidth +f.attrs["pwidth"] = pwidth +f.attrs["swidth"] = swidth +f.attrs["ifos"] = args.ifos +f.attrs["stat"] = "phasetd_newsnr_%s" % "".join(args.ifos) + +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_exclude_zerolag b/bin/all_sky_search/pycbc_exclude_zerolag index f4029387b00..66f8f9ac482 100644 --- a/bin/all_sky_search/pycbc_exclude_zerolag +++ b/bin/all_sky_search/pycbc_exclude_zerolag @@ -4,113 +4,137 @@ Remove all coincs in background_exc that contain triggers at time of zerolag coincidences from *any* coincidence type with ifar above a certain threshold """ -import numpy as np, argparse, logging, pycbc, pycbc.io -from pycbc.events import veto, significance +import argparse +import logging + +import numpy as np + +import pycbc import pycbc.conversions as conv +import pycbc.io +from pycbc.events import significance, veto parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--statmap-file', type=str, - help="Coinc statmap file to be recalculated based on foreground removal") -parser.add_argument('--other-statmap-files', nargs='+', - help="List of coinc statmap files from other coincidence types") -parser.add_argument('--censor-ifar-threshold', type=float, default=0.003, +parser.add_argument( + "--statmap-file", + type=str, + help="Coinc statmap file to be recalculated based on foreground removal", +) +parser.add_argument( + "--other-statmap-files", + nargs="+", + help="List of coinc statmap files from other coincidence types", +) +parser.add_argument( + "--censor-ifar-threshold", + type=float, + default=0.003, help="Only window out foreground triggers with IFAR (years)" - "above the threshold [default=0.003yr]") -parser.add_argument('--veto-window', type=float, default=0.1, - help="Time around each zerolag trigger to window out [default=.1s]") + "above the threshold [default=0.003yr]", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) significance.insert_significance_option_group(parser) -parser.add_argument('--output-file', help="name of output file") +parser.add_argument("--output-file", help="name of output file") args = parser.parse_args() significance.check_significance_options(args, parser) pycbc.init_logging(args.verbose) -f_in = pycbc.io.HFile(args.statmap_file,'r') +f_in = pycbc.io.HFile(args.statmap_file, "r") f_out = pycbc.io.HFile(args.output_file, "w") -f_others = [pycbc.io.HFile(fname,'r') for fname in args.other_statmap_files] +f_others = [pycbc.io.HFile(fname, "r") for fname in args.other_statmap_files] -all_ifos = f_in.attrs['ifos'].split(' ') -all_ifo_key = ''.join(all_ifos) +all_ifos = f_in.attrs["ifos"].split(" ") +all_ifo_key = "".join(all_ifos) -significance_dict = significance.digest_significance_options([all_ifo_key], - args) +significance_dict = significance.digest_significance_options([all_ifo_key], args) -logging.info('Copying attributes to %s' % args.output_file) +logging.info("Copying attributes to %s" % args.output_file) for attrk in f_in.attrs.keys(): f_out.attrs[attrk] = f_in.attrs[attrk] -logging.info('Copying unchanged datasets to %s' % args.output_file) +logging.info("Copying unchanged datasets to %s" % args.output_file) keys = pycbc.io.name_all_datasets([f_in]) for k in keys: - if 'background_exc' not in k: + if "background_exc" not in k: f_out[k] = f_in[k][:] logging.info("Collating foreground times from other coinc types") all_fg_times = [] for f in f_others: - ifar_above_thresh = np.nonzero(f['foreground/ifar'][:] > \ - args.censor_ifar_threshold)[0] - mean_times = np.mean([f['foreground/%s/time' % ifo][:] - for ifo in f.attrs['ifos'].split(' ')], axis=0) + ifar_above_thresh = np.nonzero( + f["foreground/ifar"][:] > args.censor_ifar_threshold + )[0] + mean_times = np.mean( + [f["foreground/%s/time" % ifo][:] for ifo in f.attrs["ifos"].split(" ")], axis=0 + ) all_fg_times += list(mean_times[ifar_above_thresh]) all_fg_times = np.array(all_fg_times) logging.info("Loading coinc triggers into all_trigs structure") -groups = ['decimation_factor', 'stat', 'template_id', 'timeslide_id', 'ifar'] +groups = ["decimation_factor", "stat", "template_id", "timeslide_id", "ifar"] for ifo in all_ifos: - groups += ['%s/time' % ifo] - groups += ['%s/trigger_id' % ifo] -data = {g: f_in['background_exc/%s' % g] for g in groups} + groups += ["%s/time" % ifo] + groups += ["%s/trigger_id" % ifo] +data = {g: f_in["background_exc/%s" % g] for g in groups} all_trigs = pycbc.io.DictArray(data=data) -n_triggers = f_in['background_exc/%s/time' % all_ifos[0]].size -logging.info('%d background_exc triggers' % n_triggers) +n_triggers = f_in["background_exc/%s/time" % all_ifos[0]].size +logging.info("%d background_exc triggers" % n_triggers) -logging.info('Finding background_exc triggers within {}s of any ' - 'foreground triggers'.format(args.veto_window)) +logging.info( + f"Finding background_exc triggers within {args.veto_window}s of any foreground triggers" +) remove_start_time = all_fg_times - args.veto_window remove_end_time = all_fg_times + args.veto_window all_vetoed_idx = [] for ifo in all_ifos: fg_veto_ids = veto.indices_within_times( - f_in['background_exc/%s/time' % ifo][:], - remove_start_time, remove_end_time) + f_in["background_exc/%s/time" % ifo][:], remove_start_time, remove_end_time + ) all_vetoed_idx += list(fg_veto_ids) all_vetoed_idx = np.unique(all_vetoed_idx) -logging.info('Removing {} background_exc triggers close to foreground' - ' coincs'.format(len(all_vetoed_idx))) +logging.info( + f"Removing {len(all_vetoed_idx)} background_exc triggers close to foreground coincs" +) filtered_trigs = all_trigs.remove(all_vetoed_idx) -n_triggers_new = len(filtered_trigs.data['stat']) -logging.info('%d triggers remaining' % n_triggers_new) +n_triggers_new = len(filtered_trigs.data["stat"]) +logging.info("%d triggers remaining" % n_triggers_new) -logging.info('Writing updated background_exc to file') +logging.info("Writing updated background_exc to file") for k in filtered_trigs.data: - f_out['background_exc/%s' % k] = filtered_trigs.data[k] + f_out["background_exc/%s" % k] = filtered_trigs.data[k] -logging.info('Recalculating IFARs') +logging.info("Recalculating IFARs") bg_far, fg_far, sig_info = significance.get_far( - filtered_trigs.data['stat'], - f_in['foreground/stat'][:], - filtered_trigs.data['decimation_factor'], - f_in.attrs['background_time_exc'], - **significance_dict[all_ifo_key]) + filtered_trigs.data["stat"], + f_in["foreground/stat"][:], + filtered_trigs.data["decimation_factor"], + f_in.attrs["background_time_exc"], + **significance_dict[all_ifo_key], +) fg_far = significance.apply_far_limit(fg_far, significance_dict, combo=all_ifo_key) bg_far = significance.apply_far_limit(bg_far, significance_dict, combo=all_ifo_key) -fg_ifar_exc = 1. / fg_far -bg_ifar_exc = 1. / bg_far +fg_ifar_exc = 1.0 / fg_far +bg_ifar_exc = 1.0 / bg_far -logging.info('Writing updated ifars to file') -f_out['foreground/ifar_exc'][:] = conv.sec_to_year(fg_ifar_exc) -f_out['background_exc/ifar'][:] = conv.sec_to_year(bg_ifar_exc) +logging.info("Writing updated ifars to file") +f_out["foreground/ifar_exc"][:] = conv.sec_to_year(fg_ifar_exc) +f_out["background_exc/ifar"][:] = conv.sec_to_year(bg_ifar_exc) for key, value in sig_info.items(): - f_out['foreground'].attrs[key + '_exc'] = value + f_out["foreground"].attrs[key + "_exc"] = value -fg_time_exc = conv.sec_to_year(f_in.attrs['foreground_time_exc']) -f_out['foreground/fap_exc'][:] = 1 - np.exp(-fg_time_exc / fg_ifar_exc) +fg_time_exc = conv.sec_to_year(f_in.attrs["foreground_time_exc"]) +f_out["foreground/fap_exc"][:] = 1 - np.exp(-fg_time_exc / fg_ifar_exc) logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_fit_sngls_binned b/bin/all_sky_search/pycbc_fit_sngls_binned index 126ed493fc3..f3f9c0c78e8 100644 --- a/bin/all_sky_search/pycbc_fit_sngls_binned +++ b/bin/all_sky_search/pycbc_fit_sngls_binned @@ -13,102 +13,159 @@ # Public License for more details. +import argparse +import logging import sys -import argparse, logging from matplotlib import use -use('Agg') -from matplotlib import pyplot as plt -import copy, numpy as np +use("Agg") +import copy + +import numpy as np +from matplotlib import pyplot as plt import pycbc -from pycbc import events, bin_utils, results -from pycbc.events import triggers -from pycbc.events import trigger_fits as trstats +from pycbc import bin_utils, events, results from pycbc.events import stat as pystat +from pycbc.events import trigger_fits as trstats +from pycbc.events import triggers from pycbc.io import HFile #### MAIN #### -parser = argparse.ArgumentParser(usage="", +parser = argparse.ArgumentParser( + usage="", description="Perform maximum-likelihood fits of single inspiral trigger" - " distributions to various functions") + " distributions to various functions", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--trigger-file", - help="Input hdf5 file containing single triggers. " - "Required") -parser.add_argument("--bank-file", default=None, - help="hdf file containing template parameters. Required") -parser.add_argument("--veto-file", nargs='*', default=[], action='append', - help="File(s) in .xml format with veto segments to apply " - "to triggers before fitting") -parser.add_argument("--veto-segment-name", nargs='*', default=[], action='append', - help="Name(s) of veto segments to apply. Optional, if not " - "given all segments for a given ifo will be used") -parser.add_argument("--ifo", required=True, - help="Ifo producing triggers to be fitted. Required") -parser.add_argument("--fit-function", - choices=["exponential", "rayleigh", "power"], - help="Functional form for the maximum likelihood fit") -parser.add_argument("--stat-threshold", nargs="+", type=float, - help="Only fit triggers with statistic value above this " - "threshold : can be a space-separated list, then a fit " - "will be done for each threshold. Required. Typical " - "values 6.25 6.5 6.75") -parser.add_argument("--prune-param", - help="Parameter to define bins for 'pruning' loud triggers" - " to make the fit insensitive to signals and outliers. " - "Choose from mchirp, mtotal, template_duration or a named " - "frequency cutoff in pnutils or a frequency function in " - "LALSimulation") -parser.add_argument("--prune-bins", type=int, - help="Number of bins to divide bank into when pruning") -parser.add_argument("--prune-number", type=int, - help="Number of loudest events to prune in each bin") -parser.add_argument("--log-prune-param", action='store_true', - help="Bin in the log of prune-param") -parser.add_argument("--f-lower", type=float, default=0., - help="Starting frequency for calculating template " - "duration; if not given, duration will be read from " - "single trigger files") +parser.add_argument( + "--trigger-file", help="Input hdf5 file containing single triggers. Required" +) +parser.add_argument( + "--bank-file", + default=None, + help="hdf file containing template parameters. Required", +) +parser.add_argument( + "--veto-file", + nargs="*", + default=[], + action="append", + help="File(s) in .xml format with veto segments to apply " + "to triggers before fitting", +) +parser.add_argument( + "--veto-segment-name", + nargs="*", + default=[], + action="append", + help="Name(s) of veto segments to apply. Optional, if not " + "given all segments for a given ifo will be used", +) +parser.add_argument( + "--ifo", required=True, help="Ifo producing triggers to be fitted. Required" +) +parser.add_argument( + "--fit-function", + choices=["exponential", "rayleigh", "power"], + help="Functional form for the maximum likelihood fit", +) +parser.add_argument( + "--stat-threshold", + nargs="+", + type=float, + help="Only fit triggers with statistic value above this " + "threshold : can be a space-separated list, then a fit " + "will be done for each threshold. Required. Typical " + "values 6.25 6.5 6.75", +) +parser.add_argument( + "--prune-param", + help="Parameter to define bins for 'pruning' loud triggers" + " to make the fit insensitive to signals and outliers. " + "Choose from mchirp, mtotal, template_duration or a named " + "frequency cutoff in pnutils or a frequency function in " + "LALSimulation", +) +parser.add_argument( + "--prune-bins", type=int, help="Number of bins to divide bank into when pruning" +) +parser.add_argument( + "--prune-number", type=int, help="Number of loudest events to prune in each bin" +) +parser.add_argument( + "--log-prune-param", action="store_true", help="Bin in the log of prune-param" +) +parser.add_argument( + "--f-lower", + type=float, + default=0.0, + help="Starting frequency for calculating template " + "duration; if not given, duration will be read from " + "single trigger files", +) # FIXME : allow choice of SEOBNRv2/v4 or PhenD duration formula ? -parser.add_argument("--bin-param", required=True, - help="Parameter over which to bin when fitting. Required. " - "Choose from mchirp, mtotal, template_duration or a named " - "frequency cutoff in pnutils or a frequency function in " - "LALSimulation") -parser.add_argument("--bin-spacing", choices=["linear", "log", "irregular"], - help="How to space parameter bin edges") +parser.add_argument( + "--bin-param", + required=True, + help="Parameter over which to bin when fitting. Required. " + "Choose from mchirp, mtotal, template_duration or a named " + "frequency cutoff in pnutils or a frequency function in " + "LALSimulation", +) +parser.add_argument( + "--bin-spacing", + choices=["linear", "log", "irregular"], + help="How to space parameter bin edges", +) binopt = parser.add_mutually_exclusive_group(required=True) -binopt.add_argument("--num-bins", type=int, - help="Number of regularly spaced bins to use over the " - " parameter") -binopt.add_argument("--irregular-bins", - help="Comma-separated list of parameter bin edges. " - "Required if --bin-spacing = irregular") -parser.add_argument("--bin-param-units", - help="String to display units of the binning parameter") -parser.add_argument("--approximant", default="SEOBNRv4", - help="Approximant for template duration. Default SEOBNRv4") -parser.add_argument("--min-duration", default=0., - help="Fudge factor for templates with tiny or negative " - "values of template_duration: add to duration values " - "before fitting. Units seconds") +binopt.add_argument( + "--num-bins", + type=int, + help="Number of regularly spaced bins to use over the parameter", +) +binopt.add_argument( + "--irregular-bins", + help="Comma-separated list of parameter bin edges. " + "Required if --bin-spacing = irregular", +) +parser.add_argument( + "--bin-param-units", help="String to display units of the binning parameter" +) +parser.add_argument( + "--approximant", + default="SEOBNRv4", + help="Approximant for template duration. Default SEOBNRv4", +) +parser.add_argument( + "--min-duration", + default=0.0, + help="Fudge factor for templates with tiny or negative " + "values of template_duration: add to duration values " + "before fitting. Units seconds", +) outputchoice = parser.add_mutually_exclusive_group() -outputchoice.add_argument("--plot-dir", - help="Plot the fits made, the variation of fitting " - "coefficients and the Kolmogorov-Smirnov test values " - "and save to the specified directory.") -outputchoice.add_argument("--output-file", - help="Output a plot of hists and fits made for a single " - "threshold value.") -parser.add_argument("--user-tag", default="", - help="Put a possibly informative string in the names of " - "plot files") - -pystat.insert_statistic_option_group(parser, - default_ranking_statistic='single_ranking_only') +outputchoice.add_argument( + "--plot-dir", + help="Plot the fits made, the variation of fitting " + "coefficients and the Kolmogorov-Smirnov test values " + "and save to the specified directory.", +) +outputchoice.add_argument( + "--output-file", + help="Output a plot of hists and fits made for a single threshold value.", +) +parser.add_argument( + "--user-tag", + default="", + help="Put a possibly informative string in the names of plot files", +) + +pystat.insert_statistic_option_group( + parser, default_ranking_statistic="single_ranking_only" +) args = parser.parse_args() args.veto_segment_name = sum(args.veto_segment_name, []) @@ -117,24 +174,29 @@ args.veto_file = sum(args.veto_file, []) if len(args.veto_segment_name) != len(args.veto_file): raise RuntimeError("Number of veto files much match veto file names") -if (args.prune_param or args.prune_bins or args.prune_number) and not \ - (args.prune_param and args.prune_bins and args.prune_number): - raise RuntimeError("To prune, need to specify param, number of bins and " - "nonzero number to prune in each bin!") +if (args.prune_param or args.prune_bins or args.prune_number) and not ( + args.prune_param and args.prune_bins and args.prune_number +): + raise RuntimeError( + "To prune, need to specify param, number of bins and " + "nonzero number to prune in each bin!" + ) if args.output_file is not None and len(args.stat_threshold) > 1: - raise RuntimeError("Cannot plot more than one threshold in a single " - "output file!") + raise RuntimeError("Cannot plot more than one threshold in a single output file!") pycbc.init_logging(args.verbose) -statname = "reweighted SNR" if args.sngl_ranking == "new_snr" else \ - args.sngl_ranking.replace("_", " ").replace("snr", "SNR") +statname = ( + "reweighted SNR" + if args.sngl_ranking == "new_snr" + else args.sngl_ranking.replace("_", " ").replace("snr", "SNR") +) paramname = args.bin_param.replace("_", " ") paramtag = args.bin_param.replace("_", "") if args.plot_dir: - if not args.plot_dir.endswith('/'): - args.plot_dir += '/' + if not args.plot_dir.endswith("/"): + args.plot_dir += "/" plotbase = args.plot_dir + args.ifo + "-" + args.user_tag ## Check option logic @@ -142,20 +204,20 @@ if args.bin_spacing == "irregular": if args.irregular_bins is None: raise RuntimeError("Must specify a list of irregular bin edges!") else: - args.bin_edges = [float(b) for b in args.irregular_bins.split(',')] + args.bin_edges = [float(b) for b in args.irregular_bins.split(",")] -logging.info('Opening trigger file: %s' % args.trigger_file) -trigf = HFile(args.trigger_file, 'r') -logging.info('Opening template file: %s' % args.bank_file) -templatef = HFile(args.bank_file, 'r') +logging.info("Opening trigger file: %s" % args.trigger_file) +trigf = HFile(args.trigger_file, "r") +logging.info("Opening template file: %s" % args.bank_file) +templatef = HFile(args.bank_file, "r") # get the stat values rank_method = pystat.get_statistic_from_opts(args, [args.ifo]) stat = rank_method.get_sngl_ranking(trigf[args.ifo]) # get the duration values if needed -if args.bin_param == 'template_duration' and not args.f_lower: - logging.info('Using template duration from the trigger file') +if args.bin_param == "template_duration" and not args.f_lower: + logging.info("Using template duration from the trigger file") trig_dur = True else: trig_dur = False @@ -164,45 +226,52 @@ else: minth = min(args.stat_threshold) abovethresh = stat >= minth stat = stat[abovethresh] -tid = trigf[args.ifo+'/template_id'][:][abovethresh] -time = trigf[args.ifo+'/end_time'][:][abovethresh] +tid = trigf[args.ifo + "/template_id"][:][abovethresh] +time = trigf[args.ifo + "/end_time"][:][abovethresh] if trig_dur: - tdur = trigf[args.ifo+'/template_duration'][:][abovethresh] -logging.info('%i trigs left after thresholding at %f' % (len(stat), minth)) + tdur = trigf[args.ifo + "/template_duration"][:][abovethresh] +logging.info("%i trigs left after thresholding at %f" % (len(stat), minth)) # now do vetoing for veto_file, veto_segment_name in zip(args.veto_file, args.veto_segment_name): - retain, junk = events.veto.indices_outside_segments(time, [veto_file], - ifo=args.ifo, segment_name=veto_segment_name) + retain, junk = events.veto.indices_outside_segments( + time, [veto_file], ifo=args.ifo, segment_name=veto_segment_name + ) stat = stat[retain] tid = tid[retain] time = time[retain] if trig_dur: tdur = tdur[retain] - logging.info('%i trigs left after vetoing with %s' % - (len(stat), args.veto_file)) + logging.info("%i trigs left after vetoing with %s" % (len(stat), args.veto_file)) ### Functions for doing the pruning (removal of trigs at loudest times) + def get_pars(args, tag, m1, m2, s1z, s2z): # here used for both pruning and binning params - paramarg = getattr(args, tag+'_param') + paramarg = getattr(args, tag + "_param") try: # will fail if m1 is a float rather than a sequence - logging.info('Getting %s values for %i triggers' % (paramarg, len(m1))) + logging.info("Getting %s values for %i triggers" % (paramarg, len(m1))) except: pass return triggers.get_param(paramarg, args, m1, m2, s1z, s2z) + if args.prune_param: - logging.info('Getting min and max param values') - prpars = get_pars(args, 'prune', - templatef['mass1'][:], templatef['mass2'][:], - templatef['spin1z'][:], templatef['spin2z'][:]) + logging.info("Getting min and max param values") + prpars = get_pars( + args, + "prune", + templatef["mass1"][:], + templatef["mass2"][:], + templatef["spin1z"][:], + templatef["spin2z"][:], + ) minprpar = min(prpars) maxprpar = max(prpars) del prpars - logging.info('prune param range %f %f' % (minprpar, maxprpar)) + logging.info("prune param range %f %f" % (minprpar, maxprpar)) # hard-coded time window of 0.1s args.prune_window = 0.1 @@ -222,24 +291,31 @@ if args.prune_param: # are all the bins full already? numpruned = sum([len(prunedtimes[i]) for i in range(args.prune_bins)]) if numpruned == args.prune_bins * args.prune_number: - logging.info('Finished pruning!') + logging.info("Finished pruning!") break if numpruned > args.prune_bins * args.prune_number: - logging.error('Uh-oh, we pruned too many things .. %i, to be ' - 'precise' % numpruned) + logging.error( + "Uh-oh, we pruned too many things .. %i, to be precise" % numpruned + ) raise RuntimeError loudest = np.argmax(statpruneall) lstat = statpruneall[loudest] ltid = tidpruneall[loudest] ltime = timepruneall[loudest] m1, m2, s1z, s2z = triggers.get_mass_spin(templatef, ltid) - lbin = trstats.which_bin(get_pars(args, 'prune', m1, m2, s1z, s2z), - minprpar, maxprpar, - args.prune_bins, log=args.log_prune_param) + lbin = trstats.which_bin( + get_pars(args, "prune", m1, m2, s1z, s2z), + minprpar, + maxprpar, + args.prune_bins, + log=args.log_prune_param, + ) # is the bin where the loudest trigger lives full already? if len(prunedtimes[lbin]) == args.prune_number: - logging.info('%i - Bin %i full, not pruning event with stat %f at time ' - '%.3f' % (j, lbin, lstat, ltime)) + logging.info( + "%i - Bin %i full, not pruning event with stat %f at time " + "%.3f" % (j, lbin, lstat, ltime) + ) # prune the reference trigger array retain = abs(timepruneall - ltime) > args.prune_window statpruneall = statpruneall[retain] @@ -248,13 +324,15 @@ if args.prune_param: del retain continue else: - logging.info('Pruning event with stat %f at time %.3f in bin %i' % - (lstat, ltime, lbin)) + logging.info( + "Pruning event with stat %f at time %.3f in bin %i" + % (lstat, ltime, lbin) + ) # now do the pruning retain = abs(time - ltime) > args.prune_window - logging.info('%i trigs before pruning' % len(stat)) + logging.info("%i trigs before pruning" % len(stat)) stat = stat[retain] - logging.info('%i trigs remain' % len(stat)) + logging.info("%i trigs remain" % len(stat)) tid = tid[retain] time = time[retain] if trig_dur: @@ -276,16 +354,18 @@ if trig_dur: binpars = tdur + args.min_duration else: m1, m2, s1z, s2z = triggers.get_mass_spin(templatef, tid) - binpars = get_pars(args, 'bin', m1, m2, s1z, s2z) -logging.info("Parameter range of triggers: %f - %f" % - (min(binpars), max(binpars))) + binpars = get_pars(args, "bin", m1, m2, s1z, s2z) +logging.info("Parameter range of triggers: %f - %f" % (min(binpars), max(binpars))) # remove triggers outside irregular bins if args.bin_spacing == "irregular": - logging.info("Removing triggers outside bin range %f - %f" % - (min(args.bin_edges), max(args.bin_edges))) - in_range = np.logical_and(binpars >= min(args.bin_edges), - binpars <= max(args.bin_edges)) + logging.info( + "Removing triggers outside bin range %f - %f" + % (min(args.bin_edges), max(args.bin_edges)) + ) + in_range = np.logical_and( + binpars >= min(args.bin_edges), binpars <= max(args.bin_edges) + ) binpars = binpars[in_range] stat = stat[in_range] tid = tid[in_range] @@ -329,7 +409,18 @@ stdev = {} ks_prob = {} nabove = {} -histcolors = ['r',(1.0,0.6,0),'y','g','c','b','m','k',(0.8,0.25,0),(0.25,0.8,0)] +histcolors = [ + "r", + (1.0, 0.6, 0), + "y", + "g", + "c", + "b", + "m", + "k", + (0.8, 0.25, 0), + (0.25, 0.8, 0), +] for th in args.stat_threshold: logging.info("Fitting above threshold %f" % th) @@ -352,60 +443,87 @@ for th in args.stat_threshold: logging.info("No trigs in bin %f-%f", (lower, upper)) continue # do the fit - alpha, sig_alpha = trstats.fit_above_thresh( - args.fit_function, vals_inbin, th) + alpha, sig_alpha = trstats.fit_above_thresh(args.fit_function, vals_inbin, th) alphas[th][i] = alpha stdev[th][i] = sig_alpha - _, ks_prob[th][i] = trstats.KS_test( - args.fit_function, vals_inbin, alpha, th) + _, ks_prob[th][i] = trstats.KS_test(args.fit_function, vals_inbin, alpha, th) # add histogram to plot histcounts, edges = np.histogram(vals_inbin, bins=50) cum_counts = histcounts[::-1].cumsum()[::-1] binlabel = r"%.3g - %.3g" % (lower, upper) # histogram of fitted values - plt.semilogy(edges[:-1], cum_counts, linewidth=2, - color=histcolors[i], label=binlabel, alpha=0.6) + plt.semilogy( + edges[:-1], + cum_counts, + linewidth=2, + color=histcolors[i], + label=binlabel, + alpha=0.6, + ) # fit central value - plt.semilogy(plotrange, counts[th][i] * \ - trstats.cum_fit(args.fit_function, plotrange, alpha, th), - "--", color=histcolors[i], - label=r"$\alpha = $%.2f $\pm$ %.2f" % (alpha, sig_alpha)) + plt.semilogy( + plotrange, + counts[th][i] * trstats.cum_fit(args.fit_function, plotrange, alpha, th), + "--", + color=histcolors[i], + label=r"$\alpha = $%.2f $\pm$ %.2f" % (alpha, sig_alpha), + ) # 1sigma upper deviation on alpha - plt.semilogy(plotrange, counts[th][i] * \ - trstats.cum_fit(args.fit_function, plotrange, alpha + \ - sig_alpha, th), ":", alpha=0.6, color=histcolors[i]) + plt.semilogy( + plotrange, + counts[th][i] + * trstats.cum_fit(args.fit_function, plotrange, alpha + sig_alpha, th), + ":", + alpha=0.6, + color=histcolors[i], + ) # 1sigma lower deviation - plt.semilogy(plotrange, counts[th][i] * \ - trstats.cum_fit(args.fit_function, plotrange, alpha - \ - sig_alpha, th), ":", alpha=0.6, color=histcolors[i]) + plt.semilogy( + plotrange, + counts[th][i] + * trstats.cum_fit(args.fit_function, plotrange, alpha - sig_alpha, th), + ":", + alpha=0.6, + color=histcolors[i], + ) # finish the hist plot leg = plt.legend(labelspacing=0.2) - unitstring = " (%s)" % args.bin_param_units if \ - args.bin_param_units is not None else "" - leg.set_title(paramname+unitstring) + unitstring = ( + " (%s)" % args.bin_param_units if args.bin_param_units is not None else "" + ) + leg.set_title(paramname + unitstring) plt.setp(leg.get_texts(), fontsize=11) - plt.ylim(0.7, 2*maxcount) - plt.xlim(0.9*minth, 1.1*max(plotrange)) + plt.ylim(0.7, 2 * maxcount) + plt.xlim(0.9 * minth, 1.1 * max(plotrange)) plt.grid() plt.xlabel(statname, size="large") plt.ylabel("cumulative number", size="large") if args.plot_dir: - plt.title(args.ifo + " " + statname + " distribution split by " + \ - paramname) - dest = plotbase + "_" + args.sngl_stat + "_cdf_by_" + paramtag[0:3] + \ - "_fit_thresh_" + str(th) + ".png" + plt.title(args.ifo + " " + statname + " distribution split by " + paramname) + dest = ( + plotbase + + "_" + + args.sngl_stat + + "_cdf_by_" + + paramtag[0:3] + + "_fit_thresh_" + + str(th) + + ".png" + ) logging.info("Saving cumhist to %s" % dest) plt.savefig(dest) elif args.output_file: logging.info("Saving cumhist to %s" % args.output_file) results.save_fig_with_metadata( - fig, args.output_file, - title="%s: %s histogram of single detector triggers" % (args.ifo, - statname), - caption=(r"Histogram of single detector %s values binned by %s " - "with fitted %s distribution parameterized by α"\ - % (statname, paramname, args.fit_function)), - cmd=" ".join(sys.argv) + fig, + args.output_file, + title="%s: %s histogram of single detector triggers" % (args.ifo, statname), + caption=( + r"Histogram of single detector %s values binned by %s " + "with fitted %s distribution parameterized by α" + % (statname, paramname, args.fit_function) + ), + cmd=" ".join(sys.argv), ) plt.close() @@ -415,47 +533,60 @@ if args.output_file: # make plots of alpha, trig count and KS significance for th in args.stat_threshold: - plt.errorbar(pbins.centres(), [alphas[th][i] for i in binind], - yerr=[stdev[th][i] for i in binind], fmt="+-", - label=args.ifo + " fit above %.2f" % th) -if args.bin_spacing == "log": plt.semilogx() + plt.errorbar( + pbins.centres(), + [alphas[th][i] for i in binind], + yerr=[stdev[th][i] for i in binind], + fmt="+-", + label=args.ifo + " fit above %.2f" % th, + ) +if args.bin_spacing == "log": + plt.semilogx() plt.xlim(0.03, 150) plt.ylim(0, 10) plt.grid() plt.legend(loc="best") plt.xlabel(paramname, size="large") -plt.ylabel(r"fit parameter $\alpha$", size='large') -plt.savefig(plotbase + '_alpha_vs_' + paramtag[0:3] + '.png') +plt.ylabel(r"fit parameter $\alpha$", size="large") +plt.savefig(plotbase + "_alpha_vs_" + paramtag[0:3] + ".png") plt.close() for th in args.stat_threshold: - plt.errorbar(pbins.centres(), - [float(counts[th][i])/templates[i] for i in binind], - yerr=[counts[th][i]**0.5/templates[i] for i in binind], - fmt="+-", label=args.ifo + " trigs above %.2f" % th) -if args.bin_spacing == "log": plt.semilogx() + plt.errorbar( + pbins.centres(), + [float(counts[th][i]) / templates[i] for i in binind], + yerr=[counts[th][i] ** 0.5 / templates[i] for i in binind], + fmt="+-", + label=args.ifo + " trigs above %.2f" % th, + ) +if args.bin_spacing == "log": + plt.semilogx() plt.xlim(0.03, 150) plt.grid() plt.legend(loc="best") plt.xlabel(paramname, size="large") -plt.ylabel(r"Triggers above threshold per template", size='large') -plt.savefig(plotbase + '_nabove_vs_' + paramtag[0:3] + '.png') +plt.ylabel(r"Triggers above threshold per template", size="large") +plt.savefig(plotbase + "_nabove_vs_" + paramtag[0:3] + ".png") plt.close() for th in args.stat_threshold: - plt.plot(pbins.centres(), [ks_prob[th][i] for i in binind], - '+--', label=args.ifo+' KS prob, thresh %.2f' % th) -if args.bin_spacing == 'log': + plt.plot( + pbins.centres(), + [ks_prob[th][i] for i in binind], + "+--", + label=args.ifo + " KS prob, thresh %.2f" % th, + ) +if args.bin_spacing == "log": plt.loglog() else: plt.semilogy() plt.xlim(0.03, 150) plt.grid() -leg = plt.legend(loc='best', labelspacing=0.2) +leg = plt.legend(loc="best", labelspacing=0.2) plt.setp(leg.get_texts(), fontsize=11) -plt.xlabel(paramname, size='large') -plt.ylabel('KS test p-value') -plt.savefig(plotbase + '_KS_prob_vs_' + paramtag[0:3] + '.png') +plt.xlabel(paramname, size="large") +plt.ylabel("KS test p-value") +plt.savefig(plotbase + "_KS_prob_vs_" + paramtag[0:3] + ".png") plt.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_fit_sngls_by_template b/bin/all_sky_search/pycbc_fit_sngls_by_template index c7816bb7e23..e916f7d3444 100755 --- a/bin/all_sky_search/pycbc_fit_sngls_by_template +++ b/bin/all_sky_search/pycbc_fit_sngls_by_template @@ -13,20 +13,23 @@ # Public License for more details. -import argparse, logging +import argparse +import copy +import logging -import copy, numpy as np +import numpy as np import pycbc from pycbc import events, init_logging -from pycbc.events import triggers, trigger_fits as trstats +from pycbc.events import cuts, triggers from pycbc.events import stat as statsmod -from pycbc.events import cuts -from pycbc.types.optparse import MultiDetOptionAction +from pycbc.events import trigger_fits as trstats from pycbc.io import HFile +from pycbc.types.optparse import MultiDetOptionAction #### DEFINITIONS AND FUNCTIONS #### + def get_stat(args, trigs, threshold): """ Select the triggers and calculate the single detector statistic. @@ -46,6 +49,7 @@ def get_stat(args, trigs, threshold): A boolean array that selects the triggers. np.concatenate(stat): np.ndarray The statistic values of the selected triggers. + """ # For now this is using the single detector ranking. If we want, this # could use the Stat classes in stat.py using similar code as in hdf/io.py @@ -53,16 +57,16 @@ def get_stat(args, trigs, threshold): chunk_size = 2**23 stat = [] select = [] - size = len(trigs['end_time']) + size = len(trigs["end_time"]) s = 0 trigger_cut_dict, _ = cuts.ingest_cuts_option_group(args) while s < size: - e = s + chunk_size if (s + chunk_size) <= size else size + e = min(s + chunk_size, size) # read and format chunk of data so it can be read by key # as the stat classes expect. chunk = {k: trigs[k][s:e] for k in trigs if len(trigs[k]) == size} - trigger_keep_bool = np.zeros(e-s, dtype=bool) + trigger_keep_bool = np.zeros(e - s, dtype=bool) if len(trigger_cut_dict) > 0: # Apply trigger cuts @@ -87,71 +91,117 @@ def get_stat(args, trigs, threshold): # along with the stat values above threshold return np.concatenate(select), np.concatenate(stat) + #### MAIN #### -parser = argparse.ArgumentParser(usage="", +parser = argparse.ArgumentParser( + usage="", description="Perform maximum-likelihood fits of single inspiral trigger" - " distributions to various functions") + " distributions to various functions", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--trigger-file", - help="Input hdf5 file containing single triggers. " - "Required") -parser.add_argument("--bank-file", default=None, - help="hdf file containing template parameters. Required") -parser.add_argument("--template-fraction-range", default="0/1", - help="Optional, analyze only part of template bank. " - "Format is PART/NUM_PARTS") -parser.add_argument("--veto-file", nargs='*', default=[], action='append', - help="File(s) in .xml format with veto segments to apply " - "to triggers before fitting") -parser.add_argument("--veto-segment-name", nargs='*', default=[], action='append', - help="Name(s) of veto segments to apply. Optional, if not " - "given all segments for a given ifo will be used") -parser.add_argument("--gating-veto-windows", nargs='+', - action=MultiDetOptionAction, - help="Seconds to be vetoed before and after the central time " - "of each gate. Given as detector-values pairs, e.g. " - "H1:-1,2.5 L1:-1,2.5 V1:0,0") -parser.add_argument("--output", required=True, - help="Location for output file containing fit coefficients" - ". Required") -parser.add_argument("--ifo", required=True, - help="Ifo producing triggers to be fitted. Required") -parser.add_argument("--fit-function", - choices=["exponential", "rayleigh", "power"], - help="Functional form for the maximum likelihood fit") -parser.add_argument("--stat-threshold", type=float, - help="Only fit triggers with statistic value above this " - "threshold. Required. Typically 6-6.5") -parser.add_argument("--save-trig-param", - help="For each template, save a parameter value read from " - "its trigger(s). Ex. template_duration") -parser.add_argument("--prune-param", - help="Parameter to define bins for 'pruning' loud triggers" - " to make the fit insensitive to signals and outliers. " - "Choose from mchirp, mtotal, template_duration or a named " - "frequency cutoff in pnutils or a frequency function in " - "LALSimulation") -parser.add_argument("--prune-bins", type=int, - help="Number of bins to divide bank into when pruning") -parser.add_argument("--prune-number", type=int, - help="Number of loudest events to prune in each bin") -parser.add_argument("--log-prune-param", action='store_true', - help="Bin in the log of prune-param") -parser.add_argument("--f-lower", default=-1., - help="Starting frequency for calculating template " - "duration, required if this is the prune parameter") +parser.add_argument( + "--trigger-file", help="Input hdf5 file containing single triggers. Required" +) +parser.add_argument( + "--bank-file", + default=None, + help="hdf file containing template parameters. Required", +) +parser.add_argument( + "--template-fraction-range", + default="0/1", + help="Optional, analyze only part of template bank. Format is PART/NUM_PARTS", +) +parser.add_argument( + "--veto-file", + nargs="*", + default=[], + action="append", + help="File(s) in .xml format with veto segments to apply " + "to triggers before fitting", +) +parser.add_argument( + "--veto-segment-name", + nargs="*", + default=[], + action="append", + help="Name(s) of veto segments to apply. Optional, if not " + "given all segments for a given ifo will be used", +) +parser.add_argument( + "--gating-veto-windows", + nargs="+", + action=MultiDetOptionAction, + help="Seconds to be vetoed before and after the central time " + "of each gate. Given as detector-values pairs, e.g. " + "H1:-1,2.5 L1:-1,2.5 V1:0,0", +) +parser.add_argument( + "--output", + required=True, + help="Location for output file containing fit coefficients. Required", +) +parser.add_argument( + "--ifo", required=True, help="Ifo producing triggers to be fitted. Required" +) +parser.add_argument( + "--fit-function", + choices=["exponential", "rayleigh", "power"], + help="Functional form for the maximum likelihood fit", +) +parser.add_argument( + "--stat-threshold", + type=float, + help="Only fit triggers with statistic value above this " + "threshold. Required. Typically 6-6.5", +) +parser.add_argument( + "--save-trig-param", + help="For each template, save a parameter value read from " + "its trigger(s). Ex. template_duration", +) +parser.add_argument( + "--prune-param", + help="Parameter to define bins for 'pruning' loud triggers" + " to make the fit insensitive to signals and outliers. " + "Choose from mchirp, mtotal, template_duration or a named " + "frequency cutoff in pnutils or a frequency function in " + "LALSimulation", +) +parser.add_argument( + "--prune-bins", type=int, help="Number of bins to divide bank into when pruning" +) +parser.add_argument( + "--prune-number", type=int, help="Number of loudest events to prune in each bin" +) +parser.add_argument( + "--log-prune-param", action="store_true", help="Bin in the log of prune-param" +) +parser.add_argument( + "--f-lower", + default=-1.0, + help="Starting frequency for calculating template " + "duration, required if this is the prune parameter", +) # FIXME : support using the trigger file duration as prune parameter? # FIXME : have choice of SEOBNRv2 or PhenD duration formula ? -parser.add_argument("--min-duration", default=0., - help="Fudge factor for templates with tiny or negative " - "values of template_duration: add to duration values " - "before pruning. Units seconds") -parser.add_argument("--approximant", default="SEOBNRv4", - help="Approximant for template duration. Default SEOBNRv4") - -statsmod.insert_statistic_option_group(parser, - default_ranking_statistic='single_ranking_only') +parser.add_argument( + "--min-duration", + default=0.0, + help="Fudge factor for templates with tiny or negative " + "values of template_duration: add to duration values " + "before pruning. Units seconds", +) +parser.add_argument( + "--approximant", + default="SEOBNRv4", + help="Approximant for template duration. Default SEOBNRv4", +) + +statsmod.insert_statistic_option_group( + parser, default_ranking_statistic="single_ranking_only" +) cuts.insert_cuts_option_group(parser) args = parser.parse_args() @@ -163,25 +213,28 @@ args.veto_file = sum(args.veto_file, []) if len(args.veto_segment_name) != len(args.veto_file): raise RuntimeError("Number of veto files much match veto file names") -if (args.prune_param or args.prune_bins or args.prune_number) and not \ - (args.prune_param and args.prune_bins and args.prune_number): - raise RuntimeError("To prune, need to specify param, number of bins and " - "nonzero number to prune in each bin!") +if (args.prune_param or args.prune_bins or args.prune_number) and not ( + args.prune_param and args.prune_bins and args.prune_number +): + raise RuntimeError( + "To prune, need to specify param, number of bins and " + "nonzero number to prune in each bin!" + ) -logging.info('Fitting above threshold %f' % args.stat_threshold) +logging.info("Fitting above threshold %f" % args.stat_threshold) -logging.info('Opening trigger file: %s' % args.trigger_file) -trigf = HFile(args.trigger_file, 'r') -logging.info('Opening template file: %s' % args.bank_file) -templatef = HFile(args.bank_file, 'r') +logging.info("Opening trigger file: %s" % args.trigger_file) +trigf = HFile(args.trigger_file, "r") +logging.info("Opening template file: %s" % args.bank_file) +templatef = HFile(args.bank_file, "r") -logging.info('Counting number of triggers in each template') +logging.info("Counting number of triggers in each template") # template boundaries dataset is in order of template_id -tb = trigf[args.ifo+'/template_boundaries'][:] +tb = trigf[args.ifo + "/template_boundaries"][:] tid = np.arange(len(tb)) # template boundary values ascend in the same order as template hash # hence sort by hash -hash_sort = np.argsort(templatef['template_hash'][:]) +hash_sort = np.argsort(templatef["template_hash"][:]) tb_hashorder = tb[hash_sort] # reorder template IDs in parallel to the boundary values tid_hashorder = tid[hash_sort] @@ -189,21 +242,21 @@ tid_hashorder = tid[hash_sort] # Calculate the differences between the boundary indices to get the # number in each template # adding on total number at the end to get number in the last template -total_number = len(trigf[args.ifo + '/template_id']) +total_number = len(trigf[args.ifo + "/template_id"]) count_in_template_hashorder = np.diff(np.append(tb_hashorder, total_number)) # re-reorder values from hash order to tid order tid_sort = np.argsort(tid_hashorder) count_in_template = count_in_template_hashorder[tid_sort] # get the stat values -logging.info('Calculating stat values') +logging.info("Calculating stat values") abovethresh, stat = get_stat(args, trigf[args.ifo], args.stat_threshold) -logging.info('%i trigs left after thresholding' % len(stat)) +logging.info("%i trigs left after thresholding" % len(stat)) -tid = trigf[args.ifo + '/template_id'][abovethresh] -time = trigf[args.ifo + '/end_time'][abovethresh] +tid = trigf[args.ifo + "/template_id"][abovethresh] +time = trigf[args.ifo + "/end_time"][abovethresh] if args.save_trig_param: - tparam_region = trigf[args.ifo][args.save_trig_param + '_template'][:] + tparam_region = trigf[args.ifo][args.save_trig_param + "_template"][:] tparam = [] for region in tparam_region: try: @@ -214,60 +267,68 @@ if args.save_trig_param: # Calculate total time being analysed from segments # use set() to eliminate duplicates -segment_starts = sorted(set(trigf['{}/search/start_time'.format(args.ifo)][:])) -segment_ends = sorted(set(trigf['{}/search/end_time'.format(args.ifo)][:])) +segment_starts = sorted(set(trigf[f"{args.ifo}/search/start_time"][:])) +segment_ends = sorted(set(trigf[f"{args.ifo}/search/end_time"][:])) all_segments = events.veto.start_end_to_segments(segment_starts, segment_ends) # now do vetoing for veto_file, veto_segment_name in zip(args.veto_file, args.veto_segment_name): - retain, junk = events.veto.indices_outside_segments(time, [veto_file], - ifo=args.ifo, segment_name=veto_segment_name) - all_segments -= events.veto.select_segments_by_definer(veto_file, - ifo=args.ifo, - segment_name=veto_segment_name) + retain, junk = events.veto.indices_outside_segments( + time, [veto_file], ifo=args.ifo, segment_name=veto_segment_name + ) + all_segments -= events.veto.select_segments_by_definer( + veto_file, ifo=args.ifo, segment_name=veto_segment_name + ) stat = stat[retain] tid = tid[retain] time = time[retain] - logging.info('%i trigs left after vetoing with %s' % - (len(stat), veto_file)) + logging.info("%i trigs left after vetoing with %s" % (len(stat), veto_file)) # Include gating vetoes if args.gating_veto_windows: - gating_veto = args.gating_veto_windows[args.ifo].split(',') + gating_veto = args.gating_veto_windows[args.ifo].split(",") gveto_before = float(gating_veto[0]) gveto_after = float(gating_veto[1]) if gveto_before > 0 or gveto_after < 0: - raise ValueError("Gating veto window values must be negative before " - "gates and positive after gates.") + raise ValueError( + "Gating veto window values must be negative before " + "gates and positive after gates." + ) if not (gveto_before == 0 and gveto_after == 0): - autogate_times = np.unique(trigf[args.ifo + '/gating/auto/time'][:]) - if args.ifo + '/gating/file' in trigf: - detgate_times = trigf[args.ifo + '/gating/file/time'][:] + autogate_times = np.unique(trigf[args.ifo + "/gating/auto/time"][:]) + if args.ifo + "/gating/file" in trigf: + detgate_times = trigf[args.ifo + "/gating/file/time"][:] else: detgate_times = [] gate_times = np.concatenate((autogate_times, detgate_times)) - gveto_segs = events.veto.start_end_to_segments(gate_times + gveto_before, - gate_times + gveto_after).coalesce() + gveto_segs = events.veto.start_end_to_segments( + gate_times + gveto_before, gate_times + gveto_after + ).coalesce() all_segments -= gveto_segs - gveto_retain_idx = events.veto.indices_outside_times(time, - gate_times + gveto_before, - gate_times + gveto_after) + gveto_retain_idx = events.veto.indices_outside_times( + time, gate_times + gveto_before, gate_times + gveto_after + ) stat = stat[gveto_retain_idx] tid = tid[gveto_retain_idx] time = time[gveto_retain_idx] - logging.info('%i trigs left after vetoing triggers near gates' % len(stat)) + logging.info("%i trigs left after vetoing triggers near gates" % len(stat)) total_time = abs(all_segments) # do pruning (removal of trigs at N loudest times defined over param bins) if args.prune_param: - logging.info('Getting min and max param values') - pars = triggers.get_param(args.prune_param, args, - templatef['mass1'][:], templatef['mass2'][:], - templatef['spin1z'][:], templatef['spin2z'][:]) + logging.info("Getting min and max param values") + pars = triggers.get_param( + args.prune_param, + args, + templatef["mass1"][:], + templatef["mass2"][:], + templatef["spin1z"][:], + templatef["spin2z"][:], + ) minpar = min(pars) maxpar = max(pars) del pars - logging.info('%f %f' % (minpar, maxpar)) + logging.info("%f %f" % (minpar, maxpar)) # hard-coded time window of 0.1s args.prune_window = 0.1 @@ -287,25 +348,31 @@ if args.prune_param: # are all the bins full already? numpruned = sum([len(prunedtimes[i]) for i in range(args.prune_bins)]) if numpruned == args.prune_bins * args.prune_number: - logging.info('Finished pruning!') + logging.info("Finished pruning!") break if numpruned > args.prune_bins * args.prune_number: - logging.error('Uh-oh, we pruned too many things .. %i, to be ' - 'precise' % numpruned) + logging.error( + "Uh-oh, we pruned too many things .. %i, to be precise" % numpruned + ) raise RuntimeError loudest = np.argmax(statpruneall) lstat = statpruneall[loudest] ltid = tidpruneall[loudest] ltime = timepruneall[loudest] m1, m2, s1z, s2z = triggers.get_mass_spin(templatef, ltid) - lbin = trstats.which_bin(triggers.get_param(args.prune_param, args, - m1, m2, s1z, s2z), - minpar, maxpar, args.prune_bins, - log=args.log_prune_param) + lbin = trstats.which_bin( + triggers.get_param(args.prune_param, args, m1, m2, s1z, s2z), + minpar, + maxpar, + args.prune_bins, + log=args.log_prune_param, + ) # is the bin where the loudest trigger lives full already? if len(prunedtimes[lbin]) == args.prune_number: - logging.info('%i - Bin %i full, not pruning event with stat %f at ' - 'time %.3f' % (j, lbin, lstat, ltime)) + logging.info( + "%i - Bin %i full, not pruning event with stat %f at " + "time %.3f" % (j, lbin, lstat, ltime) + ) # prune the reference trigger array retain = abs(timepruneall - ltime) > args.prune_window statpruneall = statpruneall[retain] @@ -314,13 +381,15 @@ if args.prune_param: del retain continue else: - logging.info('Pruning event with stat %f at time %.3f in bin %i' % - (lstat, ltime, lbin)) + logging.info( + "Pruning event with stat %f at time %.3f in bin %i" + % (lstat, ltime, lbin) + ) # now do the pruning retain = abs(time - ltime) > args.prune_window - logging.info('%i trigs before pruning' % len(stat)) + logging.info("%i trigs before pruning" % len(stat)) stat = stat[retain] - logging.info('%i trigs remain' % len(stat)) + logging.info("%i trigs remain" % len(stat)) tid = tid[retain] time = time[retain] # also for the reference trig arrays @@ -334,13 +403,13 @@ if args.prune_param: del statpruneall del tidpruneall del timepruneall - logging.info('%i trigs remain after pruning loop' % len(stat)) + logging.info("%i trigs remain after pruning loop" % len(stat)) # parse template range -num_templates = len(templatef['template_hash']) +num_templates = len(templatef["template_hash"]) rangestr = args.template_fraction_range -part = int(rangestr.split('/')[0]) -pieces = int(rangestr.split('/')[1]) +part = int(rangestr.split("/")[0]) +pieces = int(rangestr.split("/")[1]) tmin = int(num_templates / float(pieces) * part) tmax = int(num_templates / float(pieces) * (part + 1)) trange = range(tmin, tmax) @@ -357,40 +426,41 @@ tid = tid[tsort] stat = stat[tsort] # Get range of tid values which are the same -left = np.searchsorted(tid, trange, side='left') -right = np.searchsorted(tid, trange, side='right') +left = np.searchsorted(tid, trange, side="left") +right = np.searchsorted(tid, trange, side="right") logging.info("Fitting ...") for j, tnum in enumerate(trange): - stat_in_template = stat[left[j]:right[j]] + stat_in_template = stat[left[j] : right[j]] count_above = len(stat_in_template) count_total = count_in_template[tnum] if count_above == 0: # default/sentinel value to indicate no data, shouldn't hurt if 1/alpha is averaged - alpha = -100. + alpha = -100.0 else: alpha, sig_alpha = trstats.fit_above_thresh( - args.fit_function, stat_in_template, args.stat_threshold) + args.fit_function, stat_in_template, args.stat_threshold + ) tids.append(tnum) counts_above.append(count_above) counts_total.append(count_total) fits.append(alpha) if args.save_trig_param: tpars.append(tparam[tnum]) - if (tnum % 1000 == 0): logging.info('Fitted template %i / %i' % - (tnum - tmin, tmax - tmin)) + if tnum % 1000 == 0: + logging.info("Fitted template %i / %i" % (tnum - tmin, tmax - tmin)) logging.info("Calculating median sigma for each template") -sigma_regions = trigf[args.ifo + '/sigmasq_template'][:] +sigma_regions = trigf[args.ifo + "/sigmasq_template"][:] median_sigma = [] for reg in sigma_regions: - strigs = trigf[args.ifo + '/sigmasq'][reg] + strigs = trigf[args.ifo + "/sigmasq"][reg] if len(strigs) == 0: median_sigma.append(np.nan) continue median_sigma.append(np.median(strigs) ** 0.5) -outfile = HFile(args.output, 'w') +outfile = HFile(args.output, "w") outfile.create_dataset("template_id", data=trange) outfile.create_dataset("count_above_thresh", data=counts_above) outfile.create_dataset("fit_coeff", data=fits) @@ -408,4 +478,4 @@ outfile.attrs.create("stat_threshold", data=args.stat_threshold) outfile.attrs.create("analysis_time", data=total_time) outfile.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_fit_sngls_over_multiparam b/bin/all_sky_search/pycbc_fit_sngls_over_multiparam index 0106c9d769c..ffd0cfa1dc8 100755 --- a/bin/all_sky_search/pycbc_fit_sngls_over_multiparam +++ b/bin/all_sky_search/pycbc_fit_sngls_over_multiparam @@ -13,13 +13,17 @@ # Public License for more details. -import argparse, logging, numpy +import argparse +import logging + +import numpy from scipy.stats import norm import pycbc +from pycbc import init_logging from pycbc.events import triggers from pycbc.io import HFile -from pycbc import init_logging + def dist(i1, i2, parvals, smoothing_width): """ @@ -29,12 +33,11 @@ def dist(i1, i2, parvals, smoothing_width): """ dsq = 0 for v, s in zip(parvals, smoothing_width): - dsq += (v[i2] - v[i1]) ** 2.0 / s ** 2.0 - return dsq ** 0.5 + dsq += (v[i2] - v[i1]) ** 2.0 / s**2.0 + return dsq**0.5 -def smooth_templates(nabove, invalphan, ntotal, template_idx, - weights=None): +def smooth_templates(nabove, invalphan, ntotal, template_idx, weights=None): """ Find the smoothed values according to the specified templates, weighted appropriately. @@ -73,9 +76,11 @@ def smooth_templates(nabove, invalphan, ntotal, template_idx, ntotal_t_smoothed = numpy.average(ntotal[template_idx], weights=weights) invalphan_mean = numpy.average(invalphan[template_idx], weights=weights) - return_tuple = (nabove_t_smoothed, - nabove_t_smoothed / invalphan_mean, - ntotal_t_smoothed) + return_tuple = ( + nabove_t_smoothed, + nabove_t_smoothed / invalphan_mean, + ntotal_t_smoothed, + ) return return_tuple @@ -84,18 +89,17 @@ def smooth_tophat(nabove, invalphan, ntotal, dists): Smooth templates using a tophat function with templates within unit dists """ - idx_within_area = numpy.flatnonzero(dists < 1.) - return smooth_templates(nabove, - invalphan, - ntotal, - idx_within_area) + idx_within_area = numpy.flatnonzero(dists < 1.0) + return smooth_templates(nabove, invalphan, ntotal, idx_within_area) + # This is the default number of triggers required for n_closest smoothing _default_total_trigs = 500 -def smooth_n_closest(nabove, invalphan, ntotal, dists, - total_trigs=_default_total_trigs): +def smooth_n_closest( + nabove, invalphan, ntotal, dists, total_trigs=_default_total_trigs +): """ Smooth templates according to the closest N templates No weighting is applied @@ -107,8 +111,12 @@ def smooth_n_closest(nabove, invalphan, ntotal, dists, # starting at closest ntcs = nabove[dist_sort].cumsum() templates_required = numpy.searchsorted(ntcs, total_trigs) + 1 - logging.debug("%d template(s) required to obtain %d(>%d) triggers", - templates_required, ntcs[templates_required - 1], total_trigs) + logging.debug( + "%d template(s) required to obtain %d(>%d) triggers", + templates_required, + ntcs[templates_required - 1], + total_trigs, + ) idx_to_smooth = dist_sort[:templates_required] return smooth_templates(nabove, invalphan, ntotal, idx_to_smooth) @@ -118,15 +126,15 @@ def smooth_distance_weighted(nabove, invalphan, ntotal, dists): Smooth templates weighted according to dists in a unit-width normal distribution, truncated at three sigma """ - idx_within_area = dists < 3. + idx_within_area = dists < 3.0 weights = norm.pdf(dists[idx_within_area]) - return smooth_templates(nabove, invalphan, ntotal, - idx_within_area, weights=weights) + return smooth_templates(nabove, invalphan, ntotal, idx_within_area, weights=weights) + _smooth_dist_func = { - 'smooth_tophat': smooth_tophat, - 'n_closest': smooth_n_closest, - 'distance_weighted': smooth_distance_weighted + "smooth_tophat": smooth_tophat, + "n_closest": smooth_n_closest, + "distance_weighted": smooth_distance_weighted, } @@ -140,8 +148,10 @@ def smooth(nabove, invalphan, ntotal, dists, smoothing_method, **kwargs): dists is an array of the distances of the templates from the template of interest """ - return _smooth_dist_func[smoothing_method](nabove, invalphan, - ntotal, dists, **kwargs) + return _smooth_dist_func[smoothing_method]( + nabove, invalphan, ntotal, dists, **kwargs + ) + # Number of smoothing lengths around the current template where # distances will be calculated @@ -149,9 +159,9 @@ def smooth(nabove, invalphan, ntotal, dists, smoothing_method, **kwargs): # templates to contain n triggers, which we cannot know beforehand _smooth_cut = { - 'smooth_tophat': 1, - 'n_closest': numpy.inf, - 'distance_weighted': 3, + "smooth_tophat": 1, + "n_closest": numpy.inf, + "distance_weighted": 3, } @@ -159,12 +169,14 @@ def report_percentage(i, length): """ Convenience function - report how long through the loop we are. Every ten percent + Parameters ---------- i: integer index being looped through length : integer number of loops we will go through in total + """ pc = int(numpy.floor(i / length * 100)) pc_last = int(numpy.floor((i - 1) / length * 100)) @@ -172,64 +184,102 @@ def report_percentage(i, length): logging.info(f"Template {i} out of {length} ({pc:.0f}%)") -parser = argparse.ArgumentParser(usage="", +parser = argparse.ArgumentParser( + usage="", description="Smooth (regress) the dependence of coefficients describing " - "single-ifo background trigger distributions on a template " - "parameter, to suppress random noise in the resulting " - "background model.") + "single-ifo background trigger distributions on a template " + "parameter, to suppress random noise in the resulting " + "background model.", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--template-fit-file", required=True, nargs='+', - help="hdf5 file(s) containing fit coefficients for each " - "individual template. Can smooth over multiple " - "files provided they correspond to the same bank " - "and fitting settings. Required") -parser.add_argument("--bank-file", required=True, - help="hdf file containing template parameters. Required") -parser.add_argument("--output", required=True, - help="Location for output file containing smoothed fit " - "coefficients. Required") -parser.add_argument("--fit-param", nargs='+', - help="Parameter(s) over which to regress the background " - "fit coefficients. Required. Either read from " - "template fit file or choose from mchirp, mtotal, " - "chi_eff, eta, tau_0, tau_3, template_duration, " - "a frequency cutoff in pnutils or a frequency function" - "in LALSimulation. To regress the background over " - "multiple parameters, provide them as a list.") -parser.add_argument("--approximant", default="SEOBNRv4", - help="Approximant for template duration. Default SEOBNRv4") -parser.add_argument("--f-lower", type=float, - help="Start frequency for calculating template duration.") -parser.add_argument("--min-duration", type=float, default=0., - help="Fudge factor for templates with tiny or negative " - "values of template_duration: add to duration values" - " before fitting. Units seconds.") -parser.add_argument("--log-param", nargs='+', - help="Take the log of the fit param before smoothing. " - "Must be a list corresponding to fit params.") -parser.add_argument("--smoothing-width", type=float, nargs='+', required=True, - help="Distance in the space of fit param values (or their" - " logs) to smooth over. Required. Must be a list " - "corresponding to fit params.") -parser.add_argument("--smoothing-method", default="smooth_tophat", - choices = _smooth_dist_func.keys(), - help="Method used to smooth the fit parameters; " - "'smooth_tophat' (default) finds all templates within " - "unit distance from the template of interest " - "(distance normalised by --smoothing-width). " - "'n_closest' adds the closest templates to " - "the smoothing until 500 triggers are reached. " - "'distance_weighted' weights the closest templates " - "with a normal distribution of width smoothing-width " - "truncated at three smoothing-widths.") -parser.add_argument("--smoothing-keywords", nargs='*', - help="Keywords for the smoothing function, supplied " - "as key:value pairs, e.g. total_trigs:500 to define " - "the number of templates for n_closest smoothing.") -parser.add_argument("--output-fits-by-template", action='store_true', - help="If given, will output the input file fits to " - "fit_by_template group.") +parser.add_argument( + "--template-fit-file", + required=True, + nargs="+", + help="hdf5 file(s) containing fit coefficients for each " + "individual template. Can smooth over multiple " + "files provided they correspond to the same bank " + "and fitting settings. Required", +) +parser.add_argument( + "--bank-file", + required=True, + help="hdf file containing template parameters. Required", +) +parser.add_argument( + "--output", + required=True, + help="Location for output file containing smoothed fit coefficients. Required", +) +parser.add_argument( + "--fit-param", + nargs="+", + help="Parameter(s) over which to regress the background " + "fit coefficients. Required. Either read from " + "template fit file or choose from mchirp, mtotal, " + "chi_eff, eta, tau_0, tau_3, template_duration, " + "a frequency cutoff in pnutils or a frequency function" + "in LALSimulation. To regress the background over " + "multiple parameters, provide them as a list.", +) +parser.add_argument( + "--approximant", + default="SEOBNRv4", + help="Approximant for template duration. Default SEOBNRv4", +) +parser.add_argument( + "--f-lower", type=float, help="Start frequency for calculating template duration." +) +parser.add_argument( + "--min-duration", + type=float, + default=0.0, + help="Fudge factor for templates with tiny or negative " + "values of template_duration: add to duration values" + " before fitting. Units seconds.", +) +parser.add_argument( + "--log-param", + nargs="+", + help="Take the log of the fit param before smoothing. " + "Must be a list corresponding to fit params.", +) +parser.add_argument( + "--smoothing-width", + type=float, + nargs="+", + required=True, + help="Distance in the space of fit param values (or their" + " logs) to smooth over. Required. Must be a list " + "corresponding to fit params.", +) +parser.add_argument( + "--smoothing-method", + default="smooth_tophat", + choices=_smooth_dist_func.keys(), + help="Method used to smooth the fit parameters; " + "'smooth_tophat' (default) finds all templates within " + "unit distance from the template of interest " + "(distance normalised by --smoothing-width). " + "'n_closest' adds the closest templates to " + "the smoothing until 500 triggers are reached. " + "'distance_weighted' weights the closest templates " + "with a normal distribution of width smoothing-width " + "truncated at three smoothing-widths.", +) +parser.add_argument( + "--smoothing-keywords", + nargs="*", + help="Keywords for the smoothing function, supplied " + "as key:value pairs, e.g. total_trigs:500 to define " + "the number of templates for n_closest smoothing.", +) +parser.add_argument( + "--output-fits-by-template", + action="store_true", + help="If given, will output the input file fits to fit_by_template group.", +) args = parser.parse_args() if args.smoothing_keywords: @@ -240,14 +290,16 @@ else: kwarg_dict = {} for inputstr in smooth_kwargs: try: - key, value = inputstr.split(':') - if key == 'total_trigs': + key, value = inputstr.split(":") + if key == "total_trigs": value = int(value) kwarg_dict[key] = value except ValueError: - err_txt = "--smoothing-keywords must take input in the " \ - "form KWARG1:VALUE1 KWARG2:VALUE2 KWARG3:VALUE3 ... " \ - "Received {}".format(' '.join(args.smoothing_keywords)) + err_txt = ( + "--smoothing-keywords must take input in the " + "form KWARG1:VALUE1 KWARG2:VALUE2 KWARG3:VALUE3 ... " + "Received {}".format(" ".join(args.smoothing_keywords)) + ) raise ValueError(err_txt) assert len(args.log_param) == len(args.fit_param) == len(args.smoothing_width) @@ -270,24 +322,24 @@ logging.info("Loading input template fits") # they use the same bank num_templates = None for filename in args.template_fit_file: - with HFile(filename, 'r') as fits: + with HFile(filename, "r") as fits: if num_templates is None: - num_templates = fits['template_id'].size - elif not num_templates == fits['template_id'].size: + num_templates = fits["template_id"].size + elif not num_templates == fits["template_id"].size: raise RuntimeError( "Input fit files correspond to different banks. " "This situation is not yet supported." ) # get attributes from the template-level fit for k in fits.attrs.keys(): - if k == 'analysis_time': + if k == "analysis_time": # For this attribute only, we want the mean - analysis_time += fits.attrs['analysis_time'] + analysis_time += fits.attrs["analysis_time"] continue if k not in attr_dict: # It's the first time we encounter this attribute attr_dict[k] = fits.attrs[k] - elif k == 'ifo': + elif k == "ifo": # We don't mind if this attribute is different, however the # output attributes will only correspond to the first file's # IFO. Warn if different IFOs are being used. @@ -295,7 +347,9 @@ for filename in args.template_fit_file: logging.warning( "Fit files correspond to different IFOs: %s and %s, " "only %s is being used for output file attributes", - attr_dict[k], fits.attrs[k], attr_dict[k], + attr_dict[k], + fits.attrs[k], + attr_dict[k], ) continue elif not attr_dict[k] == fits.attrs[k]: @@ -306,19 +360,19 @@ for filename in args.template_fit_file: raise RuntimeError(err_msg) # get template id and template parameter values - tid = numpy.concatenate((tid, fits['template_id'][:])) - nabove = numpy.concatenate((nabove, fits['count_above_thresh'][:])) - ntotal = numpy.concatenate((ntotal, fits['count_in_template'][:])) - alpha = numpy.concatenate((alpha, fits['fit_coeff'][:])) + tid = numpy.concatenate((tid, fits["template_id"][:])) + nabove = numpy.concatenate((nabove, fits["count_above_thresh"][:])) + ntotal = numpy.concatenate((ntotal, fits["count_in_template"][:])) + alpha = numpy.concatenate((alpha, fits["fit_coeff"][:])) try: - median_sigma = numpy.concatenate((median_sigma, fits['median_sigma'][:])) + median_sigma = numpy.concatenate((median_sigma, fits["median_sigma"][:])) except KeyError: - logging.info('Median_sigma dataset not present in input file') + logging.info("Median_sigma dataset not present in input file") median_sigma = None # For an exponential fit 1/alpha is linear in the trigger statistic values # so calculating weighted sums or averages of 1/alpha is appropriate -invalpha = 1. / alpha +invalpha = 1.0 / alpha invalphan = invalpha * nabove # convert the sum above into a mean @@ -337,8 +391,8 @@ if len(args.template_fit_file) > 1: tidsort = tid.argsort() # For each unique template id, find the range of identical template ids - left = numpy.searchsorted(tid[tidsort], tid_unique, side='left') - right = numpy.searchsorted(tid[tidsort], tid_unique, side='right') - 1 + left = numpy.searchsorted(tid[tidsort], tid_unique, side="left") + right = numpy.searchsorted(tid[tidsort], tid_unique, side="right") - 1 # Precompute the sums so we can quickly look up differences nasum = nabove[tidsort].cumsum() @@ -355,8 +409,7 @@ if len(args.template_fit_file) > 1: # we do not mess things up when nan values are given, so we # can't use the special cumsum fast option median_sigma = [ - numpy.nanmean(median_sigma[tidsort[l:r]]) - for l, r in zip(left, right) + numpy.nanmean(median_sigma[tidsort[l:r]]) for l, r in zip(left, right) ] if args.output_fits_by_template: @@ -364,30 +417,33 @@ if args.output_fits_by_template: # For more than one input fit file, these values are averaged over the same # template in different files fbt_dict = { - 'count_above_thresh': nabove, - 'count_in_template': ntotal, + "count_above_thresh": nabove, + "count_in_template": ntotal, } - with numpy.errstate(invalid='ignore'): + with numpy.errstate(invalid="ignore"): # If n_above is zero, then we'll get an 'invalid' warning as we # are dividing zero by zero. This is normal, and we'll deal with # those properly just below, so ignore this so people don't see # a warning and panic alpha = nabove / invalphan alpha[nabove == 0] = -100 - fbt_dict['fit_coeff'] = alpha - -n_required = _default_total_trigs if 'total_trigs' not in kwarg_dict \ - else kwarg_dict['total_trigs'] -if args.smoothing_method == 'n_closest' and n_required > nabove.sum(): + fbt_dict["fit_coeff"] = alpha + +n_required = ( + _default_total_trigs + if "total_trigs" not in kwarg_dict + else kwarg_dict["total_trigs"] +) +if args.smoothing_method == "n_closest" and n_required > nabove.sum(): logging.warning( "There are %.2f triggers above threshold, not enough to give a " "total of %d for smoothing", nabove.sum(), - n_required + n_required, ) -logging.info('Calculating template parameter values') -bank = HFile(args.bank_file, 'r') +logging.info("Calculating template parameter values") +bank = HFile(args.bank_file, "r") m1, m2, s1z, s2z = triggers.get_mass_spin(bank, tid) parvals = [] @@ -395,12 +451,12 @@ parnames = [] for param, slog in zip(args.fit_param, args.log_param): data = triggers.get_param(param, args, m1, m2, s1z, s2z) - if slog in ['false', 'False', 'FALSE']: - logging.info('Using param: %s', param) + if slog in ["false", "False", "FALSE"]: + logging.info("Using param: %s", param) parvals.append(data) parnames.append(param) - elif slog in ['true', 'True', 'TRUE']: - logging.info('Using log param: %s', param) + elif slog in ["true", "True", "TRUE"]: + logging.info("Using log param: %s", param) parvals.append(numpy.log(data)) parnames.append(f"log({param})") else: @@ -417,14 +473,14 @@ smoothed_vals = numpy.zeros((num_templates, 3)) # Handle the one-dimensional case of tophat smoothing separately # as it is easier to optimize computational performance. -if len(parvals) == 1 and args.smoothing_method == 'smooth_tophat': +if len(parvals) == 1 and args.smoothing_method == "smooth_tophat": logging.info("Using efficient 1D tophat smoothing") sort = parvals[0].argsort() parvals_0 = parvals[0][sort] ntotal = ntotal[sort] nabove = nabove[sort] invalphan = invalphan[sort] - + # For each template, find the range of nearby templates which fall within # the chosen window. left = numpy.searchsorted(parvals_0, parvals[0] - args.smoothing_width[0]) @@ -439,17 +495,18 @@ if len(parvals) == 1 and args.smoothing_method == 'smooth_tophat': num = right - left logging.info("Smoothing ...") - smoothed_vals[:,0] = (nasum[right] - nasum[left]) / num + smoothed_vals[:, 0] = (nasum[right] - nasum[left]) / num invmean = (invsum[right] - invsum[left]) / num - smoothed_vals[:,1] = smoothed_vals[:, 0] / invmean - smoothed_vals[:,2] = (ntsum[right] - ntsum[left]) / num + smoothed_vals[:, 1] = smoothed_vals[:, 0] / invmean + smoothed_vals[:, 2] = (ntsum[right] - ntsum[left]) / num elif numpy.isfinite(_smooth_cut[args.smoothing_method]): c = _smooth_cut[args.smoothing_method] cut_lengths = [s * c for s in args.smoothing_width] # Find the "longest" dimension in cut lengths - sort_dim = numpy.argmax([(v.max() - v.min()) / c - for v, c in zip(parvals, cut_lengths)]) + sort_dim = numpy.argmax( + [(v.max() - v.min()) / c for v, c in zip(parvals, cut_lengths)] + ) logging.info("Sorting / Cutting on dimension %s", parnames[sort_dim]) # Sort parvals by the sort dimension @@ -458,17 +515,22 @@ elif numpy.isfinite(_smooth_cut[args.smoothing_method]): # For each template, find the range of nearby templates which fall within # the chosen window. - lefts = numpy.searchsorted(parvals[sort_dim], - parvals[sort_dim] - cut_lengths[sort_dim]) - rights = numpy.searchsorted(parvals[sort_dim], - parvals[sort_dim] + cut_lengths[sort_dim]) + lefts = numpy.searchsorted( + parvals[sort_dim], parvals[sort_dim] - cut_lengths[sort_dim] + ) + rights = numpy.searchsorted( + parvals[sort_dim], parvals[sort_dim] + cut_lengths[sort_dim] + ) n_removed = num_templates - rights + lefts - logging.info("Cutting between %d and %d templates for each smoothing", - n_removed.min(), n_removed.max()) + logging.info( + "Cutting between %d and %d templates for each smoothing", + n_removed.min(), + n_removed.max(), + ) # Sort the values to be smoothed by parameter value logging.info("Smoothing ...") - nabove_sort = nabove[par_sort] + nabove_sort = nabove[par_sort] invalphan_sort = invalphan[par_sort] ntotal_sort = ntotal[par_sort] slices = [slice(l, r) for l, r in zip(lefts, rights)] @@ -477,13 +539,13 @@ elif numpy.isfinite(_smooth_cut[args.smoothing_method]): slc = slices[i] d = dist(i, slc, parvals, args.smoothing_width) - smoothed_vals[i,:] = smooth( + smoothed_vals[i, :] = smooth( nabove_sort[slc], invalphan_sort[slc], ntotal_sort[slc], d, args.smoothing_method, - **kwarg_dict + **kwarg_dict, ) # Undo the sorts @@ -497,31 +559,26 @@ else: report_percentage(i, num_templates) d = dist(i, rang, parvals, args.smoothing_width) smoothed_vals[i, :] = smooth( - nabove, - invalphan, - ntotal, - d, - args.smoothing_method, - **kwarg_dict + nabove, invalphan, ntotal, d, args.smoothing_method, **kwarg_dict ) logging.info("Writing output") -outfile = HFile(args.output, 'w') -outfile['template_id'] = tid -outfile['count_above_thresh'] = smoothed_vals[:, 0] -outfile['fit_coeff'] = smoothed_vals[:, 1] -outfile['count_in_template'] = smoothed_vals[:, 2] +outfile = HFile(args.output, "w") +outfile["template_id"] = tid +outfile["count_above_thresh"] = smoothed_vals[:, 0] +outfile["fit_coeff"] = smoothed_vals[:, 1] +outfile["count_in_template"] = smoothed_vals[:, 2] if median_sigma is not None: - outfile['median_sigma'] = median_sigma + outfile["median_sigma"] = median_sigma for param, vals, slog in zip(args.fit_param, parvals, args.log_param): - if slog in ['false', 'False', 'FALSE']: + if slog in ["false", "False", "FALSE"]: outfile[param] = vals - elif slog in ['true', 'True', 'TRUE']: + elif slog in ["true", "True", "TRUE"]: outfile[param] = numpy.exp(vals) if args.output_fits_by_template: - fbt_group = outfile.create_group('fit_by_template') + fbt_group = outfile.create_group("fit_by_template") for k, v in fbt_dict.items(): fbt_group[k] = v @@ -529,8 +586,8 @@ if args.output_fits_by_template: for k, v in attr_dict.items(): outfile.attrs[k] = v if not analysis_time == 0: - outfile.attrs['analysis_time'] = analysis_time + outfile.attrs["analysis_time"] = analysis_time # Add a magic file attribute so that coinc_findtrigs can parse it -outfile.attrs['stat'] = attr_dict['ifo'] + '-fit_coeffs' -logging.info('Done!') +outfile.attrs["stat"] = attr_dict["ifo"] + "-fit_coeffs" +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_fit_sngls_over_param b/bin/all_sky_search/pycbc_fit_sngls_over_param index 4d77ee9559e..7f154b8a45c 100644 --- a/bin/all_sky_search/pycbc_fit_sngls_over_param +++ b/bin/all_sky_search/pycbc_fit_sngls_over_param @@ -13,196 +13,242 @@ # Public License for more details. -import argparse, logging +import argparse +import logging import numpy as np import pycbc from pycbc import init_logging -from pycbc.io import HFile from pycbc.events import triggers +from pycbc.io import HFile -parser = argparse.ArgumentParser(usage="", +parser = argparse.ArgumentParser( + usage="", description="Smooth (regress) the dependence of coefficients describing " - "single-ifo background trigger distributions on a template " - "parameter, to suppress random noise in the resulting " - "background model.") + "single-ifo background trigger distributions on a template " + "parameter, to suppress random noise in the resulting " + "background model.", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--template-fit-file", - help="Input hdf5 file containing fit coefficients for each" - " individual template. Required") -parser.add_argument("--bank-file", default=None, - help="hdf file containing template parameters. Required " - "unless reading param from template fit file") -parser.add_argument("--output", required=True, - help="Location for output file containing smoothed fit " - "coefficients. Required") -parser.add_argument("--use-template-fit-param", action="store_true", - help="Use parameter values stored in the template fit file" - "as template_param for smoothing.", default=False) -parser.add_argument("--fit-param", - help="Parameter over which to regress the background " - "fit coefficients. Required. Either read from " - "template fit file or choose from mchirp, mtotal, " - "tau_0, tau_3, template_duration, a frequency " - "cutoff in pnutils or a frequency function in " - "LALSimulation.") -parser.add_argument("--approximant", default="SEOBNRv4", - help="Approximant for template duration. Default SEOBNRv4") -parser.add_argument("--f-lower", type=float, default=0., - help="Starting frequency for calculating template " - "duration, if not reading from the template fit file") -parser.add_argument("--min-duration", type=float, default=0., - help="Fudge factor for templates with tiny or negative " - "values of template_duration: add to duration values" - " before fitting. Units seconds.") -parser.add_argument("--log-param", action='store_true', - help="Take the log of the fit param before smoothing.") -parser.add_argument("--regression-method", required=True, - choices=["nn", "tricube"], - help="Method of smoothing over the chosen fit param. " - "Required.") -parser.add_argument("--num-neighbors", type=int, default=-1, - help="Number of neighbors used in nn method. Try 2500, or " - "1/10th the total number of templates if that is " - "smaller.") -parser.add_argument("--smoothing-width", type=float, required=True, - help="Distance in the space of fit param values (or the " - "logs of them) to smooth over. Required. For log " - "template duration, try 0.2") +parser.add_argument( + "--template-fit-file", + help="Input hdf5 file containing fit coefficients for each" + " individual template. Required", +) +parser.add_argument( + "--bank-file", + default=None, + help="hdf file containing template parameters. Required " + "unless reading param from template fit file", +) +parser.add_argument( + "--output", + required=True, + help="Location for output file containing smoothed fit coefficients. Required", +) +parser.add_argument( + "--use-template-fit-param", + action="store_true", + help="Use parameter values stored in the template fit file" + "as template_param for smoothing.", + default=False, +) +parser.add_argument( + "--fit-param", + help="Parameter over which to regress the background " + "fit coefficients. Required. Either read from " + "template fit file or choose from mchirp, mtotal, " + "tau_0, tau_3, template_duration, a frequency " + "cutoff in pnutils or a frequency function in " + "LALSimulation.", +) +parser.add_argument( + "--approximant", + default="SEOBNRv4", + help="Approximant for template duration. Default SEOBNRv4", +) +parser.add_argument( + "--f-lower", + type=float, + default=0.0, + help="Starting frequency for calculating template " + "duration, if not reading from the template fit file", +) +parser.add_argument( + "--min-duration", + type=float, + default=0.0, + help="Fudge factor for templates with tiny or negative " + "values of template_duration: add to duration values" + " before fitting. Units seconds.", +) +parser.add_argument( + "--log-param", + action="store_true", + help="Take the log of the fit param before smoothing.", +) +parser.add_argument( + "--regression-method", + required=True, + choices=["nn", "tricube"], + help="Method of smoothing over the chosen fit param. Required.", +) +parser.add_argument( + "--num-neighbors", + type=int, + default=-1, + help="Number of neighbors used in nn method. Try 2500, or " + "1/10th the total number of templates if that is " + "smaller.", +) +parser.add_argument( + "--smoothing-width", + type=float, + required=True, + help="Distance in the space of fit param values (or the " + "logs of them) to smooth over. Required. For log " + "template duration, try 0.2", +) args = parser.parse_args() -if args.regression_method == 'nn' and args.num_neighbors < 1: +if args.regression_method == "nn" and args.num_neighbors < 1: raise RuntimeError("Need to give a positive number of nearest neighbors!") init_logging(args.verbose) -fits = HFile(args.template_fit_file, 'r') +fits = HFile(args.template_fit_file, "r") # get the ifo from the template-level fit -ifo = fits.attrs['ifo'] +ifo = fits.attrs["ifo"] # get template id and template parameter values -tid = fits['template_id'][:] +tid = fits["template_id"][:] if args.use_template_fit_param: - logging.info('Reading %s values from the template fit file' % - args.fit_param) + logging.info("Reading %s values from the template fit file" % args.fit_param) # check we're asking for the right param name - assert fits.attrs['save_trig_param'] == args.fit_param - parvals = fits['template_param'][:] + assert fits.attrs["save_trig_param"] == args.fit_param + parvals = fits["template_param"][:] else: - logging.info('Calculating template parameter values') - bank = HFile(args.bank_file, 'r') + logging.info("Calculating template parameter values") + bank = HFile(args.bank_file, "r") m1, m2, s1z, s2z = triggers.get_mass_spin(bank, tid) - parvals = triggers.get_param(args.fit_param, args, m1, m2, s1z, s2z) + parvals = triggers.get_param(args.fit_param, args, m1, m2, s1z, s2z) -if 'count_in_template' in fits.keys(): # recently introduced extra dataset +if "count_in_template" in fits.keys(): # recently introduced extra dataset tcount = True else: tcount = False -nabove = fits['count_above_thresh'][:] -if tcount: ntotal = fits['count_in_template'][:] +nabove = fits["count_above_thresh"][:] +if tcount: + ntotal = fits["count_in_template"][:] # for an exponential fit 1/alpha is linear in the trigger statistic values # so taking weighted sums/averages of 1/alpha is appropriate -invalpha = 1./(fits['fit_coeff'][:]) +invalpha = 1.0 / (fits["fit_coeff"][:]) # sort in ascending parameter order parsort = np.argsort(parvals) tid = tid[parsort] parvals = parvals[parsort] nabove = nabove[parsort] -if tcount: ntotal = ntotal[parsort] +if tcount: + ntotal = ntotal[parsort] invalpha = invalpha[parsort] if args.log_param: - logging.info('Using log %s to perform smoothing' % args.fit_param) + logging.info("Using log %s to perform smoothing" % args.fit_param) parvals = np.log(parvals) else: - logging.info('Using %s to perform smoothing' % args.fit_param) + logging.info("Using %s to perform smoothing" % args.fit_param) # do nearest-neighbours regression # use Gaussian weight over fitting parameter -if args.regression_method == 'nn': +if args.regression_method == "nn": # only import scikit-learn if and when needed from sklearn import neighbors - weights = lambda d:np.exp(-0.5 * (d/args.smoothing_width)**2) + + weights = lambda d: np.exp(-0.5 * (d / args.smoothing_width) ** 2) knn = neighbors.KNeighborsRegressor(args.num_neighbors, weights=weights) - logging.info('Smoothing nabove data') + logging.info("Smoothing nabove data") nabove_knn = knn.fit(parvals[:, np.newaxis], nabove) - logging.info('Evaluating smoothed nabove') + logging.info("Evaluating smoothed nabove") nabove_smoothed = [nabove_knn.predict([[p]]) for p in parvals] del nabove_knn if tcount: - logging.info('Smoothing ntotal data') + logging.info("Smoothing ntotal data") ntotal_knn = knn.fit(parvals[:, np.newaxis], ntotal) - logging.info('Evaluating smoothed ntotal') + logging.info("Evaluating smoothed ntotal") ntotal_smoothed = [ntotal_knn.predict([[p]]) for p in parvals] del ntotal_knn - logging.info('Smoothing invalpha data') + logging.info("Smoothing invalpha data") # smooth the inverse alpha values times trig count above threshold invalphan = invalpha * nabove invalphan_knn = knn.fit(parvals[:, np.newaxis], invalphan) - logging.info('Evaluating smoothed invalpha') + logging.info("Evaluating smoothed invalpha") # loop over parameter values to avoid memory error in knn predict invalphan_smoothed = [invalphan_knn.predict([[p]]) for p in parvals] del invalphan_knn # divide out by smoothed trig count invalpha_smoothed = np.array(invalphan_smoothed) / np.array(nabove_smoothed) -elif args.regression_method == 'tricube': +elif args.regression_method == "tricube": # do Nadaraya-Watson kernel regression # i.e. weighted average with parameter-dependent weights invalphan = invalpha * nabove invalpha_smoothed = [] nabove_smoothed = [] - if tcount: ntotal_smoothed = [] - logging.info('Evaluating smoothed invalpha and n_above') + if tcount: + ntotal_smoothed = [] + logging.info("Evaluating smoothed invalpha and n_above") for i, p in enumerate(parvals): if not i % 100: - logging.info('Smoothing template %i', i) + logging.info("Smoothing template %i", i) # find parameter values that go into average par_diff = parvals - p in_kernel = abs(par_diff) < args.smoothing_width # tri-cube kernel - weights = (1 - abs(par_diff[in_kernel])**3)**3 + weights = (1 - abs(par_diff[in_kernel]) ** 3) ** 3 norm = weights.sum() # weighted average n_sm = (nabove[in_kernel] * weights).sum() / norm - if tcount: nt_sm = (ntotal[in_kernel] * weights).sum() / norm + if tcount: + nt_sm = (ntotal[in_kernel] * weights).sum() / norm invalphan_smoothed = (invalphan[in_kernel] * weights).sum() / norm invalpha_smoothed.append(invalphan_smoothed / n_sm) nabove_smoothed.append(n_sm) ntotal_smoothed.append(nt_sm) # store template-dependent fit output -outfile = HFile(args.output, 'w') -outfile.create_dataset('template_id', data=tid) -outfile.create_dataset('count_above_thresh', data=nabove_smoothed) -if tcount: outfile.create_dataset('count_in_template', data=ntotal_smoothed) -outfile.create_dataset('fit_coeff', data=1. / np.array(invalpha_smoothed)) -outfile.create_dataset('param_val', data=(np.exp(parvals) if args.log_param - else parvals)) +outfile = HFile(args.output, "w") +outfile.create_dataset("template_id", data=tid) +outfile.create_dataset("count_above_thresh", data=nabove_smoothed) +if tcount: + outfile.create_dataset("count_in_template", data=ntotal_smoothed) +outfile.create_dataset("fit_coeff", data=1.0 / np.array(invalpha_smoothed)) +outfile.create_dataset( + "param_val", data=(np.exp(parvals) if args.log_param else parvals) +) # add metadata, some is inherited from template level fit -outfile.attrs.create('ifo', data=ifo) -outfile.attrs.create('stat_threshold', data=fits.attrs['stat_threshold']) -outfile.attrs.create('fit_param', data=args.fit_param) -outfile.attrs.create('regression_method', data=args.regression_method) +outfile.attrs.create("ifo", data=ifo) +outfile.attrs.create("stat_threshold", data=fits.attrs["stat_threshold"]) +outfile.attrs.create("fit_param", data=args.fit_param) +outfile.attrs.create("regression_method", data=args.regression_method) if args.num_neighbors > 0: - outfile.attrs.create('n_neighbors', data=args.num_neighbors) -outfile.attrs.create('smoothing_width', data=args.smoothing_width) -if 'analysis_time' in fits.attrs: - outfile.attrs['analysis_time'] = fits.attrs['analysis_time'] + outfile.attrs.create("n_neighbors", data=args.num_neighbors) +outfile.attrs.create("smoothing_width", data=args.smoothing_width) +if "analysis_time" in fits.attrs: + outfile.attrs["analysis_time"] = fits.attrs["analysis_time"] # add a magic file attribute so that coinc_findtrigs can parse it -outfile.attrs.create('stat', data=ifo+'-fit_coeffs') +outfile.attrs.create("stat", data=ifo + "-fit_coeffs") outfile.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_fit_sngls_split_binned b/bin/all_sky_search/pycbc_fit_sngls_split_binned index c7aa2714627..d732eac900c 100644 --- a/bin/all_sky_search/pycbc_fit_sngls_split_binned +++ b/bin/all_sky_search/pycbc_fit_sngls_split_binned @@ -13,148 +13,220 @@ # Public License for more details. +import argparse +import logging import sys -import argparse, logging from matplotlib import use -use('Agg') -from matplotlib import pyplot as plt + +use("Agg") import numpy as np +import pycbc.version +from matplotlib import pyplot as plt -from pycbc import events, bin_utils, results -from pycbc.io import SingleDetTriggers, HFile -from pycbc.events import triggers as trigs -from pycbc.events import trigger_fits as trstats +from pycbc import bin_utils, events, results from pycbc.events import stat as pystat -from pycbc.types.optparse import MultiDetOptionAction +from pycbc.events import trigger_fits as trstats +from pycbc.events import triggers as trigs +from pycbc.io import HFile, SingleDetTriggers from pycbc.tmpltbank import bank_conversions -import pycbc.version +from pycbc.types.optparse import MultiDetOptionAction -parser = argparse.ArgumentParser(usage="", - description="Plot histograms of triggers split over various parameters") +parser = argparse.ArgumentParser( + usage="", description="Plot histograms of triggers split over various parameters" +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--trigger-file", required=True, - help="Input hdf5 file containing single triggers. " - "Required") -parser.add_argument("--bank-file", default=None, required=True, - help="hdf file containing template parameters. Required") -parser.add_argument('--output-file', required=True, - help="Output image file. Required") -parser.add_argument('--bin-param', default='template_duration', - choices=bank_conversions.conversion_options, - help="Parameter for binning within plots, default " - "'template_duration'") -parser.add_argument('--bin-spacing', choices=['linear', 'log'], default='log', - help="How to space bin-param bin edges. " - "Choices=[linear, log], default log") -parser.add_argument('--num-bins', type=int, default=6, - help="Number of bins over which to split bin-param, default 6") -parser.add_argument('--max-bin-param', type=float, default=None, - help="Maximum allowed value of bin-param") -parser.add_argument('--split-param-one', default='eta', - choices=bank_conversions.conversion_options, - help="Parameter for splitting plot grid in y-direction, " - "default 'eta'") -parser.add_argument('--split-param-two', default='chi_eff', - choices=bank_conversions.conversion_options, - help="Parameter for splitting plot grid in x-direction, " - "default 'chi_eff'") -parser.add_argument('--split-one-nbins', default=3, type=int, - help="Number of split plots over split-param-one, default 3") -parser.add_argument('--split-two-nbins', default=4, type=int, - help="Number of split plots over split-param-two, default 4") -parser.add_argument('--split-one-spacing', choices=['linear', 'log'], default='linear', - help="How to space split-param-one bin edges. " - "Choices=[linear, log], default=linear") -parser.add_argument('--split-two-spacing', choices=['linear', 'log'], default='linear', - help="How to space split-param-two bin edges. " - "Choices=[linear, log], default=linear") -parser.add_argument('--ifo', help="Detector. Required") -parser.add_argument('--plot-max-x', type=float, default=None, - help="Maximum stat value to plot, if not given " - "1.05 * largest stat value will be used") -parser.add_argument("--veto-file", - help="File(s) in .xml format with veto segments to apply " - "to triggers before fitting") -parser.add_argument("--veto-segment-name", - help="Name(s) of veto segments to apply. Optional, if not " - "given all triggers for the given ifo will be used") -parser.add_argument("--gating-veto-windows", nargs='+', - action=MultiDetOptionAction, - help="Seconds to be vetoed before and after the central time " - "of each gate. Given as detector-values pairs, e.g. " - "H1:-1,2.5 L1:-1,2.5 V1:0,0") -parser.add_argument("--stat-fit-threshold", type=float, required=True, - help="Only fit triggers with statistic value above this " - "threshold. Required") -parser.add_argument("--plot-lower-stat-limit", type=float, required=True, - help="Plot triggers down to this value. Setting this too" - "low will incur huge memory usage in a full search." - "To avoid this, choose 5.5 or larger.") -parser.add_argument("--fit-function", - choices=["exponential", "rayleigh", "power"], - help="Functional form for the maximum likelihood fit") -parser.add_argument("--prune-number", type=int, default=0, - help="Number of loudest events to remove from each split " - "histogram, default 0") -parser.add_argument("--prune-window", type=float, default=0.1, - help="Time (s) to remove all triggers around a trigger " - "which is loudest in each split, default 0.1s") - -pystat.insert_statistic_option_group(parser, - default_ranking_statistic='single_ranking_only') +parser.add_argument( + "--trigger-file", + required=True, + help="Input hdf5 file containing single triggers. Required", +) +parser.add_argument( + "--bank-file", + default=None, + required=True, + help="hdf file containing template parameters. Required", +) +parser.add_argument("--output-file", required=True, help="Output image file. Required") +parser.add_argument( + "--bin-param", + default="template_duration", + choices=bank_conversions.conversion_options, + help="Parameter for binning within plots, default 'template_duration'", +) +parser.add_argument( + "--bin-spacing", + choices=["linear", "log"], + default="log", + help="How to space bin-param bin edges. Choices=[linear, log], default log", +) +parser.add_argument( + "--num-bins", + type=int, + default=6, + help="Number of bins over which to split bin-param, default 6", +) +parser.add_argument( + "--max-bin-param", + type=float, + default=None, + help="Maximum allowed value of bin-param", +) +parser.add_argument( + "--split-param-one", + default="eta", + choices=bank_conversions.conversion_options, + help="Parameter for splitting plot grid in y-direction, default 'eta'", +) +parser.add_argument( + "--split-param-two", + default="chi_eff", + choices=bank_conversions.conversion_options, + help="Parameter for splitting plot grid in x-direction, default 'chi_eff'", +) +parser.add_argument( + "--split-one-nbins", + default=3, + type=int, + help="Number of split plots over split-param-one, default 3", +) +parser.add_argument( + "--split-two-nbins", + default=4, + type=int, + help="Number of split plots over split-param-two, default 4", +) +parser.add_argument( + "--split-one-spacing", + choices=["linear", "log"], + default="linear", + help="How to space split-param-one bin edges. " + "Choices=[linear, log], default=linear", +) +parser.add_argument( + "--split-two-spacing", + choices=["linear", "log"], + default="linear", + help="How to space split-param-two bin edges. " + "Choices=[linear, log], default=linear", +) +parser.add_argument("--ifo", help="Detector. Required") +parser.add_argument( + "--plot-max-x", + type=float, + default=None, + help="Maximum stat value to plot, if not given " + "1.05 * largest stat value will be used", +) +parser.add_argument( + "--veto-file", + help="File(s) in .xml format with veto segments to apply " + "to triggers before fitting", +) +parser.add_argument( + "--veto-segment-name", + help="Name(s) of veto segments to apply. Optional, if not " + "given all triggers for the given ifo will be used", +) +parser.add_argument( + "--gating-veto-windows", + nargs="+", + action=MultiDetOptionAction, + help="Seconds to be vetoed before and after the central time " + "of each gate. Given as detector-values pairs, e.g. " + "H1:-1,2.5 L1:-1,2.5 V1:0,0", +) +parser.add_argument( + "--stat-fit-threshold", + type=float, + required=True, + help="Only fit triggers with statistic value above this threshold. Required", +) +parser.add_argument( + "--plot-lower-stat-limit", + type=float, + required=True, + help="Plot triggers down to this value. Setting this too" + "low will incur huge memory usage in a full search." + "To avoid this, choose 5.5 or larger.", +) +parser.add_argument( + "--fit-function", + choices=["exponential", "rayleigh", "power"], + help="Functional form for the maximum likelihood fit", +) +parser.add_argument( + "--prune-number", + type=int, + default=0, + help="Number of loudest events to remove from each split histogram, default 0", +) +parser.add_argument( + "--prune-window", + type=float, + default=0.1, + help="Time (s) to remove all triggers around a trigger " + "which is loudest in each split, default 0.1s", +) + +pystat.insert_statistic_option_group( + parser, default_ranking_statistic="single_ranking_only" +) args = parser.parse_args() -assert(args.stat_fit_threshold >= args.plot_lower_stat_limit) +assert args.stat_fit_threshold >= args.plot_lower_stat_limit pycbc.init_logging(args.verbose) -logging.info('Opening trigger file: %s' % args.trigger_file) -trigf = HFile(args.trigger_file, 'r') +logging.info("Opening trigger file: %s" % args.trigger_file) +trigf = HFile(args.trigger_file, "r") -logging.info('Opening template file: %s' % args.bank_file) -bank = HFile(args.bank_file, 'r') -logging.info('Getting template bank parameters') +logging.info("Opening template file: %s" % args.bank_file) +bank = HFile(args.bank_file, "r") +logging.info("Getting template bank parameters") # Calculate params needed usedparams = [args.split_param_two, args.split_param_one, args.bin_param] params = {} for par in usedparams: - if par in ['template_duration', 'duration']: - logging.info('Reading duration from trigger file') + if par in ["template_duration", "duration"]: + logging.info("Reading duration from trigger file") # Loop over template regions, return the first duration value from each region. dur = [] - for ref in trigf[args.ifo + '/template_duration_template'][:]: - dur.append(trigf[args.ifo + '/template_duration'][ref][0]) - params['template_duration'] = np.array(dur) + for ref in trigf[args.ifo + "/template_duration_template"][:]: + dur.append(trigf[args.ifo + "/template_duration"][ref][0]) + params["template_duration"] = np.array(dur) else: logging.info("Calculating %s from template parameters", par) params[par] = bank_conversions.get_bank_property( - par, bank, np.arange(bank['mass1'].size)) + par, bank, np.arange(bank["mass1"].size) + ) bank.close() -logging.info('setting up %s bins', ', '.join(usedparams)) -sp_one_bin_input = (params[args.split_param_one].min(), - params[args.split_param_one].max(), args.split_one_nbins) -sp_two_bin_input = (params[args.split_param_two].min(), - params[args.split_param_two].max(), args.split_two_nbins) +logging.info("setting up %s bins", ", ".join(usedparams)) +sp_one_bin_input = ( + params[args.split_param_one].min(), + params[args.split_param_one].max(), + args.split_one_nbins, +) +sp_two_bin_input = ( + params[args.split_param_two].min(), + params[args.split_param_two].max(), + args.split_two_nbins, +) if args.max_bin_param: - logging.info( - 'setting maximum %s value: %.3f', - args.bin_param, - args.max_bin_param - ) + logging.info("setting maximum %s value: %.3f", args.bin_param, args.max_bin_param) pbin_upper_lim = float(args.max_bin_param) else: pbin_upper_lim = params[args.bin_param].max() # For templates with no triggers, the duration will be read as zero -if args.bin_param == 'template_duration' and params[args.bin_param].min() == 0: +if args.bin_param == "template_duration" and params[args.bin_param].min() == 0: # Accessing the 0th entry of an empty region reference will return # zero due to a quirk of h5py. - logging.warning('WARNING: Some templates do not contain triggers') + logging.warning("WARNING: Some templates do not contain triggers") # Use the lowest nonzero template duration as lower limit for bins pbin_lower_lim = params[args.bin_param][params[args.bin_param] > 0].min() else: @@ -162,47 +234,57 @@ else: bb_input = (pbin_lower_lim, pbin_upper_lim, args.num_bins) -logging.info('splitting %s into bins', args.bin_param) -if args.bin_spacing == 'log': +logging.info("splitting %s into bins", args.bin_param) +if args.bin_spacing == "log": assert pbin_lower_lim > 0 pbins = bin_utils.LogarithmicBins(*bb_input) else: pbins = bin_utils.LinearBins(*bb_input) # Use sentinel value -1 for templates outside range -pind = np.array([pbins[par] if pbin_lower_lim < par < pbin_upper_lim - else -1 for par in params[args.bin_param]]) +pind = np.array( + [ + pbins[par] if pbin_lower_lim < par < pbin_upper_lim else -1 + for par in params[args.bin_param] + ] +) -if args.split_one_spacing == 'log': +if args.split_one_spacing == "log": assert params[args.split_param_one].min() > 0 sp_one_bounds = bin_utils.LogarithmicBins(*sp_one_bin_input) else: sp_one_bounds = bin_utils.LinearBins(*sp_one_bin_input) -if args.split_two_spacing == 'log': +if args.split_two_spacing == "log": assert params[args.split_param_two].min() > 0 sp_two_bounds = bin_utils.LogarithmicBins(*sp_two_bin_input) else: sp_two_bounds = bin_utils.LinearBins(*sp_two_bin_input) -logging.info('assigning template ids to different splits') +logging.info("assigning template ids to different splits") id_in_bin1 = [[]] * args.split_one_nbins id_in_bin2 = [[]] * args.split_two_nbins -for i, lower_1, upper_1 in zip(range(args.split_one_nbins), - sp_one_bounds.lower(), sp_one_bounds.upper()): - id_in_bin1[i] = np.intersect1d(np.argwhere(params[args.split_param_one] > lower_1), - np.argwhere(params[args.split_param_one] <= upper_1)) +for i, lower_1, upper_1 in zip( + range(args.split_one_nbins), sp_one_bounds.lower(), sp_one_bounds.upper() +): + id_in_bin1[i] = np.intersect1d( + np.argwhere(params[args.split_param_one] > lower_1), + np.argwhere(params[args.split_param_one] <= upper_1), + ) -for i, lower_2, upper_2 in zip(range(args.split_two_nbins), - sp_two_bounds.lower(), sp_two_bounds.upper()): - id_in_bin2[i] = np.intersect1d(np.argwhere(params[args.split_param_two] > lower_2), - np.argwhere(params[args.split_param_two] <= upper_2)) +for i, lower_2, upper_2 in zip( + range(args.split_two_nbins), sp_two_bounds.lower(), sp_two_bounds.upper() +): + id_in_bin2[i] = np.intersect1d( + np.argwhere(params[args.split_param_two] > lower_2), + np.argwhere(params[args.split_param_two] <= upper_2), + ) -logging.info('Getting template boundaries from trigger file') -boundaries = trigf[args.ifo + '/template_boundaries'][:] +logging.info("Getting template boundaries from trigger file") +boundaries = trigf[args.ifo + "/template_boundaries"][:] trigf.close() -logging.info('Calculating single stat values from trigger file') +logging.info("Calculating single stat values from trigger file") trigs = SingleDetTriggers( args.trigger_file, args.ifo, @@ -215,7 +297,7 @@ trigf = trigs.trigs_f stat = trigs.get_ranking(args.sngl_ranking) time = trigs.end_time -logging.info('Processing template boundaries') +logging.info("Processing template boundaries") max_boundary_id = np.argmax(boundaries) sorted_boundary_list = np.sort(boundaries) @@ -232,10 +314,11 @@ sorted_boundary_list = np.sort(boundaries) where_idx_end = np.zeros_like(boundaries) for idx, idx_start in enumerate(boundaries): if idx == max_boundary_id: - where_idx_end[idx] = trigf[args.ifo + '/end_time'].size + where_idx_end[idx] = trigf[args.ifo + "/end_time"].size else: where_idx_end[idx] = sorted_boundary_list[ - np.argmax(sorted_boundary_list == idx_start) + 1] + np.argmax(sorted_boundary_list == idx_start) + 1 + ] # Next we need to map these start/stop indices in the full file, to the start # stop indices in the masked list of triggers. We do this by figuring out @@ -253,83 +336,83 @@ for idx_start in sorted_boundary_list: if args.veto_file: - logging.info('Applying DQ vetoes') + logging.info("Applying DQ vetoes") remove, junk = events.veto.indices_within_segments( - time, - [args.veto_file], - ifo=args.ifo, - segment_name=args.veto_segment_name + time, [args.veto_file], ifo=args.ifo, segment_name=args.veto_segment_name ) # Set stat to zero for triggers being vetoed: given that the fit threshold # is >0 these will not be fitted or plotted. Avoids complications from # changing the number of triggers, ie changes of template boundary. - stat[remove] = 0. - time[remove] = 0. + stat[remove] = 0.0 + time[remove] = 0.0 logging.info( - '%d out of %d trigs removed after vetoing with %s from %s', + "%d out of %d trigs removed after vetoing with %s from %s", remove.size, stat.size, args.veto_segment_name, - args.veto_file + args.veto_file, ) if args.gating_veto_windows: - logging.info('Applying veto to triggers near gates') - gating_veto = args.gating_veto_windows[args.ifo].split(',') + logging.info("Applying veto to triggers near gates") + gating_veto = args.gating_veto_windows[args.ifo].split(",") gveto_before = float(gating_veto[0]) gveto_after = float(gating_veto[1]) if gveto_before > 0 or gveto_after < 0: - raise ValueError("Gating veto window values must be negative before " - "gates and positive after gates.") + raise ValueError( + "Gating veto window values must be negative before " + "gates and positive after gates." + ) if not (gveto_before == 0 and gveto_after == 0): - autogate_times = np.unique(trigf[args.ifo + '/gating/auto/time'][:]) - if args.ifo + '/gating/file' in trigf: - detgate_times = trigf[args.ifo + '/gating/file/time'][:] + autogate_times = np.unique(trigf[args.ifo + "/gating/auto/time"][:]) + if args.ifo + "/gating/file" in trigf: + detgate_times = trigf[args.ifo + "/gating/file/time"][:] else: detgate_times = [] gate_times = np.concatenate((autogate_times, detgate_times)) gveto_remove = events.veto.indices_within_times( - time, - gate_times + gveto_before, - gate_times + gveto_after + time, gate_times + gveto_before, gate_times + gveto_after ) - stat[gveto_remove] = 0. - time[gveto_remove] = 0. + stat[gveto_remove] = 0.0 + time[gveto_remove] = 0.0 logging.info( - '%d out of %d trigs removed after vetoing triggers near gates', + "%d out of %d trigs removed after vetoing triggers near gates", gveto_remove.size, - stat.size + stat.size, ) for x in range(args.split_one_nbins): if not args.prune_number: - logging.info('Not performing any pruning') + logging.info("Not performing any pruning") break elif args.prune_number and x == 0: - logging.info('Applying pruning around loudest triggers in each split') + logging.info("Applying pruning around loudest triggers in each split") id_bin1 = id_in_bin1[x] for y in range(args.split_two_nbins): id_bin2 = id_in_bin2[y] # Finding ids of templates in split id_in_both = np.intersect1d(id_bin1, id_bin2) - if len(id_in_both) == 0: continue + if len(id_in_both) == 0: + continue vals_inbin = [] time_inbin = [] # getting triggers that are in these templates for idx in id_in_both: - vals_inbin += list(stat[mask_start_idx[idx]:mask_end_idx[idx]]) - time_inbin += list(time[mask_start_idx[idx]:mask_end_idx[idx]]) + vals_inbin += list(stat[mask_start_idx[idx] : mask_end_idx[idx]]) + time_inbin += list(time[mask_start_idx[idx] : mask_end_idx[idx]]) vals_inbin = np.array(vals_inbin) time_inbin = np.array(time_inbin) count_pruned = 0 logging.info( - 'Pruning in split %s-%i %s-%i', - args.split_param_one, x, - args.split_param_two, y + "Pruning in split %s-%i %s-%i", + args.split_param_one, + x, + args.split_param_two, + y, ) - logging.info('Currently have %d triggers', len(vals_inbin)) + logging.info("Currently have %d triggers", len(vals_inbin)) while count_pruned < args.prune_number: # Getting loudest statistic value in split max_val_arg = vals_inbin.argmax() @@ -344,20 +427,20 @@ for x in range(args.split_one_nbins): abs(time_inbin[max_val_arg] - time_inbin) < args.prune_window )[0] logging.info( - 'Prune %d: removing %d triggers around %.2f, %d in this split', + "Prune %d: removing %d triggers around %.2f, %d in this split", count_pruned, remove.size, time[max_val_arg], - remove_inbin.size + remove_inbin.size, ) # Set pruned triggers' stat values to zero, as above for vetoes - vals_inbin[remove_inbin] = 0. - time_inbin[remove_inbin] = 0. - stat[remove] = 0. - time[remove] = 0. + vals_inbin[remove_inbin] = 0.0 + time_inbin[remove_inbin] = 0.0 + stat[remove] = 0.0 + time[remove] = 0.0 count_pruned += 1 -logging.info('Setting up plotting and fitting limit values') +logging.info("Setting up plotting and fitting limit values") minplot = max(stat[np.nonzero(stat)].min(), args.plot_lower_stat_limit) min_fit = max(minplot, args.stat_fit_threshold) max_fit = 1.05 * stat.max() @@ -367,35 +450,49 @@ else: maxplot = max_fit fitrange = np.linspace(min_fit, max_fit, 100) -logging.info('Setting up plotting variables') -histcolors = ['r',(1.0,0.6,0),'y','g','c','b','m','k',(0.8,0.25,0),(0.25,0.8,0)] -fig, axes = plt.subplots(args.split_one_nbins, args.split_two_nbins, - sharex=True, sharey=True, squeeze=False, - figsize=(3 * (args.split_two_nbins + 1), - 3 * args.split_one_nbins)) +logging.info("Setting up plotting variables") +histcolors = [ + "r", + (1.0, 0.6, 0), + "y", + "g", + "c", + "b", + "m", + "k", + (0.8, 0.25, 0), + (0.25, 0.8, 0), +] +fig, axes = plt.subplots( + args.split_one_nbins, + args.split_two_nbins, + sharex=True, + sharey=True, + squeeze=False, + figsize=(3 * (args.split_two_nbins + 1), 3 * args.split_one_nbins), +) # Setting up overall legend outside the split-up plots lines = [] labels = [] for i, lower, upper in zip(range(args.num_bins), pbins.lower(), pbins.upper()): binlabel = f"{lower:#.3g} - {upper:#.3g}" - line, = axes[0,0].plot([0,0], [0,0], linewidth=2, - color=histcolors[i], alpha=0.6) + (line,) = axes[0, 0].plot( + [0, 0], [0, 0], linewidth=2, color=histcolors[i], alpha=0.6 + ) lines.append(line) labels.append(binlabel) -line_fit, = axes[0,0].plot([0,0], [0,0], linestyle='--', - color='k', alpha=0.6) +(line_fit,) = axes[0, 0].plot([0, 0], [0, 0], linestyle="--", color="k", alpha=0.6) lines.append(line_fit) -labels.append(args.fit_function + ' fit to counts') -fig.legend(lines, labels, labelspacing=0.2, - loc='upper left', title=args.bin_param) +labels.append(args.fit_function + " fit to counts") +fig.legend(lines, labels, labelspacing=0.2, loc="upper left", title=args.bin_param) pidx = [] for i in range(args.num_bins): pidx.append([np.argwhere(pind == i)]) -logging.info('Starting bin, histogram and plot loop') +logging.info("Starting bin, histogram and plot loop") maxyval = 0 for x in range(args.split_one_nbins): id_bin1 = id_in_bin1[x] @@ -403,39 +500,45 @@ for x in range(args.split_one_nbins): id_bin2 = id_in_bin2[y] id_in_both = np.intersect1d(id_bin1, id_bin2) logging.info( - 'Split %s-%i %s-%i', - args.split_param_one, x, - args.split_param_two, y, + "Split %s-%i %s-%i", + args.split_param_one, + x, + args.split_param_two, + y, ) - ax = axes[x,y] - for i, lower, upper in zip(range(args.num_bins), pbins.lower(), - pbins.upper()): + ax = axes[x, y] + for i, lower, upper in zip(range(args.num_bins), pbins.lower(), pbins.upper()): indices_all_conditions = np.intersect1d(pidx[i], id_in_both) - logging.info('%s split %#.3g-%#.3g', args.bin_param, lower, upper) - if len(indices_all_conditions) == 0: continue + logging.info("%s split %#.3g-%#.3g", args.bin_param, lower, upper) + if len(indices_all_conditions) == 0: + continue vals_inbin = [] for idx in indices_all_conditions: - vals_inbin += list(stat[mask_start_idx[idx]:mask_end_idx[idx]]) + vals_inbin += list(stat[mask_start_idx[idx] : mask_end_idx[idx]]) vals_inbin = np.array(vals_inbin) vals_above_thresh = vals_inbin[vals_inbin >= args.stat_fit_threshold] if not len(vals_above_thresh): - logging.info('No triggers above threshold') + logging.info("No triggers above threshold") continue else: - logging.info('%d triggers out of %d above threshold', - len(vals_above_thresh), len(vals_inbin)) - alpha, sig_alpha = trstats.fit_above_thresh(args.fit_function, - vals_above_thresh, args.stat_fit_threshold) - fitted_cum_counts = len(vals_above_thresh) * \ - trstats.cum_fit(args.fit_function, fitrange, - alpha, args.stat_fit_threshold) + logging.info( + "%d triggers out of %d above threshold", + len(vals_above_thresh), + len(vals_inbin), + ) + alpha, sig_alpha = trstats.fit_above_thresh( + args.fit_function, vals_above_thresh, args.stat_fit_threshold + ) + fitted_cum_counts = len(vals_above_thresh) * trstats.cum_fit( + args.fit_function, fitrange, alpha, args.stat_fit_threshold + ) # upper and lower 1-sigma bounds on fit are not currently plotted - #fitted_cum_counts_plus = len(vals_above_thresh) * \ + # fitted_cum_counts_plus = len(vals_above_thresh) * \ # trstats.cum_fit(args.fit_function, # fitrange, alpha + sig_alpha, # args.stat_fit_threshold) - #fitted_cum_counts_minus = len(vals_above_thresh) * \ + # fitted_cum_counts_minus = len(vals_above_thresh) * \ # trstats.cum_fit(args.fit_function, fitrange, # alpha - sig_alpha, # args.stat_fit_threshold) @@ -444,11 +547,17 @@ for x in range(args.split_one_nbins): histcounts, edges = np.histogram(vals_inbin, bins=50) cum_counts = histcounts[::-1].cumsum()[::-1] # Plot the lines! - ax.semilogy(edges[:-1], cum_counts, linewidth=2, - color=histcolors[i], alpha=0.6) - ax.semilogy(fitrange, fitted_cum_counts, "--", color=histcolors[i], - label=f"$\\alpha = ${alpha:#.3f}" % alpha) - lgd_sub = ax.legend(fontsize='small', framealpha=0.5) + ax.semilogy( + edges[:-1], cum_counts, linewidth=2, color=histcolors[i], alpha=0.6 + ) + ax.semilogy( + fitrange, + fitted_cum_counts, + "--", + color=histcolors[i], + label=f"$\\alpha = ${alpha:#.3f}" % alpha, + ) + lgd_sub = ax.legend(fontsize="small", framealpha=0.5) del vals_inbin, vals_above_thresh @@ -457,40 +566,49 @@ for x in range(args.split_one_nbins): for i in range(args.split_one_nbins): for j in range(args.split_two_nbins): - axes[i,j].semilogy([args.stat_fit_threshold, args.stat_fit_threshold], - [1, 5 * maxyval], 'k', linestyle=':', alpha=0.2) -axes[0,0].set_ylim(1, 5 * maxyval) -axes[0,0].set_xlim(minplot, maxplot) + axes[i, j].semilogy( + [args.stat_fit_threshold, args.stat_fit_threshold], + [1, 5 * maxyval], + "k", + linestyle=":", + alpha=0.2, + ) +axes[0, 0].set_ylim(1, 5 * maxyval) +axes[0, 0].set_xlim(minplot, maxplot) for j in range(args.split_two_nbins): axes[args.split_one_nbins - 1, j].set_xlabel(args.sngl_ranking, size="large") if args.split_two_nbins == 1: break - xrange_string = f'{sp_two_bounds.lower()[j]:#.3g} to {sp_two_bounds.upper()[j]:#.3g}' - axes[0, j].set_xlabel(f'{args.split_param_two}: {xrange_string}', size="large") + xrange_string = ( + f"{sp_two_bounds.lower()[j]:#.3g} to {sp_two_bounds.upper()[j]:#.3g}" + ) + axes[0, j].set_xlabel(f"{args.split_param_two}: {xrange_string}", size="large") axes[0, j].xaxis.set_label_position("top") for i in range(args.split_one_nbins): if args.split_one_nbins == 1: - axes[0, 0].set_ylabel('Cumulative number', size='large') + axes[0, 0].set_ylabel("Cumulative number", size="large") break - yrange_string = f'{sp_one_bounds.lower()[i]:#.3g} to {sp_one_bounds.upper()[i]:#.3g}' - axes[i, 0].set_ylabel(f'{args.split_param_one}: {yrange_string}\ncumulative number', - size="large") + yrange_string = ( + f"{sp_one_bounds.lower()[i]:#.3g} to {sp_one_bounds.upper()[i]:#.3g}" + ) + axes[i, 0].set_ylabel( + f"{args.split_param_one}: {yrange_string}\ncumulative number", size="large" + ) -fig.tight_layout(rect=(1./(args.split_two_nbins+1), 0, 1, 1)) +fig.tight_layout(rect=(1.0 / (args.split_two_nbins + 1), 0, 1, 1)) -logging.info('Saving to file %s', args.output_file) +logging.info("Saving to file %s", args.output_file) results.save_fig_with_metadata( - fig, args.output_file, - title="{}: {} histogram of single detector triggers split by" - " {} and {}".format(args.ifo, args.sngl_ranking, args.split_param_one, - args.split_param_two), - caption=(r"Histogram of {} single detector {} values binned by {}, split by " - "{} and {}, with fitted {} distribution parameterized by" - " α".format(args.ifo, args.sngl_ranking, args.bin_param, - args.split_param_one, args.split_param_two, - args.fit_function)), - cmd=" ".join(sys.argv) + fig, + args.output_file, + title=f"{args.ifo}: {args.sngl_ranking} histogram of single detector triggers split by {args.split_param_one} and {args.split_param_two}", + caption=( + rf"Histogram of {args.ifo} single detector {args.sngl_ranking} values binned by {args.bin_param}, split by " + f"{args.split_param_one} and {args.split_param_two}, with fitted {args.fit_function} distribution parameterized by" + " α" + ), + cmd=" ".join(sys.argv), ) -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_followup_file b/bin/all_sky_search/pycbc_followup_file index a1820ecdaf6..e9a2b695587 100644 --- a/bin/all_sky_search/pycbc_followup_file +++ b/bin/all_sky_search/pycbc_followup_file @@ -1,80 +1,93 @@ #!/bin/env python -"""Generate the standardized file detailing the candidates/background to +""" +Generate the standardized file detailing the candidates/background to follow-up. """ -import numpy, argparse, logging, pycbc + +import argparse +import logging + +import numpy + +import pycbc from pycbc.io import HFile parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--statmap-file', - help="Statmap file containing the candidates/background to follow up") -parser.add_argument('--bank-file', - help="HDF format template bank file") -parser.add_argument('--min-stat', type=float, - help="Minimum statistic value to follow up") -parser.add_argument('--foreground-only', action='store_true', - help="Only make an followup input file which contains zerolag candidates") -parser.add_argument('--output-file') +parser.add_argument( + "--statmap-file", + help="Statmap file containing the candidates/background to follow up", +) +parser.add_argument("--bank-file", help="HDF format template bank file") +parser.add_argument( + "--min-stat", type=float, help="Minimum statistic value to follow up" +) +parser.add_argument( + "--foreground-only", + action="store_true", + help="Only make an followup input file which contains zerolag candidates", +) +parser.add_argument("--output-file") args = parser.parse_args() pycbc.init_logging(args.verbose) # Currently this expects the 2-ifo format, but the output file format # should be ok for multi-ifo as well -bank = HFile(args.bank_file, 'r') -sfile = HFile(args.statmap_file, 'r') +bank = HFile(args.bank_file, "r") +sfile = HFile(args.statmap_file, "r") -if 'detector_1' in sfile.attrs: - pivot = sfile.attrs['detector_1'] - fixed = sfile.attrs['detector_2'] +if "detector_1" in sfile.attrs: + pivot = sfile.attrs["detector_1"] + fixed = sfile.attrs["detector_2"] ifos = [pivot, fixed] else: - pivot = sfile.attrs['pivot'] - fixed = sfile.attrs['fixed'] - ifos = sfile.attrs['ifos'].split(' ') + pivot = sfile.attrs["pivot"] + fixed = sfile.attrs["fixed"] + ifos = sfile.attrs["ifos"].split(" ") -slide = sfile.attrs['timeslide_interval'] +slide = sfile.attrs["timeslide_interval"] odtype = [(ifo, numpy.float64) for ifo in ifos] -dtype = odtype + [('time', numpy.float64), - ('template_id', numpy.uint64), - ('stat', numpy.float32), - ] +dtype = odtype + [ + ("time", numpy.float64), + ("template_id", numpy.uint64), + ("stat", numpy.float32), +] # Determine which events we will include in our followup file if args.foreground_only: - sections = ['foreground'] + sections = ["foreground"] else: - sections = ['background', 'background_exc', 'foreground'] + sections = ["background", "background_exc", "foreground"] values = [] for k in sfile: if k in sections: - if 'detector_1' in sfile.attrs: - tpivot = sfile[k]['time1'][:].astype(numpy.float64) - tfixed = sfile[k]['time2'][:].astype(numpy.float64) + if "detector_1" in sfile.attrs: + tpivot = sfile[k]["time1"][:].astype(numpy.float64) + tfixed = sfile[k]["time2"][:].astype(numpy.float64) else: - tpivot = sfile[k][pivot]['time'][:].astype(numpy.float64) - tfixed = sfile[k][fixed]['time'][:].astype(numpy.float64) + tpivot = sfile[k][pivot]["time"][:].astype(numpy.float64) + tfixed = sfile[k][fixed]["time"][:].astype(numpy.float64) - template = sfile[k]['template_id'][:] + template = sfile[k]["template_id"][:] value = numpy.zeros(len(tpivot), dtype=dtype) - if k is not 'foreground': - tfixed += sfile[k]['timeslide_id'][:] * slide + if k != "foreground": + tfixed += sfile[k]["timeslide_id"][:] * slide for ifo in ifos: if ifo == pivot: continue - value[ifo] = -sfile[k]['timeslide_id'][:] * slide + value[ifo] = -sfile[k]["timeslide_id"][:] * slide - value['time'] = 0.5 * (tpivot + tfixed) - value['template_id'] = template - value['stat'] = sfile[k]['stat'][:] + value["time"] = 0.5 * (tpivot + tfixed) + value["template_id"] = template + value["stat"] = sfile[k]["stat"][:] if args.min_stat: - keep = sfile[k]['stat'][:] > args.min_stat + keep = sfile[k]["stat"][:] > args.min_stat value = value[keep] values.append(value) @@ -83,24 +96,24 @@ for k in sfile: # multiple background types) values = numpy.concatenate(values) values, invmap = numpy.unique(values, return_inverse=True) -logging.info('%s triggers to follow up', len(values)) +logging.info("%s triggers to follow up", len(values)) -f = HFile(args.output_file, 'w') -f['inverse'] = invmap -f.attrs['sections'] = sections -f['time'] = values['time'] +f = HFile(args.output_file, "w") +f["inverse"] = invmap +f.attrs["sections"] = sections +f["time"] = values["time"] offsets = numpy.zeros(len(values), dtype=odtype) for ifo in ifos: offsets[ifo] = values[ifo] -f['offsets'] = offsets -f['stat'] = values['stat'] +f["offsets"] = offsets +f["stat"] = values["stat"] wdtype = [(k, numpy.float64) for k in bank] wparam = numpy.zeros(len(values), dtype=wdtype) for k in bank: - wparam[k] = bank[k][:][values['template_id']] + wparam[k] = bank[k][:][values["template_id"]] -f['waveparams'] = wparam +f["waveparams"] = wparam -logging.info('Done') +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_foreground_censor b/bin/all_sky_search/pycbc_foreground_censor index 78e20e7d6ae..2fc8e7e0a19 100755 --- a/bin/all_sky_search/pycbc_foreground_censor +++ b/bin/all_sky_search/pycbc_foreground_censor @@ -1,55 +1,66 @@ #!/usr/bin/env python -"""Make segment file to blind the results from foreground related triggers """ +"""Make segment file to blind the results from foreground related triggers""" -import os, argparse, logging +import argparse +import logging +import os from urllib.parse import urlunparse + import pycbc.events -from pycbc.workflow import SegFile from pycbc.io import HFile +from pycbc.workflow import SegFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--foreground-triggers', - help="HDF file containing the zerolag foreground triggers " - "from the analysis") -parser.add_argument('--veto-file', - help="Baseline veto information that is added to the outptut") -parser.add_argument('--segment-name', - help="Segment name to use from the input veto file") -parser.add_argument('--output-file', help='Name of the output segment file') -parser.add_argument('--output-segment-name', - help="(optional), Name of output segment file list", - default="censor_foreground") +parser.add_argument( + "--foreground-triggers", + help="HDF file containing the zerolag foreground triggers from the analysis", +) +parser.add_argument( + "--veto-file", help="Baseline veto information that is added to the outptut" +) +parser.add_argument( + "--segment-name", help="Segment name to use from the input veto file" +) +parser.add_argument("--output-file", help="Name of the output segment file") +parser.add_argument( + "--output-segment-name", + help="(optional), Name of output segment file list", + default="censor_foreground", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Start') +logging.info("Start") -f = HFile(args.foreground_triggers, 'r') +f = HFile(args.foreground_triggers, "r") -start = f['segments/foreground_veto/start'][:] -end = f['segments/foreground_veto/end'][:] +start = f["segments/foreground_veto/start"][:] +end = f["segments/foreground_veto/end"][:] vsegs = pycbc.events.start_end_to_segments(start, end) -logging.info('Read in foreground veto segments') +logging.info("Read in foreground veto segments") # 2-ifo old style format -if 'detector_1' in f.attrs: - ifo1, ifo2 = f.attrs['detector_1'], f.attrs['detector_2'] +if "detector_1" in f.attrs: + ifo1, ifo2 = f.attrs["detector_1"], f.attrs["detector_2"] ifos = [ifo1, ifo2] # Multi-ifo format file else: - ifos = f.attrs['ifos'].split(' ') + ifos = f.attrs["ifos"].split(" ") fsegs, names = [], [] for ifo in ifos: - segs = pycbc.events.select_segments_by_definer(args.veto_file, args.segment_name, ifo) - logging.info('Read in veto segments from %s' % ifo) + segs = pycbc.events.select_segments_by_definer( + args.veto_file, args.segment_name, ifo + ) + logging.info("Read in veto segments from %s" % ifo) fsegs += [segs.coalesce() + vsegs.coalesce()] names += [args.output_segment_name] -file_url = urlunparse(['file', 'localhost', - os.path.abspath(args.output_file), None, None, None]) -SegFile.from_multi_segment_list('UNUSED', fsegs, names, ifos, file_url=file_url) -logging.info('Done') +file_url = urlunparse( + ["file", "localhost", os.path.abspath(args.output_file), None, None, None] +) +SegFile.from_multi_segment_list("UNUSED", fsegs, names, ifos, file_url=file_url) +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_get_loudest_params b/bin/all_sky_search/pycbc_get_loudest_params index b5d12002282..e13f0fdac31 100644 --- a/bin/all_sky_search/pycbc_get_loudest_params +++ b/bin/all_sky_search/pycbc_get_loudest_params @@ -1,91 +1,123 @@ #!/usr/bin/env python """ -Finds the loudest snr or newsnr event within a given time window, parses the +Finds the loudest snr or newsnr event within a given time window, parses the parameters of the template, and writes them to an hdf and/or stdout. """ -import numpy as np import argparse import logging -from pycbc import init_logging + +import numpy as np + import pycbc.events -from pycbc.pnutils import mass1_mass2_to_mchirp_eta +from pycbc import init_logging from pycbc.io import HFile +from pycbc.pnutils import mass1_mass2_to_mchirp_eta parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--single-ifo-trigs', type=str, required=True, - help='HDF file containing single IFO CBC triggers') -parser.add_argument('--tmpltbank-file', type=str, required=True, - help='HDF file containing template information for CBC search') -parser.add_argument('--ifo', type=str, required=True, - help='IFO, L1 or H1') -parser.add_argument('--central-time', type=float, required=True, - help='Central time over which to search') -parser.add_argument('--window', type=float, required=False, default=8.0, - help='Time window over which to search for loudest trigger') -parser.add_argument('--ranking-statistic', type=str, required=False, default='newsnr', - choices=['snr','newsnr'], help='Ranking statistic to use when searching for loudest events') -parser.add_argument('--output-file', type=str, required=False, - help='Output hdf file to write parameters') -parser.add_argument('--print-params', action='store_true', required=False, - help='Toggle printing parameters to stdout') +parser.add_argument( + "--single-ifo-trigs", + type=str, + required=True, + help="HDF file containing single IFO CBC triggers", +) +parser.add_argument( + "--tmpltbank-file", + type=str, + required=True, + help="HDF file containing template information for CBC search", +) +parser.add_argument("--ifo", type=str, required=True, help="IFO, L1 or H1") +parser.add_argument( + "--central-time", + type=float, + required=True, + help="Central time over which to search", +) +parser.add_argument( + "--window", + type=float, + required=False, + default=8.0, + help="Time window over which to search for loudest trigger", +) +parser.add_argument( + "--ranking-statistic", + type=str, + required=False, + default="newsnr", + choices=["snr", "newsnr"], + help="Ranking statistic to use when searching for loudest events", +) +parser.add_argument( + "--output-file", + type=str, + required=False, + help="Output hdf file to write parameters", +) +parser.add_argument( + "--print-params", + action="store_true", + required=False, + help="Toggle printing parameters to stdout", +) args = parser.parse_args() init_logging(args.verbose) -logging.info('Reading in HDF files') -trigs = HFile(args.single_ifo_trigs,'r') -template_file = HFile(args.tmpltbank_file,'r') +logging.info("Reading in HDF files") +trigs = HFile(args.single_ifo_trigs, "r") +template_file = HFile(args.tmpltbank_file, "r") if args.output_file: - outfile = HFile(args.output_file,'w') + outfile = HFile(args.output_file, "w") t_low = args.central_time - args.window t_high = args.central_time + args.window -times = trigs[args.ifo]['end_time'][:] +times = trigs[args.ifo]["end_time"][:] mask = (times > t_low) & (times < t_high) # generate vectors of sigmasq, sigma, snr, end_time, coalescence phase, and template_id -snr = trigs[args.ifo]['snr'][mask] -chisq = trigs[args.ifo]['chisq'][mask] -chisq_dof = trigs[args.ifo]['chisq_dof'][mask] -reduced_chisq = chisq/(2*chisq_dof - 2) -newsnr = pycbc.events.ranking.newsnr(snr,reduced_chisq) -template_ids = trigs[args.ifo]['template_id'][mask] -end_times = trigs[args.ifo]['end_time'][mask] - -if args.ranking_statistic == 'snr': +snr = trigs[args.ifo]["snr"][mask] +chisq = trigs[args.ifo]["chisq"][mask] +chisq_dof = trigs[args.ifo]["chisq_dof"][mask] +reduced_chisq = chisq / (2 * chisq_dof - 2) +newsnr = pycbc.events.ranking.newsnr(snr, reduced_chisq) +template_ids = trigs[args.ifo]["template_id"][mask] +end_times = trigs[args.ifo]["end_time"][mask] + +if args.ranking_statistic == "snr": idx = np.argmax(snr) else: idx = np.argmax(newsnr) tid = template_ids[idx] cbc_end_time = end_times[idx] -m1 = template_file['mass1'][tid] -m2 = template_file['mass2'][tid] -s1z = template_file['spin1z'][tid] -s2z = template_file['spin2z'][tid] +m1 = template_file["mass1"][tid] +m2 = template_file["mass2"][tid] +s1z = template_file["spin1z"][tid] +s2z = template_file["spin2z"][tid] mchirp, eta = mass1_mass2_to_mchirp_eta(m1, m2) -data = {'%s/snr' % args.ifo : [snr[idx]]} -data['%s/chisq' % args.ifo] = [chisq[idx]] -data['%s/newsnr' % args.ifo] = [newsnr[idx]] -data['%s/template_id' % args.ifo] = [tid] -data['%s/end_time' % args.ifo] = [cbc_end_time] -data['template/mass1'] = [m1] -data['template/mass2'] = [m2] -data['template/mchirp'] = [mchirp] -data['template/eta'] = [eta] -data['template/spin1z'] = [s1z] -data['template/spin2z'] = [s2z] +data = {"%s/snr" % args.ifo: [snr[idx]]} +data["%s/chisq" % args.ifo] = [chisq[idx]] +data["%s/newsnr" % args.ifo] = [newsnr[idx]] +data["%s/template_id" % args.ifo] = [tid] +data["%s/end_time" % args.ifo] = [cbc_end_time] +data["template/mass1"] = [m1] +data["template/mass2"] = [m2] +data["template/mchirp"] = [mchirp] +data["template/eta"] = [eta] +data["template/spin1z"] = [s1z] +data["template/spin2z"] = [s2z] if args.output_file: - for key in data.keys(): - outfile.create_dataset(key,data=data[key]) + for key in data: + outfile.create_dataset(key, data=data[key]) if args.print_params: - for key in data.keys(): - print(key.split('/')[1], str(data[key])) + for key in data: + print(key.split("/")[1], str(data[key])) diff --git a/bin/all_sky_search/pycbc_make_bayestar_skymap b/bin/all_sky_search/pycbc_make_bayestar_skymap index d1c7af81ffb..892a5fda12b 100644 --- a/bin/all_sky_search/pycbc_make_bayestar_skymap +++ b/bin/all_sky_search/pycbc_make_bayestar_skymap @@ -15,34 +15,45 @@ Unrecognised options will be passed straight to the bayestar subprocess """ import argparse -import subprocess -import shutil +import glob import logging import os -import glob +import shutil +import subprocess import tempfile -from igwn_ligolw import lsctables, utils as ligolw_utils +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils import pycbc -from pycbc.waveform import bank as wavebank from pycbc.io import WaveformArray from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.waveform import bank as wavebank parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--bayestar-executable', - help="The bayestar-localize-coinc executable to be run. " - "If not given, will use whatever is available in " - "the current environment.") -parser.add_argument('--event-xml', required=True, - help="XML file containing event information, SNR " - "timeseries and PSD to pass to bayestar") -parser.add_argument('--low-frequency-cutoff', type=float, default=20, - help="Low-frequency cutoff used in the matched-filtering " - "to generate the SNR timeseries") -parser.add_argument('--output-file', required=True, - help="Filename to output the fits file to.") +parser.add_argument( + "--bayestar-executable", + help="The bayestar-localize-coinc executable to be run. " + "If not given, will use whatever is available in " + "the current environment.", +) +parser.add_argument( + "--event-xml", + required=True, + help="XML file containing event information, SNR " + "timeseries and PSD to pass to bayestar", +) +parser.add_argument( + "--low-frequency-cutoff", + type=float, + default=20, + help="Low-frequency cutoff used in the matched-filtering " + "to generate the SNR timeseries", +) +parser.add_argument( + "--output-file", required=True, help="Filename to output the fits file to." +) wavebank.add_approximant_arg(parser) args, unknown = parser.parse_known_args() @@ -51,16 +62,13 @@ pycbc.init_logging(args.verbose, default_level=1) logging.info("Starting") -bayestar_exe = args.bayestar_executable or 'bayestar-localize-coincs' +bayestar_exe = args.bayestar_executable or "bayestar-localize-coincs" tmpdir = tempfile.mkdtemp() # Work out which approximant is being used # Load the file -xmldoc = ligolw_utils.load_filename( - args.event_xml, - contenthandler=LIGOLWContentHandler -) +xmldoc = ligolw_utils.load_filename(args.event_xml, contenthandler=LIGOLWContentHandler) # Grab the single inspiral table(s) which contain the template information sngl_inspiral_table = lsctables.SnglInspiralTable.get_table(xmldoc) @@ -76,30 +84,37 @@ row = WaveformArray.from_ligolw_table( waveform = wavebank.parse_approximant_arg(args.approximant, row)[0] # BAYESTAR uses TaylorF2 instead of SPAtmplt -if waveform == 'SPAtmplt': - waveform = 'TaylorF2' +if waveform == "SPAtmplt": + waveform = "TaylorF2" # Set up the command to pass to bayestar. # The XML file path being passed twice is a legacy requirement, not a mistake. -cmd = [bayestar_exe, - args.event_xml, - args.event_xml, - '--waveform', waveform, - '--f-low', str(args.low_frequency_cutoff), - '-o', tmpdir] +cmd = [ + bayestar_exe, + args.event_xml, + args.event_xml, + "--waveform", + waveform, + "--f-low", + str(args.low_frequency_cutoff), + "-o", + tmpdir, +] # Pass any unknown options straight to the subprocess cmd += unknown -logging.info("Running %s", ' '.join(cmd)) +logging.info("Running %s", " ".join(cmd)) subprocess.check_output(cmd) # Find the fits file in the temporary directory: # It would be nice to do this better! - maybe use the input xml # file to find the number which is going to be used? -fits_filenames = glob.glob(os.path.join(tmpdir, '*.fits*')) +fits_filenames = glob.glob(os.path.join(tmpdir, "*.fits*")) if len(fits_filenames) != 1: - raise ValueError(f'argh, got {len(fits_filenames)} FITS files after running BAYESTAR, I want one!') + raise ValueError( + f"argh, got {len(fits_filenames)} FITS files after running BAYESTAR, I want one!" + ) logging.info("Moving output to %s", args.output_file) shutil.move(fits_filenames[0], args.output_file) diff --git a/bin/all_sky_search/pycbc_merge_psds b/bin/all_sky_search/pycbc_merge_psds index dd8bdafde07..5ecdbc10100 100755 --- a/bin/all_sky_search/pycbc_merge_psds +++ b/bin/all_sky_search/pycbc_merge_psds @@ -15,52 +15,58 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Merge hdf psd files -""" -import logging, argparse, numpy, pycbc.types +"""Merge hdf psd files""" + +import argparse +import logging + +import numpy + +import pycbc.types from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--psd-files', nargs='+') +parser.add_argument("--psd-files", nargs="+") parser.add_argument("--output-file", required=True) args = parser.parse_args() pycbc.init_logging(args.verbose) -outf = HFile(args.output_file, 'w') +outf = HFile(args.output_file, "w") inc = {} start, end = {}, {} for psd_file in args.psd_files: - f = HFile(psd_file, 'r') + f = HFile(psd_file, "r") ifo = tuple(f.keys())[0] if ifo not in inc: inc[ifo] = 0 start[ifo], end[ifo] = [], [] - - for rkey in f['%s/psds' % ifo].keys(): + + for rkey in f["%s/psds" % ifo].keys(): int_key = int(rkey) - rkey = '%s/psds/%s' % (ifo, rkey) + rkey = "%s/psds/%s" % (ifo, rkey) psd = pycbc.types.load_frequencyseries(psd_file, group=rkey) - - key = ifo + '/psds/' + str(inc[ifo]) - outf.create_dataset(key, data=psd, - compression='gzip', compression_opts=9, shuffle=True) - - s = f['%s/start_time' % ifo][int_key] - e = f['%s/end_time' % ifo][int_key] - - outf[key].attrs['epoch'] = int(psd.epoch) - outf[key].attrs['delta_f'] = float(psd.delta_f) + + key = ifo + "/psds/" + str(inc[ifo]) + outf.create_dataset( + key, data=psd, compression="gzip", compression_opts=9, shuffle=True + ) + + s = f["%s/start_time" % ifo][int_key] + e = f["%s/end_time" % ifo][int_key] + + outf[key].attrs["epoch"] = int(psd.epoch) + outf[key].attrs["delta_f"] = float(psd.delta_f) start[ifo].append(s) end[ifo].append(e) - + inc[ifo] += 1 -for ifo in start: - outf[ifo + '/start_time'] = numpy.array(start[ifo], dtype=numpy.uint32) - outf[ifo + '/end_time'] = numpy.array(end[ifo], dtype=numpy.uint32) +for ifo in start: + outf[ifo + "/start_time"] = numpy.array(start[ifo], dtype=numpy.uint32) + outf[ifo + "/end_time"] = numpy.array(end[ifo], dtype=numpy.uint32) -outf.attrs['low_frequency_cutoff'] = f.attrs['low_frequency_cutoff'] -outf.attrs['dynamic_range_factor'] = pycbc.DYN_RANGE_FAC -logging.info('Done!') +outf.attrs["low_frequency_cutoff"] = f.attrs["low_frequency_cutoff"] +outf.attrs["dynamic_range_factor"] = pycbc.DYN_RANGE_FAC +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_plot_kde_vals b/bin/all_sky_search/pycbc_plot_kde_vals index 4318a43af9e..21ae1a913e9 100644 --- a/bin/all_sky_search/pycbc_plot_kde_vals +++ b/bin/all_sky_search/pycbc_plot_kde_vals @@ -1,59 +1,71 @@ #!/usr/bin/env python -import numpy, argparse +import argparse + import matplotlib.pyplot as plt +import numpy from matplotlib.colors import LogNorm -from pycbc import init_logging, add_common_pycbc_options + +from pycbc import add_common_pycbc_options, init_logging from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--signal-file') -parser.add_argument('--template-file', required=True) -parser.add_argument('--param', nargs='+', required=True, - help='Specify one parameter name for a kde_vs_param plot, or ' - 'two parameter names for a param_vs_param plot. Param ' - 'names must exist as datasets in the input files') -parser.add_argument('--log-axis', nargs='+', choices=['True', 'False'], required=True, - help='For each parameter, specify True for a log axis and False ' - 'for a linear axis') -parser.add_argument('--plot-type', choices=['kde_vs_param', 'param_vs_param']) -parser.add_argument('--plot-order', choices=['file', 'increasing', 'decreasing'], - default='file', - help='Choose the order to plot KDE values in: "file", "increasing"' - ' or "decreasing"') -parser.add_argument('--which-kde', choices=['signal_kde', 'template_kde', 'ratio_kde']) -parser.add_argument('--plot-dir', required=True) +parser.add_argument("--signal-file") +parser.add_argument("--template-file", required=True) +parser.add_argument( + "--param", + nargs="+", + required=True, + help="Specify one parameter name for a kde_vs_param plot, or " + "two parameter names for a param_vs_param plot. Param " + "names must exist as datasets in the input files", +) +parser.add_argument( + "--log-axis", + nargs="+", + choices=["True", "False"], + required=True, + help="For each parameter, specify True for a log axis and False for a linear axis", +) +parser.add_argument("--plot-type", choices=["kde_vs_param", "param_vs_param"]) +parser.add_argument( + "--plot-order", + choices=["file", "increasing", "decreasing"], + default="file", + help='Choose the order to plot KDE values in: "file", "increasing" or "decreasing"', +) +parser.add_argument("--which-kde", choices=["signal_kde", "template_kde", "ratio_kde"]) +parser.add_argument("--plot-dir", required=True) args = parser.parse_args() init_logging(args.verbose) -if args.plot_type == 'kde_vs_param': +if args.plot_type == "kde_vs_param": if len(args.param) != 1: - parser.error('For kde_vs_param, give exactly one parameter name') -else: - if len(args.param) != 2: - parser.error('For param_vs_param, give exactly two parameter names') + parser.error("For kde_vs_param, give exactly one parameter name") +elif len(args.param) != 2: + parser.error("For param_vs_param, give exactly two parameter names") if len(args.param) != len(args.log_axis): - parser.error('Must specify either log (True) or non-log (False) for each parameter') + parser.error("Must specify either log (True) or non-log (False) for each parameter") if args.signal_file: - signal_data = HFile(args.signal_file, 'r') - signal_kde = signal_data['data_kde'][:] -template_data = HFile(args.template_file, 'r') -template_kde = template_data['data_kde'][:] + signal_data = HFile(args.signal_file, "r") + signal_kde = signal_data["data_kde"][:] +template_data = HFile(args.template_file, "r") +template_kde = template_data["data_kde"][:] param_arrays = [template_data[param][:] for param in args.param] kde_values = { - 'signal_kde': signal_kde, - 'template_kde': template_kde, - 'ratio_kde': signal_kde / template_kde, + "signal_kde": signal_kde, + "template_kde": template_kde, + "ratio_kde": signal_kde / template_kde, }[args.which_kde] # Sort each of the parameter arrays -if args.plot_order == 'increasing': +if args.plot_order == "increasing": idx = numpy.argsort(kde_values) -elif args.plot_order == 'decreasing': +elif args.plot_order == "decreasing": idx = numpy.argsort(kde_values)[::-1] else: idx = numpy.arange(len(kde_values)) @@ -62,41 +74,43 @@ param0 = param_arrays[0] param1 = param_arrays[1] if len(param_arrays) > 1 else None kde_values = kde_values[idx] -if args.plot_type == 'kde_vs_param': - fig, ax = plt.subplots(1, figsize=(12,7), constrained_layout=True) +if args.plot_type == "kde_vs_param": + fig, ax = plt.subplots(1, figsize=(12, 7), constrained_layout=True) im = ax.scatter(kde_values, param0, marker=".", c="r", s=5) ax.set_xticklabels(args.which_kde, fontsize=13) ax.set_yticklabels(args.param[0], fontsize=13) ax.set_xlabel(args.which_kde, fontsize=15) ax.set_ylabel(args.param[0], fontsize=15) - ax.set_xscale('log') - if args.log_axis[0] == 'True': - ax.set_yscale('log') + ax.set_xscale("log") + if args.log_axis[0] == "True": + ax.set_yscale("log") else: - ax.set_yscale('linear') - plot_loc = args.plot_dir + args.which_kde + '_vs_' + args.param[0] + '.png' + ax.set_yscale("linear") + plot_loc = args.plot_dir + args.which_kde + "_vs_" + args.param[0] + ".png" plt.savefig(plot_loc) -elif args.plot_type == 'param_vs_param': - fig, ax = plt.subplots(1, figsize=(12,7), constrained_layout=True) - im = ax.scatter(param0, param1, marker=".", c=kde_values, cmap='turbo', s=5, norm=LogNorm()) +elif args.plot_type == "param_vs_param": + fig, ax = plt.subplots(1, figsize=(12, 7), constrained_layout=True) + im = ax.scatter( + param0, param1, marker=".", c=kde_values, cmap="turbo", s=5, norm=LogNorm() + ) cbar = fig.colorbar(im, ax=ax, pad=0.01) ax.set_xticklabels(args.param[0], fontsize=13) ax.set_yticklabels(args.param[1], fontsize=13) ax.set_xlabel(args.param[0], fontsize=15) ax.set_ylabel(args.param[1], fontsize=15) - if args.log_axis[0] == 'True': - ax.set_xscale('log') + if args.log_axis[0] == "True": + ax.set_xscale("log") else: - ax.set_xscale('linear') - if args.log_axis[1] == 'True': - ax.set_yscale('log') + ax.set_xscale("linear") + if args.log_axis[1] == "True": + ax.set_yscale("log") else: - ax.set_yscale('linear') + ax.set_yscale("linear") cbar.ax.set_ylabel(args.which_kde, rotation=270, fontsize=15, labelpad=15) cbar.ax.tick_params(labelsize=15) - plot_loc = f'{args.plot_dir}/{args.which_kde}_{args.plot_order}_{args.param[0]}_vs_{args.param[1]}.png' + plot_loc = f"{args.plot_dir}/{args.which_kde}_{args.plot_order}_{args.param[0]}_vs_{args.param[1]}.png" plt.savefig(plot_loc) else: - raise RuntimeError('Unknown plot type!', args.plot_type) + raise RuntimeError("Unknown plot type!", args.plot_type) diff --git a/bin/all_sky_search/pycbc_prepare_xml_for_gracedb b/bin/all_sky_search/pycbc_prepare_xml_for_gracedb index af748be75f4..087744d458d 100755 --- a/bin/all_sky_search/pycbc_prepare_xml_for_gracedb +++ b/bin/all_sky_search/pycbc_prepare_xml_for_gracedb @@ -23,9 +23,11 @@ for upload to gracedb. import argparse import logging -import numpy as np + import matplotlib -matplotlib.use('agg') +import numpy as np + +matplotlib.use("agg") import lal import lal.series @@ -40,63 +42,93 @@ from pycbc.io.ligolw import ( snr_series_to_xml, ) from pycbc.psd import interpolate -from pycbc.types import FrequencySeries, load_timeseries -from pycbc.types import MultiDetOptionAction from pycbc.results import generate_asd_plot, generate_snr_plot +from pycbc.types import FrequencySeries, MultiDetOptionAction, load_timeseries parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--psd-files", nargs='+', required=True, - help='HDF file(s) containing the PSDs to upload') -parser.add_argument("--snr-timeseries", nargs='+', required=True, - help='HDF file(s) containing the SNR timeseries to upload') -parser.add_argument('--input-file', required=True, type=str, - help='Input LIGOLW XML file of coincidences.') -parser.add_argument('--event-id', type=int, required=True, - help='ID for the event being prepared.') -parser.add_argument('--output-file', - help="Output filename for the locally stored XML. ") -parser.add_argument('--snr-timeseries-plot', - help="Output filename for the plot of SNR timeseries. " - "Must be a plot filetype as used by matplotlib.") -parser.add_argument('--psd-plot', - help="Output filename for the plot of the PSD. " - "Must be a plot filetype as used by matplotlib.") -parser.add_argument('--channel-name', action=MultiDetOptionAction, - required=True, - help="Channel name for adding to the uploaded single " - "inspiral table.") -parser.add_argument('--delta-f', type=float, default=0.25, - help="Frequency spacing of the PSD to be uploaded, Hz. " - "Default 0.25") +parser.add_argument( + "--psd-files", + nargs="+", + required=True, + help="HDF file(s) containing the PSDs to upload", +) +parser.add_argument( + "--snr-timeseries", + nargs="+", + required=True, + help="HDF file(s) containing the SNR timeseries to upload", +) +parser.add_argument( + "--input-file", + required=True, + type=str, + help="Input LIGOLW XML file of coincidences.", +) +parser.add_argument( + "--event-id", type=int, required=True, help="ID for the event being prepared." +) +parser.add_argument( + "--output-file", help="Output filename for the locally stored XML. " +) +parser.add_argument( + "--snr-timeseries-plot", + help="Output filename for the plot of SNR timeseries. " + "Must be a plot filetype as used by matplotlib.", +) +parser.add_argument( + "--psd-plot", + help="Output filename for the plot of the PSD. " + "Must be a plot filetype as used by matplotlib.", +) +parser.add_argument( + "--channel-name", + action=MultiDetOptionAction, + required=True, + help="Channel name for adding to the uploaded single inspiral table.", +) +parser.add_argument( + "--delta-f", + type=float, + default=0.25, + help="Frequency spacing of the PSD to be uploaded, Hz. Default 0.25", +) args = parser.parse_args() # Default logging level is info: --verbose adds to this pycbc.init_logging(args.verbose, default_level=1) -xmldoc = ligolw_utils.load_filename(args.input_file, - contenthandler=LIGOLWContentHandler) +xmldoc = ligolw_utils.load_filename( + args.input_file, contenthandler=LIGOLWContentHandler +) + class psd_segment(segment): def __new__(cls, psd, *args): return segment.__new__(cls, *args) + def __init__(self, psd, *args): self.psd = psd + def read_psds(psd_files): logging.info("Reading PSDs") psds = {} for psd_file in psd_files: - (ifo, group), = HFile(psd_file, "r").items() + ((ifo, group),) = HFile(psd_file, "r").items() psd = [group["psds"][str(i)] for i in range(len(group["psds"].keys()))] - psds[ifo] = segmentlist(psd_segment(*segargs) for segargs in zip( - psd, group["start_time"], group["end_time"])) + psds[ifo] = segmentlist( + psd_segment(*segargs) + for segargs in zip(psd, group["start_time"], group["end_time"]) + ) return psds + psds = read_psds(args.psd_files) + def read_snr_timeseries(snr_timeseries_files): """ Get information from the single_template snr timeseries files @@ -104,12 +136,13 @@ def read_snr_timeseries(snr_timeseries_files): logging.info("Reading SNR timeseries") snr_timeseries = {} for snr_filename in snr_timeseries_files: - with HFile(snr_filename, 'r') as snr_f: - ifo = snr_f.attrs['ifo'] - time = snr_f.attrs['event_time'] - snr_timeseries[(ifo, time)] = load_timeseries(snr_filename, 'snr') + with HFile(snr_filename, "r") as snr_f: + ifo = snr_f.attrs["ifo"] + time = snr_f.attrs["event_time"] + snr_timeseries[(ifo, time)] = load_timeseries(snr_filename, "snr") return snr_timeseries + snr_timeseries = read_snr_timeseries(args.snr_timeseries) coinc_table = lsctables.CoincTable.get_table(xmldoc) @@ -141,8 +174,10 @@ for coinc_insp in coinc_inspiral_table: if coinc_insp.coinc_event_id == event.coinc_event_id: coinc_inspiral_table_curr.append(coinc_insp) -time = coinc_inspiral_table_curr[0].end_time \ - + coinc_inspiral_table_curr[0].end_time_ns * 1e-9 +time = ( + coinc_inspiral_table_curr[0].end_time + + coinc_inspiral_table_curr[0].end_time_ns * 1e-9 +) sngl_ids = [] for coinc_map in coinc_event_map_table: @@ -158,22 +193,23 @@ for coinc_map in coinc_event_map_table: snr_ts = {} for ifo, t in snr_timeseries.keys(): if not abs(t - time) < 0.5: - raise ValueError("SNR timeseries for IFO %s does not look like it " + raise ValueError( + "SNR timeseries for IFO %s does not look like it " "corresponds to this event, event time %.3f, SNR timeseries is " - "around time %.3f" % (ifo, time, t)) + "around time %.3f" % (ifo, time, t) + ) # Convert the SNR dict to be keyed on IFO-only: snr_ts[ifo] = snr_timeseries[(ifo, t)] # IFOs from SNR timeseries: psds_event = {} psddict = {} -for ifo in snr_ts.keys(): +for ifo in snr_ts: psd = psds[ifo] psd = psd[psd.find(time)].psd # resample the psd to new spacing - psd_fs = FrequencySeries(psd, delta_f=psd.attrs["delta_f"], - dtype=np.float64) + psd_fs = FrequencySeries(psd, delta_f=psd.attrs["delta_f"], dtype=np.float64) psd_fs = interpolate(psd_fs, args.delta_f) psds_event[ifo] = psd @@ -192,7 +228,7 @@ for sngl in sngl_inspiral_table: psd = psds_event[sngl.ifo] psd_fs = psddict[sngl.ifo] - flow = psd.file.attrs['low_frequency_cutoff'] + flow = psd.file.attrs["low_frequency_cutoff"] kmin = int(flow / args.delta_f) fseries = lal.CreateREAL8FrequencySeries( @@ -201,7 +237,8 @@ for sngl in sngl_inspiral_table: kmin * args.delta_f, args.delta_f, lal.StrainUnit**2 / lal.HertzUnit, - len(psd_fs) - kmin) + len(psd_fs) - kmin, + ) fseries.data.data = psd_fs[kmin:] / np.square(pycbc.DYN_RANGE_FAC) lal_psddict[sngl.ifo] = fseries @@ -221,15 +258,12 @@ if args.psd_plot: generate_asd_plot(psddict, args.psd_plot) if args.snr_timeseries_plot: - triggers = {sngl.ifo: (sngl.end_time + sngl.end_time_ns * 1e-9, sngl.snr) - for sngl in sngl_inspiral_table_curr} + triggers = { + sngl.ifo: (sngl.end_time + sngl.end_time_ns * 1e-9, sngl.snr) + for sngl in sngl_inspiral_table_curr + } base_time = int(np.floor(time)) logging.info("Saving SNR plot %s", args.snr_timeseries_plot) - generate_snr_plot( - snr_ts, - args.snr_timeseries_plot, - triggers, - base_time - ) + generate_snr_plot(snr_ts, args.snr_timeseries_plot, triggers, base_time) -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_reduce_template_bank b/bin/all_sky_search/pycbc_reduce_template_bank index 0e5dce3ff0a..e337ec880cc 100644 --- a/bin/all_sky_search/pycbc_reduce_template_bank +++ b/bin/all_sky_search/pycbc_reduce_template_bank @@ -20,41 +20,48 @@ Reduce a template bank using some input parameter cuts """ - -import logging import argparse +import logging + import pycbc -from pycbc.io import HFile from pycbc import load_source +from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-bank", required=True, - help="Input template bank HDF file.") -parser.add_argument("--output-bank", required=True, - help="Output template bank HDF file.") -parser.add_argument("--filter-func-file", required=True, - help="This can be provided to give a function to define " - "which points are covered by the template bank " - "bounds, and which are not. The file should contain " - "a function called filter_tmpltbank, which should " - "take as call profile the template bank HDF object " - "and return a boolean (accept=1/reject=0) array.") +parser.add_argument("--input-bank", required=True, help="Input template bank HDF file.") +parser.add_argument( + "--output-bank", required=True, help="Output template bank HDF file." +) +parser.add_argument( + "--filter-func-file", + required=True, + help="This can be provided to give a function to define " + "which points are covered by the template bank " + "bounds, and which are not. The file should contain " + "a function called filter_tmpltbank, which should " + "take as call profile the template bank HDF object " + "and return a boolean (accept=1/reject=0) array.", +) opt = parser.parse_args() pycbc.init_logging(opt.verbose) -bank_fd = HFile(opt.input_bank, 'r') +bank_fd = HFile(opt.input_bank, "r") -modl = load_source('filter_func', opt.filter_func_file) +modl = load_source("filter_func", opt.filter_func_file) func = modl.filter_tmpltbank bool_arr = func(bank_fd) -logging.info("Downselecting templates. Started with", len(bool_arr), - "templates, now have ", bool_arr.sum(), "after downselecting.") -bank_ofd = HFile(opt.output_bank, 'w') +logging.info( + "Downselecting templates. Started with", + len(bool_arr), + "templates, now have ", + bool_arr.sum(), + "after downselecting.", +) +bank_ofd = HFile(opt.output_bank, "w") for name in bank_fd.keys(): bank_ofd[name] = bank_fd[name][:][bool_arr] bank_ofd.close() - diff --git a/bin/all_sky_search/pycbc_rerank_passthrough b/bin/all_sky_search/pycbc_rerank_passthrough index 5c8b44945ff..8e5ca4dc2d0 100644 --- a/bin/all_sky_search/pycbc_rerank_passthrough +++ b/bin/all_sky_search/pycbc_rerank_passthrough @@ -1,37 +1,47 @@ #!/bin/env python """Dummy script to pass through stat files and test reranking""" -import argparse, logging, pycbc + +import argparse +import logging + +import pycbc from pycbc.io import HFile parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--output-file', - help="File containing the newly assigned statistic values") +parser.add_argument( + "--output-file", help="File containing the newly assigned statistic values" +) # Options related to getting trigger information from workflow products -parser.add_argument('--input-file', - help="HDF File which gives the trigger followup information for a set") -parser.add_argument('--start-index', type=int, - help="Analyzing candidates starting from this index") -parser.add_argument('--end-index', type=int, - help="Analyzing candidates stopping at this index") -parser.add_argument('--stride', type=int, default=1, - help="Only analyze every Nth candidate") +parser.add_argument( + "--input-file", + help="HDF File which gives the trigger followup information for a set", +) +parser.add_argument( + "--start-index", type=int, help="Analyzing candidates starting from this index" +) +parser.add_argument( + "--end-index", type=int, help="Analyzing candidates stopping at this index" +) +parser.add_argument( + "--stride", type=int, default=1, help="Only analyze every Nth candidate" +) args, unknown = parser.parse_known_args() pycbc.init_logging(args.verbose) -ofile = HFile(args.output_file, 'w') +ofile = HFile(args.output_file, "w") -f = HFile(args.input_file, 'r') +f = HFile(args.input_file, "r") start = 0 if args.start_index is None else args.start_index -end = len(f['time']) if args.end_index is None else args.end_index +end = len(f["time"]) if args.end_index is None else args.end_index stride = args.stride -ofile.attrs['start_index'] = start -ofile.attrs['end_index' ] = end -ofile.attrs['stride'] = stride -stat = f['stat'][:][start:end:stride] +ofile.attrs["start_index"] = start +ofile.attrs["end_index"] = end +ofile.attrs["stride"] = stride +stat = f["stat"][:][start:end:stride] -ofile['stat'] = stat -logging.info('Done') +ofile["stat"] = stat +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_sngls_findtrigs b/bin/all_sky_search/pycbc_sngls_findtrigs index 027f833201f..70bb0d88efa 100644 --- a/bin/all_sky_search/pycbc_sngls_findtrigs +++ b/bin/all_sky_search/pycbc_sngls_findtrigs @@ -1,68 +1,96 @@ #!/usr/bin/env python -import argparse, logging, numpy as np +import argparse +import logging + +import numpy as np from numpy.random import seed, shuffle import pycbc -from pycbc.events import veto, coinc, stat -from pycbc import io -from pycbc.events import cuts +from pycbc import init_logging, io +from pycbc.events import coinc, cuts, stat, veto from pycbc.types.optparse import MultiDetOptionAction -from pycbc import init_logging parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) # Basic file input options -parser.add_argument("--trigger-files", type=str, nargs=1, - help="File containing single-detector triggers") -parser.add_argument("--template-bank", required=True, - help="Template bank file in HDF format") -parser.add_argument("--template-fraction-range", default="0/1", - help="Optional, analyze only part of template bank. Format" - " PART/NUM_PARTS") -parser.add_argument("--randomize-template-order", action="store_true", - help="Random shuffle templates with fixed seed " - "before selecting range to analyze") +parser.add_argument( + "--trigger-files", + type=str, + nargs=1, + help="File containing single-detector triggers", +) +parser.add_argument( + "--template-bank", required=True, help="Template bank file in HDF format" +) +parser.add_argument( + "--template-fraction-range", + default="0/1", + help="Optional, analyze only part of template bank. Format PART/NUM_PARTS", +) +parser.add_argument( + "--randomize-template-order", + action="store_true", + help="Random shuffle templates with fixed seed before selecting range to analyze", +) # Options to define the vetoes -parser.add_argument("--veto-files", nargs='*', action='append', default=[], - help="Optional veto file. Triggers within veto segments " - "contained in the file are ignored") -parser.add_argument("--segment-name", nargs='*', action='append', default=[], - help="Optional, name of veto segment in veto file") -parser.add_argument("--gating-veto-windows", nargs='+', - action=MultiDetOptionAction, - help="Seconds to be vetoed before and after the central time " - "of each gate. Given as detector-values pairs, e.g. " - "H1:-1,2.5 L1:-1,2.5 V1:0,0") +parser.add_argument( + "--veto-files", + nargs="*", + action="append", + default=[], + help="Optional veto file. Triggers within veto segments " + "contained in the file are ignored", +) +parser.add_argument( + "--segment-name", + nargs="*", + action="append", + default=[], + help="Optional, name of veto segment in veto file", +) +parser.add_argument( + "--gating-veto-windows", + nargs="+", + action=MultiDetOptionAction, + help="Seconds to be vetoed before and after the central time " + "of each gate. Given as detector-values pairs, e.g. " + "H1:-1,2.5 L1:-1,2.5 V1:0,0", +) # additional veto options # produces a list of lists to allow multiple invocations and multiple args -parser.add_argument('--cluster-window', type=float, - help='Window (seconds) during which to keep the trigger ' - 'with the loudest statistic value. ' - 'Default=do not cluster') -parser.add_argument("--output-file", - help="File to store the candidate triggers") -parser.add_argument("--loudest-keep-values", - default='[6:1]', - help="Apply successive multiplicative levels of" - " decimation to coincs with stat value below the" - " given thresholds. Supply as a comma-separated list" - " of threshold:decimation value pairs surrounded by" - " square brackets (no spaces!). Decimation values must" - " be positive integers." - " Ex. [15:5,10:30,5:30,0:30]." - " Default: no decimation") +parser.add_argument( + "--cluster-window", + type=float, + help="Window (seconds) during which to keep the trigger " + "with the loudest statistic value. " + "Default=do not cluster", +) +parser.add_argument("--output-file", help="File to store the candidate triggers") +parser.add_argument( + "--loudest-keep-values", + default="[6:1]", + help="Apply successive multiplicative levels of" + " decimation to coincs with stat value below the" + " given thresholds. Supply as a comma-separated list" + " of threshold:decimation value pairs surrounded by" + " square brackets (no spaces!). Decimation values must" + " be positive integers." + " Ex. [15:5,10:30,5:30,0:30]." + " Default: no decimation", +) stat.insert_statistic_option_group(parser) cuts.insert_cuts_option_group(parser) args = parser.parse_args() trigger_file = args.trigger_files[0] -if (args.veto_files and not args.segment_name) or \ - (args.segment_name and not args.veto_files): - raise RuntimeError('--veto-files and --segment-name are mutually required') +if (args.veto_files and not args.segment_name) or ( + args.segment_name and not args.veto_files +): + raise RuntimeError("--veto-files and --segment-name are mutually required") if not len(args.veto_files) == len(args.segment_name): - raise RuntimeError('--segment-name optionss are required for each --veto-files') + raise RuntimeError("--segment-name optionss are required for each --veto-files") args.segment_name = sum(args.segment_name, []) args.veto_files = sum(args.veto_files, []) @@ -71,22 +99,24 @@ init_logging(args.verbose) trigger_cut_dict, template_cut_dict = cuts.ingest_cuts_option_group(args) -logging.info('Opening trigger file: %s', trigger_file) -trigf = io.HFile(trigger_file, 'r') +logging.info("Opening trigger file: %s", trigger_file) +trigf = io.HFile(trigger_file, "r") ifo = list(trigf.keys())[0] # Set up to only load triggers from the templates of interest + def parse_template_range(num_templates, rangestr): - part = int(rangestr.split('/')[0]) - pieces = int(rangestr.split('/')[1]) + part = int(rangestr.split("/")[0]) + pieces = int(rangestr.split("/")[1]) tmin = int(num_templates / float(pieces) * part) - tmax = int(num_templates / float(pieces) * (part+1)) + tmax = int(num_templates / float(pieces) * (part + 1)) return tmin, tmax -num_templates = io.HFile(args.template_bank, "r")['template_hash'].size + +num_templates = io.HFile(args.template_bank, "r")["template_hash"].size tmin, tmax = parse_template_range(num_templates, args.template_fraction_range) -logging.info('Analyzing template %s - %s' % (tmin, tmax-1)) +logging.info("Analyzing template %s - %s" % (tmin, tmax - 1)) if args.randomize_template_order: seed(0) @@ -99,12 +129,15 @@ else: original_bank_len = len(template_ids) from pycbc.io.hdf import ReadByTemplate -trigs = ReadByTemplate(trigger_file, - args.template_bank, - args.segment_name, - args.veto_files, - args.gating_veto_windows) -logging.info("%d triggers in file", trigf[ifo + '/snr'].size) + +trigs = ReadByTemplate( + trigger_file, + args.template_bank, + args.segment_name, + args.veto_files, + args.gating_veto_windows, +) +logging.info("%d triggers in file", trigf[ifo + "/snr"].size) stat_all = [] trigger_ids_all = [] @@ -119,22 +152,27 @@ template_ids = cuts.apply_template_cuts( template_cut_dict, statistic=rank_method, ifos=[ifo], - template_ids=template_ids) + template_ids=template_ids, +) -logging.info("%d out of %d templates kept after applying template cuts", - len(template_ids), original_bank_len) +logging.info( + "%d out of %d templates kept after applying template cuts", + len(template_ids), + original_bank_len, +) if args.cluster_window is not None: - logging.info('Clustering events over %s s window within each template', - args.cluster_window) + logging.info( + "Clustering events over %s s window within each template", args.cluster_window + ) -loudest_keep_vals = args.loudest_keep_values.strip('[]').split(',') +loudest_keep_vals = args.loudest_keep_values.strip("[]").split(",") threshes = [] factors = [] tot_fac = 1 for decstr in loudest_keep_vals: - thresh, factor = decstr.split(':') + thresh, factor = decstr.split(":") if float(factor) % 1: raise RuntimeError("Non-integer decimation is not supported") if int(factor) < 1: @@ -156,22 +194,24 @@ for i, tnum in enumerate(template_ids): if i % 1000 == 0: logging.info( "Calculating statistic in template %d out of %d", - i, len(template_ids), + i, + len(template_ids), ) else: logging.debug( "Calculating statistic in template %d out of %d", - i, len(template_ids), + i, + len(template_ids), ) tids_uncut = trigs.set_template(tnum) - trigger_keep_ids = cuts.apply_trigger_cuts(trigs, trigger_cut_dict, - statistic=rank_method) + trigger_keep_ids = cuts.apply_trigger_cuts( + trigs, trigger_cut_dict, statistic=rank_method + ) tids_full = tids_uncut[trigger_keep_ids] - logging.debug('%s:%s', tnum, len(tids_uncut)) + logging.debug("%s:%s", tnum, len(tids_uncut)) if len(tids_full) < len(tids_uncut): - logging.debug("%s triggers cut", - len(tids_uncut) - len(tids_full)) + logging.debug("%s triggers cut", len(tids_uncut) - len(tids_full)) n_tot_trigs = tids_full.size if n_tot_trigs == 0: @@ -180,10 +220,9 @@ for i, tnum in enumerate(template_ids): # Stat class instance to calculate the ranking statistic sds = rank_method.single(trigs)[trigger_keep_ids] stat_t = rank_method.rank_stat_single((ifo, sds)) - trigger_times = trigs['end_time'][:][trigger_keep_ids] + trigger_times = trigs["end_time"][:][trigger_keep_ids] if args.cluster_window is not None: - cid = coinc.cluster_over_time(stat_t, trigger_times, - args.cluster_window) + cid = coinc.cluster_over_time(stat_t, trigger_times, args.cluster_window) stat_t = stat_t[cid] tids_full = tids_full[cid] trigger_times = trigger_times[cid] @@ -225,34 +264,35 @@ for t, f in zip(threshes, factors): "%d events after decimation at statistic %.3f with factor %d", dec_facs.size, t, - f + f, ) -data = {"stat": stat_all, - "decimation_factor": dec_facs, - "timeslide_id": np.zeros_like(stat_all), - "template_id": template_ids_all, - "%s/time" % ifo : trigger_times_all, - "%s/trigger_id" % ifo: trigger_ids_all} +data = { + "stat": stat_all, + "decimation_factor": dec_facs, + "timeslide_id": np.zeros_like(stat_all), + "template_id": template_ids_all, + "%s/time" % ifo: trigger_times_all, + "%s/trigger_id" % ifo: trigger_ids_all, +} logging.info("saving triggers") -f = io.HFile(args.output_file, 'w') +f = io.HFile(args.output_file, "w") for key in data: - f.create_dataset(key, data=data[key], - compression="gzip", - compression_opts=9, - shuffle=True) + f.create_dataset( + key, data=data[key], compression="gzip", compression_opts=9, shuffle=True + ) # Store segments -f['segments/%s/start' % ifo], f['segments/%s/end' % ifo] = trigs.valid +f["segments/%s/start" % ifo], f["segments/%s/end" % ifo] = trigs.valid fg_segs = veto.start_end_to_segments(*trigs.valid) fg_time = abs(fg_segs) -f.attrs['foreground_time'] = fg_time -f.attrs['background_time'] = fg_time -f.attrs['num_of_ifos'] = 1 -f.attrs['pivot'] = ifo -f.attrs['fixed'] = ifo -f.attrs['ifos'] = ifo -f.attrs['timeslide_interval'] = 0 +f.attrs["foreground_time"] = fg_time +f.attrs["background_time"] = fg_time +f.attrs["num_of_ifos"] = 1 +f.attrs["pivot"] = ifo +f.attrs["fixed"] = ifo +f.attrs["ifos"] = ifo +f.attrs["timeslide_interval"] = 0 # Do hierarchical removal # h_iterations = 0 diff --git a/bin/all_sky_search/pycbc_sngls_pastro b/bin/all_sky_search/pycbc_sngls_pastro index 84f6d179a08..f828531c0ff 100644 --- a/bin/all_sky_search/pycbc_sngls_pastro +++ b/bin/all_sky_search/pycbc_sngls_pastro @@ -10,190 +10,245 @@ coincidences. """ -import pycbc, pycbc.io, copy -import argparse, logging, numpy as np -from igwn_ligolw import lsctables, utils as ligolw_utils +import argparse +import copy +import logging + +import matplotlib +import numpy as np +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils +from igwn_segments import segment, segmentlist + +import pycbc +import pycbc.io from pycbc import conversions as conv from pycbc.events import veto from pycbc.io.ligolw import LIGOLWContentHandler -from igwn_segments import segment, segmentlist -import matplotlib -matplotlib.use('agg') + +matplotlib.use("agg") from matplotlib import pyplot as plt from scipy.stats import gaussian_kde as gk - -d_power = { - 'log': 3., - 'uniform': 2., - 'distancesquared': 1., - 'volume': 0. -} +d_power = {"log": 3.0, "uniform": 2.0, "distancesquared": 1.0, "volume": 0.0} mchirp_power = { - 'log': 0., - 'uniform': 5. / 6., - 'distancesquared': 5. / 3., - 'volume': 15. / 6. + "log": 0.0, + "uniform": 5.0 / 6.0, + "distancesquared": 5.0 / 3.0, + "volume": 15.0 / 6.0, } parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument("--single-statmap-files", nargs='+', required=True, - help="Single statmap files for which p_astro is " - "calculated.") +parser.add_argument( + "--single-statmap-files", + nargs="+", + required=True, + help="Single statmap files for which p_astro is calculated.", +) # Files to help remove foreground events -parser.add_argument('--coinc-statmap-files', nargs='+', required=True, - help="Coincident statmap files, containing coincident " - "events to be removed from the background.") -parser.add_argument('--coinc-veto-ifar-threshold', type=float, default=1.0, - help="Censor triggers around coincidences with " - "IFAR (years) above the threshold [default=1 yr") -parser.add_argument('--coinc-veto-window', type=float, default=0.1, - help="Time around each coincident event above threshold " - "to window out. Default = 0.1s") -parser.add_argument("--remove-n-loudest", type=int, - help="If given, will remove this number of triggers from " - "the background after applying foreground censor. " - "This helps prevent signal contamination, but too " - "many removals can adversely affect background " - "estimate.") +parser.add_argument( + "--coinc-statmap-files", + nargs="+", + required=True, + help="Coincident statmap files, containing coincident " + "events to be removed from the background.", +) +parser.add_argument( + "--coinc-veto-ifar-threshold", + type=float, + default=1.0, + help="Censor triggers around coincidences with " + "IFAR (years) above the threshold [default=1 yr", +) +parser.add_argument( + "--coinc-veto-window", + type=float, + default=0.1, + help="Time around each coincident event above threshold " + "to window out. Default = 0.1s", +) +parser.add_argument( + "--remove-n-loudest", + type=int, + help="If given, will remove this number of triggers from " + "the background after applying foreground censor. " + "This helps prevent signal contamination, but too " + "many removals can adversely affect background " + "estimate.", +) # Arguments for dealing with the injection trigger files for signal population -parser.add_argument('--inj-single-statmap-files', nargs='+', required=True, - help="Single statmap files for the injections. " - "Must be in same order as --inj-files") -parser.add_argument('--inj-files', nargs='+', required=True, - help="File which define injections. Must be in same " - "order as --inj-single-trigger-files") -parser.add_argument('--injection-window', type=float, default=1, - help="Window for how close a trigger is to an " - "injection to be considered to be associated " - "with it (seconds). Default=1.") +parser.add_argument( + "--inj-single-statmap-files", + nargs="+", + required=True, + help="Single statmap files for the injections. " + "Must be in same order as --inj-files", +) +parser.add_argument( + "--inj-files", + nargs="+", + required=True, + help="File which define injections. Must be in same " + "order as --inj-single-trigger-files", +) +parser.add_argument( + "--injection-window", + type=float, + default=1, + help="Window for how close a trigger is to an " + "injection to be considered to be associated " + "with it (seconds). Default=1.", +) # Arguments to decide the weighting of the injection distribution -parser.add_argument('--distance-param', choices=['distance', 'chirp_distance'], - help="Parameter used to calculate injection distribution " - "for weighting. Default='distance'") -parser.add_argument('--distribution', default='uniform', - choices=['log', 'uniform', 'distancesquared', 'volume'], - help="Form of distribution over --distance-param. " - "Default='uniform'") -parser.add_argument('--expected-signal-rate', type=float, required=True, - help="Expected rate of signals (per year) with stat " - "value above --stat-threshold for use in p_astro " - "calcation.") -parser.add_argument('--signal-ifar-threshold', type=float, default=1, - help="Coincident IFAR threshold to consider an injection " - "as 'found' (years). Default=1") +parser.add_argument( + "--distance-param", + choices=["distance", "chirp_distance"], + help="Parameter used to calculate injection distribution " + "for weighting. Default='distance'", +) +parser.add_argument( + "--distribution", + default="uniform", + choices=["log", "uniform", "distancesquared", "volume"], + help="Form of distribution over --distance-param. Default='uniform'", +) +parser.add_argument( + "--expected-signal-rate", + type=float, + required=True, + help="Expected rate of signals (per year) with stat " + "value above --stat-threshold for use in p_astro " + "calcation.", +) +parser.add_argument( + "--signal-ifar-threshold", + type=float, + default=1, + help="Coincident IFAR threshold to consider an injection " + "as 'found' (years). Default=1", +) # produces a list of lists to allow multiple invocations and multiple args -parser.add_argument('--pastro-method', required=True, - choices=['truncated_shelf', 'callister', 'background_kde'], - help="Which method to use for calculating p_astro. ") -parser.add_argument("--bg-distribution-limit", type=float, default=8, - help="The point at which the noise model will change from " - "the background distribution to --pastro-method") - -parser.add_argument('--plot-distribution', - help="If given, will plot the distribution of rate " - "densities and p_astro given ranking statistic.") -parser.add_argument('--output-file', required=True, - help="name of output file") +parser.add_argument( + "--pastro-method", + required=True, + choices=["truncated_shelf", "callister", "background_kde"], + help="Which method to use for calculating p_astro. ", +) +parser.add_argument( + "--bg-distribution-limit", + type=float, + default=8, + help="The point at which the noise model will change from " + "the background distribution to --pastro-method", +) + +parser.add_argument( + "--plot-distribution", + help="If given, will plot the distribution of rate " + "densities and p_astro given ranking statistic.", +) +parser.add_argument("--output-file", required=True, help="name of output file") args = parser.parse_args() -d_power = { - 'log': 3., - 'uniform': 2., - 'distancesquared': 1., - 'volume': 0. -}[args.distribution] +d_power = {"log": 3.0, "uniform": 2.0, "distancesquared": 1.0, "volume": 0.0}[ + args.distribution +] mchirp_power = { - 'log': 0., - 'uniform': 5. / 6., - 'distancesquared': 5. / 3., - 'volume': 15. / 6. + "log": 0.0, + "uniform": 5.0 / 6.0, + "distancesquared": 5.0 / 3.0, + "volume": 15.0 / 6.0, }[args.distribution] pycbc.init_logging(args.verbose) -sngl_ifo = pycbc.io.HFile(args.single_statmap_files[0], 'r').attrs['ifos'] -groups = ['decimation_factor', 'stat', 'template_id', 'timeslide_id', - sngl_ifo + '/time', sngl_ifo + '/trigger_id'] +sngl_ifo = pycbc.io.HFile(args.single_statmap_files[0], "r").attrs["ifos"] +groups = [ + "decimation_factor", + "stat", + "template_id", + "timeslide_id", + sngl_ifo + "/time", + sngl_ifo + "/trigger_id", +] empty_darray = {g: np.array([]) for g in groups} sngl_trigs = pycbc.io.DictArray(data=empty_darray) -logging.info('Getting single-detector triggers and segments') +logging.info("Getting single-detector triggers and segments") sngl_detector_segs = segmentlist([segment(0, 0)]) for smap_file in args.single_statmap_files: - with pycbc.io.HFile(smap_file, 'r') as f_in: - sngl_starts = f_in['segments/' + sngl_ifo + '/start'][:] - sngl_ends = f_in['segments/' + sngl_ifo + '/end'][:] + with pycbc.io.HFile(smap_file, "r") as f_in: + sngl_starts = f_in["segments/" + sngl_ifo + "/start"][:] + sngl_ends = f_in["segments/" + sngl_ifo + "/end"][:] sngl_trigs += pycbc.io.DictArray(data={g: f_in[g][:] for g in groups}) sngl_detector_segs += veto.start_end_to_segments(sngl_starts, sngl_ends) fg_time = abs(sngl_detector_segs) coinc_segs = segmentlist([segment(0, 0)]) fgveto_segs = segmentlist([segment(0, 0)]) -logging.info('Getting coinc segments and foreground vetoes') +logging.info("Getting coinc segments and foreground vetoes") for cfilename in args.coinc_statmap_files: - with pycbc.io.HFile(cfilename, 'r') as cfile: + with pycbc.io.HFile(cfilename, "r") as cfile: # Coinc segments are different depending on statmap file - if 'coinc' in cfile['segments']: - c_starts = cfile['segments/coinc/start'][:] - c_ends = cfile['segments/coinc/end'][:] + if "coinc" in cfile["segments"]: + c_starts = cfile["segments/coinc/start"][:] + c_ends = cfile["segments/coinc/end"][:] else: - ctype = cfile.attrs['ifos'].replace(' ','') - c_starts = cfile['segments/' + ctype + '/start'][:] - c_ends = cfile['segments/' + ctype + '/end'][:] + ctype = cfile.attrs["ifos"].replace(" ", "") + c_starts = cfile["segments/" + ctype + "/start"][:] + c_ends = cfile["segments/" + ctype + "/end"][:] # Test if the ifo is actually in this statmap file - if 'ifos' in cfile.attrs.keys(): - ifolist = cfile.attrs['ifos'].split(' ') - time_key = sngl_ifo + '/time' + if "ifos" in cfile.attrs.keys(): + ifolist = cfile.attrs["ifos"].split(" ") + time_key = sngl_ifo + "/time" else: - ifolist = [cfile.attrs['detector_1'], - cfile.attrs['detector_2']] - if sngl_ifo == cfile.attrs['detector_1']: - time_key = 'time1' + ifolist = [cfile.attrs["detector_1"], cfile.attrs["detector_2"]] + if sngl_ifo == cfile.attrs["detector_1"]: + time_key = "time1" else: - time_key = 'time2' + time_key = "time2" if sngl_ifo not in ifolist: logging.warning("IFO %s is not in file %s", sngl_ifo, cfilename) continue # Find the coinc events above a threshold IFAR - cifar = cfile['foreground/ifar'][:] + cifar = cfile["foreground/ifar"][:] c_above = cifar > args.coinc_veto_ifar_threshold - ctime = cfile['foreground'][time_key][:][c_above] + ctime = cfile["foreground"][time_key][:][c_above] c_fgv_starts = ctime - args.coinc_veto_window c_fgv_ends = ctime + args.coinc_veto_window coinc_segs = coinc_segs + veto.start_end_to_segments(c_starts, c_ends) - fgveto_segs = fgveto_segs + veto.start_end_to_segments(c_fgv_starts, - c_fgv_ends) + fgveto_segs = fgveto_segs + veto.start_end_to_segments(c_fgv_starts, c_fgv_ends) coinc_segs.coalesce() fgveto_segs.coalesce() logging.info("Removing triggers in single-detector time from background.") -bg_bool = np.array([t in coinc_segs - for t in sngl_trigs.data[sngl_ifo + '/time']]) +bg_bool = np.array([t in coinc_segs for t in sngl_trigs.data[sngl_ifo + "/time"]]) bg_idx = np.flatnonzero(bg_bool) bg_trigs = sngl_trigs.select(bg_idx) logging.info("Removing triggers in foreground vetoed time from background.") -fg_veto_bool_bg = np.array([t not in fgveto_segs - for t in bg_trigs.data[sngl_ifo + '/time']]) +fg_veto_bool_bg = np.array( + [t not in fgveto_segs for t in bg_trigs.data[sngl_ifo + "/time"]] +) bg_trigs = bg_trigs.select(np.flatnonzero(fg_veto_bool_bg)) bg_time = abs(coinc_segs - fgveto_segs) if args.remove_n_loudest: - logging.info("Removing %d loudest remaining triggers from background", - args.remove_n_loudest) - loudest_idx = np.argsort(bg_trigs.data['stat'])[-args.remove_n_loudest:] + logging.info( + "Removing %d loudest remaining triggers from background", args.remove_n_loudest + ) + loudest_idx = np.argsort(bg_trigs.data["stat"])[-args.remove_n_loudest :] bg_trigs = bg_trigs.remove(loudest_idx) logging.info("Getting injected signal triggers and information.") @@ -202,50 +257,52 @@ signal_stat = np.array([]) signal_weights = np.array([]) # This method requires injection definition files and injection single statmap files to be # in the same order - would prefer to not have to do this -for inj_filename, inj_trigger_filename in zip(args.inj_files, - args.inj_single_statmap_files): - - logging.info('Reading injection statistic file') - with pycbc.io.HFile(inj_trigger_filename, 'r') as inj_trig_f: - inj_trig_id = inj_trig_f[sngl_ifo]['trigger_id'][:] - inj_trig_time = inj_trig_f[sngl_ifo]['time'][:] - inj_stat = inj_trig_f['stat'][:] +for inj_filename, inj_trigger_filename in zip( + args.inj_files, args.inj_single_statmap_files +): + logging.info("Reading injection statistic file") + with pycbc.io.HFile(inj_trigger_filename, "r") as inj_trig_f: + inj_trig_id = inj_trig_f[sngl_ifo]["trigger_id"][:] + inj_trig_time = inj_trig_f[sngl_ifo]["time"][:] + inj_stat = inj_trig_f["stat"][:] time_sort = inj_trig_time.argsort() - logging.info('Reading injection file') - indoc = ligolw_utils.load_filename(inj_filename, False, - contenthandler=LIGOLWContentHandler) + logging.info("Reading injection file") + indoc = ligolw_utils.load_filename( + inj_filename, False, contenthandler=LIGOLWContentHandler + ) sim_table = lsctables.SimInspiralTable.get_table(indoc) - inj_time = np.array(sim_table.get_column('geocent_end_time') + - 1e-9 * sim_table.get_column('geocent_end_time_ns'), - dtype=np.float64) - - left = np.searchsorted(inj_trig_time[time_sort], - inj_time - args.injection_window, side='left') - right = np.searchsorted(inj_trig_time[time_sort], - inj_time + args.injection_window, side='right') + inj_time = np.array( + sim_table.get_column("geocent_end_time") + + 1e-9 * sim_table.get_column("geocent_end_time_ns"), + dtype=np.float64, + ) + + left = np.searchsorted( + inj_trig_time[time_sort], inj_time - args.injection_window, side="left" + ) + right = np.searchsorted( + inj_trig_time[time_sort], inj_time + args.injection_window, side="right" + ) found = np.flatnonzero((right - left) == 1) - found_d = np.array(sim_table.get_column('distance'), - dtype=np.float32)[found] - found_m1 = np.array(sim_table.get_column('mass1'), - dtype=np.float32)[found] - found_m2 = np.array(sim_table.get_column('mass2'), - dtype=np.float32)[found] + found_d = np.array(sim_table.get_column("distance"), dtype=np.float32)[found] + found_m1 = np.array(sim_table.get_column("mass1"), dtype=np.float32)[found] + found_m2 = np.array(sim_table.get_column("mass2"), dtype=np.float32)[found] signal_stat = np.append(signal_stat, inj_stat[left[found]]) found_mchirp = conv.mchirp_from_mass1_mass2(found_m1, found_m2) - if args.distance_param == 'chirp_distance': + if args.distance_param == "chirp_distance": found_d = conv.chirp_distance(found_d, found_mchirp) - weights = found_mchirp ** mchirp_power * found_d ** d_power + weights = found_mchirp**mchirp_power * found_d**d_power signal_weights = np.append(signal_weights, weights) -bg_stat = bg_trigs.data['stat'] -fg_stat = sngl_trigs.data['stat'] +bg_stat = bg_trigs.data["stat"] +fg_stat = sngl_trigs.data["stat"] n_sig_exp = args.expected_signal_rate * conv.sec_to_year(fg_time) logging.info("%.3f signals are expected in this amount of data", n_sig_exp) @@ -253,14 +310,12 @@ logging.info("%.3f signals are expected in this amount of data", n_sig_exp) max_bg = bg_stat.max() min_bg = bg_stat.min() -logging.info('Getting Gaussian kernel density estimates of signal ' - 'distribution') +logging.info("Getting Gaussian kernel density estimates of signal distribution") sig_kern = gk(signal_stat, weights=signal_weights, bw_method=1) sig_norm = sig_kern.integrate_box_1d(max_bg, np.inf) sig_dens = sig_kern(fg_stat) / sig_norm -logging.info('Getting Gaussian kernel density estimates of background ' - 'distribution') +logging.info("Getting Gaussian kernel density estimates of background distribution") bg_kern = gk(bg_stat, bw_method=1) bg_rate_dens = bg_kern(fg_stat) * bg_stat.size @@ -272,20 +327,22 @@ bg_dens_max_idx = bg_rate_dens.argmax() bg_dens_max_stat = fg_stat[bg_dens_max_idx] fg_stat_below_peak = fg_stat < bg_dens_max_stat -logging.info('Getting noise distribution of %s method.', args.pastro_method) +logging.info("Getting noise distribution of %s method.", args.pastro_method) # If the foreground statistic value is quieter than X # then use background density for noise pastro_noise_valid = fg_stat >= args.bg_distribution_limit noise_model = copy.deepcopy(bg_rate_dens) -if args.pastro_method == 'callister': +if args.pastro_method == "callister": noise_model[pastro_noise_valid] = sig_dens[pastro_noise_valid] -elif args.pastro_method == 'truncated_shelf': +elif args.pastro_method == "truncated_shelf": fgs = fg_stat[pastro_noise_valid] - stat_below = np.array([max(bg_stat[bg_stat < fs]) - if fs > min_bg else -np.inf for fs in fgs]) - ts_norm = np.array([sig_kern.integrate_box_1d(sbel, fs) - for sbel, fs in zip(stat_below, fgs)]) + stat_below = np.array( + [max(bg_stat[bg_stat < fs]) if fs > min_bg else -np.inf for fs in fgs] + ) + ts_norm = np.array( + [sig_kern.integrate_box_1d(sbel, fs) for sbel, fs in zip(stat_below, fgs)] + ) noise_model[pastro_noise_valid] = sig_kern(fgs) / ts_norm logging.info("Calculating pastro for foreground events") @@ -294,7 +351,7 @@ p_astro[fg_stat_below_peak] = 0 if args.plot_distribution: logging.info("Making distribution plot") - max_bin_c = min(signal_stat.max(), bg_stat.max() * 4) + max_bin_c = min(signal_stat.max(), bg_stat.max() * 4) stat_bins = np.linspace(min_bg, max_bin_c, 201) stat_bin_c = (stat_bins[1:] + stat_bins[:-1]) / 2 @@ -304,50 +361,55 @@ if args.plot_distribution: # Plot background trigger rate density bg_rate_dens_plt = bg_kern(stat_bin_c) * bg_stat.size - ax0.plot(stat_bin_c, bg_rate_dens_plt, linestyle="--", - label="Background KDE") + ax0.plot(stat_bin_c, bg_rate_dens_plt, linestyle="--", label="Background KDE") # Plot signal trigger rate density (from injections) sig_rate_dens_plt = sig_kern(stat_bin_c) / sig_norm - ax0.plot(stat_bin_c, sig_rate_dens_plt, linestyle=":", - label="Signals KDE") + ax0.plot(stat_bin_c, sig_rate_dens_plt, linestyle=":", label="Signals KDE") # Calculate noise model according to chosen method - if args.pastro_method == 'callister': + if args.pastro_method == "callister": noise_mod_plt = sig_rate_dens_plt - elif args.pastro_method == 'truncated_shelf': - stat_below_plt = [max(bg_stat[sbc >= bg_stat]) - if sbc >= min_bg else -np.inf for sbc in stat_bin_c] - ts_norm = np.array([sig_kern.integrate_box_1d(sbel, sbc) - for sbel, sbc in zip(stat_below_plt, stat_bin_c)]) + elif args.pastro_method == "truncated_shelf": + stat_below_plt = [ + max(bg_stat[sbc >= bg_stat]) if sbc >= min_bg else -np.inf + for sbc in stat_bin_c + ] + ts_norm = np.array( + [ + sig_kern.integrate_box_1d(sbel, sbc) + for sbel, sbc in zip(stat_below_plt, stat_bin_c) + ] + ) noise_mod_plt = sig_kern(stat_bin_c) / ts_norm - elif args.pastro_method == 'background_kde': + elif args.pastro_method == "background_kde": noise_mod_plt = bg_rate_dens_plt # triggers below the chosen distribution limit use the background # distribution for noise model past_noise_idx_plt = stat_bin_c < args.bg_distribution_limit noise_mod_plt[past_noise_idx_plt] = bg_rate_dens_plt[past_noise_idx_plt] - ax0.plot(stat_bin_c, noise_mod_plt, linestyle='-', c='k', - label=args.pastro_method) + ax0.plot(stat_bin_c, noise_mod_plt, linestyle="-", c="k", label=args.pastro_method) ax0.semilogy() ax0.grid() # Set axis limits according to the highest point of the background # distribution ax0.set_xlim([bg_dens_max_stat, max_bin_c]) - ax0.set_ylim([sig_kern(max_bin_c) / sig_norm / 1.5, - bg_rate_dens[bg_dens_max_idx] * 1.5]) + ax0.set_ylim( + [sig_kern(max_bin_c) / sig_norm / 1.5, bg_rate_dens[bg_dens_max_idx] * 1.5] + ) ax0.legend() ax0.set_xlabel("Ranking Statistic") ax0.set_ylabel("Rate Density") # Calculate p_astro for the plotted statistic values - p_astro_plt = \ + p_astro_plt = ( n_sig_exp * sig_rate_dens_plt / (noise_mod_plt + n_sig_exp * sig_rate_dens_plt) + ) - ax1.plot(stat_bin_c, p_astro_plt, c='k', label="p_astro distribution") - ax1.scatter(fg_stat, p_astro, c='k', marker='x', label='foreground events') + ax1.plot(stat_bin_c, p_astro_plt, c="k", label="p_astro distribution") + ax1.scatter(fg_stat, p_astro, c="k", marker="x", label="foreground events") ax1.grid() ax1.legend() ax1.set_xlim([bg_dens_max_stat, max_bin_c]) @@ -357,33 +419,40 @@ if args.plot_distribution: fig.savefig(args.plot_distribution) -f_out = pycbc.io.HFile(args.output_file, 'w') +f_out = pycbc.io.HFile(args.output_file, "w") for k in sngl_trigs.data: - f_out.create_dataset('foreground/' + k, - data=sngl_trigs.data[k], - compression='gzip', - compression_opts=9, - shuffle=True) - f_out.create_dataset('background/' + k, - data=bg_trigs.data[k], - compression='gzip', - compression_opts=9, - shuffle=True) - -f_out.create_dataset('foreground/p_astro', - data=p_astro, - compression='gzip', - compression_opts=9, - shuffle=True) + f_out.create_dataset( + "foreground/" + k, + data=sngl_trigs.data[k], + compression="gzip", + compression_opts=9, + shuffle=True, + ) + f_out.create_dataset( + "background/" + k, + data=bg_trigs.data[k], + compression="gzip", + compression_opts=9, + shuffle=True, + ) + +f_out.create_dataset( + "foreground/p_astro", + data=p_astro, + compression="gzip", + compression_opts=9, + shuffle=True, +) # Store segments -f_out['segments/%s/start' % sngl_ifo], f_out['segments/%s/end' % sngl_ifo] = \ +f_out["segments/%s/start" % sngl_ifo], f_out["segments/%s/end" % sngl_ifo] = ( veto.segments_to_start_end(sngl_detector_segs) -f_out.attrs['foreground_time'] = fg_time -f_out.attrs['background_time'] = bg_time -f_out.attrs['num_of_ifos'] = 1 -f_out.attrs['pivot'] = sngl_ifo -f_out.attrs['fixed'] = sngl_ifo -f_out.attrs['ifos'] = sngl_ifo -f_out.attrs['pastro_method'] = args.pastro_method -logging.info('Done') +) +f_out.attrs["foreground_time"] = fg_time +f_out.attrs["background_time"] = bg_time +f_out.attrs["num_of_ifos"] = 1 +f_out.attrs["pivot"] = sngl_ifo +f_out.attrs["fixed"] = sngl_ifo +f_out.attrs["ifos"] = sngl_ifo +f_out.attrs["pastro_method"] = args.pastro_method +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_sngls_statmap b/bin/all_sky_search/pycbc_sngls_statmap index 7b4f1b13313..a6b0c40228c 100755 --- a/bin/all_sky_search/pycbc_sngls_statmap +++ b/bin/all_sky_search/pycbc_sngls_statmap @@ -6,23 +6,33 @@ with producing the combined foreground and background triggers. """ import argparse -import logging, numpy, copy -from pycbc.events import veto -from pycbc.events import significance -import pycbc.pnutils, pycbc.io +import copy +import logging + +import numpy + import pycbc.conversions as conv +import pycbc.io +import pycbc.pnutils +from pycbc.events import significance, veto -class fw(object): + +class fw: def __init__(self, name): - self.f = pycbc.io.HFile(name, 'w') + self.f = pycbc.io.HFile(name, "w") self.attrs = self.f.attrs def __setitem__(self, name, data): # Make a new item if isn't in the hdf file - if not name in self.f: - self.f.create_dataset(name, data=data, compression="gzip", - compression_opts=9, shuffle=True, - maxshape=data.shape) + if name not in self.f: + self.f.create_dataset( + name, + data=data, + compression="gzip", + compression_opts=9, + shuffle=True, + maxshape=data.shape, + ) # Else reassign values else: self.f[name][:] = data @@ -30,47 +40,72 @@ class fw(object): def __getitem__(self, *args): return self.f.__getitem__(*args) + parser = argparse.ArgumentParser() # General required options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--sngls-files', nargs='+', - help='List of files containing trigger and statistic ' - 'information.') -parser.add_argument('--ifos', nargs=1, - help='List of ifos used in these coincidence files') -parser.add_argument('--cluster-window', type=float, default=10, - help='Length of time window in seconds to cluster coinc ' - 'events [default=10s]') -parser.add_argument('--veto-window', type=float, default=.1, - help='Time around each zerolag trigger to window out ' - '[default=.1s]') +parser.add_argument( + "--sngls-files", + nargs="+", + help="List of files containing trigger and statistic information.", +) +parser.add_argument( + "--ifos", nargs=1, help="List of ifos used in these coincidence files" +) +parser.add_argument( + "--cluster-window", + type=float, + default=10, + help="Length of time window in seconds to cluster coinc events [default=10s]", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) significance.insert_significance_option_group(parser) -parser.add_argument('--hierarchical-removal-ifar-threshold', - type=float, default=0.5, - help="Threshold to hierarchically remove foreground " - "triggers with IFAR (years) above this value " - "[default=0.5yr]") -parser.add_argument('--hierarchical-removal-window', type=float, default=1.0, - help='Time around each trigger to window out for a very ' - 'louder trigger in the hierarchical removal ' - 'procedure [default=1.0s]') -parser.add_argument('--max-hierarchical-removal', type=int, default=0, - help='Maximum number of hierarchical removals to carry ' - 'out. Choose -1 for unlimited hierarchical removal ' - 'until no foreground triggers are louder than the ' - '(inclusive/exclusive) background. Choose 0 to do ' - 'no hierarchical removals, choose 1 to do at most ' - '1 hierarchical removal and so on. If given, must ' - 'also provide --hierarchical-removal-against to ' - 'indicate which background to remove triggers ' - 'against. [default=0]') -parser.add_argument('--hierarchical-removal-against', type=str, - default='none', choices=['none', 'inclusive', 'exclusive'], - help='If doing hierarchical removal, remove foreground ' - 'triggers that are louder than either the "inclusive"' - ' (little-dogs-in) background, or the "exclusive" ' - '(little-dogs-out) background. [default="none"]') -parser.add_argument('--output-file') +parser.add_argument( + "--hierarchical-removal-ifar-threshold", + type=float, + default=0.5, + help="Threshold to hierarchically remove foreground " + "triggers with IFAR (years) above this value " + "[default=0.5yr]", +) +parser.add_argument( + "--hierarchical-removal-window", + type=float, + default=1.0, + help="Time around each trigger to window out for a very " + "louder trigger in the hierarchical removal " + "procedure [default=1.0s]", +) +parser.add_argument( + "--max-hierarchical-removal", + type=int, + default=0, + help="Maximum number of hierarchical removals to carry " + "out. Choose -1 for unlimited hierarchical removal " + "until no foreground triggers are louder than the " + "(inclusive/exclusive) background. Choose 0 to do " + "no hierarchical removals, choose 1 to do at most " + "1 hierarchical removal and so on. If given, must " + "also provide --hierarchical-removal-against to " + "indicate which background to remove triggers " + "against. [default=0]", +) +parser.add_argument( + "--hierarchical-removal-against", + type=str, + default="none", + choices=["none", "inclusive", "exclusive"], + help="If doing hierarchical removal, remove foreground " + 'triggers that are louder than either the "inclusive"' + ' (little-dogs-in) background, or the "exclusive" ' + '(little-dogs-out) background. [default="none"]', +) +parser.add_argument("--output-file") args = parser.parse_args() significance.check_significance_options(args, parser) @@ -78,19 +113,23 @@ significance.check_significance_options(args, parser) # Check that the user chose inclusive or exclusive background to perform # hierarchical removals of foreground triggers against. if args.max_hierarchical_removal == 0: - if args.hierarchical_removal_against != 'none': - parser.error("User Error: 0 maximum hierarchical removals chosen but " - "option for --hierarchical-removal-against was given. " - "These are conflicting options. Use with --help for more " - "information.") -else : - is_bkg_inc = (args.hierarchical_removal_against == 'inclusive') - is_bkg_exc = (args.hierarchical_removal_against == 'exclusive') - if not(is_bkg_inc or is_bkg_exc): - parser.error("--max-hierarchical-removal requires a choice of which " - "background to remove foreground triggers against, " - "inclusive or exclusive. Use with --help for more " - "information.") + if args.hierarchical_removal_against != "none": + parser.error( + "User Error: 0 maximum hierarchical removals chosen but " + "option for --hierarchical-removal-against was given. " + "These are conflicting options. Use with --help for more " + "information." + ) +else: + is_bkg_inc = args.hierarchical_removal_against == "inclusive" + is_bkg_exc = args.hierarchical_removal_against == "exclusive" + if not (is_bkg_inc or is_bkg_exc): + parser.error( + "--max-hierarchical-removal requires a choice of which " + "background to remove foreground triggers against, " + "inclusive or exclusive. Use with --help for more " + "information." + ) pycbc.init_logging(args.verbose) @@ -98,34 +137,33 @@ pycbc.init_logging(args.verbose) logging.info("Loading triggers") ifo = args.ifos[0] logging.info("IFO input: %s" % ifo) -all_trigs = pycbc.io.MultiifoStatmapData(files=args.sngls_files, - ifos=[ifo]) -assert ifo + '/time' in all_trigs.data +all_trigs = pycbc.io.MultiifoStatmapData(files=args.sngls_files, ifos=[ifo]) +assert ifo + "/time" in all_trigs.data logging.info("We have %s triggers" % len(all_trigs.stat)) logging.info("Clustering triggers") all_trigs = all_trigs.cluster(args.cluster_window) logging.info("%s triggers remain" % len(all_trigs.stat)) -fg_time = float(all_trigs.attrs['foreground_time']) +fg_time = float(all_trigs.attrs["foreground_time"]) logging.info("Dumping foreground triggers") f = fw(args.output_file) -f.attrs['num_of_ifos'] = 1 -f.attrs['ifos'] = ifo +f.attrs["num_of_ifos"] = 1 +f.attrs["ifos"] = ifo -f.attrs['timeslide_interval'] = all_trigs.attrs['timeslide_interval'] +f.attrs["timeslide_interval"] = all_trigs.attrs["timeslide_interval"] # Copy over the segment info for key in all_trigs.seg.keys(): - f['segments/%s/start' % key] = all_trigs.seg[key]['start'][:] - f['segments/%s/end' % key] = all_trigs.seg[key]['end'][:] + f["segments/%s/start" % key] = all_trigs.seg[key]["start"][:] + f["segments/%s/end" % key] = all_trigs.seg[key]["end"][:] -f['segments/foreground_veto/start'] = numpy.array([0]) -f['segments/foreground_veto/end'] = numpy.array([0]) +f["segments/foreground_veto/start"] = numpy.array([0]) +f["segments/foreground_veto/end"] = numpy.array([0]) for k in all_trigs.data: - f['foreground/' + k] = all_trigs.data[k] - f['background/' + k] = all_trigs.data[k] + f["foreground/" + k] = all_trigs.data[k] + f["background/" + k] = all_trigs.data[k] logging.info("Estimating FAN from background statistic values") # Ranking statistic of foreground and background @@ -137,11 +175,7 @@ significance_dict = significance.digest_significance_options([ifo], args) # Cumulative array of inclusive background triggers and the number of # inclusive background triggers louder than each foreground trigger bg_far, fg_far, sig_info = significance.get_far( - back_stat, - fore_stat, - bkg_dec_facs, - fg_time, - **significance_dict[ifo] + back_stat, fore_stat, bkg_dec_facs, fg_time, **significance_dict[ifo] ) fg_far = significance.apply_far_limit( @@ -155,13 +189,13 @@ bg_far = significance.apply_far_limit( combo=ifo, ) -bg_ifar = 1. / bg_far -fg_ifar = 1. / fg_far +bg_ifar = 1.0 / bg_far +fg_ifar = 1.0 / fg_far -f['background/ifar'] = conv.sec_to_year(bg_ifar) +f["background/ifar"] = conv.sec_to_year(bg_ifar) -f.attrs['background_time'] = fg_time -f.attrs['foreground_time'] = fg_time +f.attrs["background_time"] = fg_time +f.attrs["foreground_time"] = fg_time # Find foreground triggers with IFAR > the set limit and remove from # the exclusive background @@ -181,8 +215,7 @@ back_exc_locs = numpy.arange(len(all_trigs.stat)) to_keep = bg_ifar_exc <= fg_time_exc n_removed = bg_ifar_exc.size - sum(to_keep) -logging.info("Removing %s event(s) from exclusive background", - n_removed) +logging.info("Removing %s event(s) from exclusive background", n_removed) back_stat_exc = back_stat_exc[to_keep] bkg_exc_dec_facs = bkg_exc_dec_facs[to_keep] @@ -191,51 +224,42 @@ back_exc_locs = back_exc_locs[to_keep] # Cumulative array of exclusive background triggers and the number # of exclusive background triggers louder than each foreground trigger bg_far_exc, fg_far_exc, exc_sig_info = significance.get_far( - back_stat_exc, - fore_stat, - bkg_exc_dec_facs, - fg_time_exc, - **significance_dict[ifo]) - -fg_far_exc = significance.apply_far_limit( - fg_far_exc, - significance_dict, - combo=ifo) + back_stat_exc, fore_stat, bkg_exc_dec_facs, fg_time_exc, **significance_dict[ifo] +) -bg_far_exc = significance.apply_far_limit( - bg_far_exc, - significance_dict, - combo=ifo) +fg_far_exc = significance.apply_far_limit(fg_far_exc, significance_dict, combo=ifo) -bg_ifar_exc = 1. / bg_far_exc -fg_ifar_exc = 1. / fg_far_exc +bg_far_exc = significance.apply_far_limit(bg_far_exc, significance_dict, combo=ifo) + +bg_ifar_exc = 1.0 / bg_far_exc +fg_ifar_exc = 1.0 / fg_far_exc # Remove a small amount of time from the exclusive fore/background # time to account for this removal fg_time_exc -= n_removed * args.veto_window for k in all_trigs.data: - f['background_exc/' + k] = all_trigs.data[k][back_exc_locs] + f["background_exc/" + k] = all_trigs.data[k][back_exc_locs] -f['background_exc/ifar'] = conv.sec_to_year(bg_ifar_exc) -f.attrs['background_time_exc'] = fg_time_exc -f.attrs['foreground_time_exc'] = fg_time_exc +f["background_exc/ifar"] = conv.sec_to_year(bg_ifar_exc) +f.attrs["background_time_exc"] = fg_time_exc +f.attrs["foreground_time_exc"] = fg_time_exc logging.info("calculating foreground ifar/fap values") -fap = 1 - numpy.exp(- fg_time / fg_ifar) -f['foreground/ifar'] = conv.sec_to_year(fg_ifar) -f['foreground/fap'] = fap -fap_exc = 1 - numpy.exp(- fg_time_exc / fg_ifar_exc) -f['foreground/ifar_exc'] = conv.sec_to_year(fg_ifar_exc) -f['foreground/fap_exc'] = fap_exc +fap = 1 - numpy.exp(-fg_time / fg_ifar) +f["foreground/ifar"] = conv.sec_to_year(fg_ifar) +f["foreground/fap"] = fap +fap_exc = 1 - numpy.exp(-fg_time_exc / fg_ifar_exc) +f["foreground/ifar_exc"] = conv.sec_to_year(fg_ifar_exc) +f["foreground/fap_exc"] = fap_exc for key, value in sig_info.items(): - f['foreground'].attrs[key] = value + f["foreground"].attrs[key] = value for key, value in exc_sig_info.items(): - f['foreground'].attrs[f'{key}_exc'] = value + f["foreground"].attrs[f"{key}_exc"] = value -if 'name' in all_trigs.attrs: - f.attrs['name'] = all_trigs.attrs['name'] +if "name" in all_trigs.attrs: + f.attrs["name"] = all_trigs.attrs["name"] # Incorporate hierarchical removal for any other loud triggers logging.info("Beginning hierarchical removal of foreground triggers") @@ -258,9 +282,9 @@ if args.max_hierarchical_removal != 0: if is_bkg_inc: ifar_louder = fg_ifar # Otherwise user wants to remove against exclusive background - else : + else: ifar_louder = fg_ifar_exc -else : +else: # It doesn't matter if you choose inclusive or exclusive, # the while loop below will break if none are louder than # ifar_louder, or at the comparison @@ -273,32 +297,31 @@ else : # above the set threshold, or a set maximum. # Convert threshold into seconds -hier_ifar_thresh_s = args.hierarchical_removal_ifar_threshold / \ - conv.sec_to_year(1) +hier_ifar_thresh_s = args.hierarchical_removal_ifar_threshold / conv.sec_to_year(1) while numpy.any(ifar_louder > hier_ifar_thresh_s): # If the user wants to stop doing hierarchical removals after a set # number of iterations then break when that happens. - if (h_iterations == args.max_hierarchical_removal): + if h_iterations == args.max_hierarchical_removal: break # Write foreground trigger info before hierarchical removals for # downstream codes. if h_iterations == 0: - f['background_h%s/stat' % h_iterations] = back_stat - f['background_h%s/ifar' % h_iterations] = conv.sec_to_year(bg_ifar) + f["background_h%s/stat" % h_iterations] = back_stat + f["background_h%s/ifar" % h_iterations] = conv.sec_to_year(bg_ifar) for k in all_trigs.data: - f['background_h%s/' % h_iterations + k] = all_trigs.data[k] - f['foreground_h%s/stat' % h_iterations] = fore_stat - f['foreground_h%s/ifar' % h_iterations] = conv.sec_to_year(fg_ifar) - f['foreground_h%s/ifar_exc' % h_iterations] = conv.sec_to_year(fg_ifar_exc) - f['foreground_h%s/fap' % h_iterations] = fap + f["background_h%s/" % h_iterations + k] = all_trigs.data[k] + f["foreground_h%s/stat" % h_iterations] = fore_stat + f["foreground_h%s/ifar" % h_iterations] = conv.sec_to_year(fg_ifar) + f["foreground_h%s/ifar_exc" % h_iterations] = conv.sec_to_year(fg_ifar_exc) + f["foreground_h%s/fap" % h_iterations] = fap for key, value in sig_info.items(): - f['foreground_h%s' % h_iterations].attrs[key] = value + f["foreground_h%s" % h_iterations].attrs[key] = value for key, value in exc_sig_info.items(): - f['foreground_h%s' % h_iterations].attrs[key + "_exc"] = value + f["foreground_h%s" % h_iterations].attrs[key + "_exc"] = value for k in all_trigs.data: - f['foreground_h%s/' % h_iterations + k] = all_trigs.data[k] + f["foreground_h%s/" % h_iterations + k] = all_trigs.data[k] # Add the iteration number of hierarchical removals done. h_iterations += 1 @@ -315,24 +338,30 @@ while numpy.any(ifar_louder > hier_ifar_thresh_s): # Store any foreground trigger's information that we want to # hierarchically remove. - f['foreground/ifar'][orig_fore_idx] = conv.sec_to_year(fg_ifar[max_stat_idx]) - f['foreground/fap'][orig_fore_idx] = fap[max_stat_idx] + f["foreground/ifar"][orig_fore_idx] = conv.sec_to_year(fg_ifar[max_stat_idx]) + f["foreground/fap"][orig_fore_idx] = fap[max_stat_idx] - logging.info("Removing foreground trigger that is louder than the inclusive background.") + logging.info( + "Removing foreground trigger that is louder than the inclusive background." + ) # Remove the foreground trigger and all of the background triggers that # are associated with it. - ave_rm_time = all_trigs.data['%s/time' % ifo][rm_trig_idx] + ave_rm_time = all_trigs.data["%s/time" % ifo][rm_trig_idx] ind_to_rm = {} - ind_to_rm[ifo] = veto.indices_within_times(all_trigs.data['%s/time' % ifo], - [ave_rm_time - args.hierarchical_removal_window], - [ave_rm_time + args.hierarchical_removal_window]) + ind_to_rm[ifo] = veto.indices_within_times( + all_trigs.data["%s/time" % ifo], + [ave_rm_time - args.hierarchical_removal_window], + [ave_rm_time + args.hierarchical_removal_window], + ) indices_to_rm = [] indices_to_rm = numpy.concatenate([indices_to_rm, ind_to_rm[ifo]]) all_trigs = all_trigs.remove(indices_to_rm.astype(int)) - logging.info("We have %s triggers after hierarchical removal." % len(all_trigs.stat)) + logging.info( + "We have %s triggers after hierarchical removal." % len(all_trigs.stat) + ) # Step 4: Re-cluster the triggers and calculate the inclusive ifar/fap logging.info("Clustering coinc triggers (inclusive of zerolag)") @@ -344,18 +373,15 @@ while numpy.any(ifar_louder > hier_ifar_thresh_s): logging.info("Dumping foreground triggers") logging.info("Dumping background triggers (inclusive of zerolag)") for k in all_trigs.data: - f['background_h%s/' % h_iterations + k] = all_trigs.data[k] + f["background_h%s/" % h_iterations + k] = all_trigs.data[k] logging.info("Calculating FAN from background statistic values") back_stat = fore_stat = all_trigs.stat bkg_dec_facs = all_trigs.decimation_factor bg_far, fg_far, sig_info = significance.get_far( - back_stat, - fore_stat, - bkg_dec_facs, - fg_time, - **significance_dict[ifo]) + back_stat, fore_stat, bkg_dec_facs, fg_time, **significance_dict[ifo] + ) fg_far = significance.apply_far_limit( fg_far, @@ -369,8 +395,8 @@ while numpy.any(ifar_louder > hier_ifar_thresh_s): combo=ifo, ) - bg_ifar = 1. / bg_far - fg_ifar = 1. / fg_far + bg_ifar = 1.0 / bg_far + fg_ifar = 1.0 / fg_far # Update the ifar_louder criteria depending on whether foreground # triggers are being removed via inclusive or exclusive background. @@ -386,7 +412,8 @@ while numpy.any(ifar_louder > hier_ifar_thresh_s): fore_stat, bkg_exc_dec_facs, fg_time_exc, - **significance_dict[ifo]) + **significance_dict[ifo], + ) fg_far_exc = significance.apply_far_limit( fg_far_exc, @@ -394,65 +421,64 @@ while numpy.any(ifar_louder > hier_ifar_thresh_s): combo=ifo, ) - fg_ifar_exc = 1. / fg_far_exc + fg_ifar_exc = 1.0 / fg_far_exc ifar_louder = fg_ifar_exc # louder_foreground has been updated and the code can continue. logging.info("Calculating ifar/fap values") - f['background_h%s/ifar' % h_iterations] = conv.sec_to_year(bg_ifar) - f.attrs['background_time_h%s' % h_iterations] = fg_time - f.attrs['foreground_time_h%s' % h_iterations] = fg_time + f["background_h%s/ifar" % h_iterations] = conv.sec_to_year(bg_ifar) + f.attrs["background_time_h%s" % h_iterations] = fg_time + f.attrs["foreground_time_h%s" % h_iterations] = fg_time if len(all_trigs) > 0: # Write ranking statistic to file just for downstream plotting code - f['foreground_h%s/stat' % h_iterations] = fore_stat + f["foreground_h%s/stat" % h_iterations] = fore_stat for key, value in sig_info.items(): - f['foreground_h%s' % h_iterations].attrs[key] = value + f["foreground_h%s" % h_iterations].attrs[key] = value for key, value in exc_sig_info.items(): - f['foreground_h%s' % h_iterations].attrs[key + "_exc"] = value - fap = 1 - numpy.exp(- fg_time / fg_ifar) - f['foreground_h%s/ifar' % h_iterations] = conv.sec_to_year(fg_ifar) - f['foreground_h%s/fap' % h_iterations] = fap + f["foreground_h%s" % h_iterations].attrs[key + "_exc"] = value + fap = 1 - numpy.exp(-fg_time / fg_ifar) + f["foreground_h%s/ifar" % h_iterations] = conv.sec_to_year(fg_ifar) + f["foreground_h%s/fap" % h_iterations] = fap - fap_exc = 1 - numpy.exp(- fg_time / fg_ifar_exc) - f['foreground_h%s/ifar' % h_iterations] = conv.sec_to_year(fg_ifar_exc) - f['foreground_h%s/fap' % h_iterations] = fap_exc + fap_exc = 1 - numpy.exp(-fg_time / fg_ifar_exc) + f["foreground_h%s/ifar" % h_iterations] = conv.sec_to_year(fg_ifar_exc) + f["foreground_h%s/fap" % h_iterations] = fap_exc # Update ifar and fap for other foreground triggers for i in range(len(fg_ifar)): orig_fore_idx = numpy.where(orig_fore_stat == fore_stat[i])[0][0] - f['foreground/ifar'][orig_fore_idx] = conv.sec_to_year(fg_ifar[i]) - f['foreground/fap'][orig_fore_idx] = fap[i] + f["foreground/ifar"][orig_fore_idx] = conv.sec_to_year(fg_ifar[i]) + f["foreground/fap"][orig_fore_idx] = fap[i] # Save trigger ids for foreground triggers for downstream plotting code. # These don't change with the iterations but should be written at every # level. - f['foreground_h%s/template_id' % h_iterations] = all_trigs.data['template_id'] - trig_id = all_trigs.data['%s/trigger_id' % ifo] - trig_time = all_trigs.data['%s/time' % ifo] - f['foreground_h%s/%s/time' % (h_iterations,ifo)] = trig_time - f['foreground_h%s/%s/trigger_id' % (h_iterations,ifo)] = trig_id - else : - f['foreground_h%s/stat' % h_iterations] = numpy.array([]) - f['foreground_h%s/ifar' % h_iterations] = numpy.array([]) - f['foreground_h%s/fap' % h_iterations] = numpy.array([]) - f['foreground_h%s/template_id' % h_iterations] = numpy.array([]) - f['foreground_h%s/%s/time' % (h_iterations,ifo)] = numpy.array([]) - f['foreground_h%s/%s/trigger_id' % (h_iterations,ifo)] = numpy.array([]) + f["foreground_h%s/template_id" % h_iterations] = all_trigs.data["template_id"] + trig_id = all_trigs.data["%s/trigger_id" % ifo] + trig_time = all_trigs.data["%s/time" % ifo] + f["foreground_h%s/%s/time" % (h_iterations, ifo)] = trig_time + f["foreground_h%s/%s/trigger_id" % (h_iterations, ifo)] = trig_id + else: + f["foreground_h%s/stat" % h_iterations] = numpy.array([]) + f["foreground_h%s/ifar" % h_iterations] = numpy.array([]) + f["foreground_h%s/fap" % h_iterations] = numpy.array([]) + f["foreground_h%s/template_id" % h_iterations] = numpy.array([]) + f["foreground_h%s/%s/time" % (h_iterations, ifo)] = numpy.array([]) + f["foreground_h%s/%s/trigger_id" % (h_iterations, ifo)] = numpy.array([]) # Write to file how many hierarchical removals were implemented. -f.attrs['hierarchical_removal_iterations'] = h_iterations +f.attrs["hierarchical_removal_iterations"] = h_iterations # Write whether hierarchical removals were removed against the # inclusive background or the exclusive background. Have to use # numpy.bytes_ datatype. if h_iterations != 0: hrm_method = args.hierarchical_removal_against - f.attrs['hierarchical_removal_method'] = numpy.bytes_(hrm_method) + f.attrs["hierarchical_removal_method"] = numpy.bytes_(hrm_method) logging.info("Done") - diff --git a/bin/all_sky_search/pycbc_sngls_statmap_inj b/bin/all_sky_search/pycbc_sngls_statmap_inj index 47e235d7f6c..2bac65c3efa 100644 --- a/bin/all_sky_search/pycbc_sngls_statmap_inj +++ b/bin/all_sky_search/pycbc_sngls_statmap_inj @@ -6,22 +6,32 @@ with producing the combined foreground and background triggers. """ import argparse -import logging, numpy -from pycbc.events import significance -import pycbc.pnutils, pycbc.io +import logging + +import numpy + import pycbc.conversions as conv +import pycbc.io +import pycbc.pnutils +from pycbc.events import significance + -class fw(object): +class fw: def __init__(self, name): - self.f = pycbc.io.HFile(name, 'w') + self.f = pycbc.io.HFile(name, "w") self.attrs = self.f.attrs def __setitem__(self, name, data): # Make a new item if isn't in the hdf file - if not name in self.f: - self.f.create_dataset(name, data=data, compression="gzip", - compression_opts=9, shuffle=True, - maxshape=data.shape) + if name not in self.f: + self.f.create_dataset( + name, + data=data, + compression="gzip", + compression_opts=9, + shuffle=True, + maxshape=data.shape, + ) # Else reassign values else: self.f[name][:] = data @@ -29,25 +39,37 @@ class fw(object): def __getitem__(self, *args): return self.f.__getitem__(*args) + parser = argparse.ArgumentParser() # General required options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--sngls-files', nargs='+', - help='List of files containign trigger and statistic ' - 'information.') -parser.add_argument('--full-data-background', required=True, - help='background file from full data for use in analyzing ' - 'injection coincs') -parser.add_argument('--ifos', nargs=1, - help='List of ifos used in these coincidence files') -parser.add_argument('--cluster-window', type=float, default=10, - help='Length of time window in seconds to cluster coinc ' - 'events [default=10s]') -parser.add_argument('--veto-window', type=float, default=.1, - help='Time around each zerolag trigger to window out ' - '[default=.1s]') +parser.add_argument( + "--sngls-files", + nargs="+", + help="List of files containign trigger and statistic information.", +) +parser.add_argument( + "--full-data-background", + required=True, + help="background file from full data for use in analyzing injection coincs", +) +parser.add_argument( + "--ifos", nargs=1, help="List of ifos used in these coincidence files" +) +parser.add_argument( + "--cluster-window", + type=float, + default=10, + help="Length of time window in seconds to cluster coinc events [default=10s]", +) +parser.add_argument( + "--veto-window", + type=float, + default=0.1, + help="Time around each zerolag trigger to window out [default=.1s]", +) significance.insert_significance_option_group(parser) -parser.add_argument('--output-file') +parser.add_argument("--output-file") args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -58,43 +80,43 @@ logging.info("Loading triggers") logging.info("IFO input: %s" % args.ifos[0]) all_trigs = pycbc.io.MultiifoStatmapData(files=args.sngls_files, ifos=args.ifos) ifo = args.ifos[0] -assert ifo + '/time' in all_trigs.data +assert ifo + "/time" in all_trigs.data logging.info("We have %s triggers" % len(all_trigs.stat)) logging.info("Clustering triggers") all_trigs = all_trigs.cluster(args.cluster_window) -logging.info('getting background statistics') +logging.info("getting background statistics") -fb = pycbc.io.HFile(args.full_data_background,'r') -back_stat = fb['background/stat'][:] -back_stat_exc = fb['background_exc/stat'][:] +fb = pycbc.io.HFile(args.full_data_background, "r") +back_stat = fb["background/stat"][:] +back_stat_exc = fb["background_exc/stat"][:] -bkg_dec_facs = fb['background/decimation_factor'][:] -bkg_exc_dec_facs = fb['background_exc/decimation_factor'][:] +bkg_dec_facs = fb["background/decimation_factor"][:] +bkg_exc_dec_facs = fb["background_exc/decimation_factor"][:] # For now, all triggers are both in the foreground and background fore_locs = numpy.flatnonzero(all_trigs.timeslide_id == 0) -fg_time = fb.attrs['background_time'] -fg_time_exc = fb.attrs['background_time_exc'] +fg_time = fb.attrs["background_time"] +fg_time_exc = fb.attrs["background_time_exc"] logging.info("Dumping foreground triggers") f = fw(args.output_file) -f.attrs['num_of_ifos'] = 1 -f.attrs['ifos'] = ifo +f.attrs["num_of_ifos"] = 1 +f.attrs["ifos"] = ifo -f.attrs['timeslide_interval'] = all_trigs.attrs['timeslide_interval'] +f.attrs["timeslide_interval"] = all_trigs.attrs["timeslide_interval"] # Copy over the segment info for key in all_trigs.seg.keys(): - f['segments/%s/start' % key] = all_trigs.seg[key]['start'][:] - f['segments/%s/end' % key] = all_trigs.seg[key]['end'][:] + f["segments/%s/start" % key] = all_trigs.seg[key]["start"][:] + f["segments/%s/end" % key] = all_trigs.seg[key]["end"][:] -f['segments/foreground_veto/start'] = numpy.array([0]) -f['segments/foreground_veto/end'] = numpy.array([0]) +f["segments/foreground_veto/start"] = numpy.array([0]) +f["segments/foreground_veto/end"] = numpy.array([0]) for k in all_trigs.data: - f['foreground/' + k] = all_trigs.data[k] + f["foreground/" + k] = all_trigs.data[k] logging.info("Estimating FAN from background statistic values") @@ -106,37 +128,27 @@ significance_dict = significance.digest_significance_options([ifo], args) # Cumulative array of exclusive background triggers and the number # of exclusive background triggers louder than each foreground trigger bg_far_exc, fg_far_exc, sig_info = significance.get_far( - back_stat_exc, - fore_stat, - bkg_exc_dec_facs, - fg_time_exc, - **significance_dict[ifo]) - -fg_far_exc = significance.apply_far_limit( - fg_far_exc, - significance_dict, - combo=ifo) -bg_far_exc = significance.apply_far_limit( - bg_far_exc, - significance_dict, - combo=ifo) - -fg_ifar_exc = 1. / fg_far_exc -bg_ifar_exc = 1. / bg_far_exc - -f['background_exc/ifar'] = conv.sec_to_year(bg_ifar_exc) -f.attrs['background_time_exc'] = fg_time_exc -f.attrs['foreground_time_exc'] = fg_time_exc - -fap_exc = 1 - numpy.exp(- fg_time_exc / fg_ifar_exc) -f['foreground/ifar_exc'] = conv.sec_to_year(fg_ifar_exc) -f['foreground/fap_exc'] = fap_exc + back_stat_exc, fore_stat, bkg_exc_dec_facs, fg_time_exc, **significance_dict[ifo] +) -for key, value in sig_info.items(): - f['foreground'].attrs[key + '_exc'] = value +fg_far_exc = significance.apply_far_limit(fg_far_exc, significance_dict, combo=ifo) +bg_far_exc = significance.apply_far_limit(bg_far_exc, significance_dict, combo=ifo) + +fg_ifar_exc = 1.0 / fg_far_exc +bg_ifar_exc = 1.0 / bg_far_exc + +f["background_exc/ifar"] = conv.sec_to_year(bg_ifar_exc) +f.attrs["background_time_exc"] = fg_time_exc +f.attrs["foreground_time_exc"] = fg_time_exc -if 'name' in all_trigs.attrs: - f.attrs['name'] = all_trigs.attrs['name'] +fap_exc = 1 - numpy.exp(-fg_time_exc / fg_ifar_exc) +f["foreground/ifar_exc"] = conv.sec_to_year(fg_ifar_exc) +f["foreground/fap_exc"] = fap_exc + +for key, value in sig_info.items(): + f["foreground"].attrs[key + "_exc"] = value -logging.info('Done!') +if "name" in all_trigs.attrs: + f.attrs["name"] = all_trigs.attrs["name"] +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_strip_injections b/bin/all_sky_search/pycbc_strip_injections index 61d50222225..c765eec3ca7 100644 --- a/bin/all_sky_search/pycbc_strip_injections +++ b/bin/all_sky_search/pycbc_strip_injections @@ -1,55 +1,76 @@ #!/bin/env python -import numpy, argparse, pycbc.pnutils, logging +import argparse +import logging + +import numpy +from igwn_ligolw import ligolw +from igwn_ligolw import utils as ligolw_utils + +import pycbc.pnutils from pycbc.events import veto from pycbc.io.ligolw import LIGOLWContentHandler -from igwn_ligolw import ligolw, utils as ligolw_utils + +effd = {"H1": "eff_dist_h", "L1": "eff_dist_l", "V1": "eff_dist_v"} -effd = {"H1":"eff_dist_h", "L1":"eff_dist_l", "V1":"eff_dist_v"} def remove(l, i): to_remove = [l[t] for t in i] for r in to_remove: l.remove(r) + parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--injection-file') -parser.add_argument('--veto-file', - help="File containing segments used to veto injections") -parser.add_argument('--segment-name', - help="Name of segmentlist within the veto file to veto injections") -parser.add_argument('--ifos', nargs='+') -parser.add_argument('--max-effective-chirp-distance', type=float) -parser.add_argument('--output-file') +parser.add_argument("--injection-file") +parser.add_argument( + "--veto-file", help="File containing segments used to veto injections" +) +parser.add_argument( + "--segment-name", help="Name of segmentlist within the veto file to veto injections" +) +parser.add_argument("--ifos", nargs="+") +parser.add_argument("--max-effective-chirp-distance", type=float) +parser.add_argument("--output-file") args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('File: %s' % args.injection_file) -indoc = ligolw_utils.load_filename(args.injection_file, False, contenthandler=LIGOLWContentHandler) -sim_table = ligolw.Table.get_table(indoc, 'sim_inspiral') +logging.info("File: %s" % args.injection_file) +indoc = ligolw_utils.load_filename( + args.injection_file, False, contenthandler=LIGOLWContentHandler +) +sim_table = ligolw.Table.get_table(indoc, "sim_inspiral") -logging.info('%s Injections in file' % len(sim_table)) +logging.info("%s Injections in file" % len(sim_table)) if args.veto_file: - logging.info('Removing injections outside coincident time') - for ifo in args.ifos: - inj_time = numpy.array(sim_table.get_column('geocent_end_time') + 1e-9 * sim_table.get_column('geocent_end_time_ns'), dtype=numpy.float64) - idx, segs = veto.indices_outside_segments(inj_time, [args.veto_file], ifo, args.segment_name) + logging.info("Removing injections outside coincident time") + for ifo in args.ifos: + inj_time = numpy.array( + sim_table.get_column("geocent_end_time") + + 1e-9 * sim_table.get_column("geocent_end_time_ns"), + dtype=numpy.float64, + ) + idx, segs = veto.indices_outside_segments( + inj_time, [args.veto_file], ifo, args.segment_name + ) remove(sim_table, idx) - logging.info('We now have %s injections' % len(sim_table)) + logging.info("We now have %s injections" % len(sim_table)) if args.max_effective_chirp_distance: - logging.info('Removing injections that are quiet in one detector') + logging.info("Removing injections that are quiet in one detector") for ifo in args.ifos: eff_distance = numpy.array(sim_table.get_column(effd[ifo]), dtype=numpy.float32) - m1, m2 = numpy.array(sim_table.get_column('mass1')), numpy.array(sim_table.get_column('mass2')) + m1, m2 = ( + numpy.array(sim_table.get_column("mass1")), + numpy.array(sim_table.get_column("mass2")), + ) mchirp, eta = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) chirp_distance = pycbc.pnutils.chirp_distance(eff_distance, mchirp) idx = numpy.where(chirp_distance > args.max_effective_chirp_distance)[0] remove(sim_table, idx) - logging.info('We now have %s injections' % len(sim_table)) + logging.info("We now have %s injections" % len(sim_table)) outdoc = ligolw.Document() outdoc.appendChild(ligolw.LIGO_LW()).appendChild(sim_table) ligolw_utils.write_filename(outdoc, args.output_file) -logging.info('Done') +logging.info("Done") diff --git a/bin/all_sky_search/pycbc_template_kde_calc b/bin/all_sky_search/pycbc_template_kde_calc index fd1628d4a98..318d7ddbd85 100644 --- a/bin/all_sky_search/pycbc_template_kde_calc +++ b/bin/all_sky_search/pycbc_template_kde_calc @@ -12,72 +12,114 @@ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General # Public License for more details. -import numpy, operator, argparse, logging -from pycbc import init_logging, add_common_pycbc_options -from pycbc import libutils +import argparse +import logging +import operator + +import numpy + +from pycbc import add_common_pycbc_options, init_logging, libutils from pycbc.events import triggers from pycbc.io import HFile -akde = libutils.import_optional('awkde') -kf = libutils.import_optional('sklearn.model_selection') + +akde = libutils.import_optional("awkde") +kf = libutils.import_optional("sklearn.model_selection") parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--signal-file', help='File with parameters of GW signals ' - 'for KDE calculation') -parser.add_argument('--template-file', required=True, help='Hdf5 file with ' - 'template masses and spins') -parser.add_argument('--injection-file', help='Hdf5 file with masses and spins') -parser.add_argument('--min-mass', type=float, default=None, - help='Used only on signal masses: remove all' - 'signal events with mass2 < min_mass') -parser.add_argument('--min-snr', type=float, default=None, - help='Used only in injections case: remove all' - 'injection events < min_snr') -parser.add_argument('--nfold-signal', type=int, - help='Number of k-folds for signal KDE cross validation') -parser.add_argument('--nfold-template', type=int, - help='Number of k-folds for template KDE cross validation') -parser.add_argument('--nfold-injection', type=int, - help='Number of k-folds for injection KDE cross validation') -parser.add_argument('--fit-param', nargs='+', required=True, - help='Parameters over which KDE is calculated') -parser.add_argument('--log-param', nargs='+', choices=['True', 'False'], - required=True) -parser.add_argument('--output-file', required=True, help='Name of .hdf output') -parser.add_argument('--make-signal-kde', action='store_true') -parser.add_argument('--make-template-kde', action='store_true') -parser.add_argument('--make-injection-kde', action='store_true') -parser.add_argument('--fom-plot', help='Make a FOM plot for cross-validation' - ' and save it as this file') -parser.add_argument('--alpha-grid', type=float, nargs="+", - help='Grid of choices of sensitivity parameter alpha for' - ' local bandwidth') -parser.add_argument('--bw-grid', type=float, nargs='+', - help='Grid of choices of global bandwidth') -parser.add_argument('--extra-cpt-fraction', type=float, - help='Fraction of the extra component in the signal density') -parser.add_argument('--temp-volume', type=float, - help='Volume covered by the template bank') -parser.add_argument('--seed', type=int, - help='Random number generator seed') -parser.add_argument('--mchirp-downsample-power', type=float, - help='Exponent value for the power law distribution') -parser.add_argument('--min-ratio', type=float, - help='Minimum ratio for template_kde relative to the maximum') +parser.add_argument( + "--signal-file", help="File with parameters of GW signals for KDE calculation" +) +parser.add_argument( + "--template-file", required=True, help="Hdf5 file with template masses and spins" +) +parser.add_argument("--injection-file", help="Hdf5 file with masses and spins") +parser.add_argument( + "--min-mass", + type=float, + default=None, + help="Used only on signal masses: remove allsignal events with mass2 < min_mass", +) +parser.add_argument( + "--min-snr", + type=float, + default=None, + help="Used only in injections case: remove allinjection events < min_snr", +) +parser.add_argument( + "--nfold-signal", type=int, help="Number of k-folds for signal KDE cross validation" +) +parser.add_argument( + "--nfold-template", + type=int, + help="Number of k-folds for template KDE cross validation", +) +parser.add_argument( + "--nfold-injection", + type=int, + help="Number of k-folds for injection KDE cross validation", +) +parser.add_argument( + "--fit-param", + nargs="+", + required=True, + help="Parameters over which KDE is calculated", +) +parser.add_argument("--log-param", nargs="+", choices=["True", "False"], required=True) +parser.add_argument("--output-file", required=True, help="Name of .hdf output") +parser.add_argument("--make-signal-kde", action="store_true") +parser.add_argument("--make-template-kde", action="store_true") +parser.add_argument("--make-injection-kde", action="store_true") +parser.add_argument( + "--fom-plot", help="Make a FOM plot for cross-validation and save it as this file" +) +parser.add_argument( + "--alpha-grid", + type=float, + nargs="+", + help="Grid of choices of sensitivity parameter alpha for local bandwidth", +) +parser.add_argument( + "--bw-grid", type=float, nargs="+", help="Grid of choices of global bandwidth" +) +parser.add_argument( + "--extra-cpt-fraction", + type=float, + help="Fraction of the extra component in the signal density", +) +parser.add_argument( + "--temp-volume", type=float, help="Volume covered by the template bank" +) +parser.add_argument("--seed", type=int, help="Random number generator seed") +parser.add_argument( + "--mchirp-downsample-power", + type=float, + help="Exponent value for the power law distribution", +) +parser.add_argument( + "--min-ratio", + type=float, + help="Minimum ratio for template_kde relative to the maximum", +) args = parser.parse_args() init_logging(args.verbose) assert len(args.fit_param) == len(args.log_param) if args.make_signal_kde + args.make_template_kde + args.make_injection_kde != 1: - parser.error("Choose exactly one option out of --make-signal-kde, \ - --make-template-kde, or --make-injection-kde") + parser.error( + "Choose exactly one option out of --make-signal-kde, \ + --make-template-kde, or --make-injection-kde" + ) -if (args.extra_cpt_fraction is None and args.temp_volume is not None) or \ - (args.extra_cpt_fraction is not None and args.temp_volume is None): \ - parser.error("Both --extra-cpt-fraction and --temp-volume arguments \ - must be provided or neither one should be provided") +if (args.extra_cpt_fraction is None and args.temp_volume is not None) or ( + args.extra_cpt_fraction is not None and args.temp_volume is None +): + parser.error( + "Both --extra-cpt-fraction and --temp-volume arguments \ + must be provided or neither one should be provided" + ) def kde_awkde(x, x_grid, alp=0.5, gl_bandwidth=None, ret_kde=False): @@ -116,8 +158,7 @@ def optimizedparam(sampleval, bwgrid, alphagrid, nfold=2): FOM = {} for gbw in bwgrid: for alphavals in alphagrid: - FOM[(gbw, alphavals)] = kfcv_awkde(sampleval, gbw, alphavals, - k=nfold) + FOM[(gbw, alphavals)] = kfcv_awkde(sampleval, gbw, alphavals, k=nfold) optval = max(FOM.items(), key=operator.itemgetter(1))[0] optbw, optalpha = optval[0], optval[1] maxFOM = FOM[(optbw, optalpha)] @@ -125,18 +166,24 @@ def optimizedparam(sampleval, bwgrid, alphagrid, nfold=2): # Plotting FOM parameters if args.fom_plot: import matplotlib.pyplot as plt - fig = plt.figure(figsize=(12,8)) + + fig = plt.figure(figsize=(12, 8)) ax = fig.add_subplot(111) for bw in bwgrid: FOMlist = [FOM[(bw, al)] for al in alphagrid] - ax.plot(alphagrid, FOMlist, label='{0:.3f}'.format(bw)) - ax.plot(optalpha, maxFOM, 'ko', linewidth=10, label= - r'$\alpha={0:.3f},bw={1:.3f}$'.format(optalpha, optbw)) - ax.set_xlabel(r'$\alpha$', fontsize=15) - ax.set_ylabel(r'$FOM$', fontsize=15) + ax.plot(alphagrid, FOMlist, label=f"{bw:.3f}") + ax.plot( + optalpha, + maxFOM, + "ko", + linewidth=10, + label=rf"$\alpha={optalpha:.3f},bw={optbw:.3f}$", + ) + ax.set_xlabel(r"$\alpha$", fontsize=15) + ax.set_ylabel(r"$FOM$", fontsize=15) # Guess at a suitable range of FOM values to plot ax.set_ylim(maxFOM - 0.5 * npoints, maxFOM + 0.2 * npoints) - ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.135), ncol=8) + ax.legend(loc="upper center", bbox_to_anchor=(0.5, 1.135), ncol=8) plt.savefig(args.fom_plot) plt.close() @@ -144,89 +191,97 @@ def optimizedparam(sampleval, bwgrid, alphagrid, nfold=2): # Obtaining template parameters -temp_file = HFile(args.template_file, 'r') -mass1 = temp_file['mass1'][:] +temp_file = HFile(args.template_file, "r") +mass1 = temp_file["mass1"][:] tid = numpy.arange(len(mass1)) # Array of template ids mass_spin = triggers.get_mass_spin(temp_file, tid) -f_dest = HFile(args.output_file, 'w') +f_dest = HFile(args.output_file, "w") f_dest.create_dataset("template_id", data=tid) template_pars = [] for param, slog in zip(args.fit_param, args.log_param): pvals = triggers.get_param(param, args, *mass_spin) # Write the KDE param values to output file f_dest.create_dataset(param, data=pvals) - if slog in ['False']: - logging.info('Using param: %s', param) + if slog in ["False"]: + logging.info("Using param: %s", param) template_pars.append(pvals) - elif slog in ['True']: - logging.info('Using log param: %s', param) + elif slog in ["True"]: + logging.info("Using log param: %s", param) template_pars.append(numpy.log(pvals)) else: raise ValueError("invalid log param argument, use 'True', or 'False'") # Copy standard data to output file -f_dest.attrs['fit_param'] = args.fit_param -f_dest.attrs['log_param'] = args.log_param +f_dest.attrs["fit_param"] = args.fit_param +f_dest.attrs["log_param"] = args.log_param with HFile(args.template_file, "r") as f_src: f_src.copy(f_src["./"], f_dest["./"], "input_template_params") -temp_samples = numpy.vstack((template_pars)).T +temp_samples = numpy.vstack(template_pars).T if args.make_template_kde: - # Rejection sampling to reduce computational load if args.mchirp_downsample_power is not None: - logging.info('Downsampling with mchirp power ' - f'{args.mchirp_downsample_power}') - f_dest.attrs['mchirp_downsample_power'] = args.mchirp_downsample_power + logging.info(f"Downsampling with mchirp power {args.mchirp_downsample_power}") + f_dest.attrs["mchirp_downsample_power"] = args.mchirp_downsample_power try: - mchirp_index = args.fit_param.index('mchirp') + mchirp_index = args.fit_param.index("mchirp") except: raise ValueError("mchirp does not exist in args.fit_param") mc_vals = template_pars[mchirp_index] if args.log_param[mchirp_index]: mc_vals = numpy.exp(mc_vals) - power_vals = mc_vals ** args.mchirp_downsample_power + power_vals = mc_vals**args.mchirp_downsample_power probabilities = power_vals / numpy.max(power_vals) if args.seed is not None: - f_dest.attrs['seed'] = args.seed + f_dest.attrs["seed"] = args.seed numpy.random.seed(args.seed) rand_nums = numpy.random.uniform(0, 1, len(mass1)) ind = rand_nums < probabilities - logging.info(f'{ind.sum()} templates after downsampling') + logging.info(f"{ind.sum()} templates after downsampling") kde_train_samples = temp_samples[ind] f_dest.create_dataset("kde_train_samples", data=kde_train_samples) - logging.info('Starting optimization of template KDE parameters') - optbw, optalpha = optimizedparam(kde_train_samples, alphagrid=args.alpha_grid, - bwgrid=args.bw_grid, nfold=args.nfold_template) - logging.info('Bandwidth %.4f, alpha %.2f' % (optbw, optalpha)) - logging.info('Evaluating template KDE') - template_kde = kde_awkde(kde_train_samples, temp_samples, alp=optalpha, - gl_bandwidth=optbw) + logging.info("Starting optimization of template KDE parameters") + optbw, optalpha = optimizedparam( + kde_train_samples, + alphagrid=args.alpha_grid, + bwgrid=args.bw_grid, + nfold=args.nfold_template, + ) + logging.info("Bandwidth %.4f, alpha %.2f" % (optbw, optalpha)) + logging.info("Evaluating template KDE") + template_kde = kde_awkde( + kde_train_samples, temp_samples, alp=optalpha, gl_bandwidth=optbw + ) # Compensation factor for downsampling of templates - template_kde *= 1. / probabilities + template_kde *= 1.0 / probabilities if args.min_ratio is not None: - logging.info(f'Applying minimum template KDE ratio {args.min_ratio}') - f_dest.attrs['min-kde-ratio'] = args.min_ratio + logging.info(f"Applying minimum template KDE ratio {args.min_ratio}") + f_dest.attrs["min-kde-ratio"] = args.min_ratio min_val = args.min_ratio * numpy.max(template_kde) template_kde = numpy.maximum(template_kde, min_val) else: - logging.info('Starting optimization of template KDE parameters') - optbw, optalpha = optimizedparam(temp_samples, alphagrid=args.alpha_grid, - bwgrid=args.bw_grid, nfold=args.nfold_template) - logging.info('Bandwidth %.4f, alpha %.2f' % (optbw, optalpha)) - logging.info('Evaluating template KDE') - template_kde = kde_awkde(temp_samples, temp_samples, alp=optalpha, - gl_bandwidth=optbw) + logging.info("Starting optimization of template KDE parameters") + optbw, optalpha = optimizedparam( + temp_samples, + alphagrid=args.alpha_grid, + bwgrid=args.bw_grid, + nfold=args.nfold_template, + ) + logging.info("Bandwidth %.4f, alpha %.2f" % (optbw, optalpha)) + logging.info("Evaluating template KDE") + template_kde = kde_awkde( + temp_samples, temp_samples, alp=optalpha, gl_bandwidth=optbw + ) f_dest.create_dataset("data_kde", data=template_kde) - f_dest.attrs['stat'] = "template-kde_file" -f_dest.attrs['template-file'] = args.template_file + f_dest.attrs["stat"] = "template-kde_file" +f_dest.attrs["template-file"] = args.template_file def signal_kde_extra_cpt(signal_kde, frac=None, volume=None): @@ -240,64 +295,78 @@ def signal_kde_extra_cpt(signal_kde, frac=None, volume=None): # Obtaining signal parameters if args.make_signal_kde: signal_pars = [] - signal_file = numpy.genfromtxt(args.signal_file, dtype=float, - delimiter=',', names=True) - f_dest.attrs['signal-file'] = args.signal_file - mass2_sgnl = signal_file['mass2'] + signal_file = numpy.genfromtxt( + args.signal_file, dtype=float, delimiter=",", names=True + ) + f_dest.attrs["signal-file"] = args.signal_file + mass2_sgnl = signal_file["mass2"] N_original = len(mass2_sgnl) if args.min_mass: idx = mass2_sgnl > args.min_mass mass2_sgnl = mass2_sgnl[idx] - logging.info('%i triggers out of %i with MASS2 > %s' % - (len(mass2_sgnl), N_original, str(args.min_mass))) + logging.info( + "%i triggers out of %i with MASS2 > %s" + % (len(mass2_sgnl), N_original, str(args.min_mass)) + ) else: idx = numpy.full(N_original, True) - mass1_sgnl = signal_file['mass1'][idx] + mass1_sgnl = signal_file["mass1"][idx] assert min(mass1_sgnl - mass2_sgnl) > 0 - + for param, slog in zip(args.fit_param, args.log_param): pvals = signal_file[param][idx] - if slog in ['False']: - logging.info('Using param: %s', param) + if slog in ["False"]: + logging.info("Using param: %s", param) signal_pars.append(pvals) - elif slog in ['True']: - logging.info('Using log param: %s', param) + elif slog in ["True"]: + logging.info("Using log param: %s", param) signal_pars.append(numpy.log(pvals)) else: - raise ValueError("invalid log param argument, use 'True', \ - or 'False'") - - signal_samples = numpy.vstack((signal_pars)).T + raise ValueError( + "invalid log param argument, use 'True', \ + or 'False'" + ) + + signal_samples = numpy.vstack(signal_pars).T f_dest.create_dataset("kde_train_samples", data=signal_samples) - logging.info('Starting optimization of signal KDE parameters') - optbw, optalpha = optimizedparam(signal_samples, bwgrid=args.bw_grid, - alphagrid=args.alpha_grid, nfold=args.nfold_signal) - logging.info('Bandwidth %.4f, alpha %.2f' % (optbw, optalpha)) - logging.info('Evaluating signal KDE') - signal_kde = kde_awkde(signal_samples, temp_samples, - alp=optalpha, gl_bandwidth=optbw) + logging.info("Starting optimization of signal KDE parameters") + optbw, optalpha = optimizedparam( + signal_samples, + bwgrid=args.bw_grid, + alphagrid=args.alpha_grid, + nfold=args.nfold_signal, + ) + logging.info("Bandwidth %.4f, alpha %.2f" % (optbw, optalpha)) + logging.info("Evaluating signal KDE") + signal_kde = kde_awkde( + signal_samples, temp_samples, alp=optalpha, gl_bandwidth=optbw + ) if args.extra_cpt_fraction is not None and args.temp_volume is not None: - f_dest.attrs.update({'extra-fraction-used': args.extra_cpt_fraction, - 'volume': args.temp_volume}) - modified_kde = signal_kde_extra_cpt(signal_kde, frac=args.extra_cpt_fraction, - volume=args.temp_volume) + f_dest.attrs.update( + {"extra-fraction-used": args.extra_cpt_fraction, "volume": args.temp_volume} + ) + modified_kde = signal_kde_extra_cpt( + signal_kde, frac=args.extra_cpt_fraction, volume=args.temp_volume + ) else: modified_kde = signal_kde f_dest.create_dataset("data_kde", data=modified_kde) - f_dest.attrs['stat'] = "signal-kde_file" + f_dest.attrs["stat"] = "signal-kde_file" if args.make_injection_kde: inj_pars = [] - inj_file = HFile(args.injection_file, 'r') - f_dest.attrs['injection-file'] = args.injection_file + inj_file = HFile(args.injection_file, "r") + f_dest.attrs["injection-file"] = args.injection_file snr = inj_file["events"]["snr_net"][:] N_original = len(snr) if args.min_snr: idx = snr > args.min_snr snr = snr[idx] - logging.info('%i triggers out of %i with SNR > %s' % - (len(snr), N_original, str(args.min_snr))) + logging.info( + "%i triggers out of %i with SNR > %s" + % (len(snr), N_original, str(args.min_snr)) + ) else: idx = numpy.full(N_original, True) mass1 = inj_file["events"]["mass1_detector"][idx] @@ -307,36 +376,43 @@ if args.make_injection_kde: for param, slog in zip(args.fit_param, args.log_param): pvals = triggers.get_param(param, args, mass1, mass2, spin1z, spin2z) - if slog in ['False']: - logging.info('Using param: %s', param) + if slog in ["False"]: + logging.info("Using param: %s", param) inj_pars.append(pvals) - elif slog in ['True']: - logging.info('Using log param: %s', param) + elif slog in ["True"]: + logging.info("Using log param: %s", param) inj_pars.append(numpy.log(pvals)) else: raise ValueError("invalid log param argument, use 'True', or 'False'") - inj_samples = numpy.vstack((inj_pars)).T + inj_samples = numpy.vstack(inj_pars).T f_dest.create_dataset("kde_train_samples", data=inj_samples) - logging.info('Starting optimization of injection KDE parameters') - optbw, optalpha = optimizedparam(inj_samples, alphagrid=args.alpha_grid, - bwgrid=args.bw_grid, nfold=args.nfold_injection) - logging.info('Bandwidth %.4f, alpha %.2f' % (optbw, optalpha)) - logging.info('Evaluating injection KDE') - injection_kde = kde_awkde(inj_samples, temp_samples, alp=optalpha, - gl_bandwidth=optbw) + logging.info("Starting optimization of injection KDE parameters") + optbw, optalpha = optimizedparam( + inj_samples, + alphagrid=args.alpha_grid, + bwgrid=args.bw_grid, + nfold=args.nfold_injection, + ) + logging.info("Bandwidth %.4f, alpha %.2f" % (optbw, optalpha)) + logging.info("Evaluating injection KDE") + injection_kde = kde_awkde( + inj_samples, temp_samples, alp=optalpha, gl_bandwidth=optbw + ) if args.extra_cpt_fraction is not None and args.temp_volume is not None: - f_dest.attrs.update({'extra-fraction-used': args.extra_cpt_fraction, - 'volume': args.temp_volume}) - modified_kde = signal_kde_extra_cpt(injection_kde, frac=args.extra_cpt_fraction, - volume=args.temp_volume) + f_dest.attrs.update( + {"extra-fraction-used": args.extra_cpt_fraction, "volume": args.temp_volume} + ) + modified_kde = signal_kde_extra_cpt( + injection_kde, frac=args.extra_cpt_fraction, volume=args.temp_volume + ) else: modified_kde = injection_kde f_dest.create_dataset("data_kde", data=modified_kde) - f_dest.attrs['stat'] = "signal-kde_file" + f_dest.attrs["stat"] = "signal-kde_file" -f_dest.attrs['alpha'] = optalpha -f_dest.attrs['bandwidth'] = optbw +f_dest.attrs["alpha"] = optalpha +f_dest.attrs["bandwidth"] = optbw f_dest.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_template_kde_max b/bin/all_sky_search/pycbc_template_kde_max index 41ee069db79..7226e83418c 100644 --- a/bin/all_sky_search/pycbc_template_kde_max +++ b/bin/all_sky_search/pycbc_template_kde_max @@ -1,31 +1,40 @@ #!/usr/bin/env python -import numpy, h5py, argparse, logging -from pycbc import init_logging, add_common_pycbc_options +import argparse +import logging + +import h5py +import numpy + +from pycbc import add_common_pycbc_options, init_logging from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--kde-files', nargs='+', required=True, - help='HDF files with KDE values') -parser.add_argument('--output-file', required=True, help='Name of output HDF file') -parser.add_argument('--min-ratio', type=float, - help='Minimum ratio for template_kde relative to the maximum') +parser.add_argument( + "--kde-files", nargs="+", required=True, help="HDF files with KDE values" +) +parser.add_argument("--output-file", required=True, help="Name of output HDF file") +parser.add_argument( + "--min-ratio", + type=float, + help="Minimum ratio for template_kde relative to the maximum", +) args = parser.parse_args() init_logging(args.verbose) -input_kdes = [HFile(kfile, 'r') for kfile in args.kde_files] +input_kdes = [HFile(kfile, "r") for kfile in args.kde_files] if len(input_kdes) < 2: raise ValueError("At least two input files are required.") # Creating output file and save datasets and attributes from input files -f_dest = HFile(args.output_file, 'w') +f_dest = HFile(args.output_file, "w") for name in input_kdes[0]: if isinstance(input_kdes[0][name], h5py.Dataset): - if name == 'data_kde': - data_combined = [kfile['data_kde'][:] for kfile in input_kdes] + if name == "data_kde": + data_combined = [kfile["data_kde"][:] for kfile in input_kdes] template_kde = numpy.maximum.reduce(data_combined) elif all(isinstance(kfile[name], h5py.Dataset) for kfile in input_kdes): dataset_values = [kfile[name][:] for kfile in input_kdes] @@ -37,13 +46,13 @@ for attr_name in input_kdes[0].attrs.keys(): f_dest.attrs[attr_name] = attr_values[0] if args.min_ratio is not None: - logging.info(f'Applying minimum template KDE ratio {args.min_ratio}') - f_dest.attrs['min-kde-ratio'] = args.min_ratio + logging.info(f"Applying minimum template KDE ratio {args.min_ratio}") + f_dest.attrs["min-kde-ratio"] = args.min_ratio min_val = args.min_ratio * numpy.max(template_kde) template_kde = numpy.maximum(template_kde, min_val) -f_dest.create_dataset('data_kde', data=template_kde) +f_dest.create_dataset("data_kde", data=template_kde) for kfile in input_kdes: kfile.close() f_dest.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/all_sky_search/pycbc_template_recovery_hist b/bin/all_sky_search/pycbc_template_recovery_hist index 1fbbce6a110..24012f051e8 100644 --- a/bin/all_sky_search/pycbc_template_recovery_hist +++ b/bin/all_sky_search/pycbc_template_recovery_hist @@ -1,46 +1,63 @@ #!/usr/bin/python -""" Histogram of templates where injections are found. -""" -import logging, argparse, numpy as np +"""Histogram of templates where injections are found.""" + +import argparse +import logging + +import numpy as np from matplotlib import use -use('Agg') + +use("Agg") from matplotlib import pyplot as plt + +from pycbc import add_common_pycbc_options, init_logging from pycbc.events import triggers from pycbc.io import HFile -from pycbc import init_logging, add_common_pycbc_options parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--output', required=True) -parser.add_argument('--found-injection-files', dest='found', nargs='+', - help='hdf file(s) with found injections') -parser.add_argument('--inspiral-trigger-files', dest='trig', nargs='*', - default=[], - help='hdf files(s) with injection single-ifo triggers, ' - 'if supplied must be one per found injection file') -parser.add_argument('--bank-files', nargs='*', default=[], - help='hdf file(s) with template parameters') -parser.add_argument('--x-param', #choices=['template_duration'], - help='parameter to histogram over') -parser.add_argument('--num-bins', type=int, default=30) -parser.add_argument('--log-x', action='store_true', - help='use log bins in parameter') -parser.add_argument('--min-stat', type=float, - help='only plot injections above given stat value') -parser.add_argument('--min-ifar', type=float, - help='only plot injections above given ifar value') +parser.add_argument("--output", required=True) +parser.add_argument( + "--found-injection-files", + dest="found", + nargs="+", + help="hdf file(s) with found injections", +) +parser.add_argument( + "--inspiral-trigger-files", + dest="trig", + nargs="*", + default=[], + help="hdf files(s) with injection single-ifo triggers, " + "if supplied must be one per found injection file", +) +parser.add_argument( + "--bank-files", nargs="*", default=[], help="hdf file(s) with template parameters" +) +parser.add_argument( + "--x-param", # choices=['template_duration'], + help="parameter to histogram over", +) +parser.add_argument("--num-bins", type=int, default=30) +parser.add_argument("--log-x", action="store_true", help="use log bins in parameter") +parser.add_argument( + "--min-stat", type=float, help="only plot injections above given stat value" +) +parser.add_argument( + "--min-ifar", type=float, help="only plot injections above given ifar value" +) args = parser.parse_args() init_logging(args.verbose) # should be same number of inj and trig files # and either 0 or 1 bank files or 1 per inj file -if not(len(args.trig) == 0 or len(args.trig) == len(args.found)): - raise RuntimeError('If trigger files are given, must be one per injection ' - 'file!') -if not(len(args.bank_files) < 2 or len(args.bank_files) == len(args.found)): - raise RuntimeError('If multiple bank files are used, must be one per ' - 'injection file!') +if not (len(args.trig) == 0 or len(args.trig) == len(args.found)): + raise RuntimeError("If trigger files are given, must be one per injection file!") +if not (len(args.bank_files) < 2 or len(args.bank_files) == len(args.found)): + raise RuntimeError( + "If multiple bank files are used, must be one per injection file!" + ) # duplicate bank file if necessary if len(args.bank_files) == 1: @@ -54,25 +71,25 @@ foundifar = np.array([]) foundparam = np.array([]) # cycle over injection files for f, b, t in zip(args.found, args.bank_files, args.trig): - logging.info('Processing %s, %s, %s' % (f, b, t)) - f = HFile(f, 'r') - t = HFile(t, 'r') if t else None - b = HFile(b, 'r') if b else None - foundstat = np.concatenate((foundstat, f['found_after_vetoes/stat'][:])) - foundifar = np.concatenate((foundifar, f['found_after_vetoes/ifar'][:])) + logging.info("Processing %s, %s, %s" % (f, b, t)) + f = HFile(f, "r") + t = HFile(t, "r") if t else None + b = HFile(b, "r") if b else None + foundstat = np.concatenate((foundstat, f["found_after_vetoes/stat"][:])) + foundifar = np.concatenate((foundifar, f["found_after_vetoes/ifar"][:])) # acrobatics to find right ifo ifo = tuple(t.keys())[0] if t else None ## each tuple is of form ('L1', 'detector_1') - #dettuples = [(val, key) for (key, val) in f.attrs.items() + # dettuples = [(val, key) for (key, val) in f.attrs.items() # if "detector" in key] - #ifotag = [tag for (name, tag) in dettuples if name == ifo][0] + # ifotag = [tag for (name, tag) in dettuples if name == ifo][0] ## get the trigger ids - #ftrig_id = f['found_after_vetoes/trigger_id%s' % ifotag[-1]] + # ftrig_id = f['found_after_vetoes/trigger_id%s' % ifotag[-1]] fparam, found_in_ifo = triggers.get_found_param(f, b, t, args.x_param, ifo) - #if args.x_param == 'template_duration': + # if args.x_param == 'template_duration': # # duration values foundparam = np.concatenate((foundparam, fparam)) - # t['%s/template_duration' % ifo][:][ftrig_id])) + # t['%s/template_duration' % ifo][:][ftrig_id])) # apply filters if args.min_stat is not None: @@ -85,28 +102,34 @@ if args.min_ifar is not None: foundstat = foundstat[above] foundifar = foundifar[above] foundparam = foundparam[above] -logging.info('%i found injections above threshold(s)' % len(foundparam)) +logging.info("%i found injections above threshold(s)" % len(foundparam)) # plot if args.log_x: # make sure to pick up last point in either direction - hbins = np.logspace(np.log10(0.999*foundparam.min()), - np.log10(1.001*foundparam.max()), - num=args.num_bins, endpoint=True) + hbins = np.logspace( + np.log10(0.999 * foundparam.min()), + np.log10(1.001 * foundparam.max()), + num=args.num_bins, + endpoint=True, + ) lims = 0.98 * foundparam.min(), 1.02 * foundparam.max() else: extent = foundparam.max() - foundparam.min() - hbins = np.linspace(foundparam.min() - 0.001 * extent, - foundparam.max() + 0.001 * extent, - num=args.num_bins, endpoint=True) + hbins = np.linspace( + foundparam.min() - 0.001 * extent, + foundparam.max() + 0.001 * extent, + num=args.num_bins, + endpoint=True, + ) lims = foundparam.min() - 0.02 * extent, foundparam.max() + 0.02 * extent n, bins, patches = plt.hist(foundparam, bins=hbins) -if args.verbose: print(bins, n) +if args.verbose: + print(bins, n) if args.log_x: - plt.semilogx() + plt.semilogx() plt.xlim(lims) -plt.xlabel(args.x_param.replace('_', ' ')) -plt.ylabel('Number of found injections above threshold') +plt.xlabel(args.x_param.replace("_", " ")) +plt.ylabel("Number of found injections above threshold") plt.savefig(args.output) plt.close() - diff --git a/bin/all_sky_search/pycbc_upload_single_event_to_gracedb b/bin/all_sky_search/pycbc_upload_single_event_to_gracedb index ca49c48233e..61f374acba5 100755 --- a/bin/all_sky_search/pycbc_upload_single_event_to_gracedb +++ b/bin/all_sky_search/pycbc_upload_single_event_to_gracedb @@ -20,10 +20,11 @@ Take a coinc xml file containing a single event and upload to gracedb. """ -import os -import sys import argparse import logging +import os +import sys + from ligo.gracedb.rest import GraceDb import pycbc @@ -31,33 +32,51 @@ from pycbc.io.gracedb import gracedb_tag_with_version parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--xml-file-for-upload', required=True, type=str, - help='LIGOLW XML file containing the event.') -parser.add_argument('--log-message', type=str, metavar='MESSAGE', - help='Add a log entry to each upload with the given message') -parser.add_argument('--testing', action="store_true", default=False, - help="Upload event to the TEST group of gracedb.") -parser.add_argument('--production-server', action="store_true", default=False, - help="Upload event to production graceDB. If not given " - "events will be uploaded to playground server.") -parser.add_argument('--search-tag', default='AllSky', - help="Specify the search tag. Default: AllSky") -parser.add_argument('--snr-timeseries-plot', - help="SNR timeseries plot to be uploaded with the event.") -parser.add_argument('--asd-plot', - help="ASD plot to be uploaded with the event.") -parser.add_argument('--skymap-plot', - help="Skymap plot to be uploaded with the event.") -parser.add_argument('--skymap-fits-file', - help="Skymap .fits file to be uploaded with the event.") -parser.add_argument('--source-probabilities', - help="Source probabilities file to be uploaded with the " - "event.") -parser.add_argument('--source-probabilities-plot', - help="Source probabilities plot to be uploaded with the " - "event.") -parser.add_argument('--labels', nargs='+', - help="Labels to add to the event in GraceDB") +parser.add_argument( + "--xml-file-for-upload", + required=True, + type=str, + help="LIGOLW XML file containing the event.", +) +parser.add_argument( + "--log-message", + type=str, + metavar="MESSAGE", + help="Add a log entry to each upload with the given message", +) +parser.add_argument( + "--testing", + action="store_true", + default=False, + help="Upload event to the TEST group of gracedb.", +) +parser.add_argument( + "--production-server", + action="store_true", + default=False, + help="Upload event to production graceDB. If not given " + "events will be uploaded to playground server.", +) +parser.add_argument( + "--search-tag", default="AllSky", help="Specify the search tag. Default: AllSky" +) +parser.add_argument( + "--snr-timeseries-plot", help="SNR timeseries plot to be uploaded with the event." +) +parser.add_argument("--asd-plot", help="ASD plot to be uploaded with the event.") +parser.add_argument("--skymap-plot", help="Skymap plot to be uploaded with the event.") +parser.add_argument( + "--skymap-fits-file", help="Skymap .fits file to be uploaded with the event." +) +parser.add_argument( + "--source-probabilities", + help="Source probabilities file to be uploaded with the event.", +) +parser.add_argument( + "--source-probabilities-plot", + help="Source probabilities plot to be uploaded with the event.", +) +parser.add_argument("--labels", nargs="+", help="Labels to add to the event in GraceDB") args = parser.parse_args() @@ -65,12 +84,12 @@ args = parser.parse_args() pycbc.init_logging(args.verbose, default_level=1) # Make the scitokens logger a little quieter # (it is called through GraceDB) -logging.getLogger('scitokens').setLevel(logging.root.level + 10) +logging.getLogger("scitokens").setLevel(logging.root.level + 10) if args.production_server: gracedb = GraceDb() else: - gracedb = GraceDb(service_url='https://gracedb-playground.ligo.org/api/') + gracedb = GraceDb(service_url="https://gracedb-playground.ligo.org/api/") labels = [l.upper() for l in (args.labels or [])] allowed_labels = gracedb.allowed_labels @@ -81,62 +100,55 @@ if set(labels) - set(allowed_labels): err_msg += f"{','.join(allowed_labels)}." raise RuntimeError(err_msg) -group_tag = 'Test' if args.testing else 'CBC' +group_tag = "Test" if args.testing else "CBC" r = gracedb.create_event( group_tag, - 'pycbc', + "pycbc", args.xml_file_for_upload, filecontents=open(args.xml_file_for_upload, "rb").read(), search=args.search_tag, offline=True, - labels=labels + labels=labels, ).json() logging.info("Uploaded event %s.", r["graceid"]) # add info for tracking code version -gracedb_tag_with_version(gracedb, r['graceid']) +gracedb_tag_with_version(gracedb, r["graceid"]) # document the absolute path to the input file -input_file_str = 'Candidate uploaded from ' \ - + os.path.abspath(args.xml_file_for_upload) -gracedb.write_log(r['graceid'], input_file_str) +input_file_str = "Candidate uploaded from " + os.path.abspath(args.xml_file_for_upload) +gracedb.write_log(r["graceid"], input_file_str) # document the command line used in the event log -log_str = 'Upload command: ' + ' '.join(sys.argv) +log_str = "Upload command: " + " ".join(sys.argv) # add the custom log message, if provided if args.log_message is not None: - log_str += '. ' + args.log_message + log_str += ". " + args.log_message -gracedb.write_log( - r['graceid'], - log_str, - tag_name=['analyst_comments'] -) +gracedb.write_log(r["graceid"], log_str, tag_name=["analyst_comments"]) def upload_file(upload_filename, displayname, comment, tag): """ Helper function to upload files associated with the event. """ - logging.info("Uploading %s file %s to event %s.", - displayname, upload_filename, r["graceid"]) + logging.info( + "Uploading %s file %s to event %s.", displayname, upload_filename, r["graceid"] + ) gracedb.write_log( r["graceid"], comment, filename=upload_filename, tag_name=[tag], - displayName=[displayname] + displayName=[displayname], ) if args.asd_plot: upload_file( - args.asd_plot, - "ASDs", - "PyCBC ASD estimate from the time of event", - "psd" + args.asd_plot, "ASDs", "PyCBC ASD estimate from the time of event", "psd" ) if args.snr_timeseries_plot: @@ -144,39 +156,23 @@ if args.snr_timeseries_plot: args.snr_timeseries_plot, "SNR timeseries", "SNR timeseries plot upload", - "background" + "background", ) if args.skymap_plot: - upload_file( - args.skymap_plot, - "", - "Skymap plot upload", - "sky_loc" - ) + upload_file(args.skymap_plot, "", "Skymap plot upload", "sky_loc") if args.skymap_fits_file: - upload_file( - args.skymap_fits_file, - "", - "sky localization complete", - "sky_loc" - ) + upload_file(args.skymap_fits_file, "", "sky localization complete", "sky_loc") if args.source_probabilities: upload_file( - args.source_probabilities, - "", - "Source probabilities JSON file upload", - "pe" + args.source_probabilities, "", "Source probabilities JSON file upload", "pe" ) if args.source_probabilities_plot: upload_file( - args.source_probabilities_plot, - "", - "Source probabilities plot upload", - "pe" + args.source_probabilities_plot, "", "Source probabilities plot upload", "pe" ) -logging.info('Done!') +logging.info("Done!") diff --git a/bin/bank/pycbc_aligned_bank_cat b/bin/bank/pycbc_aligned_bank_cat index e983a67c82f..37acba38e95 100644 --- a/bin/bank/pycbc_aligned_bank_cat +++ b/bin/bank/pycbc_aligned_bank_cat @@ -21,51 +21,69 @@ Program for concatenating the output of the geometric aligned bank dagman. This will gather all the meta-output files and create a valid template bank xml file. """ -import glob + import argparse +import glob + import numpy +import pycbc.version from igwn_ligolw import utils -from pycbc import tmpltbank -# Old ligolw output functions no longer imported at package level -import pycbc.tmpltbank.bank_output_utils as bank_output + import pycbc import pycbc.psd import pycbc.strain -import pycbc.version -from pycbc.io.ligolw import LIGOLWContentHandler + +# Old ligolw output functions no longer imported at package level +import pycbc.tmpltbank.bank_output_utils as bank_output +from pycbc import tmpltbank from pycbc.io.hdf import HFile +from pycbc.io.ligolw import LIGOLWContentHandler from pycbc.types import positive_float - -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_aligned_bank_cat" # Read command line options -parser = argparse.ArgumentParser(description=__doc__, - formatter_class=tmpltbank.IndentedHelpFormatterWithNL) +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=tmpltbank.IndentedHelpFormatterWithNL +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("-i", "--input-glob", - help="file glob the list of paramters") -parser.add_argument("-I", "--input-files", nargs='+', - help="Explicit list of input files.") -parser.add_argument("-O", "--output-file", - help="Output file to write the bank to.") -parser.add_argument("--metadata-file", metavar="METADATA_FILE", - help="XML file containing the process and process_params " - "tables that the aligned_bank code was run with.") -parser.add_argument("--f-low", action="store", type=positive_float, - required=True, - help="Lower frequency cutoff used in computing the " - "parameter space metric. REQUIRED. UNITS=Hz") -parser.add_argument("--f-upper", action="store", type=positive_float, - required=True, - help="Upper frequency cutoff used in computing the " - "parameter space metric. REQUIRED. UNITS=Hz") -parser.add_argument('--output-f-final', action='store_true', default=False, - help="Include 'f_final' in the output hdf file.") +parser.add_argument("-i", "--input-glob", help="file glob the list of paramters") +parser.add_argument( + "-I", "--input-files", nargs="+", help="Explicit list of input files." +) +parser.add_argument("-O", "--output-file", help="Output file to write the bank to.") +parser.add_argument( + "--metadata-file", + metavar="METADATA_FILE", + help="XML file containing the process and process_params " + "tables that the aligned_bank code was run with.", +) +parser.add_argument( + "--f-low", + action="store", + type=positive_float, + required=True, + help="Lower frequency cutoff used in computing the " + "parameter space metric. REQUIRED. UNITS=Hz", +) +parser.add_argument( + "--f-upper", + action="store", + type=positive_float, + required=True, + help="Upper frequency cutoff used in computing the " + "parameter space metric. REQUIRED. UNITS=Hz", +) +parser.add_argument( + "--output-f-final", + action="store_true", + default=False, + help="Include 'f_final' in the output hdf file.", +) options = parser.parse_args() @@ -83,22 +101,22 @@ mass2 = [] spin1z = [] spin2z = [] for inp_file in input_files: - inp_fp = HFile(inp_file, 'r') - data = inp_fp['accepted_templates'][:] + inp_fp = HFile(inp_file, "r") + data = inp_fp["accepted_templates"][:] if len(data) == 0: continue - mass1.extend(data[:,0]) - mass2.extend(data[:,1]) - spin1z.extend(data[:,2]) - spin2z.extend(data[:,3]) + mass1.extend(data[:, 0]) + mass2.extend(data[:, 1]) + spin1z.extend(data[:, 2]) + spin2z.extend(data[:, 3]) inp_fp.close() -temp_bank = numpy.array([mass1,mass2,spin1z,spin2z]).T +temp_bank = numpy.array([mass1, mass2, spin1z, spin2z]).T if options.metadata_file: - outdoc = utils.load_filename(options.metadata_file, - compress='auto', - contenthandler=LIGOLWContentHandler) + outdoc = utils.load_filename( + options.metadata_file, compress="auto", contenthandler=LIGOLWContentHandler + ) else: outdoc = None @@ -109,5 +127,5 @@ bank_output.output_bank_to_file( output_duration=True, approximant="TaylorF2", optDict=options.__dict__, - outdoc=outdoc + outdoc=outdoc, ) diff --git a/bin/bank/pycbc_aligned_stoch_bank b/bin/bank/pycbc_aligned_stoch_bank index a59738b6122..7d8f1772a95 100644 --- a/bin/bank/pycbc_aligned_stoch_bank +++ b/bin/bank/pycbc_aligned_stoch_bank @@ -21,55 +21,82 @@ Stochastic aligned spin bank generator. """ import argparse -import numpy import logging +import numpy + import pycbc -from pycbc import tmpltbank -# Old ligolw output functions no longer imported at package level -import pycbc.tmpltbank.bank_output_utils as bank_output -from pycbc.types import positive_float import pycbc.psd import pycbc.strain -from pycbc.pnutils import named_frequency_cutoffs +# Old ligolw output functions no longer imported at package level +import pycbc.tmpltbank.bank_output_utils as bank_output +from pycbc import tmpltbank +from pycbc.pnutils import named_frequency_cutoffs +from pycbc.types import positive_float # Read command line option _desc = __doc__[1:] -parser = argparse.ArgumentParser(description=_desc, - formatter_class=tmpltbank.IndentedHelpFormatterWithNL) +parser = argparse.ArgumentParser( + description=_desc, formatter_class=tmpltbank.IndentedHelpFormatterWithNL +) # Begin with code specific options pycbc.add_common_pycbc_options(parser) -parser.add_argument("-V", "--vary-fupper", action="store_true", default=False, - help="Use a variable upper frequency cutoff in laying " - "out the bank. OPTIONAL.") -parser.add_argument("--bank-fupper-step", type=positive_float, default=10., - help="Size of discrete frequency steps used when varying " - "the fupper. If --calculate-ethinca-metric and " - "--ethinca-freq-step are also given, the code will use " - "the smaller of the two step values. OPTIONAL. Units=Hz") -parser.add_argument("--bank-fupper-formula", default="SchwarzISCO", - choices=named_frequency_cutoffs.keys(), - help="Frequency cutoff formula for varying fupper. " - "Frequencies will be rounded to the nearest discrete " - "step. OPTIONAL.") -parser.add_argument("-N", "--num-seeds", action="store", type=int, - default=5000000, help="Number of seed points used in " - "bank construction. OPTIONAL.") -parser.add_argument("-n", "--num-failed-cutoff", action="store", type=int, - default=1000000000, - help="Maximum number of consecutive, not-accepted test " - "points after which bank generation will be stopped. " - "OPTIONAL. Default value is really large as --num-seeds" - " is intended to provide the termination condition.") -parser.add_argument("--random-seed", action="store", type=int, - default=None, - help="Random seed to use when calling numpy.random " - "functions used in obtaining the principal components in " - "parameter space and when translating points back to " - "physical space. If given, the code should give the " - "same output when run with the same random seed.") +parser.add_argument( + "-V", + "--vary-fupper", + action="store_true", + default=False, + help="Use a variable upper frequency cutoff in laying out the bank. OPTIONAL.", +) +parser.add_argument( + "--bank-fupper-step", + type=positive_float, + default=10.0, + help="Size of discrete frequency steps used when varying " + "the fupper. If --calculate-ethinca-metric and " + "--ethinca-freq-step are also given, the code will use " + "the smaller of the two step values. OPTIONAL. Units=Hz", +) +parser.add_argument( + "--bank-fupper-formula", + default="SchwarzISCO", + choices=named_frequency_cutoffs.keys(), + help="Frequency cutoff formula for varying fupper. " + "Frequencies will be rounded to the nearest discrete " + "step. OPTIONAL.", +) +parser.add_argument( + "-N", + "--num-seeds", + action="store", + type=int, + default=5000000, + help="Number of seed points used in bank construction. OPTIONAL.", +) +parser.add_argument( + "-n", + "--num-failed-cutoff", + action="store", + type=int, + default=1000000000, + help="Maximum number of consecutive, not-accepted test " + "points after which bank generation will be stopped. " + "OPTIONAL. Default value is really large as --num-seeds" + " is intended to provide the termination condition.", +) +parser.add_argument( + "--random-seed", + action="store", + type=int, + default=None, + help="Random seed to use when calling numpy.random " + "functions used in obtaining the principal components in " + "parameter space and when translating points back to " + "physical space. If given, the code should give the " + "same output when run with the same random seed.", +) tmpltbank.insert_base_bank_options(parser) @@ -98,25 +125,25 @@ if not opts.vary_fupper: opts.bank_fupper_formula = None opts.max_mismatch = 1 - opts.min_match tmpltbank.verify_metric_calculation_options(opts, parser) -metricParams=tmpltbank.metricParameters.from_argparse(opts) +metricParams = tmpltbank.metricParameters.from_argparse(opts) tmpltbank.verify_mass_range_options(opts, parser) -massRangeParams=tmpltbank.massRangeParameters.from_argparse(opts) +massRangeParams = tmpltbank.massRangeParameters.from_argparse(opts) pycbc.psd.verify_psd_options(opts, parser) if opts.psd_estimation: pycbc.strain.verify_strain_options(opts, parser) tmpltbank.verify_ethinca_metric_options(opts, parser) -ethincaParams=tmpltbank.ethincaParameters.from_argparse(opts) +ethincaParams = tmpltbank.ethincaParameters.from_argparse(opts) # delete default ethinca frequency step if calculation is not done -if ethincaParams.doEthinca==False: +if ethincaParams.doEthinca == False: ethincaParams.freqStep = None # Ensure consistency of ethinca and bank metric parameters tmpltbank.check_ethinca_against_bank_params(ethincaParams, metricParams) # Ethinca calculation should currently only be done for non-spin templates -if ethincaParams.full_ethinca and (massRangeParams.maxNSSpinMag>0.0 or - massRangeParams.maxBHSpinMag>0.0): - parser.error("Ethinca metric calculation is currently not valid for " - "nonzero spins!") +if ethincaParams.full_ethinca and ( + massRangeParams.maxNSSpinMag > 0.0 or massRangeParams.maxBHSpinMag > 0.0 +): + parser.error("Ethinca metric calculation is currently not valid for nonzero spins!") # Decide the frequency step to be used if varying fupper / calculating ethinca if ethincaParams.doEthinca and not opts.vary_fupper: @@ -128,10 +155,13 @@ elif ethincaParams.doEthinca and opts.vary_fupper: freqStep = opts.bank_fupper_step if opts.bank_fupper_step != ethincaParams.freqStep: freqStep = min(ethincaParams.freqStep, opts.bank_fupper_step) - logging.warning("Frequency step for varying fupper was not equal " - "to ethinca frequency step! Setting freqStep to " - "the minimum of the two, "+str(freqStep)) -else: freqStep = None + logging.warning( + "Frequency step for varying fupper was not equal " + "to ethinca frequency step! Setting freqStep to " + "the minimum of the two, " + str(freqStep) + ) +else: + freqStep = None # Set random seed if needed if opts.random_seed is not None: @@ -148,11 +178,16 @@ else: logging.info("Obtaining PSD") # Want the number of samples to be a binary number and Nyquist must be above # opts.f_upper. All this assumes that 1 / deltaF is a binary number -nyquistFreq = 2**numpy.ceil(numpy.log2(opts.f_upper)) +nyquistFreq = 2 ** numpy.ceil(numpy.log2(opts.f_upper)) numSamples = int(round(nyquistFreq / opts.delta_f)) + 1 -psd = pycbc.psd.from_cli(opts, length=numSamples, delta_f=opts.delta_f, - low_frequency_cutoff=opts.f_low, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC) +psd = pycbc.psd.from_cli( + opts, + length=numSamples, + delta_f=opts.delta_f, + low_frequency_cutoff=opts.f_low, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, +) metricParams.psd = psd # Begin by calculating a metric @@ -160,13 +195,14 @@ logging.info("Calculating metric") metricParams = tmpltbank.determine_eigen_directions( metricParams, vary_fmax=(opts.vary_fupper or ethincaParams.doEthinca), - vary_density=freqStep) + vary_density=freqStep, +) logging.info("Identifying limits of frequency") # Choose the frequency values to use for metric calculation -if opts.vary_fupper==False: - if ethincaParams.doEthinca==False: +if opts.vary_fupper == False: + if ethincaParams.doEthinca == False: refFreq = metricParams.fUpper else: # use the maximum frequency for which the moments were calculated @@ -179,9 +215,7 @@ else: fs = numpy.array(list(metricParams.evals.keys()), dtype=float) fs.sort() lowEve, highEve = tmpltbank.find_max_and_min_frequencies( - opts.bank_fupper_formula, - massRangeParams, - fs + opts.bank_fupper_formula, massRangeParams, fs ) refFreq = lowEve fs = fs[fs >= lowEve] @@ -190,7 +224,8 @@ else: logging.info("Calculating covariance matrix") vals = tmpltbank.estimate_mass_range( - 1000000, massRangeParams, metricParams, refFreq, covary=False) + 1000000, massRangeParams, metricParams, refFreq, covary=False +) cov = numpy.cov(vals) evalsCV, evecsCV = numpy.linalg.eig(cov) evecsCVdict = {} @@ -200,11 +235,7 @@ metricParams.evecsCV = evecsCVdict logging.info("Initialize the PartitionedTmpltbank class") partitioned_bank_object = tmpltbank.PartitionedTmpltbank( - massRangeParams, - metricParams, - refFreq, - opts.max_mismatch ** 0.5, - bin_range_check=1 + massRangeParams, metricParams, refFreq, opts.max_mismatch**0.5, bin_range_check=1 ) # Initialise counters @@ -215,54 +246,57 @@ Nr = 0 # Map the frequency values and normalizations to idx if --vary-fupper is used if opts.vary_fupper: - partitioned_bank_object.get_freq_map_and_normalizations(fs, - opts.bank_fupper_formula) + partitioned_bank_object.get_freq_map_and_normalizations( + fs, opts.bank_fupper_formula + ) logging.info("Starting bank placement") while True: if not (Ns % 100000): # For optimization we generate points in sets of 100000 - rMass1, rMass2, rSpin1z, rSpin2z = \ - tmpltbank.get_random_mass(100000, massRangeParams) + rMass1, rMass2, rSpin1z, rSpin2z = tmpltbank.get_random_mass( + 100000, massRangeParams + ) if opts.vary_fupper: mass_dict = { - 'mass1': rMass1, - 'mass2': rMass2, - 'spin1z': rSpin1z, - 'spin2z': rSpin2z + "mass1": rMass1, + "mass2": rMass2, + "spin1z": rSpin1z, + "spin2z": rSpin2z, } refEve = tmpltbank.return_nearest_cutoff( - opts.bank_fupper_formula, mass_dict, fs) - lambdas = tmpltbank.get_chirp_params(rMass1, rMass2, rSpin1z, - rSpin2z, metricParams.f0, - metricParams.pnOrder) + opts.bank_fupper_formula, mass_dict, fs + ) + lambdas = tmpltbank.get_chirp_params( + rMass1, rMass2, rSpin1z, rSpin2z, metricParams.f0, metricParams.pnOrder + ) mus = [] idx = 0 for freq in fs: - mus.append( - tmpltbank.get_mu_params(lambdas, metricParams, freq)) + mus.append(tmpltbank.get_mu_params(lambdas, metricParams, freq)) idx += 1 mus = numpy.array(mus) - vecs = tmpltbank.get_cov_params(rMass1, rMass2, rSpin1z, rSpin2z, - metricParams, refFreq) + vecs = tmpltbank.get_cov_params( + rMass1, rMass2, rSpin1z, rSpin2z, metricParams, refFreq + ) vecs = numpy.array(vecs) Ns = 0 # Then we check each point for acceptance if not (Np % 100000): logging.info("%d seeds", Np) - vs = vecs[:,Ns] + vs = vecs[:, Ns] Np = Np + 1 # Stop if we hit break condition if Np > opts.num_seeds: break # Calculate if any existing point is too close (set store to False) if opts.vary_fupper: - reject = partitioned_bank_object.test_point_distance_vary(vs, - refEve[Ns], mus[:,:,Ns], opts.max_mismatch) + reject = partitioned_bank_object.test_point_distance_vary( + vs, refEve[Ns], mus[:, :, Ns], opts.max_mismatch + ) else: - reject = partitioned_bank_object.test_point_distance(vs, - opts.max_mismatch) + reject = partitioned_bank_object.test_point_distance(vs, opts.max_mismatch) # Increment counters, check for break condition and continue if rejected if reject: Ns = Ns + 1 @@ -273,14 +307,20 @@ while True: # Add point, increment counters and continue if accepted Nr = 0 if opts.vary_fupper: - curr_mus = mus[:,:,Ns] + curr_mus = mus[:, :, Ns] point_fupper = refEve[Ns] else: curr_mus = None point_fupper = None - partitioned_bank_object.add_point_by_chi_coords(vs, rMass1[Ns], rMass2[Ns], - rSpin1z[Ns], rSpin2z[Ns], point_fupper=point_fupper, - mus=curr_mus) + partitioned_bank_object.add_point_by_chi_coords( + vs, + rMass1[Ns], + rMass2[Ns], + rSpin1z[Ns], + rSpin2z[Ns], + point_fupper=point_fupper, + mus=curr_mus, + ) N = N + 1 if not (N % 100000): logging.info("%d templates", N) @@ -297,7 +337,7 @@ bank_output.output_bank_to_file( opts.output_file, tempBank, programName="pycbc_aligned_stoch_bank", - optDict=opts.__dict__ + optDict=opts.__dict__, ) logging.info("Done") diff --git a/bin/bank/pycbc_bank_verification b/bin/bank/pycbc_bank_verification index d0b051dca7c..9238571b7fa 100644 --- a/bin/bank/pycbc_bank_verification +++ b/bin/bank/pycbc_bank_verification @@ -21,81 +21,125 @@ Test that a template bank covers a provided parameter space using a provided metric approximation to the parameter space. """ - import argparse import logging -import numpy + import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt +import numpy -from igwn_ligolw import lsctables, utils as ligolw_utils +matplotlib.use("Agg") +import pycbc.version +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils +from matplotlib import pyplot as plt import pycbc -import pycbc.version -from pycbc import tmpltbank, psd, strain -from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc import psd, strain, tmpltbank from pycbc.io.hdf import HFile +from pycbc.io.ligolw import LIGOLWContentHandler - -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_bank_verification" # Read command line option -parser = argparse.ArgumentParser(description=__doc__, - formatter_class=tmpltbank.IndentedHelpFormatterWithNL) +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=tmpltbank.IndentedHelpFormatterWithNL +) # Begin with code specific options pycbc.add_common_pycbc_options(parser) -parser.add_argument("--histogram-output-file", action="store", default=None, - help="Output a histogram of fitting factors to the " - "supplied file. If not given no histogram is produced.") -parser.add_argument("--print-distances", action="store_true", default=False, - help="Print details about the match of each point.") -parser.add_argument("-B", "--input-bank", action="store", required=True, - help="The template bank to use an input.") -parser.add_argument("-V", "--vary-fupper", action="store_true", default=False, - help="Use a variable upper frequency cutoff in laying " - "out the bank. OPTIONAL.") -parser.add_argument("--bank-fupper-step", type=float, default=10., - help="Size of discrete frequency steps used when varying " - "the fupper. If --calculate-ethinca-metric and " - "--ethinca-freq-step are also given, the code will use " - "the smaller of the two step values. This option will do " - "nothing if the --vary-fupper flag is not given. " - "OPTIONAL, default=10. Units=Hz") -parser.add_argument("--bank-fupper-formula", default="SchwarzISCO", - choices=["SchwarzISCO","LightRing","ERD"], - help="Frequency cutoff formula for varying fupper. " - "Frequencies will be rounded to the nearest discrete " - "step. This option will do nothing if --vary-fupper is " - "not given. OPTIONAL, default='SchwarzISCO'.") -parser.add_argument("-N", "--num-points", action="store", type=int, - default=100000, help="Number of points used to test ") -parser.add_argument("-P", "--point-file", action="store", default=None, - help="List of points to test bank against. Should be a " - "space-separated ASCII file where the columns correspond " - "to mass1, mass2, spin1z and spin2z respectively.") -parser.add_argument("--random-seed", action="store", type=int, - default=None, - help="Random seed to use when calling numpy.random " - "functions used in obtaining the principal components in " - "parameter space and when translating points back to " - "physical space. If given, the code should give the " - "same output when run with the same random seed.") -parser.add_argument("--bin-spacing", action="store", type=float, - default=0.03, - help="When read in, the template bank is placed into 2D " - "bins according to the points' location in the 'natural' " - "xi coordinate system. When computing matches points are " - "only compared against points in their own bin and " - "neighbouring bins. This argument specifies, in terms of " - "mismatch, the width of these bins. Note that metric " - "distance is proportional to the square root of the " - "mismatch.") +parser.add_argument( + "--histogram-output-file", + action="store", + default=None, + help="Output a histogram of fitting factors to the " + "supplied file. If not given no histogram is produced.", +) +parser.add_argument( + "--print-distances", + action="store_true", + default=False, + help="Print details about the match of each point.", +) +parser.add_argument( + "-B", + "--input-bank", + action="store", + required=True, + help="The template bank to use an input.", +) +parser.add_argument( + "-V", + "--vary-fupper", + action="store_true", + default=False, + help="Use a variable upper frequency cutoff in laying out the bank. OPTIONAL.", +) +parser.add_argument( + "--bank-fupper-step", + type=float, + default=10.0, + help="Size of discrete frequency steps used when varying " + "the fupper. If --calculate-ethinca-metric and " + "--ethinca-freq-step are also given, the code will use " + "the smaller of the two step values. This option will do " + "nothing if the --vary-fupper flag is not given. " + "OPTIONAL, default=10. Units=Hz", +) +parser.add_argument( + "--bank-fupper-formula", + default="SchwarzISCO", + choices=["SchwarzISCO", "LightRing", "ERD"], + help="Frequency cutoff formula for varying fupper. " + "Frequencies will be rounded to the nearest discrete " + "step. This option will do nothing if --vary-fupper is " + "not given. OPTIONAL, default='SchwarzISCO'.", +) +parser.add_argument( + "-N", + "--num-points", + action="store", + type=int, + default=100000, + help="Number of points used to test ", +) +parser.add_argument( + "-P", + "--point-file", + action="store", + default=None, + help="List of points to test bank against. Should be a " + "space-separated ASCII file where the columns correspond " + "to mass1, mass2, spin1z and spin2z respectively.", +) +parser.add_argument( + "--random-seed", + action="store", + type=int, + default=None, + help="Random seed to use when calling numpy.random " + "functions used in obtaining the principal components in " + "parameter space and when translating points back to " + "physical space. If given, the code should give the " + "same output when run with the same random seed.", +) +parser.add_argument( + "--bin-spacing", + action="store", + type=float, + default=0.03, + help="When read in, the template bank is placed into 2D " + "bins according to the points' location in the 'natural' " + "xi coordinate system. When computing matches points are " + "only compared against points in their own bin and " + "neighbouring bins. This argument specifies, in terms of " + "mismatch, the width of these bins. Note that metric " + "distance is proportional to the square root of the " + "mismatch.", +) # Insert the metric calculation options tmpltbank.insert_metric_calculation_options(parser) @@ -115,9 +159,9 @@ pycbc.init_logging(opts.verbose) # Sanity check options tmpltbank.verify_metric_calculation_options(opts, parser) -metricParams=tmpltbank.metricParameters.from_argparse(opts) +metricParams = tmpltbank.metricParameters.from_argparse(opts) tmpltbank.verify_mass_range_options(opts, parser) -massRangeParams=tmpltbank.massRangeParameters.from_argparse(opts) +massRangeParams = tmpltbank.massRangeParameters.from_argparse(opts) psd.verify_psd_options(opts, parser) if opts.psd_estimation: pycbc.strain.verify_strain_options(opts, parser) @@ -137,28 +181,35 @@ else: logging.info("Obtaining PSD") # Want the number of samples to be a binary number and Nyquist must be above # opts.f_upper. All this assumes that 1 / deltaF is a binary number -nyquistFreq = 2**numpy.ceil(numpy.log2(opts.f_upper)) +nyquistFreq = 2 ** numpy.ceil(numpy.log2(opts.f_upper)) numSamples = int(round(nyquistFreq / opts.delta_f)) + 1 -psd = pycbc.psd.from_cli(opts, length=numSamples, delta_f=opts.delta_f, - low_frequency_cutoff=opts.f_low, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC) +psd = pycbc.psd.from_cli( + opts, + length=numSamples, + delta_f=opts.delta_f, + low_frequency_cutoff=opts.f_low, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, +) metricParams.psd = psd # Begin by calculating a metric logging.info("Calculating metric") -metricParams = tmpltbank.determine_eigen_directions(metricParams, - vary_fmax=opts.vary_fupper, vary_density=opts.bank_fupper_step) +metricParams = tmpltbank.determine_eigen_directions( + metricParams, vary_fmax=opts.vary_fupper, vary_density=opts.bank_fupper_step +) # Choose the frequency values to use for metric calculation -if opts.vary_fupper==False: +if opts.vary_fupper == False: refFreq = metricParams.fUpper else: - # determine upper frequency cutoffs corresponding to the min and max + # determine upper frequency cutoffs corresponding to the min and max # total masses fs = numpy.array(list(metricParams.evals.keys()), dtype=float) fs.sort() - lowEve, highEve = tmpltbank.find_max_and_min_frequencies(\ - opts.bank_fupper_formula, massRangeParams, fs) + lowEve, highEve = tmpltbank.find_max_and_min_frequencies( + opts.bank_fupper_formula, massRangeParams, fs + ) refFreq = lowEve fs = fs[fs >= lowEve] fs = fs[fs <= highEve] @@ -166,42 +217,47 @@ else: logging.info("Calculating covariance matrix") vals = tmpltbank.estimate_mass_range( - 1000000, massRangeParams, metricParams, refFreq, covary=False) + 1000000, massRangeParams, metricParams, refFreq, covary=False +) cov = numpy.cov(vals) -evalsCV,evecsCV = numpy.linalg.eig(cov) +evalsCV, evecsCV = numpy.linalg.eig(cov) metricParams.evecsCV = {} metricParams.evecsCV[refFreq] = evecsCV # Initialize the class for generating the partitioned bank logging.info("Initialize the PartitionedTmpltbank class") -partitioned_bank_object = tmpltbank.PartitionedTmpltbank(massRangeParams, - metricParams, refFreq, (opts.bin_spacing)**0.5, - bin_range_check=1) +partitioned_bank_object = tmpltbank.PartitionedTmpltbank( + massRangeParams, metricParams, refFreq, (opts.bin_spacing) ** 0.5, bin_range_check=1 +) # Map the frequency values to idx if --vary-fupper is used if opts.vary_fupper: - partitioned_bank_object.get_freq_map_and_normalizations(fs, - opts.bank_fupper_formula) + partitioned_bank_object.get_freq_map_and_normalizations( + fs, opts.bank_fupper_formula + ) # Reading in the template bank logging.info("Reading template bank.") -if opts.input_bank.endswith(('.xml','.xml.gz','.xmlgz')): - indoc = ligolw_utils.load_filename(opts.input_bank, - contenthandler=LIGOLWContentHandler) +if opts.input_bank.endswith((".xml", ".xml.gz", ".xmlgz")): + indoc = ligolw_utils.load_filename( + opts.input_bank, contenthandler=LIGOLWContentHandler + ) template_list = lsctables.SnglInspiralTable.get_table(indoc) - partitioned_bank_object.add_tmpltbank_from_xml_table\ - (template_list, vary_fupper=opts.vary_fupper) + partitioned_bank_object.add_tmpltbank_from_xml_table( + template_list, vary_fupper=opts.vary_fupper + ) -elif opts.input_bank.endswith(('.h5','.hdf','.hdf5')): - h5_fp = HFile(opts.input_bank, 'r') - partitioned_bank_object.add_tmpltbank_from_hdf_file\ - (h5_fp, vary_fupper=opts.vary_fupper) +elif opts.input_bank.endswith((".h5", ".hdf", ".hdf5")): + h5_fp = HFile(opts.input_bank, "r") + partitioned_bank_object.add_tmpltbank_from_hdf_file( + h5_fp, vary_fupper=opts.vary_fupper + ) else: - err_msg = "Don't know how to read extension {}.".format(opts.input_bank) + err_msg = f"Don't know how to read extension {opts.input_bank}." raise NotImplementedError(err_msg) logging.info("Bank read in and sorted") @@ -213,29 +269,30 @@ points_spin2z = [] points_xis = [] points_fittingfactor = [] points_closestidxes = [] -outbins = [[0,0],[0,1],[1,0],[0,-1],[-1,0],[1,1],[1,-1],[-1,1],[-1,-1]] +outbins = [[0, 0], [0, 1], [1, 0], [0, -1], [-1, 0], [1, 1], [1, -1], [-1, 1], [-1, -1]] if opts.point_file: temp_data = numpy.loadtxt(opts.point_file) - rMass1 = temp_data[:,0] - rMass2 = temp_data[:,1] - rSpin1z = temp_data[:,2] - rSpin2z = temp_data[:,3] + rMass1 = temp_data[:, 0] + rMass2 = temp_data[:, 1] + rSpin1z = temp_data[:, 2] + rSpin2z = temp_data[:, 3] opts.num_points = len(rMass1) else: - rMass1, rMass2, rSpin1z, rSpin2z =\ - tmpltbank.get_random_mass(100000, massRangeParams) + rMass1, rMass2, rSpin1z, rSpin2z = tmpltbank.get_random_mass( + 100000, massRangeParams + ) if opts.vary_fupper: mass_dict = {} - mass_dict['m1'] = rMass1 - mass_dict['m2'] = rMass2 - mass_dict['s1z'] = rSpin1z - mass_dict['s2z'] = rSpin2z - refEve = tmpltbank.return_nearest_cutoff( - opts.bank_fupper_formula, mass_dict, fs) - lambdas = tmpltbank.get_chirp_params(rMass1, rMass2, rSpin1z, rSpin2z, - metricParams.f0, metricParams.pnOrder) + mass_dict["m1"] = rMass1 + mass_dict["m2"] = rMass2 + mass_dict["s1z"] = rSpin1z + mass_dict["s2z"] = rSpin2z + refEve = tmpltbank.return_nearest_cutoff(opts.bank_fupper_formula, mass_dict, fs) + lambdas = tmpltbank.get_chirp_params( + rMass1, rMass2, rSpin1z, rSpin2z, metricParams.f0, metricParams.pnOrder + ) mus = [] for freq in fs: if freq >= lowEve and freq <= highEve: @@ -243,16 +300,16 @@ if opts.vary_fupper: mus = numpy.array(mus) else: refEve = numpy.zeros(100000) - mus = numpy.zeros([1,1,100000]) -vecs = tmpltbank.get_cov_params(rMass1, rMass2, rSpin1z, rSpin2z, - metricParams, refFreq) + mus = numpy.zeros([1, 1, 100000]) +vecs = tmpltbank.get_cov_params(rMass1, rMass2, rSpin1z, rSpin2z, metricParams, refFreq) vecs = numpy.array(vecs) for idx_curr in range(opts.num_points): - vs = vecs[:,idx_curr] + vs = vecs[:, idx_curr] if opts.vary_fupper: - min_mismatch, idxes = partitioned_bank_object.calc_point_distance_vary(\ - vs, refEve[idx_curr], mus[:,:,idx_curr]) + min_mismatch, idxes = partitioned_bank_object.calc_point_distance_vary( + vs, refEve[idx_curr], mus[:, :, idx_curr] + ) else: min_mismatch, idxes = partitioned_bank_object.calc_point_distance(vs) @@ -262,8 +319,7 @@ for idx_curr in range(opts.num_points): points_spin1z.append(rSpin1z[idx_curr]) points_spin2z.append(rSpin2z[idx_curr]) points_xis.append(vs) - if min_mismatch > 1: - min_mismatch = 1 + min_mismatch = min(min_mismatch, 1) points_fittingfactor.append(1 - min_mismatch) points_closestidxes.append(idxes) @@ -286,18 +342,27 @@ if opts.print_distances: print() for idx in range(len(points_mass1)): print("The test point with the following parameters (m1,m2,S1z,S2z)") - print("%e %e %e %e" %(points_mass1[idx], points_mass2[idx],\ - points_spin1z[idx], points_spin2z[idx])) - print("Was recovered with an overlap of %e" %(points_fittingfactor[idx])) + print( + "%e %e %e %e" + % ( + points_mass1[idx], + points_mass2[idx], + points_spin1z[idx], + points_spin2z[idx], + ) + ) + print("Was recovered with an overlap of %e" % (points_fittingfactor[idx])) if points_closestidxes[idx] is None: print("No template was within a metric distance of 1.") print() continue print("The point in the bank with this overlap had (m1,m2,S1z,S2z)") - curr_m1, curr_m2, curr_spin1, curr_spin2 = \ - partitioned_bank_object.get_point_from_bins_and_idx(\ - points_closestidxes[idx][0], points_closestidxes[idx][1], - points_closestidxes[idx][2]) - print("%e %e %e %e" %(curr_m1, curr_m2, curr_spin1, curr_spin2)) + curr_m1, curr_m2, curr_spin1, curr_spin2 = ( + partitioned_bank_object.get_point_from_bins_and_idx( + points_closestidxes[idx][0], + points_closestidxes[idx][1], + points_closestidxes[idx][2], + ) + ) + print("%e %e %e %e" % (curr_m1, curr_m2, curr_spin1, curr_spin2)) print() - diff --git a/bin/bank/pycbc_brute_bank b/bin/bank/pycbc_brute_bank index bc893e61e3e..2e6978da441 100755 --- a/bin/bank/pycbc_brute_bank +++ b/bin/bank/pycbc_brute_bank @@ -17,68 +17,121 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Generate a bank of templates using a brute force stochastic method. -""" -import numpy +"""Generate a bank of templates using a brute force stochastic method.""" + +import argparse import logging -import signal -import tqdm import os -import argparse +import signal + +import numpy import numpy.random +import tqdm from scipy.stats import gaussian_kde -import pycbc.waveform, pycbc.filter, pycbc.types, pycbc.psd, pycbc.fft, pycbc.conversions +import pycbc.conversions +import pycbc.fft +import pycbc.filter import pycbc.pool +import pycbc.psd +import pycbc.types +import pycbc.waveform from pycbc import transforms -from pycbc.waveform.spa_tmplt import spa_length_in_time from pycbc.distributions import read_params_from_config from pycbc.distributions.utils import draw_samples_from_config, prior_from_config from pycbc.io import HFile +from pycbc.waveform.spa_tmplt import spa_length_in_time parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--output-file', required=True, - help='Output file name for template bank.') -parser.add_argument('--input-file', nargs='*', default=[], - help='Bank to use as an initial set of starting samples.') -parser.add_argument('--keep-entire-input-file', help="All of the templates in the input file will be kept in the output file, even if they do not pass normal mismatch criteria.", action='store_true') -parser.add_argument('--input-config', required=True, - help='Draw parameters from the given configure file.') -parser.add_argument('--minimal-match', default=0.97, type=float) -parser.add_argument('--buffer-length', default=2, type=float, - help='size of waveform buffer in seconds') -parser.add_argument('--use-td-waveform', action='store_true', - help='Generate waveform in the time domain (default is frequency domain).') -parser.add_argument('--full-resolution-buffer-length', default=None, type=float, - help='Size of the waveform buffer in seconds for generating time-domain signals at full resolution before conversion to the frequency domain.') -parser.add_argument('--max-signal-length', type= float, - help="When specified, it cuts the maximum length of the waveform model to the lengh provided") -parser.add_argument('--sample-rate', default=2048, type=float, - help='sample rate in seconds') -parser.add_argument('--low-frequency-cutoff', default=20.0, type=float) -parser.add_argument('--enable-sigma-bound', action='store_true') -parser.add_argument('--tau0-threshold', type=float) -parser.add_argument('--permissive', action='store_true', - help='Allow waveform generator to fail.') -parser.add_argument('--placement-iterations', default=1000, type=int, - help='Specify the number of attempts the bank should make when placing points. Use this option if the bank fails to place any points.') -parser.add_argument('--seed', type=int, default=0) -parser.add_argument('--tolerance', type=float) -parser.add_argument('--use-cross', action='store_true') -parser.add_argument('--max-q', type=float) -parser.add_argument('--tau0-crawl', type=float) -parser.add_argument('--tau0-start', type=float) -parser.add_argument('--tau0-end', type=float) -parser.add_argument('--tau0-cutoff-frequency', type=float, default=15.0) -parser.add_argument('--nprocesses', type=int, default=1, - help='Number of processes to use for waveform generation parallelization. If not given then only a single core will be used.') -parser.add_argument('--parallel-check', action='store_true', help="Do bank checking parallel, note that this means that proposals WILL NOT be checked against each other.") -parser.add_argument('--max-connections', type=int, help="Maximum number of matches to store with each template", default=numpy.inf) +parser.add_argument( + "--output-file", required=True, help="Output file name for template bank." +) +parser.add_argument( + "--input-file", + nargs="*", + default=[], + help="Bank to use as an initial set of starting samples.", +) +parser.add_argument( + "--keep-entire-input-file", + help="All of the templates in the input file will be kept in the output file, even if they do not pass normal mismatch criteria.", + action="store_true", +) +parser.add_argument( + "--input-config", + required=True, + help="Draw parameters from the given configure file.", +) +parser.add_argument("--minimal-match", default=0.97, type=float) +parser.add_argument( + "--buffer-length", default=2, type=float, help="size of waveform buffer in seconds" +) +parser.add_argument( + "--use-td-waveform", + action="store_true", + help="Generate waveform in the time domain (default is frequency domain).", +) +parser.add_argument( + "--full-resolution-buffer-length", + default=None, + type=float, + help="Size of the waveform buffer in seconds for generating time-domain signals at full resolution before conversion to the frequency domain.", +) +parser.add_argument( + "--max-signal-length", + type=float, + help="When specified, it cuts the maximum length of the waveform model to the lengh provided", +) +parser.add_argument( + "--sample-rate", default=2048, type=float, help="sample rate in seconds" +) +parser.add_argument("--low-frequency-cutoff", default=20.0, type=float) +parser.add_argument("--enable-sigma-bound", action="store_true") +parser.add_argument("--tau0-threshold", type=float) +parser.add_argument( + "--permissive", action="store_true", help="Allow waveform generator to fail." +) +parser.add_argument( + "--placement-iterations", + default=1000, + type=int, + help="Specify the number of attempts the bank should make when placing points. Use this option if the bank fails to place any points.", +) +parser.add_argument("--seed", type=int, default=0) +parser.add_argument("--tolerance", type=float) +parser.add_argument("--use-cross", action="store_true") +parser.add_argument("--max-q", type=float) +parser.add_argument("--tau0-crawl", type=float) +parser.add_argument("--tau0-start", type=float) +parser.add_argument("--tau0-end", type=float) +parser.add_argument("--tau0-cutoff-frequency", type=float, default=15.0) +parser.add_argument( + "--nprocesses", + type=int, + default=1, + help="Number of processes to use for waveform generation parallelization. If not given then only a single core will be used.", +) +parser.add_argument( + "--parallel-check", + action="store_true", + help="Do bank checking parallel, note that this means that proposals WILL NOT be checked against each other.", +) +parser.add_argument( + "--max-connections", + type=int, + help="Maximum number of matches to store with each template", + default=numpy.inf, +) pycbc.psd.insert_psd_option_group(parser) -parser.add_argument('--use-trimmed-buffer', action='store_true', - help=('When specified, the match calculation will use only the first and last 100 samples ' - 'of the waveform instead of the full buffer. Enabling this option makes match computation faster.')) +parser.add_argument( + "--use-trimmed-buffer", + action="store_true", + help=( + "When specified, the match calculation will use only the first and last 100 samples " + "of the waveform instead of the full buffer. Enabling this option makes match computation faster." + ), +) args = parser.parse_args() @@ -88,43 +141,51 @@ numpy.random.seed(args.seed) config_parser = pycbc.types.config.InterpolatingConfigParser([args.input_config]) variable_args, static_args = read_params_from_config( - config_parser, prior_section='prior', - vargs_section='variable_params', - sargs_section='static_params') + config_parser, + prior_section="prior", + vargs_section="variable_params", + sargs_section="static_params", +) -if any(config_parser.get_subsections('waveform_transforms')): +if any(config_parser.get_subsections("waveform_transforms")): waveform_transforms = transforms.read_transforms_from_config( - config_parser, 'waveform_transforms') + config_parser, "waveform_transforms" + ) else: waveform_transforms = None dists_joint = prior_from_config(cp=config_parser) + def get_stable_match(p_inj, p_tmplt): """Iteratively doubles buffer length until the match stabilizes.""" frbl = args.full_resolution_buffer_length current_buflen = args.buffer_length if frbl is None else frbl last_match = -1.0 - + # We loop until convergence or a safety cap (e.g., 64s) while (last_match == -1.0) or (current_buflen <= 64): # Create a transient generator for this specific resolution - tmp_gen = GenUniformWaveform(current_buflen, args.sample_rate, args.low_frequency_cutoff) + tmp_gen = GenUniformWaveform( + current_buflen, args.sample_rate, args.low_frequency_cutoff + ) h_inj = tmp_gen.generate(**p_inj) h_tmplt = tmp_gen.generate(**p_tmplt) current_match = tmp_gen.match(h_inj, h_tmplt) # If the change is sub-leading, we are done - if (args.full_resolution_buffer_length) is not None or (abs(current_match - last_match) < 1e-4): + if (args.full_resolution_buffer_length) is not None or ( + abs(current_match - last_match) < 1e-4 + ): return current_match, current_buflen - + last_match = current_match current_buflen *= 2 - + return last_match, current_buflen - - -class Shrinker(object): + + +class Shrinker: def __init__(self, data): self.data = data @@ -135,10 +196,13 @@ class Shrinker(object): self.data = self.data[:-1] return l -class TriangleBank(object): - """ A bank of templates that uses the triangle inequality to estimate + +class TriangleBank: + """ + A bank of templates that uses the triangle inequality to estimate matches based on prior ones. """ + def __init__(self, p=None): self.waveforms = p if p is not None else [] self.tbins = {} @@ -159,9 +223,9 @@ class TriangleBank(object): for b in [hp.tbin - 1, hp.tbin, hp.tbin + 1]: if b in self.tbins: - self.tbins[b].append(len(self)-1) + self.tbins[b].append(len(self) - 1) else: - self.tbins[b] = [len(self)-1] + self.tbins[b] = [len(self) - 1] def __getitem__(self, index): return self.waveforms[index] @@ -173,14 +237,14 @@ class TriangleBank(object): return numpy.array([p.params[k] for p in self.waveforms]) def sigma_match_bound(self, sig): - if not hasattr(self, 'sigma'): + if not hasattr(self, "sigma"): self.sigma = None if self.sigma is None or len(self.sigma) != len(self): self.sigma = numpy.array([h.s for h in bank.waveforms]) return numpy.minimum(sig / self.sigma, self.sigma / sig) def range(self): - if not hasattr(self, 'r'): + if not hasattr(self, "r"): self.r = None if self.r is None or len(self.r) != len(self): self.r = numpy.arange(0, len(self)) @@ -189,8 +253,9 @@ class TriangleBank(object): def culltau0(self, threshold): cull = numpy.where(self.tau0() < threshold)[0] - class dumb(object): + class dumb: pass + for c in cull: d = dumb() d.tau0 = self.waveforms[c].tau0 @@ -199,7 +264,7 @@ class TriangleBank(object): self.waveforms[c] = d def tau0(self): - if not hasattr(self, 't0'): + if not hasattr(self, "t0"): self.t0 = None if self.t0 is None or len(self.t0) != len(self): self.t0 = numpy.array([h.tau0 for h in self]) @@ -208,7 +273,7 @@ class TriangleBank(object): def __contains__(self, hp): mmax = 0 mnum = 0 - #Apply sigmas maximal match. + # Apply sigmas maximal match. if args.enable_sigma_bound: matches = self.sigma_match_bound(hp.s) r = self.range()[matches > hp.threshold] @@ -218,12 +283,11 @@ class TriangleBank(object): msig = len(r) - #Apply tau0 threshold + # Apply tau0 threshold if args.tau0_threshold: hp.tau0 = pycbc.conversions.tau0_from_mass1_mass2( - hp.params['mass1'], - hp.params['mass2'], - args.tau0_cutoff_frequency) + hp.params["mass1"], hp.params["mass2"], args.tau0_cutoff_frequency + ) hp.tbin = int(hp.tau0 / args.tau0_threshold) if hp.tbin in self.tbins: @@ -234,25 +298,27 @@ class TriangleBank(object): mtau = len(r) # Try to do some actual matches - inc = Shrinker(r*1) + inc = Shrinker(r * 1) while 1: j = inc.pop() if j is None: msort = matches[r].argsort() - + msorted = matches[r][msort] rsorted = r[msort] keep = numpy.ones(len(msorted), dtype=bool) if args.max_connections < len(keep): - #keep[args.max_connections//2: -args.max_connections//2] = False - keep[args.max_connections:] = False - + # keep[args.max_connections//2: -args.max_connections//2] = False + keep[args.max_connections :] = False + hp.matches = msorted[keep].copy() hp.indices = rsorted[keep].copy() - logging.info("TADD MaxMatch:%0.3f Size:%i " - "AfterSigma:%i AfterTau0:%i Matches:%i" - % (mmax, len(self), msig, mtau, mnum)) + logging.info( + "TADD MaxMatch:%0.3f Size:%i " + "AfterSigma:%i AfterTau0:%i Matches:%i" + % (mmax, len(self), msig, mtau, mnum) + ) hp.max_match = mmax return False @@ -273,12 +339,17 @@ class TriangleBank(object): if m > hp.threshold: return True - if m > mmax: - mmax = m - - def check_params(self, gen, params, threshold, - force_add=False, - parallel_check=False, progress=False): + mmax = max(mmax, m) + + def check_params( + self, + gen, + params, + threshold, + force_add=False, + parallel_check=False, + progress=False, + ): num_added = 0 total_num = len(tuple(params.values())[0]) @@ -289,22 +360,28 @@ class TriangleBank(object): waveform_cache = [] pool = pycbc.pool.choose_pool(args.nprocesses) - for return_wf in tqdm.tqdm(pool.imap_unordered( + for return_wf in tqdm.tqdm( + pool.imap_unordered( wf_wrapper, - (({k: params[k][idx] for k in params}, parallel_check, threshold) for idx in chunk)), - total=total_num, disable=not progress): + ( + ({k: params[k][idx] for k in params}, parallel_check, threshold) + for idx in chunk + ), + ), + total=total_num, + disable=not progress, + ): waveform_cache += [return_wf] - for hp in waveform_cache: if hp is not None: if hp.checked is None: hp.gen = gen hp.checked = hp not in self - + if hp.checked: self.max_matches.append(hp.max_match) - + if hp.checked or force_add: num_added += 1 self.insert(hp) @@ -317,9 +394,10 @@ class TriangleBank(object): return bank, num_added / total_num + def decimate_frequency_domain(template, target_df): """ - Returns a frequency-domain waveform resampled to a lower frequency resolution + Returns a frequency-domain waveform resampled to a lower frequency resolution (delta_f) by decimation. Parameters @@ -330,43 +408,51 @@ def decimate_frequency_domain(template, target_df): The target frequency resolution (delta_f) for the decimated signal. Returns - ---------- + ------- decimated_template : pycbc.types.FrequencySeries - A new FrequencySeries object with the decimated data and the specified + A new FrequencySeries object with the decimated data and the specified target delta_f. + """ # Calculate the decimation factor decimation_factor = int(target_df / template.delta_f) if decimation_factor < 1: - raise ValueError("Target delta_f must be greater than or equal to the original delta_f.") + raise ValueError( + "Target delta_f must be greater than or equal to the original delta_f." + ) # Decimate the data by selecting every 'decimation_factor'-th point decimated_signal = template.data[::decimation_factor] # Create a new FrequencySeries object with the decimated data and the target delta_f - decimated_template = pycbc.types.FrequencySeries(decimated_signal, delta_f=target_df) + decimated_template = pycbc.types.FrequencySeries( + decimated_signal, delta_f=target_df + ) return decimated_template + def handle_exit_signal(signum, frame): logging.warning(f"Signal {signum} received. Triggering emergency save...") save_bank() # Force exit so the script doesn't try to resume the loops os._exit(0) -def save_bank(): # Use a local name + +def save_bank(): # Use a local name logging.info("Saving current bank to %s", args.output_file) - with HFile(args.output_file, 'w') as o: - o.attrs['minimal_match'] = args.minimal_match + with HFile(args.output_file, "w") as o: + o.attrs["minimal_match"] = args.minimal_match for k in bank.keys(): val = bank.key(k) - if val.dtype.char == 'U': - val = val.astype('bytes') + if val.dtype.char == "U": + val = val.astype("bytes") o[k] = val - o['max_matches'] = numpy.array(bank.max_matches) + o["max_matches"] = numpy.array(bank.max_matches) logging.info("Save complete.") -class GenUniformWaveform(object): + +class GenUniformWaveform: def __init__(self, buffer_length, sample_rate, f_lower): self.f_lower = f_lower self.delta_f = 1.0 / buffer_length @@ -374,10 +460,10 @@ class GenUniformWaveform(object): self.flen = tlen // 2 + 1 psd = pycbc.psd.from_cli(args, self.flen, self.delta_f, self.f_lower) self.kmin = int(f_lower * buffer_length) - self.w = ((1.0 / psd[self.kmin:-1]) ** 0.5).astype(numpy.float32) + self.w = ((1.0 / psd[self.kmin : -1]) ** 0.5).astype(numpy.float32) qtilde = pycbc.types.zeros(tlen, numpy.complex64) q = pycbc.types.zeros(tlen, numpy.complex64) - self.qtilde_view = qtilde[self.kmin:self.flen - 1] + self.qtilde_view = qtilde[self.kmin : self.flen - 1] self.ifft = pycbc.fft.IFFT(qtilde, q) self.md = q @@ -389,52 +475,58 @@ class GenUniformWaveform(object): def generate(self, **kwds): if args.max_signal_length is not None: - flow = numpy.arange(self.f_lower, 100, .1)[::-1] - length = spa_length_in_time(mass1=kwds['mass1'], mass2=kwds['mass2'], f_lower=flow, phase_order=-1) - maxlen = args.max_signal_length - x = numpy.searchsorted(length, maxlen) - 1 - l = length[x] - f = flow[x] + flow = numpy.arange(self.f_lower, 100, 0.1)[::-1] + length = spa_length_in_time( + mass1=kwds["mass1"], mass2=kwds["mass2"], f_lower=flow, phase_order=-1 + ) + maxlen = args.max_signal_length + x = numpy.searchsorted(length, maxlen) - 1 + l = length[x] + f = flow[x] else: - f = self.f_lower + f = self.f_lower - if 'f_lower' not in kwds: - kwds['f_lower'] = f + if "f_lower" not in kwds: + kwds["f_lower"] = f - if hasattr(kwds['approximant'], 'decode'): - kwds['approximant'] = kwds['approximant'].decode() + if hasattr(kwds["approximant"], "decode"): + kwds["approximant"] = kwds["approximant"].decode() - if args.full_resolution_buffer_length is not None: - buff_len = args.full_resolution_buffer_length + if args.full_resolution_buffer_length is not None: + buff_len = args.full_resolution_buffer_length else: buff_len = 1.0 / self.delta_f - if args.use_td_waveform and kwds['approximant'] in pycbc.waveform.td_approximants(): + if ( + args.use_td_waveform + and kwds["approximant"] in pycbc.waveform.td_approximants() + ): + hp, hc = pycbc.waveform.get_td_waveform( + delta_t=1.0 / args.sample_rate, **kwds + ) - hp, hc = pycbc.waveform.get_td_waveform(delta_t=1.0 / args.sample_rate, - **kwds) + hp = hp.to_frequencyseries(delta_f=1.0 / buff_len) + hc = hc.to_frequencyseries(delta_f=1.0 / buff_len) - hp = hp.to_frequencyseries(delta_f = 1.0 / buff_len) - hc = hc.to_frequencyseries(delta_f = 1.0 / buff_len) + elif kwds["approximant"] in pycbc.waveform.fd_approximants(): + hp, hc = pycbc.waveform.get_fd_waveform(delta_f=1.0 / buff_len, **kwds) - elif kwds['approximant'] in pycbc.waveform.fd_approximants(): - - hp, hc = pycbc.waveform.get_fd_waveform(delta_f = 1.0 / buff_len, - **kwds) - - if args.use_cross: + if args.use_cross: hp = hc - if 'fratio' in kwds: - hp = hc * kwds['fratio'] + hp * (1 - kwds['fratio']) + if "fratio" in kwds: + hp = hc * kwds["fratio"] + hp * (1 - kwds["fratio"]) else: dt = 1.0 / args.sample_rate hp = pycbc.waveform.get_waveform_filter( - pycbc.types.zeros(buff_len * args.sample_rate // 2 + 1, - dtype=numpy.complex64), - delta_f=1.0 / buff_len, delta_t=dt, - **kwds) + pycbc.types.zeros( + buff_len * args.sample_rate // 2 + 1, dtype=numpy.complex64 + ), + delta_f=1.0 / buff_len, + delta_t=dt, + **kwds, + ) if args.full_resolution_buffer_length is not None: # Decimate the generated signal to a reduced frequency resolution @@ -442,13 +534,12 @@ class GenUniformWaveform(object): hp.resize(self.flen) hp = hp.astype(numpy.complex64) - - hp[self.kmin:-1] *= self.w - s = float(1.0 / pycbc.filter.sigmasq(hp, - low_frequency_cutoff=f) ** 0.5) + + hp[self.kmin : -1] *= self.w + s = float(1.0 / pycbc.filter.sigmasq(hp, low_frequency_cutoff=f) ** 0.5) hp *= s hp.params = kwds - hp.view = hp[self.kmin:-1] + hp.view = hp[self.kmin : -1] hp.s = (1.0 / s) ** 2.0 return hp @@ -458,6 +549,7 @@ class GenUniformWaveform(object): m = max(abs(self.md).max(), abs(self.md2).max()) return m * 4.0 * self.delta_f + r = 0 if not args.tolerance: tolerance = (1 - args.minimal_match) / 10 @@ -466,63 +558,75 @@ else: size = int(1.0 / tolerance) -gen = GenUniformWaveform(args.buffer_length, - args.sample_rate, args.low_frequency_cutoff) +gen = GenUniformWaveform( + args.buffer_length, args.sample_rate, args.low_frequency_cutoff +) bank = TriangleBank() + def silent_child_exit(signum, frame): # The child workers execute this branch # os._exit(0) ensures they die instantly and quietly os._exit(0) + def wf_wrapper(args): p, parallel_check, threshold = args - + signal.signal(signal.SIGINT, silent_child_exit) signal.signal(signal.SIGTERM, silent_child_exit) try: hp = gen.generate(**p) hp.checked = None - hp.threshold = threshold + hp.threshold = threshold if parallel_check: hp.gen = gen - hp.checked = hp not in bank + hp.checked = hp not in bank hp.gen = None return hp except Exception as e: print(e) return None + if len(args.input_file) > 0: total_raw_templates = 0 for input_file in args.input_file: logging.info("Loading: %s", input_file) - with HFile(input_file, 'r') as f: - params = {k: f[k][:] for k in f.keys() if k not in ['max_matches']} + with HFile(input_file, "r") as f: + params = {k: f[k][:] for k in f.keys() if k not in ["max_matches"]} total_raw_templates += len(params[list(params.keys())[0]]) - bank, _ = bank.check_params(gen, params, args.minimal_match, - force_add=args.keep_entire_input_file, - parallel_check=args.parallel_check, - progress=True) + bank, _ = bank.check_params( + gen, + params, + args.minimal_match, + force_add=args.keep_entire_input_file, + parallel_check=args.parallel_check, + progress=True, + ) reduction = total_raw_templates - len(bank) - logging.info("Combined Size: %i (Naive Total: %i, Removed: %i)", - len(bank), total_raw_templates, reduction) + logging.info( + "Combined Size: %i (Naive Total: %i, Removed: %i)", + len(bank), + total_raw_templates, + reduction, + ) + def draw(rtype): - if rtype == 'uniform': + if rtype == "uniform": # `draw_samples_from_config` has its own fixed seed, so must overwrite it. - random_seed = numpy.random.randint(low=0, high=2**32-1) + random_seed = numpy.random.randint(low=0, high=2**32 - 1) samples = draw_samples_from_config(args.input_config, size, random_seed) params = {name: samples[name] for name in samples.fieldnames} # Add `static_args` back. if static_args is not None: for k in static_args.keys(): - params[k] = numpy.array([static_args[k]]*size) + params[k] = numpy.array([static_args[k]] * size) - elif rtype == 'kde': + elif rtype == "kde": trail = 300 - if trail > len(bank): - trail = len(bank) + trail = min(trail, len(bank)) p = variable_args bdata = numpy.array([bank.key(k)[-trail:] for k in p]) kde = gaussian_kde(bdata) @@ -532,7 +636,7 @@ def draw(rtype): # Add `static_args` back, some transformations may need them. if static_args is not None: for k in static_args.keys(): - params[k] = numpy.array([static_args[k]]*size) + params[k] = numpy.array([static_args[k]] * size) # Apply `waveform_transforms` defined in the .ini file to samples. if waveform_transforms is not None: @@ -543,26 +647,26 @@ def draw(rtype): params = {k: params[k][l] for k in params} return params + def cdraw(rtype, ts, te): from pycbc.conversions import tau0_from_mass1_mass2 p = draw(rtype) - if len(p[list(p.keys())[0]]) > 0: - t = tau0_from_mass1_mass2(p['mass1'], p['mass2'], - args.tau0_cutoff_frequency) + if len(p[list(p.keys())[0]]) > 0: + t = tau0_from_mass1_mass2(p["mass1"], p["mass2"], args.tau0_cutoff_frequency) l = (t < te) & (t > ts) p = {k: p[k][l] for k in p} i = 0 while len(p[list(p.keys())[0]]) < size: - tp = draw(rtype) - if len(tp[list(tp.keys())[0]]) > 0: - t = tau0_from_mass1_mass2(tp['mass1'], tp['mass2'], - args.tau0_cutoff_frequency) + if len(tp[list(tp.keys())[0]]) > 0: + t = tau0_from_mass1_mass2( + tp["mass1"], tp["mass2"], args.tau0_cutoff_frequency + ) l = (t < te) & (t > ts) tp = {k: tp[k][l] for k in tp} - + p = {k: numpy.concatenate([p[k], tp[k]]) for k in p} i += 1 @@ -574,6 +678,7 @@ def cdraw(rtype, ts, te): return p + signal.signal(signal.SIGINT, handle_exit_signal) signal.signal(signal.SIGTERM, handle_exit_signal) @@ -584,7 +689,6 @@ go = True region = 0 while tau0s < args.tau0_end: - tau0e = min(tau0e, args.tau0_end) conv = 1 @@ -592,39 +696,56 @@ while tau0s < args.tau0_end: while conv > tolerance: # Standard Round r += 1 - params = cdraw('uniform', tau0s, tau0e) + params = cdraw("uniform", tau0s, tau0e) if params is None: if len(bank) > 0: go = False break blen = len(bank) - bank, uconv = bank.check_params(gen, params, args.minimal_match, - parallel_check=args.parallel_check) - logging.info("%s: Round (U): %s Size: %s conv: %s added: %s", - region, r, len(bank), uconv, len(bank) - blen) + bank, uconv = bank.check_params( + gen, params, args.minimal_match, parallel_check=args.parallel_check + ) + logging.info( + "%s: Round (U): %s Size: %s conv: %s added: %s", + region, + r, + len(bank), + uconv, + len(bank) - blen, + ) if r > 10: conv = uconv kloop = 0 - while ((kloop == 0) or (kconv / okconv) > .5) and len(bank) > 10: + while ((kloop == 0) or (kconv / okconv) > 0.5) and len(bank) > 10: r += 1 kloop += 1 - params = cdraw('kde', tau0s, tau0e) + params = cdraw("kde", tau0s, tau0e) blen = len(bank) - bank, kconv = bank.check_params(gen, params, args.minimal_match, - parallel_check=args.parallel_check) + bank, kconv = bank.check_params( + gen, params, args.minimal_match, parallel_check=args.parallel_check + ) - trail_matches = numpy.array(bank.max_matches[int(len(bank.max_matches)*0.9):]) + trail_matches = numpy.array( + bank.max_matches[int(len(bank.max_matches) * 0.9) :] + ) ave = numpy.mean(trail_matches) - logging.info("%s: Round (K) (%s): %s Size: %s conv: %0.4f added: %s Trail Ave: %0.4f Trail Min: %0.4f", - region, kloop, r, len(bank), kconv, len(bank) - blen, - ave, trail_matches.min()) - + logging.info( + "%s: Round (K) (%s): %s Size: %s conv: %0.4f added: %s Trail Ave: %0.4f Trail Min: %0.4f", + region, + kloop, + r, + len(bank), + kconv, + len(bank) - blen, + ave, + trail_matches.min(), + ) if uconv: - logging.info('Ratio of convergences: %2.3f' % (kconv / (uconv))) - logging.info('Progress: {:.0%} completed'.format(tau0e/args.tau0_end)) + logging.info("Ratio of convergences: %2.3f" % (kconv / (uconv))) + logging.info(f"Progress: {tau0e / args.tau0_end:.0%} completed") if kloop == 1: okconv = kconv @@ -634,8 +755,7 @@ while tau0s < args.tau0_end: break bank.culltau0(tau0s - args.tau0_threshold * 2.0) - logging.info("Region Done %3.1f-%3.1f, %s stored", - tau0s, tau0e, bank.activelen()) + logging.info("Region Done %3.1f-%3.1f, %s stored", tau0s, tau0e, bank.activelen()) region += 1 tau0s += args.tau0_crawl / 2 tau0e += args.tau0_crawl / 2 diff --git a/bin/bank/pycbc_coinc_bank2hdf b/bin/bank/pycbc_coinc_bank2hdf index caaa905746c..a68a84a7932 100644 --- a/bin/bank/pycbc_coinc_bank2hdf +++ b/bin/bank/pycbc_coinc_bank2hdf @@ -16,10 +16,12 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This program converts a standard sngl_inspiral table based template bank +""" +This program converts a standard sngl_inspiral table based template bank into an hdf format that includes a template hash used to associate triggers with their template. """ + import argparse import logging @@ -29,13 +31,17 @@ from pycbc.waveform import bank # the following are the default parameters that will be loaded from the # xml file (and what they are called in the xml file) default_parameters = [ - "mass1", "mass2", - "spin1z", "spin2z", + "mass1", + "mass2", + "spin1z", + "spin2z", "f_lower:alpha6", - ] +] + def parse_parameters(parameters): - """Parses the parameters argument into names to write to and columns + """ + Parses the parameters argument into names to write to and columns to read from. """ outnames = [] @@ -47,64 +53,82 @@ def parse_parameters(parameters): elif len(ps) == 2: outname, column = ps else: - raise ValueError("parameter %s not formatted correctly; " %(p) + - "see help") + raise ValueError( + "parameter %s not formatted correctly; " % (p) + "see help" + ) outnames.append(outname) columns.append(column) return outnames, columns + parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--bank-file', required=True, - help="The bank file to load. Must end in '.xml[.gz]' " - "and must contain a SnglInspiral table or must end " - "in .hdf/.h5 and contain template bank parameters.") -parser.add_argument('--output-file', required=True, - help="The ouput file name. Must end in '.hdf'.") -parser.add_argument("--parameters", metavar="parameter[:xml_column]", - nargs="+", default=None, - help="The parameters to load from the xml file and to " - "write to the hdf file. The given name will be the " - "dataset's name in the hdf file. If this is " - "different than the column name in the xml file, " - "provide the column name after a colon, e.g., " - "'inclination:alpha3'. Otherwise, the given name " - "will assumed to be the same as the column name. " - "Not used with HDF input files. " - "Default is '%s'." %(' '.join(default_parameters))) +parser.add_argument( + "--bank-file", + required=True, + help="The bank file to load. Must end in '.xml[.gz]' " + "and must contain a SnglInspiral table or must end " + "in .hdf/.h5 and contain template bank parameters.", +) +parser.add_argument( + "--output-file", required=True, help="The ouput file name. Must end in '.hdf'." +) +parser.add_argument( + "--parameters", + metavar="parameter[:xml_column]", + nargs="+", + default=None, + help="The parameters to load from the xml file and to " + "write to the hdf file. The given name will be the " + "dataset's name in the hdf file. If this is " + "different than the column name in the xml file, " + "provide the column name after a colon, e.g., " + "'inclination:alpha3'. Otherwise, the given name " + "will assumed to be the same as the column name. " + "Not used with HDF input files. " + "Default is '%s'." % (" ".join(default_parameters)), +) # add the ability to specify the approximant to use -bank.add_approximant_arg(parser, - help="Specify the approximant to use with each " - "template. See pycbc_inspiral's help message " - "for syntax details. If provided, 'approximant'" - "will be added to the list of parameters.") -parser.add_argument("--force", action="store_true", default=False, - help="Overwrite the given hdf file if it exists. " - "Otherwise, an error is raised.") +bank.add_approximant_arg( + parser, + help="Specify the approximant to use with each " + "template. See pycbc_inspiral's help message " + "for syntax details. If provided, 'approximant'" + "will be added to the list of parameters.", +) +parser.add_argument( + "--force", + action="store_true", + default=False, + help="Overwrite the given hdf file if it exists. Otherwise, an error is raised.", +) args = parser.parse_args() pycbc.init_logging(args.verbose) # parse the parameters -if args.bank_file.endswith(('.xml','.xml.gz','.xmlgz')) or \ - args.parameters is not None: +if ( + args.bank_file.endswith((".xml", ".xml.gz", ".xmlgz")) + or args.parameters is not None +): if args.parameters is None: args.parameters = default_parameters outnames, columns = parse_parameters(args.parameters) name_map = dict(zip(columns, outnames)) else: - columns=None - name_map=None + columns = None + name_map = None # load the file -logging.info("Loading %s" %(args.bank_file)) -bankf = bank.TemplateBank(args.bank_file, approximant=args.approximant, - parameters=columns) +logging.info("Loading %s" % (args.bank_file)) +bankf = bank.TemplateBank( + args.bank_file, approximant=args.approximant, parameters=columns +) # rename the columns to the outnames if name_map is not None: params = list(bankf.table.fieldnames) - for ii,p in enumerate(params): + for ii, p in enumerate(params): try: params[ii] = name_map[p] except KeyError: @@ -116,5 +140,5 @@ logging.info("Getting template hashes") bankf.ensure_hash() # write to output -logging.info("Writing to %s" %(args.output_file)) +logging.info("Writing to %s" % (args.output_file)) bankf.write_to_hdf(args.output_file, force=args.force) diff --git a/bin/bank/pycbc_geom_aligned_2dstack b/bin/bank/pycbc_geom_aligned_2dstack index 547450844b9..0aee514cebf 100644 --- a/bin/bank/pycbc_geom_aligned_2dstack +++ b/bin/bank/pycbc_geom_aligned_2dstack @@ -28,73 +28,110 @@ It also prints information about the discarded points for debugging. import argparse import copy -import numpy import logging -import pycbc.tmpltbank +import numpy import pycbc.version + +import pycbc.tmpltbank from pycbc import pnutils -from pycbc.tmpltbank.lambda_mapping import pycbcValidOrdersHelpDescriptions from pycbc.io import HFile +from pycbc.tmpltbank.lambda_mapping import pycbcValidOrdersHelpDescriptions - -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_geom_aligned_2dstack" # Feed in command line options usage = """usage: %prog [options]""" _desc = __doc__[1:] -parser = argparse.ArgumentParser(usage, description=_desc, - formatter_class=pycbc.tmpltbank.IndentedHelpFormatterWithNL) +parser = argparse.ArgumentParser( + usage, + description=_desc, + formatter_class=pycbc.tmpltbank.IndentedHelpFormatterWithNL, +) pycbc.add_common_pycbc_options(parser) # Code specific options -parser.add_argument("--pn-order", action="store", type=str,\ - default=None,\ - help="Determines the PN order to use. Note that if you "+\ - "placing a bank of non-spinning templates, any "+\ - "spin-related terms in the metric will always "+\ - "be zero. REQUIRED ARGUMENT: "+\ - "choices are: %s" %(pycbcValidOrdersHelpDescriptions)) -parser.add_argument("--f0", action="store", type=float,\ - default=70.,\ - help="f0 is used as a dynamic rescaling factor when "+\ - "calculating the integrals used in metric "+\ - "construction. IE. instead of integrating F(f) we "+\ - "integrate F(f/f0) and then remove f0 after the fact."+\ - "The default option should be fine here for most "+\ - "applications. OPTIONAL"+\ - "WARNING: If using ethinca calculation this must be "+\ - "equal to f-low. UNITS=Hz") -parser.add_argument("-I", "--input-file", action="store", type=str, - required=True, - help="HDF file containing the input bank and the " - "various metric information. " - "REQUIRED ARGUMENT.") -parser.add_argument("--split-bank-num", action="store", type=int, - required=True, - help="Which splitbank do I read out of the HDF file?.") -parser.add_argument("--stack-distance", action="store", type=float,\ - default=0.2, help="Minimum metric spacing before we stack"+\ - "OPTIONAL") -parser.add_argument("--threed-lattice", action="store_true", default=False,\ - help="Set this to use a 3D lattice. "+\ - "OPTIONAL") -parser.add_argument("--skip-vec4-depth",action="store_true", default=False,\ - help="Assume the 4th direction to have negligible depth. "+\ - "OPTIONAL") -parser.add_argument("--skip-vec5-depth",action="store_true", default=False,\ - help="Assume the 5th direction to have negligible depth. "+\ - "OPTIONAL") -parser.add_argument("--random-seed", action="store", type=int,\ - default=None, - help="""Random seed to use whenever the numpy random +parser.add_argument( + "--pn-order", + action="store", + type=str, + default=None, + help="Determines the PN order to use. Note that if you " + "placing a bank of non-spinning templates, any " + "spin-related terms in the metric will always " + "be zero. REQUIRED ARGUMENT: " + + "choices are: %s" % (pycbcValidOrdersHelpDescriptions), +) +parser.add_argument( + "--f0", + action="store", + type=float, + default=70.0, + help="f0 is used as a dynamic rescaling factor when " + "calculating the integrals used in metric " + "construction. IE. instead of integrating F(f) we " + "integrate F(f/f0) and then remove f0 after the fact." + "The default option should be fine here for most " + "applications. OPTIONAL" + "WARNING: If using ethinca calculation this must be " + "equal to f-low. UNITS=Hz", +) +parser.add_argument( + "-I", + "--input-file", + action="store", + type=str, + required=True, + help="HDF file containing the input bank and the " + "various metric information. " + "REQUIRED ARGUMENT.", +) +parser.add_argument( + "--split-bank-num", + action="store", + type=int, + required=True, + help="Which splitbank do I read out of the HDF file?.", +) +parser.add_argument( + "--stack-distance", + action="store", + type=float, + default=0.2, + help="Minimum metric spacing before we stack" + "OPTIONAL", +) +parser.add_argument( + "--threed-lattice", + action="store_true", + default=False, + help="Set this to use a 3D lattice. " + "OPTIONAL", +) +parser.add_argument( + "--skip-vec4-depth", + action="store_true", + default=False, + help="Assume the 4th direction to have negligible depth. " + "OPTIONAL", +) +parser.add_argument( + "--skip-vec5-depth", + action="store_true", + default=False, + help="Assume the 5th direction to have negligible depth. " + "OPTIONAL", +) +parser.add_argument( + "--random-seed", + action="store", + type=int, + default=None, + help="""Random seed to use whenever the numpy random functions are called when doing the monte-carlo for obtaining the principal components and when translating all points back to physical space. If this is used the code should give the same - output if run with the same random seed.""") + output if run with the same random seed.""", +) pycbc.tmpltbank.insert_base_bank_options(parser) @@ -112,7 +149,7 @@ if not opts.pn_order: parser.error("Must supply --pn-order") opts.max_mismatch = 1 - opts.min_match pycbc.tmpltbank.verify_mass_range_options(opts, parser) -massRangeParams=pycbc.tmpltbank.massRangeParameters.from_argparse(opts) +massRangeParams = pycbc.tmpltbank.massRangeParameters.from_argparse(opts) # Set random seed if needed if opts.random_seed is not None: numpy.random.seed(opts.random_seed) @@ -122,24 +159,25 @@ if opts.random_seed is not None: # function and are not used. f0 is used in place of fUpper, but this is only # used as a key so doesn't matter. We have, however, removed information # that might be needed for something like ethinca. Maybe restore this? -metricParams = pycbc.tmpltbank.metricParameters(opts.pn_order, 0, opts.f0, \ - 0, f0=opts.f0) +metricParams = pycbc.tmpltbank.metricParameters( + opts.pn_order, 0, opts.f0, 0, f0=opts.f0 +) # Load the list of points from file -h5file = HFile(opts.input_file, 'r') -v1s = h5file['split_banks/split_bank_%05d/v1s' % opts.split_bank_num][:] -v2s = h5file['split_banks/split_bank_%05d/v2s' % opts.split_bank_num][:] -if 'v3s' in h5file['split_banks/split_bank_%05d' % opts.split_bank_num].keys(): - v3s = h5file['split_banks/split_bank_%05d/v3s' % opts.split_bank_num][:] - temp_bank = numpy.array([v1s,v2s,v3s]).T +h5file = HFile(opts.input_file, "r") +v1s = h5file["split_banks/split_bank_%05d/v1s" % opts.split_bank_num][:] +v2s = h5file["split_banks/split_bank_%05d/v2s" % opts.split_bank_num][:] +if "v3s" in h5file["split_banks/split_bank_%05d" % opts.split_bank_num].keys(): + v3s = h5file["split_banks/split_bank_%05d/v3s" % opts.split_bank_num][:] + temp_bank = numpy.array([v1s, v2s, v3s]).T else: - temp_bank = numpy.array([v1s,v2s]).T + temp_bank = numpy.array([v1s, v2s]).T # Load the files giving the information needed to define the xi_i # parameter space -evals = h5file['metric_evals'] -evecs = h5file['metric_evecs'] -evecsCV = h5file['cov_evecs'] +evals = h5file["metric_evals"] +evecs = h5file["metric_evecs"] +evecsCV = h5file["cov_evecs"] metricParams.evals = {} metricParams.evecs = {} @@ -157,12 +195,14 @@ if len(evals) < 5: # Create a large set of points and map to xi_i to give a starting point when # mapping from xi_i to masses and spins -rMass1, rMass2, rSpin1z, rSpin2z = \ - pycbc.tmpltbank.get_random_mass(2000000, massRangeParams) +rMass1, rMass2, rSpin1z, rSpin2z = pycbc.tmpltbank.get_random_mass( + 2000000, massRangeParams +) rTotmass, rEta = pnutils.mass1_mass2_to_mtotal_eta(rMass1, rMass2) -rXis = pycbc.tmpltbank.get_cov_params(rMass1, rMass2, rSpin1z, rSpin2z, - metricParams, opts.f0) +rXis = pycbc.tmpltbank.get_cov_params( + rMass1, rMass2, rSpin1z, rSpin2z, metricParams, opts.f0 +) xis = (numpy.array(rXis)).T physMasses = numpy.array([rTotmass, rEta, rSpin1z, rSpin2z]) @@ -185,52 +225,74 @@ numtemps = len(temp_bank) for entry in temp_bank: temp_number += 1 logging.info("Analysing template %d" % temp_number) - # First find the closest point in our set of 2000000 defined above + # First find the closest point in our set of 2000000 defined above # This is used as the starting point xi1_des = entry[0] xi2_des = entry[1] - xis_des = [xi1_des,xi2_des] + xis_des = [xi1_des, xi2_des] if opts.threed_lattice: xi3_des = entry[2] xis_des.append(xi3_des) req_match = 0.0001 - dist = (xi1_des - xis[:,0])**2 + (xi2_des - xis[:,1])**2 + dist = (xi1_des - xis[:, 0]) ** 2 + (xi2_des - xis[:, 1]) ** 2 if opts.threed_lattice: - dist += (xi3_des - xis[:,2])**2 + dist += (xi3_des - xis[:, 2]) ** 2 xis_close = xis[dist < 0.03] masses_close = physMasses[dist < 0.03] bestMasses = physMasses[dist.argmin()] bestXis = xis[dist.argmin()] - logging.info("Template %d has initial distance of %e" \ - % (temp_number, dist.min())) + logging.info("Template %d has initial distance of %e" % (temp_number, dist.min())) # Reject point if it is too far away from *any* of these points - if dist.min() > 2.: + if dist.min() > 2.0: logging.info("Template %d rejected as too far away" % temp_number) # Print info to the rejected points file if opts.threed_lattice: - reject_info.append([xi1_des, xi2_des, xi3_des, 0, 0, - 0, 0, dist.min()]) + reject_info.append([xi1_des, xi2_des, xi3_des, 0, 0, 0, 0, dist.min()]) else: reject_info.append([xi1_des, xi2_des, 0, 0, 0, 0, dist.min()]) continue # This function will use the starting point and iteratively find a # physical point has a mismatch < 0.0001 with the desired one - masses = pycbc.tmpltbank.get_physical_covaried_masses(xis_des,\ - copy.deepcopy(bestMasses), copy.deepcopy(bestXis), req_match,\ - massRangeParams, metricParams, opts.f0) + masses = pycbc.tmpltbank.get_physical_covaried_masses( + xis_des, + copy.deepcopy(bestMasses), + copy.deepcopy(bestXis), + req_match, + massRangeParams, + metricParams, + opts.f0, + ) # Now how close is it? - logging.info("Template %d has corrected distance of %e" \ - % (temp_number, masses[5])) + logging.info("Template %d has corrected distance of %e" % (temp_number, masses[5])) if masses[5] > opts.max_mismatch: # Reject point, it is too far away logging.info("Template %d rejected as too far away" % temp_number) if opts.threed_lattice: - reject_info.append([xi1_des, xi2_des, xi3_des, masses[0], - masses[1], masses[2], masses[3], masses[5]]) + reject_info.append( + [ + xi1_des, + xi2_des, + xi3_des, + masses[0], + masses[1], + masses[2], + masses[3], + masses[5], + ] + ) else: - reject_info.append([xi1_des, xi2_des, masses[0], masses[1], - masses[2], masses[3], masses[5]]) + reject_info.append( + [ + xi1_des, + xi2_des, + masses[0], + masses[1], + masses[2], + masses[3], + masses[5], + ] + ) continue # If we got this far the point will be accepted. # Now we figure out if the depth of the *other* directions are wide enough @@ -238,30 +300,37 @@ for entry in temp_bank: # We begin by evaluating the depth of the third direction, this is not # needed if a 3D lattice is being employed tmpTotMass = masses[0] + masses[1] - tmpEta = masses[0] * masses[1] / (tmpTotMass*tmpTotMass) + tmpEta = masses[0] * masses[1] / (tmpTotMass * tmpTotMass) if not opts.threed_lattice: # If point is close enough, determine depth of xi3 direction - vec3_min, vec3_max=\ - pycbc.tmpltbank.stack_xi_direction_brute(\ - [masses[6][0],masses[6][1]],\ - [tmpTotMass,tmpEta,masses[2],masses[3]],\ - copy.deepcopy(bestXis), 2, opts.max_mismatch,\ - massRangeParams, metricParams, opts.f0) + vec3_min, vec3_max = pycbc.tmpltbank.stack_xi_direction_brute( + [masses[6][0], masses[6][1]], + [tmpTotMass, tmpEta, masses[2], masses[3]], + copy.deepcopy(bestXis), + 2, + opts.max_mismatch, + massRangeParams, + metricParams, + opts.f0, + ) vec3_depth = vec3_max - vec3_min # Double check that no points appear outside what was calculated above if len(xis_close): - if vec3_min > xis_close[:,2].min(): - logging.warning( - "WARNING: Numerical placement fails, trying again" - ) - temp_idx = xis_close[:,2].argmin() + if vec3_min > xis_close[:, 2].min(): + logging.warning("WARNING: Numerical placement fails, trying again") + temp_idx = xis_close[:, 2].argmin() temmpBestMasses = masses_close[temp_idx] temmpBestXis = xis_close[temp_idx] - temmpvec3_min, temmpvec3_max =\ - pycbc.tmpltbank.stack_xi_direction_brute([xi1_des,xi2_des],\ - copy.deepcopy(temmpBestMasses),\ - copy.deepcopy(temmpBestXis), 2, opts.max_mismatch, \ - massRangeParams, metricParams, opts.f0) + temmpvec3_min, temmpvec3_max = pycbc.tmpltbank.stack_xi_direction_brute( + [xi1_des, xi2_des], + copy.deepcopy(temmpBestMasses), + copy.deepcopy(temmpBestXis), + 2, + opts.max_mismatch, + massRangeParams, + metricParams, + opts.f0, + ) temmpvec3_depth = temmpvec3_max - temmpvec3_min if temmpvec3_min < vec3_min: vec3_min = temmpvec3_min @@ -269,18 +338,21 @@ for entry in temp_bank: if temmpvec3_max > vec3_max: vec3_max = temmpvec3_max vec3_depth = vec3_max - vec3_min - if vec3_max < xis_close[:,2].max(): - logging.warning( - "WARNING: Numerical placement fails, trying again" - ) - temp_idx = xis_close[:,2].argmax() + if vec3_max < xis_close[:, 2].max(): + logging.warning("WARNING: Numerical placement fails, trying again") + temp_idx = xis_close[:, 2].argmax() temmpBestMasses = physMasses[temp_idx] temmpBestXis = xis[temp_idx] - temmpvec3_min, temmpvec3_max =\ - pycbc.tmpltbank.stack_xi_direction_brute([xi1_des,xi2_des],\ - copy.deepcopy(temmpBestMasses),\ - copy.deepcopy(temmpBestXis), 2, opts.max_mismatch, \ - massRangeParams, metricParams, opts.f0) + temmpvec3_min, temmpvec3_max = pycbc.tmpltbank.stack_xi_direction_brute( + [xi1_des, xi2_des], + copy.deepcopy(temmpBestMasses), + copy.deepcopy(temmpBestXis), + 2, + opts.max_mismatch, + massRangeParams, + metricParams, + opts.f0, + ) temmpvec3_depth = temmpvec3_max - temmpvec3_min if temmpvec3_max > vec3_max: vec3_max = temmpvec3_max @@ -288,26 +360,34 @@ for entry in temp_bank: if temmpvec3_min < vec3_min: vec3_min = temmpvec3_min vec3_depth = vec3_max - vec3_min - # Determine depth of xi4 direction (is this needed is xi3 depth is found to - # be small?) + # Determine depth of xi4 direction (is this needed is xi3 depth is found to + # be small?) if opts.eval_vec4_depth: - vec4_min, vec4_max =\ - pycbc.tmpltbank.stack_xi_direction_brute(\ - [masses[6][0],masses[6][1]],\ - [tmpTotMass,tmpEta,masses[2],masses[3]],\ - copy.deepcopy(bestXis), 3, opts.max_mismatch,\ - massRangeParams, metricParams, opts.f0) + vec4_min, vec4_max = pycbc.tmpltbank.stack_xi_direction_brute( + [masses[6][0], masses[6][1]], + [tmpTotMass, tmpEta, masses[2], masses[3]], + copy.deepcopy(bestXis), + 3, + opts.max_mismatch, + massRangeParams, + metricParams, + opts.f0, + ) vec4_depth = vec4_max - vec4_min else: vec4_min = vec4_max = vec4_depth = 0 - # Determine depth of xi5 direction + # Determine depth of xi5 direction if opts.eval_vec5_depth: - vec5_min, vec5_max =\ - pycbc.tmpltbank.stack_xi_direction_brute(\ - [masses[6][0],masses[6][1]],\ - [tmpTotMass,tmpEta,masses[2],masses[3]],\ - copy.deepcopy(bestXis), 4, opts.max_mismatch,\ - massRangeParams, metricParams, opts.f0) + vec5_min, vec5_max = pycbc.tmpltbank.stack_xi_direction_brute( + [masses[6][0], masses[6][1]], + [tmpTotMass, tmpEta, masses[2], masses[3]], + copy.deepcopy(bestXis), + 4, + opts.max_mismatch, + massRangeParams, + metricParams, + opts.f0, + ) vec5_depth = vec5_max - vec5_min else: vec5_min = vec5_max = vec5_depth = 0 @@ -315,9 +395,8 @@ for entry in temp_bank: if opts.threed_lattice: depths_info.append([xi1_des, xi2_des, xi3_des, vec4_depth, vec5_depth]) else: - depths_info.append([xi1_des, xi2_des, vec3_depth, vec4_depth, - vec5_depth]) - # Figure out how many templates we need to stack in 3rd direction + depths_info.append([xi1_des, xi2_des, vec3_depth, vec4_depth, vec5_depth]) + # Figure out how many templates we need to stack in 3rd direction vec3DepthVal = opts.stack_distance if opts.threed_lattice: numV3Temps = 1 @@ -325,54 +404,77 @@ for entry in temp_bank: numV3Temps = int(round(vec3_depth // vec3DepthVal)) + 1 for ite in range(numV3Temps): if not opts.threed_lattice: - xi3_des = vec3_min + \ - (vec3_depth) * (2 * ite + 1) / (2. * (numV3Temps)) - dist = (xi1_des - xis[:,0])**2 + (xi2_des - xis[:,1])**2 + \ - (xi3_des - xis[:,2])**2 + xi3_des = vec3_min + (vec3_depth) * (2 * ite + 1) / (2.0 * (numV3Temps)) + dist = ( + (xi1_des - xis[:, 0]) ** 2 + + (xi2_des - xis[:, 1]) ** 2 + + (xi3_des - xis[:, 2]) ** 2 + ) bestMasses = physMasses[dist.argmin()] bestXis = xis[dist.argmin()] # Find close point to this 3d position - masses = pycbc.tmpltbank.get_physical_covaried_masses(\ - [xi1_des,xi2_des,xi3_des], copy.deepcopy(bestMasses),\ - copy.deepcopy(bestXis), req_match,\ - massRangeParams, metricParams, opts.f0) + masses = pycbc.tmpltbank.get_physical_covaried_masses( + [xi1_des, xi2_des, xi3_des], + copy.deepcopy(bestMasses), + copy.deepcopy(bestXis), + req_match, + massRangeParams, + metricParams, + opts.f0, + ) # If vec4 depth is negligible or we didn't get close, stop here if vec4_depth < vec3DepthVal or masses[5] > opts.max_mismatch: if masses[5]: # Write point to file - points_info.append([masses[0], masses[1], masses[2], masses[3], - masses[5]]) + points_info.append( + [masses[0], masses[1], masses[2], masses[3], masses[5]] + ) else: # OR we need to estimate the depth in the 4th direction at this # 3d point tmpTotMass = masses[0] + masses[1] - tmpEta = masses[0] * masses[1] / (tmpTotMass*tmpTotMass) - vec4_minT, vec4_maxT = \ - pycbc.tmpltbank.stack_xi_direction_brute(\ - [masses[6][0],masses[6][1],masses[6][2]],\ - [tmpTotMass,tmpEta,masses[2],masses[3]],\ - copy.deepcopy(bestXis), 3, opts.max_mismatch,\ - massRangeParams, metricParams, opts.f0) + tmpEta = masses[0] * masses[1] / (tmpTotMass * tmpTotMass) + vec4_minT, vec4_maxT = pycbc.tmpltbank.stack_xi_direction_brute( + [masses[6][0], masses[6][1], masses[6][2]], + [tmpTotMass, tmpEta, masses[2], masses[3]], + copy.deepcopy(bestXis), + 3, + opts.max_mismatch, + massRangeParams, + metricParams, + opts.f0, + ) vec4_depthT = vec4_maxT - vec4_minT # Then loop over necessary templates in 4th direction numV4Temps = int(round(vec4_depthT // vec3DepthVal)) + 1 for ite in range(numV4Temps): - xi4_des = vec4_minT + \ - (vec4_depthT) * (2 * ite + 1) / (2. * (numV4Temps)) - dist = (xi1_des - xis[:,0])**2 + (xi2_des - xis[:,1])**2 + \ - (xi3_des - xis[:,2])**2 + (xi4_des - xis[:,3])**2 + xi4_des = vec4_minT + (vec4_depthT) * (2 * ite + 1) / ( + 2.0 * (numV4Temps) + ) + dist = ( + (xi1_des - xis[:, 0]) ** 2 + + (xi2_des - xis[:, 1]) ** 2 + + (xi3_des - xis[:, 2]) ** 2 + + (xi4_des - xis[:, 3]) ** 2 + ) bestMasses = physMasses[dist.argmin()] bestXis = xis[dist.argmin()] - masses = pycbc.tmpltbank.get_physical_covaried_masses(\ - [xi1_des,xi2_des,xi3_des,xi4_des],\ - copy.deepcopy(bestMasses), copy.deepcopy(bestXis), \ - req_match, massRangeParams, metricParams, opts.f0) + masses = pycbc.tmpltbank.get_physical_covaried_masses( + [xi1_des, xi2_des, xi3_des, xi4_des], + copy.deepcopy(bestMasses), + copy.deepcopy(bestXis), + req_match, + massRangeParams, + metricParams, + opts.f0, + ) if masses[5]: - points_info.append([masses[0], masses[1], masses[2], - masses[3], masses[5]]) + points_info.append( + [masses[0], masses[1], masses[2], masses[3], masses[5]] + ) -h5outfile = HFile(opts.output_file, 'w') -h5outfile['reject_points'] = numpy.array(reject_info) -h5outfile['point_depths'] = numpy.array(depths_info) -h5outfile['accepted_templates'] = numpy.array(points_info) +h5outfile = HFile(opts.output_file, "w") +h5outfile["reject_points"] = numpy.array(reject_info) +h5outfile["point_depths"] = numpy.array(depths_info) +h5outfile["accepted_templates"] = numpy.array(points_info) h5outfile.close() diff --git a/bin/bank/pycbc_geom_aligned_bank b/bin/bank/pycbc_geom_aligned_bank index 48fa2a25fd1..cf9f3801c66 100644 --- a/bin/bank/pycbc_geom_aligned_bank +++ b/bin/bank/pycbc_geom_aligned_bank @@ -23,73 +23,109 @@ translation between points in the chirp parameters parameter space is done in a workflow. """ -import copy import argparse -import numpy +import copy import logging -from scipy import spatial +import numpy +import pycbc.version from igwn_ligolw import ligolw from igwn_ligolw import utils as ligolw_utils +from scipy import spatial import pycbc import pycbc.psd import pycbc.strain -import pycbc.version import pycbc.tmpltbank - -from pycbc.io.ligolw import create_process_table from pycbc.io import HFile +from pycbc.io.ligolw import create_process_table -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_geom_aligned_bank" # Read command line options _desc = __doc__[1:] -parser = argparse.ArgumentParser(description=_desc, - formatter_class=pycbc.tmpltbank.IndentedHelpFormatterWithNL) +parser = argparse.ArgumentParser( + description=_desc, formatter_class=pycbc.tmpltbank.IndentedHelpFormatterWithNL +) # Begin with code specific options pycbc.add_common_pycbc_options(parser) -parser.add_argument("-s", "--stack-distance", action="store", type=float,\ - default=0.2, help="Minimum metric spacing before we "+\ - "stack.") -parser.add_argument("-3", "--threed-lattice", action="store_true", default=False,\ - help="Set this to use a 3D lattice. "+\ - "OPTIONAL") -parser.add_argument("-S", "--num-split-jobs", action="store", type=int,\ - default=100,\ - help="Number of parallel jobs to split the bank into. "+\ - "OPTIONAL") -parser.add_argument("-F", "--filter-points", action="store_true", default=False, - help="Remove nearby points before generating the bank.") -parser.add_argument("--random-seed", action="store", type=int,\ - default=None, - help="""Random seed to use whenever the numpy random +parser.add_argument( + "-s", + "--stack-distance", + action="store", + type=float, + default=0.2, + help="Minimum metric spacing before we " + "stack.", +) +parser.add_argument( + "-3", + "--threed-lattice", + action="store_true", + default=False, + help="Set this to use a 3D lattice. " + "OPTIONAL", +) +parser.add_argument( + "-S", + "--num-split-jobs", + action="store", + type=int, + default=100, + help="Number of parallel jobs to split the bank into. " + "OPTIONAL", +) +parser.add_argument( + "-F", + "--filter-points", + action="store_true", + default=False, + help="Remove nearby points before generating the bank.", +) +parser.add_argument( + "--random-seed", + action="store", + type=int, + default=None, + help="""Random seed to use whenever the numpy random functions are called when doing the monte-carlo for obtaining the principal components and when translating all points back to physical space. If this is used the code should give the same - output if run with the same random seed.""") -parser.add_argument("--print-chi-points", action="store", default=None, - metavar="FILENAME", - help="Add a node to print off an ASCII list of mass " - "parameters and corresponding location in the xi space " - "using pycbc_tmpltbank_to_chi_params. This will be " - "written to FILENAME. If this argument is not given, no " - "chi points file will be written.") -parser.add_argument("--metadata-file", type=str, required=True, - help="Location of the output file containing the metadata " - "that will be added to the final XML file.") -parser.add_argument("--storage-path-base", default=None, - help="If running this code as a sub-workflow then this " - "path is pretended to all storage directories.") -parser.add_argument("--supplement-config-file", action="append", - help="This can be used to add additional options to those " - "that this code will supply. If there are conflicts the " - "code will fail. Can be supplied multiple times.") + output if run with the same random seed.""", +) +parser.add_argument( + "--print-chi-points", + action="store", + default=None, + metavar="FILENAME", + help="Add a node to print off an ASCII list of mass " + "parameters and corresponding location in the xi space " + "using pycbc_tmpltbank_to_chi_params. This will be " + "written to FILENAME. If this argument is not given, no " + "chi points file will be written.", +) +parser.add_argument( + "--metadata-file", + type=str, + required=True, + help="Location of the output file containing the metadata " + "that will be added to the final XML file.", +) +parser.add_argument( + "--storage-path-base", + default=None, + help="If running this code as a sub-workflow then this " + "path is pretended to all storage directories.", +) +parser.add_argument( + "--supplement-config-file", + action="append", + help="This can be used to add additional options to those " + "that this code will supply. If there are conflicts the " + "code will fail. Can be supplied multiple times.", +) pycbc.tmpltbank.insert_base_bank_options(parser) @@ -126,9 +162,9 @@ pycbc.init_logging(opts.verbose) opts.max_mismatch = 1 - opts.min_match pycbc.tmpltbank.verify_metric_calculation_options(opts, parser) -metricParams=pycbc.tmpltbank.metricParameters.from_argparse(opts) +metricParams = pycbc.tmpltbank.metricParameters.from_argparse(opts) pycbc.tmpltbank.verify_mass_range_options(opts, parser) -massRangeParams=pycbc.tmpltbank.massRangeParameters.from_argparse(opts) +massRangeParams = pycbc.tmpltbank.massRangeParameters.from_argparse(opts) pycbc.psd.verify_psd_options(opts, parser) if opts.psd_estimation: pycbc.strain.verify_strain_options(opts, parser) @@ -137,10 +173,10 @@ ethincaParams = pycbc.tmpltbank.ethincaParameters.from_argparse(opts) # Ensure consistency of ethinca and bank metric parameters pycbc.tmpltbank.check_ethinca_against_bank_params(ethincaParams, metricParams) # Ethinca calculation should currently only be done for non-spin templates -if ethincaParams.full_ethinca and (massRangeParams.maxNSSpinMag>0.0 or - massRangeParams.maxBHSpinMag>0.0): - parser.error("Ethinca metric calculation is currently not valid for " - "nonzero spins!") +if ethincaParams.full_ethinca and ( + massRangeParams.maxNSSpinMag > 0.0 or massRangeParams.maxBHSpinMag > 0.0 +): + parser.error("Ethinca metric calculation is currently not valid for nonzero spins!") # Set random seed if needed if opts.random_seed is not None: @@ -158,11 +194,16 @@ else: logging.info("Obtaining PSD") # Want the number of samples to be a binary number and Nyquist must be above # opts.f_upper. All this assumes that 1 / deltaF is a binary number -nyquistFreq = 2**numpy.ceil(numpy.log2(opts.f_upper)) +nyquistFreq = 2 ** numpy.ceil(numpy.log2(opts.f_upper)) numSamples = int(round(nyquistFreq / opts.delta_f)) + 1 -psd = pycbc.psd.from_cli(opts, length=numSamples, delta_f=opts.delta_f, \ - low_frequency_cutoff=opts.f_low, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC) +psd = pycbc.psd.from_cli( + opts, + length=numSamples, + delta_f=opts.delta_f, + low_frequency_cutoff=opts.f_low, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, +) metricParams.psd = psd # Begin by calculating a metric @@ -171,8 +212,9 @@ metricParams = pycbc.tmpltbank.determine_eigen_directions(metricParams) logging.info("Calculating covariance matrix") -vals = pycbc.tmpltbank.estimate_mass_range(1000000, massRangeParams, \ - metricParams, metricParams.fUpper, covary=False) +vals = pycbc.tmpltbank.estimate_mass_range( + 1000000, massRangeParams, metricParams, metricParams.fUpper, covary=False +) cov = numpy.cov(vals) evalsCV, evecsCV = numpy.linalg.eig(cov) evecsCVdict = {} @@ -182,8 +224,9 @@ metricParams.evecsCV = evecsCVdict logging.info("Determining parameter space extent") -vals = pycbc.tmpltbank.estimate_mass_range(1000000, massRangeParams, \ - metricParams, metricParams.fUpper, covary=True) +vals = pycbc.tmpltbank.estimate_mass_range( + 1000000, massRangeParams, metricParams, metricParams.fUpper, covary=True +) chi1Max = vals[0].max() chi1Min = vals[0].min() @@ -195,18 +238,26 @@ chi2Diff = chi2Max - chi2Min logging.info("Calculating lattice") if not opts.threed_lattice: - v1s,v2s = pycbc.tmpltbank.generate_hexagonal_lattice(\ - chi1Max+(0.02*chi1Diff), chi1Min-(0.02*chi1Diff),\ - chi2Max+(0.02*chi2Diff), chi2Min-(0.02*chi2Diff),\ - opts.max_mismatch) + v1s, v2s = pycbc.tmpltbank.generate_hexagonal_lattice( + chi1Max + (0.02 * chi1Diff), + chi1Min - (0.02 * chi1Diff), + chi2Max + (0.02 * chi2Diff), + chi2Min - (0.02 * chi2Diff), + opts.max_mismatch, + ) else: chi3Max = vals[2].max() chi3Min = vals[2].min() chi3Diff = chi3Max - chi3Min - v1s, v2s, v3s = pycbc.tmpltbank.generate_anstar_3d_lattice(\ - chi1Max+(0.02*chi1Diff), chi1Min-(0.02*chi1Diff),\ - chi2Max+(0.02*chi2Diff), chi2Min-(0.02*chi2Diff),\ - chi3Max+(0.02*chi3Diff), chi3Min-(0.02*chi3Diff), opts.max_mismatch) + v1s, v2s, v3s = pycbc.tmpltbank.generate_anstar_3d_lattice( + chi1Max + (0.02 * chi1Diff), + chi1Min - (0.02 * chi1Diff), + chi2Max + (0.02 * chi2Diff), + chi2Min - (0.02 * chi2Diff), + chi3Max + (0.02 * chi3Diff), + chi3Min - (0.02 * chi3Diff), + opts.max_mismatch, + ) chi3Max = vals[2].max() chi3Min = vals[2].min() chi3Diff = chi3Max - chi3Min @@ -220,10 +271,13 @@ if opts.filter_points: # Create a large set of points and map to xi_i to give a starting point when # mapping from xi_i to masses and spins # Use the EM constraint only if asked to do so - rMass1, rMass2, rSpin1z, rSpin2z = \ - pycbc.tmpltbank.get_random_mass(2000000, massRangeParams) + rMass1, rMass2, rSpin1z, rSpin2z = pycbc.tmpltbank.get_random_mass( + 2000000, massRangeParams + ) - rXis = pycbc.tmpltbank.get_cov_params(rMass1, rMass2, rSpin1z, rSpin2z, metricParams, metricParams.fUpper) + rXis = pycbc.tmpltbank.get_cov_params( + rMass1, rMass2, rSpin1z, rSpin2z, metricParams, metricParams.fUpper + ) xis = (numpy.array(rXis)).T f0 = opts.f0 @@ -243,11 +297,11 @@ if opts.filter_points: # Use scipy's KDtree to quickly calculate Euclidean distances logging.info("Setting up KDtree to compute distances.") if opts.threed_lattice: - tree = spatial.KDTree(xis[:,:3]) - xi_points = list(zip(v1s,v2s,v3s)) + tree = spatial.KDTree(xis[:, :3]) + xi_points = list(zip(v1s, v2s, v3s)) else: - tree = spatial.KDTree(xis[:,:2]) - xi_points = list(zip(v1s,v2s)) + tree = spatial.KDTree(xis[:, :2]) + xi_points = list(zip(v1s, v2s)) logging.info("Computing distances using KDtree.") dists, pointargs = tree.query(xi_points) @@ -255,7 +309,7 @@ if opts.filter_points: logging.info("Removing far-away points.") for i in range(len(v1s)): - if dists[i] < 2.: + if dists[i] < 2.0: newV1s.append(v1s[i]) newV2s.append(v2s[i]) if opts.threed_lattice: @@ -270,12 +324,12 @@ else: # Now begin to generate the dag -h5file = HFile(opts.output_file, 'w') +h5file = HFile(opts.output_file, "w") # Dump the full bank in \xi_i coordinates -h5file['full_bank/v1s'] = newV1s -h5file['full_bank/v2s'] = newV2s +h5file["full_bank/v1s"] = newV1s +h5file["full_bank/v2s"] = newV2s if opts.threed_lattice: - h5file['full_bank/v3s'] = newV3s + h5file["full_bank/v3s"] = newV3s # Now store split banks bank_num = 0 @@ -285,33 +339,37 @@ v1s = [] v2s = [] v3s = [] total_points = len(newV1s) -points_per_job = int(numpy.ceil(total_points / float(opts.num_split_jobs))) if total_points > 0 and opts.num_split_jobs > 0 else 1 +points_per_job = ( + int(numpy.ceil(total_points / float(opts.num_split_jobs))) + if total_points > 0 and opts.num_split_jobs > 0 + else 1 +) for i in range(len(newV1s)): v1s.append(newV1s[i]) v2s.append(newV2s[i]) if opts.threed_lattice: v3s.append(newV3s[i]) - if not (i+1) % points_per_job: - h5file['split_banks/split_bank_%05d/v1s' % bank_num] = v1s - h5file['split_banks/split_bank_%05d/v2s' % bank_num] = v2s + if not (i + 1) % points_per_job: + h5file["split_banks/split_bank_%05d/v1s" % bank_num] = v1s + h5file["split_banks/split_bank_%05d/v2s" % bank_num] = v2s v1s = [] v2s = [] if opts.threed_lattice: - h5file['split_banks/split_bank_%05d/v3s' % bank_num] = v3s + h5file["split_banks/split_bank_%05d/v3s" % bank_num] = v3s v3s = [] bank_num = bank_num + 1 -if len(v1s): +if v1s: # There are still templates to dump - h5file['split_banks/split_bank_%05d/v1s' % bank_num] = v1s - h5file['split_banks/split_bank_%05d/v2s' % bank_num] = v2s + h5file["split_banks/split_bank_%05d/v1s" % bank_num] = v1s + h5file["split_banks/split_bank_%05d/v2s" % bank_num] = v2s if opts.threed_lattice: - h5file['split_banks/split_bank_%05d/v3s' % bank_num] = v3s + h5file["split_banks/split_bank_%05d/v3s" % bank_num] = v3s -logging.info('Adding metric information to HDF file.') -h5file['cov_evecs'] = metricParams.evecsCV[metricParams.fUpper] -h5file['metric_evals'] = metricParams.evals[metricParams.fUpper] -h5file['metric_evecs'] = metricParams.evecs[metricParams.fUpper] +logging.info("Adding metric information to HDF file.") +h5file["cov_evecs"] = metricParams.evecsCV[metricParams.fUpper] +h5file["metric_evals"] = metricParams.evals[metricParams.fUpper] +h5file["metric_evecs"] = metricParams.evecs[metricParams.fUpper] h5file.close() diff --git a/bin/bank/pycbc_geom_nonspinbank b/bin/bank/pycbc_geom_nonspinbank index b778f7fafb0..e9f369e0da4 100644 --- a/bin/bank/pycbc_geom_nonspinbank +++ b/bin/bank/pycbc_geom_nonspinbank @@ -21,42 +21,46 @@ Template bank generator for placing a bank of non-spinning templates. """ import argparse -import math import copy -import numpy import logging +import math -import pycbc +import numpy import pycbc.version + +import pycbc import pycbc.psd import pycbc.strain -from pycbc import tmpltbank + # Old ligolw output functions no longer imported at package level import pycbc.tmpltbank.bank_output_utils as bank_output -from pycbc import pnutils - +from pycbc import pnutils, tmpltbank -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_geom_nonspinbank" # Read command line options _desc = __doc__[1:] parser = argparse.ArgumentParser( - description=_desc, - formatter_class=tmpltbank.IndentedHelpFormatterWithNL) + description=_desc, formatter_class=tmpltbank.IndentedHelpFormatterWithNL +) # Begin with code specific options pycbc.add_common_pycbc_options(parser) -parser.add_argument("--random-seed", action="store", type=int, - default=None, - help="""Random seed to use when calling numpy.random +parser.add_argument( + "--random-seed", + action="store", + type=int, + default=None, + help="""Random seed to use when calling numpy.random functions used in obtaining the principal components and when translating points back to physical space. If given, the code should give the same output - when run with the same random seed.""") + when run with the same random seed.""", +) tmpltbank.insert_base_bank_options(parser) @@ -81,14 +85,14 @@ pycbc.init_logging(opts.verbose) opts.max_mismatch = 1 - opts.min_match tmpltbank.verify_metric_calculation_options(opts, parser) -metricParams=tmpltbank.metricParameters.from_argparse(opts) +metricParams = tmpltbank.metricParameters.from_argparse(opts) tmpltbank.verify_mass_range_options(opts, parser, nonSpin=True) -massRangeParams=tmpltbank.massRangeParameters.from_argparse(opts,nonSpin=True) +massRangeParams = tmpltbank.massRangeParameters.from_argparse(opts, nonSpin=True) pycbc.psd.verify_psd_options(opts, parser) if opts.psd_estimation: pycbc.strain.verify_strain_options(opts, parser) tmpltbank.verify_ethinca_metric_options(opts, parser) -ethincaParams=tmpltbank.ethincaParameters.from_argparse(opts) +ethincaParams = tmpltbank.ethincaParameters.from_argparse(opts) # Ensure consistency of ethinca and bank metric parameters tmpltbank.check_ethinca_against_bank_params(ethincaParams, metricParams) @@ -109,27 +113,33 @@ logging.info("Obtaining PSD") # FIXME: Revisit the following! # Want the number of samples to be a binary number and Nyquist must be above # opts.f_upper. All this assumes that 1 / deltaF is a binary number -nyquistFreq = 2**numpy.ceil(numpy.log2(opts.f_upper)) +nyquistFreq = 2 ** numpy.ceil(numpy.log2(opts.f_upper)) numSamples = int(round(nyquistFreq / opts.delta_f)) + 1 -psd = pycbc.psd.from_cli(opts, length=numSamples, delta_f=opts.delta_f, - low_frequency_cutoff=opts.f_low, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC) +psd = pycbc.psd.from_cli( + opts, + length=numSamples, + delta_f=opts.delta_f, + low_frequency_cutoff=opts.f_low, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, +) metricParams.psd = psd logging.info("Calculating metric") # Begin by calculating a metric metricParams = tmpltbank.determine_eigen_directions( - metricParams, vary_fmax=ethincaParams.doEthinca, - vary_density=ethincaParams.freqStep) + metricParams, vary_fmax=ethincaParams.doEthinca, vary_density=ethincaParams.freqStep +) # This is used to calculate evalsCV and evecsCV with describe the rotations # needed to move into the principal component directions. evalsCV will be 1s. logging.info("Calculating covariance matrix") -vals = tmpltbank.estimate_mass_range(1000000, massRangeParams, metricParams, - metricParams.fUpper, covary=False) +vals = tmpltbank.estimate_mass_range( + 1000000, massRangeParams, metricParams, metricParams.fUpper, covary=False +) cov = numpy.cov(vals) evalsCV, evecsCV = numpy.linalg.eig(cov) evecsCVdict = {} @@ -140,8 +150,9 @@ metricParams.evecsCV = evecsCVdict logging.info("Estimating extent of parameter space") # This is to get an estimate of the largest values of \chi1 and \chi2 -vals = tmpltbank.estimate_mass_range(5000000, massRangeParams, metricParams, - metricParams.fUpper, covary=True) +vals = tmpltbank.estimate_mass_range( + 5000000, massRangeParams, metricParams, metricParams.fUpper, covary=True +) chi1Max = vals[0].max() chi1Min = vals[0].min() @@ -167,8 +178,8 @@ spacing1D = opts.max_mismatch**0.5 * 2 v1s = [] v2s = [] for i in range(100): - tempChi1Lower = chi1Min + i * chi1Diff/100. - tempChi1Upper = chi1Min + (i+1) * chi1Diff/100. + tempChi1Lower = chi1Min + i * chi1Diff / 100.0 + tempChi1Upper = chi1Min + (i + 1) * chi1Diff / 100.0 lgc = (vals[0] > tempChi1Lower) & (vals[0] < tempChi1Upper) tempChi1 = vals[0][lgc] tempChi2 = vals[1][lgc] @@ -176,17 +187,15 @@ for i in range(100): tempChi2Min = tempChi2.min() if (tempChi2Max - tempChi2Min) < 0.2: # 1D lattice through here - chi2Loc = (tempChi2Max + tempChi2Min) / 2. + chi2Loc = (tempChi2Max + tempChi2Min) / 2.0 currIter = 0 startPoint = chi1Min - 0.02 * chi1Diff - while(1): + while 1: currChi1 = startPoint + currIter * spacing1D if (currChi1 < tempChi1Lower) and (i != 0): currIter = currIter + 1 continue - elif (currChi1 > tempChi1Upper) and (i != 99): - break - elif (currChi1 > tempChi1Upper + 0.02*chi1Diff): + elif ((currChi1 > tempChi1Upper) and (i != 99)) or currChi1 > tempChi1Upper + 0.02 * chi1Diff: break v1s.append(currChi1) v2s.append(chi2Loc) @@ -197,14 +206,14 @@ for i in range(100): break # Next we start from the end and work backwards -for i in range(99,-1,-1): +for i in range(99, -1, -1): # Only need to do this if there is a lower boundary if lower1DBoundary is None: break # FIXME: Move this to a function, duplicated above! # FIXME: Maybe move all this to a generate_NS_lattice function - tempChi1Lower = chi1Min + i * chi1Diff/100. - tempChi1Upper = chi1Min + (i+1) * chi1Diff/100. + tempChi1Lower = chi1Min + i * chi1Diff / 100.0 + tempChi1Upper = chi1Min + (i + 1) * chi1Diff / 100.0 lgc = (vals[0] > tempChi1Lower) & (vals[0] < tempChi1Upper) tempChi1 = vals[0][lgc] tempChi2 = vals[1][lgc] @@ -212,17 +221,15 @@ for i in range(99,-1,-1): tempChi2Min = tempChi2.min() if (tempChi2Max - tempChi2Min) < 0.2: # 1D lattice through here - chi2Loc = (tempChi2Max + tempChi2Min) / 2. + chi2Loc = (tempChi2Max + tempChi2Min) / 2.0 currIter = 0 startPoint = chi1Min - 0.02 * chi1Diff - while(1): + while 1: currChi1 = startPoint + currIter * spacing1D if (currChi1 < tempChi1Lower) and (i != 0): currIter = currIter + 1 continue - elif (currChi1 > tempChi1Upper) and (i != 99): - break - elif (currChi1 > tempChi1Upper + 0.02*chi1Diff): + elif ((currChi1 > tempChi1Upper) and (i != 99)) or currChi1 > tempChi1Upper + 0.02 * chi1Diff: break v1s.append(currChi1) v2s.append(chi2Loc) @@ -233,25 +240,29 @@ for i in range(99,-1,-1): break # TESTING CODE: USE THIS TO TURN OFF THE 1D LATTICE, IE DO 2D EVERYWHERE -#lower1DBoundary = 0 -#upper1DBoundary = 100 -#v1s = [] -#v2s = [] +# lower1DBoundary = 0 +# upper1DBoundary = 100 +# v1s = [] +# v2s = [] # Anything left is covered with a 2D hexagonal array if lower1DBoundary is not None: # Need some 2D parts in the bank if lower1DBoundary == 0: - lowerChi1 = chi1Min - 0.02*chi1Diff + lowerChi1 = chi1Min - 0.02 * chi1Diff else: - lowerChi1 = chi1Min + lower1DBoundary*chi1Diff/100. + lowerChi1 = chi1Min + lower1DBoundary * chi1Diff / 100.0 if upper1DBoundary == 100: - upperChi1 = chi1Max + 0.02*chi1Diff + upperChi1 = chi1Max + 0.02 * chi1Diff else: - upperChi1 = chi1Min + upper1DBoundary*chi1Diff/100. + upperChi1 = chi1Min + upper1DBoundary * chi1Diff / 100.0 tempv1s, tempv2s = tmpltbank.generate_hexagonal_lattice( - upperChi1, lowerChi1, chi2Max + 0.02*chi2Diff, - chi2Min - 0.02*chi2Diff, opts.max_mismatch) + upperChi1, + lowerChi1, + chi2Max + 0.02 * chi2Diff, + chi2Min - 0.02 * chi2Diff, + opts.max_mismatch, + ) v1s.extend(tempv1s) v2s.extend(tempv2s) @@ -271,12 +282,12 @@ logging.info("%d points in the lattice", len(v1s)) # FIXME: Most of this should probably move to a function in pycbc.tmpltbank # Choose an initial set of points to begin comparing points to physical space -rMass1, rMass2, rSpin1z, rSpin2z = \ - tmpltbank.get_random_mass(2000000, massRangeParams) +rMass1, rMass2, rSpin1z, rSpin2z = tmpltbank.get_random_mass(2000000, massRangeParams) rTotmass, rEta = pnutils.mass1_mass2_to_mtotal_eta(rMass1, rMass2) # Here we have a set of m1s,m2s and mapped to Xis below -rXis = tmpltbank.get_cov_params(rMass1, rMass2, rSpin1z, rSpin2z, metricParams, - metricParams.fUpper) +rXis = tmpltbank.get_cov_params( + rMass1, rMass2, rSpin1z, rSpin2z, metricParams, metricParams.fUpper +) xis = (numpy.array(rXis)).T physMasses = numpy.array([rTotmass, rEta, rSpin1z, rSpin2z]) @@ -286,15 +297,15 @@ physMasses = physMasses.T v1smin = v1s.min() v1smax = v1s.max() v1sdiff = v1smax - v1smin -numv1bins = int(math.ceil(v1sdiff / 1.)) +numv1bins = int(math.ceil(v1sdiff / 1.0)) v2smin = v2s.min() v2smax = v2s.max() v2sdiff = v2smax - v2smin -numv2bins = int(math.ceil(v2sdiff / 1.)) +numv2bins = int(math.ceil(v2sdiff / 1.0)) sortedBins = {} -#for i in range(numv1bins): +# for i in range(numv1bins): # print i # sortedBins[i] = {} # for j in range(numv2bins): @@ -302,19 +313,19 @@ sortedBins = {} logging.info("Sorting guide points") -bin_size = (opts.max_mismatch)**0.5 * 6. +bin_size = (opts.max_mismatch) ** 0.5 * 6.0 for iter, currXi in enumerate(xis): xi1 = currXi[0] xi2 = currXi[1] - x1bin = int((xi1 - v1smin)/bin_size) - x2bin = int((xi2 - v2smin)/bin_size) + x1bin = int((xi1 - v1smin) / bin_size) + x2bin = int((xi2 - v2smin) / bin_size) if x1bin not in sortedBins: sortedBins[x1bin] = {} if x2bin not in sortedBins[x1bin]: sortedBins[x1bin][x2bin] = [] if len(sortedBins[x1bin][x2bin]) < 100: - sortedBins[x1bin][x2bin].append( (iter, xi1, xi2) ) + sortedBins[x1bin][x2bin].append((iter, xi1, xi2)) logging.info("Converting to physical points") @@ -326,37 +337,43 @@ tempBank = [] temp_number = 0 req_match = 0.0001 # This can be used for debugging -#fileP = open('bank_file.dat','w') +# fileP = open('bank_file.dat','w') for iter, (v1, v2) in enumerate(zip(v1s, v2s)): # First check if point is within the physical space # and find some example nearby physical points if it is. - v1bin = int((v1 - v1smin)/bin_size) - v2bin = int((v2 - v2smin)/bin_size) + v1bin = int((v1 - v1smin) / bin_size) + v2bin = int((v2 - v2smin) / bin_size) physBin = None - for i in [v1bin, v1bin+1, v1bin-1]: + for i in [v1bin, v1bin + 1, v1bin - 1]: if i not in sortedBins: continue - for j in [v2bin, v2bin+1, v2bin-1]: + for j in [v2bin, v2bin + 1, v2bin - 1]: if j not in sortedBins[i]: continue - physBin = [i,j] + physBin = [i, j] points = numpy.array(sortedBins[physBin[0]][physBin[1]]) - dist = (v1 - points[:,1])**2 + (v2 - points[:,2])**2 + dist = (v1 - points[:, 1]) ** 2 + (v2 - points[:, 2]) ** 2 break if physBin: break else: # No nearby physical points found: continue continue - iters = points[:,0] - bestXis = points[dist.argmin(),1:] - bestMasses = physMasses[int(points[dist.argmin(),0] + 0.5)] + iters = points[:, 0] + bestXis = points[dist.argmin(), 1:] + bestMasses = physMasses[int(points[dist.argmin(), 0] + 0.5)] # Reject point if it is too far from physical space - masses = tmpltbank.get_physical_covaried_masses([v1,v2], - copy.deepcopy(bestMasses), copy.deepcopy(bestXis), req_match, - massRangeParams, metricParams, metricParams.fUpper, - giveUpThresh=20) + masses = tmpltbank.get_physical_covaried_masses( + [v1, v2], + copy.deepcopy(bestMasses), + copy.deepcopy(bestXis), + req_match, + massRangeParams, + metricParams, + metricParams.fUpper, + giveUpThresh=20, + ) # If point is still not close enough to desired space, remove it if masses[5] > opts.max_mismatch: @@ -374,10 +391,7 @@ for iter, (v1, v2) in enumerate(zip(v1s, v2s)): logging.info("Writing output to file %s", opts.output_file) bank_output.output_bank_to_file( - opts.output_file, - tempBank, - programName=__program__, - optDict=opts.__dict__ + opts.output_file, tempBank, programName=__program__, optDict=opts.__dict__ ) logging.info("Done") diff --git a/bin/bank/pycbc_tmpltbank_to_chi_params b/bin/bank/pycbc_tmpltbank_to_chi_params index 817d2fbbe8c..7d097000c91 100644 --- a/bin/bank/pycbc_tmpltbank_to_chi_params +++ b/bin/bank/pycbc_tmpltbank_to_chi_params @@ -20,59 +20,72 @@ Read in a tmpltbank and create a file containing the list of chi parameters """ +import pycbc.version import pycbc -import pycbc.version -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_bank_verification" import argparse import logging -import numpy -from igwn_ligolw import lsctables, utils as ligolw_utils +import numpy +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils -from pycbc import tmpltbank, psd, strain -from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc import psd, strain, tmpltbank from pycbc.io.hdf import HFile - +from pycbc.io.ligolw import LIGOLWContentHandler # Read command line option -parser = argparse.ArgumentParser(description=__doc__, - formatter_class=tmpltbank.IndentedHelpFormatterWithNL) +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=tmpltbank.IndentedHelpFormatterWithNL +) # Begin with code specific options pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-bank", action="store", required=True, - help="The template bank to use an input.") -parser.add_argument("--output-file", action="store", required=True, - help="The ASCII file to create as output. The columns in " - "this file will hold: mass1, mass2, spin1z, spin2z, chi_1, " - "chi_2, ....") +parser.add_argument( + "--input-bank", + action="store", + required=True, + help="The template bank to use an input.", +) +parser.add_argument( + "--output-file", + action="store", + required=True, + help="The ASCII file to create as output. The columns in " + "this file will hold: mass1, mass2, spin1z, spin2z, chi_1, " + "chi_2, ....", +) # Not sure that vary-fupper makes sense in this code. -#parser.add_argument("-V", "--vary-fupper", action="store_true", default=False, +# parser.add_argument("-V", "--vary-fupper", action="store_true", default=False, # help="Use a variable upper frequency cutoff in laying " # "out the bank. OPTIONAL.") -#parser.add_argument("--bank-fupper-step", type=float, default=10., +# parser.add_argument("--bank-fupper-step", type=float, default=10., # help="Size of discrete frequency steps used when varying " # "the fupper. If --calculate-ethinca-metric and " # "--ethinca-freq-step are also given, the code will use " # "the smaller of the two step values. OPTIONAL. Units=Hz") -#parser.add_argument("--bank-fupper-formula", default="SchwarzISCO", +# parser.add_argument("--bank-fupper-formula", default="SchwarzISCO", # choices=["SchwarzISCO","LightRing","ERD"], # help="Frequency cutoff formula for varying fupper. " # "Frequencies will be rounded to the nearest discrete " # "step. OPTIONAL.") -parser.add_argument("--random-seed", action="store", type=int, - default=None, - help="Random seed to use when calling numpy.random " - "functions used in obtaining the principal components in " - "parameter space and when translating points back to " - "physical space. If given, the code should give the " - "same output when run with the same random seed.") +parser.add_argument( + "--random-seed", + action="store", + type=int, + default=None, + help="Random seed to use when calling numpy.random " + "functions used in obtaining the principal components in " + "parameter space and when translating points back to " + "physical space. If given, the code should give the " + "same output when run with the same random seed.", +) # Insert the metric calculation options tmpltbank.insert_metric_calculation_options(parser) @@ -92,9 +105,9 @@ pycbc.init_logging(opts.verbose) # Sanity check options tmpltbank.verify_metric_calculation_options(opts, parser) -metricParams=tmpltbank.metricParameters.from_argparse(opts) +metricParams = tmpltbank.metricParameters.from_argparse(opts) tmpltbank.verify_mass_range_options(opts, parser) -massRangeParams=tmpltbank.massRangeParameters.from_argparse(opts) +massRangeParams = tmpltbank.massRangeParameters.from_argparse(opts) psd.verify_psd_options(opts, parser) if opts.psd_estimation: pycbc.strain.verify_strain_options(opts, parser) @@ -114,17 +127,23 @@ else: logging.info("Obtaining PSD") # Want the number of samples to be a binary number and Nyquist must be above # opts.f_upper. All this assumes that 1 / deltaF is a binary number -nyquistFreq = 2**numpy.ceil(numpy.log2(opts.f_upper)) +nyquistFreq = 2 ** numpy.ceil(numpy.log2(opts.f_upper)) numSamples = int(round(nyquistFreq / opts.delta_f)) + 1 -psd = pycbc.psd.from_cli(opts, length=numSamples, delta_f=opts.delta_f, - low_frequency_cutoff=opts.f_low, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC) +psd = pycbc.psd.from_cli( + opts, + length=numSamples, + delta_f=opts.delta_f, + low_frequency_cutoff=opts.f_low, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, +) metricParams.psd = psd # Begin by calculating a metric logging.info("Calculating metric") -metricParams = tmpltbank.determine_eigen_directions(metricParams, - vary_fmax=False, vary_density=None) +metricParams = tmpltbank.determine_eigen_directions( + metricParams, vary_fmax=False, vary_density=None +) # Choose the frequency values to use for metric calculation refFreq = metricParams.fUpper @@ -132,38 +151,42 @@ refFreq = metricParams.fUpper logging.info("Calculating covariance matrix") vals = tmpltbank.estimate_mass_range( - 1000000, massRangeParams, metricParams, refFreq, covary=False) + 1000000, massRangeParams, metricParams, refFreq, covary=False +) cov = numpy.cov(vals) -evalsCV,evecsCV = numpy.linalg.eig(cov) +evalsCV, evecsCV = numpy.linalg.eig(cov) metricParams.evecsCV = {} metricParams.evecsCV[refFreq] = evecsCV logging.info("Reading template bank.") -if opts.input_bank.endswith(('.xml','.xml.gz','.xmlgz')): - indoc = ligolw_utils.load_filename(opts.input_bank, - contenthandler=LIGOLWContentHandler) +if opts.input_bank.endswith((".xml", ".xml.gz", ".xmlgz")): + indoc = ligolw_utils.load_filename( + opts.input_bank, contenthandler=LIGOLWContentHandler + ) template_list = lsctables.SnglInspiralTable.get_table(indoc) mass1 = [template.mass1 for template in template_list] mass2 = [template.mass2 for template in template_list] spin1z = [template.spin1z for template in template_list] spin2z = [template.spin2z for template in template_list] -elif opts.input_bank.endswith(('.h5','.hdf','.hdf5')): - h5_fp = HFile(opts.input_bank, 'r') - mass1 = h5_fp['mass1'][:] - mass2 = h5_fp['mass2'][:] - spin1z = h5_fp['spin1z'][:] - spin2z = h5_fp['spin2z'][:] +elif opts.input_bank.endswith((".h5", ".hdf", ".hdf5")): + h5_fp = HFile(opts.input_bank, "r") + mass1 = h5_fp["mass1"][:] + mass2 = h5_fp["mass2"][:] + spin1z = h5_fp["spin1z"][:] + spin2z = h5_fp["spin2z"][:] else: - err_msg = "Don't know how to read extension {}.".format(opts.input_bank) + err_msg = f"Don't know how to read extension {opts.input_bank}." raise NotImplementedError(err_msg) fP = open(opts.output_file, "w") for idx in range(len(mass1)): - chi_coords = tmpltbank.get_cov_params(mass1[idx], mass2[idx], spin1z[idx], - spin2z[idx], metricParams, refFreq) - vals = " ".join([str(mass1[idx]), str(mass2[idx]), str(spin1z[idx]), - str(spin2z[idx])]) + chi_coords = tmpltbank.get_cov_params( + mass1[idx], mass2[idx], spin1z[idx], spin2z[idx], metricParams, refFreq + ) + vals = " ".join( + [str(mass1[idx]), str(mass2[idx]), str(spin1z[idx]), str(spin2z[idx])] + ) vals += " " vals += " ".join([str(i) for i in chi_coords]) - fP.write("%s\n" %(vals)) + fP.write("%s\n" % (vals)) fP.close() diff --git a/bin/hwinj/pycbc_generate_hwinj b/bin/hwinj/pycbc_generate_hwinj index 66b0b6af451..3d003a2aca3 100644 --- a/bin/hwinj/pycbc_generate_hwinj +++ b/bin/hwinj/pycbc_generate_hwinj @@ -18,35 +18,31 @@ import argparse import logging -import numpy import os import sys -from igwn_ligolw import ligolw -from igwn_ligolw import lsctables -from igwn_ligolw import utils +import numpy +import pycbc.version +from igwn_ligolw import ligolw, lsctables, utils import pycbc -import pycbc.version -from pycbc import DYN_RANGE_FAC -from pycbc import fft -from pycbc import pnutils +from pycbc import DYN_RANGE_FAC, fft, pnutils from pycbc import psd as _psd from pycbc import strain as _strain from pycbc.detector import Detector +from pycbc.filter import make_frequency_series, sigmasq from pycbc.inject import InjectionSet, legacy_approximant_name -from pycbc.filter import make_frequency_series -from pycbc.filter import sigmasq +from pycbc.io.ligolw import create_process_table from pycbc.types import TimeSeries from pycbc.types.optparse import convert_to_process_params_dict -from pycbc.io.ligolw import create_process_table from pycbc.waveform import get_td_waveform, td_approximants + def _empty_row(obj): - """Create an empty sim_inspiral or sngl_inspiral row where the columns have - default values of 0.0 for a float, 0 for an int, '' for a string. """ - + Create an empty sim_inspiral or sngl_inspiral row where the columns have + default values of 0.0 for a float, 0 for an int, '' for a string. + """ # check if sim_inspiral or sngl_inspiral if obj == lsctables.SimInspiral: row = lsctables.SimInspiral() @@ -57,32 +53,33 @@ def _empty_row(obj): # populate columns with default values for entry in cols.keys(): - if cols[entry] in ['real_4','real_8']: - setattr(row,entry,0.) - elif cols[entry] == 'int_4s': - setattr(row,entry,0) - elif cols[entry] == 'lstring': - setattr(row,entry,'') - elif entry in ['process_id', 'process:process_id']: + if cols[entry] in ["real_4", "real_8"]: + setattr(row, entry, 0.0) + elif cols[entry] == "int_4s": + setattr(row, entry, 0) + elif cols[entry] == "lstring": + setattr(row, entry, "") + elif entry in ["process_id", "process:process_id"]: row.process_id = 0 - elif entry == 'simulation_id': + elif entry == "simulation_id": row.simulation_id = 0 - elif entry == 'event_id': + elif entry == "event_id": row.event_id = 0 else: - raise ValueError("Column %s not recognized." %(entry) ) + raise ValueError("Column %s not recognized." % (entry)) return row + def pad_timeseries_to_integer_length(timeseries, sample_rate): - ''' This function zero pads a time series so that its length is an integer + """ + This function zero pads a time series so that its length is an integer multiple of the sampling rate. Padding is adding symmetically to the start and end of the time series. If the number of samples to pad is odd then the end zero padding will have one more sample than the start zero padding. - ''' - + """ # calculate how many sample points needed to pad to get # integer second time series remainder = sample_rate - len(timeseries) % sample_rate @@ -94,134 +91,241 @@ def pad_timeseries_to_integer_length(timeseries, sample_rate): end_array = numpy.zeros(end_pad) # pad waveform with arrays of zeroes - initial_array = numpy.concatenate([start_array,timeseries,end_array]) - return TimeSeries(initial_array, delta_t=timeseries.delta_t, - epoch=timeseries.start_time, dtype=timeseries.dtype) + initial_array = numpy.concatenate([start_array, timeseries, end_array]) + return TimeSeries( + initial_array, + delta_t=timeseries.delta_t, + epoch=timeseries.start_time, + dtype=timeseries.dtype, + ) + # map order integer to a string that can be parsed by lalsimulation pn_orders = { - 'default' : -1, - 'zeroPN' : 0, - 'onePN' : 2, - 'onePointFivePN' : 3, - 'twoPN' : 4, - 'twoPointFivePN' : 5, - 'threePN' : 6, - 'threePointFivePN' : 7, - 'pseudoFourPN' : 8, + "default": -1, + "zeroPN": 0, + "onePN": 2, + "onePointFivePN": 3, + "twoPN": 4, + "twoPointFivePN": 5, + "threePN": 6, + "threePointFivePN": 7, + "pseudoFourPN": 8, } # command line usage parser = argparse.ArgumentParser( - usage=__name__ + ' [--options]', - description="Generates a hardware injection waveform using " - "a time-domain waveform.") + usage=__name__ + " [--options]", + description="Generates a hardware injection waveform using a time-domain waveform.", +) pycbc.add_common_pycbc_options(parser) # IFO network options -parser.add_argument('--network-snr', type=float, required=True, - help='The network SNR of the injection.') +parser.add_argument( + "--network-snr", type=float, required=True, help="The network SNR of the injection." +) # sky location options -parser.add_argument('--ra', type=float, required=True, - help='The right ascension of the injection in radians.') -parser.add_argument('--dec', type=float, required=True, - help='The declination of the injection in radians.') -parser.add_argument('--polarization', type=float, required=True, - help='The polarization of the injection in radians.') +parser.add_argument( + "--ra", + type=float, + required=True, + help="The right ascension of the injection in radians.", +) +parser.add_argument( + "--dec", + type=float, + required=True, + help="The declination of the injection in radians.", +) +parser.add_argument( + "--polarization", + type=float, + required=True, + help="The polarization of the injection in radians.", +) # waveform parameter options -parser.add_argument('--approximant', type=str, required=True, - choices=td_approximants(), - help='Approximant to use for generating waveform.') -parser.add_argument("--order", type=str, default='default', - choices = pn_orders.keys(), - help='The integer half-PN order at which to generate ' - 'the approximant.') -parser.add_argument('--mass1', type=float, required=True, - help='First mass of the binary in solar masses.') -parser.add_argument('--mass2', type=float, required=True, - help='Second mass of the binary in solar masses.') -parser.add_argument('--inclination', type=float, required=True, - help='Inclination of the binary in radians.') -parser.add_argument('--coa-phase', type=float, default=0.0, - help='Reference orbital phase parameter in radians. ' - 'Note, this is not the same as the constant ' - 'frequency-domain phase shift that is maximized ' - 'over in the standard matched-filter search and ' - 'which is also commonly called coa-phase.' - 'Called phiRef in LALSimulation. ' - 'Called coa_phase in sim_inspiral tables.') -parser.add_argument('--taper', required=True, - choices=['TAPER_NONE', 'TAPER_START', - 'TAPER_END', 'TAPER_STARTEND'], - help='Taper the wavform before FFT.') -parser.add_argument('--waveform-low-frequency-cutoff', type=float, - required=True, - help='Frequency to begin generating the waveform in Hz.') +parser.add_argument( + "--approximant", + type=str, + required=True, + choices=td_approximants(), + help="Approximant to use for generating waveform.", +) +parser.add_argument( + "--order", + type=str, + default="default", + choices=pn_orders.keys(), + help="The integer half-PN order at which to generate the approximant.", +) +parser.add_argument( + "--mass1", + type=float, + required=True, + help="First mass of the binary in solar masses.", +) +parser.add_argument( + "--mass2", + type=float, + required=True, + help="Second mass of the binary in solar masses.", +) +parser.add_argument( + "--inclination", + type=float, + required=True, + help="Inclination of the binary in radians.", +) +parser.add_argument( + "--coa-phase", + type=float, + default=0.0, + help="Reference orbital phase parameter in radians. " + "Note, this is not the same as the constant " + "frequency-domain phase shift that is maximized " + "over in the standard matched-filter search and " + "which is also commonly called coa-phase." + "Called phiRef in LALSimulation. " + "Called coa_phase in sim_inspiral tables.", +) +parser.add_argument( + "--taper", + required=True, + choices=["TAPER_NONE", "TAPER_START", "TAPER_END", "TAPER_STARTEND"], + help="Taper the wavform before FFT.", +) +parser.add_argument( + "--waveform-low-frequency-cutoff", + type=float, + required=True, + help="Frequency to begin generating the waveform in Hz.", +) # waveform spin parameter options -parser.add_argument('--spin1z', type=float, default=0.0, - help='(optional) Spin in z direction for mass1.') -parser.add_argument('--spin1y', type=float, default=0.0, - help='(optional) Spin in y direction for mass1.') -parser.add_argument('--spin1x', type=float, default=0.0, - help='(optional) Spin in x direction for mass1.') -parser.add_argument('--spin2z', type=float, default=0.0, - help='(optional) Spin in z direction for mass2.') -parser.add_argument('--spin2y', type=float, default=0.0, - help='(optional) Spin in y direction for mass2.') -parser.add_argument('--spin2x', type=float, default=0.0, - help='(optional) Spin in x direction for mass2.') +parser.add_argument( + "--spin1z", + type=float, + default=0.0, + help="(optional) Spin in z direction for mass1.", +) +parser.add_argument( + "--spin1y", + type=float, + default=0.0, + help="(optional) Spin in y direction for mass1.", +) +parser.add_argument( + "--spin1x", + type=float, + default=0.0, + help="(optional) Spin in x direction for mass1.", +) +parser.add_argument( + "--spin2z", + type=float, + default=0.0, + help="(optional) Spin in z direction for mass2.", +) +parser.add_argument( + "--spin2y", + type=float, + default=0.0, + help="(optional) Spin in y direction for mass2.", +) +parser.add_argument( + "--spin2x", + type=float, + default=0.0, + help="(optional) Spin in x direction for mass2.", +) # tidal options -parser.add_argument('--lambda1', type=float, default=None, - help='(optional) Tidal lambda term for mass1.' - 'WARNING: Giving this option will produce an XML ' - 'file containing a lambda1 column. This file will ' - 'not be readable by default methods.') -parser.add_argument('--lambda2', type=float, default=None, - help='(optional) Tidal lambda term for mass2. ' - 'WARNING: Giving this option will produce an XML ' - 'file containing a lambda2 column. This file will ' - 'not be readable by default methods.') -parser.add_argument('--dquad-mon1', type=float, default=None, - help='(optional) Tidal self-spin term for mass1, ' - 'for BHs this is 0 (its the deformation relative to ' - 'Kerr. ' - 'WARNING: Giving this option will produce an XML ' - 'file containing a dquad_mon1 column. This file will ' - 'not be readable by default methods.') -parser.add_argument('--dquad-mon2', type=float, default=None, - help='(optional) Tidal self-spin term for mass2, ' - 'for BHs this is 0 (its the deformation relative to ' - 'Kerr. ' - 'WARNING: Giving this option will produce an XML ' - 'file containing a dquad_mon2 column. This file will ' - 'not be readable by default methods.') +parser.add_argument( + "--lambda1", + type=float, + default=None, + help="(optional) Tidal lambda term for mass1." + "WARNING: Giving this option will produce an XML " + "file containing a lambda1 column. This file will " + "not be readable by default methods.", +) +parser.add_argument( + "--lambda2", + type=float, + default=None, + help="(optional) Tidal lambda term for mass2. " + "WARNING: Giving this option will produce an XML " + "file containing a lambda2 column. This file will " + "not be readable by default methods.", +) +parser.add_argument( + "--dquad-mon1", + type=float, + default=None, + help="(optional) Tidal self-spin term for mass1, " + "for BHs this is 0 (its the deformation relative to " + "Kerr. " + "WARNING: Giving this option will produce an XML " + "file containing a dquad_mon1 column. This file will " + "not be readable by default methods.", +) +parser.add_argument( + "--dquad-mon2", + type=float, + default=None, + help="(optional) Tidal self-spin term for mass2, " + "for BHs this is 0 (its the deformation relative to " + "Kerr. " + "WARNING: Giving this option will produce an XML " + "file containing a dquad_mon2 column. This file will " + "not be readable by default methods.", +) # Use this for NR injections -parser.add_argument('--numrel-data', type=str, default='', - help="Location of NR data file if using NR injections.") +parser.add_argument( + "--numrel-data", + type=str, + default="", + help="Location of NR data file if using NR injections.", +) # end time options -parser.add_argument('--geocentric-end-time', type=float, required=True, - help='The geocentric GPS end time of the injection.') +parser.add_argument( + "--geocentric-end-time", + type=float, + required=True, + help="The geocentric GPS end time of the injection.", +) # data conditioning options -parser.add_argument('--low-frequency-cutoff', type=float, required=True, - help='Frequency to begin generating the PSD in Hz. This ' - 'is the start frequency of the SNR calculation.') -parser.add_argument('--high-frequency-cutoff', type=float, - help='(optional) Upper frequency to terminate the SNR ' - 'calculation. Default will be Nyquist frequency, ' - 'ie. int(sample_rate/2).') +parser.add_argument( + "--low-frequency-cutoff", + type=float, + required=True, + help="Frequency to begin generating the PSD in Hz. This " + "is the start frequency of the SNR calculation.", +) +parser.add_argument( + "--high-frequency-cutoff", + type=float, + help="(optional) Upper frequency to terminate the SNR " + "calculation. Default will be Nyquist frequency, " + "ie. int(sample_rate/2).", +) # output options -parser.add_argument("--tag", type=str, default='hwinjcbc', - help="Prefix added to output filenames.") -parser.add_argument("--instruments", nargs="+", type=str, required=True, - help="List of instruments to analyze.") +parser.add_argument( + "--tag", type=str, default="hwinjcbc", help="Prefix added to output filenames." +) +parser.add_argument( + "--instruments", + nargs="+", + type=str, + required=True, + help="List of instruments to analyze.", +) # add option groups fft.insert_fft_option_group(parser) @@ -234,11 +338,14 @@ opts = parser.parse_args() # verify options are sane if using strain options if opts.psd_estimation: _strain.verify_strain_options_multi_ifo(opts, parser, opts.instruments) -if not opts.psd_estimation and (opts.frame_files or opts.frame_type - or opts.frame_cache or opts.fake_strain): - raise KeyError("Must use --psd-estimation with frame options" - "(--frame-files, --frame-type, --frame-cache, " - "and --fake-strain).") +if not opts.psd_estimation and ( + opts.frame_files or opts.frame_type or opts.frame_cache or opts.fake_strain +): + raise KeyError( + "Must use --psd-estimation with frame options" + "(--frame-files, --frame-type, --frame-cache, " + "and --fake-strain)." + ) # setup log: default is DEBUG (2) log_level = 2 if opts.verbose is None else opts.verbose + 2 @@ -247,7 +354,7 @@ pycbc.init_logging(opts.verbose) # check that sample rates are the same for ifo in opts.instruments: if opts.sample_rate[ifo] != opts.sample_rate[opts.instruments[0]]: - logging.warning('Sample rates must be equal for all IFOs.') + logging.warning("Sample rates must be equal for all IFOs.") sys.exit() sample_rate = opts.sample_rate[opts.instruments[0]] @@ -259,7 +366,7 @@ else: # check that frame types are not lists if opts.frame_type: - for key,value in opts.frame_type.items(): + for key, value in opts.frame_type.items(): if type(opts.frame_type[key]) == list: opts.frame_type[key] = value[0] @@ -270,14 +377,15 @@ distance = 40.0 network_snr = 0.0 # create output XML file -logging.info('Creating XML file') +logging.info("Creating XML file") outdoc = ligolw.Document() outdoc.appendChild(ligolw.LIGO_LW()) # create process table llw_opts = convert_to_process_params_dict(opts) -create_process_table(outdoc, sys.argv[0], options=llw_opts, - detectors=[''.join(opts.instruments)]) +create_process_table( + outdoc, sys.argv[0], options=llw_opts, detectors=["".join(opts.instruments)] +) # create sim_inspiral row for injection # and populate non-IFO-specific columns in XML output file @@ -285,18 +393,23 @@ create_process_table(outdoc, sys.argv[0], options=llw_opts, # If using tidal terms some hacking is currently needed. Hopefully this can # be resolved in the future by using HDF injection files! if opts.lambda1 is not None: - lsctables.SimInspiralTable.validcolumns['lambda1'] = 'real_4' + lsctables.SimInspiralTable.validcolumns["lambda1"] = "real_4" if opts.lambda2 is not None: - lsctables.SimInspiralTable.validcolumns['lambda2'] = 'real_4' + lsctables.SimInspiralTable.validcolumns["lambda2"] = "real_4" if opts.dquad_mon1 is not None: - lsctables.SimInspiralTable.validcolumns['dquad_mon1'] = 'real_4' + lsctables.SimInspiralTable.validcolumns["dquad_mon1"] = "real_4" if opts.dquad_mon2 is not None: - lsctables.SimInspiralTable.validcolumns['dquad_mon2'] = 'real_4' + lsctables.SimInspiralTable.validcolumns["dquad_mon2"] = "real_4" + + # You'd think that would it to redefine the table columns, but wait, it gets # worse! class SimInspiralNew(lsctables.SimInspiral): - __slots__ = tuple(map(ligolw.Column.ColumnName, - lsctables.SimInspiralTable.validcolumns)) + __slots__ = tuple( + map(ligolw.Column.ColumnName, lsctables.SimInspiralTable.validcolumns) + ) + + lsctables.SimInspiral = SimInspiralNew lsctables.SimInspiralTable.RowType = SimInspiralNew @@ -354,8 +467,7 @@ outdoc.childNodes[0].appendChild(sngl_table) sngl = _empty_row(lsctables.SnglInspiral) sngl.mass1 = opts.mass1 sngl.mass2 = opts.mass2 -sngl.mchirp, sngl.eta = pnutils.mass1_mass2_to_mchirp_eta(sngl.mass1, - sngl.mass2) +sngl.mchirp, sngl.eta = pnutils.mass1_mass2_to_mchirp_eta(sngl.mass1, sngl.mass2) sngl.mtotal = sngl.mass1 + sngl.mass2 sngl.spin1z = opts.spin1z sngl.spin1y = opts.spin1y @@ -365,12 +477,18 @@ sngl.spin2y = opts.spin2y sngl.spin2x = opts.spin2x # generate waveform -logging.info('Generating waveform at %.3fMpc beginning at %.3fHz for ' - 'SNR calculation', sim.distance, opts.waveform_low_frequency_cutoff) -h_plus, h_cross = get_td_waveform(sim, approximant=name, - phase_order=phase_order, - f_lower=opts.waveform_low_frequency_cutoff, - delta_t=1.0 / sample_rate) +logging.info( + "Generating waveform at %.3fMpc beginning at %.3fHz for SNR calculation", + sim.distance, + opts.waveform_low_frequency_cutoff, +) +h_plus, h_cross = get_td_waveform( + sim, + approximant=name, + phase_order=phase_order, + f_lower=opts.waveform_low_frequency_cutoff, + delta_t=1.0 / sample_rate, +) # zero pad polarizations to get integer second time series h_plus = pad_timeseries_to_integer_length(h_plus, sample_rate) @@ -381,8 +499,9 @@ h_cross = pad_timeseries_to_integer_length(h_cross, sample_rate) # then genreate a zeroNoise TimeSeries to get length, delta_f, etc. if not opts.psd_estimation: opts.fake_strain = "zeroNoise" -strain_dict = _strain.from_cli_multi_ifos(opts, opts.instruments, - dyn_range_fac=DYN_RANGE_FAC) +strain_dict = _strain.from_cli_multi_ifos( + opts, opts.instruments, dyn_range_fac=DYN_RANGE_FAC +) # organize options for multi-IFO PSD # if not generating strain then set those related options to None @@ -397,71 +516,86 @@ for ifo in opts.instruments: low_frequency_cutoff_dict[ifo] = opts.low_frequency_cutoff # get PSD -logging.info('Generating PSDs') +logging.info("Generating PSDs") psd_dict = _psd.from_cli_multi_ifos( - opts, length_dict, delta_f_dict, - low_frequency_cutoff_dict, opts.instruments, - strain_dict=strain_dict, - dyn_range_factor=DYN_RANGE_FAC, - precision="double") + opts, + length_dict, + delta_f_dict, + low_frequency_cutoff_dict, + opts.instruments, + strain_dict=strain_dict, + dyn_range_factor=DYN_RANGE_FAC, + precision="double", +) # loop over IFOs to calculate sigma for ifo in opts.instruments: - # get Detector instance for IFO det = Detector(ifo) # get time delay to detector from center of the Earth - time_delay = det.time_delay_from_earth_center(sim.longitude, sim.latitude, - sim.geocent_end_time) + time_delay = det.time_delay_from_earth_center( + sim.longitude, sim.latitude, sim.geocent_end_time + ) end_time = sim.geocent_end_time + time_delay # get antenna pattern - f_plus, f_cross = det.antenna_pattern(sim.longitude, sim.latitude, - sim.polarization, - sim.geocent_end_time) + f_plus, f_cross = det.antenna_pattern( + sim.longitude, sim.latitude, sim.polarization, sim.geocent_end_time + ) # calculate strain - logging.info('Calculating strain for %s', ifo) + logging.info("Calculating strain for %s", ifo) strain = f_plus * h_plus + f_cross * h_cross # taper waveform - logging.info('Tapering strain for %s', ifo) + logging.info("Tapering strain for %s", ifo) strain = strain.taper_timeseries(location=sim.taper) # FFT strain - logging.info('FFT strain for %s', ifo) + logging.info("FFT strain for %s", ifo) strain_tilde = make_frequency_series(strain) # interpolate PSD to waveform delta_f if psd_dict[ifo].delta_f != strain_tilde.delta_f: - logging.info('Interpolating PSD for %s from %fHz to %fHz', - ifo, psd_dict[ifo].delta_f, strain_tilde.delta_f) - psd_dict[ifo] = _psd.interpolate( - psd_dict[ifo], strain_tilde.delta_f) + logging.info( + "Interpolating PSD for %s from %fHz to %fHz", + ifo, + psd_dict[ifo].delta_f, + strain_tilde.delta_f, + ) + psd_dict[ifo] = _psd.interpolate(psd_dict[ifo], strain_tilde.delta_f) # calculate sigma-squared SNR - logging.info('Calculating sigma for %s', ifo) + logging.info("Calculating sigma for %s", ifo) sigma_squared = sigmasq( - DYN_RANGE_FAC * strain_tilde, - psd=psd_dict[ifo], - low_frequency_cutoff=opts.low_frequency_cutoff, - high_frequency_cutoff=f_high) - logging.info('Sigma integrated from %.3f to %.3fHz for %s is %.3f', - opts.low_frequency_cutoff, f_high, ifo, - numpy.sqrt(sigma_squared)) + DYN_RANGE_FAC * strain_tilde, + psd=psd_dict[ifo], + low_frequency_cutoff=opts.low_frequency_cutoff, + high_frequency_cutoff=f_high, + ) + logging.info( + "Sigma integrated from %.3f to %.3fHz for %s is %.3f", + opts.low_frequency_cutoff, + f_high, + ifo, + numpy.sqrt(sigma_squared), + ) # populate IFO end time columns - setattr(sim, ifo[0].lower()+'_end_time', int(end_time)) - setattr(sim, ifo[0].lower()+'_end_time_ns', int(end_time % 1 * 1e9)) + setattr(sim, ifo[0].lower() + "_end_time", int(end_time)) + setattr(sim, ifo[0].lower() + "_end_time_ns", int(end_time % 1 * 1e9)) # populate IFO distance columns - eff_distance = det.effective_distance(sim.distance, - sim.longitude, sim.latitude, - sim.polarization, - sim.geocent_end_time, - sim.inclination) - setattr(sim, 'eff_dist_'+ifo[0].lower(), eff_distance) + eff_distance = det.effective_distance( + sim.distance, + sim.longitude, + sim.latitude, + sim.polarization, + sim.geocent_end_time, + sim.inclination, + ) + setattr(sim, "eff_dist_" + ifo[0].lower(), eff_distance) # populate IFO end time columns sngl.end_time = int(end_time) @@ -473,19 +607,25 @@ for ifo in opts.instruments: # distance scaling factor to get target snr network_snr = numpy.sqrt(network_snr) scale = network_snr / opts.network_snr -sim.distance = scale*sim.distance +sim.distance = scale * sim.distance for ifo in opts.instruments: - attrname='eff_dist_'+ifo[0].lower() - effdist=getattr(sim,attrname) - setattr(sim,attrname,effdist*scale) + attrname = "eff_dist_" + ifo[0].lower() + effdist = getattr(sim, attrname) + setattr(sim, attrname, effdist * scale) # generate waveform -logging.info('Generating waveform at %.3fMpc beginning at %.3fHz for ' - 'SNR calculation', sim.distance, opts.waveform_low_frequency_cutoff) -h_plus, h_cross = get_td_waveform(sim, approximant=name, - phase_order=phase_order, - f_lower=opts.waveform_low_frequency_cutoff, - delta_t=1.0 / sample_rate) +logging.info( + "Generating waveform at %.3fMpc beginning at %.3fHz for SNR calculation", + sim.distance, + opts.waveform_low_frequency_cutoff, +) +h_plus, h_cross = get_td_waveform( + sim, + approximant=name, + phase_order=phase_order, + f_lower=opts.waveform_low_frequency_cutoff, + delta_t=1.0 / sample_rate, +) # zero pad polarizations to get integer second time series h_plus = pad_timeseries_to_integer_length(h_plus, sample_rate) @@ -493,11 +633,12 @@ h_cross = pad_timeseries_to_integer_length(h_cross, sample_rate) # figure out length of time series to inject waveform into -logging.info('Calculating number of sample points in output file') -h_plus, _ = get_td_waveform(sim, approximant=name, phase_order=phase_order, - delta_t=1.0/sample_rate) +logging.info("Calculating number of sample points in output file") +h_plus, _ = get_td_waveform( + sim, approximant=name, phase_order=phase_order, delta_t=1.0 / sample_rate +) pad_seconds = 5 -template_duration_seconds = int( len(h_plus) / sample_rate ) + 1 +template_duration_seconds = int(len(h_plus) / sample_rate) + 1 start_time = int(sim.geocent_end_time) - template_duration_seconds - pad_seconds end_time = int(sim.geocent_end_time) + 1 + pad_seconds num_samples = int((end_time - start_time) * sample_rate) @@ -507,43 +648,45 @@ sim_table.append(sim) sngl_table.append(sngl) # construct filenames prefix -prefix = opts.tag + '_' + str(start_time) +prefix = opts.tag + "_" + str(start_time) # save XML output file if it does not exist -logging.info('Writing XML file') -xml_filename = prefix + '.xml.gz' +logging.info("Writing XML file") +xml_filename = prefix + ".xml.gz" if os.path.exists(xml_filename): - logging.warning('Filename %s already exists and will not be overwritten', - xml_filename) + logging.warning( + "Filename %s already exists and will not be overwritten", xml_filename + ) sys.exit() else: utils.write_filename(outdoc, xml_filename) # loop over IFOs for writing waveforms to file for ifo in opts.instruments: - # create a time series of zeroes to inject waveform into initial_array = numpy.zeros(num_samples, dtype=strain.dtype) - output = TimeSeries(initial_array, delta_t=1.0 / sample_rate, - epoch=start_time, dtype=strain.dtype) + output = TimeSeries( + initial_array, delta_t=1.0 / sample_rate, epoch=start_time, dtype=strain.dtype + ) # inject waveform - logging.info('Injecting %s waveform into timeseries of zeroes', ifo) + logging.info("Injecting %s waveform into timeseries of zeroes", ifo) injections = InjectionSet(xml_filename) injections.apply(output, ifo) # set output filename - txt_filename = prefix + '_' + ifo + '.txt' + txt_filename = prefix + "_" + ifo + ".txt" # check if filename exists if os.path.exists(txt_filename): - logging.warning('Filename %s already exists and will not be overwritten', - txt_filename) + logging.warning( + "Filename %s already exists and will not be overwritten", txt_filename + ) sys.exit() # save waveform as single-column ASCII for awgstream to use - logging.info('Writing strain for %s', ifo) + logging.info("Writing strain for %s", ifo) numpy.savetxt(txt_filename, output) # finish -logging.info('Done') +logging.info("Done") diff --git a/bin/hwinj/pycbc_generate_hwinj_from_xml b/bin/hwinj/pycbc_generate_hwinj_from_xml index 09349eff4ad..2d2c7311bb0 100644 --- a/bin/hwinj/pycbc_generate_hwinj_from_xml +++ b/bin/hwinj/pycbc_generate_hwinj_from_xml @@ -19,21 +19,21 @@ import argparse import logging -import numpy import os import sys +import numpy + from pycbc import add_common_pycbc_options, init_logging +from pycbc.detector import get_available_detectors from pycbc.inject import InjectionSet, legacy_approximant_name from pycbc.types import TimeSeries from pycbc.waveform import get_td_waveform -from pycbc.detector import get_available_detectors - # command line usage parser = argparse.ArgumentParser( - usage='pycbc_generate_hwinj_from_xml --injection-file [INJECTION_FILE] --sample-rate [SAMPLE_RATE]', - description='Generates all the time-domain injections from a LIGOLW \ + usage="pycbc_generate_hwinj_from_xml --injection-file [INJECTION_FILE] --sample-rate [SAMPLE_RATE]", + description="Generates all the time-domain injections from a LIGOLW \ sim_inspiral table for H1 and L1. \ To run the executable you have to specify a LIGOLW XML file \ that contains a sim_inspiral table (--injection-file) and the sample rate \ @@ -43,20 +43,35 @@ parser = argparse.ArgumentParser( current working directory with the standard LIGO naming convetion, \ eg. H1-HWINJ_CBC_FROM_SIMULATION_ID_X-Y-Z.txt where X is the simulation_id \ from the sim_inspiral row, Y is the GPS start time of the ASCII waveform file, \ - and Z is the duration of the file in seconds.') + and Z is the duration of the file in seconds.", +) add_common_pycbc_options(parser) # add command line options -parser.add_argument('--injection-file', type=str, required=True, - help='Path to the LIGOLW XML file that contains a sim_inspiral table.') -parser.add_argument('--sample-rate', type=int, required=True, - help='Sample rate that waveforms will be generated.') -parser.add_argument("--tag", type=str, default='hwinjcbcsimid', - help="Prefix added to output filenames.") -parser.add_argument('--ifos', nargs='+', default=['H1', 'L1'], required=True, - choices=get_available_detectors(), - help='List of IFOs to generate injections for.') +parser.add_argument( + "--injection-file", + type=str, + required=True, + help="Path to the LIGOLW XML file that contains a sim_inspiral table.", +) +parser.add_argument( + "--sample-rate", + type=int, + required=True, + help="Sample rate that waveforms will be generated.", +) +parser.add_argument( + "--tag", type=str, default="hwinjcbcsimid", help="Prefix added to output filenames." +) +parser.add_argument( + "--ifos", + nargs="+", + default=["H1", "L1"], + required=True, + choices=get_available_detectors(), + help="List of IFOs to generate injections for.", +) # parse command line opts = parser.parse_args() @@ -64,14 +79,13 @@ opts = parser.parse_args() init_logging(args.verbose, default_level=2) # read in injection LIGOLW XML file -logging.info('Reading injection file') +logging.info("Reading injection file") injections = InjectionSet(opts.injection_file) # loop over rows in sim_inspiral table for sim in injections.table: - # print statement - logging.info('Begin generating waveform for %s', sim.simulation_id) + logging.info("Begin generating waveform for %s", sim.simulation_id) # parse the sim_inspiral waveform column name, phase_order = legacy_approximant_name(sim.waveform) @@ -82,46 +96,46 @@ for sim in injections.table: # generate the waveform # we generate the waveform so we know how many sample points the ASCII file # will need to have to contain the full waveform - logging.info('Getting number of sample points in waveform') - h_plus, _ = get_td_waveform(sim, approximant=name, phase_order=phase_order, - delta_t=1.0/opts.sample_rate) + logging.info("Getting number of sample points in waveform") + h_plus, _ = get_td_waveform( + sim, approximant=name, phase_order=phase_order, delta_t=1.0 / opts.sample_rate + ) # figure out length of the time series to inject waveform into pad_seconds = 5 - template_duration_seconds = int( len(h_plus) / opts.sample_rate ) + 1 + template_duration_seconds = int(len(h_plus) / opts.sample_rate) + 1 start_time = int(sim.geocent_end_time) - template_duration_seconds - pad_seconds end_time = int(sim.geocent_end_time) + 1 + pad_seconds num_samples = (end_time - start_time) * opts.sample_rate # loop over IFOs for writing waveforms to file for ifo in opts.ifos: - # create a time series of zeroes to inject waveform into initial_array = numpy.zeros(num_samples, dtype=h_plus.dtype) output = TimeSeries( initial_array, - delta_t=1.0/opts.sample_rate, + delta_t=1.0 / opts.sample_rate, epoch=start_time, - dtype=h_plus.dtype) + dtype=h_plus.dtype, + ) # inject waveform into time series of zeroes - logging.info('Injecting %s waveform into timeseries of zeroes', ifo) + logging.info("Injecting %s waveform into timeseries of zeroes", ifo) injections.apply(output, ifo, simulation_ids=[sim.simulation_id]) # set output filename - txt_filename = f'{opts.tag}{sim_id:d}_{start_time:d}_{ifo}.txt' + txt_filename = f"{opts.tag}{sim_id:d}_{start_time:d}_{ifo}.txt" # check if filename does not exist if os.path.exists(txt_filename): logging.warning( - 'Filename %s already exists and will not be overwritten', - txt_filename + "Filename %s already exists and will not be overwritten", txt_filename ) sys.exit() # save waveform as single column ASCII for awgstream to use - logging.info('Writing strain for %s', ifo) + logging.info("Writing strain for %s", ifo) numpy.savetxt(txt_filename, output) # finish -logging.info('Done') +logging.info("Done") diff --git a/bin/hwinj/pycbc_insert_frame_hwinj b/bin/hwinj/pycbc_insert_frame_hwinj index 458f426e2e5..61f860a530d 100644 --- a/bin/hwinj/pycbc_insert_frame_hwinj +++ b/bin/hwinj/pycbc_insert_frame_hwinj @@ -18,38 +18,57 @@ import argparse import logging + import numpy from pycbc import add_common_pycbc_options, init_logging -from pycbc.frame import write_frame from pycbc import strain as _strain +from pycbc.frame import write_frame # command line usage -parser = argparse.ArgumentParser(usage='pycbc_insert_frame_hwinj [--options]', - description="Inserts a single-column ASCII " - "file into frame data.") +parser = argparse.ArgumentParser( + usage="pycbc_insert_frame_hwinj [--options]", + description="Inserts a single-column ASCII file into frame data.", +) add_common_pycbc_options(parser) # injection options -parser.add_argument('--hwinj-file', type=str, required=True, - help='Path to single-column ASCII file.') -parser.add_argument('--hwinj-start-time', type=int, required=True, - help='Start time of the single-column ASCII file.') -parser.add_argument('--scale-factor', type=float, default=1.0, - help='Scales the waveform amplitude by a float.') +parser.add_argument( + "--hwinj-file", type=str, required=True, help="Path to single-column ASCII file." +) +parser.add_argument( + "--hwinj-start-time", + type=int, + required=True, + help="Start time of the single-column ASCII file.", +) +parser.add_argument( + "--scale-factor", + type=float, + default=1.0, + help="Scales the waveform amplitude by a float.", +) # frame options -parser.add_argument('--ifo', type=str, required=True, - help='IFO.') -parser.add_argument('--output-file', type=str, required=True, - help='Path to output frame file.') -parser.add_argument('--precision', type=str, default='double', - choices=['single', 'double'], - help='Store as a float32 or float64.') +parser.add_argument("--ifo", type=str, required=True, help="IFO.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output frame file." +) +parser.add_argument( + "--precision", + type=str, + default="double", + choices=["single", "double"], + help="Store as a float32 or float64.", +) # fake data options -parser.add_argument('--low-frequency-cutoff', type=int, required=False, - help='Low-frequency cutoff for generating fake PSD.') +parser.add_argument( + "--low-frequency-cutoff", + type=int, + required=False, + help="Low-frequency cutoff for generating fake PSD.", +) # add option groups _strain.insert_strain_option_group(parser) @@ -64,22 +83,22 @@ init_logging(opts.verbose, default_level=2) strain = _strain.from_cli(opts, precision=opts.precision) # load data -logging.info('Reading the hardware injection data') +logging.info("Reading the hardware injection data") initial_array = numpy.loadtxt(opts.hwinj_file) initial_array *= opts.scale_factor # figure out how much to pad -start_pad = (opts.hwinj_start_time-strain.start_time) * opts.sample_rate +start_pad = (opts.hwinj_start_time - strain.start_time) * opts.sample_rate start_pad = int(start_pad + 0.5) # add the two time series -logging.info('Summing the two time series') +logging.info("Summing the two time series") for i in range(len(initial_array)): - strain[start_pad+1+i] += initial_array[i] + strain[start_pad + 1 + i] += initial_array[i] # write frame -logging.info('Writing data') -write_frame(opts.output_file, opts.ifo+':HWINJ_INJECTED', strain) +logging.info("Writing data") +write_frame(opts.output_file, opts.ifo + ":HWINJ_INJECTED", strain) # exit -logging.info('Done.') +logging.info("Done.") diff --git a/bin/hwinj/pycbc_plot_hwinj b/bin/hwinj/pycbc_plot_hwinj index 8c34959300b..69e8f7eb2e2 100644 --- a/bin/hwinj/pycbc_plot_hwinj +++ b/bin/hwinj/pycbc_plot_hwinj @@ -15,43 +15,45 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plots a single-column ASCII file data. -""" +"""Plots a single-column ASCII file data.""" import argparse import logging -import matplotlib as mpl; mpl.use("Agg") + +import matplotlib as mpl + +mpl.use("Agg") import matplotlib.pyplot as plt import numpy from pycbc import add_common_pycbc_options, init_logging # command line usage -parser = argparse.ArgumentParser(usage=__file__ + " [--options]", - description=__doc__) +parser = argparse.ArgumentParser(usage=__file__ + " [--options]", description=__doc__) add_common_pycbc_options(parser) # I/O options -parser.add_argument("--input-file", required=True, - help="Path to single column ASCII file.") -parser.add_argument("--output-file", required=True, - help="Path to output plot.") +parser.add_argument( + "--input-file", required=True, help="Path to single column ASCII file." +) +parser.add_argument("--output-file", required=True, help="Path to output plot.") # plotting options -parser.add_argument("--x-min", type=float, default=None, - help="Minimum x-value to plot.") -parser.add_argument("--x-max", type=float, default=None, - help="Maximum x-value to plot.") -parser.add_argument("--y-min", type=float, default=None, - help="Minimum y-value to plot.") -parser.add_argument("--y-max", type=float, default=None, - help="Maximum y-value to plot.") -parser.add_argument("--x-label", default="Samples", - help="Label for x-axis.") -parser.add_argument("--y-label", default="Counts", - help="Label for y-axis.") -parser.add_argument("--title", default="", - help="Title shown above plot.") +parser.add_argument( + "--x-min", type=float, default=None, help="Minimum x-value to plot." +) +parser.add_argument( + "--x-max", type=float, default=None, help="Maximum x-value to plot." +) +parser.add_argument( + "--y-min", type=float, default=None, help="Minimum y-value to plot." +) +parser.add_argument( + "--y-max", type=float, default=None, help="Maximum y-value to plot." +) +parser.add_argument("--x-label", default="Samples", help="Label for x-axis.") +parser.add_argument("--y-label", default="Counts", help="Label for y-axis.") +parser.add_argument("--title", default="", help="Title shown above plot.") # parse command line opts = parser.parse_args() diff --git a/bin/inference/pycbc_inference b/bin/inference/pycbc_inference index 1706e14e699..d224cc03aec 100644 --- a/bin/inference/pycbc_inference +++ b/bin/inference/pycbc_inference @@ -15,54 +15,70 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Runs a sampler to find the posterior distributions. -""" +"""Runs a sampler to find the posterior distributions.""" -import os -import sys import argparse import logging +import os +import sys import numpy import pycbc -from pycbc import fft, opt, scheme, pool - -from pycbc import inference +from pycbc import fft, inference, opt, pool, scheme from pycbc.inference import models from pycbc.inference.io import loadfile from pycbc.workflow import configuration # command line usage -parser = argparse.ArgumentParser(usage=__file__ + " [--options]", - description=__doc__) +parser = argparse.ArgumentParser(usage=__file__ + " [--options]", description=__doc__) pycbc.add_common_pycbc_options(parser) # output options -parser.add_argument("--output-file", type=str, required=True, - help="Output file path.") -parser.add_argument("--force", action="store_true", default=False, - help="If the output-file already exists, overwrite it. " - "Otherwise, an OSError is raised.") -parser.add_argument("--save-backup", action="store_true", - default=False, - help="Don't delete the backup file after the run has " - "completed.") +parser.add_argument("--output-file", type=str, required=True, help="Output file path.") +parser.add_argument( + "--force", + action="store_true", + default=False, + help="If the output-file already exists, overwrite it. " + "Otherwise, an OSError is raised.", +) +parser.add_argument( + "--save-backup", + action="store_true", + default=False, + help="Don't delete the backup file after the run has completed.", +) # parallelization options -parser.add_argument("--nprocesses", type=int, default=1, - help="Number of processes to use. If not given then only " - "a single core will be used.") -parser.add_argument("--use-mpi", action='store_true', default=False, - help="Use MPI to parallelize the sampler") -parser.add_argument("--samples-file", default=None, - help="Use an iteration from an InferenceFile as the " - "initial proposal distribution. The same " - "number of walkers and the same [variable_params] " - "section in the configuration file should be used. " - "The priors must allow encompass the initial " - "positions from the InferenceFile being read.") -parser.add_argument("--seed", type=int, default=0, - help="Seed to use for the random number generator that " - "initially distributes the walkers. Default is 0.") +parser.add_argument( + "--nprocesses", + type=int, + default=1, + help="Number of processes to use. If not given then only " + "a single core will be used.", +) +parser.add_argument( + "--use-mpi", + action="store_true", + default=False, + help="Use MPI to parallelize the sampler", +) +parser.add_argument( + "--samples-file", + default=None, + help="Use an iteration from an InferenceFile as the " + "initial proposal distribution. The same " + "number of walkers and the same [variable_params] " + "section in the configuration file should be used. " + "The priors must allow encompass the initial " + "positions from the InferenceFile being read.", +) +parser.add_argument( + "--seed", + type=int, + default=0, + help="Seed to use for the random number generator that " + "initially distributes the walkers. Default is 0.", +) # add config options configuration.add_workflow_command_line_group(parser) # add module pre-defined options @@ -86,8 +102,7 @@ scheme.verify_processing_options(opts, parser) # check that the output file doesn't already exist if os.path.exists(opts.output_file) and not opts.force: - raise OSError("output-file already exists; use force if you " - "wish to overwrite it.") + raise OSError("output-file already exists; use force if you wish to overwrite it.") # set seed @@ -96,14 +111,13 @@ logging.info("Using seed %i", opts.seed) # we'll silence numpy warnings since they are benign and make for confusing # logging output -numpy.seterr(divide='ignore', invalid='ignore') +numpy.seterr(divide="ignore", invalid="ignore") # get scheme ctx = scheme.from_cli(opts) fft.from_cli(opts) with ctx: - # read configuration file cp_original = configuration.WorkflowConfigParser.from_cli(opts) # some models will interally modify original cp for sampling, @@ -112,13 +126,15 @@ with ctx: cp = cp_original.__deepcopy__(cp_original) # create an empty checkpoint file, if needed - condor_ckpt = cp.has_option('sampler', 'checkpoint-signal') + condor_ckpt = cp.has_option("sampler", "checkpoint-signal") if condor_ckpt: logging.info( "Sampler will exit with signal {} after checkpointing".format( - cp.get('sampler', 'checkpoint-signal'))) + cp.get("sampler", "checkpoint-signal") + ) + ) # create an empty output file to keep condor happy - open(opts.output_file, 'a').close() + open(opts.output_file, "a").close() logging.info("Setting up model") @@ -132,13 +148,17 @@ with ctx: # unless you enjoy angering your cluster admins, # NO SAMPLES FILE IO SHOULD BE DONE PRIOR TO THIS POINT!!! sampler = inference.sampler.load_from_config( - cp, model, output_file=opts.output_file, nprocesses=opts.nprocesses, - use_mpi=opts.use_mpi) + cp, + model, + output_file=opts.output_file, + nprocesses=opts.nprocesses, + use_mpi=opts.use_mpi, + ) # store the checkpoint file if pool.is_main_process(): for fn in [sampler.checkpoint_file, sampler.backup_file]: - with loadfile(fn, 'a') as fp: + with loadfile(fn, "a") as fp: fp.write_config_file(cp_original) # Run the sampler @@ -151,11 +171,11 @@ with ctx: sampler.finalize() if condor_ckpt: - # unlink the empty output file - try: - os.unlink(opts.output_file) - except: - pass + # unlink the empty output file + try: + os.unlink(opts.output_file) + except: + pass # rename checkpoint to output and delete backup logging.info("Moving checkpoint to output") @@ -165,12 +185,12 @@ if not opts.save_backup: os.remove(sampler.backup_file) # write the end time -with sampler.io(opts.output_file, 'a') as fp: +with sampler.io(opts.output_file, "a") as fp: fp.write_run_end_time() if condor_ckpt: - # create an empty checkpoint file - open(sampler.checkpoint_file, 'a').close() + # create an empty checkpoint file + open(sampler.checkpoint_file, "a").close() # exit logging.info("Done") diff --git a/bin/inference/pycbc_inference_create_calibration_config b/bin/inference/pycbc_inference_create_calibration_config index 8d767d85bf5..4b3c1e39fbd 100644 --- a/bin/inference/pycbc_inference_create_calibration_config +++ b/bin/inference/pycbc_inference_create_calibration_config @@ -16,80 +16,115 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Generate the calibration configuration file for a given event/GPS time. +""" +Generate the calibration configuration file for a given event/GPS time. Also plots the calibration envelop for each detector for sanity checks. """ -import os, argparse -import pycbc -import numpy -import matplotlib.pyplot as plt +import argparse import logging +import os -from pycbc.types import MultiDetOptionAction -from pycbc.strain.recalibrate import get_calibration_files, read_calibration_envelop_file +import matplotlib.pyplot as plt +import numpy +import pycbc +from pycbc.strain.recalibrate import ( + get_calibration_files, + read_calibration_envelop_file, +) +from pycbc.types import MultiDetOptionAction parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--calibration-files-path', - help='Calibration envelope file path. It can be pre-' - 'downloaded or downloaded on the fly', - type=str) - - -parser.add_argument('--calibration-files', nargs='+', action=MultiDetOptionAction, - metavar = 'DETECTOR:COLUMN', type=str, - help='Calibration file for each detector',) - -parser.add_argument('--ifos', nargs='+', - help='Specify ifos for the analysis') - -parser.add_argument('--minimum-frequency', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', type=float, - help='For each detector, provide minimum frequency for' - 'the analysis. The minimum frequency should be higher than' - 'the minimum frequency in calibration envelop file so that' - 'we do not extrapolate.') - - -parser.add_argument('--maximum-frequency', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', type=float, - help='For each detector, provide maximum frequency for' - 'the analysis. The minimum frequency should be lower than' - 'the maximum frequency in calibration envelop file so that' - 'we do not extrapolate.') - - -parser.add_argument('--gps-time', - help='GPS time at which calibration configuration files' - 'are needed', - required=True, type=float) - -parser.add_argument('--output-dir', - help="Path for the output calibration config files", - default='.', type=str) - -parser.add_argument('--plots-dir', - help="Path for the 'sanity check' plots", - default='.', type=str) - -parser.add_argument('--n-nodes', - help='Number of frequency nodes to be used', - default=10, type=int) - -parser.add_argument('--correction-type', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', - help="Provide the correction type for each detector", - required=True) - -parser.add_argument('--plot-sanity-checks', help='Make some plots for sanity checks', - default=True, type=bool) - -parser.add_argument('--tag', help='Provide your tag for naming calibration file', - default=None) +parser.add_argument( + "--calibration-files-path", + help="Calibration envelope file path. It can be pre-" + "downloaded or downloaded on the fly", + type=str, +) + + +parser.add_argument( + "--calibration-files", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=str, + help="Calibration file for each detector", +) + +parser.add_argument("--ifos", nargs="+", help="Specify ifos for the analysis") + +parser.add_argument( + "--minimum-frequency", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=float, + help="For each detector, provide minimum frequency for" + "the analysis. The minimum frequency should be higher than" + "the minimum frequency in calibration envelop file so that" + "we do not extrapolate.", +) + + +parser.add_argument( + "--maximum-frequency", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=float, + help="For each detector, provide maximum frequency for" + "the analysis. The minimum frequency should be lower than" + "the maximum frequency in calibration envelop file so that" + "we do not extrapolate.", +) + + +parser.add_argument( + "--gps-time", + help="GPS time at which calibration configuration filesare needed", + required=True, + type=float, +) + +parser.add_argument( + "--output-dir", + help="Path for the output calibration config files", + default=".", + type=str, +) + +parser.add_argument( + "--plots-dir", help="Path for the 'sanity check' plots", default=".", type=str +) + +parser.add_argument( + "--n-nodes", help="Number of frequency nodes to be used", default=10, type=int +) + +parser.add_argument( + "--correction-type", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + help="Provide the correction type for each detector", + required=True, +) + +parser.add_argument( + "--plot-sanity-checks", + help="Make some plots for sanity checks", + default=True, + type=bool, +) + +parser.add_argument( + "--tag", help="Provide your tag for naming calibration file", default=None +) # parse the command line @@ -97,8 +132,10 @@ opts = parser.parse_args() if opts.calibration_files_path and opts.calibration_files: - raise ValueError("Only one of the arguments should be provided among:" - "--calibration-files-path and --calibration-files") + raise ValueError( + "Only one of the arguments should be provided among:" + "--calibration-files-path and --calibration-files" + ) calibration_files_path = opts.calibration_files_path @@ -115,7 +152,7 @@ else: tag = opts.tag for ifo in opts.ifos: - logging.info("Minimum %s frequency: %.3f", ifo, min_freq_list[ifo]) + logging.info("Minimum %s frequency: %.3f", ifo, min_freq_list[ifo]) # Read calibration files if opts.calibration_files: @@ -123,17 +160,21 @@ if opts.calibration_files: for ifo in opts.ifos: calibration_files_dict[ifo] = opts.calibration_files[ifo] else: - calibration_files_dict = get_calibration_files(opts.ifos, gps_time, calibration_files_path) + calibration_files_dict = get_calibration_files( + opts.ifos, gps_time, calibration_files_path + ) logging.debug("Calibration files dictionary: %s", calibration_files_dict) input_dict = {} for ii in range(len(opts.ifos)): ifo = str(opts.ifos[ii]) - input_dict['min_freq_%s'%ifo]=min_freq_list[ifo] - input_dict['max_freq_%s'%ifo]=max_freq_list[ifo] - input_dict['correction_type_%s'%ifo]=correction_type_list[ifo] - input_dict['calibration_envelop_file_%s'%ifo] = calibration_files_dict[ifo.upper()] + input_dict["min_freq_%s" % ifo] = min_freq_list[ifo] + input_dict["max_freq_%s" % ifo] = max_freq_list[ifo] + input_dict["correction_type_%s" % ifo] = correction_type_list[ifo] + input_dict["calibration_envelop_file_%s" % ifo] = calibration_files_dict[ + ifo.upper() + ] logging.debug("Input dictionary: %s", input_dict) # Upper and Lower indices of detectors seems unnecessary # but they are there to make sure the correct convention @@ -141,109 +182,148 @@ logging.debug("Input dictionary: %s", input_dict) # file as well as in detector notation. prior_dict = {} for ifo in opts.ifos: - #ifo = ifo.lower() - min_freq = input_dict['min_freq_%s'%ifo] - max_freq = input_dict['max_freq_%s'%ifo] - correction_type = input_dict['correction_type_%s'%ifo] - calib_env_file = input_dict['calibration_envelop_file_%s'%ifo] - log_nodes, amplitude_median_nodes, amplitude_sigma_nodes, phase_median_nodes, phase_sigma_nodes = \ - read_calibration_envelop_file(calib_env_file, correction_type, - min_freq, max_freq, n_nodes) - prior_dict['log_nodes_%s'%ifo] = log_nodes - prior_dict['amplitude_median_nodes_%s'%ifo] = amplitude_median_nodes - prior_dict['amplitude_sigma_nodes_%s'%ifo] = amplitude_sigma_nodes - prior_dict['phase_median_nodes_%s'%ifo] = phase_median_nodes - prior_dict['phase_sigma_nodes_%s'%ifo] =phase_sigma_nodes - logging.info('Done for detector %s', ifo) - - -calibration_filename = '%s/calibration-%s.ini'%(output_dir, tag) + # ifo = ifo.lower() + min_freq = input_dict["min_freq_%s" % ifo] + max_freq = input_dict["max_freq_%s" % ifo] + correction_type = input_dict["correction_type_%s" % ifo] + calib_env_file = input_dict["calibration_envelop_file_%s" % ifo] + ( + log_nodes, + amplitude_median_nodes, + amplitude_sigma_nodes, + phase_median_nodes, + phase_sigma_nodes, + ) = read_calibration_envelop_file( + calib_env_file, correction_type, min_freq, max_freq, n_nodes + ) + prior_dict["log_nodes_%s" % ifo] = log_nodes + prior_dict["amplitude_median_nodes_%s" % ifo] = amplitude_median_nodes + prior_dict["amplitude_sigma_nodes_%s" % ifo] = amplitude_sigma_nodes + prior_dict["phase_median_nodes_%s" % ifo] = phase_median_nodes + prior_dict["phase_sigma_nodes_%s" % ifo] = phase_sigma_nodes + logging.info("Done for detector %s", ifo) + + +calibration_filename = "%s/calibration-%s.ini" % (output_dir, tag) if os.path.isfile(calibration_filename) != True: text_file = open(calibration_filename, "w") text_file.write("[calibration] \n") for ifo in opts.ifos: - min_freq = input_dict['min_freq_%s'%ifo] - max_freq = input_dict['max_freq_%s'%ifo] - text_file.write("%s_model = cubic_spline \n"%ifo.lower()) - text_file.write("%s_minimum_frequency = %d \n"%(ifo.lower(),min_freq)) - text_file.write("%s_maximum_frequency = %d \n"%(ifo.lower(),max_freq)) - text_file.write("%s_n_points = %d \n"%(ifo.lower(),n_nodes)) + min_freq = input_dict["min_freq_%s" % ifo] + max_freq = input_dict["max_freq_%s" % ifo] + text_file.write("%s_model = cubic_spline \n" % ifo.lower()) + text_file.write("%s_minimum_frequency = %d \n" % (ifo.lower(), min_freq)) + text_file.write("%s_maximum_frequency = %d \n" % (ifo.lower(), max_freq)) + text_file.write("%s_n_points = %d \n" % (ifo.lower(), n_nodes)) text_file.write(" \n") text_file.write("[variable_params] \n") for ifo in opts.ifos: for ii in range(n_nodes): - text_file.write("recalib_amplitude_%s_%d = \n"%(ifo.lower(),ii)) - text_file.write("recalib_phase_%s_%d = \n"%(ifo.lower(),ii)) + text_file.write("recalib_amplitude_%s_%d = \n" % (ifo.lower(), ii)) + text_file.write("recalib_phase_%s_%d = \n" % (ifo.lower(), ii)) text_file.write(" \n") for ifo in opts.ifos: for ii in range(n_nodes): - amplitude_median_nodes = prior_dict['amplitude_median_nodes_%s'%ifo][ii] - amplitude_sigma_nodes = prior_dict['amplitude_sigma_nodes_%s'%ifo][ii] - text_file.write("[prior-recalib_amplitude_%s_%d] \n"%(ifo.lower(),ii)) + amplitude_median_nodes = prior_dict["amplitude_median_nodes_%s" % ifo][ii] + amplitude_sigma_nodes = prior_dict["amplitude_sigma_nodes_%s" % ifo][ii] + text_file.write("[prior-recalib_amplitude_%s_%d] \n" % (ifo.lower(), ii)) text_file.write("name = gaussian \n") - text_file.write("recalib_amplitude_%s_%d_mean = %.3g \n"%(ifo.lower(),ii,amplitude_median_nodes)) - text_file.write("recalib_amplitude_%s_%d_var = %.3g \n"%(ifo.lower(),ii,amplitude_sigma_nodes**2)) + text_file.write( + "recalib_amplitude_%s_%d_mean = %.3g \n" + % (ifo.lower(), ii, amplitude_median_nodes) + ) + text_file.write( + "recalib_amplitude_%s_%d_var = %.3g \n" + % (ifo.lower(), ii, amplitude_sigma_nodes**2) + ) text_file.write(" \n") for ifo in opts.ifos: for ii in range(n_nodes): - phase_median_nodes = prior_dict['phase_median_nodes_%s'%ifo][ii] - phase_sigma_nodes = prior_dict['phase_sigma_nodes_%s'%ifo][ii] - text_file.write("[prior-recalib_phase_%s_%d] \n"%(ifo.lower(),ii)) + phase_median_nodes = prior_dict["phase_median_nodes_%s" % ifo][ii] + phase_sigma_nodes = prior_dict["phase_sigma_nodes_%s" % ifo][ii] + text_file.write("[prior-recalib_phase_%s_%d] \n" % (ifo.lower(), ii)) text_file.write("name = gaussian \n") - text_file.write("recalib_phase_%s_%d_mean = %.3g \n"%(ifo.lower(),ii,phase_median_nodes)) - text_file.write("recalib_phase_%s_%d_var = %.3g \n"%(ifo.lower(),ii,phase_sigma_nodes**2)) + text_file.write( + "recalib_phase_%s_%d_mean = %.3g \n" + % (ifo.lower(), ii, phase_median_nodes) + ) + text_file.write( + "recalib_phase_%s_%d_var = %.3g \n" + % (ifo.lower(), ii, phase_sigma_nodes**2) + ) text_file.write(" \n") text_file.write(" \n") text_file.write(" \n") text_file.close() else: raise ValueError( - 'Calibration file already present. ' - 'Either delete it or rename it. Overwriting ' - 'this file is not allowed!!!' + "Calibration file already present. " + "Either delete it or rename it. Overwriting " + "this file is not allowed!!!" ) # Sanity check plots -plots_dir = '%s/plots'%output_dir -os.system('mkdir -p %s'%plots_dir) +plots_dir = "%s/plots" % output_dir +os.system("mkdir -p %s" % plots_dir) if opts.plot_sanity_checks: for ifo in opts.ifos: - calib_env_path = input_dict['calibration_envelop_file_%s'%ifo] - min_freq = input_dict['min_freq_%s'%ifo] - max_freq = input_dict['max_freq_%s'%ifo] - log_nodes = prior_dict['log_nodes_%s'%ifo] - amp_median_nodes = prior_dict['amplitude_median_nodes_%s'%ifo] - amp_sigma_nodes = prior_dict['amplitude_sigma_nodes_%s'%ifo] + calib_env_path = input_dict["calibration_envelop_file_%s" % ifo] + min_freq = input_dict["min_freq_%s" % ifo] + max_freq = input_dict["max_freq_%s" % ifo] + log_nodes = prior_dict["log_nodes_%s" % ifo] + amp_median_nodes = prior_dict["amplitude_median_nodes_%s" % ifo] + amp_sigma_nodes = prior_dict["amplitude_sigma_nodes_%s" % ifo] d1 = numpy.loadtxt(calib_env_path) - plt.figure(figsize=(8,5)) - plt.plot(d1[:,0],d1[:,1],label=r'$\mu$') - plt.plot(d1[:,0],d1[:,3],label=r'$\mu-\sigma$') - plt.plot(d1[:,0],d1[:,5],label=r'$\mu+\sigma$') - plt.plot(numpy.exp(log_nodes),amp_median_nodes+1,'.',label=r'$\mu$(Prior)') - plt.plot(numpy.exp(log_nodes),amp_median_nodes+1-amp_sigma_nodes,'.',label=r'$\mu-\sigma$ (Prior)') - plt.plot(numpy.exp(log_nodes),amp_median_nodes+1+amp_sigma_nodes,'.',label=r'$\mu+\sigma$ (Prior)') - plt.xscale('log') - plt.xlabel('Frequency (Hz)') - plt.ylabel('Amplitude') - plt.title('Calibration envelop (Amplitude) for %s'%ifo) + plt.figure(figsize=(8, 5)) + plt.plot(d1[:, 0], d1[:, 1], label=r"$\mu$") + plt.plot(d1[:, 0], d1[:, 3], label=r"$\mu-\sigma$") + plt.plot(d1[:, 0], d1[:, 5], label=r"$\mu+\sigma$") + plt.plot(numpy.exp(log_nodes), amp_median_nodes + 1, ".", label=r"$\mu$(Prior)") + plt.plot( + numpy.exp(log_nodes), + amp_median_nodes + 1 - amp_sigma_nodes, + ".", + label=r"$\mu-\sigma$ (Prior)", + ) + plt.plot( + numpy.exp(log_nodes), + amp_median_nodes + 1 + amp_sigma_nodes, + ".", + label=r"$\mu+\sigma$ (Prior)", + ) + plt.xscale("log") + plt.xlabel("Frequency (Hz)") + plt.ylabel("Amplitude") + plt.title("Calibration envelop (Amplitude) for %s" % ifo) plt.legend() - plt.savefig('%s/calibration_envelop_amplitude_%s_%s.png'%(plots_dir,tag,ifo)) + plt.savefig( + "%s/calibration_envelop_amplitude_%s_%s.png" % (plots_dir, tag, ifo) + ) plt.clf() - phase_median_nodes = prior_dict['phase_median_nodes_%s'%ifo] - phase_sigma_nodes = prior_dict['phase_sigma_nodes_%s'%ifo] - plt.figure(figsize=(8,5)) - plt.plot(d1[:,0],d1[:,2],label=r'$\mu$') - plt.plot(d1[:,0],d1[:,4],label=r'$\mu-\sigma$') - plt.plot(d1[:,0],d1[:,6],label=r'$\mu+\sigma$') - plt.plot(numpy.exp(log_nodes),phase_median_nodes,'.',label=r'$\mu$ (Prior)') - plt.plot(numpy.exp(log_nodes),phase_median_nodes-phase_sigma_nodes,'.',label=r'$\mu-\sigma$ (Prior)') - plt.plot(numpy.exp(log_nodes),phase_median_nodes+phase_sigma_nodes,'.',label=r'$\mu+\sigma$') - plt.xscale('log') - plt.xlabel('Frequency (Hz)') - plt.ylabel('Phase') - plt.title('Calibration envelop (Phase) for %s'%ifo) + phase_median_nodes = prior_dict["phase_median_nodes_%s" % ifo] + phase_sigma_nodes = prior_dict["phase_sigma_nodes_%s" % ifo] + plt.figure(figsize=(8, 5)) + plt.plot(d1[:, 0], d1[:, 2], label=r"$\mu$") + plt.plot(d1[:, 0], d1[:, 4], label=r"$\mu-\sigma$") + plt.plot(d1[:, 0], d1[:, 6], label=r"$\mu+\sigma$") + plt.plot(numpy.exp(log_nodes), phase_median_nodes, ".", label=r"$\mu$ (Prior)") + plt.plot( + numpy.exp(log_nodes), + phase_median_nodes - phase_sigma_nodes, + ".", + label=r"$\mu-\sigma$ (Prior)", + ) + plt.plot( + numpy.exp(log_nodes), + phase_median_nodes + phase_sigma_nodes, + ".", + label=r"$\mu+\sigma$", + ) + plt.xscale("log") + plt.xlabel("Frequency (Hz)") + plt.ylabel("Phase") + plt.title("Calibration envelop (Phase) for %s" % ifo) plt.legend() - plt.savefig('%s/calibration_envelop_phase_%s_%s.png'%(plots_dir,tag,ifo)) + plt.savefig("%s/calibration_envelop_phase_%s_%s.png" % (plots_dir, tag, ifo)) plt.clf() - diff --git a/bin/inference/pycbc_inference_create_fits b/bin/inference/pycbc_inference_create_fits index 83bd29109bc..deed0567c8b 100644 --- a/bin/inference/pycbc_inference_create_fits +++ b/bin/inference/pycbc_inference_create_fits @@ -16,7 +16,8 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Wrapper around ligo-skymap-from-samples that creates fits files from a +""" +Wrapper around ligo-skymap-from-samples that creates fits files from a inference or posterior hdf file. Requires ligo.skymap to be installed, which requires python 3. @@ -30,49 +31,67 @@ same as what is used for the --parameters option in pycbc_inference_extract_samples (see that program for details). """ -import os -import numpy import argparse +import os import subprocess +import numpy + from pycbc import add_common_pycbc_options, init_logging from pycbc.inference import io parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--input-file', required=True, - help='The inference or posterior hdf file to load.') -parser.add_argument('--output-file', required=True, - help='The output file name; it must end in ".fits".') -parser.add_argument('--maxpts', type=int, - help='The number of posterior samples to extract for ' - 'making the fits file. Default is to load all. Using ' - 'fewer points will result in faster runtime, but ' - 'more error in the result sky map.') -parser.add_argument('--ra', default='ra', - help='The name of the right ascension parameter in the ' - 'input file. May also provide a constant, or math ' - 'operations on multiple parameters. Default is "ra".') -parser.add_argument('--dec', default='dec', - help='The name of the declination parameter in the ' - 'input file. May also provide a constant, or math ' - 'operations on multiple parameters. Default ' - 'is "dec".') -parser.add_argument('--distance', default='distance', - help='The name of the distance parameter in the input ' - 'file. May also provide a constant, or math ' - 'operations on multiple parameters. Default is ' - '"distance".') -parser.add_argument('--tc', default='tc', - help='The name of the coalesence time parameter in the ' - 'input file. May also provide a constant, or math ' - 'operations on multiple parameters. Default is ' - '"tc".') +parser.add_argument( + "--input-file", required=True, help="The inference or posterior hdf file to load." +) +parser.add_argument( + "--output-file", required=True, help='The output file name; it must end in ".fits".' +) +parser.add_argument( + "--maxpts", + type=int, + help="The number of posterior samples to extract for " + "making the fits file. Default is to load all. Using " + "fewer points will result in faster runtime, but " + "more error in the result sky map.", +) +parser.add_argument( + "--ra", + default="ra", + help="The name of the right ascension parameter in the " + "input file. May also provide a constant, or math " + 'operations on multiple parameters. Default is "ra".', +) +parser.add_argument( + "--dec", + default="dec", + help="The name of the declination parameter in the " + "input file. May also provide a constant, or math " + "operations on multiple parameters. Default " + 'is "dec".', +) +parser.add_argument( + "--distance", + default="distance", + help="The name of the distance parameter in the input " + "file. May also provide a constant, or math " + "operations on multiple parameters. Default is " + '"distance".', +) +parser.add_argument( + "--tc", + default="tc", + help="The name of the coalesence time parameter in the " + "input file. May also provide a constant, or math " + "operations on multiple parameters. Default is " + '"tc".', +) opts = parser.parse_args() init_logging(opts.verbose) -fp = io.loadfile(opts.input_file, 'r') +fp = io.loadfile(opts.input_file, "r") samples = fp.read_samples([opts.ra, opts.dec, opts.distance, opts.tc]) fp.close() @@ -81,20 +100,19 @@ out[:, 0] = samples[opts.ra] out[:, 1] = samples[opts.dec] out[:, 2] = samples[opts.distance] out[:, 3] = samples[opts.tc] -basename = opts.output_file.replace('.fits', '') -txtfile = '{}.dat'.format(basename) -numpy.savetxt(txtfile, out, header='ra dec distance time') +basename = opts.output_file.replace(".fits", "") +txtfile = f"{basename}.dat" +numpy.savetxt(txtfile, out, header="ra dec distance time") if opts.maxpts is not None: - maxptsarg = '--maxpts {} '.format(opts.maxpts) + maxptsarg = f"--maxpts {opts.maxpts} " else: - maxptsarg = ' ' -cmd = 'ligo-skymap-from-samples {}{} --fitsoutname {}'.format( - maxptsarg, txtfile, opts.output_file) + maxptsarg = " " +cmd = f"ligo-skymap-from-samples {maxptsarg}{txtfile} --fitsoutname {opts.output_file}" ret = subprocess.run(cmd.split()) os.remove(txtfile) ret.check_returncode() try: - os.remove('skypost.obj') + os.remove("skypost.obj") except: pass diff --git a/bin/inference/pycbc_inference_extract_samples b/bin/inference/pycbc_inference_extract_samples index 1137d41edaf..7aab24c9d12 100644 --- a/bin/inference/pycbc_inference_extract_samples +++ b/bin/inference/pycbc_inference_extract_samples @@ -24,7 +24,8 @@ # # ============================================================================= # -"""Extracts posterior samples from a Sampler HDF file, writing to a +""" +Extracts posterior samples from a Sampler HDF file, writing to a posterior HDF file. Parameters can be renamed in the output file using the --parameters option. @@ -38,18 +39,23 @@ across combined files. """ import os +import warnings + import numpy +from scipy.special import logsumexp + import pycbc -from pycbc.inference.io import (ResultsArgumentParser, results_from_cli, - PosteriorFile, loadfile) +from pycbc.inference.io import ( + PosteriorFile, + ResultsArgumentParser, + loadfile, + results_from_cli, +) from pycbc.inference.io.base_hdf import format_attr -from scipy.special import logsumexp -import warnings def isthesame(current_val, val): - """Checks that attrs values from two hdf files are the same (or nearly so). - """ + """Checks that attrs values from two hdf files are the same (or nearly so).""" if isinstance(current_val, float): issame = numpy.isclose(current_val, val, equal_nan=True) elif isinstance(current_val, numpy.ndarray): @@ -69,51 +75,68 @@ def isthesame(current_val, val): # we didn't have an array, assume it was a single value return issame -parser = ResultsArgumentParser(defaultparams='all', autoparamlabels=False, - description=__doc__) + +parser = ResultsArgumentParser( + defaultparams="all", autoparamlabels=False, description=__doc__ +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--output-file", type=str, required=True, - help="Output file to create.") -parser.add_argument("--force", action="store_true", default=False, - help="If the output-file already exists, overwrite it. " - "Otherwise, an IOError is raised.") -parser.add_argument("--skip-groups", default=None, nargs="+", - help="Don't write the specified groups in the output " - "(aside from samples; samples are always written), " - "for example, 'sampler_info'. If 'all' skip " - "all groups, only writing the samples. Default is " - "to write all groups if only one file is provided, " - "and all groups from the first file except " - "sampler_info if multiple files are provided.") -parser.add_argument("--combine-via-sampling", action='store_true', - default=False, - help="Specify whether to combine the posteriors by " - "sampling. By default, extract_samples will dump all " - "samples from multiple inputs into one output file. " - "If this option is specified, the samples will " - "be randomly sampled with weighting based on the " - "evidence in each input. For example, if the " - "evidence of input 1 is twice that of input 2, " - "the resulting posterior file with this option " - "specified will have twice as many points from input " - "1 than input 2. The output will have the same " - "number of samples as the smallest input file. " - "An error is thrown if this option is specified and " - "any of the input files does not have a log_evidence " - "attribute (e.g., the file used a sampler like emcee " - "that does not report evidence). This option " - "assumes that the priors of all files do not " - "overlap; we cannot properly combine posteriors if " - "the priors do overlap.") -parser.add_argument("--mutually-exclusive-priors", action='store_true', - default=False, - help="If specifying --combine-via-sampling, specify " - "whether to treat priors as mutually exclusive. By " - "default, the provided input files are assumed to " - "have identical priors, and their evidences will be " - "averaged. If this option is specified, the priors " - "are assumed to be non-overlapping across files, and " - "their evidences will be summed together.") +parser.add_argument( + "--output-file", type=str, required=True, help="Output file to create." +) +parser.add_argument( + "--force", + action="store_true", + default=False, + help="If the output-file already exists, overwrite it. " + "Otherwise, an IOError is raised.", +) +parser.add_argument( + "--skip-groups", + default=None, + nargs="+", + help="Don't write the specified groups in the output " + "(aside from samples; samples are always written), " + "for example, 'sampler_info'. If 'all' skip " + "all groups, only writing the samples. Default is " + "to write all groups if only one file is provided, " + "and all groups from the first file except " + "sampler_info if multiple files are provided.", +) +parser.add_argument( + "--combine-via-sampling", + action="store_true", + default=False, + help="Specify whether to combine the posteriors by " + "sampling. By default, extract_samples will dump all " + "samples from multiple inputs into one output file. " + "If this option is specified, the samples will " + "be randomly sampled with weighting based on the " + "evidence in each input. For example, if the " + "evidence of input 1 is twice that of input 2, " + "the resulting posterior file with this option " + "specified will have twice as many points from input " + "1 than input 2. The output will have the same " + "number of samples as the smallest input file. " + "An error is thrown if this option is specified and " + "any of the input files does not have a log_evidence " + "attribute (e.g., the file used a sampler like emcee " + "that does not report evidence). This option " + "assumes that the priors of all files do not " + "overlap; we cannot properly combine posteriors if " + "the priors do overlap.", +) +parser.add_argument( + "--mutually-exclusive-priors", + action="store_true", + default=False, + help="If specifying --combine-via-sampling, specify " + "whether to treat priors as mutually exclusive. By " + "default, the provided input files are assumed to " + "have identical priors, and their evidences will be " + "averaged. If this option is specified, the priors " + "are assumed to be non-overlapping across files, and " + "their evidences will be summed together.", +) opts = parser.parse_args() @@ -121,8 +144,7 @@ pycbc.init_logging(opts.verbose) # check that the output doesn't already exist if os.path.exists(opts.output_file) and not opts.force: - raise IOError("output file already exists; use --force if you wish to " - "overwrite.") + raise OSError("output file already exists; use --force if you wish to overwrite.") # load the samples fps, params, labels, samples = results_from_cli(opts) @@ -133,25 +155,28 @@ if len(opts.input_file) == 1: # also stack results if multiple files were provided if len(opts.input_file) > 1: if opts.combine_via_sampling: - raw_samples = {labels[p]: numpy.concatenate([s[p] for s in samples]) - for p in params} + raw_samples = { + labels[p]: numpy.concatenate([s[p] for s in samples]) for p in params + } logz_list = [] dlogz_list = [] len_list = [] raw_samps_list = [] weights_list = [] for file in opts.input_file: - fp = loadfile(file, 'r') + fp = loadfile(file, "r") # get evidence from each file if possible try: logz, dlogz = fp.log_evidence except KeyError: - raise ValueError(f"Cannot combine evidences; file {file} " - "does not have a log_evidence attr") + raise ValueError( + f"Cannot combine evidences; file {file} " + "does not have a log_evidence attr" + ) logz_list.append(logz) dlogz_list.append(dlogz) # get samples from each file - file_samps = fp.read_samples(list(fp['samples'].keys())) + file_samps = fp.read_samples(list(fp["samples"].keys())) raw_samps_list.append(file_samps) # get the number of samples from each file len_list.append(len(file_samps)) @@ -162,26 +187,28 @@ if len(opts.input_file) > 1: for i in range(len(opts.input_file)): # weight each file's samples according to logz logwt = logz_list[i] - logz_net - weights_list.append([numpy.exp(logwt)/len_list[i] for j in - range(len_list[i])]) + weights_list.append( + [numpy.exp(logwt) / len_list[i] for j in range(len_list[i])] + ) # randomly sample indices from all samples weights = numpy.concatenate(weights_list) - idx = numpy.random.choice(int(len_net), size=out_size, replace=True, - p=weights) - samples = {param: raw_samples[param][idx] for param in - raw_samples.keys()} + idx = numpy.random.choice(int(len_net), size=out_size, replace=True, p=weights) + samples = {param: raw_samples[param][idx] for param in raw_samples} else: - samples = {labels[p]: numpy.concatenate([s[p] for s in samples]) - for p in params} + samples = { + labels[p]: numpy.concatenate([s[p] for s in samples]) for p in params + } else: samples = {labels[p]: samples[p] for p in params} if opts.combine_via_sampling: - warnings.warn("Specified combine_via_sampling with only one input " - "file. This option will have no effect.") + warnings.warn( + "Specified combine_via_sampling with only one input " + "file. This option will have no effect." + ) # create the file outtype = PosteriorFile.name -out = loadfile(opts.output_file, 'w', filetype=outtype) +out = loadfile(opts.output_file, "w", filetype=outtype) # write the samples out.write_samples(samples) @@ -197,9 +224,10 @@ for fp in fps: current_val = val # enforce that the metadata must be the same across multiple files if not isthesame(current_val, val): - raise ValueError("cannot combine all files; samples group attr {} " - "is not the same across all files. ({} vs {})" - .format(key, current_val, val)) + raise ValueError( + f"cannot combine all files; samples group attr {key} " + f"is not the same across all files. ({current_val} vs {val})" + ) # Preserve top-level metadata... # ...except for: the filetype, since that's set when the file was loaded; @@ -207,14 +235,23 @@ for fp in fps: # cmd (we will replace it with this program's command) # resume points (may not be the same across multiple files) # effective nsamples (that can just be obtained from the samples size) -skip_attrs = ['filetype', 'thin_start', 'thin_interval', 'thin_end', - 'thinned_by', 'cmd', 'resume_points', 'effective_nsamples', - 'run_start_time', 'run_end_time'] +skip_attrs = [ + "filetype", + "thin_start", + "thin_interval", + "thin_end", + "thinned_by", + "cmd", + "resume_points", + "effective_nsamples", + "run_start_time", + "run_end_time", +] # also skip evidence if multiple files are being combined; this will be handled # via sampling if specified if len(opts.input_file) > 1: - skip_attrs += ['log_evidence', 'dlog_evidence'] - + skip_attrs += ["log_evidence", "dlog_evidence"] + # make sure attrs are the same between files... cat_params = False for fp in fps: @@ -227,48 +264,50 @@ for fp in fps: out.attrs[key] = val current_val = format_attr(out.attrs[key]) if not isthesame(current_val, val): - if key == 'remapped_params': + if key == "remapped_params": # ...unless it's remapped_params; just save first file's # entries if they don't match between files - warnings.warn("WARNING: remapped_params metadata does not " - "match between files; saving metadata from " - "first file") + warnings.warn( + "WARNING: remapped_params metadata does not " + "match between files; saving metadata from " + "first file" + ) else: - raise ValueError("cannot combine all files; file attr {} is " - "not the same across all files ({} vs {})" - .format(key, current_val, val)) + raise ValueError( + f"cannot combine all files; file attr {key} is " + f"not the same across all files ({current_val} vs {val})" + ) # store what parameters were renamed (if not already saved above) if not cat_params: - out.attrs['remapped_params'] = list(labels.items()) + out.attrs["remapped_params"] = list(labels.items()) # write combined evidence and dlog evidence if combining via sampling if opts.combine_via_sampling: if len(opts.input_file) == 1: # there's only one file; just return what's in the input - out.attrs['log_evidence'] = fps[0].log_evidence[0] - out.attrs['dlog_evidence'] = fps[0].log_evidence[1] + out.attrs["log_evidence"] = fps[0].log_evidence[0] + out.attrs["dlog_evidence"] = fps[0].log_evidence[1] elif opts.mutually_exclusive_priors: # add together the evidences; quadrature sum the residuals - out.attrs['log_evidence'] = logsumexp(logz_list) - out.attrs['dlog_evidence'] = numpy.sqrt(sum([i**2 for i in dlogz_list])) + out.attrs["log_evidence"] = logsumexp(logz_list) + out.attrs["dlog_evidence"] = numpy.sqrt(sum([i**2 for i in dlogz_list])) else: # average the evidences; quadrature sum and scale the residuals n = len(opts.input_file) - out.attrs['log_evidence'] = logsumexp(logz_list) - numpy.log(n) - out.attrs['dlog_evidence'] = numpy.sqrt(sum([i**2 for i in dlogz_list])) / n + out.attrs["log_evidence"] = logsumexp(logz_list) - numpy.log(n) + out.attrs["dlog_evidence"] = numpy.sqrt(sum([i**2 for i in dlogz_list])) / n # write the other groups using the first file fp = fps[0] skip_groups = opts.skip_groups -if skip_groups is not None and 'all' in opts.skip_groups: - skip_groups = [group for group in fp.keys() - if group != fp.samples_group] +if skip_groups is not None and "all" in opts.skip_groups: + skip_groups = [group for group in fp.keys() if group != fp.samples_group] # don't write the sampler info if more than one file was provided if len(opts.input_file) > 1: if skip_groups is None: skip_groups = [] - skip_groups.append('sampler_info') + skip_groups.append("sampler_info") fp.copy_info(out, ignore=skip_groups) diff --git a/bin/inference/pycbc_inference_model_stats b/bin/inference/pycbc_inference_model_stats index ba152ceb021..48d1aee6979 100644 --- a/bin/inference/pycbc_inference_model_stats +++ b/bin/inference/pycbc_inference_model_stats @@ -16,44 +16,64 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Recalculates log likelihood and prior for points in a given inference or +""" +Recalculates log likelihood and prior for points in a given inference or posterior file and writes them to a new file. Also records auxillary model stats that may have been ignored by the sampler. """ -import os import argparse -import shutil import logging +import os +import shutil + import numpy import tqdm import pycbc -from pycbc.pool import (use_mpi, choose_pool) -from pycbc.io import FieldArray -from pycbc.inference.io import loadfile from pycbc.inference import models +from pycbc.inference.io import loadfile +from pycbc.io import FieldArray +from pycbc.pool import choose_pool, use_mpi from pycbc.workflow import WorkflowConfigParser - parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-file", type=str, required=True, - help="Input HDF file to read. Can be either an inference " - "file or a posterior file.") -parser.add_argument("--output-file", type=str, required=True, - help="Output file to create.") -parser.add_argument("--force", action="store_true", default=False, - help="If the output-file already exists, overwrite it. " - "Otherwise, an OSError is raised.") -parser.add_argument("--nprocesses", type=int, default=1, - help="Number of processes to use. If not given then only " - "a single core will be used.") -parser.add_argument("--config-file", nargs="+", type=str, default=None, - help="Override the config file stored in the input file " - "with the given file(s).") -parser.add_argument('--reconstruct-parameters', action="store_true", - help="Reconstruct marginalized parameters") +parser.add_argument( + "--input-file", + type=str, + required=True, + help="Input HDF file to read. Can be either an inference file or a posterior file.", +) +parser.add_argument( + "--output-file", type=str, required=True, help="Output file to create." +) +parser.add_argument( + "--force", + action="store_true", + default=False, + help="If the output-file already exists, overwrite it. " + "Otherwise, an OSError is raised.", +) +parser.add_argument( + "--nprocesses", + type=int, + default=1, + help="Number of processes to use. If not given then only " + "a single core will be used.", +) +parser.add_argument( + "--config-file", + nargs="+", + type=str, + default=None, + help="Override the config file stored in the input file with the given file(s).", +) +parser.add_argument( + "--reconstruct-parameters", + action="store_true", + help="Reconstruct marginalized parameters", +) # parse command line @@ -61,8 +81,7 @@ opts = parser.parse_args() # check that the output file doesn't already exist if os.path.exists(opts.output_file) and not opts.force: - raise OSError("output-file already exists; use force if you " - "wish to overwrite it.") + raise OSError("output-file already exists; use force if you wish to overwrite it.") # setup log # If we're running in MPI mode, only allow the parent to print @@ -76,11 +95,13 @@ logging.info("Loading config file") if opts.config_file is None: # try to load the config file from the input file try: - with loadfile(opts.input_file, 'r') as fp: + with loadfile(opts.input_file, "r") as fp: cp = fp.read_config_file() except ValueError: - raise ValueError("no config file found in {}; please provide one " - "using the --config-file".format(opts.input_file)) + raise ValueError( + f"no config file found in {opts.input_file}; please provide one " + "using the --config-file" + ) else: cp = WorkflowConfigParser(opts.config_file) @@ -91,6 +112,7 @@ model = models.read_from_config(cp) # from the variable parameter space model.sampling_transforms = None + # create function for calling the model to get the stats def callmodel(arg): iteration, paramvals = arg @@ -106,14 +128,15 @@ def callmodel(arg): rec = model.reconstruct(seed=iteration) return stats, rec + # these help for parallelization for MPI models._global_instance = callmodel model_call = models._call_global_model pool = choose_pool(processes=opts.nprocesses) -logging.info('Getting samples') -with loadfile(opts.input_file, 'r') as fp: +logging.info("Getting samples") +with loadfile(opts.input_file, "r") as fp: # we'll need the shape; all the arrays in the samples group should have the # same shape pick = list(fp[fp.samples_group].keys())[0] @@ -128,8 +151,7 @@ logging.info("Loaded %i samples", samples.size) # get the stats logging.info("Calculating stats") -data = list(tqdm.tqdm(pool.imap(model_call, enumerate(samples)), - total=len(samples))) +data = list(tqdm.tqdm(pool.imap(model_call, enumerate(samples)), total=len(samples))) stats = [x[0] for x in data] rec = [x[1] for x in data] @@ -138,7 +160,7 @@ logging.info("Copying input to output") shutil.copy(opts.input_file, opts.output_file) logging.info("Writing stats to output") -out = loadfile(opts.output_file, 'a') +out = loadfile(opts.output_file, "a") idx = range(len(stats)) for pi, p in enumerate(model.default_stats): vals = numpy.array([stats[ii][pi] for ii in idx]).reshape(shape) diff --git a/bin/inference/pycbc_inference_monitor b/bin/inference/pycbc_inference_monitor index 090bdf6a9ff..b960a276334 100644 --- a/bin/inference/pycbc_inference_monitor +++ b/bin/inference/pycbc_inference_monitor @@ -1,45 +1,60 @@ #!/bin/env python -""" Monitor and create updating results page of an inference run -""" +"""Monitor and create updating results page of an inference run""" -import os, sys, time, glob, shutil, logging, argparse, pycbc +import argparse +import glob +import logging +import os +import shutil +import sys +import time + +import pycbc import pycbc.workflow as wf parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--inference-file', help="The name of the inference file") -parser.add_argument('--check-interval', default=10, type=int, - help="Polling interval to check for file changes (s)") -parser.add_argument('--output-dir', help="Output plots / results directory") -parser.add_argument('--allow-failure', action='store_true', - help="Allow for a failure in plot generation") +parser.add_argument("--inference-file", help="The name of the inference file") +parser.add_argument( + "--check-interval", + default=10, + type=int, + help="Polling interval to check for file changes (s)", +) +parser.add_argument("--output-dir", help="Output plots / results directory") +parser.add_argument( + "--allow-failure", + action="store_true", + help="Allow for a failure in plot generation", +) wf.add_workflow_command_line_group(parser) args = parser.parse_args() pycbc.init_logging(args.verbose) -c = wf.Workflow(args, 'results') -exes = c.cp.options('executables') -exes = [e for e in exes if e != 'results_page'] +c = wf.Workflow(args, "results") +exes = c.cp.options("executables") +exes = [e for e in exes if e != "results_page"] wf.makedir(args.output_dir) -wf.makedir(args.output_dir + '/logs') +wf.makedir(args.output_dir + "/logs") + +pnode = wf.Executable(c.cp, "results_page", out_dir=args.output_dir).create_node() +pnode.add_opt("--plots-dir", "./plots") +pnode.add_opt("--output-path", "./html") -pnode = wf.Executable(c.cp, 'results_page', out_dir=args.output_dir).create_node() -pnode.add_opt('--plots-dir', './plots') -pnode.add_opt('--output-path', './html') def process_plots(fname, subdir): - outdir = args.output_dir + '/plots/checkpoints/' + subdir + outdir = args.output_dir + "/plots/checkpoints/" + subdir wf.makedir(outdir) - wf.makedir(outdir + '/logs') + wf.makedir(outdir + "/logs") nodes = [] for exe in exes: secs = c.cp.get_subsections(exe) for sec in secs: e = wf.Executable(c.cp, exe, out_dir=outdir, tags=[sec]) node = e.create_node() - node.add_opt('--input-file', os.path.abspath(fname)) - node.add_opt('--output-file', exe + '-' + sec + '.png') + node.add_opt("--input-file", os.path.abspath(fname)) + node.add_opt("--output-file", exe + "-" + sec + ".png") try: c.execute_node(node) except Exception as e: @@ -47,13 +62,14 @@ def process_plots(fname, subdir): pass else: raise e - fmade = glob.glob(outdir + '/*.png') + fmade = glob.glob(outdir + "/*.png") for fn in fmade: - shutil.copy(fn, args.output_dir + '/plots/' + os.path.basename(fn)) + shutil.copy(fn, args.output_dir + "/plots/" + os.path.basename(fn)) c.execute_node(pnode) + # Look for checkpoint / bkup file using the standard naming convention -fname = args.inference_file + '.bkup' +fname = args.inference_file + ".bkup" last_changed = time.time() iteration = 0 @@ -62,16 +78,15 @@ while 1: if not os.path.isfile(fname): if os.path.isfile(args.inference_file): fname = args.inference_file - logging.info('Run appears to have finished making final plots') + logging.info("Run appears to have finished making final plots") exit = True else: - raise RuntimeError('Could not located inference files, exiting...') + raise RuntimeError("Could not located inference files, exiting...") tmod = os.path.getmtime(fname) if tmod > last_changed: - tstr = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(tmod)) - logging.info("File update %.1f seconds ago: %s", - time.time() - tmod, tstr) + tstr = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(tmod)) + logging.info("File update %.1f seconds ago: %s", time.time() - tmod, tstr) last_changed = tmod process_plots(fname, subdir=str(iteration)) iteration += 1 diff --git a/bin/inference/pycbc_inference_plot_acceptance_rate b/bin/inference/pycbc_inference_plot_acceptance_rate index 8723464b693..1bbfeaeb510 100644 --- a/bin/inference/pycbc_inference_plot_acceptance_rate +++ b/bin/inference/pycbc_inference_plot_acceptance_rate @@ -18,40 +18,60 @@ import argparse import logging + from matplotlib import use -use('agg') + +use("agg") +import sys + import matplotlib.pyplot as plt + import pycbc from pycbc import results from pycbc.inference import io -import sys # add options to command line parser = argparse.ArgumentParser( - usage="pycbc_inference_plot_acceptance_rate [--options]", - description="Plots histogram of the fractions of steps " - "accepted by walkers.") + usage="pycbc_inference_plot_acceptance_rate [--options]", + description="Plots histogram of the fractions of steps accepted by walkers.", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-file", type=str, required=True, - help="Path to input HDF file.") +parser.add_argument( + "--input-file", type=str, required=True, help="Path to input HDF file." +) # output plot options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) # add walkers to plot -parser.add_argument("--walkers", type=int, nargs='+', default=None, - help="Specify walkers whose acceptance fraction would " - "be plotted. The acceptance fraction for a walker is the " - "fraction of steps accepted by it. Default is plot for " - "all walkers.") -parser.add_argument("--temps", type=int, nargs="+", default=None, - help="Specify temperatures whose acceptance fraction " - "would be plotted (if available). [default=None] " - "specifies that all temperatures will be plotted. " - "This will create N histograms on the plot for N " - "temperatures in file.") +parser.add_argument( + "--walkers", + type=int, + nargs="+", + default=None, + help="Specify walkers whose acceptance fraction would " + "be plotted. The acceptance fraction for a walker is the " + "fraction of steps accepted by it. Default is plot for " + "all walkers.", +) +parser.add_argument( + "--temps", + type=int, + nargs="+", + default=None, + help="Specify temperatures whose acceptance fraction " + "would be plotted (if available). [default=None] " + "specifies that all temperatures will be plotted. " + "This will create N histograms on the plot for N " + "temperatures in file.", +) # add number of bins for histogram -parser.add_argument("--bins", type=int, default=10, - help="Specify number of bins for the histogram plot.") +parser.add_argument( + "--bins", + type=int, + default=10, + help="Specify number of bins for the histogram plot.", +) # parse the command line opts = parser.parse_args() @@ -67,10 +87,11 @@ fp = io.loadfile(opts.input_file, "r") # add the temperature arguments if it is specified additional_args = {} if opts.temps is not None: - additional_args['temps'] = opts.temps + additional_args["temps"] = opts.temps -acceptance_fraction = fp.read_acceptance_fraction(walkers=opts.walkers, - **additional_args) +acceptance_fraction = fp.read_acceptance_fraction( + walkers=opts.walkers, **additional_args +) # Close the file fp.close() @@ -87,10 +108,13 @@ plt.xlabel("Mean Acceptance Rate") caption = """This plot shows a histogram of the acceptance rate of the walkers. The acceptance rate of a walker is the fraction of steps accepted by it.""" -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title="Acceptance Rate", - caption=caption) +results.save_fig_with_metadata( + fig, + opts.output_file, + cmd=" ".join(sys.argv), + title="Acceptance Rate", + caption=caption, +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_acf b/bin/inference/pycbc_inference_plot_acf index 3453f4c450e..e2ca22b1260 100644 --- a/bin/inference/pycbc_inference_plot_acf +++ b/bin/inference/pycbc_inference_plot_acf @@ -21,36 +21,44 @@ import sys import numpy from matplotlib import use -use('agg') + +use("agg") from matplotlib import pyplot as plt import pycbc from pycbc import results from pycbc.inference import io - from pycbc.inference.sampler import samplers # command line usage # turn off thin-interval since it won't be used -parser = io.ResultsArgumentParser(skip_args='thin-interval', - description="Plots autocorrelation function " - "from inference samples.") +parser = io.ResultsArgumentParser( + skip_args="thin-interval", + description="Plots autocorrelation function from inference samples.", +) pycbc.add_common_pycbc_options(parser) # output plot options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) # add plotting options -parser.add_argument("--ymin", type=float, - help="Minimum value to plot on y-axis.") -parser.add_argument("--ymax", type=float, - help="Maximum value to plot on y-axis.") -parser.add_argument("--per-walker", action="store_true", default=False, +parser.add_argument("--ymin", type=float, help="Minimum value to plot on y-axis.") +parser.add_argument("--ymax", type=float, help="Maximum value to plot on y-axis.") +parser.add_argument( + "--per-walker", + action="store_true", + default=False, help="Plot the ACF for each walker separately. Default is to average " - "parameter values over the walkers, then compute the ACF for the " - "averaged chain. Warning: turning this option on can significantly " - "increase the run time.") -parser.add_argument("--no-legend", action="store_true", default=False, - help="Do not add a legend to the plot.") + "parameter values over the walkers, then compute the ACF for the " + "averaged chain. Warning: turning this option on can significantly " + "increase the run time.", +) +parser.add_argument( + "--no-legend", + action="store_true", + default=False, + help="Do not add a legend to the plot.", +) # parse the command line opts = parser.parse_args() @@ -61,24 +69,26 @@ pycbc.init_logging(opts.verbose) # add the temperature args if it exists additional_args = {} try: - additional_args['temps'] = opts.temps + additional_args["temps"] = opts.temps except AttributeError: pass # load the results -logging.info('Loading parameters') +logging.info("Loading parameters") fp, parameters, labels, _ = io.results_from_cli(opts, load_samples=False) # calculate autocorrelation function logging.info("Calculating autocorrelation functions") -acfs = samplers[fp.sampler].compute_acf(fp.filename, - start_index=opts.thin_start, - end_index=opts.thin_end, - per_walker=opts.per_walker, - walkers=opts.walkers, - parameters=parameters, - **additional_args) +acfs = samplers[fp.sampler].compute_acf( + fp.filename, + start_index=opts.thin_start, + end_index=opts.thin_end, + per_walker=opts.per_walker, + walkers=opts.walkers, + parameters=parameters, + **additional_args, +) # check if we have multiple temperatures if opts.per_walker: @@ -86,42 +96,42 @@ if opts.per_walker: else: multi_tempered = tuple(acfs.values())[0].ndim == 2 try: - temps = additional_args['temps'] - if temps == 'all': + temps = additional_args["temps"] + if temps == "all": temps = range(tuple(acfs.values())[0].shape[0]) if isinstance(temps, int): temps = [temps] except KeyError: temps = [0] -fig_height = max(6, 3*len(temps)) +fig_height = max(6, 3 * len(temps)) # plot autocorrelation logging.info("Plotting autocorrelation functions") -fig = plt.figure(figsize=(8,fig_height)) -for kk,tk in enumerate(temps): +fig = plt.figure(figsize=(8, fig_height)) +for kk, tk in enumerate(temps): if multi_tempered: - logging.info("Temperature {}".format(tk)) + logging.info(f"Temperature {tk}") axnum = len(temps) - kk ax = fig.add_subplot(len(temps), 1, axnum) for param in parameters: - logging.info("Parameter {}".format(param)) + logging.info(f"Parameter {param}") if multi_tempered: - acf = acfs[param][kk,...] + acf = acfs[param][kk, ...] else: acf = acfs[param] if opts.per_walker: - lbl = '{}, walker {}' + lbl = "{}, walker {}" # account for thinning that was done on the fly - xvals = numpy.arange(acf.shape[1])*fp.thinned_by + xvals = numpy.arange(acf.shape[1]) * fp.thinned_by for wi in range(acf.shape[0]): wnum = opts.walkers[wi] if opts.walkers is not None else wi - ax.plot(xvals, acf[wi,:], label=lbl.format(labels[param], wnum)) + ax.plot(xvals, acf[wi, :], label=lbl.format(labels[param], wnum)) else: - xvals = numpy.arange(len(acf))*fp.thinned_by + xvals = numpy.arange(len(acf)) * fp.thinned_by ax.plot(xvals, acf, label=labels[param]) if multi_tempered: - t = 1./fp[fp.sampler_group].attrs['betas'][tk] - ax.set_title(r'$T = {}$ (chain {})'.format(t, tk)) + t = 1.0 / fp[fp.sampler_group].attrs["betas"][tk] + ax.set_title(rf"$T = {t}$ (chain {tk})") if axnum == len(temps): ax.set_xlabel("iteration") ax.set_ylabel("autocorrelation function") @@ -137,14 +147,13 @@ plt.tight_layout() # save figure with meta-data caption_kwargs = { - "parameters" : ", ".join(labels), + "parameters": ", ".join(labels), } caption = """Autocorrelation function (ACF).""" title = "Autocorrelation Function for {parameters}".format(**caption_kwargs) -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title=title, - caption=caption) +results.save_fig_with_metadata( + fig, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_acl b/bin/inference/pycbc_inference_plot_acl index 20a8a692745..d7ef731ed9a 100644 --- a/bin/inference/pycbc_inference_plot_acl +++ b/bin/inference/pycbc_inference_plot_acl @@ -20,29 +20,28 @@ import logging import sys import numpy - from matplotlib import use -use('agg') + +use("agg") from matplotlib import pyplot as plt import pycbc from pycbc import results from pycbc.filter import autocorrelation - from pycbc.inference import io # command line usage # turn off thin-interval and temps since they won't be used -parser = io.ResultsArgumentParser(skip_args=['thin-interval', 'temps'], - description="Histograms autocorrelation " - "length per walker from an MCMC " - "sampler.") +parser = io.ResultsArgumentParser( + skip_args=["thin-interval", "temps"], + description="Histograms autocorrelation length per walker from an MCMC sampler.", +) pycbc.add_common_pycbc_options(parser) # output plot options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") -parser.add_argument("--bins", type=int, default=10, - help="Number of bins in histogram.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) +parser.add_argument("--bins", type=int, default=10, help="Number of bins in histogram.") # parse the command line opts = parser.parse_args() @@ -61,8 +60,13 @@ for param_name in parameters: # for each walker acls = [] for i in range(fp.nwalkers): - y = fp.read_samples(param_name, walkers=i, thin_start=opts.thin_start, - thin_end=opts.thin_end, thin_interval=1) + y = fp.read_samples( + param_name, + walkers=i, + thin_start=opts.thin_start, + thin_end=opts.thin_end, + thin_interval=1, + ) acl = autocorrelation.calculate_acl(y[param_name], dtype=int) if acl == numpy.inf: acl = fp.niterations @@ -74,24 +78,23 @@ for param_name in parameters: plt.hist(acls, opts.bins, histtype="step", label=labels[param_name]) plt.xlabel("Autocorrelation time") -plt.ylabel(r'Number of walkers') +plt.ylabel(r"Number of walkers") # plot autocorrelation time saved in hdf file -plt.axvline(fp.act, linestyle='--') +plt.axvline(fp.act, linestyle="--") plt.legend() # save figure with meta-data caption_kwargs = { - "parameters" : ", ".join(labels), + "parameters": ", ".join(labels), } caption = """ Histograms of the autocorrelation time (ACT) from all the walker chains for the parameters. The vertical dashed line is the ACT read from the input file.""" title = "Autocorrelation time for {parameters}".format(**caption_kwargs) -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title=title, - caption=caption) +results.save_fig_with_metadata( + fig, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_dynesty_run b/bin/inference/pycbc_inference_plot_dynesty_run index 3e1aa6795b2..6e0109d94a7 100644 --- a/bin/inference/pycbc_inference_plot_dynesty_run +++ b/bin/inference/pycbc_inference_plot_dynesty_run @@ -18,27 +18,31 @@ import argparse import logging + from matplotlib import use -use('agg') -from matplotlib import pyplot as plt + +use("agg") +import sys from dynesty import plotting as dyplot +from matplotlib import pyplot as plt -from pycbc.inference import io import pycbc from pycbc import results -import sys +from pycbc.inference import io parser = argparse.ArgumentParser( - usage="pycbc_inference_plot_dynesty_run [--options]", - description="Plots various figures showing evolution of " - "dynesty run.") + usage="pycbc_inference_plot_dynesty_run [--options]", + description="Plots various figures showing evolution of dynesty run.", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-file", type=str, required=True, - help="Path to input HDF file.") +parser.add_argument( + "--input-file", type=str, required=True, help="Path to input HDF file." +) # output plot options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) # parse the command line opts = parser.parse_args() @@ -56,16 +60,20 @@ fp.close() # plot evolution of dynesty run logging.info("Plotting dynesty runplot") -fig, axes = dyplot.runplot(sampler_state.results, - lnz_truth=sampler_state.results['logz'][-1]) -#plt.savefig() +fig, axes = dyplot.runplot( + sampler_state.results, lnz_truth=sampler_state.results["logz"][-1] +) +# plt.savefig() # save figure with meta-data caption = """Set of plots shows evolution of Evidence, nlive points, likelihood values, and importance weight pdf.""" -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title="Dynesty runplots", - caption=caption) +results.save_fig_with_metadata( + fig, + opts.output_file, + cmd=" ".join(sys.argv), + title="Dynesty runplots", + caption=caption, +) plt.close() diff --git a/bin/inference/pycbc_inference_plot_dynesty_traceplot b/bin/inference/pycbc_inference_plot_dynesty_traceplot index 2d6c9a968b6..4584f0ecdfc 100644 --- a/bin/inference/pycbc_inference_plot_dynesty_traceplot +++ b/bin/inference/pycbc_inference_plot_dynesty_traceplot @@ -18,29 +18,32 @@ import argparse import logging + from matplotlib import use -use('agg') -from matplotlib import pyplot as plt + +use("agg") +import sys from dynesty import plotting as dyplot +from matplotlib import pyplot as plt -from pycbc.inference import io import pycbc from pycbc import results -import sys - +from pycbc.inference import io parser = argparse.ArgumentParser( - usage="pycbc_inference_plot_dynesty_traceplots [--options]", - description="Plots various figures showing evolution of " - "dynesty run.") + usage="pycbc_inference_plot_dynesty_traceplots [--options]", + description="Plots various figures showing evolution of dynesty run.", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-file", type=str, required=True, - help="Path to input HDF file.") +parser.add_argument( + "--input-file", type=str, required=True, help="Path to input HDF file." +) # output plot options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) # parse the command line opts = parser.parse_args() @@ -57,17 +60,24 @@ fp.close() # plot evolution of dynesty run logging.info("Plotting dynesty tracelots") -fig, axes = dyplot.traceplot(sampler_state.results, truth_color='black', - show_titles=True, trace_cmap='viridis', - connect=True) -#plt.savefig() +fig, axes = dyplot.traceplot( + sampler_state.results, + truth_color="black", + show_titles=True, + trace_cmap="viridis", + connect=True, +) +# plt.savefig() # save figure with meta-data caption = """Set of plots shows 'evolution of particles (and their marginal posterior distributions) in 1-D projections'.""" -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title="Dynesty traceplots", - caption=caption) +results.save_fig_with_metadata( + fig, + opts.output_file, + cmd=" ".join(sys.argv), + title="Dynesty traceplots", + caption=caption, +) plt.close() diff --git a/bin/inference/pycbc_inference_plot_gelman_rubin b/bin/inference/pycbc_inference_plot_gelman_rubin index b9fe8cfb49a..dca14878086 100644 --- a/bin/inference/pycbc_inference_plot_gelman_rubin +++ b/bin/inference/pycbc_inference_plot_gelman_rubin @@ -14,39 +14,54 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plots the Gelman-Rubin convergence diagnositic statistic. -""" +"""Plots the Gelman-Rubin convergence diagnositic statistic.""" import logging + import matplotlib + matplotlib.use("Agg") -import matplotlib.pyplot as plt import sys -from pycbc import ( - results, init_logging, add_common_pycbc_options -) -from pycbc.inference import (gelman_rubin, io) +import matplotlib.pyplot as plt + +from pycbc import add_common_pycbc_options, init_logging, results +from pycbc.inference import gelman_rubin, io # add options to command line -parser = io.ResultsArgumentParser(skip_args=['walkers']) +parser = io.ResultsArgumentParser(skip_args=["walkers"]) add_common_pycbc_options(parser) # output options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") -parser.add_argument("--walkers", type=int, nargs="+", default=None, - help="Specific walkers to plot. Default is plot " - "all walkers.") -parser.add_argument("--hline", type=float, default=None, - help="Plots a horizontal line.") -parser.add_argument("--segment-start", type=int, required=True, - help="Index in chain to start calculation.") -parser.add_argument("--segment-end", type=int, required=True, - help="Index in chain to end calculation.") -parser.add_argument("--segment-step", type=int, required=True, - help="Step size in chain to next calculation.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) +parser.add_argument( + "--walkers", + type=int, + nargs="+", + default=None, + help="Specific walkers to plot. Default is plot all walkers.", +) +parser.add_argument( + "--hline", type=float, default=None, help="Plots a horizontal line." +) +parser.add_argument( + "--segment-start", + type=int, + required=True, + help="Index in chain to start calculation.", +) +parser.add_argument( + "--segment-end", type=int, required=True, help="Index in chain to end calculation." +) +parser.add_argument( + "--segment-step", + type=int, + required=True, + help="Step size in chain to next calculation.", +) # parse the command line @@ -60,9 +75,7 @@ if opts.iteration is not None: raise ValueError("Cannot use --iteration") # load the results -fp, params, labels, _ = io.results_from_cli( - opts, load_samples=False, - walkers=None) +fp, params, labels, _ = io.results_from_cli(opts, load_samples=False, walkers=None) # if use wants specific walkers walkers = range(fp.nwalkers) if opts.walkers is None else opts.walkers @@ -76,16 +89,22 @@ for param, label in zip(params, labels): logging.info("Plotting parameter %s", param) # get samples for each chain - chains = [fp.read_samples(param, walkers=i, - thin_start=opts.thin_start, - thin_interval=opts.thin_interval, - thin_end=opts.thin_end)[param] for i in walkers] + chains = [ + fp.read_samples( + param, + walkers=i, + thin_start=opts.thin_start, + thin_interval=opts.thin_interval, + thin_end=opts.thin_end, + )[param] + for i in walkers + ] # calculate the Gelman-Rubin convergence diagnostic statistic chains = [chain.reshape(1, len(chain)) for chain in chains] - starts, ends, stats = gelman_rubin.walk(chains, opts.segment_start, - opts.segment_end, - opts.segment_step) + starts, ends, stats = gelman_rubin.walk( + chains, opts.segment_start, opts.segment_end, opts.segment_step + ) # plot plt.plot(ends, stats[0, :], label=label) @@ -101,15 +120,14 @@ if opts.hline: # save figure with meta-data caption_kwargs = { - "parameters" : ", ".join(labels), + "parameters": ", ".join(labels), } caption = """The Gelman-Rubin convergence diagnostic statistic for {parameters} read from the input file.""".format(**caption_kwargs) title = "Gelman-Rubin Convergence for {parameters}".format(**caption_kwargs) -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title=title, - caption=caption) +results.save_fig_with_metadata( + fig, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_geweke b/bin/inference/pycbc_inference_plot_geweke index 49a79474e0e..cf8e5228b44 100644 --- a/bin/inference/pycbc_inference_plot_geweke +++ b/bin/inference/pycbc_inference_plot_geweke @@ -15,48 +15,77 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plots the Geweke convergence diagnositic statistic. -""" +"""Plots the Geweke convergence diagnositic statistic.""" import logging -import matplotlib as mpl; mpl.use("Agg") + +import matplotlib as mpl + +mpl.use("Agg") +import sys + import matplotlib.pyplot as plt + import pycbc from pycbc import results -import sys - -from pycbc.inference import (io, geweke) +from pycbc.inference import geweke, io # add options to command line -parser = io.ResultsArgumentParser(skip_args=['walkers']) +parser = io.ResultsArgumentParser(skip_args=["walkers"]) pycbc.add_common_pycbc_options(parser) # program-specific # output options -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") -parser.add_argument("--walkers", type=int, nargs="+", default=None, - help="Specific walkers to plot. Default is plot " - "all walkers.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) +parser.add_argument( + "--walkers", + type=int, + nargs="+", + default=None, + help="Specific walkers to plot. Default is plot all walkers.", +) # Geweke calculation options -parser.add_argument("--segment-length", type=int, required=True, - help="Number of iterations to use in calculation.") -parser.add_argument("--segment-stride", type=int, required=True, - help="How many iterations to advance for next calculation " - "along chain.") -parser.add_argument("--segment-start", type=int, required=True, - help="Start iteration for calculating statistic.") -parser.add_argument("--segment-end", type=int, required=True, - help="End iteration for calculating statistic.") -parser.add_argument("--reference-start", type=int, required=True, - help="Start of reference segment for " - "calculating statistic.") -parser.add_argument("--reference-end", type=int, required=True, - help="End of reference segment for " - "calculating statistic.") +parser.add_argument( + "--segment-length", + type=int, + required=True, + help="Number of iterations to use in calculation.", +) +parser.add_argument( + "--segment-stride", + type=int, + required=True, + help="How many iterations to advance for next calculation along chain.", +) +parser.add_argument( + "--segment-start", + type=int, + required=True, + help="Start iteration for calculating statistic.", +) +parser.add_argument( + "--segment-end", + type=int, + required=True, + help="End iteration for calculating statistic.", +) +parser.add_argument( + "--reference-start", + type=int, + required=True, + help="Start of reference segment for calculating statistic.", +) +parser.add_argument( + "--reference-end", + type=int, + required=True, + help="End of reference segment for calculating statistic.", +) # parse the command line @@ -71,8 +100,7 @@ if opts.iteration is not None: raise ValueError("Cannot use --iteration") # load the results -fp, params, labels, _ = io.results_from_cli( - opts, load_samples=False) +fp, params, labels, _ = io.results_from_cli(opts, load_samples=False) # get walkers to plot walkers = range(fp.nwalkers) if opts.walkers is None else opts.walkers @@ -96,10 +124,14 @@ for param, label in zip(params, labels): # calculate the Geweke convergence statistic starts, ends, stats = geweke.geweke( - vals, opts.segment_length, opts.segment_stride, - opts.segment_end, opts.reference_start, - ref_end=opts.reference_end, - seg_start=opts.segment_start) + vals, + opts.segment_length, + opts.segment_stride, + opts.segment_end, + opts.reference_start, + ref_end=opts.reference_end, + seg_start=opts.segment_start, + ) # plot walker for start, end, stat in zip(starts, ends, stats): @@ -114,15 +146,14 @@ plt.hlines([-1, 1], 0, len(vals), "r", linestyles="dashed") # save figure with meta-data caption_kwargs = { - "parameters" : ", ".join(labels), + "parameters": ", ".join(labels), } caption = """The Geweke convergence diagnostic statistic for {parameters} read from the input file.""".format(**caption_kwargs) title = "Geweke Convergence for {parameters}".format(**caption_kwargs) -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title=title, - caption=caption) +results.save_fig_with_metadata( + fig, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_inj_recovery b/bin/inference/pycbc_inference_plot_inj_recovery index b82b4d0f9fe..d561aac7315 100644 --- a/bin/inference/pycbc_inference_plot_inj_recovery +++ b/bin/inference/pycbc_inference_plot_inj_recovery @@ -1,27 +1,38 @@ #!/usr/bin/env python -"""Plots the recovered versus injected parameter values from a population +""" +Plots the recovered versus injected parameter values from a population of injections. """ -import sys import logging -import matplotlib as mpl; mpl.use("Agg") +import sys + +import matplotlib as mpl + +mpl.use("Agg") import matplotlib.colorbar as cbar import matplotlib.pyplot as plt import numpy -import pycbc from matplotlib import cm -from pycbc.inference import (option_utils, io) + +import pycbc +from pycbc.inference import io, option_utils from pycbc.results import save_fig_with_metadata # parse command line parser = io.ResultsArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--output-file", required=True, type=str, - help="Path to save output plot.") -parser.add_argument("--percentiles", nargs=2, type=float, default=[5, 95], - help="Percentiles to use as limits.") +parser.add_argument( + "--output-file", required=True, type=str, help="Path to save output plot." +) +parser.add_argument( + "--percentiles", + nargs=2, + type=float, + default=[5, 95], + help="Percentiles to use as limits.", +) option_utils.add_scatter_option_group(parser) option_utils.add_injsamples_map_opt(parser) opts = parser.parse_args() @@ -34,11 +45,11 @@ fp, parameters, labels, samples = io.results_from_cli(opts) # only plot one parameter assert len(parameters) == 1 -parameter = parameters[0] -label = labels[parameters[0]] +parameter = parameters[0] +label = labels[parameters[0]] # create figure -fig = plt.figure(figsize=(6,6)) +fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(111) # typecast to list for iteratation @@ -64,20 +75,22 @@ if opts.z_arg is not None: # create colormap cmap = cm.get_cmap(opts.scatter_cmap) - vmin = opts.vmin if opts.vmin else min_zval - vmax = opts.vmax if opts.vmax else max_zval + vmin = opts.vmin or min_zval + vmax = opts.vmax or max_zval norm = mpl.colors.Normalize(vmin, vmax) # loop over input files and its samples logging.info("Plotting") -for i, (input_file, input_fp, input_samples) in enumerate(zip(opts.input_file, - fp, samples)): +for i, (input_file, input_fp, input_samples) in enumerate( + zip(opts.input_file, fp, samples) +): # get paramter values sampled_vals = input_samples[parameter] injected_vals = injs[parameter][i] # compute percentiles of sampled results - percentiles = numpy.array([numpy.percentile(sampled_vals, p) - for p in opts.percentiles]) + percentiles = numpy.array( + [numpy.percentile(sampled_vals, p) for p in opts.percentiles] + ) # get median and lowest and highest quntiles for plotting med = numpy.median(sampled_vals) @@ -91,10 +104,14 @@ for i, (input_file, input_fp, input_samples) in enumerate(zip(opts.input_file, color = "black" # plot a point for each injection - ax.errorbar([injected_vals], - [med - injected_vals], - yerr=[[(med - low)], [(high - med)]], - ecolor=color, linestyle="None", zorder=10) + ax.errorbar( + [injected_vals], + [med - injected_vals], + yerr=[[(med - low)], [(high - med)]], + ecolor=color, + linestyle="None", + zorder=10, + ) # create a colorbar if opts.z_arg: @@ -104,7 +121,7 @@ if opts.z_arg: # set labels ax.set_ylabel(r"Recovered " + label + r"- Injected " + label) -ax.set_xlabel(r"Injected " + r"{}".format(label)) +ax.set_xlabel(r"Injected " + rf"{label}") # add grid to plot ax.grid() @@ -113,15 +130,18 @@ ax.grid() ax.axhline(0, linestyle="dashed", color="gray", zorder=9) # save plot -caption = ('Difference in recovered value vs injected value vs injected ' - 'value. The vertical lines show the width of the {} to {} ' - 'percentile credible interval.' - .format(opts.percentiles[0], opts.percentiles[1])) -save_fig_with_metadata(fig, opts.output_file, - caption=caption, - cmd=' '.join(sys.argv), - fig_kwds={'bbox_inches': 'tight'}) +caption = ( + "Difference in recovered value vs injected value vs injected " + f"value. The vertical lines show the width of the {opts.percentiles[0]} to {opts.percentiles[1]} " + "percentile credible interval." +) +save_fig_with_metadata( + fig, + opts.output_file, + caption=caption, + cmd=" ".join(sys.argv), + fig_kwds={"bbox_inches": "tight"}, +) # done logging.info("Done") - diff --git a/bin/inference/pycbc_inference_plot_mcmc_history b/bin/inference/pycbc_inference_plot_mcmc_history index e66c6b33fe4..e2c540eb9e6 100644 --- a/bin/inference/pycbc_inference_plot_mcmc_history +++ b/bin/inference/pycbc_inference_plot_mcmc_history @@ -17,13 +17,14 @@ """Makes a plot of MCMC parameters saved over checkpoint history..""" -import sys -import logging import argparse +import logging +import sys -import numpy import matplotlib -matplotlib.use('Agg') +import numpy + +matplotlib.use("Agg") from matplotlib import pyplot import pycbc @@ -32,44 +33,64 @@ from pycbc.results import metadata parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input-file', type=str, required=True, - help='Path to the input HDF file.') -parser.add_argument('--output-file', type=str, required=True, - help='Name of the output plot.') -parser.add_argument('-t', '--plot-checkpoint-dt', action='store_true', - help='Plot the wall-clock time between checkpoints.') -parser.add_argument('-a', '--plot-act', action='store_true', - help='Plot ACT vs checkpoint iteration.') -parser.add_argument('-n', '--plot-effective-nsamples', action='store_true', - help='Plot the number of effective samples versus ' - 'checkpoint iteration.') -parser.add_argument('-b', '--plot-nchains-burned-in', action='store_true', - help='Plot the number of chains that were burned in ' - 'versus checkpoint iteration. Note that for ' - 'ensemble samplers, this will be all or nothing.') +parser.add_argument( + "--input-file", type=str, required=True, help="Path to the input HDF file." +) +parser.add_argument( + "--output-file", type=str, required=True, help="Name of the output plot." +) +parser.add_argument( + "-t", + "--plot-checkpoint-dt", + action="store_true", + help="Plot the wall-clock time between checkpoints.", +) +parser.add_argument( + "-a", "--plot-act", action="store_true", help="Plot ACT vs checkpoint iteration." +) +parser.add_argument( + "-n", + "--plot-effective-nsamples", + action="store_true", + help="Plot the number of effective samples versus checkpoint iteration.", +) +parser.add_argument( + "-b", + "--plot-nchains-burned-in", + action="store_true", + help="Plot the number of chains that were burned in " + "versus checkpoint iteration. Note that for " + "ensemble samplers, this will be all or nothing.", +) opts = parser.parse_args() # Default logging level is info: --verbose adds to this pycbc.init_logging(opts.verbose, default_level=1) -nplots = sum([opts.plot_act, opts.plot_effective_nsamples, - opts.plot_nchains_burned_in, opts.plot_checkpoint_dt]) +nplots = sum( + [ + opts.plot_act, + opts.plot_effective_nsamples, + opts.plot_nchains_burned_in, + opts.plot_checkpoint_dt, + ] +) if nplots == 0: raise ValueError("nothing to plot") # load the data -logging.info('Loading data') +logging.info("Loading data") -fp = io.loadfile(opts.input_file, 'r') -history = fp['sampler_info/checkpoint_history'] -iterations = history['niterations'][()] +fp = io.loadfile(opts.input_file, "r") +history = fp["sampler_info/checkpoint_history"] +iterations = history["niterations"][()] if opts.plot_checkpoint_dt: - checkpoint_dt = history['checkpoint_dt'][()] + checkpoint_dt = history["checkpoint_dt"][()] if opts.plot_act: try: - raw_acts = history['act'][()] + raw_acts = history["act"][()] except KeyError: raise ValueError("ACT history was not saved") if raw_acts.ndim == 2: @@ -84,11 +105,11 @@ if opts.plot_act: acts = raw_acts if opts.plot_effective_nsamples: - nsamples = history['effective_nsamples'][()] + nsamples = history["effective_nsamples"][()] if opts.plot_nchains_burned_in: try: - burn_in_iter = history['burn_in_iteration'][()] + burn_in_iter = history["burn_in_iteration"][()] except KeyError: raise ValueError("Burn-in history not saved") nchains_burned_in = numpy.zeros(iterations.size, dtype=int) @@ -96,88 +117,97 @@ if opts.plot_nchains_burned_in: for ii in range(nchains_burned_in.size): if burn_in_iter.ndim == 1: # ensemble sampler; all or none - nchains_burned_in[ii] = nchains*(burn_in_iter[ii] > 0) + nchains_burned_in[ii] = nchains * (burn_in_iter[ii] > 0) else: nchains_burned_in[ii] = (burn_in_iter[:, ii] > 0).sum() fp.close() # plot logging.info("Plotting") -fig, axes = pyplot.subplots(nrows=nplots, figsize=(8, 3*nplots)) +fig, axes = pyplot.subplots(nrows=nplots, figsize=(8, 3 * nplots)) if nplots == 1: axes = [axes] pi = -1 dx = iterations.max() - iterations.min() if dx == 0: dx = 1 -xmin = iterations.min() - 0.025*dx -xmax = iterations.max() + 0.025*dx +xmin = iterations.min() - 0.025 * dx +xmax = iterations.max() + 0.025 * dx -caption = ("Status of the sampler at each checkpoint iteration (indicated by " - "the x markers). Shown are:") +caption = ( + "Status of the sampler at each checkpoint iteration (indicated by " + "the x markers). Shown are:" +) if opts.plot_checkpoint_dt: pi += 1 ax = axes[pi] - ax.plot(iterations, checkpoint_dt/60., lw=2, marker='x') - ax.set_ylabel('wallclock dt (m)') + ax.plot(iterations, checkpoint_dt / 60.0, lw=2, marker="x") + ax.set_ylabel("wallclock dt (m)") ax.set_xlim(xmin, xmax) - caption += (" ({}) the amount of wall-clock time between checkpoints " - "(in minutes);".format(pi+1)) + caption += ( + f" ({pi + 1}) the amount of wall-clock time between checkpoints (in minutes);" + ) if opts.plot_act: pi += 1 ax = axes[pi] - ax.plot(iterations, acts, lw=2, marker='x', zorder=1) + ax.plot(iterations, acts, lw=2, marker="x", zorder=1) if raw_acts.ndim == 2: # plot each of the chains separately for ii in range(raw_acts.shape[0]): - ax.plot(iterations, raw_acts[ii, :], lw=1, color='C1', - alpha=0.3, - zorder=0) - ax.set_ylabel('ACT') + ax.plot(iterations, raw_acts[ii, :], lw=1, color="C1", alpha=0.3, zorder=0) + ax.set_ylabel("ACT") ax.set_xlim(xmin, xmax) - caption += (" ({}) the autocorrelation time (ACT) as computed from the " - "estimated burn-in iteration to the end of the chain(s) at " - "the checkpoint;".format(pi+1)) + caption += ( + f" ({pi + 1}) the autocorrelation time (ACT) as computed from the " + "estimated burn-in iteration to the end of the chain(s) at " + "the checkpoint;" + ) if opts.plot_effective_nsamples: pi += 1 ax = axes[pi] - ax.plot(iterations, nsamples, lw=2, marker='x', zorder=1) - ax.set_ylabel(r'eff. N samples') + ax.plot(iterations, nsamples, lw=2, marker="x", zorder=1) + ax.set_ylabel(r"eff. N samples") ax.set_xlim(xmin, xmax) - caption += (" ({}) the estimateed effective number of samples post " - "burn-in at that checkpoint;".format(pi+1)) + caption += ( + f" ({pi + 1}) the estimateed effective number of samples post " + "burn-in at that checkpoint;" + ) if opts.plot_nchains_burned_in: pi += 1 ax = axes[pi] - ax.plot(iterations, nchains_burned_in, lw=2, marker='x', zorder=1) + ax.plot(iterations, nchains_burned_in, lw=2, marker="x", zorder=1) # put a horizontal line at the total number of chains - ax.axhline(nchains, ls='--', color='C3', zorder=0) - ax.set_ylabel(r'N chains burned in') + ax.axhline(nchains, ls="--", color="C3", zorder=0) + ax.set_ylabel(r"N chains burned in") ax.set_xlim(xmin, xmax) - caption += (" ({}) the number of chains that are burned in at the " - "checkpoint (for ensemble samplers, this will either be all " - "or nothing), and the total number of chains (red dashed " - "line);".format(pi+1)) - -ax.set_xlabel('iteration') + caption += ( + f" ({pi + 1}) the number of chains that are burned in at the " + "checkpoint (for ensemble samplers, this will either be all " + "or nothing), and the total number of chains (red dashed " + "line);" + ) + +ax.set_xlabel("iteration") # replace the traling semicolon with a period caption = caption[:-1] + "." # common settings for ii, ax in enumerate(axes): - ax.grid(ls=':', zorder=-1) + ax.grid(ls=":", zorder=-1) # turn off x ticks for all but the bottom if ii < len(axes) - 1: ax.set_xticklabels([]) # save metadata.save_fig_with_metadata( - fig, opts.output_file, + fig, + opts.output_file, cmd=" ".join(sys.argv), title="MCMC history", caption=caption, - fig_kwds={'bbox_inches': 'tight'}) + fig_kwds={"bbox_inches": "tight"}, +) diff --git a/bin/inference/pycbc_inference_plot_movie b/bin/inference/pycbc_inference_plot_movie index 607731e79db..9b4c08542ab 100644 --- a/bin/inference/pycbc_inference_plot_movie +++ b/bin/inference/pycbc_inference_plot_movie @@ -15,7 +15,8 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Creates a movie showing how a sampler evolves. +""" +Creates a movie showing how a sampler evolves. To determine the length of the movie, you can either specify a --frame-number or a --frame-step. The former specifies the number of frames to use in the @@ -38,29 +39,29 @@ the movie. # ============================================================================= # +import glob import logging -import subprocess import os -import glob -from pycbc.pool import BroadcastPool as Pool +import subprocess +import matplotlib import numpy -import matplotlib -matplotlib.use('agg') +from pycbc.pool import BroadcastPool as Pool + +matplotlib.use("agg") from matplotlib import pyplot import pycbc.results from pycbc import conversions -from pycbc.inference import (option_utils, io) - -from pycbc.results.scatter_histograms import (create_multidim_plot, - get_scale_fac) -from pycbc.results.plot import (add_style_opt_to_parser, set_style_from_cli) +from pycbc.inference import io, option_utils +from pycbc.results.plot import add_style_opt_to_parser, set_style_from_cli +from pycbc.results.scatter_histograms import create_multidim_plot, get_scale_fac def integer_logspace(start, end, num): - """Generates a list of integers that are spaced approximately uniformly + """ + Generates a list of integers that are spaced approximately uniformly in log10 space between `start` and `end`. This is done such that the length of the output array is guaranteed to have length equal to num. @@ -77,12 +78,17 @@ def integer_logspace(start, end, num): ------- array The output array of integers. + """ start += 1 end += 1 out = numpy.zeros(num, dtype=int) - x = numpy.round(numpy.logspace(numpy.log10(start), numpy.log10(end), - num=num)).astype(int) - 1 + x = ( + numpy.round( + numpy.logspace(numpy.log10(start), numpy.log10(end), num=num) + ).astype(int) + - 1 + ) dx = numpy.diff(x) start_idx = 0 while (dx == 0).any(): @@ -98,57 +104,89 @@ def integer_logspace(start, end, num): # regenerate, starting from the new starting point num -= len(keep) start = keep[-1] + 2 - x = numpy.round(numpy.logspace(numpy.log10(start), numpy.log10(end), - num=num)).astype(int) - 1 + x = ( + numpy.round( + numpy.logspace(numpy.log10(start), numpy.log10(end), num=num) + ).astype(int) + - 1 + ) dx = numpy.diff(x) - out[start_idx:len(x)+start_idx] = x + out[start_idx : len(x) + start_idx] = x return out + # we won't add thinning arguments nor iteration, since this is determined by # the frame number/step options -skip_args = ['thin-start', 'thin-interval', 'thin-end', 'iteration'] +skip_args = ["thin-start", "thin-interval", "thin-end", "iteration"] parser = io.ResultsArgumentParser(description=__doc__, skip_args=skip_args) pycbc.add_common_pycbc_options(parser) # make frame number and frame step mutually exclusive group = parser.add_mutually_exclusive_group(required=True) -group.add_argument("--frame-number", type=int, - help="Maximum number of frames for the movie.") -group.add_argument("--frame-step", type=int, - help="Number of sample indices to skip between frames.") -parser.add_argument("--start-index", type=int, default=0, - help="The starting index of the samples to load. Must be " - "< the number of samples that are stored in the " - "file. Default is 0.") -parser.add_argument("--end-index", type=int, default=None, - help="The ending index of the samples to load. Must be < " - "the number of samples that are stored in the file. " - "Default is to load everything to the end.") -parser.add_argument("--log-steps", action="store_true", default=False, - help="If frame-number is specified, make the number of " - "samples between frames uniform in log10. This " - "provides more detail of the early iterations, when " - "the sampler is changing most rapidly. An error will " - "be raised if frame-number is not provided.") -parser.add_argument("--output-prefix", type=str, required=True, - help="Output path and prefix for the frame files " - "(without extension).") -parser.add_argument('--dpi', type=int, default=200, - help='Set the dpi for each frame; default is 200') -parser.add_argument("--nprocesses", type=int, default=None, - help="Number of processes to use. If not given then " - "use maximum.") -parser.add_argument("--movie-file", type=str, - help="Path for creating the movie automatically after " - "generating all the frames. If a movie with the same " - "name already exists, it will remove it. Format: mp4. " - "Installation of FFMPEG required.") -parser.add_argument("--cleanup", action='store_true', - help="Delete all plots generated after creating the movie." - " Only works together with option make-movie.") +group.add_argument( + "--frame-number", type=int, help="Maximum number of frames for the movie." +) +group.add_argument( + "--frame-step", type=int, help="Number of sample indices to skip between frames." +) +parser.add_argument( + "--start-index", + type=int, + default=0, + help="The starting index of the samples to load. Must be " + "< the number of samples that are stored in the " + "file. Default is 0.", +) +parser.add_argument( + "--end-index", + type=int, + default=None, + help="The ending index of the samples to load. Must be < " + "the number of samples that are stored in the file. " + "Default is to load everything to the end.", +) +parser.add_argument( + "--log-steps", + action="store_true", + default=False, + help="If frame-number is specified, make the number of " + "samples between frames uniform in log10. This " + "provides more detail of the early iterations, when " + "the sampler is changing most rapidly. An error will " + "be raised if frame-number is not provided.", +) +parser.add_argument( + "--output-prefix", + type=str, + required=True, + help="Output path and prefix for the frame files (without extension).", +) +parser.add_argument( + "--dpi", type=int, default=200, help="Set the dpi for each frame; default is 200" +) +parser.add_argument( + "--nprocesses", + type=int, + default=None, + help="Number of processes to use. If not given then use maximum.", +) +parser.add_argument( + "--movie-file", + type=str, + help="Path for creating the movie automatically after " + "generating all the frames. If a movie with the same " + "name already exists, it will remove it. Format: mp4. " + "Installation of FFMPEG required.", +) +parser.add_argument( + "--cleanup", + action="store_true", + help="Delete all plots generated after creating the movie." + " Only works together with option make-movie.", +) # add options for what plots to create option_utils.add_plot_posterior_option_group(parser) # add scatter and density configuration options -option_utils.add_scatter_option_group(parser) +option_utils.add_scatter_option_group(parser) option_utils.add_density_option_group(parser) add_style_opt_to_parser(parser) @@ -163,7 +201,7 @@ if len(opts.input_file) > 1: raise ValueError("this program can only plot one file at a time") # Get data -logging.info('Loading parameters') +logging.info("Loading parameters") fp, parameters, labels, _ = io.results_from_cli(opts, load_samples=False) # get the total number of samples on disk @@ -171,16 +209,20 @@ thinned_by = fp.thinned_by nsamples = fp.niterations // thinned_by start_index = opts.start_index if start_index >= nsamples: - raise ValueError("given start-index {} is >= the number of samples in the " - "input file {}".format(start_index, nsamples)) + raise ValueError( + f"given start-index {start_index} is >= the number of samples in the input file {nsamples}" + ) end_index = opts.end_index if end_index is not None and end_index >= nsamples: - raise ValueError("given end-index {} is >= the number of samples in the " - "input file {}".format(end_index, nsamples)) + raise ValueError( + f"given end-index {end_index} is >= the number of samples in the input file {nsamples}" + ) if opts.log_steps and not opts.frame_number: - raise ValueError("log-steps requires a non-zero frame-number to be " - "provided; see help for details.") + raise ValueError( + "log-steps requires a non-zero frame-number to be " + "provided; see help for details." + ) if opts.frame_number: thinint = 1 @@ -188,9 +230,14 @@ else: # frame step was provided thinint = opts.frame_step # get the samples -samples = fp.samples_from_cli(opts, parameters, thin_start=start_index, - thin_interval=thinint, thin_end=end_index, - flatten=False) +samples = fp.samples_from_cli( + opts, + parameters, + thin_start=start_index, + thin_interval=thinint, + thin_end=end_index, + flatten=False, +) if samples.ndim > 2: # multi-tempered samplers will return 3 dims, so flatten _, ii, jj = samples.shape @@ -199,11 +246,13 @@ if samples.ndim > 2: # pull out the samples we want if opts.frame_number: if opts.log_steps: - indices = integer_logspace(start_index, samples.shape[1]-1, - opts.frame_number) + indices = integer_logspace(start_index, samples.shape[1] - 1, opts.frame_number) else: - indices = numpy.unique(numpy.linspace(start_index, samples.shape[1]-1, - num=opts.frame_number).astype(int)) + indices = numpy.unique( + numpy.linspace( + start_index, samples.shape[1] - 1, num=opts.frame_number + ).astype(int) + ) samples = samples[:, indices] else: # set the index counter based on what was loaded @@ -213,14 +262,19 @@ else: # Get z-values if opts.z_arg is not None: logging.info("Getting samples for colorbar") - if opts.z_arg == 'snr': - z_arg = 'loglikelihood' + if opts.z_arg == "snr": + z_arg = "loglikelihood" else: z_arg = opts.z_arg - zsamples = fp.samples_from_cli(opts, z_arg, thin_start=start_index, - thin_interval=thinint, thin_end=end_index, - flatten=False) - if opts.z_arg == 'snr': + zsamples = fp.samples_from_cli( + opts, + z_arg, + thin_start=start_index, + thin_interval=thinint, + thin_end=end_index, + flatten=False, + ) + if opts.z_arg == "snr": loglr = zsamples[z_arg] - zsamples.lognl zsamples[z_arg] = conversions.snr_from_loglr(loglr) zlbl = opts.z_arg_labels[opts.z_arg] @@ -262,10 +316,12 @@ if opts.plot_injection_parameters: continue unique_vals = numpy.unique(vals) if unique_vals.size != 1: - raise ValueError("More than one injection found! To use " + raise ValueError( + "More than one injection found! To use " "plot-injection-parameters, there must be a single unique " "injection in all input files. Use the expected-parameters " - "option to specify an expected parameter instead.") + "option to specify an expected parameter instead." + ) # passed: use the value for the expected expected_parameters[p] = unique_vals[0] @@ -273,7 +329,7 @@ if opts.plot_injection_parameters: expected_parameters.update(option_utils.expected_parameters_from_cli(opts)) expected_parameters_color = opts.expected_parameters_color -logging.info('Choosing common characteristics for all figures') +logging.info("Choosing common characteristics for all figures") # Set common min and max for axis in all plots mins, maxs = option_utils.plot_ranges_from_cli(opts) # add any missing parameters @@ -286,16 +342,16 @@ for p in parameters: # set colors: # the 1d marginal colors will just be the first color in the style's cycle -linecolor = list(matplotlib.rcParams['axes.prop_cycle'])[0]['color'] +linecolor = list(matplotlib.rcParams["axes.prop_cycle"])[0]["color"] # make the hist color black or white, depending on if the # dark background is used -if opts.mpl_style == 'dark_background': - hist_color = 'white' +if opts.mpl_style == "dark_background": + hist_color = "white" else: - hist_color = 'black' + hist_color = "black" # make the default contour color white if plot density is on if not opts.contour_color and opts.plot_density: - contour_color = 'white' + contour_color = "white" # otherwise, make the default be the same as the hist color elif not opts.contour_color: contour_color = hist_color @@ -305,36 +361,44 @@ else: # Make each figure # for sorting purposes, we will need to zero-pad the sample number with the # appriopriate number of 0's -max_iter_num = indices[-1]*thinned_by + 1 +max_iter_num = indices[-1] * thinned_by + 1 + def _make_frame(frame): - """Wrapper for making the plot in a pooled environment. - """ - plotargs = samples[:,frame] + """Wrapper for making the plot in a pooled environment.""" + plotargs = samples[:, frame] if zvals is not None: - z = zvals[:,frame] + z = zvals[:, frame] else: z = None - iter_num = str(indices[frame]*thinned_by + 1) + iter_num = str(indices[frame] * thinned_by + 1) iter_num = iter_num.zfill(len(str(max_iter_num))) - output = opts.output_prefix + '-{}.png'.format(iter_num) - - fig, axis_dict = create_multidim_plot(parameters, plotargs, labels=labels, - mins=mins, maxs=maxs, - plot_marginal=opts.plot_marginal, - line_color=linecolor, - hist_color=hist_color, - plot_scatter=opts.plot_scatter, - zvals=z, show_colorbar=show_colorbar, - cbar_label=zlbl, vmin=vmin, vmax=vmax, - scatter_cmap=opts.scatter_cmap, - plot_density=opts.plot_density, - plot_contours=opts.plot_contours, - density_cmap=opts.density_cmap, - contour_color=contour_color, - use_kombine=opts.use_kombine_kde, - expected_parameters=expected_parameters, - expected_parameters_color=expected_parameters_color) + output = opts.output_prefix + f"-{iter_num}.png" + + fig, axis_dict = create_multidim_plot( + parameters, + plotargs, + labels=labels, + mins=mins, + maxs=maxs, + plot_marginal=opts.plot_marginal, + line_color=linecolor, + hist_color=hist_color, + plot_scatter=opts.plot_scatter, + zvals=z, + show_colorbar=show_colorbar, + cbar_label=zlbl, + vmin=vmin, + vmax=vmax, + scatter_cmap=opts.scatter_cmap, + plot_density=opts.plot_density, + plot_contours=opts.plot_contours, + density_cmap=opts.density_cmap, + contour_color=contour_color, + use_kombine=opts.use_kombine_kde, + expected_parameters=expected_parameters, + expected_parameters_color=expected_parameters_color, + ) # Write sample number if show_colorbar: @@ -343,14 +407,20 @@ def _make_frame(frame): xtxt = 0.9 ytxt = 0.95 scale_fac = get_scale_fac(fig) - fontsize = 8*scale_fac - pyplot.annotate('Iteration {}'.format(iter_num), xy=(xtxt, ytxt), - xycoords='figure fraction', horizontalalignment='right', - verticalalignment='top', fontsize=fontsize) - - fig.savefig(output, bbox_inches='tight', dpi=opts.dpi) + fontsize = 8 * scale_fac + pyplot.annotate( + f"Iteration {iter_num}", + xy=(xtxt, ytxt), + xycoords="figure fraction", + horizontalalignment="right", + verticalalignment="top", + fontsize=fontsize, + ) + + fig.savefig(output, bbox_inches="tight", dpi=opts.dpi) pyplot.close() - return fig.get_figheight()/fig.get_figwidth() + return fig.get_figheight() / fig.get_figwidth() + # create the pool if opts.nprocesses is None or opts.nprocesses > 1: @@ -369,15 +439,26 @@ if opts.movie_file: logging.info("Making movie") frame_files = opts.output_prefix + "*.png" # set the aspect ratio - aspect_ratio = "1024x{}".format(int(1024*aspect_ratio)) + aspect_ratio = f"1024x{int(1024 * aspect_ratio)}" if os.path.isfile(opts.movie_file): os.remove(opts.movie_file) - subprocess.call(["ffmpeg", "-pix_fmt", "yuv420p", "-s", aspect_ratio, - "-pattern_type", "glob", "-i", - frame_files, opts.movie_file]) + subprocess.call( + [ + "ffmpeg", + "-pix_fmt", + "yuv420p", + "-s", + aspect_ratio, + "-pattern_type", + "glob", + "-i", + frame_files, + opts.movie_file, + ] + ) if opts.cleanup: logging.info("Removing frames") for frame in glob.glob(frame_files): os.remove(frame) -logging.info('Done') +logging.info("Done") diff --git a/bin/inference/pycbc_inference_plot_posterior b/bin/inference/pycbc_inference_plot_posterior index a47da4bfa31..25fa05dcba9 100644 --- a/bin/inference/pycbc_inference_plot_posterior +++ b/bin/inference/pycbc_inference_plot_posterior @@ -29,63 +29,91 @@ import itertools import logging import sys -import numpy - import matplotlib -from matplotlib import (patches, use) +import numpy +from matplotlib import patches, use import pycbc -from pycbc.results.plot import (add_style_opt_to_parser, set_style_from_cli) -from pycbc.results import metadata -from pycbc.io import FieldArray from pycbc import conversions -from pycbc.workflow import WorkflowConfigParser -from pycbc.inference import (option_utils, io) from pycbc.distributions.utils import prior_from_config - +from pycbc.inference import io, option_utils +from pycbc.io import FieldArray +from pycbc.results import metadata +from pycbc.results.plot import add_style_opt_to_parser, set_style_from_cli from pycbc.results.scatter_histograms import create_multidim_plot +from pycbc.workflow import WorkflowConfigParser -use('agg') +use("agg") # add options to command line parser = io.ResultsArgumentParser() pycbc.add_common_pycbc_options(parser) # program-specific -parser.add_argument("--output-file", type=str, required=True, - help="Output plot path.") -parser.add_argument("--plot-prior", nargs="+", type=str, - help="Plot the prior on the 1D marginal plots using the " - "given config file(s).") -parser.add_argument("--prior-nsamples", type=int, default=10000, - help="The number of samples to use for plotting the " - "prior. Default is 10000.") -parser.add_argument("--colors-multi-run", nargs="+", type=str, - help="For multiple runs, provide colours to be used for successively. Default setting is to use the successive colours specified in matplotlib color cycle.") -parser.add_argument("--fill-hist", action="store_true", default=False, - help="Fill the 1D marginalized histograms") -parser.add_argument("--hist-color", - help="Provide color for histogram outline. Default is black") -parser.add_argument("--hist-fill-color", default='gray', - help="Provide the fill_color for filled histograms. Default is gray") +parser.add_argument("--output-file", type=str, required=True, help="Output plot path.") +parser.add_argument( + "--plot-prior", + nargs="+", + type=str, + help="Plot the prior on the 1D marginal plots using the given config file(s).", +) +parser.add_argument( + "--prior-nsamples", + type=int, + default=10000, + help="The number of samples to use for plotting the prior. Default is 10000.", +) +parser.add_argument( + "--colors-multi-run", + nargs="+", + type=str, + help="For multiple runs, provide colours to be used for successively. Default setting is to use the successive colours specified in matplotlib color cycle.", +) +parser.add_argument( + "--fill-hist", + action="store_true", + default=False, + help="Fill the 1D marginalized histograms", +) +parser.add_argument( + "--hist-color", help="Provide color for histogram outline. Default is black" +) +parser.add_argument( + "--hist-fill-color", + default="gray", + help="Provide the fill_color for filled histograms. Default is gray", +) # add options for what plots to create option_utils.add_plot_posterior_option_group(parser) # scatter configuration option_utils.add_scatter_option_group(parser) # density configuration option_utils.add_density_option_group(parser) -parser.add_argument("--plot-maxl", action="store_true", default=False, - help="Put a marker on the 2D marginal where the maxL " - "point is.") -parser.add_argument('--legend-location', default='upper right', - help='Where to put the legend (if multiple files are ' - 'provided). Default is "upper right".') -parser.add_argument('--legend-bbox-to-anchor', default='(1.0, 1.0)', - help='Box that is used to position the legend in conjunction with loc.') -parser.add_argument('--legend-fontsize', default='medium', - help="The font size of legend. Default is 'medium'." - "Choose from int or {'xx-small', 'x-small', 'small', 'medium', 'large', 'x-large', 'xx-large'}.") -parser.add_argument('--dpi', type=int, default=200, - help="Set the DPI of the plot. Default is 200.") +parser.add_argument( + "--plot-maxl", + action="store_true", + default=False, + help="Put a marker on the 2D marginal where the maxL point is.", +) +parser.add_argument( + "--legend-location", + default="upper right", + help="Where to put the legend (if multiple files are " + 'provided). Default is "upper right".', +) +parser.add_argument( + "--legend-bbox-to-anchor", + default="(1.0, 1.0)", + help="Box that is used to position the legend in conjunction with loc.", +) +parser.add_argument( + "--legend-fontsize", + default="medium", + help="The font size of legend. Default is 'medium'." + "Choose from int or {'xx-small', 'x-small', 'small', 'medium', 'large', 'x-large', 'xx-large'}.", +) +parser.add_argument( + "--dpi", type=int, default=200, help="Set the DPI of the plot. Default is 200." +) # style option add_style_opt_to_parser(parser) @@ -96,11 +124,11 @@ opts = parser.parse_args() set_style_from_cli(opts) # get any kdeargs; we'll always pass through the kde max samples -kdeargs = {'max_kde_samples': opts.max_kde_samples} +kdeargs = {"max_kde_samples": opts.max_kde_samples} # add any other options that were specified if opts.kde_args is not None: for opt in opts.kde_args: - opt, val = opt.split(':') + opt, val = opt.split(":") try: val = float(val) # convert to int if no remainder @@ -112,12 +140,13 @@ if opts.kde_args is not None: if opts.plot_maxl: # add loglikelihood to list of parameters - add_logl = 'loglikelihood' not in opts.parameters + add_logl = "loglikelihood" not in opts.parameters if add_logl: - opts.parameters.append('loglikelihood') + opts.parameters.append("loglikelihood") else: add_logl = False + # parse legend bbox def parse_tuple(value): try: @@ -125,6 +154,7 @@ def parse_tuple(value): except Exception as e: raise argparse.ArgumentTypeError(f"Invalid tuple format: {value}") from e + opts.legend_bbox_to_anchor = parse_tuple(opts.legend_bbox_to_anchor) # set logging @@ -134,7 +164,7 @@ pycbc.init_logging(opts.verbose) fps, parameters, labels, samples = io.results_from_cli(opts) if add_logl: - parameters = [p for p in parameters if p != 'loglikelihood'] + parameters = [p for p in parameters if p != "loglikelihood"] # typecast to list so the input files can be iterated over fps = fps if isinstance(fps, list) else [fps] @@ -143,12 +173,12 @@ samples = samples if isinstance(samples, list) else [samples] # if a z-arg is specified, load samples for it if opts.z_arg is not None: logging.info("Getting samples for colorbar") - z_arg = 'loglikelihood' if opts.z_arg == 'snr' else opts.z_arg + z_arg = "loglikelihood" if opts.z_arg == "snr" else opts.z_arg zlbl = opts.z_arg_labels[opts.z_arg] zvals = [] for fp in fps: zsamples = fp.samples_from_cli(opts, parameters=z_arg) - if opts.z_arg == 'snr': + if opts.z_arg == "snr": loglr = zsamples[z_arg] - zsamples.lognl zsamples[z_arg] = conversions.snr_from_loglr(loglr) zvals.append(zsamples[z_arg]) @@ -173,9 +203,11 @@ if not numpy.any(plot_options): if opts.plot_prior is not None: # check that we're plotting 1D marginals if not opts.plot_marginal: - raise ValueError("prior may only be plotted on 1D marginal plot; " - "either turn on --plot-marginal, or turn off " - "--plot-prior") + raise ValueError( + "prior may only be plotted on 1D marginal plot; " + "either turn on --plot-marginal, or turn off " + "--plot-prior" + ) logging.info("Loading prior") cp = WorkflowConfigParser(opts.plot_prior) prior = prior_from_config(cp) @@ -184,18 +216,18 @@ if opts.plot_prior is not None: # we'll just use the first file for metadata fp = fps[0] # add the static params - for param in fp.attrs['static_params']: + for param in fp.attrs["static_params"]: setattr(prior_samples, param, fp.attrs[param]) # remap any parameters - if 'remapped_params' in fp.attrs: + if "remapped_params" in fp.attrs: remapped_params = {} - for func, param in fp.attrs['remapped_params']: + for func, param in fp.attrs["remapped_params"]: try: remapped_params[param] = prior_samples[func] except (NameError, TypeError, AttributeError): continue prior_samples = FieldArray.from_kwargs(**remapped_params) - for param in fp.attrs['static_params']: + for param in fp.attrs["static_params"]: setattr(prior_samples, param, fp.attrs[param]) # get minimum and maximum ranges for each parameter from command line @@ -214,20 +246,22 @@ if opts.plot_injection_parameters: injections = io.injections_from_cli(opts) if opts.pick_injection_by_time: - if 'tc' not in injections: + if "tc" not in injections: raise ValueError("Couldn't determine injection time, tried tc") - inj_time = injections['tc'] + inj_time = injections["tc"] - if 'tc' in samples[0]: - pos_time = samples[0]['tc'].mean() - elif 'trigger_time' in fps[0].attrs: - pos_time = fps[0].attrs['trigger_time'] - elif 'tc_ref' in fps[0].attrs: - pos_time = fps[0].attrs['tc_ref'] + if "tc" in samples[0]: + pos_time = samples[0]["tc"].mean() + elif "trigger_time" in fps[0].attrs: + pos_time = fps[0].attrs["trigger_time"] + elif "tc_ref" in fps[0].attrs: + pos_time = fps[0].attrs["tc_ref"] else: - raise ValueError("Couldn't find posterior time, " - "tried tc, tc_ref, and trigger_time attribute") + raise ValueError( + "Couldn't find posterior time, " + "tried tc, tc_ref, and trigger_time attribute" + ) pick = abs(inj_time - pos_time).argmin() for p in parameters: @@ -244,10 +278,12 @@ if opts.plot_injection_parameters: # check that all of the injections are the same unique_vals = numpy.unique(vals) if unique_vals.size != 1: - raise ValueError("More than one injection found! To use " + raise ValueError( + "More than one injection found! To use " "plot-injection-parameters, there must be a single unique " "injection in all input files. Use the expected-parameters" - " option to specify an expected parameter instead.") + " option to specify an expected parameter instead." + ) # passed: use the value for the expected expected_parameters[p] = unique_vals[0] @@ -260,7 +296,7 @@ for fp in fps: expected_parameters.update(option_utils.expected_parameters_from_cli(opts)) # get the color cycle to use -color_cycle = [c['color'] for c in matplotlib.rcParams['axes.prop_cycle']] +color_cycle = [c["color"] for c in matplotlib.rcParams["axes.prop_cycle"]] if opts.colors_multi_run is not None: colors = itertools.cycle(opts.colors_multi_run) else: @@ -269,8 +305,7 @@ else: # plot each input file logging.info("Plotting") hist_colors = [] -for (i, s) in enumerate(samples): - +for i, s in enumerate(samples): # on first iteration create figure otherwise update old figure if i == 0: fig = None @@ -285,13 +320,12 @@ for (i, s) in enumerate(samples): if len(opts.input_file) == 1: # make the hist color black or white, depending on if dark background # is used - if opts.mpl_style == 'dark_background': - hist_color = 'white' + if opts.mpl_style == "dark_background": + hist_color = "white" + elif opts.hist_color: + hist_color = opts.hist_color else: - if opts.hist_color: - hist_color = opts.hist_color - else: - hist_color = 'black' + hist_color = "black" # fill histogram if fill_hist is True if opts.fill_hist: fill_color = opts.hist_fill_color @@ -299,7 +333,7 @@ for (i, s) in enumerate(samples): fill_color = None # make the default contour color white if plot density is on if not opts.contour_color and opts.plot_density: - contour_color = 'white' + contour_color = "white" # otherwise, make the default be the same as the hist color elif not opts.contour_color: contour_color = hist_color @@ -316,46 +350,63 @@ for (i, s) in enumerate(samples): # plot fig, axis_dict = create_multidim_plot( - parameters, s, labels=labels, fig=fig, axis_dict=axis_dict, - plot_marginal=opts.plot_marginal, - plot_marginal_lines=not opts.no_marginal_lines, - plot_maxl=opts.plot_maxl, - marginal_percentiles=opts.marginal_percentiles, - marginal_title=not opts.no_marginal_titles, - plot_scatter=opts.plot_scatter, - zvals=zvals[i] if zvals is not None else None, - show_colorbar=opts.z_arg is not None, - cbar_label=zlbl, - vmin=opts.vmin, vmax=opts.vmax, - scatter_cmap=opts.scatter_cmap, - plot_density=opts.plot_density, - plot_contours=opts.plot_contours, - contour_percentiles=opts.contour_percentiles, - density_cmap=opts.density_cmap, - contour_color=contour_color, - contour_linestyles=opts.contour_linestyles, - label_contours=not opts.no_contour_labels, - hist_color=hist_color, - line_color=linecolor, - fill_color=fill_color, - use_kombine=opts.use_kombine_kde, - kdeargs=kdeargs, - mins=mins, maxs=maxs, - expected_parameters=expected_parameters, - expected_parameters_color=opts.expected_parameters_color) + parameters, + s, + labels=labels, + fig=fig, + axis_dict=axis_dict, + plot_marginal=opts.plot_marginal, + plot_marginal_lines=not opts.no_marginal_lines, + plot_maxl=opts.plot_maxl, + marginal_percentiles=opts.marginal_percentiles, + marginal_title=not opts.no_marginal_titles, + plot_scatter=opts.plot_scatter, + zvals=zvals[i] if zvals is not None else None, + show_colorbar=opts.z_arg is not None, + cbar_label=zlbl, + vmin=opts.vmin, + vmax=opts.vmax, + scatter_cmap=opts.scatter_cmap, + plot_density=opts.plot_density, + plot_contours=opts.plot_contours, + contour_percentiles=opts.contour_percentiles, + density_cmap=opts.density_cmap, + contour_color=contour_color, + contour_linestyles=opts.contour_linestyles, + label_contours=not opts.no_contour_labels, + hist_color=hist_color, + line_color=linecolor, + fill_color=fill_color, + use_kombine=opts.use_kombine_kde, + kdeargs=kdeargs, + mins=mins, + maxs=maxs, + expected_parameters=expected_parameters, + expected_parameters_color=opts.expected_parameters_color, + ) # plot the prior if opts.plot_prior: if len(opts.input_file) > 1: hist_color = next(colors) fig, axis_dict = create_multidim_plot( - parameters, prior_samples, fig=fig, axis_dict=axis_dict, - labels=labels, plot_marginal=True, marginal_percentiles=[], - plot_scatter=False, plot_density=False, plot_contours=False, + parameters, + prior_samples, + fig=fig, + axis_dict=axis_dict, + labels=labels, + plot_marginal=True, + marginal_percentiles=[], + plot_scatter=False, + plot_density=False, + plot_contours=False, fill_color=None, - marginal_title=False, marginal_linestyle=':', + marginal_title=False, + marginal_linestyle=":", hist_color=hist_color, - mins=mins, maxs=maxs) + mins=mins, + maxs=maxs, + ) # add legend to upper right for input files if len(opts.input_file) > 1: @@ -371,20 +422,26 @@ if len(opts.input_file) > 1: addto = axis_dict[parameters[0], parameters[0]][0] else: addto = fig - addto.legend(loc=opts.legend_location, handles=handles, - labels=labels, fontsize=opts.legend_fontsize, - bbox_to_anchor=opts.legend_bbox_to_anchor) + addto.legend( + loc=opts.legend_location, + handles=handles, + labels=labels, + fontsize=opts.legend_fontsize, + bbox_to_anchor=opts.legend_bbox_to_anchor, + ) # set DPI fig.set_dpi(opts.dpi) # save metadata.save_fig_with_metadata( - fig, opts.output_file, - cmd=" ".join(sys.argv), - title="Posteriors", - caption="Posterior probability density functions.", - fig_kwds={'bbox_inches': 'tight'}) + fig, + opts.output_file, + cmd=" ".join(sys.argv), + title="Posteriors", + caption="Posterior probability density functions.", + fig_kwds={"bbox_inches": "tight"}, +) # finish logging.info("Done") diff --git a/bin/inference/pycbc_inference_plot_pp b/bin/inference/pycbc_inference_plot_pp index faba65261d7..696cc217089 100644 --- a/bin/inference/pycbc_inference_plot_pp +++ b/bin/inference/pycbc_inference_plot_pp @@ -15,42 +15,48 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Plots the fraction of injections with their parameter value recovered +""" +Plots the fraction of injections with their parameter value recovered within a credible interval versus credible interval. """ -import sys -import logging import itertools +import logging +import sys +import matplotlib import numpy - from scipy import stats -import matplotlib -matplotlib.use('agg') +matplotlib.use("agg") from matplotlib import pyplot as plt import pycbc import pycbc.results.plot +from pycbc.inference import io, option_utils from pycbc.results import save_fig_with_metadata -from pycbc.inference import (option_utils, io) # parse command line parser = io.ResultsArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--output-file", required=True, type=str, - help="Path to save output plot.") -parser.add_argument("--injection-hdf-group", default="injections", - help="HDF group that contains injection values. " - "Default is 'injections'.") -parser.add_argument("--do-ks-test", action="store_true", default=False, - help="Perform a KS test between the percentile-percentile " - "plot for each parameter and the (expected) uniform " - "distribution. Results are printed in the legend.") +parser.add_argument( + "--output-file", required=True, type=str, help="Path to save output plot." +) +parser.add_argument( + "--injection-hdf-group", + default="injections", + help="HDF group that contains injection values. Default is 'injections'.", +) +parser.add_argument( + "--do-ks-test", + action="store_true", + default=False, + help="Perform a KS test between the percentile-percentile " + "plot for each parameter and the (expected) uniform " + "distribution. Results are printed in the legend.", +) option_utils.add_injsamples_map_opt(parser) -pycbc.results.plot.add_style_opt_to_parser(parser, - default='seaborn-colorblind') +pycbc.results.plot.add_style_opt_to_parser(parser, default="seaborn-colorblind") opts = parser.parse_args() # set style @@ -60,7 +66,7 @@ pycbc.results.plot.set_style_from_cli(opts) pycbc.init_logging(opts.verbose) # read results -logging.info('Loading parameters') +logging.info("Loading parameters") _, parameters, labels, samples = io.results_from_cli(opts) # typecast to list for iteration @@ -78,7 +84,7 @@ for input_file, input_samples in zip(opts.input_file, samples): for p in parameters: inj_val = inj_parameters[p] sample_vals = input_samples[p] - measured = stats.percentileofscore(sample_vals, inj_val, kind='weak') + measured = stats.percentileofscore(sample_vals, inj_val, kind="weak") try: measured_percentiles[p].append(measured) except KeyError: @@ -88,32 +94,38 @@ for input_file, input_samples in zip(opts.input_file, samples): # set the color and line styles; total number of unique combinations is 3 * # the number of colors in the style's cycle (seaborn-colorblind has 6 colors, # so the default is 18 unique combinations) -color_cycle = matplotlib.rcParams['axes.prop_cycle'] -colors = itertools.cycle([x['color'] for x in color_cycle]) +color_cycle = matplotlib.rcParams["axes.prop_cycle"] +colors = itertools.cycle([x["color"] for x in color_cycle]) ncolors = len(color_cycle) -ls_cycle = ['-']*ncolors + [':']*ncolors + ['-.']*ncolors +ls_cycle = ["-"] * ncolors + [":"] * ncolors + ["-."] * ncolors line_styles = itertools.cycle(ls_cycle) # create figure for plotting -fig = plt.figure(figsize=(6,6)) +fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(111) # calculate the expected percentile for each injection and plot for param in parameters: label = labels[param] meas = numpy.array(measured_percentiles[param]) meas.sort() - expected = numpy.array([stats.percentileofscore(meas, x, kind='weak') - for x in meas]) + expected = numpy.array( + [stats.percentileofscore(meas, x, kind="weak") for x in meas] + ) # perform ks test if opts.do_ks_test: - ks, p = stats.kstest(meas/100., 'uniform') - label = '{} $D_{{KS}}$: {:.3f} p-value: {:.3f}'.format(label, ks, p) - ax.plot(meas/100., expected/100., c=next(colors), ls=next(line_styles), - label=label) + ks, p = stats.kstest(meas / 100.0, "uniform") + label = f"{label} $D_{{KS}}$: {ks:.3f} p-value: {p:.3f}" + ax.plot( + meas / 100.0, + expected / 100.0, + c=next(colors), + ls=next(line_styles), + label=label, + ) # set legend breaknum = 9 -if len(parameters) > 2*breaknum: +if len(parameters) > 2 * breaknum: ncols = 3 elif len(parameters) > breaknum: ncols = 2 @@ -133,22 +145,26 @@ ax.grid() ax.plot([0, 1], [0, 1], linestyle="dashed", color="gray", zorder=9) # save plot -caption = ('Percentile-percentile plot. The value of the KS statistic ' - '$D_{KS}$ is given in the legend. This gives the maximum distance ' - 'between the observed line and the expected (dashed) line; i.e., ' - 'it gives the maximum distance between the measured CDF and the ' - 'expected (uniform) CDF. The associated two-tailed p-value gives ' - 'the probability of getting a maximum distance (either above ' - 'or below the expected line) larger than the observed $D_{KS}$ ' - 'assuming that the measured CDF is the same as the expected. ' - 'In other words, the larger (smaller) the p-value ($D_{KS}$), the ' - 'more likely the measured distribution is the same as the ' - 'expected.') -save_fig_with_metadata(fig, opts.output_file, +caption = ( + "Percentile-percentile plot. The value of the KS statistic " + "$D_{KS}$ is given in the legend. This gives the maximum distance " + "between the observed line and the expected (dashed) line; i.e., " + "it gives the maximum distance between the measured CDF and the " + "expected (uniform) CDF. The associated two-tailed p-value gives " + "the probability of getting a maximum distance (either above " + "or below the expected line) larger than the observed $D_{KS}$ " + "assuming that the measured CDF is the same as the expected. " + "In other words, the larger (smaller) the p-value ($D_{KS}$), the " + "more likely the measured distribution is the same as the " + "expected." +) +save_fig_with_metadata( + fig, + opts.output_file, caption=caption, - cmd=' '.join(sys.argv), - fig_kwds={'bbox_inches': 'tight'}) + cmd=" ".join(sys.argv), + fig_kwds={"bbox_inches": "tight"}, +) # done logging.info("Done") - diff --git a/bin/inference/pycbc_inference_plot_prior b/bin/inference/pycbc_inference_plot_prior index 25cabc06085..0fc57701177 100644 --- a/bin/inference/pycbc_inference_plot_prior +++ b/bin/inference/pycbc_inference_plot_prior @@ -22,59 +22,74 @@ import logging import sys import numpy - from matplotlib import use -use('agg') + +use("agg") from matplotlib import pyplot as plt import pycbc -from pycbc import (results, waveform) -from pycbc.inference.option_utils import ParseParametersArg +from pycbc import results, waveform from pycbc.distributions.utils import prior_from_config -from pycbc.workflow import WorkflowConfigParser +from pycbc.inference.option_utils import ParseParametersArg from pycbc.results.scatter_histograms import create_multidim_plot +from pycbc.workflow import WorkflowConfigParser def cartesian(arrays): - """ Returns a cartesian product from a list of iterables. - """ - return numpy.array([numpy.array(element) - for element in itertools.product(*arrays)]) + """Returns a cartesian product from a list of iterables.""" + return numpy.array([numpy.array(element) for element in itertools.product(*arrays)]) + # command line usage parser = argparse.ArgumentParser( usage="pycbc_inference_plot_prior [--options]", - description="Plots prior distributions.") + description="Plots prior distributions.", +) pycbc.add_common_pycbc_options(parser) # add input options -parser.add_argument("--config-files", type=str, nargs="+", required=True, - help="A file parsable by " - "pycbc.workflow.WorkflowConfigParser.") -parser.add_argument("--parameters", nargs="+", action=ParseParametersArg, - metavar="PARAM[:LABEL]", - help="Only plot the given parameters. May also provide a " - "label for each parameter. If provided, the " - "parameters can be any of parameters in the " - "[variable_params] section, derived parameters from " - "them, or any function of them. Syntax for " - "functions is python; any math functions in the " - "numpy libary may be used. Can optionally also " - "specify a LABEL for each parameter. If no LABEL is " - "provided, PARAM will used as the LABEL. If LABEL " - "is the same as a parameter in " - "pycbc.waveform.parameters, the label " - "property of that parameter will be used. If not " - "provided, will plot all of the parameters in the " - "[variable_params] section of the config file.") -parser.add_argument("--prior-section", type=str, default="prior", - help="Name of section with distribution configurations. " - "Default is 'prior'.") +parser.add_argument( + "--config-files", + type=str, + nargs="+", + required=True, + help="A file parsable by pycbc.workflow.WorkflowConfigParser.", +) +parser.add_argument( + "--parameters", + nargs="+", + action=ParseParametersArg, + metavar="PARAM[:LABEL]", + help="Only plot the given parameters. May also provide a " + "label for each parameter. If provided, the " + "parameters can be any of parameters in the " + "[variable_params] section, derived parameters from " + "them, or any function of them. Syntax for " + "functions is python; any math functions in the " + "numpy libary may be used. Can optionally also " + "specify a LABEL for each parameter. If no LABEL is " + "provided, PARAM will used as the LABEL. If LABEL " + "is the same as a parameter in " + "pycbc.waveform.parameters, the label " + "property of that parameter will be used. If not " + "provided, will plot all of the parameters in the " + "[variable_params] section of the config file.", +) +parser.add_argument( + "--prior-section", + type=str, + default="prior", + help="Name of section with distribution configurations. Default is 'prior'.", +) # add output options -parser.add_argument("--nsamples", type=int, default=10000, - help="The number of samples to draw from the prior for " - "plotting. Default is 10000.") -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--nsamples", + type=int, + default=10000, + help="The number of samples to draw from the prior for plotting. Default is 10000.", +) +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) # parse the command line opts = parser.parse_args() @@ -107,14 +122,16 @@ else: logging.info("Plotting") fig, axis_dict = create_multidim_plot( - parameters, samples, + parameters, + samples, labels=labels, plot_marginal=True, marginal_percentiles=[5, 50, 95], plot_scatter=False, plot_density=True, plot_contours=True, - contour_percentiles=[50, 90],) + contour_percentiles=[50, 90], +) # set DPI @@ -124,10 +141,9 @@ fig.set_dpi(200) caption = """This plot shows the probability density function (PDF) from the prior distributions.""" title = "Prior Distribution" -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title=title, - caption=caption) +results.save_fig_with_metadata( + fig, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_samples b/bin/inference/pycbc_inference_plot_samples index 6f407087bf7..027ae699058 100644 --- a/bin/inference/pycbc_inference_plot_samples +++ b/bin/inference/pycbc_inference_plot_samples @@ -15,31 +15,38 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plots samples from inference sampler. -""" +"""Plots samples from inference sampler.""" import argparse import logging + from matplotlib import use -use('agg') -from matplotlib import pyplot as plt + +use("agg") +import sys + import numpy +from matplotlib import pyplot as plt + import pycbc from pycbc import results from pycbc.inference import io -import sys # command line usage -parser = argparse.parser = io.ResultsArgumentParser( - skip_args=['chains', 'iteration']) +parser = argparse.parser = io.ResultsArgumentParser(skip_args=["chains", "iteration"]) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--chains", nargs='+', default=None, - help="Chain/walker indices to plot. Options are 'all' or " - "one or more chain indices. Default is to plot the " - "average of all chains for the input " - "`--parameters`.") -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--chains", + nargs="+", + default=None, + help="Chain/walker indices to plot. Options are 'all' or " + "one or more chain indices. Default is to plot the " + "average of all chains for the input " + "`--parameters`.", +) +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) # parse the command line opts = parser.parse_args() @@ -54,7 +61,7 @@ fp, parameters, labels, _ = io.results_from_cli(opts, load_samples=False) ndim = len(parameters) # get chain indices -if opts.chains == ['all'] or opts.chains == None: +if opts.chains == ["all"] or opts.chains == None: chains = range(fp.nchains) else: chains = list(map(int, opts.chains)) @@ -62,7 +69,7 @@ else: # plot samples # plot each parameter as a different subplot logging.info("Plotting samples") -fig, axs = plt.subplots(ndim, figsize=(10,8), sharex=True) +fig, axs = plt.subplots(ndim, figsize=(10, 8), sharex=True) plt.xlabel("Iteration") # loop over parameters @@ -72,7 +79,7 @@ if opts.thin_start is not None: xmin = opts.thin_start else: try: - xmin = fp.attrs['thin_start'] + xmin = fp.attrs["thin_start"] except KeyError: xmin = 0 @@ -80,62 +87,65 @@ if opts.thin_interval is not None: xint = opts.thin_interval else: try: - xint = fp.attrs['thin_interval'] + xint = fp.attrs["thin_interval"] except KeyError: xint = 1 -thinned_by = fp.thinned_by*xint -xmin = xmin*fp.thinned_by +thinned_by = fp.thinned_by * xint +xmin = xmin * fp.thinned_by # create the kwargs to load samples -kwargs = {'thin_start': opts.thin_start, - 'thin_interval': opts.thin_interval, - 'thin_end': opts.thin_end} +kwargs = { + "thin_start": opts.thin_start, + "thin_interval": opts.thin_interval, + "thin_end": opts.thin_end, +} # add the temperature args if it exists try: - kwargs['temps'] = opts.temps + kwargs["temps"] = opts.temps except AttributeError: pass for i, arg in enumerate(parameters): chains_arg = [] for cidx in chains: - kwargs['chains'] = cidx + kwargs["chains"] = cidx try: chain = fp.read_samples(parameters, **kwargs) except TypeError: # will get this if ensemble sampler; change "chains" to "walkers" - kwargs['walkers'] = kwargs.pop('chains') + kwargs["walkers"] = kwargs.pop("chains") chain = fp.read_samples(parameters, **kwargs) chains_arg.append(chain[arg]) if opts.chains is not None: for chain in chains_arg: # plot each chain as a different line on the subplot - axs[i].plot((numpy.arange(len(chain)))*thinned_by + xmin, chain, - alpha=0.6) + axs[i].plot( + (numpy.arange(len(chain))) * thinned_by + xmin, chain, alpha=0.6 + ) else: # plot the average of all chains for the parameter on the subplot chains_arg = numpy.array(chains_arg) - avg_chain = [chains_arg[:, j].sum()/fp.nchains - for j in range(len(chains_arg[0]))] - axs[i].plot((numpy.arange(len(avg_chain)))*thinned_by + xmin, avg_chain) + avg_chain = [ + chains_arg[:, j].sum() / fp.nchains for j in range(len(chains_arg[0])) + ] + axs[i].plot((numpy.arange(len(avg_chain))) * thinned_by + xmin, avg_chain) # Set y labels axs[i].set_ylabel(labels[arg]) fp.close() # save figure with meta-data caption_kwargs = { - "parameters" : ", ".join(sorted(list(labels.values()))), + "parameters": ", ".join(sorted(list(labels.values()))), } caption = r"""Parameter samples from the chains whose indices were provided as inputs. Each line is a different chain of samples in that case. If no chain indices were provided, the plot shows the variation of the parameter sample values averaged over all chains.""" title = "Samples for {parameters}".format(**caption_kwargs) -results.save_fig_with_metadata(fig, opts.output_file, - cmd=" ".join(sys.argv), - title=title, - caption=caption) +results.save_fig_with_metadata( + fig, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) plt.close() # exit diff --git a/bin/inference/pycbc_inference_plot_skymap b/bin/inference/pycbc_inference_plot_skymap index a2660e48c78..52e63e698a4 100644 --- a/bin/inference/pycbc_inference_plot_skymap +++ b/bin/inference/pycbc_inference_plot_skymap @@ -16,7 +16,8 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Wrapper around ligo-skymap-plot that creates a skymap plot from a fits file. +""" +Wrapper around ligo-skymap-plot that creates a skymap plot from a fits file. The main purpose of this is to add metadata to the resulting png plot, so that it can be used in an inference results page. If you don't care about that, and want more control on the plot options, run ligo-skaymap-plot instead. @@ -26,39 +27,39 @@ pycbc_inference_create_fits. This requires ligo.skymap to be installed, which requires python 3. """ + import argparse -import sys import subprocess +import sys + from PIL import Image, PngImagePlugin from pycbc import add_common_pycbc_options, init_logging parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--input-file', required=True, - help='Input fits file"') -parser.add_argument('--output-file', required=True, - help='Output png file."') -parser.add_argument('--colormap', - help='Specify the colormap to use.') +parser.add_argument("--input-file", required=True, help='Input fits file"') +parser.add_argument("--output-file", required=True, help='Output png file."') +parser.add_argument("--colormap", help="Specify the colormap to use.") opts = parser.parse_args() init_logging(opts.verbose) -cmd = 'ligo-skymap-plot {} -o {} --annotate --contour 50 90'.format( - opts.input_file, opts.output_file) +cmd = f"ligo-skymap-plot {opts.input_file} -o {opts.output_file} --annotate --contour 50 90" if opts.colormap is not None: - cmd += ' --colormap {}'.format(opts.colormap) + cmd += f" --colormap {opts.colormap}" ret = subprocess.run(cmd.split()) ret.check_returncode() im = Image.open(opts.output_file) meta = PngImagePlugin.PngInfo() -kwds = {'cmd': " ".join(sys.argv), - 'title': "Sky map", - 'caption': "Sky location posterior probability."} +kwds = { + "cmd": " ".join(sys.argv), + "title": "Sky map", + "caption": "Sky location posterior probability.", +} for key in kwds: meta.add_text(str(key), str(kwds[key])) im.save(opts.output_file, "png", pnginfo=meta) diff --git a/bin/inference/pycbc_inference_plot_thermodynamic_integrand b/bin/inference/pycbc_inference_plot_thermodynamic_integrand index c48ec990001..bf594057ddc 100644 --- a/bin/inference/pycbc_inference_plot_thermodynamic_integrand +++ b/bin/inference/pycbc_inference_plot_thermodynamic_integrand @@ -1,6 +1,6 @@ #!/usr/bin/env python -# Copyright (C) 2019 Steven Reyes +# Copyright (C) 2019 Steven Reyes # # This program is free software; you can redistribute it and/or modify it # under the terms of the GNU General Public License as published by the @@ -26,7 +26,9 @@ # import argparse + import matplotlib + matplotlib.use("agg") import matplotlib.pyplot as plt import numpy @@ -36,20 +38,37 @@ from pycbc.inference import io parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument("--inference-file", type=str, - help="The PyCBC multi-temper inference file.") -parser.add_argument("--thin-start", type=int, default=None, - help="MCMC iteration to begin drawing samples from.") -parser.add_argument("--thin-end", type=int, default=None, - help="MCMC iteration to end drawing samples from.") -parser.add_argument("--thin-interval", type=int, default=None, - help="MCMC intervals to downselect samples between" - " thin-start and thin-end.") -parser.add_argument("--beta-log-scale", action="store_true", - help="Put inverse-temperature axis in logarithm" - " base 10 scale.") -parser.add_argument("--integrand-logarithmic", action="store_true", - help="Plot the integrand as integrated in log beta.") +parser.add_argument( + "--inference-file", type=str, help="The PyCBC multi-temper inference file." +) +parser.add_argument( + "--thin-start", + type=int, + default=None, + help="MCMC iteration to begin drawing samples from.", +) +parser.add_argument( + "--thin-end", + type=int, + default=None, + help="MCMC iteration to end drawing samples from.", +) +parser.add_argument( + "--thin-interval", + type=int, + default=None, + help="MCMC intervals to downselect samples between thin-start and thin-end.", +) +parser.add_argument( + "--beta-log-scale", + action="store_true", + help="Put inverse-temperature axis in logarithm base 10 scale.", +) +parser.add_argument( + "--integrand-logarithmic", + action="store_true", + help="Plot the integrand as integrated in log beta.", +) parser.add_argument("--output-file", type=str) args = parser.parse_args() @@ -57,10 +76,13 @@ init_logging(args.verbose) # Read in the necessary data fp = io.loadfile(args.inference_file, "r") -logl = fp.read_samples("loglikelihood", thin_start=args.thin_start, - thin_end=args.thin_end, - thin_interval=args.thin_interval, - flatten=False)["loglikelihood"] +logl = fp.read_samples( + "loglikelihood", + thin_start=args.thin_start, + thin_end=args.thin_end, + thin_interval=args.thin_interval, + flatten=False, +)["loglikelihood"] betas = fp["sampler_info"].attrs["betas"] fp.close() @@ -70,21 +92,21 @@ betas = betas[order] logl = logl[order] if args.integrand_logarithmic: - y = betas * numpy.average(logl, axis=(1,2)) + y = betas * numpy.average(logl, axis=(1, 2)) plt.ylabel(r"$\beta \, \langle ln \, \mathcal{L} \rangle_\beta$") plt.xscale("log") else: - y = numpy.average(logl, axis=(1,2)) + y = numpy.average(logl, axis=(1, 2)) plt.ylabel(r"$\langle ln \, \mathcal{L} \rangle_\beta$") if args.beta_log_scale: plt.xscale("log") plt.minorticks_on() -plt.grid(True, which='minor', axis='y', ls="--", alpha=0.2) -plt.grid(True, which='major', axis='y', ls="-", alpha=0.4) -plt.grid(True, which='major', axis='x', ls="-", alpha=0.4) +plt.grid(True, which="minor", axis="y", ls="--", alpha=0.2) +plt.grid(True, which="major", axis="y", ls="-", alpha=0.4) +plt.grid(True, which="major", axis="x", ls="-", alpha=0.4) plt.plot(betas, y, marker="v") plt.xlabel(r"$\beta$") diff --git a/bin/inference/pycbc_inference_pp_table_summary b/bin/inference/pycbc_inference_pp_table_summary index 90c4fe2a217..ef03dfacf2f 100644 --- a/bin/inference/pycbc_inference_pp_table_summary +++ b/bin/inference/pycbc_inference_pp_table_summary @@ -15,24 +15,31 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Performs a PP test on a collection of parameters, writing results out to +""" +Performs a PP test on a collection of parameters, writing results out to an html table. """ import logging -import numpy import sys + +import numpy from scipy import stats + import pycbc from pycbc import results +from pycbc.inference.io import ( + ResultsArgumentParser, + injections_from_cli, + results_from_cli, +) from pycbc.inference.option_utils import add_injsamples_map_opt -from pycbc.inference.io import (ResultsArgumentParser, results_from_cli, - injections_from_cli) parser = ResultsArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) add_injsamples_map_opt(parser) # parse the command line @@ -41,7 +48,7 @@ opts = parser.parse_args() pycbc.init_logging(opts.verbose) # read results -logging.info('Loading parameters') +logging.info("Loading parameters") _, parameters, labels, samples = results_from_cli(opts) # typecast to list for iteration @@ -58,7 +65,7 @@ for input_file, input_samples in zip(opts.input_file, samples): for p in parameters: inj_val = inj_parameters[p] sample_vals = input_samples[p] - measured = stats.percentileofscore(sample_vals, inj_val, kind='weak') + measured = stats.percentileofscore(sample_vals, inj_val, kind="weak") try: measured_percentiles[p].append(measured) except KeyError: @@ -68,63 +75,73 @@ for input_file, input_samples in zip(opts.input_file, samples): logging.info("Performing percentile-percentile test") p_values = numpy.zeros(len(parameters)) table = [] -fmt = '{:.3f}' +fmt = "{:.3f}" for ii, param in enumerate(parameters): row = [labels[param]] meas = numpy.array(measured_percentiles[param]) meas.sort() - expected = numpy.array([stats.percentileofscore(meas, x, kind='weak') - for x in meas]) + expected = numpy.array( + [stats.percentileofscore(meas, x, kind="weak") for x in meas] + ) # perform ks test - ks, p = stats.kstest(meas/100., 'uniform') + ks, p = stats.kstest(meas / 100.0, "uniform") p_values[ii] = p row += [fmt.format(ks), fmt.format(p)] table.append(row) # do the p-value of p-values test if more than one parameter provided if len(parameters) > 1: - summary_ks, summary_p = stats.kstest(p_values, 'uniform') - table.append(['p-values', - ''+fmt.format(summary_ks)+'', - ''+fmt.format(summary_p)+'']) + summary_ks, summary_p = stats.kstest(p_values, "uniform") + table.append( + [ + "p-values", + "" + fmt.format(summary_ks) + "", + "" + fmt.format(summary_p) + "", + ] + ) # for the caption extra_caption = ( - ' If all parameters ' - 'satisfy the percentile-percentile test, then the distribution of ' - 'p-values should be uniform. The p-value of p-values is obtained ' - 'by applying the KS test to the distribution of p-values. The ' - 'smaller this value, the lower the probability of obtaining the ' - 'observed distribution of parameters assuming that the analysis ' - 'did satsify the expectation that $X$ percent of injections fall ' - 'in the $X$ percent credible interval.') + " If all parameters " + "satisfy the percentile-percentile test, then the distribution of " + "p-values should be uniform. The p-value of p-values is obtained " + "by applying the KS test to the distribution of p-values. The " + "smaller this value, the lower the probability of obtaining the " + "observed distribution of parameters assuming that the analysis " + "did satsify the expectation that $X$ percent of injections fall " + "in the $X$ percent credible interval." + ) else: - extra_caption = '' + extra_caption = "" # create table header headers = ["Parameter", "$D_{{KS}}$", "p-value"] # add mathjax header to display latex -html = results.mathjax_html_header() + '\n%s'%( - str(results.static_table(table, headers) )) +html = results.mathjax_html_header() + "\n%s" % ( + str(results.static_table(table, headers)) +) # add the number of injections that were used -mdatatmplt = '

{}: {}

' -metadata = [mdatatmplt.format('Number of injections', ninjections)] +mdatatmplt = "

{}: {}

" +metadata = [mdatatmplt.format("Number of injections", ninjections)] -html += '\n'.join(metadata) + '
\n
\n' +html += "\n".join(metadata) + "
\n
\n" # save HTML table results.save_fig_with_metadata( - html, opts.output_file, {}, + html, + opts.output_file, + {}, cmd=" ".join(sys.argv), title="Percentile-Percentile Test", - caption='Percentile-percentile test results. The value of the KS ' - 'statistic $D_{KS}$ gives the maximum distance ' - 'between the the measured CDF and the ' - 'expected (uniform) CDF. The associated two-tailed p-value gives ' - 'the probability of getting a maximum distance larger than the ' - 'observed $D_{KS}$ ' - 'assuming that the measured CDF is the same as the expected. ' - 'In other words, the larger (smaller) the p-value ($D_{KS}$), the ' - 'more likely the measured distribution is the same as the ' - 'expected.'+extra_caption) + caption="Percentile-percentile test results. The value of the KS " + "statistic $D_{KS}$ gives the maximum distance " + "between the the measured CDF and the " + "expected (uniform) CDF. The associated two-tailed p-value gives " + "the probability of getting a maximum distance larger than the " + "observed $D_{KS}$ " + "assuming that the measured CDF is the same as the expected. " + "In other words, the larger (smaller) the p-value ($D_{KS}$), the " + "more likely the measured distribution is the same as the " + "expected." + extra_caption, +) diff --git a/bin/inference/pycbc_inference_savage_dickey b/bin/inference/pycbc_inference_savage_dickey index 87732b71002..b37ce01d27a 100644 --- a/bin/inference/pycbc_inference_savage_dickey +++ b/bin/inference/pycbc_inference_savage_dickey @@ -24,8 +24,9 @@ # # ============================================================================= # -"""Estimates the Bayes factor using the Savage-Dickey density ratio from a -given posterior file. +""" +Estimates the Bayes factor using the Savage-Dickey density ratio from a +given posterior file. The Savage-Dickey ratio is defined with respect to a single parameter. The ratio compares two models: a null hypothesis, where the parameter of interest @@ -48,86 +49,131 @@ required for the given parameter. """ import os + import numpy +from scipy.stats import gaussian_kde + import pycbc -from pycbc.inference.io import (ResultsArgumentParser, loadfile) from pycbc import distributions -from pycbc.distributions.utils import prior_from_config from pycbc.boundaries import Bounds -from scipy.stats import gaussian_kde +from pycbc.distributions.utils import prior_from_config +from pycbc.inference.io import ResultsArgumentParser, loadfile -parser = ResultsArgumentParser(defaultparams='all', autoparamlabels=False, - description=__doc__) +parser = ResultsArgumentParser( + defaultparams="all", autoparamlabels=False, description=__doc__ +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--parameter-null-value", required=True, - help="Value at which to evaluate the posterior and prior " - "when calculating the Savage-Dickey ratio.") -parser.add_argument("--kde-samples", type=int, default=50000, - help="The number of samples to use for generating the " - "prior KDE. Default 50000.") -parser.add_argument("--kde-bandwidth", type=float, default=0.01, - help="The bandwidth of the KDEs, which determine how " - "'smooth' the KDEs are. Lower values lead to smoother " - "KDEs; higher values lead to more 'jagged' KDEs but " - "will potentially capture more structure. Use " - "'--plot' to aid in fine-tuning this.") -parser.add_argument("--prior-name", type=str, default=None, - help="The type of prior distribution to generate. " - "Accepts valid prior names from the distributions " - "module (e.g. 'uniform', 'uniform_log10', etc.). " - "If specifying a name, a min and max must also be " - "specified using --prior-min and --prior-max " - "respectively. By default, the prior is read from " - "the input posterior file(s).") -parser.add_argument("--prior-min", type=float, default=None, - help="The minimum value of the prior if specifying " - "--prior-name. By default (i.e. if prior-name is not " - "specified), this is read from the input file(s).") -parser.add_argument("--prior-max", type=float, default=None, - help="The maximum value of the prior if specifying " - "--prior-name. By default (i.e. if prior-name is not " - "specified), this is read from the input file(s).") -parser.add_argument("--prior-cyclic", type=bool, default=False, - help="Specify whether the prior has cyclic bounds if " - "specifying --prior-name. By default (i.e. if " - "prior-name is not specified), this is read from the " - "input file(s).") -parser.add_argument("--reflect-kde-on-left", action="store_true", default=False, - help="Specify whether to reflect samples on the left bound " - "when generating KDEs. This is done to prevent KDE " - "sampling issues at the edges of a histogram. This is " - "not done by default.") -parser.add_argument("--reflect-kde-on-right", action="store_true", default=False, - help="Specify whether to reflect samples on the right bound " - "when generating KDEs. This is done to prevent KDE " - "sampling issues at the edges of a histogram. This is " - "not done by default.") -parser.add_argument("--plot-name", type=str, default=None, nargs='+', - help="Specify the name(s) of the output plot. If a name is " - "specified, a plot will be generated showing the " - "prior and posterior KDEs and samples for each file. " - "By default, a plot is not generated. One name is " - "required per file specified by '--input-file'.") +parser.add_argument( + "--parameter-null-value", + required=True, + help="Value at which to evaluate the posterior and prior " + "when calculating the Savage-Dickey ratio.", +) +parser.add_argument( + "--kde-samples", + type=int, + default=50000, + help="The number of samples to use for generating the prior KDE. Default 50000.", +) +parser.add_argument( + "--kde-bandwidth", + type=float, + default=0.01, + help="The bandwidth of the KDEs, which determine how " + "'smooth' the KDEs are. Lower values lead to smoother " + "KDEs; higher values lead to more 'jagged' KDEs but " + "will potentially capture more structure. Use " + "'--plot' to aid in fine-tuning this.", +) +parser.add_argument( + "--prior-name", + type=str, + default=None, + help="The type of prior distribution to generate. " + "Accepts valid prior names from the distributions " + "module (e.g. 'uniform', 'uniform_log10', etc.). " + "If specifying a name, a min and max must also be " + "specified using --prior-min and --prior-max " + "respectively. By default, the prior is read from " + "the input posterior file(s).", +) +parser.add_argument( + "--prior-min", + type=float, + default=None, + help="The minimum value of the prior if specifying " + "--prior-name. By default (i.e. if prior-name is not " + "specified), this is read from the input file(s).", +) +parser.add_argument( + "--prior-max", + type=float, + default=None, + help="The maximum value of the prior if specifying " + "--prior-name. By default (i.e. if prior-name is not " + "specified), this is read from the input file(s).", +) +parser.add_argument( + "--prior-cyclic", + type=bool, + default=False, + help="Specify whether the prior has cyclic bounds if " + "specifying --prior-name. By default (i.e. if " + "prior-name is not specified), this is read from the " + "input file(s).", +) +parser.add_argument( + "--reflect-kde-on-left", + action="store_true", + default=False, + help="Specify whether to reflect samples on the left bound " + "when generating KDEs. This is done to prevent KDE " + "sampling issues at the edges of a histogram. This is " + "not done by default.", +) +parser.add_argument( + "--reflect-kde-on-right", + action="store_true", + default=False, + help="Specify whether to reflect samples on the right bound " + "when generating KDEs. This is done to prevent KDE " + "sampling issues at the edges of a histogram. This is " + "not done by default.", +) +parser.add_argument( + "--plot-name", + type=str, + default=None, + nargs="+", + help="Specify the name(s) of the output plot. If a name is " + "specified, a plot will be generated showing the " + "prior and posterior KDEs and samples for each file. " + "By default, a plot is not generated. One name is " + "required per file specified by '--input-file'.", +) opts = parser.parse_args() pycbc.init_logging(opts.verbose) # read in the posterior of given parameter if len(opts.parameters) > 1: - raise ValueError("Multiple parameters specified. Only one parameter is " - "accepted at a time") + raise ValueError( + "Multiple parameters specified. Only one parameter is accepted at a time" + ) param = opts.parameters[0] null = float(opts.parameter_null_value) -print(f"Calculating Bayes factor for selected parameter {param} versus {param} " - f"at null value {null}") +print( + f"Calculating Bayes factor for selected parameter {param} versus {param} " + f"at null value {null}" +) print("======================================================================") priors = [] posteriors = [] for file in opts.input_file: - fp = loadfile(file, 'r') + fp = loadfile(file, "r") cp = fp.read_config_file() # read in the prior of given parameter @@ -137,9 +183,11 @@ for file in opts.input_file: elif opts.prior_max is None: raise ValueError("Must specify --prior-max when using --prior-name") else: - bounds = Bounds(min_bound = opts.prior_min, - max_bound = opts.prior_max, - cyclic = opts.prior_cyclic) + bounds = Bounds( + min_bound=opts.prior_min, + max_bound=opts.prior_max, + cyclic=opts.prior_cyclic, + ) prior_name = opts.prior_name else: if opts.prior_min is not None: @@ -148,7 +196,7 @@ for file in opts.input_file: raise ValueError("Must specify --prior-name when using --prior-max") full_prior = prior_from_config(cp) bounds = full_prior.bounds[param] - prior_name = cp.get_opt_tag('prior', 'name', param) + prior_name = cp.get_opt_tag("prior", "name", param) kws = {param: bounds} prior = distributions.distribs[prior_name](**kws) @@ -158,13 +206,15 @@ for file in opts.input_file: pos_samples = fp.read_samples([param])[param] # if in log scale, rescale the distributions - if 'log' in prior.name: + if "log" in prior.name: pos = numpy.log10(pos_samples) pri = numpy.log10(prior_samples) null = numpy.log10(null) - bounds = Bounds(min_bound = numpy.log10(bounds.min), - max_bound = numpy.log10(bounds.max), - cyclic = bounds.cyclic) + bounds = Bounds( + min_bound=numpy.log10(bounds.min), + max_bound=numpy.log10(bounds.max), + cyclic=bounds.cyclic, + ) else: pos = pos_samples pri = prior_samples @@ -173,14 +223,14 @@ for file in opts.input_file: net_pos = pos net_pri = pri if opts.reflect_kde_on_left: - left_pos = 2*bounds.min - pos + left_pos = 2 * bounds.min - pos net_pos = numpy.append(net_pos, left_pos) - left_pri = 2*bounds.min - pri + left_pri = 2 * bounds.min - pri net_pri = numpy.append(net_pri, left_pri) if opts.reflect_kde_on_right: - right_pos = 2*bounds.max - pos + right_pos = 2 * bounds.max - pos net_pos = numpy.append(net_pos, right_pos) - right_pri = 2*bounds.max - pri + right_pri = 2 * bounds.max - pri net_pri = numpy.append(net_pri, right_pri) # generate KDEs @@ -200,29 +250,31 @@ for file in opts.input_file: # plotting if opts.plot_name is not None: if len(opts.plot_name) != len(opts.input_file): - raise KeyError(f"Number of plot names ({opts.plot_name}) " - f"does not match number of input files " - f"({opts.input_file})") + raise KeyError( + f"Number of plot names ({opts.plot_name}) " + f"does not match number of input files " + f"({opts.input_file})" + ) import matplotlib.pyplot as plt if "log" not in prior.name: # non-log distributions - vals, bins, _ = plt.hist(pos, bins=100, density=True, alpha=0.4, - label="Posterior Samples") - plt.hist(pri, bins=bins, density=True, alpha=0.4, - label="Prior Samples") - plt.plot(bins, pos_pdf(bins)/pos_norm, color='green', - label="Posterior KDE") - plt.plot(bins, prior_pdf(bins)/prior_norm, color='red', - label="Prior KDE") + vals, bins, _ = plt.hist( + pos, bins=100, density=True, alpha=0.4, label="Posterior Samples" + ) + plt.hist(pri, bins=bins, density=True, alpha=0.4, label="Prior Samples") + plt.plot( + bins, pos_pdf(bins) / pos_norm, color="green", label="Posterior KDE" + ) + plt.plot(bins, prior_pdf(bins) / prior_norm, color="red", label="Prior KDE") else: # special handling for log distributions bins = numpy.logspace(bounds.min, bounds.max, 100) log_bins = numpy.log10(bins) - plt.hist(10**pos, bins=bins, density=True, alpha=0.4, - label="Posterior Samples") - plt.hist(10**pri, bins=bins, density=True, alpha=0.4, - label="Prior Samples") + plt.hist( + 10**pos, bins=bins, density=True, alpha=0.4, label="Posterior Samples" + ) + plt.hist(10**pri, bins=bins, density=True, alpha=0.4, label="Prior Samples") # convert log scale KDEs to linear scale # p(x) = p(y) / (x ln 10), where y = log x jac = 1 / bins / numpy.log(10) @@ -230,19 +282,19 @@ for file in opts.input_file: pri_plot = prior_pdf(log_bins) * jac pos_plot_norm = numpy.trapezoid(pos_plot, bins) pri_plot_norm = numpy.trapezoid(pri_plot, bins) - plt.plot(bins, pos_plot/pos_plot_norm, color='green', - label="Posterior KDE") - plt.plot(bins, pri_plot/pri_plot_norm, color='red', - label="Prior KDE") - plt.xscale('log') + plt.plot( + bins, pos_plot / pos_plot_norm, color="green", label="Posterior KDE" + ) + plt.plot(bins, pri_plot / pri_plot_norm, color="red", label="Prior KDE") + plt.xscale("log") f = os.path.basename(file) plt.title(f"Posterior and Prior of {param} from {f}") plt.xlabel(f"{param}") plt.legend() idx = opts.input_file.index(file) - plt.savefig(f'{opts.plot_name[idx]}') - plt.close() + plt.savefig(f"{opts.plot_name[idx]}") + plt.close() # print output print(f"File: {file}") @@ -250,5 +302,4 @@ for file in opts.input_file: print(f"Bayes factor: {B[0]:.3f}") else: print(f"Bayes factor: {B[0]:.5e}") - print("") - + print() diff --git a/bin/inference/pycbc_inference_start_from_samples b/bin/inference/pycbc_inference_start_from_samples index 5e99e2dfd9c..7a7c31b75ad 100644 --- a/bin/inference/pycbc_inference_start_from_samples +++ b/bin/inference/pycbc_inference_start_from_samples @@ -1,8 +1,8 @@ #!/bin/env python -""" Convert inference file to parallel temperred compatible start format -""" +"""Convert inference file to parallel temperred compatible start format""" import argparse + from numpy.random import choice from pycbc import add_common_pycbc_options, init_logging @@ -11,12 +11,15 @@ from pycbc.inference.sampler import samplers parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--input-file') -parser.add_argument('--output-file') -parser.add_argument('--sampler', default='emcee_pt', - help="The output sampler type, if none, we assume emcee_pt") -parser.add_argument('--ntemps', type=int) -parser.add_argument('--nwalkers', type=int) +parser.add_argument("--input-file") +parser.add_argument("--output-file") +parser.add_argument( + "--sampler", + default="emcee_pt", + help="The output sampler type, if none, we assume emcee_pt", +) +parser.add_argument("--ntemps", type=int) +parser.add_argument("--nwalkers", type=int) args = parser.parse_args() init_logging(opts.verbose) @@ -30,22 +33,22 @@ ntemps = args.ntemps nwalkers = args.nwalkers f = loadfile(args.input_file) -params = list(f.variable_params) + ['loglikelihood'] +params = list(f.variable_params) + ["loglikelihood"] samples = f.read_samples(params) nsample = len(samples) # These are the ids we'll use for the temps / walkers use = choice(nsample, replace=False, size=ntemps * nwalkers) -o = loadfile(args.output_file, 'w', filetype=samplers[args.sampler]._io.name) +o = loadfile(args.output_file, "w", filetype=samplers[args.sampler]._io.name) for k in params: data = samples[k][use] - o['samples/' + k] = data.reshape(ntemps, nwalkers, 1) + o["samples/" + k] = data.reshape(ntemps, nwalkers, 1) -o.attrs['static_params'] = [] -o.attrs['variable_params'] = f.variable_params -o.create_group('sampler_info') -o['sampler_info'].attrs['nchains'] = nwalkers +o.attrs["static_params"] = [] +o.attrs["variable_params"] = f.variable_params +o.create_group("sampler_info") +o["sampler_info"].attrs["nchains"] = nwalkers o.close() f.close() diff --git a/bin/inference/pycbc_inference_table_summary b/bin/inference/pycbc_inference_table_summary index 3c41d0df18c..3d61d0b0afd 100644 --- a/bin/inference/pycbc_inference_table_summary +++ b/bin/inference/pycbc_inference_table_summary @@ -17,35 +17,45 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import os -import numpy import sys + +import numpy + import pycbc from pycbc import results -from pycbc.inference.option_utils import ParseLabelArg from pycbc.inference.io import ResultsArgumentParser, results_from_cli +from pycbc.inference.option_utils import ParseLabelArg -parser = ResultsArgumentParser( - description="Makes a table of posterior results.") +parser = ResultsArgumentParser(description="Makes a table of posterior results.") pycbc.add_common_pycbc_options(parser) -parser.add_argument("--output-file", type=str, required=True, - help="Path to output plot.") -parser.add_argument("--print-metadata", nargs="+", metavar="PARAM[:LABEL]", - default=[], - action=ParseLabelArg, - help="Add metadata information after the parameter table. " - "Any parameter stored in the any of the file's attrs " - "may be printed. To specify an attribute in a " - "sub-group, prepend the parameter with the group " - "name in the file. A label may be provided for the " - "html output; otherwise, the parameter name will be " - "used. If nothing provided, will just print the " - "number of posterior samples that were used (this is " - "always printed).") -parser.add_argument("--percentiles", type=float, nargs=3, - default=[5, 50, 95], - help="Percentiles to calculate. Must provide 3 values. " - "Default is 5 50 95 (the median with 90%% credible " - "interval).") +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) +parser.add_argument( + "--print-metadata", + nargs="+", + metavar="PARAM[:LABEL]", + default=[], + action=ParseLabelArg, + help="Add metadata information after the parameter table. " + "Any parameter stored in the any of the file's attrs " + "may be printed. To specify an attribute in a " + "sub-group, prepend the parameter with the group " + "name in the file. A label may be provided for the " + "html output; otherwise, the parameter name will be " + "used. If nothing provided, will just print the " + "number of posterior samples that were used (this is " + "always printed).", +) +parser.add_argument( + "--percentiles", + type=float, + nargs=3, + default=[5, 50, 95], + help="Percentiles to calculate. Must provide 3 values. " + "Default is 5 50 95 (the median with 90%% credible " + "interval).", +) # parse the command line opts = parser.parse_args() @@ -59,10 +69,10 @@ if ninputs > 1: # make sure the loglikelihood and logprior are included in the parameters # so that we can get the maxL and MAP values requested_params = [p for p in opts.parameters] -if 'loglikelihood' not in requested_params: - opts.parameters.append('loglikelihood') -if 'logprior' not in requested_params: - opts.parameters.append('logprior') +if "loglikelihood" not in requested_params: + opts.parameters.append("loglikelihood") +if "logprior" not in requested_params: + opts.parameters.append("logprior") # load the results fp, parameters, labels, samples = results_from_cli(opts) @@ -70,10 +80,10 @@ fp, parameters, labels, samples = results_from_cli(opts) parameters = requested_params # get the index of the map and maxl -if 'loglikelihood' in samples.fieldnames: - maxlidx = samples['loglikelihood'].argmax() - if 'logprior' in samples.fieldnames: - mapidx = samples['loglikelihood+logprior'].argmax() +if "loglikelihood" in samples.fieldnames: + maxlidx = samples["loglikelihood"].argmax() + if "logprior" in samples.fieldnames: + mapidx = samples["loglikelihood+logprior"].argmax() else: # the sampler didn't save the logprior mapidx = None @@ -89,13 +99,17 @@ for param in parameters: row = [labels[param]] # calculate the score at a given percentile x = samples[param] - percentiles = numpy.array([numpy.percentile(x, q) - for q in sorted(opts.percentiles)]) + percentiles = numpy.array( + [numpy.percentile(x, q) for q in sorted(opts.percentiles)] + ) values_min, values_med, values_max = percentiles negerror = values_med - values_min poserror = values_max - values_med - fmt = '${0}$'.format(results.format_value( - values_med, negerror, plus_error=poserror, use_scientific_notation=5)) + fmt = "${0}$".format( + results.format_value( + values_med, negerror, plus_error=poserror, use_scientific_notation=5 + ) + ) row.append(fmt) # get the maxl and map values mapval = x[mapidx] @@ -106,57 +120,68 @@ for param in parameters: for idx in [mapidx, maxlidx]: if idx is not None: val = x[idx] - fmt = '${0}$'.format(results.format_value( - val, error, use_scientific_notation=5, include_error=False)) + fmt = "${0}$".format( + results.format_value( + val, error, use_scientific_notation=5, include_error=False + ) + ) else: # sampler didn't provide a loglikelihood or logprior, so just # enter nothing - fmt = '--' + fmt = "--" row.append(fmt) # add to the table table.append(row) # create table header -interval = opts.percentiles[-1]-opts.percentiles[0] -headers = ["Parameter", - "{0:d}% Credible Interval".format(int(interval)), - "Maximum Posterior", - "Maximum Likelihood" - ] +interval = opts.percentiles[-1] - opts.percentiles[0] +headers = [ + "Parameter", + f"{int(interval):d}% Credible Interval", + "Maximum Posterior", + "Maximum Likelihood", +] # add mathjax header to display latex -html = results.mathjax_html_header() + '\n%s'%( - str(results.static_table(table, headers) )) +html = results.mathjax_html_header() + "\n%s" % ( + str(results.static_table(table, headers)) +) # add extra metadata -mdatatmplt = '

{}: {}

' +mdatatmplt = "

{}: {}

" + + def formatattr(fp, attr): group = os.path.dirname(attr) - if not group.startswith('/'): - group = '/' + group + if not group.startswith("/"): + group = "/" + group attr = os.path.basename(attr) val = fp[group].attrs[attr] if isinstance(val, numpy.ndarray): - val = ', '.join(['{}'.format(x) for x in val]) + val = ", ".join([f"{x}" for x in val]) return val -metadata = [mdatatmplt.format(opts.print_metadata_labels[p], - formatattr(fp, p)) - for p in opts.print_metadata] + + +metadata = [ + mdatatmplt.format(opts.print_metadata_labels[p], formatattr(fp, p)) + for p in opts.print_metadata +] # add the number of posterior samples -metadata.append(mdatatmplt.format('Number of posterior samples', samples.size)) +metadata.append(mdatatmplt.format("Number of posterior samples", samples.size)) -html += '\n'.join(metadata) + '
\n
\n' +html += "\n".join(metadata) + "
\n
\n" # save HTML table results.save_fig_with_metadata( - html, opts.output_file, {}, + html, + opts.output_file, + {}, cmd=" ".join(sys.argv), title="Parameter Estimates", - caption="Summary of parameter estimates. The {0:d}-percent credible " - "interval is the {1:d}th +/- {2:d}th/{3:d}th percentiles. " - "The maximum posterior parameters are the parameters " - "with the largest likelihood * prior. The maximum " - "likelihood parameters are the parameters with the " - "largest likelihood." - .format(int(interval), int(opts.percentiles[1]), - int(opts.percentiles[2]), int(opts.percentiles[0]))) + caption=f"Summary of parameter estimates. The {int(interval):d}-percent credible " + f"interval is the {int(opts.percentiles[1]):d}th +/- {int(opts.percentiles[2]):d}th/{int(opts.percentiles[0]):d}th percentiles. " + "The maximum posterior parameters are the parameters " + "with the largest likelihood * prior. The maximum " + "likelihood parameters are the parameters with the " + "largest likelihood.", +) diff --git a/bin/inference/pycbc_validate_test_posterior b/bin/inference/pycbc_validate_test_posterior index aae4ed3e093..b86dc41b12b 100644 --- a/bin/inference/pycbc_validate_test_posterior +++ b/bin/inference/pycbc_validate_test_posterior @@ -1,40 +1,43 @@ #!/usr/bin/env python -""" Validate and generate diagnostic plots for a inference file using the +""" +Validate and generate diagnostic plots for a inference file using the test posterior model. """ + +import argparse import sys + import numpy -import argparse from matplotlib import use -use('Agg') -from matplotlib import pyplot as plt +use("Agg") +from matplotlib import pyplot as plt from scipy.stats import ks_2samp +from pycbc import add_common_pycbc_options, init_logging from pycbc.distributions.utils import prior_from_config -from pycbc.inference import models, io +from pycbc.inference import io, models from pycbc.io import FieldArray -from pycbc import add_common_pycbc_options, init_logging numpy.random.seed(0) parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--input-file', help='inference posterior file') -parser.add_argument('--output-file', help='diagnostic plot') -parser.add_argument('--p-value-threshold', help='minimum ks test p-value', - type=float) -parser.add_argument('--ind-samples', help='use only this number of samples', - default=1000, type=int) +parser.add_argument("--input-file", help="inference posterior file") +parser.add_argument("--output-file", help="diagnostic plot") +parser.add_argument("--p-value-threshold", help="minimum ks test p-value", type=float) +parser.add_argument( + "--ind-samples", help="use only this number of samples", default=1000, type=int +) args = parser.parse_args() init_logging(args.verbose) size = int(1e6) -d1 = io.loadfile(args.input_file, 'r') +d1 = io.loadfile(args.input_file, "r") -#We directly recreate the model and prior from the stored -#config to ensure the same configuration +# We directly recreate the model and prior from the stored +# config to ensure the same configuration config = d1.read_config_file() prior = prior_from_config(config) @@ -52,20 +55,19 @@ for dist in prior.distributions: ref = ref[(bound.min < ref[param]) & (ref[param] < bound.max)] nparam = len(model.variable_params) -fig, axs = plt.subplots(1, nparam, figsize=[6*nparam, 4], dpi=100) +fig, axs = plt.subplots(1, nparam, figsize=[6 * nparam, 4], dpi=100) result = d1.read_samples(model.variable_params) failed = False for param, ax in zip(model.variable_params, axs): - rpart = numpy.random.choice(result[param], replace=False, - size=args.ind_samples) + rpart = numpy.random.choice(result[param], replace=False, size=args.ind_samples) kv, pvalue = ks_2samp(ref[param], rpart) - print("{}, p-value={:.3f}".format(param, pvalue)) + print(f"{param}, p-value={pvalue:.3f}") plt.sca(ax) - plt.hist(ref[param], density=True, bins=30, label='reference') - plt.hist(result[param], density=True, bins=30, alpha=0.5, label='sampler') - plt.title('KS p-value = {:.4f}'.format(pvalue)) + plt.hist(ref[param], density=True, bins=30, label="reference") + plt.hist(result[param], density=True, bins=30, alpha=0.5, label="sampler") + plt.title(f"KS p-value = {pvalue:.4f}") plt.xlabel(param) plt.legend() ax.get_yaxis().set_visible(False) diff --git a/bin/live/pycbc_live_collate_triggers b/bin/live/pycbc_live_collate_triggers index bfabb0172f1..a0a79aadf0f 100644 --- a/bin/live/pycbc_live_collate_triggers +++ b/bin/live/pycbc_live_collate_triggers @@ -14,17 +14,17 @@ """Find trigger files and combine them into a single hdf trigger merge file.""" -import numpy import argparse -import h5py -import os import logging +import os -from igwn_segments import segmentlist, segment +import h5py +import numpy +from igwn_segments import segment, segmentlist import pycbc -from pycbc.io import live as liveio from pycbc.events import cuts, veto +from pycbc.io import live as liveio # Set up the command line argument parser parser = argparse.ArgumentParser(description=__doc__) @@ -33,21 +33,17 @@ liveio.add_live_trigger_selection_options(parser) cuts.insert_cuts_option_group(parser) parser.add_argument( - '--ifos', - nargs='+', + "--ifos", + nargs="+", required=True, - help="The list of detectors to include triggers in the merged file" + help="The list of detectors to include triggers in the merged file", ) parser.add_argument( - '--output-file', - required=True, - help='The output file containing merged triggers.' + "--output-file", required=True, help="The output file containing merged triggers." ) parser.add_argument( - "--bank-file", - required=True, - help="The bank file used in the search" + "--bank-file", required=True, help="The bank file used in the search" ) args = parser.parse_args() @@ -68,22 +64,19 @@ args.trigger_cuts.append(f"end_time:{args.gps_end_time}:upper_inc") trigger_cut_dict, template_cut_dict = cuts.ingest_cuts_option_group(args) -logging.info( - "Collating triggers to %s", - args.output_file -) +logging.info("Collating triggers to %s", args.output_file) # Some tracking objects file_count = 0 -n_triggers = {ifo: 0 for ifo in args.ifos} -n_triggers_cut = {ifo: 0 for ifo in args.ifos} +n_triggers = dict.fromkeys(args.ifos, 0) +n_triggers_cut = dict.fromkeys(args.ifos, 0) segs = {ifo: segmentlist([]) for ifo in args.ifos} -with h5py.File(args.bank_file,'r') as bank_file: +with h5py.File(args.bank_file, "r") as bank_file: # Count the number of templates - n_templates = bank_file['template_hash'].size + n_templates = bank_file["template_hash"].size -with h5py.File(args.output_file, 'w') as destination: +with h5py.File(args.output_file, "w") as destination: # Create the ifo groups in the trigger file for ifo in args.ifos: if ifo not in destination: @@ -92,21 +85,17 @@ with h5py.File(args.output_file, 'w') as destination: for file_count, source_file in enumerate(trigger_files): trigger_file = os.path.basename(source_file) - start_time = float(trigger_file.split('-')[2]) - duration = float(trigger_file.split('-')[3][:-4]) + start_time = float(trigger_file.split("-")[2]) + duration = float(trigger_file.split("-")[3][:-4]) end_time = start_time + duration if file_count % 100 == 0: - logging.info( - "Files appended: %d/%d", - file_count, - len(trigger_files) - ) + logging.info("Files appended: %d/%d", file_count, len(trigger_files)) - with h5py.File(source_file, 'r') as source: + with h5py.File(source_file, "r") as source: for ifo in args.ifos: try: - n_trigs_ifo = source[ifo]['snr'].size + n_trigs_ifo = source[ifo]["snr"].size n_triggers[ifo] += n_trigs_ifo except KeyError: # No triggers in this IFO in this file @@ -116,14 +105,15 @@ with h5py.File(args.output_file, 'w') as destination: segs[ifo].append(segment(start_time, end_time)) triggers = { - k: source[ifo][k][:] for k in source[ifo].keys() - if k not in ('loudest', 'stat', 'gates', 'psd') + k: source[ifo][k][:] + for k in source[ifo].keys() + if k not in ("loudest", "stat", "gates", "psd") and source[ifo][k].size == n_trigs_ifo } # The stored chisq is actually reduced chisq, so convert back # to unreduced chisq using chisq_dof - tmpchisq = triggers['chisq'][:] * (2 * triggers['chisq_dof'][:] - 2) - triggers['chisq'][:] = tmpchisq + tmpchisq = triggers["chisq"][:] * (2 * triggers["chisq_dof"][:] - 2) + triggers["chisq"][:] = tmpchisq # Apply the cuts to triggers keep_idx = cuts.apply_trigger_cuts(triggers, trigger_cut_dict) @@ -131,39 +121,30 @@ with h5py.File(args.output_file, 'w') as destination: # triggers contains the datasets that we want to use for # the template cuts, so here it can be used as the template bank keep_idx = cuts.apply_template_cuts( - triggers, - template_cut_dict, - template_ids=keep_idx + triggers, template_cut_dict, template_ids=keep_idx ) if not any(keep_idx): # No triggers kept after cuts in this ifo for this file continue n_triggers_cut[ifo] += keep_idx.size - triggers = { - k: triggers[k][keep_idx] - for k in triggers.keys() - } + triggers = {k: triggers[k][keep_idx] for k in triggers} - if ('approximant' not in triggers) or (len(triggers['approximant']) == 0): + if ("approximant" not in triggers) or ( + len(triggers["approximant"]) == 0 + ): continue for name, dataset in triggers.items(): if name in destination[ifo]: # Append new data to existing dataset in destination - if triggers[name].shape[0] == 0: + if dataset.shape[0] == 0: continue - destination[ifo][name].resize( - n_triggers_cut[ifo], - axis=0 - ) - destination[ifo][name][-keep_idx.size:] = dataset + destination[ifo][name].resize(n_triggers_cut[ifo], axis=0) + destination[ifo][name][-keep_idx.size :] = dataset else: destination[ifo].create_dataset( - name, - data=dataset, - chunks=True, - maxshape=(None,) + name, data=dataset, chunks=True, maxshape=(None,) ) for attr_name, attr_value in source.attrs.items(): @@ -171,22 +152,19 @@ with h5py.File(args.output_file, 'w') as destination: for ifo in args.ifos: # Collect the segments, and output to the file - search_grp = destination[ifo].create_group('search') + search_grp = destination[ifo].create_group("search") segs[ifo].coalesce() seg_starts, seg_ends = veto.segments_to_start_end(segs[ifo]) search_grp.create_dataset( - 'end_time', + "end_time", data=seg_ends, ) search_grp.create_dataset( - 'start_time', + "start_time", data=seg_starts, ) logging.info( - "Found %d %s triggers in %s seconds", - n_triggers[ifo], - ifo, - abs(segs[ifo]) + "Found %d %s triggers in %s seconds", n_triggers[ifo], ifo, abs(segs[ifo]) ) if n_triggers[ifo] != n_triggers_cut[ifo]: logging.info( @@ -199,7 +177,7 @@ with h5py.File(args.output_file, 'w') as destination: logging.info("Processing %s", ifo) triggers = destination[ifo] try: - template_ids = triggers['template_id'][:] + template_ids = triggers["template_id"][:] except KeyError: logging.info("No triggers for %s, skipping", ifo) continue @@ -208,14 +186,12 @@ with h5py.File(args.output_file, 'w') as destination: sorted_indices = numpy.argsort(template_ids) sorted_template_ids = template_ids[sorted_indices] unique_template_ids, template_id_counts = numpy.unique( - sorted_template_ids, - return_counts=True + sorted_template_ids, return_counts=True ) template_boundaries = numpy.searchsorted( - sorted_template_ids, - numpy.arange(n_templates) + sorted_template_ids, numpy.arange(n_templates) ) - triggers['template_boundaries'] = template_boundaries + triggers["template_boundaries"] = template_boundaries # Sort other datasets by template_id so it makes sense # Datasets with the same length as the number of triggers: @@ -224,29 +200,33 @@ with h5py.File(args.output_file, 'w') as destination: # template_duration, template_hash, template_id for key in triggers.keys(): if len(triggers[key]) == len(template_ids): - logging.info('Sorting %s by template id', key) + logging.info("Sorting %s by template id", key) sorted_key = triggers[key][:][sorted_indices] triggers[key][:] = sorted_key logging.info("Setting up region references") # Datasets which need region references: - region_ref_datasets = ('chisq_dof', 'chisq', 'coa_phase', - 'end_time', 'sg_chisq', 'snr', - 'template_duration', 'sigmasq') - if 'psd_var_val' in triggers.keys(): - region_ref_datasets += ('psd_var_val',) - if 'dq_state' in triggers.keys(): - region_ref_datasets += ('dq_state',) + region_ref_datasets = ( + "chisq_dof", + "chisq", + "coa_phase", + "end_time", + "sg_chisq", + "snr", + "template_duration", + "sigmasq", + ) + if "psd_var_val" in triggers.keys(): + region_ref_datasets += ("psd_var_val",) + if "dq_state" in triggers.keys(): + region_ref_datasets += ("dq_state",) start_boundaries = template_boundaries end_boundaries = numpy.roll(start_boundaries, -1) end_boundaries[-1] = len(template_ids) for dataset in region_ref_datasets: - logging.info( - "Region references for %s", - dataset - ) + logging.info("Region references for %s", dataset) refs = [ triggers[dataset].regionref[l:r] for l, r in zip(start_boundaries, end_boundaries) @@ -254,9 +234,9 @@ with h5py.File(args.output_file, 'w') as destination: logging.info("Adding to file") triggers.create_dataset( - dataset + '_template', + dataset + "_template", data=refs, - dtype=h5py.special_dtype(ref=h5py.RegionReference) + dtype=h5py.special_dtype(ref=h5py.RegionReference), ) logging.info("Done!") diff --git a/bin/live/pycbc_live_collated_dq_trigger_rates b/bin/live/pycbc_live_collated_dq_trigger_rates index d16fb4d73d7..9bd138b06a5 100644 --- a/bin/live/pycbc_live_collated_dq_trigger_rates +++ b/bin/live/pycbc_live_collated_dq_trigger_rates @@ -17,20 +17,19 @@ Calculate the noise trigger rate adjustment as a function of data-quality state for a day of PyCBC Live triggers. """ -import logging import argparse +import glob import hashlib +import logging import os -import glob import numpy as np - -from igwn_segments import segmentlist, segment +from igwn_segments import segment, segmentlist import pycbc import pycbc.dq +from pycbc.frame.frame import query_and_read_frame, read_frame from pycbc.io import HFile -from pycbc.frame.frame import read_frame, query_and_read_frame def find_frames(ifo, frame_dir, frame_type, start, end): @@ -38,8 +37,8 @@ def find_frames(ifo, frame_dir, frame_type, start, end): Find the frame files for a given time range. """ ifo_frame_type = frame_type.format(ifo=ifo) - ifo_pattern = f'{ifo[0]}-{ifo_frame_type}_llhoft-' - folder_pattern = ifo_pattern + '{supertime}' + ifo_pattern = f"{ifo[0]}-{ifo_frame_type}_llhoft-" + folder_pattern = ifo_pattern + "{supertime}" path_pattern = os.path.join( frame_dir, ifo_frame_type, @@ -51,7 +50,7 @@ def find_frames(ifo, frame_dir, frame_type, start, end): files = [] for t in range(start_round, end_round + 1): framedir = path_pattern.format(supertime=t) - files += glob.glob(os.path.join(framedir, '*.gwf')) + files += glob.glob(os.path.join(framedir, "*.gwf")) return files @@ -62,12 +61,16 @@ parser.add_argument("--ifo", required=True) parser.add_argument("--gps-start-time", required=True, type=int) parser.add_argument("--gps-end-time", required=True, type=int) parser.add_argument("--template-bin-file", required=True) -parser.add_argument("--use-gwdatafind", action='store_true', - help='Use gwdatafind to find frame files. If provided, ' - 'frame-directory argument will be ignored.') -parser.add_argument("--frame-directory", - help='Directory containing frame files. Required if ' - 'not using gwdatafind.') +parser.add_argument( + "--use-gwdatafind", + action="store_true", + help="Use gwdatafind to find frame files. If provided, " + "frame-directory argument will be ignored.", +) +parser.add_argument( + "--frame-directory", + help="Directory containing frame files. Required if not using gwdatafind.", +) parser.add_argument("--frame-type", required=True) parser.add_argument("--analysis-flag-name") parser.add_argument("--dq-channel", required=True) @@ -82,7 +85,7 @@ args = parser.parse_args() pycbc.init_logging(args.verbose) if not args.use_gwdatafind and args.frame_directory is None: - raise ValueError('Must provide frame-directory if not using gwdatafind.') + raise ValueError("Must provide frame-directory if not using gwdatafind.") if args.analysis_flag_name is not None: # Get observing segs @@ -93,15 +96,13 @@ if args.analysis_flag_name is not None: else: # If we are not querying the database for segments, just assume that the # full time is in observing mode - observing_flag = segmentlist( - [segment(args.gps_start_time, args.gps_end_time)] - ) + observing_flag = segmentlist([segment(args.gps_start_time, args.gps_end_time)]) # shift observing flag to current time observing_flag.protract(args.replay_offset) livetime = abs(observing_flag) -logging.info(f'Found {livetime} seconds of observing time at {args.ifo}.') +logging.info(f"Found {livetime} seconds of observing time at {args.ifo}.") # for each segment, check how much time was dq flagged flagged_time = 0 @@ -113,7 +114,7 @@ for seg in observing_flag: args.frame_type.format(ifo=args.ifo), [dq_channel, dq_ok_channel], start_time=seg[0], - end_time=seg[1] + end_time=seg[1], ) else: frames = find_frames( @@ -125,8 +126,7 @@ for seg in observing_flag: ) ts_dq_channel, ts_dq_ok_channel = read_frame( - frames, [dq_channel, dq_ok_channel], - start_time=seg[0], end_time=seg[1] + frames, [dq_channel, dq_ok_channel], start_time=seg[0], end_time=seg[1] ) # build segmentlists: @@ -135,7 +135,7 @@ for seg in observing_flag: dq_ok_segs = ts_dq_ok_channel.bool_to_segmentlist() # Is the channel above threshold? - dq_mask = (ts_dq_channel >= args.dq_thresh) + dq_mask = ts_dq_channel >= args.dq_thresh # If not, set it to zero / one accordingly ts_dq_channel._data = dq_mask dq_segs = ts_dq_channel.bool_to_segmentlist() @@ -151,82 +151,88 @@ for seg in observing_flag: # intersection with observing_flag and add livetime flagged_time += abs(valid_bad_dq & observing_flag) -logging.info(f'Found {flagged_time} seconds of dq flagged time at {args.ifo}.') +logging.info(f"Found {flagged_time} seconds of dq flagged time at {args.ifo}.") bg_livetime = livetime - flagged_time state_time = np.array([bg_livetime, flagged_time]) # read in template bins -logging.info(f'Reading template bins from {args.template_bin_file}') +logging.info(f"Reading template bins from {args.template_bin_file}") template_bins = {} num_templates = 0 -with HFile(args.template_bin_file, 'r') as bin_file: - f_lower = bin_file.attrs['f_lower'] - bank_file = bin_file.attrs['bank_file'] - bin_string = bin_file.attrs['background_bins'] +with HFile(args.template_bin_file, "r") as bin_file: + f_lower = bin_file.attrs["f_lower"] + bank_file = bin_file.attrs["bank_file"] + bin_string = bin_file.attrs["background_bins"] # only ever one group in this file grp = bin_file[list(bin_file.keys())[0]] num_bins = len(grp.keys()) for k in grp.keys(): - template_bins[k] = grp[k]['tids'][:] + template_bins[k] = grp[k]["tids"][:] num_templates += len(template_bins[k]) # for each bin, get total number of triggers and number of dq flagged triggers bin_total_triggers = np.zeros(num_bins) bin_dq_triggers = np.zeros(num_bins) -logging.info(f'Reading triggers from {args.trigger_file}') -with HFile(args.trigger_file, 'r') as trigf: - for bin_name in template_bins.keys(): +logging.info(f"Reading triggers from {args.trigger_file}") +with HFile(args.trigger_file, "r") as trigf: + for bin_name in template_bins: # bins are named as 'bin{bin_num}' bin_num = int(bin_name[3:]) template_ids = template_bins[bin_name] for template_id in template_ids: # get dq states for all triggers with this template # dq state is either 0 or 1 for each trigger - dq_key = trigf[f'{args.ifo}/dq_state_template'][template_id] - dq_states = trigf[f'{args.ifo}/dq_state'][dq_key] + dq_key = trigf[f"{args.ifo}/dq_state_template"][template_id] + dq_states = trigf[f"{args.ifo}/dq_state"][dq_key] # update trigger counts bin_total_triggers[bin_num] += len(dq_states) bin_dq_triggers[bin_num] += np.sum(dq_states) # write outputs to file -logging.info(f'Writing results to {args.output}') -with HFile(args.output, 'w') as f: +logging.info(f"Writing results to {args.output}") +with HFile(args.output, "w") as f: ifo_group = f.create_group(args.ifo) - ifo_group.create_dataset('observing_livetime', data=livetime) - ifo_group.create_dataset('dq_flag_livetime', data=flagged_time) - bin_group = ifo_group.create_group('bins') - for bin_name in template_bins.keys(): + ifo_group.create_dataset("observing_livetime", data=livetime) + ifo_group.create_dataset("dq_flag_livetime", data=flagged_time) + bin_group = ifo_group.create_group("bins") + for bin_name in template_bins: bin_num = int(bin_name[3:]) bgrp = bin_group.create_group(bin_name) - bgrp.create_dataset('tids', data=template_bins[bin_name]) - bgrp.create_dataset('total_triggers', data=bin_total_triggers[bin_num]) - bgrp.create_dataset('dq_triggers', data=bin_dq_triggers[bin_num]) + bgrp.create_dataset("tids", data=template_bins[bin_name]) + bgrp.create_dataset("total_triggers", data=bin_total_triggers[bin_num]) + bgrp.create_dataset("dq_triggers", data=bin_dq_triggers[bin_num]) bg_triggers = bin_total_triggers[bin_num] - bin_dq_triggers[bin_num] num_trigs = np.array([bg_triggers, bin_dq_triggers[bin_num]]) trig_rates = num_trigs / state_time mean_rate = bin_total_triggers[bin_num] / livetime normalized_rates = trig_rates / mean_rate - bgrp.create_dataset('dq_rates', data=normalized_rates) + bgrp.create_dataset("dq_rates", data=normalized_rates) - f.attrs['dq_thresh'] = args.dq_thresh - f.attrs['dq_channel'] = dq_channel - f.attrs['dq_ok_channel'] = dq_ok_channel - f.attrs['gps_start_time'] = args.gps_start_time - f.attrs['gps_end_time'] = args.gps_end_time - f.attrs['f_lower'] = f_lower - f.attrs['bank_file'] = bank_file - f.attrs['background_bins'] = bin_string + f.attrs["dq_thresh"] = args.dq_thresh + f.attrs["dq_channel"] = dq_channel + f.attrs["dq_ok_channel"] = dq_ok_channel + f.attrs["gps_start_time"] = args.gps_start_time + f.attrs["gps_end_time"] = args.gps_end_time + f.attrs["f_lower"] = f_lower + f.attrs["bank_file"] = bank_file + f.attrs["background_bins"] = bin_string # hash is used to check if different files have compatible settings - settings_to_hash = [args.dq_thresh, dq_channel, dq_ok_channel, - f_lower, bank_file, bin_string] - setting_str = ' '.join([str(s) for s in settings_to_hash]) + settings_to_hash = [ + args.dq_thresh, + dq_channel, + dq_ok_channel, + f_lower, + bank_file, + bin_string, + ] + setting_str = " ".join([str(s) for s in settings_to_hash]) hash_object = hashlib.sha256(setting_str.encode()) - f.attrs['settings_hash'] = hash_object.hexdigest() + f.attrs["settings_hash"] = hash_object.hexdigest() -logging.info('Done') \ No newline at end of file +logging.info("Done") diff --git a/bin/live/pycbc_live_combine_dq_trigger_rates b/bin/live/pycbc_live_combine_dq_trigger_rates index 7bf0c43e360..b70ce449fd7 100644 --- a/bin/live/pycbc_live_combine_dq_trigger_rates +++ b/bin/live/pycbc_live_combine_dq_trigger_rates @@ -14,8 +14,8 @@ """Combine the data-quality adjusted trigger rates from multiple days.""" -import logging import argparse +import logging import numpy as np @@ -24,8 +24,12 @@ from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--daily-dq-files", nargs="+", required=True, - help="Files containing daily dq trigger rates") +parser.add_argument( + "--daily-dq-files", + nargs="+", + required=True, + help="Files containing daily dq trigger rates", +) parser.add_argument("--ifo", required=True) parser.add_argument("--output", required=True) args = parser.parse_args() @@ -38,17 +42,18 @@ daily_files.sort() # need all files to use compatible settings # get settings hash from the last file # we will only use files that have the same hash -with HFile(daily_files[-1], 'r') as last_file: - settings_hash = last_file.attrs['settings_hash'] - bin_str = last_file.attrs['background_bins'] - bank_file = last_file.attrs['bank_file'] - f_lower = last_file.attrs['f_lower'] - dq_thresh = last_file.attrs['dq_thresh'] - dq_channel = last_file.attrs['dq_channel'] - dq_ok_channel = last_file.attrs['dq_ok_channel'] - bin_group = last_file[f'{args.ifo}/bins'] - template_bins = {bin_name: bin_group[bin_name]['tids'][:] - for bin_name in bin_group.keys()} +with HFile(daily_files[-1], "r") as last_file: + settings_hash = last_file.attrs["settings_hash"] + bin_str = last_file.attrs["background_bins"] + bank_file = last_file.attrs["bank_file"] + f_lower = last_file.attrs["f_lower"] + dq_thresh = last_file.attrs["dq_thresh"] + dq_channel = last_file.attrs["dq_channel"] + dq_ok_channel = last_file.attrs["dq_ok_channel"] + bin_group = last_file[f"{args.ifo}/bins"] + template_bins = { + bin_name: bin_group[bin_name]["tids"][:] for bin_name in bin_group.keys() + } num_bins = len(bin_str.split()) total_livetime = 0 @@ -56,20 +61,20 @@ flagged_livetime = 0 total_triggers = np.zeros(num_bins) flagged_triggers = np.zeros(num_bins) for fpath in daily_files: - with HFile(fpath, 'r') as f: - if f.attrs['settings_hash'] != settings_hash: - warning_str = f'File {fpath} has incompatible settings, skipping' + with HFile(fpath, "r") as f: + if f.attrs["settings_hash"] != settings_hash: + warning_str = f"File {fpath} has incompatible settings, skipping" logging.warning(warning_str) continue - total_livetime += f[f'{args.ifo}/observing_livetime'][()] - flagged_livetime += f[f'{args.ifo}/dq_flag_livetime'][()] - bin_group = f[f'{args.ifo}/bins'] + total_livetime += f[f"{args.ifo}/observing_livetime"][()] + flagged_livetime += f[f"{args.ifo}/dq_flag_livetime"][()] + bin_group = f[f"{args.ifo}/bins"] for bin_name in bin_group.keys(): # bins are named as 'bin{bin_num}' bin_num = int(bin_name[3:]) bgrp = bin_group[bin_name] - total_triggers[bin_num] += bgrp['total_triggers'][()] - flagged_triggers[bin_num] += bgrp['dq_triggers'][()] + total_triggers[bin_num] += bgrp["total_triggers"][()] + flagged_triggers[bin_num] += bgrp["dq_triggers"][()] total_trigger_rate = total_triggers / total_livetime flag_trigger_rate = flagged_triggers / flagged_livetime @@ -78,26 +83,26 @@ bg_livetime = total_livetime - flagged_livetime bg_trigger_rate = bg_triggers / bg_livetime # save results -with HFile(args.output, 'w') as f: +with HFile(args.output, "w") as f: ifo_grp = f.create_group(args.ifo) - all_bin_grp = ifo_grp.create_group('bins') + all_bin_grp = ifo_grp.create_group("bins") for bin_name, bin_tids in template_bins.items(): bin_grp = all_bin_grp.create_group(bin_name) - bin_grp['tids'] = bin_tids + bin_grp["tids"] = bin_tids bin_num = int(bin_name[3:]) bin_trig_rates = [bg_trigger_rate[bin_num], flag_trigger_rate[bin_num]] bin_trig_rates /= total_trigger_rate[bin_num] - bin_grp['dq_rates'] = bin_trig_rates - bin_grp['num_triggers'] = total_triggers[bin_num] + bin_grp["dq_rates"] = bin_trig_rates + bin_grp["num_triggers"] = total_triggers[bin_num] - f.attrs['settings_hash'] = settings_hash - f.attrs['stat'] = f'{args.ifo}-dq_stat_info' - f.attrs['total_livetime'] = total_livetime - f.attrs['flagged_livetime'] = flagged_livetime - f.attrs['dq_thresh'] = dq_thresh - f.attrs['dq_channel'] = dq_channel - f.attrs['dq_ok_channel'] = dq_ok_channel - f.attrs['background_bins'] = bin_str - f.attrs['bank_file'] = bank_file - f.attrs['f_lower'] = f_lower + f.attrs["settings_hash"] = settings_hash + f.attrs["stat"] = f"{args.ifo}-dq_stat_info" + f.attrs["total_livetime"] = total_livetime + f.attrs["flagged_livetime"] = flagged_livetime + f.attrs["dq_thresh"] = dq_thresh + f.attrs["dq_channel"] = dq_channel + f.attrs["dq_ok_channel"] = dq_ok_channel + f.attrs["background_bins"] = bin_str + f.attrs["bank_file"] = bank_file + f.attrs["f_lower"] = f_lower diff --git a/bin/live/pycbc_live_combine_single_significance_fits b/bin/live/pycbc_live_combine_single_significance_fits index 21aaa1afbb4..1e2d393b537 100644 --- a/bin/live/pycbc_live_combine_single_significance_fits +++ b/bin/live/pycbc_live_combine_single_significance_fits @@ -12,13 +12,15 @@ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General # Public License for more details. -"""Combine PyCBC Live single-detector trigger fitting parameters from several -different files.""" +""" +Combine PyCBC Live single-detector trigger fitting parameters from several +different files. +""" import argparse import logging -import numpy as np +import numpy as np from igwn_segments import segment, segmentlist import pycbc @@ -26,15 +28,25 @@ from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--trfits-files", nargs="+", required=True, - help="Files containing daily trigger fits") -parser.add_argument("--conservative-percentile", type=int, default=95, - help="What percentile to use for the conservative " - "combined fit. Integer in range 50-99. Default=95") -parser.add_argument("--output", required=True, - help="Output file for combined fit parameters") -parser.add_argument("--ifos", required=True, nargs="+", - help="list of ifos fo collect info for") +parser.add_argument( + "--trfits-files", + nargs="+", + required=True, + help="Files containing daily trigger fits", +) +parser.add_argument( + "--conservative-percentile", + type=int, + default=95, + help="What percentile to use for the conservative " + "combined fit. Integer in range 50-99. Default=95", +) +parser.add_argument( + "--output", required=True, help="Output file for combined fit parameters" +) +parser.add_argument( + "--ifos", required=True, nargs="+", help="list of ifos fo collect info for" +) args = parser.parse_args() @@ -42,11 +54,12 @@ pycbc.init_logging(args.verbose) # Assert some sensible limits on the arguments -if args.conservative_percentile < 50 or \ - args.conservative_percentile > 99: - parser.error("--conservative-percentile must be between 50 and 99, " - "otherwise it is either not a percentile, or not " - "conservative.") +if args.conservative_percentile < 50 or args.conservative_percentile > 99: + parser.error( + "--conservative-percentile must be between 50 and 99, " + "otherwise it is either not a percentile, or not " + "conservative." + ) logging.info("%d input files", len(args.trfits_files)) @@ -58,16 +71,14 @@ logging.info("Determining the most recent configuration parameters") latest_date = None for f in args.trfits_files: - with HFile(f, 'r') as fit_f: - if latest_date is None or fit_f.attrs['fit_end_gps_time'] > latest_date: - latest_date = fit_f.attrs['fit_end_gps_time'] - bl = fit_f['bins_lower'][:] - bu = fit_f['bins_upper'][:] - sngl_rank = fit_f.attrs['sngl_ranking'] - fit_thresh = {ifo: fit_f[ifo].attrs['fit_threshold'] - for ifo in args.ifos} - fit_func = {ifo: fit_f[ifo].attrs['fit_function'] - for ifo in args.ifos} + with HFile(f, "r") as fit_f: + if latest_date is None or fit_f.attrs["fit_end_gps_time"] > latest_date: + latest_date = fit_f.attrs["fit_end_gps_time"] + bl = fit_f["bins_lower"][:] + bu = fit_f["bins_upper"][:] + sngl_rank = fit_f.attrs["sngl_ranking"] + fit_thresh = {ifo: fit_f[ifo].attrs["fit_threshold"] for ifo in args.ifos} + fit_func = {ifo: fit_f[ifo].attrs["fit_function"] for ifo in args.ifos} # Now go back through the fit files and read the actual information. Skip the # files that have fit parameters inconsistent with what we found earlier. @@ -83,29 +94,26 @@ alphas_all = {ifo: [] for ifo in args.ifos} # Check whether the fit files overlap one another seg_all = segmentlist([]) for f in args.trfits_files: - with HFile(f, 'r') as fits_f: + with HFile(f, "r") as fits_f: # Check that the file uses the same setup as file 0, to make sure # all coefficients are comparable - new_fit_func = {ifo: fits_f[ifo].attrs['fit_function'] - for ifo in args.ifos} - new_fit_thresh = {ifo: fits_f[ifo].attrs['fit_threshold'] - for ifo in args.ifos} - same_conf = (fits_f.attrs['sngl_ranking'] == sngl_rank - and new_fit_thresh == fit_thresh - and new_fit_func == fit_func - and fits_f['bins_lower'].size == bl.size - and all(fits_f['bins_lower'][:] == bl) - and all(fits_f['bins_upper'][:] == bu)) + new_fit_func = {ifo: fits_f[ifo].attrs["fit_function"] for ifo in args.ifos} + new_fit_thresh = {ifo: fits_f[ifo].attrs["fit_threshold"] for ifo in args.ifos} + same_conf = ( + fits_f.attrs["sngl_ranking"] == sngl_rank + and new_fit_thresh == fit_thresh + and new_fit_func == fit_func + and fits_f["bins_lower"].size == bl.size + and all(fits_f["bins_lower"][:] == bl) + and all(fits_f["bins_upper"][:] == bu) + ) if not same_conf: - logging.warning( - "Found a change in the fit configuration, skipping %s", - f - ) + logging.warning("Found a change in the fit configuration, skipping %s", f) continue # Get the start/end times of the trigger_fits file - fit_start = fits_f.attrs['fit_start_gps_time'] - fit_end = fits_f.attrs['fit_end_gps_time'] + fit_start = fits_f.attrs["fit_start_gps_time"] + fit_end = fits_f.attrs["fit_end_gps_time"] trigger_file_starts.append(fit_start) trigger_file_ends.append(fit_end) seg = segmentlist([segment(fit_start, fit_end)]) @@ -126,12 +134,12 @@ for f in args.trfits_files: alphas_all[ifo].append(-1 * np.ones_like(bl)) else: ffi = fits_f[ifo] - live_times[ifo].append(ffi.attrs['live_time']) - counts_all[ifo].append(ffi['counts'][:]) - alphas_all[ifo].append(ffi['fit_coeff'][:]) - if any(np.isnan(ffi['fit_coeff'][:])): + live_times[ifo].append(ffi.attrs["live_time"]) + counts_all[ifo].append(ffi["counts"][:]) + alphas_all[ifo].append(ffi["fit_coeff"][:]) + if any(np.isnan(ffi["fit_coeff"][:])): logging.warning("nan in %s, %s", f, ifo) - logging.warning(ffi['fit_coeff'][:]) + logging.warning(ffi["fit_coeff"][:]) # Set up the date array, this is stored as an offset from the first trigger time of # the first file to the last trigger of the file @@ -153,27 +161,26 @@ cons_counts_out = {ifo: np.inf * np.ones(len(alphas_bin[ifo])) for ifo in args.i logging.info("Writing results") -fout = HFile(args.output, 'w') -fout.attrs['conservative_percentile'] = args.conservative_percentile -fout.attrs['ifos'] = args.ifos -fout['bins_edges'] = list(bl) + [bu[-1]] -fout['fits_dates'] = ad + start_time_n +fout = HFile(args.output, "w") +fout.attrs["conservative_percentile"] = args.conservative_percentile +fout.attrs["ifos"] = args.ifos +fout["bins_edges"] = list(bl) + [bu[-1]] +fout["fits_dates"] = ad + start_time_n for ifo in args.ifos: logging.info(ifo) fout_ifo = fout.create_group(ifo) - fout_ifo.attrs['fit_threshold'] = fit_thresh[ifo] - fout_ifo.attrs['fit_function'] = fit_func[ifo] + fout_ifo.attrs["fit_threshold"] = fit_thresh[ifo] + fout_ifo.attrs["fit_function"] = fit_func[ifo] l_times = np.array(live_times[ifo]) total_time = l_times.sum() - fout_ifo.attrs['live_time'] = total_time + fout_ifo.attrs["live_time"] = total_time - fout_ifo['separate_fits/live_times'] = l_times[ad_order] - fout_ifo['separate_fits/start_time'] = trigger_file_starts[ad_order] - fout_ifo['separate_fits/end_time'] = trigger_file_ends[ad_order] + fout_ifo["separate_fits/live_times"] = l_times[ad_order] + fout_ifo["separate_fits/start_time"] = trigger_file_starts[ad_order] + fout_ifo["separate_fits/end_time"] = trigger_file_ends[ad_order] - for counter, (a, c) in enumerate(zip(alphas_bin[ifo], - counts_bin[ifo])): + for counter, (a, c) in enumerate(zip(alphas_bin[ifo], counts_bin[ifo])): # Sort alpha and counts by date a = np.array(a)[ad_order] c = np.array(c)[ad_order] @@ -181,8 +188,8 @@ for ifo in args.ifos: # ignore anything with the 'invalid' salient values valid = c > 0 - fout_ifo[f'separate_fits/bin_{counter:d}/fit_coeff'] = a - fout_ifo[f'separate_fits/bin_{counter:d}/counts'] = c + fout_ifo[f"separate_fits/bin_{counter:d}/fit_coeff"] = a + fout_ifo[f"separate_fits/bin_{counter:d}/counts"] = c if not any(valid): cons_alphas_out[ifo][counter] = np.nan @@ -193,45 +200,38 @@ for ifo in args.ifos: invalphan = c[valid] / a[valid] mean_alpha = c[valid].mean() / invalphan.mean() - cons_count = np.percentile( - c[valid], - args.conservative_percentile - ) - cons_invalphan = np.percentile( - invalphan, - args.conservative_percentile - ) + cons_count = np.percentile(c[valid], args.conservative_percentile) + cons_invalphan = np.percentile(invalphan, args.conservative_percentile) # Conservative alpha estimate using percentile rather than mean - cons_alphas_out[ifo][counter] = cons_count / cons_invalphan + cons_alphas_out[ifo][counter] = cons_count / cons_invalphan alphas_out[ifo][counter] = mean_alpha # To get the count values, we need to convert to rates and back again - r = c[valid]/ l_times[ad_order][valid] + r = c[valid] / l_times[ad_order][valid] cons_rate = np.percentile(r, args.conservative_percentile) cons_counts_out[ifo][counter] = cons_rate * total_time counts_out[ifo][counter] = np.mean(r) * total_time # Output the mean average values - fout_ifo['mean/fit_coeff'] = alphas_out[ifo] - fout_ifo['mean/counts'] = counts_out[ifo] + fout_ifo["mean/fit_coeff"] = alphas_out[ifo] + fout_ifo["mean/counts"] = counts_out[ifo] # Output the conservative values - fout_ifo['conservative/fit_coeff'] = cons_alphas_out[ifo] - fout_ifo['conservative/counts'] = cons_counts_out[ifo] + fout_ifo["conservative/fit_coeff"] = cons_alphas_out[ifo] + fout_ifo["conservative/counts"] = cons_counts_out[ifo] # For the fixed version, we just set this to 1 - fout_ifo['fixed/counts'] = [1] * len(counts_out[ifo]) - fout_ifo['fixed/fit_coeff'] = [0] * len(alphas_out[ifo]) + fout_ifo["fixed/counts"] = [1] * len(counts_out[ifo]) + fout_ifo["fixed/fit_coeff"] = [0] * len(alphas_out[ifo]) # Take some averages for plotting and summary values overall_invalphan = counts_out[ifo] / alphas_out[ifo] - overall_meanalpha = np.nanmean(counts_out[ifo]) \ - / np.nanmean(overall_invalphan) + overall_meanalpha = np.nanmean(counts_out[ifo]) / np.nanmean(overall_invalphan) # Add some useful info to the output file - fout_ifo.attrs['mean_alpha'] = overall_meanalpha - fout_ifo.attrs['total_counts'] = np.nansum(counts_out[ifo]) + fout_ifo.attrs["mean_alpha"] = overall_meanalpha + fout_ifo.attrs["total_counts"] = np.nansum(counts_out[ifo]) fout.close() -logging.info('Done') +logging.info("Done") diff --git a/bin/live/pycbc_live_plot_combined_single_significance_fits b/bin/live/pycbc_live_plot_combined_single_significance_fits index 71be63c3562..e178f8284ee 100644 --- a/bin/live/pycbc_live_plot_combined_single_significance_fits +++ b/bin/live/pycbc_live_plot_combined_single_significance_fits @@ -12,54 +12,69 @@ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General # Public License for more details. -"""Plot the time evolution of fit parameters of PyCBC Live triggers. -""" +"""Plot the time evolution of fit parameters of PyCBC Live triggers.""" import argparse import logging -import numpy as np + import matplotlib -matplotlib.use('agg') +import numpy as np + +matplotlib.use("agg") from matplotlib import pyplot as plt import pycbc from pycbc.io.hdf import HFile from pycbc.time import strip_time_from_gps - parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--combined-fits-file", required=True, - help="File containing information on the combined trigger " - "fits.") +parser.add_argument( + "--combined-fits-file", + required=True, + help="File containing information on the combined trigger fits.", +) default_plot_format = "{ifo}-TRIGGER-FITS-{type}.png" -parser.add_argument("--output-plot-name-format", - default=default_plot_format, - help="Format to save plots, must contain '{ifo}' and " - "'{type}' to indicate ifo and 'alphas' or 'counts' " - " in filename. Default: " + - default_plot_format) -parser.add_argument("--ifos", nargs="+", - help="List of ifos to plot. If not given, use all " - "in the combined file.") -parser.add_argument("--colormap", default="rainbow_r", choices=plt.colormaps(), - help="Colormap to use for choosing the colours of the " - "duration bin lines. Default=rainbow_r") -parser.add_argument("--log-colormap", action='store_true', - help="Use log spacing for choosing colormap values " - "based on duration bins.") parser.add_argument( - '--log-counts', - action='store_true', - help="Plot the trigger rate above threshold on a log scale." + "--output-plot-name-format", + default=default_plot_format, + help="Format to save plots, must contain '{ifo}' and " + "'{type}' to indicate ifo and 'alphas' or 'counts' " + " in filename. Default: " + default_plot_format, +) +parser.add_argument( + "--ifos", + nargs="+", + help="List of ifos to plot. If not given, use all in the combined file.", +) +parser.add_argument( + "--colormap", + default="rainbow_r", + choices=plt.colormaps(), + help="Colormap to use for choosing the colours of the " + "duration bin lines. Default=rainbow_r", +) +parser.add_argument( + "--log-colormap", + action="store_true", + help="Use log spacing for choosing colormap values based on duration bins.", +) +parser.add_argument( + "--log-counts", + action="store_true", + help="Plot the trigger rate above threshold on a log scale.", ) args = parser.parse_args() -if '{ifo}' not in args.output_plot_name_format or \ - '{type}' not in args.output_plot_name_format: - parser.error("--output-plot-name-format must contain '{ifo}' and " - "'{type}' to indicate ifo and 'alphas' or 'counts' " - " in filename.") +if ( + "{ifo}" not in args.output_plot_name_format + or "{type}" not in args.output_plot_name_format +): + parser.error( + "--output-plot-name-format must contain '{ifo}' and " + "'{type}' to indicate ifo and 'alphas' or 'counts' " + " in filename." + ) pycbc.init_logging(args.verbose) @@ -76,28 +91,30 @@ separate_ends = {} separate_times = {} logging.info("Loading Data") -with HFile(args.combined_fits_file, 'r') as cff: - ifos = args.ifos or cff.attrs['ifos'] - bins_edges = cff['bins_edges'][:] - conservative_percentile = cff.attrs['conservative_percentile'] +with HFile(args.combined_fits_file, "r") as cff: + ifos = args.ifos or cff.attrs["ifos"] + bins_edges = cff["bins_edges"][:] + conservative_percentile = cff.attrs["conservative_percentile"] n_bins = len(bins_edges) - 1 for ifo in ifos: logging.info(ifo) - live_total[ifo] = cff[ifo].attrs['live_time'] - mean_count[ifo] = cff[ifo]['mean']['counts'][:] - mean_alpha[ifo] = cff[ifo]['mean']['fit_coeff'][:] - cons_count[ifo] = cff[ifo]['conservative']['counts'][:] - cons_alpha[ifo] = cff[ifo]['conservative']['fit_coeff'][:] - - separate_starts[ifo] = cff[ifo]['separate_fits']['start_time'][:] - separate_ends[ifo] = cff[ifo]['separate_fits']['end_time'][:] - separate_times[ifo] = cff[ifo]['separate_fits']['live_times'][:] - - separate_data = cff[ifo]['separate_fits'] - separate_alphas[ifo] = np.array([separate_data[f'bin_{i}']['fit_coeff'][:] - for i in range(n_bins)]) - separate_counts[ifo] = np.array([separate_data[f'bin_{i}']['counts'][:] - for i in range(n_bins)]) + live_total[ifo] = cff[ifo].attrs["live_time"] + mean_count[ifo] = cff[ifo]["mean"]["counts"][:] + mean_alpha[ifo] = cff[ifo]["mean"]["fit_coeff"][:] + cons_count[ifo] = cff[ifo]["conservative"]["counts"][:] + cons_alpha[ifo] = cff[ifo]["conservative"]["fit_coeff"][:] + + separate_starts[ifo] = cff[ifo]["separate_fits"]["start_time"][:] + separate_ends[ifo] = cff[ifo]["separate_fits"]["end_time"][:] + separate_times[ifo] = cff[ifo]["separate_fits"]["live_times"][:] + + separate_data = cff[ifo]["separate_fits"] + separate_alphas[ifo] = np.array( + [separate_data[f"bin_{i}"]["fit_coeff"][:] for i in range(n_bins)] + ) + separate_counts[ifo] = np.array( + [separate_data[f"bin_{i}"]["counts"][:] for i in range(n_bins)] + ) bin_starts = bins_edges[:-1] bin_ends = bins_edges[1:] @@ -147,7 +164,7 @@ for ifo in ifos: continue l_times = separate_times[ifo] - with np.errstate(divide='ignore', invalid='ignore'): + with np.errstate(divide="ignore", invalid="ignore"): rate = counts / l_times ma = mean_alpha[ifo][i] @@ -158,24 +175,39 @@ for ifo in ifos: bin_prop = i / len(bin_starts) bin_colour = plt.get_cmap(args.colormap)(bin_prop) bin_label = f"duration {bl:.2f}-{bu:.2f}" - alpha_lines += ax_alpha.plot(separate_starts[ifo], alphas, c=bin_colour, - label=bin_label, marker='.', - markersize=10) - alpha_lines.append(ax_alpha.axhline(ma, - label="total fit = %.2f" % ma, - c=bin_colour, linestyle='--',)) + alpha_lines += ax_alpha.plot( + separate_starts[ifo], + alphas, + c=bin_colour, + label=bin_label, + marker=".", + markersize=10, + ) + alpha_lines.append( + ax_alpha.axhline( + ma, + label="total fit = %.2f" % ma, + c=bin_colour, + linestyle="--", + ) + ) alpha_lab = f"{conservative_percentile:d}th %ile = {ca:.2f}" - alpha_lines.append(ax_alpha.axhline(ca, - c=bin_colour, linestyle=':', - label=alpha_lab)) + alpha_lines.append( + ax_alpha.axhline(ca, c=bin_colour, linestyle=":", label=alpha_lab) + ) # Invalid counts inv_counts = rate <= 0 rate[inv_counts] = None - count_lines += ax_count.plot(separate_starts[ifo], rate, c=bin_colour, - label=bin_label, marker='.', - markersize=10) + count_lines += ax_count.plot( + separate_starts[ifo], + rate, + c=bin_colour, + label=bin_label, + marker=".", + markersize=10, + ) if mr < 1e-3: mlab = f"mean = {mr:.3e}" @@ -184,42 +216,47 @@ for ifo in ifos: mlab = f"mean = {mr:.3f}" clab = f"{conservative_percentile:d}th %ile = {cr:.3f}" - count_lines.append(ax_count.axhline(mr, - c=bin_colour, linestyle='--', - label=mlab)) - count_lines.append(ax_count.axhline(cr, - c=bin_colour, linestyle=':', - label=clab)) + count_lines.append( + ax_count.axhline(mr, c=bin_colour, linestyle="--", label=mlab) + ) + count_lines.append( + ax_count.axhline(cr, c=bin_colour, linestyle=":", label=clab) + ) alpha_labels = [l.get_label() for l in alpha_lines] - ax_alpha.legend(alpha_lines, alpha_labels, loc='lower center', - ncol=5, bbox_to_anchor=(0.5, 1.01)) - ax_alpha.set_ylabel('Fit coefficient') + ax_alpha.legend( + alpha_lines, + alpha_labels, + loc="lower center", + ncol=5, + bbox_to_anchor=(0.5, 1.01), + ) + ax_alpha.set_ylabel("Fit coefficient") count_labels = [l.get_label() for l in count_lines] if args.log_counts: ax_count.semilogy() - ax_count.legend(count_lines, count_labels, loc='lower center', - ncol=5, bbox_to_anchor=(0.5, 1.01)) - ax_count.set_ylabel('Rate of triggers above fit threshold [s$^{-1}$]') + ax_count.legend( + count_lines, + count_labels, + loc="lower center", + ncol=5, + bbox_to_anchor=(0.5, 1.01), + ) + ax_count.set_ylabel("Rate of triggers above fit threshold [s$^{-1}$]") for ax in [ax_count, ax_alpha]: ax.set_xticks(xtix) ax.set_xticklabels(xtix_labels, rotation=90) # Add 1/4 day padding either side of the lines ax.set_xlim( - min(separate_starts[ifo]) - 21600, - max(separate_starts[ifo]) + 21600 + min(separate_starts[ifo]) - 21600, max(separate_starts[ifo]) + 21600 ) ax.grid(zorder=-30) fig_count.tight_layout() - fig_count.savefig( - args.output_plot_name_format.format(ifo=ifo, type='counts') - ) + fig_count.savefig(args.output_plot_name_format.format(ifo=ifo, type="counts")) fig_alpha.tight_layout() - fig_alpha.savefig( - args.output_plot_name_format.format(ifo=ifo, type='fit_coeffs') - ) + fig_alpha.savefig(args.output_plot_name_format.format(ifo=ifo, type="fit_coeffs")) logging.info("Done") diff --git a/bin/live/pycbc_live_plot_single_significance_fits b/bin/live/pycbc_live_plot_single_significance_fits index e79e0a9b643..b54aa2bd5ec 100644 --- a/bin/live/pycbc_live_plot_single_significance_fits +++ b/bin/live/pycbc_live_plot_single_significance_fits @@ -12,15 +12,18 @@ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General # Public License for more details. -"""Plot histograms of PyCBC Live triggers split over various parameters, and +""" +Plot histograms of PyCBC Live triggers split over various parameters, and the corresponding fits. """ import argparse import logging -import numpy as np + import matplotlib -matplotlib.use('agg') +import numpy as np + +matplotlib.use("agg") from matplotlib import pyplot as plt import pycbc @@ -31,54 +34,67 @@ from pycbc.time import gps_to_utc_str parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--trigger-fits-file", required=True, - help="Trigger fits file to plot") +parser.add_argument( + "--trigger-fits-file", required=True, help="Trigger fits file to plot" +) default_plot_format = "{ifo}-TRIGGER-FITS.png" -parser.add_argument("--output-plot-name-format", - default=default_plot_format, - help="Format to save plots, must contain '{ifo}' to " - "indicate ifo in filename. Default: " + - default_plot_format) -parser.add_argument("--colormap", default="rainbow_r", choices=plt.colormaps(), - help="Colormap to use for choosing the colours of the " - "duration bin lines. Default=rainbow_r") -parser.add_argument("--log-colormap", action='store_true', - help="Use log spacing for choosing colormap values " - "based on duration bins.") -parser.add_argument("--x-lim-lower", type=float, - help="Add a lower limit to the x-axis of the plot") -parser.add_argument("--x-lim-upper", type=float, - help="Add an upper limit to the x-axis of the plot") -parser.add_argument("--y-lim-lower", type=float, - help="Add a lower limit to the y-axis of the plot") -parser.add_argument("--y-lim-upper", type=float, - help="Add an upper limit to the y-axis of the plot") - -#Add some input sanitisation +parser.add_argument( + "--output-plot-name-format", + default=default_plot_format, + help="Format to save plots, must contain '{ifo}' to " + "indicate ifo in filename. Default: " + default_plot_format, +) +parser.add_argument( + "--colormap", + default="rainbow_r", + choices=plt.colormaps(), + help="Colormap to use for choosing the colours of the " + "duration bin lines. Default=rainbow_r", +) +parser.add_argument( + "--log-colormap", + action="store_true", + help="Use log spacing for choosing colormap values based on duration bins.", +) +parser.add_argument( + "--x-lim-lower", type=float, help="Add a lower limit to the x-axis of the plot" +) +parser.add_argument( + "--x-lim-upper", type=float, help="Add an upper limit to the x-axis of the plot" +) +parser.add_argument( + "--y-lim-lower", type=float, help="Add a lower limit to the y-axis of the plot" +) +parser.add_argument( + "--y-lim-upper", type=float, help="Add an upper limit to the y-axis of the plot" +) + +# Add some input sanitisation args = parser.parse_args() -if '{ifo}' not in args.output_plot_name_format: - parser.error("--output-plot-name-format must contain '{ifo}' " - "to indicate ifo in filename.") +if "{ifo}" not in args.output_plot_name_format: + parser.error( + "--output-plot-name-format must contain '{ifo}' to indicate ifo in filename." + ) pycbc.init_logging(args.verbose) logging.info("Getting trigger fits file information") -with HFile(args.trigger_fits_file, 'r') as trfit_f: +with HFile(args.trigger_fits_file, "r") as trfit_f: # Get the ifos to plot from the file # Check that all the ifos we want to plot are in the file - all_ifos = trfit_f.attrs['ifos'].split(',') - ifos = [k for k in trfit_f.keys() if not k.startswith('bins')] + all_ifos = trfit_f.attrs["ifos"].split(",") + ifos = [k for k in trfit_f.keys() if not k.startswith("bins")] # Grab some info from the attributes - if trfit_f.attrs['ranking_statistic'] in ['quadsum','single_ranking_only']: - x_label = trfit_f.attrs['sngl_ranking'] + if trfit_f.attrs["ranking_statistic"] in ["quadsum", "single_ranking_only"]: + x_label = trfit_f.attrs["sngl_ranking"] else: x_label = "Ranking statistic" - fit_threshold = {ifo: trfit_f[ifo].attrs['fit_threshold'] for ifo in ifos} - fit_function = {ifo: trfit_f[ifo].attrs['fit_function'] for ifo in ifos} - start_time = trfit_f.attrs['fit_start_gps_time'] - end_time = trfit_f.attrs['fit_end_gps_time'] + fit_threshold = {ifo: trfit_f[ifo].attrs["fit_threshold"] for ifo in ifos} + fit_function = {ifo: trfit_f[ifo].attrs["fit_function"] for ifo in ifos} + start_time = trfit_f.attrs["fit_start_gps_time"] + end_time = trfit_f.attrs["fit_end_gps_time"] start_str = gps_to_utc_str(start_time) end_str = gps_to_utc_str(end_time) @@ -87,15 +103,15 @@ with HFile(args.trigger_fits_file, 'r') as trfit_f: stats = {ifo: {} for ifo in ifos} durations = {ifo: {} for ifo in ifos} for ifo in ifos: - if 'triggers' not in trfit_f[ifo]: + if "triggers" not in trfit_f[ifo]: continue - stats[ifo] = trfit_f[ifo]['triggers']['stat'][:] - durations[ifo] = trfit_f[ifo]['triggers']['template_duration'][:] - live_time = {ifo: trfit_f[ifo].attrs['live_time'] for ifo in ifos} - alphas = {ifo: trfit_f[ifo]['fit_coeff'][:] for ifo in ifos} - counts = {ifo: trfit_f[ifo]['counts'][:] for ifo in ifos} - bu = trfit_f['bins_upper'][:] - bl = trfit_f['bins_lower'][:] + stats[ifo] = trfit_f[ifo]["triggers"]["stat"][:] + durations[ifo] = trfit_f[ifo]["triggers"]["template_duration"][:] + live_time = {ifo: trfit_f[ifo].attrs["live_time"] for ifo in ifos} + alphas = {ifo: trfit_f[ifo]["fit_coeff"][:] for ifo in ifos} + counts = {ifo: trfit_f[ifo]["counts"][:] for ifo in ifos} + bu = trfit_f["bins_upper"][:] + bl = trfit_f["bins_lower"][:] duration_bin_edges = list(bl) + [bu[-1]] tbins = IrregularBins(duration_bin_edges) @@ -112,13 +128,15 @@ for ifo in all_ifos: if ifo not in ifos or not len(stats[ifo]): # Plot a blank plot with a message to show it worked, but there # weren't any triggers - plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, - right=False) + plt.tick_params( + labelcolor="none", top=False, bottom=False, left=False, right=False + ) ax.text( - 0.5, 0.5, + 0.5, + 0.5, "No triggers above threshold in this detector", - horizontalalignment='center', - verticalalignment='center', + horizontalalignment="center", + verticalalignment="center", ) logging.info(f"Saving {oput_plot}") # Save initial logging level @@ -150,58 +168,52 @@ for ifo in all_ifos: # Skip if there are no triggers in this bin in this IFO if not any(inbin) or alphas[ifo][bin_num] == -1: - ax.plot( - [], - [], - linewidth=2, - color=bin_colour, - label=binlabel, - alpha=0.6 - ) - ax.plot( - [], - [], - "--", - color=bin_colour, - label="No triggers" - ) + ax.plot([], [], linewidth=2, color=bin_colour, label=binlabel, alpha=0.6) + ax.plot([], [], "--", color=bin_colour, label="No triggers") continue binned_sngl_stats = stats[ifo][event_bin == bin_num] # Histogram the triggers - histcounts, edges = np.histogram(binned_sngl_stats, - bins=plotbins) + histcounts, edges = np.histogram(binned_sngl_stats, bins=plotbins) cum_rate = histcounts[::-1].cumsum()[::-1] / live_time[ifo] max_rate = max(max_rate, cum_rate[0]) ecf = eval_cum_fit( - fit_function[ifo], - plotbins, - alphas[ifo][bin_num], - fit_threshold[ifo] + fit_function[ifo], plotbins, alphas[ifo][bin_num], fit_threshold[ifo] ) cum_fit = counts[ifo][bin_num] / live_time[ifo] * ecf - ax.plot(edges[:-1], cum_rate, linewidth=2, - color=bin_colour, label=binlabel, alpha=0.6) - ax.plot(plotbins, cum_fit, "--", color=bin_colour, - label=r"$\alpha = $%.2f" % alphas[ifo][bin_num]) + ax.plot( + edges[:-1], + cum_rate, + linewidth=2, + color=bin_colour, + label=binlabel, + alpha=0.6, + ) + ax.plot( + plotbins, + cum_fit, + "--", + color=bin_colour, + label=r"$\alpha = $%.2f" % alphas[ifo][bin_num], + ) ax.semilogy() ax.grid() ax.set_xlim( fit_threshold[ifo] if args.x_lim_lower is None else args.x_lim_lower, - plotmax if args.x_lim_upper is None else args.x_lim_upper + plotmax if args.x_lim_upper is None else args.x_lim_upper, ) ax.set_ylim( 0.5 / live_time[ifo] if args.y_lim_lower is None else args.y_lim_lower, - 1.5 * max_rate if args.y_lim_upper is None else args.y_lim_upper + 1.5 * max_rate if args.y_lim_upper is None else args.y_lim_upper, ) ax.set_xlabel(x_label) ax.set_ylabel("Number of louder triggers per live time") title = f"{ifo} singles significance fits from\n{start_str} to {end_str}" ax.set_title(title) - ax.legend(loc='center left', bbox_to_anchor=(1.01, 0.5)) + ax.legend(loc="center left", bbox_to_anchor=(1.01, 0.5)) logging.info(f"Saving {oput_plot}") # Save initial logging level logger.setLevel(logging.WARNING) diff --git a/bin/live/pycbc_live_single_significance_fits b/bin/live/pycbc_live_single_significance_fits index 6d80bdda911..a1dd36cf9bf 100644 --- a/bin/live/pycbc_live_single_significance_fits +++ b/bin/live/pycbc_live_single_significance_fits @@ -12,52 +12,69 @@ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General # Public License for more details. -"""Fit a background model to single-detector triggers from PyCBC Live. +""" +Fit a background model to single-detector triggers from PyCBC Live. -See https://arxiv.org/abs/2008.07494 for a description of the method.""" +See https://arxiv.org/abs/2008.07494 for a description of the method. +""" -import sys import argparse import logging +import sys + import numpy as np import pycbc from pycbc.bin_utils import IrregularBins -from pycbc.events import cuts, trigger_fits as trstats, stat +from pycbc.events import cuts, stat +from pycbc.events import trigger_fits as trstats +from pycbc.events.coinc import cluster_over_time from pycbc.io import DictArray, HFile from pycbc.io import live as liveio from pycbc.live import significance_fits as sngls_io -from pycbc.events.coinc import cluster_over_time from pycbc.types import MultiDetOptionAction parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--ifos", nargs="+", required=True, - help="Which ifo(s) are we fitting the triggers for? " - "Required") +parser.add_argument( + "--ifos", + nargs="+", + required=True, + help="Which ifo(s) are we fitting the triggers for? Required", +) liveio.add_live_trigger_selection_options(parser) sngls_io.add_live_significance_trigger_pruning_options(parser) sngls_io.add_live_significance_duration_bin_options(parser) -parser.add_argument("--fit-function", default="exponential", - action=MultiDetOptionAction, - choices=["exponential", "rayleigh", "power"], - help="Functional form for the maximum likelihood fit. " - "Choose from exponential, rayleigh or power. " - "Default: exponential") -parser.add_argument("--fit-threshold", type=float, default=5, - action=MultiDetOptionAction, - help="Lower threshold used in fitting the triggers." - "Default 5. Can be supplied as a single value, " - "or as a set of IFO:value pairs, e.g. H1:0:, L1:-1") -parser.add_argument("--cluster", action='store_true', - help="Only use maximum of the --sngl-ranking value " - "from each file.") -parser.add_argument("--output", required=True, - help="File in which to save the output trigger fit " - "parameters.") +parser.add_argument( + "--fit-function", + default="exponential", + action=MultiDetOptionAction, + choices=["exponential", "rayleigh", "power"], + help="Functional form for the maximum likelihood fit. " + "Choose from exponential, rayleigh or power. " + "Default: exponential", +) +parser.add_argument( + "--fit-threshold", + type=float, + default=5, + action=MultiDetOptionAction, + help="Lower threshold used in fitting the triggers." + "Default 5. Can be supplied as a single value, " + "or as a set of IFO:value pairs, e.g. H1:0:, L1:-1", +) +parser.add_argument( + "--cluster", + action="store_true", + help="Only use maximum of the --sngl-ranking value from each file.", +) +parser.add_argument( + "--output", + required=True, + help="File in which to save the output trigger fit parameters.", +) stat.insert_statistic_option_group( - parser, - default_ranking_statistic='single_ranking_only' + parser, default_ranking_statistic="single_ranking_only" ) cuts.insert_cuts_option_group(parser) @@ -71,8 +88,7 @@ sngls_io.verify_live_significance_duration_bin_options(args, parser) duration_bin_edges = sngls_io.duration_bins_from_cli(args) logging.info( - "Duration bin edges: %s", - ', '.join([f'{e:.3e}' for e in duration_bin_edges]) + "Duration bin edges: %s", ", ".join([f"{e:.3e}" for e in duration_bin_edges]) ) logging.info("Finding files") @@ -103,7 +119,7 @@ tbins = IrregularBins(duration_bin_edges) # Live time is not immediately obvious - get an approximation with 8 second # granularity by adding 8 seconds per 'valid' file -live_time = {ifo: 0 for ifo in args.ifos} +live_time = dict.fromkeys(args.ifos, 0) logging.info("Getting events which meet criteria") @@ -122,31 +138,35 @@ for counter, filename in enumerate(files): # In case of no triggers for an extended period logging.info("%s: No data", ifo) else: - logging.info("%s: %d triggers in %.0f s", ifo, - events[ifo].data['snr'].size, live_time[ifo]) + logging.info( + "%s: %d triggers in %.0f s", + ifo, + events[ifo].data["snr"].size, + live_time[ifo], + ) # If there is an IOerror with the file, don't fail, just carry on try: - HFile(filename, 'r') - except IOError: - logging.warning('IOError with file %s', f) + HFile(filename, "r") + except OSError: + logging.warning("IOError with file %s", f) continue # Triggers for this file triggers = {} - with HFile(filename, 'r') as fin: + with HFile(filename, "r") as fin: # Open the file: does it have the ifo group and snr dataset? for ifo in args.ifos: - if not (ifo in fin and 'snr' in fin[ifo]): + if not (ifo in fin and "snr" in fin[ifo]): continue # Eventual FIX ME: live output files should (soon) have the live time # added, but for now, extract from the filename # Format of the filename is to have the live time as a dash, # followed by '.hdf' at the end of the filename - lt = int(filename.split('-')[-1][:-4]) + lt = int(filename.split("-")[-1][:-4]) live_time[ifo] += lt - n_triggers = fin[ifo]['snr'].size + n_triggers = fin[ifo]["snr"].size # Skip if there are no triggers if not n_triggers: continue @@ -155,34 +175,34 @@ for counter, filename in enumerate(files): # Get all datasets with the same size as the trigger SNRs, # except for edge cases where the number of loudest, gates etc. # happens to be the same as the trigger count - triggers[ifo] = {k: fin[ifo][k][:] for k in fin[ifo].keys() - if k not in ('loudest', 'stat', 'gates', 'psd') - and fin[ifo][k].size == n_triggers} + triggers[ifo] = { + k: fin[ifo][k][:] + for k in fin[ifo].keys() + if k not in ("loudest", "stat", "gates", "psd") + and fin[ifo][k].size == n_triggers + } # The stored chisq is actually reduced chisq, so hack the # chisq_dof dataset to use the standard conversions. # chisq_dof of 1.5 gives the right number (2 * 1.5 - 2 = 1) - triggers[ifo]['chisq_dof'] = \ - 1.5 * np.ones_like(triggers[ifo]['snr']) - + triggers[ifo]["chisq_dof"] = 1.5 * np.ones_like(triggers[ifo]["snr"]) for ifo, trigs_ifo in triggers.items(): - # Apply the cuts to triggers keep_idx = cuts.apply_trigger_cuts(trigs_ifo, trigger_cut_dict) # triggers contains the datasets that we want to use for # the template cuts, so here it can be used as the template bank - keep_idx = cuts.apply_template_cuts(trigs_ifo, template_cut_dict, - template_ids=keep_idx) + keep_idx = cuts.apply_template_cuts( + trigs_ifo, template_cut_dict, template_ids=keep_idx + ) # Skip if no triggers survive the cuts if not keep_idx.size: continue # Apply the cuts - triggers_cut = {k: trigs_ifo[k][keep_idx] - for k in trigs_ifo.keys()} + triggers_cut = {k: trigs_ifo[k][keep_idx] for k in trigs_ifo.keys()} # Calculate the sngl_ranking values sds = rank_method[ifo].single(triggers_cut) @@ -190,7 +210,7 @@ for counter, filename in enumerate(files): (ifo, sds), ) - triggers_cut['stat'] = sngls_value + triggers_cut["stat"] = sngls_value triggers_da = DictArray(data=triggers_cut) @@ -212,10 +232,9 @@ for ifo in args.ifos: if ifo not in events: logging.info("%s: No data", ifo) else: - logging.info("%s: %d in %.2fs", - ifo, len(events[ifo]), live_time[ifo]) + logging.info("%s: %d in %.2fs", ifo, len(events[ifo]), live_time[ifo]) -logging.info('Sorting events into template duration bins') +logging.info("Sorting events into template duration bins") # Set up bins and prune loud events in each bin n_bins = duration_bin_edges.size - 1 @@ -226,20 +245,22 @@ times_to_prune = {ifo: [] for ifo in args.ifos} for ifo in events: # Sort the events into their bins - event_bins[ifo] = np.array([tbins[d] - for d in events[ifo].data['template_duration']]) + event_bins[ifo] = np.array( + [tbins[d] for d in events[ifo].data["template_duration"]] + ) if args.prune_loudest: for bin_num in range(n_bins): inbin = event_bins[ifo] == bin_num - binned_events = events[ifo].data['stat'][inbin] - binned_event_times = events[ifo].data['end_time'][inbin] + binned_events = events[ifo].data["stat"][inbin] + binned_event_times = events[ifo].data["end_time"][inbin] # Cluster triggers in time with the pruning window to ensure # that clusters are independent - cidx = cluster_over_time(binned_events, binned_event_times, - args.prune_window) + cidx = cluster_over_time( + binned_events, binned_event_times, args.prune_window + ) # Find clusters at/above the statistic threshold above_stat_min = binned_events[cidx] >= args.prune_stat_threshold @@ -252,18 +273,22 @@ for ifo in events: continue # Find the loudest of the triggers in this bin - argloudest = np.argsort(binned_events[cidx])[-args.prune_loudest:] + argloudest = np.argsort(binned_events[cidx])[-args.prune_loudest :] times_to_prune[ifo] += list(binned_event_times[cidx][argloudest]) n_pruned = {ifo: [] for ifo in args.ifos} pruned_trigger_times = {} if args.prune_loudest: - logging.info("Pruning triggers %.2fs either side of the loudest %d " - "triggers in each bin if %s > %.2f", args.prune_window, - args.prune_loudest, args.sngl_ranking, - args.prune_stat_threshold) + logging.info( + "Pruning triggers %.2fs either side of the loudest %d " + "triggers in each bin if %s > %.2f", + args.prune_window, + args.prune_loudest, + args.sngl_ranking, + args.prune_stat_threshold, + ) for ifo in events: - times = events[ifo].data['end_time'][:] + times = events[ifo].data["end_time"][:] outwith_window = np.ones_like(times, dtype=bool) for t in times_to_prune[ifo]: outwith_window &= abs(times - t) > args.prune_window @@ -284,13 +309,13 @@ if args.prune_loudest: pruned_inbin = pruned_trigger_bins == bin_num n_pruned_thisbin = np.count_nonzero(pruned_inbin) n_pruned[ifo].append(n_pruned_thisbin) - logging.info("Pruned %d triggers from %s bin %d", - n_pruned_thisbin, ifo, bin_num) + logging.info( + "Pruned %d triggers from %s bin %d", n_pruned_thisbin, ifo, bin_num + ) # Do the fitting for each bin for ifo in events: for bin_num in range(n_bins): - inbin = event_bins[ifo] == bin_num if not np.count_nonzero(inbin): @@ -299,54 +324,51 @@ for ifo in events: alphas[ifo][bin_num] = -1 continue - stat_inbin = events[ifo].data['stat'][inbin] - counts[ifo][bin_num] = \ - np.count_nonzero(stat_inbin > args.fit_threshold[ifo]) + stat_inbin = events[ifo].data["stat"][inbin] + counts[ifo][bin_num] = np.count_nonzero(stat_inbin > args.fit_threshold[ifo]) alphas[ifo][bin_num], _ = trstats.fit_above_thresh( - args.fit_function[ifo], - stat_inbin, - args.fit_threshold[ifo] + args.fit_function[ifo], stat_inbin, args.fit_threshold[ifo] ) logging.info("Writing results") -with HFile(args.output, 'w') as fout: +with HFile(args.output, "w") as fout: for ifo in args.ifos: fout_ifo = fout.create_group(ifo) - fout_ifo.attrs['fit_function'] = args.fit_function[ifo] - fout_ifo.attrs['fit_threshold'] = args.fit_threshold[ifo] + fout_ifo.attrs["fit_function"] = args.fit_function[ifo] + fout_ifo.attrs["fit_threshold"] = args.fit_threshold[ifo] if ifo not in events: # There were no triggers, but we should still produce some # information - fout_ifo['fit_coeff'] = -1 * np.zeros(n_bins) - fout_ifo['counts'] = np.zeros(n_bins) - fout_ifo.attrs['live_time'] = live_time[ifo] - fout_ifo.attrs['pruned_times'] = [] - fout_ifo.attrs['n_pruned'] = 0 + fout_ifo["fit_coeff"] = -1 * np.zeros(n_bins) + fout_ifo["counts"] = np.zeros(n_bins) + fout_ifo.attrs["live_time"] = live_time[ifo] + fout_ifo.attrs["pruned_times"] = [] + fout_ifo.attrs["n_pruned"] = 0 continue # Save the triggers we have used for the fits - fout_ifo_trigs = fout_ifo.create_group('triggers') + fout_ifo_trigs = fout_ifo.create_group("triggers") for key in events[ifo].data: fout_ifo_trigs[key] = events[ifo].data[key] if ifo in pruned_trigger_times: - fout_ifo['pruned_trigger_times'] = pruned_trigger_times[ifo] - - fout_ifo['fit_coeff'] = alphas[ifo] - fout_ifo['counts'] = counts[ifo] - fout_ifo.attrs['live_time'] = live_time[ifo] - fout_ifo.attrs['pruned_times'] = times_to_prune[ifo] - fout_ifo.attrs['n_pruned'] = n_pruned[ifo] - - fout['bins_upper'] = tbins.upper() - fout['bins_lower'] = tbins.lower() - - fout.attrs['ifos'] = ','.join(args.ifos) - fout.attrs['fit_start_gps_time'] = args.gps_start_time - fout.attrs['fit_end_gps_time'] = args.gps_end_time - fout.attrs['input'] = sys.argv - fout.attrs['cuts'] = args.template_cuts + args.trigger_cuts - fout.attrs['sngl_ranking'] = args.sngl_ranking - fout.attrs['ranking_statistic'] = args.ranking_statistic + fout_ifo["pruned_trigger_times"] = pruned_trigger_times[ifo] + + fout_ifo["fit_coeff"] = alphas[ifo] + fout_ifo["counts"] = counts[ifo] + fout_ifo.attrs["live_time"] = live_time[ifo] + fout_ifo.attrs["pruned_times"] = times_to_prune[ifo] + fout_ifo.attrs["n_pruned"] = n_pruned[ifo] + + fout["bins_upper"] = tbins.upper() + fout["bins_lower"] = tbins.lower() + + fout.attrs["ifos"] = ",".join(args.ifos) + fout.attrs["fit_start_gps_time"] = args.gps_start_time + fout.attrs["fit_end_gps_time"] = args.gps_end_time + fout.attrs["input"] = sys.argv + fout.attrs["cuts"] = args.template_cuts + args.trigger_cuts + fout.attrs["sngl_ranking"] = args.sngl_ranking + fout.attrs["ranking_statistic"] = args.ranking_statistic logging.info("Done") diff --git a/bin/live/pycbc_live_supervise_collated_trigger_fits b/bin/live/pycbc_live_supervise_collated_trigger_fits index 43aa1b24b64..da09c660b81 100755 --- a/bin/live/pycbc_live_supervise_collated_trigger_fits +++ b/bin/live/pycbc_live_supervise_collated_trigger_fits @@ -1,47 +1,46 @@ #!/usr/bin/env python -"""Supervise the periodic re-fitting of PyCBC Live single-detector triggers, +""" +Supervise the periodic re-fitting of PyCBC Live single-detector triggers, and the associated plots. """ -import re -import logging import argparse -import shutil +import logging import os +import re +import shutil +from datetime import datetime, timedelta import numpy as np -from datetime import datetime, timedelta - import pycbc -from pycbc.live import supervision as sv -from pycbc.types.config import InterpolatingConfigParser as icp from pycbc.io import HFile +from pycbc.live import supervision as sv from pycbc.time import gps_to_utc_datetime, utc_datetime_to_gps +from pycbc.types.config import InterpolatingConfigParser as icp def read_options(args): """ - read the options into a dictionary + Read the options into a dictionary """ logging.info("Reading config file") cp = icp(configFiles=[args.config_file]) config_opts = { - section: {k: v for k, v in cp[section].items()} - for section in cp.sections() + section: {k: v for k, v in cp[section].items()} for section in cp.sections() } - del config_opts['environment'] + del config_opts["environment"] return config_opts def get_true_date(replay_day_dt, controls): - if 'replay-start-time' not in controls: + if "replay-start-time" not in controls: return replay_day_dt - replay_start_time = int(controls['replay-start-time']) - true_start_time = int(controls['true-start-time']) - replay_duration = int(controls['replay-duration']) + replay_start_time = int(controls["replay-start-time"]) + true_start_time = int(controls["true-start-time"]) + replay_duration = int(controls["replay-duration"]) dt_replay_start = gps_to_utc_datetime(replay_start_time) td = (replay_day_dt - dt_replay_start).total_seconds() @@ -58,19 +57,14 @@ def get_true_date(replay_day_dt, controls): def trigger_collation( - day_dt, - day_str, - collation_control_options, - collation_options, - output_dir, - controls + day_dt, day_str, collation_control_options, collation_options, output_dir, controls ): """ Perform the trigger collation as specified """ logging.info("Performing trigger collation") collate_args = [ - 'pycbc_live_collate_triggers', + "pycbc_live_collate_triggers", ] collate_args += sv.dict_to_args(collation_options) gps_start = int(utc_datetime_to_gps(day_dt)) @@ -78,16 +72,19 @@ def trigger_collation( trig_merge_file = os.path.join( output_dir, - collation_control_options['collated-triggers-format'].format( - ifos=''.join(controls['ifos'].split()), + collation_control_options["collated-triggers-format"].format( + ifos="".join(controls["ifos"].split()), start=gps_start, - duration=(gps_end - gps_start) - ) + duration=(gps_end - gps_start), + ), ) collate_args += [ - '--gps-start-time', f'{gps_start:d}', - '--gps-end-time', f'{gps_end:d}', - '--output-file', trig_merge_file, + "--gps-start-time", + f"{gps_start:d}", + "--gps-end-time", + f"{gps_end:d}", + "--output-file", + trig_merge_file, ] sv.run_and_error(collate_args, controls) @@ -99,44 +96,41 @@ def get_active_ifos(trigger_file): """ Get the active ifos from the trigger file """ - with HFile(trigger_file, 'r') as f: - ifos = [ifo for ifo in f.keys() if 'snr' in f[ifo].keys()] - logging.info("Interferometers with data: %s", ' '.join(ifos)) + with HFile(trigger_file, "r") as f: + ifos = [ifo for ifo in f.keys() if "snr" in f[ifo].keys()] + logging.info("Interferometers with data: %s", " ".join(ifos)) return ifos def fit_by_template( - trigger_merge_file, - day_str, - fbt_control_options, - fbt_options, - output_dir, - ifo, - controls + trigger_merge_file, + day_str, + fbt_control_options, + fbt_options, + output_dir, + ifo, + controls, ): """ Supervise the running of pycbc_fit_sngls_by_template on live triggers """ logging.info("Performing daily fit_by_template") - fbt_out_fname = fbt_control_options['fit-by-template-format'].format( + fbt_out_fname = fbt_control_options["fit-by-template-format"].format( date=day_str, ifo=ifo, ) fbt_out_full = os.path.join(output_dir, fbt_out_fname) - fit_by_args = ['pycbc_fit_sngls_by_template'] - fit_by_args += ['--trigger-file', trigger_merge_file] + fit_by_args = ["pycbc_fit_sngls_by_template"] + fit_by_args += ["--trigger-file", trigger_merge_file] fit_by_args += sv.dict_to_args(fbt_options) - fit_by_args += ['--output', fbt_out_full, '--ifo', ifo] + fit_by_args += ["--output", fbt_out_full, "--ifo", ifo] sv.run_and_error(fit_by_args, controls) return fbt_out_full, day_str def find_daily_files( - combined_control_options, - daily_fname_format, - daily_files_dir, - ifo=None + combined_control_options, daily_fname_format, daily_files_dir, ifo=None ): """ Find files which match the specified formats @@ -145,7 +139,7 @@ def find_daily_files( if ifo is not None: log_str += f" in detector {ifo}" logging.info(log_str) - combined_days = int(combined_control_options['combined-days']) + combined_days = int(combined_control_options["combined-days"]) current_date = get_true_date(day_dt, combined_control_options) @@ -155,7 +149,7 @@ def find_daily_files( missed_files = 0 # Maximum consecutive number of days between files before a warning is raised # 10 days of the detector being off would be unusual for current detectors - max_nmissed = combined_control_options.get('maximum_missed_files', 10) + max_nmissed = combined_control_options.get("maximum_missed_files", 10) found_files = 0 while found_files < combined_days and missed_files < max_nmissed: # Loop through the possible file locations and see if the file exists @@ -166,20 +160,14 @@ def find_daily_files( ifo=ifo, ) - output_dir = os.path.join( - daily_files_dir, - date_out - ) - daily_full = os.path.join( - output_dir, - daily_fname - ) + output_dir = os.path.join(daily_files_dir, date_out) + daily_full = os.path.join(output_dir, daily_fname) # Check that the file exists: if not os.path.exists(daily_full): missed_files += 1 logging.info("File %s does not exist - skipping", daily_full) continue - if not len(daily_files): + if not daily_files: end_date = date_out # This is now the oldest file first_date = date_out @@ -196,177 +184,149 @@ def find_daily_files( # is wrong with the analysis. Warn about this and use fewer # files logging.warning( - f'More than {max_nmissed} days between files, only using ' - f'{found_files} files!' + f"More than {max_nmissed} days between files, only using " + f"{found_files} files!" ) return daily_files, first_date, end_date -def fit_over_multiparam( - fit_over_controls, - fit_over_options, - ifo, - output_dir, - controls -): +def fit_over_multiparam(fit_over_controls, fit_over_options, ifo, output_dir, controls): """ Supervise the smoothing of live trigger fits using pycbc_fit_sngls_over_multiparam """ daily_files, first_date, end_date = find_daily_files( fit_over_controls, - fit_over_controls['fit-by-format'], - controls['output-directory'], - ifo=ifo + fit_over_controls["fit-by-format"], + controls["output-directory"], + ifo=ifo, ) logging.info( "Smoothing fits using fit_over_multiparam with %d files and " "specified parameters", - len(daily_files) + len(daily_files), ) - file_id_str = f'{first_date}-{end_date}' - out_fname = fit_over_controls['fit-over-format'].format( + file_id_str = f"{first_date}-{end_date}" + out_fname = fit_over_controls["fit-over-format"].format( dates=file_id_str, ifo=ifo, ) - fit_over_args = ['pycbc_fit_sngls_over_multiparam', '--template-fit-file'] + fit_over_args = ["pycbc_fit_sngls_over_multiparam", "--template-fit-file"] fit_over_args += daily_files fit_over_args += sv.dict_to_args(fit_over_options) fit_over_full = os.path.join(output_dir, out_fname) - fit_over_args += ['--output', fit_over_full] + fit_over_args += ["--output", fit_over_full] sv.run_and_error(fit_over_args, controls) - if 'variable-fit-over-param' in fit_over_controls: - variable_fits = fit_over_controls['variable-fit-over-param'].format( - ifo=ifo - ) + if "variable-fit-over-param" in fit_over_controls: + variable_fits = fit_over_controls["variable-fit-over-param"].format(ifo=ifo) shutil.copyfile(fit_over_full, variable_fits) return fit_over_full, file_id_str def plot_fits( - fits_file, - ifo, - day_title_str, - plot_fit_options, - controls, - smoothed=False + fits_file, ifo, day_title_str, plot_fit_options, controls, smoothed=False ): """Plotting for fit_by files, and linking to the public directory""" - fits_plot_output = fits_file[:-3] + 'png' - logging.info( - "Plotting template fits %s to %s", - fits_file, - fits_plot_output - ) + fits_plot_output = fits_file[:-3] + "png" + logging.info("Plotting template fits %s to %s", fits_file, fits_plot_output) fits_plot_arguments = [ - 'pycbc_plot_bank_corner', - '--fits-file', + "pycbc_plot_bank_corner", + "--fits-file", fits_file, - '--output-plot-file', + "--output-plot-file", fits_plot_output, ] fits_plot_arguments += sv.dict_to_args(plot_fit_options) - title = "Fit parameters for pycbc-live, triggers from {}, {}".format( - ifo, - day_title_str - ) + title = f"Fit parameters for pycbc-live, triggers from {ifo}, {day_title_str}" if smoothed: - title += ', smoothed' - fits_plot_arguments += ['--title', title] + title += ", smoothed" + fits_plot_arguments += ["--title", title] sv.run_and_error(fits_plot_arguments, controls) - if 'public-dir' in controls: - public_dir = os.path.abspath(os.path.join( - controls['public-dir'], - *day_str.split('_') - )) + if "public-dir" in controls: + public_dir = os.path.abspath( + os.path.join(controls["public-dir"], *day_str.split("_")) + ) sv.symlink(fits_plot_output, public_dir) def single_significance_fits( - daily_fit_controls, - daily_fit_options, - output_dir, - day_str, - day_dt, - controls, - stat_files=None, + daily_fit_controls, + daily_fit_options, + output_dir, + day_str, + day_dt, + controls, + stat_files=None, ): """ Supervise the significance fits for live triggers using pycbc_live_single_significance_fits """ - daily_fit_options['output'] = os.path.join( + daily_fit_options["output"] = os.path.join( output_dir, - daily_fit_controls['sig-daily-format'].format( - ifos=''.join(sorted(controls['ifos'].split())), - date=day_str + daily_fit_controls["sig-daily-format"].format( + ifos="".join(sorted(controls["ifos"].split())), date=day_str ), ) - daily_args = ['pycbc_live_single_significance_fits'] + daily_args = ["pycbc_live_single_significance_fits"] gps_start_time = int(utc_datetime_to_gps(day_dt)) - gps_end_time = int(utc_datetime_to_gps(day_dt)+ timedelta(days=1)) + gps_end_time = int(utc_datetime_to_gps(day_dt) + timedelta(days=1)) - daily_fit_options['gps-start-time'] = f'{gps_start_time:d}' - daily_fit_options['gps-end-time'] = f'{gps_end_time:d}' + daily_fit_options["gps-start-time"] = f"{gps_start_time:d}" + daily_fit_options["gps-end-time"] = f"{gps_end_time:d}" daily_args += sv.dict_to_args(daily_fit_options) if stat_files is not None: - daily_args += ['--statistic-files'] + stat_files + daily_args += ["--statistic-files"] + stat_files sv.run_and_error(daily_args, controls) - return daily_fit_options['output'] + return daily_fit_options["output"] def plot_single_significance_fits(daily_output, daily_plot_options, controls): """ Plotting daily significance fits, and link to public directory if wanted """ - daily_plot_output = daily_output[:-4].replace( - ''.join(sorted(controls['ifos'].split())), - '{ifo}' - ) + '.png' + daily_plot_output = ( + daily_output[:-4].replace("".join(sorted(controls["ifos"].split())), "{ifo}") + + ".png" + ) logging.info( "Plotting daily significance fits from %s to %s", daily_output, - daily_plot_output + daily_plot_output, ) daily_plot_arguments = [ - 'pycbc_live_plot_single_significance_fits', - '--trigger-fits-file', + "pycbc_live_plot_single_significance_fits", + "--trigger-fits-file", daily_output, - '--output-plot-name-format', + "--output-plot-name-format", daily_plot_output, ] daily_plot_arguments += sv.dict_to_args(daily_plot_options) sv.run_and_error(daily_plot_arguments, controls) # Link the plots to the public-dir if wanted - if 'public-dir' in controls: + if "public-dir" in controls: daily_plot_outputs = [ - daily_plot_output.format(ifo=ifo) - for ifo in controls['ifos'].split() + daily_plot_output.format(ifo=ifo) for ifo in controls["ifos"].split() ] logging.info("Linking daily fits plots") for dpo in daily_plot_outputs: - public_dir = os.path.abspath(os.path.join( - controls['public-dir'], - *day_str.split('_') - )) + public_dir = os.path.abspath( + os.path.join(controls["public-dir"], *day_str.split("_")) + ) sv.symlink(dpo, public_dir) def combine_significance_fits( - combined_fit_options, - combined_fit_controls, - output_dir, - day_str, - controls + combined_fit_options, combined_fit_controls, output_dir, day_str, controls ): """ Supervise the smoothing of live trigger significance fits using @@ -376,64 +336,57 @@ def combine_significance_fits( # string, but not the date daily_files, first_date, end_date = find_daily_files( combined_fit_controls, - combined_fit_controls['daily-format'].format( - ifos=''.join(sorted(controls['ifos'].split())), - date='{date}' + combined_fit_controls["daily-format"].format( + ifos="".join(sorted(controls["ifos"].split())), date="{date}" ), - controls['output-directory'], - ) - logging.info( - "Smoothing significance fits over %d files", - len(daily_files) + controls["output-directory"], ) - date_range = f'{first_date}-{end_date}' - outfile_name = combined_fit_controls['outfile-format'].format( + logging.info("Smoothing significance fits over %d files", len(daily_files)) + date_range = f"{first_date}-{end_date}" + outfile_name = combined_fit_controls["outfile-format"].format( date=day_str, date_range=date_range, ) - combined_fit_options['output'] = os.path.join(output_dir, outfile_name) - combined_fit_options['trfits-files'] = ' '.join(daily_files) + combined_fit_options["output"] = os.path.join(output_dir, outfile_name) + combined_fit_options["trfits-files"] = " ".join(daily_files) - combined_fit_args = ['pycbc_live_combine_single_significance_fits'] + combined_fit_args = ["pycbc_live_combine_single_significance_fits"] combined_fit_args += sv.dict_to_args(combined_fit_options) sv.run_and_error(combined_fit_args, controls) - if 'variable-significance-fits' in combined_fit_controls: + if "variable-significance-fits" in combined_fit_controls: logging.info("Linking to variable significance fits file") shutil.copyfile( - combined_fit_options['output'], - combined_fit_controls['variable-significance-fits'] + combined_fit_options["output"], + combined_fit_controls["variable-significance-fits"], ) - return combined_fit_options['output'], date_range + return combined_fit_options["output"], date_range def plot_combined_significance_fits( - csf_file, - date_range, - output_dir, - combined_plot_options, - combined_plot_control_options, - controls + csf_file, + date_range, + output_dir, + combined_plot_options, + combined_plot_control_options, + controls, ): """ Plotting combined significance fits, and link to public directory if wanted """ - - oput_fmt = combined_plot_control_options['output-plot-name-format'] - if '{date_range}' not in oput_fmt: - raise RuntimeError( - "Must specify {date_range} in output-plot-name-format" - ) - oput_fmt = oput_fmt.replace('{date_range}', date_range) + oput_fmt = combined_plot_control_options["output-plot-name-format"] + if "{date_range}" not in oput_fmt: + raise RuntimeError("Must specify {date_range} in output-plot-name-format") + oput_fmt = oput_fmt.replace("{date_range}", date_range) oput_full = os.path.join(output_dir, oput_fmt) combined_plot_arguments = [ - 'pycbc_live_plot_combined_single_significance_fits', - '--combined-fits-file', + "pycbc_live_plot_combined_single_significance_fits", + "--combined-fits-file", csf_file, - '--output-plot-name-format', - oput_full + "--output-plot-name-format", + oput_full, ] combined_plot_arguments += sv.dict_to_args(combined_plot_options) @@ -441,33 +394,30 @@ def plot_combined_significance_fits( # Get the list of combined plotting output files: combined_plot_outputs = [ - oput_full.format(ifo=ifo, type='fit_coeffs') for ifo in - controls['ifos'].split() + oput_full.format(ifo=ifo, type="fit_coeffs") for ifo in controls["ifos"].split() ] combined_plot_outputs += [ - oput_full.format(ifo=ifo, type='counts') for ifo in - controls['ifos'].split() + oput_full.format(ifo=ifo, type="counts") for ifo in controls["ifos"].split() ] - if 'public-dir' in controls: + if "public-dir" in controls: logging.info("Linking combined fits to public dir") - public_dir = os.path.abspath(os.path.join( - controls['public-dir'], - *day_str.split('_') - )) + public_dir = os.path.abspath( + os.path.join(controls["public-dir"], *day_str.split("_")) + ) for cpo in combined_plot_outputs: sv.symlink(cpo, public_dir) def daily_dq_trigger_rates( - trigger_merge_file, - day_str, - day_dt, - daily_dq_controls, - daily_dq_options, - output_dir, - ifo, - controls + trigger_merge_file, + day_str, + day_dt, + daily_dq_controls, + daily_dq_options, + output_dir, + ifo, + controls, ): """ For a given day, record the number of idq flagged triggers @@ -475,7 +425,7 @@ def daily_dq_trigger_rates( of time flagged by idq and the total observing time. """ logging.info(f"Calculating daily dq trigger rates for {ifo}") - fname_format = daily_dq_controls['daily-dq-format'] + fname_format = daily_dq_controls["daily-dq-format"] ddtr_out_fname = fname_format.format(date=day_str, ifo=ifo) start_time = get_true_date(day_dt, daily_dq_controls) @@ -485,14 +435,14 @@ def daily_dq_trigger_rates( offset = int(utc_datetime_to_gps(day_dt) - gps_start_time) ddtr_out_full = os.path.join(output_dir, ddtr_out_fname) - daily_dq_args = ['pycbc_live_collated_dq_trigger_rates'] - daily_dq_args += ['--trigger-file', trigger_merge_file] - daily_dq_args += ['--output', ddtr_out_full] - daily_dq_args += ['--ifo', ifo] - - daily_dq_options['gps-start-time'] = f'{gps_start_time:d}' - daily_dq_options['gps-end-time'] = f'{gps_end_time:d}' - daily_dq_options['replay-offset'] = f'{offset:d}' + daily_dq_args = ["pycbc_live_collated_dq_trigger_rates"] + daily_dq_args += ["--trigger-file", trigger_merge_file] + daily_dq_args += ["--output", ddtr_out_full] + daily_dq_args += ["--ifo", ifo] + + daily_dq_options["gps-start-time"] = f"{gps_start_time:d}" + daily_dq_options["gps-end-time"] = f"{gps_end_time:d}" + daily_dq_options["replay-offset"] = f"{offset:d}" daily_dq_args += sv.dict_to_args(daily_dq_options) sv.run_and_error(daily_dq_args, controls) @@ -501,24 +451,20 @@ def daily_dq_trigger_rates( def combine_dq_trigger_rates( - combined_dq_controls, - combined_dq_options, - ifo, - output_dir, - controls + combined_dq_controls, combined_dq_options, ifo, output_dir, controls ): """ Calculate the dq trigger rate adjustment from multiple days of data """ daily_files, first_date, end_date = find_daily_files( combined_dq_controls, - combined_dq_controls['daily-dq-format'], - controls['output-directory'], - ifo + combined_dq_controls["daily-dq-format"], + controls["output-directory"], + ifo, ) - date_range = f'{first_date}-{end_date}' - outfile_name = combined_dq_controls['combined-dq-format'].format( + date_range = f"{first_date}-{end_date}" + outfile_name = combined_dq_controls["combined-dq-format"].format( ifo=ifo, dates=date_range, ) @@ -528,67 +474,51 @@ def combine_dq_trigger_rates( f"from {date_range}" ) - combined_dq_options['output'] = os.path.join(output_dir, outfile_name) - combined_dq_options['daily-dq-files'] = ' '.join(daily_files) - combined_dq_options['ifo'] = ifo + combined_dq_options["output"] = os.path.join(output_dir, outfile_name) + combined_dq_options["daily-dq-files"] = " ".join(daily_files) + combined_dq_options["ifo"] = ifo - combine_dq_args = ['pycbc_live_combine_dq_trigger_rates'] + combine_dq_args = ["pycbc_live_combine_dq_trigger_rates"] combine_dq_args += sv.dict_to_args(combined_dq_options) sv.run_and_error(combine_dq_args, controls) - if 'variable-dq-trigger-rates' in combined_dq_controls: + if "variable-dq-trigger-rates" in combined_dq_controls: logging.info("Linking to variable significance fits file") shutil.copyfile( - combined_dq_options['output'], - combined_dq_controls['variable-dq-trigger-rates'].format( - ifo=ifo - ) + combined_dq_options["output"], + combined_dq_controls["variable-dq-trigger-rates"].format(ifo=ifo), ) - return combined_dq_options['output'], date_range + return combined_dq_options["output"], date_range -def plot_dq_trigger_rates( - dq_file, - ifo, - date_str, - plot_dq_options, - controls -): +def plot_dq_trigger_rates(dq_file, ifo, date_str, plot_dq_options, controls): """ Plotting for dq trigger rate files, and linking to the public directory """ - dq_plot_output = dq_file[:-3] + 'png' - logging.info( - "Plotting dq trigger rates %s to %s", - dq_file, - dq_plot_output - ) + dq_plot_output = dq_file[:-3] + "png" + logging.info("Plotting dq trigger rates %s to %s", dq_file, dq_plot_output) dq_plot_arguments = [ - 'pycbc_plot_dq_flag_likelihood', - '--dq-file', + "pycbc_plot_dq_flag_likelihood", + "--dq-file", dq_file, - '--output-file', + "--output-file", dq_plot_output, - '--ifo', + "--ifo", ifo, ] dq_plot_arguments += sv.dict_to_args(plot_dq_options) - title = "Trigger rate adjustment for iDQ flagged times from {}, {}".format( - ifo, - date_str - ) - dq_plot_arguments += ['--title', title] + title = f"Trigger rate adjustment for iDQ flagged times from {ifo}, {date_str}" + dq_plot_arguments += ["--title", title] sv.run_and_error(dq_plot_arguments, controls) - if 'public-dir' in controls: - public_dir = os.path.abspath(os.path.join( - controls['public-dir'], - *day_str.split('_') - )) + if "public-dir" in controls: + public_dir = os.path.abspath( + os.path.join(controls["public-dir"], *day_str.split("_")) + ) sv.symlink(dq_plot_output, public_dir) @@ -598,48 +528,41 @@ def supervise_collation_fits_dq(args, day_dt, day_str): """ # Read in the config file and pack into appropriate dictionaries config_opts = read_options(args) - controls = config_opts['control'] - collation_options = config_opts['collation'] - collation_control_options = config_opts['collation_control'] - fit_by_template_options = config_opts['fit_by_template'] - fit_by_template_control_options = config_opts['fit_by_template_control'] - fit_over_options = config_opts['fit_over_multiparam'] - fit_over_control_options = config_opts['fit_over_multiparam_control'] - plot_fit_options = config_opts['plot_fit'] - daily_significance_options = config_opts['significance_daily_fits'] - daily_significance_control_options = config_opts['significance_daily_fits_control'] - daily_significance_plot_options = config_opts['plot_significance_daily'] - combined_significance_fit_options = config_opts['significance_combined_fits'] - combined_control_options = config_opts['significance_combined_fits_control'] - combined_plot_options = config_opts['plot_significance_combined'] - combined_plot_control_options = config_opts['plot_significance_combined_control'] - daily_dq_options = config_opts['daily_dq_trigger_rates'] - daily_dq_control_options = config_opts['daily_dq_trigger_rates_control'] - combined_dq_options = config_opts['combined_dq_trigger_rates'] - combined_dq_control_options = config_opts['combined_dq_trigger_rates_control'] - plot_dq_options = config_opts['plot_dq_trigger_rates'] + controls = config_opts["control"] + collation_options = config_opts["collation"] + collation_control_options = config_opts["collation_control"] + fit_by_template_options = config_opts["fit_by_template"] + fit_by_template_control_options = config_opts["fit_by_template_control"] + fit_over_options = config_opts["fit_over_multiparam"] + fit_over_control_options = config_opts["fit_over_multiparam_control"] + plot_fit_options = config_opts["plot_fit"] + daily_significance_options = config_opts["significance_daily_fits"] + daily_significance_control_options = config_opts["significance_daily_fits_control"] + daily_significance_plot_options = config_opts["plot_significance_daily"] + combined_significance_fit_options = config_opts["significance_combined_fits"] + combined_control_options = config_opts["significance_combined_fits_control"] + combined_plot_options = config_opts["plot_significance_combined"] + combined_plot_control_options = config_opts["plot_significance_combined_control"] + daily_dq_options = config_opts["daily_dq_trigger_rates"] + daily_dq_control_options = config_opts["daily_dq_trigger_rates_control"] + combined_dq_options = config_opts["combined_dq_trigger_rates"] + combined_dq_control_options = config_opts["combined_dq_trigger_rates_control"] + plot_dq_options = config_opts["plot_dq_trigger_rates"] # The main output directory will have a date subdirectory which we # put the output into sv.ensure_directories(controls, day_str) - output_dir = os.path.join( - controls['output-directory'], - day_str - ) + output_dir = os.path.join(controls["output-directory"], day_str) logging.info("Outputs to %s", output_dir) - if 'public_dir' in controls: - public_dir = os.path.abspath(os.path.join( - controls['public-dir'], - *day_str.split('_') - )) + if "public_dir" in controls: + public_dir = os.path.abspath( + os.path.join(controls["public-dir"], *day_str.split("_")) + ) logging.info("Outputs to be linked to % ", public_dir) if args.collated_trigger_file is not None: - logging.info( - "Using collated trigger file %s", - args.collated_trigger_file - ) + logging.info("Using collated trigger file %s", args.collated_trigger_file) merged_triggers = args.collated_trigger_file else: merged_triggers = trigger_collation( @@ -648,14 +571,14 @@ def supervise_collation_fits_dq(args, day_dt, day_str): collation_control_options, collation_options, output_dir, - controls + controls, ) # Store the locations of files needed for the statistic stat_files = [] # daily fits should only be done for IFOs that had data # and were included in the config file - all_ifos = controls['ifos'].split() + all_ifos = controls["ifos"].split() active_ifos = get_active_ifos(merged_triggers) active_ifos = [ifo for ifo in active_ifos if ifo in all_ifos] for ifo in active_ifos: @@ -670,13 +593,7 @@ def supervise_collation_fits_dq(args, day_dt, day_str): ifo, controls, ) - plot_fits( - fbt_file, - ifo, - date_str, - plot_fit_options, - controls - ) + plot_fits(fbt_file, ifo, date_str, plot_fit_options, controls) if args.daily_dq_trigger_rates: # compute and plot daily dq trigger rates @@ -688,26 +605,16 @@ def supervise_collation_fits_dq(args, day_dt, day_str): daily_dq_options, output_dir, ifo, - controls - ) - plot_dq_trigger_rates( - ddtr_file, - ifo, - day_str, - plot_dq_options, - controls + controls, ) + plot_dq_trigger_rates(ddtr_file, ifo, day_str, plot_dq_options, controls) for ifo in all_ifos: if args.fit_over_multiparam: # compute and plot template fits smoothed over parameter space # and combining multiple days of triggers fom_file, date_str = fit_over_multiparam( - fit_over_control_options, - fit_over_options, - ifo, - output_dir, - controls + fit_over_control_options, fit_over_options, ifo, output_dir, controls ) stat_files.append(fom_file) plot_fits( @@ -726,15 +633,11 @@ def supervise_collation_fits_dq(args, day_dt, day_str): combined_dq_options, ifo, output_dir, - controls + controls, ) stat_files.append(cdtr_file) plot_dq_trigger_rates( - cdtr_file, - ifo, - dq_date_str, - plot_dq_options, - controls + cdtr_file, ifo, dq_date_str, plot_dq_options, controls ) if args.single_significance_fits: @@ -748,9 +651,7 @@ def supervise_collation_fits_dq(args, day_dt, day_str): stat_files=stat_files, ) plot_single_significance_fits( - ssf_file, - daily_significance_plot_options, - controls + ssf_file, daily_significance_plot_options, controls ) if args.combine_significance_fits: csf_file, date_str = combine_significance_fits( @@ -758,7 +659,7 @@ def supervise_collation_fits_dq(args, day_dt, day_str): combined_control_options, output_dir, date_str, - controls + controls, ) plot_combined_significance_fits( csf_file, @@ -766,91 +667,86 @@ def supervise_collation_fits_dq(args, day_dt, day_str): output_dir, combined_plot_options, combined_plot_control_options, - controls + controls, ) def get_yesterday_date(): - """ Get the date string for yesterday's triggers """ + """Get the date string for yesterday's triggers""" day_dt = datetime.utcnow() - timedelta(days=1) day_dt = datetime.combine(day_dt, datetime.min.time()) - day_str = day_dt.strftime('%Y_%m_%d') + day_str = day_dt.strftime("%Y_%m_%d") return day_dt, day_str parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) +parser.add_argument("--config-file", required=True) parser.add_argument( - '--config-file', - required=True -) -parser.add_argument( - '--date', - help='Date to analyse, if not given, will analyse yesterday (UTC). ' - 'Format YYYY_MM_DD. Do not use if using --run-daily-at.' + "--date", + help="Date to analyse, if not given, will analyse yesterday (UTC). " + "Format YYYY_MM_DD. Do not use if using --run-daily-at.", ) parser.add_argument( - '--collated-trigger-file', - help='Collated trigger file to use, if not given, will be created. ' - 'Useful for testing.' + "--collated-trigger-file", + help="Collated trigger file to use, if not given, will be created. " + "Useful for testing.", ) parser.add_argument( - '--fit-by-template', - action='store_true', - help="Perform template fits calculation." + "--fit-by-template", action="store_true", help="Perform template fits calculation." ) parser.add_argument( - '--fit-over-multiparam', - action='store_true', - help="Perform template fits smoothing." + "--fit-over-multiparam", + action="store_true", + help="Perform template fits smoothing.", ) parser.add_argument( - '--single-significance-fits', - action='store_true', - help="Perform daily singles significance fits." + "--single-significance-fits", + action="store_true", + help="Perform daily singles significance fits.", ) parser.add_argument( - '--combine-significance-fits', - action='store_true', - help="Do combination of singles significance fits." + "--combine-significance-fits", + action="store_true", + help="Do combination of singles significance fits.", ) parser.add_argument( - '--daily-dq-trigger-rates', - action='store_true', - help="Calculate daily dq trigger rates." + "--daily-dq-trigger-rates", + action="store_true", + help="Calculate daily dq trigger rates.", ) parser.add_argument( - '--combine-dq-trigger-rates', - action='store_true', - help="Combine dq trigger rates over multiple days." + "--combine-dq-trigger-rates", + action="store_true", + help="Combine dq trigger rates over multiple days.", ) parser.add_argument( - '--run-daily-at', - metavar='HH:MM:SS', - help='Stay running and repeat the fitting daily at the given UTC hour.' + "--run-daily-at", + metavar="HH:MM:SS", + help="Stay running and repeat the fitting daily at the given UTC hour.", ) args = parser.parse_args() pycbc.init_logging(args.verbose, default_level=1) if args.run_daily_at is not None and args.date is not None: - parser.error('Cannot take --run-daily-at and --date at the same time') + parser.error("Cannot take --run-daily-at and --date at the same time") if args.run_daily_at is not None: # keep running and repeat the fitting every day at the given hour - if not re.match('[0-9][0-9]:[0-9][0-9]:[0-9][0-9]', args.run_daily_at): - parser.error('--run-daily-at takes a UTC time in the format HH:MM:SS') - logging.info('Starting in daily run mode') + if not re.match("[0-9][0-9]:[0-9][0-9]:[0-9][0-9]", args.run_daily_at): + parser.error("--run-daily-at takes a UTC time in the format HH:MM:SS") + logging.info("Starting in daily run mode") while True: sv.wait_for_utc_time(args.run_daily_at) day_dt, day_str = get_yesterday_date() - logging.info('==== Time to update the single fits, waking up ====') + logging.info("==== Time to update the single fits, waking up ====") supervise_collation_fits_dq(args, day_dt, day_str) else: # run just once if args.date: day_str = args.date - day_dt = datetime.strptime(args.date, '%Y_%m_%d') + day_dt = datetime.strptime(args.date, "%Y_%m_%d") else: day_dt, day_str = get_yesterday_date() supervise_collation_fits_dq(args, day_dt, day_str) diff --git a/bin/minifollowups/pycbc_foreground_minifollowup b/bin/minifollowups/pycbc_foreground_minifollowup index d547e011fd1..de7b78a79b7 100644 --- a/bin/minifollowups/pycbc_foreground_minifollowup +++ b/bin/minifollowups/pycbc_foreground_minifollowup @@ -14,51 +14,67 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Followup foreground events -""" -import os +"""Followup foreground events""" + import argparse import logging +import os import re import pycbc.workflow as wf -from pycbc import init_logging, add_common_pycbc_options +import pycbc.workflow.minifollowups as mini +from pycbc import add_common_pycbc_options, init_logging +from pycbc.events import coinc, select_segments_by_definer +from pycbc.io import HFile, get_all_subkeys from pycbc.results import layout from pycbc.types import MultiDetOptionAction -from pycbc.events import select_segments_by_definer, coinc -from pycbc.io import get_all_subkeys, HFile -import pycbc.workflow.minifollowups as mini from pycbc.workflow.core import resolve_url_to_file parser = argparse.ArgumentParser(description=__doc__[1:]) add_common_pycbc_options(parser) -parser.add_argument('--bank-file', - help="HDF format template bank file") -parser.add_argument('--statmap-file', - help="HDF format clustered coincident trigger result file") -parser.add_argument('--single-detector-triggers', nargs='+', action=MultiDetOptionAction, - help="HDF format merged single detector trigger files") -parser.add_argument('--inspiral-segments', - help="xml segment files containing the inspiral analysis times") -parser.add_argument('--inspiral-data-read-name', - help="Name of inspiral segmentlist containing data read in " - "by each analysis job.") -parser.add_argument('--inspiral-data-analyzed-name', - help="Name of inspiral segmentlist containing data " - "analyzed by each analysis job.") -parser.add_argument('--analysis-category', type=str, required=False, - default='background_exc', - choices = ['foreground', 'background', 'background_exc'], - help='Designates whether to look at foreground triggers ' - 'background triggers (including "little dogs") ' - 'or background triggers (with "little dogs" removed)') -parser.add_argument('--sort-variable', default='ifar', - help='Which subgroup of --analysis-category to use for ' - 'sorting. Default=ifar') -parser.add_argument('--sort-order', default='descending', - choices=['ascending','descending'], - help='Which direction to use when sorting on ' - '--sort-variable. Default=descending') +parser.add_argument("--bank-file", help="HDF format template bank file") +parser.add_argument( + "--statmap-file", help="HDF format clustered coincident trigger result file" +) +parser.add_argument( + "--single-detector-triggers", + nargs="+", + action=MultiDetOptionAction, + help="HDF format merged single detector trigger files", +) +parser.add_argument( + "--inspiral-segments", + help="xml segment files containing the inspiral analysis times", +) +parser.add_argument( + "--inspiral-data-read-name", + help="Name of inspiral segmentlist containing data read in by each analysis job.", +) +parser.add_argument( + "--inspiral-data-analyzed-name", + help="Name of inspiral segmentlist containing data analyzed by each analysis job.", +) +parser.add_argument( + "--analysis-category", + type=str, + required=False, + default="background_exc", + choices=["foreground", "background", "background_exc"], + help="Designates whether to look at foreground triggers " + 'background triggers (including "little dogs") ' + 'or background triggers (with "little dogs" removed)', +) +parser.add_argument( + "--sort-variable", + default="ifar", + help="Which subgroup of --analysis-category to use for sorting. Default=ifar", +) +parser.add_argument( + "--sort-order", + default="descending", + choices=["ascending", "descending"], + help="Which direction to use when sorting on --sort-variable. Default=descending", +) wf.add_workflow_command_line_group(parser) wf.add_workflow_settings_cli(parser, include_subdax_opts=True) @@ -84,34 +100,34 @@ insp_data_seglists = {} insp_analysed_seglists = {} for ifo in args.single_detector_triggers: fname = args.single_detector_triggers[ifo] - strig_file = resolve_url_to_file(os.path.abspath(fname), - attrs={'ifos': ifo}) + strig_file = resolve_url_to_file(os.path.abspath(fname), attrs={"ifos": ifo}) single_triggers.append(strig_file) - fsdt[ifo] = HFile(args.single_detector_triggers[ifo], 'r') + fsdt[ifo] = HFile(args.single_detector_triggers[ifo], "r") insp_data_seglists[ifo] = select_segments_by_definer( - args.inspiral_segments, - segment_name=args.inspiral_data_read_name, - ifo=ifo) + args.inspiral_segments, segment_name=args.inspiral_data_read_name, ifo=ifo + ) insp_analysed_seglists[ifo] = select_segments_by_definer( - args.inspiral_segments, - segment_name=args.inspiral_data_analyzed_name, - ifo=ifo) + args.inspiral_segments, segment_name=args.inspiral_data_analyzed_name, ifo=ifo + ) # NOTE: make_singles_timefreq needs a coalesced set of segments. If this is # being used to determine command-line options for other codes, # please think if that code requires coalesced, or not, segments. insp_data_seglists[ifo].coalesce() insp_analysed_seglists[ifo].coalesce() -num_events = int(workflow.cp.get_opt_tags('workflow-minifollowups', 'num-events', '')) -f = HFile(args.statmap_file, 'r') +num_events = int(workflow.cp.get_opt_tags("workflow-minifollowups", "num-events", "")) +f = HFile(args.statmap_file, "r") file_val = args.analysis_category -stat = f['{}/stat'.format(file_val)][:] +stat = f[f"{file_val}/stat"][:] if args.sort_variable not in f[file_val]: - all_datasets = [re.sub(file_val, '', ds).strip('/') - for ds in get_all_subkeys(f, file_val)] - raise KeyError(f'Sort variable {args.sort_variable} not in {file_val}: sort' - f'choices in {file_val} are ' + ', '.join(all_datasets)) + all_datasets = [ + re.sub(file_val, "", ds).strip("/") for ds in get_all_subkeys(f, file_val) + ] + raise KeyError( + f"Sort variable {args.sort_variable} not in {file_val}: sort" + f"choices in {file_val} are " + ", ".join(all_datasets) + ) # In case we are doing background minifollowup with repeated events, # we must include the ordering / template / trigger / time info for @@ -119,14 +135,12 @@ if args.sort_variable not in f[file_val]: events_to_read = num_events * 100 # We've asked for more events than there are! -if len(stat) < num_events: - num_events = len(stat) -if len(stat) < events_to_read: - events_to_read = len(stat) +num_events = min(num_events, len(stat)) +events_to_read = min(events_to_read, len(stat)) # Get the indices of the events we are considering in the order specified -sorting = f[file_val + '/' + args.sort_variable][:].argsort() -if args.sort_order == 'descending': +sorting = f[file_val + "/" + args.sort_variable][:].argsort() +if args.sort_order == "descending": sorting = sorting[::-1] event_idx = sorting[0:events_to_read] stat = stat[event_idx] @@ -134,15 +148,15 @@ stat = stat[event_idx] # Save the time / trigger / template ids for the events times = {} tids = {} -ifo_list = f.attrs['ifos'].split(' ') +ifo_list = f.attrs["ifos"].split(" ") f_cat = f[file_val] for ifo in ifo_list: - times[ifo] = f_cat[ifo]['time'][:][event_idx] - tids[ifo] = f_cat[ifo]['trigger_id'][:][event_idx] -bank_ids = f_cat['template_id'][:][event_idx] + times[ifo] = f_cat[ifo]["time"][:][event_idx] + tids[ifo] = f_cat[ifo]["trigger_id"][:][event_idx] +bank_ids = f_cat["template_id"][:][event_idx] f.close() -bank_data = HFile(args.bank_file, 'r') +bank_data = HFile(args.bank_file, "r") # loop over number of loudest events to be followed up event_times = {} @@ -159,8 +173,7 @@ while event_count < num_events and curr_idx < (events_to_read - 1): ifo_tids_strings = [] # Mean time to use as reference for zerolag - mtime = coinc.mean_if_greater_than_zero( - [times[ifo][curr_idx] for ifo in times])[0] + mtime = coinc.mean_if_greater_than_zero([times[ifo][curr_idx] for ifo in times])[0] for ifo in times: plottime = times[ifo][curr_idx] @@ -168,17 +181,19 @@ while event_count < num_events and curr_idx < (events_to_read - 1): # Use mean time instead and don't plot special trigger plottime = mtime else: # Do plot special trigger - ifo_tids_strings += ['%s:%s' % (ifo, tids[ifo][curr_idx])] - ifo_times_strings += ['%s:%s' % (ifo, plottime)] + ifo_tids_strings += ["%s:%s" % (ifo, tids[ifo][curr_idx])] + ifo_times_strings += ["%s:%s" % (ifo, plottime)] # For background do not want to follow up 10 background coincs with the # same event in ifo 1 and different events in ifo 2, so record times if ifo not in event_times: event_times[ifo] = [] # Only do this for background coincident triggers & not sentinel time -1 - if 'background' in args.analysis_category and \ - times[ifo][curr_idx] != -1 and \ - int(times[ifo][curr_idx]) in event_times[ifo]: + if ( + "background" in args.analysis_category + and times[ifo][curr_idx] != -1 + and int(times[ifo][curr_idx]) in event_times[ifo] + ): skipped_data.append((ifo, int(times[ifo][curr_idx]))) duplicate = True break @@ -188,35 +203,44 @@ while event_count < num_events and curr_idx < (events_to_read - 1): for ifo in times: event_times[ifo].append(int(times[ifo][curr_idx])) - ifo_times = ' '.join(ifo_times_strings) - ifo_tids = ' '.join(ifo_tids_strings) + ifo_times = " ".join(ifo_times_strings) + ifo_tids = " ".join(ifo_tids_strings) event_count += 1 if skipped_data: - layouts += (mini.make_skipped_html( - workflow, - skipped_data, - args.output_dir, - tags=[f'SKIP_{event_count}']),) + layouts += ( + mini.make_skipped_html( + workflow, skipped_data, args.output_dir, tags=[f"SKIP_{event_count}"] + ), + ) skipped_data = [] bank_id = bank_ids[curr_idx] - layouts += (mini.make_coinc_info(workflow, single_triggers, tmpltbank_file, - coinc_file, args.output_dir, n_loudest=curr_idx, - sort_order=args.sort_order, sort_var=args.sort_variable, - tags=args.tags + [str(event_count)]),) - files += mini.make_trigger_timeseries(workflow, single_triggers, - ifo_times, args.output_dir, special_tids=ifo_tids, - tags=args.tags + [str(event_count)]) + layouts += ( + mini.make_coinc_info( + workflow, + single_triggers, + tmpltbank_file, + coinc_file, + args.output_dir, + n_loudest=curr_idx, + sort_order=args.sort_order, + sort_var=args.sort_variable, + tags=args.tags + [str(event_count)], + ), + ) + files += mini.make_trigger_timeseries( + workflow, + single_triggers, + ifo_times, + args.output_dir, + special_tids=ifo_tids, + tags=args.tags + [str(event_count)], + ) params = mini.get_single_template_params( - curr_idx, - times, - bank_data, - bank_ids[curr_idx], - fsdt, - tids + curr_idx, times, bank_data, bank_ids[curr_idx], fsdt, tids ) _, sngl_tmplt_plots = mini.make_single_template_plots( @@ -227,32 +251,43 @@ while event_count < num_events and curr_idx < (events_to_read - 1): params, args.output_dir, data_segments=insp_data_seglists, - tags=args.tags + [str(event_count)] + tags=args.tags + [str(event_count)], ) files += sngl_tmplt_plots - for single in single_triggers: time = times[single.ifo][curr_idx] - if time==-1: + if time == -1: # If this detector did not trigger, still make the plot, but use # the average time of detectors which did trigger - time = coinc.mean_if_greater_than_zero([times[sngl.ifo][curr_idx] - for sngl in single_triggers])[0] + time = coinc.mean_if_greater_than_zero( + [times[sngl.ifo][curr_idx] for sngl in single_triggers] + )[0] for seg in insp_analysed_seglists[single.ifo]: if time in seg: - files += mini.make_singles_timefreq(workflow, single, tmpltbank_file, - time, args.output_dir, - data_segments=insp_data_seglists[single.ifo], - tags=args.tags + [str(event_count)]) - files += mini.make_qscan_plot\ - (workflow, single.ifo, time, args.output_dir, - data_segments=insp_data_seglists[single.ifo], - tags=args.tags + [str(event_count)]) + files += mini.make_singles_timefreq( + workflow, + single, + tmpltbank_file, + time, + args.output_dir, + data_segments=insp_data_seglists[single.ifo], + tags=args.tags + [str(event_count)], + ) + files += mini.make_qscan_plot( + workflow, + single.ifo, + time, + args.output_dir, + data_segments=insp_data_seglists[single.ifo], + tags=args.tags + [str(event_count)], + ) break else: - logging.info(f'Trigger time {time} is not valid in ' - f'{single.ifo}, skipping singles plots') + logging.info( + f"Trigger time {time} is not valid in " + f"{single.ifo}, skipping singles plots" + ) layouts += list(layout.grouper(files, 2)) diff --git a/bin/minifollowups/pycbc_injection_minifollowup b/bin/minifollowups/pycbc_injection_minifollowup index 871b81f51ab..bcbf48d0397 100644 --- a/bin/minifollowups/pycbc_injection_minifollowup +++ b/bin/minifollowups/pycbc_injection_minifollowup @@ -17,32 +17,33 @@ """Create a workflow for following up missed loud injections.""" -import os import argparse -import logging import copy +import logging +import os + import numpy -from pycbc import init_logging, add_common_pycbc_options import pycbc.workflow as wf import pycbc.workflow.minifollowups as mini -from pycbc.types import MultiDetOptionAction -from pycbc.events import select_segments_by_definer, coinc -from pycbc.results import layout +from pycbc import add_common_pycbc_options, init_logging from pycbc.detector import Detector +from pycbc.events import coinc, select_segments_by_definer +from pycbc.io.hdf import HFile, SingleDetTriggers +from pycbc.results import layout +from pycbc.types import MultiDetOptionAction from pycbc.workflow.core import resolve_url_to_file -from pycbc.io.hdf import SingleDetTriggers, HFile - legal_distance_types = [ - 'decisive_optimal_snr', - 'comb_optimal_snr', - 'dec_chirp_distance' + "decisive_optimal_snr", + "comb_optimal_snr", + "dec_chirp_distance", ] def sort_injections(args, inj_group, missed): - """Return an array of indices to sort the missed injections from most to + """ + Return an array of indices to sort the missed injections from most to least likely to be detected, according to a metric of choice. Parameters @@ -60,59 +61,58 @@ def sort_injections(args, inj_group, missed): ------- missed_sorted : array Array of indices of missed injections sorted as requested. + """ - if not hasattr(args, 'distance_type'): - raise ValueError('Distance type not provided') + if not hasattr(args, "distance_type"): + raise ValueError("Distance type not provided") if args.distance_type not in legal_distance_types: raise ValueError( f'Invalid distance type "{args.distance_type}", ' - f'allowed types are {", ".join(legal_distance_types)}' + f"allowed types are {', '.join(legal_distance_types)}" ) - if 'optimal_snr' in args.distance_type: + if "optimal_snr" in args.distance_type: optimal_snrs = [ - inj_group[dsn][:][missed] for dsn in inj_group.keys() - if dsn.startswith('optimal_snr_') + inj_group[dsn][:][missed] + for dsn in inj_group.keys() + if dsn.startswith("optimal_snr_") ] - assert optimal_snrs, 'These injections do not have optimal SNRs' + assert optimal_snrs, "These injections do not have optimal SNRs" - if args.distance_type == 'decisive_optimal_snr': + if args.distance_type == "decisive_optimal_snr": # descending order of decisive (2nd largest) optimal SNR - dec_snr = numpy.array([ - sorted(snrs)[-2] for snrs in zip(*optimal_snrs) - ]) + dec_snr = numpy.array([sorted(snrs)[-2] for snrs in zip(*optimal_snrs)]) if args.maximum_decisive_snr is not None: # By setting to 0, these injections will not be considered dec_snr[dec_snr > args.maximum_decisive_snr] = 0 sorter = dec_snr.argsort()[::-1] return missed[sorter] - if args.distance_type == 'comb_optimal_snr': + if args.distance_type == "comb_optimal_snr": # descending order of network optimal SNR optimal_snrs = numpy.vstack(optimal_snrs) - net_opt_snrs_squared = (optimal_snrs ** 2).sum(axis=0) + net_opt_snrs_squared = (optimal_snrs**2).sum(axis=0) sorter = net_opt_snrs_squared.argsort()[::-1] return missed[sorter] - if args.distance_type == 'dec_chirp_distance': + if args.distance_type == "dec_chirp_distance": # ascending order of decisive (2nd smallest) chirp distance - from pycbc.conversions import mchirp_from_mass1_mass2, chirp_distance + from pycbc.conversions import chirp_distance, mchirp_from_mass1_mass2 eff_dists = [] for ifo in args.single_detector_triggers: eff_dist = Detector(ifo).effective_distance( - inj_group['distance'][:][missed], - inj_group['ra'][:][missed], - inj_group['dec'][:][missed], - inj_group['polarization'][:][missed], - inj_group['tc'][:][missed], - inj_group['inclination'][:][missed] + inj_group["distance"][:][missed], + inj_group["ra"][:][missed], + inj_group["dec"][:][missed], + inj_group["polarization"][:][missed], + inj_group["tc"][:][missed], + inj_group["inclination"][:][missed], ) eff_dists.append(eff_dist) dec_eff_dist = sorted(eff_dists)[-2] mchirp = mchirp_from_mass1_mass2( - inj_group['mass1'][:][missed], - inj_group['mass2'][:][missed] + inj_group["mass1"][:][missed], inj_group["mass2"][:][missed] ) dec_chirp_dist = chirp_distance(dec_dist, mchirp) sorter = dec_chirp_dist.argsort() @@ -121,39 +121,60 @@ def sort_injections(args, inj_group, missed): parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--bank-file', - help="HDF format template bank file") -parser.add_argument('--injection-file', - help="HDF format injection results file") -parser.add_argument('--injection-xml-file', - help="XML format injection file") -parser.add_argument('--single-detector-triggers', nargs='+', action=MultiDetOptionAction, - help="HDF format merged single detector trigger files") -parser.add_argument('--inspiral-segments', - help="xml segment files containing the inspiral analysis times") -parser.add_argument('--inspiral-data-read-name', - help="Name of inspiral segmentlist containing data read in " - "by each analysis job.") -parser.add_argument('--inspiral-data-analyzed-name', - help="Name of inspiral segmentlist containing data " - "analyzed by each analysis job.") -parser.add_argument('--inj-window', type=int, default=0.5, - help="Time window in which to look for injection triggers") -parser.add_argument('--ifar-threshold', type=float, default=None, - help="If given also followup injections with ifar smaller " - "than this threshold.") -parser.add_argument('--maximum-decisive-snr', type=float, default=None, - help="If given, only followup injections where the " - "decisive SNR is smaller than this value.") -parser.add_argument('--nearby-triggers-window', type=float, default=0.05, - help="Maximum time difference between the missed " - "injection and the loudest SNR nearby trigger to " - "display, seconds. Default=0.05") -parser.add_argument('--distance-type', - required=True, - choices=legal_distance_types, - help="How to sort missed injections from most to least " - "likely to be detected") +parser.add_argument("--bank-file", help="HDF format template bank file") +parser.add_argument("--injection-file", help="HDF format injection results file") +parser.add_argument("--injection-xml-file", help="XML format injection file") +parser.add_argument( + "--single-detector-triggers", + nargs="+", + action=MultiDetOptionAction, + help="HDF format merged single detector trigger files", +) +parser.add_argument( + "--inspiral-segments", + help="xml segment files containing the inspiral analysis times", +) +parser.add_argument( + "--inspiral-data-read-name", + help="Name of inspiral segmentlist containing data read in by each analysis job.", +) +parser.add_argument( + "--inspiral-data-analyzed-name", + help="Name of inspiral segmentlist containing data analyzed by each analysis job.", +) +parser.add_argument( + "--inj-window", + type=int, + default=0.5, + help="Time window in which to look for injection triggers", +) +parser.add_argument( + "--ifar-threshold", + type=float, + default=None, + help="If given also followup injections with ifar smaller than this threshold.", +) +parser.add_argument( + "--maximum-decisive-snr", + type=float, + default=None, + help="If given, only followup injections where the " + "decisive SNR is smaller than this value.", +) +parser.add_argument( + "--nearby-triggers-window", + type=float, + default=0.05, + help="Maximum time difference between the missed " + "injection and the loudest SNR nearby trigger to " + "display, seconds. Default=0.05", +) +parser.add_argument( + "--distance-type", + required=True, + choices=legal_distance_types, + help="How to sort missed injections from most to least likely to be detected", +) wf.add_workflow_command_line_group(parser) wf.add_workflow_settings_cli(parser, include_subdax_opts=True) args = parser.parse_args() @@ -179,18 +200,13 @@ insp_data_seglists = {} insp_analysed_seglists = {} for ifo in args.single_detector_triggers: fname = args.single_detector_triggers[ifo] - strig_file = resolve_url_to_file(os.path.abspath(fname), - attrs={'ifos': ifo}) + strig_file = resolve_url_to_file(os.path.abspath(fname), attrs={"ifos": ifo}) single_triggers.append(strig_file) insp_data_seglists[ifo] = select_segments_by_definer( - args.inspiral_segments, - segment_name=args.inspiral_data_read_name, - ifo=ifo + args.inspiral_segments, segment_name=args.inspiral_data_read_name, ifo=ifo ) insp_analysed_seglists[ifo] = select_segments_by_definer( - args.inspiral_segments, - segment_name=args.inspiral_data_analyzed_name, - ifo=ifo + args.inspiral_segments, segment_name=args.inspiral_data_analyzed_name, ifo=ifo ) # NOTE: make_singles_timefreq needs a coalesced set of segments. If this is # being used to determine command-line options for other codes, @@ -198,27 +214,27 @@ for ifo in args.single_detector_triggers: insp_data_seglists[ifo].coalesce() insp_analysed_seglists[ifo].coalesce() -f = HFile(args.injection_file, 'r') -inj_def = f['injections'] -missed = f['missed/after_vetoes'][:] +f = HFile(args.injection_file, "r") +inj_def = f["injections"] +missed = f["missed/after_vetoes"][:] if args.ifar_threshold is not None: try: # injections may not have (inclusive) IFAR present - ifars = f['found_after_vetoes']['ifar'][:] + ifars = f["found_after_vetoes"]["ifar"][:] except KeyError: - ifars = f['found_after_vetoes']['ifar_exc'][:] - logging.warning('Inclusive IFAR not found, using exclusive') + ifars = f["found_after_vetoes"]["ifar_exc"][:] + logging.warning("Inclusive IFAR not found, using exclusive") lgc_arr = ifars < args.ifar_threshold - missed = numpy.append(missed, - f['found_after_vetoes']['injection_index'][lgc_arr]) + missed = numpy.append(missed, f["found_after_vetoes"]["injection_index"][lgc_arr]) # Get the trigger SNRs and times # But only ones which are within a small window of the missed injection -missed_inj_times = numpy.sort(inj_def['tc'][:][missed]) +missed_inj_times = numpy.sort(inj_def["tc"][:][missed]) # Note: Adding Earth diameter in light seconds to the window here # to allow for different IFO's arrival times of the injection safe_window = args.nearby_triggers_window + 0.0425 + def nearby_missedinj(endtime, snr): """ Convenience function to check if trigger times are within a small @@ -237,51 +253,49 @@ def nearby_missedinj(endtime, snr): ------- boolean array True for triggers which are close to any missed injection + """ left = numpy.searchsorted(missed_inj_times - safe_window, endtime) right = numpy.searchsorted(missed_inj_times + safe_window, endtime) return left != right + trigger_idx = {} trigger_snrs = {} trigger_times = {} # This finds the triggers near to _any_ missed injection for trig in single_triggers: ifo = trig.ifo - with HFile(trig.lfn, 'r') as trig_f: - trigger_idx[ifo], data_tuple = \ - trig_f.select( - nearby_missedinj, - f'{ifo}/end_time', - f'{ifo}/snr', - ) + with HFile(trig.lfn, "r") as trig_f: + trigger_idx[ifo], data_tuple = trig_f.select( + nearby_missedinj, + f"{ifo}/end_time", + f"{ifo}/snr", + ) trigger_times[ifo], trigger_snrs[ifo] = data_tuple # figure out how many injections to follow up -num_events = int(workflow.cp.get_opt_tags( - 'workflow-injection_minifollowups', - 'num-events', - '' -)) -if len(missed) < num_events: - num_events = len(missed) +num_events = int( + workflow.cp.get_opt_tags("workflow-injection_minifollowups", "num-events", "") +) +num_events = min(num_events, len(missed)) # sort the injections missed = sort_injections(args, inj_def, missed) # loop over sorted missed injections to be followed up -found_inj_idxes = f['found_after_vetoes/injection_index'][:] +found_inj_idxes = f["found_after_vetoes/injection_index"][:] for num_event in range(num_events): files = wf.FileList([]) injection_index = missed[num_event] - time = inj_def['tc'][injection_index] - lon = inj_def['ra'][injection_index] - lat = inj_def['dec'][injection_index] + time = inj_def["tc"][injection_index] + lon = inj_def["ra"][injection_index] + lat = inj_def["dec"][injection_index] - ifo_times = '' + ifo_times = "" inj_params = {} - for val in ['mass1', 'mass2', 'spin1z', 'spin2z', 'tc']: + for val in ["mass1", "mass2", "spin1z", "spin2z", "tc"]: inj_params[val] = inj_def[val][injection_index] for single in single_triggers: ifo = single.ifo @@ -293,26 +307,45 @@ for num_event in range(num_events): else: ifo_time = -1.0 - ifo_times += ' %s:%s ' % (ifo, ifo_time) - inj_params[ifo + '_end_time'] = ifo_time - all_times = [inj_params[sngl.ifo + '_end_time'] for sngl in single_triggers] - inj_params['mean_time'] = coinc.mean_if_greater_than_zero(all_times)[0] - - layouts += [(mini.make_inj_info(workflow, injection_file, injection_index, num_event, - args.output_dir, tags=args.tags + [str(num_event)])[0],)] + ifo_times += " %s:%s " % (ifo, ifo_time) + inj_params[ifo + "_end_time"] = ifo_time + all_times = [inj_params[sngl.ifo + "_end_time"] for sngl in single_triggers] + inj_params["mean_time"] = coinc.mean_if_greater_than_zero(all_times)[0] + + layouts += [ + ( + mini.make_inj_info( + workflow, + injection_file, + injection_index, + num_event, + args.output_dir, + tags=args.tags + [str(num_event)], + )[0], + ) + ] if injection_index in found_inj_idxes: trig_id = numpy.where(found_inj_idxes == injection_index)[0][0] - layouts += [(mini.make_coinc_info - (workflow, single_triggers, tmpltbank_file, - injection_file, args.output_dir, trig_id=trig_id, - file_substring='found_after_vetoes', - title="Details of closest event", - tags=args.tags + [str(num_event)])[0],)] + layouts += [ + ( + mini.make_coinc_info( + workflow, + single_triggers, + tmpltbank_file, + injection_file, + args.output_dir, + trig_id=trig_id, + file_substring="found_after_vetoes", + title="Details of closest event", + tags=args.tags + [str(num_event)], + )[0], + ) + ] for sngl in single_triggers: # Find the triggers close to _this_ injection at this IFO ifo = sngl.ifo - trig_tdiff = abs(inj_params[ifo + '_end_time'] - trigger_times[ifo]) + trig_tdiff = abs(inj_params[ifo + "_end_time"] - trigger_times[ifo]) nearby = trig_tdiff < args.nearby_triggers_window if not any(nearby): # If there are no triggers in the defined window, @@ -323,68 +356,101 @@ for num_event in range(num_events): # Convert to the index in the trigger file nearby_trigger_idx = trigger_idx[ifo][nearby][loudest] # Make the info snippet - sngl_info = mini.make_sngl_ifo(workflow, sngl, tmpltbank_file, - nearby_trigger_idx, args.output_dir, ifo, + sngl_info = mini.make_sngl_ifo( + workflow, + sngl, + tmpltbank_file, + nearby_trigger_idx, + args.output_dir, + ifo, title=f"Parameters of loudest SNR nearby trigger in {ifo}", - tags=args.tags + [str(num_event)])[0] + tags=args.tags + [str(num_event)], + )[0] layouts += [(sngl_info,)] - files += mini.make_trigger_timeseries(workflow, single_triggers, - ifo_times, args.output_dir, - tags=args.tags + [str(num_event)]) + files += mini.make_trigger_timeseries( + workflow, + single_triggers, + ifo_times, + args.output_dir, + tags=args.tags + [str(num_event)], + ) for single in single_triggers: checkedtime = time - if (inj_params[single.ifo + '_end_time'] == -1.0): - checkedtime = inj_params['mean_time'] + if inj_params[single.ifo + "_end_time"] == -1.0: + checkedtime = inj_params["mean_time"] for seg in insp_analysed_seglists[single.ifo]: if checkedtime in seg: - files += mini.make_singles_timefreq(workflow, single, tmpltbank_file, - checkedtime, args.output_dir, - data_segments=insp_data_seglists[single.ifo], - tags=args.tags + [str(num_event)]) - files += mini.make_qscan_plot\ - (workflow, single.ifo, checkedtime, args.output_dir, - data_segments=insp_data_seglists[single.ifo], - injection_file=injection_xml_file, - tags=args.tags + [str(num_event)]) + files += mini.make_singles_timefreq( + workflow, + single, + tmpltbank_file, + checkedtime, + args.output_dir, + data_segments=insp_data_seglists[single.ifo], + tags=args.tags + [str(num_event)], + ) + files += mini.make_qscan_plot( + workflow, + single.ifo, + checkedtime, + args.output_dir, + data_segments=insp_data_seglists[single.ifo], + injection_file=injection_xml_file, + tags=args.tags + [str(num_event)], + ) break else: logging.info( - 'Trigger time %s is not valid in %s, skipping singles plots', - checkedtime, single.ifo + "Trigger time %s is not valid in %s, skipping singles plots", + checkedtime, + single.ifo, ) - _, norm_plot = mini.make_single_template_plots(workflow, insp_segs, - args.inspiral_data_read_name, - args.inspiral_data_analyzed_name, inj_params, - args.output_dir, inj_file=injection_xml_file, - tags=args.tags+['INJ_PARAMS',str(num_event)], - params_str='injection parameters as template, ' +\ - 'here the injection is made as normal', - use_exact_inj_params=True) + _, norm_plot = mini.make_single_template_plots( + workflow, + insp_segs, + args.inspiral_data_read_name, + args.inspiral_data_analyzed_name, + inj_params, + args.output_dir, + inj_file=injection_xml_file, + tags=args.tags + ["INJ_PARAMS", str(num_event)], + params_str="injection parameters as template, " + "here the injection is made as normal", + use_exact_inj_params=True, + ) files += norm_plot - _, inv_plot = mini.make_single_template_plots(workflow, insp_segs, - args.inspiral_data_read_name, - args.inspiral_data_analyzed_name, inj_params, - args.output_dir, inj_file=injection_xml_file, - tags=args.tags + ['INJ_PARAMS_INVERTED', - str(num_event)], - params_str='injection parameters as template, ' +\ - 'here the injection is made inverted', - use_exact_inj_params=True) + _, inv_plot = mini.make_single_template_plots( + workflow, + insp_segs, + args.inspiral_data_read_name, + args.inspiral_data_analyzed_name, + inj_params, + args.output_dir, + inj_file=injection_xml_file, + tags=args.tags + ["INJ_PARAMS_INVERTED", str(num_event)], + params_str="injection parameters as template, " + "here the injection is made inverted", + use_exact_inj_params=True, + ) files += inv_plot - _, noinj_plot = mini.make_single_template_plots(workflow, insp_segs, - args.inspiral_data_read_name, - args.inspiral_data_analyzed_name, inj_params, - args.output_dir, inj_file=injection_xml_file, - tags=args.tags + ['INJ_PARAMS_NOINJ', - str(num_event)], - params_str='injection parameters, here no ' +\ - 'injection was actually performed', - use_exact_inj_params=True) + _, noinj_plot = mini.make_single_template_plots( + workflow, + insp_segs, + args.inspiral_data_read_name, + args.inspiral_data_analyzed_name, + inj_params, + args.output_dir, + inj_file=injection_xml_file, + tags=args.tags + ["INJ_PARAMS_NOINJ", str(num_event)], + params_str="injection parameters, here no " + "injection was actually performed", + use_exact_inj_params=True, + ) files += noinj_plot for curr_ifo in args.single_detector_triggers: @@ -392,8 +458,8 @@ for num_event in range(num_events): # First, find triggers close to the missed injection single_fname = args.single_detector_triggers[curr_ifo] idx, _ = HFile(single_fname).select( - lambda t: abs(t - inj_params['tc']) < args.inj_window, - f'{curr_ifo}/end_time', + lambda t: abs(t - inj_params["tc"]) < args.inj_window, + f"{curr_ifo}/end_time", return_data=False, ) @@ -402,10 +468,7 @@ for num_event in range(num_events): continue hd_sngl = SingleDetTriggers( - single_fname, - curr_ifo, - bank_file=args.bank_file, - premask=idx + single_fname, curr_ifo, bank_file=args.bank_file, premask=idx ) # Next, find the loudest within this set of triggers # Use SNR here or NewSNR, or other?? @@ -414,34 +477,39 @@ for num_event in range(num_events): # What are the parameters of this trigger? curr_params = copy.deepcopy(inj_params) - curr_params['mass1'] = hd_sngl.mass1[0] - curr_params['mass2'] = hd_sngl.mass2[0] - curr_params['spin1z'] = hd_sngl.spin1z[0] - curr_params['spin2z'] = hd_sngl.spin2z[0] - curr_params['f_lower'] = hd_sngl.f_lower[0] + curr_params["mass1"] = hd_sngl.mass1[0] + curr_params["mass2"] = hd_sngl.mass2[0] + curr_params["spin1z"] = hd_sngl.spin1z[0] + curr_params["spin2z"] = hd_sngl.spin2z[0] + curr_params["f_lower"] = hd_sngl.f_lower[0] # don't require precessing template info if not present try: - curr_params['spin1x'] = hd_sngl.spin1x[0] - curr_params['spin2x'] = hd_sngl.spin2x[0] - curr_params['spin1y'] = hd_sngl.spin1y[0] - curr_params['spin2y'] = hd_sngl.spin2y[0] - curr_params['inclination'] = hd_sngl.inclination[0] + curr_params["spin1x"] = hd_sngl.spin1x[0] + curr_params["spin2x"] = hd_sngl.spin2x[0] + curr_params["spin1y"] = hd_sngl.spin1y[0] + curr_params["spin2y"] = hd_sngl.spin2y[0] + curr_params["inclination"] = hd_sngl.inclination[0] except KeyError: pass try: # Only present for precessing search - curr_params['u_vals'] = hd_sngl.u_vals[0] + curr_params["u_vals"] = hd_sngl.u_vals[0] except: pass - curr_tags = ['TMPLT_PARAMS_%s' %(curr_ifo,)] + curr_tags = ["TMPLT_PARAMS_%s" % (curr_ifo,)] curr_tags += [str(num_event)] - _, loudest_plot = mini.make_single_template_plots(workflow, insp_segs, - args.inspiral_data_read_name, - args.inspiral_data_analyzed_name, curr_params, - args.output_dir, inj_file=injection_xml_file, - tags=args.tags + curr_tags, - params_str='loudest template in %s' % curr_ifo ) + _, loudest_plot = mini.make_single_template_plots( + workflow, + insp_segs, + args.inspiral_data_read_name, + args.inspiral_data_analyzed_name, + curr_params, + args.output_dir, + inj_file=injection_xml_file, + tags=args.tags + curr_tags, + params_str="loudest template in %s" % curr_ifo, + ) files += loudest_plot layouts += list(layout.grouper(files, 2)) diff --git a/bin/minifollowups/pycbc_page_coincinfo b/bin/minifollowups/pycbc_page_coincinfo index 2f495ef88a1..93bcc6d31ea 100644 --- a/bin/minifollowups/pycbc_page_coincinfo +++ b/bin/minifollowups/pycbc_page_coincinfo @@ -14,116 +14,153 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Make tables describing a foreground event""" +"""Make tables describing a foreground event""" import argparse import logging import sys + import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") import numpy -from pycbc import add_common_pycbc_options -import pycbc.results import pycbc.pnutils -from pycbc.io.hdf import HFile +import pycbc.results +from pycbc import add_common_pycbc_options from pycbc.events import stat as pystat +from pycbc.io.hdf import HFile from pycbc.results import followup from pycbc.time import gps_to_utc_datetime - parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--single-trigger-files', nargs='+', - help="HDF format single detector trigger files for the full data run") -parser.add_argument('--bank-file', - help="HDF format template bank file") -parser.add_argument('--output-file') -parser.add_argument('--statmap-file', required=True, - help="HDF format clustered statmap file containing the result " - "triggers. Required") -parser.add_argument('--statmap-file-subspace-name', default='background_exc', +parser.add_argument( + "--single-trigger-files", + nargs="+", + help="HDF format single detector trigger files for the full data run", +) +parser.add_argument("--bank-file", help="HDF format template bank file") +parser.add_argument("--output-file") +parser.add_argument( + "--statmap-file", + required=True, + help="HDF format clustered statmap file containing the result triggers. Required", +) +parser.add_argument( + "--statmap-file-subspace-name", + default="background_exc", help="If given look in this 'sub-directory' of the HDF file for triggers, " - "takes a default value of 'background_exc'.") + "takes a default value of 'background_exc'.", +) trig_input = parser.add_mutually_exclusive_group(required=True) -trig_input.add_argument('--n-loudest', type=int, - help="Examine the n'th loudest trigger, use with statmap file") -trig_input.add_argument('--trigger-id', type=int, +trig_input.add_argument( + "--n-loudest", + type=int, + help="Examine the n'th loudest trigger, use with statmap file", +) +trig_input.add_argument( + "--trigger-id", + type=int, help="Examine the trigger with specified ID, use with statmap file. An " - "alternative to --n-loudest. Cannot be used together") -parser.add_argument('--sort-variable', default='ifar', - help='Which subgroup of --analysis-category to use for ' - 'sorting if using --n-loudest. Default=ifar') -parser.add_argument('--sort-order', default='descending', - choices=['ascending','descending'], - help='Which direction to use when sorting on ' - '--sort-variable with --n-loudest. Default=descending') -parser.add_argument('--title', + "alternative to --n-loudest. Cannot be used together", +) +parser.add_argument( + "--sort-variable", + default="ifar", + help="Which subgroup of --analysis-category to use for " + "sorting if using --n-loudest. Default=ifar", +) +parser.add_argument( + "--sort-order", + default="descending", + choices=["ascending", "descending"], + help="Which direction to use when sorting on " + "--sort-variable with --n-loudest. Default=descending", +) +parser.add_argument( + "--title", help="Supply a title for the event details. Defaults are " - "'Parameters of event ranked N' if --n-loudest is given, or " - "'Details of trigger' for --trigger-id.") -parser.add_argument('--include-summary-page-link', action='store_true', + "'Parameters of event ranked N' if --n-loudest is given, or " + "'Details of trigger' for --trigger-id.", +) +parser.add_argument( + "--include-summary-page-link", + action="store_true", help="If given, will include a link to the DQ summary page on the " - "single detector trigger tables.") -parser.add_argument('--include-gracedb-link', action='store_true', + "single detector trigger tables.", +) +parser.add_argument( + "--include-gracedb-link", + action="store_true", help="If given, will provide a link to search GraceDB for events " - "within a 3s window around the coincidence time.") -parser.add_argument('--max-columns', type=int, - help="Maximum number of columns allowed in the table (not including detector names)") -pystat.insert_statistic_option_group(parser, - default_ranking_statistic='single_ranking_only') + "within a 3s window around the coincidence time.", +) +parser.add_argument( + "--max-columns", + type=int, + help="Maximum number of columns allowed in the table (not including detector names)", +) +pystat.insert_statistic_option_group( + parser, default_ranking_statistic="single_ranking_only" +) args = parser.parse_args() pycbc.init_logging(args.verbose) -if args.ranking_statistic not in ['quadsum', 'single_ranking_only']: +if args.ranking_statistic not in ["quadsum", "single_ranking_only"]: logging.warning( "For the coincident info table, we only use single ranking, not %s, " "this option will be ignored", - args.ranking_statistic + args.ranking_statistic, ) - args.ranking_statistic = 'quadsum' + args.ranking_statistic = "quadsum" # Get the nth loudest trigger from the output of pycbc_coinc_statmap -f = HFile(args.statmap_file, 'r') +f = HFile(args.statmap_file, "r") d = f[args.statmap_file_subspace_name] if args.n_loudest is not None: sorting = d[args.sort_variable][:].argsort() - if args.sort_order == 'descending': + if args.sort_order == "descending": sorting = sorting[::-1] n = sorting[args.n_loudest] - title = 'Parameters of event ranked %s' % (args.n_loudest + 1) - caption = ('Parameters of event ranked %s by %s %s in the search. The figures below' - ' show the mini-followup data for this event.' % - (args.n_loudest + 1, args.sort_order, args.sort_variable)) + title = "Parameters of event ranked %s" % (args.n_loudest + 1) + caption = ( + "Parameters of event ranked %s by %s %s in the search. The figures below" + " show the mini-followup data for this event." + % (args.n_loudest + 1, args.sort_order, args.sort_variable) + ) elif args.trigger_id is not None: n = args.trigger_id - title = 'Details of trigger' - caption = ('Parameters of event. The figures below show the ' - 'mini-followup data for this event.') + title = "Details of trigger" + caption = ( + "Parameters of event. The figures below show the " + "mini-followup data for this event." + ) else: # It shouldn't be possible to get here! - raise ValueError() + raise ValueError # Make a table for the event information ################################# -hdrs = ["Ranking statistic", - "Inclusive IFAR (yr)", - "Inclusive FAP", - "Exclusive IFAR (yr)", - "Exclusive FAP" - ] +hdrs = [ + "Ranking statistic", + "Inclusive IFAR (yr)", + "Inclusive FAP", + "Exclusive IFAR (yr)", + "Exclusive FAP", +] -dsets = ['stat', 'ifar', 'fap', 'ifar_exc', 'fap_exc'] -formats = ['%5.2f', '%5.2f', '%5.2e', '%5.2f', '%5.2e'] +dsets = ["stat", "ifar", "fap", "ifar_exc", "fap_exc"] +formats = ["%5.2f", "%5.2f", "%5.2e", "%5.2f", "%5.2e"] tbl = [[h, fmt % d[dst][n]] for fmt, dst, h in zip(formats, dsets, hdrs) if dst in d] headers, table = zip(*tbl) headers = list(headers) table = list(table) if args.include_gracedb_link: # Get the time from the single detectors - times = [d['%s/time' % ifo][n] for ifo in f.attrs['ifos'].split(' ')] + times = [d["%s/time" % ifo][n] for ifo in f.attrs["ifos"].split(" ")] times = numpy.array(times) time = numpy.mean(times[times > 0]) gdb_search_link = followup.get_gracedb_search_link(time) @@ -132,25 +169,26 @@ if args.include_gracedb_link: table = numpy.array([table], dtype=str) -html = pycbc.results.dq.redirect_javascript + \ - str(pycbc.results.static_table(table, headers)) +html = pycbc.results.dq.redirect_javascript + str( + pycbc.results.static_table(table, headers) +) # Make a table for the single detector information ############################ idx = {} -ifo_list = f.attrs['ifos'].split(' ') +ifo_list = f.attrs["ifos"].split(" ") for ifo in ifo_list: - idx[ifo] = d['%s/trigger_id' % ifo][n] + idx[ifo] = d["%s/trigger_id" % ifo][n] # Store the single detector trigger files keyed by ifo in a dictionary files = {} for fname in args.single_trigger_files: - f2 = HFile(fname, 'r') + f2 = HFile(fname, "r") ifos = f2.keys() for ifo in ifos: files[ifo] = f2[ifo] -bank = HFile(args.bank_file, 'r') +bank = HFile(args.bank_file, "r") statmapfile = d # Data will store the values that will appear in the resulting single-detector # table. Each entry in data corresponds to each row in the final table and @@ -159,28 +197,28 @@ data = [] row_labels = [] rank_method = pystat.get_statistic_from_opts(args, list(files.keys())) -for ifo in files.keys(): - +for ifo in files: # ignore ifo if coinc didn't participate (only for multi-ifo workflow) - if (statmapfile['%s/time' % ifo][n] == -1.0): + if statmapfile["%s/time" % ifo][n] == -1.0: continue row_labels.append(ifo) d = files[ifo] i = idx[ifo] - tid = d['template_id'][i] - rchisq = d['chisq'][i] / (d['chisq_dof'][i] * 2 - 2) - mchirp = (pycbc.pnutils.mass1_mass2_to_mchirp_eta(bank['mass1'][tid], - bank['mass2'][tid]))[0] + tid = d["template_id"][i] + rchisq = d["chisq"][i] / (d["chisq_dof"][i] * 2 - 2) + mchirp = ( + pycbc.pnutils.mass1_mass2_to_mchirp_eta(bank["mass1"][tid], bank["mass2"][tid]) + )[0] - time = d['end_time'][i] + time = d["end_time"][i] utc = gps_to_utc_datetime(time) trig_dict = { k: numpy.array([d[k][i]]) for k in d.keys() - if not k.endswith('_template') - and k not in ['gating', 'search', 'template_boundaries'] + if not k.endswith("_template") + and k not in ["gating", "search", "template_boundaries"] } # Headers will store the headers that will appear in the table. headers = [] @@ -193,11 +231,11 @@ for ifo in files.keys(): # End times data[-1].append(str(utc)) - data[-1].append('%.3f' % time) + data[-1].append("%.3f" % time) headers.append("UTC End Time") headers.append("GPS End time") - #headers.append("Stat") + # headers.append("Stat") # Determine statistic naming if args.sngl_ranking == "newsnr": sngl_stat_name = "Reweighted SNR" @@ -212,36 +250,36 @@ for ifo in files.keys(): stat = rank_method.get_sngl_ranking(trig_dict) headers.append(sngl_stat_name) - data[-1].append('%5.2f' % stat[0]) + data[-1].append("%5.2f" % stat[0]) # SNR and phase (not showing any single-det stat here) - data[-1].append('%5.2f' % d['snr'][i]) - data[-1].append('%5.2f' % d['coa_phase'][i]) + data[-1].append("%5.2f" % d["snr"][i]) + data[-1].append("%5.2f" % d["coa_phase"][i]) headers.append("ρ") headers.append("Phase") # Signal-glitch discrimators - data[-1].append('%5.2f' % rchisq) - data[-1].append('%i' % d['chisq_dof'][i]) + data[-1].append("%5.2f" % rchisq) + data[-1].append("%i" % d["chisq_dof"][i]) headers.append("χ2r") headers.append("χ2 bins") try: - data[-1].append('%5.2f' % d['sg_chisq'][i]) + data[-1].append("%5.2f" % d["sg_chisq"][i]) headers.append("sg χ2") except: pass try: - data[-1].append('%5.2f' % d['psd_var_val'][i]) + data[-1].append("%5.2f" % d["psd_var_val"][i]) headers.append("PSD var") except: pass # Template parameters - data[-1].append('%5.2f' % bank['mass1'][tid]) - data[-1].append('%5.2f' % bank['mass2'][tid]) - data[-1].append('%5.2f' % mchirp) - data[-1].append('%5.2f' % bank['spin1z'][tid]) - data[-1].append('%5.2f' % bank['spin2z'][tid]) - data[-1].append('%5.2f' % d['template_duration'][i]) + data[-1].append("%5.2f" % bank["mass1"][tid]) + data[-1].append("%5.2f" % bank["mass2"][tid]) + data[-1].append("%5.2f" % mchirp) + data[-1].append("%5.2f" % bank["spin1z"][tid]) + data[-1].append("%5.2f" % bank["spin2z"][tid]) + data[-1].append("%5.2f" % d["template_duration"][i]) headers.append("m1") headers.append("m2") headers.append("Mc") @@ -250,22 +288,25 @@ for ifo in files.keys(): headers.append("Duration") # display eccentricity and relative anomaly if they exist try: - data[-1].append('%5.2f' % bank['eccentricity'][tid]) - data[-1].append('%5.2f' % bank['rel_anomaly'][tid]) + data[-1].append("%5.2f" % bank["eccentricity"][tid]) + data[-1].append("%5.2f" % bank["rel_anomaly"][tid]) headers.append("eccentricity") headers.append("anomaly") except: pass -html += str(pycbc.results.static_table( - data, - headers, - columns_max=args.max_columns, - row_labels=row_labels -)) +html += str( + pycbc.results.static_table( + data, headers, columns_max=args.max_columns, row_labels=row_labels + ) +) ############################################################################### -pycbc.results.save_fig_with_metadata(html, args.output_file, {}, - cmd=' '.join(sys.argv), - title=args.title if args.title else title, - caption=caption) +pycbc.results.save_fig_with_metadata( + html, + args.output_file, + {}, + cmd=" ".join(sys.argv), + title=args.title or title, + caption=caption, +) diff --git a/bin/minifollowups/pycbc_page_injinfo b/bin/minifollowups/pycbc_page_injinfo index 6e01912ec0f..3399b3e1226 100644 --- a/bin/minifollowups/pycbc_page_injinfo +++ b/bin/minifollowups/pycbc_page_injinfo @@ -14,35 +14,44 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Make tables describing a missed injection""" +"""Make tables describing a missed injection""" + import argparse import sys + import numpy +import pycbc.pnutils import pycbc.results +from pycbc import add_common_pycbc_options, init_logging from pycbc import conversions as conv -import pycbc.pnutils -from pycbc import init_logging, add_common_pycbc_options from pycbc.detector import Detector from pycbc.io.hdf import HFile parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--output-file') -parser.add_argument('--injection-file', required=True, - help="The HDF format injection file. Required") -parser.add_argument('--injection-index', type=int, required=True, - help="The index of the injection to print out. Required") -parser.add_argument('--n-nearest', type=int, - help="Optional, used in the title") -parser.add_argument('--max-columns', type=int, - help="Optional, maximum number of columns used for the table") +parser.add_argument("--output-file") +parser.add_argument( + "--injection-file", required=True, help="The HDF format injection file. Required" +) +parser.add_argument( + "--injection-index", + type=int, + required=True, + help="The index of the injection to print out. Required", +) +parser.add_argument("--n-nearest", type=int, help="Optional, used in the title") +parser.add_argument( + "--max-columns", + type=int, + help="Optional, maximum number of columns used for the table", +) args = parser.parse_args() init_logging(args.verbose) -f = HFile(args.injection_file, 'r') +f = HFile(args.injection_file, "r") iidx = args.injection_index # make a table for the coincident information ################################# @@ -51,85 +60,102 @@ headers = [] data = [] labels = { - 'tc': 'End time', - 'dec_chirp_dist': 'Dec. chirp dist (HL)', - 'eff_dist_h' : 'Deff H', - 'eff_dist_l' : 'Deff L', - 'eff_dist_v' : 'Deff V', - 'mass1' : 'm1', - 'mass2' : 'm2', - 'mchirp' : 'Mc', - 'eta' : 'η', - 'ra' : 'RA', - 'dec' : 'Dec', - 'inclination': 'ι', - 'spin1x': 's1x', - 'spin1y': 's1y', - 'spin1z': 's1z', - 'spin2x': 's2x', - 'spin2y': 's2y', - 'spin2z': 's2z', - 'chieff': 'χeff', - 'chip': 'χp', - 'eccentricity': 'eccentricity', - 'rel_anomaly': 'anomaly', + "tc": "End time", + "dec_chirp_dist": "Dec. chirp dist (HL)", + "eff_dist_h": "Deff H", + "eff_dist_l": "Deff L", + "eff_dist_v": "Deff V", + "mass1": "m1", + "mass2": "m2", + "mchirp": "Mc", + "eta": "η", + "ra": "RA", + "dec": "Dec", + "inclination": "ι", + "spin1x": "s1x", + "spin1y": "s1y", + "spin1z": "s1z", + "spin2x": "s2x", + "spin2y": "s2y", + "spin2z": "s2z", + "chieff": "χeff", + "chip": "χp", + "eccentricity": "eccentricity", + "rel_anomaly": "anomaly", } -params += ['tc'] +params += ["tc"] -m1, m2 = f['injections']['mass1'][iidx], f['injections']['mass2'][iidx] -s1x, s2x = f['injections']['spin1x'][iidx], f['injections']['spin2x'][iidx] -s1y, s2y = f['injections']['spin1y'][iidx], f['injections']['spin2y'][iidx] -s1z, s2z = f['injections']['spin1z'][iidx], f['injections']['spin2z'][iidx] +m1, m2 = f["injections"]["mass1"][iidx], f["injections"]["mass2"][iidx] +s1x, s2x = f["injections"]["spin1x"][iidx], f["injections"]["spin2x"][iidx] +s1y, s2y = f["injections"]["spin1y"][iidx], f["injections"]["spin2y"][iidx] +s1z, s2z = f["injections"]["spin1z"][iidx], f["injections"]["spin2z"][iidx] try: - eccentricity, rel_anomaly = f['injections']['eccentricity'][iidx], f['injections']['rel_anomaly'][iidx] + eccentricity, rel_anomaly = ( + f["injections"]["eccentricity"][iidx], + f["injections"]["rel_anomaly"][iidx], + ) except KeyError: eccentricity, rel_anomaly = 0, 0 derived = {} -derived['mchirp'], derived['eta'] = \ - pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) -derived['mtotal'] = conv.mtotal_from_mass1_mass2(m1, m2) -derived['chieff'] = conv.chi_eff(m1, m2, s1z, s2z) -derived['chip'] = conv.chi_p(m1, m2, s1x, s1y, s2x, s2y) - -if 'optimal_snr' in ' '.join(list(f['injections'].keys())): - ifolist = f.attrs['ifos'].split(' ') +derived["mchirp"], derived["eta"] = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) +derived["mtotal"] = conv.mtotal_from_mass1_mass2(m1, m2) +derived["chieff"] = conv.chi_eff(m1, m2, s1z, s2z) +derived["chip"] = conv.chi_p(m1, m2, s1x, s1y, s2x, s2y) + +if "optimal_snr" in " ".join(list(f["injections"].keys())): + ifolist = f.attrs["ifos"].split(" ") for ifo in ifolist: - labels['optimal_snr_%s' % ifo] = 'Opt. SNR %s' % ifo - params += ['optimal_snr_%s' % ifo] + labels["optimal_snr_%s" % ifo] = "Opt. SNR %s" % ifo + params += ["optimal_snr_%s" % ifo] else: eff_dist = {} - for ifo in ['H1', 'L1', 'V1']: + for ifo in ["H1", "L1", "V1"]: eff_dist[ifo] = Detector(ifo).effective_distance( - f['injections/distance'][iidx], - f['injections/ra'][iidx], - f['injections/dec'][iidx], - f['injections/polarization'][iidx], - f['injections/tc'][iidx], - f['injections/inclination'][iidx] + f["injections/distance"][iidx], + f["injections/ra"][iidx], + f["injections/dec"][iidx], + f["injections/polarization"][iidx], + f["injections/tc"][iidx], + f["injections/inclination"][iidx], ) - params += ['dec_chirp_dist', 'eff_dist_h', 'eff_dist_l', 'eff_dist_v'] - dec_dist = max(eff_dist['H1'], eff_dist['L1']) + params += ["dec_chirp_dist", "eff_dist_h", "eff_dist_l", "eff_dist_v"] + dec_dist = max(eff_dist["H1"], eff_dist["L1"]) dec_chirp_dist = pycbc.pnutils.chirp_distance(dec_dist, mchirp) -params += ['mass1', 'mass2', 'mchirp', 'eta', 'ra', 'dec', - 'inclination', 'spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y', - 'spin2z', 'chieff', 'chip','eccentricity', 'rel_anomaly'] +params += [ + "mass1", + "mass2", + "mchirp", + "eta", + "ra", + "dec", + "inclination", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + "chieff", + "chip", + "eccentricity", + "rel_anomaly", +] for p in params: - - if p in f['injections']: - data += ["%.2f" % f['injections'][p][iidx]] - elif p in derived.keys(): - data += [f'{derived[p]:.2f}'] - elif 'eff_dist' in p: - ifo = '%s1' % p.split('_')[-1] + if p in f["injections"]: + data += ["%.2f" % f["injections"][p][iidx]] + elif p in derived: + data += [f"{derived[p]:.2f}"] + elif "eff_dist" in p: + ifo = "%s1" % p.split("_")[-1] data += ["%.2f" % eff_dist[ifo.upper()]] - elif p == 'dec_chirp_dist': + elif p == "dec_chirp_dist": data += ["%.2f" % dec_chirp_dist] else: - print("No data for {}, so skipping".format(p)) + print(f"No data for {p}, so skipping") continue headers += [labels[p]] @@ -137,11 +163,15 @@ for p in params: table = numpy.array([data], dtype=str) html = str(pycbc.results.static_table(table, headers, columns_max=args.max_columns)) -tag = '' +tag = "" if args.n_nearest is not None: - tag = ':%s' % (args.n_nearest + 1) - -pycbc.results.save_fig_with_metadata(html, args.output_file, {}, - cmd = ' '.join(sys.argv), - title = 'Parameters of missed injection' + tag, - caption = "Parameters of this injection") + tag = ":%s" % (args.n_nearest + 1) + +pycbc.results.save_fig_with_metadata( + html, + args.output_file, + {}, + cmd=" ".join(sys.argv), + title="Parameters of missed injection" + tag, + caption="Parameters of this injection", +) diff --git a/bin/minifollowups/pycbc_page_snglinfo b/bin/minifollowups/pycbc_page_snglinfo index 5678dd8a7d0..2ac4226818f 100644 --- a/bin/minifollowups/pycbc_page_snglinfo +++ b/bin/minifollowups/pycbc_page_snglinfo @@ -15,60 +15,90 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Make tables describing a sngl event""" -import sys -import numpy +"""Make tables describing a sngl event""" + import argparse +import sys + import matplotlib -matplotlib.use('Agg') +import numpy -import pycbc.events, pycbc.results, pycbc.pnutils -from pycbc.results import followup +matplotlib.use("Agg") + +import pycbc.events +import pycbc.pnutils +import pycbc.results +from pycbc import add_common_pycbc_options, init_logging from pycbc.events import stat as pystat from pycbc.io import hdf -from pycbc import init_logging, add_common_pycbc_options +from pycbc.results import followup from pycbc.time import gps_to_utc_datetime - parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--single-trigger-file', required=True, - help="HDF format single detector trigger files for the full " - "data run") -parser.add_argument('--bank-file', required=True, - help="HDF format template bank file") -parser.add_argument('--output-file') +parser.add_argument( + "--single-trigger-file", + required=True, + help="HDF format single detector trigger files for the full data run", +) +parser.add_argument("--bank-file", required=True, help="HDF format template bank file") +parser.add_argument("--output-file") # CURRENTLY UNUSED, but may be needed if using foreground censor -#parser.add_argument('--statmap-file', +# parser.add_argument('--statmap-file', # help="The HDF format clustered coincident statmap file containing the " # "result triggers.") -parser.add_argument('--veto-file', - help="The veto file to be used if vetoing triggers. Optional") -parser.add_argument('--veto-segment-name', - help="If using veto file, the name of the segments to use as a veto.") +parser.add_argument( + "--veto-file", help="The veto file to be used if vetoing triggers. Optional" +) +parser.add_argument( + "--veto-segment-name", + help="If using veto file, the name of the segments to use as a veto.", +) parser.add_argument("--instrument", help="Name of ifo (e.g. H1)") trig_args = parser.add_mutually_exclusive_group(required=True) -trig_args.add_argument("--n-loudest", type=int, default=None, +trig_args.add_argument( + "--n-loudest", + type=int, + default=None, help="Examine the n'th loudest trigger (loudest is n=0). Must supply " - "either this or --trigger-id.") -trig_args.add_argument("--trigger-id", type=int, default=None, + "either this or --trigger-id.", +) +trig_args.add_argument( + "--trigger-id", + type=int, + default=None, help="Use the trigger with this ID in the HDF file. Must supply either " - "this option or --n-loudest.") -parser.add_argument('--title', + "this option or --n-loudest.", +) +parser.add_argument( + "--title", help="Title for the produced snippet. If not given, will default to " - "pre-sets according to the options given") -parser.add_argument('--include-summary-page-link', action='store_true', - help="If given, will include a link to the DQ summary page") -parser.add_argument('--include-gracedb-link', action='store_true', + "pre-sets according to the options given", +) +parser.add_argument( + "--include-summary-page-link", + action="store_true", + help="If given, will include a link to the DQ summary page", +) +parser.add_argument( + "--include-gracedb-link", + action="store_true", help="If given, will provide a link to search GraceDB for events " - "within a 3s window around the coincidence time.") -parser.add_argument('--significance-file', + "within a 3s window around the coincidence time.", +) +parser.add_argument( + "--significance-file", help="If given, will search for this trigger's id in the file to see if " - "stat and p_astro values exists for this trigger.") -parser.add_argument('--max-columns', type=int, - help="Optional. Set a maximum number of columns to be used in the output table") -pystat.insert_statistic_option_group(parser, - default_ranking_statistic='single_ranking_only') + "stat and p_astro values exists for this trigger.", +) +parser.add_argument( + "--max-columns", + type=int, + help="Optional. Set a maximum number of columns to be used in the output table", +) +pystat.insert_statistic_option_group( + parser, default_ranking_statistic="single_ranking_only" +) args = parser.parse_args() @@ -100,8 +130,7 @@ elif args.n_loudest is not None: # Cluster by a ranking statistic and retain only the loudest n clustered # triggers sngl_file.mask_to_n_loudest_clustered_events( - rank_method, - n_loudest=args.n_loudest+1 + rank_method, n_loudest=args.n_loudest + 1 ) else: raise ValueError("Must give --n-loudest or --trigger-id.") @@ -137,15 +166,15 @@ if args.include_summary_page_link: # End times data[0].append(str(utc)) -data[0].append('%.3f' % time) +data[0].append("%.3f" % time) headers.append("UTC") headers.append("End time") # SNR and statistic headers.append("ρ") -data[0].append('%5.2f' % sngl_file.snr[0]) +data[0].append("%5.2f" % sngl_file.snr[0]) headers.append("Phase") -data[0].append('%5.2f' % sngl_file.get_column('coa_phase')[0]) +data[0].append("%5.2f" % sngl_file.get_column("coa_phase")[0]) # Determine statistic naming if args.sngl_ranking == "newsnr": sngl_stat_name = "Reweighted SNR" @@ -163,37 +192,35 @@ if args.ranking_statistic in ["quadsum", "single_ranking_only"]: stat_name_long = sngl_stat_name else: # Name would be too long - just call it ranking statistic - stat_name = 'Ranking Statistic' - stat_name_long = ' with '.join( - [args.ranking_statistic, args.sngl_ranking] - ) + stat_name = "Ranking Statistic" + stat_name_long = " with ".join([args.ranking_statistic, args.sngl_ranking]) headers.append(stat_name) -data[0].append('%5.2f' % stat[0]) +data[0].append("%5.2f" % stat[0]) # Signal-glitch discrimators -data[0].append('%5.2f' % sngl_file.rchisq[0]) -data[0].append('%i' % sngl_file.get_column('chisq_dof')[0]) +data[0].append("%5.2f" % sngl_file.rchisq[0]) +data[0].append("%i" % sngl_file.get_column("chisq_dof")[0]) headers.append("χ2r") headers.append("χ2 bins") try: - data[0].append('%5.2f' % sngl_file.sgchisq[0]) + data[0].append("%5.2f" % sngl_file.sgchisq[0]) headers.append("sgχ2") except: pass try: - data[0].append('%5.2f' % sngl_file.psd_var_val[0]) + data[0].append("%5.2f" % sngl_file.psd_var_val[0]) headers.append("PSD var") except: pass # Template parameters -data[0].append('%5.2f' % sngl_file.mass1[0]) -data[0].append('%5.2f' % sngl_file.mass2[0]) -data[0].append('%5.2f' % sngl_file.mchirp[0]) -data[0].append('%5.2f' % sngl_file.spin1z[0]) -data[0].append('%5.2f' % sngl_file.spin2z[0]) -data[0].append('%5.2f' % sngl_file.template_duration[0]) +data[0].append("%5.2f" % sngl_file.mass1[0]) +data[0].append("%5.2f" % sngl_file.mass2[0]) +data[0].append("%5.2f" % sngl_file.mchirp[0]) +data[0].append("%5.2f" % sngl_file.spin1z[0]) +data[0].append("%5.2f" % sngl_file.spin2z[0]) +data[0].append("%5.2f" % sngl_file.template_duration[0]) headers.append("m1") headers.append("m2") headers.append("Mc") @@ -202,18 +229,18 @@ headers.append("s2z") headers.append("Duration") # display eccentricity and relative anomaly if they exist try: - data[0].append('%5.2f' % sngl_file.eccentricity[0]) - data[0].append('%5.2f' % sngl_file.rel_anomaly[0]) + data[0].append("%5.2f" % sngl_file.eccentricity[0]) + data[0].append("%5.2f" % sngl_file.rel_anomaly[0]) headers.append("eccentricity") headers.append("anomaly") except KeyError: pass if args.significance_file and not args.n_loudest: - with hdf.HFile(args.significance_file, 'r') as sig_f: - trigger_ids = sig_f['foreground'][args.instrument]['trigger_id'][:] - sngl_stat = sig_f['foreground/stat'][:] - sngl_pastro = sig_f['foreground/p_astro_exc'][:] + with hdf.HFile(args.significance_file, "r") as sig_f: + trigger_ids = sig_f["foreground"][args.instrument]["trigger_id"][:] + sngl_stat = sig_f["foreground/stat"][:] + sngl_pastro = sig_f["foreground/p_astro_exc"][:] if trig_id in trigger_ids: trig_idx = numpy.nonzero(trigger_ids == trig_id)[0][0] @@ -229,20 +256,29 @@ if args.include_gracedb_link: headers.append("GraceDB Search Link") data[0].append(gdb_search_link) -html = pycbc.results.dq.redirect_javascript + \ - str(pycbc.results.static_table(data, headers, row_labels=row_labels, columns_max=args.max_columns)) +html = pycbc.results.dq.redirect_javascript + str( + pycbc.results.static_table( + data, headers, row_labels=row_labels, columns_max=args.max_columns + ) +) ############################################################################### # Set up default titles and the captions for the file if args.n_loudest: - title = 'Parameters of single-detector event ranked %s' \ - % (args.n_loudest + 1) - caption = 'Parameters of the single-detector event ranked number %s by %s. The figures below show the mini-followup data for this event.' % (args.n_loudest + 1, stat_name_long) + title = "Parameters of single-detector event ranked %s" % (args.n_loudest + 1) + caption = ( + "Parameters of the single-detector event ranked number %s by %s. The figures below show the mini-followup data for this event." + % (args.n_loudest + 1, stat_name_long) + ) else: - title = 'Parameters of single-detector event' - caption = 'Parameters of the single-detector event. The figures below show the mini-followup data for this event.' + title = "Parameters of single-detector event" + caption = "Parameters of the single-detector event. The figures below show the mini-followup data for this event." -pycbc.results.save_fig_with_metadata(html, args.output_file, {}, - cmd = ' '.join(sys.argv), - title = args.title if args.title else title, - caption = caption) +pycbc.results.save_fig_with_metadata( + html, + args.output_file, + {}, + cmd=" ".join(sys.argv), + title=args.title or title, + caption=caption, +) diff --git a/bin/minifollowups/pycbc_plot_chigram b/bin/minifollowups/pycbc_plot_chigram index 5f31c2d8ec9..b25a169376b 100644 --- a/bin/minifollowups/pycbc_plot_chigram +++ b/bin/minifollowups/pycbc_plot_chigram @@ -1,71 +1,77 @@ #!/bin/env python -import numpy import argparse import sys + import matplotlib -matplotlib.use('Agg') +import numpy + +matplotlib.use("Agg") from matplotlib import pyplot as plt -from pycbc import init_logging, add_common_pycbc_options -import pycbc.types import pycbc.results +import pycbc.types +from pycbc import add_common_pycbc_options, init_logging from pycbc.io.hdf import HFile -parser=argparse.ArgumentParser() +parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--single-template-file', - help="HDF output of pycbc_single_template") -parser.add_argument('--central-time', type=float, - help="Time to center the plot, optional", - default=0) -parser.add_argument('--window', type=float, - help="Time around central time to plot, optional. " - "Used with the central time option") -parser.add_argument('--plot-type', choices=['snr', 'chisq'], default='chisq') -parser.add_argument('--output-file') +parser.add_argument( + "--single-template-file", help="HDF output of pycbc_single_template" +) +parser.add_argument( + "--central-time", type=float, help="Time to center the plot, optional", default=0 +) +parser.add_argument( + "--window", + type=float, + help="Time around central time to plot, optional. " + "Used with the central time option", +) +parser.add_argument("--plot-type", choices=["snr", "chisq"], default="chisq") +parser.add_argument("--output-file") args = parser.parse_args() init_logging(args.verbose) -f = HFile(args.single_template_file, 'r') -y = f['chisq_boundaries'][:] +f = HFile(args.single_template_file, "r") +y = f["chisq_boundaries"][:] fig = plt.figure() ax = plt.gca() -snr = pycbc.types.load_timeseries(args.single_template_file, group='snr') -chisq = pycbc.types.load_timeseries(args.single_template_file, group='chisq') +snr = pycbc.types.load_timeseries(args.single_template_file, group="snr") +chisq = pycbc.types.load_timeseries(args.single_template_file, group="chisq") s = None -for i in range(len(f['chisq_bins'].keys())): - ts = pycbc.types.load_timeseries(args.single_template_file, group='chisq_bins/%s' % i) +for i in range(len(f["chisq_bins"].keys())): + ts = pycbc.types.load_timeseries( + args.single_template_file, group="chisq_bins/%s" % i + ) delta_t = ts.delta_t - st = ts.sample_times.numpy() - - - if args.plot_type == 'chisq': - ts = ts.squared_norm() - snr.squared_norm() / (len(y) -1) - elif args.plot_type == 'snr': + st = ts.sample_times.numpy() + + if args.plot_type == "chisq": + ts = ts.squared_norm() - snr.squared_norm() / (len(y) - 1) + elif args.plot_type == "snr": ts = ts.squared_norm() - - + ts = ts.numpy() x = numpy.append(st - delta_t / 2.0, [st[-1] + delta_t / 2.0]) - args.central_time - l = y[i:i+2] - ax.pcolorfast(x, numpy.array([i, i+1]), ts.reshape(1, len(ts))) + l = y[i : i + 2] + ax.pcolorfast(x, numpy.array([i, i + 1]), ts.reshape(1, len(ts))) -plt.ylabel('Frequency (Hz)') +plt.ylabel("Frequency (Hz)") -xlabel = 'Time (s)' +xlabel = "Time (s)" if args.central_time: - xlabel += ' - %.2f' % args.central_time + xlabel += " - %.2f" % args.central_time if args.window: plt.xlim(xmin=-args.window, xmax=args.window) plt.xlabel(xlabel) c = plt.colorbar(ax.get_children()[2], ax=ax) -if args.plot_type == 'chisq': +if args.plot_type == "chisq": c.set_label("$\\rho_l^2 - \\rho^2/p$") else: c.set_label("$\\rho_l^2$") @@ -73,10 +79,12 @@ else: # Set the frequency label fig.canvas.draw() labels = [item.get_text() for item in ax.get_yticklabels()] -ax.set_yticklabels(['%.0f' % y[int(label)] for label in labels]) +ax.set_yticklabels(["%.0f" % y[int(label)] for label in labels]) -pycbc.results.save_fig_with_metadata(fig, args.output_file, - cmd=' '.join(sys.argv), - title="chisq timeseries for each bin", - caption="Plot of the time series of each chisq bin") - +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + cmd=" ".join(sys.argv), + title="chisq timeseries for each bin", + caption="Plot of the time series of each chisq bin", +) diff --git a/bin/minifollowups/pycbc_plot_trigger_timeseries b/bin/minifollowups/pycbc_plot_trigger_timeseries index a248a315316..2ef4531e2de 100644 --- a/bin/minifollowups/pycbc_plot_trigger_timeseries +++ b/bin/minifollowups/pycbc_plot_trigger_timeseries @@ -15,41 +15,65 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plot the single detector trigger timeseries """ +"""Plot the single detector trigger timeseries""" + import argparse import logging import sys + import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt + +matplotlib.use("Agg") import numpy +from matplotlib import pyplot as plt -from pycbc import init_logging, add_common_pycbc_options import pycbc.results -from pycbc.types import MultiDetOptionAction +from pycbc import add_common_pycbc_options, init_logging from pycbc.events import ranking from pycbc.io import HFile, SingleDetTriggers +from pycbc.types import MultiDetOptionAction parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--single-trigger-files', nargs='+', - action=MultiDetOptionAction, metavar="IFO:FILE", +parser.add_argument( + "--single-trigger-files", + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:FILE", help="The HDF format single detector merged trigger files, in " - "multi-ifo argument format, H1:file1.hdf L1:file2.hdf, etc") -parser.add_argument('--window', type=float, default=10, - help="Time in seconds around the coincident trigger to plot") -parser.add_argument('--times', nargs='+', type=float, - action=MultiDetOptionAction, metavar="IFO:GPS_TIME", + "multi-ifo argument format, H1:file1.hdf L1:file2.hdf, etc", +) +parser.add_argument( + "--window", + type=float, + default=10, + help="Time in seconds around the coincident trigger to plot", +) +parser.add_argument( + "--times", + nargs="+", + type=float, + action=MultiDetOptionAction, + metavar="IFO:GPS_TIME", help="The gps times to plot around in multi-ifo argument format, " - "H1:132341323 L1:132423422") -parser.add_argument('--special-trigger-ids', nargs='+', type=int, - action=MultiDetOptionAction, metavar="IFO:GPS_TIME", - help="The set of special trigger ids to plot a star at") -parser.add_argument('--plot-type', - choices=ranking.sngls_ranking_function_dict, default='snr', - help="Which single-detector ranking statistic to plot.") -parser.add_argument('--output-file') -parser.add_argument('--log-y-axis', action='store_true') + "H1:132341323 L1:132423422", +) +parser.add_argument( + "--special-trigger-ids", + nargs="+", + type=int, + action=MultiDetOptionAction, + metavar="IFO:GPS_TIME", + help="The set of special trigger ids to plot a star at", +) +parser.add_argument( + "--plot-type", + choices=ranking.sngls_ranking_function_dict, + default="snr", + help="Which single-detector ranking statistic to plot.", +) +parser.add_argument("--output-file") +parser.add_argument("--log-y-axis", action="store_true") args = parser.parse_args() init_logging(args.verbose) @@ -64,30 +88,29 @@ for ifo in args.single_trigger_files.keys(): t = args.times[ifo] # Identify trigger idxs within window of trigger time - with HFile(args.single_trigger_files[ifo], 'r') as data: + with HFile(args.single_trigger_files[ifo], "r") as data: idx, _ = data.select( lambda endtime: abs(endtime - t) < args.window, - 'end_time', + "end_time", group=ifo, return_data=False, ) - data_mask = numpy.zeros(data[ifo]['snr'].size, dtype=bool) + data_mask = numpy.zeros(data[ifo]["snr"].size, dtype=bool) data_mask[idx] = True if not len(idx): # No triggers in this window, add to the legend and continue # Make sure it isnt on the plot - plt.scatter(-2 * args.window, 0, - color=pycbc.results.ifo_color(ifo), - marker='x', - label=ifo) + plt.scatter( + -2 * args.window, + 0, + color=pycbc.results.ifo_color(ifo), + marker="x", + label=ifo, + ) continue - trigs = SingleDetTriggers( - args.single_trigger_files[ifo], - ifo, - premask=data_mask - ) + trigs = SingleDetTriggers(args.single_trigger_files[ifo], ifo, premask=data_mask) any_data = True logging.info("Keeping %d triggers in the window", len(idx)) @@ -95,9 +118,13 @@ for ifo in args.single_trigger_files.keys(): logging.info("Getting %s", args.plot_type) rank = ranking.get_sngls_ranking_from_trigs(trigs, args.plot_type) - plt.scatter(trigs['end_time'] - t, rank, - color=pycbc.results.ifo_color(ifo), marker='x', - label=ifo) + plt.scatter( + trigs["end_time"] - t, + rank, + color=pycbc.results.ifo_color(ifo), + marker="x", + label=ifo, + ) min_rank = min(min_rank, rank.min()) @@ -106,21 +133,25 @@ for ifo in args.single_trigger_files.keys(): if special_idx == None: # No special trigger for this ifo continue elif special_idx not in idx: - logging.info("IDX %d not in kept list", - args.special_trigger_ids[ifo]) + logging.info("IDX %d not in kept list", args.special_trigger_ids[ifo]) continue special_red_idx = numpy.where(idx == special_idx)[0] - plt.scatter(trigs.trigs_f[f'{ifo}/end_time'][special_idx] - t, - rank[special_red_idx], marker='*', s=50, color='yellow') + plt.scatter( + trigs.trigs_f[f"{ifo}/end_time"][special_idx] - t, + rank[special_red_idx], + marker="*", + s=50, + color="yellow", + ) if args.log_y_axis and any_data: - plt.yscale('log') + plt.yscale("log") if not numpy.isinf(min_rank): plt.ylim(ymin=min_rank) -plt.xlabel('time (s)') +plt.xlabel("time (s)") plt.ylabel(args.plot_type) plt.xlim(xmin=-args.window, xmax=args.window) @@ -128,13 +159,15 @@ plt.legend() plt.grid() logging.info("Saving figure") -pycbc.results.save_fig_with_metadata(fig, args.output_file, - cmd = ' '.join(sys.argv), - title = 'Single Detector Trigger Timeseries (%s)' % args.plot_type, - caption = 'Time series showing the single-detector triggers ' - 'centered around the time of the trigger of interest. ' - 'Triggers with ranking 1 have been downweighted beyond ' - 'consideration, but may still form insignificant ' - 'events.', - ) +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + cmd=" ".join(sys.argv), + title="Single Detector Trigger Timeseries (%s)" % args.plot_type, + caption="Time series showing the single-detector triggers " + "centered around the time of the trigger of interest. " + "Triggers with ranking 1 have been downweighted beyond " + "consideration, but may still form insignificant " + "events.", +) logging.info("Done!") diff --git a/bin/minifollowups/pycbc_single_template_plot b/bin/minifollowups/pycbc_single_template_plot index cb3a74a00d3..db0d26cc56c 100644 --- a/bin/minifollowups/pycbc_single_template_plot +++ b/bin/minifollowups/pycbc_single_template_plot @@ -14,13 +14,15 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plot the output of pycbc_single_template """ +"""Plot the output of pycbc_single_template""" import argparse import sys -import numpy + import matplotlib -matplotlib.use('Agg') +import numpy + +matplotlib.use("Agg") from matplotlib import pyplot as plt import pycbc.results @@ -29,86 +31,92 @@ from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--single-template-file', required=True, +parser.add_argument( + "--single-template-file", + required=True, help="HDF file containing the SNR and CHISQ timeseries. " - " The output of pycbc_single_template") -parser.add_argument('--window', type=float, required=True, - help="The time in seconds to plot") -parser.add_argument('--event-time', type=float, - help='GPS time to use as the center of the plot') -parser.add_argument('--output-file', required=True) -parser.add_argument('--log-y', action='store_true', - help='Use log scale for y-axis') -parser.add_argument('--plot-title', - help="If given, use this as the plot title") -parser.add_argument('--plot-caption', - help="If given, use this as the plot caption") + " The output of pycbc_single_template", +) +parser.add_argument( + "--window", type=float, required=True, help="The time in seconds to plot" +) +parser.add_argument( + "--event-time", type=float, help="GPS time to use as the center of the plot" +) +parser.add_argument("--output-file", required=True) +parser.add_argument("--log-y", action="store_true", help="Use log scale for y-axis") +parser.add_argument("--plot-title", help="If given, use this as the plot title") +parser.add_argument("--plot-caption", help="If given, use this as the plot caption") args = parser.parse_args() pycbc.init_logging(args.verbose) # The event is chosen from the input of pycbc_single template, so we # just extract the parameters here -f = HFile(args.single_template_file, 'r') +f = HFile(args.single_template_file, "r") try: - delta_t = f['snr'].attrs['delta_t'] - start_time = f['snr'].attrs['start_time'] + delta_t = f["snr"].attrs["delta_t"] + start_time = f["snr"].attrs["start_time"] except: - plt.text(0.5, 0.5, 'no triggers found') + plt.text(0.5, 0.5, "no triggers found") plt.savefig(args.output_file) sys.exit() if args.event_time is not None: time = args.event_time -elif 'event_time' in f.attrs: - time = f.attrs['event_time'] +elif "event_time" in f.attrs: + time = f.attrs["event_time"] else: - raise ValueError('Event time is neither specified nor found in the HDF file') + raise ValueError("Event time is neither specified nor found in the HDF file") -ifo = f.attrs['ifo'] +ifo = f.attrs["ifo"] center = int((time - start_time) / delta_t) # Determine where in the timeseries we need to plot left = center - int(args.window / delta_t) right = center + int(args.window / delta_t) -left = max(left,0) -right = min(right,len(f['snr'])) +left = max(left, 0) +right = min(right, len(f["snr"])) -snr = abs(f['snr'][left:right][:]) -chisq = f['chisq'][left:right][:] +snr = abs(f["snr"][left:right][:]) +chisq = f["chisq"][left:right][:] rang = (numpy.arange(0, len(snr), 1) - (center - left)) * delta_t newsnr = ranking.newsnr(snr, chisq) fig, ax1 = plt.subplots() -ax1.plot(rang, snr, color='blue', label='SNR') -ax1.plot(rang, newsnr, color='purple', label='NewSNR') -ax1.set_ylabel('SNR') +ax1.plot(rang, snr, color="blue", label="SNR") +ax1.plot(rang, newsnr, color="purple", label="NewSNR") +ax1.set_ylabel("SNR") if args.log_y: - ax1.set_yscale('log') + ax1.set_yscale("log") ax1.legend(loc="upper left") -ax1.set_xlabel('Time - %.3f (s)' % time) +ax1.set_xlabel("Time - %.3f (s)" % time) ax2 = ax1.twinx() -ax2.plot(rang, chisq, color='green', label='$\\chi^2_{\\rm r}$') -ax2.set_ylabel('Reduced $\\chi^2$') +ax2.plot(rang, chisq, color="green", label="$\\chi^2_{\\rm r}$") +ax2.set_ylabel("Reduced $\\chi^2$") if args.log_y: - ax2.set_yscale('log') + ax2.set_yscale("log") ax2.legend(loc="upper right") if args.plot_title is None: - args.plot_title = '%s: SNR and chi^2 time series' % ifo + args.plot_title = "%s: SNR and chi^2 time series" % ifo if args.plot_caption is None: - args.plot_caption = '' + args.plot_caption = "" -if 'command_line' in f.attrs: - cmd = f.attrs['command_line'] + '\\n\\n' + ' '.join(sys.argv) +if "command_line" in f.attrs: + cmd = f.attrs["command_line"] + "\\n\\n" + " ".join(sys.argv) else: - cmd = ' '.join(sys.argv) - -pycbc.results.save_fig_with_metadata(fig, args.output_file, - cmd=cmd, fig_kwds={'dpi': 150}, - title=args.plot_title, - caption=args.plot_caption) + cmd = " ".join(sys.argv) + +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + cmd=cmd, + fig_kwds={"dpi": 150}, + title=args.plot_title, + caption=args.plot_caption, +) diff --git a/bin/minifollowups/pycbc_sngl_minifollowup b/bin/minifollowups/pycbc_sngl_minifollowup index 00fe4727143..ffb357bd427 100644 --- a/bin/minifollowups/pycbc_sngl_minifollowup +++ b/bin/minifollowups/pycbc_sngl_minifollowup @@ -17,82 +17,117 @@ """ Followup single-detector triggers which do not contribute to foreground events """ -import os + import argparse import logging -import numpy +import os +import numpy from igwn_ligolw import ligolw from igwn_ligolw import utils as ligolw_utils -from pycbc import init_logging, add_common_pycbc_options +import pycbc.events +import pycbc.workflow as wf +import pycbc.workflow.minifollowups as mini +from pycbc import add_common_pycbc_options, init_logging +from pycbc.events import select_segments_by_definer, stat, veto +from pycbc.io import hdf from pycbc.results import layout from pycbc.types.optparse import MultiDetOptionAction -from pycbc.events import select_segments_by_definer -import pycbc.workflow.minifollowups as mini -import pycbc.workflow as wf -import pycbc.events from pycbc.workflow.core import resolve_url_to_file -from pycbc.events import stat, veto -from pycbc.io import hdf parser = argparse.ArgumentParser(description=__doc__[1:]) add_common_pycbc_options(parser) -parser.add_argument('--bank-file', - help="HDF format template bank file") -parser.add_argument('--single-detector-file', - help="HDF format merged single detector trigger files") -parser.add_argument('--instrument', help="Name of interferometer e.g. H1") -parser.add_argument('--foreground-censor-file', - help="The censor file to be used if vetoing triggers " - "in the foreground of the search (optional).") -parser.add_argument('--foreground-segment-name', - help="If using foreground censor file must also provide " - "the name of the segment to use as a veto.") -parser.add_argument('--veto-file', - help="The veto file to be used if vetoing triggers " - "(optional).") -parser.add_argument('--veto-segment-name', - help="If using veto file must also provide the name of " - "the segment to use as a veto.") -parser.add_argument("--gating-veto-windows", nargs='+', - action=MultiDetOptionAction, - help="Seconds to be vetoed before and after the central time " - "of each gate. Given as detector-values pairs, e.g. " - "H1:-1,2.5 L1:-1,2.5 V1:0,0") -parser.add_argument('--inspiral-segments', - help="xml segment file containing the inspiral analysis " - "times") -parser.add_argument('--inspiral-data-read-name', - help="Name of inspiral segmentlist containing data read in " - "by each analysis job.") -parser.add_argument('--inspiral-data-analyzed-name', - help="Name of inspiral segmentlist containing data " - "analyzed by each analysis job.") -parser.add_argument('--min-sngl-ranking', type=float, default=6.5, - help="Minimum sngl-ranking to consider for loudest " - "triggers. Useful for efficiency savings. " - "Default=6.5.") -parser.add_argument('--non-coinc-time-only', action='store_true', - help="If given remove (veto) single-detector triggers " - "that occur during a time when at least one other " - "instrument is taking science data.") -parser.add_argument('--vetoed-time-only', action='store_true', - help="If given, only report on single-detector triggers " - "that occur during vetoed times.") -parser.add_argument('--minimum-duration', default=None, type=float, - help="If given only consider single-detector triggers " - "with template duration larger than this.") -parser.add_argument('--maximum-duration', default=None, type=float, - help="If given only consider single-detector triggers " - "with template duration smaller than this.") -parser.add_argument('--cluster-window', type=float, default=10, - help="Window (seconds) over which to cluster triggers " - "when finding the loudest-ranked. Default=10") +parser.add_argument("--bank-file", help="HDF format template bank file") +parser.add_argument( + "--single-detector-file", help="HDF format merged single detector trigger files" +) +parser.add_argument("--instrument", help="Name of interferometer e.g. H1") +parser.add_argument( + "--foreground-censor-file", + help="The censor file to be used if vetoing triggers " + "in the foreground of the search (optional).", +) +parser.add_argument( + "--foreground-segment-name", + help="If using foreground censor file must also provide " + "the name of the segment to use as a veto.", +) +parser.add_argument( + "--veto-file", help="The veto file to be used if vetoing triggers (optional)." +) +parser.add_argument( + "--veto-segment-name", + help="If using veto file must also provide the name of " + "the segment to use as a veto.", +) +parser.add_argument( + "--gating-veto-windows", + nargs="+", + action=MultiDetOptionAction, + help="Seconds to be vetoed before and after the central time " + "of each gate. Given as detector-values pairs, e.g. " + "H1:-1,2.5 L1:-1,2.5 V1:0,0", +) +parser.add_argument( + "--inspiral-segments", + help="xml segment file containing the inspiral analysis times", +) +parser.add_argument( + "--inspiral-data-read-name", + help="Name of inspiral segmentlist containing data read in by each analysis job.", +) +parser.add_argument( + "--inspiral-data-analyzed-name", + help="Name of inspiral segmentlist containing data analyzed by each analysis job.", +) +parser.add_argument( + "--min-sngl-ranking", + type=float, + default=6.5, + help="Minimum sngl-ranking to consider for loudest " + "triggers. Useful for efficiency savings. " + "Default=6.5.", +) +parser.add_argument( + "--non-coinc-time-only", + action="store_true", + help="If given remove (veto) single-detector triggers " + "that occur during a time when at least one other " + "instrument is taking science data.", +) +parser.add_argument( + "--vetoed-time-only", + action="store_true", + help="If given, only report on single-detector triggers " + "that occur during vetoed times.", +) +parser.add_argument( + "--minimum-duration", + default=None, + type=float, + help="If given only consider single-detector triggers " + "with template duration larger than this.", +) +parser.add_argument( + "--maximum-duration", + default=None, + type=float, + help="If given only consider single-detector triggers " + "with template duration smaller than this.", +) +parser.add_argument( + "--cluster-window", + type=float, + default=10, + help="Window (seconds) over which to cluster triggers " + "when finding the loudest-ranked. Default=10", +) wf.add_workflow_command_line_group(parser) wf.add_workflow_settings_cli(parser, include_subdax_opts=True) -stat.insert_statistic_option_group(parser, - default_ranking_statistic='single_ranking_only') +stat.insert_statistic_option_group( + parser, default_ranking_statistic="single_ranking_only" +) args = parser.parse_args() # Default logging level is info: --verbose adds to this @@ -109,32 +144,33 @@ layouts = [] tmpltbank_file = resolve_url_to_file(os.path.abspath(args.bank_file)) sngl_file = resolve_url_to_file( - os.path.abspath(args.single_detector_file), - attrs={'ifos': args.instrument} + os.path.abspath(args.single_detector_file), attrs={"ifos": args.instrument} ) # Flatten the statistic_files option: statfiles = [] for f in sum(args.statistic_files, []): statfiles.append(resolve_url_to_file(os.path.abspath(f))) -statfiles = wf.FileList(statfiles) if statfiles is not [] else None +statfiles = wf.FileList(statfiles) if statfiles != [] else None if args.veto_file is not None: veto_file = resolve_url_to_file( - os.path.abspath(args.veto_file), - attrs={'ifos': args.instrument} + os.path.abspath(args.veto_file), attrs={"ifos": args.instrument} ) else: veto_file = None insp_segs = resolve_url_to_file(os.path.abspath(args.inspiral_segments)) -insp_data_seglists = select_segments_by_definer\ - (args.inspiral_segments, segment_name=args.inspiral_data_read_name, - ifo=args.instrument) +insp_data_seglists = select_segments_by_definer( + args.inspiral_segments, + segment_name=args.inspiral_data_read_name, + ifo=args.instrument, +) insp_data_seglists.coalesce() -num_events = int(workflow.cp.get_opt_tags('workflow-sngl_minifollowups', - 'num-sngl-events', '')) +num_events = int( + workflow.cp.get_opt_tags("workflow-sngl_minifollowups", "num-sngl-events", "") +) trigs = hdf.SingleDetTriggers( args.single_detector_file, @@ -143,48 +179,47 @@ trigs = hdf.SingleDetTriggers( veto_file=args.foreground_censor_file, segment_name=args.foreground_segment_name, filter_rank=args.sngl_ranking, - filter_threshold=args.min_sngl_ranking + filter_threshold=args.min_sngl_ranking, ) # Include gating vetoes if args.gating_veto_windows: logging.info("Getting gating vetoes") - gating_veto = args.gating_veto_windows[args.instrument].split(',') + gating_veto = args.gating_veto_windows[args.instrument].split(",") gveto_before = float(gating_veto[0]) gveto_after = float(gating_veto[1]) if gveto_before > 0 or gveto_after < 0: - raise ValueError("Gating veto window values must be negative before " - "gates and positive after gates.") + raise ValueError( + "Gating veto window values must be negative before " + "gates and positive after gates." + ) if not (gveto_before == 0 and gveto_after == 0): - gate_group = f[args.instrument + '/gating/'] - autogate_times = numpy.unique(gate_group['auto/time'][:]) - if 'file' in gate_group: - detgate_times = gate_group['file/time'][:] + gate_group = f[args.instrument + "/gating/"] + autogate_times = numpy.unique(gate_group["auto/time"][:]) + if "file" in gate_group: + detgate_times = gate_group["file/time"][:] else: detgate_times = [] gate_times = numpy.concatenate((autogate_times, detgate_times)) gveto_idx = veto.indices_within_times( - trigs.end_time, - gate_times + gveto_before, - gate_times + gveto_after + trigs.end_time, gate_times + gveto_before, gate_times + gveto_after ) - logging.info('%i triggers in gating vetoes', gveto_idx.size) + logging.info("%i triggers in gating vetoes", gveto_idx.size) else: gveto_idx = numpy.array([], dtype=numpy.uint64) if args.veto_file: - logging.info('Getting file vetoes') + logging.info("Getting file vetoes") # veto_mask is an array of indices into the trigger arrays # giving the surviving triggers veto_file_idx, _ = pycbc.events.veto.indices_within_segments( trigs.end_time, [args.veto_file], ifo=args.instrument, - segment_name=args.veto_segment_name + segment_name=args.veto_segment_name, ) - logging.info('%i triggers in file-vetoed segments', - veto_file_idx.size) + logging.info("%i triggers in file-vetoed segments", veto_file_idx.size) else: veto_file_idx = numpy.array([], dtype=numpy.uint64) @@ -207,35 +242,37 @@ elif args.vetoed_time_only and vetoed_idx.size == 0: if args.non_coinc_time_only: from pycbc.io.ligolw import LIGOLWContentHandler as h - segs_doc = ligolw_utils.load_filename(args.inspiral_segments, - contenthandler=h) - seg_def_table = ligolw.Table.get_table(segs_doc, 'segment_definer') - def_ifos = seg_def_table.getColumnByName('ifos') + segs_doc = ligolw_utils.load_filename(args.inspiral_segments, contenthandler=h) + seg_def_table = ligolw.Table.get_table(segs_doc, "segment_definer") + def_ifos = seg_def_table.getColumnByName("ifos") def_ifos = [str(ifo) for ifo in def_ifos] ifo_list = list(set(def_ifos)) ifo_list.remove(args.instrument) for ifo in ifo_list: curr_veto_mask, segs = pycbc.events.veto.indices_outside_segments( - trigs.end_time, [args.inspiral_segments], - ifo=ifo, segment_name=args.inspiral_data_analyzed_name) + trigs.end_time, + [args.inspiral_segments], + ifo=ifo, + segment_name=args.inspiral_data_analyzed_name, + ) curr_veto_mask.sort() trigs.apply_mask(curr_veto_mask) if args.minimum_duration is not None: - logging.info('applying minimum duration') + logging.info("applying minimum duration") durations = trigs.template_duration lgc_mask = durations > args.minimum_duration trigs.apply_mask(lgc_mask) - logging.info('remaining triggers: %s', trigs.mask.sum()) + logging.info("remaining triggers: %s", trigs.mask.sum()) if args.maximum_duration is not None: - logging.info('applying maximum duration') + logging.info("applying maximum duration") durations = trigs.template_duration lgc_mask = durations < args.maximum_duration trigs.apply_mask(lgc_mask) - logging.info('remaining triggers: %s', trigs.mask.sum()) + logging.info("remaining triggers: %s", trigs.mask.sum()) -logging.info('Finding loudest clustered events') +logging.info("Finding loudest clustered events") rank_method = stat.get_statistic_from_opts(args, [args.instrument]) trigs.mask_to_n_loudest_clustered_events( @@ -255,73 +292,102 @@ trig_stat = trigs.stat # loop over number of loudest events to be followed up order = trig_stat.argsort()[::-1] for rank, num_event in enumerate(order): - logging.info('Processing event: %s', rank) + logging.info("Processing event: %s", rank) files = wf.FileList([]) time = times[num_event] - ifo_time = '%s:%s' %(args.instrument, str(time)) + ifo_time = "%s:%s" % (args.instrument, str(time)) tid = trigger_ids[num_event] - ifo_tid = '%s:%s' %(args.instrument, str(tid)) - - layouts += (mini.make_sngl_ifo(workflow, sngl_file, tmpltbank_file, - tid, args.output_dir, args.instrument, - statfiles=statfiles, - tags=args.tags + [str(rank)]),) - files += mini.make_trigger_timeseries(workflow, [sngl_file], - ifo_time, args.output_dir, special_tids=ifo_tid, - tags=args.tags + [str(rank)]) + ifo_tid = "%s:%s" % (args.instrument, str(tid)) + + layouts += ( + mini.make_sngl_ifo( + workflow, + sngl_file, + tmpltbank_file, + tid, + args.output_dir, + args.instrument, + statfiles=statfiles, + tags=args.tags + [str(rank)], + ), + ) + files += mini.make_trigger_timeseries( + workflow, + [sngl_file], + ifo_time, + args.output_dir, + special_tids=ifo_tid, + tags=args.tags + [str(rank)], + ) curr_params = {} - curr_params['mass1'] = trigs.mass1[num_event] - curr_params['mass2'] = trigs.mass2[num_event] - curr_params['spin1z'] = trigs.spin1z[num_event] - curr_params['spin2z'] = trigs.spin2z[num_event] - curr_params['f_lower'] = trigs.f_lower[num_event] - curr_params[args.instrument + '_end_time'] = time - curr_params['mean_time'] = time + curr_params["mass1"] = trigs.mass1[num_event] + curr_params["mass2"] = trigs.mass2[num_event] + curr_params["spin1z"] = trigs.spin1z[num_event] + curr_params["spin2z"] = trigs.spin2z[num_event] + curr_params["f_lower"] = trigs.f_lower[num_event] + curr_params[args.instrument + "_end_time"] = time + curr_params["mean_time"] = time # optional eccentric parameters, only present when using SEOBNRv4/5E try: - curr_params['eccentricity'] = trigs.eccentricity[num_event] - curr_params['rel_anomaly'] = trigs.rel_anomaly[num_event] + curr_params["eccentricity"] = trigs.eccentricity[num_event] + curr_params["rel_anomaly"] = trigs.rel_anomaly[num_event] except KeyError: pass # don't require precessing template info if not present try: - curr_params['spin1x'] = trigs.spin1x[num_event] - curr_params['spin2x'] = trigs.spin2x[num_event] - curr_params['spin1y'] = trigs.spin1y[num_event] - curr_params['spin2y'] = trigs.spin2y[num_event] - curr_params['inclination'] = trigs.inclination[num_event] + curr_params["spin1x"] = trigs.spin1x[num_event] + curr_params["spin2x"] = trigs.spin2x[num_event] + curr_params["spin1y"] = trigs.spin1y[num_event] + curr_params["spin2y"] = trigs.spin2y[num_event] + curr_params["inclination"] = trigs.inclination[num_event] except KeyError: pass try: # Only present for precessing search - curr_params['u_vals'] = trigs.u_vals[num_event] + curr_params["u_vals"] = trigs.u_vals[num_event] except: pass - _, sngl_plot = mini.make_single_template_plots(workflow, insp_segs, - args.inspiral_data_read_name, - args.inspiral_data_analyzed_name, - curr_params, - args.output_dir, - data_segments={args.instrument : insp_data_seglists}, - tags=args.tags+[str(rank)]) + _, sngl_plot = mini.make_single_template_plots( + workflow, + insp_segs, + args.inspiral_data_read_name, + args.inspiral_data_analyzed_name, + curr_params, + args.output_dir, + data_segments={args.instrument: insp_data_seglists}, + tags=args.tags + [str(rank)], + ) files += sngl_plot - files += mini.make_plot_waveform_plot(workflow, curr_params, - args.output_dir, [args.instrument], - tags=args.tags + [str(rank)]) + files += mini.make_plot_waveform_plot( + workflow, + curr_params, + args.output_dir, + [args.instrument], + tags=args.tags + [str(rank)], + ) - files += mini.make_singles_timefreq(workflow, sngl_file, tmpltbank_file, - time, args.output_dir, - data_segments=insp_data_seglists, - tags=args.tags + [str(rank)]) + files += mini.make_singles_timefreq( + workflow, + sngl_file, + tmpltbank_file, + time, + args.output_dir, + data_segments=insp_data_seglists, + tags=args.tags + [str(rank)], + ) - files += mini.make_qscan_plot(workflow, args.instrument, time, - args.output_dir, - data_segments=insp_data_seglists, - tags=args.tags + [str(rank)]) + files += mini.make_qscan_plot( + workflow, + args.instrument, + time, + args.output_dir, + data_segments=insp_data_seglists, + tags=args.tags + [str(rank)], + ) layouts += list(layout.grouper(files, 2)) diff --git a/bin/minifollowups/pycbc_upload_prep_minifollowup b/bin/minifollowups/pycbc_upload_prep_minifollowup index a30cc347a3d..05362d7d867 100644 --- a/bin/minifollowups/pycbc_upload_prep_minifollowup +++ b/bin/minifollowups/pycbc_upload_prep_minifollowup @@ -14,46 +14,62 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Prepare files for upload to GraceDB for foreground events -""" -import os +"""Prepare files for upload to GraceDB for foreground events""" + import argparse import logging -import numpy as np +import os import igwn_segments as segments +import numpy as np -from pycbc import init_logging, add_common_pycbc_options import pycbc.workflow as wf -from pycbc.types import MultiDetOptionAction +import pycbc.workflow.minifollowups as mini +from pycbc import add_common_pycbc_options, init_logging from pycbc.events import select_segments_by_definer from pycbc.io import HFile -import pycbc.workflow.minifollowups as mini -from pycbc.workflow.core import resolve_url_to_file, resolve_td_option +from pycbc.types import MultiDetOptionAction +from pycbc.workflow.core import resolve_td_option, resolve_url_to_file parser = argparse.ArgumentParser(description=__doc__[1:]) add_common_pycbc_options(parser) -parser.add_argument('--bank-file', - help="HDF format template bank file") -parser.add_argument('--statmap-file', - help="HDF format clustered coincident trigger result file") -parser.add_argument('--xml-all-file', - help="XML format result file containing all events") -parser.add_argument('--single-detector-triggers', nargs='+', action=MultiDetOptionAction, - help="HDF format merged single detector trigger files") -parser.add_argument('--inspiral-segments', - help="xml segment files containing the inspiral analysis times") -parser.add_argument('--inspiral-data-read-name', - help="Name of inspiral segmentlist containing data read in " - "by each analysis job.") -parser.add_argument('--inspiral-data-analyzed-name', - help="Name of inspiral segmentlist containing data " - "analyzed by each analysis job.") -parser.add_argument('--psd-files', nargs='+', action=MultiDetOptionAction, - help="HDF format merged single detector PSD files") -parser.add_argument('--ifar-thresh', type=float, - help="IFAR threshold for preparing SNR timeseries " - "files for upload. Default=No upload prep") +parser.add_argument("--bank-file", help="HDF format template bank file") +parser.add_argument( + "--statmap-file", help="HDF format clustered coincident trigger result file" +) +parser.add_argument( + "--xml-all-file", help="XML format result file containing all events" +) +parser.add_argument( + "--single-detector-triggers", + nargs="+", + action=MultiDetOptionAction, + help="HDF format merged single detector trigger files", +) +parser.add_argument( + "--inspiral-segments", + help="xml segment files containing the inspiral analysis times", +) +parser.add_argument( + "--inspiral-data-read-name", + help="Name of inspiral segmentlist containing data read in by each analysis job.", +) +parser.add_argument( + "--inspiral-data-analyzed-name", + help="Name of inspiral segmentlist containing data analyzed by each analysis job.", +) +parser.add_argument( + "--psd-files", + nargs="+", + action=MultiDetOptionAction, + help="HDF format merged single detector PSD files", +) +parser.add_argument( + "--ifar-thresh", + type=float, + help="IFAR threshold for preparing SNR timeseries " + "files for upload. Default=No upload prep", +) wf.add_workflow_command_line_group(parser) wf.add_workflow_settings_cli(parser, include_subdax_opts=True) @@ -69,16 +85,12 @@ wf.makedir(args.output_dir) channel_opts = {} for ifo in workflow.ifos: channel_opts[ifo] = workflow.cp.get_opt_tags( - "workflow", - "%s-channel-name" % ifo.lower(), - "" + "workflow", "%s-channel-name" % ifo.lower(), "" ) # create a FileList that will contain all output files layouts = [] -logging.info( - "Grabbing inputs: template bank, insp segs and XML with all events" -) +logging.info("Grabbing inputs: template bank, insp segs and XML with all events") tmpltbank_file = resolve_url_to_file(os.path.abspath(args.bank_file)) insp_segs = resolve_url_to_file(os.path.abspath(args.inspiral_segments)) xml_all = resolve_url_to_file(os.path.abspath(args.xml_all_file)) @@ -91,42 +103,36 @@ insp_analysed_seglists = {} for ifo in args.single_detector_triggers: strig_fname = args.single_detector_triggers[ifo] logging.info("Getting %s single-detector trigger file", ifo) - strig_file = resolve_url_to_file(os.path.abspath(strig_fname), - attrs={'ifos': ifo}) + strig_file = resolve_url_to_file(os.path.abspath(strig_fname), attrs={"ifos": ifo}) single_triggers.append(strig_file) logging.info("Getting %s PSD file", ifo) psd_fname = args.psd_files[ifo] - psd_file = resolve_url_to_file(os.path.abspath(psd_fname), - attrs={'ifos': ifo}) + psd_file = resolve_url_to_file(os.path.abspath(psd_fname), attrs={"ifos": ifo}) psd_files.append(psd_file) - fsdt[ifo] = HFile(args.single_detector_triggers[ifo], 'r') + fsdt[ifo] = HFile(args.single_detector_triggers[ifo], "r") logging.info("Loading inspiral segments information") insp_data_seglists[ifo] = select_segments_by_definer( - args.inspiral_segments, - segment_name=args.inspiral_data_read_name, - ifo=ifo) + args.inspiral_segments, segment_name=args.inspiral_data_read_name, ifo=ifo + ) insp_analysed_seglists[ifo] = select_segments_by_definer( - args.inspiral_segments, - segment_name=args.inspiral_data_analyzed_name, - ifo=ifo) + args.inspiral_segments, segment_name=args.inspiral_data_analyzed_name, ifo=ifo + ) insp_data_seglists[ifo].coalesce() insp_analysed_seglists[ifo].coalesce() -f = HFile(args.statmap_file, 'r') -stat = f['foreground/stat'][:] +f = HFile(args.statmap_file, "r") +stat = f["foreground/stat"][:] -bank_data = HFile(args.bank_file, 'r') +bank_data = HFile(args.bank_file, "r") ifar_limit = args.ifar_thresh # Get indices of all events which pass the IFAR threshold -event_ifars = f['foreground/ifar'][:] +event_ifars = f["foreground/ifar"][:] events_to_read = np.count_nonzero(event_ifars > ifar_limit) logging.info( - "%d events exceed the IFAR threshold of %.3f years", - events_to_read, - ifar_limit + "%d events exceed the IFAR threshold of %.3f years", events_to_read, ifar_limit ) # Sort by IFAR, descending event_idx = event_ifars.argsort()[::-1][:events_to_read] @@ -136,11 +142,11 @@ tids = {} bank_ids = {} logging.info("Getting event information") -ifo_list = f.attrs['ifos'].split(' ') +ifo_list = f.attrs["ifos"].split(" ") for ifo in ifo_list: - times[ifo] = f[f'foreground/{ifo}/time'][:][event_idx] - tids[ifo] = f[f'foreground/{ifo}/trigger_id'][:][event_idx] -bank_ids = f['foreground/template_id'][:][event_idx] + times[ifo] = f[f"foreground/{ifo}/time"][:][event_idx] + tids[ifo] = f[f"foreground/{ifo}/trigger_id"][:][event_idx] +bank_ids = f["foreground/template_id"][:][event_idx] f.close() @@ -148,17 +154,12 @@ for curr_idx in range(event_idx.size): logging.info("Event number %d", curr_idx) logging.info("Getting template parameters") params = mini.get_single_template_params( - curr_idx, - times, - bank_data, - bank_ids[curr_idx], - fsdt, - tids + curr_idx, times, bank_data, bank_ids[curr_idx], fsdt, tids ) # Extract approximant try: - appx = params.pop('approximant') + appx = params.pop("approximant") except KeyError: # approximant not stored in params, use default appx = None @@ -167,19 +168,16 @@ for curr_idx in range(event_idx.size): for ifo in ifo_list: ifo_chname = resolve_td_option( channel_opts[ifo], - segments.segment(params['mean_time'], params['mean_time']) + segments.segment(params["mean_time"], params["mean_time"]), ) channel_name += ifo_chname + " " single_temp_files = [] for ifo in ifo_list: - if params['mean_time'] not in insp_analysed_seglists[ifo]: - logging.info("Mean time %.3f not in segment list", - params['mean_time']) + if params["mean_time"] not in insp_analysed_seglists[ifo]: + logging.info("Mean time %.3f not in segment list", params["mean_time"]) continue - logging.info( - "Making single-template files" - ) + logging.info("Making single-template files") single_temp_files += mini.make_single_template_files( workflow, insp_segs, @@ -189,7 +187,7 @@ for curr_idx in range(event_idx.size): params, args.output_dir, store_file=True, - tags=args.tags+['upload', str(curr_idx)], + tags=args.tags + ["upload", str(curr_idx)], ) mini.make_upload_files( @@ -201,7 +199,7 @@ for curr_idx in range(event_idx.size): appx, args.output_dir, channel_name, - tags=args.tags+['upload', str(curr_idx)] + tags=args.tags + ["upload", str(curr_idx)], ) workflow.save() diff --git a/bin/plotting/pycbc_banksim_plot_eff_fitting_factor b/bin/plotting/pycbc_banksim_plot_eff_fitting_factor index 3d7e8cb8d6c..f810ab81693 100644 --- a/bin/plotting/pycbc_banksim_plot_eff_fitting_factor +++ b/bin/plotting/pycbc_banksim_plot_eff_fitting_factor @@ -19,44 +19,54 @@ Plot effective fitting factor vs mass1 and mass2 from various point-source files. """ -import sys -import numpy import argparse +import sys + import matplotlib -matplotlib.use('Agg') -import matplotlib.pyplot as plt +import numpy -from pycbc import init_logging, add_common_pycbc_options +matplotlib.use("Agg") +import matplotlib.pyplot as plt import pycbc.version -from pycbc import results + +from pycbc import add_common_pycbc_options, init_logging, results from pycbc.io.hdf import HFile -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_banksim_plot_eff_fitting_factor" -parser = argparse.ArgumentParser(usage='', - description=__doc__) +parser = argparse.ArgumentParser(usage="", description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--input-files', nargs='+', default=None, required=True, - help="List of input files.") -parser.add_argument('--output-file', default=None, required=True, - help="Output file.") -parser.add_argument('--filter-injections', action='store_true', default=False, - help="If true only consider and plot injections that are " - "marked in the input HDF file as passing the filter " - "that was supplied when generating the HDF file.") +parser.add_argument( + "--input-files", nargs="+", default=None, required=True, help="List of input files." +) +parser.add_argument("--output-file", default=None, required=True, help="Output file.") +parser.add_argument( + "--filter-injections", + action="store_true", + default=False, + help="If true only consider and plot injections that are " + "marked in the input HDF file as passing the filter " + "that was supplied when generating the HDF file.", +) # Plotting options -parser.add_argument('--plot-title', - help="If given, use this as the plot title") -parser.add_argument('--plot-caption', - help="If given, use this as the plot caption") -parser.add_argument('--log-axes', action='store_true', default=False, - help="If given, use logarithmic axes instead of linear.") -parser.add_argument('--log-colorbar', action='store_true', default=False, - help="If given, use logarithmic values for the colorbar, " - "showing 1 - fitting factor.") +parser.add_argument("--plot-title", help="If given, use this as the plot title") +parser.add_argument("--plot-caption", help="If given, use this as the plot caption") +parser.add_argument( + "--log-axes", + action="store_true", + default=False, + help="If given, use logarithmic axes instead of linear.", +) +parser.add_argument( + "--log-colorbar", + action="store_true", + default=False, + help="If given, use logarithmic values for the colorbar, " + "showing 1 - fitting factor.", +) opt = parser.parse_args() @@ -66,13 +76,13 @@ m1 = [] m2 = [] eff_ff = [] for file_name in opt.input_files: - curr_fp = HFile(file_name, 'r') - m1.append(curr_fp['inj_params/mass1'][0]) - m2.append(curr_fp['inj_params/mass2'][0]) + curr_fp = HFile(file_name, "r") + m1.append(curr_fp["inj_params/mass1"][0]) + m2.append(curr_fp["inj_params/mass2"][0]) if opt.filter_injections: - eff_ff.append(curr_fp['filtered_eff_fitting_factor'][()]) + eff_ff.append(curr_fp["filtered_eff_fitting_factor"][()]) else: - eff_ff.append(curr_fp['eff_fitting_factor'][()]) + eff_ff.append(curr_fp["eff_fitting_factor"][()]) curr_fp.close() if opt.filter_injections: @@ -89,31 +99,39 @@ if opt.filter_injections: fig = plt.figure() ax = fig.gca() -cmap = plt.get_cmap('viridis') +cmap = plt.get_cmap("viridis") if not opt.log_colorbar: - sctr = ax.scatter(m1, m2, c=eff_ff, cmap=cmap, edgecolors='none', - marker='o', s=40) + sctr = ax.scatter(m1, m2, c=eff_ff, cmap=cmap, edgecolors="none", marker="o", s=40) cb = fig.colorbar(sctr) cb.set_label("Effective fitting factor") if opt.filter_injections: - sctr = ax.scatter(m1_empty, m2_empty, c='k', cmap=cmap, - edgecolors='k', marker='x', s=40) + sctr = ax.scatter( + m1_empty, m2_empty, c="k", cmap=cmap, edgecolors="k", marker="x", s=40 + ) else: - sctr = ax.scatter(m1, m2, c=numpy.log10(1 - numpy.array(eff_ff)), - cmap=cmap, edgecolors='none', marker='o', s=40) + sctr = ax.scatter( + m1, + m2, + c=numpy.log10(1 - numpy.array(eff_ff)), + cmap=cmap, + edgecolors="none", + marker="o", + s=40, + ) cb = fig.colorbar(sctr) cb.set_label("log10(1 - Effective fitting factor)") if opt.filter_injections: - sctr = ax.scatter(m1_empty, m2_empty, c='k', cmap=cmap, - edgecolors='k', marker='x', s=40) + sctr = ax.scatter( + m1_empty, m2_empty, c="k", cmap=cmap, edgecolors="k", marker="x", s=40 + ) -ax.set_xlabel('Mass 1 (solar masses)') -ax.set_ylabel('Mass 2 (solar masses)') +ax.set_xlabel("Mass 1 (solar masses)") +ax.set_ylabel("Mass 2 (solar masses)") if opt.log_axes: - ax.set_yscale('log') - ax.set_xscale('log') + ax.set_yscale("log") + ax.set_xscale("log") ax.grid() if opt.plot_title is None: @@ -122,24 +140,32 @@ if opt.plot_title is None: else: opt.plot_title = "Effective fitting factor" if opt.plot_caption is None: - opt.plot_caption = ("The effective fitting factor is the signal-strength " - "weighted average fitting factor. This is shown, as " - "a function of masses, for the input injection sets.") + opt.plot_caption = ( + "The effective fitting factor is the signal-strength " + "weighted average fitting factor. This is shown, as " + "a function of masses, for the input injection sets." + ) if opt.filter_injections: - opt.plot_caption += (" Injections are filtered so that only those " - "passing the filter function used when creating " - "the input files are considered. Usually this is " - "used to restrict to only signals that are " - "the parameter space used to create the template " - "bank. Any point marked with a black " - "cross contain no injections that pass the " - "filter function.") + opt.plot_caption += ( + " Injections are filtered so that only those " + "passing the filter function used when creating " + "the input files are considered. Usually this is " + "used to restrict to only signals that are " + "the parameter space used to create the template " + "bank. Any point marked with a black " + "cross contain no injections that pass the " + "filter function." + ) fig_kwds = {} -if '.png' in opt.output_file: - fig_kwds['dpi'] = 200 +if ".png" in opt.output_file: + fig_kwds["dpi"] = 200 -results.save_fig_with_metadata(fig, opt.output_file, - fig_kwds=fig_kwds, title=opt.plot_title, - cmd=' '.join(sys.argv), - caption=opt.plot_caption) +results.save_fig_with_metadata( + fig, + opt.output_file, + fig_kwds=fig_kwds, + title=opt.plot_title, + cmd=" ".join(sys.argv), + caption=opt.plot_caption, +) diff --git a/bin/plotting/pycbc_banksim_plot_fitting_factors b/bin/plotting/pycbc_banksim_plot_fitting_factors index e071ecd30a4..89057383250 100644 --- a/bin/plotting/pycbc_banksim_plot_fitting_factors +++ b/bin/plotting/pycbc_banksim_plot_fitting_factors @@ -14,157 +14,167 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Make Alex Nielsen's fitting factor plots for a single banksim run. -""" +"""Make Alex Nielsen's fitting factor plots for a single banksim run.""" -import sys import argparse -import numpy +import sys + import matplotlib -matplotlib.use('Agg') -import matplotlib.pyplot as plt +import numpy +matplotlib.use("Agg") +import matplotlib.pyplot as plt import pycbc.version + from pycbc import results from pycbc.io.hdf import HFile __author__ = "Alex Nielsen , " -__author__ += "Ian Harry " +__author__ += "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_banksim_plot_fitting_factors" -parser = argparse.ArgumentParser(usage='', - description="Plot fitting factor distribution.") +parser = argparse.ArgumentParser( + usage="", description="Plot fitting factor distribution." +) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input-file', default=None, required=True, - help="List of input files.") -parser.add_argument('--output-file', default=None, required=True, - help="Output file.") -parser.add_argument('--filter-injections', action='store_true', default=False, - help="If true only consider and plot injections that are " - "marked in the input HDF file as passing the filter " - "that was supplied when generating the HDF file.") +parser.add_argument( + "--input-file", default=None, required=True, help="List of input files." +) +parser.add_argument("--output-file", default=None, required=True, help="Output file.") +parser.add_argument( + "--filter-injections", + action="store_true", + default=False, + help="If true only consider and plot injections that are " + "marked in the input HDF file as passing the filter " + "that was supplied when generating the HDF file.", +) # Plotting options -parser.add_argument('--plot-title', - help="If given, use this as the plot title") -parser.add_argument('--plot-caption', - help="If given, use this as the plot caption") +parser.add_argument("--plot-title", help="If given, use this as the plot title") +parser.add_argument("--plot-caption", help="If given, use this as the plot caption") + - opt = parser.parse_args() pycbc.init_logging(opt.verbose) -curr_fp = HFile(opt.input_file, 'r') -m1 = curr_fp['inj_params/mass1'][:] -m2 = curr_fp['inj_params/mass2'][:] -s1z = curr_fp['inj_params/spin1z'][:] -s2z = curr_fp['inj_params/spin2z'][:] -match = curr_fp['trig_params/match'][:] +curr_fp = HFile(opt.input_file, "r") +m1 = curr_fp["inj_params/mass1"][:] +m2 = curr_fp["inj_params/mass2"][:] +s1z = curr_fp["inj_params/spin1z"][:] +s2z = curr_fp["inj_params/spin2z"][:] +match = curr_fp["trig_params/match"][:] if opt.filter_injections: - bool_arr = curr_fp['filtered_points'][:] + bool_arr = curr_fp["filtered_points"][:] m1 = m1[bool_arr] m2 = m2[bool_arr] s1z = s1z[bool_arr] s2z = s2z[bool_arr] match = match[bool_arr] mtot = m1 + m2 -eta = m1*m2 / (mtot*mtot) -effspin = (s1z*m1 + s2z*m2) / mtot +eta = m1 * m2 / (mtot * mtot) +effspin = (s1z * m1 + s2z * m2) / mtot curr_fp.close() -cmap = plt.get_cmap('magma') +cmap = plt.get_cmap("magma") + def adjust_axes(ax_instance): """Adjust axes so that there is always a range of >= 0.1 on the plots.""" low, upp = ax_instance.get_xlim() if upp - low < 0.1: - cax.set_xlim([low-0.1, upp+0.1]) + cax.set_xlim([low - 0.1, upp + 0.1]) low, upp = ax_instance.get_ylim() if upp - low < 0.1: - cax.set_ylim([low-0.1, upp+0.1]) + cax.set_ylim([low - 0.1, upp + 0.1]) -pointsize=4 -fig, axarr = plt.subplots(nrows=3, ncols=2, figsize=(10,8)) -fig.subplots_adjust(left=0.1, bottom=0.1, right=0.8, top=0.96, wspace=0.3, - hspace=0.3) +pointsize = 4 +fig, axarr = plt.subplots(nrows=3, ncols=2, figsize=(10, 8)) +fig.subplots_adjust(left=0.1, bottom=0.1, right=0.8, top=0.96, wspace=0.3, hspace=0.3) cax = axarr[0, 0] -cax.autoscale_view('tight') +cax.autoscale_view("tight") cax.scatter(s1z, s2z, c=match, zorder=10, s=pointsize, linewidth=0, cmap=cmap) -cax.set_xlabel('Injected spin1') -cax.set_ylabel('Injected spin2') +cax.set_xlabel("Injected spin1") +cax.set_ylabel("Injected spin2") adjust_axes(cax) cax = axarr[0, 1] -cax.autoscale_view('tight') -cax.scatter(eta, effspin, c=match, zorder=10, s=pointsize, linewidth=0, - cmap=cmap) -cax.set_xlabel('Injected eta') -cax.set_ylabel('Injected effective spin') +cax.autoscale_view("tight") +cax.scatter(eta, effspin, c=match, zorder=10, s=pointsize, linewidth=0, cmap=cmap) +cax.set_xlabel("Injected eta") +cax.set_ylabel("Injected effective spin") adjust_axes(cax) -cax = axarr[1,0] -cax.autoscale_view('tight') +cax = axarr[1, 0] +cax.autoscale_view("tight") cax.scatter(m1, m2, c=match, zorder=10, s=pointsize, linewidth=0, cmap=cmap) -cax.set_xlabel('Injected mass1') -cax.set_ylabel('Injected mass2') +cax.set_xlabel("Injected mass1") +cax.set_ylabel("Injected mass2") adjust_axes(cax) -cax = axarr[1,1] -cax.autoscale_view('tight') -image = cax.scatter(mtot, effspin, c=match, zorder=10, s=pointsize, - linewidth=0, cmap=cmap) -cax.set_xlabel('Injected total mass') -cax.set_ylabel('Injected effective spin') +cax = axarr[1, 1] +cax.autoscale_view("tight") +image = cax.scatter( + mtot, effspin, c=match, zorder=10, s=pointsize, linewidth=0, cmap=cmap +) +cax.set_xlabel("Injected total mass") +cax.set_ylabel("Injected effective spin") adjust_axes(cax) -cax = axarr[2,0] -cax.autoscale_view('tight') +cax = axarr[2, 0] +cax.autoscale_view("tight") hist, bins = numpy.histogram(match, bins=50) width = 1.0 * (bins[1] - bins[0]) center = (bins[:-1] + bins[1:]) / 2 -cax.bar(center, hist, align='center', width=width, edgecolor="none") -cax.set_xlabel('fitting factor') -cax.set_ylabel('number') -cax.set_yscale('log') -cax.set_ylim([0.9,1000]) - -cax = axarr[2,1] -cax.autoscale_view('tight') +cax.bar(center, hist, align="center", width=width, edgecolor="none") +cax.set_xlabel("fitting factor") +cax.set_ylabel("number") +cax.set_yscale("log") +cax.set_ylim([0.9, 1000]) + +cax = axarr[2, 1] +cax.autoscale_view("tight") if len(m1): normed_hist = hist / hist.sum() else: normed_hist = hist cumulative = numpy.cumsum(normed_hist) -cax.bar(center, cumulative, align='center', width=width, edgecolor="none") -cax.set_xlabel('fitting factor') -cax.set_ylabel('cumulative fraction') +cax.bar(center, cumulative, align="center", width=width, edgecolor="none") +cax.set_xlabel("fitting factor") +cax.set_ylabel("cumulative fraction") if len(m1) > 50: - cax.set_yscale('log') -cax.set_ylim([0,1]) + cax.set_yscale("log") +cax.set_ylim([0, 1]) fig.subplots_adjust(right=0.8) cbar_ax = fig.add_axes([0.85, 0.05, 0.05, 0.91]) if len(mtot): cbar = fig.colorbar(image, cax=cbar_ax) - cbar.set_label('Recovered fitting factor') + cbar.set_label("Recovered fitting factor") if opt.plot_title is None: - opt.plot_title = 'Fitting factor plots' + opt.plot_title = "Fitting factor plots" if opt.plot_caption is None: - opt.plot_caption = ("A sequence of plots showing the fitting factor of " - "the input injections as a function of both spins and " - "masses (top 2 rows), as well as showing cumulative " - "and non-cumulative distributions of fitting factor.") + opt.plot_caption = ( + "A sequence of plots showing the fitting factor of " + "the input injections as a function of both spins and " + "masses (top 2 rows), as well as showing cumulative " + "and non-cumulative distributions of fitting factor." + ) fig_kwds = {} -if '.png' in opt.output_file: - fig_kwds['dpi'] = 200 - -results.save_fig_with_metadata(fig, opt.output_file, - fig_kwds=fig_kwds, title=opt.plot_title, - cmd=' '.join(sys.argv), - caption=opt.plot_caption) +if ".png" in opt.output_file: + fig_kwds["dpi"] = 200 + +results.save_fig_with_metadata( + fig, + opt.output_file, + fig_kwds=fig_kwds, + title=opt.plot_title, + cmd=" ".join(sys.argv), + caption=opt.plot_caption, +) diff --git a/bin/plotting/pycbc_banksim_table_point_injs b/bin/plotting/pycbc_banksim_table_point_injs index 5d6ec5192b5..06878f94de5 100644 --- a/bin/plotting/pycbc_banksim_table_point_injs +++ b/bin/plotting/pycbc_banksim_table_point_injs @@ -14,43 +14,52 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Make summary table for a set of point injection runs. -""" +"""Make summary table for a set of point injection runs.""" -import sys import argparse -import numpy +import sys +import numpy import pycbc.version + from pycbc import results from pycbc.io.hdf import HFile -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_banksim_table_point_injs" -parser = argparse.ArgumentParser(usage='', - description="Plot effective fitting factor vs mass1 and mass2.") +parser = argparse.ArgumentParser( + usage="", description="Plot effective fitting factor vs mass1 and mass2." +) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input-files', nargs='+', default=None, required=True, - help="List of input files.") -parser.add_argument('--directory-links', nargs='+', default=None, - help="Relative directory paths to corresponding input") -parser.add_argument('--output-file', default=None, required=True, - help="Output file.") - +parser.add_argument( + "--input-files", nargs="+", default=None, required=True, help="List of input files." +) +parser.add_argument( + "--directory-links", + nargs="+", + default=None, + help="Relative directory paths to corresponding input", +) +parser.add_argument("--output-file", default=None, required=True, help="Output file.") + opt = parser.parse_args() pycbc.init_logging(opt.verbose) -col_names = ['Mass 1', 'Mass 2', 'Signal recovery fraction', - 'Effective fitting factor', - 'Maximum fitting factor', - 'Minimum fitting factor'] -format_strings = ['##.##', '##.##', '#.###', '#.###', '#.###', '#.###'] +col_names = [ + "Mass 1", + "Mass 2", + "Signal recovery fraction", + "Effective fitting factor", + "Maximum fitting factor", + "Minimum fitting factor", +] +format_strings = ["##.##", "##.##", "#.###", "#.###", "#.###", "#.###"] if opt.directory_links is not None: - col_names = ['More details link'] + col_names + col_names = ["More details link"] + col_names format_strings = [None] + format_strings dir_names = [] @@ -63,59 +72,72 @@ min_ff = [] for idx, file_name in enumerate(opt.input_files): if opt.directory_links is not None: d = opt.directory_links[idx] - dir_names.append('LINK'.format(d)) - curr_fp = HFile(file_name, 'r') - m1.append(curr_fp['inj_params/mass1'][0]) - m2.append(curr_fp['inj_params/mass2'][0]) - eff_ff.append(curr_fp['eff_fitting_factor'][()]) - srf.append(curr_fp['sig_rec_fac'][()]) - max_ff.append(max(curr_fp['trig_params/match'][:])) - min_ff.append(min(curr_fp['trig_params/match'][:])) + dir_names.append(f'LINK') + curr_fp = HFile(file_name, "r") + m1.append(curr_fp["inj_params/mass1"][0]) + m2.append(curr_fp["inj_params/mass2"][0]) + eff_ff.append(curr_fp["eff_fitting_factor"][()]) + srf.append(curr_fp["sig_rec_fac"][()]) + max_ff.append(max(curr_fp["trig_params/match"][:])) + min_ff.append(min(curr_fp["trig_params/match"][:])) curr_fp.close() -columns = [numpy.array(m1), numpy.array(m2), numpy.array(srf), - numpy.array(eff_ff), numpy.array(max_ff), numpy.array(min_ff)] +columns = [ + numpy.array(m1), + numpy.array(m2), + numpy.array(srf), + numpy.array(eff_ff), + numpy.array(max_ff), + numpy.array(min_ff), +] if opt.directory_links is not None: columns = [numpy.array(dir_names)] + columns -test_fp = HFile(opt.input_files[0], 'r') -if 'filtered_points' in test_fp.keys(): +test_fp = HFile(opt.input_files[0], "r") +if "filtered_points" in test_fp.keys(): test_fp.close() - cn1 = 'Fraction of points within template bank' - cn2 = 'Filtered points signal recoveryfraction' - cn3 = 'Filtered points effective fittingfactor' - cn4 = 'Filtered points maximum fittingfactor' - cn5 = 'Filtered points minimum fittingfactor' + cn1 = "Fraction of points within template bank" + cn2 = "Filtered points signal recoveryfraction" + cn3 = "Filtered points effective fittingfactor" + cn4 = "Filtered points maximum fittingfactor" + cn5 = "Filtered points minimum fittingfactor" col_names += [cn1, cn2, cn3, cn4, cn5] - format_strings += ['#.###', '#.###', '#.###', '#.###', '#.###'] + format_strings += ["#.###", "#.###", "#.###", "#.###", "#.###"] point_frac = [] filt_srf = [] filt_eff_ff = [] filt_max_ff = [] filt_min_ff = [] for file_name in opt.input_files: - curr_fp = HFile(file_name, 'r') - point_frac.append(curr_fp['frac_points_within_bank'][()]) - filt_srf.append(curr_fp['filtered_sig_rec_fac'][()]) - filt_eff_ff.append(curr_fp['filtered_eff_fitting_factor'][()]) + curr_fp = HFile(file_name, "r") + point_frac.append(curr_fp["frac_points_within_bank"][()]) + filt_srf.append(curr_fp["filtered_sig_rec_fac"][()]) + filt_eff_ff.append(curr_fp["filtered_eff_fitting_factor"][()]) if point_frac[-1] > 0: - bool_arr = curr_fp['filtered_points'][:] - filt_max_ff.append(max(curr_fp['trig_params/match'][:][bool_arr])) - filt_min_ff.append(min(curr_fp['trig_params/match'][:][bool_arr])) + bool_arr = curr_fp["filtered_points"][:] + filt_max_ff.append(max(curr_fp["trig_params/match"][:][bool_arr])) + filt_min_ff.append(min(curr_fp["trig_params/match"][:][bool_arr])) else: filt_max_ff.append(-1) filt_min_ff.append(-1) curr_fp.close() - columns += [numpy.array(point_frac), numpy.array(filt_srf), - numpy.array(filt_eff_ff), numpy.array(filt_max_ff), - numpy.array(filt_min_ff)] + columns += [ + numpy.array(point_frac), + numpy.array(filt_srf), + numpy.array(filt_eff_ff), + numpy.array(filt_max_ff), + numpy.array(filt_min_ff), + ] else: test_fp.close() -html_table = results.html_table(columns, col_names, format_strings=format_strings, - page_size=len(m1)) +html_table = results.html_table( + columns, col_names, format_strings=format_strings, page_size=len(m1) +) -kwds = {'title' : 'Point Injection Results', - 'cmd' :' '.join(sys.argv), } +kwds = { + "title": "Point Injection Results", + "cmd": " ".join(sys.argv), +} results.save_fig_with_metadata(str(html_table), opt.output_file, **kwds) diff --git a/bin/plotting/pycbc_create_html_snippet b/bin/plotting/pycbc_create_html_snippet index 47b046eb90a..81e18894e0f 100644 --- a/bin/plotting/pycbc_create_html_snippet +++ b/bin/plotting/pycbc_create_html_snippet @@ -16,22 +16,22 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -import sys import argparse +import sys -from pycbc import init_logging, add_common_pycbc_options import pycbc.version + import pycbc.results +from pycbc import add_common_pycbc_options, init_logging # parse command line parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--output-file', type=str, - help='Path of the output HTML file.') -parser.add_argument('--html-text', type=str, - help="Contents of output HTML file.") -parser.add_argument('--title', type=str, default="Title", - help="Title of figure for results webpage.") +parser.add_argument("--output-file", type=str, help="Path of the output HTML file.") +parser.add_argument("--html-text", type=str, help="Contents of output HTML file.") +parser.add_argument( + "--title", type=str, default="Title", help="Title of figure for results webpage." +) opts = parser.parse_args() init_logging(opts.verbose) @@ -40,7 +40,10 @@ init_logging(opts.verbose) caption = opts.html_text # save the plot as an interactive HTML -pycbc.results.save_fig_with_metadata(opts.html_text, opts.output_file, - title=opts.title, - cmd=' '.join(sys.argv), - caption=caption) +pycbc.results.save_fig_with_metadata( + opts.html_text, + opts.output_file, + title=opts.title, + cmd=" ".join(sys.argv), + caption=caption, +) diff --git a/bin/plotting/pycbc_faithsim_plots b/bin/plotting/pycbc_faithsim_plots index 70c9ec5fad4..db3747fc05b 100755 --- a/bin/plotting/pycbc_faithsim_plots +++ b/bin/plotting/pycbc_faithsim_plots @@ -6,20 +6,22 @@ that compare two approximants and compute the match between them. """ import argparse + import matplotlib + matplotlib.use("Agg") import matplotlib.cm -from matplotlib.ticker import MultipleLocator -from matplotlib import pyplot as plt import numpy as np +from matplotlib import pyplot as plt +from matplotlib.ticker import MultipleLocator -from pycbc import init_logging, add_common_pycbc_options +from pycbc import add_common_pycbc_options, init_logging from pycbc.conversions import ( - mtotal_from_mass1_mass2, - q_from_mass1_mass2, + chi_eff, eta_from_mass1_mass2, mchirp_from_mass1_mass2, - chi_eff, + mtotal_from_mass1_mass2, + q_from_mass1_mass2, ) @@ -140,11 +142,11 @@ derived_func_map = { "mass_ratio": lambda d: q_from_mass1_mass2(d["mass1"], d["mass2"]), "mchirp": lambda d: mchirp_from_mass1_mass2(d["mass1"], d["mass2"]), "spin1_magnitude": lambda d: ( - d["spin1x"] ** 2 + d["spin1y"] ** 2 + d["spin1z"] ** 2 - ) - ** 0.5, - "spin2_magnitude": lambda d: (d["spin2x"] ** 2 + d["spin2y"] ** 2 + d["spin2z"] ** 2) - ** 0.5, + (d["spin1x"] ** 2 + d["spin1y"] ** 2 + d["spin1z"] ** 2) ** 0.5 + ), + "spin2_magnitude": lambda d: ( + (d["spin2x"] ** 2 + d["spin2y"] ** 2 + d["spin2z"] ** 2) ** 0.5 + ), "eta": lambda d: eta_from_mass1_mass2(d["mass1"], d["mass2"]), "chi_eff": lambda d: chi_eff(d["mass1"], d["mass2"], d["spin1z"], d["spin2z"]), "horizon_distance_1": lambda d: d["sigma1"] / 8, diff --git a/bin/plotting/pycbc_ifar_catalog b/bin/plotting/pycbc_ifar_catalog index f74a418e667..7eef21ed015 100644 --- a/bin/plotting/pycbc_ifar_catalog +++ b/bin/plotting/pycbc_ifar_catalog @@ -16,81 +16,94 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import argparse -import numpy -import sys import logging -import matplotlib as mpl; mpl.use('Agg') -from matplotlib import pyplot as plt +import sys + +import matplotlib as mpl +import numpy +mpl.use("Agg") +from matplotlib import pyplot as plt from scipy.stats import norm, poisson import pycbc.results -from pycbc import conversions -from pycbc import init_logging, add_common_pycbc_options +from pycbc import add_common_pycbc_options, conversions, init_logging from pycbc.io.hdf import HFile -parser = argparse.ArgumentParser(usage='pycbc_ifar_catalog [--options]', - description='Plots cumulative IFAR vs count for' - ' foreground triggers') +parser = argparse.ArgumentParser( + usage="pycbc_ifar_catalog [--options]", + description="Plots cumulative IFAR vs count for foreground triggers", +) add_common_pycbc_options(parser) -parser.add_argument('--trigger-files', nargs='+', - help='Path to coincident trigger HDF file(s)') -parser.add_argument('--output-file', required=True, - help='Path to output plot') -parser.add_argument('--truncate-threshold', type=float, - help='Truncate plot above IFAR threshold (units: years)') -parser.add_argument('--remove-threshold', type=float, - help='Remove detected events with IFAR above threshold' - ' (units years)') -parser.add_argument('--open-box', action='store_true', - help='Are we putting open box results onto the plot? ' - 'Default=False.') -parser.add_argument('--use-tex', action='store_true', - help="Render using latex") -parser.add_argument('--use-hierarchical-level', type=int, default=None, - help='Indicate which inclusive background and FARs of ' - 'foreground triggers to plot if there were any ' - 'hierarchical removals done. Choosing None plots ' - 'the inclusive backgrounds prior to any ' - 'hierarchical removals with the updated FARs for ' - 'foreground triggers after hierarchical removal(s). ' - 'Choosing 0 means plotting inclusive background ' - 'from prior to any hierarchical removals with FARs ' - 'for foreground triggers prior to hierarchical ' - 'removal. Choosing 1 means plotting the inclusive ' - 'background after doing 1 hierarchical removal, and ' - 'includes updated FARs from after 1 hierarchical ' - 'removal. [default=plot after application of all ' - 'hierarchical removal iterations]') -parser.add_argument('--use-exclusive-ifar', action='store_true', - help='Whether to use exclusive background for ' - 'IFAR calculations. Default=False.') +parser.add_argument( + "--trigger-files", nargs="+", help="Path to coincident trigger HDF file(s)" +) +parser.add_argument("--output-file", required=True, help="Path to output plot") +parser.add_argument( + "--truncate-threshold", + type=float, + help="Truncate plot above IFAR threshold (units: years)", +) +parser.add_argument( + "--remove-threshold", + type=float, + help="Remove detected events with IFAR above threshold (units years)", +) +parser.add_argument( + "--open-box", + action="store_true", + help="Are we putting open box results onto the plot? Default=False.", +) +parser.add_argument("--use-tex", action="store_true", help="Render using latex") +parser.add_argument( + "--use-hierarchical-level", + type=int, + default=None, + help="Indicate which inclusive background and FARs of " + "foreground triggers to plot if there were any " + "hierarchical removals done. Choosing None plots " + "the inclusive backgrounds prior to any " + "hierarchical removals with the updated FARs for " + "foreground triggers after hierarchical removal(s). " + "Choosing 0 means plotting inclusive background " + "from prior to any hierarchical removals with FARs " + "for foreground triggers prior to hierarchical " + "removal. Choosing 1 means plotting the inclusive " + "background after doing 1 hierarchical removal, and " + "includes updated FARs from after 1 hierarchical " + "removal. [default=plot after application of all " + "hierarchical removal iterations]", +) +parser.add_argument( + "--use-exclusive-ifar", + action="store_true", + help="Whether to use exclusive background for IFAR calculations. Default=False.", +) opts = parser.parse_args() init_logging(opts.verbose) if opts.use_tex: - plt.rc('text', usetex=True) - plt.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']}) + plt.rc("text", usetex=True) + plt.rc("font", family="serif", serif=["Computer Modern"]) -trigf = [HFile(f, 'r') for f in opts.trigger_files] +trigf = [HFile(f, "r") for f in opts.trigger_files] # Parse which inclusive background to use for the plotting h_inc_back_num = opts.use_hierarchical_level if opts.use_exclusive_ifar: - ifar_str = 'ifar_exc' + ifar_str = "ifar_exc" else: - ifar_str = 'ifar' + ifar_str = "ifar" try: - h_iterations = max([f.attrs['hierarchical_removal_iterations'] - for f in trigf]) + h_iterations = max([f.attrs["hierarchical_removal_iterations"] for f in trigf]) except KeyError: h_iterations = 0 if h_inc_back_num is None: - logging.info('Using %d hierarchical removal iterations', h_iterations) + logging.info("Using %d hierarchical removal iterations", h_iterations) h_inc_back_num = h_iterations if h_inc_back_num > h_iterations: @@ -101,18 +114,29 @@ if h_inc_back_num > h_iterations: ax.set_xlim(0, 1) ax.set_ylim(0, 1) - output_message = "No more foreground events louder than all background\n" \ - "at this removal level.\nAttempted to show " + \ - str(h_inc_back_num) + " removal(s),\n" \ - "but only " + str(h_iterations) + " removal(s) done." - - ax.text(0.5, 0.5, output_message, horizontalalignment='center', - verticalalignment='center') - - pycbc.results.save_fig_with_metadata(fig, opts.output_file, - title='Cumulative Number vs. IFAR', - caption=output_message, - cmd=' '.join(sys.argv)) + output_message = ( + "No more foreground events louder than all background\n" + "at this removal level.\nAttempted to show " + + str(h_inc_back_num) + + " removal(s),\n" + "but only " + str(h_iterations) + " removal(s) done." + ) + + ax.text( + 0.5, + 0.5, + output_message, + horizontalalignment="center", + verticalalignment="center", + ) + + pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + title="Cumulative Number vs. IFAR", + caption=output_message, + cmd=" ".join(sys.argv), + ) # Exit the code successfully and bypass the rest of the plotting code. sys.exit(0) @@ -122,9 +146,9 @@ if h_inc_back_num >= 0 and h_iterations is not None and h_iterations != 0: fore_ifar = numpy.array([]) supp_fore_ifar = numpy.array([]) for f in trigf: - ifar = f['foreground_h%s/%s' % (h_inc_back_num, ifar_str)][:] + ifar = f["foreground_h%s/%s" % (h_inc_back_num, ifar_str)][:] # Get ifar of hierarchically removed foreground triggers - supp_ifar = f['foreground/' + ifar_str][:] + supp_ifar = f["foreground/" + ifar_str][:] # increment arrays to be plotted fore_ifar = numpy.append(fore_ifar, ifar) supp_fore_ifar = numpy.append(supp_fore_ifar, supp_ifar) @@ -139,7 +163,7 @@ if h_inc_back_num >= 0 and h_iterations is not None and h_iterations != 0: else: fore_ifar = numpy.array([]) for f in trigf: - fore_ifar = numpy.append(fore_ifar, f['foreground/' + ifar_str][:]) + fore_ifar = numpy.append(fore_ifar, f["foreground/" + ifar_str][:]) if opts.remove_threshold is not None and opts.truncate_threshold is not None: raise RuntimeError("Can't both remove and truncate foreground events!") @@ -153,19 +177,26 @@ fore_cumnum = numpy.arange(len(fore_ifar), 0, -1) # get expected (noise-only) foreground IFAR values and cumulative number for # each IFAR value -expected_ifar = numpy.logspace(-4., numpy.log10(opts.truncate_threshold or 10000), - num=1000, base=10.0) +expected_ifar = numpy.logspace( + -4.0, numpy.log10(opts.truncate_threshold or 10000), num=1000, base=10.0 +) fg_time = 0 for f in trigf: - fg_time += f.attrs['foreground_time'] + fg_time += f.attrs["foreground_time"] expected_cumnum = conversions.sec_to_year(fg_time) / expected_ifar # make figure fig = plt.figure(1) # plot the expected background -plt.loglog(expected_ifar, expected_cumnum, linestyle='--', linewidth=1, - color='black', label='Expected Background') +plt.loglog( + expected_ifar, + expected_cumnum, + linestyle="--", + linewidth=1, + color="black", + label="Expected Background", +) # plot the counting error lower = {} @@ -204,36 +235,59 @@ for sigma in [1, 2, 3, 4]: upper[sigma].append(nup + 1.5) plotifar = expected_ifar[::-1] -plt.fill_between(plotifar, lower[4], lower[3], facecolor='k', alpha=0.15, - label=r'$<4\sigma$') -plt.fill_between(plotifar, lower[3], lower[2], facecolor='k', alpha=0.3, - label=r'$<3\sigma$') -plt.fill_between(plotifar, lower[2], lower[1], facecolor='k', alpha=0.45, - label=r'$<2\sigma$') -plt.fill_between(plotifar, lower[1], upper[1], facecolor='k', alpha=0.6, - label=r'$<1\sigma$') -plt.fill_between(plotifar, upper[1], upper[2], facecolor='k', alpha=0.45) -plt.fill_between(plotifar, upper[2], upper[3], facecolor='k', alpha=0.3) -plt.fill_between(plotifar, upper[3], upper[4], facecolor='k', alpha=0.15) +plt.fill_between( + plotifar, lower[4], lower[3], facecolor="k", alpha=0.15, label=r"$<4\sigma$" +) +plt.fill_between( + plotifar, lower[3], lower[2], facecolor="k", alpha=0.3, label=r"$<3\sigma$" +) +plt.fill_between( + plotifar, lower[2], lower[1], facecolor="k", alpha=0.45, label=r"$<2\sigma$" +) +plt.fill_between( + plotifar, lower[1], upper[1], facecolor="k", alpha=0.6, label=r"$<1\sigma$" +) +plt.fill_between(plotifar, upper[1], upper[2], facecolor="k", alpha=0.45) +plt.fill_between(plotifar, upper[2], upper[3], facecolor="k", alpha=0.3) +plt.fill_between(plotifar, upper[3], upper[4], facecolor="k", alpha=0.15) # plot the foreground triggers if opts.open_box: if opts.truncate_threshold: over_trunc = fore_ifar > opts.truncate_threshold - fore_ifar[over_trunc] = numpy.ones(over_trunc.sum()) * \ - opts.truncate_threshold + fore_ifar[over_trunc] = numpy.ones(over_trunc.sum()) * opts.truncate_threshold for i in fore_cumnum[over_trunc]: - plt.arrow(opts.truncate_threshold, i, opts.truncate_threshold, 0, - head_width=0.1 * i, head_length=0.4 * \ - opts.truncate_threshold, ec='b', fc='b') - plt.loglog(fore_ifar, fore_cumnum, linestyle='None', color='blue', - marker='^', ms=6, label='Foreground') + plt.arrow( + opts.truncate_threshold, + i, + opts.truncate_threshold, + 0, + head_width=0.1 * i, + head_length=0.4 * opts.truncate_threshold, + ec="b", + fc="b", + ) + plt.loglog( + fore_ifar, + fore_cumnum, + linestyle="None", + color="blue", + marker="^", + ms=6, + label="Foreground", + ) max_ifar = max(fore_ifar) if h_inc_back_num > 0: max_ifar = max(max_ifar, max(h_rm_ifar)) - plt.loglog(h_rm_ifar, h_rm_cumnum, linestyle='None', color='#b66dff', - marker='v', label='Hierarchically Removed Foreground') + plt.loglog( + h_rm_ifar, + h_rm_cumnum, + linestyle="None", + color="#b66dff", + marker="v", + label="Hierarchically Removed Foreground", + ) # format plot if opts.open_box: @@ -243,20 +297,24 @@ if opts.open_box: else: plt.ylim(0.7, max(expected_cumnum)) plt.grid() -plt.legend(loc='upper right', fontsize=13) -plt.ylabel('Cumulative Number', size='large') -ifar_label = 'Inverse False Alarm Rate (yr)' +plt.legend(loc="upper right", fontsize=13) +plt.ylabel("Cumulative Number", size="large") +ifar_label = "Inverse False Alarm Rate (yr)" if opts.use_exclusive_ifar: - ifar_label = 'Exclusive ' + ifar_label -plt.xlabel(ifar_label, size='large') + ifar_label = "Exclusive " + ifar_label +plt.xlabel(ifar_label, size="large") # save -caption = 'This is a cumulative histogram of triggers. The blue triangles ' \ - + 'represent foreground triggers. The dashed line ' \ - + 'represents the expected background given the analysis time. ' \ - + 'The shaded regions represent counting errors.' -pycbc.results.save_fig_with_metadata(fig, opts.output_file, - title='Cumulative Number vs. IFAR', - caption=caption, - cmd=' '.join(sys.argv)) - +caption = ( + "This is a cumulative histogram of triggers. The blue triangles " + "represent foreground triggers. The dashed line " + "represents the expected background given the analysis time. " + "The shaded regions represent counting errors." +) +pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + title="Cumulative Number vs. IFAR", + caption=caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_mass_area_plot b/bin/plotting/pycbc_mass_area_plot index 486cd3d7741..b35b799bad3 100644 --- a/bin/plotting/pycbc_mass_area_plot +++ b/bin/plotting/pycbc_mass_area_plot @@ -4,19 +4,22 @@ # By A. Curiel Barroso # August 2019 -"""This script computes the area corresponding to different CBC on +""" +This script computes the area corresponding to different CBC on the m1 & m2 plane when given a central mchirp value and uncertainty. """ import argparse + import numpy -from matplotlib import use; use("Agg") +from matplotlib import use + +use("Agg") from matplotlib import pyplot -from pycbc import init_logging, add_common_pycbc_options -from pycbc.mchirp_area import calc_areas -from pycbc.mchirp_area import src_mass_from_z_det_mass +from pycbc import add_common_pycbc_options, init_logging from pycbc.conversions import mass2_from_mchirp_mass1 as m2mcm1 +from pycbc.mchirp_area import calc_areas, src_mass_from_z_det_mass # ARGUMENT PARSER parser = argparse.ArgumentParser() @@ -68,8 +71,7 @@ print("agns = " + str(areas["gns"])) print("abns = " + str(areas["bns"])) # PLOT GENERATION -src_mchirp = src_mass_from_z_det_mass(central_z, delta_z, - central_mc, delta_mc) +src_mchirp = src_mass_from_z_det_mass(central_z, delta_z, central_mc, delta_mc) mcb = src_mchirp[0] + src_mchirp[1] mcs = src_mchirp[0] - src_mchirp[1] @@ -77,8 +79,8 @@ mcs = src_mchirp[0] - src_mchirp[1] # The points where the equal mass line and a chirp mass # curve intersect is m1 = m2 = (2**0.2)*mchirp -mib = (2**0.2)*mcb -mis = (2**0.2)*mcs +mib = (2**0.2) * mcb +mis = (2**0.2) * mcs lim_m1b = min(m1_max, m2mcm1(mcb, m2_min)) m1b = numpy.linspace(mib, lim_m1b, num=100) @@ -92,8 +94,7 @@ if mib > m1_max: pyplot.plot((m1_max, m1_max), (m2mcm1(mcs, lim_m1s), m1_max), "b") else: pyplot.plot(m1b, m2b, "b") - pyplot.plot((m1_max, m1_max), (m2mcm1(mcs, lim_m1s), - m2mcm1(mcb, lim_m1b)),"b") + pyplot.plot((m1_max, m1_max), (m2mcm1(mcs, lim_m1s), m2mcm1(mcb, lim_m1b)), "b") if mis >= m2_min: pyplot.plot(m1s, m2s, "b") @@ -109,6 +110,5 @@ pyplot.plot((gap_max, m1_max), (gap_max, gap_max), "k:") pyplot.xlabel("M1") pyplot.ylabel("M2") -pyplot.title("MChirp = " + str(0.5 * (mcb + mcs)) + " +/- " - + str((mcb - mcs) * 0.5)) +pyplot.title("MChirp = " + str(0.5 * (mcb + mcs)) + " +/- " + str((mcb - mcs) * 0.5)) pyplot.savefig("mass_plot.png") diff --git a/bin/plotting/pycbc_mchirp_plots b/bin/plotting/pycbc_mchirp_plots index 49003497a6b..7affe639d34 100644 --- a/bin/plotting/pycbc_mchirp_plots +++ b/bin/plotting/pycbc_mchirp_plots @@ -4,16 +4,20 @@ # By A. Curiel Barroso # August 2019 -"""This script computes the area corresponding to different CBC on the m1 & m2 +""" +This script computes the area corresponding to different CBC on the m1 & m2 plane as a function of central mchirp value. """ import argparse + import numpy -from matplotlib import use; use("Agg") +from matplotlib import use + +use("Agg") from matplotlib import pyplot -from pycbc import init_logging, add_common_pycbc_options +from pycbc import add_common_pycbc_options, init_logging from pycbc.mchirp_area import calc_areas # ARGUMENT PARSER @@ -57,8 +61,8 @@ y_ansbh = numpy.zeros(n, float) y_agns = numpy.zeros(n, float) y_abns = numpy.zeros(n, float) -for i in range(0, n): - central_mc = 0.8 + i*(10.0 - 0.8)/(n - 1) +for i in range(n): + central_mc = 0.8 + i * (10.0 - 0.8) / (n - 1) delta_mc = central_mc * 0.01 trig_mc = {"central": central_mc, "delta": delta_mc} x_mc[i] = central_mc diff --git a/bin/plotting/pycbc_page_banktriggerrate b/bin/plotting/pycbc_page_banktriggerrate index 464bcb681ed..507d92b64ac 100644 --- a/bin/plotting/pycbc_page_banktriggerrate +++ b/bin/plotting/pycbc_page_banktriggerrate @@ -1,31 +1,35 @@ #!/usr/bin/python -""" Plot the rate of triggers across the template bank -""" +"""Plot the rate of triggers across the template bank""" + import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") +import argparse + +import numpy from matplotlib import pyplot as plt -import numpy, argparse, pycbc.pnutils -from pycbc.io.hdf import HFile -from pycbc import init_logging, add_common_pycbc_options +import pycbc.pnutils +from pycbc import add_common_pycbc_options, init_logging +from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--trigger-files', nargs='+') -parser.add_argument('--bank-file') -parser.add_argument('--output-file') -parser.add_argument('--chisq-bins') +parser.add_argument("--trigger-files", nargs="+") +parser.add_argument("--bank-file") +parser.add_argument("--output-file") +parser.add_argument("--chisq-bins") args = parser.parse_args() init_logging(args.verbose) bf = HFile(args.bank_file) -m1 = bf['mass1'][:] -m2 = bf['mass2'][:] +m1 = bf["mass1"][:] +m2 = bf["mass2"][:] mc, et = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) -num_templates = len(m1) +num_templates = len(m1) template_num = numpy.zeros(num_templates) max_snr = numpy.zeros(num_templates) @@ -35,23 +39,23 @@ chisqs = [] for trig_filename in args.trigger_files: f = HFile(trig_filename) - tid = f['template_id'][:] - + tid = f["template_id"][:] + # Count triggers produced by each template tsort = tid.argsort() tid = tid[tsort] - snr = f['snr'][:][tsort] - chisq = f['chisq'][:][tsort] - + snr = f["snr"][:][tsort] + chisq = f["chisq"][:][tsort] + u = numpy.unique(tid) - l = numpy.searchsorted(tid, u, side='left') - r = numpy.searchsorted(tid, u, side='right') + l = numpy.searchsorted(tid, u, side="left") + r = numpy.searchsorted(tid, u, side="right") n = r - l template_num[u] += n - + snrs += [snr] chisqs += [chisq] - + num_trigs += len(tsort) chisq = numpy.concatenate(chisqs) / (float(args.chisq_bins) * 2 - 2) @@ -60,12 +64,12 @@ snr = numpy.concatenate(snrs) plt.figure() plt.scatter(snr[0:1000000], chisq[0:1000000]) plt.xlim(6, 8) -plt.ylim(.8, 3) -plt.savefig('snrchi.png') - -plt.figure() -plt.scatter(et, m1+m2, c=template_num, s=template_num/template_num.max()*50) -plt.ylabel('Total Mass') -plt.xlabel('Eta') +plt.ylim(0.8, 3) +plt.savefig("snrchi.png") + +plt.figure() +plt.scatter(et, m1 + m2, c=template_num, s=template_num / template_num.max() * 50) +plt.ylabel("Total Mass") +plt.xlabel("Eta") plt.colorbar() plt.savefig(args.output_file) diff --git a/bin/plotting/pycbc_page_coinc_snrchi b/bin/plotting/pycbc_page_coinc_snrchi index b7d8b636d1b..3c41e0b7116 100644 --- a/bin/plotting/pycbc_page_coinc_snrchi +++ b/bin/plotting/pycbc_page_coinc_snrchi @@ -1,69 +1,97 @@ #!/usr/bin/env python +import argparse import sys -import numpy, argparse, matplotlib + +import matplotlib +import numpy from matplotlib import colors -matplotlib.use('Agg') + +matplotlib.use("Agg") from matplotlib import pyplot as plt + import pycbc.results -from pycbc.io import ( - get_chisq_from_file_choice, chisq_choices, SingleDetTriggers, HFile -) -from pycbc import conversions, init_logging, add_common_pycbc_options +from pycbc import add_common_pycbc_options, conversions, init_logging from pycbc.detector import Detector +from pycbc.io import HFile, SingleDetTriggers, chisq_choices, get_chisq_from_file_choice -def snr_from_chisq(chisq, newsnr, q=6.): + +def snr_from_chisq(chisq, newsnr, q=6.0): snr = numpy.zeros(len(chisq)) + float(newsnr) - ind = numpy.where(chisq > 1.)[0] - snr[ind] = float(newsnr) / ( 0.5 * (1. + chisq[ind] ** (q/2.)) ) ** (-1./q) + ind = numpy.where(chisq > 1.0)[0] + snr[ind] = float(newsnr) / (0.5 * (1.0 + chisq[ind] ** (q / 2.0))) ** (-1.0 / q) return snr + parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--found-injection-file', required=True, - help='HDF format found injection file. Required') -parser.add_argument('--single-injection-file', required=True, - help='Single detector trigger files from the injection set' - ': one file per ifo') -parser.add_argument('--coinc-statistic-file', required=True, - help='HDF format statistic file. Required') -parser.add_argument('--single-trigger-file', required=True, - help='Single detector trigger files from the zero lag run.' - ' Required') -parser.add_argument('--newsnr-contours', nargs='*', default=[], - help="List of newsnr values to draw contours. Optional") -parser.add_argument('--background-front', action='store_true', default=False, - help='If set, plot background on top of injections rather ' - 'than vice versa') -parser.add_argument('--colorbar-choice', choices=('effective_spin', 'mchirp', - 'eta', 'effective_distance', 'mtotal', 'optimal_snr', - 'redshift'), default='effective_distance', - help='Parameter to use for the colorbar. ' - 'Default=effective_distance') -parser.add_argument('--chisq-choice', choices=chisq_choices, - default='traditional', - help='Which chisquared to plot. Default=traditional') -parser.add_argument('--output-file', required=True) +parser.add_argument( + "--found-injection-file", + required=True, + help="HDF format found injection file. Required", +) +parser.add_argument( + "--single-injection-file", + required=True, + help="Single detector trigger files from the injection set: one file per ifo", +) +parser.add_argument( + "--coinc-statistic-file", required=True, help="HDF format statistic file. Required" +) +parser.add_argument( + "--single-trigger-file", + required=True, + help="Single detector trigger files from the zero lag run. Required", +) +parser.add_argument( + "--newsnr-contours", + nargs="*", + default=[], + help="List of newsnr values to draw contours. Optional", +) +parser.add_argument( + "--background-front", + action="store_true", + default=False, + help="If set, plot background on top of injections rather than vice versa", +) +parser.add_argument( + "--colorbar-choice", + choices=( + "effective_spin", + "mchirp", + "eta", + "effective_distance", + "mtotal", + "optimal_snr", + "redshift", + ), + default="effective_distance", + help="Parameter to use for the colorbar. Default=effective_distance", +) +parser.add_argument( + "--chisq-choice", + choices=chisq_choices, + default="traditional", + help="Which chisquared to plot. Default=traditional", +) +parser.add_argument("--output-file", required=True) args = parser.parse_args() init_logging(args.verbose) # First - check the IFO being used -with HFile(args.single_trigger_file, 'r') as stf: +with HFile(args.single_trigger_file, "r") as stf: ifo = tuple(stf.keys())[0] # Add the background triggers -with HFile(args.coinc_statistic_file, 'r') as csf: - b_tids = csf['background_exc'][ifo]['trigger_id'][:] +with HFile(args.coinc_statistic_file, "r") as csf: + b_tids = csf["background_exc"][ifo]["trigger_id"][:] # Remove trigger ids == -1, as this indicates that it was found in # other detector(s) b_tids = b_tids[b_tids >= 0] - ifos = csf.attrs['ifos'].split(' ') + ifos = csf.attrs["ifos"].split(" ") -trigs = SingleDetTriggers( - args.single_trigger_file, - ifo, - premask=b_tids -) +trigs = SingleDetTriggers(args.single_trigger_file, ifo, premask=b_tids) bkg_snr = trigs.snr bkg_chisq = get_chisq_from_file_choice(trigs, args.chisq_choice) @@ -81,51 +109,62 @@ bkg_pos = bkg_chisq > 0 bkg_snr = bkg_snr[bkg_pos] bkg_chisq = bkg_chisq[bkg_pos] fig = plt.figure() -plt.scatter(bkg_snr, bkg_chisq, marker='o', color='black', - linewidth=0, s=4, label='Background', alpha=0.6, - zorder=args.background_front) +plt.scatter( + bkg_snr, + bkg_chisq, + marker="o", + color="black", + linewidth=0, + s=4, + label="Background", + alpha=0.6, + zorder=args.background_front, +) # Add the found injection points -f = HFile(args.found_injection_file, 'r') -inj_tid = f['found_after_vetoes'][ifo]['trigger_id'][:] +f = HFile(args.found_injection_file, "r") +inj_tid = f["found_after_vetoes"][ifo]["trigger_id"][:] # Remove trigger ids == -1, as this indicates that it was found in # other detector(s) valid_inj = inj_tid >= 0 inj_tid = inj_tid[valid_inj] -eff_dists = Detector(ifo).effective_distance(f['injections/distance'][:], - f['injections/ra'][:], - f['injections/dec'][:], - f['injections/polarization'][:], - f['injections/tc'][:], - f['injections/inclination'][:]) +eff_dists = Detector(ifo).effective_distance( + f["injections/distance"][:], + f["injections/ra"][:], + f["injections/dec"][:], + f["injections/polarization"][:], + f["injections/tc"][:], + f["injections/inclination"][:], +) -inj_idx = f['found_after_vetoes/injection_index'][valid_inj] +inj_idx = f["found_after_vetoes/injection_index"][valid_inj] eff_dist = eff_dists[inj_idx] -m1, m2 = f['injections/mass1'][:][inj_idx], f['injections/mass2'][:][inj_idx] -s1, s2 = f['injections/spin1z'][:][inj_idx], f['injections/spin2z'][:][inj_idx] +m1, m2 = f["injections/mass1"][:][inj_idx], f["injections/mass2"][:][inj_idx] +s1, s2 = f["injections/spin1z"][:][inj_idx], f["injections/spin2z"][:][inj_idx] mchirp = conversions.mchirp_from_mass1_mass2(m1, m2) eta = conversions.eta_from_mass1_mass2(m1, m2) weighted_spin = conversions.chi_eff(m1, m2, s1, s2) -redshift = f['injections/redshift'][:][inj_idx] if \ - args.colorbar_choice == 'redshift' else None +redshift = ( + f["injections/redshift"][:][inj_idx] if args.colorbar_choice == "redshift" else None +) # choices to color the found injections -coloring = {'effective_distance': (eff_dist, "Effective Distance (Mpc)", - colors.LogNorm()), - 'mchirp': (mchirp, "Chirp Mass", colors.LogNorm()), - 'eta': (eta, "Symmetric Mass Ratio", colors.LogNorm()), - 'effective_spin': (weighted_spin, "Weighted Aligned Spin", None), - 'mtotal': (m1 + m2, "Total Mass", colors.LogNorm()), - 'redshift': (redshift, "Redshift", None) - } - -if 'optimal_snr_{}'.format(ifo) in f['injections']: - opt_snr_str = 'injections/optimal_snr_{}'.format(ifo) +coloring = { + "effective_distance": (eff_dist, "Effective Distance (Mpc)", colors.LogNorm()), + "mchirp": (mchirp, "Chirp Mass", colors.LogNorm()), + "eta": (eta, "Symmetric Mass Ratio", colors.LogNorm()), + "effective_spin": (weighted_spin, "Weighted Aligned Spin", None), + "mtotal": (m1 + m2, "Total Mass", colors.LogNorm()), + "redshift": (redshift, "Redshift", None), +} + +if f"optimal_snr_{ifo}" in f["injections"]: + opt_snr_str = f"injections/optimal_snr_{ifo}" opt_snr = f[opt_snr_str][:][inj_idx] - coloring['optimal_snr'] = (opt_snr, 'Optimal SNR', colors.LogNorm()) + coloring["optimal_snr"] = (opt_snr, "Optimal SNR", colors.LogNorm()) -f = HFile(args.single_injection_file, 'r')[ifo] +f = HFile(args.single_injection_file, "r")[ifo] if len(inj_tid): inj_trigs = SingleDetTriggers( args.single_injection_file, @@ -137,10 +176,7 @@ if len(inj_tid): # Again, the SingleDetTriggers will have discarded order etc. of tids, # so we need to do index juggling - _, inj_tid_idx_inverse = numpy.unique( - inj_tid, - return_inverse=True - ) + _, inj_tid_idx_inverse = numpy.unique(inj_tid, return_inverse=True) inj_snr = inj_snr[inj_tid_idx_inverse] inj_chisq = inj_chisq[inj_tid_idx_inverse] @@ -153,64 +189,85 @@ inj_pos = inj_chisq > 0 if len(coloring[args.colorbar_choice][0]) == 0: coloring[args.colorbar_choice] = (None, None, None) else: # Only plot positive chisq - plt.scatter(inj_snr[inj_pos], inj_chisq[inj_pos], - c=coloring[args.colorbar_choice][0][inj_pos], - norm=coloring[args.colorbar_choice][2], s=20, - marker='^', linewidth=0, label="Injections", - zorder=(not args.background_front)) + plt.scatter( + inj_snr[inj_pos], + inj_chisq[inj_pos], + c=coloring[args.colorbar_choice][0][inj_pos], + norm=coloring[args.colorbar_choice][2], + s=20, + marker="^", + linewidth=0, + label="Injections", + zorder=(not args.background_front), + ) try: - r = numpy.logspace(numpy.log(min(bkg_chisq.min(), inj_chisq[inj_pos].min()) - * 0.9), - numpy.log(max(bkg_chisq.max(), inj_chisq.max()) * 1.1), 200) + r = numpy.logspace( + numpy.log(min(bkg_chisq.min(), inj_chisq[inj_pos].min()) * 0.9), + numpy.log(max(bkg_chisq.max(), inj_chisq.max()) * 1.1), + 200, + ) except ValueError: # Allow code to continue in the absence of injection triggers - r = numpy.logspace(numpy.log(bkg_chisq.min() * 0.9), - numpy.log(bkg_chisq.max() * 1.1), 200) + r = numpy.logspace( + numpy.log(bkg_chisq.min() * 0.9), numpy.log(bkg_chisq.max() * 1.1), 200 + ) if args.newsnr_contours: for cval in args.newsnr_contours: snrv = snr_from_chisq(r, cval) - plt.plot(snrv, r, '--', color='grey', linewidth=1) + plt.plot(snrv, r, "--", color="grey", linewidth=1) ax = plt.gca() -ax.set_xscale('log') -ax.set_yscale('log') +ax.set_xscale("log") +ax.set_yscale("log") try: cb = plt.colorbar() - cb.set_label(coloring[args.colorbar_choice][1], size='large') + cb.set_label(coloring[args.colorbar_choice][1], size="large") except (TypeError, ZeroDivisionError): # Catch case of no injection triggers if len(inj_chisq): raise -plt.title('%s Coincident Triggers' % ifo, size='large') -plt.xlabel('SNR', size='large') -plt.ylabel('Reduced $\chi^2$', size='large') +plt.title("%s Coincident Triggers" % ifo, size="large") +plt.xlabel("SNR", size="large") +plt.ylabel(r"Reduced $\chi^2$", size="large") try: - plt.xlim(min(inj_snr.min(), bkg_snr.min()) * 0.99, - max(inj_snr.max(), bkg_snr.max()) * 1.4) - plt.ylim(min(bkg_chisq.min(), inj_chisq[inj_pos].min()) * 0.7, - max(bkg_chisq.max(), inj_chisq.max()) * 1.4) + plt.xlim( + min(inj_snr.min(), bkg_snr.min()) * 0.99, + max(inj_snr.max(), bkg_snr.max()) * 1.4, + ) + plt.ylim( + min(bkg_chisq.min(), inj_chisq[inj_pos].min()) * 0.7, + max(bkg_chisq.max(), inj_chisq.max()) * 1.4, + ) except ValueError: # Raised if no injection triggers pass -plt.legend(loc='lower right', prop={'size': 10}) -plt.grid(which='major', ls='solid', alpha=0.7, linewidth=.5) -plt.grid(which='minor', ls='solid', alpha=0.7, linewidth=.1) +plt.legend(loc="lower right", prop={"size": 10}) +plt.grid(which="major", ls="solid", alpha=0.7, linewidth=0.5) +plt.grid(which="minor", ls="solid", alpha=0.7, linewidth=0.1) -title = '%s %s chisq vs SNR. %s background with injections %s' \ - % (ifo.upper(), args.chisq_choice, ''.join(ifos).upper(), - 'behind' if args.background_front else 'ontop') +title = "%s %s chisq vs SNR. %s background with injections %s" % ( + ifo.upper(), + args.chisq_choice, + "".join(ifos).upper(), + "behind" if args.background_front else "ontop", +) caption = """Distribution of SNR and %s chi-squared veto for single detector triggers. Black points are %s background triggers. Triangles are injection triggers colored by %s of the injection. Dashed lines show contours of -constant NewSNR.""" % (args.chisq_choice, ''.join(ifos).upper(), - coloring[args.colorbar_choice][1]) -pycbc.results.save_fig_with_metadata(fig, - args.output_file, - title=title, - caption=caption, - cmd=' '.join(sys.argv), - fig_kwds={'dpi':200}) +constant NewSNR.""" % ( + args.chisq_choice, + "".join(ifos).upper(), + coloring[args.colorbar_choice][1], +) +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title=title, + caption=caption, + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200}, +) diff --git a/bin/plotting/pycbc_page_dq_table b/bin/plotting/pycbc_page_dq_table index e500ab9e038..90e7d9fc057 100644 --- a/bin/plotting/pycbc_page_dq_table +++ b/bin/plotting/pycbc_page_dq_table @@ -1,8 +1,9 @@ #!/usr/bin/env python -""" Make a table of dq state information -""" -import sys +"""Make a table of dq state information""" + import argparse +import sys + import h5py as h5 import numpy as np @@ -11,44 +12,48 @@ import pycbc.results parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--ifo', required=True) -parser.add_argument('--dq-file', required=True) -parser.add_argument('--output-file') +parser.add_argument("--ifo", required=True) +parser.add_argument("--dq-file", required=True) +parser.add_argument("--output-file") args = parser.parse_args() pycbc.init_logging(args.verbose) dq_states = { - 'dq_state_0': 'Clean', - 'dq_state_1': 'DQ Flag', - 'dq_state_2': 'Autogating', + "dq_state_0": "Clean", + "dq_state_1": "DQ Flag", + "dq_state_2": "Autogating", } -f = h5.File(args.dq_file, 'r') -grp = f[args.ifo]['dq_segments'] +f = h5.File(args.dq_file, "r") +grp = f[args.ifo]["dq_segments"] livetimes = [] total_livetime = 0 for dq_state in dq_states: - livetime = grp[dq_state]['livetime'][()] + livetime = grp[dq_state]["livetime"][()] livetimes.append(livetime) total_livetime += livetime livetimes.append(total_livetime) frac_livetimes = [lt / total_livetime for lt in livetimes] -state_names = list(dq_states.values()) + ['Total'] +state_names = list(dq_states.values()) + ["Total"] columns = [state_names, livetimes, frac_livetimes] columns = [np.array(c) for c in columns] -col_names = ['DQ State', 'Livetime', '% of Livetime'] +col_names = ["DQ State", "Livetime", "% of Livetime"] -format_strings = [None, '0.0', '0.00%'] +format_strings = [None, "0.0", "0.00%"] -html_table = pycbc.results.html_table(columns, col_names, - page_size=len(state_names), - format_strings=format_strings) -title = f'{args.ifo} DQ State Livetimes' -caption = 'Table of DQ state livetimes' +html_table = pycbc.results.html_table( + columns, col_names, page_size=len(state_names), format_strings=format_strings +) +title = f"{args.ifo} DQ State Livetimes" +caption = "Table of DQ state livetimes" pycbc.results.save_fig_with_metadata( - str(html_table), args.output_file, title=title, - caption=caption, cmd=' '.join(sys.argv)) + str(html_table), + args.output_file, + title=title, + caption=caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_page_fars_vs_stat b/bin/plotting/pycbc_page_fars_vs_stat index a83979468b1..f8bed12a81b 100644 --- a/bin/plotting/pycbc_page_fars_vs_stat +++ b/bin/plotting/pycbc_page_fars_vs_stat @@ -1,52 +1,60 @@ #!/usr/bin/env python -import sys, h5py +import sys + +import h5py import matplotlib -matplotlib.use('agg') -from matplotlib import pyplot as plt + +matplotlib.use("agg") +import argparse +import logging + import numpy as np -import argparse, logging -from pycbc import init_logging, add_common_pycbc_options, results -from pycbc.events import significance +from matplotlib import pyplot as plt + +from pycbc import add_common_pycbc_options, init_logging, results from pycbc.conversions import sec_to_year as convert_s_to_y +from pycbc.events import significance -parser = argparse.ArgumentParser(usage='pycbc_page_ifar_vs_stat [--options]', - description='Plots cumulative IFAR vs stat for' - 'backgrounds of different event types.') +parser = argparse.ArgumentParser( + usage="pycbc_page_ifar_vs_stat [--options]", + description="Plots cumulative IFAR vs stat for" + "backgrounds of different event types.", +) add_common_pycbc_options(parser) # get singles fit options significance.insert_significance_option_group(parser) -parser.add_argument('--trigger-files', nargs='+', - help='Paths to separate-event-type statmap files.') -parser.add_argument('--ifo-combos', nargs='+', - help='Which event types (detector combinations) to plot.') -parser.add_argument('--min-x', type=float, default=-15., - help='X axis limit.') -parser.add_argument('--max-x', type=float, default=25., - help='X axis limit.') -parser.add_argument('--output-file', required=True, - help='Path to output plot.') +parser.add_argument( + "--trigger-files", nargs="+", help="Paths to separate-event-type statmap files." +) +parser.add_argument( + "--ifo-combos", nargs="+", help="Which event types (detector combinations) to plot." +) +parser.add_argument("--min-x", type=float, default=-15.0, help="X axis limit.") +parser.add_argument("--max-x", type=float, default=25.0, help="X axis limit.") +parser.add_argument("--output-file", required=True, help="Path to output plot.") opts = parser.parse_args() init_logging(opts.verbose) sig_dict = significance.digest_significance_options(opts.ifo_combos, opts) coinc_color_map = { - 'H': '#ee0000', # red - 'L': '#4ba6ff', # blue - 'V': '#9b59b6', # magenta/purple - 'HL': '#0000DD', - 'LV': '#FFCC33', - 'HV': '#990000', - 'HLV': np.array([20, 120, 20]) / 255.} + "H": "#ee0000", # red + "L": "#4ba6ff", # blue + "V": "#9b59b6", # magenta/purple + "HL": "#0000DD", + "LV": "#FFCC33", + "HV": "#990000", + "HLV": np.array([20, 120, 20]) / 255.0, +} coinc_line_map = { - 'H': '--', - 'L': '--', - 'V': '--', - 'HL': '-', - 'LV': '-', - 'HV': '-', - 'HLV':'-' + "H": "--", + "L": "--", + "V": "--", + "HL": "-", + "LV": "-", + "HV": "-", + "HLV": "-", } fig = plt.figure(figsize=(9, 6)) @@ -57,71 +65,90 @@ ifar_calc_points = np.linspace(opts.min_x, opts.max_x, 200) fitted = [] # List of combos for which a fit is used for f in opts.trigger_files: logging.info(f"Opening {f}") - with h5py.File(f, 'r') as exc_zlag_f: + with h5py.File(f, "r") as exc_zlag_f: logging.info(f"Ifos {exc_zlag_f.attrs['ifos']}") # eg 'H1 L1' - sig_key = exc_zlag_f.attrs['ifos'].replace(' ', '') # eg 'H1L1' - coinc_key = sig_key.replace('1', '') # eg 'HL' - stat_exc = exc_zlag_f['background_exc']['stat'][:] - dec_fac_exc = exc_zlag_f['background_exc']['decimation_factor'][:] - bg_time = convert_s_to_y(exc_zlag_f.attrs['background_time_exc']) - ifar_file = exc_zlag_f['background_exc']['ifar'][:] + sig_key = exc_zlag_f.attrs["ifos"].replace(" ", "") # eg 'H1L1' + coinc_key = sig_key.replace("1", "") # eg 'HL' + stat_exc = exc_zlag_f["background_exc"]["stat"][:] + dec_fac_exc = exc_zlag_f["background_exc"]["decimation_factor"][:] + bg_time = convert_s_to_y(exc_zlag_f.attrs["background_time_exc"]) + ifar_file = exc_zlag_f["background_exc"]["ifar"][:] logging.info(f"{coinc_key} background time (y): {bg_time}") # If fitting / extrapolation is done - if sig_dict[sig_key]['method'] == 'trigger_fit': + if sig_dict[sig_key]["method"] == "trigger_fit": fitted.append(sig_key) - fit_thresh = sig_dict[sig_key]['fit_threshold'] - far_tuple = significance.get_far(stat_exc, - ifar_calc_points, - dec_fac_exc, - bg_time, - **sig_dict[sig_key]) + fit_thresh = sig_dict[sig_key]["fit_threshold"] + far_tuple = significance.get_far( + stat_exc, ifar_calc_points, dec_fac_exc, bg_time, **sig_dict[sig_key] + ) far = far_tuple[1] # Version-proof: get_far returns 2 or 3 values # Plot n-louder points with a solid line and sngl fits with dashed - ax.plot(ifar_calc_points[ifar_calc_points < fit_thresh], - far[ifar_calc_points < fit_thresh], - '-', c=coinc_color_map[coinc_key], zorder=-5) - ax.plot(ifar_calc_points[ifar_calc_points > fit_thresh], - far[ifar_calc_points > fit_thresh], - coinc_line_map[coinc_key], - c=coinc_color_map[coinc_key], label=coinc_key, zorder=-5) + ax.plot( + ifar_calc_points[ifar_calc_points < fit_thresh], + far[ifar_calc_points < fit_thresh], + "-", + c=coinc_color_map[coinc_key], + zorder=-5, + ) + ax.plot( + ifar_calc_points[ifar_calc_points > fit_thresh], + far[ifar_calc_points > fit_thresh], + coinc_line_map[coinc_key], + c=coinc_color_map[coinc_key], + label=coinc_key, + zorder=-5, + ) del fit_thresh # avoid variable hanging around else: - far_tuple = significance.get_far(stat_exc, - ifar_calc_points, - dec_fac_exc, - bg_time, - method='n_louder') - - ax.plot(ifar_calc_points, far_tuple[1], '-', - c=coinc_color_map[coinc_key], label=coinc_key, zorder=-5) + far_tuple = significance.get_far( + stat_exc, ifar_calc_points, dec_fac_exc, bg_time, method="n_louder" + ) + + ax.plot( + ifar_calc_points, + far_tuple[1], + "-", + c=coinc_color_map[coinc_key], + label=coinc_key, + zorder=-5, + ) del sig_key, coinc_key, far_tuple # avoid variables hanging around # Plot the thresholds for combo in fitted: - ax.axvline(sig_dict[combo]['fit_threshold'], ls='-.', zorder=-10, - c=coinc_color_map[combo.replace('1', '')]) + ax.axvline( + sig_dict[combo]["fit_threshold"], + ls="-.", + zorder=-10, + c=coinc_color_map[combo.replace("1", "")], + ) ax.semilogy() ax.legend(ncol=2, fontsize=12) ax.grid() ax.set_xlim([opts.min_x, opts.max_x]) ax.set_ylim([1e-3, 1e6]) -ax.tick_params(axis='both', labelsize=14) -ax.set_xlabel('Ranking Statistic', fontsize=16) -ax.set_ylabel('Cumulative Event Rate [y$^{-1}$]', fontsize=16) +ax.tick_params(axis="both", labelsize=14) +ax.set_xlabel("Ranking Statistic", fontsize=16) +ax.set_ylabel("Cumulative Event Rate [y$^{-1}$]", fontsize=16) # Save -caption = 'Cumulative rates of noise events for separate event types. Solid ' \ - 'lines represent estimates from counts of louder (zerolag or time ' \ - 'shifted) events. Dashed lines represent estimates from fitting / ' \ - 'extrapolation.' - -results.save_fig_with_metadata(fig, opts.output_file, - fig_kwds={'bbox_inches': 'tight'}, - title='Cumulative Rate vs. Statistic', - caption=caption, - cmd=' '.join(sys.argv)) +caption = ( + "Cumulative rates of noise events for separate event types. Solid " + "lines represent estimates from counts of louder (zerolag or time " + "shifted) events. Dashed lines represent estimates from fitting / " + "extrapolation." +) + +results.save_fig_with_metadata( + fig, + opts.output_file, + fig_kwds={"bbox_inches": "tight"}, + title="Cumulative Rate vs. Statistic", + caption=caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_page_foreground b/bin/plotting/pycbc_page_foreground index 0144db45e8f..4ee24626664 100755 --- a/bin/plotting/pycbc_page_foreground +++ b/bin/plotting/pycbc_page_foreground @@ -1,12 +1,14 @@ #!/usr/bin/env python -"""Make table of the foreground events. Also includes ability to write FARs and +""" +Make table of the foreground events. Also includes ability to write FARs and p-values for intermediary hierarchical removal steps. """ -import sys import argparse import logging +import sys + import numpy import pycbc @@ -14,32 +16,34 @@ import pycbc.results from pycbc.io import hdf from pycbc.pnutils import mass1_mass2_to_mchirp_eta - parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-file', required=True) -parser.add_argument('--bank-file', required=True) -parser.add_argument('--single-detector-triggers', nargs='+') -parser.add_argument('--output-file', required=True) -parser.add_argument('--num-to-write', type=int) -parser.add_argument('--use-hierarchical-level', type=int, - help='Indicate which FARs to write to the table ' - 'based on the number of hierarchical removals done. ' - 'Choosing None defaults to giving the FARs after ' - 'all hierarchical removals were done depending on ' - 'previous configuration choices. Choosing 0 selects ' - 'writing the FARs prior to any hierarchical ' - 'removals. Choosing 1 means writing the FARs after ' - 'doing 1 hierarchical removal. The program will ' - 'fail if the user selects a number above the number ' - 'of hierarchical removals done. [default=None]') +parser.add_argument("--trigger-file", required=True) +parser.add_argument("--bank-file", required=True) +parser.add_argument("--single-detector-triggers", nargs="+") +parser.add_argument("--output-file", required=True) +parser.add_argument("--num-to-write", type=int) +parser.add_argument( + "--use-hierarchical-level", + type=int, + help="Indicate which FARs to write to the table " + "based on the number of hierarchical removals done. " + "Choosing None defaults to giving the FARs after " + "all hierarchical removals were done depending on " + "previous configuration choices. Choosing 0 selects " + "writing the FARs prior to any hierarchical " + "removals. Choosing 1 means writing the FARs after " + "doing 1 hierarchical removal. The program will " + "fail if the user selects a number above the number " + "of hierarchical removals done. [default=None]", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -with hdf.HFile(args.trigger_file, 'r') as f: +with hdf.HFile(args.trigger_file, "r") as f: try: - h_iterations = f.attrs['hierarchical_removal_iterations'] + h_iterations = f.attrs["hierarchical_removal_iterations"] except KeyError: h_iterations = None @@ -51,21 +55,19 @@ if h_num_rm is not None and (h_iterations is None or h_num_rm > h_iterations): col_two = numpy.array([h_iterations]) columns = [col_one, col_two] names = ["Hierarchical Removals Requested", "Hierarchical Removals Performed"] - format_strings = ['#', '#'] - html_table = pycbc.results.html_table(columns, names, - format_strings=format_strings, page_size=10) - - kwds = { 'title' : - 'No more events louder than all background at this removal level.', - 'cmd' :' '.join(sys.argv), } - pycbc.results.save_fig_with_metadata( - str(html_table), - args.output_file, - **kwds + format_strings = ["#", "#"] + html_table = pycbc.results.html_table( + columns, names, format_strings=format_strings, page_size=10 ) + + kwds = { + "title": "No more events louder than all background at this removal level.", + "cmd": " ".join(sys.argv), + } + pycbc.results.save_fig_with_metadata(str(html_table), args.output_file, **kwds) sys.exit(0) -if args.output_file.endswith('.xml') or args.output_file.endswith('.xml.gz'): +if args.output_file.endswith(".xml") or args.output_file.endswith(".xml.gz"): if args.single_detector_triggers is None: err_msg = "If creating xml files must provide the single detector " err_msg += "trigger lists with --single-detector-triggers." @@ -74,36 +76,53 @@ if args.output_file.endswith('.xml') or args.output_file.endswith('.xml.gz'): # Ensure that the user uses a trigger file that is hierarchical # removal compatible. Also default to regular behavior if hierarchical # removals were not performed. -if h_num_rm is not None and h_iterations is not None \ - and h_num_rm >= 0 and h_iterations != 0: +if ( + h_num_rm is not None + and h_iterations is not None + and h_num_rm >= 0 + and h_iterations != 0 +): # read in the triggers from the group corresponding to the removal stage - fortrigs = hdf.ForegroundTriggers(args.trigger_file, args.bank_file, - sngl_files=args.single_detector_triggers, - n_loudest=args.num_to_write, - group='foreground_h%s' % h_num_rm) + fortrigs = hdf.ForegroundTriggers( + args.trigger_file, + args.bank_file, + sngl_files=args.single_detector_triggers, + n_loudest=args.num_to_write, + group="foreground_h%s" % h_num_rm, + ) # read in the regular triggers to get the exclusive ifar which is # independent of any removal - fortrigs_exclusive = hdf.ForegroundTriggers(args.trigger_file, args.bank_file, - sngl_files=args.single_detector_triggers, - n_loudest=args.num_to_write) - -else : - fortrigs = hdf.ForegroundTriggers(args.trigger_file, args.bank_file, - sngl_files=args.single_detector_triggers, - n_loudest=args.num_to_write) - -if args.output_file.endswith('.html'): - ifar = fortrigs.get_coincfile_array('ifar') - fap = fortrigs.get_coincfile_array('fap') - stat = fortrigs.get_coincfile_array('stat') - mass1 = fortrigs.get_bankfile_array('mass1') - mass2 = fortrigs.get_bankfile_array('mass2') - spin1z = fortrigs.get_bankfile_array('spin1z') - spin2z = fortrigs.get_bankfile_array('spin2z') + fortrigs_exclusive = hdf.ForegroundTriggers( + args.trigger_file, + args.bank_file, + sngl_files=args.single_detector_triggers, + n_loudest=args.num_to_write, + ) + +else: + fortrigs = hdf.ForegroundTriggers( + args.trigger_file, + args.bank_file, + sngl_files=args.single_detector_triggers, + n_loudest=args.num_to_write, + ) + +if args.output_file.endswith(".html"): + ifar = fortrigs.get_coincfile_array("ifar") + fap = fortrigs.get_coincfile_array("fap") + stat = fortrigs.get_coincfile_array("stat") + mass1 = fortrigs.get_bankfile_array("mass1") + mass2 = fortrigs.get_bankfile_array("mass2") + spin1z = fortrigs.get_bankfile_array("spin1z") + spin2z = fortrigs.get_bankfile_array("spin2z") time = fortrigs.get_end_time() - if h_num_rm is not None and h_iterations is not None \ - and h_num_rm >= 0 and h_iterations != 0: + if ( + h_num_rm is not None + and h_iterations is not None + and h_num_rm >= 0 + and h_iterations != 0 + ): # FIXME: Currently hierarchical removal is not in the multi-detector # workflow. If this functionality is added, and we want to make # this table for hierarchically removed triggers, we would need @@ -112,18 +131,18 @@ if args.output_file.endswith('.html'): # This information is not stored at all levels of the trigger file # so grab from the top level and populate the variables for output later - ifar_exc_orig = fortrigs_exclusive.get_coincfile_array('ifar_exc') - fap_exc_orig = fortrigs_exclusive.get_coincfile_array('fap_exc') + ifar_exc_orig = fortrigs_exclusive.get_coincfile_array("ifar_exc") + fap_exc_orig = fortrigs_exclusive.get_coincfile_array("fap_exc") - trig_id1 = fortrigs_exclusive.get_coincfile_array('trigger_id1') - trig_id1_hrm = fortrigs.get_coincfile_array('trigger_id1') + trig_id1 = fortrigs_exclusive.get_coincfile_array("trigger_id1") + trig_id1_hrm = fortrigs.get_coincfile_array("trigger_id1") # Get the list of id's that map from one to the other set idx1_list = [] for i in range(len(trig_id1)): for j in range(len(trig_id1_hrm)): - if trig_id1[i] == trig_id1_hrm[j]: - idx1_list.append(i) + if trig_id1[i] == trig_id1_hrm[j]: + idx1_list.append(i) ifar_exc = numpy.zeros(len(ifar), dtype=numpy.float64) fap_exc = numpy.zeros(len(ifar), dtype=numpy.float64) @@ -131,63 +150,95 @@ if args.output_file.endswith('.html'): ifar_exc[k] = ifar_exc_orig[idx1_list[k]] fap_exc[k] = fap_exc_orig[idx1_list[k]] - else : - ifar_exc = fortrigs.get_coincfile_array('ifar_exc') - fap_exc = fortrigs.get_coincfile_array('fap_exc') + else: + ifar_exc = fortrigs.get_coincfile_array("ifar_exc") + fap_exc = fortrigs.get_coincfile_array("fap_exc") mchirp, eta = mass1_mass2_to_mchirp_eta(mass1, mass2) - columns = [ifar_exc, ifar, fap_exc, fap, stat, time, - mchirp, mass1, mass2, spin1z, spin2z] + columns = [ + ifar_exc, + ifar, + fap_exc, + fap, + stat, + time, + mchirp, + mass1, + mass2, + spin1z, + spin2z, + ] # Not supposed to use FAP anymore, nomenclature is p-value. Should be # fixed consistently accross PyCBC. - names = ['Exc. IFAR (YR)', 'Inc. IFAR (YR)', 'Exc. FAP', - 'Inc. FAP', 'Ranking Statistic', 'End Time', - 'mchirp', 'm1','m2', 's1z', 's2z'] - format_strings = ['#.###E0', '#.###E0', '#.##E0', '#.##E0', '##.###', - None, '##.##', '##.##', '##.##', '##.##', '##.##'] + names = [ + "Exc. IFAR (YR)", + "Inc. IFAR (YR)", + "Exc. FAP", + "Inc. FAP", + "Ranking Statistic", + "End Time", + "mchirp", + "m1", + "m2", + "s1z", + "s2z", + ] + format_strings = [ + "#.###E0", + "#.###E0", + "#.##E0", + "#.##E0", + "##.###", + None, + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + ] if args.single_detector_triggers: - single_snr = fortrigs.get_snglfile_array_dict('snr') - single_chisq = fortrigs.get_snglfile_array_dict('chisq') - single_chisq_dof = fortrigs.get_snglfile_array_dict('chisq_dof') + single_snr = fortrigs.get_snglfile_array_dict("snr") + single_chisq = fortrigs.get_snglfile_array_dict("chisq") + single_chisq_dof = fortrigs.get_snglfile_array_dict("chisq_dof") for ifo in sorted(single_snr.keys()): snrs = [] chisqs = [] for idx in range(len(single_snr[ifo][0])): if single_snr[ifo][1][idx]: - snrs.append('%.2f' % single_snr[ifo][0][idx]) - rchisq = single_chisq[ifo][0][idx] / \ - (single_chisq_dof[ifo][0][idx]*2 - 2) - chisqs.append('%.2f' % rchisq) + snrs.append("%.2f" % single_snr[ifo][0][idx]) + rchisq = single_chisq[ifo][0][idx] / ( + single_chisq_dof[ifo][0][idx] * 2 - 2 + ) + chisqs.append("%.2f" % rchisq) else: - snrs.append(' ') - chisqs.append(' ') + snrs.append(" ") + chisqs.append(" ") columns.append(numpy.array(snrs)) - names.append("%s SNR" %(ifo)) + names.append("%s SNR" % (ifo)) format_strings.append("##.##") columns.append(numpy.array(chisqs)) - names.append("%s Red. Chisq" %(ifo)) + names.append("%s Red. Chisq" % (ifo)) format_strings.append("##.##") - logging.info('Making table of foreground triggers') - html_table = pycbc.results.html_table(columns, names, - format_strings=format_strings, page_size=10) - - kwds = { 'title' : 'Loudest Event Table', - 'cmd' :' '.join(sys.argv), } - pycbc.results.save_fig_with_metadata( - str(html_table), - args.output_file, - **kwds + logging.info("Making table of foreground triggers") + html_table = pycbc.results.html_table( + columns, names, format_strings=format_strings, page_size=10 ) -elif args.output_file.endswith('.xml') or args.output_file.endswith('.xml.gz'): + kwds = { + "title": "Loudest Event Table", + "cmd": " ".join(sys.argv), + } + pycbc.results.save_fig_with_metadata(str(html_table), args.output_file, **kwds) + +elif args.output_file.endswith(".xml") or args.output_file.endswith(".xml.gz"): fortrigs.to_coinc_xml_object(args.output_file) -elif args.output_file.endswith('.hdf'): +elif args.output_file.endswith(".hdf"): fortrigs.to_coinc_hdf_object(args.output_file) else: diff --git a/bin/plotting/pycbc_page_foundmissed b/bin/plotting/pycbc_page_foundmissed index f1d16519641..42ec019af98 100644 --- a/bin/plotting/pycbc_page_foundmissed +++ b/bin/plotting/pycbc_page_foundmissed @@ -1,83 +1,110 @@ #!/usr/bin/env python -""" Plot found and missed injections. -""" -import numpy -import logging +"""Plot found and missed injections.""" + import argparse +import logging import sys + import matplotlib -matplotlib.use('Agg') +import numpy + +matplotlib.use("Agg") import matplotlib.pyplot as plot -import pycbc.results.followup, pycbc.pnutils, pycbc.results import pycbc.pnutils +import pycbc.results +import pycbc.results.followup from pycbc.detector import Detector from pycbc.io.hdf import HFile -labels={'mchirp': 'Chirp Mass', - 'mtotal': 'Total Mass', - 'mass_ratio': 'Mass Ratio', - 'decisive_distance': 'Decisive Distance (Mpc)', - 'dec_chirp_distance': 'Decisive Chirp Distance (Mpc)', - 'min_eff_distance': 'Minimum Effective Distance (Mpc)', - 'min_eff_chirp_distance': 'Minimum Effective Chirp Distance (Mpc)', - 'chirp_distance': 'Chirp Distance (Mpc)', - 'comb_optimal_snr': 'Combined Optimal SNR', - 'decisive_optimal_snr': 'Decisive Optimal SNR', - 'max_optimal_snr': 'Maximum Optimal SNR', - 'redshift': 'Redshift', - 'time': 'Time (s)', - 'effective_spin': 'Effective Inspiral Spin', - } +labels = { + "mchirp": "Chirp Mass", + "mtotal": "Total Mass", + "mass_ratio": "Mass Ratio", + "decisive_distance": "Decisive Distance (Mpc)", + "dec_chirp_distance": "Decisive Chirp Distance (Mpc)", + "min_eff_distance": "Minimum Effective Distance (Mpc)", + "min_eff_chirp_distance": "Minimum Effective Chirp Distance (Mpc)", + "chirp_distance": "Chirp Distance (Mpc)", + "comb_optimal_snr": "Combined Optimal SNR", + "decisive_optimal_snr": "Decisive Optimal SNR", + "max_optimal_snr": "Maximum Optimal SNR", + "redshift": "Redshift", + "time": "Time (s)", + "effective_spin": "Effective Inspiral Spin", +} parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--injection-file', - help="The hdf injection file to plot", required=True) -parser.add_argument('--axis-type', default='mchirp', choices=['mchirp', - 'effective_spin', 'time', 'mtotal', 'mass_ratio']) -parser.add_argument('--log-x', action='store_true', default=False) -parser.add_argument('--distance-type', default='decisive_optimal_snr', - choices=list(labels), - help="Variable related to injected distance. Decisive " - "distance and dec chirp distance only available for " - "2-ifo search") -parser.add_argument('--plot-all-distance', action='store_true', default=False, - help="Plot all values of distance or SNR. If not given, " - "the plot will be truncated below 1") -parser.add_argument('--colormap',default='cividis_r', - help="Type of colormap to be used for the plots.") -parser.add_argument('--log-distance', action='store_true', default=False) -parser.add_argument('--dynamic', action='store_true', default=False) -parser.add_argument('--gradient-far', action='store_true', - help="Show far of found injections as a gradient") -parser.add_argument('--output-file', required=True) -parser.add_argument('--ifar-limits', nargs=2, type=float, - help="Supply upper and lower limits for IFAR colors. " - "0 indicates no limit.") -parser.add_argument('--far-type', choices=('inclusive', 'exclusive'), - default='inclusive', - help="Type of far to plot for the color. Choices are " - "'inclusive' or 'exclusive'. Default = 'inclusive'") -parser.add_argument('--missed-on-top', action='store_true', - help="Plot missed injections on top of found ones and " - "high FAR on top of low FAR") +parser.add_argument( + "--injection-file", help="The hdf injection file to plot", required=True +) +parser.add_argument( + "--axis-type", + default="mchirp", + choices=["mchirp", "effective_spin", "time", "mtotal", "mass_ratio"], +) +parser.add_argument("--log-x", action="store_true", default=False) +parser.add_argument( + "--distance-type", + default="decisive_optimal_snr", + choices=list(labels), + help="Variable related to injected distance. Decisive " + "distance and dec chirp distance only available for " + "2-ifo search", +) +parser.add_argument( + "--plot-all-distance", + action="store_true", + default=False, + help="Plot all values of distance or SNR. If not given, " + "the plot will be truncated below 1", +) +parser.add_argument( + "--colormap", default="cividis_r", help="Type of colormap to be used for the plots." +) +parser.add_argument("--log-distance", action="store_true", default=False) +parser.add_argument("--dynamic", action="store_true", default=False) +parser.add_argument( + "--gradient-far", + action="store_true", + help="Show far of found injections as a gradient", +) +parser.add_argument("--output-file", required=True) +parser.add_argument( + "--ifar-limits", + nargs=2, + type=float, + help="Supply upper and lower limits for IFAR colors. 0 indicates no limit.", +) +parser.add_argument( + "--far-type", + choices=("inclusive", "exclusive"), + default="inclusive", + help="Type of far to plot for the color. Choices are " + "'inclusive' or 'exclusive'. Default = 'inclusive'", +) +parser.add_argument( + "--missed-on-top", + action="store_true", + help="Plot missed injections on top of found ones and high FAR on top of low FAR", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Read in the data') -f = HFile(args.injection_file, 'r') -time = f['injections/tc'][:] -found = f['found_after_vetoes/injection_index'][:] -missed = f['missed/after_vetoes'][:] +logging.info("Read in the data") +f = HFile(args.injection_file, "r") +time = f["injections/tc"][:] +found = f["found_after_vetoes/injection_index"][:] +missed = f["missed/after_vetoes"][:] -if args.far_type == 'inclusive': - ifar_found = f['found_after_vetoes/ifar'][:] - far_title = 'Inclusive' -elif args.far_type == 'exclusive': - ifar_found = f['found_after_vetoes/ifar_exc'][:] - far_title = 'Exclusive' +if args.far_type == "inclusive": + ifar_found = f["found_after_vetoes/ifar"][:] + far_title = "Inclusive" +elif args.far_type == "exclusive": + ifar_found = f["found_after_vetoes/ifar_exc"][:] + far_title = "Exclusive" upper_lim = False lower_lim = False @@ -93,92 +120,101 @@ if args.ifar_limits: upper_lim = True ifar_found = numpy.minimum(ifar_found, args.ifar_limits[1]) -s1z = f['injections/spin1z'][:] -s2z = f['injections/spin2z'][:] -dist = f['injections/distance'][:] -m1, m2 = f['injections/mass1'][:], f['injections/mass2'][:] +s1z = f["injections/spin1z"][:] +s2z = f["injections/spin2z"][:] +dist = f["injections/distance"][:] +m1, m2 = f["injections/mass1"][:], f["injections/mass2"][:] vals = {} -vals['mchirp'], eta = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) +vals["mchirp"], eta = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) # GRRR should use pnutils formula -vals['effective_spin'] = (m1 * s1z + m2 * s2z) / (m1 + m2) -vals['time'] = time -vals['mtotal'] = m1 + m2 +vals["effective_spin"] = (m1 * s1z + m2 * s2z) / (m1 + m2) +vals["time"] = time +vals["mtotal"] = m1 + m2 # use convention that q>1 -vals['mass_ratio'] = numpy.maximum(m1/m2, m2/m1) +vals["mass_ratio"] = numpy.maximum(m1 / m2, m2 / m1) dvals = {} -if 'ifos' in f.attrs: - ifos = f.attrs['ifos'].split(' ') +if "ifos" in f.attrs: + ifos = f.attrs["ifos"].split(" ") else: logging.warning("Ifos not found in input file, assuming H-L") - ifos = ['H1', 'L1'] + ifos = ["H1", "L1"] # NOTE: Effective distance is hardcoded to these values. It's also normally # meaningless for precessing injections. try: eff_dists = [] - for ifo in ['H1', 'L1', 'V1']: - eff_dists.append(Detector(ifo).effective_distance( - f['injections/distance'][:], - f['injections/ra'][:], - f['injections/dec'][:], - f['injections/polarization'][:], - f['injections/tc'][:], - f['injections/inclination'][:])) + for ifo in ["H1", "L1", "V1"]: + eff_dists.append( + Detector(ifo).effective_distance( + f["injections/distance"][:], + f["injections/ra"][:], + f["injections/dec"][:], + f["injections/polarization"][:], + f["injections/tc"][:], + f["injections/inclination"][:], + ) + ) eff_dists = numpy.array(eff_dists).T # "Decisive" distance for coincidences is the second smallest # effective distance - dvals['decisive_distance'] = numpy.sort(eff_dists)[:,1] - dvals['dec_chirp_distance'] = \ - pycbc.pnutils.chirp_distance(dvals['decisive_distance'], - vals['mchirp']) + dvals["decisive_distance"] = numpy.sort(eff_dists)[:, 1] + dvals["dec_chirp_distance"] = pycbc.pnutils.chirp_distance( + dvals["decisive_distance"], vals["mchirp"] + ) # When single-detector triggers are included in the analysis, # minimum distance is the relevant quantity for injection efficiency - dvals['min_eff_distance'] = numpy.min(eff_dists, axis=1) - dvals['min_eff_chirp_distance'] = \ - pycbc.pnutils.chirp_distance(dvals['min_eff_distance'], - vals['mchirp']) + dvals["min_eff_distance"] = numpy.min(eff_dists, axis=1) + dvals["min_eff_chirp_distance"] = pycbc.pnutils.chirp_distance( + dvals["min_eff_distance"], vals["mchirp"] + ) except KeyError: # If the ifo isn't in the effective distance columns you can't get this. # But you can still use other values. pass -dvals['chirp_distance'] = pycbc.pnutils.chirp_distance(dist, vals['mchirp']) -if args.distance_type == 'redshift': - dvals['redshift'] = f['injections/redshift'][:] +dvals["chirp_distance"] = pycbc.pnutils.chirp_distance(dist, vals["mchirp"]) +if args.distance_type == "redshift": + dvals["redshift"] = f["injections/redshift"][:] -if 'snr' in args.distance_type: # only evaluate SNRs if needed - opt_snrsq_arr = \ - [f['injections/optimal_snr_%s' % ifo][:] ** 2. for ifo in ifos] - dvals['comb_optimal_snr'] = \ - numpy.array([numpy.sqrt(sum(opt_snrsq)) - for opt_snrsq in zip(*opt_snrsq_arr)]) +if "snr" in args.distance_type: # only evaluate SNRs if needed + opt_snrsq_arr = [f["injections/optimal_snr_%s" % ifo][:] ** 2.0 for ifo in ifos] + dvals["comb_optimal_snr"] = numpy.array( + [numpy.sqrt(sum(opt_snrsq)) for opt_snrsq in zip(*opt_snrsq_arr)] + ) # Decisive optimal SNR is the 2nd largest optimal SNR - dvals['decisive_optimal_snr'] = \ - numpy.array([numpy.sqrt(sorted(opt_snrsq)[-2]) - for opt_snrsq in zip(*opt_snrsq_arr)]) - dvals['max_optimal_snr'] = \ - numpy.array([numpy.sqrt(max(opt_snrsq)) - for opt_snrsq in zip(*opt_snrsq_arr)]) + dvals["decisive_optimal_snr"] = numpy.array( + [numpy.sqrt(sorted(opt_snrsq)[-2]) for opt_snrsq in zip(*opt_snrsq_arr)] + ) + dvals["max_optimal_snr"] = numpy.array( + [numpy.sqrt(max(opt_snrsq)) for opt_snrsq in zip(*opt_snrsq_arr)] + ) fdvals = dvals[args.distance_type][found] mdvals = dvals[args.distance_type][missed] if args.missed_on_top: - fig_title = 'Missed and Found Injections' + fig_title = "Missed and Found Injections" else: - fig_title = 'Found and Missed Injections' + fig_title = "Found and Missed Injections" fig = plot.figure() zmissed = args.missed_on_top zfound = not args.missed_on_top -mpoints = plot.scatter(vals[args.axis_type][missed], mdvals, s=16, - linewidth=0.5, marker='x', color='red', - label='missed', zorder=zmissed) +mpoints = plot.scatter( + vals[args.axis_type][missed], + mdvals, + s=16, + linewidth=0.5, + marker="x", + color="red", + label="missed", + zorder=zmissed, +) fvals = vals[args.axis_type][found] @@ -199,34 +235,47 @@ if not args.gradient_far: color[thousand] = 0 norm = matplotlib.colors.Normalize() - caption = (fig_title + ": Red x's are missed injections. " - "Blue circles are found with IFAR < 100 years, gray are < " - "1000 years, and yellow are found with IFAR >=1000 years. ") + caption = ( + fig_title + ": Red x's are missed injections. " + "Blue circles are found with IFAR < 100 years, gray are < " + "1000 years, and yellow are found with IFAR >=1000 years. " + ) else: color = 1.0 / ifsorted if len(color) < 2: - color=None + color = None norm = matplotlib.colors.LogNorm() - caption = (fig_title + ": Red x's are missed injections. " - "Circles are found injections. The color indicates the value " - "of the false alarm rate." ) - -points = plot.scatter(fvals, fdval, c=color, linewidth=0, s=16, norm=norm, - marker='o', label='found', zorder=zfound, - cmap=args.colormap) + caption = ( + fig_title + ": Red x's are missed injections. " + "Circles are found injections. The color indicates the value " + "of the false alarm rate." + ) + +points = plot.scatter( + fvals, + fdval, + c=color, + linewidth=0, + s=16, + norm=norm, + marker="o", + label="found", + zorder=zfound, + cmap=args.colormap, +) if args.gradient_far: try: if upper_lim and lower_lim: - ext = 'both' + ext = "both" elif lower_lim and not upper_lim: - ext = 'max' + ext = "max" elif upper_lim and not lower_lim: - ext = 'min' + ext = "min" else: - ext = 'neither' + ext = "neither" c = plot.colorbar(extend=ext) - c.set_label('False Alarm Rate $(yr^{-1})$, %s' % far_title) + c.set_label("False Alarm Rate $(yr^{-1})$, %s" % far_title) # Set up tick labels - there will be 5 if possible min_tick = numpy.ceil(min(numpy.log10(color))) @@ -240,9 +289,9 @@ if args.gradient_far: raise if args.missed_on_top: - caption += "Missed injections are shown on top of found injections." + caption += "Missed injections are shown on top of found injections." else: - caption += "Found injections are shown on top of missed injections." + caption += "Found injections are shown on top of missed injections." ax = plot.gca() plot.xlabel(labels[args.axis_type]) @@ -251,14 +300,14 @@ plot.grid() if args.log_x: # log x axis may fail for some choices, eg effective spin - ax.set_xscale('log') + ax.set_xscale("log") tmpxvals = list(vals[args.axis_type][missed]) tmpxvals += list(vals[args.axis_type][found]) xmax = 1.4 * max(tmpxvals) xmin = 0.7 * min(tmpxvals) plot.xlim(xmin, xmax) if args.log_distance: - ax.set_yscale('log') + ax.set_yscale("log") tmpyvals = list(fdvals) + list(mdvals) ymax = 1.2 * max(tmpyvals) @@ -267,7 +316,7 @@ if args.plot_all_distance: # note: ymin=0 will clash with args.log_distance # in that case it *should* throw an error! plot.ylim(ymin, ymax) -elif args.distance_type == 'redshift': +elif args.distance_type == "redshift": # default y limit: min redshift 0 plot.ylim(ymin=0, ymax=ymax) else: @@ -275,19 +324,23 @@ else: plot.ylim(ymin=1, ymax=ymax) fig_kwds = {} -if '.png' in args.output_file: - fig_kwds['dpi'] = 200 +if ".png" in args.output_file: + fig_kwds["dpi"] = 200 -if ('.html' in args.output_file): +if ".html" in args.output_file: plot.subplots_adjust(left=0.1, right=0.8, top=0.9, bottom=0.1) - import mpld3, mpld3.plugins, mpld3.utils - mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fmt='.5g')) - legend = mpld3.plugins.InteractiveLegendPlugin([mpoints, points], - ['missed', 'found'], - alpha_unsel=0.1) + import mpld3 + import mpld3.plugins + import mpld3.utils + + mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fmt=".5g")) + legend = mpld3.plugins.InteractiveLegendPlugin( + [mpoints, points], ["missed", "found"], alpha_unsel=0.1 + ) mpld3.plugins.connect(fig, legend) -title = '%s: %s vs %s' % (fig_title, args.axis_type, args.distance_type) -cmd = ' '.join(sys.argv) -pycbc.results.save_fig_with_metadata(fig, args.output_file, fig_kwds=fig_kwds, - title=title, cmd=cmd, caption=caption) +title = "%s: %s vs %s" % (fig_title, args.axis_type, args.distance_type) +cmd = " ".join(sys.argv) +pycbc.results.save_fig_with_metadata( + fig, args.output_file, fig_kwds=fig_kwds, title=title, cmd=cmd, caption=caption +) diff --git a/bin/plotting/pycbc_page_ifar b/bin/plotting/pycbc_page_ifar index a93321829fa..8052086d8f1 100644 --- a/bin/plotting/pycbc_page_ifar +++ b/bin/plotting/pycbc_page_ifar @@ -16,88 +16,112 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import argparse -import numpy -import sys import copy +import sys + import matplotlib as mpl -mpl.use('Agg') +import numpy + +mpl.use("Agg") from matplotlib import pyplot as plt -from pycbc import init_logging, add_common_pycbc_options import pycbc.results -from pycbc.events import veto +from pycbc import add_common_pycbc_options, init_logging from pycbc import conversions as conv +from pycbc.events import veto from pycbc.io import HFile -def calculate_time_slide_duration(pifo_segments, fifo_segments, offset=0): - ''' Returns the amount of coincident time between two segmentlists. - ''' - # dertermine how much coincident time - overlap = pifo_segments.coalesce() & fifo_segments.coalesce().shift(offset) - duration = abs(overlap.coalesce()) +def calculate_time_slide_duration(pifo_segments, fifo_segments, offset=0): + """Returns the amount of coincident time between two segmentlists.""" + # dertermine how much coincident time + overlap = pifo_segments.coalesce() & fifo_segments.coalesce().shift(offset) + duration = abs(overlap.coalesce()) - # unshift the shifted segmentlist - fifo_segments.shift(-offset) + # unshift the shifted segmentlist + fifo_segments.shift(-offset) - return duration + return duration # parser command line -parser = argparse.ArgumentParser(usage='pycbc_page_ifar [--options]', - description='Plots a cumulative histogram of IFAR for' - 'coincident foreground triggers and a subset of' - 'the coincident time slide triggers.') +parser = argparse.ArgumentParser( + usage="pycbc_page_ifar [--options]", + description="Plots a cumulative histogram of IFAR for" + "coincident foreground triggers and a subset of" + "the coincident time slide triggers.", +) add_common_pycbc_options(parser) -parser.add_argument('--trigger-file', type=str, required=True, - help='Path to coincident trigger HDF file.') -parser.add_argument('--output-file', type=str, required=True, - help='Path to output plot.') -parser.add_argument('--decimation-factor', type=int, required=True, - help='Decimation factor used in background estimation.' - 'Decimation factor means that every nth time' - 'slide kept all of its coincident triggers in' - 'the HDF file.') -parser.add_argument('--all-decimated-background-min-ifar', type=float, - default=0.01, metavar='VAL', - help='Threshold the IFARs of decimated background events ' - 'before combining them and rescaling to zerolag ' - '(green line). Plotting all decimated background ' - 'can be memory intensive and is not necessary for ' - 'checking the background estimate at interesting ' - 'IFAR levels. Slide coincs with IFARs below VAL ' - 'will be cut, the green line will then be valid for ' - 'zerolag down to VAL*(zerolag coinc time/total ' - 'decimated slide time). Recommended value equal ' - 'to default value 0.01yr.') -parser.add_argument('--use-hierarchical-level', type=int, default=None, - help='Indicate which inclusive background and FARs of ' - 'foreground triggers to plot if there were any ' - 'hierarchical removals done. Choosing None plots ' - 'the inclusive backgrounds prior to any ' - 'hierarchical removals with the updated FARs for ' - 'foreground triggers after hierarchical removal(s). ' - 'Choosing 0 means plotting inclusive background ' - 'from prior to any hierarchical removals with FARs ' - 'for foreground triggers prior to hierarchical ' - 'removal. Choosing 1 means plotting the inclusive ' - 'background after doing 1 hierarchical removal, and ' - 'includes updated FARs from after 1 hierarchical ' - 'removal. [default=None]') -parser.add_argument('--open-box', action='store_true', default=False, - help='Show the foreground triggers on the output plot. ') +parser.add_argument( + "--trigger-file", + type=str, + required=True, + help="Path to coincident trigger HDF file.", +) +parser.add_argument( + "--output-file", type=str, required=True, help="Path to output plot." +) +parser.add_argument( + "--decimation-factor", + type=int, + required=True, + help="Decimation factor used in background estimation." + "Decimation factor means that every nth time" + "slide kept all of its coincident triggers in" + "the HDF file.", +) +parser.add_argument( + "--all-decimated-background-min-ifar", + type=float, + default=0.01, + metavar="VAL", + help="Threshold the IFARs of decimated background events " + "before combining them and rescaling to zerolag " + "(green line). Plotting all decimated background " + "can be memory intensive and is not necessary for " + "checking the background estimate at interesting " + "IFAR levels. Slide coincs with IFARs below VAL " + "will be cut, the green line will then be valid for " + "zerolag down to VAL*(zerolag coinc time/total " + "decimated slide time). Recommended value equal " + "to default value 0.01yr.", +) +parser.add_argument( + "--use-hierarchical-level", + type=int, + default=None, + help="Indicate which inclusive background and FARs of " + "foreground triggers to plot if there were any " + "hierarchical removals done. Choosing None plots " + "the inclusive backgrounds prior to any " + "hierarchical removals with the updated FARs for " + "foreground triggers after hierarchical removal(s). " + "Choosing 0 means plotting inclusive background " + "from prior to any hierarchical removals with FARs " + "for foreground triggers prior to hierarchical " + "removal. Choosing 1 means plotting the inclusive " + "background after doing 1 hierarchical removal, and " + "includes updated FARs from after 1 hierarchical " + "removal. [default=None]", +) +parser.add_argument( + "--open-box", + action="store_true", + default=False, + help="Show the foreground triggers on the output plot. ", +) opts = parser.parse_args() init_logging(opts.verbose) # read file -fp = HFile(opts.trigger_file, 'r') +fp = HFile(opts.trigger_file, "r") # Parse which inclusive background to use for the plotting h_inc_back_num = opts.use_hierarchical_level try: - h_iterations = fp.attrs['hierarchical_removal_iterations'] + h_iterations = fp.attrs["hierarchical_removal_iterations"] except KeyError: h_iterations = 0 @@ -107,35 +131,44 @@ if h_inc_back_num is None: if h_inc_back_num > h_iterations: # Produce a null plot saying no hierarchical removals can be plotted import sys + fig = plt.figure() ax = fig.add_subplot(111) ax.set_xlim(0, 1) ax.set_ylim(0, 1) - output_message = "No more foreground events louder than all background\n" \ - "at this removal level.\nAttempted to show %d " \ - "removal(s),\nbut only %d removal(s) done." % \ - (h_inc_back_num, h_iterations) - - ax.text(0.5, 0.5, output_message, horizontalalignment='center', - verticalalignment='center') - - pycbc.results.save_fig_with_metadata(fig, opts.output_file, - title='Cumulative Number vs. IFAR', - caption=output_message, - cmd=' '.join(sys.argv)) + output_message = ( + "No more foreground events louder than all background\n" + "at this removal level.\nAttempted to show %d " + "removal(s),\nbut only %d removal(s) done." % (h_inc_back_num, h_iterations) + ) + + ax.text( + 0.5, + 0.5, + output_message, + horizontalalignment="center", + verticalalignment="center", + ) + + pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + title="Cumulative Number vs. IFAR", + caption=output_message, + cmd=" ".join(sys.argv), + ) # Exit the code successfully and bypass the rest of the plotting code. sys.exit(0) # get foreground IFAR values and cumulative number for each IFAR value if opts.open_box: - if h_inc_back_num >= 0 and h_iterations is not None and h_iterations != 0: - fore_ifar = fp['foreground_h%s/ifar' % h_inc_back_num][:] + fore_ifar = fp["foreground_h%s/ifar" % h_inc_back_num][:] # Get ifar of hierarchically removed foreground triggers - supp_fore_ifar = fp['foreground/ifar'][:] + supp_fore_ifar = fp["foreground/ifar"][:] # Use to plot hierarchically removed foreground triggers # on the plot and correct the cumulative number at IFAR level @@ -145,24 +178,24 @@ if opts.open_box: h_rm_ifar = hrm_sorted[idx_start:] h_rm_cumnum = numpy.arange(len(h_rm_ifar), 0, -1) - else : - fore_ifar = fp['foreground/ifar'][:] + else: + fore_ifar = fp["foreground/ifar"][:] fore_ifar.sort() fore_cumnum = numpy.arange(len(fore_ifar), 0, -1) # get expected foreground IFAR values and cumulative number for each IFAR value expected_ifar = numpy.logspace(-8, 2, num=100, endpoint=True, base=10.0) -expected_cumnum = conv.sec_to_year(fp.attrs['foreground_time'] / expected_ifar) +expected_cumnum = conv.sec_to_year(fp.attrs["foreground_time"] / expected_ifar) # get background timeslide IDs and IFAR values # Use these backgrounds for n-th stage of hierarchical removals -# and test for backwards compatibility and that stage was written. +# and test for backwards compatibility and that stage was written. if opts.open_box: if h_inc_back_num >= 0 and h_iterations is not None and h_iterations != 0: - back_tsid = fp['background_h%s/timeslide_id' % h_inc_back_num][:] - back_ifar = fp['background_h%s/ifar' % h_inc_back_num][:] + back_tsid = fp["background_h%s/timeslide_id" % h_inc_back_num][:] + back_ifar = fp["background_h%s/ifar" % h_inc_back_num][:] # If the user wants to plot foreground events at h_inc_back_num == 0 # and h_iterations == 0 then no hierarchical removals were done and @@ -170,32 +203,31 @@ if opts.open_box: # is before hierarchical removals so the else below catches that. # The third case is requesting to plot hierarchical removal level # greater than was done and an empty plot is made earlier in the code. - else : - back_tsid = fp['background/timeslide_id'][:] - back_ifar = fp['background/ifar'][:] - -else : - # If the box is closed use the background from no hierarchical - # removals. Results should be agnostic of judgments about - # whether a hierarchical removal was done. - if h_iterations > 0 and h_iterations is not None: - back_tsid = fp['background_h0/timeslide_id'][:] - back_ifar = fp['background_h0/ifar'][:] - - # Otherwise the box is closed and no hierarchical removals - # were done so use top level background - else : - back_tsid = fp['background/timeslide_id'][:] - back_ifar = fp['background/ifar'][:] + else: + back_tsid = fp["background/timeslide_id"][:] + back_ifar = fp["background/ifar"][:] + +# If the box is closed use the background from no hierarchical +# removals. Results should be agnostic of judgments about +# whether a hierarchical removal was done. +elif h_iterations > 0 and h_iterations is not None: + back_tsid = fp["background_h0/timeslide_id"][:] + back_ifar = fp["background_h0/ifar"][:] + +# Otherwise the box is closed and no hierarchical removals +# were done so use top level background +else: + back_tsid = fp["background/timeslide_id"][:] + back_ifar = fp["background/ifar"][:] # make figure fig = plt.figure(1) # get a unique list of timeslide_ids and loop over them -interval = fp.attrs['timeslide_interval'] -ifo_joined = fp.attrs['ifos'].replace(' ','') -p_starts = fp['segments'][ifo_joined]['start'][:] -p_ends = fp['segments'][ifo_joined]['end'][:] +interval = fp.attrs["timeslide_interval"] +ifo_joined = fp.attrs["ifos"].replace(" ", "") +p_starts = fp["segments"][ifo_joined]["start"][:] +p_ends = fp["segments"][ifo_joined]["end"][:] pifo_segments = veto.start_end_to_segments(p_starts, p_ends) fifo_segments = copy.deepcopy(pifo_segments) min_tsid = (p_starts.min() - p_ends.max()) / interval @@ -209,8 +241,9 @@ else: min_tsid = numpy.rint(min_tsid) // opts.decimation_factor max_tsid = numpy.rint(max_tsid) // opts.decimation_factor - tsids = numpy.arange(min_tsid, max_tsid, 1).astype(numpy.int64) * \ - opts.decimation_factor + tsids = ( + numpy.arange(min_tsid, max_tsid, 1).astype(numpy.int64) * opts.decimation_factor + ) allbkg_dur = 0 allbkg_ifars = [] @@ -225,23 +258,25 @@ for tsid in tsids: ts_indx = numpy.where(back_tsid == tsid) # calculate the amount of coincident time in this time slide - offset = tsid*fp.attrs['timeslide_interval'] - back_dur = calculate_time_slide_duration(pifo_segments, fifo_segments, - offset=offset) + offset = tsid * fp.attrs["timeslide_interval"] + back_dur = calculate_time_slide_duration( + pifo_segments, fifo_segments, offset=offset + ) if back_dur == 0: continue plotted_slide_count += 1 allbkg_dur = allbkg_dur + back_dur # Limit how far back the decimated background gets plotted - red_trigs = back_ifar[ts_indx][back_ifar[ts_indx] > \ - opts.all_decimated_background_min_ifar] + red_trigs = back_ifar[ts_indx][ + back_ifar[ts_indx] > opts.all_decimated_background_min_ifar + ] allbkg_ifars.extend(red_trigs) # apply the correction factor for this time slide to its IFAR # you need a correction factor because the analyzed time of the time slide # is not the same as the analyzed time of the foreground - ts_ifar = back_ifar[ts_indx] * (fp.attrs['foreground_time'] / back_dur) + ts_ifar = back_ifar[ts_indx] * (fp.attrs["foreground_time"] / back_dur) ts_ifar.sort() # get the cumulative number of triggers in this time slide @@ -249,45 +284,80 @@ for tsid in tsids: # plot the time slide triggers if len(ts_ifar) > 1: - plt.loglog(ts_ifar, ts_cumnum, color='gray', alpha=0.4) + plt.loglog(ts_ifar, ts_cumnum, color="gray", alpha=0.4) elif len(ts_ifar) == 1: - plt.plot(ts_ifar, ts_cumnum, color='gray', marker='.', alpha=0.4) + plt.plot(ts_ifar, ts_cumnum, color="gray", marker=".", alpha=0.4) else: empty_slide_count += 1 allbkg_ifars = numpy.array(allbkg_ifars) -allbkg_ifars = allbkg_ifars * (fp.attrs['foreground_time'] / allbkg_dur ) +allbkg_ifars = allbkg_ifars * (fp.attrs["foreground_time"] / allbkg_dur) allbkg_ifars.sort() allbkg_cumnum = numpy.arange(len(allbkg_ifars), 0, -1) -plt.loglog(allbkg_ifars, allbkg_cumnum, color='green', linewidth=1.5, - label="All decimated background") +plt.loglog( + allbkg_ifars, + allbkg_cumnum, + color="green", + linewidth=1.5, + label="All decimated background", +) # plot the expected background -plt.loglog(expected_ifar, expected_cumnum, linestyle='--', linewidth=2, - color='black', label='Expected Background') +plt.loglog( + expected_ifar, + expected_cumnum, + linestyle="--", + linewidth=2, + color="black", + label="Expected Background", +) # plot the counting error error_plus = expected_cumnum + numpy.sqrt(expected_cumnum) error_minus = expected_cumnum - numpy.sqrt(expected_cumnum) -error_minus = numpy.where(error_minus<=0, 1e-5, error_minus) -plt.fill_between(expected_ifar, error_minus, error_plus, facecolor='y', - alpha=0.4, label='$N^{1/2}$ Errors') +error_minus = numpy.where(error_minus <= 0, 1e-5, error_minus) +plt.fill_between( + expected_ifar, + error_minus, + error_plus, + facecolor="y", + alpha=0.4, + label="$N^{1/2}$ Errors", +) # plot the counting error error_plus = expected_cumnum + 2 * numpy.sqrt(expected_cumnum) -error_minus = expected_cumnum - 2*numpy.sqrt(expected_cumnum) -error_minus = numpy.where(error_minus<=0, 1e-5, error_minus) -plt.fill_between(expected_ifar, error_minus, error_plus, facecolor='y', - alpha=0.2, label='$2N^{1/2}$ Errors') +error_minus = expected_cumnum - 2 * numpy.sqrt(expected_cumnum) +error_minus = numpy.where(error_minus <= 0, 1e-5, error_minus) +plt.fill_between( + expected_ifar, + error_minus, + error_plus, + facecolor="y", + alpha=0.2, + label="$2N^{1/2}$ Errors", +) # plot the foreground triggers if opts.open_box: - plt.loglog(fore_ifar, fore_cumnum, linestyle='None', color='blue', - marker='^', label='Foreground') + plt.loglog( + fore_ifar, + fore_cumnum, + linestyle="None", + color="blue", + marker="^", + label="Foreground", + ) if h_inc_back_num > 0: - plt.loglog(h_rm_ifar, h_rm_cumnum, linestyle='None', color='#b66dff', - marker='v', label='Hierarchically Removed Foreground') + plt.loglog( + h_rm_ifar, + h_rm_cumnum, + linestyle="None", + color="#b66dff", + marker="v", + label="Hierarchically Removed Foreground", + ) # format plot if opts.open_box and len(fore_cumnum) > 100: @@ -301,22 +371,27 @@ elif len(allbkg_cumnum) > 0: plt.ylim(0.8, 1.1 * len(allbkg_cumnum)) plt.xlim(0.9 * min(allbkg_ifars)) plt.grid() -plt.legend(loc='upper right', fontsize=9) -plt.ylabel('Cumulative Number') -plt.xlabel('Inverse False Alarm Rate (yr)') +plt.legend(loc="upper right", fontsize=9) +plt.ylabel("Cumulative Number") +plt.xlabel("Inverse False Alarm Rate (yr)") # save -caption = 'This is a cumulative histogram of triggers. The blue triangles ' \ - 'represent coincident foreground triggers. The dashed line ' \ - 'represents the expected background given the analysis time. The ' \ - 'shaded regions represent counting errors. The gray lines are ' \ - 'time slides treated as zero lag, here there are %d time slides ' \ - 'plotted. Gray dots are time slides with only one event. %d of ' \ - 'the plotted slides have zero events. The green line represents ' \ - 'all decimated time slides rescaled to the analysis time.' \ - % (plotted_slide_count, empty_slide_count) - -pycbc.results.save_fig_with_metadata(fig, opts.output_file, - title='Cumulative Number vs. IFAR', - caption=caption, - cmd=' '.join(sys.argv)) +caption = ( + "This is a cumulative histogram of triggers. The blue triangles " + "represent coincident foreground triggers. The dashed line " + "represents the expected background given the analysis time. The " + "shaded regions represent counting errors. The gray lines are " + "time slides treated as zero lag, here there are %d time slides " + "plotted. Gray dots are time slides with only one event. %d of " + "the plotted slides have zero events. The green line represents " + "all decimated time slides rescaled to the analysis time." + % (plotted_slide_count, empty_slide_count) +) + +pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + title="Cumulative Number vs. IFAR", + caption=caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_page_injtable b/bin/plotting/pycbc_page_injtable index 5cafa75009c..6434b67b763 100644 --- a/bin/plotting/pycbc_page_injtable +++ b/bin/plotting/pycbc_page_injtable @@ -1,137 +1,189 @@ #!/usr/bin/env python -""" Make a table of found injection information -""" +"""Make a table of found injection information""" + import argparse -import numpy as np import sys from itertools import combinations -import pycbc.results +import numpy as np + import pycbc.detector -import pycbc.pnutils import pycbc.events -from pycbc.io.hdf import HFile +import pycbc.pnutils +import pycbc.results from pycbc import add_common_pycbc_options, init_logging +from pycbc.io.hdf import HFile from pycbc.types import MultiDetOptionAction - parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--injection-file', - help='HDF File containing the matched injections') -parser.add_argument('--single-trigger-files', nargs='*', - action=MultiDetOptionAction, - help="HDF format single detector trigger files") -parser.add_argument('--show-missed', action='store_true') -parser.add_argument('--output-file') +parser.add_argument( + "--injection-file", help="HDF File containing the matched injections" +) +parser.add_argument( + "--single-trigger-files", + nargs="*", + action=MultiDetOptionAction, + help="HDF format single detector trigger files", +) +parser.add_argument("--show-missed", action="store_true") +parser.add_argument("--output-file") args = parser.parse_args() init_logging(args.verbose) -f = HFile(args.injection_file,'r') -inj = f['injections'] +f = HFile(args.injection_file, "r") +inj = f["injections"] found_cols, found_names, found_formats = [], [], [] -ifos = f.attrs['ifos'].split(' ') +ifos = f.attrs["ifos"].split(" ") if args.show_missed: title = "Missed Injections" - idx = f['missed/after_vetoes'][:] + idx = f["missed/after_vetoes"][:] else: title = "Found Injections" - found = f['found_after_vetoes'] - idx = found['injection_index'][:] - detectors = f.attrs['ifos'].split(' ') - keys = f['found_after_vetoes'].keys() + found = f["found_after_vetoes"] + idx = found["injection_index"][:] + detectors = f.attrs["ifos"].split(" ") + keys = f["found_after_vetoes"].keys() detectors_used = [] - found = f['found_after_vetoes'] + found = f["found_after_vetoes"] for det in detectors: - if(det in keys): + if det in keys: detectors_used.append(det) - det_two_combo= np.array(list(combinations(detectors_used,2))) + det_two_combo = np.array(list(combinations(detectors_used, 2))) tdiff = [] tdiff_str = [] - tdiff_format =[] + tdiff_format = [] for i in range(len(det_two_combo)): - time_1 = np.array(found[det_two_combo[i,0]+'/time'][:]) - time_2 = np.array(found[det_two_combo[i,1]+'/time'][:]) + time_1 = np.array(found[det_two_combo[i, 0] + "/time"][:]) + time_2 = np.array(found[det_two_combo[i, 1] + "/time"][:]) tdiff_vals = (time_1 - time_2) * 1000 tdiff_vals[np.logical_or(time_1 < 0, time_2 < 0)] = np.nan - tdiff_1 = ['%.2f' % td if not np.isnan(td) else ' ' for td in tdiff_vals] + tdiff_1 = ["%.2f" % td if not np.isnan(td) else " " for td in tdiff_vals] tdiff.append(tdiff_1) - tdiff_head= '%s - %s time (ms)' % (det_two_combo[i,0], det_two_combo[i,1]) + tdiff_head = "%s - %s time (ms)" % (det_two_combo[i, 0], det_two_combo[i, 1]) tdiff_str.append(tdiff_head) - tdiff_format.append('##.##') - ids = {detector:found[detector+'/trigger_id'][:] for detector in detectors_used} - - found_cols = [found['stat'], found['ifar_exc']] + tdiff - found_names = ['Ranking Stat.', 'Exc. IFAR'] + tdiff_str - found_formats = ['##.##', '##.##'] + tdiff_format + tdiff_format.append("##.##") + ids = {detector: found[detector + "/trigger_id"][:] for detector in detectors_used} + found_cols = [found["stat"], found["ifar_exc"]] + tdiff + found_names = ["Ranking Stat.", "Exc. IFAR"] + tdiff_str + found_formats = ["##.##", "##.##"] + tdiff_format if args.single_trigger_files: for ifo in args.single_trigger_files: - f = HFile(args.single_trigger_files[ifo], 'r')[ifo] + f = HFile(args.single_trigger_files[ifo], "r")[ifo] ids_ifo = np.array(ids[ifo]) ids_na = np.argwhere(ids_ifo == -1) - snr_vals = f['snr'][:][ids_ifo] + snr_vals = f["snr"][:][ids_ifo] snr_vals[ids_ifo == -1] = np.nan - chisq_vals = f['chisq'][:][ids_ifo] / (2 * f['chisq_dof'][:][ids_ifo] - 2) + chisq_vals = f["chisq"][:][ids_ifo] / (2 * f["chisq_dof"][:][ids_ifo] - 2) chisq_vals[ids_ifo == -1] = np.nan newsnr_vals = pycbc.events.ranking.newsnr(snr_vals, chisq_vals) - snr = ['%.2f' % s if not np.isnan(s) else ' ' for s in snr_vals] - chisq = ['%.2f' % c if not np.isnan(c) else ' ' for c in chisq_vals] - newsnr = ['%.2f' % s if not np.isnan(s) else ' ' for s in newsnr_vals] + snr = ["%.2f" % s if not np.isnan(s) else " " for s in snr_vals] + chisq = ["%.2f" % c if not np.isnan(c) else " " for c in chisq_vals] + newsnr = ["%.2f" % s if not np.isnan(s) else " " for s in newsnr_vals] found_names += [ifo + " SNR", ifo + " CHISQ", ifo + " NewSNR"] found_cols += [snr, chisq, newsnr] - found_formats += ['##.##', '##.##', '##.##'] + found_formats += ["##.##", "##.##", "##.##"] -eff_dist = {'eff_dist_%s' % i[0].lower() : 'Eff Dist (%s)' % i for i in ifos} +eff_dist = {"eff_dist_%s" % i[0].lower(): "Eff Dist (%s)" % i for i in ifos} keys = inj.keys() eff_dist_str = [] eff_distance = [] eff_dist_format = [] -for dist in eff_dist : - ifo = ('%s1' % dist.split('_')[-1]).upper() +for dist in eff_dist: + ifo = ("%s1" % dist.split("_")[-1]).upper() d = pycbc.detector.Detector(ifo) edist = d.effective_distance( - inj['distance'][:][idx], - inj['ra'][:][idx], - inj['dec'][:][idx], - inj['polarization'][:][idx], - inj['tc'][:][idx], - inj['inclination'][:][idx]) + inj["distance"][:][idx], + inj["ra"][:][idx], + inj["dec"][:][idx], + inj["polarization"][:][idx], + inj["tc"][:][idx], + inj["inclination"][:][idx], + ) eff_distance.append(edist) eff_dist_str.append(eff_dist[dist]) - eff_dist_format.append('##.##') + eff_dist_format.append("##.##") dec_dist = np.max(eff_distance, 0) -m1, m2 = inj['mass1'][:][idx], inj['mass2'][:][idx] +m1, m2 = inj["mass1"][:][idx], inj["mass2"][:][idx] mchirp, eta = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2) dec_chirp_dist = pycbc.pnutils.chirp_distance(dec_dist, mchirp) -columns = [dec_chirp_dist, inj['tc'][:][idx], m1, m2, mchirp, eta, - inj['spin1x'][:][idx], inj['spin1y'][:][idx], inj['spin1z'][:][idx], - inj['spin2x'][:][idx], inj['spin2y'][:][idx], inj['spin2z'][:][idx], - inj['distance'][:][idx]] + eff_distance + found_cols - -names = ['DChirp Dist', 'Inj Time', 'Mass1', 'Mass2', 'Mchirp', 'Eta', - 's1x', 's1y', 's1z', - 's2x', 's2y', 's2z', - 'Dist'] + eff_dist_str + found_names - -format_strings = ['##.##', '##.##', '##.##', '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', - '##.##'] + eff_dist_format + found_formats +columns = ( + [ + dec_chirp_dist, + inj["tc"][:][idx], + m1, + m2, + mchirp, + eta, + inj["spin1x"][:][idx], + inj["spin1y"][:][idx], + inj["spin1z"][:][idx], + inj["spin2x"][:][idx], + inj["spin2y"][:][idx], + inj["spin2z"][:][idx], + inj["distance"][:][idx], + ] + + eff_distance + + found_cols +) + +names = ( + [ + "DChirp Dist", + "Inj Time", + "Mass1", + "Mass2", + "Mchirp", + "Eta", + "s1x", + "s1y", + "s1z", + "s2x", + "s2y", + "s2z", + "Dist", + ] + + eff_dist_str + + found_names +) + +format_strings = ( + [ + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + ] + + eff_dist_format + + found_formats +) columns = [np.array(col) for col in columns] -html_table = pycbc.results.html_table(columns, names, - format_strings=format_strings, - page_size=20) - -kwds = { 'title' : title, - 'caption' : "A table of %s and their coincident statistic information." % title.lower(), - 'cmd' :' '.join(sys.argv), } +html_table = pycbc.results.html_table( + columns, names, format_strings=format_strings, page_size=20 +) + +kwds = { + "title": title, + "caption": "A table of %s and their coincident statistic information." + % title.lower(), + "cmd": " ".join(sys.argv), +} pycbc.results.save_fig_with_metadata(str(html_table), args.output_file, **kwds) diff --git a/bin/plotting/pycbc_page_recovery b/bin/plotting/pycbc_page_recovery index f960f325e88..808d42f6840 100644 --- a/bin/plotting/pycbc_page_recovery +++ b/bin/plotting/pycbc_page_recovery @@ -1,7 +1,13 @@ #!/usr/bin/python -''' Make plots of recovered injection parameters -''' -import numpy, logging, argparse, sys, matplotlib +"""Make plots of recovered injection parameters""" + +import argparse +import logging +import sys + +import matplotlib +import numpy + matplotlib.use("Agg") import matplotlib.pyplot as plot @@ -12,40 +18,68 @@ from pycbc.io.hdf import HFile parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument("--injection-file", required=True, - help="hdf injection file containing found injections. " - "Required") -parser.add_argument("--bank-file", - help="hdf bank file containing template parameters. " - "Required for some parameters, ex. chirp mass") -parser.add_argument("--trigger-file", - help="hdf trigger merge file containing single triggers. " - "Required for some parameters, ex. end_time") -parser.add_argument("--ifo", required=True, - help="Detector where injections are recovered. Required") -parser.add_argument("--min-ifar", default=0, - help="Minimum IFAR for which to plot injections. Units:years") -parser.add_argument("--error-param", required=True, - help="Property to be compared between injected and recovered " - "event, ex. 'mchirp'. Required") -parser.add_argument("--x-param", - help="Injected parameter to plot on x-axis. Required if " - "--plot-type is err_v_param or fracerr_v_param") -parser.add_argument("--log-x", action="store_true", - help="Use a logarithmic x-axis.") -parser.add_argument("--log-y", action="store_true", - help="Use a logarithmic y-axis.") -parser.add_argument("--plot-type", choices=["scatter", "err", "fracerr", - "err_v_param", "fracerr_v_param", "errhist", - "fracerrhist", "ratio", "ratio_v_param"], - default="scatter", - help="Type of plot to show accuracy of recovery. " - "Default='scatter'") -parser.add_argument("--gradient-far", action="store_true", - help="Show FAR of found injections as a color gradient") -parser.add_argument("--f-lower", type=int, - help="lower frequency value for e.g. calculation of " - "template duration") +parser.add_argument( + "--injection-file", + required=True, + help="hdf injection file containing found injections. Required", +) +parser.add_argument( + "--bank-file", + help="hdf bank file containing template parameters. " + "Required for some parameters, ex. chirp mass", +) +parser.add_argument( + "--trigger-file", + help="hdf trigger merge file containing single triggers. " + "Required for some parameters, ex. end_time", +) +parser.add_argument( + "--ifo", required=True, help="Detector where injections are recovered. Required" +) +parser.add_argument( + "--min-ifar", + default=0, + help="Minimum IFAR for which to plot injections. Units:years", +) +parser.add_argument( + "--error-param", + required=True, + help="Property to be compared between injected and recovered " + "event, ex. 'mchirp'. Required", +) +parser.add_argument( + "--x-param", + help="Injected parameter to plot on x-axis. Required if " + "--plot-type is err_v_param or fracerr_v_param", +) +parser.add_argument("--log-x", action="store_true", help="Use a logarithmic x-axis.") +parser.add_argument("--log-y", action="store_true", help="Use a logarithmic y-axis.") +parser.add_argument( + "--plot-type", + choices=[ + "scatter", + "err", + "fracerr", + "err_v_param", + "fracerr_v_param", + "errhist", + "fracerrhist", + "ratio", + "ratio_v_param", + ], + default="scatter", + help="Type of plot to show accuracy of recovery. Default='scatter'", +) +parser.add_argument( + "--gradient-far", + action="store_true", + help="Show FAR of found injections as a color gradient", +) +parser.add_argument( + "--f-lower", + type=int, + help="lower frequency value for e.g. calculation of template duration", +) parser.add_argument("--output-file", required=True) args = parser.parse_args() @@ -80,10 +114,10 @@ ifar_found = ifar_found[above_min_ifar] # need to hack end time and eff dist as different between detectors param = args.error_param if args.error_param == "end_time": - param = "end_time_"+site + param = "end_time_" + site if args.error_param == "effective_distance": - param = "eff_dist_"+site -# calculate the parameter for all injections, then select those found and + param = "eff_dist_" + site +# calculate the parameter for all injections, then select those found and # above the IFAR threshold inj_params = triggers.get_inj_param(injs, param, args.ifo, args)[found] @@ -91,9 +125,9 @@ inj_params = triggers.get_inj_param(injs, param, args.ifo, args)[found] if args.x_param: xparam = args.x_param if args.x_param == "end_time": - xparam = "end_time_"+site + xparam = "end_time_" + site if args.x_param == "effective_distance": - xparam = "eff_dist_"+site + xparam = "eff_dist_" + site inj_xparams = triggers.get_inj_param(injs, xparam, args.ifo, args)[found] # need a string value when defining axis labels even if not used args.x_param = args.x_param or "badger" @@ -103,8 +137,9 @@ args.x_param = args.x_param or "badger" param = args.error_param # calculate parameters for all recovered injection triggers -rec_params, found_in_ifo = triggers.get_found_param(injs, bank, trig, param, - args.ifo, args) +rec_params, found_in_ifo = triggers.get_found_param( + injs, bank, trig, param, args.ifo, args +) # select values above IFAR threshold rec_params = rec_params[above_min_ifar] found_in_ifo = found_in_ifo[above_min_ifar] @@ -115,8 +150,9 @@ if not all(found_in_ifo): inj_params = inj_params[found_in_ifo] inj_xparams = inj_xparams[found_in_ifo] ifar_found = ifar_found[found_in_ifo] - logging.info("Removing %i inj not found in %s" % - ((found_in_ifo == False).sum(), args.ifo)) + logging.info( + "Removing %i inj not found in %s" % ((found_in_ifo == False).sum(), args.ifo) + ) # calculate needed values if "err" in args.plot_type: @@ -127,49 +163,49 @@ if "ratio" in args.plot_type: errrat = rec_params / inj_params x_dict = { - "scatter" : "inj_params", - "err" : "inj_params", - "fracerr" : "inj_params", - "ratio" : "inj_params", - "err_v_param" : "inj_xparams", - "fracerr_v_param" : "inj_xparams", - "ratio_v_param" : "inj_xparams", - "errhist" : "diff", - "fracerrhist" : "reldiff" + "scatter": "inj_params", + "err": "inj_params", + "fracerr": "inj_params", + "ratio": "inj_params", + "err_v_param": "inj_xparams", + "fracerr_v_param": "inj_xparams", + "ratio_v_param": "inj_xparams", + "errhist": "diff", + "fracerrhist": "reldiff", } xlabels = { - "scatter" : "Injected "+args.error_param, - "err" : "Injected "+args.error_param, - "fracerr" : "Injected "+args.error_param, - "ratio" : "Injected "+args.error_param, - "err_v_param" : "Injected "+args.x_param, - "fracerr_v_param" : "Injected "+args.x_param, - "ratio_v_param" : "Injected "+args.x_param, - "errhist" : "Error (rec-inj) in "+args.error_param, - "fracerrhist" : "Fractional error (rec-inj)/inj in "+args.error_param + "scatter": "Injected " + args.error_param, + "err": "Injected " + args.error_param, + "fracerr": "Injected " + args.error_param, + "ratio": "Injected " + args.error_param, + "err_v_param": "Injected " + args.x_param, + "fracerr_v_param": "Injected " + args.x_param, + "ratio_v_param": "Injected " + args.x_param, + "errhist": "Error (rec-inj) in " + args.error_param, + "fracerrhist": "Fractional error (rec-inj)/inj in " + args.error_param, } y_dict = { - "scatter" : "rec_params", - "err" : "diff", - "fracerr" : "reldiff", - "ratio" : "errrat", - "err_v_param" : "diff", - "fracerr_v_param" : "reldiff", - "ratio_v_param" : "errrat", - "errhist" : "None", - "fracerrhist" : "None" + "scatter": "rec_params", + "err": "diff", + "fracerr": "reldiff", + "ratio": "errrat", + "err_v_param": "diff", + "fracerr_v_param": "reldiff", + "ratio_v_param": "errrat", + "errhist": "None", + "fracerrhist": "None", } ylabels = { - "scatter" : "Recovered "+args.error_param, - "err" : "Error (rec-inj) in "+args.error_param, - "fracerr" : "Fractional error (rec-inj)/inj in "+args.error_param, - "ratio" : "Ratio rec/inj in "+args.error_param, - "err_v_param" : "Error (rec-inj) in "+args.error_param, - "fracerr_v_param" : "Fractional error (rec-inj)/inj in "+args.error_param, - "ratio_v_param" : "Ratio rec/inj in "+args.error_param, - "errhist" : "Number of injections", - "fracerrhist" : "Number of injections" + "scatter": "Recovered " + args.error_param, + "err": "Error (rec-inj) in " + args.error_param, + "fracerr": "Fractional error (rec-inj)/inj in " + args.error_param, + "ratio": "Ratio rec/inj in " + args.error_param, + "err_v_param": "Error (rec-inj) in " + args.error_param, + "fracerr_v_param": "Fractional error (rec-inj)/inj in " + args.error_param, + "ratio_v_param": "Ratio rec/inj in " + args.error_param, + "errhist": "Number of injections", + "fracerrhist": "Number of injections", } xvals = eval(x_dict[args.plot_type]) @@ -187,11 +223,22 @@ if "hist" not in args.plot_type: thousand = numpy.where(ifar_found > 1000)[0] color[hundred] = 0.5 color[thousand] = 1.0 - caption = ("Found injections: blue circles are found with IFAR<100yr, " - "green are IFAR<1000yr, red are IFAR >=1000yr") - points = plot.scatter(xvals, yvals, c=color, linewidth=0, vmin=0, - vmax=1, s=12, marker="o", label="found", - alpha=0.6) + caption = ( + "Found injections: blue circles are found with IFAR<100yr, " + "green are IFAR<1000yr, red are IFAR >=1000yr" + ) + points = plot.scatter( + xvals, + yvals, + c=color, + linewidth=0, + vmin=0, + vmax=1, + s=12, + marker="o", + label="found", + alpha=0.6, + ) else: # make a pretty rainbow coloured plot @@ -201,11 +248,17 @@ if "hist" not in args.plot_type: xvals = xvals[csort] yvals = yvals[csort] color = color[csort] - caption = ("Found injections: color indicates the estimated false " - "alarm rate") - points = plot.scatter(xvals, yvals, c=color, linewidth=0, - norm=matplotlib.colors.LogNorm(), - s=16, marker="o", label="found") + caption = "Found injections: color indicates the estimated false alarm rate" + points = plot.scatter( + xvals, + yvals, + c=color, + linewidth=0, + norm=matplotlib.colors.LogNorm(), + s=16, + marker="o", + label="found", + ) plot.subplots_adjust(right=0.99) try: c = plot.colorbar() @@ -216,33 +269,38 @@ if "hist" not in args.plot_type: raise else: # it's a histogram! - plot.hist(xvals, bins=int(len(xvals)**0.5), histtype="step", label="found") - caption = ("Found injections") + plot.hist(xvals, bins=int(len(xvals) ** 0.5), histtype="step", label="found") + caption = "Found injections" ax = plot.gca() if args.log_x: ax.set_xscale("log") if args.log_y: ax.set_yscale("log") -plot.xlabel(xlabels[args.plot_type], size='large') -plot.ylabel(ylabels[args.plot_type], size='large') +plot.xlabel(xlabels[args.plot_type], size="large") +plot.ylabel(ylabels[args.plot_type], size="large") plot.grid(True) fig_kwds = {} if ".png" in args.output_file: fig_kwds["dpi"] = 150 -if (".html" in args.output_file): +if ".html" in args.output_file: plot.subplots_adjust(left=0.1, right=0.8, top=0.9, bottom=0.1) - import mpld3, mpld3.plugins, mpld3.utils + import mpld3 + import mpld3.plugins + import mpld3.utils + mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fmt=".5g")) - legend = mpld3.plugins.InteractiveLegendPlugin([points], ["found"], - alpha_unsel=0.1) + legend = mpld3.plugins.InteractiveLegendPlugin([points], ["found"], alpha_unsel=0.1) mpld3.plugins.connect(fig, legend) -results.save_fig_with_metadata(fig, args.output_file, - fig_kwds=fig_kwds, - title="Injection %s recovery" % (args.error_param), - cmd=" ".join(sys.argv), - caption=caption) +results.save_fig_with_metadata( + fig, + args.output_file, + fig_kwds=fig_kwds, + title="Injection %s recovery" % (args.error_param), + cmd=" ".join(sys.argv), + caption=caption, +) logging.info("Done!") diff --git a/bin/plotting/pycbc_page_segments b/bin/plotting/pycbc_page_segments index 75b1d4ab7c2..66fbf6f811c 100644 --- a/bin/plotting/pycbc_page_segments +++ b/bin/plotting/pycbc_page_segments @@ -1,29 +1,30 @@ #!/usr/bin/env python -""" Make interactive visualization of segments -""" +"""Make interactive visualization of segments""" + import argparse from itertools import cycle + import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt -import numpy + +matplotlib.use("Agg") import mpld3 import mpld3.plugins +import numpy +from matplotlib import pyplot as plt from matplotlib.patches import Rectangle import pycbc.events from pycbc.results.mpld3_utils import MPLSlide, Tooltip - parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--segment-files', nargs='+', - help="List of segment files to plot") -parser.add_argument('--output-file', help="output html file") +parser.add_argument("--segment-files", nargs="+", help="List of segment files to plot") +parser.add_argument("--output-file", help="output html file") args = parser.parse_args() pycbc.init_logging(args.verbose) + def timestr(s): t = "" s = int(s) @@ -39,31 +40,37 @@ def timestr(s): t += "%ss " % s return t + def get_name(segment_file): - from igwn_ligolw import ligolw, utils as ligolw_utils + from igwn_ligolw import ligolw + from igwn_ligolw import utils as ligolw_utils + from pycbc.io.ligolw import LIGOLWContentHandler indoc = ligolw_utils.load_filename( - segment_file, False, contenthandler=LIGOLWContentHandler) - n = ligolw.Table.get_table(indoc, 'segment_definer')[0] + segment_file, False, contenthandler=LIGOLWContentHandler + ) + n = ligolw.Table.get_table(indoc, "segment_definer")[0] return "%s:%s:%s" % (n.ifos, n.name, n.version) + def plot_segs(start, end, color=None, y=0, h=1): patches = [] - if not hasattr(plot_segs, 'colors'): - plot_segs.colors = cycle(['red', 'blue', 'green', 'yellow', 'cyan', 'violet']) + if not hasattr(plot_segs, "colors"): + plot_segs.colors = cycle(["red", "blue", "green", "yellow", "cyan", "violet"]) if color is None: color = next(plot_segs.colors) for s, e in zip(start, end): ax = plt.gca() - patch = Rectangle((s, y), (e-s), h, facecolor=color) + patch = Rectangle((s, y), (e - s), h, facecolor=color) ax.add_patch(patch) patches.append(patch) return patches + # Define some CSS to control our custom labels css = """ table @@ -93,15 +100,15 @@ ax = fig.gca() names = [] smin, smax = -1, 1 for i, seg_file in enumerate(sorted(args.segment_files)[::-1]): - y = i + .05 - h = .7 + y = i + 0.05 + h = 0.7 name = get_name(seg_file) start, end = pycbc.events.start_end_from_segments(seg_file) if len(start) > 0: # remove duplicate segments, e.g. inspiral jobs with split banks unique_startend = numpy.array(list(set(zip(start, end)))) - start, end = unique_startend[:,0], unique_startend[:,1] + start, end = unique_startend[:, 0], unique_startend[:, 1] dur = end - start total = timestr(abs(pycbc.events.start_end_to_segments(start, end).coalesce())) @@ -112,10 +119,10 @@ for i, seg_file in enumerate(sorted(args.segment_files)[::-1]): Duration%s """ - smin = start.min() if len(start) and ( start.min() < smin or smin == -1) else smin + smin = start.min() if len(start) and (start.min() < smin or smin == -1) else smin smax = end.max() if len(end) and end.max() > smax else smax - names += [(name, total, y + h + .1)] + names += [(name, total, y + h + 0.1)] patches = plot_segs(start, end, y=y, h=h) for i, p in enumerate(patches): @@ -131,10 +138,10 @@ for name, total, h in names: ax.set_ylim(0, h + 0.2) ax.set_xlim(smin, smax) ax.set_yticks([]) -ax.set_xlabel('GPS Time (s)') +ax.set_xlabel("GPS Time (s)") -mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fontsize=14, fmt='10f')) +mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fontsize=14, fmt="10f")) mpld3.plugins.connect(fig, mpld3.plugins.BoxZoom()) mpld3.plugins.connect(fig, MPLSlide()) mpld3.plugins.connect(fig, mpld3.plugins.Reset()) -mpld3.save_html(fig, open(args.output_file, 'w')) +mpld3.save_html(fig, open(args.output_file, "w")) diff --git a/bin/plotting/pycbc_page_segplot b/bin/plotting/pycbc_page_segplot index e50383a2aa6..314b51fafe7 100644 --- a/bin/plotting/pycbc_page_segplot +++ b/bin/plotting/pycbc_page_segplot @@ -17,12 +17,19 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import argparse -import matplotlib; matplotlib.use('Agg') -import matplotlib.pyplot as plt -import numpy, pycbc.events, mpld3, mpld3.plugins + +import matplotlib + +matplotlib.use("Agg") import sys from itertools import cycle +import matplotlib.pyplot as plt +import mpld3 +import mpld3.plugins +import numpy + +import pycbc.events from pycbc import add_common_pycbc_options, init_logging from pycbc.events.veto import get_segment_definer_comments from pycbc.results.mpld3_utils import LineTooltip @@ -31,23 +38,36 @@ from pycbc.workflow import SegFile # parse command line parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--segment-files', type=str, nargs="+", - help='XML files with a segment definer table to read.') -parser.add_argument('--segment-names', type=str, nargs="+", required=False, - help='Names of segments in the segment definer table.') -parser.add_argument('--output-file', type=str, - help='Path of the output HTML file.') -parser.add_argument('--ifos', nargs='+', default=['H1', 'L1'], - help='Space-separated list of IFOs to plot, default H1 L1.') +parser.add_argument( + "--segment-files", + type=str, + nargs="+", + help="XML files with a segment definer table to read.", +) +parser.add_argument( + "--segment-names", + type=str, + nargs="+", + required=False, + help="Names of segments in the segment definer table.", +) +parser.add_argument("--output-file", type=str, help="Path of the output HTML file.") +parser.add_argument( + "--ifos", + nargs="+", + default=["H1", "L1"], + help="Space-separated list of IFOs to plot, default H1 L1.", +) opts = parser.parse_args() init_logging(opts.verbose) + def timestr(s): - """ Takes seconds and returns a human-readable string for the amount + """ + Takes seconds and returns a human-readable string for the amount of time. """ - t = "" s = int(s) d = s / 86400 @@ -62,14 +82,15 @@ def timestr(s): t += "%ss " % s return t + # FIXME set colors color_cycle = {} -color_cycle['H1'] = cycle(['#ee0000', '#8c000f', '#840000', '#653700']) -color_cycle['L1'] = cycle(['#4ba6ff', '#d0fefe', '#0165fc', '#001146']) -color_cycle['V1'] = cycle(['#9b59b6', '#c20078', '#9a0eea', '#7e1e9c']) -color_cycle['K1'] = cycle(['#ffb200', '#fdde6c', '#fdb915', '#d8863b']) -color_cycle['I1'] = cycle(['#b0dd8b', 'ForestGreen', 'DarkGreen', 'DarkOliveGreen']) -color_cycle['G1'] = cycle(['#222222', '#d8dcd6', '#363737', '#000000']) +color_cycle["H1"] = cycle(["#ee0000", "#8c000f", "#840000", "#653700"]) +color_cycle["L1"] = cycle(["#4ba6ff", "#d0fefe", "#0165fc", "#001146"]) +color_cycle["V1"] = cycle(["#9b59b6", "#c20078", "#9a0eea", "#7e1e9c"]) +color_cycle["K1"] = cycle(["#ffb200", "#fdde6c", "#fdb915", "#d8863b"]) +color_cycle["I1"] = cycle(["#b0dd8b", "ForestGreen", "DarkGreen", "DarkOliveGreen"]) +color_cycle["G1"] = cycle(["#222222", "#d8dcd6", "#363737", "#000000"]) # set default plugins mpld3.plugins.DEFAULT_PLUGINS = [] @@ -81,7 +102,7 @@ names = [] smin, smax = numpy.inf, -numpy.inf # set height of rectangles -h = .7 +h = 0.7 ifos = opts.ifos @@ -92,33 +113,31 @@ caption = "Visualization of the input and output segments for the search. Shown line_collections = [] # create a figure -fig, ax = plt.subplots(figsize=(16,9)) +fig, ax = plt.subplots(figsize=(16, 9)) # loop over segment XML files i = 0 seg_list = [] s = [] for segment_file in opts.segment_files: - # read segment definer table seg_dict = SegFile.from_segment_xml(segment_file).segment_dict # read comments from segment definer table - comment_dict = get_segment_definer_comments(open(segment_file, 'rb'), - include_version=False) + comment_dict = get_segment_definer_comments( + open(segment_file, "rb"), include_version=False + ) # loop over segment names for segment_name in opts.segment_names: - # get segments for ifo in ifos: - # get new color for this segment name color = next(color_cycle[ifo]) for key in seg_dict.keys(): - seg_key = ifo+':'+segment_name - if key == seg_key or key.startswith(ifo+':'+segment_name+':'): + seg_key = ifo + ":" + segment_name + if key == seg_key or key.startswith(ifo + ":" + segment_name + ":"): segs = seg_dict[key] # increment y position of bits @@ -128,9 +147,9 @@ for segment_file in opts.segment_files: continue # put comment in caption - caption += ifo + ':' + segment_name + caption += ifo + ":" + segment_name if comment_dict[key] != None: - caption += " ("+comment_dict[key]+")" + caption += " (" + comment_dict[key] + ")" caption += " " # get a start time and end time array @@ -146,8 +165,8 @@ for segment_file in opts.segment_files: # plot segments sub_line_collections = [] - for s,e in zip(start, end): - l = ax.plot([s,e], [y,y], '-', lw=120, color=color, alpha=0.3) + for s, e in zip(start, end): + l = ax.plot([s, e], [y, y], "-", lw=120, color=color, alpha=0.3) sub_line_collections += l # set HTML table string @@ -157,53 +176,57 @@ for segment_file in opts.segment_files: End%.0f Duration%s - """ % (segment_name, s, e, e-s) + """ % (segment_name, s, e, e - s) # add tooltip for segment if len(line_collections) < 1: - mpld3.plugins.connect(fig, mpld3.plugins.LineHTMLTooltip(l[0], label)) + mpld3.plugins.connect( + fig, mpld3.plugins.LineHTMLTooltip(l[0], label) + ) else: mpld3.plugins.connect(fig, LineTooltip(l[0], label)) # add name to list - names += [ifo+':'+segment_name] + names += [ifo + ":" + segment_name] # add list of lines to list line_collections.append(sub_line_collections) # loop over IFOs for ifo in ifos: - # increment y position of bits y = ifos.index(ifo) + 2 * 0.33 # set y position of text - ax.text(smin, y, ifo+' Segments', size=24) + ax.text(smin, y, ifo + " Segments", size=24) # add interactive legend -interactive_legend = mpld3.plugins.InteractiveLegendPlugin(line_collections, - names, - alpha_unsel=0.1) +interactive_legend = mpld3.plugins.InteractiveLegendPlugin( + line_collections, names, alpha_unsel=0.1 +) mpld3.plugins.connect(fig, interactive_legend) # format the plot plt.ylim(0, len(ifos) * (h + 0.2)) plt.xlim(smin, smax) -plt.xlabel('GPS Time (s)') +plt.xlabel("GPS Time (s)") # add whitespace for the legend fig.subplots_adjust(left=0.0, right=0.6, top=0.9, bottom=0.1) # add plugins to the plot -mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fontsize=14, fmt='10f')) +mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fontsize=14, fmt="10f")) mpld3.plugins.connect(fig, mpld3.plugins.BoxZoom()) mpld3.plugins.connect(fig, mpld3.plugins.Zoom()) mpld3.plugins.connect(fig, mpld3.plugins.Reset()) # save the plot as an interactive HTML fig_kwds = {} -pycbc.results.save_fig_with_metadata(fig, opts.output_file, - fig_kwds=fig_kwds, - title='Data Input and Search Output Segments', - cmd=' '.join(sys.argv), - caption=caption) +pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + fig_kwds=fig_kwds, + title="Data Input and Search Output Segments", + cmd=" ".join(sys.argv), + caption=caption, +) diff --git a/bin/plotting/pycbc_page_segtable b/bin/plotting/pycbc_page_segtable index 4819c483886..35a4f3c0b7b 100644 --- a/bin/plotting/pycbc_page_segtable +++ b/bin/plotting/pycbc_page_segtable @@ -17,38 +17,59 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import argparse -import numpy -import pycbc.results -import sys import itertools +import sys import igwn_segments as segments +import numpy +import pycbc.results from pycbc.events.veto import get_segment_definer_comments from pycbc.results import save_fig_with_metadata from pycbc.workflow import SegFile + def powerset_ifos(ifo_set): combo_set = [] for n_ifos in range(1, len(ifo_set) + 1): combo_set += itertools.combinations(ifo_set, n_ifos) return combo_set + # parse command line parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--segment-files', type=str, nargs="+", - help='XML files with a segment definer table to read.') -parser.add_argument('--segment-names', type=str, nargs="+", required=False, default="", - help='Names of segments in the segment definer table.') -parser.add_argument('--description', type=str, required=False, default="", - help='Additional descriptive text for caption.') -parser.add_argument('--title-text', type=str, required=False, - help='Additional text to append to title.') -parser.add_argument('--output-file', type=str, - help='Path of the output HTML file.') -parser.add_argument('--ifos', nargs='+', default=['H1', 'L1'], - help='Space-separated list of detectors: default H1 L1.') +parser.add_argument( + "--segment-files", + type=str, + nargs="+", + help="XML files with a segment definer table to read.", +) +parser.add_argument( + "--segment-names", + type=str, + nargs="+", + required=False, + default="", + help="Names of segments in the segment definer table.", +) +parser.add_argument( + "--description", + type=str, + required=False, + default="", + help="Additional descriptive text for caption.", +) +parser.add_argument( + "--title-text", type=str, required=False, help="Additional text to append to title." +) +parser.add_argument("--output-file", type=str, help="Path of the output HTML file.") +parser.add_argument( + "--ifos", + nargs="+", + default=["H1", "L1"], + help="Space-separated list of detectors: default H1 L1.", +) opts = parser.parse_args() # setup log @@ -58,11 +79,11 @@ pycbc.init_logging(opts.verbose) ifo_combinations = powerset_ifos(opts.ifos) # set column names -columns = (('Name', []),) +columns = (("Name", []),) # columns for detectors and combinations for combo in ifo_combinations: - det_combo_string = ''.join(combo).upper() - columns += (('%s Time (s)' % det_combo_string, []),) + det_combo_string = "".join(combo).upper() + columns += (("%s Time (s)" % det_combo_string, []),) caption = "This table shows the cumulative amount of time for each segment. Shown are:" @@ -70,17 +91,16 @@ caption = "This table shows the cumulative amount of time for each segment. Show seg_dict = {} comment_dict = {} for segment_file in opts.segment_files: - # read segment definer table seg_dict.update(SegFile.from_segment_xml(segment_file).segment_dict) - comment_dict.update(get_segment_definer_comments(open(segment_file, 'rb'), - include_version=False)) + comment_dict.update( + get_segment_definer_comments(open(segment_file, "rb"), include_version=False) + ) # loop over segment names for segment_name in opts.segment_names: - # allow user to find the and of segments if they use a "&" on the command line - names = segment_name.split('&') + names = segment_name.split("&") # create a dict that will hold first segment_name in list for each IFO base_segs = segments.segmentlistdict({}) @@ -90,14 +110,12 @@ for segment_name in opts.segment_names: # loop over IFOs for ifo in opts.ifos: - # make an empty list for each IFO base_segs[ifo] = segments.segmentlist([]) comp_segs[ifo] = segments.segmentlist([]) # loop over names from string split on "&" for name in names: - # construct the first part of the segment flag # note that if you don't include a version on the command line then # it will select what ever version it finds first for that @@ -106,17 +124,18 @@ for segment_name in opts.segment_names: # loop over segments names from segment_definer table # from the segment XML files - for key in seg_dict.keys(): - + for key in seg_dict: # if you find the segment flag for this IFO then add those # segments to the list and exit this for loop # this is a for loop because the version of the flag may not # be defined by the command line - if key == segment_flag or key.startswith(segment_flag+':'): + if key == segment_flag or key.startswith(segment_flag + ":"): if not len(base_segs[ifo]) and name == names[0]: base_segs[ifo] = seg_dict[key].coalesce() else: - comp_segs[ifo] = comp_segs[ifo].coalesce() + seg_dict[key].coalesce() + comp_segs[ifo] = ( + comp_segs[ifo].coalesce() + seg_dict[key].coalesce() + ) comp_segs[ifo].coalesce() break @@ -125,16 +144,18 @@ for segment_name in opts.segment_names: # FIXME: This only worked by fluke before. Not sure what it's supposed to # do in the case where there are two segments anyway. # Commenting for now - #if comment_dict[key] != None: + # if comment_dict[key] != None: # caption += " ("+comment_dict[key]+")" # set up dictionary of length of time of single-ifo segments in seconds ifo_len = {} for ifo in opts.ifos: if len(names) > 1: - ifo_len[ifo] = abs( ( base_segs[ifo].coalesce() & comp_segs[ifo].coalesce() ).coalesce() ) + ifo_len[ifo] = abs( + (base_segs[ifo].coalesce() & comp_segs[ifo].coalesce()).coalesce() + ) else: - ifo_len[ifo] = abs( base_segs[ifo].coalesce() ) + ifo_len[ifo] = abs(base_segs[ifo].coalesce()) columns[0][1].append(segment_name) counter = 1 @@ -147,34 +168,43 @@ for segment_name in opts.segment_names: first_ifo = False else: for ifo in combo: - combo_base_segs = combo_base_segs.coalesce() & base_segs[ifo].coalesce() - combo_comp_segs = combo_comp_segs.coalesce() + comp_segs[ifo].coalesce() + combo_base_segs = ( + combo_base_segs.coalesce() & base_segs[ifo].coalesce() + ) + combo_comp_segs = ( + combo_comp_segs.coalesce() + comp_segs[ifo].coalesce() + ) combo_base_segs.coalesce() combo_comp_segs.coalesce() if len(names) > 1: - combo_len = abs( (combo_base_segs.coalesce() & combo_comp_segs.coalesce() ).coalesce() ) + combo_len = abs( + (combo_base_segs.coalesce() & combo_comp_segs.coalesce()).coalesce() + ) else: - combo_len = abs( combo_base_segs.coalesce() ) + combo_len = abs(combo_base_segs.coalesce()) columns[counter][1].append(float(combo_len)) counter += 1 # cast columns into arrays -keys = [numpy.array(key, dtype=type(key[0])) for key,_ in columns] -vals = [numpy.array(val, dtype=type(val[0])) for _,val in columns] +keys = [numpy.array(key, dtype=type(key[0])) for key, _ in columns] +vals = [numpy.array(val, dtype=type(val[0])) for _, val in columns] # write HTML table -title = 'Segment Summary' -caption += '.' +title = "Segment Summary" +caption += "." if opts.title_text: title = title + ": " + opts.title_text if opts.description: - caption = caption + " " + opts.description + caption = caption + " " + opts.description fig_kwds = {} html_table = pycbc.results.html_table(vals, keys, page_size=10) -save_fig_with_metadata(str(html_table), opts.output_file, - fig_kwds=fig_kwds, - title=title, - cmd=' '.join(sys.argv), - caption=caption) +save_fig_with_metadata( + str(html_table), + opts.output_file, + fig_kwds=fig_kwds, + title=title, + cmd=" ".join(sys.argv), + caption=caption, +) diff --git a/bin/plotting/pycbc_page_sensitivity b/bin/plotting/pycbc_page_sensitivity index f3777c88b88..a0a4ff58117 100755 --- a/bin/plotting/pycbc_page_sensitivity +++ b/bin/plotting/pycbc_page_sensitivity @@ -1,264 +1,313 @@ #!/usr/bin/python -""" Plot search sensitivity as a function of significance. -""" +"""Plot search sensitivity as a function of significance.""" + import argparse -import numpy import logging -import matplotlib import sys -matplotlib.use('Agg') + +import matplotlib +import numpy + +matplotlib.use("Agg") from matplotlib import pyplot as plt +import pycbc import pycbc.pnutils import pycbc.results -import pycbc -from pycbc import sensitivity from pycbc import conversions as conv +from pycbc import sensitivity from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--injection-file', nargs='+', - help="Required. HDF format injection result file or space " - "separated list of files") -parser.add_argument('--output-file', required=True, - help='Destination file for the plot') -parser.add_argument('--bin-type', choices=['spin', 'mchirp', 'total_mass', - 'max_mass', 'eta', - 'template_duration'], - default='mchirp', - help="Parameter used to bin injections. Default 'mchirp'") +parser.add_argument( + "--injection-file", + nargs="+", + help="Required. HDF format injection result file or space separated list of files", +) +parser.add_argument( + "--output-file", required=True, help="Destination file for the plot" +) +parser.add_argument( + "--bin-type", + choices=["spin", "mchirp", "total_mass", "max_mass", "eta", "template_duration"], + default="mchirp", + help="Parameter used to bin injections. Default 'mchirp'", +) # FIXME: Make bins options properly optional, if no bins are specified then # plot everything together -parser.add_argument('--bins', nargs='+', - help="Required. Parameter bin boundaries, ex. 1.3 2.6 4.1") -parser.add_argument('--sig-type', choices=['ifar', 'fap', 'stat'], - default='ifar', - help="x-axis significance measure. Default 'ifar'") -parser.add_argument('--exclusive-sig', action='store_true', - help="Plot only exclusive injection FAR, not inclusive") -parser.add_argument('--sig-bins', nargs='*', - help="Boundaries of x-axis significance bins. If not given" - ", hard-coded defaults will be used"), -parser.add_argument('--dist-type', choices=['distance', 'volume', 'vt'], - default='distance', - help="y-axis sensitivity measure. Default 'distance'") -parser.add_argument('--log-dist', action='store_true', - help='Plot the sensitivity axis in log scale') -parser.add_argument('--min-dist', type=float, - help="Lower y-axis limit for sensitive distance") -parser.add_argument('--max-dist', type=float, - help="Upper y-axis limit for sensitive distance") -parser.add_argument('--integration-method', default='pylal', - choices=['pylal', 'shell', 'mc', 'vchirp_mc'], - help="Sensitive volume estimation method. Default 'pylal'") -parser.add_argument('--dist-bins', type=int, default=100, - help="Number of distance bins for 'pylal' volume " - "estimation. Default 100") -parser.add_argument('--distance-param', choices=['distance', 'chirp_distance'], - help="Parameter D used to generate injection distribution " - "over distance, required for 'mc' volume estimation") -parser.add_argument('--distribution', - choices=['log', 'uniform', 'distancesquared', 'volume'], - help="Form of distribution over D, required by 'mc' method") -parser.add_argument('--limits-param', choices=['distance', 'chirp_distance'], - help="Parameter Dlim specifying limits of injection " - "distribution, used by 'mc' method. If not given, " - "will be set equal to --distance-param") -parser.add_argument('--max-param', type=float, - help="Maximum value of Dlim, used by 'mc' method. If not " - "given, the maximum injected value will be used") -parser.add_argument('--min-param', type=float, - help="Minimum value of Dlim, used by 'mc' method with log " - "distribution. If not given, min injected value will " - "be used") -parser.add_argument('--spin-frame', choices=['line-of-sight', 'orbit'], - default='orbit', help='Frame convention used by injections ' - 'for specifying spin vectors. LAL versions after summer ' - '2015 should use the orbit convention, which is the ' - 'default choice here.') -parser.add_argument('--hdf-out', help='HDF file to save curve data') -parser.add_argument('--f-lower', default=20, type=float, help='Low frequency ' - 'cutoff for calculating template durations') +parser.add_argument( + "--bins", nargs="+", help="Required. Parameter bin boundaries, ex. 1.3 2.6 4.1" +) +parser.add_argument( + "--sig-type", + choices=["ifar", "fap", "stat"], + default="ifar", + help="x-axis significance measure. Default 'ifar'", +) +parser.add_argument( + "--exclusive-sig", + action="store_true", + help="Plot only exclusive injection FAR, not inclusive", +) +( + parser.add_argument( + "--sig-bins", + nargs="*", + help="Boundaries of x-axis significance bins. If not given" + ", hard-coded defaults will be used", + ), +) +parser.add_argument( + "--dist-type", + choices=["distance", "volume", "vt"], + default="distance", + help="y-axis sensitivity measure. Default 'distance'", +) +parser.add_argument( + "--log-dist", action="store_true", help="Plot the sensitivity axis in log scale" +) +parser.add_argument( + "--min-dist", type=float, help="Lower y-axis limit for sensitive distance" +) +parser.add_argument( + "--max-dist", type=float, help="Upper y-axis limit for sensitive distance" +) +parser.add_argument( + "--integration-method", + default="pylal", + choices=["pylal", "shell", "mc", "vchirp_mc"], + help="Sensitive volume estimation method. Default 'pylal'", +) +parser.add_argument( + "--dist-bins", + type=int, + default=100, + help="Number of distance bins for 'pylal' volume estimation. Default 100", +) +parser.add_argument( + "--distance-param", + choices=["distance", "chirp_distance"], + help="Parameter D used to generate injection distribution " + "over distance, required for 'mc' volume estimation", +) +parser.add_argument( + "--distribution", + choices=["log", "uniform", "distancesquared", "volume"], + help="Form of distribution over D, required by 'mc' method", +) +parser.add_argument( + "--limits-param", + choices=["distance", "chirp_distance"], + help="Parameter Dlim specifying limits of injection " + "distribution, used by 'mc' method. If not given, " + "will be set equal to --distance-param", +) +parser.add_argument( + "--max-param", + type=float, + help="Maximum value of Dlim, used by 'mc' method. If not " + "given, the maximum injected value will be used", +) +parser.add_argument( + "--min-param", + type=float, + help="Minimum value of Dlim, used by 'mc' method with log " + "distribution. If not given, min injected value will " + "be used", +) +parser.add_argument( + "--spin-frame", + choices=["line-of-sight", "orbit"], + default="orbit", + help="Frame convention used by injections " + "for specifying spin vectors. LAL versions after summer " + "2015 should use the orbit convention, which is the " + "default choice here.", +) +parser.add_argument("--hdf-out", help="HDF file to save curve data") +parser.add_argument( + "--f-lower", + default=20, + type=float, + help="Low frequency cutoff for calculating template durations", +) args = parser.parse_args() if len(args.bins) < 2: raise RuntimeError("At least 2 injection bin boundaries are required!") -if args.integration_method == 'mc' and (args.distance_param is None or \ - args.distribution is None): - raise RuntimeError("The 'mc' method requires --distance-param and " - "--distribution !") -if args.integration_method == 'mc' and args.limits_param is None: +if args.integration_method == "mc" and ( + args.distance_param is None or args.distribution is None +): + raise RuntimeError("The 'mc' method requires --distance-param and --distribution !") +if args.integration_method == "mc" and args.limits_param is None: args.limits_param = args.distance_param pycbc.init_logging(args.verbose) -logging.info('Read in the data') +logging.info("Read in the data") + def get_spin(frame, inc, m1, m2, s1x, s1z, s2x, s2z): # 'spin' means effective spin along orbital angular momentum - if frame == 'line-of-sight': + if frame == "line-of-sight": s1 = s1x * numpy.sin(inc) + s1z * numpy.cos(inc) s2 = s2x * numpy.sin(inc) + s2z * numpy.cos(inc) - elif frame == 'orbit': + elif frame == "orbit": s1, s2 = s1z, s2z else: - raise RuntimeError('Unknown spin frame!') + raise RuntimeError("Unknown spin frame!") return (m1 * s1 + m2 * s2) / (m1 + m2) + # initialize injection arrays and duration # mchirp is required for Monte-Carlo method for dchirp distributed inj missed = { - 'dist' : numpy.array([]), - 'param' : numpy.array([]), - 'mchirp': numpy.array([]), + "dist": numpy.array([]), + "param": numpy.array([]), + "mchirp": numpy.array([]), } found = { - 'dist' : numpy.array([]), - 'param' : numpy.array([]), - 'mchirp': numpy.array([]), - 'sig' : numpy.array([]), - 'sig_exc' : numpy.array([]), + "dist": numpy.array([]), + "param": numpy.array([]), + "mchirp": numpy.array([]), + "sig": numpy.array([]), + "sig_exc": numpy.array([]), } -t = 0. +t = 0.0 for fi in args.injection_file: - with HFile(fi, 'r') as f: - + with HFile(fi, "r") as f: # Get the found (at any FAR)/missed injection indices - foundi = f['found_after_vetoes/injection_index'][:] - missedi = f['missed/after_vetoes'][:] + foundi = f["found_after_vetoes/injection_index"][:] + missedi = f["missed/after_vetoes"][:] # retrieve injection parameters - dist = f['injections/distance'][:] - m1, m2 = f['injections/mass1'][:], f['injections/mass2'][:] - s1x, s2x = f['injections/spin1x'][:], f['injections/spin2x'][:] - s1z, s2z = f['injections/spin1z'][:], f['injections/spin2z'][:] + dist = f["injections/distance"][:] + m1, m2 = f["injections/mass1"][:], f["injections/mass2"][:] + s1x, s2x = f["injections/spin1x"][:], f["injections/spin2x"][:] + s1z, s2z = f["injections/spin1z"][:], f["injections/spin2z"][:] # y-components not used but read them in for symmetry - s1y, s2y = f['injections/spin1y'][:], f['injections/spin2y'][:] - inc = f['injections/inclination'][:] + s1y, s2y = f["injections/spin1y"][:], f["injections/spin2y"][:] + inc = f["injections/inclination"][:] mchirp = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2)[0] - if args.bin_type == 'mchirp': + if args.bin_type == "mchirp": pvals = mchirp - elif args.bin_type == 'eta': + elif args.bin_type == "eta": pvals = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2)[1] - elif args.bin_type == 'total_mass': + elif args.bin_type == "total_mass": pvals = m1 + m2 - elif args.bin_type == 'max_mass': + elif args.bin_type == "max_mass": pvals = numpy.maximum(m1, m2) - elif args.bin_type == 'spin': - pvals = get_spin(args.spin_frame, inc, - m1, m2, s1x, s1z, s2x, s2z) - elif args.bin_type == 'template_duration': + elif args.bin_type == "spin": + pvals = get_spin(args.spin_frame, inc, m1, m2, s1x, s1z, s2x, s2z) + elif args.bin_type == "template_duration": # Will default to SEOBNRv4 approximant value # Only valid/useful for non-spin or aligned-spin signals - pvals = pycbc.pnutils.get_imr_duration(m1, m2, s1z, s2z, - f_low=args.f_lower) + pvals = pycbc.pnutils.get_imr_duration(m1, m2, s1z, s2z, f_low=args.f_lower) else: - raise RuntimeError('Unrecognized --bin-type value!') + raise RuntimeError("Unrecognized --bin-type value!") - if args.sig_type == 'stat': - sig_exc = f['found_after_vetoes/stat'][:] + if args.sig_type == "stat": + sig_exc = f["found_after_vetoes/stat"][:] sig = None - elif args.sig_type == 'ifar': - sig_exc = f['found_after_vetoes/ifar_exc'][:] + elif args.sig_type == "ifar": + sig_exc = f["found_after_vetoes/ifar_exc"][:] try: - sig = f['found_after_vetoes/ifar'][:] + sig = f["found_after_vetoes/ifar"][:] except KeyError: # multiifo inj files may not have inclusive FAR sig = numpy.array([]) - elif args.sig_type == 'fap': - sig_exc = f['found_after_vetoes/fap_exc'][:] + elif args.sig_type == "fap": + sig_exc = f["found_after_vetoes/fap_exc"][:] try: - sig = f['found_after_vetoes/fap'][:] + sig = f["found_after_vetoes/fap"][:] except KeyError: # multiifo inj files may not have inclusive FAP sig = numpy.array([]) else: - raise RuntimeError('Unrecognized --sig-type value!') + raise RuntimeError("Unrecognized --sig-type value!") # Add the current values to the arrays - missed['dist'] = numpy.append(missed['dist'], dist[missedi]) - missed['param'] = numpy.append(missed['param'], pvals[missedi]) - missed['mchirp']= numpy.append(missed['mchirp'], mchirp[missedi]) - found['dist'] = numpy.append(found['dist'], dist[foundi]) - found['param'] = numpy.append(found['param'], pvals[foundi]) - found['mchirp'] = numpy.append(found['mchirp'], mchirp[foundi]) - found['sig'] = numpy.append(found['sig'], sig) - found['sig_exc']= numpy.append(found['sig_exc'], sig_exc) + missed["dist"] = numpy.append(missed["dist"], dist[missedi]) + missed["param"] = numpy.append(missed["param"], pvals[missedi]) + missed["mchirp"] = numpy.append(missed["mchirp"], mchirp[missedi]) + found["dist"] = numpy.append(found["dist"], dist[foundi]) + found["param"] = numpy.append(found["param"], pvals[foundi]) + found["mchirp"] = numpy.append(found["mchirp"], mchirp[foundi]) + found["sig"] = numpy.append(found["sig"], sig) + found["sig_exc"] = numpy.append(found["sig_exc"], sig_exc) # Time in years - if args.dist_type == 'vt': - t += conv.sec_to_year(f.attrs['foreground_time_exc']) + if args.dist_type == "vt": + t += conv.sec_to_year(f.attrs["foreground_time_exc"]) # Parameter bin legend labels labels = { - 'mchirp' : "$ M_{\\rm chirp} \in [%5.2f, %5.2f] M_\odot $", - 'eta' : "$ \\eta \in [%5.2f, %5.2f] $", - 'total_mass' : "$ M_{\\rm total} \in [%5.2f, %5.2f] M_\odot $", - 'max_mass' : "$ {\\rm max}(m_1, m_2) \in [%5.2f, %5.2f] M_\odot $", - 'spin' : "$ \\chi_{\\rm eff} \in [%5.2f, %5.2f] $", - 'template_duration' : "$ \\tau \in [%5.2f, %5.2f] s $" + "mchirp": "$ M_{\\rm chirp} \\in [%5.2f, %5.2f] M_\\odot $", + "eta": "$ \\eta \\in [%5.2f, %5.2f] $", + "total_mass": "$ M_{\\rm total} \\in [%5.2f, %5.2f] M_\\odot $", + "max_mass": "$ {\\rm max}(m_1, m_2) \\in [%5.2f, %5.2f] M_\\odot $", + "spin": "$ \\chi_{\\rm eff} \\in [%5.2f, %5.2f] $", + "template_duration": "$ \\tau \\in [%5.2f, %5.2f] s $", } ylabel = xlabel = "" # Set up significance values for x-axis -if args.sig_type == 'stat': - xlabel = 'Ranking Statistic Value' +if args.sig_type == "stat": + xlabel = "Ranking Statistic Value" if args.sig_bins: x_values = [float(v) for v in args.sig_bins] else: - x_values = numpy.arange(9, 14, .05) -elif args.sig_type == 'ifar': - xlabel = 'Inverse False Alarm Rate (years)' + x_values = numpy.arange(9, 14, 0.05) +elif args.sig_type == "ifar": + xlabel = "Inverse False Alarm Rate (years)" if args.sig_bins: x_values = [float(v) for v in args.sig_bins] else: # From approx 2 per day to 1 per 10,000yr - x_values = 10. ** (numpy.arange(-3, 4, .1)) -elif args.sig_type == 'fap': - xlabel = 'False Alarm Probability' + x_values = 10.0 ** (numpy.arange(-3, 4, 0.1)) +elif args.sig_type == "fap": + xlabel = "False Alarm Probability" if args.sig_bins: x_values = [float(v) for v in args.sig_bins] else: # From ~0.3 down to 1e-6 - x_values = 10. ** numpy.arange(-0.5, -6, -0.5) + x_values = 10.0 ** numpy.arange(-0.5, -6, -0.5) if args.hdf_out: plotdict = {} - plotdict['xvals'] = x_values + plotdict["xvals"] = x_values # Switches for plotting inclusive/exclusive significance -color = iter(plt.cm.rainbow(numpy.linspace(0, 1, len(args.bins)-1))) +color = iter(plt.cm.rainbow(numpy.linspace(0, 1, len(args.bins) - 1))) if not args.exclusive_sig: - fvalues = [found['sig'], found['sig_exc']] + fvalues = [found["sig"], found["sig_exc"]] do_labels = [True, False] - alphas = [.8, .3] + alphas = [0.8, 0.3] else: - fvalues = [found['sig_exc']] + fvalues = [found["sig_exc"]] do_labels = [True] - alphas = [.6] + alphas = [0.6] fig = plt.figure() # Cycle over parameter bins plotting each in turn -for j in range(len(args.bins)-1): +for j in range(len(args.bins) - 1): c = next(color) # cycle over inclusive / exclusive significance if available for sig_val, do_label, alpha in zip(fvalues, do_labels, alphas): if sig_val[0] is None: - logging.info('Skipping exclusive significance') + logging.info("Skipping exclusive significance") continue - left = float(args.bins[j]) - right = float(args.bins[j+1]) - logging.info('Injections between param values %5.2f and %5.2f' % - (left, right)) + left = float(args.bins[j]) + right = float(args.bins[j + 1]) + logging.info("Injections between param values %5.2f and %5.2f" % (left, right)) # Get distance of missed injections within parameter bin - binm = numpy.logical_and(missed['param'] >= left, - missed['param'] < right) - m_dist = missed['dist'][binm] + binm = numpy.logical_and(missed["param"] >= left, missed["param"] < right) + m_dist = missed["dist"][binm] # Abort if the bin has too few triggers if len(m_dist) < 2: @@ -267,51 +316,66 @@ for j in range(len(args.bins)-1): vols, vol_errors = [], [] # Slice up found injections in parameter bin - binf = numpy.logical_and(found['param'] >= left, - found['param'] < right) - binfsig = sig_val[binf] + binf = numpy.logical_and(found["param"] >= left, found["param"] < right) + binfsig = sig_val[binf] # Calculate each sensitive distance at a given significance threshold for x_val in x_values: - logging.info('Thresholding on significance at %5.2f' % x_val) + logging.info("Thresholding on significance at %5.2f" % x_val) # Count found inj towards sensitivity if IFAR/stat exceeds threshold # or if FAP value is less than threshold - if args.sig_type == 'ifar' or args.sig_type == 'stat': + if args.sig_type == "ifar" or args.sig_type == "stat": loud = binfsig >= x_val quiet = binfsig < x_val - elif args.sig_type == 'fap': + elif args.sig_type == "fap": loud = binfsig <= x_val quiet = binfsig > x_val # Distances of inj found above threshold - f_dist = found['dist'][binf][loud] + f_dist = found["dist"][binf][loud] # Distances of inj found below threshold - fm_dist = found['dist'][binf][quiet] + fm_dist = found["dist"][binf][quiet] # Add distances of 'quiet' found injections to the missed list m_dist_full = numpy.append(m_dist, fm_dist) # Choose the volume estimation method - if args.integration_method == 'shell': + if args.integration_method == "shell": vol, vol_err = sensitivity.volume_shell(f_dist, m_dist_full) - elif args.integration_method == 'pylal': - vol, vol_err = sensitivity.volume_binned_pylal(f_dist, - m_dist_full, bins=args.dist_bins) - elif args.integration_method in ['mc', 'vchirp_mc']: - found_mchirp = found['mchirp'][binf][loud] - missed_mchirp = numpy.append(missed['mchirp'][binm], - found['mchirp'][binf][quiet]) - - if args.integration_method == 'mc': - vol, vol_err = sensitivity.volume_montecarlo(f_dist, - m_dist_full, found_mchirp, missed_mchirp, - args.distance_param, args.distribution, - args.limits_param, args.min_param, args.max_param) - else: # vchirp_mc + elif args.integration_method == "pylal": + vol, vol_err = sensitivity.volume_binned_pylal( + f_dist, m_dist_full, bins=args.dist_bins + ) + elif args.integration_method in ["mc", "vchirp_mc"]: + found_mchirp = found["mchirp"][binf][loud] + missed_mchirp = numpy.append( + missed["mchirp"][binm], found["mchirp"][binf][quiet] + ) + + if args.integration_method == "mc": + vol, vol_err = sensitivity.volume_montecarlo( + f_dist, + m_dist_full, + found_mchirp, + missed_mchirp, + args.distance_param, + args.distribution, + args.limits_param, + args.min_param, + args.max_param, + ) + else: # vchirp_mc vol, vol_err = sensitivity.chirp_volume_montecarlo( - f_dist, m_dist_full, found_mchirp, missed_mchirp, - args.distance_param, args.distribution, - args.limits_param, args.min_param, args.max_param) + f_dist, + m_dist_full, + found_mchirp, + missed_mchirp, + args.distance_param, + args.distribution, + args.limits_param, + args.min_param, + args.max_param, + ) vols.append(vol) vol_errors.append(vol_err) @@ -319,42 +383,45 @@ for j in range(len(args.bins)-1): vols = numpy.array(vols) vol_errors = numpy.array(vol_errors) - if args.dist_type == 'distance': - ylabel = 'Sensitive Distance (Mpc)' - reach, ehigh, elow = sensitivity.volume_to_distance_with_errors(vols, vol_errors) - elif args.dist_type == 'volume': + if args.dist_type == "distance": + ylabel = "Sensitive Distance (Mpc)" + reach, ehigh, elow = sensitivity.volume_to_distance_with_errors( + vols, vol_errors + ) + elif args.dist_type == "volume": ylabel = "Sensitive Volume (Mpc$^3$)" reach, ehigh, elow = vols, vol_errors, vol_errors - elif args.dist_type == 'vt': + elif args.dist_type == "vt": ylabel = "Volume $\\times$ Time (yr Mpc$^3$)" - plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0)) + plt.ticklabel_format(style="sci", axis="y", scilimits=(0, 0)) reach, ehigh, elow = vols * t, vol_errors * t, vol_errors * t label = labels[args.bin_type] % (left, right) if do_label else None plt.plot(x_values, reach, label=label, c=c) - plt.plot(x_values, reach, alpha=alpha, c='black') - plt.fill_between(x_values, reach - elow, reach + ehigh, - facecolor=c, edgecolor=c, alpha=alpha) + plt.plot(x_values, reach, alpha=alpha, c="black") + plt.fill_between( + x_values, reach - elow, reach + ehigh, facecolor=c, edgecolor=c, alpha=alpha + ) if label and args.hdf_out: - plotdict['data/%s' % label] = reach - plotdict['errorhigh/%s' % label] = ehigh - plotdict['errorlow/%s' % label] = elow + plotdict["data/%s" % label] = reach + plotdict["errorhigh/%s" % label] = ehigh + plotdict["errorlow/%s" % label] = elow if args.hdf_out: - outfile = HFile(args.hdf_out,'w') - for key in plotdict.keys(): + outfile = HFile(args.hdf_out, "w") + for key in plotdict: outfile.create_dataset(key, data=plotdict[key]) ax = plt.gca() if args.log_dist: - ax.set_yscale('log') + ax.set_yscale("log") -if args.sig_type != 'stat': - ax.set_xscale('log') +if args.sig_type != "stat": + ax.set_xscale("log") -if args.sig_type == 'fap': +if args.sig_type == "fap": ax.invert_xaxis() if args.min_dist is not None: @@ -366,18 +433,23 @@ plt.ylabel(ylabel) plt.xlabel(xlabel) plt.grid() -plt.legend(loc='lower left') - -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title="Sensitive %s vs %s: binned by %s using %s method" - % (args.dist_type.title(), args.sig_type.upper(), - args.bin_type, args.integration_method), - caption="Sensitive %s as a function of Significance:" - "Lighter lines represent the significance without" - " including injections in their own background," - " while darker lines include each injection" - " individually in the background. The integration" - " method used is based on %s." - % (args.dist_type.title(), args.integration_method), - cmd=' '.join(sys.argv)) - +plt.legend(loc="lower left") + +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title="Sensitive %s vs %s: binned by %s using %s method" + % ( + args.dist_type.title(), + args.sig_type.upper(), + args.bin_type, + args.integration_method, + ), + caption="Sensitive %s as a function of Significance:" + "Lighter lines represent the significance without" + " including injections in their own background," + " while darker lines include each injection" + " individually in the background. The integration" + " method used is based on %s." % (args.dist_type.title(), args.integration_method), + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_page_snrchi b/bin/plotting/pycbc_page_snrchi index 42f8e15fb2e..7ca57668560 100644 --- a/bin/plotting/pycbc_page_snrchi +++ b/bin/plotting/pycbc_page_snrchi @@ -1,70 +1,71 @@ #!/usr/bin/env python -import logging -import numpy import argparse -import matplotlib +import logging import sys -matplotlib.use('Agg') + +import matplotlib +import numpy + +matplotlib.use("Agg") from matplotlib import pyplot as plt import pycbc.results -from pycbc.io import ( - get_chisq_from_file_choice, chisq_choices, SingleDetTriggers, HFile -) +from pycbc.io import HFile, SingleDetTriggers, chisq_choices, get_chisq_from_file_choice parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-file', help='Single ifo trigger file') -parser.add_argument('--veto-file', - help='Optional, file of veto segments to remove triggers') -parser.add_argument('--segment-name', default=None, type=str, - help='Optional, name of segment list to use for vetoes') -parser.add_argument('--min-snr', type=float, - help='Optional, Minimum SNR to plot') -parser.add_argument('--output-file') +parser.add_argument("--trigger-file", help="Single ifo trigger file") +parser.add_argument( + "--veto-file", help="Optional, file of veto segments to remove triggers" +) parser.add_argument( - '--newsnr-contours', - nargs='*', + "--segment-name", + default=None, + type=str, + help="Optional, name of segment list to use for vetoes", +) +parser.add_argument("--min-snr", type=float, help="Optional, Minimum SNR to plot") +parser.add_argument("--output-file") +parser.add_argument( + "--newsnr-contours", + nargs="*", help="List of newsnr values to draw contours at. Only valid for " - "--chisq-choice traditional (which is default)" + "--chisq-choice traditional (which is default)", +) +parser.add_argument( + "--chisq-choice", + choices=chisq_choices, + default="traditional", + help="Which chisquared to plot. Default=traditional", ) -parser.add_argument('--chisq-choice', choices=chisq_choices, - default='traditional', - help='Which chisquared to plot. Default=traditional') args = parser.parse_args() pycbc.init_logging(args.verbose, default_level=1) -if args.newsnr_contours is not None and not args.chisq_choice == 'traditional': +if args.newsnr_contours is not None and not args.chisq_choice == "traditional": parser.error( "Newsnr contours are being plotted assuming a traditional " - "chisquared, but this is not being plotted! We are plotting %s" % - args.chisq_choice + "chisquared, but this is not being plotted! We are plotting %s" + % args.chisq_choice ) -args.newsnr_contours = \ - args.newsnr_contours if args.newsnr_contours is not None else [] +args.newsnr_contours = args.newsnr_contours if args.newsnr_contours is not None else [] # First - get the IFO from the file -with HFile(args.trigger_file, 'r') as f: +with HFile(args.trigger_file, "r") as f: ifo = tuple(f.keys())[0] - n_triggers = f[ifo]['snr'].size + n_triggers = f[ifo]["snr"].size -logging.info( - "Plotting SNR vs %s chisq for %s triggers", - args.chisq_choice, - ifo -) +logging.info("Plotting SNR vs %s chisq for %s triggers", args.chisq_choice, ifo) if args.min_snr is not None: - filter_rank = 'snr' + filter_rank = "snr" else: filter_rank = None if n_triggers > 1e7: logging.warning( - "Getting %d triggers with no SNR cut - this will be expensive!", - n_triggers + "Getting %d triggers with no SNR cut - this will be expensive!", n_triggers ) trigs = SingleDetTriggers( @@ -77,20 +78,19 @@ trigs = SingleDetTriggers( ) logging.info( - "Reading SNR and %s chisq for %d triggers", - args.chisq_choice, - trigs.mask_size + "Reading SNR and %s chisq for %d triggers", args.chisq_choice, trigs.mask_size ) snr = trigs.snr chisq = get_chisq_from_file_choice(trigs, args.chisq_choice) -def snr_from_chisq(chisq, newsnr, q=6.): +def snr_from_chisq(chisq, newsnr, q=6.0): snr = numpy.zeros(len(chisq)) + float(newsnr) - ind = numpy.where(chisq > 1.)[0] - snr[ind] = float(newsnr) / ( 0.5 * (1. + chisq[ind] ** (q/2.)) ) ** (-1./q) + ind = numpy.where(chisq > 1.0)[0] + snr[ind] = float(newsnr) / (0.5 * (1.0 + chisq[ind] ** (q / 2.0))) ** (-1.0 / q) return snr + fig = plt.figure(1) r = numpy.logspace(numpy.log(chisq.min()), numpy.log(chisq.max()), 300) @@ -98,7 +98,7 @@ r = numpy.logspace(numpy.log(chisq.min()), numpy.log(chisq.max()), 300) for i, cval in enumerate(args.newsnr_contours): logging.info("Plotting newsnr %s contour", cval) snrv = snr_from_chisq(r, cval) - plt.plot(snrv, r, color='black', lw=0.5) + plt.plot(snrv, r, color="black", lw=0.5) if i == 0: label = "$\\hat{\\rho} = %s$" % cval else: @@ -107,26 +107,42 @@ for i, cval in enumerate(args.newsnr_contours): label_pos_idx = numpy.where(snrv > snr.max() * 0.8)[0][0] except IndexError: label_pos_idx = 0 - plt.text(snrv[label_pos_idx], r[label_pos_idx], label, fontsize=6, - horizontalalignment='center', verticalalignment='center', - bbox=dict(facecolor='white', lw=0, pad=0, alpha=0.9)) + plt.text( + snrv[label_pos_idx], + r[label_pos_idx], + label, + fontsize=6, + horizontalalignment="center", + verticalalignment="center", + bbox=dict(facecolor="white", lw=0, pad=0, alpha=0.9), + ) -plt.hexbin(snr, chisq, gridsize=300, xscale='log', yscale='log', lw=0.04, - mincnt=1, norm=matplotlib.colors.LogNorm()) +plt.hexbin( + snr, + chisq, + gridsize=300, + xscale="log", + yscale="log", + lw=0.04, + mincnt=1, + norm=matplotlib.colors.LogNorm(), +) ax = plt.gca() -plt.grid() -ax.set_xscale('log') -cb = plt.colorbar() +plt.grid() +ax.set_xscale("log") +cb = plt.colorbar() plt.xlim(snr.min(), snr.max() * 1.1) plt.ylim(chisq.min(), chisq.max() * 1.1) -cb.set_label('Trigger Density') -plt.xlabel('Signal-to-Noise Ratio') -plt.ylabel('Reduced $\\chi^2$') -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title="%s :SNR vs Reduced %s χ2" % (ifo, args.chisq_choice), - caption="Distribution of SNR and %s χ² for single detector triggers: " - "Black lines show contours of constant NewSNR." \ - %(args.chisq_choice,), - cmd=' '.join(sys.argv), - fig_kwds={'dpi':300}) +cb.set_label("Trigger Density") +plt.xlabel("Signal-to-Noise Ratio") +plt.ylabel("Reduced $\\chi^2$") +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title="%s :SNR vs Reduced %s χ2" % (ifo, args.chisq_choice), + caption="Distribution of SNR and %s χ² for single detector triggers: " + "Black lines show contours of constant NewSNR." % (args.chisq_choice,), + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 300}, +) diff --git a/bin/plotting/pycbc_page_snrifar b/bin/plotting/pycbc_page_snrifar index 02c08cc4b66..8857f6e55ea 100644 --- a/bin/plotting/pycbc_page_snrifar +++ b/bin/plotting/pycbc_page_snrifar @@ -1,97 +1,125 @@ #!/usr/bin/python -""" Make cumulative histogram of foreground coincident events via rate vs - ranking statistic, or make statistical significance vs ranking statistic - cumulative histograms. """ -import argparse, numpy, logging, sys +Make cumulative histogram of foreground coincident events via rate vs +ranking statistic, or make statistical significance vs ranking statistic +cumulative histograms. +""" + +import argparse +import logging +import sys + import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt +import numpy +matplotlib.use("Agg") +from matplotlib import pyplot as plt from scipy.special import erfc, erfinv -from pycbc.io.hdf import HFile import pycbc.results from pycbc import conversions as conv +from pycbc.io.hdf import HFile + def sigma_from_p(p): - return - erfinv(1 - (1 - p) * 2) * 2**0.5 + return -erfinv(1 - (1 - p) * 2) * 2**0.5 + def p_from_sigma(sig): - return erfc((sig)/numpy.sqrt(2))/2. + return erfc((sig) / numpy.sqrt(2)) / 2.0 + def p_from_far(far, livetime): return 1 - numpy.exp(-far * livetime) + def _far_from_p(p, livetime, max_far): if p < 0.0001: - lmbda = p + p**2./2. + lmbda = p + p**2.0 / 2.0 elif p == 1: return max_far else: - lmbda = -numpy.log(1-p) - far = lmbda/livetime + lmbda = -numpy.log(1 - p) + far = lmbda / livetime if far > max_far: return max_far return far + far_from_p = numpy.vectorize(_far_from_p) parser = argparse.ArgumentParser() # General required options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-file') -parser.add_argument('--output-file') -parser.add_argument('--not-cumulative', action='store_true') -parser.add_argument('--trials-factor', type=int, default=1, - help='Trials factor to divide from p-value and ' - 'divide from IFAR to account for look-elsewhere ' - 'effect. [default=1]') -parser.add_argument('--use-hierarchical-level', type=int, default=None, - help='Indicate which inclusive background and FARs of ' - 'foreground triggers to plot if there were any ' - 'hierarchical removals done. Choosing None plots ' - 'the inclusive backgrounds after all ' - 'hierarchical removals with the updated FARs for ' - 'foreground triggers after hierarchical removal(s). ' - 'Choosing 0 means plotting inclusive background ' - 'from prior to any hierarchical removals with FARs ' - 'for foreground triggers prior to hierarchical ' - 'removal. Choosing 1 means plotting the inclusive ' - 'background after doing 1 hierarchical removal, and ' - 'includes updated FARs from after 1 hierarchical ' - 'removal. [default=None]') -parser.add_argument('--fg-marker', default='^', - help='Marker to use for the foreground triggers, ' - '[default = ^]') -parser.add_argument('--fg-marker-h-rm', default='v', - help='Marker to use for the hierarchically removed ' - 'foreground triggers. [default = v]') -parser.add_argument('--closed-box', action='store_true', - help="Make a closed box version that excludes foreground " - "triggers") -parser.add_argument('--xmin', type=float, - help='Set the minimum value of the x-axis') -parser.add_argument('--xmax', type=float, - help='Set the maximum value of the x-axis') -parser.add_argument('--ymin', type=float, - help='Set the minimum value of the y-axis ' - '(in units of 1/years)') -parser.add_argument('--ymax', type=float, - help='Set the maximum value of the y-axis ' - '(in units of 1/years)') +parser.add_argument("--trigger-file") +parser.add_argument("--output-file") +parser.add_argument("--not-cumulative", action="store_true") +parser.add_argument( + "--trials-factor", + type=int, + default=1, + help="Trials factor to divide from p-value and " + "divide from IFAR to account for look-elsewhere " + "effect. [default=1]", +) +parser.add_argument( + "--use-hierarchical-level", + type=int, + default=None, + help="Indicate which inclusive background and FARs of " + "foreground triggers to plot if there were any " + "hierarchical removals done. Choosing None plots " + "the inclusive backgrounds after all " + "hierarchical removals with the updated FARs for " + "foreground triggers after hierarchical removal(s). " + "Choosing 0 means plotting inclusive background " + "from prior to any hierarchical removals with FARs " + "for foreground triggers prior to hierarchical " + "removal. Choosing 1 means plotting the inclusive " + "background after doing 1 hierarchical removal, and " + "includes updated FARs from after 1 hierarchical " + "removal. [default=None]", +) +parser.add_argument( + "--fg-marker", + default="^", + help="Marker to use for the foreground triggers, [default = ^]", +) +parser.add_argument( + "--fg-marker-h-rm", + default="v", + help="Marker to use for the hierarchically removed " + "foreground triggers. [default = v]", +) +parser.add_argument( + "--closed-box", + action="store_true", + help="Make a closed box version that excludes foreground triggers", +) +parser.add_argument("--xmin", type=float, help="Set the minimum value of the x-axis") +parser.add_argument("--xmax", type=float, help="Set the maximum value of the x-axis") +parser.add_argument( + "--ymin", + type=float, + help="Set the minimum value of the y-axis (in units of 1/years)", +) +parser.add_argument( + "--ymax", + type=float, + help="Set the maximum value of the y-axis (in units of 1/years)", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Read in the data') -f = HFile(args.trigger_file, 'r') +logging.info("Read in the data") +f = HFile(args.trigger_file, "r") # Parse which inclusive background to use for the plotting h_inc_back_num = args.use_hierarchical_level try: - h_iterations = f.attrs['hierarchical_removal_iterations'] + h_iterations = f.attrs["hierarchical_removal_iterations"] except KeyError: h_iterations = 0 @@ -101,53 +129,68 @@ if h_inc_back_num is None: if h_inc_back_num > h_iterations: # Produce a null plot saying no hierarchical removals can be plotted import sys + fig = plt.figure() ax = fig.add_subplot(111) ax.set_xlim(0, 1) ax.set_ylim(0, 1) - output_message = "No more foreground events louder than all background\n" \ - "at this removal level.\n" \ - "Attempted to show " + str(h_inc_back_num) + " removal(s),\n" \ - "but only " + str(h_iterations) + " removal(s) done." - - ax.text(0.5, 0.5, output_message, horizontalalignment='center', - verticalalignment='center') - - figure_title = "%s: " % f.attrs['ifos'] if 'ifos' in f.attrs else "" - figure_title += "%s bin, Cumulative Rate vs Rank" % f.attrs['name'] if 'name' in f.attrs else "FAR vs Rank" - - pycbc.results.save_fig_with_metadata(fig, args.output_file, - title=figure_title, - caption=output_message, - cmd=' '.join(sys.argv)) + output_message = ( + "No more foreground events louder than all background\n" + "at this removal level.\n" + "Attempted to show " + str(h_inc_back_num) + " removal(s),\n" + "but only " + str(h_iterations) + " removal(s) done." + ) + + ax.text( + 0.5, + 0.5, + output_message, + horizontalalignment="center", + verticalalignment="center", + ) + + figure_title = "%s: " % f.attrs["ifos"] if "ifos" in f.attrs else "" + figure_title += ( + "%s bin, Cumulative Rate vs Rank" % f.attrs["name"] + if "name" in f.attrs + else "FAR vs Rank" + ) + + pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title=figure_title, + caption=output_message, + cmd=" ".join(sys.argv), + ) # Exit the code successfully and bypass the rest of the plotting code. sys.exit(0) args.cumulative = not args.not_cumulative -foreground_livetime = conv.sec_to_year(f.attrs['foreground_time']) - +foreground_livetime = conv.sec_to_year(f.attrs["foreground_time"]) + if args.cumulative: - cstat_fore = f['foreground/stat'][:] + cstat_fore = f["foreground/stat"][:] cstat_fore.sort() cstat_rate = numpy.arange(len(cstat_fore), 0, -1) / foreground_livetime else: - cstat_fore = f['foreground/stat'][:] + cstat_fore = f["foreground/stat"][:] ssort = cstat_fore.argsort() - cstat_rate = 1.0 / f['foreground/ifar'][:] * args.trials_factor - # Correct the foreground ifars depending on the values calculated - # under a particular inclusive background distribution. + cstat_rate = 1.0 / f["foreground/ifar"][:] * args.trials_factor + # Correct the foreground ifars depending on the values calculated + # under a particular inclusive background distribution. try: - ifar_h_inc = f['foreground_h%s/ifar' % h_inc_back_num][:] + ifar_h_inc = f["foreground_h%s/ifar" % h_inc_back_num][:] except KeyError: ifar_h_inc = [] try: - stat_h_inc = f['foreground_h%s/stat' % h_inc_back_num][:] + stat_h_inc = f["foreground_h%s/stat" % h_inc_back_num][:] except KeyError: stat_h_inc = [] @@ -160,21 +203,21 @@ else: cstat_fore = cstat_fore[ssort] cstat_fap = p_from_far(cstat_rate, foreground_livetime) -logging.info('Found %s foreground triggers' % len(cstat_fore)) +logging.info("Found %s foreground triggers" % len(cstat_fore)) if cstat_fore is not None and len(cstat_fore) == 0: cstat_fore = None # Make choosing the background backwards compatible try: - back_ifar = f['background_h%s/ifar' % h_inc_back_num][:] + back_ifar = f["background_h%s/ifar" % h_inc_back_num][:] except KeyError: - back_ifar = f['background/ifar'][:] + back_ifar = f["background/ifar"][:] try: - cstat_back = f['background_h%s/stat' % h_inc_back_num][:] + cstat_back = f["background_h%s/stat" % h_inc_back_num][:] except KeyError: - cstat_back = f['background/stat'][:] + cstat_back = f["background/stat"][:] back_sort = cstat_back.argsort() cstat_back = cstat_back[back_sort] @@ -186,10 +229,10 @@ if not args.cumulative: fap_back = p_from_far(far_back, foreground_livetime) -logging.info('Found %s background (inclusive zerolag) triggers' % len(cstat_back)) +logging.info("Found %s background (inclusive zerolag) triggers" % len(cstat_back)) -back_ifar_exc = f['background_exc/ifar'][:] -cstat_back_exc = f['background_exc/stat'][:] +back_ifar_exc = f["background_exc/ifar"][:] +cstat_back_exc = f["background_exc/stat"][:] back_sort_exc = cstat_back_exc.argsort() cstat_back_exc = cstat_back_exc[back_sort_exc] far_back_exc = 1.0 / back_ifar_exc[back_sort_exc] @@ -197,16 +240,16 @@ far_back_exc = 1.0 / back_ifar_exc[back_sort_exc] if not args.cumulative: far_back_exc *= args.trials_factor -logging.info('Found %s background (exclusive zerolag) triggers' % len(cstat_back_exc)) +logging.info("Found %s background (exclusive zerolag) triggers" % len(cstat_back_exc)) # We'll use the background far for the ylimits if closed box, foreground # otherwise if args.ymin is not None: plot_ymin = args.ymin elif args.closed_box or cstat_fore is None: - plot_ymin = far_back_exc.min()/10. + plot_ymin = far_back_exc.min() / 10.0 else: - plot_ymin = far_back.min()/10. + plot_ymin = far_back.min() / 10.0 if args.ymax is not None: plot_ymax = args.ymax @@ -234,15 +277,28 @@ else: plot_xmax = cstat_back_exc.max() else: plot_xmax = max(cstat_back.max(), cstat_fore.max()) - plot_xmax += (plot_xmax - plot_xmin)/10. + plot_xmax += (plot_xmax - plot_xmin) / 10.0 fig = plt.figure(1) -back_marker = 'x' -plt.scatter(cstat_back_exc, far_back_exc, color='gray', marker=back_marker, s=10, label='Closed Box Background') +back_marker = "x" +plt.scatter( + cstat_back_exc, + far_back_exc, + color="gray", + marker=back_marker, + s=10, + label="Closed Box Background", +) if not args.closed_box: - plt.scatter(cstat_back, far_back, color='black', marker=back_marker, s=10, - label='Open Box Background') + plt.scatter( + cstat_back, + far_back, + color="black", + marker=back_marker, + s=10, + label="Open Box Background", + ) if cstat_fore is not None and len(cstat_fore): # Remove hierarchically removed foreground triggers from the list @@ -256,36 +312,55 @@ if not args.closed_box: # to lower ranking statistic. for i in range(h_inc_back_num): rm_idx = cstat_fore.argmax() - cstat_fore_h_rm = numpy.append(cstat_fore_h_rm, - cstat_fore[rm_idx]) - cstat_rate_h_rm = numpy.append(cstat_rate_h_rm, - cstat_rate[rm_idx]) + cstat_fore_h_rm = numpy.append(cstat_fore_h_rm, cstat_fore[rm_idx]) + cstat_rate_h_rm = numpy.append(cstat_rate_h_rm, cstat_rate[rm_idx]) cstat_fore = numpy.delete(cstat_fore, rm_idx) cstat_rate = numpy.delete(cstat_rate, rm_idx) - plt.scatter(cstat_fore_h_rm, cstat_rate_h_rm, s=60, color='#b66dff', - marker=args.fg_marker_h_rm, - label='Hierarchically Removed Foreground', zorder=100, - linewidth=0.5, edgecolors='white') - - plt.scatter(cstat_fore, cstat_rate, s=60, color='#ff6600', - marker=args.fg_marker, label='Foreground', zorder=100, - linewidth=0.5, edgecolors='white') + plt.scatter( + cstat_fore_h_rm, + cstat_rate_h_rm, + s=60, + color="#b66dff", + marker=args.fg_marker_h_rm, + label="Hierarchically Removed Foreground", + zorder=100, + linewidth=0.5, + edgecolors="white", + ) + + plt.scatter( + cstat_fore, + cstat_rate, + s=60, + color="#ff6600", + marker=args.fg_marker, + label="Foreground", + zorder=100, + linewidth=0.5, + edgecolors="white", + ) if args.not_cumulative: # add arrows to any points > the loudest background - louder_pts = numpy.where(cstat_fore > cstat_back.max())[0] + louder_pts = numpy.where(cstat_fore > cstat_back.max())[0] for ii in louder_pts: r = cstat_fore[ii] arr_start = cstat_rate[ii] # make the arrow length 1/15 the height of the plot - arr_end = arr_start * (plot_ymin / plot_ymax) ** (1./15) - plt.plot([r, r], [arr_start, arr_end], lw=2, color='black', - zorder=99) - plt.plot([r, r], [arr_start, arr_end], lw=2.6, color='white', - zorder=97) - plt.scatter([r], [arr_end], marker='v', c='black', - edgecolors='white', lw=0.5, s=40, zorder=98) + arr_end = arr_start * (plot_ymin / plot_ymax) ** (1.0 / 15) + plt.plot([r, r], [arr_start, arr_end], lw=2, color="black", zorder=99) + plt.plot([r, r], [arr_start, arr_end], lw=2.6, color="white", zorder=97) + plt.scatter( + [r], + [arr_end], + marker="v", + c="black", + edgecolors="white", + lw=0.5, + s=40, + zorder=98, + ) if h_inc_back_num > 0: louder_pts = numpy.where(cstat_fore_h_rm > cstat_back.max())[0] @@ -295,62 +370,77 @@ if not args.closed_box: if arr_start > far_back_exc.min(): continue # make the arrow length 1/15 the height of the plot - arr_end = arr_start * (plot_ymin / plot_ymax) ** (1./15) - plt.plot([r, r], [arr_start, arr_end], lw=2, - color='black', zorder=99) - plt.plot([r, r], [arr_start, arr_end], lw=2.6, - color='white', zorder=97) - plt.scatter([r], [arr_end], marker='v', c='black', - edgecolors='white', lw=0.5, s=40, zorder=98) + arr_end = arr_start * (plot_ymin / plot_ymax) ** (1.0 / 15) + plt.plot( + [r, r], [arr_start, arr_end], lw=2, color="black", zorder=99 + ) + plt.plot( + [r, r], [arr_start, arr_end], lw=2.6, color="white", zorder=97 + ) + plt.scatter( + [r], + [arr_end], + marker="v", + c="black", + edgecolors="white", + lw=0.5, + s=40, + zorder=98, + ) if not args.cumulative: # add second y-axis for probabilities, sigmas - sigmas = numpy.arange(6)+1 + sigmas = numpy.arange(6) + 1 ax1 = plt.gca() - if hasattr(ax1, 'set_facecolor'): - ax1.set_facecolor('none') + if hasattr(ax1, "set_facecolor"): + ax1.set_facecolor("none") else: - ax1.set_axis_bgcolor('none') + ax1.set_axis_bgcolor("none") ax2 = ax1.twinx() - ax1.set_zorder(ax2.get_zorder()+1) # put axis1 on top + ax1.set_zorder(ax2.get_zorder() + 1) # put axis1 on top plt.sca(ax2) # where to stick the sigma lables; we'll put them 1/25th from the # right axis - anntx = plot_xmax - (plot_xmax - plot_xmin)/25. + anntx = plot_xmax - (plot_xmax - plot_xmin) / 25.0 sigps = p_from_sigma(sigmas) - for ii,p in enumerate(sigps[:-1]): - nextp = sigps[ii+1] - plt.axhspan(far_from_p(nextp, foreground_livetime, far_back.max()), - far_from_p(p, foreground_livetime, far_back.max()), - linewidth=0, - color=plt.cm.Blues(float(sigmas[ii+1]) / sigmas.size), - alpha=0.3, zorder=-1) - # add sigma label - plt.annotate('%1.0f$\sigma$' % sigmas[ii], - (anntx, far_from_p(p, foreground_livetime, - far_back.max())), zorder=100) - ax2.plot([],[]) + for ii, p in enumerate(sigps[:-1]): + nextp = sigps[ii + 1] + plt.axhspan( + far_from_p(nextp, foreground_livetime, far_back.max()), + far_from_p(p, foreground_livetime, far_back.max()), + linewidth=0, + color=plt.cm.Blues(float(sigmas[ii + 1]) / sigmas.size), + alpha=0.3, + zorder=-1, + ) + # add sigma label + plt.annotate( + r"%1.0f$\sigma$" % sigmas[ii], + (anntx, far_from_p(p, foreground_livetime, far_back.max())), + zorder=100, + ) + ax2.plot([], []) plt.sca(ax1) -plt.xlabel(r'Ranking Statistic') -plt.yscale('log') +plt.xlabel(r"Ranking Statistic") +plt.yscale("log") plt.ylim(plot_ymin, plot_ymax * 10.0) plt.xlim(plot_xmin, plot_xmax) plt.legend(loc="upper right", fontsize=9) plt.grid() - + if args.cumulative: - plt.ylabel('Cumulative Rate (yr$^{-1}$)') + plt.ylabel("Cumulative Rate (yr$^{-1}$)") else: if args.trials_factor == 1: - plt.ylabel('False Alarm Rate (yr$^{-1}$)') + plt.ylabel("False Alarm Rate (yr$^{-1}$)") elif args.trials_factor >= 1: - plt.ylabel('Combined False Alarm Rate (yr$^{-1}$)') - ax2.set_ylabel('p-value') - ax2.set_yscale('log') + plt.ylabel("Combined False Alarm Rate (yr$^{-1}$)") + ax2.set_ylabel("p-value") + ax2.set_yscale("log") ymin, ymax = ax1.get_ylim() ax2.set_ylim(ymin, ymax) - + if args.cumulative: figure_caption = "Cumulative histogram of foreground and background triggers." else: @@ -364,24 +454,33 @@ else: tick_min = numpy.ceil(numpy.log10(pymin)) tick_max = numpy.floor(numpy.log10(pymax)) # Tick ranges - pticks = numpy.arange(tick_min, tick_max+1) + pticks = numpy.arange(tick_min, tick_max + 1) # Convert back to FAR fticks = far_from_p(10**pticks, foreground_livetime, far_back.max()) ax2.yaxis.set_major_locator(plt.NullLocator()) ax2.set_yticks(fticks) # Set the labels - ax2.set_yticklabels(['$10^{%i}$' %(val) for val in pticks.astype(int)]) + ax2.set_yticklabels(["$10^{%i}$" % (val) for val in pticks.astype(int)]) plt.tight_layout() -figure_title = "%s: " % f.attrs['ifos'] if 'ifos' in f.attrs else "" -figure_title += "%s bin, Cumulative Rate vs Rank" % f.attrs['name'] if 'name' in f.attrs else "FAR vs Rank" - -figure_caption = figure_caption + "Orange triangle (if present) represent triggers from the " \ -"zero-lag (foreground) analysis. Solid crosses show " \ -"the background inclusive of zerolag events, and grey crosses show the " \ -"background constructed without triggers that are " \ -"coincident in the zero-lag data.", - -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title=figure_title, - caption=figure_caption, - cmd=' '.join(sys.argv)) +figure_title = "%s: " % f.attrs["ifos"] if "ifos" in f.attrs else "" +figure_title += ( + "%s bin, Cumulative Rate vs Rank" % f.attrs["name"] + if "name" in f.attrs + else "FAR vs Rank" +) + +figure_caption = ( + figure_caption + "Orange triangle (if present) represent triggers from the " + "zero-lag (foreground) analysis. Solid crosses show " + "the background inclusive of zerolag events, and grey crosses show the " + "background constructed without triggers that are " + "coincident in the zero-lag data.", +) + +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title=figure_title, + caption=figure_caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_page_snrratehist b/bin/plotting/pycbc_page_snrratehist index 4b5e7a32a0f..692cbeaa721 100755 --- a/bin/plotting/pycbc_page_snrratehist +++ b/bin/plotting/pycbc_page_snrratehist @@ -1,24 +1,29 @@ #!/usr/bin/python -""" Make SNR vs rate of triggers histogram for foreground coincident events. - Also has the ability to plot inclusive backgrounds from different stages - of hierarchical removal. """ +Make SNR vs rate of triggers histogram for foreground coincident events. +Also has the ability to plot inclusive backgrounds from different stages +of hierarchical removal. +""" + import argparse -import numpy import logging import sys + import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt +import numpy +matplotlib.use("Agg") +from matplotlib import pyplot as plt from scipy.special import erf, erfinv -from pycbc.io.hdf import HFile import pycbc.results from pycbc import conversions as conv +from pycbc.io.hdf import HFile + def sigma_from_p(p): - return - erfinv(1 - (1 - p) * 2) * 2**0.5 + return -erfinv(1 - (1 - p) * 2) * 2**0.5 + def p_from_sigma(sig): return 1 - (1 - erf(sig / 2**0.5)) / 2 @@ -27,38 +32,44 @@ def p_from_sigma(sig): parser = argparse.ArgumentParser() # General required options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-file') -parser.add_argument('--output-file') -parser.add_argument('--bin-size', type=float) -parser.add_argument('--x-min', type=float) -parser.add_argument('--trials-factor', type=int, default=1) -parser.add_argument('--use-hierarchical-level', type=int, default=None, - help='Indicate which inclusive background to plot ' - 'with the foreground triggers if there were ' - 'any hierarchical removals done. Choosing 0 ' - 'indicates plotting the inclusive background prior ' - 'to any hierarchical removals. Choosing 1 indicates ' - 'the inclusive background after the loudest ' - 'foreground trigger was removed. Choosing another ' - 'integer gives the inclusive background after that ' - 'integer number of hierarchical removals. If not ' - 'provided, default to the background from after all ' - 'hierarchical removals performed. [default=None]') -parser.add_argument('--closed-box', action='store_true', - help='Make a closed box version that excludes ' - 'foreground triggers') +parser.add_argument("--trigger-file") +parser.add_argument("--output-file") +parser.add_argument("--bin-size", type=float) +parser.add_argument("--x-min", type=float) +parser.add_argument("--trials-factor", type=int, default=1) +parser.add_argument( + "--use-hierarchical-level", + type=int, + default=None, + help="Indicate which inclusive background to plot " + "with the foreground triggers if there were " + "any hierarchical removals done. Choosing 0 " + "indicates plotting the inclusive background prior " + "to any hierarchical removals. Choosing 1 indicates " + "the inclusive background after the loudest " + "foreground trigger was removed. Choosing another " + "integer gives the inclusive background after that " + "integer number of hierarchical removals. If not " + "provided, default to the background from after all " + "hierarchical removals performed. [default=None]", +) +parser.add_argument( + "--closed-box", + action="store_true", + help="Make a closed box version that excludes foreground triggers", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Read in the data') -f = HFile(args.trigger_file, 'r') +logging.info("Read in the data") +f = HFile(args.trigger_file, "r") # Determine which inclusive background to plot. h_inc_back_num = args.use_hierarchical_level try: - h_iterations = f.attrs['hierarchical_removal_iterations'] + h_iterations = f.attrs["hierarchical_removal_iterations"] except KeyError: h_iterations = 0 @@ -68,24 +79,37 @@ if h_inc_back_num is None: if h_inc_back_num > h_iterations: # Produce a null plot saying no hierarchical removals can be plotted import sys + fig = plt.figure() ax = fig.add_subplot(111) ax.set_xlim(0, 1) ax.set_ylim(0, 1) - output_message = "No more foreground events louder than all background\n" \ - "at this removal level.\n" \ - "Attempted to show " + str(h_inc_back_num) + " removal(s),\n" \ - "but only " + str(h_iterations) + " removal(s) done." + output_message = ( + "No more foreground events louder than all background\n" + "at this removal level.\n" + "Attempted to show " + str(h_inc_back_num) + " removal(s),\n" + "but only " + str(h_iterations) + " removal(s) done." + ) - ax.text(0.5, 0.5, output_message, horizontalalignment='center', - verticalalignment='center') + ax.text( + 0.5, + 0.5, + output_message, + horizontalalignment="center", + verticalalignment="center", + ) - pycbc.results.save_fig_with_metadata(fig, args.output_file, - title="%s bin, Count vs Rank" % f.attrs['name'] if 'name' in f.attrs else "Count vs Rank", + pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title="%s bin, Count vs Rank" % f.attrs["name"] + if "name" in f.attrs + else "Count vs Rank", caption=output_message, - cmd=' '.join(sys.argv)) + cmd=" ".join(sys.argv), + ) # Exit the code successfully and bypass the rest of the plotting code. sys.exit(0) @@ -94,7 +118,7 @@ if args.closed_box: fstat = None else: try: - fstat = f['foreground/stat'][:] + fstat = f["foreground/stat"][:] fstat.sort() except: fstat = None @@ -102,24 +126,29 @@ else: fstat = None if h_inc_back_num == 0: - bstat = f['background/stat'][:] - fap = 1 - numpy.exp(- conv.sec_to_year(f.attrs['foreground_time']) / f['background/ifar'][:]) - dec = f['background/decimation_factor'][:] -else : - bstat = f['background_h%s/stat' % h_inc_back_num][:] - fap = 1 - numpy.exp(- conv.sec_to_year(f.attrs['foreground_time_h%s' % h_inc_back_num]) / f['background_h%s/ifar' % h_inc_back_num][:]) - dec = f['background_h%s/decimation_factor' % h_inc_back_num][:] + bstat = f["background/stat"][:] + fap = 1 - numpy.exp( + -conv.sec_to_year(f.attrs["foreground_time"]) / f["background/ifar"][:] + ) + dec = f["background/decimation_factor"][:] +else: + bstat = f["background_h%s/stat" % h_inc_back_num][:] + fap = 1 - numpy.exp( + -conv.sec_to_year(f.attrs["foreground_time_h%s" % h_inc_back_num]) + / f["background_h%s/ifar" % h_inc_back_num][:] + ) + dec = f["background_h%s/decimation_factor" % h_inc_back_num][:] s = bstat.argsort() dec, bstat, fap = dec[s], bstat[s], fap[s] -logging.info('Found %s background (inclusive zerolag) triggers' % len(bstat)) +logging.info("Found %s background (inclusive zerolag) triggers" % len(bstat)) -dec_exc = f['background_exc/decimation_factor'][:] -bstat_exc = f['background_exc/stat'][:] +dec_exc = f["background_exc/decimation_factor"][:] +bstat_exc = f["background_exc/stat"][:] s = bstat_exc.argsort() dec_exc, bstat_exc = dec_exc[s], bstat_exc[s] -logging.info('Found %s background (exclusive zerolag) triggers' % len(bstat_exc)) +logging.info("Found %s background (exclusive zerolag) triggers" % len(bstat_exc)) fig = plt.figure() @@ -133,47 +162,55 @@ else: minimum = bstat.min() maximum = bstat.max() -bin_size = args.bin_size if args.bin_size else (maximum - minimum) / 100. +bin_size = args.bin_size or (maximum - minimum) / 100.0 bins = numpy.arange(minimum, maximum + bin_size, bin_size) # plot background minus foreground -exc_binweights = dec_exc / conv.sec_to_year(f.attrs['background_time_exc']) -exc_binvals = plt.hist(bstat_exc, bins=bins, histtype='step', - linewidth=2, - color='grey', log=True, - label= 'Background Uncorrelated with Foreground', - weights=exc_binweights) +exc_binweights = dec_exc / conv.sec_to_year(f.attrs["background_time_exc"]) +exc_binvals = plt.hist( + bstat_exc, + bins=bins, + histtype="step", + linewidth=2, + color="grey", + log=True, + label="Background Uncorrelated with Foreground", + weights=exc_binweights, +) histpeak = bins[exc_binvals[0].argmax()] # plot full background if not args.closed_box: - bg_key = 'background_time' if h_inc_back_num == 0 \ - else 'background_time_h%s' % h_inc_back_num + bg_key = ( + "background_time" + if h_inc_back_num == 0 + else "background_time_h%s" % h_inc_back_num + ) binweights = dec / conv.sec_to_year(f.attrs[bg_key]) plt.hist( bstat, bins=bins, - histtype='step', + histtype="step", linewidth=2, - color='black', + color="black", log=True, - label='Full Background', - weights=binweights + label="Full Background", + weights=binweights, ) if fstat is not None and not args.closed_box: le, re = bins[:-1], bins[1:] # We need to plot a histogram "errorbar" with the hierarchically # removed foreground triggers in purple. - if h_iterations > 0 : + if h_iterations > 0: fstat_h_rm = numpy.array([], dtype=float) # Since there is only one background bin we can just remove # hierarchically removed triggers from highest ranking statistic # to lower ranking statistic. - for i in range(0, h_inc_back_num): + for i in range(h_inc_back_num): rm_idx = fstat.argmax() fstat_h_rm = numpy.append(fstat_h_rm, fstat[rm_idx]) fstat = numpy.delete(fstat, rm_idx) @@ -189,34 +226,52 @@ if fstat is not None and not args.closed_box: # Just use the top level foreground time if you want background # before h-removal. if h_inc_back_num == 0: - count_h_rm = (right_h_rm - left_h_rm) / \ - conv.sec_to_year(f.attrs['foreground_time']) + count_h_rm = (right_h_rm - left_h_rm) / conv.sec_to_year( + f.attrs["foreground_time"] + ) # Or use the foreground time after h-removal else: - count_h_rm = (right_h_rm - left_h_rm) / \ - conv.sec_to_year(f.attrs['foreground_time_h%s' % h_inc_back_num]) - - plt.errorbar(bins[:-1] + bin_size / 2, count_h_rm, - xerr=bin_size/2, - label='Hierarchically Removed Foreground', mec='none', - fmt='s', ms=1, capthick=0, elinewidth=4, - color='#b66dff') + count_h_rm = (right_h_rm - left_h_rm) / conv.sec_to_year( + f.attrs["foreground_time_h%s" % h_inc_back_num] + ) + + plt.errorbar( + bins[:-1] + bin_size / 2, + count_h_rm, + xerr=bin_size / 2, + label="Hierarchically Removed Foreground", + mec="none", + fmt="s", + ms=1, + capthick=0, + elinewidth=4, + color="#b66dff", + ) left = numpy.searchsorted(fstat, le) right = numpy.searchsorted(fstat, re) - count = (right - left) / conv.sec_to_year(f.attrs['foreground_time']) - plt.errorbar(bins[:-1] + bin_size / 2, count, xerr=bin_size/2, - label='Foreground', mec='none', fmt='o', ms=1, capthick=0, - elinewidth=4, color='#ff6600') + count = (right - left) / conv.sec_to_year(f.attrs["foreground_time"]) + plt.errorbar( + bins[:-1] + bin_size / 2, + count, + xerr=bin_size / 2, + label="Foreground", + mec="none", + fmt="o", + ms=1, + capthick=0, + elinewidth=4, + color="#ff6600", + ) -plt.xlabel('Ranking statistic (bin size = %.2f)' % bin_size) -plt.ylabel('Trigger Rate (yr$^{-1})$') +plt.xlabel("Ranking statistic (bin size = %.2f)" % bin_size) +plt.ylabel("Trigger Rate (yr$^{-1})$") if args.x_min is not None: plt.xlim(xmin=args.x_min) else: plt.xlim(xmin=numpy.floor(histpeak)) -plt.ylim(ymin=0.5 / conv.sec_to_year(f.attrs['background_time_exc'])) +plt.ylim(ymin=0.5 / conv.sec_to_year(f.attrs["background_time_exc"])) plt.grid() leg = plt.legend(fontsize=9) @@ -245,41 +300,48 @@ if not args.closed_box: x = [bstat[::-1][x1], bstat[::-1][x2]] except IndexError: break - plt.fill_between(x, ymin, ymax, zorder=-1, - color=plt.cm.Blues(next_sig / 8.0)) + plt.fill_between(x, ymin, ymax, zorder=-1, color=plt.cm.Blues(next_sig / 8.0)) if next_sig == end: - next_sig = '%.1f' % next_sig - - plt.text(bstat[::-1][x2] - .1, ymax, r"$%s \sigma$" % next_sig, - fontsize=10, horizontalalignment='center', - verticalalignment='bottom') - -ax1 = plt.gca() + next_sig = "%.1f" % next_sig + + plt.text( + bstat[::-1][x2] - 0.1, + ymax, + r"$%s \sigma$" % next_sig, + fontsize=10, + horizontalalignment="center", + verticalalignment="bottom", + ) + +ax1 = plt.gca() ax2 = ax1.twinx() if h_inc_back_num == 0: - fac = conv.sec_to_year(f.attrs['foreground_time']) -else : - fac = conv.sec_to_year(f.attrs['foreground_time_h%s' % h_inc_back_num]) + fac = conv.sec_to_year(f.attrs["foreground_time"]) +else: + fac = conv.sec_to_year(f.attrs["foreground_time_h%s" % h_inc_back_num]) ymin = ax1.get_ylim()[0] * fac -ymax = ax1.get_ylim()[1] * fac +ymax = ax1.get_ylim()[1] * fac ax2.set_ylim(ymin=ymin, ymax=ymax) -ax2.set_yscale('log') -ax2.set_ylabel('Number per Experiment') - -if 'name' in f.attrs: - title = "%s bin, Count vs Rank" % f.attrs['name'] - caption="Histogram of FAR vs the ranking statistic in the search." -elif 'ifos' in f.attrs: - title = "%s coincidences, Count vs Rank" % f.attrs['ifos'] - caption="Histogram of the FAR vs the ranking statistic " \ - "for %s coincidences only" % f.attrs['ifos'] +ax2.set_yscale("log") +ax2.set_ylabel("Number per Experiment") + +if "name" in f.attrs: + title = "%s bin, Count vs Rank" % f.attrs["name"] + caption = "Histogram of FAR vs the ranking statistic in the search." +elif "ifos" in f.attrs: + title = "%s coincidences, Count vs Rank" % f.attrs["ifos"] + caption = ( + "Histogram of the FAR vs the ranking statistic " + "for %s coincidences only" % f.attrs["ifos"] + ) else: title = "Count vs Rank" - caption="Histogram of the FAR vs the ranking statistic in the search." + caption = "Histogram of the FAR vs the ranking statistic in the search." -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title=title, caption=caption, cmd=' '.join(sys.argv)) +pycbc.results.save_fig_with_metadata( + fig, args.output_file, title=title, caption=caption, cmd=" ".join(sys.argv) +) diff --git a/bin/plotting/pycbc_page_template_bin_table b/bin/plotting/pycbc_page_template_bin_table index 7e4c3f662b8..35e7ba94640 100644 --- a/bin/plotting/pycbc_page_template_bin_table +++ b/bin/plotting/pycbc_page_template_bin_table @@ -1,8 +1,9 @@ #!/usr/bin/env python -""" Make a table of template bin information -""" -import sys +"""Make a table of template bin information""" + import argparse +import sys + import h5py as h5 import numpy as np @@ -11,24 +12,24 @@ import pycbc.results parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--ifo', required=True) -parser.add_argument('--dq-file', required=True) -parser.add_argument('--output-file') +parser.add_argument("--ifo", required=True) +parser.add_argument("--dq-file", required=True) +parser.add_argument("--output-file") args = parser.parse_args() pycbc.init_logging(args.verbose) -f = h5.File(args.dq_file, 'r') -grp = f[args.ifo]['bins'] +f = h5.File(args.dq_file, "r") +grp = f[args.ifo]["bins"] bin_names = list(grp.keys()) -sngl_ranking = f.attrs['sngl_ranking'] -sngl_thresh = f.attrs['sngl_ranking_threshold'] +sngl_ranking = f.attrs["sngl_ranking"] +sngl_thresh = f.attrs["sngl_ranking_threshold"] livetime = 0 -seg_grp = f[args.ifo]['dq_segments'] +seg_grp = f[args.ifo]["dq_segments"] for k in seg_grp.keys(): - livetime += seg_grp[k]['livetime'][()] + livetime += seg_grp[k]["livetime"][()] num_templates = [] num_triggers = [] @@ -37,39 +38,56 @@ total_triggers = 0 for bin_name in bin_names: bin_grp = grp[bin_name] - n_tmp = len(bin_grp['tids'][:]) + n_tmp = len(bin_grp["tids"][:]) num_templates.append(n_tmp) total_templates += n_tmp - n_trig = bin_grp['num_triggers'][()] + n_trig = bin_grp["num_triggers"][()] num_triggers.append(n_trig) total_triggers += n_trig -bin_names.append('Total') +bin_names.append("Total") num_triggers.append(total_triggers) num_templates.append(total_templates) frac_triggers = [n / total_triggers for n in num_triggers] frac_templates = [n / total_templates for n in num_templates] trigger_rate = [n / livetime for n in num_triggers] -col_names = ['Template Bin', 'Number of Templates', '% of Templates', - 'Number of Loud Triggers', '% of Loud Triggers', - 'Loud Trigger Rate (Hz)'] -columns = [bin_names, num_templates, frac_templates, - num_triggers, frac_triggers, trigger_rate] +col_names = [ + "Template Bin", + "Number of Templates", + "% of Templates", + "Number of Loud Triggers", + "% of Loud Triggers", + "Loud Trigger Rate (Hz)", +] +columns = [ + bin_names, + num_templates, + frac_templates, + num_triggers, + frac_triggers, + trigger_rate, +] columns = [np.array(c) for c in columns] -format_strings = [None, '#', '0.000%', '#', '0.0%', '0.00E0'] +format_strings = [None, "#", "0.000%", "#", "0.0%", "0.00E0"] -html_table = pycbc.results.html_table(columns, col_names, - page_size=len(bin_names), - format_strings=format_strings) -title = f'{args.ifo} DQ Template Bin Information' -caption = 'Table of information about template bins ' \ - + 'used for DQ trigger rate calculations. ' \ - + 'Loud triggers are defined as those with ' \ - + f'{sngl_ranking} > {sngl_thresh}.' +html_table = pycbc.results.html_table( + columns, col_names, page_size=len(bin_names), format_strings=format_strings +) +title = f"{args.ifo} DQ Template Bin Information" +caption = ( + "Table of information about template bins " + "used for DQ trigger rate calculations. " + "Loud triggers are defined as those with " + f"{sngl_ranking} > {sngl_thresh}." +) pycbc.results.save_fig_with_metadata( - str(html_table), args.output_file, title=title, - caption=caption, cmd=' '.join(sys.argv)) + str(html_table), + args.output_file, + title=title, + caption=caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_page_versioning b/bin/plotting/pycbc_page_versioning index 2bb083437b0..ae28873610b 100755 --- a/bin/plotting/pycbc_page_versioning +++ b/bin/plotting/pycbc_page_versioning @@ -10,57 +10,64 @@ import argparse import logging import pycbc -from pycbc.results import (save_fig_with_metadata, html_escape, - get_library_version_info, get_code_version_numbers) +from pycbc.results import ( + get_code_version_numbers, + get_library_version_info, + html_escape, + save_fig_with_metadata, +) parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--executables', nargs='+', required=True, - help="List of executables to provide version " - "information for") -parser.add_argument('--executables-names', nargs='+', required=True, - help="Names of the executables, must be in the " - "same order as --executables-files") -parser.add_argument("--output-file", required=True, - help="The directory for output html snippets") +parser.add_argument( + "--executables", + nargs="+", + required=True, + help="List of executables to provide version information for", +) +parser.add_argument( + "--executables-names", + nargs="+", + required=True, + help="Names of the executables, must be in the same order as --executables-files", +) +parser.add_argument( + "--output-file", required=True, help="The directory for output html snippets" +) args = parser.parse_args() pycbc.init_logging(args.verbose) if not len(args.executables) == len(args.executables_names): - raise parser.error("--executables-files and executables-names must be " - "the same number of arguments") + raise parser.error( + "--executables-files and executables-names must be the same number of arguments" + ) logging.info("Getting version information for libraries") library_list = get_library_version_info() -html_text = '' +html_text = "" for curr_lib in library_list: - lib_name = curr_lib['Name'] + lib_name = curr_lib["Name"] logging.info(f"Getting {lib_name} information") - html_text += f'

{lib_name} Version Information

:
\n' + html_text += f"

{lib_name} Version Information

:
\n" for key, value in curr_lib.items(): - html_text += '
  • %s : %s
    • \n' % (key, value) + html_text += "
  • %s : %s
    • \n" % (key, value) -code_version_dict = get_code_version_numbers( - args.executables_names, - args.executables -) +code_version_dict = get_code_version_numbers(args.executables_names, args.executables) -html_text += f'

    Version Information from Executables

    :
    \n' +html_text += "

    Version Information from Executables

    :
    \n" for key, value in code_version_dict.items(): - html_text += '
  • %s:
    %s


    • \n' \ - % (key, str(value).replace('@', '@')) + html_text += "
  • %s:
    %s


    • \n" % ( + key, + str(value).replace("@", "@"), + ) kwds = { - 'render-function' : 'render_text', - 'title' : 'Version Information', + "render-function": "render_text", + "title": "Version Information", } -save_fig_with_metadata( - html_escape(html_text), - args.output_file, - **kwds -) +save_fig_with_metadata(html_escape(html_text), args.output_file, **kwds) logging.info("Done") diff --git a/bin/plotting/pycbc_page_vetotable b/bin/plotting/pycbc_page_vetotable index e3c6d9077e3..40feed4f23b 100644 --- a/bin/plotting/pycbc_page_vetotable +++ b/bin/plotting/pycbc_page_vetotable @@ -16,29 +16,27 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"Writes veto_definer table contents to HTML table." +"""Writes veto_definer table contents to HTML table.""" import argparse import logging -import numpy import sys -from igwn_ligolw import lsctables -from igwn_ligolw import utils +import numpy +from igwn_ligolw import lsctables, utils import pycbc.results -from pycbc.results import save_fig_with_metadata from pycbc.io.ligolw import LIGOLWContentHandler - +from pycbc.results import save_fig_with_metadata parser = argparse.ArgumentParser(description=__doc__) # add command line options pycbc.add_common_pycbc_options(parser) -parser.add_argument('--veto-definer-file', type=str, - help='XML files with a veto_definer table to read.') -parser.add_argument('--output-file', type=str, - help='Path of the output HTML file.') +parser.add_argument( + "--veto-definer-file", type=str, help="XML files with a veto_definer table to read." +) +parser.add_argument("--output-file", type=str, help="Path of the output HTML file.") # parse command line opts = parser.parse_args() @@ -47,32 +45,33 @@ opts = parser.parse_args() pycbc.init_logging(opts.verbose, default_level=1) # set column names -columns = (('Category', []), - ('IFO', []), - ('Name', []), - ('Version', []), - ('Start Padding (s)', []), - ('End Padding (s)', []), - ('Start Time', []), - ('End Time', []), - ('Comment', []), +columns = ( + ("Category", []), + ("IFO", []), + ("Name", []), + ("Version", []), + ("Start Padding (s)", []), + ("End Padding (s)", []), + ("Start Time", []), + ("End Time", []), + ("Comment", []), ) # set caption name caption = "This table shows each veto in the veto definer table." # read input file -logging.info('Reading veto definer XML file') -vetodef_xml = utils.load_filename(opts.veto_definer_file, - contenthandler=LIGOLWContentHandler) +logging.info("Reading veto definer XML file") +vetodef_xml = utils.load_filename( + opts.veto_definer_file, contenthandler=LIGOLWContentHandler +) # get veto_definer table vetodef_table = lsctables.VetoDefTable.get_table(vetodef_xml) # loop over rows in veto_definer table -logging.info('Looping through veto_definer table and saving information') +logging.info("Looping through veto_definer table and saving information") for vetodef in vetodef_table: - # put values into columns columns[0][1].append(vetodef.category) columns[1][1].append(str(vetodef.ifo)) @@ -85,17 +84,20 @@ for vetodef in vetodef_table: columns[8][1].append(str(vetodef.comment)) # cast columns into arrays -keys = [numpy.array(key, dtype=type(key[0])) for key,_ in columns] -vals = [numpy.array(val, dtype=type(val[0])) for _,val in columns] +keys = [numpy.array(key, dtype=type(key[0])) for key, _ in columns] +vals = [numpy.array(val, dtype=type(val[0])) for _, val in columns] # write HTML table -logging.info('Writing HTML table') +logging.info("Writing HTML table") fig_kwds = {} html_table = pycbc.results.html_table(vals, keys, page_size=25) -save_fig_with_metadata(str(html_table), opts.output_file, - fig_kwds=fig_kwds, - title='Veto Definer Table', - cmd=' '.join(sys.argv), - caption=caption) +save_fig_with_metadata( + str(html_table), + opts.output_file, + fig_kwds=fig_kwds, + title="Veto Definer Table", + cmd=" ".join(sys.argv), + caption=caption, +) -logging.info('Done.') +logging.info("Done.") diff --git a/bin/plotting/pycbc_plot_background_coincs b/bin/plotting/pycbc_plot_background_coincs index 7faeb01f4bc..7e879745bbd 100644 --- a/bin/plotting/pycbc_plot_background_coincs +++ b/bin/plotting/pycbc_plot_background_coincs @@ -1,66 +1,71 @@ #!/usr/bin/env python -""" Plot PyCBC's background coinc triggers -""" +"""Plot PyCBC's background coinc triggers""" + import argparse + import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from matplotlib import pyplot as plt from matplotlib.colors import LogNorm from matplotlib.ticker import LogLocator -from pycbc.io.hdf import HFile from pycbc import add_common_pycbc_options, init_logging +from pycbc.io.hdf import HFile + def get_var(data, name): if name in data: return data[name][:] + parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--coinc-file', help="Coincident trigger file. The result" - " of pycbc_coinc_statmap ") -parser.add_argument('--x-var', type=str, required=True, - help='Parameter to plot on the x-axis') -parser.add_argument('--y-var', type=str, required=True, - help='Parameter to plot on the y-axis') -parser.add_argument('--z-var', required=True, - help='Quantity to plot on the color scale', - choices=['density', 'ranking_stat']) -parser.add_argument('--min-z', type=float, help='Optional minimum z value') -parser.add_argument('--max-z', type=float, help='Optional maximum z value') -parser.add_argument('--grid-size', default=100, help="Number of hexbins", type=int) -parser.add_argument('--dpi', type=int, default=200) -parser.add_argument('--output-file') -args = parser.parse_args() +parser.add_argument( + "--coinc-file", help="Coincident trigger file. The result of pycbc_coinc_statmap " +) +parser.add_argument( + "--x-var", type=str, required=True, help="Parameter to plot on the x-axis" +) +parser.add_argument( + "--y-var", type=str, required=True, help="Parameter to plot on the y-axis" +) +parser.add_argument( + "--z-var", + required=True, + help="Quantity to plot on the color scale", + choices=["density", "ranking_stat"], +) +parser.add_argument("--min-z", type=float, help="Optional minimum z value") +parser.add_argument("--max-z", type=float, help="Optional maximum z value") +parser.add_argument("--grid-size", default=100, help="Number of hexbins", type=int) +parser.add_argument("--dpi", type=int, default=200) +parser.add_argument("--output-file") +args = parser.parse_args() init_logging(args.verbose) f = HFile(args.coinc_file) -bdata = f['background_exc'] +bdata = f["background_exc"] x = get_var(bdata, args.x_var) y = get_var(bdata, args.y_var) -hexbin_style = { - 'gridsize': args.grid_size, - 'mincnt': 1, - 'linewidths': 0.02 -} +hexbin_style = {"gridsize": args.grid_size, "mincnt": 1, "linewidths": 0.02} if args.min_z is not None: - hexbin_style['vmin'] = args.min_z + hexbin_style["vmin"] = args.min_z if args.max_z is not None: - hexbin_style['vmax'] = args.max_z + hexbin_style["vmax"] = args.max_z fig = plt.figure() ax = fig.gca() -if args.z_var == 'density': +if args.z_var == "density": hb = ax.hexbin(x, y, norm=LogNorm(), vmin=1, **hexbin_style) - fig.colorbar(hb, ticks=LogLocator(subs=range(10))) -elif args.z_var == 'ranking_stat': - hb = ax.hexbin(x, y, C=bdata['stat'][:], reduce_C_function=max, **hexbin_style) + fig.colorbar(hb, ticks=LogLocator(subs=range(10))) +elif args.z_var == "ranking_stat": + hb = ax.hexbin(x, y, C=bdata["stat"][:], reduce_C_function=max, **hexbin_style) fig.colorbar(hb) ax.set_xlabel(args.x_var) ax.set_ylabel(args.y_var) ax.set_title("Coincident Background Triggers, %s" % args.z_var) fig.savefig(args.output_file, dpi=args.dpi) - diff --git a/bin/plotting/pycbc_plot_bank_bins b/bin/plotting/pycbc_plot_bank_bins index e40e95ef9ce..f44b2f04eaf 100644 --- a/bin/plotting/pycbc_plot_bank_bins +++ b/bin/plotting/pycbc_plot_bank_bins @@ -1,140 +1,174 @@ #!/bin/env python -""" plot a an hdf bank file based on background binning -""" -import sys +"""plot a an hdf bank file based on background binning""" + import argparse +import sys + import h5py -import numpy import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt +import numpy + +matplotlib.use("Agg") import inspect from itertools import cycle -import pycbc.events, pycbc.pnutils, pycbc.conversions, pycbc.results +from matplotlib import pyplot as plt + +import pycbc.conversions +import pycbc.events +import pycbc.pnutils +import pycbc.results class H5BankFile(h5py.File): - "Convenience class for getting CBC parameters out of an HDF5 bank." + """Convenience class for getting CBC parameters out of an HDF5 bank.""" @classmethod def get_param_names(cls): - """Returns a list of CBC parameters which can be obtained - from the class instance.""" - return sorted([m[0].replace('_param', '') for m in inspect.getmembers(cls) - if m[0].endswith('_param')]) + """ + Returns a list of CBC parameters which can be obtained + from the class instance. + """ + return sorted( + [ + m[0].replace("_param", "") + for m in inspect.getmembers(cls) + if m[0].endswith("_param") + ] + ) def __len__(self): - return len(self['mass1']) + return len(self["mass1"]) def mass1_param(self): - 'Mass 1 $M_\odot$' - return self['mass1'][:] + r"""Mass 1 $M_\odot$""" + return self["mass1"][:] def mass2_param(self): - 'Mass 2 $M_\odot$' - return self['mass2'][:] + r"""Mass 2 $M_\odot$""" + return self["mass2"][:] def spin1z_param(self): - 'Spin 1 z-component' - return self['spin1z'][:] + """Spin 1 z-component""" + return self["spin1z"][:] def spin2z_param(self): - 'Spin 2 z-component' - return self['spin2z'][:] + """Spin 2 z-component""" + return self["spin2z"][:] def chirp_mass_param(self): - 'Chirp mass $M_\odot$' + r"""Chirp mass $M_\odot$""" return pycbc.pnutils.mass1_mass2_to_mchirp_eta( - self.mass1_param(), self.mass2_param())[0] + self.mass1_param(), self.mass2_param() + )[0] def total_mass_param(self): - 'Total mass $M_\odot$' + r"""Total mass $M_\odot$""" return self.mass1_param() + self.mass2_param() def mass_ratio_param(self): - 'Mass ratio' + """Mass ratio""" return self.mass1_param() / self.mass2_param() def eta_param(self): - 'Symmetric mass ratio' + """Symmetric mass ratio""" return pycbc.conversions.eta_from_mass1_mass2( - self.mass1_param(), self.mass2_param()) + self.mass1_param(), self.mass2_param() + ) def effective_spin_param(self): - 'Effective spin z-component' - return (self.spin1z_param() * self.mass1_param() + self.spin2z_param() * self.mass2_param()) / self.total_mass_param() + """Effective spin z-component""" + return ( + self.spin1z_param() * self.mass1_param() + + self.spin2z_param() * self.mass2_param() + ) / self.total_mass_param() def tau0_param(self): - '$\\tau_0$' + """$\\tau_0$""" return pycbc.pnutils.mass1_mass2_to_tau0_tau3( - self.mass1_param(), self.mass2_param(), self.f_lower)[0] + self.mass1_param(), self.mass2_param(), self.f_lower + )[0] def tau3_param(self): - '$\\tau_3$' + """$\\tau_3$""" return pycbc.pnutils.mass1_mass2_to_tau0_tau3( - self.mass1_param(), self.mass2_param(), self.f_lower)[1] + self.mass1_param(), self.mass2_param(), self.f_lower + )[1] parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--bank-file', help='hdf format template bank file', - required=True) -parser.add_argument('--background-bins', nargs='+', - help='list of background bin format strings') -parser.add_argument('--f-lower', type=float, - help="Lower frequency cutoff for evaluating template " - "duration. Should be equal to the lower cutoff " - "used in inspiral jobs") -parser.add_argument('--output-file', help='output file', required=True) -parser.add_argument('--x-var', type=str, choices=H5BankFile.get_param_names(), - default='mass1', - help='Template parameter to plot on the x-axis') -parser.add_argument('--y-var', type=str, choices=H5BankFile.get_param_names(), - default='mass2', - help='Template parameter to plot on the y-axis') -parser.add_argument('--log-x', action='store_true', - help='Make x-axis logarithmic') -parser.add_argument('--log-y', action='store_true', - help='Make y-axis logarithmic') +parser.add_argument("--bank-file", help="hdf format template bank file", required=True) +parser.add_argument( + "--background-bins", nargs="+", help="list of background bin format strings" +) +parser.add_argument( + "--f-lower", + type=float, + help="Lower frequency cutoff for evaluating template " + "duration. Should be equal to the lower cutoff " + "used in inspiral jobs", +) +parser.add_argument("--output-file", help="output file", required=True) +parser.add_argument( + "--x-var", + type=str, + choices=H5BankFile.get_param_names(), + default="mass1", + help="Template parameter to plot on the x-axis", +) +parser.add_argument( + "--y-var", + type=str, + choices=H5BankFile.get_param_names(), + default="mass2", + help="Template parameter to plot on the y-axis", +) +parser.add_argument("--log-x", action="store_true", help="Make x-axis logarithmic") +parser.add_argument("--log-y", action="store_true", help="Make y-axis logarithmic") args = parser.parse_args() pycbc.init_logging(args.verbose) -bank = H5BankFile(args.bank_file, 'r') -f_lower = args.f_lower or bank['f_lower'][:] +bank = H5BankFile(args.bank_file, "r") +f_lower = args.f_lower or bank["f_lower"][:] if args.background_bins: - data = {'mass1': bank['mass1'][:], 'mass2': bank['mass2'][:], - 'spin1z': bank['spin1z'][:], 'spin2z': bank['spin2z'][:], - 'f_lower': f_lower} + data = { + "mass1": bank["mass1"][:], + "mass2": bank["mass2"][:], + "spin1z": bank["spin1z"][:], + "spin2z": bank["spin2z"][:], + "f_lower": f_lower, + } locs_dict = pycbc.events.background_bin_from_string(args.background_bins, data) else: - locs_dict = {'Template Bank': numpy.arange(0, len(bank), 1)} + locs_dict = {"Template Bank": numpy.arange(0, len(bank), 1)} -color = cycle(['red', 'green', 'blue', 'purple']) +color = cycle(["red", "green", "blue", "purple"]) -x_var = getattr(bank, args.x_var + '_param')() -x_var_name = getattr(bank, args.x_var + '_param').__doc__ -y_var = getattr(bank, args.y_var + '_param')() -y_var_name = getattr(bank, args.y_var + '_param').__doc__ +x_var = getattr(bank, args.x_var + "_param")() +x_var_name = getattr(bank, args.x_var + "_param").__doc__ +y_var = getattr(bank, args.y_var + "_param")() +y_var_name = getattr(bank, args.y_var + "_param").__doc__ fig = plt.figure() plt.grid() for name in locs_dict: locs = locs_dict[name] - plt.scatter(x_var[locs], y_var[locs], label=name, edgecolor='none', s=1, - c=next(color)) + plt.scatter( + x_var[locs], y_var[locs], label=name, edgecolor="none", s=1, c=next(color) + ) -plt.legend(loc='upper left', markerscale=5) +plt.legend(loc="upper left", markerscale=5) plt.xlabel(x_var_name) plt.ylabel(y_var_name) plt.xlim(x_var.min(), x_var.max()) plt.ylim(y_var.min(), y_var.max()) if args.log_x: - plt.xscale('log') + plt.xscale("log") if args.log_y: - plt.yscale('log') + plt.yscale("log") title = "Template Bank and Bins Used to Compute Background" caption = """This plot shows the template bank in the {x_var}-{y_var} plane. @@ -143,5 +177,6 @@ compute the search background. Note that the bins may be chosen in a space higher than two dimensions for spinning templates, causing apparent overlap in the {x_var}-{y_var} plane.""" caption = caption.format(x_var=args.x_var, y_var=args.y_var) -pycbc.results.save_fig_with_metadata(fig, args.output_file, title=title, - caption=caption, cmd=' '.join(sys.argv)) +pycbc.results.save_fig_with_metadata( + fig, args.output_file, title=title, caption=caption, cmd=" ".join(sys.argv) +) diff --git a/bin/plotting/pycbc_plot_bank_compression b/bin/plotting/pycbc_plot_bank_compression index 3afd4eaa88e..46bf159033c 100644 --- a/bin/plotting/pycbc_plot_bank_compression +++ b/bin/plotting/pycbc_plot_bank_compression @@ -6,18 +6,21 @@ with compressed waveforms """ import argparse + import matplotlib + matplotlib.use("agg") -from matplotlib import pyplot as plt -import numpy as np import logging import sys +import numpy as np +from matplotlib import pyplot as plt + import pycbc +from pycbc import tmpltbank +from pycbc.inference import option_utils from pycbc.io import HFile from pycbc.results import save_fig_with_metadata -from pycbc.inference import option_utils -import pycbc.tmpltbank as tmpltbank parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) @@ -57,37 +60,35 @@ parser.add_argument( "--histogram-density", action="store_true", help="Flag to indicate that the histogram should be a density " - "rather than a count per bin", + "rather than a count per bin", ) parser.add_argument( "--log-histogram", action="store_true", - help="Flag to indicate that histogram values should be plotted " - "on a log scale" + help="Flag to indicate that histogram values should be plotted on a log scale", ) default_param = "template_duration" -parser.add_argument("--comparison-parameter", +parser.add_argument( + "--comparison-parameter", action=option_utils.ParseParametersArg, metavar="PARAM[:LABEL]", help="Plot the scatter plot of compressin factor versus the given " - "parameter. Optionally provide a LABEL for use in the plot. " - "Choose from " + ", ".join(tmpltbank.conversion_options) + ", " - "though some options may not be buildable from bank parameters. " - "If no LABEL is provided, PARAM will used as the LABEL. If LABEL " - "is the same as a parameter in pycbc.waveform.parameters, the label " - "property of that parameter will be used. Default: " + default_param + "parameter. Optionally provide a LABEL for use in the plot. " + "Choose from " + ", ".join(tmpltbank.conversion_options) + ", " + "though some options may not be buildable from bank parameters. " + "If no LABEL is provided, PARAM will used as the LABEL. If LABEL " + "is the same as a parameter in pycbc.waveform.parameters, the label " + "property of that parameter will be used. Default: " + default_param, ) args = parser.parse_args() if args.comparison_parameter is None: args.comparison_parameter = default_param - args.comparison_parameter_labels = { - default_param: "Template Duration (s)" - } + args.comparison_parameter_labels = {default_param: "Template Duration (s)"} elif args.comparison_parameter not in tmpltbank.conversion_options: raise parser.error( "--comparison-parameter %s not in conversion options %s, see help" - % (args.comparison_parameter, ', '.join(tmpltbank.conversion_options)) + % (args.comparison_parameter, ", ".join(tmpltbank.conversion_options)) ) pycbc.init_logging(args.verbose) @@ -96,7 +97,7 @@ comp_label = args.comparison_parameter_labels[args.comparison_parameter] # Quieten the matplotlib logger plt.set_loglevel("info" if args.verbose else "warning") -logging.getLogger('matplotlib.font_manager').setLevel(logging.ERROR) +logging.getLogger("matplotlib.font_manager").setLevel(logging.ERROR) logging.info("Getting information from the bank") @@ -107,12 +108,7 @@ approximants = [] mismatch = [] total_templates = 0 for i, bank_fname in enumerate(args.bank_files): - logging.debug( - "Bank %d out of %d: %s", - i, - len(args.bank_files), - bank_fname - ) + logging.debug("Bank %d out of %d: %s", i, len(args.bank_files), bank_fname) with HFile(bank_fname, "r") as bank_f: compressed_grp = bank_f["compressed_waveforms"] @@ -123,7 +119,7 @@ for i, bank_fname in enumerate(args.bank_files): mismatch_thisbank = np.zeros(thashes.size, dtype=float) logging.debug( "Bank contains %s approximant(s)", - ','.join(apx.decode() for apx in np.unique(bank_f["approximant"][:])) + ",".join(apx.decode() for apx in np.unique(bank_f["approximant"][:])), ) logging.debug("Getting compression factors and mismatches") @@ -133,38 +129,33 @@ for i, bank_fname in enumerate(args.bank_files): except KeyError: continue - compression_thisbank[i] = this_grp.attrs['compression_factor'] - mismatch_thisbank[i] = this_grp.attrs['mismatch'] + compression_thisbank[i] = this_grp.attrs["compression_factor"] + mismatch_thisbank[i] = this_grp.attrs["mismatch"] valid_idx[i] = True - + logging.info( "%d out of %d compressed waveforms in %s", np.count_nonzero(valid_idx), thashes.size, - bank_fname + bank_fname, ) if not any(valid_idx): continue - + compression_factor += list(compression_thisbank[valid_idx]) mismatch += list(mismatch_thisbank[valid_idx]) logging.debug("Getting approximants") - approximants += [ - apx.decode() for apx in - bank_f["approximant"][valid_idx] - ] + approximants += [apx.decode() for apx in bank_f["approximant"][valid_idx]] logging.debug("Getting %s", args.comparison_parameter) - comparison_values += list(tmpltbank.get_bank_property( - args.comparison_parameter, - bank_f, - template_ids=valid_idx - )) + comparison_values += list( + tmpltbank.get_bank_property( + args.comparison_parameter, bank_f, template_ids=valid_idx + ) + ) if not len(compression_factor): - raise ValueError( - "No compressed waveforms found in any of the given banks" - ) + raise ValueError("No compressed waveforms found in any of the given banks") approximants = np.array(approximants) @@ -173,9 +164,7 @@ compression_factor = np.array(compression_factor) mismatch = np.array(mismatch) logging.info( - "%d out of %d templates have compressed waveforms", - mismatch.size, - total_templates + "%d out of %d templates have compressed waveforms", mismatch.size, total_templates ) # Store the max/min factors, as these are used for setting @@ -193,58 +182,47 @@ if args.log_compression: factor_range = np.log10(max_factor) - np.log10(min_factor) factor_max = np.log10(max_factor) + 0.05 * factor_range factor_min = max(np.log10(min_factor) - 0.05 * factor_range, 0) - bin_comp_edges = np.logspace( - factor_min, - factor_max, - args.n_bins - ) + bin_comp_edges = np.logspace(factor_min, factor_max, args.n_bins) else: factor_range = max_factor - min_factor factor_min = max(min_factor - 0.05 * factor_range, 1) factor_max = max_factor + 0.05 * factor_range - bin_comp_edges = np.linspace( - factor_min, - factor_max, - args.n_bins - ) + bin_comp_edges = np.linspace(factor_min, factor_max, args.n_bins) if args.log_mismatch: mmatch_range = np.log10(max_mmatch) - np.log10(min_mmatch) mmatch_max = np.log10(max_mmatch) + 0.05 * mmatch_range mmatch_min = np.log10(min_mmatch) - 0.05 * mmatch_range - bin_mmatch_edges = np.logspace( - mmatch_min, - mmatch_max, - args.n_bins - ) + bin_mmatch_edges = np.logspace(mmatch_min, mmatch_max, args.n_bins) else: mmatch_range = max_mmatch - min_mmatch mmatch_min = min_mmatch - 0.05 * mmatch_range mmatch_max = max_mmatch + 0.05 * mmatch_range - bin_mmatch_edges = np.linspace( - mmatch_min, - mmatch_max, - args.n_bins - ) + bin_mmatch_edges = np.linspace(mmatch_min, mmatch_max, args.n_bins) # These are used as the x values in the histogram plot; # this will be each bin edge repeated twice in order, # except the first and last which will appear once bin_comp_step = np.concatenate(tuple(zip(bin_comp_edges[:-1], bin_comp_edges[1:]))) -bin_mmatch_step = np.concatenate(tuple(zip(bin_mmatch_edges[:-1], bin_mmatch_edges[1:]))) +bin_mmatch_step = np.concatenate( + tuple(zip(bin_mmatch_edges[:-1], bin_mmatch_edges[1:])) +) fig, axes = plt.subplots( - 2, 2, - figsize=(8,8), + 2, + 2, + figsize=(8, 8), layout="constrained", - sharex='col', + sharex="col", ) # Make the histogram and scatter points for each approximant separately apx_names, apx_count = np.unique(approximants, return_counts=True) apx_colors = { - apx_name: col for apx_name, col in - zip(apx_names, plt.rcParams['axes.prop_cycle'].by_key()['color']) + apx_name: col + for apx_name, col in zip( + apx_names, plt.rcParams["axes.prop_cycle"].by_key()["color"] + ) } for apx_name, apx_count in zip(apx_names, apx_count): # Filter to just this approximant @@ -261,11 +239,11 @@ for apx_name, apx_count in zip(apx_names, apx_count): comp_hist *= apx_count / compression_factor.size logging.debug("Plotting compression factor histogram") - axes[0,0].plot( + axes[0, 0].plot( bin_comp_step, np.repeat(comp_hist, 2), c=apx_colors[apx_name], - label=f'{apx_name}: {apx_count}', + label=f"{apx_name}: {apx_count}", ) logging.info("Making %s mismatch histogram", apx_name) @@ -278,35 +256,26 @@ for apx_name, apx_count in zip(apx_names, apx_count): mmatch_hist *= apx_count / mismatch.size logging.debug("Plotting mismatch histogram") - axes[0,1].plot( + axes[0, 1].plot( bin_mmatch_step, np.repeat(mmatch_hist, 2), c=apx_colors[apx_name], - label=f'{apx_name}: {apx_count}', + label=f"{apx_name}: {apx_count}", ) - logging.debug( - "Plotting compression vs %s scatter", - args.comparison_parameter - ) + logging.debug("Plotting compression vs %s scatter", args.comparison_parameter) # This makes it so that if there aren't that many points, # they are visible, but more points become more cloud-like - scatter_alpha = max( - min(100 / len(compression_factor), 1), - 0.05 - ) - axes[1,0].scatter( + scatter_alpha = max(min(100 / len(compression_factor), 1), 0.05) + axes[1, 0].scatter( compression_factor[this_appx], comparison_values[this_appx], alpha=scatter_alpha, c=apx_colors[apx_name], s=5, ) - logging.debug( - "Plotting mismatch vs %s scatter", - args.comparison_parameter - ) - axes[1,1].scatter( + logging.debug("Plotting mismatch vs %s scatter", args.comparison_parameter) + axes[1, 1].scatter( mismatch[this_appx], comparison_values[this_appx], alpha=scatter_alpha, @@ -314,62 +283,58 @@ for apx_name, apx_count in zip(apx_names, apx_count): s=5, ) # This one is for use in the legend: - axes[1,0].scatter( - [],[], - c=apx_colors[apx_name], - label=f'{apx_name}: {apx_count}' - ) + axes[1, 0].scatter([], [], c=apx_colors[apx_name], label=f"{apx_name}: {apx_count}") logging.debug("Setting scales and limits of axes") if args.log_compression: - axes[1,0].set_xlim(10 ** factor_min, 10 ** factor_max) - axes[1,0].set_xscale("log") + axes[1, 0].set_xlim(10**factor_min, 10**factor_max) + axes[1, 0].set_xscale("log") else: - axes[0,0].set_xlim(factor_min, factor_max) - axes[0,1].set_ylim(factor_min, factor_max) + axes[0, 0].set_xlim(factor_min, factor_max) + axes[0, 1].set_ylim(factor_min, factor_max) if args.log_mismatch: - axes[1,1].set_xscale("log") - axes[1,1].set_xlim(10 ** mmatch_min, 10 ** mmatch_max) + axes[1, 1].set_xscale("log") + axes[1, 1].set_xlim(10**mmatch_min, 10**mmatch_max) else: - axes[1,1].set_xlim(mmatch_min, mmatch_max) + axes[1, 1].set_xlim(mmatch_min, mmatch_max) if args.log_comparison: - axes[1,0].set_yscale("log") - axes[1,1].set_yscale("log") + axes[1, 0].set_yscale("log") + axes[1, 1].set_yscale("log") if args.log_histogram: - axes[0,0].set_yscale("log") - axes[0,1].set_yscale("log") + axes[0, 0].set_yscale("log") + axes[0, 1].set_yscale("log") else: - axes[0,0].set_ylim(bottom=0) - axes[0,1].set_ylim(bottom=0) + axes[0, 0].set_ylim(bottom=0) + axes[0, 1].set_ylim(bottom=0) logging.info("Setting axes labels") if args.histogram_density: - axes[0,0].set_ylabel("Template Density") - axes[0,1].set_ylabel("Template Density") + axes[0, 0].set_ylabel("Template Density") + axes[0, 1].set_ylabel("Template Density") else: - axes[0,0].set_ylabel("Number of Templates") - axes[0,1].set_ylabel("Number of Templates") -axes[1,0].set_xlabel("Compression Factor") + axes[0, 0].set_ylabel("Number of Templates") + axes[0, 1].set_ylabel("Number of Templates") +axes[1, 0].set_xlabel("Compression Factor") -axes[1,0].set_ylabel(comp_label) -axes[1,1].set_ylabel(comp_label) -axes[1,1].set_xlabel("Mismatch") +axes[1, 0].set_ylabel(comp_label) +axes[1, 1].set_ylabel(comp_label) +axes[1, 1].set_xlabel("Mismatch") -axes[0,0].legend(loc='upper left') -axes[1,0].legend(loc='upper left') +axes[0, 0].legend(loc="upper left") +axes[1, 0].legend(loc="upper left") for ax in axes.flatten(): ax.grid(zorder=-100) caption = ( "Plot showing the a histogram of compression factor (left), a " - "scatter plot of compression factor vs {label} ({parameter}) (middle) " - "and mismatch vs {label} ({parameter}) (right)." + f"scatter plot of compression factor vs {comp_label} ({args.comparison_parameter}) (middle) " + f"and mismatch vs {comp_label} ({args.comparison_parameter}) (right)." "Legend entries indicate the number of templates per approximant. " -).format(label=comp_label, parameter=args.comparison_parameter) +) if args.histogram_density: caption += "Density for each histogram is weighted by the number of templates " @@ -380,6 +345,6 @@ save_fig_with_metadata( args.output, title="Bank compression vs %s" % comp_label, caption=caption, - cmd=' '.join(sys.argv) + cmd=" ".join(sys.argv), ) logging.info("Done!") diff --git a/bin/plotting/pycbc_plot_bank_corner b/bin/plotting/pycbc_plot_bank_corner index 522ed1e2d89..888c72e9865 100644 --- a/bin/plotting/pycbc_plot_bank_corner +++ b/bin/plotting/pycbc_plot_bank_corner @@ -19,119 +19,133 @@ Plots various parameters against one another for pycbc banks in an hdf file. """ - +import argparse +import logging import os import sys +from textwrap import wrap + import h5py import numpy as np -import argparse -import logging -from textwrap import wrap import pycbc -from pycbc.results.plot import (add_style_opt_to_parser, set_style_from_cli) -from pycbc.io import FieldArray, HFile from pycbc.inference import option_utils -from pycbc.tmpltbank import bank_conversions as bconv -from pycbc.results.scatter_histograms import create_multidim_plot +from pycbc.io import FieldArray, HFile from pycbc.results import metadata +from pycbc.results.plot import add_style_opt_to_parser, set_style_from_cli +from pycbc.results.scatter_histograms import create_multidim_plot +from pycbc.tmpltbank import bank_conversions as bconv conversion_options = bconv.conversion_options _fit_parameters = [ - 'count_above_thresh', - 'count_in_template', - 'fit_coeff', - 'median_sigma' + "count_above_thresh", + "count_in_template", + "fit_coeff", + "median_sigma", ] parameter_options = conversion_options + _fit_parameters -parser = argparse.ArgumentParser(usage='pycbc_plot_bank_corner [--options]', - description=__doc__) +parser = argparse.ArgumentParser( + usage="pycbc_plot_bank_corner [--options]", description=__doc__ +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--bank-file", - required=True, - help="The bank file to read in and plot") -parser.add_argument("--output-plot-file", - required=True, - help="Name of the output file") -parser.add_argument("--parameters", +parser.add_argument( + "--bank-file", required=True, help="The bank file to read in and plot" +) +parser.add_argument("--output-plot-file", required=True, help="Name of the output file") +parser.add_argument( + "--parameters", nargs="+", default=[], action=option_utils.ParseParametersArg, metavar="PARAM[:LABEL]", help="Only plot the given parameters. May also provide a " - "label for each parameter. Choose from a combination of " + \ - ", ".join(parameter_options) + ", though some " - "may not be buildable from bank parameters. Can optionally " - "also specify a LABEL for each parameter. If no LABEL is " - "provided, PARAM will used as the LABEL. If LABEL " - "is the same as a parameter in " - "pycbc.waveform.parameters, the label " - "property of that parameter will be used. If not " - "provided, will plot all of the parameters in the " - "bank.") + "label for each parameter. Choose from a combination of " + + ", ".join(parameter_options) + + ", though some " + "may not be buildable from bank parameters. Can optionally " + "also specify a LABEL for each parameter. If no LABEL is " + "provided, PARAM will used as the LABEL. If LABEL " + "is the same as a parameter in " + "pycbc.waveform.parameters, the label " + "property of that parameter will be used. If not " + "provided, will plot all of the parameters in the " + "bank.", +) parser.add_argument( - '--log-parameters', - nargs='+', + "--log-parameters", + nargs="+", default=[], help="Which parameters are to be plotted on a log scale? " - "Must be also given in parameters" + "Must be also given in parameters", +) +parser.add_argument( + "--plot-histogram", + action="store_true", + help="Plot 1D histograms of parameters on the diagonal axes.", ) -parser.add_argument('--plot-histogram', - action='store_true', - help="Plot 1D histograms of parameters on the " - "diagonal axes.") # add mins, maxs options -parser.add_argument('--mins', - nargs='+', - metavar='PARAM:VAL', +parser.add_argument( + "--mins", + nargs="+", + metavar="PARAM:VAL", default=[], help="Specify minimum parameter values to plot. This " - "should be done by specifying the parameter name " - "followed by the value. Parameter names must be " - "the same as the PARAM argument in --parameters " - "(or, if no parameters are provided, the same as " - "the parameter name specified in the variable " - "args in the input bank. If none provided, " - "the smallest parameter value in the posterior " - "will be used.") -parser.add_argument('--maxs', - nargs='+', - metavar='PARAM:VAL', + "should be done by specifying the parameter name " + "followed by the value. Parameter names must be " + "the same as the PARAM argument in --parameters " + "(or, if no parameters are provided, the same as " + "the parameter name specified in the variable " + "args in the input bank. If none provided, " + "the smallest parameter value in the posterior " + "will be used.", +) +parser.add_argument( + "--maxs", + nargs="+", + metavar="PARAM:VAL", default=[], - help="Same as mins, but for the maximum values to " - "plot.") -parser.add_argument("--color-parameter", + help="Same as mins, but for the maximum values to plot.", +) +parser.add_argument( + "--color-parameter", nargs=1, action=option_utils.ParseParametersArg, - metavar='PARAM[:LABEL]', + metavar="PARAM[:LABEL]", help="Color scatter points according to the parameter given. " - "May optionally provide a label in the same way as for " - "--parameters. Default=No scatter point coloring.") + "May optionally provide a label in the same way as for " + "--parameters. Default=No scatter point coloring.", +) parser.add_argument( - '--log-colormap', - action='store_true', - help="Should the colorbar be plotted on a log scale?" + "--log-colormap", + action="store_true", + help="Should the colorbar be plotted on a log scale?", ) -parser.add_argument('--dpi', - type=int, - default=200, - help="Set the DPI of the plot. Default is 200.") -parser.add_argument('--fits-file', - help="Provide a fits file to plot parameters from. Required if any of " - ', '.join(_fit_parameters) + " are given as parameters.") -parser.add_argument('--title', +parser.add_argument( + "--dpi", type=int, default=200, help="Set the DPI of the plot. Default is 200." +) +parser.add_argument( + "--fits-file", + help="Provide a fits file to plot parameters from. Required if any of , ".join( + _fit_parameters + ) + + " are given as parameters.", +) +parser.add_argument( + "--title", help="A title for the plot. If not given, the files supplied " - "and the number of templates will be used") -parser.add_argument('--no-suptitle', action='store_true', - help="If true no suptitle is shown.") + "and the number of templates will be used", +) +parser.add_argument( + "--no-suptitle", action="store_true", help="If true no suptitle is shown." +) add_style_opt_to_parser(parser) args = parser.parse_args() for lp in args.log_parameters: - if not lp in args.parameters: + if lp not in args.parameters: parser.error( "--log-parameters should be in --parameters. " f"{lp} not in [{', '.join(args.parameters)}]" @@ -144,18 +158,20 @@ mins, maxs = option_utils.plot_ranges_from_cli(args) logging.info("Reading in the bank") bank = {} -with HFile(args.bank_file, 'r') as bankf: +with HFile(args.bank_file, "r") as bankf: for p in bankf.keys(): - if not isinstance(bankf[p], h5py.Dataset): continue + if not isinstance(bankf[p], h5py.Dataset): + continue # Ignore things which aren't numbers, which cannot # be histogrammed - e.g. approximant - if not np.issubdtype(bankf[p].dtype, np.number): continue + if not np.issubdtype(bankf[p].dtype, np.number): + continue bank[p] = bankf[p][:] banklen = bankf[p].size if args.fits_file is not None: # Add fit parameters to the bank object - with HFile(args.fits_file, 'r') as fits_f: + with HFile(args.fits_file, "r") as fits_f: for p in _fit_parameters: if not fits_f[p].size == banklen: raise RuntimeError( @@ -174,8 +190,7 @@ logging.info("Got %d templates from the bank", banklen) # If no parameters are given - just plot whatever is in the bank if not args.parameters: # Plot anything which is not singular - args.parameters = [k for k in bank - if len(np.unique(bank[k])) > 1] + args.parameters = [k for k in bank if len(np.unique(bank[k])) > 1] args.parameters_labels = {k: k for k in bank} # Check through the conversion options to see if any are used in @@ -191,46 +206,47 @@ for co in parameter_options: required_params.append(co) # Check for possible double-counting of duration parameter: -if 'duration' in required_params and any([p.endswith('_duration') - for p in required_params]): +if "duration" in required_params and any( + [p.endswith("_duration") for p in required_params] +): # Check whether duration was specified directly: - dur_pars = [par for par in args.parameters if 'duration' in par] - _dur_pars = [par for par in args.parameters if '_duration' in par] + dur_pars = [par for par in args.parameters if "duration" in par] + _dur_pars = [par for par in args.parameters if "_duration" in par] if len(dur_pars) == len(_dur_pars): # It looks like duration has been added where it didnt need to be - required_params.remove('duration') + required_params.remove("duration") # Do the same with the duration functions, but here we need to make # sure we have some other keys in order to calculate duration -duration_required_keys = ['mass1', 'mass2', 'spin1z', - 'spin2z', 'f_lower'] +duration_required_keys = ["mass1", "mass2", "spin1z", "spin2z", "f_lower"] -if any(['duration' in par for par in args.parameters]): +if any(["duration" in par for par in args.parameters]): required_params += duration_required_keys # Some things may have been double counted - undo this required_params = np.unique(required_params) -logging.info("Required parameters to get from bank: %s", - ', '.join(required_params)) +logging.info("Required parameters to get from bank: %s", ", ".join(required_params)) # Get parameters not directly in the bank: for p in required_params: - if p in bank: continue + if p in bank: + continue if p not in parameter_options: - raise KeyError(f"Parameter {p} not in bank, fits or conversion " - "options, choose from bank parameters or " + \ - ', '.join(parameter_options)) + raise KeyError( + f"Parameter {p} not in bank, fits or conversion " + "options, choose from bank parameters or " + ", ".join(parameter_options) + ) if p in _fit_parameters and args.fits_file is None: parser.error(f"Parameter {p} needs a --fits-file, but none is given") logging.info("Converting %s", p) bank[p] = bconv.get_bank_property(p, bank, np.arange(banklen)) -if args.mpl_style == 'dark_background': - hist_color = 'white' +if args.mpl_style == "dark_background": + hist_color = "white" else: - hist_color = 'black' + hist_color = "black" # All parameters should be in the bank now, check they are the right size: assert all([len(bank[p]) == banklen for p in bank]) @@ -279,7 +295,7 @@ fig, axis_dict = create_multidim_plot( scatter_log_cmap=args.log_colormap, marginal_title=False, marginal_percentiles=[], - fill_color='g', + fill_color="g", zvals=zvals, show_colorbar=cpar is not False, cbar_label=cpar_label, @@ -295,7 +311,7 @@ title_text = f"{os.path.basename(args.bank_file)}" if args.fits_file is not None: title_text += f", {os.path.basename(args.fits_file)}" title_text += f" - {banklen}\u00a0templates" -title_text = '\n'.join(wrap(args.title if args.title is not None else title_text, 60)) +title_text = "\n".join(wrap(args.title if args.title is not None else title_text, 60)) if not args.no_suptitle: fig.suptitle(title_text) for k, v in axis_dict.items(): @@ -343,15 +359,16 @@ logging.info("Plot generated") fig.set_dpi(args.dpi) # save -caption = ("Template bank as a corner plot with " - "scatter points for each waveform.") +caption = "Template bank as a corner plot with scatter points for each waveform." metadata.save_fig_with_metadata( - fig, args.output_plot_file, - cmd=" ".join(sys.argv), - title=title_text, - caption=caption, - fig_kwds={'bbox_inches': 'tight'}) + fig, + args.output_plot_file, + cmd=" ".join(sys.argv), + title=title_text, + caption=caption, + fig_kwds={"bbox_inches": "tight"}, +) # finish logging.info("Done") diff --git a/bin/plotting/pycbc_plot_dq_flag_likelihood b/bin/plotting/pycbc_plot_dq_flag_likelihood index f8d459d02b6..f8b73c12c02 100644 --- a/bin/plotting/pycbc_plot_dq_flag_likelihood +++ b/bin/plotting/pycbc_plot_dq_flag_likelihood @@ -1,13 +1,16 @@ #!/usr/bin/env python -""" Plot the log likelihood percentiles for a DQ bin -""" -import sys +"""Plot the log likelihood percentiles for a DQ bin""" + import argparse +import sys + import numpy -import pycbc -from matplotlib import use as matplotlib_use from matplotlib import pyplot -matplotlib_use('Agg') +from matplotlib import use as matplotlib_use + +import pycbc + +matplotlib_use("Agg") import pycbc.results from pycbc.io.hdf import HFile @@ -17,13 +20,16 @@ pycbc.add_common_pycbc_options(parser) parser.add_argument("--dq-file", required=True) parser.add_argument("--ifo", type=str, required=True) parser.add_argument("--output-file", required=True) -parser.add_argument("--low-latency", action='store_true') +parser.add_argument("--low-latency", action="store_true") title_grp = parser.add_mutually_exclusive_group(required=True) -title_grp.add_argument("--dq-label", type=str, - help="Name of dq flag. Used in plot title," - " mutually exclusive with --title") -title_grp.add_argument("--title", type=str, - help="Title of plot, mutually exclusive with --dq-label") +title_grp.add_argument( + "--dq-label", + type=str, + help="Name of dq flag. Used in plot title, mutually exclusive with --title", +) +title_grp.add_argument( + "--title", type=str, help="Title of plot, mutually exclusive with --dq-label" +) args = parser.parse_args() @@ -31,34 +37,30 @@ pycbc.init_logging(args.verbose) ifo = args.ifo -f = HFile(args.dq_file, 'r') +f = HFile(args.dq_file, "r") ifo_grp = f[ifo] -bin_names = ifo_grp['bins'].keys() +bin_names = ifo_grp["bins"].keys() x = numpy.arange(len(bin_names)) if args.low_latency: dq_states = { - 0: 'Clean', - 1: 'DQ Flag', + 0: "Clean", + 1: "DQ Flag", } x_shift = 0.5 width = 0.33 else: dq_states = { - 0: 'Clean', - 1: 'DQ Flag', - 2: 'Autogating', + 0: "Clean", + 1: "DQ Flag", + 2: "Autogating", } x_shift = 1 width = 0.25 -colors = { - 0: 'green', - 1: 'gold', - 2: 'red' -} +colors = {0: "green", 1: "gold", 2: "red"} fig, ax = pyplot.subplots(figsize=(9, 6)) ax2 = ax.twinx() @@ -66,8 +68,7 @@ ax2 = ax.twinx() ymax = 1 ymin = 0.9 for n, dqstate_name in dq_states.items(): - dq_rates = numpy.array( - [ifo_grp['bins'][b]['dq_rates'][n] for b in bin_names]) + dq_rates = numpy.array([ifo_grp["bins"][b]["dq_rates"][n] for b in bin_names]) dq_rates = numpy.maximum(dq_rates, 1) logrates = numpy.log(dq_rates) @@ -78,27 +79,31 @@ for n, dqstate_name in dq_states.items(): ymax = ymax**1.05 ax.set_ylim(ymin, ymax) -ax.set_yscale('log') -ax.set_ylabel('(Trigger rate during DQ state)/(Mean Rate) [Min 1]') +ax.set_yscale("log") +ax.set_ylabel("(Trigger rate during DQ state)/(Mean Rate) [Min 1]") ax.set_xticks(x) -ax.set_xlabel('Template Bin Number (Longer duration -> Lower number)') +ax.set_xlabel("Template Bin Number (Longer duration -> Lower number)") ax.legend() ax.grid() -ax2.set_ylabel('DQ Log Likelihood Penalty') +ax2.set_ylabel("DQ Log Likelihood Penalty") ax2.set_ylim(numpy.log(ymin), numpy.log(ymax)) # add meta data and save figure if args.title is not None: plot_title = args.title else: - plot_title = f'{ifo}:{args.dq_label} DQ Trigger Rates' + plot_title = f"{ifo}:{args.dq_label} DQ Trigger Rates" ax.set_title(plot_title) -plot_caption = 'The log likelihood correction \ - during during each dq state for each template bin.' +plot_caption = "The log likelihood correction \ + during during each dq state for each template bin." pycbc.results.save_fig_with_metadata( - fig, args.output_file, title=plot_title, - caption=plot_caption, cmd=' '.join(sys.argv)) + fig, + args.output_file, + title=plot_title, + caption=plot_caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_plot_dq_likelihood_vs_time b/bin/plotting/pycbc_plot_dq_likelihood_vs_time index 803bed718ad..8d0ae2b90b2 100644 --- a/bin/plotting/pycbc_plot_dq_likelihood_vs_time +++ b/bin/plotting/pycbc_plot_dq_likelihood_vs_time @@ -1,12 +1,15 @@ #!/usr/bin/env python -"""Plot the DQ log likelihood versus time for a specific background bin -""" -import sys +"""Plot the DQ log likelihood versus time for a specific background bin""" + import argparse +import sys + import numpy -import pycbc from matplotlib import use -use('Agg') + +import pycbc + +use("Agg") from matplotlib import pyplot import pycbc.results @@ -16,7 +19,7 @@ parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) parser.add_argument("--ifo", type=str, required=True) parser.add_argument("--dq-file", required=True) -parser.add_argument('--background-bin', default='all_bin') +parser.add_argument("--background-bin", default="all_bin") parser.add_argument("--output-file", required=True) args = parser.parse_args() @@ -25,20 +28,20 @@ pycbc.init_logging(args.verbose) ifo = args.ifo -f = HFile(args.dq_file, 'r') +f = HFile(args.dq_file, "r") bin_name = args.background_bin -if bin_name not in f[f'{ifo}/dq_vals'].keys(): +if bin_name not in f[f"{ifo}/dq_vals"].keys(): raise ValueError("Background bin name not found in DQ file") else: - yvals = f['%s/dq_vals/%s'%(ifo,bin_name)][:] + yvals = f["%s/dq_vals/%s" % (ifo, bin_name)][:] -xvals = f[f'{ifo}/times'][:] -dq_name = f.attrs['stat'].split('-')[0] +xvals = f[f"{ifo}/times"][:] +dq_name = f.attrs["stat"].split("-")[0] -xmax = max(xvals) + 0.2 * (max(xvals)-min(xvals)) -xmin = min(xvals) - 0.2 * (max(xvals)-min(xvals)) +xmax = max(xvals) + 0.2 * (max(xvals) - min(xvals)) +xmin = min(xvals) - 0.2 * (max(xvals) - min(xvals)) ymax = 1.2 * max(yvals) @@ -48,14 +51,19 @@ color = pycbc.results.ifo_color(ifo) fig = pyplot.figure(0) ax = fig.add_subplot(111) -ax.scatter(xvals[yvals>0],yvals[yvals>0], - label=' '.join([ifo,dq_name,bin_name]), edgecolor='none', - s=1, color=color) -ax.legend(loc='upper left', markerscale=5) +ax.scatter( + xvals[yvals > 0], + yvals[yvals > 0], + label=" ".join([ifo, dq_name, bin_name]), + edgecolor="none", + s=1, + color=color, +) +ax.legend(loc="upper left", markerscale=5) # format plot -ax.set_ylabel('Data quality log likelihood') -ax.set_xlabel('Time (s)') +ax.set_ylabel("Data quality log likelihood") +ax.set_xlabel("Time (s)") ax.set_ylim(ymin=0, ymax=ymax) ax.set_xlim(xmin=xmin, xmax=xmax) @@ -69,17 +77,19 @@ ax2 = ax.twinx() ax2_ymax = numpy.exp(ymax) ax2.set_ylim(1, ax2_ymax) ax2.plot(xmax + 1e9, 100) -new_ticks = range(0, int(numpy.ceil(numpy.log10(ax2_ymax)))) +new_ticks = range(int(numpy.ceil(numpy.log10(ax2_ymax)))) ax2.set_yticks([10**t for t in new_ticks]) -ax2.set_ylabel('Relative Trigger Rate') +ax2.set_ylabel("Relative Trigger Rate") ax2.set_xlim(xmin, xmax) -ax2.set_yscale('log') +ax2.set_yscale("log") # add meta data and save figure fig_kwds = {} -pycbc.results.save_fig_with_metadata(fig, args.output_file, - fig_kwds=fig_kwds, - title = '%s %s log likelihood versus time' % (ifo, dq_name), - caption = 'The log likelihood verus time of a DQ product.', - cmd = ' '.join(sys.argv)) - +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + fig_kwds=fig_kwds, + title="%s %s log likelihood versus time" % (ifo, dq_name), + caption="The log likelihood verus time of a DQ product.", + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_plot_dq_percentiles b/bin/plotting/pycbc_plot_dq_percentiles index 845cfc91fe7..3a331a664ba 100644 --- a/bin/plotting/pycbc_plot_dq_percentiles +++ b/bin/plotting/pycbc_plot_dq_percentiles @@ -1,12 +1,15 @@ #!/usr/bin/env python -""" Plot the log likelihood percentiles for a DQ bin -""" -import sys +"""Plot the log likelihood percentiles for a DQ bin""" + import argparse +import sys + import numpy -import pycbc from matplotlib import use -use('Agg') + +import pycbc + +use("Agg") from matplotlib import pyplot import pycbc.results @@ -14,12 +17,12 @@ from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--ifo", type=str,required=True) +parser.add_argument("--ifo", type=str, required=True) parser.add_argument("--dq-file", required=True) -parser.add_argument('--background-bin', default='all_bin') -parser.add_argument("--x-max",default=100) -parser.add_argument("--x-min",default=0) -parser.add_argument('--log-x', action="store_true") +parser.add_argument("--background-bin", default="all_bin") +parser.add_argument("--x-max", default=100) +parser.add_argument("--x-min", default=0) +parser.add_argument("--log-x", action="store_true") parser.add_argument("--output-file", required=True) args = parser.parse_args() @@ -27,23 +30,23 @@ pycbc.init_logging(args.verbose) ifo = args.ifo -f = HFile(args.dq_file, 'r') +f = HFile(args.dq_file, "r") bin_name = args.background_bin -if bin_name not in f['%s/dq_percentiles'%ifo].keys(): +if bin_name not in f["%s/dq_percentiles" % ifo].keys(): raise ValueError("Background bin name not found in DQ file") else: - yvals = f['%s/dq_percentiles/%s'%(ifo,bin_name)][:] + yvals = f["%s/dq_percentiles/%s" % (ifo, bin_name)][:] -xvals = numpy.linspace(0,100,len(yvals)+1) +xvals = numpy.linspace(0, 100, len(yvals) + 1) centers = (xvals[1:] + xvals[:-1]) / 2 -dq_name = f.attrs['stat'].split('-')[0] +dq_name = f.attrs["stat"].split("-")[0] xmax = float(args.x_max) xmin = float(args.x_min) -if args.log_x and xmin==0: +if args.log_x and xmin == 0: raise ValueError("Cannot set the x axis to log scale when xmin=0") ymax = 1.2 * max(yvals) @@ -54,19 +57,23 @@ color = pycbc.results.ifo_color(ifo) fig = pyplot.figure(0) ax = fig.add_subplot(111) -counts_, bins_, _ = ax.hist(centers, bins=len(yvals), - weights=yvals, color=color, - label=' '.join([ifo,dq_name,bin_name]), - range=(min(xvals), max(xvals))) -ax.legend(loc='upper left', markerscale=5) +counts_, bins_, _ = ax.hist( + centers, + bins=len(yvals), + weights=yvals, + color=color, + label=" ".join([ifo, dq_name, bin_name]), + range=(min(xvals), max(xvals)), +) +ax.legend(loc="upper left", markerscale=5) # format plot -ax.set_ylabel('Data quality log likelihood') -ax.set_xlabel('Percentile') -ax.set_ylim(ymin=0,ymax=ymax) -ax.set_xlim(xmin=xmin,xmax=xmax) +ax.set_ylabel("Data quality log likelihood") +ax.set_xlabel("Percentile") +ax.set_ylim(ymin=0, ymax=ymax) +ax.set_xlim(xmin=xmin, xmax=xmax) if args.log_x: - ax.set_xscale('log') + ax.set_xscale("log") # add a grid to the plot ax.grid() @@ -74,21 +81,23 @@ ax.grid() yticks = ax.get_yticks() # add second axis -ax2=ax.twinx() +ax2 = ax.twinx() ax2_ymax = numpy.exp(ymax) -ax2.set_ylim(1,ax2_ymax) -ax2.plot(xmax+1e9,100) -new_ticks = range(0,int(numpy.ceil(numpy.log10(ax2_ymax)))) +ax2.set_ylim(1, ax2_ymax) +ax2.plot(xmax + 1e9, 100) +new_ticks = range(int(numpy.ceil(numpy.log10(ax2_ymax)))) ax2.set_yticks([10**t for t in new_ticks]) -ax2.set_ylabel('Relative Trigger Rate') +ax2.set_ylabel("Relative Trigger Rate") ax2.set_xlim(xmin, xmax) -ax2.set_yscale('log') +ax2.set_yscale("log") # add meta data and save figure -plot_title = '%s: %s log likelihood versus percentile' % (ifo, dq_name) -plot_caption = 'The log likelihood verus percentile of a DQ product.' -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title = plot_title, - caption = plot_caption, - cmd = ' '.join(sys.argv)) - +plot_title = "%s: %s log likelihood versus percentile" % (ifo, dq_name) +plot_caption = "The log likelihood verus percentile of a DQ product." +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title=plot_title, + caption=plot_caption, + cmd=" ".join(sys.argv), +) diff --git a/bin/plotting/pycbc_plot_gate_triggers b/bin/plotting/pycbc_plot_gate_triggers index 261229ea8f1..5c84f6cd335 100755 --- a/bin/plotting/pycbc_plot_gate_triggers +++ b/bin/plotting/pycbc_plot_gate_triggers @@ -16,66 +16,86 @@ Plot histograms of triggers around gated times. """ -import numpy -import logging import argparse +import logging + +import numpy from matplotlib import use -use('Agg') + +use("Agg") import matplotlib.pyplot as plt -from pycbc.io.hdf import HFile from pycbc import add_common_pycbc_options, init_logging from pycbc.events import ranking +from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--single-trigger-files', nargs='+', - help='HDF format single detector merged trigger file(s) path. ' - 'All triggers are combined in the histogram') -parser.add_argument('--ifo-tag', default='', - help='String to label plots, eg H1, H1L1') -parser.add_argument('--gating-type', choices=['auto', 'detchar'], - help='Type of gate to analyze.') -parser.add_argument('--window', default=5, type=float, - help='Time in seconds around the gate to plot (default 5s).') -parser.add_argument('--snr-cut-type', choices=['snr', 'newsnr'], - help='Type of snr threshold to apply.') -parser.add_argument('--snr-cut-vals', nargs='+', type=float, - help='Threshold snr/newsnr value(s) for histogram(s).') -parser.add_argument('--template-duration', nargs='+', - help='Triggers with a template duration greater than ' - '(GT) or less than (LT) some time value in seconds. ' - 'Usage: --template-duration GT/LT time value') -parser.add_argument('--bin-duration', type=float, default=0.25, - help='Time in seconds represented by each bin in the histogram ' - '(default 0.25s).') -parser.add_argument('--log-y', action='store_true', - help='Set y-axis scale logarithmic.') -parser.add_argument('--output-file', - help='Plot file name (optional).') +parser.add_argument( + "--single-trigger-files", + nargs="+", + help="HDF format single detector merged trigger file(s) path. " + "All triggers are combined in the histogram", +) +parser.add_argument("--ifo-tag", default="", help="String to label plots, eg H1, H1L1") +parser.add_argument( + "--gating-type", choices=["auto", "detchar"], help="Type of gate to analyze." +) +parser.add_argument( + "--window", + default=5, + type=float, + help="Time in seconds around the gate to plot (default 5s).", +) +parser.add_argument( + "--snr-cut-type", choices=["snr", "newsnr"], help="Type of snr threshold to apply." +) +parser.add_argument( + "--snr-cut-vals", + nargs="+", + type=float, + help="Threshold snr/newsnr value(s) for histogram(s).", +) +parser.add_argument( + "--template-duration", + nargs="+", + help="Triggers with a template duration greater than " + "(GT) or less than (LT) some time value in seconds. " + "Usage: --template-duration GT/LT time value", +) +parser.add_argument( + "--bin-duration", + type=float, + default=0.25, + help="Time in seconds represented by each bin in the histogram (default 0.25s).", +) +parser.add_argument( + "--log-y", action="store_true", help="Set y-axis scale logarithmic." +) +parser.add_argument("--output-file", help="Plot file name (optional).") args = parser.parse_args() init_logging(args.verbose) if args.gating_type: - if args.gating_type == 'auto': - gating_type = 'auto' - elif args.gating_type == 'detchar': - gating_type = 'file' + if args.gating_type == "auto": + gating_type = "auto" + elif args.gating_type == "detchar": + gating_type = "file" else: - raise ValueError('A gating type argument must be provided.') + raise ValueError("A gating type argument must be provided.") if args.snr_cut_type: if args.snr_cut_vals: snr_cut_vals = sorted(args.snr_cut_vals) else: - raise ValueError('At least one SNR cut value must be provided!') + raise ValueError("At least one SNR cut value must be provided!") else: snr_cut_vals = numpy.array([0]) if args.template_duration: args.duration_ltgt = args.template_duration[0] - assert args.duration_ltgt in ['LT', 'GT'] + assert args.duration_ltgt in ["LT", "GT"] args.duration_cut = float(args.template_duration[1]) # Binning along trigger time @@ -85,44 +105,45 @@ ifo_trigger_times = [[] for thr in snr_cut_vals] n_gates = 0 for single_file in args.single_trigger_files: - logging.info("Reading trigger file %s...", single_file.split('/')[-1]) + logging.info("Reading trigger file %s...", single_file.split("/")[-1]) file_ = HFile(single_file, "r") ifo = str(list(file_.keys())[0]) data = file_[ifo] - if gating_type in data['gating'].keys(): + if gating_type in data["gating"].keys(): logging.info("Extracting gated times...") - gate_time = numpy.unique(data[f'gating/{gating_type}/time'][:]) + gate_time = numpy.unique(data[f"gating/{gating_type}/time"][:]) n_gates += len(gate_time) logging.info("%i %s gates found", len(gate_time), args.gating_type) else: - logging.info('Trigger file does not contain %s gates. ' - 'Skipping file...', args.gating_type) + logging.info( + "Trigger file does not contain %s gates. Skipping file...", args.gating_type + ) continue logging.info("Extracting trigger times...") - times = data['end_time'][:] + times = data["end_time"][:] if args.template_duration: logging.info("Cutting on template duration...") - template_duration = data['template_duration'][:] - if args.duration_ltgt == 'LT': + template_duration = data["template_duration"][:] + if args.duration_ltgt == "LT": times = times[template_duration <= args.duration_cut] else: times = times[template_duration >= args.duration_cut] if args.snr_cut_type: logging.info("Cutting on %s...", args.snr_cut_type) - if args.snr_cut_type == 'snr': - snr_val = data['snr'][:] - if args.snr_cut_type == 'newsnr': - rchisq = data['chisq'][:] / (2 * data['chisq_dof'][:] - 2) + if args.snr_cut_type == "snr": + snr_val = data["snr"][:] + if args.snr_cut_type == "newsnr": + rchisq = data["chisq"][:] / (2 * data["chisq_dof"][:] - 2) if len(rchisq) > 0: - snr_val = ranking.newsnr(data['snr'][:], rchisq) + snr_val = ranking.newsnr(data["snr"][:], rchisq) else: snr_val = numpy.array([]) del rchisq if args.template_duration: - if args.duration_ltgt == 'LT': + if args.duration_ltgt == "LT": snr_val = snr_val[template_duration <= args.duration_cut] else: snr_val = snr_val[template_duration >= args.duration_cut] @@ -144,54 +165,67 @@ for single_file in args.single_trigger_files: if args.output_file: out_name = args.output_file else: - out_name = '%s_%s' % (str(args.gating_type).upper(), args.ifo_tag) + out_name = "%s_%s" % (str(args.gating_type).upper(), args.ifo_tag) if args.template_duration: templ_cut_name = ( - str("{:.1f}".format(args.duration_cut)).replace('.', '-') - if '.' in args.template_duration[1] - else args.template_duration[1] + str(f"{args.duration_cut:.1f}").replace(".", "-") + if "." in args.template_duration[1] + else args.template_duration[1] + ) + out_name = "_".join( + [out_name, "TD-%s-%s" % (args.duration_ltgt, templ_cut_name)] ) - out_name = '_'.join([out_name, 'TD-%s-%s' % (args.duration_ltgt, templ_cut_name)]) if args.snr_cut_type: snr_cut_name = (args.snr_cut_type).upper() snr_cuts = [ - str("{:.1f}".format(cut_val)).replace('.', '-') if '.' in str(cut_val) - else str(cut_val) for cut_val in snr_cut_vals + str(f"{cut_val:.1f}").replace(".", "-") + if "." in str(cut_val) + else str(cut_val) + for cut_val in snr_cut_vals ] - snr_cut_value = '_'.join(snr_cuts) - out_name = '_'.join([out_name, '%s-%s' % (snr_cut_name, snr_cut_value)]) + snr_cut_value = "_".join(snr_cuts) + out_name = "_".join([out_name, "%s-%s" % (snr_cut_name, snr_cut_value)]) - window_name = (str("{:.1f}".format(args.window)).replace('.', '-') if '.' in - str(args.window) else str(args.window)) - out_name = '_'.join([out_name, 'WINDOW-%s' % window_name]) + window_name = ( + str(f"{args.window:.1f}").replace(".", "-") + if "." in str(args.window) + else str(args.window) + ) + out_name = "_".join([out_name, "WINDOW-%s" % window_name]) if args.log_y: - out_name = '_'.join([out_name, 'LOG']) + out_name = "_".join([out_name, "LOG"]) - out_name = out_name + '.png' # Filenames via command line may have any extension + out_name = out_name + ".png" # Filenames via command line may have any extension if n_gates > 0: for i, thr in enumerate(snr_cut_vals): plt.hist( ifo_trigger_times[i], bins=numpy.linspace(-args.window, args.window, nbins), - label=['%s > %.1f' % (args.snr_cut_type, thr) if args.snr_cut_type else ''] + label=["%s > %.1f" % (args.snr_cut_type, thr) if args.snr_cut_type else ""], ) if args.log_y: - plt.yscale('log') + plt.yscale("log") plt.grid(True) - plt.xlabel('Time relative to gate centre (s)') - plt.ylabel('Number of triggers') + plt.xlabel("Time relative to gate centre (s)") + plt.ylabel("Number of triggers") plt.title( # Complicated recipe - '%s - %i %s GATES %s %s' % - (args.ifo_tag, n_gates, str(args.gating_type).upper(), - 'IN %i CHUNKS' % len(args.single_trigger_files) if \ - len(args.single_trigger_files) > 1 else '', - '\n Triggers with template duration %s %s s' % - ('<' if args.duration_ltgt == 'LT' else '>', str(args.duration_cut)) - if args.template_duration else '') + "%s - %i %s GATES %s %s" + % ( + args.ifo_tag, + n_gates, + str(args.gating_type).upper(), + "IN %i CHUNKS" % len(args.single_trigger_files) + if len(args.single_trigger_files) > 1 + else "", + "\n Triggers with template duration %s %s s" + % ("<" if args.duration_ltgt == "LT" else ">", str(args.duration_cut)) + if args.template_duration + else "", + ) ) if args.snr_cut_type: plt.legend() @@ -200,10 +234,15 @@ else: # Nothing to plot, make an empty figure fig = plt.figure() ax = fig.add_subplot(111) output_message = "Sorry, no gates to plot!" - ax.text(0.5, 0.5, output_message, horizontalalignment='center', - verticalalignment='center') + ax.text( + 0.5, + 0.5, + output_message, + horizontalalignment="center", + verticalalignment="center", + ) -logging.info('Saving histogram...') +logging.info("Saving histogram...") plt.savefig(out_name) plt.close() diff --git a/bin/plotting/pycbc_plot_gating b/bin/plotting/pycbc_plot_gating index a4ab6756e57..a48f7b6bfa6 100644 --- a/bin/plotting/pycbc_plot_gating +++ b/bin/plotting/pycbc_plot_gating @@ -1,29 +1,34 @@ #!/usr/bin/env python -"Plot gated time segments from inspiral HDF5 files." +"""Plot gated time segments from inspiral HDF5 files.""" import argparse import logging -import numpy as np + import matplotlib -matplotlib.use('agg') -from matplotlib import pyplot as plt -from matplotlib.patches import Rectangle +import numpy as np + +matplotlib.use("agg") import mpld3 import mpld3.plugins +from matplotlib import pyplot as plt +from matplotlib.patches import Rectangle import pycbc -from pycbc.results.color import ifo_color from pycbc.io.hdf import HFile - +from pycbc.results.color import ifo_color parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input-file', nargs='+', required=True, - help='Single-detector inspiral HDF5 files to take gating ' - 'data from.') -parser.add_argument('--output-file', required=True, - help='Destination file for the plot.') +parser.add_argument( + "--input-file", + nargs="+", + required=True, + help="Single-detector inspiral HDF5 files to take gating data from.", +) +parser.add_argument( + "--output-file", required=True, help="Destination file for the plot." +) args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -31,18 +36,18 @@ pycbc.init_logging(args.verbose) gate_data = {} have_gates = False for fn in args.input_file: - logging.info('Reading gates from %s', fn) - f = HFile(fn, 'r') + logging.info("Reading gates from %s", fn) + f = HFile(fn, "r") ifo = tuple(f.keys())[0] - for gate_type in ['file', 'auto']: + for gate_type in ["file", "auto"]: try: - gate_time = f[ifo + '/gating/' + gate_type + '/time'][:] - gate_width = f[ifo + '/gating/' + gate_type + '/width'][:] - gate_pad = f[ifo + '/gating/' + gate_type + '/pad'][:] + gate_time = f[ifo + "/gating/" + gate_type + "/time"][:] + gate_width = f[ifo + "/gating/" + gate_type + "/width"][:] + gate_pad = f[ifo + "/gating/" + gate_type + "/pad"][:] except KeyError: continue key = (ifo, gate_type) - if not key in gate_data: + if key not in gate_data: gate_data[key] = [] gate_data[key] += [g for g in zip(gate_time, gate_width, gate_pad)] have_gates = have_gates or (len(gate_data[key]) > 0) @@ -57,42 +62,50 @@ if have_gates: total_time = {} total_time_pad = {} for i, (ifo, gate_type) in enumerate(sorted(gate_data.keys())): - logging.info('Plotting %s gates for %s', gate_type, ifo) + logging.info("Plotting %s gates for %s", gate_type, ifo) gd = np.array(sorted(frozenset(gate_data[(ifo, gate_type)]))) if len(gd) > 0: - for s, e in zip(gd[:,0] - gd[:,1] - gd[:,2], - gd[:,0] + gd[:,1] + gd[:,2]): - patch = Rectangle((s, i), e - s, 0.7, color=ifo_color(ifo), - alpha=0.5) + for s, e in zip( + gd[:, 0] - gd[:, 1] - gd[:, 2], gd[:, 0] + gd[:, 1] + gd[:, 2] + ): + patch = Rectangle((s, i), e - s, 0.7, color=ifo_color(ifo), alpha=0.5) ax.add_patch(patch) if t_min is None or s < t_min: t_min = s if t_max is None or e > t_max: t_max = e - total_time[(ifo, gate_type)] = sum(gd[:,1]) * 2 - total_time_pad[(ifo, gate_type)] = sum(gd[:,1] + gd[:,2]) * 2 + total_time[(ifo, gate_type)] = sum(gd[:, 1]) * 2 + total_time_pad[(ifo, gate_type)] = sum(gd[:, 1] + gd[:, 2]) * 2 else: - total_time[(ifo, gate_type)] = 0. - total_time_pad[(ifo, gate_type)] = 0. + total_time[(ifo, gate_type)] = 0.0 + total_time_pad[(ifo, gate_type)] = 0.0 for i, (ifo, gate_type) in enumerate(sorted(gate_data.keys())): - label = '%s %s gates: %.1f s (zeroes), %.1f s (zeroes + pad)' \ - % (ifo, gate_type, total_time[(ifo, gate_type)], - total_time_pad[(ifo, gate_type)]) + label = "%s %s gates: %.1f s (zeroes), %.1f s (zeroes + pad)" % ( + ifo, + gate_type, + total_time[(ifo, gate_type)], + total_time_pad[(ifo, gate_type)], + ) ax.text(t_min, i + 0.75, label) ax.set_xlim(t_min, t_max) ax.set_ylim(0, len(gate_data.keys())) - ax.set_xlabel('GPS Time (s)') + ax.set_xlabel("GPS Time (s)") else: - ax.text(0.5, 0.5, 'No gating data to plot', horizontalalignment='center', - verticalalignment='center') + ax.text( + 0.5, + 0.5, + "No gating data to plot", + horizontalalignment="center", + verticalalignment="center", + ) ax.set_xticks([]) ax.set_xlim(0, 1) ax.set_ylim(0, 1) -mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fontsize=14, fmt='10.1f')) +mpld3.plugins.connect(fig, mpld3.plugins.MousePosition(fontsize=14, fmt="10.1f")) mpld3.plugins.connect(fig, mpld3.plugins.BoxZoom()) mpld3.plugins.connect(fig, mpld3.plugins.Zoom()) mpld3.plugins.connect(fig, mpld3.plugins.Reset()) -mpld3.save_html(fig, open(args.output_file, 'w')) +mpld3.save_html(fig, open(args.output_file, "w")) diff --git a/bin/plotting/pycbc_plot_hist b/bin/plotting/pycbc_plot_hist index 3b5ce12f917..c8f43d85212 100644 --- a/bin/plotting/pycbc_plot_hist +++ b/bin/plotting/pycbc_plot_hist @@ -1,43 +1,44 @@ #!/bin/env python -""" Make histograms of single detector triggers -""" +"""Make histograms of single detector triggers""" -import numpy import argparse import sys from itertools import cycle + +import numpy from matplotlib import use -use('Agg') + +use("Agg") from matplotlib import pyplot import pycbc -import pycbc.results import pycbc.io +import pycbc.results from pycbc.events import background_bin_from_string, ranking parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-file', required=True, - help="Combined single detector hdf trigger file") -parser.add_argument('--veto-file', - help="segment xml file, indicates which triggers to ignore") -parser.add_argument('--segment-name', - help="name of segment list in the veto file") -parser.add_argument('--x-var', - help="name of value to histogram") -parser.add_argument('--output-file') -parser.add_argument('--bank-file', default="", - help='template bank hdf file') -parser.add_argument('--background-bins', nargs='+', - help='list of background bin format strings') -parser.add_argument('--bins', default=100, type=int, - help="number of bins in histogram") -parser.add_argument('--x-max', type=float) -parser.add_argument('--x-min', type=float, default=6, - help="Minimum x-value. Default 6") -parser.add_argument('--special-time', type=float, - help="plot triggers within +-1s of a given time in a " - "different color (black)") +parser.add_argument( + "--trigger-file", required=True, help="Combined single detector hdf trigger file" +) +parser.add_argument( + "--veto-file", help="segment xml file, indicates which triggers to ignore" +) +parser.add_argument("--segment-name", help="name of segment list in the veto file") +parser.add_argument("--x-var", help="name of value to histogram") +parser.add_argument("--output-file") +parser.add_argument("--bank-file", default="", help="template bank hdf file") +parser.add_argument( + "--background-bins", nargs="+", help="list of background bin format strings" +) +parser.add_argument("--bins", default=100, type=int, help="number of bins in histogram") +parser.add_argument("--x-max", type=float) +parser.add_argument("--x-min", type=float, default=6, help="Minimum x-value. Default 6") +parser.add_argument( + "--special-time", + type=float, + help="plot triggers within +-1s of a given time in a different color (black)", +) args = parser.parse_args() # sanity check command line options @@ -48,7 +49,7 @@ if args.special_time and args.background_bins: pycbc.init_logging(args.verbose) # read trigger file to determine IFO -f = pycbc.io.HFile(args.trigger_file, 'r') +f = pycbc.io.HFile(args.trigger_file, "r") ifo = tuple(f.keys())[0] # Apply presort if we can; if we can't, load everything but it may be @@ -69,75 +70,91 @@ trigs = pycbc.io.SingleDetTriggers( # get x values and find the maximum x value val = getattr(trigs, args.x_var) -x_max = args.x_max if args.x_max else val.max() * 1.1 +x_max = args.x_max or val.max() * 1.1 # create a figure to add a histogram fig = pyplot.figure(0) # command line says to plot triggers around a certain time if args.special_time: - time = getattr(trigs, 'end_time') - val_special = val[abs(time-args.special_time) < 1] - val_boring = val[abs(time-args.special_time) >= 1] + time = trigs.end_time + val_special = val[abs(time - args.special_time) < 1] + val_boring = val[abs(time - args.special_time) >= 1] binvals = numpy.linspace(args.x_min, x_max, args.bins, endpoint=True) - pyplot.hist([val_boring, val_special], bins=binvals, histtype='stepfilled', - stacked=True, color=[pycbc.results.ifo_color(ifo), 'k'], - label=['Other times', 'Triggers near %i' % (int(args.special_time))]) - pyplot.legend(loc='upper right') + pyplot.hist( + [val_boring, val_special], + bins=binvals, + histtype="stepfilled", + stacked=True, + color=[pycbc.results.ifo_color(ifo), "k"], + label=["Other times", "Triggers near %i" % (int(args.special_time))], + ) + pyplot.legend(loc="upper right") # command line says to plot triggers in each background bin elif args.background_bins: - # get a dict where key is each bin's name and value is a list of indexes # for the triggers in that bin - bank_data = {'mass1' : trigs.mass1, - 'mass2' : trigs.mass2, - 'spin1z' : trigs.spin1z, - 'spin2z' : trigs.spin2z, + bank_data = { + "mass1": trigs.mass1, + "mass2": trigs.mass2, + "spin1z": trigs.spin1z, + "spin2z": trigs.spin2z, } locs_dict = background_bin_from_string(args.background_bins, bank_data) - + # get a list of bin names and a corresponding list for x values loc_bin_keys = [key for key in locs_dict.keys()] loc_bin_vals = [val[locs_dict[key]] for key in loc_bin_keys] # assign a color for each bin - color_cycle = cycle(['red', 'green', 'blue', 'black', 'magenta', 'cyan']) + color_cycle = cycle(["red", "green", "blue", "black", "magenta", "cyan"]) loc_bin_colors = [next(color_cycle) for key in loc_bin_keys] # get number of overflows for each background bin - loc_bin_overflows = [len(vals[vals>=x_max]) for vals in loc_bin_vals] + loc_bin_overflows = [len(vals[vals >= x_max]) for vals in loc_bin_vals] num_bins = len(loc_bin_overflows) # remove overflow triggers from plotting - loc_bin_vals = [vals[vals low.min(): y_min = low.min() @@ -125,18 +150,25 @@ for psd_file in args.psd_files: ax.set_xlim(flow, samples[-1]) if args.psd_model or args.psd_file or args.asd_file: - reference_psd = pycbc.psd.from_cli(args, 2048, 1., 10., None) + reference_psd = pycbc.psd.from_cli(args, 2048, 1.0, 10.0, None) if reference_psd is not None: - ax.loglog(reference_psd.sample_frequencies, reference_psd ** 0.5, - '-k', lw=0.3, label='Reference') + ax.loglog( + reference_psd.sample_frequencies, + reference_psd**0.5, + "-k", + lw=0.3, + label="Reference", + ) ax.set_ylim(y_min * 0.5, y_min * 100) -ax.legend(loc='upper right', fontsize='small') -pycbc.results.save_fig_with_metadata(fig, args.output_file, +ax.legend(loc="upper right", fontsize="small") +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, title="Noise Amplitude Spectral Density", - caption="Median amplitude spectral density plotted with a shaded region " - "between the 5th and 95th perentiles. ", - cmd=' '.join(sys.argv), - fig_kwds={'dpi': 200, - 'bbox_inches': 'tight'}) + caption="Median amplitude spectral density plotted with a shaded region " + "between the 5th and 95th perentiles. ", + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200, "bbox_inches": "tight"}, +) diff --git a/bin/plotting/pycbc_plot_psd_timefreq b/bin/plotting/pycbc_plot_psd_timefreq index f7dc23bebac..8c4d8813986 100644 --- a/bin/plotting/pycbc_plot_psd_timefreq +++ b/bin/plotting/pycbc_plot_psd_timefreq @@ -21,62 +21,74 @@ Plot the variation of the amplitude spectral density (ASD) over time using the PSD hdf file generated by the pipeline """ -import logging import argparse -import numpy +import logging import sys + import matplotlib -matplotlib.use('agg') +import numpy + +matplotlib.use("agg") from matplotlib import pyplot as plt from matplotlib.colors import LogNorm -from pycbc.io.hdf import HFile import pycbc import pycbc.results +from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--psd-file', required=True, - help='HDF file containing the PSDs.') -parser.add_argument('--output-file', required=True, - help='Name for the output plot.') -parser.add_argument('--normalize', action='store_true', - help='Select this option if you want to plot the amplitude ' - 'relative to median.') -parser.add_argument('--bins', type=int, default=1000, - help='Number of (logarithmic-spaced) frequency bins to plot' - ' (default %(default)d).') -parser.add_argument('--reduction-method', choices=['min','max','mean','median'], - default='mean', - help='Method of reducing the ASD data into fewer' - ' frequency bins to plot (default %(default)s).') +parser.add_argument("--psd-file", required=True, help="HDF file containing the PSDs.") +parser.add_argument("--output-file", required=True, help="Name for the output plot.") +parser.add_argument( + "--normalize", + action="store_true", + help="Select this option if you want to plot the amplitude relative to median.", +) +parser.add_argument( + "--bins", + type=int, + default=1000, + help="Number of (logarithmic-spaced) frequency bins to plot (default %(default)d).", +) +parser.add_argument( + "--reduction-method", + choices=["min", "max", "mean", "median"], + default="mean", + help="Method of reducing the ASD data into fewer" + " frequency bins to plot (default %(default)s).", +) opts = parser.parse_args() pycbc.init_logging(opts.verbose) -fig=plt.figure() -ax=fig.gca() +fig = plt.figure() +ax = fig.gca() -logging.info('Reading %s', opts.psd_file) +logging.info("Reading %s", opts.psd_file) -f = HFile(opts.psd_file, 'r') +f = HFile(opts.psd_file, "r") ifo = tuple(f.keys())[0] -df = f[ifo + '/psds/0'].attrs['delta_f'] -keys = f[ifo + '/psds'].keys() -psds = [f[ifo + '/psds/' + str(key)][:] for key in range(len(f[ifo + '/psds']))] +df = f[ifo + "/psds/0"].attrs["delta_f"] +keys = f[ifo + "/psds"].keys() +psds = [f[ifo + "/psds/" + str(key)][:] for key in range(len(f[ifo + "/psds"]))] -flow = f.attrs['low_frequency_cutoff'] +flow = f.attrs["low_frequency_cutoff"] kmin = int(flow / df) fac = 1.0 / pycbc.DYN_RANGE_FAC -freqs = numpy.arange( 0 , len(psds[0]) )[kmin:] * df -start, end = f[ifo + '/start_time'][:], f[ifo + '/end_time'][:] - -logging.info('Starting the plot PSD over time') - -reduction_function = { 'min' : numpy.min , 'max' : numpy.max , 'mean' : numpy.mean , - 'median' : numpy.median } -asd_median = numpy.median( psds, axis = 0 )[kmin:] ** 0.5 * fac +freqs = numpy.arange(0, len(psds[0]))[kmin:] * df +start, end = f[ifo + "/start_time"][:], f[ifo + "/end_time"][:] + +logging.info("Starting the plot PSD over time") + +reduction_function = { + "min": numpy.min, + "max": numpy.max, + "mean": numpy.mean, + "median": numpy.median, +} +asd_median = numpy.median(psds, axis=0)[kmin:] ** 0.5 * fac # Select limits for the colorbar if opts.normalize: asdmin = 0.5 @@ -86,8 +98,8 @@ else: asdmax = 1e-20 for psd_segment in range(len(psds)): - logging.info('Plotting PSD segment number %d', psd_segment) - time = numpy.array( [ start[psd_segment] , end[psd_segment] ] ) + logging.info("Plotting PSD segment number %d", psd_segment) + time = numpy.array([start[psd_segment], end[psd_segment]]) if opts.normalize: asd = psds[psd_segment][kmin:] ** 0.5 * fac / asd_median else: @@ -95,53 +107,62 @@ for psd_segment in range(len(psds)): # Get the frequency bins equally spaced in a logarithmic scale and the # indices needed for the ASDs - freqsegs = numpy.logspace( numpy.log10(flow) , numpy.log10(freqs.max()) , - num = opts.bins ) - indices = numpy.searchsorted( freqs , freqsegs ) + freqsegs = numpy.logspace( + numpy.log10(flow), numpy.log10(freqs.max()), num=opts.bins + ) + indices = numpy.searchsorted(freqs, freqsegs) # Take the ASD of each frequency bin using the method given asdless = [] - for index in range( 1 , len(indices) ): - asdless.append( reduction_function[ opts.reduction_method ] - ( asd[indices[index-1] : indices[index]] )) + for index in range(1, len(indices)): + asdless.append( + reduction_function[opts.reduction_method]( + asd[indices[index - 1] : indices[index]] + ) + ) - Y , X = numpy.meshgrid( freqsegs[:-1] , time ) - im = ax.pcolormesh( X , Y , numpy.array([asdless]) , norm = LogNorm( vmin = asdmin , vmax = asdmax ) ) + Y, X = numpy.meshgrid(freqsegs[:-1], time) + im = ax.pcolormesh( + X, Y, numpy.array([asdless]), norm=LogNorm(vmin=asdmin, vmax=asdmax) + ) -logging.info('Saving results') +logging.info("Saving results") -cb = fig.colorbar( im ) +cb = fig.colorbar(im) if opts.normalize: - cbLabel = 'Amplitude relative to median' - cb.set_ticks( [0.5, 1, 2] ) - cb.ax.set_yticklabels( ['0.5', '1', '2'] ) + cbLabel = "Amplitude relative to median" + cb.set_ticks([0.5, 1, 2]) + cb.ax.set_yticklabels(["0.5", "1", "2"]) else: - cbLabel = 'Amplitude Spectral Density (Strain / $\\sqrt{\\rm Hz}$)' -cb.set_label( cbLabel, fontsize=20 ) -ax.set_yscale('log') -ax.set_xlim( start.min() - 5000 , end.max() + 5000 ) -ax.set_ylim( numpy.min(freqs) , numpy.max(freqs) ) + cbLabel = "Amplitude Spectral Density (Strain / $\\sqrt{\\rm Hz}$)" +cb.set_label(cbLabel, fontsize=20) +ax.set_yscale("log") +ax.set_xlim(start.min() - 5000, end.max() + 5000) +ax.set_ylim(numpy.min(freqs), numpy.max(freqs)) ticks = [] -for days in range( (end.max() - start.min()) / 86400 + 1 ): - ticks.append( start.min() + 86400 * days ) -ticklabels = numpy.arange( (end.max() - start.min()) / 86400 + 1) -ax.set_xticks( ticks ) -ax.set_xticklabels( ticklabels ) - -ax.set_xlabel('Days since GPS %d' % start.min(), fontsize=18 ) -ax.set_ylabel('Frequency (Hz)', fontsize=18 ) +for days in range((end.max() - start.min()) / 86400 + 1): + ticks.append(start.min() + 86400 * days) +ticklabels = numpy.arange((end.max() - start.min()) / 86400 + 1) +ax.set_xticks(ticks) +ax.set_xticklabels(ticklabels) + +ax.set_xlabel("Days since GPS %d" % start.min(), fontsize=18) +ax.set_ylabel("Frequency (Hz)", fontsize=18) fig.set_size_inches(18.5, 10.5) plt.tight_layout() -title = ('Evolution of the noise spectral density over time in %s' % ifo) -caption = ('Variation of the amplitude spectral density over time. The original' - 'frequency bins are reduced to %d logarithmic frequency bins by taking' - 'the %s. Each time segment contains %d seconds.') % \ - (opts.bins , opts.reduction_method , end[0] - start[0] ) -pycbc.results.save_fig_with_metadata(fig, opts.output_file, - title = title, - caption = caption, - cmd=' '.join(sys.argv), - fig_kwds={'dpi': 200}) - +title = "Evolution of the noise spectral density over time in %s" % ifo +caption = ( + "Variation of the amplitude spectral density over time. The original" + "frequency bins are reduced to %d logarithmic frequency bins by taking" + "the %s. Each time segment contains %d seconds." +) % (opts.bins, opts.reduction_method, end[0] - start[0]) +pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + title=title, + caption=caption, + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200}, +) diff --git a/bin/plotting/pycbc_plot_qscan b/bin/plotting/pycbc_plot_qscan index 105bca313a9..066d6fa6811 100644 --- a/bin/plotting/pycbc_plot_qscan +++ b/bin/plotting/pycbc_plot_qscan @@ -23,89 +23,152 @@ See https://iopscience.iop.org/article/10.1088/0264-9381/21/20/024 for information on the Q-transform and its parameters. """ -import sys import argparse +import sys import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") from matplotlib import pyplot as plt from matplotlib.colors import LogNorm -import pycbc.strain import pycbc.results +import pycbc.strain + # https://stackoverflow.com/questions/9978880/python-argument-parser-list-of-list-or-tuple-of-tuples def t_window(s): try: - start, end = map(float, s.split(',')) + start, end = map(float, s.split(",")) return [start, end] except: raise argparse.ArgumentTypeError("Input must be start,end start,end") + parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--output-file', required=True, help='Output plot') -parser.add_argument('--center-time', type=float, - help='Center plot on the given GPS time. If omitted, use ' - 'center of interval set via --gps-start-time and ' - '--gps-end-time') -parser.add_argument('--time-windows', required=True, type=t_window, - nargs='+', metavar='START,END', - help='Use these set of times for time windows. Each ' - 'window produces a different panel. Should ' - 'be provided as start1,end1 start2,end2 ... ' - 'where start/end times are relative to --center-time ' - 'and both positive') -parser.add_argument("--low-frequency-cutoff", type=float, - help="The low frequency cutoff to use for fake strain generation") - -parser.add_argument('--qtransform-delta-t', default=0.001, type=float, - help='The time resolution to interpolate to (optional)') -parser.add_argument('--qtransform-delta-f', type=float, - help='Frequency resolution to interpolate to (optional)') -parser.add_argument('--qtransform-logfsteps', type=int, default=200, - help='Do a log interpolation (incompatible with ' - '--qtransform-delta-f option) and set the number ' - 'of steps to take') -parser.add_argument('--qtransform-frange-lower', type=float, - help='Lower frequency (in Hz) at which to compute ' - 'qtransform. Optional, default=20 Hz') -parser.add_argument('--qtransform-frange-upper', type=float, - help='Upper frequency (in Hz) at which to compute ' - 'qtransform. Optional, default=Half of Nyquist') -parser.add_argument('--qtransform-qrange-lower', default=4, type=float, - help='Lower limit of the range of q to consider, ' - 'default=%(default)f') -parser.add_argument('--qtransform-qrange-upper', default=64, type=float, - help='Upper limit of the range of q to consider, ' - 'default=%(default)f') -parser.add_argument('--qtransform-mismatch', default=0.2, type=float, - help='Mismatch between frequency tiles, default=%(default)f') - -parser.add_argument('--linear-y-axis', dest='log_y', default=True, - action='store_false', - help='Use a linear y-axis. By default a log axis is used.') -parser.add_argument('--linear-colorbar', dest='log_colorbar', default=True, - action='store_false', - help='Use a linear colorbar scale.') -parser.add_argument('--plot-title', - help="If given, use this as the plot title") -parser.add_argument('--plot-caption', - help="If given, use this as the plot caption") -parser.add_argument('--colormap', choices=plt.colormaps(), default='viridis', - help='Colormap to use (default %(default)s)') -parser.add_argument('--mass1', type=float, - help='Provide masses to plot the waveform track on top of ' - 'the qscan.') -parser.add_argument('--mass2', type=float, - help='Provide masses to plot the waveform track on top of ' - 'the qscan.') -parser.add_argument('--spin1z', type=float, default=0, - help='If masses provided, provide spins to plot the ' - 'waveform track on top of the qscan (default=0).') -parser.add_argument('--spin2z', type=float, default=0, - help='If masses provided, provide spins to plot the ' - 'waveform track on top of the qscan (default=0).') +parser.add_argument("--output-file", required=True, help="Output plot") +parser.add_argument( + "--center-time", + type=float, + help="Center plot on the given GPS time. If omitted, use " + "center of interval set via --gps-start-time and " + "--gps-end-time", +) +parser.add_argument( + "--time-windows", + required=True, + type=t_window, + nargs="+", + metavar="START,END", + help="Use these set of times for time windows. Each " + "window produces a different panel. Should " + "be provided as start1,end1 start2,end2 ... " + "where start/end times are relative to --center-time " + "and both positive", +) +parser.add_argument( + "--low-frequency-cutoff", + type=float, + help="The low frequency cutoff to use for fake strain generation", +) + +parser.add_argument( + "--qtransform-delta-t", + default=0.001, + type=float, + help="The time resolution to interpolate to (optional)", +) +parser.add_argument( + "--qtransform-delta-f", + type=float, + help="Frequency resolution to interpolate to (optional)", +) +parser.add_argument( + "--qtransform-logfsteps", + type=int, + default=200, + help="Do a log interpolation (incompatible with " + "--qtransform-delta-f option) and set the number " + "of steps to take", +) +parser.add_argument( + "--qtransform-frange-lower", + type=float, + help="Lower frequency (in Hz) at which to compute " + "qtransform. Optional, default=20 Hz", +) +parser.add_argument( + "--qtransform-frange-upper", + type=float, + help="Upper frequency (in Hz) at which to compute " + "qtransform. Optional, default=Half of Nyquist", +) +parser.add_argument( + "--qtransform-qrange-lower", + default=4, + type=float, + help="Lower limit of the range of q to consider, default=%(default)f", +) +parser.add_argument( + "--qtransform-qrange-upper", + default=64, + type=float, + help="Upper limit of the range of q to consider, default=%(default)f", +) +parser.add_argument( + "--qtransform-mismatch", + default=0.2, + type=float, + help="Mismatch between frequency tiles, default=%(default)f", +) + +parser.add_argument( + "--linear-y-axis", + dest="log_y", + default=True, + action="store_false", + help="Use a linear y-axis. By default a log axis is used.", +) +parser.add_argument( + "--linear-colorbar", + dest="log_colorbar", + default=True, + action="store_false", + help="Use a linear colorbar scale.", +) +parser.add_argument("--plot-title", help="If given, use this as the plot title") +parser.add_argument("--plot-caption", help="If given, use this as the plot caption") +parser.add_argument( + "--colormap", + choices=plt.colormaps(), + default="viridis", + help="Colormap to use (default %(default)s)", +) +parser.add_argument( + "--mass1", + type=float, + help="Provide masses to plot the waveform track on top of the qscan.", +) +parser.add_argument( + "--mass2", + type=float, + help="Provide masses to plot the waveform track on top of the qscan.", +) +parser.add_argument( + "--spin1z", + type=float, + default=0, + help="If masses provided, provide spins to plot the " + "waveform track on top of the qscan (default=0).", +) +parser.add_argument( + "--spin2z", + type=float, + default=0, + help="If masses provided, provide spins to plot the " + "waveform track on top of the qscan (default=0).", +) pycbc.strain.insert_strain_option_group(parser) opts = parser.parse_args() @@ -113,39 +176,43 @@ opts = parser.parse_args() pycbc.init_logging(opts.verbose, default_level=1) if opts.center_time is None: - center_time = (opts.gps_start_time + opts.gps_end_time) / 2. + center_time = (opts.gps_start_time + opts.gps_end_time) / 2.0 else: center_time = opts.center_time if center_time == -1.0: - fig, axes = plt.subplots(1,1) + fig, axes = plt.subplots(1, 1) fig.add_subplot(111, frameon=False) - plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, - right=False) + plt.tick_params(labelcolor="none", top=False, bottom=False, left=False, right=False) plt.grid(False) - plt.xlabel('Time from {:.3f} (s)'.format(opts.center_time)) - plt.ylabel('Frequency (Hz)') - title = 'No data for this detector' + plt.xlabel(f"Time from {opts.center_time:.3f} (s)") + plt.ylabel("Frequency (Hz)") + title = "No data for this detector" pycbc.results.save_fig_with_metadata( - fig, opts.output_file, cmd=' '.join(sys.argv), - fig_kwds={'dpi': 150}, title=title, caption=opts.plot_caption) + fig, + opts.output_file, + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 150}, + title=title, + caption=opts.plot_caption, + ) sys.exit(0) plot_chirp = False if opts.mass1 is not None or opts.mass2 is not None: if None in [opts.mass1, opts.mass2]: - parser.error('Must provide either both --mass1 and --mass2 or neither') + parser.error("Must provide either both --mass1 and --mass2 or neither") plot_chirp = True strain = pycbc.strain.from_cli(opts, pycbc.DYN_RANGE_FAC) -if opts.qtransform_frange_upper is None and \ - opts.qtransform_frange_lower is None: - curr_frange = (20, opts.sample_rate / 4.) -elif opts.qtransform_frange_upper is None or \ - opts.qtransform_frange_lower is None: - parser.error('Must provide either both --qtransform-frange-upper and ' - '--qtransform-frange-lower or neither option.') +if opts.qtransform_frange_upper is None and opts.qtransform_frange_lower is None: + curr_frange = (20, opts.sample_rate / 4.0) +elif opts.qtransform_frange_upper is None or opts.qtransform_frange_lower is None: + parser.error( + "Must provide either both --qtransform-frange-upper and " + "--qtransform-frange-lower or neither option." + ) else: curr_frange = (opts.qtransform_frange_lower, opts.qtransform_frange_upper) @@ -170,13 +237,18 @@ for curr_idx in range(len(wins)): curr_win[0] = float(opts.center_time - strain.start_time - 0.01) if opts.center_time + curr_win[1] > strain.end_time: curr_win[1] = float(strain.end_time - opts.center_time - 0.01) - strain_zoom = strain.time_slice(opts.center_time - curr_win[0], - opts.center_time + curr_win[1]) + strain_zoom = strain.time_slice( + opts.center_time - curr_win[0], opts.center_time + curr_win[1] + ) times, freqs, qvals = strain_zoom.qtransform( - delta_t=opts.qtransform_delta_t, delta_f=opts.qtransform_delta_f, - logfsteps=opts.qtransform_logfsteps, frange=curr_frange, - qrange=qrange, mismatch=opts.qtransform_mismatch) + delta_t=opts.qtransform_delta_t, + delta_f=opts.qtransform_delta_f, + logfsteps=opts.qtransform_logfsteps, + frange=curr_frange, + qrange=qrange, + mismatch=opts.qtransform_mismatch, + ) times_all.append(times) freqs_all.append(freqs) qvals_all.append(qvals) @@ -194,55 +266,70 @@ for curr_idx in range(len(wins)): if opts.log_colorbar: norm = LogNorm(vmin=1, vmax=max_qval) - im = ax.pcolormesh(times - opts.center_time, freqs, qvals, norm=norm, - cmap=opts.colormap) + im = ax.pcolormesh( + times - opts.center_time, freqs, qvals, norm=norm, cmap=opts.colormap + ) ax.set_xlim(-curr_win[0], curr_win[1]) ax.set_ylim(curr_frange[0], curr_frange[1]) if opts.log_y: - ax.set_yscale('log') + ax.set_yscale("log") # https://stackoverflow.com/questions/6963035/pyplot-axes-labels-for-subplots fig.add_subplot(111, frameon=False) -plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, - right=False) +plt.tick_params(labelcolor="none", top=False, bottom=False, left=False, right=False) plt.grid(False) -plt.xlabel('Time from {:.3f} (s)'.format(opts.center_time)) -plt.ylabel('Frequency (Hz)') +plt.xlabel(f"Time from {opts.center_time:.3f} (s)") +plt.ylabel("Frequency (Hz)") # https://stackoverflow.com/questions/13784201/matplotlib-2-subplots-1-colorbar cb = fig.colorbar(im, ax=(axes.ravel().tolist() if len(wins) > 1 else axes)) -cb.set_label('Normalized power') +cb.set_label("Normalized power") if plot_chirp: from pycbc.pnutils import get_inspiral_tf - f_low = 20. - approximant = 'SPAtmplt' if opts.mass1+opts.mass2<4 else 'SEOBNRv4_ROM' - track_t, track_f = get_inspiral_tf(opts.center_time, opts.mass1, - opts.mass2, opts.spin1z, opts.spin2z, - f_low, approximant=approximant) + + f_low = 20.0 + approximant = "SPAtmplt" if opts.mass1 + opts.mass2 < 4 else "SEOBNRv4_ROM" + track_t, track_f = get_inspiral_tf( + opts.center_time, + opts.mass1, + opts.mass2, + opts.spin1z, + opts.spin2z, + f_low, + approximant=approximant, + ) for curr_idx in range(len(wins)): ax = axes[curr_idx] if len(wins) > 1 else axes - ax.plot(track_t - opts.center_time, track_f, 'r-', lw=1.5) + ax.plot(track_t - opts.center_time, track_f, "r-", lw=1.5) if opts.channel_name is not None: plt.suptitle(opts.channel_name) if opts.plot_title is None: - opts.plot_title = 'Q-transform plot around {:.3f}'.format(opts.center_time) + opts.plot_title = f"Q-transform plot around {opts.center_time:.3f}" if opts.plot_caption is None: - opts.plot_caption = ("This shows the Q-transform as a function of time and " - "frequency.") + opts.plot_caption = ( + "This shows the Q-transform as a function of time and frequency." + ) if opts.channel_name is not None: - opts.plot_caption += (' The strain channel is ' + opts.channel_name - + '.') + opts.plot_caption += " The strain channel is " + opts.channel_name + "." if plot_chirp: - chirp_caption = (' The red curve is the time-frequency curve of a' - ' quadrupole-order quasicircular aligned-spin' - ' inspiral with mass1={:.3f}, mass2={:.3f},' - ' spin1z={:.3f}, spin1z={:.3f}.') - opts.plot_caption += chirp_caption.format(opts.mass1, opts.mass2, - opts.spin1z, opts.spin2z) + chirp_caption = ( + " The red curve is the time-frequency curve of a" + " quadrupole-order quasicircular aligned-spin" + " inspiral with mass1={:.3f}, mass2={:.3f}," + " spin1z={:.3f}, spin1z={:.3f}." + ) + opts.plot_caption += chirp_caption.format( + opts.mass1, opts.mass2, opts.spin1z, opts.spin2z + ) pycbc.results.save_fig_with_metadata( - fig, opts.output_file, cmd=' '.join(sys.argv), fig_kwds={'dpi': 150}, - title=opts.plot_title, caption=opts.plot_caption) + fig, + opts.output_file, + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 150}, + title=opts.plot_title, + caption=opts.plot_caption, +) diff --git a/bin/plotting/pycbc_plot_range b/bin/plotting/pycbc_plot_range index f6671652c09..da808ba2d77 100644 --- a/bin/plotting/pycbc_plot_range +++ b/bin/plotting/pycbc_plot_range @@ -1,33 +1,36 @@ #!/usr/bin/env python -""" Plot variation in PSD -""" +"""Plot variation in PSD""" + import matplotlib -matplotlib.use('Agg'); + +matplotlib.use("Agg") +import argparse import logging +import sys + import numpy -import argparse from matplotlib import pyplot as plt -import sys +import pycbc.filter import pycbc.results import pycbc.types import pycbc.waveform -import pycbc.filter +from pycbc.fft.fftw import set_measure_level from pycbc.io.hdf import HFile -from pycbc.fft.fftw import set_measure_level set_measure_level(0) parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--psd-files", nargs='+', help='HDF file of psds') -parser.add_argument("--output-file", help='output file name') -parser.add_argument("--mass1", nargs="+", type=float, - help="Mass of first component in solar masses") -parser.add_argument("--mass2", nargs="+", type=float, - help="Mass of second component in solar masses") -parser.add_argument("--approximant", nargs="+", - help="approximant to use for range") +parser.add_argument("--psd-files", nargs="+", help="HDF file of psds") +parser.add_argument("--output-file", help="output file name") +parser.add_argument( + "--mass1", nargs="+", type=float, help="Mass of first component in solar masses" +) +parser.add_argument( + "--mass2", nargs="+", type=float, help="Mass of second component in solar masses" +) +parser.add_argument("--approximant", nargs="+", help="approximant to use for range") args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -35,30 +38,35 @@ pycbc.init_logging(args.verbose) canonical_snr = 8.0 fig = plt.figure(0) -plt.xlabel('Time (s)') -plt.ylabel('Inspiral Range (Mpc)') +plt.xlabel("Time (s)") +plt.ylabel("Inspiral Range (Mpc)") plt.grid() for psd_file in args.psd_files: - f = HFile(psd_file, 'r') + f = HFile(psd_file, "r") ifo = tuple(f.keys())[0] - flow = f.attrs['low_frequency_cutoff'] - keys = list(f[ifo + '/psds'].keys()) - start, end = f[ifo + '/start_time'][:], f[ifo + '/end_time'][:] + flow = f.attrs["low_frequency_cutoff"] + keys = list(f[ifo + "/psds"].keys()) + start, end = f[ifo + "/start_time"][:], f[ifo + "/end_time"][:] f.close() ranges = {} for i in range(len(keys)): - name = ifo + '/psds/' + str(i) + name = ifo + "/psds/" + str(i) psd = pycbc.types.load_frequencyseries(psd_file, group=name) delta_t = 1.0 / ((len(psd) - 1) * 2 * psd.delta_f) out = pycbc.types.zeros(len(psd), dtype=numpy.complex64) for m1, m2, apx in zip(args.mass1, args.mass2, args.approximant): - htilde = pycbc.waveform.get_waveform_filter(out, - mass1=m1,mass2=m2, approximant=apx, - f_lower=flow, delta_f=psd.delta_f, - delta_t=delta_t, - distance = 1.0/pycbc.DYN_RANGE_FAC) + htilde = pycbc.waveform.get_waveform_filter( + out, + mass1=m1, + mass2=m2, + approximant=apx, + f_lower=flow, + delta_f=psd.delta_f, + delta_t=delta_t, + distance=1.0 / pycbc.DYN_RANGE_FAC, + ) htilde = htilde.astype(numpy.complex64) sigma = pycbc.filter.sigma(htilde, psd=psd, low_frequency_cutoff=flow) horizon_distance = sigma / canonical_snr @@ -72,25 +80,32 @@ for psd_file in args.psd_files: for m1, m2, apx in zip(args.mass1, args.mass2, args.approximant): if len(args.approximant) > 1: - label = '%s: $%sM_{\odot}-%sM_{\odot}$ (%s)' % (ifo, m1, m2, apx) + label = r"%s: $%sM_{\odot}-%sM_{\odot}$ (%s)" % (ifo, m1, m2, apx) else: label = str(ifo) wf_key = (m1, m2, apx) - plt.errorbar((start+end)/2, ranges[wf_key], xerr=(end-start)/2, - ecolor=pycbc.results.ifo_color(ifo), label=label, - fmt='none') + plt.errorbar( + (start + end) / 2, + ranges[wf_key], + xerr=(end - start) / 2, + ecolor=pycbc.results.ifo_color(ifo), + label=label, + fmt="none", + ) -plt.legend(loc="best", fontsize='small') +plt.legend(loc="best", fontsize="small") if len(args.approximant) == 1: - fig.suptitle('$%sM_{\odot}-%sM_{\odot}$ %s' % (m1, m2, apx)) - -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title = "Inspiral Range", - caption = "The canonical sky- and orientation-averaged inspiral range for a single " - "detector at SNR 8. This range is comparable to the SenseMon range and a factor of 2.26 smaller than the horizon distance.", - cmd = ' '.join(sys.argv), - fig_kwds={'dpi':200} - ) + fig.suptitle(r"$%sM_{\odot}-%sM_{\odot}$ %s" % (m1, m2, apx)) + +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title="Inspiral Range", + caption="The canonical sky- and orientation-averaged inspiral range for a single " + "detector at SNR 8. This range is comparable to the SenseMon range and a factor of 2.26 smaller than the horizon distance.", + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200}, +) logging.info("Done!") diff --git a/bin/plotting/pycbc_plot_range_vs_mtot b/bin/plotting/pycbc_plot_range_vs_mtot index f03c7a90b72..1f067ba2290 100644 --- a/bin/plotting/pycbc_plot_range_vs_mtot +++ b/bin/plotting/pycbc_plot_range_vs_mtot @@ -1,36 +1,39 @@ #!/usr/bin/env python -""" Plot variation in PSD -""" +"""Plot variation in PSD""" + import matplotlib -matplotlib.use('Agg') -import matplotlib.pyplot as plt -import numpy + +matplotlib.use("Agg") import argparse -import sys import math +import sys +import matplotlib.pyplot as plt +import numpy + +import pycbc.filter import pycbc.results import pycbc.types import pycbc.waveform -import pycbc.filter +from pycbc.fft.fftw import set_measure_level from pycbc.io.hdf import HFile -from pycbc.fft.fftw import set_measure_level set_measure_level(0) parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--psd-files", nargs='+', - help='HDF file of psds') -parser.add_argument("--output-file", help='output file name') -parser.add_argument("--min_mtot", nargs="+", type=float, - help="Minimum total mass for range") -parser.add_argument("--max_mtot", nargs="+", type=float, - help="Maximum total mass for range") -parser.add_argument("--d_mtot", nargs="+", type=float, - help="Delta total mass for range ") -parser.add_argument("--approximant", nargs="+", - help="approximant to use for range") +parser.add_argument("--psd-files", nargs="+", help="HDF file of psds") +parser.add_argument("--output-file", help="output file name") +parser.add_argument( + "--min_mtot", nargs="+", type=float, help="Minimum total mass for range" +) +parser.add_argument( + "--max_mtot", nargs="+", type=float, help="Maximum total mass for range" +) +parser.add_argument( + "--d_mtot", nargs="+", type=float, help="Delta total mass for range " +) +parser.add_argument("--approximant", nargs="+", help="approximant to use for range") args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -38,16 +41,16 @@ pycbc.init_logging(args.verbose) canonical_snr = 8.0 fig = plt.figure(0) -plt.xlabel('Total Mass (M$_{\odot}$)') -plt.ylabel('Inspiral Range (Mpc)') -plt.grid() +plt.xlabel(r"Total Mass (M$_{\odot}$)") +plt.ylabel("Inspiral Range (Mpc)") +plt.grid() for psd_file in args.psd_files: - f = HFile(psd_file, 'r') + f = HFile(psd_file, "r") ifo = tuple(f.keys())[0] - flow = f.attrs['low_frequency_cutoff'] - keys = f[ifo + '/psds'].keys() - start, end = f[ifo + '/start_time'][:], f[ifo + '/end_time'][:] + flow = f.attrs["low_frequency_cutoff"] + keys = f[ifo + "/psds"].keys() + start, end = f[ifo + "/start_time"][:], f[ifo + "/end_time"][:] seglen = numpy.subtract(end, start) tott = sum(seglen) f.close() @@ -55,22 +58,28 @@ for psd_file in args.psd_files: avg_range, rangerr, mbin = [], [], [] for i in range(len(keys)): - name = ifo + '/psds/' + str(i) + name = ifo + "/psds/" + str(i) psd = pycbc.types.load_frequencyseries(psd_file, group=name) delta_t = 1.0 / ((len(psd) - 1) * 2 * psd.delta_f) out = pycbc.types.zeros(len(psd), dtype=numpy.complex64) - for mi, mf, dm, apx in zip(args.min_mtot, args.max_mtot, args.d_mtot, args.approximant): + for mi, mf, dm, apx in zip( + args.min_mtot, args.max_mtot, args.d_mtot, args.approximant + ): for M in numpy.arange(mi, mf, dm): - htilde = pycbc.waveform.get_waveform_filter(out, - mass1=M/2.,mass2=M/2., approximant=apx, - f_lower=flow, delta_f=psd.delta_f, - delta_t=delta_t, - distance = 1.0/pycbc.DYN_RANGE_FAC) + htilde = pycbc.waveform.get_waveform_filter( + out, + mass1=M / 2.0, + mass2=M / 2.0, + approximant=apx, + f_lower=flow, + delta_f=psd.delta_f, + delta_t=delta_t, + distance=1.0 / pycbc.DYN_RANGE_FAC, + ) htilde = htilde.astype(numpy.complex64) - sigma = pycbc.filter.sigma(htilde, psd=psd, - low_frequency_cutoff=flow) - horizon_distance = sigma / canonical_snr + sigma = pycbc.filter.sigma(htilde, psd=psd, low_frequency_cutoff=flow) + horizon_distance = sigma / canonical_snr inspiral_range = horizon_distance / 2.26 if M in ranges: @@ -80,21 +89,30 @@ for psd_file in args.psd_files: for M in numpy.arange(mi, mf, dm): mean = numpy.average(ranges[M], weights=seglen) - variance = numpy.average((ranges[M]-mean)**2, weights=seglen) + variance = numpy.average((ranges[M] - mean) ** 2, weights=seglen) stddev = math.sqrt(variance) avg_range.append(mean), rangerr.append(stddev), mbin.append(M) for apx in args.approximant: - label = '%s-%s' % (ifo, apx) - plt.errorbar(mbin, avg_range, yerr=rangerr, ecolor=pycbc.results.ifo_color(ifo), label=label, fmt='none') + label = "%s-%s" % (ifo, apx) + plt.errorbar( + mbin, + avg_range, + yerr=rangerr, + ecolor=pycbc.results.ifo_color(ifo), + label=label, + fmt="none", + ) plt.plot(mbin, avg_range, color=pycbc.results.ifo_color(ifo)) plt.legend(loc="upper left") -pycbc.results.save_fig_with_metadata(fig, args.output_file, - title = "Inspiral Range", - caption = "The canonical sky-averaged inspiral range for a single " - "detector at SNR 8 vs total mass:equal mass binary", - cmd = ' '.join(sys.argv), - fig_kwds={'dpi':200} - ) +pycbc.results.save_fig_with_metadata( + fig, + args.output_file, + title="Inspiral Range", + caption="The canonical sky-averaged inspiral range for a single " + "detector at SNR 8 vs total mass:equal mass binary", + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200}, +) diff --git a/bin/plotting/pycbc_plot_singles_timefreq b/bin/plotting/pycbc_plot_singles_timefreq index 35b40ddaf4d..adf5464b154 100644 --- a/bin/plotting/pycbc_plot_singles_timefreq +++ b/bin/plotting/pycbc_plot_singles_timefreq @@ -21,71 +21,83 @@ Plot single-detector inspiral triggers in the time-frequency plane along with a spectrogram of the strain data. """ -import sys -import logging import argparse -import numpy as np +import logging +import sys + import matplotlib -matplotlib.use('agg') +import numpy as np + +matplotlib.use("agg") +from matplotlib import mlab from matplotlib import pyplot as plt -import matplotlib.mlab as mlab from matplotlib.colors import LogNorm from matplotlib.ticker import LogLocator import pycbc -from pycbc.io import HFile import pycbc.events import pycbc.pnutils -import pycbc.strain import pycbc.results +import pycbc.strain import pycbc.waveform - +from pycbc.io import HFile parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trig-file', required=True, - help='HDF5 file containing single triggers') -parser.add_argument('--output-file', required=True, help='Output plot') -parser.add_argument('--bank-file', required=True, - help='HDF5 file containing template bank') -parser.add_argument('--veto-file', help='LIGOLW file containing veto segments') -parser.add_argument('--f-low', type=float, default=20, - help='Low-frequency cutoff') -parser.add_argument('--rank', choices=['snr', 'newsnr'], default='newsnr', - help='Ranking statistic for sorting triggers') -parser.add_argument('--num-loudest', type=int, default=1000, - help='Number of loudest triggers to plot') -parser.add_argument('--interesting-trig', type=int, - help='Index of interesting trigger to highlight') -parser.add_argument('--detector', type=str, required=True) -parser.add_argument('--center-time', type=float, - help='Center plot on the given GPS time') +parser.add_argument( + "--trig-file", required=True, help="HDF5 file containing single triggers" +) +parser.add_argument("--output-file", required=True, help="Output plot") +parser.add_argument( + "--bank-file", required=True, help="HDF5 file containing template bank" +) +parser.add_argument("--veto-file", help="LIGOLW file containing veto segments") +parser.add_argument("--f-low", type=float, default=20, help="Low-frequency cutoff") +parser.add_argument( + "--rank", + choices=["snr", "newsnr"], + default="newsnr", + help="Ranking statistic for sorting triggers", +) +parser.add_argument( + "--num-loudest", type=int, default=1000, help="Number of loudest triggers to plot" +) +parser.add_argument( + "--interesting-trig", type=int, help="Index of interesting trigger to highlight" +) +parser.add_argument("--detector", type=str, required=True) +parser.add_argument( + "--center-time", type=float, help="Center plot on the given GPS time" +) # add the approximant argument -pycbc.waveform.bank.add_approximant_arg(parser, - help='Waveform model used for computing inspiral tracks.' - ' Python expressions can be used to specify a ' - 'parameter-dependent approximant as done in ' - 'pycbc_inspiral. Default is to use the approximant ' - 'specified in the bank file.') +pycbc.waveform.bank.add_approximant_arg( + parser, + help="Waveform model used for computing inspiral tracks." + " Python expressions can be used to specify a " + "parameter-dependent approximant as done in " + "pycbc_inspiral. Default is to use the approximant " + "specified in the bank file.", +) pycbc.strain.insert_strain_option_group(parser) opts = parser.parse_args() pycbc.init_logging(opts.verbose) if opts.center_time is None: - center_time = (opts.gps_start_time + opts.gps_end_time) / 2. + center_time = (opts.gps_start_time + opts.gps_end_time) / 2.0 else: center_time = opts.center_time -fig = plt.figure(figsize=(11,5.5)) +fig = plt.figure(figsize=(11, 5.5)) fig.subplots_adjust(left=0.06, right=0.95, bottom=0.09, top=0.95) ax = fig.gca() opts.low_frequency_cutoff = opts.f_low strain = pycbc.strain.from_cli(opts, pycbc.DYN_RANGE_FAC) -logging.info('Plotting strain spectrogram') -Pxx, freq, t = mlab.specgram(strain, NFFT=1024, noverlap=1000, - Fs=opts.sample_rate, mode='psd') +logging.info("Plotting strain spectrogram") +Pxx, freq, t = mlab.specgram( + strain, NFFT=1024, noverlap=1000, Fs=opts.sample_rate, mode="psd" +) del strain median_psd = np.median(Pxx, axis=1) median_psd_tile = np.tile(np.array([median_psd]).T, (1, len(t))) @@ -93,37 +105,46 @@ del median_psd Pxx /= median_psd_tile del median_psd_tile norm = LogNorm(vmin=1, vmax=1000) -pc = ax.pcolormesh(t + opts.gps_start_time - center_time, freq, Pxx, - norm=norm, cmap='afmhot_r', shading='gouraud') +pc = ax.pcolormesh( + t + opts.gps_start_time - center_time, + freq, + Pxx, + norm=norm, + cmap="afmhot_r", + shading="gouraud", +) del freq, Pxx -logging.info('Loading trigs') -trig_f = HFile(opts.trig_file, 'r') +logging.info("Loading trigs") +trig_f = HFile(opts.trig_file, "r") trigs = trig_f[opts.detector] + def rough_filter(snr, chisq, chisq_dof, end_time, tmp_id, tmp_dur): return np.logical_and( end_time > opts.gps_start_time, - end_time < opts.gps_end_time + tmp_dur.astype('float64') + end_time < opts.gps_end_time + tmp_dur.astype("float64"), ) + indices, data_tuple = trig_f.select( rough_filter, - 'snr', - 'chisq', - 'chisq_dof', - 'end_time', - 'template_id', - 'template_duration', - group=opts.detector + "snr", + "chisq", + "chisq_dof", + "end_time", + "template_id", + "template_duration", + group=opts.detector, ) snr, chisq, chisq_dof, end_time, template_ids, template_duration = data_tuple if len(indices) > 0: if opts.veto_file: - logging.info('Loading veto segments') + logging.info("Loading veto segments") locs, _ = pycbc.events.veto.indices_outside_segments( - end_time, [opts.veto_file], ifo=opts.detector) + end_time, [opts.veto_file], ifo=opts.detector + ) end_time = end_time[locs] snr = snr[locs] chisq = chisq[locs] @@ -132,9 +153,9 @@ if len(indices) > 0: indices = indices[locs] del locs - if opts.rank == 'snr': + if opts.rank == "snr": rank = snr - elif opts.rank == 'newsnr': + elif opts.rank == "newsnr": rchisq = chisq / (chisq_dof * 2 - 2) rank = pycbc.events.ranking.newsnr(snr, rchisq) if type(rank) in [np.float32, np.float64]: @@ -142,34 +163,46 @@ if len(indices) > 0: del rchisq del snr, chisq, chisq_dof - sorter = np.argsort(rank)[::-1][:opts.num_loudest] + sorter = np.argsort(rank)[::-1][: opts.num_loudest] sorted_end_time = end_time[sorter] sorted_rank = rank[sorter] sorted_template_ids = template_ids[sorter] try: - max_rank = max([sorted_rank[i] for i in range(len(sorted_rank)) - if sorted_end_time[i] <= opts.gps_end_time]) + max_rank = max( + [ + sorted_rank[i] + for i in range(len(sorted_rank)) + if sorted_end_time[i] <= opts.gps_end_time + ] + ) except ValueError: max_rank = None - logging.info('Loading bank') - bank = pycbc.waveform.bank.TemplateBank(opts.bank_file, - approximant=opts.approximant, - parameters=['mass1', 'mass2', 'spin1z', 'spin2z', 'approximant'], - load_compressed=False) + logging.info("Loading bank") + bank = pycbc.waveform.bank.TemplateBank( + opts.bank_file, + approximant=opts.approximant, + parameters=["mass1", "mass2", "spin1z", "spin2z", "approximant"], + load_compressed=False, + ) tmplts = bank.table - logging.info('Plotting %d trigs', len(sorted_end_time)) + logging.info("Plotting %d trigs", len(sorted_end_time)) for tc, rho, tid in zip(sorted_end_time, sorted_rank, sorted_template_ids): track_t, track_f = pycbc.pnutils.get_inspiral_tf( - tc - center_time, tmplts.mass1[tid], tmplts.mass2[tid], - tmplts.spin1z[tid], tmplts.spin2z[tid], - opts.f_low, approximant=bank.approximant(tid)) + tc - center_time, + tmplts.mass1[tid], + tmplts.mass2[tid], + tmplts.spin1z[tid], + tmplts.spin2z[tid], + opts.f_low, + approximant=bank.approximant(tid), + ) if max_rank and rho == max_rank: - ax.plot(track_t, track_f, '-', color='#0080ff', zorder=3, lw=2) + ax.plot(track_t, track_f, "-", color="#0080ff", zorder=3, lw=2) else: - ax.plot(track_t, track_f, '-', color='#0000ff', zorder=2, lw=0.5, alpha=0.5) + ax.plot(track_t, track_f, "-", color="#0000ff", zorder=2, lw=0.5, alpha=0.5) if opts.interesting_trig is not None and opts.interesting_trig in indices: interesting_id = np.where(indices == opts.interesting_trig)[0] @@ -178,31 +211,46 @@ if len(indices) > 0: interesting_trig_rank = rho tid = template_ids[interesting_id] track_t, track_f = pycbc.pnutils.get_inspiral_tf( - tc - center_time, tmplts.mass1[tid], tmplts.mass2[tid], - tmplts.spin1z[tid], tmplts.spin2z[tid], - opts.f_low, approximant=bank.approximant(tid)) - ax.plot(track_t, track_f, '-', color='#00f000', zorder=3, lw=2) + tc - center_time, + tmplts.mass1[tid], + tmplts.mass2[tid], + tmplts.spin1z[tid], + tmplts.spin2z[tid], + opts.f_low, + approximant=bank.approximant(tid), + ) + ax.plot(track_t, track_f, "-", color="#00f000", zorder=3, lw=2) else: interesting_trig_rank = None if max_rank: - title = '%s - loudest %d triggers by %s - max %s = %.2f (light blue curve)' \ - % (opts.channel_name, opts.num_loudest, opts.rank, opts.rank, max_rank) + title = "%s - loudest %d triggers by %s - max %s = %.2f (light blue curve)" % ( + opts.channel_name, + opts.num_loudest, + opts.rank, + opts.rank, + max_rank, + ) else: - title = '%s - loudest %d triggers by %s' \ - % (opts.channel_name, opts.num_loudest, opts.rank) + title = "%s - loudest %d triggers by %s" % ( + opts.channel_name, + opts.num_loudest, + opts.rank, + ) if interesting_trig_rank is not None: - title += ' - selected %s = %.2f (green curve)' \ - % (opts.rank, interesting_trig_rank) + title += " - selected %s = %.2f (green curve)" % ( + opts.rank, + interesting_trig_rank, + ) else: - title = '%s - no triggers' % opts.channel_name + title = "%s - no triggers" % opts.channel_name -logging.info('Loading and plotting gates') -for gate_type, hatch_style in [('file', '\\'), ('auto', '/')]: +logging.info("Loading and plotting gates") +for gate_type, hatch_style in [("file", "\\"), ("auto", "/")]: try: - gate_time = trigs['gating/' + gate_type + '/time'][:] - gate_width = trigs['gating/' + gate_type + '/width'][:] - gate_pad = trigs['gating/' + gate_type + '/pad'][:] + gate_time = trigs["gating/" + gate_type + "/time"][:] + gate_width = trigs["gating/" + gate_type + "/width"][:] + gate_pad = trigs["gating/" + gate_type + "/pad"][:] except KeyError: continue gate_unique = list(frozenset(zip(gate_time, gate_width, gate_pad))) @@ -210,46 +258,61 @@ for gate_type, hatch_style in [('file', '\\'), ('auto', '/')]: # only plot gates within the plot time window if gt + gw < opts.gps_start_time or gt - gw > opts.gps_end_time: continue - ax.axvspan(gt - gw - center_time, gt + gw - center_time, - hatch=hatch_style, facecolor='none', edgecolor='#00ff00') + ax.axvspan( + gt - gw - center_time, + gt + gw - center_time, + hatch=hatch_style, + facecolor="none", + edgecolor="#00ff00", + ) if opts.veto_file: - logging.info('Loading and plotting veto segments') + logging.info("Loading and plotting veto segments") veto_segs = pycbc.events.veto.select_segments_by_definer( - opts.veto_file, ifo=opts.detector) + opts.veto_file, ifo=opts.detector + ) veto_segs.coalesce() for seg in veto_segs: if seg[0] > opts.gps_end_time or seg[1] < opts.gps_start_time: continue - ax.axvspan(seg[0] - center_time, seg[1] - center_time, - hatch='x', facecolor='none', edgecolor='#ff0000') + ax.axvspan( + seg[0] - center_time, + seg[1] - center_time, + hatch="x", + facecolor="none", + edgecolor="#ff0000", + ) -half_width = max(opts.gps_end_time - center_time, - center_time - opts.gps_start_time) +half_width = max(opts.gps_end_time - center_time, center_time - opts.gps_start_time) ax.set_xlim(-half_width, half_width) ax.set_ylim(opts.f_low, opts.sample_rate / 2) -ax.set_yscale('log') -ax.grid(ls='solid', alpha=0.2) -ax.set_xlabel('Time - %.3f (s)' % center_time) -ax.set_ylabel('Frequency (Hz)') +ax.set_yscale("log") +ax.grid(ls="solid", alpha=0.2) +ax.set_xlabel("Time - %.3f (s)" % center_time) +ax.set_ylabel("Frequency (Hz)") ax.set_title(title) -cb = fig.colorbar(pc, fraction=0.04, pad=0.01, - ticks=LogLocator(subs=range(10))) -cb.set_label('Power density normalized to its median over time') - -caption = ('This plot shows the power spectrogram of the strain data, ' - 'normalized to its median over time, as a heatmap. The ' - 'time-frequency evolution of each single trigger is shown as a ' - 'blue curve. The light-blue curve is the loudest trigger by {0}. ' - 'Only the loudest {1} triggers by {0} are shown. Green hatched ' - 'areas denote gated strain data (\\\\ = externally-provided ' - 'gates, // = autogates). Red hatched areas denote vetoed time.') +cb = fig.colorbar(pc, fraction=0.04, pad=0.01, ticks=LogLocator(subs=range(10))) +cb.set_label("Power density normalized to its median over time") + +caption = ( + "This plot shows the power spectrogram of the strain data, " + "normalized to its median over time, as a heatmap. The " + "time-frequency evolution of each single trigger is shown as a " + "blue curve. The light-blue curve is the loudest trigger by {0}. " + "Only the loudest {1} triggers by {0} are shown. Green hatched " + "areas denote gated strain data (\\\\ = externally-provided " + "gates, // = autogates). Red hatched areas denote vetoed time." +) caption = caption.format(opts.rank, opts.num_loudest) -md_title = 'Strain spectrogram and inspiral tracks for ' + opts.detector +md_title = "Strain spectrogram and inspiral tracks for " + opts.detector pycbc.results.save_fig_with_metadata( - fig, opts.output_file, cmd=' '.join(sys.argv), + fig, + opts.output_file, + cmd=" ".join(sys.argv), title=md_title, - caption=caption, fig_kwds={'dpi': 100}) + caption=caption, + fig_kwds={"dpi": 100}, +) -logging.info('Done') +logging.info("Done") diff --git a/bin/plotting/pycbc_plot_singles_vs_params b/bin/plotting/pycbc_plot_singles_vs_params index 05f4197787b..9fca1c82c17 100644 --- a/bin/plotting/pycbc_plot_singles_vs_params +++ b/bin/plotting/pycbc_plot_singles_vs_params @@ -20,72 +20,96 @@ Plot PyCBC's single-detector triggers over the search parameter space. """ -import logging import argparse -import numpy as np +import logging + import matplotlib -matplotlib.use('agg') +import numpy as np + +matplotlib.use("agg") +import sys + from matplotlib import pyplot as plt from matplotlib.colors import LogNorm from matplotlib.ticker import LogLocator -import sys from packaging.version import Version import pycbc -import pycbc.pnutils import pycbc.events -import pycbc.results import pycbc.io +import pycbc.pnutils +import pycbc.results from pycbc.events import ranking parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--single-trig-file', required=True, - help='Path to file containing single-detector triggers in ' - 'HDF5 format. Required') -parser.add_argument('--bank-file', required=True, - help='Path to file containing template bank in HDF5 format' - '. Required') -parser.add_argument('--veto-file', type=str, - help='Optional path to file containing veto segments') -parser.add_argument('--segment-name', default=None, type=str, - help='Optional, name of segment list to use for vetoes') -parser.add_argument('--filter-string', default=None, type=str, - help='Optional, boolean expression for filtering triggers.' - 'Example: "self.snr >=6"' - 'Parameter in "self.parameter" should be either from' - 'trigger or bank params.') -parser.add_argument('--output-file', type=str, required=True, - help='Destination path for plot') -parser.add_argument('--x-var', required=True, - choices=pycbc.io.SingleDetTriggers.get_param_names(), - help='Parameter to plot on the x-axis. Required') -parser.add_argument('--y-var', required=True, - choices=pycbc.io.SingleDetTriggers.get_param_names(), - help='Parameter to plot on the y-axis. Required') +parser.add_argument( + "--single-trig-file", + required=True, + help="Path to file containing single-detector triggers in HDF5 format. Required", +) +parser.add_argument( + "--bank-file", + required=True, + help="Path to file containing template bank in HDF5 format. Required", +) +parser.add_argument( + "--veto-file", type=str, help="Optional path to file containing veto segments" +) +parser.add_argument( + "--segment-name", + default=None, + type=str, + help="Optional, name of segment list to use for vetoes", +) +parser.add_argument( + "--filter-string", + default=None, + type=str, + help="Optional, boolean expression for filtering triggers." + 'Example: "self.snr >=6"' + 'Parameter in "self.parameter" should be either from' + "trigger or bank params.", +) +parser.add_argument( + "--output-file", type=str, required=True, help="Destination path for plot" +) +parser.add_argument( + "--x-var", + required=True, + choices=pycbc.io.SingleDetTriggers.get_param_names(), + help="Parameter to plot on the x-axis. Required", +) +parser.add_argument( + "--y-var", + required=True, + choices=pycbc.io.SingleDetTriggers.get_param_names(), + help="Parameter to plot on the y-axis. Required", +) ranking_keys = list(ranking.sngls_ranking_function_dict.keys()) -parser.add_argument('--z-var', required=True, - choices=['density'] + ranking_keys, - help='Quantity to plot on the color scale. Required') -parser.add_argument('--detector', required=True, - help='Detector. Required') -parser.add_argument('--grid-size', type=int, default=80, - help='Bin resolution (larger = smaller bins)') -parser.add_argument('--log-x', action='store_true', - help='Use log scale for x-axis') -parser.add_argument('--log-y', action='store_true', - help='Use log scale for y-axis') -parser.add_argument('--min-x', type=float, help='Optional minimum x value') -parser.add_argument('--max-x', type=float, help='Optional maximum x value') -parser.add_argument('--min-y', type=float, help='Optional minimum y value') -parser.add_argument('--max-y', type=float, help='Optional maximum y value') -parser.add_argument('--min-z', type=float, help='Optional minimum z value') -parser.add_argument('--max-z', type=float, help='Optional maximum z value') +parser.add_argument( + "--z-var", + required=True, + choices=["density"] + ranking_keys, + help="Quantity to plot on the color scale. Required", +) +parser.add_argument("--detector", required=True, help="Detector. Required") +parser.add_argument( + "--grid-size", type=int, default=80, help="Bin resolution (larger = smaller bins)" +) +parser.add_argument("--log-x", action="store_true", help="Use log scale for x-axis") +parser.add_argument("--log-y", action="store_true", help="Use log scale for y-axis") +parser.add_argument("--min-x", type=float, help="Optional minimum x value") +parser.add_argument("--max-x", type=float, help="Optional maximum x value") +parser.add_argument("--min-y", type=float, help="Optional minimum y value") +parser.add_argument("--max-y", type=float, help="Optional maximum y value") +parser.add_argument("--min-z", type=float, help="Optional minimum z value") +parser.add_argument("--max-z", type=float, help="Optional maximum z value") opts = parser.parse_args() pycbc.init_logging(opts.verbose) -if opts.z_var == 'density' or opts.min_z is None: +if opts.z_var == "density" or opts.min_z is None: filter_rank = None filter_thresh = None else: @@ -101,7 +125,7 @@ trigs = pycbc.io.SingleDetTriggers( segment_name=opts.segment_name, filter_func=opts.filter_string, filter_rank=filter_rank, - filter_threshold=filter_thresh + filter_threshold=filter_thresh, ) # Can this be folded into the SingleDetTriggers call? Maybe, but we can @@ -121,52 +145,53 @@ if opts.max_y is not None: x = x[mask] y = y[mask] -title = f'{opts.z_var.title()} of {opts.detector} triggers ' + \ - f'over {opts.x_var.title()} and {opts.y_var.title()}' -fig_caption = f"This plot shows the {opts.z_var} of single detector " + \ - f"triggers for the {opts.detector} detector. " + \ - f"{opts.z_var.title()} is shown on the colorbar axis " + \ - f"against {opts.x_var} and {opts.y_var} on the x- and y-axes." +title = ( + f"{opts.z_var.title()} of {opts.detector} triggers " + f"over {opts.x_var.title()} and {opts.y_var.title()}" +) +fig_caption = ( + f"This plot shows the {opts.z_var} of single detector " + f"triggers for the {opts.detector} detector. " + f"{opts.z_var.title()} is shown on the colorbar axis " + f"against {opts.x_var} and {opts.y_var} on the x- and y-axes." +) if not any(mask): # All triggers removed - make a blank plot which says so: fig = plt.figure() ax = fig.gca() - plt.text(0.5, 0.5, 'no triggers in the range') + plt.text(0.5, 0.5, "no triggers in the range") pycbc.results.save_fig_with_metadata( fig, opts.output_file, title=title, caption=fig_caption, - cmd=' '.join(sys.argv), - fig_kwds={'dpi': 200} + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200}, ) sys.exit() -hexbin_style = { - 'gridsize': opts.grid_size, - 'linewidths': 0.03 -} +hexbin_style = {"gridsize": opts.grid_size, "linewidths": 0.03} # In earlier versions mpl will try to take the max over bins with 0 triggers # and fail, unless we tell it to leave these blank by setting mincnt -if Version(matplotlib.__version__) < Version('3.8.1'): - hexbin_style['mincnt'] = 0 +if Version(matplotlib.__version__) < Version("3.8.1"): + hexbin_style["mincnt"] = 0 if opts.log_x: - hexbin_style['xscale'] = 'log' + hexbin_style["xscale"] = "log" if opts.log_y: - hexbin_style['yscale'] = 'log' -minz = opts.min_z if opts.min_z else 1 + hexbin_style["yscale"] = "log" +minz = opts.min_z or 1 maxz = opts.max_z -hexbin_style['norm'] = LogNorm(vmin=minz, vmax=maxz) +hexbin_style["norm"] = LogNorm(vmin=minz, vmax=maxz) -logging.info('Plotting') +logging.info("Plotting") fig = plt.figure() ax = fig.gca() -if opts.z_var == 'density': +if opts.z_var == "density": hb = ax.hexbin(x, y, **hexbin_style) fig.colorbar(hb, ticks=LogLocator(subs=range(10))) elif opts.z_var in ranking.sngls_ranking_function_dict: @@ -177,17 +202,22 @@ elif opts.z_var in ranking.sngls_ranking_function_dict: min_z = z.min() if opts.min_z is None else opts.min_z max_z = z.max() if opts.max_z is None else opts.max_z if max_z / min_z > 10: - cb_style['ticks'] = LogLocator(subs=range(10)) + cb_style["ticks"] = LogLocator(subs=range(10)) hb = ax.hexbin(x, y, C=z, reduce_C_function=np.max, **hexbin_style) fig.colorbar(hb, **cb_style) else: - raise RuntimeError('z_var = %s is not recognized!' % (opts.z_var)) + raise RuntimeError("z_var = %s is not recognized!" % (opts.z_var)) ax.set_xlabel(opts.x_var) ax.set_ylabel(opts.y_var) -ax.set_title(opts.z_var.title() + ' of %s triggers ' % (opts.detector)) -pycbc.results.save_fig_with_metadata(fig, opts.output_file, title=title, - caption=fig_caption, cmd=' '.join(sys.argv), - fig_kwds={'dpi': 200}) +ax.set_title(opts.z_var.title() + " of %s triggers " % (opts.detector)) +pycbc.results.save_fig_with_metadata( + fig, + opts.output_file, + title=title, + caption=fig_caption, + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 200}, +) -logging.info('Done') +logging.info("Done") diff --git a/bin/plotting/pycbc_plot_throughput b/bin/plotting/pycbc_plot_throughput index 4b8438a7fd5..296c1072908 100755 --- a/bin/plotting/pycbc_plot_throughput +++ b/bin/plotting/pycbc_plot_throughput @@ -1,83 +1,88 @@ #!/usr/bin/env python import argparse + import matplotlib -matplotlib.use('Agg') -from matplotlib import pyplot as plt +matplotlib.use("Agg") +from matplotlib import pyplot as plt from scipy.stats import hmean import pycbc -from pycbc.results.color import ifo_color from pycbc.io.hdf import HFile +from pycbc.results.color import ifo_color parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input-file', nargs='+', required=True, - help='Single-detector inspiral HDF5 files to get ' - 'templates per core.') -parser.add_argument('--output-file', required=True, - help='Destination file for the plot.') -parser.add_argument('--duration-weighted', action="store_true") +parser.add_argument( + "--input-file", + nargs="+", + required=True, + help="Single-detector inspiral HDF5 files to get templates per core.", +) +parser.add_argument( + "--output-file", required=True, help="Destination file for the plot." +) +parser.add_argument("--duration-weighted", action="store_true") args = parser.parse_args() pycbc.init_logging(args.verbose) -fig, (ax1, ax2, ax3) = plt.subplots(3,1,figsize=(10,10)) +fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 10)) for pa in args.input_file: - f = HFile(pa, 'r') + f = HFile(pa, "r") ifo = tuple(f.keys())[0] - dur = f['%s/search/end_time' % ifo][:] - f['%s/search/start_time' % ifo][:] + dur = f["%s/search/end_time" % ifo][:] - f["%s/search/start_time" % ifo][:] dur /= dur.mean() if args.duration_weighted: w = dur else: w = None - if 'templates_per_core' in f['%s/search' % ifo].keys(): - tpc = f['%s/search/templates_per_core' % ifo][:] + if "templates_per_core" in f["%s/search" % ifo].keys(): + tpc = f["%s/search/templates_per_core" % ifo][:] else: tpc = None - if 'filter_rate_per_core' in f['%s/search' % ifo].keys(): - fpc = f['%s/search/filter_rate_per_core' % ifo][:] + if "filter_rate_per_core" in f["%s/search" % ifo].keys(): + fpc = f["%s/search/filter_rate_per_core" % ifo][:] else: fpc = None - if 'setup_time_fraction' in f['%s/search' % ifo].keys(): - stf = f['%s/search/setup_time_fraction' % ifo][:] + if "setup_time_fraction" in f["%s/search" % ifo].keys(): + stf = f["%s/search/setup_time_fraction" % ifo][:] else: stf = None if tpc is not None: - label = str(ifo) + ': Harmonic mean - ' + str(hmean(tpc)) + label = str(ifo) + ": Harmonic mean - " + str(hmean(tpc)) if args.duration_weighted: - avg_tpc = 1.0 / ( w / tpc).mean() + avg_tpc = 1.0 / (w / tpc).mean() else: avg_tpc = hmean(tpc) - label = str(ifo) + ': Harmonic mean - ' + str(avg_tpc) - ax1.hist(tpc, 100, color=ifo_color(ifo), alpha = 0.65, label = label, weights=w) - #ax1.set_title('Templates per Core') - ax1.set_xlabel('Templates per Core') - ax1.legend(loc = 'upper right') + label = str(ifo) + ": Harmonic mean - " + str(avg_tpc) + ax1.hist(tpc, 100, color=ifo_color(ifo), alpha=0.65, label=label, weights=w) + # ax1.set_title('Templates per Core') + ax1.set_xlabel("Templates per Core") + ax1.legend(loc="upper right") ax1.grid(True) if fpc is not None: if w is not None: fpc *= w - label = str(ifo) + ': Mean average - ' + str(fpc.mean()) - ax2.hist(fpc, 100, color=ifo_color(ifo), alpha = 0.65, label = label, weights=w) - #ax2.set_title('Filter rate per core') - ax2.set_xlabel('Filter rate per core (# FFTs per second per core)') - ax2.legend(loc = 'upper right') + label = str(ifo) + ": Mean average - " + str(fpc.mean()) + ax2.hist(fpc, 100, color=ifo_color(ifo), alpha=0.65, label=label, weights=w) + # ax2.set_title('Filter rate per core') + ax2.set_xlabel("Filter rate per core (# FFTs per second per core)") + ax2.legend(loc="upper right") ax2.grid(True) if stf is not None: if w is not None: stf *= w - label = str(ifo) + ': Mean average - ' + str(stf.mean()) - ax3.hist(stf, 100, color=ifo_color(ifo), alpha = 0.65, label = label, weights=w) - #ax2.set_title('Filter rate per core') - ax3.set_xlabel('Fraction of time doing setup operations') - ax3.legend(loc = 'upper right') + label = str(ifo) + ": Mean average - " + str(stf.mean()) + ax3.hist(stf, 100, color=ifo_color(ifo), alpha=0.65, label=label, weights=w) + # ax2.set_title('Filter rate per core') + ax3.set_xlabel("Fraction of time doing setup operations") + ax3.legend(loc="upper right") ax3.grid(True) fig.tight_layout() diff --git a/bin/plotting/pycbc_plot_trigrate b/bin/plotting/pycbc_plot_trigrate index 12ce70bf7bb..d3bfe08ddd8 100644 --- a/bin/plotting/pycbc_plot_trigrate +++ b/bin/plotting/pycbc_plot_trigrate @@ -13,95 +13,150 @@ # Public License for more details. -import sys import argparse import logging +import sys + from matplotlib import use -use('Agg') -from matplotlib import pyplot as plt + +use("Agg") import numpy as np +from matplotlib import pyplot as plt from scipy import stats as scistats import pycbc -from pycbc import io, events, bin_utils, results -from pycbc.events import triggers -from pycbc.events import ranking +from pycbc import bin_utils, events, io, results +from pycbc.events import ranking, triggers #### DEFINITIONS AND FUNCTIONS #### + def get_stat(statchoice, trigs): # For now this is using the single detector ranking. If we want, this # could use the Stat classes in stat.py using similar code as in hdf/io.py # This requires additional options, so only change this if it's useful! return ranking.get_sngls_ranking_from_trigs(trigs, statname) + #### MAIN #### parser = argparse.ArgumentParser(usage="", description="Plot trigger rates") pycbc.add_common_pycbc_options(parser) -parser.add_argument("--trigger-file", - help="Input hdf5 file containing single triggers. " - "Required") -parser.add_argument("--bank-file", default=None, - help="hdf file containing template parameters. Required") -parser.add_argument("--veto-file", nargs='*', default=[], action='append', - help="File(s) in .xml format with veto segments to apply " - "to triggers before fitting") -parser.add_argument("--veto-segment-name", nargs='*', default=[], action='append', - help="Name(s) of veto segments to apply. Optional, if not " - "given all segments for a given ifo will be used") -parser.add_argument("--ifo", required=True, - help="Ifo producing triggers to be fitted. Required") -parser.add_argument("--gps-start-time", type=float, - help="Time from which to load and plot triggers") -parser.add_argument("--gps-end-time", type=float, - help="End time up to which to load/plot triggers") -parser.add_argument("--sngl-stat", default="new_snr", - choices=ranking.sngls_ranking_function_dict.keys(), - help="Function of SNR and chisq to threshold on") -parser.add_argument("--stat-factor", type=float, - help="Adjustable magic number used in some sngl " - "statistics. Values commonly used: 6 for new_snr, 250 " - "or 50 for effective_snr") -parser.add_argument("--stat-threshold", type=float, - help="Only plot triggers with statistic value above this " - "threshold") -parser.add_argument("--thorne-limit", action="store_true", - help="Remove triggers from templates with one or both " - "spins above 0.998") -parser.add_argument("--f-lower", type=float, default=0., - help="Starting frequency for calculating template " - "duration; if not given, duration will be read from " - "single trigger files") -parser.add_argument("--bin-param", required=True, - help="Parameter over which to bin. Required. " - "Choose from mchirp, mtotal, template_duration or a named " - "frequency cutoff in pnutils or a frequency function in " - "LALSimulation") -parser.add_argument("--bin-spacing", choices=["linear", "log", "irreg"], - help="How to space parameter bin edges") +parser.add_argument( + "--trigger-file", help="Input hdf5 file containing single triggers. Required" +) +parser.add_argument( + "--bank-file", + default=None, + help="hdf file containing template parameters. Required", +) +parser.add_argument( + "--veto-file", + nargs="*", + default=[], + action="append", + help="File(s) in .xml format with veto segments to apply " + "to triggers before fitting", +) +parser.add_argument( + "--veto-segment-name", + nargs="*", + default=[], + action="append", + help="Name(s) of veto segments to apply. Optional, if not " + "given all segments for a given ifo will be used", +) +parser.add_argument( + "--ifo", required=True, help="Ifo producing triggers to be fitted. Required" +) +parser.add_argument( + "--gps-start-time", type=float, help="Time from which to load and plot triggers" +) +parser.add_argument( + "--gps-end-time", type=float, help="End time up to which to load/plot triggers" +) +parser.add_argument( + "--sngl-stat", + default="new_snr", + choices=ranking.sngls_ranking_function_dict.keys(), + help="Function of SNR and chisq to threshold on", +) +parser.add_argument( + "--stat-factor", + type=float, + help="Adjustable magic number used in some sngl " + "statistics. Values commonly used: 6 for new_snr, 250 " + "or 50 for effective_snr", +) +parser.add_argument( + "--stat-threshold", + type=float, + help="Only plot triggers with statistic value above this threshold", +) +parser.add_argument( + "--thorne-limit", + action="store_true", + help="Remove triggers from templates with one or both spins above 0.998", +) +parser.add_argument( + "--f-lower", + type=float, + default=0.0, + help="Starting frequency for calculating template " + "duration; if not given, duration will be read from " + "single trigger files", +) +parser.add_argument( + "--bin-param", + required=True, + help="Parameter over which to bin. Required. " + "Choose from mchirp, mtotal, template_duration or a named " + "frequency cutoff in pnutils or a frequency function in " + "LALSimulation", +) +parser.add_argument( + "--bin-spacing", + choices=["linear", "log", "irreg"], + help="How to space parameter bin edges", +) binspec = parser.add_mutually_exclusive_group() -binspec.add_argument("--num-bins", type=int, - help="Number of regularly spaced bins to use over the " - " parameter") -binspec.add_argument("--irregular-bins", type=float, nargs="*", - help="Boundaries of irregular bins") -parser.add_argument("--bin-param-units", - help="String to display units of the binning parameter") -parser.add_argument("--approximant", default="SEOBNRv4", - help="Approximant for template duration. Default SEOBNRv4") -parser.add_argument("--min-duration", default=0., - help="Fudge factor for templates with tiny or negative " - "values of template_duration: add to duration values " - "before fitting. Units seconds") -parser.add_argument("--kde-bandwidth", type=float, default=1., - help="Width of the smoothing kernel as compared to total " - "plot duration. 0.02 usually works OK") -parser.add_argument("--raw-rate", action="store_true", - help="Plot rates without normalizing to template count") +binspec.add_argument( + "--num-bins", + type=int, + help="Number of regularly spaced bins to use over the parameter", +) +binspec.add_argument( + "--irregular-bins", type=float, nargs="*", help="Boundaries of irregular bins" +) +parser.add_argument( + "--bin-param-units", help="String to display units of the binning parameter" +) +parser.add_argument( + "--approximant", + default="SEOBNRv4", + help="Approximant for template duration. Default SEOBNRv4", +) +parser.add_argument( + "--min-duration", + default=0.0, + help="Fudge factor for templates with tiny or negative " + "values of template_duration: add to duration values " + "before fitting. Units seconds", +) +parser.add_argument( + "--kde-bandwidth", + type=float, + default=1.0, + help="Width of the smoothing kernel as compared to total " + "plot duration. 0.02 usually works OK", +) +parser.add_argument( + "--raw-rate", + action="store_true", + help="Plot rates without normalizing to template count", +) parser.add_argument("--log-y", action="store_true") -parser.add_argument("--output-file", - help="Name of file to save plot") +parser.add_argument("--output-file", help="Name of file to save plot") args = parser.parse_args() @@ -111,28 +166,32 @@ args.veto_file = sum(args.veto_file, []) if len(args.veto_segment_name) != len(args.veto_file): raise RuntimeError("Number of veto files much match veto file names") -if (args.gps_start_time or args.gps_end_time) and not (args.gps_start_time \ - and args.gps_end_time): +if (args.gps_start_time or args.gps_end_time) and not ( + args.gps_start_time and args.gps_end_time +): raise RuntimeError("I need both gps start time and end time!") pycbc.init_logging(opts.verbose) -statname = "reweighted SNR" if args.sngl_stat == "new_snr" else \ - args.sngl_stat.replace("_", " ").replace("snr", "SNR") +statname = ( + "reweighted SNR" + if args.sngl_stat == "new_snr" + else args.sngl_stat.replace("_", " ").replace("snr", "SNR") +) paramname = args.bin_param.replace("_", " ") paramtag = args.bin_param.replace("_", "") -logging.info('Opening trigger file: %s' % args.trigger_file) -trigf = io.HFile(args.trigger_file, 'r') -logging.info('Opening template file: %s' % args.bank_file) -templatef = io.HFile(args.bank_file, 'r') +logging.info("Opening trigger file: %s" % args.trigger_file) +trigf = io.HFile(args.trigger_file, "r") +logging.info("Opening template file: %s" % args.bank_file) +templatef = io.HFile(args.bank_file, "r") # get the stat values stat = get_stat(args.sngl_stat, trigf[args.ifo]) # get the duration values if needed -if args.bin_param == 'template_duration' and not args.f_lower: - logging.info('Using template duration from the trigger file') +if args.bin_param == "template_duration" and not args.f_lower: + logging.info("Using template duration from the trigger file") trig_dur = True else: trig_dur = False @@ -140,33 +199,35 @@ else: # stat threshold to reduce trigger numbers abovethresh = stat >= args.stat_threshold stat = stat[abovethresh] -tid = trigf[args.ifo+'/template_id'][:][abovethresh] -time = trigf[args.ifo+'/end_time'][:][abovethresh] +tid = trigf[args.ifo + "/template_id"][:][abovethresh] +time = trigf[args.ifo + "/end_time"][:][abovethresh] if trig_dur: - tdur = trigf[args.ifo+'/template_duration'][:][abovethresh] -logging.info('%i trigs left after thresholding at %f' % (len(stat), args.stat_threshold)) + tdur = trigf[args.ifo + "/template_duration"][:][abovethresh] +logging.info( + "%i trigs left after thresholding at %f" % (len(stat), args.stat_threshold) +) del stat if args.gps_start_time and args.gps_end_time: inside = np.logical_and(time >= args.gps_start_time, time < args.gps_end_time) - #stat = stat[inside] + # stat = stat[inside] tid = tid[inside] time = time[inside] if trig_dur: tdur = tdur[inside] - logging.info('%i trigs left after restricting gps times' % len(time)) + logging.info("%i trigs left after restricting gps times" % len(time)) # now do vetoing for vfile, vsegmentname in zip(args.veto_file, args.veto_segment_name): - retain, junk = events.veto.indices_outside_segments(time, [vfile], - ifo=args.ifo, segment_name=vsegmentname) - #stat = stat[retain] + retain, junk = events.veto.indices_outside_segments( + time, [vfile], ifo=args.ifo, segment_name=vsegmentname + ) + # stat = stat[retain] tid = tid[retain] time = time[retain] if trig_dur: tdur = tdur[retain] - logging.info('%i trigs left after vetoing with %s' % - (len(time), args.veto_file)) + logging.info("%i trigs left after vetoing with %s" % (len(time), args.veto_file)) # get a minimum time for plotting purposes if not args.gps_start_time: @@ -177,38 +238,48 @@ plottime = time - args.gps_start_time if args.thorne_limit: m1, m2, s1z, s2z = triggers.get_mass_spin(templatef, tid) inside = np.logical_and(abs(s1z) < 0.998, abs(s2z) < 0.998) - #stat = stat[inside] + # stat = stat[inside] tid = tid[inside] time = time[inside] if trig_dur: tdur = tdur[inside] + def get_pars(args, tag, m1, m2, s1z, s2z): # used for binning params - paramarg = getattr(args, tag+'_param') + paramarg = getattr(args, tag + "_param") try: # will fail if m1 is a float rather than a sequence - logging.info('Getting %s values for %i triggers' % (paramarg, len(m1))) + logging.info("Getting %s values for %i triggers" % (paramarg, len(m1))) except: pass return triggers.get_param(paramarg, args, m1, m2, s1z, s2z) + # get binning params if trig_dur: binpars = tdur + args.min_duration else: m1, m2, s1z, s2z = triggers.get_mass_spin(templatef, tid) - binpars = get_pars(args, 'bin', m1, m2, s1z, s2z) -logging.info("Parameter range of triggers: %f - %f" % - (min(binpars), max(binpars))) + binpars = get_pars(args, "bin", m1, m2, s1z, s2z) +logging.info("Parameter range of triggers: %f - %f" % (min(binpars), max(binpars))) # get the bins # we assume that parvals are all positive assert min(binpars) >= 0 pmin = 0.999 * min(binpars) pmax = 1.001 * max(binpars) -bincolors = ['r',(1.0,0.65,0),#'y', - 'g','c','b','m','k',(0.8,0.25,0),(0.25,0.8,0)] +bincolors = [ + "r", + (1.0, 0.65, 0), #'y', + "g", + "c", + "b", + "m", + "k", + (0.8, 0.25, 0), + (0.25, 0.8, 0), +] if args.bin_spacing == "linear": pbins = bin_utils.LinearBins(pmin, pmax, args.num_bins) elif args.bin_spacing == "log": @@ -217,7 +288,7 @@ elif args.bin_spacing == "irreg": # allow bins in reverse order! if args.irregular_bins[1] < args.irregular_bins[0]: args.irregular_bins = args.irregular_bins[::-1] - #bincolors = bincolors[::-1] + # bincolors = bincolors[::-1] pbins = bin_utils.IrregularBins(args.irregular_bins) # list of bin indices @@ -254,21 +325,20 @@ for i, lower, upper in zip(binind, pbins.lower(), pbins.upper()): yplot = kd(xplot) * bincounts[i] else: yplot = kd(xplot) * bincounts[i] / bintemplates[i] - minrate = min(minrate, yplot.max()/5e2) + minrate = min(minrate, yplot.max() / 5e2) print(minrate) maxrate = max(maxrate, yplot.max()) binlabel = r"%.3g - %.3g" % (lower, upper) - plt.plot(xplot, yplot, '-', c=bincolors[i], label=binlabel) + plt.plot(xplot, yplot, "-", c=bincolors[i], label=binlabel) # finish the plot -leg = plt.legend(labelspacing=0.2, loc='lower center') -unitstring = " (%s)" % args.bin_param_units if \ - args.bin_param_units is not None else "" -leg.set_title(paramname+unitstring) +leg = plt.legend(labelspacing=0.2, loc="lower center") +unitstring = " (%s)" % args.bin_param_units if args.bin_param_units is not None else "" +leg.set_title(paramname + unitstring) plt.setp(leg.get_texts(), fontsize=11) if args.log_y: - plt.yscale('log') -plt.ylim(minrate, 1.4*maxrate) + plt.yscale("log") +plt.ylim(minrate, 1.4 * maxrate) plt.xlim(0, maxtime) plt.grid() plt.xlabel("GPS time after %i" % args.gps_start_time, size="large") @@ -278,13 +348,15 @@ else: plt.ylabel(r"Rate per template (s$^{-1}$)", size="large") logging.info("Saving to %s" % args.output_file) results.save_fig_with_metadata( - fig, args.output_file, - title="%s: rate of triggers above a %s threshold %f" % (args.ifo, - statname, args.stat_threshold), - caption=(r"Rate of %s single detector triggers thresholded on %s" \ - % (args.ifo, statname)), - cmd=" ".join(sys.argv) - ) + fig, + args.output_file, + title="%s: rate of triggers above a %s threshold %f" + % (args.ifo, statname, args.stat_threshold), + caption=( + r"Rate of %s single detector triggers thresholded on %s" % (args.ifo, statname) + ), + cmd=" ".join(sys.argv), +) plt.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/plotting/pycbc_plot_vt_ratio b/bin/plotting/pycbc_plot_vt_ratio index 66a26815c30..1c7b7e36012 100644 --- a/bin/plotting/pycbc_plot_vt_ratio +++ b/bin/plotting/pycbc_plot_vt_ratio @@ -6,60 +6,72 @@ sets of injections. It reads two HDF files produced by pycbc_page_sensitivity's --hdf-out option, and plots the ratios of their VTs at various IFARs. """ -import sys import argparse import logging +import sys + import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np -from pycbc import init_logging, add_common_pycbc_options -from pycbc.results import save_fig_with_metadata +from pycbc import add_common_pycbc_options, init_logging from pycbc.io.hdf import HFile +from pycbc.results import save_fig_with_metadata parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--vt-files-one', nargs='+', - help='HDF files containing VT curves, data for ' - 'the numerator (top) of the ratio') -parser.add_argument('--vt-files-two', nargs='+', - help='HDF files containing VT curves, data for ' - 'the denominator (bottom) of the ratio') -parser.add_argument('--desc-one', required=True, - help='Descriptor tag for first set of data ' - '(short, for use in subscript)') -parser.add_argument('--desc-two', type=str, required=True, - help='Descriptor tag for second set of data ' - '(short, for use in subscript)') -parser.add_argument('--outfile', type=str, required=True, - help='Output file to save to') -parser.add_argument('--ifars', type=float, required=True, nargs='+', - help='IFAR values to plot VT ratio for. Note that the ' - 'plotted values will be the closest values available ' - 'from the VT files') -parser.add_argument('--log-x', action='store_true', - help='Use logarithmic x-axis') -parser.add_argument('--log-y', action='store_true', - help='Use logarithmic y-axis') +parser.add_argument( + "--vt-files-one", + nargs="+", + help="HDF files containing VT curves, data for the numerator (top) of the ratio", +) +parser.add_argument( + "--vt-files-two", + nargs="+", + help="HDF files containing VT curves, data for " + "the denominator (bottom) of the ratio", +) +parser.add_argument( + "--desc-one", + required=True, + help="Descriptor tag for first set of data (short, for use in subscript)", +) +parser.add_argument( + "--desc-two", + type=str, + required=True, + help="Descriptor tag for second set of data (short, for use in subscript)", +) +parser.add_argument("--outfile", type=str, required=True, help="Output file to save to") +parser.add_argument( + "--ifars", + type=float, + required=True, + nargs="+", + help="IFAR values to plot VT ratio for. Note that the " + "plotted values will be the closest values available " + "from the VT files", +) +parser.add_argument("--log-x", action="store_true", help="Use logarithmic x-axis") +parser.add_argument("--log-y", action="store_true", help="Use logarithmic y-axis") args = parser.parse_args() init_logging(args.verbose) # Warn user if different numbers of files in numerator vs denominator if len(args.vt_files_one) != len(args.vt_files_two): - logging.warning( - 'WATCH OUT! You gave different numbers of One and Two files!') + logging.warning("WATCH OUT! You gave different numbers of One and Two files!") # Load in the first numerator file -with HFile(args.vt_files_one[0], 'r') as ftop_init: +with HFile(args.vt_files_one[0], "r") as ftop_init: # Find the index closest to the given IFAR value - idxs = [np.argmin(np.abs(ftop_init['xvals'][:] - ifv)) - for ifv in args.ifars] - plot_ifars = ftop_init['xvals'][idxs] + idxs = [np.argmin(np.abs(ftop_init["xvals"][:] - ifv)) for ifv in args.ifars] + plot_ifars = ftop_init["xvals"][idxs] # Get binning keys for reference - keys = list(ftop_init['data'].keys()) + keys = list(ftop_init["data"].keys()) # Dicts holding data for total VT and variances vt_top = {k: np.zeros_like(plot_ifars) for k in keys} @@ -71,50 +83,52 @@ vt_bot_errsqlow = {k: np.zeros_like(plot_ifars) for k in keys} # Cycle over inputs for numerator for ftop in args.vt_files_one: - with HFile(ftop, 'r') as f: + with HFile(ftop, "r") as f: # Check the input bins - if list(f['data'].keys()) != keys: + if list(f["data"].keys()) != keys: raise ValueError( - f'keys do not match for the given input files - ' - '{keys} v {list(f["data"].keys())}') + "keys do not match for the given input files - " + '{keys} v {list(f["data"].keys())}' + ) # Add the data for k in keys: - vt_top[k] += f['data'][k][idxs] + vt_top[k] += f["data"][k][idxs] # Variances add over files - vt_top_errsqhi[k] += f['errorhigh'][k][idxs] ** 2. - vt_top_errsqlow[k] += f['errorlow'][k][idxs] ** 2. + vt_top_errsqhi[k] += f["errorhigh"][k][idxs] ** 2.0 + vt_top_errsqlow[k] += f["errorlow"][k][idxs] ** 2.0 # Same for denominator for fbot in args.vt_files_two: - with HFile(fbot, 'r') as f: - if list(f['data'].keys()) != keys: + with HFile(fbot, "r") as f: + if list(f["data"].keys()) != keys: raise ValueError( - f'keys do not match for the given input files - ' - '{keys} v {list(f["data"].keys())}') + "keys do not match for the given input files - " + '{keys} v {list(f["data"].keys())}' + ) for k in keys: - vt_bot[k] += f['data'][k][idxs] - vt_bot_errsqhi[k] += f['errorhigh'][k][idxs] ** 2. - vt_bot_errsqlow[k] += f['errorlow'][k][idxs] ** 2. + vt_bot[k] += f["data"][k][idxs] + vt_bot_errsqhi[k] += f["errorhigh"][k][idxs] ** 2.0 + vt_bot_errsqlow[k] += f["errorlow"][k][idxs] ** 2.0 # make the plot pretty -plt.rc('axes.formatter', limits=[-3, 4]) -plt.rc('figure', dpi=300) +plt.rc("axes.formatter", limits=[-3, 4]) +plt.rc("figure", dpi=300) fig_mi = plt.figure(figsize=(10, 4)) ax_mi = fig_mi.gca() ax_mi.grid(True, zorder=1) # read in labels for the different plotting points -labels = ['$ ' + label.split('\\in')[-1] for label in keys] +labels = ["$ " + label.split("\\in")[-1] for label in keys] # read in the splitting parameter name from the first data set -x_param = r'$' + tuple(keys)[0].split('\\in')[0].strip('$').strip() + r'$' +x_param = r"$" + tuple(keys)[0].split("\\in")[0].strip("$").strip() + r"$" # read in the positions from the labels -xpos = np.array([float(l.split('[')[1].split(',')[0]) for l in labels]) +xpos = np.array([float(l.split("[")[1].split(",")[0]) for l in labels]) # offset different ifars by 1/20th of the mean distance between parameters -try: +try: if args.log_x: xpos_logdiffmean = np.diff(np.log(xpos)).mean() xpos_add_dx = 0.05 * np.ones_like(xpos) * xpos_logdiffmean @@ -126,61 +140,73 @@ except IndexError: xpos_add_dx = 0.05 # set the x ticks to be the positions given in the labels -plt.xticks(xpos, labels, rotation='horizontal') +plt.xticks(xpos, labels, rotation="horizontal") -colors = ['#7b85d4', '#f37738', '#83c995', '#d7369e', '#c4c9d8', '#859795'] +colors = ["#7b85d4", "#f37738", "#83c995", "#d7369e", "#c4c9d8", "#859795"] # loop through each IFAR and plot the VT ratio with error bars for j in range(len(idxs)): data1 = np.array([vt_top[key][j] for key in keys]) errsqhi1 = np.array([vt_top_errsqhi[key][j] for key in keys]) errsqlow1 = np.array([vt_top_errsqlow[key][j] for key in keys]) - + data2 = np.array([vt_bot[key][j] for key in keys]) errsqhi2 = np.array([vt_bot_errsqhi[key][j] for key in keys]) errsqlow2 = np.array([vt_bot_errsqlow[key][j] for key in keys]) ys = data1 / data2 # fractional error propagation - yerr_low = (errsqlow1 / (data1**2.) + errsqlow2 / (data2**2.)) ** 0.5 * ys - yerr_hi = (errsqhi1 / (data1**2.) + errsqhi2 / (data2**2.)) ** 0.5 * ys + yerr_low = (errsqlow1 / (data1**2.0) + errsqlow2 / (data2**2.0)) ** 0.5 * ys + yerr_hi = (errsqhi1 / (data1**2.0) + errsqhi2 / (data2**2.0)) ** 0.5 * ys if args.log_x: - xvals = np.exp(np.log(xpos) + - xpos_add_dx * (j - float(len(args.ifars) - 1) / 2.)) + xvals = np.exp( + np.log(xpos) + xpos_add_dx * (j - float(len(args.ifars) - 1) / 2.0) + ) else: - xvals = xpos + xpos_add_dx * (j - float(len(args.ifars) - 1) / 2.) - ax_mi.errorbar(xvals, ys, - yerr=[yerr_low, yerr_hi], fmt='o', markersize=7, linewidth=5, - label='IFAR = %d yr' % plot_ifars[j], capsize=5, - capthick=2, mec='k', color=colors[j % len(colors)]) + xvals = xpos + xpos_add_dx * (j - float(len(args.ifars) - 1) / 2.0) + ax_mi.errorbar( + xvals, + ys, + yerr=[yerr_low, yerr_hi], + fmt="o", + markersize=7, + linewidth=5, + label="IFAR = %d yr" % plot_ifars[j], + capsize=5, + capthick=2, + mec="k", + color=colors[j % len(colors)], + ) if args.log_x: - plt.xscale('log') + plt.xscale("log") if args.log_y: - plt.yscale('log') -plt.xticks(xpos, labels, rotation='horizontal') + plt.yscale("log") +plt.xticks(xpos, labels, rotation="horizontal") # get the limit of the x axes, and draw a black line in order to highlight # equal comparison xlimits = plt.xlim() -plt.plot(xlimits, [1, 1], 'k', lw=2, zorder=0) -plt.xlim(xlimits) # reassert the x limits so that the plot doesn't expand +plt.plot(xlimits, [1, 1], "k", lw=2, zorder=0) +plt.xlim(xlimits) # reassert the x limits so that the plot doesn't expand -ax_mi.legend(bbox_to_anchor=(0.5, 1.01), ncol=len(args.ifars), - loc='lower center') -ax_mi.get_legend().get_title().set_fontsize('14') +ax_mi.legend(bbox_to_anchor=(0.5, 1.01), ncol=len(args.ifars), loc="lower center") +ax_mi.get_legend().get_title().set_fontsize("14") ax_mi.get_legend().get_frame().set_alpha(0.7) -ax_mi.set_xlabel(x_param, size='large') -ax_mi.set_ylabel(r'$\frac{VT(\mathrm{' + args.desc_one +'})}\ - {VT(\mathrm{' + args.desc_two +'})}$', - size='large') +ax_mi.set_xlabel(x_param, size="large") +ax_mi.set_ylabel( + r"$\frac{VT(\mathrm{" + + args.desc_one + + "})}\ + {VT(\\mathrm{" + + args.desc_two + + "})}$", + size="large", +) plt.tight_layout() -title = f'VT sensitivity comparison between {args.desc_one} and ' \ - f'{args.desc_two}' -save_fig_with_metadata(fig_mi, args.outfile, cmd=' '.join(sys.argv), - title=title) +title = f"VT sensitivity comparison between {args.desc_one} and {args.desc_two}" +save_fig_with_metadata(fig_mi, args.outfile, cmd=" ".join(sys.argv), title=title) plt.close() - diff --git a/bin/plotting/pycbc_plot_waveform b/bin/plotting/pycbc_plot_waveform index 626cdd25c05..311ce6d4c60 100644 --- a/bin/plotting/pycbc_plot_waveform +++ b/bin/plotting/pycbc_plot_waveform @@ -14,117 +14,183 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Plot a waveform in both time and frequency domain. -""" +"""Plot a waveform in both time and frequency domain.""" # imports -import sys, argparse, matplotlib, numpy -matplotlib.use('Agg') +import argparse +import sys + +import matplotlib +import numpy + +matplotlib.use("Agg") import matplotlib.pyplot as plt from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes -from pycbc import waveform, io -from pycbc import results -from pycbc import init_logging, add_common_pycbc_options -from pycbc.types import zeros, complex64 +from pycbc import add_common_pycbc_options, init_logging, io, results, waveform +from pycbc.types import complex64, zeros -parser = argparse.ArgumentParser(usage='', description=__doc__) +parser = argparse.ArgumentParser(usage="", description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--output-file', required=True) -parser.add_argument("--low-frequency-cutoff", type=float, - help="The low frequency cutoff to use for generation.") +parser.add_argument("--output-file", required=True) +parser.add_argument( + "--low-frequency-cutoff", + type=float, + help="The low frequency cutoff to use for generation.", +) # add the approximant argument -waveform.bank.add_approximant_arg(parser, - help="The name of the approximant to use for filtering. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--mass1", type=float, required=True, - help="The mass of the first component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--mass2", type=float, required=True, - help="The mass of the second component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--spin1z", type=float, required=True, - help="The aligned spin of the first component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--spin2z", type=float, required=True, - help="The aligned pin of the second component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--sample-rate", type=float, required=True, - help="Sample rate to use when generating waveform.") -parser.add_argument("--waveform-length", type=float, required=True, - help="Used to set the value of delta-f when either generating" - "in frequency domain, or when FFTing.") +waveform.bank.add_approximant_arg( + parser, + help="The name of the approximant to use for filtering. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--mass1", + type=float, + required=True, + help="The mass of the first component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--mass2", + type=float, + required=True, + help="The mass of the second component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--spin1z", + type=float, + required=True, + help="The aligned spin of the first component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--spin2z", + type=float, + required=True, + help="The aligned pin of the second component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--sample-rate", + type=float, + required=True, + help="Sample rate to use when generating waveform.", +) +parser.add_argument( + "--waveform-length", + type=float, + required=True, + help="Used to set the value of delta-f when either generating" + "in frequency domain, or when FFTing.", +) # Optional orientation arguments -parser.add_argument("--inclination", type=float, default=0, - help="Latitude of the observer w.r.t the source. ") -parser.add_argument("--coa-phase", type=float, default=0, - help="Longitude of the observer w.r.t these source. ") +parser.add_argument( + "--inclination", + type=float, + default=0, + help="Latitude of the observer w.r.t the source. ", +) +parser.add_argument( + "--coa-phase", + type=float, + default=0, + help="Longitude of the observer w.r.t these source. ", +) # Optional arguments for precession -parser.add_argument("--spin1x", type=float, default=0, - help="Non-aligned spin of the first component object. ") -parser.add_argument("--spin1y", type=float, default=0, - help="Non-aligned spin of the first component object. ") -parser.add_argument("--spin2x", type=float, default=0, - help="Non-aligned spin of the second component object. ") -parser.add_argument("--spin2y", type=float, default=0, - help="Non-aligned spin of the second component object. ") -parser.add_argument("--u-val", type=float, default=0, - help="Ratio between h_+ and h_x according to " - "h(t) = h_+ * u_val + h_x. Only needed in the case of " - "waveforms where h_+ and h_x are not related by a " - "simple phase shift.") +parser.add_argument( + "--spin1x", + type=float, + default=0, + help="Non-aligned spin of the first component object. ", +) +parser.add_argument( + "--spin1y", + type=float, + default=0, + help="Non-aligned spin of the first component object. ", +) +parser.add_argument( + "--spin2x", + type=float, + default=0, + help="Non-aligned spin of the second component object. ", +) +parser.add_argument( + "--spin2y", + type=float, + default=0, + help="Non-aligned spin of the second component object. ", +) +parser.add_argument( + "--u-val", + type=float, + default=0, + help="Ratio between h_+ and h_x according to " + "h(t) = h_+ * u_val + h_x. Only needed in the case of " + "waveforms where h_+ and h_x are not related by a " + "simple phase shift.", +) # Optional arguments for eccentricity -parser.add_argument("--eccentricity", type=float, default=0, - help="Eccentricity of the orbit. ") -parser.add_argument("--rel-anomaly", type=float, default=0, - help="Relativistic anomaly. ") +parser.add_argument( + "--eccentricity", type=float, default=0, help="Eccentricity of the orbit. " +) +parser.add_argument( + "--rel-anomaly", type=float, default=0, help="Relativistic anomaly. " +) # Other optional arguments -parser.add_argument("--taper-template", choices=["start","end","startend"], - help="For time-domain approximants, taper the start and/or" - " end of the waveform before FFTing.") +parser.add_argument( + "--taper-template", + choices=["start", "end", "startend"], + help="For time-domain approximants, taper the start and/or" + " end of the waveform before FFTing.", +) # Plotting options -parser.add_argument('--plot-title', - help="If given, use this as the plot title") -parser.add_argument('--plot-caption', - help="If given, use this as the plot caption") +parser.add_argument("--plot-title", help="If given, use this as the plot title") +parser.add_argument("--plot-caption", help="If given, use this as the plot caption") opt = parser.parse_args() init_logging(opt.verbose) -delta_f = 1. / opt.waveform_length -delta_t = 1. / opt.sample_rate +delta_f = 1.0 / opt.waveform_length +delta_t = 1.0 / opt.sample_rate tlen = int(opt.waveform_length * opt.sample_rate) flen = tlen // 2 + 1 tmp_params = io.WaveformArray.from_kwargs( - mass1=opt.mass1, - mass2=opt.mass2, - spin1x=opt.spin1x, - spin1y=opt.spin1y, - spin1z=opt.spin1z, - spin2x=opt.spin2x, - spin2y=opt.spin2y, - spin2z=opt.spin2z, - inclination=opt.inclination, - coa_phase=opt.coa_phase, - eccentricity=opt.eccentricity, - rel_anomaly=opt.rel_anomaly) + mass1=opt.mass1, + mass2=opt.mass2, + spin1x=opt.spin1x, + spin1y=opt.spin1y, + spin1z=opt.spin1z, + spin2x=opt.spin2x, + spin2y=opt.spin2y, + spin2z=opt.spin2z, + inclination=opt.inclination, + coa_phase=opt.coa_phase, + eccentricity=opt.eccentricity, + rel_anomaly=opt.rel_anomaly, +) # Deal with parameter-dependent approximant string -approximant = waveform.bank.parse_approximant_arg(opt.approximant, - tmp_params)[0] +approximant = waveform.bank.parse_approximant_arg(opt.approximant, tmp_params)[0] # This is a hack. We don't have a two-pol SPATmplt, so use TaylorF2 -if approximant == 'SPAtmplt': - approximant = 'TaylorF2' +if approximant == "SPAtmplt": + approximant = "TaylorF2" -hp, hc = waveform.get_two_pol_waveform_filter(zeros(flen, dtype=complex64), - zeros(flen, dtype=complex64), - tmp_params[0], approximant=approximant, - taper=opt.taper_template, - f_lower=opt.low_frequency_cutoff, - delta_f=delta_f, delta_t=delta_t) +hp, hc = waveform.get_two_pol_waveform_filter( + zeros(flen, dtype=complex64), + zeros(flen, dtype=complex64), + tmp_params[0], + approximant=approximant, + taper=opt.taper_template, + f_lower=opt.low_frequency_cutoff, + delta_f=delta_f, + delta_t=delta_t, +) # 0.1s (post-merger) added for safety tmplt_length = hp.length_in_time + 0.1 @@ -142,46 +208,67 @@ template_td = template.to_timeseries() template_td.roll(int(-post_merger_length * opt.sample_rate)) fig = plt.figure() -ax = plt.subplot(2,1,1) +ax = plt.subplot(2, 1, 1) ax.plot(template_td.sample_times, template_td) -ax.set_xlim([template_td.sample_times.max() - tmplt_length*1.1, - template_td.sample_times.max()]) -ax.set_xlabel('Time (s)') -ax.set_ylabel('$h(t)$') +ax.set_xlim( + [ + template_td.sample_times.max() - tmplt_length * 1.1, + template_td.sample_times.max(), + ] +) +ax.set_xlabel("Time (s)") +ax.set_ylabel("$h(t)$") -if tmplt_length > 1.: +if tmplt_length > 1.0: # Make inset zoom on merger - axins = zoomed_inset_axes(ax, 0.5, loc=3,bbox_to_anchor=(0.12, 0.69), - bbox_transform=ax.figure.transFigure) + axins = zoomed_inset_axes( + ax, + 0.5, + loc=3, + bbox_to_anchor=(0.12, 0.69), + bbox_transform=ax.figure.transFigure, + ) x_zoom_fac = tmplt_length / 0.2 axins.plot(template_td.sample_times * x_zoom_fac, template_td) - axins.set_xlim([(template_td.sample_times.max() - 0.2)*x_zoom_fac, - x_zoom_fac*template_td.sample_times.max()]) + axins.set_xlim( + [ + (template_td.sample_times.max() - 0.2) * x_zoom_fac, + x_zoom_fac * template_td.sample_times.max(), + ] + ) axins.get_xaxis().set_ticks([]) axins.get_yaxis().set_ticks([]) -plt.subplot(2,1,2) +plt.subplot(2, 1, 2) plt.plot(template.sample_frequencies, abs(template)) -plt.xlabel('Frequency (Hz)') -plt.xlim([0,f_length*1.1]) -plt.ylabel('abs( $\\tilde{h}(f)$ )') +plt.xlabel("Frequency (Hz)") +plt.xlim([0, f_length * 1.1]) +plt.ylabel("abs( $\\tilde{h}(f)$ )") plt.tight_layout() if opt.plot_title is None: - opt.plot_title = 'Waveform plot' -if opt.plot_caption is None and tmplt_length > 1.: - opt.plot_caption = ("The first plot represents the template waveform" - "in the time domain with the inset showing a" - "zoom-up at merger. The second plot represents" - "the template waveform in the frequency domain.") -if opt.plot_caption is None and tmplt_length <=1 : - opt.plot_caption = ("The first plot represents the template waveform" - "in the time domain. The second plot represents" - "the template waveform in the frequency domain.") - -results.save_fig_with_metadata(fig, opt.output_file, - cmd=' '.join(sys.argv), fig_kwds={'dpi': 150}, - title=opt.plot_title, - caption=opt.plot_caption) + opt.plot_title = "Waveform plot" +if opt.plot_caption is None and tmplt_length > 1.0: + opt.plot_caption = ( + "The first plot represents the template waveform" + "in the time domain with the inset showing a" + "zoom-up at merger. The second plot represents" + "the template waveform in the frequency domain." + ) +if opt.plot_caption is None and tmplt_length <= 1: + opt.plot_caption = ( + "The first plot represents the template waveform" + "in the time domain. The second plot represents" + "the template waveform in the frequency domain." + ) + +results.save_fig_with_metadata( + fig, + opt.output_file, + cmd=" ".join(sys.argv), + fig_kwds={"dpi": 150}, + title=opt.plot_title, + caption=opt.plot_caption, +) diff --git a/bin/population/pycbc_multiifo_pastro b/bin/population/pycbc_multiifo_pastro index 71c17dcfdfa..bcbbf72c905 100644 --- a/bin/population/pycbc_multiifo_pastro +++ b/bin/population/pycbc_multiifo_pastro @@ -8,108 +8,185 @@ # option) any later version. -import os import argparse import logging -import numpy as np +import os +import numpy as np from matplotlib import use -use('Agg') -from matplotlib import rcParams + +use("Agg") from matplotlib import pyplot as plt +from matplotlib import rcParams -from pycbc import init_logging, add_common_pycbc_options +from pycbc import add_common_pycbc_options, init_logging from pycbc.population import fgmc_functions as utils from pycbc.population import fgmc_laguerre as fgmcl from pycbc.population import fgmc_plots - -rcParams.update({'axes.labelsize': 12, - 'font.size': 12, - 'legend.fontsize': 12, - 'xtick.labelsize': 12, - 'ytick.labelsize': 12, - 'text.usetex': True, - }) +rcParams.update( + { + "axes.labelsize": 12, + "font.size": 12, + "legend.fontsize": 12, + "xtick.labelsize": 12, + "ytick.labelsize": 12, + "text.usetex": True, + } +) parser = argparse.ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--diagnostic-plots', action='store_true', default=False, - help='Make extra diagnostic scatter plots') -parser.add_argument('--full-data-files', nargs='+', - help='hdf5 file(s) (space-separated) containing zerolag events' - 'and (decimated/undecimated) timeslide background') -parser.add_argument('--bank-files', nargs='+', - help='hdf5 file(s) (space-separated) containing bank parameters') -parser.add_argument('--injection-files', nargs='+', - help='hdf5 file(s) (space-separated) containing injection event ' - 'data in hdfinjfind format') -parser.add_argument('--rank-injections', nargs='*', - help='hdf5 file(s) (space-separated) containing injection events ' - 'in statmap format, to treat as extra zerolag events. ' - 'Optional') -parser.add_argument('--prior-event-info', help='.txt data file with gps, stat, ifar, ' - 'ln bayes factor for prior events. Only ln BF is used in FGMC ' - 'calculation') -parser.add_argument('--network', nargs='*', default=['H1', 'L1', 'V1'], - help='All ifos to consider in determining coinc times/types ' - '(space-separated)') -parser.add_argument('--coinc-times', nargs='+', - help='Types of coincident analysis time to read in ' - '(space-separated). Ex. H1L1 H1L1V1') -parser.add_argument('--stat-threshold', type=float, - help='Only analyze events and injections above this ' - 'value') -parser.add_argument('--bg-bin-width', default=None, - help='Width of stat bins to estimate background density, eg 0.25') -parser.add_argument('--inj-bin-width', default=None, - help='Width of bins to estimate injection density, eg 0.4') -parser.add_argument('--bg-log-bin-width', type=float, default=None, - help='Width of log spaced bins to estimate background' - ' density, eg 0.02') -parser.add_argument('--inj-log-bin-width', type=float, default=None, - help='Width of log spaced bins to estimate injection' - ' density, eg 0.04') -parser.add_argument('--min-mchirp', type=float, required=True, - help='Minimum template chirp mass to accept') -parser.add_argument('--max-mchirp', type=float, required=True, - help='Maximum template chirp mass to accept') -parser.add_argument('--power-law-prior', default=-0.5, type=float, - help='Power-law prior on foreground count parameter, ' - 'default -0.5 ("Jeffreys")') -parser.add_argument('--laguerre-degree', type=int, default=20, - help='Degree of polynomial used for generalized L-G quadrature') -parser.add_argument('--ntop', type=int, default=15, - help='Number of loudest events listed, default 15') -parser.add_argument('--p-astro-txt', help='Save GPS and p_astro in txt file') -parser.add_argument('--p-astro-inj', help='Save injection GPS / p_astro in txt file.' - ' Optional') -parser.add_argument('--json-tag', help='If given, save event json files with tag ' - 'in name') -parser.add_argument('--json-min-ifar', type=float, help='IFAR threshold to save ' - 'json files') -parser.add_argument('--vt-mean', type=float, - help='Sensitivity estimate conjugate to merger rate: if ' - 'provided, rate estimate will be output in the same units') -parser.add_argument('--vt-err', type=float, - help='Standard error on vt-mean value, assumed in same units') -parser.add_argument('--rate-units', default=r'Gpc^-3 yr^-1', - help=r'Sensitivity units for rate output, default 1/(Gpc^3 yr)') -parser.add_argument('--plot-max-stat', type=float, default=30., - help='Maximum background stat value to plot') -parser.add_argument('--plot-bg-hist', action='store_true', - help='Make diagnostic background PDF plots') -parser.add_argument('--plot-inj-hist', action='store_true', - help='Make diagnostic injection PDF plots') -parser.add_argument('--plot-dir', required=True, help='Destination for plots, ' - 'must include final "/"') -parser.add_argument('--plot-tag', help='String for plot filenames') +parser.add_argument( + "--diagnostic-plots", + action="store_true", + default=False, + help="Make extra diagnostic scatter plots", +) +parser.add_argument( + "--full-data-files", + nargs="+", + help="hdf5 file(s) (space-separated) containing zerolag events" + "and (decimated/undecimated) timeslide background", +) +parser.add_argument( + "--bank-files", + nargs="+", + help="hdf5 file(s) (space-separated) containing bank parameters", +) +parser.add_argument( + "--injection-files", + nargs="+", + help="hdf5 file(s) (space-separated) containing injection event " + "data in hdfinjfind format", +) +parser.add_argument( + "--rank-injections", + nargs="*", + help="hdf5 file(s) (space-separated) containing injection events " + "in statmap format, to treat as extra zerolag events. " + "Optional", +) +parser.add_argument( + "--prior-event-info", + help=".txt data file with gps, stat, ifar, " + "ln bayes factor for prior events. Only ln BF is used in FGMC " + "calculation", +) +parser.add_argument( + "--network", + nargs="*", + default=["H1", "L1", "V1"], + help="All ifos to consider in determining coinc times/types (space-separated)", +) +parser.add_argument( + "--coinc-times", + nargs="+", + help="Types of coincident analysis time to read in " + "(space-separated). Ex. H1L1 H1L1V1", +) +parser.add_argument( + "--stat-threshold", + type=float, + help="Only analyze events and injections above this value", +) +parser.add_argument( + "--bg-bin-width", + default=None, + help="Width of stat bins to estimate background density, eg 0.25", +) +parser.add_argument( + "--inj-bin-width", + default=None, + help="Width of bins to estimate injection density, eg 0.4", +) +parser.add_argument( + "--bg-log-bin-width", + type=float, + default=None, + help="Width of log spaced bins to estimate background density, eg 0.02", +) +parser.add_argument( + "--inj-log-bin-width", + type=float, + default=None, + help="Width of log spaced bins to estimate injection density, eg 0.04", +) +parser.add_argument( + "--min-mchirp", + type=float, + required=True, + help="Minimum template chirp mass to accept", +) +parser.add_argument( + "--max-mchirp", + type=float, + required=True, + help="Maximum template chirp mass to accept", +) +parser.add_argument( + "--power-law-prior", + default=-0.5, + type=float, + help='Power-law prior on foreground count parameter, default -0.5 ("Jeffreys")', +) +parser.add_argument( + "--laguerre-degree", + type=int, + default=20, + help="Degree of polynomial used for generalized L-G quadrature", +) +parser.add_argument( + "--ntop", type=int, default=15, help="Number of loudest events listed, default 15" +) +parser.add_argument("--p-astro-txt", help="Save GPS and p_astro in txt file") +parser.add_argument( + "--p-astro-inj", help="Save injection GPS / p_astro in txt file. Optional" +) +parser.add_argument( + "--json-tag", help="If given, save event json files with tag in name" +) +parser.add_argument( + "--json-min-ifar", type=float, help="IFAR threshold to save json files" +) +parser.add_argument( + "--vt-mean", + type=float, + help="Sensitivity estimate conjugate to merger rate: if " + "provided, rate estimate will be output in the same units", +) +parser.add_argument( + "--vt-err", + type=float, + help="Standard error on vt-mean value, assumed in same units", +) +parser.add_argument( + "--rate-units", + default=r"Gpc^-3 yr^-1", + help=r"Sensitivity units for rate output, default 1/(Gpc^3 yr)", +) +parser.add_argument( + "--plot-max-stat", + type=float, + default=30.0, + help="Maximum background stat value to plot", +) +parser.add_argument( + "--plot-bg-hist", action="store_true", help="Make diagnostic background PDF plots" +) +parser.add_argument( + "--plot-inj-hist", action="store_true", help="Make diagnostic injection PDF plots" +) +parser.add_argument( + "--plot-dir", required=True, help='Destination for plots, must include final "/"' +) +parser.add_argument("--plot-tag", help="String for plot filenames") args = parser.parse_args() -if not args.plot_dir.endswith('/'): - args.plot_dir += '/' +if not args.plot_dir.endswith("/"): + args.plot_dir += "/" if not os.path.exists(args.plot_dir): raise RuntimeError("Output dir %s doesn't exist!" % args.plot_dir) if args.bg_bin_width is not None and args.bg_log_bin_width is not None: @@ -125,8 +202,11 @@ MANY_BANKS = True if len(args.bank_files) != nchunks: MANY_BANKS = False if len(args.bank_files) != 1: - raise RuntimeError('Either need the same number of banks as chunks, '\ - + str(nchunks) + ', or exactly 1 bank!') + raise RuntimeError( + "Either need the same number of banks as chunks, " + + str(nchunks) + + ", or exactly 1 bank!" + ) else: args.bank_files = args.bank_files * nchunks @@ -134,35 +214,37 @@ if len(args.bank_files) != nchunks: tot_exp_bg = 0 zl_inputs = zip(args.full_data_files, args.bank_files) -logging.info('Setting up zerolag') -fg = utils.ForegroundEvents(args, args.coinc_times, - bin_lo=args.min_mchirp, bin_hi=args.max_mchirp) +logging.info("Setting up zerolag") +fg = utils.ForegroundEvents( + args, args.coinc_times, bin_lo=args.min_mchirp, bin_hi=args.max_mchirp +) -logging.info('Setting up background') -bg = utils.BackgroundEventRate(args, args.coinc_times, - bin_lo=args.min_mchirp, bin_hi=args.max_mchirp) +logging.info("Setting up background") +bg = utils.BackgroundEventRate( + args, args.coinc_times, bin_lo=args.min_mchirp, bin_hi=args.max_mchirp +) for fd, b in zl_inputs: - logging.info('Adding bank info from ' + b) + logging.info("Adding bank info from " + b) fg.add_bank(b) fg.filter_templates() bg.add_bank(b) bg.filter_templates() - logging.info('Adding event info from ' + fd) + logging.info("Adding event info from " + fd) fg.add_zerolag(fd) bg.add_background(fd) if args.diagnostic_plots: - plt.plot(fg.stat, fg.masspars, '+b', ms=3) + plt.plot(fg.stat, fg.masspars, "+b", ms=3) plt.grid(True) plt.ylim(ymax=1.1 * fg.masspars.max()) - plt.xlabel('Rank statistic') - plt.ylabel('Template param (mchirp)') - plt.savefig(args.plot_dir + 'fg_stat_vs_mchirp.png') + plt.xlabel("Rank statistic") + plt.ylabel("Template param (mchirp)") + plt.savefig(args.plot_dir + "fg_stat_vs_mchirp.png") plt.semilogy() plt.ylim(0.9 * fg.masspars.min(), 1.1 * fg.masspars.max()) - plt.savefig(args.plot_dir + 'fg_stat_vs_logmchirp.png') + plt.savefig(args.plot_dir + "fg_stat_vs_logmchirp.png") plt.close() if args.plot_bg_hist: @@ -170,16 +252,17 @@ if args.plot_bg_hist: # normalize and count expected bg events bg.get_norms() -logging.info('Total expected bg count ' + str(bg.norm)) -logging.info('Actual zerolag event count ' + str(len(fg.stat))) +logging.info("Total expected bg count " + str(bg.norm)) +logging.info("Actual zerolag event count " + str(len(fg.stat))) # Read signal (injection) data inj_inputs = zip(args.injection_files, args.bank_files, args.full_data_files) -sg = utils.SignalEventRate(args, args.coinc_times, - bin_lo=args.min_mchirp, bin_hi=args.max_mchirp) +sg = utils.SignalEventRate( + args, args.coinc_times, bin_lo=args.min_mchirp, bin_hi=args.max_mchirp +) for jf, b, fd in inj_inputs: - logging.info('Adding inj info from' + jf) + logging.info("Adding inj info from" + jf) sg.add_bank(b) sg.filter_templates() sg.add_injections(jf, fd) @@ -188,7 +271,7 @@ if args.plot_inj_hist: sg.plot_inj() sg.get_norms() -logging.info('Total number of inj used ' + str(sg.norm)) +logging.info("Total number of inj used " + str(sg.norm)) fg.get_bg_pdf(bg) fg.get_sg_pdf(sg) @@ -197,119 +280,153 @@ allfgbg = fg.sg_pdf - fg.bg_pdf # make diagnostic plots if args.diagnostic_plots: - print('ln Bayes factors from', min(allfgbg), max(allfgbg)) - col = {'H1L1': 'r', 'H1L1V1': 'b', 'H1V1': 'm', 'L1V1': 'g'} + print("ln Bayes factors from", min(allfgbg), max(allfgbg)) + col = {"H1L1": "r", "H1L1V1": "b", "H1V1": "m", "L1V1": "g"} for cty in col: in_type = fg.ctype == cty - plt.semilogx(fg.stat[in_type], fg.sg_pdf[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.semilogx( + fg.stat[in_type], + fg.sg_pdf[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('Signal ln PDF') - plt.savefig(args.plot_dir + 'fg_stat_vs_sgpdf.png') + plt.xlabel("Rank statistic") + plt.ylabel("Signal ln PDF") + plt.savefig(args.plot_dir + "fg_stat_vs_sgpdf.png") plt.close() for cty in col: in_type = fg.ctype == cty - plt.semilogx(fg.stat[in_type], fg.bg_pdf[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.semilogx( + fg.stat[in_type], + fg.bg_pdf[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('Noise ln PDF') - plt.savefig(args.plot_dir + 'fg_stat_vs_bgpdf.png') + plt.xlabel("Rank statistic") + plt.ylabel("Noise ln PDF") + plt.savefig(args.plot_dir + "fg_stat_vs_bgpdf.png") plt.close() for cty in col: in_type = fg.ctype == cty - plt.plot(fg.stat[in_type], allfgbg[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.plot( + fg.stat[in_type], + allfgbg[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('fg/bg Bayes factor') + plt.xlabel("Rank statistic") + plt.ylabel("fg/bg Bayes factor") plt.xlim(xmax=args.plot_max_stat) # zoom to see not-super-loud events - plt.savefig(args.plot_dir + 'fg_stat_vs_fgbg_bf.png') + plt.savefig(args.plot_dir + "fg_stat_vs_fgbg_bf.png") plt.semilogx() plt.xlim(args.stat_threshold, 1.1 * fg.stat.max()) - plt.savefig(args.plot_dir + 'fg_logstat_vs_fgbg_bf.png') + plt.savefig(args.plot_dir + "fg_logstat_vs_fgbg_bf.png") plt.close() for cty in col: in_type = fg.ctype == cty - plt.loglog(fg.stat[in_type], fg.ifar[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.loglog( + fg.stat[in_type], + fg.ifar[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('IFAR (yr)') - plt.savefig(args.plot_dir + 'fg_stat_vs_ifar.png') + plt.xlabel("Rank statistic") + plt.ylabel("IFAR (yr)") + plt.savefig(args.plot_dir + "fg_stat_vs_ifar.png") plt.close() # add in prior event info if args.prior_event_info: - logging.warning('Adding events from', args.prior_event_info) + logging.warning("Adding events from", args.prior_event_info) prior_info = np.genfromtxt(args.prior_event_info, names=True) - allfgbg = np.append(allfgbg, prior_info['ln_bayes_factor']) + allfgbg = np.append(allfgbg, prior_info["ln_bayes_factor"]) else: # no prior events! - prior_info = {'stat': np.array([]), - 'ifar': np.array([]), - 'gps': np.array([]), - 'bayes_factor': np.array([])} + prior_info = { + "stat": np.array([]), + "ifar": np.array([]), + "gps": np.array([]), + "bayes_factor": np.array([]), + } # do the calculation -foreground_count_dist = \ - fgmcl.count_posterior(allfgbg, laguerre_n=args.laguerre_degree, - Lambda0=bg.norm, name='foreground count posterior', - prior=args.power_law_prior) +foreground_count_dist = fgmcl.count_posterior( + allfgbg, + laguerre_n=args.laguerre_degree, + Lambda0=bg.norm, + name="foreground count posterior", + prior=args.power_law_prior, +) # Roulet et al. argue that information on rate scales with sum (p_astro^2) -logging.info('sum of p_astro^2', - ((1 - foreground_count_dist.p_bg(allfgbg)) ** 2.).sum()) +logging.info( + "sum of p_astro^2", ((1 - foreground_count_dist.p_bg(allfgbg)) ** 2.0).sum() +) fgmc_plots.odds_summary( args, - np.append(fg.stat, prior_info['stat']), - np.append(fg.ifar, prior_info['ifar']), + np.append(fg.stat, prior_info["stat"]), + np.append(fg.ifar, prior_info["ifar"]), foreground_count_dist.p_bg(allfgbg), args.ntop, - times=np.append(fg.gpst, prior_info['gps']), - name='foreground events', - plot_extensions=['.png'], + times=np.append(fg.gpst, prior_info["gps"]), + name="foreground events", + plot_extensions=[".png"], ) args.plot_limits = None cmed, cminus, cplus = fgmc_plots.dist_summary( args, foreground_count_dist, - plot_extensions=['.png'], - middle='median', + plot_extensions=[".png"], + middle="median", credible_intervals=[0.90], ) # if VT and error are provided, convert signal PDF to rate PDF if args.vt_mean: if args.vt_err is None: - logging.warning('NB : no VT error provided, will use 0!') - args.vt_err = 0. - logging.info('Rate in units of ' + args.rate_units + ' ...') - print('Rate without VT uncertainty :') - print('%.1f_{-%.1f}^{+%.1f}' % ( - cmed / args.vt_mean, abs(cminus) / args.vt_mean, cplus / args.vt_mean)) + logging.warning("NB : no VT error provided, will use 0!") + args.vt_err = 0.0 + logging.info("Rate in units of " + args.rate_units + " ...") + print("Rate without VT uncertainty :") + print( + "%.1f_{-%.1f}^{+%.1f}" + % (cmed / args.vt_mean, abs(cminus) / args.vt_mean, cplus / args.vt_mean) + ) import rate_functions + fracerr = args.vt_err / args.vt_mean # attr to arbitrarily rescale counts in rate calculation : here, don't rescale - foreground_count_dist.scale = 1. + foreground_count_dist.scale = 1.0 # we want a Jeffreys type prior: this is already supplied by the default value # of the --power-law-prior option (see above) so leave a uniform prior on rate - rate_distribution = \ - rate_functions.rate_posterior(foreground_count_dist, args.vt_mean, - fracerr, unit=args.rate_units, texunit='') - median, minus, plus = fgmc_plots.dist_summary(args, rate_distribution, - plot_extensions=None, credible_intervals=[0.9]) - print('Rate with VT uncertainty via rate_posterior :') - print('%.1f_{-%.1f}^{+%.1f}' % (median, abs(minus), plus)) + rate_distribution = rate_functions.rate_posterior( + foreground_count_dist, args.vt_mean, fracerr, unit=args.rate_units, texunit="" + ) + median, minus, plus = fgmc_plots.dist_summary( + args, rate_distribution, plot_extensions=None, credible_intervals=[0.9] + ) + print("Rate with VT uncertainty via rate_posterior :") + print("%.1f_{-%.1f}^{+%.1f}" % (median, abs(minus), plus)) elif args.vt_err: raise RuntimeError("Can't supply VT error without a VT mean!") @@ -317,8 +434,9 @@ elif args.vt_err: # if injections are to be evaluated 'like zerolag', read them in if args.rank_injections: assert args.p_astro_inj is not None - inj = utils.ForegroundEvents(args, args.coinc_times, - bin_lo=args.min_mchirp, bin_hi=args.max_mchirp) + inj = utils.ForegroundEvents( + args, args.coinc_times, bin_lo=args.min_mchirp, bin_hi=args.max_mchirp + ) # if using the same single input file for all results if not MANY_BANKS: @@ -328,11 +446,11 @@ if args.rank_injections: assert len(args.rank_injections) == len(args.bank_files) banks = args.bank_files for i, b in zip(args.rank_injections, banks): - logging.info('Adding bank info from ' + b) + logging.info("Adding bank info from " + b) inj.add_bank(b) inj.filter_templates() - logging.info('Adding event info from ' + i) + logging.info("Adding event info from " + i) inj.add_zerolag(i) # get densities for injections as if zerolag @@ -342,66 +460,82 @@ if args.rank_injections: injfgbg = inj.sg_pdf - inj.bg_pdf if args.diagnostic_plots: - col = {'H1L1': 'r', 'H1L1V1': 'b', 'H1V1': 'm', 'L1V1': 'g'} + col = {"H1L1": "r", "H1L1V1": "b", "H1V1": "m", "L1V1": "g"} for cty in col: in_type = inj.ctype == cty - plt.semilogx(inj.stat[in_type], inj.sg_pdf[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.semilogx( + inj.stat[in_type], + inj.sg_pdf[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('Signal ln PDF') - plt.savefig(args.plot_dir + 'inj_stat_vs_sgpdf.png') + plt.xlabel("Rank statistic") + plt.ylabel("Signal ln PDF") + plt.savefig(args.plot_dir + "inj_stat_vs_sgpdf.png") plt.close() for cty in col: in_type = inj.ctype == cty - plt.semilogx(inj.stat[in_type], inj.bg_pdf[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.semilogx( + inj.stat[in_type], + inj.bg_pdf[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('Noise ln PDF') - plt.savefig(args.plot_dir + 'inj_stat_vs_bgpdf.png') + plt.xlabel("Rank statistic") + plt.ylabel("Noise ln PDF") + plt.savefig(args.plot_dir + "inj_stat_vs_bgpdf.png") plt.close() for cty in col: in_type = inj.ctype == cty - plt.plot(inj.stat[in_type], injfgbg[in_type], col[cty]+'+', - ms=3, alpha=0.6, label=cty) + plt.plot( + inj.stat[in_type], + injfgbg[in_type], + col[cty] + "+", + ms=3, + alpha=0.6, + label=cty, + ) plt.grid(True) plt.legend() - plt.xlabel('Rank statistic') - plt.ylabel('fg/bg Bayes factor') + plt.xlabel("Rank statistic") + plt.ylabel("fg/bg Bayes factor") # zoom to see not-super-loud events plt.xlim(args.stat_threshold, args.plot_max_stat) - plt.savefig(args.plot_dir + 'inj_stat_vs_fgbg_bf.png') + plt.savefig(args.plot_dir + "inj_stat_vs_fgbg_bf.png") plt.semilogx() plt.xlim(args.stat_threshold, 1.1 * inj.stat.max()) - plt.savefig(args.plot_dir + 'inj_logstat_vs_fgbg_bf.png') + plt.savefig(args.plot_dir + "inj_logstat_vs_fgbg_bf.png") plt.close() # evaluate p_astro using zerolag signal rate inj_pbg = foreground_count_dist.p_bg(injfgbg) # make a special txt format output file - np.savetxt(args.p_astro_inj, - np.column_stack((inj.gpst, - 1. / inj.ifar, - inj.ifar, - inj.stat, - 1. - inj_pbg)), - fmt=['%.3f', '%e', '%e', '%6f', '%.6f'], - delimiter=',', - header='geocent_end_time,FAR,IFAR,detection_statistic,p_astro') + np.savetxt( + args.p_astro_inj, + np.column_stack((inj.gpst, 1.0 / inj.ifar, inj.ifar, inj.stat, 1.0 - inj_pbg)), + fmt=["%.3f", "%e", "%e", "%6f", "%.6f"], + delimiter=",", + header="geocent_end_time,FAR,IFAR,detection_statistic,p_astro", + ) if args.diagnostic_plots: - plt.loglog(inj.stat, (1. - inj_pbg) / inj_pbg, 'k.', ms=3) + plt.loglog(inj.stat, (1.0 - inj_pbg) / inj_pbg, "k.", ms=3) plt.grid(True) - plt.xlabel('Ranking statistic') - plt.ylabel(r'$P_1/P_0$') - plt.xlim(0.99 * inj.stat.min(), 2. * args.plot_max_stat) - plt.savefig(args.plot_dir + 'injection_events_odds.png') + plt.xlabel("Ranking statistic") + plt.ylabel(r"$P_1/P_0$") + plt.xlim(0.99 * inj.stat.min(), 2.0 * args.plot_max_stat) + plt.savefig(args.plot_dir + "injection_events_odds.png") plt.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/population/pycbc_population_plots b/bin/population/pycbc_population_plots index ec43dca514b..3360a4c473b 100644 --- a/bin/population/pycbc_population_plots +++ b/bin/population/pycbc_population_plots @@ -26,108 +26,136 @@ __version__ = "0.0" __date__ = "31.10.2017" import argparse -from numpy import logaddexp, log, newaxis, expm1 + import numpy as np +import pycbc.version +import scipy.stats as ss from matplotlib import cm from matplotlib import pyplot as plt - -import scipy.stats as ss +from numpy import expm1, log, logaddexp, newaxis import pycbc -import pycbc.version from pycbc.io.hdf import HFile from pycbc.population import rates_functions as rf # Parse command line -parser = argparse.ArgumentParser( - description=__doc__ -) +parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--prior-samples', required=True, - help="File containing rate prior samples") -parser.add_argument('--posterior-samples', nargs='+', required=True, - help="Files(s) containing samples for the rate posterior.") -parser.add_argument('--population-models', nargs='+', required=True, - help="Models to which posteriors belong ('lnm', 'imf', 'bns').") -parser.add_argument('--output-rates', dest='rate_file', required=True, - help="File saving all the rate posteriors into one.") -parser.add_argument('--rates-figure', required=True, - help="Name of file to draw rates prior and posterior.") -parser.add_argument('--pastro-figure', required=True, - help="Name of file to save p_astro figure.") -parser.add_argument('--plot-labels', nargs='+', - required=True, help="Labels for the population models.") +parser.add_argument( + "--prior-samples", required=True, help="File containing rate prior samples" +) +parser.add_argument( + "--posterior-samples", + nargs="+", + required=True, + help="Files(s) containing samples for the rate posterior.", +) +parser.add_argument( + "--population-models", + nargs="+", + required=True, + help="Models to which posteriors belong ('lnm', 'imf', 'bns').", +) +parser.add_argument( + "--output-rates", + dest="rate_file", + required=True, + help="File saving all the rate posteriors into one.", +) +parser.add_argument( + "--rates-figure", + required=True, + help="Name of file to draw rates prior and posterior.", +) +parser.add_argument( + "--pastro-figure", required=True, help="Name of file to save p_astro figure." +) +parser.add_argument( + "--plot-labels", nargs="+", required=True, help="Labels for the population models." +) opts = parser.parse_args() pycbc.init_logging(opts.verbose) -assert len(opts.posterior_samples) == len(opts.population_models), \ - "Unequal number of posterior files and population models!" +assert len(opts.posterior_samples) == len(opts.population_models), ( + "Unequal number of posterior files and population models!" +) -#Save rate posteriors in one file -with HFile(opts.rate_file, 'w') as out: +# Save rate posteriors in one file +with HFile(opts.rate_file, "w") as out: for fname, model in zip(opts.posterior_samples, opts.population_models): - f = HFile(fname, "r") - Rf = f[model + '/Rf'][:] - Lf = f[model + '/Lf'][:] + Rf = f[model + "/Rf"][:] + Lf = f[model + "/Lf"][:] pl = out.create_group(model) - pl.create_dataset('Rf', data=Rf, compression='gzip') - pl.create_dataset('Lf', data=Lf, compression='gzip') + pl.create_dataset("Rf", data=Rf, compression="gzip") + pl.create_dataset("Lf", data=Lf, compression="gzip") f.close() # Make prior/posterior plot -- estimate p_astro p_astro = [] plt.figure() -color=iter(cm.rainbow(np.linspace(0, 1, len(opts.population_models)))) +color = iter(cm.rainbow(np.linspace(0, 1, len(opts.population_models)))) mods = zip(opts.posterior_samples, opts.population_models, opts.plot_labels) f = HFile(opts.prior_samples, "r") for fpost, model, lbl in mods: - c = next(color) fpo = HFile(fpost, "r") - Rfpr, Rfpo = f[model + '/Rf'][:], fpo[model + '/Rf'][:] + Rfpr, Rfpo = f[model + "/Rf"][:], fpo[model + "/Rf"][:] prior_alpha, prior_mu, prior_sigma = rf.fit(Rfpr) post_alpha, post_mu, post_sigma = rf.fit(Rfpo) log_R = np.log(Rfpr) xs = np.linspace(min(log_R), max(log_R), 200) - plt.plot(np.exp(xs), ss.skewnorm.pdf(xs, prior_alpha, prior_mu, - prior_sigma), '--', label=lbl + ' Prior', color=c) - plt.plot(np.exp(xs), ss.skewnorm.pdf(xs, post_alpha, post_m, - post_sigma), label=lbl + ' Posterior', color=c) - - Lfpo, Lbpo = fpo[model + '/Lf'][:], fpo[model + '/Lb'][:] - log_fg_ratios = fpo['data/log_fg_bg_ratio'][:] - - log_pastros = logaddexp.reduce(log(Lfpo[:, newaxis]) +\ - log_fg_ratios[newaxis,:] - logaddexp(log(Lfpo[:, newaxis]) +\ - log_fg_ratios[newaxis,:], log(Lbpo[:, newaxis])), axis=0) -\ - log(Lfpo.shape[0]) + plt.plot( + np.exp(xs), + ss.skewnorm.pdf(xs, prior_alpha, prior_mu, prior_sigma), + "--", + label=lbl + " Prior", + color=c, + ) + plt.plot( + np.exp(xs), + ss.skewnorm.pdf(xs, post_alpha, post_m, post_sigma), + label=lbl + " Posterior", + color=c, + ) + + Lfpo, Lbpo = fpo[model + "/Lf"][:], fpo[model + "/Lb"][:] + log_fg_ratios = fpo["data/log_fg_bg_ratio"][:] + + log_pastros = logaddexp.reduce( + log(Lfpo[:, newaxis]) + + log_fg_ratios[newaxis, :] + - logaddexp( + log(Lfpo[:, newaxis]) + log_fg_ratios[newaxis, :], log(Lbpo[:, newaxis]) + ), + axis=0, + ) - log(Lfpo.shape[0]) p_astro.append(1 + expm1(np.sort(log_pastros)[::-1])) fpo.close() f.close() -plt.xscale('log') -plt.xlabel(r'$R$ ($\mathrm{Gpc}^{-3} \, \mathrm{yr}^{-1}$)') -plt.ylabel(r'$RP(R)$') -plt.legend(loc='best') +plt.xscale("log") +plt.xlabel(r"$R$ ($\mathrm{Gpc}^{-3} \, \mathrm{yr}^{-1}$)") +plt.ylabel(r"$RP(R)$") +plt.legend(loc="best") plt.savefig(opts.rates_figure) plt.figure() -color=iter(cm.rainbow(np.linspace(0, 1, len(opts.population_models)))) +color = iter(cm.rainbow(np.linspace(0, 1, len(opts.population_models)))) for pas, lbl in zip(p_astro, opts.plot_labels): c = next(color) - plt.plot(log_fg_ratios, 1 - pas, '.', label = lbl, color = c) + plt.plot(log_fg_ratios, 1 - pas, ".", label=lbl, color=c) -plt.xlabel(r'$\log p(x\mid f)/p(x\mid b)$') -plt.ylabel(r'$1-p_\mathrm{astro}$') -plt.legend(loc='best') -plt.yscale('log') +plt.xlabel(r"$\log p(x\mid f)/p(x\mid b)$") +plt.ylabel(r"$1-p_\mathrm{astro}$") +plt.legend(loc="best") +plt.yscale("log") plt.savefig(opts.pastro_figure) print(p_astro) diff --git a/bin/population/pycbc_population_rates b/bin/population/pycbc_population_rates index 004099af44e..8b1f445855d 100644 --- a/bin/population/pycbc_population_rates +++ b/bin/population/pycbc_population_rates @@ -33,55 +33,104 @@ __version__ = "0.0" __date__ = "31.10.2017" import argparse + import numpy as np +import pycbc.version import pycbc -from pycbc.population import scale_injections as si -from pycbc.population import rates_functions as rf from pycbc.io.hdf import HFile -import pycbc.version - +from pycbc.population import rates_functions as rf +from pycbc.population import scale_injections as si # Parse command line parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--sim-files', nargs='+', required=True, help="List of " - "simulation files to estimate the sensitive volume") -parser.add_argument('--m_dist', nargs='+', required=True, help="Specify " - "the mass distribution for the simulations") -parser.add_argument('--s_dist', nargs='+',required=True, help="Specify " - "the spin distribution for the simulations") -parser.add_argument('--d_dist', nargs='+', required=True, help="Specify " - "the distance distribution for the simulations") -parser.add_argument('--bank-file', required=True, help="File containing " - "template bank used in the search.") -parser.add_argument('--statmap-file', required=True, help="File containing " - "trigger information") -parser.add_argument('--prior-samples', required=True, help="File storing " - "samples of prior for the analysis - posterior from the previous analysis") -parser.add_argument('--output-file', required=True, help="Name of the " - "output file for saving rate posteriors.") -parser.add_argument('--thr-var', required=False, default='stat', - help="Variable used to define threshold.") -parser.add_argument('--thr-val', type=float, required=False, default=8.0, - help="Value of threshold.") -parser.add_argument('--population-model', required=True, - help = 'Population model defined in rates_functions') -parser.add_argument('--min-mass', type=float, required=True, - help="Minimum mass of the compact object.") -parser.add_argument('--max-mass', type=float, required=True, - help="Maximum mass of the compact object.") -parser.add_argument('--max-mtotal', type=float, required=True, - help="Maximum total mass of the binary.") -parser.add_argument('--min-tmplt-mchirp', type=float, required=True, - help="Minimum chirp mass of the template considered " - "for trigger identification.") -parser.add_argument('--max-tmplt-mchirp', type=float, required=True, - help="Maximum chirp mass in the template considered " - "for trigger identification.") -parser.add_argument('--calibration-error', dest='cal_err', type=float, - required=False, default=3.0, help="Percentage calibration" - " errors in measurement of distance.") +parser.add_argument( + "--sim-files", + nargs="+", + required=True, + help="List of simulation files to estimate the sensitive volume", +) +parser.add_argument( + "--m_dist", + nargs="+", + required=True, + help="Specify the mass distribution for the simulations", +) +parser.add_argument( + "--s_dist", + nargs="+", + required=True, + help="Specify the spin distribution for the simulations", +) +parser.add_argument( + "--d_dist", + nargs="+", + required=True, + help="Specify the distance distribution for the simulations", +) +parser.add_argument( + "--bank-file", + required=True, + help="File containing template bank used in the search.", +) +parser.add_argument( + "--statmap-file", required=True, help="File containing trigger information" +) +parser.add_argument( + "--prior-samples", + required=True, + help="File storing " + "samples of prior for the analysis - posterior from the previous analysis", +) +parser.add_argument( + "--output-file", + required=True, + help="Name of the output file for saving rate posteriors.", +) +parser.add_argument( + "--thr-var", + required=False, + default="stat", + help="Variable used to define threshold.", +) +parser.add_argument( + "--thr-val", type=float, required=False, default=8.0, help="Value of threshold." +) +parser.add_argument( + "--population-model", + required=True, + help="Population model defined in rates_functions", +) +parser.add_argument( + "--min-mass", type=float, required=True, help="Minimum mass of the compact object." +) +parser.add_argument( + "--max-mass", type=float, required=True, help="Maximum mass of the compact object." +) +parser.add_argument( + "--max-mtotal", type=float, required=True, help="Maximum total mass of the binary." +) +parser.add_argument( + "--min-tmplt-mchirp", + type=float, + required=True, + help="Minimum chirp mass of the template considered for trigger identification.", +) +parser.add_argument( + "--max-tmplt-mchirp", + type=float, + required=True, + help="Maximum chirp mass in the template considered for trigger identification.", +) +parser.add_argument( + "--calibration-error", + dest="cal_err", + type=float, + required=False, + default=3.0, + help="Percentage calibration errors in measurement of distance.", +) opts = parser.parse_args() @@ -89,81 +138,90 @@ pycbc.init_logging(opts.verbose) path = opts.output_file -assert opts.min_tmplt_mchirp < opts.max_tmplt_mchirp, \ - "Minimum chirp mass should be less than the maximum chirp mass" +assert opts.min_tmplt_mchirp < opts.max_tmplt_mchirp, ( + "Minimum chirp mass should be less than the maximum chirp mass" +) # Read the simulation files -injections = si.read_injections(opts.sim_files, - opts.m_dist, opts.s_dist, opts.d_dist) +injections = si.read_injections(opts.sim_files, opts.m_dist, opts.s_dist, opts.d_dist) # Read the chirp-mass samples -- Imported from rates_function -if opts.population_model == 'imf': +if opts.population_model == "imf": mchirp_sampler = rf.mchirp_sampler_imf prob = rf.prob_imf -elif opts.population_model == 'lnm': +elif opts.population_model == "lnm": mchirp_sampler = rf.mchirp_sampler_lnm prob = rf.prob_lnm -elif opts.population_model == 'bns': +elif opts.population_model == "bns": mchirp_sampler = rf.mchirp_sampler_flat prob = rf.prob_flat # Estimate the rates and make supporting plots -vt = si.estimate_vt(injections, mchirp_sampler, prob, - thr_var = opts.thr_var, - thr_val = opts.thr_val, - min_mass = opts.min_mass, - max_mass = opts.max_mass, - max_mtotal = opts.max_mtotal) - -vol_time, sig_vt = vt['VT'], vt['VT_err'] -inj_falloff = vt['thr_falloff'] +vt = si.estimate_vt( + injections, + mchirp_sampler, + prob, + thr_var=opts.thr_var, + thr_val=opts.thr_val, + min_mass=opts.min_mass, + max_mass=opts.max_mass, + max_mtotal=opts.max_mtotal, +) + +vol_time, sig_vt = vt["VT"], vt["VT_err"] +inj_falloff = vt["thr_falloff"] # Include the calibration uncertainity vol, vol_err, cal_err = vol_time, sig_vt, opts.cal_err -sigma_w_cal_uncrt = np.sqrt((3*cal_err/100.)**2 + (vol_err/vol)**2) - -#Sabe background data and coincidences -all_bkg, coincs = rf.save_bkg_falloff(opts.statmap_file, opts.bank_file, - path, opts.thr_val, opts.min_tmplt_mchirp, - opts.max_tmplt_mchirp) - -#Load background data and coincidences/ make some plots +sigma_w_cal_uncrt = np.sqrt((3 * cal_err / 100.0) ** 2 + (vol_err / vol) ** 2) + +# Sabe background data and coincidences +all_bkg, coincs = rf.save_bkg_falloff( + opts.statmap_file, + opts.bank_file, + path, + opts.thr_val, + opts.min_tmplt_mchirp, + opts.max_tmplt_mchirp, +) + +# Load background data and coincidences/ make some plots bg_l, bg_h, bg_counts = all_bkg bg_bins = np.append(bg_l, bg_h[-1]) -#fg_stats = np.concatenate([inj_falloff[dist] for dist in distrs]) +# fg_stats = np.concatenate([inj_falloff[dist] for dist in distrs]) fg_stats = inj_falloff[inj_falloff > opts.thr_val] fg_bins = np.logspace(np.log10(opts.thr_val), np.log10(np.max(fg_stats)), 101) log_fg_ratios = rf.log_rho_fgmc(coincs, fg_stats, fg_bins) log_fg_ratios -= rf.log_rho_bg(coincs, bg_bins, bg_counts) -#Load prior samples and fit a skew-log-normal to it +# Load prior samples and fit a skew-log-normal to it with HFile(opts.prior_samples, "r") as f: - R = np.array(f[opts.population_model+'/Rf']) + R = np.array(f[opts.population_model + "/Rf"]) alpha, mu, sigma = rf.fit(R) -#Estimate rates +# Estimate rates rate_samples = {} log_R = np.log(R) -mu_log_vt = np.log(vol_time/1e9) +mu_log_vt = np.log(vol_time / 1e9) sigma_log_vt = sigma_w_cal_uncrt Rf_samp = rf.skew_lognormal_samples(alpha, mu, sigma, min(log_R), max(log_R)) -rate_samples['Rf'], rate_samples['Lf'], rate_samples['Lb'] = \ -rf.fgmc(log_fg_ratios, mu_log_vt, sigma_log_vt, Rf_samp, max(fg_stats)) +rate_samples["Rf"], rate_samples["Lf"], rate_samples["Lb"] = rf.fgmc( + log_fg_ratios, mu_log_vt, sigma_log_vt, Rf_samp, max(fg_stats) +) -rate_post = rate_samples['Rf'] +rate_post = rate_samples["Rf"] r50, r95, r05 = np.percentile(rate_post, [50, 95, 5]) -#Save rate posteriors -with HFile(opts.output_file, 'w') as out: - +# Save rate posteriors +with HFile(opts.output_file, "w") as out: pl = out.create_group(opts.population_model) - pl.create_dataset('Lf', data=rate_samples['Lf'], compression='gzip') - pl.create_dataset('Lb', data=rate_samples['Lb'], compression='gzip') - pl.create_dataset('Rf', data=rate_samples['Rf'], compression='gzip') + pl.create_dataset("Lf", data=rate_samples["Lf"], compression="gzip") + pl.create_dataset("Lb", data=rate_samples["Lb"], compression="gzip") + pl.create_dataset("Rf", data=rate_samples["Rf"], compression="gzip") - d = out.create_group('data') - d.create_dataset('log_fg_bg_ratio', data=log_fg_ratios, compression='gzip') - d.create_dataset('newsnr', data=coincs, compression='gzip') + d = out.create_group("data") + d.create_dataset("log_fg_bg_ratio", data=log_fg_ratios, compression="gzip") + d.create_dataset("newsnr", data=coincs, compression="gzip") diff --git a/bin/pycbc_banksim b/bin/pycbc_banksim index 55320f5f517..fc7552fc992 100644 --- a/bin/pycbc_banksim +++ b/bin/pycbc_banksim @@ -16,49 +16,60 @@ """Calculate the fitting factors of simulated signals with a template bank.""" - import logging -from tqdm import tqdm -from numpy import complex64, array from argparse import ArgumentParser from math import ceil, log -from igwn_ligolw import utils as ligolw_utils from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils +from numpy import array, complex64 +from tqdm import tqdm -from pycbc.pnutils import mass1_mass2_to_mchirp_eta -from pycbc.pnutils import mass1_mass2_to_tau0_tau3 -from pycbc.pnutils import nearest_larger_binary_number -from pycbc.waveform import get_td_waveform, get_fd_waveform, td_approximants, fd_approximants +import pycbc.fft +import pycbc.psd +import pycbc.scheme +import pycbc.strain from pycbc import DYN_RANGE_FAC -from pycbc.types import FrequencySeries, TimeSeries, zeros, complex_same_precision_as +from pycbc.detector import overhead_antenna_pattern as generate_fplus_fcross from pycbc.filter import match, sigmasq from pycbc.io.ligolw import LIGOLWContentHandler -import pycbc.psd, pycbc.scheme, pycbc.fft, pycbc.strain -from pycbc.detector import overhead_antenna_pattern as generate_fplus_fcross -from pycbc.waveform import TemplateBank - +from pycbc.pnutils import ( + mass1_mass2_to_mchirp_eta, + mass1_mass2_to_tau0_tau3, + nearest_larger_binary_number, +) +from pycbc.types import FrequencySeries, TimeSeries, complex_same_precision_as, zeros +from pycbc.waveform import ( + TemplateBank, + fd_approximants, + get_fd_waveform, + get_td_waveform, + td_approximants, +) ## Remove the need for these functions ######################################## - + + def generate_detector_strain(template_params, h_plus, h_cross): - latitude = 0 - longitude = 0 - polarization = 0 + latitude = 0 + longitude = 0 + polarization = 0 - if hasattr(template_params, 'latitude'): + if hasattr(template_params, "latitude"): latitude = template_params.latitude - if hasattr(template_params, 'longitude'): + if hasattr(template_params, "longitude"): longitude = template_params.longitude - if hasattr(template_params, 'polarization'): + if hasattr(template_params, "polarization"): polarization = template_params.polarization f_plus, f_cross = generate_fplus_fcross(longitude, latitude, polarization) return h_plus * f_plus + h_cross * f_cross + def make_padded_frequency_series(vec, filter_N=None, delta_f=None): - """Convert vec (TimeSeries or FrequencySeries) to a FrequencySeries. If + """ + Convert vec (TimeSeries or FrequencySeries) to a FrequencySeries. If filter_N and/or delta_f are given, the output will take those values. If not told otherwise the code will attempt to pad a timeseries first such that the waveform will not wraparound. However, if delta_f is specified to be @@ -66,14 +77,15 @@ def make_padded_frequency_series(vec, filter_N=None, delta_f=None): """ if filter_N is None: power = ceil(log(len(vec), 2)) + 1 - N = 2 ** power + N = 2**power else: N = filter_N n = N // 2 + 1 if isinstance(vec, FrequencySeries): - vectilde = FrequencySeries(zeros(n, dtype=complex_same_precision_as(vec)), - delta_f=1.0, copy=False) + vectilde = FrequencySeries( + zeros(n, dtype=complex_same_precision_as(vec)), delta_f=1.0, copy=False + ) if len(vectilde) < len(vec): cplen = len(vectilde) else: @@ -86,7 +98,7 @@ def make_padded_frequency_series(vec, filter_N=None, delta_f=None): # and increase if necessary curr_length = len(vec) new_length = int(nearest_larger_binary_number(curr_length)) - while new_length * vec.delta_t < 1./delta_f: + while new_length * vec.delta_t < 1.0 / delta_f: new_length = new_length * 2 vec.resize(new_length) # Then convert to frequencyseries @@ -99,108 +111,196 @@ def make_padded_frequency_series(vec, filter_N=None, delta_f=None): i_delta_f = v_tilde.get_delta_f() v_tilde = v_tilde.numpy() df_ratio = int(delta_f / i_delta_f) - n_freq_len = int((n-1) * df_ratio +1) - assert(n <= len(v_tilde)) + n_freq_len = int((n - 1) * df_ratio + 1) + assert n <= len(v_tilde) df_ratio = int(delta_f / i_delta_f) v_tilde = v_tilde[:n_freq_len:df_ratio] vectilde = FrequencySeries(v_tilde, delta_f=delta_f, dtype=complex64) - return FrequencySeries(vectilde * DYN_RANGE_FAC, delta_f=delta_f, - dtype=complex64) + return FrequencySeries(vectilde * DYN_RANGE_FAC, delta_f=delta_f, dtype=complex64) + -def get_waveform(approximant, phase_order, amplitude_order, spin_order, - wf_params, start_frequency, sample_rate, length, - filter_rate): +def get_waveform( + approximant, + phase_order, + amplitude_order, + spin_order, + wf_params, + start_frequency, + sample_rate, + length, + filter_rate, +): if type(approximant) is not str: - approximant = approximant.decode('utf-8') + approximant = approximant.decode("utf-8") delta_f = filter_rate / length if approximant in fd_approximants(): - hp, hc = get_fd_waveform(wf_params, approximant=approximant, - phase_order=phase_order, delta_f=delta_f, - spin_order=spin_order, - f_lower=start_frequency, - amplitude_order=amplitude_order) + hp, hc = get_fd_waveform( + wf_params, + approximant=approximant, + phase_order=phase_order, + delta_f=delta_f, + spin_order=spin_order, + f_lower=start_frequency, + amplitude_order=amplitude_order, + ) hvec = generate_detector_strain(wf_params, hp, hc) elif approximant in td_approximants(): - hp, hc = get_td_waveform(wf_params, - approximant=approximant, - phase_order=phase_order, - spin_order=spin_order, - delta_t=1./sample_rate, - f_lower=start_frequency, - amplitude_order=amplitude_order) - if hasattr(wf_params, 'taper'): - hp = hp.taper_timeseries(location=wf_params.taper, - tapermethod=wf_params.get('taper_method', 'lal'), - taper_window=wf_params.get('taper_window')) - hc = hc.taper_timeseries(location=wf_params.taper, - tapermethod=wf_params.get('taper_method', 'lal'), - taper_window=wf_params.get('taper_window')) + hp, hc = get_td_waveform( + wf_params, + approximant=approximant, + phase_order=phase_order, + spin_order=spin_order, + delta_t=1.0 / sample_rate, + f_lower=start_frequency, + amplitude_order=amplitude_order, + ) + if hasattr(wf_params, "taper"): + hp = hp.taper_timeseries( + location=wf_params.taper, + tapermethod=wf_params.get("taper_method", "lal"), + taper_window=wf_params.get("taper_window"), + ) + hc = hc.taper_timeseries( + location=wf_params.taper, + tapermethod=wf_params.get("taper_method", "lal"), + taper_window=wf_params.get("taper_window"), + ) hvec = generate_detector_strain(wf_params, hp, hc) return make_padded_frequency_series(hvec, filter_N, delta_f=delta_f) + aprs = sorted(list(set(td_approximants() + fd_approximants()))) -#File output Settings +# File output Settings parser = ArgumentParser(description=__doc__) -parser.add_argument("--match-file", dest="out_file", metavar="FILE", - required=True, help="File to output match results") +parser.add_argument( + "--match-file", + dest="out_file", + metavar="FILE", + required=True, + help="File to output match results", +) pycbc.add_common_pycbc_options(parser) -#Template Settings -parser.add_argument("--template-file", dest="bank_file", metavar="FILE", - required=True, help="File specifying the template parameters, " - "either in HDF or LIGOLW XML format.") -parser.add_argument("--total-mass-divide", type=float, - help="Total mass to switch from --template-approximant to " - "--highmass-approximant.") -parser.add_argument("--highmass-approximant", choices=aprs, - help="Waveform approximant for highmass templates.") -parser.add_argument("--template-approximant", choices=aprs, required=True, - help="Waveform approximant for templates") -parser.add_argument("--template-phase-order", default=-1, type=int, - help="PN order to use for the template phase") -parser.add_argument("--template-amplitude-order", default=-1, type=int, - help="PN order to use for the template amplitude") -parser.add_argument("--template-spin-order", default=-1, type=int, - help="PN order to use for the template spin terms") -parser.add_argument("--template-start-frequency", type=float, - help="Starting frequency for templates [Hz]") -parser.add_argument("--template-sample-rate", type=float, - help="Sample rate for templates [Hz]") - -#Signal Settings -parser.add_argument("--signal-file", dest="sim_file", metavar="FILE", - required=True, help="SimInspiral or SnglInspiral XML file " - "containing the signal parameters") -parser.add_argument("--signal-approximant", choices=aprs, required=True, - help="Waveform approximant for signals") -parser.add_argument("--signal-phase-order", default=-1, type=int, - help="PN order to use for the signal phase") -parser.add_argument("--signal-spin-order", default=-1, type=int, - help="PN order to use for the signal spin terms") -parser.add_argument("--signal-amplitude-order", default=-1, type=int, - help="PN order to use for the signal amplitude") -parser.add_argument("--signal-start-frequency", type=float, - help="Starting frequency for signals [Hz]") -parser.add_argument("--signal-sample-rate", type=float, - help="Sample rate for signals [Hz]") -parser.add_argument("--use-sky-location", action='store_true', - help="Inject into a theoretical detector at the celestial " - "North pole of a non-rotating Earth rather than overhead") - -#Filtering Settings -parser.add_argument('--filter-low-frequency-cutoff', metavar='FREQ', type=float, - required=True, help='low frequency cutoff of matched filter') -parser.add_argument("--filter-sample-rate", type=float, required=True, - help="Filter sample rate [Hz]") -parser.add_argument("--filter-signal-length", type=int, required=True, - help="Length of signal for filtering, shoud be longer " - "than all waveforms and include some padding") +# Template Settings +parser.add_argument( + "--template-file", + dest="bank_file", + metavar="FILE", + required=True, + help="File specifying the template parameters, either in HDF or LIGOLW XML format.", +) +parser.add_argument( + "--total-mass-divide", + type=float, + help="Total mass to switch from --template-approximant to --highmass-approximant.", +) +parser.add_argument( + "--highmass-approximant", + choices=aprs, + help="Waveform approximant for highmass templates.", +) +parser.add_argument( + "--template-approximant", + choices=aprs, + required=True, + help="Waveform approximant for templates", +) +parser.add_argument( + "--template-phase-order", + default=-1, + type=int, + help="PN order to use for the template phase", +) +parser.add_argument( + "--template-amplitude-order", + default=-1, + type=int, + help="PN order to use for the template amplitude", +) +parser.add_argument( + "--template-spin-order", + default=-1, + type=int, + help="PN order to use for the template spin terms", +) +parser.add_argument( + "--template-start-frequency", + type=float, + help="Starting frequency for templates [Hz]", +) +parser.add_argument( + "--template-sample-rate", type=float, help="Sample rate for templates [Hz]" +) + +# Signal Settings +parser.add_argument( + "--signal-file", + dest="sim_file", + metavar="FILE", + required=True, + help="SimInspiral or SnglInspiral XML file containing the signal parameters", +) +parser.add_argument( + "--signal-approximant", + choices=aprs, + required=True, + help="Waveform approximant for signals", +) +parser.add_argument( + "--signal-phase-order", + default=-1, + type=int, + help="PN order to use for the signal phase", +) +parser.add_argument( + "--signal-spin-order", + default=-1, + type=int, + help="PN order to use for the signal spin terms", +) +parser.add_argument( + "--signal-amplitude-order", + default=-1, + type=int, + help="PN order to use for the signal amplitude", +) +parser.add_argument( + "--signal-start-frequency", type=float, help="Starting frequency for signals [Hz]" +) +parser.add_argument( + "--signal-sample-rate", type=float, help="Sample rate for signals [Hz]" +) +parser.add_argument( + "--use-sky-location", + action="store_true", + help="Inject into a theoretical detector at the celestial " + "North pole of a non-rotating Earth rather than overhead", +) + +# Filtering Settings +parser.add_argument( + "--filter-low-frequency-cutoff", + metavar="FREQ", + type=float, + required=True, + help="low frequency cutoff of matched filter", +) +parser.add_argument( + "--filter-sample-rate", type=float, required=True, help="Filter sample rate [Hz]" +) +parser.add_argument( + "--filter-signal-length", + type=int, + required=True, + help="Length of signal for filtering, shoud be longer " + "than all waveforms and include some padding", +) # add PSD options pycbc.psd.insert_psd_option_group(parser, output=False) @@ -208,25 +308,34 @@ pycbc.psd.insert_psd_option_group(parser, output=False) # Insert the data reading options pycbc.strain.insert_strain_option_group(parser) -#hardware support +# hardware support pycbc.scheme.insert_processing_option_group(parser) pycbc.fft.insert_fft_option_group(parser) -#Restricted maximization -parser.add_argument("--mchirp-window", type=str, metavar="FRACTION", - help="Ignore templates whose chirp mass deviates from " - "signal's one more than given fraction. Provide two " - "comma separated numbers to have different bounds " - "above and below the signal's, with below bound " - "listed first.") -parser.add_argument("--tau0-window", type=float, metavar="TIME", default=None, - help="Ignore templates whose Newtonian order chirp time " - "(tau0) varies from the signals by more than the " - "supplied amount. If option is not provided no " - "window on tau0 is used. The " - "filter-low-frequency-cutoff is used to calculate " - "the value of tau0 for all cases. Provided in units " - "of seconds.") +# Restricted maximization +parser.add_argument( + "--mchirp-window", + type=str, + metavar="FRACTION", + help="Ignore templates whose chirp mass deviates from " + "signal's one more than given fraction. Provide two " + "comma separated numbers to have different bounds " + "above and below the signal's, with below bound " + "listed first.", +) +parser.add_argument( + "--tau0-window", + type=float, + metavar="TIME", + default=None, + help="Ignore templates whose Newtonian order chirp time " + "(tau0) varies from the signals by more than the " + "supplied amount. If option is not provided no " + "window on tau0 is used. The " + "filter-low-frequency-cutoff is used to calculate " + "the value of tau0 for all cases. Provided in units " + "of seconds.", +) options = parser.parse_args() @@ -238,32 +347,40 @@ if options.psd_estimation: pycbc.strain.verify_strain_options(options, parser) if options.total_mass_divide and options.highmass_approximant is None: - parser.error("You must provide a highmass-approximant if you want total-mass-divide.") + parser.error( + "You must provide a highmass-approximant if you want total-mass-divide." + ) if options.mchirp_window is None: + def outside_mchirp_window(template_mchirp, signal_mchirp): return False -elif ',' in options.mchirp_window: +elif "," in options.mchirp_window: # asymmetric chirp mass window mchirp_window_lower = float(options.mchirp_window.split(",")[0]) mchirp_window_upper = float(options.mchirp_window.split(",")[1]) + def outside_mchirp_window(template_mchirp, signal_mchirp): delta = (template_mchirp - signal_mchirp) / signal_mchirp return delta > mchirp_window_upper or -delta > mchirp_window_lower else: # symmetric chirp mass window mchirp_window = float(options.mchirp_window) + def outside_mchirp_window(template_mchirp, signal_mchirp): - return abs(signal_mchirp - template_mchirp) > \ - (mchirp_window * signal_mchirp) + return abs(signal_mchirp - template_mchirp) > (mchirp_window * signal_mchirp) + if options.tau0_window is None: + def outside_tau0_window(template_tau0, signal_tau0, window): return False else: + def outside_tau0_window(template_tau0, signal_tau0, window): return abs(signal_tau0 - template_tau0) > window + # If we are going to use h(t) to estimate a PSD we need h(t) if options.psd_estimation: logging.info("Obtaining h(t) for PSD generation") @@ -282,14 +399,15 @@ else: ctx = pycbc.scheme.from_cli(options) -logging.info('Reading template bank') +logging.info("Reading template bank") temp_bank = TemplateBank(options.bank_file) template_table = temp_bank.table logging.info(" %d templates", len(template_table)) -logging.info('Reading simulation list') -indoc = ligolw_utils.load_filename(options.sim_file, False, - contenthandler=LIGOLWContentHandler) +logging.info("Reading simulation list") +indoc = ligolw_utils.load_filename( + options.sim_file, False, contenthandler=LIGOLWContentHandler +) try: signal_table = lsctables.SimInspiralTable.get_table(indoc) except ValueError: @@ -303,12 +421,17 @@ filter_n = filter_N // 2 + 1 filter_delta_f = 1.0 / float(options.filter_signal_length) logging.info("Reading and Interpolating PSD") -psd = pycbc.psd.from_cli(options, filter_n, filter_delta_f, - options.filter_low_frequency_cutoff, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC, - precision='single') - -with ctx: +psd = pycbc.psd.from_cli( + options, + filter_n, + filter_delta_f, + options.filter_low_frequency_cutoff, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, + precision="single", +) + +with ctx: pycbc.fft.from_cli(options) logging.info("Pregenerating Signals") @@ -320,19 +443,23 @@ with ctx: prog = tqdm(total=len(signal_table), disable=(not options.verbose)) for index, signal_params in enumerate(signal_table): prog.update(1) - if not options.use_sky_location and hasattr(signal_params, 'latitude'): - signal_params.latitude = 0. - signal_params.longitude = 0. - stilde = get_waveform(options.signal_approximant, - options.signal_phase_order, - options.signal_amplitude_order, - options.signal_spin_order, - signal_params, - options.signal_start_frequency, - signal_sample_rate, - filter_N, options.filter_sample_rate) - s_norm = sigmasq(stilde, psd=psd, - low_frequency_cutoff=options.filter_low_frequency_cutoff) + if not options.use_sky_location and hasattr(signal_params, "latitude"): + signal_params.latitude = 0.0 + signal_params.longitude = 0.0 + stilde = get_waveform( + options.signal_approximant, + options.signal_phase_order, + options.signal_amplitude_order, + options.signal_spin_order, + signal_params, + options.signal_start_frequency, + signal_sample_rate, + filter_N, + options.filter_sample_rate, + ) + s_norm = sigmasq( + stilde, psd=psd, low_frequency_cutoff=options.filter_low_frequency_cutoff + ) stilde /= psd signals.append((stilde, s_norm, [], signal_params)) sig_m1.append(signal_params.mass1) @@ -340,15 +467,17 @@ with ctx: prog.close() sig_m1 = array(sig_m1) sig_m2 = array(sig_m2) - sig_tau0, _ = mass1_mass2_to_tau0_tau3(sig_m1, sig_m2, - options.filter_low_frequency_cutoff) + sig_tau0, _ = mass1_mass2_to_tau0_tau3( + sig_m1, sig_m2, options.filter_low_frequency_cutoff + ) sig_mchirp, _ = mass1_mass2_to_mchirp_eta(sig_m1, sig_m2) logging.info("Calculating Mchirp and Tau0") template_m1 = array([tp.mass1 for tp in template_table]) template_m2 = array([tp.mass2 for tp in template_table]) - template_tau0, _ = mass1_mass2_to_tau0_tau3(template_m1, template_m2, - options.filter_low_frequency_cutoff) + template_tau0, _ = mass1_mass2_to_tau0_tau3( + template_m1, template_m2, options.filter_low_frequency_cutoff + ) template_mchirp, _ = mass1_mass2_to_mchirp_eta(template_m1, template_m2) logging.info("Calculating Overlaps") @@ -364,21 +493,24 @@ with ctx: if f_lower < options.filter_low_frequency_cutoff: # Not entirely clear what to do here? if not flow_warned: - logging.warning("Template's flower is smaller than " - "--filter-low-frequency-cutoff. Raising " - "flower of template to match.") - flow_warned=True + logging.warning( + "Template's flower is smaller than " + "--filter-low-frequency-cutoff. Raising " + "flower of template to match." + ) + flow_warned = True f_lower = options.filter_low_frequency_cutoff h_norm = htilde = None for sidx, (stilde, s_norm, matches, signal_params) in enumerate(signals): # Check if we need to look at this check_logic = stilde is None - check_logic |= outside_tau0_window(template_tau0[index], - sig_tau0[sidx], - options.tau0_window) - check_logic |= outside_mchirp_window(template_mchirp[index], - sig_mchirp[sidx]) + check_logic |= outside_tau0_window( + template_tau0[index], sig_tau0[sidx], options.tau0_window + ) + check_logic |= outside_mchirp_window( + template_mchirp[index], sig_mchirp[sidx] + ) if check_logic: matches.append(0) continue @@ -390,25 +522,37 @@ with ctx: # However, while we are still using the high-mass divide # in XML banks, this must be retained. try: - this_approximant = template_params['approximant'] + this_approximant = template_params["approximant"] except: this_approximant = options.template_approximant - if options.total_mass_divide is not None and (template_params.mass1+template_params.mass2) >= options.total_mass_divide: + if ( + options.total_mass_divide is not None + and (template_params.mass1 + template_params.mass2) + >= options.total_mass_divide + ): this_approximant = options.highmass_approximant - htilde = get_waveform(this_approximant, - options.template_phase_order, - options.template_amplitude_order, - options.template_spin_order, - template_params, - options.template_start_frequency, - template_sample_rate, - filter_N, options.filter_sample_rate) + htilde = get_waveform( + this_approximant, + options.template_phase_order, + options.template_amplitude_order, + options.template_spin_order, + template_params, + options.template_start_frequency, + template_sample_rate, + filter_N, + options.filter_sample_rate, + ) h_norm = sigmasq(htilde, psd=psd, low_frequency_cutoff=f_lower) - o, i = match(htilde, stilde, v1_norm=h_norm, v2_norm=s_norm, - low_frequency_cutoff=f_lower) + o, i = match( + htilde, + stilde, + v1_norm=h_norm, + v2_norm=s_norm, + low_frequency_cutoff=f_lower, + ) matches.append(o) prog.close() diff --git a/bin/pycbc_banksim_combine_banks b/bin/pycbc_banksim_combine_banks index 4d82c5ebdd7..8bc963000c7 100644 --- a/bin/pycbc_banksim_combine_banks +++ b/bin/pycbc_banksim_combine_banks @@ -23,11 +23,12 @@ Concatenation of injections is done separately. """ import argparse + import numpy as np import pycbc -__author__ = "Ian Harry " +__author__ = "Ian Harry " __program__ = "pycbc_banksim_combine_banks" # Read command line options @@ -35,34 +36,45 @@ _desc = __doc__[1:] parser = argparse.ArgumentParser(description=_desc) pycbc.add_common_pycbc_options(parser) -parser.add_argument("-I", "--input-files", nargs='+', - help="Explicit list of input files.") -parser.add_argument("-o", "--output-file", required=True, - help="Output file name") +parser.add_argument( + "-I", "--input-files", nargs="+", help="Explicit list of input files." +) +parser.add_argument("-o", "--output-file", required=True, help="Output file name") options = parser.parse_args() pycbc.init_logging(options.verbose) -dtypef = np.dtype([('match', np.float64), ('bank', np.str_, 256), - ('bank_i', np.int32), ('sim', np.str_, 256), - ('sim_i', np.int32), ('sigmasq', np.float64)]) +dtypef = np.dtype( + [ + ("match", np.float64), + ("bank", np.str_, 256), + ("bank_i", np.int32), + ("sim", np.str_, 256), + ("sim_i", np.int32), + ("sigmasq", np.float64), + ] +) -matches=[] +matches = [] maxmatch = [] for fil in options.input_files: matches.append(np.loadtxt(fil, dtype=dtypef)) # It is possible for the input files to only contain a single injection # if the user has split the injections many times. -if np.array(matches, dtype=dtypef)['match'].ndim == 1: - index = np.array(matches, dtype=dtypef)['match'].argmax() +if np.array(matches, dtype=dtypef)["match"].ndim == 1: + index = np.array(matches, dtype=dtypef)["match"].argmax() maxmatch.append(matches[index]) else: - indices = np.array(matches, dtype=dtypef)['match'].argmax(0) + indices = np.array(matches, dtype=dtypef)["match"].argmax(0) for i, j in enumerate(indices): maxmatch.append(matches[j][i]) maxmatch = np.array(maxmatch, dtype=dtypef) -np.savetxt(options.output_file, maxmatch, - fmt=('%5.5f', '%s', '%i', '%s', '%i', '%5.5f'), delimiter=' ') +np.savetxt( + options.output_file, + maxmatch, + fmt=("%5.5f", "%s", "%i", "%s", "%i", "%5.5f"), + delimiter=" ", +) diff --git a/bin/pycbc_banksim_match_combine b/bin/pycbc_banksim_match_combine index 922ca04f0d6..4554f147447 100644 --- a/bin/pycbc_banksim_match_combine +++ b/bin/pycbc_banksim_match_combine @@ -23,20 +23,19 @@ bank files, and the number of injections in each must correspond one-to-one. """ import argparse -import numpy as np -from igwn_ligolw import utils, ligolw +import numpy as np +from igwn_ligolw import ligolw, utils import pycbc -from pycbc import pnutils -from pycbc.waveform import TemplateBank -from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc import load_source, pnutils from pycbc.io.hdf import HFile -from pycbc import load_source +from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.waveform import TemplateBank -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_banksim_match_combine" @@ -44,28 +43,36 @@ __program__ = "pycbc_banksim_match_combine" parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--match-files", nargs='+', - help="Explicit list of match files.") -#parser.add_argument("--inj-files", nargs='+', +parser.add_argument("--match-files", nargs="+", help="Explicit list of match files.") +# parser.add_argument("--inj-files", nargs='+', # help="Explicit list of injection files. These must be in " # "the same order, and match one-to-one with the " # "match-files.") -parser.add_argument("-o", "--output-file", required=True, - help="Output file name") -parser.add_argument("--filter-func-file", default=None, - help="This can be provided to give a function to define " - "which points are covered by the template bank " - "bounds, and which are not. The file should contain " - "a function called filter_injections, which should " - "take as call profile, mass1, mass2, spin1z, spin2z, " - "as numpy arrays.") +parser.add_argument("-o", "--output-file", required=True, help="Output file name") +parser.add_argument( + "--filter-func-file", + default=None, + help="This can be provided to give a function to define " + "which points are covered by the template bank " + "bounds, and which are not. The file should contain " + "a function called filter_injections, which should " + "take as call profile, mass1, mass2, spin1z, spin2z, " + "as numpy arrays.", +) options = parser.parse_args() pycbc.init_logging(options.verbose) -dtypem = np.dtype([('match', np.float64), ('bank', np.str_, 256), - ('bank_i', np.int32), ('sim', np.str_, 256), - ('sim_i', np.int32), ('sigmasq', np.float64)]) +dtypem = np.dtype( + [ + ("match", np.float64), + ("bank", np.str_, 256), + ("bank_i", np.int32), + ("sim", np.str_, 256), + ("sim_i", np.int32), + ("sigmasq", np.float64), + ] +) # Collect the results res = None @@ -76,7 +83,7 @@ for fil in options.match_files: res = np.loadtxt(fil, dtype=dtypem) btables = {} -itables = {} +itables = {} f = HFile(options.output_file, "w") @@ -84,38 +91,59 @@ f = HFile(options.output_file, "w") # why I'm using a dictionary now. bank_params = {} -bank_par_list = ['mass1', 'mass2', 'spin1x', 'spin1y', 'spin1z', 'spin2x', - 'spin2y', 'spin2z'] +bank_par_list = [ + "mass1", + "mass2", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", +] for val in bank_par_list: bank_params[val] = np.zeros(len(res), dtype=np.float64) inj_params = {} -inj_par_list = ['mass1', 'mass2', 'spin1x', 'spin1y', 'spin1z', 'spin2x', - 'spin2y', 'spin2z', 'coa_phase', 'inclination', 'latitude', - 'longitude', 'polarization'] +inj_par_list = [ + "mass1", + "mass2", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + "coa_phase", + "inclination", + "latitude", + "longitude", + "polarization", +] for val in inj_par_list: inj_params[val] = np.zeros(len(res), dtype=np.float64) trig_params = {} -trig_par_list = ['match', 'sigmasq'] +trig_par_list = ["match", "sigmasq"] for val in trig_par_list: trig_params[val] = np.zeros(len(res), dtype=np.float64) -for idx, row in enumerate(res): +for idx, row in enumerate(res): outstr = "" - if row['bank'] not in btables: - bank_path = row['bank'] + if row["bank"] not in btables: + bank_path = row["bank"] temp_bank = TemplateBank(bank_path) - btables[row['bank']] = temp_bank.table - - if row['sim'] not in itables: - indoc = utils.load_filename(row['sim'], False, - contenthandler=LIGOLWContentHandler) - itables[row['sim']] = ligolw.Table.get_table(indoc, "sim_inspiral") - - bt = btables[row['bank']][row['bank_i']] - it = itables[row['sim']][row['sim_i']] - + btables[row["bank"]] = temp_bank.table + + if row["sim"] not in itables: + indoc = utils.load_filename( + row["sim"], False, contenthandler=LIGOLWContentHandler + ) + itables[row["sim"]] = ligolw.Table.get_table(indoc, "sim_inspiral") + + bt = btables[row["bank"]][row["bank_i"]] + it = itables[row["sim"]][row["sim_i"]] + for val in trig_par_list: trig_params[val][idx] = row[val] for val in bank_par_list: @@ -124,65 +152,70 @@ for idx, row in enumerate(res): except AttributeError: # If not present set to 0. # For example spin1x is not always stored in aligned-spin banks - bank_params[val][idx] = 0. + bank_params[val][idx] = 0.0 for val in inj_par_list: inj_params[val][idx] = getattr(it, val) for val in bank_par_list: - f['bank_params/{}'.format(val)] = bank_params[val] + f[f"bank_params/{val}"] = bank_params[val] for val in inj_par_list: - f['inj_params/{}'.format(val)] = inj_params[val] + f[f"inj_params/{val}"] = inj_params[val] for val in trig_par_list: - f['trig_params/{}'.format(val)] = trig_params[val] + f[f"trig_params/{val}"] = trig_params[val] if options.filter_func_file: - modl = load_source('filter_func', options.filter_func_file) + modl = load_source("filter_func", options.filter_func_file) func = modl.filter_injections - bool_arr = func(inj_params['mass1'], inj_params['mass2'], - inj_params['spin1z'], inj_params['spin2z']) + bool_arr = func( + inj_params["mass1"], + inj_params["mass2"], + inj_params["spin1z"], + inj_params["spin2z"], + ) bool_arr = np.array(bool_arr) # Also consider values over the whole set # Signal recovery fraction -sigma = trig_params['sigmasq']**0.5 -srfn = np.sum((trig_params['match'] * sigma)**3.) -srfd = np.sum(sigma**3.) - -f['sig_rec_fac'] = srfn / srfd -f['eff_fitting_factor'] = (srfn / srfd)**(1./3.) -mchirp, _ = pnutils.mass1_mass2_to_mchirp_eta(inj_params['mass1'], - inj_params['mass2']) -srfn_mcweighted = np.sum((trig_params['match'] * mchirp**(-5./6.) *\ - sigma)**3.) -srfd_mcweighted = np.sum((mchirp**(-5./6.) * sigma)**3.) -f['sig_rec_fac_chirp_mass_weighted'] = srfn_mcweighted / srfd_mcweighted -f['eff_fitting_factor_chirp_mass_weighted'] = \ - (srfn_mcweighted / srfd_mcweighted)**(1./3.) +sigma = trig_params["sigmasq"] ** 0.5 +srfn = np.sum((trig_params["match"] * sigma) ** 3.0) +srfd = np.sum(sigma**3.0) + +f["sig_rec_fac"] = srfn / srfd +f["eff_fitting_factor"] = (srfn / srfd) ** (1.0 / 3.0) +mchirp, _ = pnutils.mass1_mass2_to_mchirp_eta(inj_params["mass1"], inj_params["mass2"]) +srfn_mcweighted = np.sum((trig_params["match"] * mchirp ** (-5.0 / 6.0) * sigma) ** 3.0) +srfd_mcweighted = np.sum((mchirp ** (-5.0 / 6.0) * sigma) ** 3.0) +f["sig_rec_fac_chirp_mass_weighted"] = srfn_mcweighted / srfd_mcweighted +f["eff_fitting_factor_chirp_mass_weighted"] = (srfn_mcweighted / srfd_mcweighted) ** ( + 1.0 / 3.0 +) if options.filter_func_file: - num_filtered = len(inj_params['mass1'][bool_arr]) + num_filtered = len(inj_params["mass1"][bool_arr]) if num_filtered == 0: - f['frac_points_within_bank'] = 0 - f['filtered_sig_rec_fac'] = -1 - f['filtered_eff_fitting_factor'] = -1 + f["frac_points_within_bank"] = 0 + f["filtered_sig_rec_fac"] = -1 + f["filtered_eff_fitting_factor"] = -1 else: - f['frac_points_within_bank'] = \ - num_filtered / float(len(inj_params['mass1'])) - filt_match = trig_params['match'][bool_arr] - filt_sigma = trig_params['sigmasq'][bool_arr]**0.5 - srfn_filt = np.sum((filt_match * filt_sigma)**3.) + f["frac_points_within_bank"] = num_filtered / float(len(inj_params["mass1"])) + filt_match = trig_params["match"][bool_arr] + filt_sigma = trig_params["sigmasq"][bool_arr] ** 0.5 + srfn_filt = np.sum((filt_match * filt_sigma) ** 3.0) srfd_filt = np.sum(filt_sigma**3) - f['filtered_sig_rec_fac'] = srfn_filt / srfd_filt - f['filtered_eff_fitting_factor'] = (srfn_filt / srfd_filt)**(1./3.) + f["filtered_sig_rec_fac"] = srfn_filt / srfd_filt + f["filtered_eff_fitting_factor"] = (srfn_filt / srfd_filt) ** (1.0 / 3.0) mchirp = mchirp[bool_arr] - srfn_mcweighted = np.sum((filt_match * mchirp**(-5./6.) *\ - filt_sigma)**3.) - srfd_mcweighted = np.sum((mchirp**(-5./6.) * filt_sigma)**3.) - f['filtered_sig_rec_fac_chirp_mass_weighted'] = \ + srfn_mcweighted = np.sum( + (filt_match * mchirp ** (-5.0 / 6.0) * filt_sigma) ** 3.0 + ) + srfd_mcweighted = np.sum((mchirp ** (-5.0 / 6.0) * filt_sigma) ** 3.0) + f["filtered_sig_rec_fac_chirp_mass_weighted"] = ( + srfn_mcweighted / srfd_mcweighted + ) + f["filtered_eff_fitting_factor_chirp_mass_weighted"] = ( srfn_mcweighted / srfd_mcweighted - f['filtered_eff_fitting_factor_chirp_mass_weighted'] = \ - (srfn_mcweighted / srfd_mcweighted)**(1./3.) + ) ** (1.0 / 3.0) - f['filtered_points'] = bool_arr + f["filtered_points"] = bool_arr f.close() diff --git a/bin/pycbc_banksim_skymax b/bin/pycbc_banksim_skymax index f544e9e4b8e..6c801f750af 100644 --- a/bin/pycbc_banksim_skymax +++ b/bin/pycbc_banksim_skymax @@ -16,51 +16,66 @@ """Calculate the fitting factors of simulated signals with a template bank.""" - import logging -from numpy import complex64, sqrt, argmax, real, array from argparse import ArgumentParser from math import ceil, log -from tqdm import tqdm -from igwn_ligolw import utils as ligolw_utils from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils +from numpy import argmax, array, complex64, real, sqrt +from tqdm import tqdm -from pycbc.pnutils import mass1_mass2_to_mchirp_eta -from pycbc.pnutils import nearest_larger_binary_number -from pycbc.pnutils import mass1_mass2_to_tau0_tau3 -from pycbc.waveform import get_td_waveform, get_fd_waveform, td_approximants, fd_approximants +import pycbc.fft +import pycbc.psd +import pycbc.scheme +import pycbc.strain from pycbc import DYN_RANGE_FAC -from pycbc.types import FrequencySeries, TimeSeries, zeros, complex_same_precision_as -from pycbc.filter import sigmasq -from pycbc.filter import overlap_cplx, matched_filter -from pycbc.filter import compute_max_snr_over_sky_loc_stat -from pycbc.filter import compute_max_snr_over_sky_loc_stat_no_phase -from pycbc.io.ligolw import LIGOLWContentHandler -import pycbc.psd, pycbc.scheme, pycbc.fft, pycbc.strain from pycbc.detector import overhead_antenna_pattern as generate_fplus_fcross -from pycbc.waveform import TemplateBank +from pycbc.filter import ( + compute_max_snr_over_sky_loc_stat, + compute_max_snr_over_sky_loc_stat_no_phase, + matched_filter, + overlap_cplx, + sigmasq, +) +from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.pnutils import ( + mass1_mass2_to_mchirp_eta, + mass1_mass2_to_tau0_tau3, + nearest_larger_binary_number, +) +from pycbc.types import FrequencySeries, TimeSeries, complex_same_precision_as, zeros +from pycbc.waveform import ( + TemplateBank, + fd_approximants, + get_fd_waveform, + get_td_waveform, + td_approximants, +) ## Remove the need for these functions ######################################## - + + def generate_detector_strain(template_params, h_plus, h_cross): - latitude = 0 - longitude = 0 - polarization = 0 + latitude = 0 + longitude = 0 + polarization = 0 - if hasattr(template_params, 'latitude'): + if hasattr(template_params, "latitude"): latitude = template_params.latitude - if hasattr(template_params, 'longitude'): + if hasattr(template_params, "longitude"): longitude = template_params.longitude - if hasattr(template_params, 'polarization'): + if hasattr(template_params, "polarization"): polarization = template_params.polarization f_plus, f_cross = generate_fplus_fcross(longitude, latitude, polarization) return h_plus * f_plus + h_cross * f_cross + def make_padded_frequency_series(vec, filter_N=None, delta_f=None): - """Convert vec (TimeSeries or FrequencySeries) to a FrequencySeries. If + """ + Convert vec (TimeSeries or FrequencySeries) to a FrequencySeries. If filter_N and/or delta_f are given, the output will take those values. If not told otherwise the code will attempt to pad a timeseries first such that the waveform will not wraparound. However, if delta_f is specified to be @@ -68,14 +83,15 @@ def make_padded_frequency_series(vec, filter_N=None, delta_f=None): """ if filter_N is None: power = ceil(log(len(vec), 2)) + 1 - N = 2 ** power + N = 2**power else: N = filter_N n = N // 2 + 1 if isinstance(vec, FrequencySeries): - vectilde = FrequencySeries(zeros(n, dtype=complex_same_precision_as(vec)), - delta_f=1.0, copy=False) + vectilde = FrequencySeries( + zeros(n, dtype=complex_same_precision_as(vec)), delta_f=1.0, copy=False + ) if len(vectilde) < len(vec): cplen = len(vectilde) else: @@ -88,7 +104,7 @@ def make_padded_frequency_series(vec, filter_N=None, delta_f=None): # and increase if necessary curr_length = len(vec) new_length = int(nearest_larger_binary_number(curr_length)) - while new_length * vec.delta_t < 1./delta_f: + while new_length * vec.delta_t < 1.0 / delta_f: new_length = new_length * 2 vec.resize(new_length) # Then convert to frequencyseries @@ -101,120 +117,214 @@ def make_padded_frequency_series(vec, filter_N=None, delta_f=None): i_delta_f = v_tilde.get_delta_f() v_tilde = v_tilde.numpy() df_ratio = int(delta_f / i_delta_f) - n_freq_len = int((n-1) * df_ratio +1) - assert(n <= len(v_tilde)) + n_freq_len = int((n - 1) * df_ratio + 1) + assert n <= len(v_tilde) df_ratio = int(delta_f / i_delta_f) v_tilde = v_tilde[:n_freq_len:df_ratio] vectilde = FrequencySeries(v_tilde, delta_f=delta_f, dtype=complex64) - return FrequencySeries(vectilde * DYN_RANGE_FAC, delta_f=delta_f, - dtype=complex64) + return FrequencySeries(vectilde * DYN_RANGE_FAC, delta_f=delta_f, dtype=complex64) -def get_waveform(approximant, phase_order, amplitude_order, spin_order, - wf_params, start_frequency, sample_rate, length, - filter_rate, sky_max_template=False): + +def get_waveform( + approximant, + phase_order, + amplitude_order, + spin_order, + wf_params, + start_frequency, + sample_rate, + length, + filter_rate, + sky_max_template=False, +): if type(approximant) is not str: - approximant = approximant.decode('utf-8') + approximant = approximant.decode("utf-8") delta_f = filter_rate / length if approximant in fd_approximants(): - hp, hc = get_fd_waveform(wf_params, approximant=approximant, - phase_order=phase_order, delta_f=delta_f, - spin_order=spin_order, - f_lower=start_frequency, - amplitude_order=amplitude_order) + hp, hc = get_fd_waveform( + wf_params, + approximant=approximant, + phase_order=phase_order, + delta_f=delta_f, + spin_order=spin_order, + f_lower=start_frequency, + amplitude_order=amplitude_order, + ) elif approximant in td_approximants(): - hp, hc = get_td_waveform(wf_params, - approximant=approximant, - phase_order=phase_order, - spin_order=spin_order, - delta_t=1./sample_rate, - f_lower=start_frequency, - amplitude_order=amplitude_order) - if hasattr(wf_params, 'taper'): - hp = hp.taper_timeseries(location=wf_params.taper, - tapermethod=wf_params.get('taper_method', 'lal'), - taper_window=wf_params.get('taper_window')) - hc = hc.taper_timeseries(location=wf_params.taper, - tapermethod=wf_params.get('taper_method', 'lal'), - taper_window=wf_params.get('taper_window')) + hp, hc = get_td_waveform( + wf_params, + approximant=approximant, + phase_order=phase_order, + spin_order=spin_order, + delta_t=1.0 / sample_rate, + f_lower=start_frequency, + amplitude_order=amplitude_order, + ) + if hasattr(wf_params, "taper"): + hp = hp.taper_timeseries( + location=wf_params.taper, + tapermethod=wf_params.get("taper_method", "lal"), + taper_window=wf_params.get("taper_window"), + ) + hc = hc.taper_timeseries( + location=wf_params.taper, + tapermethod=wf_params.get("taper_method", "lal"), + taper_window=wf_params.get("taper_window"), + ) if not sky_max_template: hvec = generate_detector_strain(wf_params, hp, hc) return make_padded_frequency_series(hvec, filter_N, delta_f=delta_f) - else: - return make_padded_frequency_series(hp, filter_N, delta_f=delta_f), \ - make_padded_frequency_series(hc, filter_N, delta_f=delta_f) + return make_padded_frequency_series( + hp, filter_N, delta_f=delta_f + ), make_padded_frequency_series(hc, filter_N, delta_f=delta_f) + aprs = sorted(list(set(td_approximants() + fd_approximants()))) -#File output Settings +# File output Settings parser = ArgumentParser(description=__doc__) -parser.add_argument("--match-file", dest="out_file", metavar="FILE", - required=True, help="File to output match results") +parser.add_argument( + "--match-file", + dest="out_file", + metavar="FILE", + required=True, + help="File to output match results", +) pycbc.add_common_pycbc_options(parser) -#Template Settings -parser.add_argument("--template-file", dest="bank_file", metavar="FILE", - required=True, help="SimInspiral or SnglInspiral XML file " - "containing the template parameters") -parser.add_argument("--total-mass-divide", type=float, - help="Total mass to switch from --template-approximant to " - "--highmass-approximant.") -parser.add_argument("--highmass-approximant", choices=aprs, - help="Waveform approximant for highmass templates.") -parser.add_argument("--template-approximant", choices=aprs, required=True, - help="Waveform approximant for templates") -parser.add_argument("--template-phase-order", default=-1, type=int, - help="PN order to use for the template phase") -parser.add_argument("--template-amplitude-order", default=-1, type=int, - help="PN order to use for the template amplitude") -parser.add_argument("--template-spin-order", default=-1, type=int, - help="PN order to use for the template spin terms") -parser.add_argument("--template-start-frequency", type=float, - help="Starting frequency for templates [Hz]") -parser.add_argument("--template-sample-rate", type=float, - help="Sample rate for templates [Hz]") - -#Signal Settings -parser.add_argument("--signal-file", dest="sim_file", metavar="FILE", - required=True, help="SimInspiral or SnglInspiral XML file " - "containing the signal parameters") -parser.add_argument("--signal-approximant", choices=aprs, required=True, - help="Waveform approximant for signals") -parser.add_argument("--signal-phase-order", default=-1, type=int, - help="PN order to use for the signal phase") -parser.add_argument("--signal-spin-order", default=-1, type=int, - help="PN order to use for the signal spin terms") -parser.add_argument("--signal-amplitude-order", default=-1, type=int, - help="PN order to use for the signal amplitude") -parser.add_argument("--signal-start-frequency", type=float, - help="Starting frequency for signals [Hz]") -parser.add_argument("--signal-sample-rate", type=float, - help="Sample rate for signals [Hz]") -parser.add_argument("--use-sky-location", action='store_true', - help="Inject into a theoretical detector at the celestial " - "North pole of a non-rotating Earth rather than overhead") - -#Filtering Settings -parser.add_argument('--filter-low-frequency-cutoff', metavar='FREQ', type=float, - required=True, help='low frequency cutoff of matched filter') -parser.add_argument('--filter-low-freq-cutoff-column', metavar='NAME', type=str, - help='If given, use a per-template low-frequency cutoff ' - 'from column NAME of the template table instead of ' - 'the value given via --filter-low-frequency-cutoff. ' - 'Signals are still normalized using ' - '--filter-low-frequency-cutoff and the value in ' - 'column NAME must be larger than or equal to it.') -parser.add_argument("--filter-sample-rate", type=float, required=True, - help="Filter sample rate [Hz]") -parser.add_argument("--filter-signal-length", type=int, required=True, - help="Length of signal for filtering, shoud be longer " - "than all waveforms and include some padding") -parser.add_argument("--sky-maximization-method", required=True, - choices=["precessing", "hom"]) +# Template Settings +parser.add_argument( + "--template-file", + dest="bank_file", + metavar="FILE", + required=True, + help="SimInspiral or SnglInspiral XML file containing the template parameters", +) +parser.add_argument( + "--total-mass-divide", + type=float, + help="Total mass to switch from --template-approximant to --highmass-approximant.", +) +parser.add_argument( + "--highmass-approximant", + choices=aprs, + help="Waveform approximant for highmass templates.", +) +parser.add_argument( + "--template-approximant", + choices=aprs, + required=True, + help="Waveform approximant for templates", +) +parser.add_argument( + "--template-phase-order", + default=-1, + type=int, + help="PN order to use for the template phase", +) +parser.add_argument( + "--template-amplitude-order", + default=-1, + type=int, + help="PN order to use for the template amplitude", +) +parser.add_argument( + "--template-spin-order", + default=-1, + type=int, + help="PN order to use for the template spin terms", +) +parser.add_argument( + "--template-start-frequency", + type=float, + help="Starting frequency for templates [Hz]", +) +parser.add_argument( + "--template-sample-rate", type=float, help="Sample rate for templates [Hz]" +) + +# Signal Settings +parser.add_argument( + "--signal-file", + dest="sim_file", + metavar="FILE", + required=True, + help="SimInspiral or SnglInspiral XML file containing the signal parameters", +) +parser.add_argument( + "--signal-approximant", + choices=aprs, + required=True, + help="Waveform approximant for signals", +) +parser.add_argument( + "--signal-phase-order", + default=-1, + type=int, + help="PN order to use for the signal phase", +) +parser.add_argument( + "--signal-spin-order", + default=-1, + type=int, + help="PN order to use for the signal spin terms", +) +parser.add_argument( + "--signal-amplitude-order", + default=-1, + type=int, + help="PN order to use for the signal amplitude", +) +parser.add_argument( + "--signal-start-frequency", type=float, help="Starting frequency for signals [Hz]" +) +parser.add_argument( + "--signal-sample-rate", type=float, help="Sample rate for signals [Hz]" +) +parser.add_argument( + "--use-sky-location", + action="store_true", + help="Inject into a theoretical detector at the celestial " + "North pole of a non-rotating Earth rather than overhead", +) + +# Filtering Settings +parser.add_argument( + "--filter-low-frequency-cutoff", + metavar="FREQ", + type=float, + required=True, + help="low frequency cutoff of matched filter", +) +parser.add_argument( + "--filter-low-freq-cutoff-column", + metavar="NAME", + type=str, + help="If given, use a per-template low-frequency cutoff " + "from column NAME of the template table instead of " + "the value given via --filter-low-frequency-cutoff. " + "Signals are still normalized using " + "--filter-low-frequency-cutoff and the value in " + "column NAME must be larger than or equal to it.", +) +parser.add_argument( + "--filter-sample-rate", type=float, required=True, help="Filter sample rate [Hz]" +) +parser.add_argument( + "--filter-signal-length", + type=int, + required=True, + help="Length of signal for filtering, shoud be longer " + "than all waveforms and include some padding", +) +parser.add_argument( + "--sky-maximization-method", required=True, choices=["precessing", "hom"] +) # add PSD options @@ -223,24 +333,33 @@ pycbc.psd.insert_psd_option_group(parser, output=False) # Insert the data reading options pycbc.strain.insert_strain_option_group(parser) -#hardware support +# hardware support pycbc.scheme.insert_processing_option_group(parser) pycbc.fft.insert_fft_option_group(parser) -#Restricted maximization -parser.add_argument("--mchirp-window", type=str, metavar="FRACTION", - help="Ignore templates whose chirp mass deviates from " - "signal's one more than given fraction. Provide two " - "comma separated numbers to have different bounds " - "above and below the signal's, with below bound " - "listed first.") -parser.add_argument("--tau0-window", type=float, metavar="TIME", default=None, - help="Ignore templates whose Newtonian order chirp time " - "(tau0) varies from the signals by more than the " - "supplied amount. If option is not provided no " - "window on tau0 is used. The " - "filter-low-frequency-cutoff is used to calculate " - "the value of tau0 for all cases.") +# Restricted maximization +parser.add_argument( + "--mchirp-window", + type=str, + metavar="FRACTION", + help="Ignore templates whose chirp mass deviates from " + "signal's one more than given fraction. Provide two " + "comma separated numbers to have different bounds " + "above and below the signal's, with below bound " + "listed first.", +) +parser.add_argument( + "--tau0-window", + type=float, + metavar="TIME", + default=None, + help="Ignore templates whose Newtonian order chirp time " + "(tau0) varies from the signals by more than the " + "supplied amount. If option is not provided no " + "window on tau0 is used. The " + "filter-low-frequency-cutoff is used to calculate " + "the value of tau0 for all cases.", +) options = parser.parse_args() @@ -252,29 +371,36 @@ if options.psd_estimation: pycbc.strain.verify_strain_options(options, parser) if options.total_mass_divide and options.highmass_approximant is None: - parser.error("You must provide a highmass-approximant if you want total-mass-divide.") + parser.error( + "You must provide a highmass-approximant if you want total-mass-divide." + ) if options.mchirp_window is None: + def outside_mchirp_window(template_mchirp, signal_mchirp): return False -elif ',' in options.mchirp_window: +elif "," in options.mchirp_window: # asymmetric chirp mass window mchirp_window_lower = float(options.mchirp_window.split(",")[0]) mchirp_window_upper = float(options.mchirp_window.split(",")[1]) + def outside_mchirp_window(template_mchirp, signal_mchirp): delta = (template_mchirp - signal_mchirp) / signal_mchirp return delta > mchirp_window_upper or -delta > mchirp_window_lower else: # symmetric chirp mass window mchirp_window = float(options.mchirp_window) + def outside_mchirp_window(template_mchirp, signal_mchirp): - return abs(signal_mchirp - template_mchirp) > \ - (mchirp_window * signal_mchirp) + return abs(signal_mchirp - template_mchirp) > (mchirp_window * signal_mchirp) + if options.tau0_window is None: + def outside_tau0_window(template_tau0, signal_tau0, window): return False else: + def outside_tau0_window(template_tau0, signal_tau0, window): return abs(signal_tau0 - template_tau0) > window @@ -297,15 +423,16 @@ else: ctx = pycbc.scheme.from_cli(options) -logging.info('Reading template bank') +logging.info("Reading template bank") temp_bank = TemplateBank(options.bank_file) template_table = temp_bank.table - + logging.info(" %d templates", len(template_table)) -logging.info('Reading simulation list') -indoc = ligolw_utils.load_filename(options.sim_file, False, - contenthandler=LIGOLWContentHandler) +logging.info("Reading simulation list") +indoc = ligolw_utils.load_filename( + options.sim_file, False, contenthandler=LIGOLWContentHandler +) try: signal_table = lsctables.SimInspiralTable.get_table(indoc) except ValueError: @@ -319,12 +446,17 @@ filter_n = filter_N // 2 + 1 filter_delta_f = 1.0 / float(options.filter_signal_length) logging.info("Reading and Interpolating PSD") -psd = pycbc.psd.from_cli(options, filter_n, filter_delta_f, - options.filter_low_frequency_cutoff, strain=strain, - dyn_range_factor=pycbc.DYN_RANGE_FAC, - precision='single') - -with ctx: +psd = pycbc.psd.from_cli( + options, + filter_n, + filter_delta_f, + options.filter_low_frequency_cutoff, + strain=strain, + dyn_range_factor=pycbc.DYN_RANGE_FAC, + precision="single", +) + +with ctx: pycbc.fft.from_cli(options) logging.info("Pregenerating Signals") @@ -333,22 +465,26 @@ with ctx: # Used for getting mchirp/tau0 later sig_m1 = [] sig_m2 = [] - prog = tqdm(total=len(signal_table), disable=(not options.verbose)) + prog = tqdm(total=len(signal_table), disable=(not options.verbose)) for index, signal_params in enumerate(signal_table): prog.update(1) if not options.use_sky_location: - signal_params.latitude = 0. - signal_params.longitude = 0. - stilde = get_waveform(options.signal_approximant, - options.signal_phase_order, - options.signal_amplitude_order, - options.signal_spin_order, - signal_params, - options.signal_start_frequency, - signal_sample_rate, - filter_N, options.filter_sample_rate) - s_norm = sigmasq(stilde, psd=psd, - low_frequency_cutoff=options.filter_low_frequency_cutoff) + signal_params.latitude = 0.0 + signal_params.longitude = 0.0 + stilde = get_waveform( + options.signal_approximant, + options.signal_phase_order, + options.signal_amplitude_order, + options.signal_spin_order, + signal_params, + options.signal_start_frequency, + signal_sample_rate, + filter_N, + options.filter_sample_rate, + ) + s_norm = sigmasq( + stilde, psd=psd, low_frequency_cutoff=options.filter_low_frequency_cutoff + ) stilde /= sqrt(float(s_norm)) stilde /= psd signals.append((stilde, s_norm, [], signal_params)) @@ -357,8 +493,9 @@ with ctx: prog.close() sig_m1 = array(sig_m1) sig_m2 = array(sig_m2) - sig_tau0, _ = mass1_mass2_to_tau0_tau3(sig_m1, sig_m2, - options.filter_low_frequency_cutoff) + sig_tau0, _ = mass1_mass2_to_tau0_tau3( + sig_m1, sig_m2, options.filter_low_frequency_cutoff + ) sig_mchirp, _ = mass1_mass2_to_mchirp_eta(sig_m1, sig_m2) logging.info("Calculating Mchirp and Tau0") @@ -369,14 +506,15 @@ with ctx: template_m2.append(template_params.mass2) template_m1 = array(template_m1) template_m2 = array(template_m2) - template_tau0, _ = mass1_mass2_to_tau0_tau3(template_m1, template_m2, - options.filter_low_frequency_cutoff) + template_tau0, _ = mass1_mass2_to_tau0_tau3( + template_m1, template_m2, options.filter_low_frequency_cutoff + ) template_mchirp, _ = mass1_mass2_to_mchirp_eta(template_m1, template_m2) logging.info("Calculating Overlaps") flow_warned = False - prog = tqdm(total=len(template_table), disable=(not options.verbose)) + prog = tqdm(total=len(template_table), disable=(not options.verbose)) for index, template_params in enumerate(template_table): prog.update(1) f_lower = template_params.f_lower @@ -386,22 +524,24 @@ with ctx: if f_lower < options.filter_low_frequency_cutoff: # Not entirely clear what to do here? if not flow_warned: - logging.warning("Template's flower is smaller than " - "--filter-low-frequency-cutoff. Raising " - "flower of template to match.") - flow_warned=True + logging.warning( + "Template's flower is smaller than " + "--filter-low-frequency-cutoff. Raising " + "flower of template to match." + ) + flow_warned = True f_lower = options.filter_low_frequency_cutoff - h_norm = htilde = None for sidx, (stilde, s_norm, matches, signal_params) in enumerate(signals): # Check if we need to look at this check_logic = stilde is None - check_logic |= outside_tau0_window(template_tau0[index], - sig_tau0[sidx], - options.tau0_window) - check_logic |= outside_mchirp_window(template_mchirp[index], - sig_mchirp[sidx]) + check_logic |= outside_tau0_window( + template_tau0[index], sig_tau0[sidx], options.tau0_window + ) + check_logic |= outside_mchirp_window( + template_mchirp[index], sig_mchirp[sidx] + ) if check_logic: matches.append(0) continue @@ -413,47 +553,76 @@ with ctx: # However, while we are still using the high-mass divide # in XML banks, this must be retained. try: - this_approximant = template_params['approximant'] + this_approximant = template_params["approximant"] except: this_approximant = options.template_approximant - if options.total_mass_divide is not None and (template_params.mass1+template_params.mass2) >= options.total_mass_divide: + if ( + options.total_mass_divide is not None + and (template_params.mass1 + template_params.mass2) + >= options.total_mass_divide + ): this_approximant = options.highmass_approximant - hplus, hcross = get_waveform(this_approximant, - options.template_phase_order, - options.template_amplitude_order, - options.template_spin_order, - template_params, - options.template_start_frequency, - template_sample_rate, - filter_N, options.filter_sample_rate, - sky_max_template=True) + hplus, hcross = get_waveform( + this_approximant, + options.template_phase_order, + options.template_amplitude_order, + options.template_spin_order, + template_params, + options.template_start_frequency, + template_sample_rate, + filter_N, + options.filter_sample_rate, + sky_max_template=True, + ) hp_norm = sigmasq(hplus, psd=psd, low_frequency_cutoff=f_lower) hc_norm = sigmasq(hcross, psd=psd, low_frequency_cutoff=f_lower) hplus /= sqrt(float(hp_norm)) hcross /= sqrt(float(hc_norm)) - hpc_corr = overlap_cplx(hplus, hcross, psd=psd, - low_frequency_cutoff=options.filter_low_frequency_cutoff, - normalized=False) + hpc_corr = overlap_cplx( + hplus, + hcross, + psd=psd, + low_frequency_cutoff=options.filter_low_frequency_cutoff, + normalized=False, + ) hpc_corr_R = real(hpc_corr) - htilde=1 - - I_plus = matched_filter(hplus, stilde, - low_frequency_cutoff=options.filter_low_frequency_cutoff, - sigmasq=1.) - - I_cross = matched_filter(hcross, stilde, - low_frequency_cutoff=options.filter_low_frequency_cutoff, - sigmasq=1.) - - if options.sky_maximization_method == 'precessing': - det_stat = compute_max_snr_over_sky_loc_stat\ - (I_plus, I_cross, hpc_corr_R, hpnorm=1., hcnorm=1., - thresh=0.1, analyse_slice=slice(0,len(I_plus.data))) - elif options.sky_maximization_method == 'hom': - det_stat = compute_max_snr_over_sky_loc_stat_no_phase\ - (I_plus, I_cross, hpc_corr_R, hpnorm=1., hcnorm=1., - thresh=0.1, analyse_slice=slice(0,len(I_plus.data))) + htilde = 1 + + I_plus = matched_filter( + hplus, + stilde, + low_frequency_cutoff=options.filter_low_frequency_cutoff, + sigmasq=1.0, + ) + + I_cross = matched_filter( + hcross, + stilde, + low_frequency_cutoff=options.filter_low_frequency_cutoff, + sigmasq=1.0, + ) + + if options.sky_maximization_method == "precessing": + det_stat = compute_max_snr_over_sky_loc_stat( + I_plus, + I_cross, + hpc_corr_R, + hpnorm=1.0, + hcnorm=1.0, + thresh=0.1, + analyse_slice=slice(0, len(I_plus.data)), + ) + elif options.sky_maximization_method == "hom": + det_stat = compute_max_snr_over_sky_loc_stat_no_phase( + I_plus, + I_cross, + hpc_corr_R, + hpnorm=1.0, + hcnorm=1.0, + thresh=0.1, + analyse_slice=slice(0, len(I_plus.data)), + ) else: err_msg = "I really shouldn't be here! Who gone broked me?" raise ValueError(err_msg) diff --git a/bin/pycbc_compress_bank b/bin/pycbc_compress_bank index 0a1ddf3e230..303afc052f0 100755 --- a/bin/pycbc_compress_bank +++ b/bin/pycbc_compress_bank @@ -15,137 +15,203 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -__description__ = \ -""""Loads a bank of waveforms and compresses them using the specified +__description__ = """"Loads a bank of waveforms and compresses them using the specified compression algorithm. The resulting compressed waveforms are saved to an hdf file.""" import argparse -import numpy import logging import multiprocessing +import numpy + import pycbc import pycbc.psd -from pycbc.waveform import compress from pycbc import waveform from pycbc.types import real_same_precision_as +from pycbc.waveform import compress + # --- WORKER FUNCTION --- def _worker(ii): """Worker function to generate and compress a single template.""" tmplt = bank_table[ii] - if tmplt['approximant'] in do_not_compress: + if tmplt["approximant"] in do_not_compress: return None - + # Generate the waveform htilde = bank_obj[ii] fmin = tmplt.f_lower template_duration = htilde.chirp_length - + if template_duration is None: template_duration = seg_len // 2 - logging.info("Warning: was not able to determine template duration, " - "setting half segment length for template %s", tmplt.template_hash) - + logging.info( + "Warning: was not able to determine template duration, " + "setting half segment length for template %s", + tmplt.template_hash, + ) + kmin = int(numpy.ceil(fmin / df_val)) - htilde[:kmin] = 0 # This ensures proper phase unwrapping + htilde[:kmin] = 0 # This ensures proper phase unwrapping if numpy.abs(htilde[kmin]) == 0: - raise ValueError("The amplitude of the waveform at the " - "low_frequency_cutoff is zero. A non-zero value " - "is required.") - + raise ValueError( + "The amplitude of the waveform at the " + "low_frequency_cutoff is zero. A non-zero value " + "is required." + ) + kmax = numpy.nonzero(abs(htilde))[0][-1] - + # Get the compressed sample points - if alg == 'mchirp': - sample_points = compress.mchirp_compression(tmplt.mass1, tmplt.mass2, - fmin, kmax*df_val, min_seglen=t_pad, df_multiple=df_val, - scale=scale_val).astype(real_same_precision_as(htilde)) - elif alg == 'spa': - sample_points = compress.spa_compression(htilde, fmin, kmax*df_val, - min_seglen=t_pad, scale=scale_val).astype(real_same_precision_as(htilde)) - elif alg == 'bounds': - sample_points = numpy.array([fmin, kmax*df_val]) + if alg == "mchirp": + sample_points = compress.mchirp_compression( + tmplt.mass1, + tmplt.mass2, + fmin, + kmax * df_val, + min_seglen=t_pad, + df_multiple=df_val, + scale=scale_val, + ).astype(real_same_precision_as(htilde)) + elif alg == "spa": + sample_points = compress.spa_compression( + htilde, fmin, kmax * df_val, min_seglen=t_pad, scale=scale_val + ).astype(real_same_precision_as(htilde)) + elif alg == "bounds": + sample_points = numpy.array([fmin, kmax * df_val]) else: raise ValueError("unrecognized compression algorithm %s" % alg) # Compress hcompressed = compress.compress_waveform( - htilde, sample_points, tol, interp, - 'double', decomp_scratch=None, psd=psd_obj) + htilde, sample_points, tol, interp, "double", decomp_scratch=None, psd=psd_obj + ) return ii, tmplt.template_hash, hcompressed, template_duration + # --- MAIN EXECUTION --- if __name__ == "__main__": parser = argparse.ArgumentParser(description=__description__) pycbc.add_common_pycbc_options(parser) - parser.add_argument("--bank-file", type=str, required=True, - help="Bank hdf file to load.") - parser.add_argument("--output", type=str, required=True, - help="The hdf file to save the templates and " - "compressed waveforms to.") - parser.add_argument("--low-frequency-cutoff", type=float, default=None, - help="The low frequency cutoff to use for generating " - "the waveforms (Hz). If this is not provided, the code " - "will look for a low frequency cutoff corressponding " - "to each template stored in the bank. If a " - "--low-frequency-cutoff is provided and the bank stores " - "a low frequency cutoff for each template as well, the " - "former is used to generate the waveforms.") + parser.add_argument( + "--bank-file", type=str, required=True, help="Bank hdf file to load." + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="The hdf file to save the templates and compressed waveforms to.", + ) + parser.add_argument( + "--low-frequency-cutoff", + type=float, + default=None, + help="The low frequency cutoff to use for generating " + "the waveforms (Hz). If this is not provided, the code " + "will look for a low frequency cutoff corressponding " + "to each template stored in the bank. If a " + "--low-frequency-cutoff is provided and the bank stores " + "a low frequency cutoff for each template as well, the " + "former is used to generate the waveforms.", + ) pycbc.waveform.bank.add_approximant_arg(parser) - parser.add_argument("--sample-rate", type=int, required=True, - help="Half this value sets the maximum frequency of " - "the compressed waveforms. Must be a power of 2.") - parser.add_argument("--segment-length", type=int, required=True, - help="The segment length to use for generating the " - "templates for compression, and for calculating " - "overlap for tolerance. Must be a power of 2, " - "and at least twice as long as the longest template " - "in the bank.") - parser.add_argument("--tmplt-index", nargs=2, type=int, default=None, - help="Only generate compressed waveforms for the given " - "indices in the bank file. Must provide both a start " - "and a stop index. Default is to compress all of the " - "templates in the bank file.") - parser.add_argument("--compression-algorithm", type=str, required=True, - choices=list(compress.compression_algorithms.keys()) + ['bounds'], - help="The compression algorithm to use for selecting " - "frequency points.") - parser.add_argument("--nprocesses", type=int, default=1, - help="Number of parallel processes to use.") - parser.add_argument("--t-pad", type=float, default=0.001, - help="The minimum duration used for t(f) in seconds. " - "The inverse of this gives the maximum frequency " - "step that will be used in the compressed waveforms. " - "Default is 0.001.") - parser.add_argument("--tolerance", type=float, default=0.001, - help="The maximum mismatch to allow between the " - "interpolated waveform must and the full waveform. " - "Points will be added to the compressed waveform " - "until its interpolation has a mismatch <= this value. " - "Default is 0.001.") - parser.add_argument("--interpolation", type=str, default="inline_linear", - help="The interpolation to use for decompressing the " - "waveforms for checking tolerance. Options are " - "'inline_linear', or any interpolation recognized by " - "scipy's interp1d kind argument. Default is inline_linear.") - parser.add_argument("--precision", type=str, choices=["double", "single"], - default="single", - help="What precision to generate and store the " - "waveforms with; default is double.") - parser.add_argument("--force", action="store_true", default=False, - help="Overwrite the given hdf file if it exists. " - "Otherwise, an error is raised.") - parser.add_argument("--scale", type=float, default=1, - help="Scale factor to apply to the courseness of the initial " - "compression points. >1 means a courser set of initial " - "points. The algorithm will still add points to meet" - "the tolerance goals.") - parser.add_argument("--do-not-compress", nargs="+", - help="If given, will not compress waveforms using " - "the given approximant. Recommended 'SPAtmplt TaylorF2'.") + parser.add_argument( + "--sample-rate", + type=int, + required=True, + help="Half this value sets the maximum frequency of " + "the compressed waveforms. Must be a power of 2.", + ) + parser.add_argument( + "--segment-length", + type=int, + required=True, + help="The segment length to use for generating the " + "templates for compression, and for calculating " + "overlap for tolerance. Must be a power of 2, " + "and at least twice as long as the longest template " + "in the bank.", + ) + parser.add_argument( + "--tmplt-index", + nargs=2, + type=int, + default=None, + help="Only generate compressed waveforms for the given " + "indices in the bank file. Must provide both a start " + "and a stop index. Default is to compress all of the " + "templates in the bank file.", + ) + parser.add_argument( + "--compression-algorithm", + type=str, + required=True, + choices=list(compress.compression_algorithms.keys()) + ["bounds"], + help="The compression algorithm to use for selecting frequency points.", + ) + parser.add_argument( + "--nprocesses", type=int, default=1, help="Number of parallel processes to use." + ) + parser.add_argument( + "--t-pad", + type=float, + default=0.001, + help="The minimum duration used for t(f) in seconds. " + "The inverse of this gives the maximum frequency " + "step that will be used in the compressed waveforms. " + "Default is 0.001.", + ) + parser.add_argument( + "--tolerance", + type=float, + default=0.001, + help="The maximum mismatch to allow between the " + "interpolated waveform must and the full waveform. " + "Points will be added to the compressed waveform " + "until its interpolation has a mismatch <= this value. " + "Default is 0.001.", + ) + parser.add_argument( + "--interpolation", + type=str, + default="inline_linear", + help="The interpolation to use for decompressing the " + "waveforms for checking tolerance. Options are " + "'inline_linear', or any interpolation recognized by " + "scipy's interp1d kind argument. Default is inline_linear.", + ) + parser.add_argument( + "--precision", + type=str, + choices=["double", "single"], + default="single", + help="What precision to generate and store the " + "waveforms with; default is double.", + ) + parser.add_argument( + "--force", + action="store_true", + default=False, + help="Overwrite the given hdf file if it exists. " + "Otherwise, an error is raised.", + ) + parser.add_argument( + "--scale", + type=float, + default=1, + help="Scale factor to apply to the courseness of the initial " + "compression points. >1 means a courser set of initial " + "points. The algorithm will still add points to meet" + "the tolerance goals.", + ) + parser.add_argument( + "--do-not-compress", + nargs="+", + help="If given, will not compress waveforms using " + "the given approximant. Recommended 'SPAtmplt TaylorF2'.", + ) pycbc.psd.insert_psd_option_group(parser, include_data_options=False) args = parser.parse_args() @@ -167,19 +233,24 @@ if __name__ == "__main__": if args.do_not_compress is None: args.do_not_compress = [] - df_val = 1./args.segment_length - fmax = args.sample_rate/2. + df_val = 1.0 / args.segment_length + fmax = args.sample_rate / 2.0 N = args.sample_rate * args.segment_length dtype = numpy.complex128 logging.info("loading bank") # we'll do everything in double precision; if single is desired, we'll # cast to single when saving the waveforms - bank_obj = waveform.FilterBank(args.bank_file, N//2+1, df_val, dtype, - low_frequency_cutoff=args.low_frequency_cutoff, - approximant=args.approximant, - enable_compressed_waveforms=False) - + bank_obj = waveform.FilterBank( + args.bank_file, + N // 2 + 1, + df_val, + dtype, + low_frequency_cutoff=args.low_frequency_cutoff, + approximant=args.approximant, + enable_compressed_waveforms=False, + ) + if args.tmplt_index is not None: imin, imax = args.tmplt_index bank_obj.table = bank_obj.table[imin:imax] @@ -193,20 +264,25 @@ if __name__ == "__main__": alg = args.compression_algorithm t_pad = args.t_pad scale_val = args.scale - + logging.info("getting psd") - psd_obj = pycbc.psd.from_cli(args, length=N//2+1, delta_f=df_val, - low_frequency_cutoff=bank_table.f_lower.min(), - dyn_range_factor=pycbc.DYN_RANGE_FAC, - precision='double') + psd_obj = pycbc.psd.from_cli( + args, + length=N // 2 + 1, + delta_f=df_val, + low_frequency_cutoff=bank_table.f_lower.min(), + dyn_range_factor=pycbc.DYN_RANGE_FAC, + precision="double", + ) logging.info("writing template info to output") - output = bank_obj.write_to_hdf(args.output, force=args.force, - write_compressed_waveforms=False) + output = bank_obj.write_to_hdf( + args.output, force=args.force, write_compressed_waveforms=False + ) output.create_group("compressed_waveforms") logging.info("Starting compression with %d processes", args.nprocesses) - + pool = multiprocessing.Pool(processes=args.nprocesses) results = pool.imap_unordered(_worker, range(bank_table.size)) @@ -214,13 +290,14 @@ if __name__ == "__main__": if result is None: continue idx, tmplt_hash, hcompressed, duration = result - + # Verify segment length against duration one last time in main thread - if args.segment_length < 2*duration: - raise ValueError("segment length is < twice the duration " - "({}) of template {}".format(duration, tmplt_hash)) - - output['template_duration'][idx] = duration + if args.segment_length < 2 * duration: + raise ValueError( + f"segment length is < twice the duration ({duration}) of template {tmplt_hash}" + ) + + output["template_duration"][idx] = duration hcompressed.write_to_hdf(output, tmplt_hash, precision=args.precision) logging.info("Saved compressed template %s", tmplt_hash) diff --git a/bin/pycbc_condition_strain b/bin/pycbc_condition_strain index 4f65cc717a3..44175a4f9b0 100644 --- a/bin/pycbc_condition_strain +++ b/bin/pycbc_condition_strain @@ -24,55 +24,75 @@ also be used to generate frames of simulated strain data, with or without injections. """ -import logging import argparse +import logging -import pycbc.strain -import pycbc.frame import pycbc.fft +import pycbc.frame +import pycbc.strain from pycbc.types import float32, float64 def write_strain(file_name, channel, data): - logging.info('Writing output strain to %s', file_name) + logging.info("Writing output strain to %s", file_name) - if file_name.endswith('.gwf'): + if file_name.endswith(".gwf"): pycbc.frame.write_frame(file_name, channel, data) - elif file_name.endswith(('.hdf', '.h5')): + elif file_name.endswith((".hdf", ".h5")): data.save(file_name, group=channel) else: - raise ValueError('Unknown extension for ' + file_name) + raise ValueError("Unknown extension for " + file_name) parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--output-strain-file', required=True, - help='Name of output frame file. The file format is ' - 'selected based on the extension (.gwf, .npy, .hdf ' - 'and .txt accepted)') -parser.add_argument('--output-channel-name', - help='Name of channel in output frame file (default: same ' - 'as input channel)') -parser.add_argument('--output-gates-file', - help='Save gating info to specified file, in the same ' - 'format as accepted by the --gating-file option') -parser.add_argument('--output-precision', type=str, - choices=['single', 'double'], default='double', - help='Precision of output strain, %(default)s by default') -parser.add_argument('--dyn-range-factor', action='store_true', - help='Scale the output strain by a large factor (%f) ' - 'to avoid underflows in subsequent ' - 'calculations' % pycbc.DYN_RANGE_FAC) -parser.add_argument('--low-frequency-cutoff', type=float, - help='Provide a low-frequency-cutoff for fake strain. ' - 'This is only needed if fake-strain or ' - 'fake-strain-from-file is used') -parser.add_argument('--frame-duration', metavar='SECONDS', type=int, - help='Split the produced data into different frame files ' - 'of the given duration. The output file name should ' - 'contain the strings {start} and {duration}, which ' - 'will be replaced by the start GPS time and duration ' - 'in seconds') +parser.add_argument( + "--output-strain-file", + required=True, + help="Name of output frame file. The file format is " + "selected based on the extension (.gwf, .npy, .hdf " + "and .txt accepted)", +) +parser.add_argument( + "--output-channel-name", + help="Name of channel in output frame file (default: same as input channel)", +) +parser.add_argument( + "--output-gates-file", + help="Save gating info to specified file, in the same " + "format as accepted by the --gating-file option", +) +parser.add_argument( + "--output-precision", + type=str, + choices=["single", "double"], + default="double", + help="Precision of output strain, %(default)s by default", +) +parser.add_argument( + "--dyn-range-factor", + action="store_true", + help="Scale the output strain by a large factor (%f) " + "to avoid underflows in subsequent " + "calculations" % pycbc.DYN_RANGE_FAC, +) +parser.add_argument( + "--low-frequency-cutoff", + type=float, + help="Provide a low-frequency-cutoff for fake strain. " + "This is only needed if fake-strain or " + "fake-strain-from-file is used", +) +parser.add_argument( + "--frame-duration", + metavar="SECONDS", + type=int, + help="Split the produced data into different frame files " + "of the given duration. The output file name should " + "contain the strings {start} and {duration}, which " + "will be replaced by the start GPS time and duration " + "in seconds", +) pycbc.strain.insert_strain_option_group(parser) pycbc.fft.insert_fft_option_group(parser) @@ -85,23 +105,27 @@ pycbc.fft.verify_fft_options(args, parser) pycbc.fft.from_cli(args) if args.frame_duration is not None and args.frame_duration <= 0: - parser.error('Frame duration should be positive integer, {} given'.format(args.frame_duration)) + parser.error( + f"Frame duration should be positive integer, {args.frame_duration} given" + ) # read and condition strain as pycbc_inspiral would do -out_strain = pycbc.strain.from_cli(args, dyn_range_fac=pycbc.DYN_RANGE_FAC, - precision=args.output_precision) +out_strain = pycbc.strain.from_cli( + args, dyn_range_fac=pycbc.DYN_RANGE_FAC, precision=args.output_precision +) # if requested, save the gates while we have them if args.output_gates_file: - logging.info('Writing output gates') - with file(args.output_gates_file, 'wb') as gate_f: + logging.info("Writing output gates") + with file(args.output_gates_file, "wb") as gate_f: for k, v in out_strain.gating_info.items(): for t, w, p in v: - gate_f.write('%.4f %.2f %.2f\n' % (t, w, p)) + gate_f.write("%.4f %.2f %.2f\n" % (t, w, p)) # force strain precision to be as requested out_strain = out_strain.astype( - float32 if args.output_precision == 'single' else float64) + float32 if args.output_precision == "single" else float64 +) # unless asked otherwise, revert the dynamic range factor if not args.dyn_range_factor: @@ -116,11 +140,12 @@ if args.frame_duration: # Last frame duration can be shorter than duration if stop doesn't allow for s in range(start, stop, step): - ts = out_strain.time_slice(s, s+step if s+step < stop else stop) + ts = out_strain.time_slice(s, min(stop, s + step)) complete_fn = args.output_strain_file.format( - start=s, duration=step if s+step < stop else stop - s) + start=s, duration=step if s + step < stop else stop - s + ) write_strain(complete_fn, output_channel_name, ts) else: write_strain(args.output_strain_file, output_channel_name, out_strain) -logging.info('Done') +logging.info("Done") diff --git a/bin/pycbc_convertinjfiletohdf b/bin/pycbc_convertinjfiletohdf index 259bc5162b2..8950834cf39 100755 --- a/bin/pycbc_convertinjfiletohdf +++ b/bin/pycbc_convertinjfiletohdf @@ -16,30 +16,34 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" This program converts standard injection file formats (e.g. XML format +""" +This program converts standard injection file formats (e.g. XML format or the LIGO HDF format) into a PyCBC hdf injection format """ import argparse -import numpy import shutil +import numpy +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils + import pycbc from pycbc.inject import CBCHDFInjectionSet -from pycbc.io.record import FieldArray from pycbc.io.hdf import HFile - from pycbc.io.ligolw import LIGOLWContentHandler -from igwn_ligolw import utils as ligolw_utils, lsctables - +from pycbc.io.record import FieldArray # Handlers for alternative injection formats + def legacy_approximant_name(apx): - """Convert the old style xml approximant name to a name + """ + Convert the old style xml approximant name to a name and phase_order. """ import lalsimulation as lalsim + apx = str(apx) try: order = lalsim.GetOrderFromString(apx) @@ -50,9 +54,9 @@ def legacy_approximant_name(apx): return name, order -class XMLInjectionSet(object): - - """Reads injections from LIGOLW XML files +class XMLInjectionSet: + """ + Reads injections from LIGOLW XML files Parameters ---------- @@ -64,11 +68,13 @@ class XMLInjectionSet(object): ---------- indoc table + """ - def __init__(self, sim_file): + def __init__(self, sim_file): self.indoc = ligolw_utils.load_filename( - sim_file, False, contenthandler=LIGOLWContentHandler) + sim_file, False, contenthandler=LIGOLWContentHandler + ) self.table = lsctables.SimInspiralTable.get_table(self.indoc) def end_times(self): @@ -81,54 +87,49 @@ class XMLInjectionSet(object): # Some XML files can have empty columns which are read as None in # python. For these cases we ignore the columns so they do not # appear in the output HDF file. (Such XML files cannot currently - # be read by LALSuite C code.) + # be read by LALSuite C code.) if getattr(self.table[0], key) is not None: - data[str(key)] = numpy.array( - [getattr(t, key) for t in self.table] - ) + data[str(key)] = numpy.array([getattr(t, key) for t in self.table]) - for k in ['simulation_id', 'process_id']: + for k in ["simulation_id", "process_id"]: a = data.pop(k) - data['approximant'], data['phase_order'] = numpy.array( - [legacy_approximant_name(wf) for wf in data['waveform']] + data["approximant"], data["phase_order"] = numpy.array( + [legacy_approximant_name(wf) for wf in data["waveform"]] ).T - data['tc'] = ( - data['geocent_end_time'] + 1e-9 * data['geocent_end_time_ns'] - ) - data['dec'] = data['latitude'] - data['ra'] = data['longitude'] + data["tc"] = data["geocent_end_time"] + 1e-9 * data["geocent_end_time_ns"] + data["dec"] = data["latitude"] + data["ra"] = data["longitude"] return data -class LVKNewStyleInjectionSet(object): - """Reads injections from new-style LVK injection files - """ +class LVKNewStyleInjectionSet: + """Reads injections from new-style LVK injection files""" translation_dict = { - 'mass1_det': 'mass1', - 'mass2_det': 'mass2', - 'f22_ref_spin': 'f_ref', - 't_co_gps': 'tc', - 'd_lum': 'distance', - 'f22_start': 'f_lower', - 'cbc_model': 'approximant', - 'longAscNodes': 'long_asc_nodes', - 'eccentricity': 'eccentricity', - 'meanPerAno': 'mean_per_ano', - 't_co_gps_add': None, - 'ModeArray': 'FIXME', # FIXME (Will need special handling. Need an example) - 'ModeArrayJframe': 'FIXME', # FIXME (See ModeArray above) + "mass1_det": "mass1", + "mass2_det": "mass2", + "f22_ref_spin": "f_ref", + "t_co_gps": "tc", + "d_lum": "distance", + "f22_start": "f_lower", + "cbc_model": "approximant", + "longAscNodes": "long_asc_nodes", + "eccentricity": "eccentricity", + "meanPerAno": "mean_per_ano", + "t_co_gps_add": None, + "ModeArray": "FIXME", # FIXME (Will need special handling. Need an example) + "ModeArrayJframe": "FIXME", # FIXME (See ModeArray above) } - subdir = 'cbc_waveform_params' - + subdir = "cbc_waveform_params" + def __init__(self, sim_file): - self.inj_file = HFile(sim_file, 'r') + self.inj_file = HFile(sim_file, "r") def columns_in_pycbc_format(self): # Loop over columns in input file - + for field_name in self.inj_file[self.subdir].keys(): if field_name in self.translation_dict: yield (field_name, self.translation_dict[field_name]) @@ -136,23 +137,23 @@ class LVKNewStyleInjectionSet(object): yield (field_name, field_name) def get_coalescence_time(self): - tcs = self.inj_file[f'{self.subdir}/t_co_gps'][:] - if 't_co_gps_add' in self.inj_file[self.subdir]: - tcs += self.inj_file[f'{self.subdir}/t_co_gps_add'][:] + tcs = self.inj_file[f"{self.subdir}/t_co_gps"][:] + if "t_co_gps_add" in self.inj_file[self.subdir]: + tcs += self.inj_file[f"{self.subdir}/t_co_gps_add"][:] return tcs def pack_data_into_pycbc_format_input(self): data = {} for lvk_name, pycbc_name in self.columns_in_pycbc_format(): - if lvk_name == 't_co_gps': + if lvk_name == "t_co_gps": # Special case data[pycbc_name] = self.get_coalescence_time() - elif lvk_name == 't_co_gps_add': + elif lvk_name == "t_co_gps_add": continue else: - lvk_file_dset = self.inj_file[f'{self.subdir}/{lvk_name}'] - if lvk_file_dset.dtype.char in ['U', 'O']: - data[pycbc_name] = lvk_file_dset[:].astype('S') + lvk_file_dset = self.inj_file[f"{self.subdir}/{lvk_name}"] + if lvk_file_dset.dtype.char in ["U", "O"]: + data[pycbc_name] = lvk_file_dset[:].astype("S") else: data[pycbc_name] = lvk_file_dset[:] return data @@ -160,22 +161,26 @@ class LVKNewStyleInjectionSet(object): parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--injection-file', required=True, - help="The injection file to load. Must end in '.xml[.gz]' " - "and must contain a SimInspiral table") -parser.add_argument('--output-file', required=True, - help="The ouput file name. Must end in '.hdf'.") +parser.add_argument( + "--injection-file", + required=True, + help="The injection file to load. Must end in '.xml[.gz]' " + "and must contain a SimInspiral table", +) +parser.add_argument( + "--output-file", required=True, help="The ouput file name. Must end in '.hdf'." +) args = parser.parse_args() pycbc.init_logging(args.verbose) injclass = None -if args.injection_file.endswith(('.xml', '.xml.gz', '.xmlgz')): +if args.injection_file.endswith((".xml", ".xml.gz", ".xmlgz")): injclass = XMLInjectionSet(args.injection_file) else: # Assume a HDF file - check if new LVK format first - inj_file = HFile(args.injection_file, 'r') - if 'file_format' in inj_file.attrs or b'file_format' in inj_file.attrs: + inj_file = HFile(args.injection_file, "r") + if "file_format" in inj_file.attrs or b"file_format" in inj_file.attrs: # Assume LVK as PyCBC doesn't have this attribute injclass = LVKNewStyleInjectionSet(args.injection_file) inj_file.close() diff --git a/bin/pycbc_create_injections b/bin/pycbc_create_injections index 604a0f62935..1a97eca71e3 100644 --- a/bin/pycbc_create_injections +++ b/bin/pycbc_create_injections @@ -16,7 +16,8 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Generates injections drawn from a distribution read from a config file. +""" +Generates injections drawn from a distribution read from a config file. Config file syntax ------------------ @@ -108,71 +109,102 @@ along with any numpy ufunc. So, in the above example, mass ratio (q) is constrained to be <= 4 by using a function from the conversions module. """ -import os -import sys import argparse import logging +import os +import sys + import numpy from numpy.random import uniform import pycbc -from pycbc.inject import InjectionSet -from pycbc import distributions -from pycbc import transforms +from pycbc import distributions, transforms from pycbc.distributions import JointDistribution -from pycbc.workflow import configuration -from pycbc.workflow import WorkflowConfigParser +from pycbc.inject import InjectionSet +from pycbc.workflow import WorkflowConfigParser, configuration -parser = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter +) configuration.add_workflow_command_line_group(parser) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--ninjections', type=int, - help='Number of injections to create.') -parser.add_argument('--gps-start-time', type=int, help="Alternative to " - "ninjections argument. " - "The number will be chosen to fill the chosen time range. ") -parser.add_argument('--gps-end-time', type=int, help="Alternative to " - "ninjections argument. " - "The number will be chosen to fill the chosen time range. ") -parser.add_argument('--time-step', type=float, - help="Minimum separation between injections") -parser.add_argument('--time-window', type=float, - help="Uniform window of time to place injection into") -parser.add_argument('--seed', type=int, default=0, - help='Seed to use for the random number generator. ' - 'Default is 0.') -parser.add_argument('--output-file', required=True, - help='Output file to save to. If ends in ".xml[.gz]", ' - 'injections will be written to a sim_inspiral table ' - 'in an xml file. Otherwise, results will be written ' - 'to an hdf file.') -parser.add_argument('--dist-section', default='prior', - help='What section in the config file to load ' - 'distributions from. Default is prior.') -parser.add_argument('--variable-params-section', default='variable_params', - help='What section in the config file to load the ' - 'parameters to vary. Default is variable_params.') -parser.add_argument('--static-params-section', default="static_params", - help='What section to load the static params from. Default ' - 'is static_params.') -parser.add_argument("--force", action="store_true", default=False, - help="If the output-file already exists, overwrite it. " - "Otherwise, an OSError is raised.") +parser.add_argument("--ninjections", type=int, help="Number of injections to create.") +parser.add_argument( + "--gps-start-time", + type=int, + help="Alternative to " + "ninjections argument. " + "The number will be chosen to fill the chosen time range. ", +) +parser.add_argument( + "--gps-end-time", + type=int, + help="Alternative to " + "ninjections argument. " + "The number will be chosen to fill the chosen time range. ", +) +parser.add_argument( + "--time-step", type=float, help="Minimum separation between injections" +) +parser.add_argument( + "--time-window", type=float, help="Uniform window of time to place injection into" +) +parser.add_argument( + "--seed", + type=int, + default=0, + help="Seed to use for the random number generator. Default is 0.", +) +parser.add_argument( + "--output-file", + required=True, + help='Output file to save to. If ends in ".xml[.gz]", ' + "injections will be written to a sim_inspiral table " + "in an xml file. Otherwise, results will be written " + "to an hdf file.", +) +parser.add_argument( + "--dist-section", + default="prior", + help="What section in the config file to load " + "distributions from. Default is prior.", +) +parser.add_argument( + "--variable-params-section", + default="variable_params", + help="What section in the config file to load the " + "parameters to vary. Default is variable_params.", +) +parser.add_argument( + "--static-params-section", + default="static_params", + help="What section to load the static params from. Default is static_params.", +) +parser.add_argument( + "--force", + action="store_true", + default=False, + help="If the output-file already exists, overwrite it. " + "Otherwise, an OSError is raised.", +) opts = parser.parse_args() pycbc.init_logging(opts.verbose) if os.path.exists(opts.output_file) and not opts.force: - raise OSError("output-file already exists; use --force if you wish to " - "overwrite it.") - -if opts.ninjections and (opts.gps_start_time is not None or - opts.gps_end_time is not None or - opts.time_step is not None or - opts.time_window is not None): - raise ValueError("Cannot provide both ninjections and " - " start/end/step/window time options.") + raise OSError( + "output-file already exists; use --force if you wish to overwrite it." + ) + +if opts.ninjections and ( + opts.gps_start_time is not None + or opts.gps_end_time is not None + or opts.time_step is not None + or opts.time_window is not None +): + raise ValueError( + "Cannot provide both ninjections and start/end/step/window time options." + ) numpy.random.seed(opts.seed) @@ -180,16 +212,18 @@ logging.info("Loading config file") cp = WorkflowConfigParser.from_cli(opts) # get the vairable and static arguments from the config file -variable_params, static_params = distributions.read_params_from_config(cp, +variable_params, static_params = distributions.read_params_from_config( + cp, prior_section=opts.dist_section, vargs_section=opts.variable_params_section, - sargs_section=opts.static_params_section) -constraints = distributions.read_constraints_from_config( - cp, static_args=static_params) - -if any(cp.get_subsections('waveform_transforms')): - waveform_transforms = transforms.read_transforms_from_config(cp, - 'waveform_transforms') + sargs_section=opts.static_params_section, +) +constraints = distributions.read_constraints_from_config(cp, static_args=static_params) + +if any(cp.get_subsections("waveform_transforms")): + waveform_transforms = transforms.read_transforms_from_config( + cp, "waveform_transforms" + ) else: waveform_transforms = None write_args = variable_params @@ -199,13 +233,14 @@ logging.info("Reading distributions") dists = distributions.read_distributions_from_config(cp, opts.dist_section) # construct class that will draw the samples -randomsampler = JointDistribution(variable_params, *dists, - **{"constraints": constraints}) +randomsampler = JointDistribution( + variable_params, *dists, constraints=constraints +) if opts.ninjections: draw_size = opts.ninjections else: - draw_size = 4000 # Just a default so it's not super slow drawing large sets + draw_size = 4000 # Just a default so it's not super slow drawing large sets old_samples = None while True: @@ -216,36 +251,43 @@ while True: logging.info("Transforming to waveform transform parameters") for t in waveform_transforms: if not set(t.inputs).isdisjoint(set(static_params.keys())): - for item in list((set(t.inputs) & set(static_params.keys()) - - set(samples.fieldnames))): - samples = samples.add_fields([numpy.repeat(static_params[item], - draw_size).astype(float)], - [item]) + for item in list( + + set(t.inputs) + & set(static_params.keys()) - set(samples.fieldnames) + + ): + samples = samples.add_fields( + [numpy.repeat(static_params[item], draw_size).astype(float)], + [item], + ) samples = transforms.apply_transforms(samples, waveform_transforms) # We are drawing until we hit a time so check if we've reached it if opts.gps_start_time is not None: - if 'tc' in samples.fieldnames or 'tc' in static_params: - raise RuntimeError("A value or distribution was given for 'tc'. " - "This is not compatible with providing " - "start/end times.") + if "tc" in samples.fieldnames or "tc" in static_params: + raise RuntimeError( + "A value or distribution was given for 'tc'. " + "This is not compatible with providing " + "start/end times." + ) if old_samples is None: tstart = opts.gps_start_time else: - tstart = old_samples['tc'].max() + opts.time_step + tstart = old_samples["tc"].max() + opts.time_step tc = numpy.arange(0, draw_size) * opts.time_step + tstart tc += uniform(0, high=opts.time_window, size=draw_size).cumsum() - samples = samples.add_fields([tc], ['tc']) - + samples = samples.add_fields([tc], ["tc"]) + if old_samples is not None: samples = old_samples.append(samples) old_samples = samples if tc.max() >= opts.gps_end_time: - samples = samples[samples['tc'] < opts.gps_end_time] - logging.info('Total Injections: %s', len(samples)) + samples = samples[samples["tc"] < opts.gps_end_time] + logging.info("Total Injections: %s", len(samples)) break # We got as many samples as we needed so we can stop @@ -254,7 +296,7 @@ while True: # write results logging.info("Writing results") -write_args = [arg for arg in samples.fieldnames - if arg not in static_params.keys()] -InjectionSet.write(opts.output_file, samples, write_args, static_params, - cmd=" ".join(sys.argv)) +write_args = [arg for arg in samples.fieldnames if arg not in static_params.keys()] +InjectionSet.write( + opts.output_file, samples, write_args, static_params, cmd=" ".join(sys.argv) +) diff --git a/bin/pycbc_data_store b/bin/pycbc_data_store index 69f373b29e9..63b8f80fd20 100755 --- a/bin/pycbc_data_store +++ b/bin/pycbc_data_store @@ -1,14 +1,14 @@ #!/usr/bin/env python -""" Create HDF strain cache file -""" -import logging +"""Create HDF strain cache file""" + import argparse +import logging import pycbc -import pycbc.strain import pycbc.dq -from pycbc.fft.fftw import set_measure_level +import pycbc.strain from pycbc.events.veto import segments_to_start_end +from pycbc.fft.fftw import set_measure_level from pycbc.io.hdf import HFile set_measure_level(0) @@ -25,30 +25,37 @@ pycbc.strain.insert_strain_option_group(parser) args = parser.parse_args() pycbc.init_logging(args.verbose) -logging.info('Querying science segemnts') -segs = pycbc.dq.query_str(args.instrument, - args.science_name, - args.gps_start_time, - args.gps_end_time, - server=args.segment_server, - veto_definer=args.veto_definer_file) -logging.info('Found %s segments, %ss total', len(segs), abs(segs)) - -f = HFile(args.output_file, 'w') +logging.info("Querying science segemnts") +segs = pycbc.dq.query_str( + args.instrument, + args.science_name, + args.gps_start_time, + args.gps_end_time, + server=args.segment_server, + veto_definer=args.veto_definer_file, +) +logging.info("Found %s segments, %ss total", len(segs), abs(segs)) + +f = HFile(args.output_file, "w") starts, ends = segments_to_start_end(segs) pad = 0 if args.pad_data is None else args.pad_data -f['{}/segments/start'.format(args.channel_name)] = starts + pad -f['{}/segments/end'.format(args.channel_name)] = ends - pad +f[f"{args.channel_name}/segments/start"] = starts + pad +f[f"{args.channel_name}/segments/end"] = ends - pad for i, seg in enumerate(segs): - logging.info('Processing science segment %s/%s of duration %ss', - i, len(segs), abs(seg)) + logging.info( + "Processing science segment %s/%s of duration %ss", i, len(segs), abs(seg) + ) args.gps_start_time = seg[0] + pad args.gps_end_time = seg[1] - pad - logging.info('Reading %s-%s', seg[0], seg[1]) + logging.info("Reading %s-%s", seg[0], seg[1]) ht = pycbc.strain.from_cli(args) - f.create_dataset("{}/{}".format(args.channel_name, i), data=ht.data[:], - compression_opts=9, compression='gzip') -logging.info('Done!') + f.create_dataset( + f"{args.channel_name}/{i}", + data=ht.data[:], + compression_opts=9, + compression="gzip", + ) +logging.info("Done!") diff --git a/bin/pycbc_faithsim b/bin/pycbc_faithsim index 01a0a915146..106bc884b4f 100644 --- a/bin/pycbc_faithsim +++ b/bin/pycbc_faithsim @@ -23,122 +23,192 @@ # ============================================================================= # -import logging -from numpy import complex64 import argparse +import logging import sys -from igwn_ligolw import utils as ligolw_utils from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils +from numpy import complex64 -import pycbc.strain import pycbc.psd -from pycbc.waveform import td_approximants, fd_approximants -from pycbc.waveform import get_two_pol_waveform_filter +import pycbc.strain from pycbc import DYN_RANGE_FAC -from pycbc.types import FrequencySeries, zeros from pycbc.filter import match, overlap, sigma -from pycbc.scheme import CPUScheme, CUDAScheme from pycbc.io.ligolw import LIGOLWContentHandler +from pycbc.scheme import CPUScheme, CUDAScheme +from pycbc.types import FrequencySeries, zeros +from pycbc.waveform import fd_approximants, get_two_pol_waveform_filter, td_approximants + def update_progress(progress): - print('Progress: {}/{} members processed'.format(progress, 100)) + print(f"Progress: {progress}/{100} members processed") if progress == 100: print("Done") sys.stdout.flush() -def get_waveform(approximant, phase_order, amplitude_order, spin_order, tapering, template_params, start_frequency, sample_rate, length): +def get_waveform( + approximant, + phase_order, + amplitude_order, + spin_order, + tapering, + template_params, + start_frequency, + sample_rate, + length, +): delta_f = sample_rate / length - vecplus = FrequencySeries(zeros(filter_n), delta_f=delta_f, - dtype=complex64) - veccross = FrequencySeries(zeros(filter_n), delta_f=delta_f, - dtype=complex64) + vecplus = FrequencySeries(zeros(filter_n), delta_f=delta_f, dtype=complex64) + veccross = FrequencySeries(zeros(filter_n), delta_f=delta_f, dtype=complex64) if tapering is not None: - curr_taper = tapering - elif hasattr(template_params, 'taper'): - curr_taper = template_params.taper + curr_taper = tapering + elif hasattr(template_params, "taper"): + curr_taper = template_params.taper else: - curr_taper = None + curr_taper = None # NOTE: for now only hplus is used! For precessing faithsims one would want # to also specify a polarization phase, or "u_val". - hplus, hcross = get_two_pol_waveform_filter(vecplus, veccross, - template_params, approximant=approximant, spin_order=spin_order, - phase_order=phase_order, delta_t=1.0 / sample_rate, delta_f=delta_f, - f_lower=start_frequency, amplitude_order=amplitude_order, - taper=curr_taper) + hplus, hcross = get_two_pol_waveform_filter( + vecplus, + veccross, + template_params, + approximant=approximant, + spin_order=spin_order, + phase_order=phase_order, + delta_t=1.0 / sample_rate, + delta_f=delta_f, + f_lower=start_frequency, + amplitude_order=amplitude_order, + taper=curr_taper, + ) + + return hplus * DYN_RANGE_FAC - return hplus*DYN_RANGE_FAC ############################################################################### aprs = list(set(td_approximants() + fd_approximants())) psd_names = pycbc.psd.get_lalsim_psd_list() -#File I/O Settings -taper_choices = ["start","end","startend"] -parser = argparse.ArgumentParser(usage='', - description="Calculate faithfulness for a set of waveforms.") +# File I/O Settings +taper_choices = ["start", "end", "startend"] +parser = argparse.ArgumentParser( + usage="", description="Calculate faithfulness for a set of waveforms." +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--param-file", dest="bank_file", metavar="FILE", - help="Sngl or Sim Inspiral Table containing waveform " - "parameters.") -parser.add_argument("--match-file", dest="out_file", metavar="FILE", - help="File to output match results to.") - -#Waveform generation Settings -parser.add_argument("--waveform1-approximant", choices=aprs, - help="Waveform 1's approximant.") -parser.add_argument("--waveform1-phase-order", type=int, default=-1, - help="PN order to use for the phase") -parser.add_argument("--waveform1-amplitude-order", default=-1, type=int, - help="PN order to use for the amplitude.") -parser.add_argument("--waveform1-spin-order", default=-1, type=int, - help="Spin terms up to the given pN order are included") -parser.add_argument("--waveform1-start-frequency", type=float, - help="Starting frequency for waveform generation.") -parser.add_argument("--waveform1-taper-template", choices=taper_choices, - default=None, - help="For time-domain approximants, taper the start and/or" - " end of the waveform before FFTing. This can also be" - " provided in the sim_inspiral table. Providing this" - " option will override any entry in that table.") - -parser.add_argument("--waveform2-approximant", choices=aprs, - help="Waveform 2's approximant.") -parser.add_argument("--waveform2-phase-order", type=int, default=-1, - help="PN order to use for the phase") -parser.add_argument("--waveform2-amplitude-order", default=-1, type=int, - help="PN order to use for the amplitude.") -parser.add_argument("--waveform2-spin-order", default=-1, type=int, - help="Spin terms up to the given pN order are included") -parser.add_argument("--waveform2-start-frequency", type=float, - help="Starting frequency for waveform generation.") -parser.add_argument("--waveform2-taper-template", choices=taper_choices, - default=None, - help="For time-domain approximants, taper the start and/or" - " end of the waveform before FFTing. This can also be" - " provided in the sim_inspiral table. Providing this" - " option will override any entry in that table.") - -#Filter Settings -parser.add_argument('--filter-low-frequency-cutoff', metavar='FREQ', - type=float, - help='Low frequency cutoff of matched filter.') -parser.add_argument('--filter-high-frequency-cutoff', metavar='FREQ', - type=float, - help='High frequency cutoff of matched filter.') -parser.add_argument("--filter-sample-rate", type=float, - help="Filter Sample Rate [Hz]") -parser.add_argument("--filter-waveform-length", type=int, - help="Length of waveform for filtering, should be longer " - "than all waveforms and include some padding.") - -parser.add_argument("--cuda", action="store_true", - help="Use CUDA for calculations.") +parser.add_argument( + "--param-file", + dest="bank_file", + metavar="FILE", + help="Sngl or Sim Inspiral Table containing waveform parameters.", +) +parser.add_argument( + "--match-file", + dest="out_file", + metavar="FILE", + help="File to output match results to.", +) + +# Waveform generation Settings +parser.add_argument( + "--waveform1-approximant", choices=aprs, help="Waveform 1's approximant." +) +parser.add_argument( + "--waveform1-phase-order", + type=int, + default=-1, + help="PN order to use for the phase", +) +parser.add_argument( + "--waveform1-amplitude-order", + default=-1, + type=int, + help="PN order to use for the amplitude.", +) +parser.add_argument( + "--waveform1-spin-order", + default=-1, + type=int, + help="Spin terms up to the given pN order are included", +) +parser.add_argument( + "--waveform1-start-frequency", + type=float, + help="Starting frequency for waveform generation.", +) +parser.add_argument( + "--waveform1-taper-template", + choices=taper_choices, + default=None, + help="For time-domain approximants, taper the start and/or" + " end of the waveform before FFTing. This can also be" + " provided in the sim_inspiral table. Providing this" + " option will override any entry in that table.", +) + +parser.add_argument( + "--waveform2-approximant", choices=aprs, help="Waveform 2's approximant." +) +parser.add_argument( + "--waveform2-phase-order", + type=int, + default=-1, + help="PN order to use for the phase", +) +parser.add_argument( + "--waveform2-amplitude-order", + default=-1, + type=int, + help="PN order to use for the amplitude.", +) +parser.add_argument( + "--waveform2-spin-order", + default=-1, + type=int, + help="Spin terms up to the given pN order are included", +) +parser.add_argument( + "--waveform2-start-frequency", + type=float, + help="Starting frequency for waveform generation.", +) +parser.add_argument( + "--waveform2-taper-template", + choices=taper_choices, + default=None, + help="For time-domain approximants, taper the start and/or" + " end of the waveform before FFTing. This can also be" + " provided in the sim_inspiral table. Providing this" + " option will override any entry in that table.", +) + +# Filter Settings +parser.add_argument( + "--filter-low-frequency-cutoff", + metavar="FREQ", + type=float, + help="Low frequency cutoff of matched filter.", +) +parser.add_argument( + "--filter-high-frequency-cutoff", + metavar="FREQ", + type=float, + help="High frequency cutoff of matched filter.", +) +parser.add_argument("--filter-sample-rate", type=float, help="Filter Sample Rate [Hz]") +parser.add_argument( + "--filter-waveform-length", + type=int, + help="Length of waveform for filtering, should be longer " + "than all waveforms and include some padding.", +) + +parser.add_argument("--cuda", action="store_true", help="Use CUDA for calculations.") # Insert the PSD options pycbc.psd.insert_psd_option_group(parser) @@ -156,14 +226,15 @@ else: ctx = CPUScheme() # Load in the waveform1 bank file -indoc = ligolw_utils.load_filename(options.bank_file, False, - contenthandler=LIGOLWContentHandler) -try : +indoc = ligolw_utils.load_filename( + options.bank_file, False, contenthandler=LIGOLWContentHandler +) +try: waveform_table = lsctables.SnglInspiralTable.get_table(indoc) except ValueError: waveform_table = lsctables.SimInspiralTable.get_table(indoc) -# open the output file where the max overlaps over the bank are stored +# open the output file where the max overlaps over the bank are stored fout = open(options.out_file, "w") logging.info("Writing matches to " + options.out_file) @@ -181,9 +252,15 @@ if options.psd_estimation: else: strain = None -psd = pycbc.psd.from_cli(options, length=filter_n, delta_f=delta_f, - low_frequency_cutoff=options.filter_low_frequency_cutoff, strain=strain, - dyn_range_factor=DYN_RANGE_FAC, precision='single') +psd = pycbc.psd.from_cli( + options, + length=filter_n, + delta_f=delta_f, + low_frequency_cutoff=options.filter_low_frequency_cutoff, + strain=strain, + dyn_range_factor=DYN_RANGE_FAC, + precision="single", +) matches = [] overlaps = [] @@ -192,53 +269,71 @@ s1s = [] s2s = [] logging.info("Calculating Overlaps") with ctx: - index = 0 + index = 0 # Calculate the overlaps for waveform_params in waveform_table: index += 1 if options.verbose: - update_progress(index*100/len(waveform_table)) + update_progress(index * 100 / len(waveform_table)) try: - htilde1 = get_waveform(options.waveform1_approximant, - options.waveform1_phase_order, - options.waveform1_amplitude_order, - options.waveform1_spin_order, - options.waveform1_taper_template, - waveform_params, - options.waveform1_start_frequency, - options.filter_sample_rate, - filter_N) - - htilde2 = get_waveform(options.waveform2_approximant, - options.waveform2_phase_order, - options.waveform2_amplitude_order, - options.waveform2_spin_order, - options.waveform2_taper_template, - waveform_params, - options.waveform2_start_frequency, - options.filter_sample_rate, - filter_N) - - m,i = match(htilde1, htilde2, psd=psd, + htilde1 = get_waveform( + options.waveform1_approximant, + options.waveform1_phase_order, + options.waveform1_amplitude_order, + options.waveform1_spin_order, + options.waveform1_taper_template, + waveform_params, + options.waveform1_start_frequency, + options.filter_sample_rate, + filter_N, + ) + + htilde2 = get_waveform( + options.waveform2_approximant, + options.waveform2_phase_order, + options.waveform2_amplitude_order, + options.waveform2_spin_order, + options.waveform2_taper_template, + waveform_params, + options.waveform2_start_frequency, + options.filter_sample_rate, + filter_N, + ) + + m, i = match( + htilde1, + htilde2, + psd=psd, low_frequency_cutoff=options.filter_low_frequency_cutoff, - high_frequency_cutoff=options.filter_high_frequency_cutoff) + high_frequency_cutoff=options.filter_high_frequency_cutoff, + ) - o = overlap(htilde1, htilde2, psd=psd, + o = overlap( + htilde1, + htilde2, + psd=psd, low_frequency_cutoff=options.filter_low_frequency_cutoff, - high_frequency_cutoff=options.filter_high_frequency_cutoff) + high_frequency_cutoff=options.filter_high_frequency_cutoff, + ) - s1 = sigma(htilde1, psd=psd, + s1 = sigma( + htilde1, + psd=psd, low_frequency_cutoff=options.filter_low_frequency_cutoff, - high_frequency_cutoff=options.filter_high_frequency_cutoff) - s2 = sigma(htilde2, psd=psd, + high_frequency_cutoff=options.filter_high_frequency_cutoff, + ) + s2 = sigma( + htilde2, + psd=psd, low_frequency_cutoff=options.filter_low_frequency_cutoff, - high_frequency_cutoff=options.filter_high_frequency_cutoff) + high_frequency_cutoff=options.filter_high_frequency_cutoff, + ) matches.append(m) overlaps.append(o) if i > filter_n: - i = i - filter_N - time_offsets.append(i * 1./options.filter_sample_rate) + i = i - filter_N + time_offsets.append(i * 1.0 / options.filter_sample_rate) s1s.append(s1) s2s.append(s2) except Exception as e: @@ -250,7 +345,7 @@ with ctx: s1s.append(-1) s2s.append(-1) -#Output the overlaps to a file +# Output the overlaps to a file for m, o, i, s1, s2 in zip(matches, overlaps, time_offsets, s1s, s2s): - match_str= "%5.5f %5.5f %5.5f %5.5f %5.5f\n" % (m, o, i, s1, s2) + match_str = "%5.5f %5.5f %5.5f %5.5f %5.5f\n" % (m, o, i, s1, s2) fout.write(match_str) diff --git a/bin/pycbc_faithsim_collect_results b/bin/pycbc_faithsim_collect_results index 60ca40d41a3..fcf63a90d3e 100755 --- a/bin/pycbc_faithsim_collect_results +++ b/bin/pycbc_faithsim_collect_results @@ -2,13 +2,13 @@ """ Program for collecting the results of pycbc_faithsim comparing two approximants -computing the match between them and creating a .dat file with the results. +computing the match between them and creating a .dat file with the results. """ import argparse -import numpy as np -from igwn_ligolw import utils, lsctables +import numpy as np +from igwn_ligolw import lsctables, utils from pycbc import add_common_pycbc_options, init_logging from pycbc.io.ligolw import LIGOLWContentHandler @@ -97,7 +97,7 @@ for i in range(len(match_files)): md = np.loadtxt(match, dtype=dtypem) if md.size == 0: continue - except IOError: + except OSError: continue pdata = np.zeros(len(bt), dtype=dtypeo) diff --git a/bin/pycbc_fit_sngl_trigs b/bin/pycbc_fit_sngl_trigs index 7acde3384d3..4c6531376ee 100644 --- a/bin/pycbc_fit_sngl_trigs +++ b/bin/pycbc_fit_sngl_trigs @@ -13,121 +13,175 @@ # Public License for more details. -import argparse, logging +import argparse +import logging from matplotlib import use -use('Agg') -from matplotlib import pyplot as plt + +use("Agg") import numpy as np +from matplotlib import pyplot as plt import pycbc -from pycbc import io, events, bin_utils +from pycbc import bin_utils, events, io from pycbc.events import ranking from pycbc.events import trigger_fits as trstats #### DEFINITIONS AND FUNCTIONS #### stat_dict = { - "new_snr" : ranking.newsnr, - "effective_snr" : ranking.effsnr, - "snr" : lambda snr, rchisq : snr, - "snronchi" : lambda snr, rchisq : snr / (rchisq ** 0.5) + "new_snr": ranking.newsnr, + "effective_snr": ranking.effsnr, + "snr": lambda snr, rchisq: snr, + "snronchi": lambda snr, rchisq: snr / (rchisq**0.5), } + def get_stat(statchoice, snr, rchisq, fac): if fac is not None: if statchoice not in ["new_snr", "effective_snr"]: raise RuntimeError("Can't use --stat-factor with this statistic!") return stat_dict[statchoice](snr, rchisq, fac) - else: - return stat_dict[statchoice](snr, rchisq) + return stat_dict[statchoice](snr, rchisq) + def get_bins(opt, pmin, pmax): if opt.bin_spacing == "linear": return bin_utils.LinearBins(pmin, pmax, opt.num_bins) - elif opt.bin_spacing == "log": + if opt.bin_spacing == "log": return bin_utils.LogarithmicBins(pmin, pmax, opt.num_bins) - elif opt.bin_spacing == "irregular": + if opt.bin_spacing == "irregular": return bin_utils.IrregularBins(opt.bin_edges) + #### MAIN #### -parser = argparse.ArgumentParser(usage="", +parser = argparse.ArgumentParser( + usage="", description="Perform maximum-likelihood fits of single inspiral trigger" - "distributions to various functions") + "distributions to various functions", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--inputs", nargs="+", - help="Input file or space-separated list of input files " - "containing single triggers. Currently .xml(.gz) " - "and .hdf supported. Required") -parser.add_argument("--bank-file", default=None, - help="hdf file containing template parameters; required " - "if fitting over trigger masses and/or spins with hdf " - "format triggers") -parser.add_argument("--veto-file", default=None, - help="File in .xml format with veto segments to apply to " - "triggers before fitting") -parser.add_argument("--veto-segment-name", default=None, - help="Name of veto segments to apply. Optional: if not " - "given, all segments for a given ifo will be used") -parser.add_argument("--output", required=True, - help="Location for output file containing fit coefficients" - ". Required") -parser.add_argument("--plot-dir", default=None, - help="Plot the fits made, the variation of fitting " - "coefficients and the Kolmogorov-Smirnov test values " - "and save to the specified directory") -parser.add_argument("--user-tag", default="", - help="Put a possibly informative string in the names of " - "plot files") -parser.add_argument("--ifos", nargs="+", - help="Ifo or space-separated list of ifos to select " - "triggers to be fit. Required") -parser.add_argument("--fit-function", - choices=["exponential", "rayleigh", "power"], - help="Functional form for the maximum likelihood fit") -parser.add_argument("--sngl-stat", default="new_snr", - choices=["snr", "snronchi", "effective_snr", "new_snr"], - help="Function of SNR and chisq to perform fits with") +parser.add_argument( + "--inputs", + nargs="+", + help="Input file or space-separated list of input files " + "containing single triggers. Currently .xml(.gz) " + "and .hdf supported. Required", +) +parser.add_argument( + "--bank-file", + default=None, + help="hdf file containing template parameters; required " + "if fitting over trigger masses and/or spins with hdf " + "format triggers", +) +parser.add_argument( + "--veto-file", + default=None, + help="File in .xml format with veto segments to apply to triggers before fitting", +) +parser.add_argument( + "--veto-segment-name", + default=None, + help="Name of veto segments to apply. Optional: if not " + "given, all segments for a given ifo will be used", +) +parser.add_argument( + "--output", + required=True, + help="Location for output file containing fit coefficients. Required", +) +parser.add_argument( + "--plot-dir", + default=None, + help="Plot the fits made, the variation of fitting " + "coefficients and the Kolmogorov-Smirnov test values " + "and save to the specified directory", +) +parser.add_argument( + "--user-tag", + default="", + help="Put a possibly informative string in the names of plot files", +) +parser.add_argument( + "--ifos", + nargs="+", + help="Ifo or space-separated list of ifos to select triggers to be fit. Required", +) +parser.add_argument( + "--fit-function", + choices=["exponential", "rayleigh", "power"], + help="Functional form for the maximum likelihood fit", +) +parser.add_argument( + "--sngl-stat", + default="new_snr", + choices=["snr", "snronchi", "effective_snr", "new_snr"], + help="Function of SNR and chisq to perform fits with", +) # FIXME: how to parse various possible mixtures of chisq -#parser.add_argument("--chisq-choice", default="power", +# parser.add_argument("--chisq-choice", default="power", # choices=["power","bank","auto","maxpowerbank"], # help="Which chisq value or values to form the " # "single-trigger statistic with") -parser.add_argument("--stat-factor", type=float, - help="Adjustable magic number used in some sngl " - "statistics. Values commonly used: 6 for new_snr, 250 " - "or 50 for effective_snr") +parser.add_argument( + "--stat-factor", + type=float, + help="Adjustable magic number used in some sngl " + "statistics. Values commonly used: 6 for new_snr, 250 " + "or 50 for effective_snr", +) default_thresh = 5.5 cuts = parser.add_mutually_exclusive_group() -cuts.add_argument("--snr-threshold", type=float, - help="Only fit triggers with SNR above this threshold. " - "Default %.1f" % default_thresh) -cuts.add_argument("--trig-filter", default=None, - help="Filter function applied to trigger properties: only " - "implemented for hdf format. May include names from math, " - "numpy as 'np', pycbc.events and pycbc.pnutils. Ex. " - "'ranking.newsnr(self.snr, self.bank_chisq/self.bank_chisq_" - "dof) > 6.)'") -parser.add_argument("--stat-threshold", nargs="+", type=float, - help="Only fit triggers with statistic value above this " - "threshold : can be a space-separated list, then a fit " - "will be done for each threshold. Required. Typical " - "values 6.5 6.75 7") -parser.add_argument("--fit-param", required=True, - help="Parameter over which to estimate variation of the " - "fits. Required. Must be a SnglInspiralTable column for " - ".xml input or the name of a dataset for .hdf") +cuts.add_argument( + "--snr-threshold", + type=float, + help="Only fit triggers with SNR above this threshold. " + "Default %.1f" % default_thresh, +) +cuts.add_argument( + "--trig-filter", + default=None, + help="Filter function applied to trigger properties: only " + "implemented for hdf format. May include names from math, " + "numpy as 'np', pycbc.events and pycbc.pnutils. Ex. " + "'ranking.newsnr(self.snr, self.bank_chisq/self.bank_chisq_" + "dof) > 6.)'", +) +parser.add_argument( + "--stat-threshold", + nargs="+", + type=float, + help="Only fit triggers with statistic value above this " + "threshold : can be a space-separated list, then a fit " + "will be done for each threshold. Required. Typical " + "values 6.5 6.75 7", +) +parser.add_argument( + "--fit-param", + required=True, + help="Parameter over which to estimate variation of the " + "fits. Required. Must be a SnglInspiralTable column for " + ".xml input or the name of a dataset for .hdf", +) # FIXME: allow for math functions of columns. Ex. 1./mtotal") -parser.add_argument("--bin-spacing", choices=["linear", "log", "irregular"], - help="How to space parameter bin edges") +parser.add_argument( + "--bin-spacing", + choices=["linear", "log", "irregular"], + help="How to space parameter bin edges", +) binopt = parser.add_mutually_exclusive_group(required=True) -binopt.add_argument("--num-bins", type=int, - help="Number of regularly spaced bins to use over the " - " parameter") -binopt.add_argument("--irregular-bins", - help="Comma-separated list of parameter bin edges. " - "Required if --bin-spacing = irregular") +binopt.add_argument( + "--num-bins", + type=int, + help="Number of regularly spaced bins to use over the parameter", +) +binopt.add_argument( + "--irregular-bins", + help="Comma-separated list of parameter bin edges. " + "Required if --bin-spacing = irregular", +) opt = parser.parse_args() @@ -138,26 +192,28 @@ paramname = opt.fit_param.replace("_", " ") paramtag = opt.fit_param.replace("_", "") if opt.plot_dir is not None: - outdir = opt.plot_dir if opt.plot_dir.endswith('/') else opt.plot_dir+'/' + outdir = opt.plot_dir if opt.plot_dir.endswith("/") else opt.plot_dir + "/" -outfile = open(opt.output, 'w') +outfile = open(opt.output, "w") ## Check option logic if opt.bin_spacing == "irregular": if opt.irregular_bins is None: raise RuntimeError("Must specify a list of irregular bin edges!") else: - opt.bin_edges = [float(b) for b in opt.irregular_bins.split(',')] + opt.bin_edges = [float(b) for b in opt.irregular_bins.split(",")] -if opt.inputs[0].split('.')[-1] == "xml" or \ - opt.inputs[0].split('.')[-2:] == ["xml", "gz"]: +if opt.inputs[0].split(".")[-1] == "xml" or opt.inputs[0].split(".")[-2:] == [ + "xml", + "gz", +]: trigformat = "xml" elif "hdf" in opt.inputs[0]: trigformat = "hdf" # columns for reading triggers without template file info cols = ["snr", "chisq", "template_duration", "template_id", "end_time"] else: - print('inputs:', opt.inputs) + print("inputs:", opt.inputs) raise RuntimeError("I didn't recognize the file format for the inputs!") if opt.trig_filter and trigformat != "hdf": @@ -165,9 +221,9 @@ if opt.trig_filter and trigformat != "hdf": # hack default values of SNR threshold and filter function if not opt.snr_threshold and not opt.trig_filter: opt.snr_threshold = default_thresh - filter_func = 'self.snr > %f' % opt.snr_threshold + filter_func = "self.snr > %f" % opt.snr_threshold elif opt.snr_threshold: - filter_func = 'self.snr > %f' % opt.snr_threshold + filter_func = "self.snr > %f" % opt.snr_threshold else: filter_func = opt.trig_filter @@ -179,10 +235,23 @@ fits = {} stdev = {} ks_prob = {} -histcolors = ['r',(1.0,0.6,0),'y','g','c','b','m','k',(0.8,0.25,0),(0.25,0.8,0)] +histcolors = [ + "r", + (1.0, 0.6, 0), + "y", + "g", + "c", + "b", + "m", + "k", + (0.8, 0.25, 0), + (0.25, 0.8, 0), +] # header -outfile.write("# ifo threshold lower upper templates triggers alpha sig_alpha ks_prob\n") +outfile.write( + "# ifo threshold lower upper templates triggers alpha sig_alpha ks_prob\n" +) done_vetoing = False for ifo in opt.ifos: @@ -190,9 +259,12 @@ for ifo in opt.ifos: # to the 'get_column' syntax used by SnglInspiralUtils if trigformat == "xml": from pylal import SnglInspiralUtils as sniuls - sngls = sniuls.ReadSnglInspiralFromFiles(opt.inputs, - filterFunc=lambda s: s.snr>opt.snr_threshold, - verbose=opt.verbose) + + sngls = sniuls.ReadSnglInspiralFromFiles( + opt.inputs, + filterFunc=lambda s: s.snr > opt.snr_threshold, + verbose=opt.verbose, + ) sngls.ifocut(ifo, inplace=True) elif trigformat == "hdf": if opt.bank_file: @@ -207,9 +279,9 @@ for ifo in opt.ifos: ) done_vetoing = True else: - sngls = io.hdf.DataFromFiles(opt.inputs, group=ifo, - columnlist=cols, - filter_func=filter_func) + sngls = io.hdf.DataFromFiles( + opt.inputs, group=ifo, columnlist=cols, filter_func=filter_func + ) if opt.plot_dir: plotbase = outdir + ifo + "-" + opt.user_tag @@ -230,34 +302,40 @@ for ifo in opt.ifos: mass2 = sngls.get_column("mass2") parvals = sngls.get_column(opt.fit_param) rchisq = chisq / (2 * chisq_dof - 2) - #bank_rchisq = sngls.get_column('bank_chisq') / sngls.get_column('bank_chisq_dof') - #rchisq = np.maximum(rchisq, bank_rchisq) + # bank_rchisq = sngls.get_column('bank_chisq') / sngls.get_column('bank_chisq_dof') + # rchisq = np.maximum(rchisq, bank_rchisq) if opt.veto_file and not done_vetoing: - logging.info("Applying vetoes from %s for ifo %s" % - (opt.veto_file, ifo)) + logging.info("Applying vetoes from %s for ifo %s" % (opt.veto_file, ifo)) keep_idx, segs = events.veto.indices_outside_segments( - sngls.get_column("end_time").astype(int), - [opt.veto_file], ifo=ifo, - segment_name=opt.veto_segment_name) + sngls.get_column("end_time").astype(int), + [opt.veto_file], + ifo=ifo, + segment_name=opt.veto_segment_name, + ) snr = snr[keep_idx] rchisq = rchisq[keep_idx] - if mass1 is not None: mass1 = mass1[keep_idx] - if mass2 is not None: mass2 = mass2[keep_idx] - if tid is not None: tid = tid[keep_idx] + if mass1 is not None: + mass1 = mass1[keep_idx] + if mass2 is not None: + mass2 = mass2[keep_idx] + if tid is not None: + tid = tid[keep_idx] parvals = parvals[keep_idx] logging.info("%i trigs remain" % len(snr)) logging.info("Calculating ranking statistic values") statv = get_stat(opt.sngl_stat, snr, rchisq, opt.stat_factor) - logging.info("Parameter range of triggers: %f - %f" % - (min(parvals), max(parvals))) + logging.info("Parameter range of triggers: %f - %f" % (min(parvals), max(parvals))) if opt.bin_spacing == "irregular": - logging.info("Removing triggers outside bin range %f - %f" % - (min(opt.bin_edges), max(opt.bin_edges))) - in_range = np.logical_and(parvals >= min(opt.bin_edges), - parvals <= max(opt.bin_edges)) + logging.info( + "Removing triggers outside bin range %f - %f" + % (min(opt.bin_edges), max(opt.bin_edges)) + ) + in_range = np.logical_and( + parvals >= min(opt.bin_edges), parvals <= max(opt.bin_edges) + ) statv = statv[in_range] parvals = parvals[in_range] logging.info("%i remain" % len(statv)) @@ -302,8 +380,7 @@ for ifo in opt.ifos: else: mass1_inbin = mass1[pind == i] mass2_inbin = mass2[pind == i] - mass_tuples = [(m1, m2) for m1, m2 in - zip(mass1_inbin, mass2_inbin)] + mass_tuples = [(m1, m2) for m1, m2 in zip(mass1_inbin, mass2_inbin)] numtmpl = len(set(mass_tuples)) templates[ifo][i] = numtmpl vals_inbin = statv[pind == i] @@ -313,82 +390,137 @@ for ifo in opt.ifos: continue # FIXME: Also allow to use dynamically determined threshold # calculated via Nth loudest trig - #thresh = trstats.tail_threshold(vals_inbin, N=nmax) - #counts[ifo][th][i] = nmax + # thresh = trstats.tail_threshold(vals_inbin, N=nmax) + # counts[ifo][th][i] = nmax # do the fit alpha, sig_alpha = trstats.fit_above_thresh( - opt.fit_function, vals_inbin, th) - #alpha, sig_alpha = trstats.fit_above_thresh( + opt.fit_function, vals_inbin, th + ) + # alpha, sig_alpha = trstats.fit_above_thresh( # opt.fit_function, vals_inbin, thresh) fits[ifo][th][i] = alpha stdev[ifo][th][i] = sig_alpha _, ks_prob[ifo][th][i] = trstats.KS_test( - opt.fit_function, vals_inbin, alpha, th) - outfile.write("%s %.2f %.3g %.3g %d %d %.3f %.3f %.3g\n" % - (ifo, th, lower, upper, numtmpl, counts[ifo][th][i], alpha, - sig_alpha, ks_prob[ifo][th][i])) + opt.fit_function, vals_inbin, alpha, th + ) + outfile.write( + "%s %.2f %.3g %.3g %d %d %.3f %.3f %.3g\n" + % ( + ifo, + th, + lower, + upper, + numtmpl, + counts[ifo][th][i], + alpha, + sig_alpha, + ks_prob[ifo][th][i], + ) + ) # add histogram to plot if opt.plot_dir: histcounts, edges = np.histogram(vals_inbin, bins=50) cum_counts = histcounts[::-1].cumsum()[::-1] binlabel = r"%.2g - %.2g" % (lower, upper) - plt.semilogy(edges[:-1], cum_counts, linewidth=2, - color=histcolors[i], label=binlabel, alpha=0.6) - - plt.semilogy(plotrange, counts[ifo][th][i] * \ - trstats.cum_fit(opt.fit_function, plotrange, alpha, th), - "--", color=histcolors[i], - label=r"$\alpha = $%.2f $\pm$ %.2f" % (alpha, sig_alpha)) - plt.semilogy(plotrange, counts[ifo][th][i] * \ - trstats.cum_fit(opt.fit_function, plotrange, alpha + \ - sig_alpha, th), ":", alpha=0.6, color=histcolors[i]) - plt.semilogy(plotrange, counts[ifo][th][i] * \ - trstats.cum_fit(opt.fit_function, plotrange, alpha - \ - sig_alpha, th), ":", alpha=0.6, color=histcolors[i]) + plt.semilogy( + edges[:-1], + cum_counts, + linewidth=2, + color=histcolors[i], + label=binlabel, + alpha=0.6, + ) + + plt.semilogy( + plotrange, + counts[ifo][th][i] + * trstats.cum_fit(opt.fit_function, plotrange, alpha, th), + "--", + color=histcolors[i], + label=r"$\alpha = $%.2f $\pm$ %.2f" % (alpha, sig_alpha), + ) + plt.semilogy( + plotrange, + counts[ifo][th][i] + * trstats.cum_fit( + opt.fit_function, plotrange, alpha + sig_alpha, th + ), + ":", + alpha=0.6, + color=histcolors[i], + ) + plt.semilogy( + plotrange, + counts[ifo][th][i] + * trstats.cum_fit( + opt.fit_function, plotrange, alpha - sig_alpha, th + ), + ":", + alpha=0.6, + color=histcolors[i], + ) if opt.plot_dir: leg = plt.legend(labelspacing=0.2) plt.setp(leg.get_texts(), fontsize=11) - plt.ylim(0.7, 2*maxcount) - plt.xlim(0.9*min(opt.stat_threshold), 1.1*max(plotrange)) + plt.ylim(0.7, 2 * maxcount) + plt.xlim(0.9 * min(opt.stat_threshold), 1.1 * max(plotrange)) plt.grid() - plt.title(ifo + " " + statname + " distribution split by " + \ - paramname) + plt.title(ifo + " " + statname + " distribution split by " + paramname) plt.xlabel(statname, size="large") plt.ylabel("Cumulative number", size="large") - dest = plotbase + "_" + opt.sngl_stat + "_cdf_by_" + \ - paramtag[0:3] + "_fit_thresh_" + str(th) + ".png" + dest = ( + plotbase + + "_" + + opt.sngl_stat + + "_cdf_by_" + + paramtag[0:3] + + "_fit_thresh_" + + str(th) + + ".png" + ) logging.info("Saving plot to %s" % dest) plt.savefig(dest) plt.close() # make plots of alpha and KS significance for ifos having triggers if opt.plot_dir: - for ifo in fits.keys(): + for ifo in fits: for th in opt.stat_threshold: - plt.errorbar(pbins.centres(), [fits[ifo][th][i] for i in binind], - yerr=[stdev[ifo][th][i] for i in binind], fmt="+-", - label=ifo + " fit above %.2f" % th) - if opt.bin_spacing == "log": plt.semilogx() + plt.errorbar( + pbins.centres(), + [fits[ifo][th][i] for i in binind], + yerr=[stdev[ifo][th][i] for i in binind], + fmt="+-", + label=ifo + " fit above %.2f" % th, + ) + if opt.bin_spacing == "log": + plt.semilogx() plt.grid() plt.legend(loc="best") plt.xlabel(paramname, size="large") - plt.ylabel(r"fit parameter $\alpha$", size='large') - plt.savefig(plotbase + '_alpha_vs_' + paramtag[0:3] + '.png') + plt.ylabel(r"fit parameter $\alpha$", size="large") + plt.savefig(plotbase + "_alpha_vs_" + paramtag[0:3] + ".png") plt.close() for th in opt.stat_threshold: - plt.plot(pbins.centres(), [ks_prob[ifo][th][i] for i in binind], - '+--', label=ifo+' KS prob, thresh %.2f' % th) - if opt.bin_spacing == 'log': plt.loglog() - else : plt.semilogy() + plt.plot( + pbins.centres(), + [ks_prob[ifo][th][i] for i in binind], + "+--", + label=ifo + " KS prob, thresh %.2f" % th, + ) + if opt.bin_spacing == "log": + plt.loglog() + else: + plt.semilogy() plt.grid() - leg = plt.legend(loc='best', labelspacing=0.2) + leg = plt.legend(loc="best", labelspacing=0.2) plt.setp(leg.get_texts(), fontsize=11) - plt.xlabel(paramname, size='large') - plt.ylabel('KS test p-value') - plt.savefig(plotbase + '_KS_prob_vs_' + paramtag[0:3] + '.png') + plt.xlabel(paramname, size="large") + plt.ylabel("KS test p-value") + plt.savefig(plotbase + "_KS_prob_vs_" + paramtag[0:3] + ".png") plt.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/pycbc_generate_mock_data b/bin/pycbc_generate_mock_data index ceed90f68c7..e90aa024cdf 100644 --- a/bin/pycbc_generate_mock_data +++ b/bin/pycbc_generate_mock_data @@ -15,94 +15,129 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -__description__ = \ -"""Generate the mock data frame files for a given network of detectors. +__description__ = """Generate the mock data frame files for a given network of detectors. It contains Gaussian noise, population injections, and glitches. """ -import os import argparse +import logging +import os + import numpy + import pycbc from pycbc.types import MultiDetOptionAction -import logging - parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument("--ifo-list", nargs="+", type=str, - help="List of detector names, e.g. H1 L1 V1") - -parser.add_argument('--low-frequency-cutoff', nargs='+', - action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', type=float, - help='For each detector, provide minimum frequency for' - 'the data generation.') - - -parser.add_argument('--channel-name', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', type=str, - help='For each detector, provide channel names for the' - 'frame files.') - -parser.add_argument('--psd-model', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', type=str, - help='For each detector, provide PSD option for' - 'the noise generation.') - -parser.add_argument('--fake-strain-seed', nargs='+', - action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', type=int, - help='For each detector, provide a random number seed') - -parser.add_argument('--fake-strain-from-file', nargs='+', - action=MultiDetOptionAction, - metavar='DETECTOR:FILE', type=str, - help='For each detector, provide a file (e.g. txt) to ' - 'load the fake strain from, instead of generating it.') - -parser.add_argument('--gps-start-time', type=int, - help='GPS start time for frame files') - -parser.add_argument('--gps-end-time', type=int, - help='GPS end time for frame files') - -parser.add_argument('--sample-rate', type=float, - help='Sample rate of the frames to be generated') - -parser.add_argument('--fake-strain-sample-rate', type=float, - help='Sample rate for the fake strain generation (defaults to sample-rate if not set)') - -parser.add_argument('--fake-strain-filter-duration', type=float, - help='Duration of the filter used for coloring the noise') - -parser.add_argument('--injection-file', type=str, - help='Injection file containing parameters of' - 'injection population') - -parser.add_argument('--output-path', - help="Path for the output frame files", - default='.', type=str) - -parser.add_argument('--tag', help='Provide your tag for naming frame file', - default=None) - -parser.add_argument('--len-arm', type=float, default=None, - help="The arm length of LISA, in the unit of 'm'") - -parser.add_argument('--acc-noise-level', type=float, - help="The level of acceleration noise") - -parser.add_argument('--oms-noise-level', type=float, - help="The level of OMS noise") - -parser.add_argument('--tdi', type=str, - help="The version of TDI. Choose from '1.5' or '2.0'") - -parser.add_argument('--duration', type=float, - help="The duration of observation, between 0 and 10," - "in the unit of years") +parser.add_argument( + "--ifo-list", nargs="+", type=str, help="List of detector names, e.g. H1 L1 V1" +) + +parser.add_argument( + "--low-frequency-cutoff", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=float, + help="For each detector, provide minimum frequency forthe data generation.", +) + + +parser.add_argument( + "--channel-name", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=str, + help="For each detector, provide channel names for theframe files.", +) + +parser.add_argument( + "--psd-model", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=str, + help="For each detector, provide PSD option forthe noise generation.", +) + +parser.add_argument( + "--fake-strain-seed", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + type=int, + help="For each detector, provide a random number seed", +) + +parser.add_argument( + "--fake-strain-from-file", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:FILE", + type=str, + help="For each detector, provide a file (e.g. txt) to " + "load the fake strain from, instead of generating it.", +) + +parser.add_argument("--gps-start-time", type=int, help="GPS start time for frame files") + +parser.add_argument("--gps-end-time", type=int, help="GPS end time for frame files") + +parser.add_argument( + "--sample-rate", type=float, help="Sample rate of the frames to be generated" +) + +parser.add_argument( + "--fake-strain-sample-rate", + type=float, + help="Sample rate for the fake strain generation (defaults to sample-rate if not set)", +) + +parser.add_argument( + "--fake-strain-filter-duration", + type=float, + help="Duration of the filter used for coloring the noise", +) + +parser.add_argument( + "--injection-file", + type=str, + help="Injection file containing parameters ofinjection population", +) + +parser.add_argument( + "--output-path", help="Path for the output frame files", default=".", type=str +) + +parser.add_argument( + "--tag", help="Provide your tag for naming frame file", default=None +) + +parser.add_argument( + "--len-arm", + type=float, + default=None, + help="The arm length of LISA, in the unit of 'm'", +) + +parser.add_argument( + "--acc-noise-level", type=float, help="The level of acceleration noise" +) + +parser.add_argument("--oms-noise-level", type=float, help="The level of OMS noise") + +parser.add_argument( + "--tdi", type=str, help="The version of TDI. Choose from '1.5' or '2.0'" +) + +parser.add_argument( + "--duration", + type=float, + help="The duration of observation, between 0 and 10,in the unit of years", +) # parse the command line opts = parser.parse_args() @@ -110,90 +145,99 @@ opts = parser.parse_args() # Function to generate command line script def create_command_opts(ifo, rseed, output_file): - if 'LISA' in ifo: - detector_tag = 'space' + if "LISA" in ifo: + detector_tag = "space" else: - detector_tag = 'ground' + detector_tag = "ground" # Determine the source of the strain (PSD Model OR File) strain_source_args = [] - + # Priority: Check if file is provided first if opts.fake_strain_from_file and ifo in opts.fake_strain_from_file: - strain_source_args.append('--fake-strain-from-file %s' % opts.fake_strain_from_file[ifo]) + strain_source_args.append( + "--fake-strain-from-file %s" % opts.fake_strain_from_file[ifo] + ) # Otherwise, check if PSD model is provided elif opts.psd_model and ifo in opts.psd_model: - strain_source_args.append('--fake-strain %s' % opts.psd_model[ifo]) + strain_source_args.append("--fake-strain %s" % opts.psd_model[ifo]) else: # This shouldn't happen if sanity check passes, but good for safety raise ValueError(f"No strain source (PSD or File) provided for detector {ifo}") - if detector_tag=='ground': + if detector_tag == "ground": args_opt = [ - 'pycbc_condition_strain', - '--fake-strain-seed %d'%rseed, - '--sample-rate %d'%(opts.sample_rate), - '--gps-start-time %d'%(opts.gps_start_time), - '--gps-end-time %d'%(opts.gps_end_time), - '--channel-name %s:%s'%(ifo, opts.channel_name[ifo]), - '--output-strain-file %s'%(output_file)] - + "pycbc_condition_strain", + "--fake-strain-seed %d" % rseed, + "--sample-rate %d" % (opts.sample_rate), + "--gps-start-time %d" % (opts.gps_start_time), + "--gps-end-time %d" % (opts.gps_end_time), + "--channel-name %s:%s" % (ifo, opts.channel_name[ifo]), + "--output-strain-file %s" % (output_file), + ] + # Add the determined source arguments args_opt.extend(strain_source_args) - + # Pass explicit extended args if provided (no auto-fill for ground) if opts.fake_strain_sample_rate is not None: - args_opt.append('--fake-strain-sample-rate %f' % opts.fake_strain_sample_rate) + args_opt.append( + "--fake-strain-sample-rate %f" % opts.fake_strain_sample_rate + ) if opts.fake_strain_filter_duration is not None: - args_opt.append('--fake-strain-filter-duration %f' % opts.fake_strain_filter_duration) + args_opt.append( + "--fake-strain-filter-duration %f" % opts.fake_strain_filter_duration + ) - elif detector_tag=='space': + elif detector_tag == "space": # Sanity Check if opts.len_arm is None: raise ValueError("Space-borne missions require len-arm") - # If user didn't specify fake-strain-sample-rate, + # If user didn't specify fake-strain-sample-rate, # force it to match the global sample-rate. fake_rate = opts.fake_strain_sample_rate if fake_rate is None and opts.sample_rate is not None: fake_rate = opts.sample_rate - + # Build extra args for LISA extra_args = [] if opts.len_arm is not None: - extra_args.append('len_arm:%s' % opts.len_arm) + extra_args.append("len_arm:%s" % opts.len_arm) if opts.acc_noise_level is not None: - extra_args.append('acc_noise_level:%s' % opts.acc_noise_level) + extra_args.append("acc_noise_level:%s" % opts.acc_noise_level) if opts.oms_noise_level is not None: - extra_args.append('oms_noise_level:%s' % opts.oms_noise_level) + extra_args.append("oms_noise_level:%s" % opts.oms_noise_level) if opts.tdi is not None: - extra_args.append('tdi:%s' % opts.tdi) + extra_args.append("tdi:%s" % opts.tdi) if opts.duration is not None: - extra_args.append('duration:%s' % opts.duration) - + extra_args.append("duration:%s" % opts.duration) + args_opt = [ - 'pycbc_condition_strain', - '--fake-strain-seed %d'%rseed, - '--sample-rate %f'%(opts.sample_rate), - '--gps-start-time %d'%(opts.gps_start_time), - '--gps-end-time %d'%(opts.gps_end_time), - '--channel-name %s:%s'%(ifo, opts.channel_name[ifo]), - '--output-strain-file %s'%(output_file) + "pycbc_condition_strain", + "--fake-strain-seed %d" % rseed, + "--sample-rate %f" % (opts.sample_rate), + "--gps-start-time %d" % (opts.gps_start_time), + "--gps-end-time %d" % (opts.gps_end_time), + "--channel-name %s:%s" % (ifo, opts.channel_name[ifo]), + "--output-strain-file %s" % (output_file), ] # Add the determined source arguments args_opt.extend(strain_source_args) if fake_rate is not None: - args_opt.append('--fake-strain-sample-rate %f' % fake_rate) - + args_opt.append("--fake-strain-sample-rate %f" % fake_rate) + if opts.fake_strain_filter_duration is not None: - args_opt.append('--fake-strain-filter-duration %f' % opts.fake_strain_filter_duration) - + args_opt.append( + "--fake-strain-filter-duration %f" % opts.fake_strain_filter_duration + ) + # Add extra args if any if extra_args: - args_opt.append('--fake-strain-extra-args %s' % ' '.join(extra_args)) - + args_opt.append("--fake-strain-extra-args %s" % " ".join(extra_args)) + return args_opt @@ -203,11 +247,15 @@ if set(opts.ifo_list) != set(opts.low_frequency_cutoff.keys()): # New Logic: Ensure each IFO has EITHER a PSD model OR a File provided_psds = set(opts.psd_model.keys()) if opts.psd_model else set() -provided_files = set(opts.fake_strain_from_file.keys()) if opts.fake_strain_from_file else set() +provided_files = ( + set(opts.fake_strain_from_file.keys()) if opts.fake_strain_from_file else set() +) # Check if IFOs are subset of (PSDs U Files) if not set(opts.ifo_list).issubset(provided_psds.union(provided_files)): - raise ValueError("For each detector, you must provide either a --psd-model or a --fake-strain-from-file.") + raise ValueError( + "For each detector, you must provide either a --psd-model or a --fake-strain-from-file." + ) for ifo in opts.ifo_list: @@ -216,7 +264,7 @@ for ifo in opts.ifo_list: if set(opts.ifo_list) != set(opts.fake_strain_seed.keys()): raise ValueError("Fake strain seed for each detector is required") - logging.info("Generating strain for the Detector:%s"%ifo) + logging.info("Generating strain for the Detector:%s" % ifo) if opts.fake_strain_seed[ifo] is None: rseed = numpy.random.randint(1e8) else: @@ -224,22 +272,28 @@ for ifo in opts.ifo_list: time_duration = opts.gps_end_time - opts.gps_start_time if opts.tag is None: - output_file = '%s/%s-SIMULATED_STRAIN-%d-%d.gwf'%( - opts.output_path, ifo, opts.gps_start_time, time_duration) + output_file = "%s/%s-SIMULATED_STRAIN-%d-%d.gwf" % ( + opts.output_path, + ifo, + opts.gps_start_time, + time_duration, + ) else: - output_file = '%s/%s-SIMULATED_STRAIN-%s-%d-%d.gwf'%( - opts.output_path, ifo, opts.tag, opts.gps_start_time, - time_duration) + output_file = "%s/%s-SIMULATED_STRAIN-%s-%d-%d.gwf" % ( + opts.output_path, + ifo, + opts.tag, + opts.gps_start_time, + time_duration, + ) # Get the base command options args_opt = create_command_opts(ifo, rseed, output_file) if opts.injection_file is not None: - args_opt.append('--injection-file %s'%(opts.injection_file)) + args_opt.append("--injection-file %s" % (opts.injection_file)) if opts.low_frequency_cutoff is not None: - args_opt.append('--fake-strain-flow %f'%( - opts.low_frequency_cutoff[ifo])) - + args_opt.append("--fake-strain-flow %f" % (opts.low_frequency_cutoff[ifo])) - cmd = ' '.join(args_opt) + cmd = " ".join(args_opt) print(cmd) os.system(cmd) diff --git a/bin/pycbc_get_ffinal b/bin/pycbc_get_ffinal index 46ca2e7853b..a9c91a68c6b 100644 --- a/bin/pycbc_get_ffinal +++ b/bin/pycbc_get_ffinal @@ -1,90 +1,139 @@ #! /usr/bin/env python -"""Finds the maximum frequency of waveforms by generating them. Will also report +""" +Finds the maximum frequency of waveforms by generating them. Will also report the duration of time domain waveforms. """ -__prog__ = 'pycbc_get_ffinal' -__author__ = 'Collin Capano ' +__prog__ = "pycbc_get_ffinal" +__author__ = "Collin Capano " +import argparse import sys + import numpy -import argparse +from igwn_ligolw import ligolw +from igwn_ligolw import utils as ligolw_utils import pycbc from pycbc import waveform -from igwn_ligolw import utils as ligolw_utils -from igwn_ligolw import ligolw - - appx_options = waveform.td_approximants() + waveform.fd_approximants() parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("-i", "--input", - help="Input file. If specified, any single waveform " - "parameters given by the below options will be " - "ignored.") -parser.add_argument("-o", "--output", - help="Output file. Required if specifying an input " - "file.") -parser.add_argument("-a", "--approximant", required=True, - help=("What approximant to use to generate the " - "waveform(s). Options are: TD approximants: %s." - % ', '.join(waveform.td_approximants()) + - ". FD approximants: %s." % - ', '.join(waveform.fd_approximants())), - choices=appx_options) -parser.add_argument('-f', '--f-min', type=float, - help="Frequency at which to start the waveform " - "generation (in Hz).") -parser.add_argument('-r', '--sample-rate', type=int, - help="Sample rate to use (in Hz). Required for TD " - "approximants.") -parser.add_argument('-s', '--segment-length', type=int, - help = "The inverse of deltaF (in s). Required for FD " - "approximants.") -parser.add_argument('-m', '--max-sample-rate', type=int, - help="Optional. Maximum sample rate to use (in Hz). " - "If the Nyquist frequency of the given sample rate " - "is lower than the ringdown frequency, an error " - "will occur for some waveform approximants. If " - "max-sample-rate is specified, the code will try " - "increasing the sample-rate by a factor of 2 until " - "it finds a frequency that works or until it " - "exceeds the specified maximum rate.") -waveform_opts = parser.add_argument_group("Waveform Options", +parser.add_argument( + "-i", + "--input", + help="Input file. If specified, any single waveform " + "parameters given by the below options will be " + "ignored.", +) +parser.add_argument( + "-o", "--output", help="Output file. Required if specifying an input file." +) +parser.add_argument( + "-a", + "--approximant", + required=True, + help=( + "What approximant to use to generate the " + "waveform(s). Options are: TD approximants: %s." + % ", ".join(waveform.td_approximants()) + + ". FD approximants: %s." % ", ".join(waveform.fd_approximants()) + ), + choices=appx_options, +) +parser.add_argument( + "-f", + "--f-min", + type=float, + help="Frequency at which to start the waveform generation (in Hz).", +) +parser.add_argument( + "-r", + "--sample-rate", + type=int, + help="Sample rate to use (in Hz). Required for TD approximants.", +) +parser.add_argument( + "-s", + "--segment-length", + type=int, + help="The inverse of deltaF (in s). Required for FD approximants.", +) +parser.add_argument( + "-m", + "--max-sample-rate", + type=int, + help="Optional. Maximum sample rate to use (in Hz). " + "If the Nyquist frequency of the given sample rate " + "is lower than the ringdown frequency, an error " + "will occur for some waveform approximants. If " + "max-sample-rate is specified, the code will try " + "increasing the sample-rate by a factor of 2 until " + "it finds a frequency that works or until it " + "exceeds the specified maximum rate.", +) +waveform_opts = parser.add_argument_group( + "Waveform Options", "Optional arguments for specifying a single waveform. These will be " - "ignored if an input file is given.") -waveform_opts.add_argument('--mass1', type=float, - help="Specify mass1 of a single waveform.") -waveform_opts.add_argument('--mass2', type=float, - help="Specify mass2 of a single waveform.") -waveform_opts.add_argument('--spin1x', type=float, default=0, - help='Specify spin1x of a single waveform.') -waveform_opts.add_argument('--spin1y', type=float, default=0, - help='Specify spin1y of a single waveform.') -waveform_opts.add_argument('--spin1z', type=float, default=0, - help='Specify spin1z of a single waveform.') -waveform_opts.add_argument('--spin2x', type=float, default=0, - help='Specify spin2x of a single waveform.') -waveform_opts.add_argument('--spin2y', type=float, default=0, - help='Specify spin2y of a single waveform.') -waveform_opts.add_argument('--spin2z', type=float, default=0, - help='Specify spin2z of a single waveform.') -waveform_opts.add_argument('--lambda1', type=float, default=0, - help='Specify lambda1 of a single waveform.') -waveform_opts.add_argument('--lambda2', type=float, default=0, - help='Specify lambda2 of a single waveform.') -waveform_opts.add_argument('--phase-order', type=int, default=-1, - help='Specify phase-order of a single waveform.') -waveform_opts.add_argument('--spin-order', type=int, default=-1, - help='Specify spin-order of a single waveform.') -waveform_opts.add_argument('--tidal-order', type=int, default=-1, - help='Specify tidal-order of a single waveform.') -waveform_opts.add_argument('--amplitude-order', type=int, default=-1, - help='Specify amplitude-order of a single waveform.') + "ignored if an input file is given.", +) +waveform_opts.add_argument( + "--mass1", type=float, help="Specify mass1 of a single waveform." +) +waveform_opts.add_argument( + "--mass2", type=float, help="Specify mass2 of a single waveform." +) +waveform_opts.add_argument( + "--spin1x", type=float, default=0, help="Specify spin1x of a single waveform." +) +waveform_opts.add_argument( + "--spin1y", type=float, default=0, help="Specify spin1y of a single waveform." +) +waveform_opts.add_argument( + "--spin1z", type=float, default=0, help="Specify spin1z of a single waveform." +) +waveform_opts.add_argument( + "--spin2x", type=float, default=0, help="Specify spin2x of a single waveform." +) +waveform_opts.add_argument( + "--spin2y", type=float, default=0, help="Specify spin2y of a single waveform." +) +waveform_opts.add_argument( + "--spin2z", type=float, default=0, help="Specify spin2z of a single waveform." +) +waveform_opts.add_argument( + "--lambda1", type=float, default=0, help="Specify lambda1 of a single waveform." +) +waveform_opts.add_argument( + "--lambda2", type=float, default=0, help="Specify lambda2 of a single waveform." +) +waveform_opts.add_argument( + "--phase-order", + type=int, + default=-1, + help="Specify phase-order of a single waveform.", +) +waveform_opts.add_argument( + "--spin-order", + type=int, + default=-1, + help="Specify spin-order of a single waveform.", +) +waveform_opts.add_argument( + "--tidal-order", + type=int, + default=-1, + help="Specify tidal-order of a single waveform.", +) +waveform_opts.add_argument( + "--amplitude-order", + type=int, + default=-1, + help="Specify amplitude-order of a single waveform.", +) opts = parser.parse_args() @@ -92,8 +141,7 @@ pycbc.init_logging(opts.verbose) # check options if opts.input is None and (opts.mass1 is None or opts.mass2 is None): - parser.error("Must specify input file or at least mass1,mass2 " - "of a single waveform") + parser.error("Must specify input file or at least mass1,mass2 of a single waveform") infile = opts.input if opts.input is not None and opts.output is None: @@ -104,7 +152,7 @@ if opts.approximant is None: fd_approx = opts.approximant in waveform.fd_approximants() if not fd_approx and opts.approximant not in waveform.td_approximants(): - raise ValueError("Unrecognized approximant {}".format(opts.approximant)) + raise ValueError(f"Unrecognized approximant {opts.approximant}") if not fd_approx and opts.sample_rate is None: raise ValueError("TD approximants require a sample-rate") min_sample_rate = opts.sample_rate @@ -116,25 +164,34 @@ if fd_approx and opts.segment_length is None: raise ValueError("FD approximants require a segment-length") seg_length = opts.segment_length wFmin = opts.f_min -if not fd_approx and wFmin >= min_sample_rate / 2.: +if not fd_approx and wFmin >= min_sample_rate / 2.0: raise ValueError("f-min must be < half the sample-rate") -phi0 = 0. +phi0 = 0.0 # load the xml file if infile is not None: - xmldoc = ligolw_utils.load_filename( - infile, - compress='auto', - verbose=opts.verbose - ) + xmldoc = ligolw_utils.load_filename(infile, compress="auto", verbose=opts.verbose) this_process = xmldoc.register_process(__prog__, opts.__dict__) - sngl_insp_table = ligolw.Table.get_table(xmldoc, 'sngl_inspiral') + sngl_insp_table = ligolw.Table.get_table(xmldoc, "sngl_inspiral") else: # FIXME: try to get this from the waveform_opts group - tmplt_args = ['mass1', 'mass2', 'spin1x', 'spin1y', 'spin1z', 'spin2x', - 'spin2y', 'spin2z', 'lambda1', 'lambda2', 'phase_order', 'spin_order', - 'tidal_order', 'amplitude_order'] - tmplt = dict([ [arg, getattr(opts, arg)] for arg in tmplt_args]) + tmplt_args = [ + "mass1", + "mass2", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + "lambda1", + "lambda2", + "phase_order", + "spin_order", + "tidal_order", + "amplitude_order", + ] + tmplt = dict([[arg, getattr(opts, arg)] for arg in tmplt_args]) sngl_insp_table = [tmplt] if opts.verbose and infile is not None: @@ -142,21 +199,29 @@ if opts.verbose and infile is not None: for cache_id, tmplt in enumerate(sngl_insp_table): if opts.verbose and infile is not None: - print("Template %i/%i\r" %(cache_id+1, - len(sngl_insp_table)), end=' ', file=sys.stdout) + print( + "Template %i/%i\r" % (cache_id + 1, len(sngl_insp_table)), + end=" ", + file=sys.stdout, + ) sys.stdout.flush() - if fd_approx: if infile is None: hplus, hcross = waveform.get_fd_waveform( - approximant=opts.approximant, delta_f=1./seg_length, - f_lower=wFmin, **tmplt) + approximant=opts.approximant, + delta_f=1.0 / seg_length, + f_lower=wFmin, + **tmplt, + ) else: hplus, hcross = waveform.get_fd_waveform( - template=tmplt, approximant=opts.approximant, - delta_f=1./seg_length, f_lower=wFmin) - + template=tmplt, + approximant=opts.approximant, + delta_f=1.0 / seg_length, + f_lower=wFmin, + ) + freq = hplus.delta_f * numpy.nonzero(abs(hplus.data))[0][-1] else: @@ -167,34 +232,40 @@ for cache_id, tmplt in enumerate(sngl_insp_table): try: if infile is None: hplus, hcross = waveform.get_td_waveform( - approximant=opts.approximant, delta_t=1./sample_rate, - f_lower=wFmin, **tmplt) + approximant=opts.approximant, + delta_t=1.0 / sample_rate, + f_lower=wFmin, + **tmplt, + ) else: hplus, hcross = waveform.get_td_waveform( - template=tmplt, approximant=opts.approximant, - delta_t=1./sample_rate, f_lower=wFmin) + template=tmplt, + approximant=opts.approximant, + delta_t=1.0 / sample_rate, + f_lower=wFmin, + ) # success break success = True except: sample_rate *= 2 if sample_rate > max_sample_rate: - err_msg = "Unable to generate template %i;" %(cache_id) + err_msg = "Unable to generate template %i;" % (cache_id) err_msg += " try increasing max sample rate" raise ValueError(err_msg) - + # Remove possible zero padding at the end of the waveforms - amp = hplus**2 + hcross**2 + amp = hplus**2 + hcross**2 end_pnt = numpy.nonzero(amp.data)[0][-1] + 1 hplus = hplus[:end_pnt] hcross = hcross[:end_pnt] - + # find f_final freq = waveform.frequency_from_polarizations(hplus, hcross).max() if infile is None: - print("f Final: %f Hz" %(freq), file=sys.stdout) + print("f Final: %f Hz" % (freq), file=sys.stdout) if not fd_approx: - print("Duration: %f s" %(hplus.duration), file=sys.stdout) + print("Duration: %f s" % (hplus.duration), file=sys.stdout) else: sngl_insp_table[cache_id].f_final = freq @@ -207,7 +278,7 @@ if infile is not None: sys.stdout.flush() # write the xml file - ligolw_utils.write_filename(xmldoc, outfile, gz = outfile.endswith('.gz')) + ligolw_utils.write_filename(xmldoc, outfile, gz=outfile.endswith(".gz")) if opts.verbose: print("Finished!", file=sys.stdout) diff --git a/bin/pycbc_gwosc_segment_query b/bin/pycbc_gwosc_segment_query index b7d33e5a324..6421bb95952 100644 --- a/bin/pycbc_gwosc_segment_query +++ b/bin/pycbc_gwosc_segment_query @@ -1,9 +1,9 @@ #!/usr/bin/env python -import os -import logging -import json import argparse +import json +import logging +import os import shutil from urllib.request import urlopen @@ -30,43 +30,46 @@ def query_gwosc(ifo, segment_name, gps_start_time, duration): The amount of time in seconds after the gps start time. Returns - --------- + ------- segment_list : igwn_segments.segmentlist The interval returned by GWOSC segment_summary : igwn_segments.segmentlist The segments returned by GWOSC - """ + """ response = urlopen( - f'https://www.gwosc.org/timeline/segments/json/O1/{ifo}_{segment_name}/{gps_start_time}/{duration}/' + f"https://www.gwosc.org/timeline/segments/json/O1/{ifo}_{segment_name}/{gps_start_time}/{duration}/" ) logging.info(response.info()) json_segment_data = json.loads(response.read()) - summary_segment = igwn_segments.segmentlist([igwn_segments.segment( - json_segment_data['start'], - json_segment_data['end'])]) + summary_segment = igwn_segments.segmentlist( + [igwn_segments.segment(json_segment_data["start"], json_segment_data["end"])] + ) - segments = igwn_segments.segmentlist([igwn_segments.segment( - x[0],x[1]) for x in json_segment_data['segments']]) + segments = igwn_segments.segmentlist( + [igwn_segments.segment(x[0], x[1]) for x in json_segment_data["segments"]] + ) return summary_segment, segments + def write_xml_file(ifo, summary_segment, segments, filename): - file_url = 'file://' + os.path.abspath(filename) - sf = SegFile.from_segment_list('GWOSC segments', segments, 'RESULT', ifo, - summary_segment, file_url=file_url) + file_url = "file://" + os.path.abspath(filename) + sf = SegFile.from_segment_list( + "GWOSC segments", segments, "RESULT", ifo, summary_segment, file_url=file_url + ) sf.to_segment_xml() parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) -parser.add_argument('--gps-start-time', type=int, required=True) -parser.add_argument('--gps-end-time', type=int, required=True) -parser.add_argument('--include-segments', type=str, required=True) -parser.add_argument('--output-file', type=str, required=True) -parser.add_argument('--protract-hw-inj', type=int, default=0) +parser.add_argument("--gps-start-time", type=int, required=True) +parser.add_argument("--gps-end-time", type=int, required=True) +parser.add_argument("--include-segments", type=str, required=True) +parser.add_argument("--output-file", type=str, required=True) +parser.add_argument("--protract-hw-inj", type=int, default=0) args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -75,9 +78,8 @@ gps_start_time = args.gps_start_time gps_end_time = args.gps_end_time duration = gps_end_time - gps_start_time -logging.info("Reading in GWOSC files from %s to %s.", - gps_start_time, gps_end_time) -detector = args.include_segments.split(':')[0] +logging.info("Reading in GWOSC files from %s to %s.", gps_start_time, gps_end_time) +detector = args.include_segments.split(":")[0] logging.info("Querying for %s", detector) file_list = [] @@ -86,39 +88,40 @@ logging.info("Querying science segments") sci_summ, sci_segs = query_gwosc(detector, "DATA", gps_start_time, duration) sci_segs.coalesce() -sci_file_name = "{}-SCIENCE_SEGMENTS.xml".format(detector) +sci_file_name = f"{detector}-SCIENCE_SEGMENTS.xml" write_xml_file(detector, sci_summ, sci_segs, sci_file_name) file_list.append(sci_file_name) logging.info("Calculating CAT1 veto time") -not_cat1_summ, not_cat1_segs = query_gwosc(detector, "CBC_CAT1", - gps_start_time, duration) +not_cat1_summ, not_cat1_segs = query_gwosc( + detector, "CBC_CAT1", gps_start_time, duration +) not_cat1_segs.coalesce() cat1_segs = ~not_cat1_segs cat1_segs &= sci_segs -cat1_file_name = "{}-VETOTIME_CAT1-{}-{}.xml".format(detector, - gps_start_time, duration) +cat1_file_name = f"{detector}-VETOTIME_CAT1-{gps_start_time}-{duration}.xml" write_xml_file(detector, not_cat1_summ, cat1_segs, cat1_file_name) file_list.append(cat1_file_name) logging.info("Calculating CAT2 veto time") -not_cat2_summ, not_cat2_segs = query_gwosc(detector, "CBC_CAT2", - gps_start_time, duration) +not_cat2_summ, not_cat2_segs = query_gwosc( + detector, "CBC_CAT2", gps_start_time, duration +) not_cat2_segs.coalesce() cat2_segs = ~not_cat2_segs cat2_segs &= sci_segs -cat2_file_name = "{}-VETOTIME_CAT2-{}-{}.xml".format(detector, - gps_start_time, duration) +cat2_file_name = f"{detector}-VETOTIME_CAT2-{gps_start_time}-{duration}.xml" write_xml_file(detector, not_cat2_summ, cat2_segs, cat2_file_name) file_list.append(cat2_file_name) logging.info("Calculating HW injection veto time") -not_hw_inj_summ, not_hw_inj_segs = query_gwosc(detector, "NO_CBC_HW_INJ", - gps_start_time, duration) +not_hw_inj_summ, not_hw_inj_segs = query_gwosc( + detector, "NO_CBC_HW_INJ", gps_start_time, duration +) not_hw_inj_segs.coalesce() hw_inj_segs = ~not_hw_inj_segs @@ -126,15 +129,14 @@ hw_inj_segs.protract(args.protract_hw_inj) hw_inj_segs.coalesce() hw_inj_segs &= sci_segs -hw_inj_file_name = "{}-VETOTIME_CAT3-{}-{}.xml".format(detector, - gps_start_time, duration) +hw_inj_file_name = f"{detector}-VETOTIME_CAT3-{gps_start_time}-{duration}.xml" write_xml_file(detector, not_hw_inj_summ, hw_inj_segs, hw_inj_file_name) file_list.append(hw_inj_file_name) destination_path = os.path.dirname(os.path.abspath(args.output_file)) for f in file_list: - d = os.path.join(destination_path,f) + d = os.path.join(destination_path, f) logging.info("Copying %s to %s", f, d) shutil.copy2(f, os.path.join(destination_path, f)) os.unlink(f) diff --git a/bin/pycbc_hdf5_splitbank b/bin/pycbc_hdf5_splitbank index d0f41ab8723..d549daddd56 100755 --- a/bin/pycbc_hdf5_splitbank +++ b/bin/pycbc_hdf5_splitbank @@ -23,44 +23,67 @@ of templates per bank can be specified. """ import argparse -import numpy import logging + +import numpy from numpy import random import pycbc from pycbc.waveform import bank -__author__ = "Soumi De " +__author__ = "Soumi De " parser = argparse.ArgumentParser(description=__doc__[1:]) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--bank-file", type=str, required=True, - help="Bank hdf file to load.") +parser.add_argument( + "--bank-file", type=str, required=True, help="Bank hdf file to load." +) outbanks = parser.add_mutually_exclusive_group(required=True) -outbanks.add_argument("--templates-per-bank", type=int, - help="Number of templates in each output sub-banks. " - "Either specify this or --number-of-banks, not both.") -outbanks.add_argument("--number-of-banks", type=int, - help="Number of output sub-banks. Either specify this " - "or --templates-per-bank, not both.") -outbanks.add_argument("--output-filenames", nargs='*', - action="store", - help="Directly specify the names of the output files. " - "The bank will be split equally between files.") -parser.add_argument("--output-prefix", - help="Prefix to add to the output template bank names, " - "for example 'H1L1-BANK'. Output file names would then be " - "'H1L1-BANK{x}.hdf' where {x} is 1,2,...") +outbanks.add_argument( + "--templates-per-bank", + type=int, + help="Number of templates in each output sub-banks. " + "Either specify this or --number-of-banks, not both.", +) +outbanks.add_argument( + "--number-of-banks", + type=int, + help="Number of output sub-banks. Either specify this " + "or --templates-per-bank, not both.", +) +outbanks.add_argument( + "--output-filenames", + nargs="*", + action="store", + help="Directly specify the names of the output files. " + "The bank will be split equally between files.", +) +parser.add_argument( + "--output-prefix", + help="Prefix to add to the output template bank names, " + "for example 'H1L1-BANK'. Output file names would then be " + "'H1L1-BANK{x}.hdf' where {x} is 1,2,...", +) sortopt = parser.add_mutually_exclusive_group() -sortopt.add_argument("--mchirp-sort", action="store_true", default=False, - help="Sort templates by chirp mass before splitting") -sortopt.add_argument("--random-sort", action="store_true", default=False, - help="Sort templates randomly before splitting") -parser.add_argument("--random-seed", type=int, - help="Random seed for --random-sort") -parser.add_argument("--force", action="store_true", default=False, - help="Overwrite the given hdf file if it exists. " - "Otherwise, an error is raised.") +sortopt.add_argument( + "--mchirp-sort", + action="store_true", + default=False, + help="Sort templates by chirp mass before splitting", +) +sortopt.add_argument( + "--random-sort", + action="store_true", + default=False, + help="Sort templates randomly before splitting", +) +parser.add_argument("--random-seed", type=int, help="Random seed for --random-sort") +parser.add_argument( + "--force", + action="store_true", + default=False, + help="Overwrite the given hdf file if it exists. Otherwise, an error is raised.", +) args = parser.parse_args() @@ -128,17 +151,18 @@ for ii in range(num_files): if args.output_filenames: outname = args.output_filenames[ii] elif args.output_prefix: - outname = args.output_prefix + str(ii) + '.hdf' + outname = args.output_prefix + str(ii) + ".hdf" else: - raise RuntimeError("I shouldn't be able to reach this point. One out " - "of --output-filenames and --output-prefix must " - "have been supplied!") + raise RuntimeError( + "I shouldn't be able to reach this point. One out " + "of --output-filenames and --output-prefix must " + "have been supplied!" + ) # Generate the hdf5 output file for the ii'th sub-bank, which would # be a slice of the input template bank having a start index and # end index as calculated above - output = tmplt_bank.write_to_hdf(outname, start_idx, end_idx, - force=args.force) + output = tmplt_bank.write_to_hdf(outname, start_idx, end_idx, force=args.force) output.close() start_idx = end_idx diff --git a/bin/pycbc_hdf_splitinj b/bin/pycbc_hdf_splitinj index 288ecee39a6..821ff6a6add 100644 --- a/bin/pycbc_hdf_splitinj +++ b/bin/pycbc_hdf_splitinj @@ -6,27 +6,29 @@ Split sets are organized to maximize time between injections. """ import argparse + import numpy as np import pycbc from pycbc.inject import InjectionSet from pycbc.io.hdf import HFile - # Parse command line parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("-f", "--output-files", nargs='*', required=True, - help="Names of output files") -parser.add_argument("-i", "--input-file", required=True, - help="Injection file to be split") +parser.add_argument( + "-f", "--output-files", nargs="*", required=True, help="Names of output files" +) +parser.add_argument( + "-i", "--input-file", required=True, help="Injection file to be split" +) args = parser.parse_args() pycbc.init_logging(args.verbose) # Read in input file as both an hdf file and an InjectionSet object -inj_file = HFile(args.input_file, 'r') +inj_file = HFile(args.input_file, "r") inj_set = InjectionSet(args.input_file) # Define table of injection info @@ -35,12 +37,12 @@ inj_table = inj_set.table # InjectionSet.write() requires static params as a dictionary, # so get that from file object. # Ignore the "static_params" copy. -static_params = {key : inj_file.attrs[key] for key in inj_file.attrs - if key != 'static_args'} +static_params = { + key: inj_file.attrs[key] for key in inj_file.attrs if key != "static_args" +} # Also get the names of variable params as write_args -write_args = [arg for arg in inj_table.fieldnames - if arg not in static_params] +write_args = [arg for arg in inj_table.fieldnames if arg not in static_params] num_injs = len(inj_table) num_splits = len(args.output_files) @@ -60,14 +62,15 @@ if remainder > 0: assert sum(injs_per_split) == num_injs, "Not all injections were accounted for!" # Sort injections by time -inj_table.sort(order='tc') +inj_table.sort(order="tc") # Split injections into a list of smaller sets for i in range(num_splits): # Number of injections in this split injs_in_split = injs_per_split[i] # Spread injections by time so they don't overlap - injs_to_get = [i+(num_splits*j) for j in range(injs_in_split)] + injs_to_get = [i + (num_splits * j) for j in range(injs_in_split)] # Write to file - InjectionSet.write(args.output_files[i], inj_table[injs_to_get], - write_args, static_params) + InjectionSet.write( + args.output_files[i], inj_table[injs_to_get], write_args, static_params + ) diff --git a/bin/pycbc_inj_cut b/bin/pycbc_inj_cut index 3b17466c07d..dd70d5534c1 100644 --- a/bin/pycbc_inj_cut +++ b/bin/pycbc_inj_cut @@ -20,7 +20,7 @@ Given an xml(.gz) table with CBC injections, this script separates them into: (1) (potentially) found injections (2) injections that we expect to miss -The two sets are stored into two separate output files. +The two sets are stored into two separate output files. """ __author__ = "Stephen Fairhurst" @@ -28,10 +28,10 @@ __email__ = "stephen.fairhurst@ligo.org" import argparse import logging -import numpy -from igwn_ligolw import utils as ligolw_utils +import numpy from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils import pycbc import pycbc.inject @@ -39,21 +39,43 @@ from pycbc.types import MultiDetOptionAction parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--input', dest='inj_xml', required=True, help='Input LIGOLW injections file.') -parser.add_argument('--output-missed', dest='output_missed', required=False, - help="Output LIGOLW file containing injections we expect to miss.") -parser.add_argument('--output-file', dest='output_inject', required=True, - help="Output LIGOLW file containing injections that we might find.") -parser.add_argument('--snr-threshold', dest='snr_thresh', required=True, type=float, - help="Select the SNR threshold which is required to keep the injections") -parser.add_argument('--snr-columns', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', required=True, - help='Defines in which column of the sim_inspiral table' \ - ' the optimal SNR for each detector has been stored.' \ - ' COLUMN should be an existing sim_inspiral column with ' \ - ' no useful data in it, good candidates are usually' \ - ' alpha1, alpha2 etc.') -parser.add_argument("-z", "--write-compress", action="store_true", help="Write compressed xml.gz files.") +parser.add_argument( + "--input", dest="inj_xml", required=True, help="Input LIGOLW injections file." +) +parser.add_argument( + "--output-missed", + dest="output_missed", + required=False, + help="Output LIGOLW file containing injections we expect to miss.", +) +parser.add_argument( + "--output-file", + dest="output_inject", + required=True, + help="Output LIGOLW file containing injections that we might find.", +) +parser.add_argument( + "--snr-threshold", + dest="snr_thresh", + required=True, + type=float, + help="Select the SNR threshold which is required to keep the injections", +) +parser.add_argument( + "--snr-columns", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + required=True, + help="Defines in which column of the sim_inspiral table" + " the optimal SNR for each detector has been stored." + " COLUMN should be an existing sim_inspiral column with " + " no useful data in it, good candidates are usually" + " alpha1, alpha2 etc.", +) +parser.add_argument( + "-z", "--write-compress", action="store_true", help="Write compressed xml.gz files." +) opts = parser.parse_args() @@ -64,19 +86,19 @@ injections = pycbc.inject.InjectionSet(opts.inj_xml) # injections that we should do: table and name of the file that will store it output_inject = opts.output_inject if opts.write_compress: - if not output_inject.endswith('gz'): - output_inject = output_inject+'.gz' + if not output_inject.endswith("gz"): + output_inject = output_inject + ".gz" out_sim_inspiral_inject = lsctables.SimInspiralTable.new( columns=injections.table.columnnames ) -# missed injections table +# missed injections table out_sim_inspiral_missed = lsctables.SimInspiralTable.new( columns=injections.table.columnnames ) if opts.write_compress: - if not output_missed.endswith('gz'): - output_missed = output_missed+'.gz' + if not output_missed.endswith("gz"): + output_missed = output_missed + ".gz" for i, inj in enumerate(injections.table): snrs = numpy.array([getattr(inj, column) for column in opts.snr_columns.values()]) @@ -84,36 +106,36 @@ for i, inj in enumerate(injections.table): out_sim_inspiral_missed.append(inj) else: out_sim_inspiral_inject.append(inj) - + logging.info( - 'Found %d/%d (potentially) found injections. Storing them to %s.', + "Found %d/%d (potentially) found injections. Storing them to %s.", len(out_sim_inspiral_inject), len(injections.table), - output_inject + output_inject, ) logging.info( - 'Found %d/%d missed injections.', + "Found %d/%d missed injections.", len(out_sim_inspiral_missed), - len(injections.table) + len(injections.table), ) -logging.info('Writing output') +logging.info("Writing output") llw_doc = injections.indoc llw_root = llw_doc.childNodes[0] llw_root.removeChild(injections.table) llw_root.appendChild(out_sim_inspiral_inject) -ligolw_utils.write_filename(llw_doc, output_inject, compress='auto') +ligolw_utils.write_filename(llw_doc, output_inject, compress="auto") if opts.output_missed: output_missed = opts.output_missed if opts.write_compress: - if not output_missed.endswith('gz'): - output_missed = output_missed+'.gz' + if not output_missed.endswith("gz"): + output_missed = output_missed + ".gz" llw_root.removeChild(out_sim_inspiral_inject) llw_root.appendChild(out_sim_inspiral_missed) - logging.info('Storing them to %s.', output_missed) + logging.info("Storing them to %s.", output_missed) - ligolw_utils.write_filename(llw_doc, output_missed, compress='auto') + ligolw_utils.write_filename(llw_doc, output_missed, compress="auto") -logging.info('Done') +logging.info("Done") diff --git a/bin/pycbc_inspiral b/bin/pycbc_inspiral index e3bd5f7731e..2cc437e40dd 100644 --- a/bin/pycbc_inspiral +++ b/bin/pycbc_inspiral @@ -16,196 +16,315 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -import sys -import os +import argparse import copy import logging -import argparse -import numpy +import os +import sys import time -from pycbc.pool import BroadcastPool as Pool + +import numpy import pycbc -from pycbc import vetoes, psd, waveform, strain, scheme, fft, DYN_RANGE_FAC, events -from pycbc.vetoes.sgchisq import SingleDetSGChisq -from pycbc.filter import MatchedFilterControl, qtransform -from pycbc.types import zeros, float32, complex64 -import pycbc.opt import pycbc.inject +import pycbc.opt +from pycbc import DYN_RANGE_FAC, events, fft, psd, scheme, strain, vetoes, waveform +from pycbc.filter import MatchedFilterControl, qtransform +from pycbc.pool import BroadcastPool as Pool +from pycbc.types import complex64, float32, zeros +from pycbc.vetoes.sgchisq import SingleDetSGChisq last_progress_update = -1.0 -def update_progress(p,u,n): - """ updates a file 'progress.txt' with a value 0 .. 1.0 when enough (filtering) progress was made - """ + +def update_progress(p, u, n): + """Updates a file 'progress.txt' with a value 0 .. 1.0 when enough (filtering) progress was made""" global last_progress_update if p > last_progress_update + u: - f = open(n,"w") + f = open(n, "w") if f: f.write("%.4f" % p) f.close() last_progress_update = p + tstart = time.time() -parser = argparse.ArgumentParser(usage='', - description="Find single detector gravitational-wave triggers.") +parser = argparse.ArgumentParser( + usage="", description="Find single detector gravitational-wave triggers." +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--update-progress", - help="updates a file 'progress.txt' with a value 0 .. 1.0 when this amount of (filtering) progress was made", - type=float, default=0) -parser.add_argument("--update-progress-file", - help="name of the file to write the amount of (filtering) progress to", - type=str, default="progress.txt") +parser.add_argument( + "--update-progress", + help="updates a file 'progress.txt' with a value 0 .. 1.0 when this amount of (filtering) progress was made", + type=float, + default=0, +) +parser.add_argument( + "--update-progress-file", + help="name of the file to write the amount of (filtering) progress to", + type=str, + default="progress.txt", +) parser.add_argument("--output", type=str, help="FIXME: ADD") parser.add_argument("--bank-file", type=str, help="FIXME: ADD") -parser.add_argument("--snr-threshold", - help="SNR threshold for trigger generation", type=float) -parser.add_argument("--newsnr-threshold", type=float, metavar='THRESHOLD', - help="Cut triggers with NewSNR less than THRESHOLD") -parser.add_argument("--low-frequency-cutoff", type=float, - help="The low frequency cutoff to use for filtering (Hz)") -parser.add_argument("--enable-bank-start-frequency", action='store_true', - help="Read the starting frequency of template waveforms" - " from the template bank.") -parser.add_argument("--max-template-length", type=float, - help="The maximum length of a template is seconds. The " - "starting frequency of the template is modified to " - "ensure the proper length") -parser.add_argument("--enable-q-transform", action='store_true', - help="compute the q-transform for each segment of a " - "given analysis run. (default = False)") +parser.add_argument( + "--snr-threshold", help="SNR threshold for trigger generation", type=float +) +parser.add_argument( + "--newsnr-threshold", + type=float, + metavar="THRESHOLD", + help="Cut triggers with NewSNR less than THRESHOLD", +) +parser.add_argument( + "--low-frequency-cutoff", + type=float, + help="The low frequency cutoff to use for filtering (Hz)", +) +parser.add_argument( + "--enable-bank-start-frequency", + action="store_true", + help="Read the starting frequency of template waveforms from the template bank.", +) +parser.add_argument( + "--max-template-length", + type=float, + help="The maximum length of a template is seconds. The " + "starting frequency of the template is modified to " + "ensure the proper length", +) +parser.add_argument( + "--enable-q-transform", + action="store_true", + help="compute the q-transform for each segment of a " + "given analysis run. (default = False)", +) # add approximant arg pycbc.waveform.bank.add_approximant_arg(parser) -parser.add_argument("--order", type=int, - help="The integer half-PN order at which to generate" - " the approximant. Default is -1 which indicates to use" - " approximant defined default.", default=-1, - choices = numpy.arange(-1, 9, 1)) -taper_choices = ["start","end","startend"] -parser.add_argument("--taper-template", choices=taper_choices, - help="For time-domain approximants, taper the start and/or" - " end of the waveform before FFTing.") -parser.add_argument("--cluster-function", choices=["findchirp", "symmetric"], - help="How to cluster together triggers within a window. " - "'findchirp' uses a forward sliding window; 'symmetric' " - "will compare each window to the one before and after, keeping " - "only a local maximum.", default="findchirp") -parser.add_argument("--cluster-window", type=float, default=0, - help="Length of clustering window in seconds." - " Set to 0 to disable clustering.") +parser.add_argument( + "--order", + type=int, + help="The integer half-PN order at which to generate" + " the approximant. Default is -1 which indicates to use" + " approximant defined default.", + default=-1, + choices=numpy.arange(-1, 9, 1), +) +taper_choices = ["start", "end", "startend"] +parser.add_argument( + "--taper-template", + choices=taper_choices, + help="For time-domain approximants, taper the start and/or" + " end of the waveform before FFTing.", +) +parser.add_argument( + "--cluster-function", + choices=["findchirp", "symmetric"], + help="How to cluster together triggers within a window. " + "'findchirp' uses a forward sliding window; 'symmetric' " + "will compare each window to the one before and after, keeping " + "only a local maximum.", + default="findchirp", +) +parser.add_argument( + "--cluster-window", + type=float, + default=0, + help="Length of clustering window in seconds. Set to 0 to disable clustering.", +) parser.add_argument("--bank-veto-bank-file", type=str, help="FIXME: ADD") -parser.add_argument("--chisq-snr-threshold", type=float, - help="Minimum SNR to calculate the power chisq") -parser.add_argument("--chisq-bins", default=0, help= - "Number of frequency bins to use for power chisq. Specify" - " an integer for a constant number of bins, or a function " - "of template attributes. Math functions are " - "allowed, ex. " - "'10./math.sqrt((params.mass1+params.mass2)/100.)'. " - "Non-integer values will be rounded down.") -parser.add_argument("--chisq-threshold", type=float, default=0, - help="FIXME: ADD") +parser.add_argument( + "--chisq-snr-threshold", type=float, help="Minimum SNR to calculate the power chisq" +) +parser.add_argument( + "--chisq-bins", + default=0, + help="Number of frequency bins to use for power chisq. Specify" + " an integer for a constant number of bins, or a function " + "of template attributes. Math functions are " + "allowed, ex. " + "'10./math.sqrt((params.mass1+params.mass2)/100.)'. " + "Non-integer values will be rounded down.", +) +parser.add_argument("--chisq-threshold", type=float, default=0, help="FIXME: ADD") parser.add_argument("--chisq-delta", type=float, default=0, help="FIXME: ADD") -parser.add_argument("--autochi-number-points", type=int, default=0, - help="The number of points to use, in both directions if" - "doing a two-sided auto-chisq, to calculate the" - "auto-chisq statistic.") -parser.add_argument("--autochi-stride", type=int, default=0, - help="The gap, in sample points, between the points at" - "which to calculate auto-chisq.") -parser.add_argument("--autochi-two-phase", action="store_true", - default=False, - help="If given auto-chisq will be calculated by testing " - "against both phases of the SNR time-series. " - "If not given, only the phase matching the trigger " - "will be used.") -parser.add_argument("--autochi-onesided", action='store', default=None, - choices=['left','right'], - help="Decide whether to calculate auto-chisq using" - "points on both sides of the trigger or only on one" - "side. If not given points on both sides will be" - "used. If given, with either 'left' or 'right'," - "only points on that side (right = forward in time," - "left = back in time) will be used.") -parser.add_argument("--autochi-reverse-template", action="store_true", - default=False, - help="If given, time-reverse the template before" - "calculating the auto-chisq statistic. This will" - "come at additional computational cost as the SNR" - "time-series will need recomputing for the time-" - "reversed template.") -parser.add_argument("--autochi-max-valued", action="store_true", - default=False, - help="If given, store only the maximum value of the auto-" - "chisq over all points tested. A disadvantage of this " - "is that the mean value will not be known " - "analytically.") -parser.add_argument("--autochi-max-valued-dof", action="store", metavar="INT", - type=int, - help="If using --autochi-max-valued this value denotes " - "the pre-calculated mean value that will be stored " - "as the auto-chisq degrees-of-freedom value.") -parser.add_argument("--downsample-factor", type=int, - help="Factor that determines the interval between the " - "initial SNR sampling. If not set (or 1) no sparse sample " - "is created, and the standard full SNR is calculated.", default=1) -parser.add_argument("--upsample-threshold", type=float, - help="The fraction of the SNR threshold to check the sparse SNR sample.") -parser.add_argument("--upsample-method", choices=["pruned_fft"], - help="The method to find the SNR points between the sparse SNR sample.", - default='pruned_fft') -parser.add_argument("--user-tag", type=str, metavar="TAG", help=""" +parser.add_argument( + "--autochi-number-points", + type=int, + default=0, + help="The number of points to use, in both directions if" + "doing a two-sided auto-chisq, to calculate the" + "auto-chisq statistic.", +) +parser.add_argument( + "--autochi-stride", + type=int, + default=0, + help="The gap, in sample points, between the points at" + "which to calculate auto-chisq.", +) +parser.add_argument( + "--autochi-two-phase", + action="store_true", + default=False, + help="If given auto-chisq will be calculated by testing " + "against both phases of the SNR time-series. " + "If not given, only the phase matching the trigger " + "will be used.", +) +parser.add_argument( + "--autochi-onesided", + action="store", + default=None, + choices=["left", "right"], + help="Decide whether to calculate auto-chisq using" + "points on both sides of the trigger or only on one" + "side. If not given points on both sides will be" + "used. If given, with either 'left' or 'right'," + "only points on that side (right = forward in time," + "left = back in time) will be used.", +) +parser.add_argument( + "--autochi-reverse-template", + action="store_true", + default=False, + help="If given, time-reverse the template before" + "calculating the auto-chisq statistic. This will" + "come at additional computational cost as the SNR" + "time-series will need recomputing for the time-" + "reversed template.", +) +parser.add_argument( + "--autochi-max-valued", + action="store_true", + default=False, + help="If given, store only the maximum value of the auto-" + "chisq over all points tested. A disadvantage of this " + "is that the mean value will not be known " + "analytically.", +) +parser.add_argument( + "--autochi-max-valued-dof", + action="store", + metavar="INT", + type=int, + help="If using --autochi-max-valued this value denotes " + "the pre-calculated mean value that will be stored " + "as the auto-chisq degrees-of-freedom value.", +) +parser.add_argument( + "--downsample-factor", + type=int, + help="Factor that determines the interval between the " + "initial SNR sampling. If not set (or 1) no sparse sample " + "is created, and the standard full SNR is calculated.", + default=1, +) +parser.add_argument( + "--upsample-threshold", + type=float, + help="The fraction of the SNR threshold to check the sparse SNR sample.", +) +parser.add_argument( + "--upsample-method", + choices=["pruned_fft"], + help="The method to find the SNR points between the sparse SNR sample.", + default="pruned_fft", +) +parser.add_argument( + "--user-tag", + type=str, + metavar="TAG", + help=""" This is used to identify FULL_DATA jobs for compatibility with pipedown post-processing. - Option will be removed when no longer needed.""") -parser.add_argument("--keep-loudest-log-chirp-window", type=float, - help="Keep loudest triggers within ln chirp mass window") -parser.add_argument("--keep-loudest-interval", type=float, - help="Window in seconds to maximize triggers over bank") -parser.add_argument("--keep-loudest-num", type=int, - help="Number of triggers to keep from each maximization interval") -parser.add_argument("--keep-loudest-stat", default="newsnr", - choices=events.ranking.sngls_ranking_function_dict.keys(), - help="Statistic used to determine loudest to keep") -parser.add_argument("--finalize-events-template-rate", default=None, - type=int, metavar="NUM TEMPLATES", - help="After NUM TEMPLATES perform the various clustering " - "and rejection tests that would be performed at the " - "end of this job. Default is to only do those things " - "at the end of the job. This can help control memory " - "usage if a lot of triggers that would be rejected " - "are being retained. A suggested value for this is " - "512, but a good number may depend on other settings " - "and your specific use-case.") -parser.add_argument("--gpu-callback-method", default='none') + Option will be removed when no longer needed.""", +) +parser.add_argument( + "--keep-loudest-log-chirp-window", + type=float, + help="Keep loudest triggers within ln chirp mass window", +) +parser.add_argument( + "--keep-loudest-interval", + type=float, + help="Window in seconds to maximize triggers over bank", +) +parser.add_argument( + "--keep-loudest-num", + type=int, + help="Number of triggers to keep from each maximization interval", +) +parser.add_argument( + "--keep-loudest-stat", + default="newsnr", + choices=events.ranking.sngls_ranking_function_dict.keys(), + help="Statistic used to determine loudest to keep", +) +parser.add_argument( + "--finalize-events-template-rate", + default=None, + type=int, + metavar="NUM TEMPLATES", + help="After NUM TEMPLATES perform the various clustering " + "and rejection tests that would be performed at the " + "end of this job. Default is to only do those things " + "at the end of the job. This can help control memory " + "usage if a lot of triggers that would be rejected " + "are being retained. A suggested value for this is " + "512, but a good number may depend on other settings " + "and your specific use-case.", +) +parser.add_argument("--gpu-callback-method", default="none") parser.add_argument( "--use-compressed-waveforms", action="store_true", default=False, - help='Use compressed waveforms from the bank file (if available).' -) -parser.add_argument("--waveform-decompression-method", action='store', default=None, - help='Method to be used decompress waveforms from the bank file.') -parser.add_argument("--checkpoint-interval", type=int, - help="Save results to checkpoint file every X seconds. " - "Default is no checkpointing.") -parser.add_argument('--require-valid-checkpoint', default=False, action="store_true", - help="If the checkpoint file is invalid, raise an error. " - "Default is to ignore invalid checkpoint files and to " - "delete the broken file.") -parser.add_argument("--checkpoint-exit-maxtime", type=int, - help="Checkpoint and exit if X seconds of execution" - " time is exceeded. Default is no checkpointing.") -parser.add_argument("--checkpoint-exit-code", type=int, default=77, - help="Exit code returned if exiting after a checkpoint") -parser.add_argument("--multiprocessing-nprocesses", type=int, - help="Parallelize over multiple processes, note this is " - "separate from threading using the proc. scheme. " - "Used in conjunction with the option" - "--finalize-events-template-rate which should be set" - "to a multiple of the number of processes.") + help="Use compressed waveforms from the bank file (if available).", +) +parser.add_argument( + "--waveform-decompression-method", + action="store", + default=None, + help="Method to be used decompress waveforms from the bank file.", +) +parser.add_argument( + "--checkpoint-interval", + type=int, + help="Save results to checkpoint file every X seconds. " + "Default is no checkpointing.", +) +parser.add_argument( + "--require-valid-checkpoint", + default=False, + action="store_true", + help="If the checkpoint file is invalid, raise an error. " + "Default is to ignore invalid checkpoint files and to " + "delete the broken file.", +) +parser.add_argument( + "--checkpoint-exit-maxtime", + type=int, + help="Checkpoint and exit if X seconds of execution" + " time is exceeded. Default is no checkpointing.", +) +parser.add_argument( + "--checkpoint-exit-code", + type=int, + default=77, + help="Exit code returned if exiting after a checkpoint", +) +parser.add_argument( + "--multiprocessing-nprocesses", + type=int, + help="Parallelize over multiple processes, note this is " + "separate from threading using the proc. scheme. " + "Used in conjunction with the option" + "--finalize-events-template-rate which should be set" + "to a multiple of the number of processes.", +) # Add options groups psd.insert_psd_option_group(parser) @@ -223,7 +342,7 @@ psd.verify_psd_options(opt, parser) strain.verify_strain_options(opt, parser) strain.StrainSegments.verify_segment_options(opt, parser) scheme.verify_processing_options(opt, parser) -fft.verify_fft_options(opt,parser) +fft.verify_fft_options(opt, parser) pycbc.opt.verify_optimization_options(opt, parser) pycbc.init_logging(opt.verbose) @@ -232,14 +351,15 @@ fft.from_cli(opt) inj_filter_rejector = pycbc.inject.InjFilterRejector.from_cli(opt) ctx = scheme.from_cli(opt) -gwstrain = strain.from_cli(opt, dyn_range_fac=DYN_RANGE_FAC, - inj_filter_rejector=inj_filter_rejector) +gwstrain = strain.from_cli( + opt, dyn_range_fac=DYN_RANGE_FAC, inj_filter_rejector=inj_filter_rejector +) strain_segments = strain.StrainSegments.from_cli(opt, gwstrain) + def template_triggers(t_num): - """ Get the triggers for a specific template - """ + """Get the triggers for a specific template""" template = None tparam = None out_vals_all = [] @@ -247,67 +367,78 @@ def template_triggers(t_num): # Filter check checks the 'inj_filter_rejector' options to # determine whether # to filter this template/segment if injections are present. - if not inj_filter_rejector.template_segment_checker( - bank, t_num, stilde): + if not inj_filter_rejector.template_segment_checker(bank, t_num, stilde): continue if template is None: template = bank[t_num] tparam = template.params if opt.update_progress: - update_progress((t_num + (s_num / float(len(segments))) ) / len(bank), - opt.update_progress, opt.update_progress_file) - logging.info("Filtering template %d/%d segment %d/%d" % - (t_num + 1, len(bank), s_num + 1, len(segments))) + update_progress( + (t_num + (s_num / float(len(segments)))) / len(bank), + opt.update_progress, + opt.update_progress_file, + ) + logging.info( + "Filtering template %d/%d segment %d/%d" + % (t_num + 1, len(bank), s_num + 1, len(segments)) + ) sigmasq = template.sigmasq(stilde.psd) - snr, norm, corr, idx, snrv = \ - matched_filter.matched_filter_and_cluster(s_num, - sigmasq, - cluster_window, - epoch=stilde._epoch) + snr, norm, corr, idx, snrv = matched_filter.matched_filter_and_cluster( + s_num, sigmasq, cluster_window, epoch=stilde._epoch + ) if not len(idx): continue out_vals = out_vals_ref.copy() - out_vals['bank_chisq'], out_vals['bank_chisq_dof'] = \ - bank_chisq.values(template, stilde.psd, stilde, snrv, norm, - idx+stilde.analyze.start) - - out_vals['chisq'], out_vals['chisq_dof'] = \ - power_chisq.values(corr, snrv, norm, stilde.psd, - idx+stilde.analyze.start, template) - - out_vals['sg_chisq'] = sg_chisq.values(stilde, template, stilde.psd, - snrv, norm, - out_vals['chisq'], - out_vals['chisq_dof'], - idx+stilde.analyze.start) - - out_vals['cont_chisq'], _ = \ - autochisq.values(snr, idx+stilde.analyze.start, template, - stilde.psd, norm, stilde=stilde, - low_frequency_cutoff=flow) + out_vals["bank_chisq"], out_vals["bank_chisq_dof"] = bank_chisq.values( + template, stilde.psd, stilde, snrv, norm, idx + stilde.analyze.start + ) + + out_vals["chisq"], out_vals["chisq_dof"] = power_chisq.values( + corr, snrv, norm, stilde.psd, idx + stilde.analyze.start, template + ) + + out_vals["sg_chisq"] = sg_chisq.values( + stilde, + template, + stilde.psd, + snrv, + norm, + out_vals["chisq"], + out_vals["chisq_dof"], + idx + stilde.analyze.start, + ) + + out_vals["cont_chisq"], _ = autochisq.values( + snr, + idx + stilde.analyze.start, + template, + stilde.psd, + norm, + stilde=stilde, + low_frequency_cutoff=flow, + ) idx += stilde.cumulative_index - out_vals['time_index'] = idx - out_vals['snr'] = snrv * norm - out_vals['sigmasq'] = numpy.zeros(len(snrv), dtype=float32) + sigmasq + out_vals["time_index"] = idx + out_vals["snr"] = snrv * norm + out_vals["sigmasq"] = numpy.zeros(len(snrv), dtype=float32) + sigmasq if opt.psdvar_short_segment is not None: - out_vals['psd_var_val'] = \ - pycbc.psd.find_trigger_value(psd_var, - out_vals['time_index'], - opt.gps_start_time, opt.sample_rate) - #print(idx, out_vals['time_index']) + out_vals["psd_var_val"] = pycbc.psd.find_trigger_value( + psd_var, out_vals["time_index"], opt.gps_start_time, opt.sample_rate + ) + # print(idx, out_vals['time_index']) out_vals_all.append(copy.deepcopy(out_vals)) - #print(out_vals_all) + # print(out_vals_all) return out_vals_all, tparam -with ctx: - if opt.fft_backends == 'fftw': +with ctx: + if opt.fft_backends == "fftw": # The following FFTW specific options needed to wait until # we were inside the scheme context. @@ -317,43 +448,58 @@ with ctx: # Read specified user-provided wisdom files if opt.fftw_input_float_wisdom_file is not None: - fft.fftw.import_single_wisdom_from_filename(opt.fftw_input_float_wisdom_file) + fft.fftw.import_single_wisdom_from_filename( + opt.fftw_input_float_wisdom_file + ) if opt.fftw_input_double_wisdom_file is not None: - fft.fftw.import_double_wisdom_from_filename(opt.fftw_input_double_wisdom_file) + fft.fftw.import_double_wisdom_from_filename( + opt.fftw_input_double_wisdom_file + ) flow = opt.low_frequency_cutoff flen = strain_segments.freq_len tlen = strain_segments.time_len delta_f = strain_segments.delta_f - logging.info("Making frequency-domain data segments") segments = strain_segments.fourier_segments() - psd.associate_psds_to_segments(opt, segments, gwstrain, flen, delta_f, - flow, dyn_range_factor=DYN_RANGE_FAC, precision='single') + psd.associate_psds_to_segments( + opt, + segments, + gwstrain, + flen, + delta_f, + flow, + dyn_range_factor=DYN_RANGE_FAC, + precision="single", + ) # storage for values and types to be passed to event manager out_types = { - 'time_index' : int, - 'snr' : complex64, - 'chisq' : float32, - 'chisq_dof' : int, - 'bank_chisq' : float32, - 'bank_chisq_dof' : int, - 'cont_chisq' : float32, - 'psd_var_val' : float32, - 'sigmasq' : float32, - } + "time_index": int, + "snr": complex64, + "chisq": float32, + "chisq_dof": int, + "bank_chisq": float32, + "bank_chisq_dof": int, + "cont_chisq": float32, + "psd_var_val": float32, + "sigmasq": float32, + } out_types.update(SingleDetSGChisq.returns) - out_vals_ref = {key: None for key in out_types} + out_vals_ref = dict.fromkeys(out_types) names = sorted(out_vals_ref.keys()) if len(strain_segments.segment_slices) == 0: logging.info("--filter-inj-only specified and no injections in analysis time") event_mgr = events.EventManager( - opt, names, [out_types[n] for n in names], psd=None, - gating_info=gwstrain.gating_info) + opt, + names, + [out_types[n] for n in names], + psd=None, + gating_info=gwstrain.gating_info, + ) event_mgr.finalize_template_events() event_mgr.write_events(opt.output) logging.info("Finished") @@ -362,11 +508,17 @@ with ctx: # FIXME: Maybe we should use the PSD corresponding to each trigger if opt.psdvar_segment is not None: logging.info("Calculating PSD variation") - psd_var = pycbc.psd.calc_filt_psd_variation(gwstrain, opt.psdvar_segment, - opt.psdvar_short_segment, opt.psdvar_long_segment, - opt.psdvar_psd_duration, opt.psdvar_psd_stride, - opt.psd_estimation, opt.psdvar_low_freq, opt.psdvar_high_freq) - + psd_var = pycbc.psd.calc_filt_psd_variation( + gwstrain, + opt.psdvar_segment, + opt.psdvar_short_segment, + opt.psdvar_long_segment, + opt.psdvar_psd_duration, + opt.psdvar_psd_stride, + opt.psd_estimation, + opt.psdvar_low_freq, + opt.psdvar_high_freq, + ) if opt.enable_q_transform: logging.info("Performing q-transform on analysis segments") @@ -376,28 +528,33 @@ with ctx: q_trans = {} if opt.checkpoint_interval: - checkpoint_file = opt.output + '.checkpoint' + checkpoint_file = opt.output + ".checkpoint" checkpoint_exists = False tnum_start = 0 - if (opt.checkpoint_interval and os.path.isfile(checkpoint_file)): + if opt.checkpoint_interval and os.path.isfile(checkpoint_file): try: tnum_start, event_mgr = events.EventManager.restore_state(checkpoint_file) checkpoint_exists = True except Exception as e: if args.require_valid_checkpoint: logging.info("Failed to load checkpoint file") - raise(e) + raise (e) - logging.info('Failed to load checkpoint file, starting anew') + logging.info("Failed to load checkpoint file, starting anew") logging.info(e) if not checkpoint_exists: event_mgr = events.EventManager( - opt, names, [out_types[n] for n in names], psd=segments[0].psd, - gating_info=gwstrain.gating_info, q_trans=q_trans) - - template_mem = zeros(tlen, dtype = complex64) + opt, + names, + [out_types[n] for n in names], + psd=segments[0].psd, + gating_info=gwstrain.gating_info, + q_trans=q_trans, + ) + + template_mem = zeros(tlen, dtype=complex64) cluster_window = int(opt.cluster_window * gwstrain.sample_rate) if opt.cluster_window == 0.0: @@ -406,50 +563,72 @@ with ctx: use_cluster = True if hasattr(ctx, "num_threads"): - ncores = ctx.num_threads + ncores = ctx.num_threads else: - ncores = 1 + ncores = 1 if opt.multiprocessing_nprocesses: ncores *= opt.multiprocessing_nprocesses - - matched_filter = MatchedFilterControl(opt.low_frequency_cutoff, None, - opt.snr_threshold, tlen, delta_f, complex64, - segments, template_mem, use_cluster, - downsample_factor=opt.downsample_factor, - upsample_threshold=opt.upsample_threshold, - upsample_method=opt.upsample_method, - gpu_callback_method=opt.gpu_callback_method, - cluster_function=opt.cluster_function) - - bank_chisq = vetoes.SingleDetBankVeto(opt.bank_veto_bank_file, - flen, delta_f, flow, complex64, - phase_order=opt.order, - approximant=opt.approximant) + matched_filter = MatchedFilterControl( + opt.low_frequency_cutoff, + None, + opt.snr_threshold, + tlen, + delta_f, + complex64, + segments, + template_mem, + use_cluster, + downsample_factor=opt.downsample_factor, + upsample_threshold=opt.upsample_threshold, + upsample_method=opt.upsample_method, + gpu_callback_method=opt.gpu_callback_method, + cluster_function=opt.cluster_function, + ) + + bank_chisq = vetoes.SingleDetBankVeto( + opt.bank_veto_bank_file, + flen, + delta_f, + flow, + complex64, + phase_order=opt.order, + approximant=opt.approximant, + ) power_chisq = vetoes.SingleDetPowerChisq(opt.chisq_bins, opt.chisq_snr_threshold) - autochisq = vetoes.SingleDetAutoChisq(opt.autochi_stride, - opt.autochi_number_points, - onesided=opt.autochi_onesided, - twophase=opt.autochi_two_phase, - reverse_template=opt.autochi_reverse_template, - take_maximum_value=opt.autochi_max_valued, - maximal_value_dof=opt.autochi_max_valued_dof) + autochisq = vetoes.SingleDetAutoChisq( + opt.autochi_stride, + opt.autochi_number_points, + onesided=opt.autochi_onesided, + twophase=opt.autochi_two_phase, + reverse_template=opt.autochi_reverse_template, + take_maximum_value=opt.autochi_max_valued, + maximal_value_dof=opt.autochi_max_valued_dof, + ) logging.info("Overwhitening frequency-domain data segments") for seg in segments: seg /= seg.psd logging.info("Read in template bank") - bank = waveform.FilterBank(opt.bank_file, flen, delta_f, + bank = waveform.FilterBank( + opt.bank_file, + flen, + delta_f, low_frequency_cutoff=None if opt.enable_bank_start_frequency else flow, - dtype=complex64, phase_order=opt.order, - taper=opt.taper_template, approximant=opt.approximant, - out=template_mem, max_template_length=opt.max_template_length, + dtype=complex64, + phase_order=opt.order, + taper=opt.taper_template, + approximant=opt.approximant, + out=template_mem, + max_template_length=opt.max_template_length, enable_compressed_waveforms=True if opt.use_compressed_waveforms else False, - waveform_decompression_method= - opt.waveform_decompression_method if opt.use_compressed_waveforms else None) + waveform_decompression_method=opt.waveform_decompression_method + if opt.use_compressed_waveforms + else None, + ) sg_chisq = SingleDetSGChisq.from_cli(opt, bank, opt.chisq_bins) @@ -463,7 +642,7 @@ with ctx: tanalyze = list(range(tnum_start, len(bank))) n = opt.finalize_events_template_rate n = 1 if n is None else n - tchunks = [tanalyze[i:i + n] for i in range(0, len(tanalyze), n)] + tchunks = [tanalyze[i : i + n] for i in range(0, len(tanalyze), n)] mmap = map if opt.multiprocessing_nprocesses: @@ -486,14 +665,16 @@ with ctx: if opt.finalize_events_template_rate is not None: event_mgr.consolidate_events(opt, gwstrain=gwstrain) - if opt.checkpoint_interval and \ - (time.time() - tcheckpoint > opt.checkpoint_interval): - event_mgr.save_state(max(tchunk), opt.output + '.checkpoint') + if opt.checkpoint_interval and ( + time.time() - tcheckpoint > opt.checkpoint_interval + ): + event_mgr.save_state(max(tchunk), opt.output + ".checkpoint") tcheckpoint = time.time() - if opt.checkpoint_exit_maxtime and \ - (time.time() - tstart > opt.checkpoint_exit_maxtime): - event_mgr.save_state(max(tchunk), opt.output + '.checkpoint') + if opt.checkpoint_exit_maxtime and ( + time.time() - tstart > opt.checkpoint_exit_maxtime + ): + event_mgr.save_state(max(tchunk), opt.output + ".checkpoint") sys.exit(opt.checkpoint_exit_code) event_mgr.consolidate_events(opt, gwstrain=gwstrain) @@ -507,7 +688,7 @@ event_mgr.save_performance(ncores, len(segments), len(bank), run_time, tsetup) logging.info("Writing out triggers") event_mgr.write_events(opt.output) -if opt.fft_backends == 'fftw': +if opt.fft_backends == "fftw": if opt.fftw_output_float_wisdom_file: fft.fftw.export_single_wisdom_to_filename(opt.fftw_output_float_wisdom_file) diff --git a/bin/pycbc_inspiral_skymax b/bin/pycbc_inspiral_skymax index bc1972e7968..07981c4a3ae 100755 --- a/bin/pycbc_inspiral_skymax +++ b/bin/pycbc_inspiral_skymax @@ -16,151 +16,248 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -import logging import argparse +import logging import sys + import numpy -from pycbc import vetoes, psd, waveform, events, strain, scheme, fft -from pycbc.vetoes.sgchisq import SingleDetSGChisq -from pycbc import DYN_RANGE_FAC -from pycbc.filter import MatchedFilterSkyMaxControl, MatchedFilterSkyMaxControlNoPhase -from pycbc.types import zeros, float32, complex64 import pycbc.fft.fftw import pycbc.opt +from pycbc import DYN_RANGE_FAC, events, fft, psd, scheme, strain, vetoes, waveform +from pycbc.filter import MatchedFilterSkyMaxControl, MatchedFilterSkyMaxControlNoPhase +from pycbc.types import complex64, float32, zeros +from pycbc.vetoes.sgchisq import SingleDetSGChisq -parser = argparse.ArgumentParser(usage='', - description="Find single detector precessing gravitational-wave triggers.") +parser = argparse.ArgumentParser( + usage="", description="Find single detector precessing gravitational-wave triggers." +) pycbc.add_common_pycbc_options(parser) parser.add_argument("--output", type=str, help="FIXME: ADD") parser.add_argument("--bank-file", type=str, help="FIXME: ADD") -parser.add_argument("--snr-threshold", - help="SNR threshold for trigger generation", type=float) -parser.add_argument("--newsnr-threshold", type=float, metavar='THRESHOLD', - help="Cut triggers with NewSNR less than THRESHOLD") -parser.add_argument("--low-frequency-cutoff", type=float, - help="The low frequency cutoff to use for filtering (Hz)") -parser.add_argument("--enable-bank-start-frequency", action='store_true', - help="Read the starting frequency of template waveforms" - " from the template bank.") -parser.add_argument("--max-template-length", type=float, - help="The maximum length of a template is seconds. The " - "starting frequency of the template is modified to " - "ensure the proper length") -parser.add_argument("--enable-q-transform", action='store_true', - help="compute the q-transform for each segment of a " - "given analysis run. (default = False)") +parser.add_argument( + "--snr-threshold", help="SNR threshold for trigger generation", type=float +) +parser.add_argument( + "--newsnr-threshold", + type=float, + metavar="THRESHOLD", + help="Cut triggers with NewSNR less than THRESHOLD", +) +parser.add_argument( + "--low-frequency-cutoff", + type=float, + help="The low frequency cutoff to use for filtering (Hz)", +) +parser.add_argument( + "--enable-bank-start-frequency", + action="store_true", + help="Read the starting frequency of template waveforms from the template bank.", +) +parser.add_argument( + "--max-template-length", + type=float, + help="The maximum length of a template is seconds. The " + "starting frequency of the template is modified to " + "ensure the proper length", +) +parser.add_argument( + "--enable-q-transform", + action="store_true", + help="compute the q-transform for each segment of a " + "given analysis run. (default = False)", +) # add approximant arg waveform.bank.add_approximant_arg(parser) -parser.add_argument("--order", type=int, - help="The integer half-PN order at which to generate" - " the approximant. Default is -1 which indicates to use" - " approximant defined default.", default=-1, - choices = numpy.arange(-1, 9, 1)) -taper_choices = ["start","end","startend"] -parser.add_argument("--taper-template", choices=taper_choices, - help="For time-domain approximants, taper the start and/or " - "end of the waveform before FFTing.") -parser.add_argument("--cluster-method", choices=["template", "window"], - help="FIXME: ADD") -parser.add_argument("--cluster-function", choices=["findchirp", "symmetric"], - help="How to cluster together triggers within a window. " - "'findchirp' uses a forward sliding window; 'symmetric' " - "will compare each window to the one before and after, keeping " - "only a local maximum.", default="findchirp") -parser.add_argument("--cluster-window", type=float, default = -1, - help="Length of clustering window in seconds." - " Set to 0 to disable clustering.") -parser.add_argument("--maximization-interval", type=float, default=0, - help="Maximize triggers over the template bank (ms)") +parser.add_argument( + "--order", + type=int, + help="The integer half-PN order at which to generate" + " the approximant. Default is -1 which indicates to use" + " approximant defined default.", + default=-1, + choices=numpy.arange(-1, 9, 1), +) +taper_choices = ["start", "end", "startend"] +parser.add_argument( + "--taper-template", + choices=taper_choices, + help="For time-domain approximants, taper the start and/or " + "end of the waveform before FFTing.", +) +parser.add_argument( + "--cluster-method", choices=["template", "window"], help="FIXME: ADD" +) +parser.add_argument( + "--cluster-function", + choices=["findchirp", "symmetric"], + help="How to cluster together triggers within a window. " + "'findchirp' uses a forward sliding window; 'symmetric' " + "will compare each window to the one before and after, keeping " + "only a local maximum.", + default="findchirp", +) +parser.add_argument( + "--cluster-window", + type=float, + default=-1, + help="Length of clustering window in seconds. Set to 0 to disable clustering.", +) +parser.add_argument( + "--maximization-interval", + type=float, + default=0, + help="Maximize triggers over the template bank (ms)", +) parser.add_argument("--bank-veto-bank-file", type=str, help="FIXME: ADD") -parser.add_argument("--chisq-snr-threshold", type=float, - help="Minimum SNR to calculate the power chisq") -parser.add_argument("--chisq-bins", default=0, help= - "Number of frequency bins to use for power chisq. Specify" - " an integer for a constant number of bins, or a function " - "of template attributes. Math functions are " - "allowed, ex. " - "'10./math.sqrt((params.mass1+params.mass2)/100.)'. " - "Non-integer values will be rounded down.") -parser.add_argument("--chisq-threshold", type=float, default=0, - help="FIXME: ADD") +parser.add_argument( + "--chisq-snr-threshold", type=float, help="Minimum SNR to calculate the power chisq" +) +parser.add_argument( + "--chisq-bins", + default=0, + help="Number of frequency bins to use for power chisq. Specify" + " an integer for a constant number of bins, or a function " + "of template attributes. Math functions are " + "allowed, ex. " + "'10./math.sqrt((params.mass1+params.mass2)/100.)'. " + "Non-integer values will be rounded down.", +) +parser.add_argument("--chisq-threshold", type=float, default=0, help="FIXME: ADD") parser.add_argument("--chisq-delta", type=float, default=0, help="FIXME: ADD") -parser.add_argument("--autochi-number-points", type=int, default=0, - help="The number of points to use, in both directions if" - "doing a two-sided auto-chisq, to calculate the" - "auto-chisq statistic.") -parser.add_argument("--autochi-stride", type=int, default=0, - help="The gap, in sample points, between the points at" - "which to calculate auto-chisq.") -parser.add_argument("--autochi-two-phase", action="store_true", - default=False, - help="If given auto-chisq will be calculated by testing " - "against both phases of the SNR time-series. " - "If not given, only the phase matching the trigger " - "will be used.") -parser.add_argument("--autochi-onesided", action='store', default=None, - choices=['left','right'], - help="Decide whether to calculate auto-chisq using" - "points on both sides of the trigger or only on one" - "side. If not given points on both sides will be" - "used. If given, with either 'left' or 'right'," - "only points on that side (right = forward in time," - "left = back in time) will be used.") -parser.add_argument("--autochi-reverse-template", action="store_true", - default=False, - help="If given, time-reverse the template before" - "calculating the auto-chisq statistic. This will" - "come at additional computational cost as the SNR" - "time-series will need recomputing for the time-" - "reversed template.") -parser.add_argument("--autochi-max-valued", action="store_true", - default=False, - help="If given, store only the maximum value of the auto-" - "chisq over all points tested. A disadvantage of this " - "is that the mean value will not be known " - "analytically.") -parser.add_argument("--autochi-max-valued-dof", action="store", metavar="INT", - type=int, - help="If using --autochi-max-valued this value denotes " - "the pre-calculated mean value that will be stored " - "as the auto-chisq degrees-of-freedom value.") -parser.add_argument("--downsample-factor", type=int, - help="Factor that determines the interval between the " - "initial SNR sampling. If not set (or 1) no sparse sample " - "is created, and the standard full SNR is calculated.", default=1) -parser.add_argument("--upsample-threshold", type=float, - help="The fraction of the SNR threshold to check the sparse SNR sample.") -parser.add_argument("--upsample-method", choices=["pruned_fft"], - help="The method to find the SNR points between the sparse SNR sample.", default='pruned_fft') - -parser.add_argument("--user-tag", type=str, metavar="TAG", help=""" +parser.add_argument( + "--autochi-number-points", + type=int, + default=0, + help="The number of points to use, in both directions if" + "doing a two-sided auto-chisq, to calculate the" + "auto-chisq statistic.", +) +parser.add_argument( + "--autochi-stride", + type=int, + default=0, + help="The gap, in sample points, between the points at" + "which to calculate auto-chisq.", +) +parser.add_argument( + "--autochi-two-phase", + action="store_true", + default=False, + help="If given auto-chisq will be calculated by testing " + "against both phases of the SNR time-series. " + "If not given, only the phase matching the trigger " + "will be used.", +) +parser.add_argument( + "--autochi-onesided", + action="store", + default=None, + choices=["left", "right"], + help="Decide whether to calculate auto-chisq using" + "points on both sides of the trigger or only on one" + "side. If not given points on both sides will be" + "used. If given, with either 'left' or 'right'," + "only points on that side (right = forward in time," + "left = back in time) will be used.", +) +parser.add_argument( + "--autochi-reverse-template", + action="store_true", + default=False, + help="If given, time-reverse the template before" + "calculating the auto-chisq statistic. This will" + "come at additional computational cost as the SNR" + "time-series will need recomputing for the time-" + "reversed template.", +) +parser.add_argument( + "--autochi-max-valued", + action="store_true", + default=False, + help="If given, store only the maximum value of the auto-" + "chisq over all points tested. A disadvantage of this " + "is that the mean value will not be known " + "analytically.", +) +parser.add_argument( + "--autochi-max-valued-dof", + action="store", + metavar="INT", + type=int, + help="If using --autochi-max-valued this value denotes " + "the pre-calculated mean value that will be stored " + "as the auto-chisq degrees-of-freedom value.", +) +parser.add_argument( + "--downsample-factor", + type=int, + help="Factor that determines the interval between the " + "initial SNR sampling. If not set (or 1) no sparse sample " + "is created, and the standard full SNR is calculated.", + default=1, +) +parser.add_argument( + "--upsample-threshold", + type=float, + help="The fraction of the SNR threshold to check the sparse SNR sample.", +) +parser.add_argument( + "--upsample-method", + choices=["pruned_fft"], + help="The method to find the SNR points between the sparse SNR sample.", + default="pruned_fft", +) + +parser.add_argument( + "--user-tag", + type=str, + metavar="TAG", + help=""" This is used to identify FULL_DATA jobs for compatibility with pipedown post-processing. - Option will be removed when no longer needed.""") -parser.add_argument("--keep-loudest-interval", type=float, - help="Window in seconds to maximize triggers over bank") -parser.add_argument("--keep-loudest-num", type=int, - help="Number of triggers to keep from each maximization interval") -parser.add_argument("--keep-loudest-stat", default="newsnr", - choices=events.ranking.sngls_ranking_function_dict.keys(), - help="Statistic used to determine loudest to keep") -parser.add_argument("--keep-loudest-log-chirp-window", type=float, - help="Keep loudest triggers within ln chirp mass window") -parser.add_argument("--finalize-events-template-rate", default=None, - type=int, metavar="NUM TEMPLATES", - help="After NUM TEMPLATES perform the various clustering " - "and rejection tests that would be performed at the " - "end of this job. Default is to only do those things " - "at the end of the job. This can help control memory " - "usage if a lot of triggers that would be rejected " - "are being retained. A suggested value for this is " - "500, but a good number may depend on other settings " - "and your specific use-case.") -parser.add_argument("--gpu-callback-method", default='none') -parser.add_argument("--sky-maximization-method", required=True, - choices=["precessing", "hom"]) + Option will be removed when no longer needed.""", +) +parser.add_argument( + "--keep-loudest-interval", + type=float, + help="Window in seconds to maximize triggers over bank", +) +parser.add_argument( + "--keep-loudest-num", + type=int, + help="Number of triggers to keep from each maximization interval", +) +parser.add_argument( + "--keep-loudest-stat", + default="newsnr", + choices=events.ranking.sngls_ranking_function_dict.keys(), + help="Statistic used to determine loudest to keep", +) +parser.add_argument( + "--keep-loudest-log-chirp-window", + type=float, + help="Keep loudest triggers within ln chirp mass window", +) +parser.add_argument( + "--finalize-events-template-rate", + default=None, + type=int, + metavar="NUM TEMPLATES", + help="After NUM TEMPLATES perform the various clustering " + "and rejection tests that would be performed at the " + "end of this job. Default is to only do those things " + "at the end of the job. This can help control memory " + "usage if a lot of triggers that would be rejected " + "are being retained. A suggested value for this is " + "500, but a good number may depend on other settings " + "and your specific use-case.", +) +parser.add_argument("--gpu-callback-method", default="none") +parser.add_argument( + "--sky-maximization-method", required=True, choices=["precessing", "hom"] +) # Add options groups psd.insert_psd_option_group(parser) @@ -178,7 +275,7 @@ psd.verify_psd_options(opt, parser) strain.verify_strain_options(opt, parser) strain.StrainSegments.verify_segment_options(opt, parser) scheme.verify_processing_options(opt, parser) -fft.verify_fft_options(opt,parser) +fft.verify_fft_options(opt, parser) pycbc.opt.verify_optimization_options(opt, parser) pycbc.init_logging(opt.verbose) @@ -197,33 +294,44 @@ with ctx: logging.info("Making frequency-domain data segments") segments = strain_segments.fourier_segments() - psd.associate_psds_to_segments(opt, segments, gwstrain, flen, delta_f, - flow, dyn_range_factor=DYN_RANGE_FAC, precision='single') + psd.associate_psds_to_segments( + opt, + segments, + gwstrain, + flen, + delta_f, + flow, + dyn_range_factor=DYN_RANGE_FAC, + precision="single", + ) # storage for values and types to be passed to event manager out_types = { - 'time_index' : int, - 'snr' : float32, - 'chisq' : float32, - 'chisq_dof' : int, - 'bank_chisq' : float32, - 'bank_chisq_dof' : int, - 'cont_chisq' : float32, - 'psd_var_val' : float32, - 'u_vals' : float32, - 'coa_phase' : float32, - 'hplus_cross_corr' : float32 - } + "time_index": int, + "snr": float32, + "chisq": float32, + "chisq_dof": int, + "bank_chisq": float32, + "bank_chisq_dof": int, + "cont_chisq": float32, + "psd_var_val": float32, + "u_vals": float32, + "coa_phase": float32, + "hplus_cross_corr": float32, + } out_types.update(SingleDetSGChisq.returns) - out_vals = {key: None for key in out_types} + out_vals = dict.fromkeys(out_types) names = sorted(out_vals.keys()) - if len(strain_segments.segment_slices) == 0: logging.info("--filter-inj-only specified and no injections in analysis time") event_mgr = events.EventManager( - opt, names, [out_types[n] for n in names], psd=None, - gating_info=gwstrain.gating_info) + opt, + names, + [out_types[n] for n in names], + psd=None, + gating_info=gwstrain.gating_info, + ) event_mgr.finalize_template_events() event_mgr.write_events(opt.output) logging.info("Finished") @@ -232,55 +340,78 @@ with ctx: # FIXME: Maybe we should use the PSD corresponding to each trigger if opt.psdvar_segment is not None: logging.info("Calculating PSD variation") - psd_var = pycbc.psd.calc_filt_psd_variation(gwstrain, opt.psdvar_segment, - opt.psdvar_short_segment, opt.psdvar_long_segment, - opt.psdvar_psd_duration, opt.psdvar_psd_stride, - opt.psd_estimation, opt.psdvar_low_freq, opt.psdvar_high_freq) - - if opt.sky_maximization_method == 'hom': - - matched_filter = MatchedFilterSkyMaxControlNoPhase(opt.low_frequency_cutoff, - None, opt.snr_threshold, tlen, delta_f, - complex64) + psd_var = pycbc.psd.calc_filt_psd_variation( + gwstrain, + opt.psdvar_segment, + opt.psdvar_short_segment, + opt.psdvar_long_segment, + opt.psdvar_psd_duration, + opt.psdvar_psd_stride, + opt.psd_estimation, + opt.psdvar_low_freq, + opt.psdvar_high_freq, + ) + + if opt.sky_maximization_method == "hom": + matched_filter = MatchedFilterSkyMaxControlNoPhase( + opt.low_frequency_cutoff, None, opt.snr_threshold, tlen, delta_f, complex64 + ) else: - - matched_filter = MatchedFilterSkyMaxControl(opt.low_frequency_cutoff, - None, opt.snr_threshold, tlen, delta_f, - complex64) - - bank_chisq = vetoes.SingleDetSkyMaxBankVeto(opt.bank_veto_bank_file, - flen, delta_f, flow, complex64, - phase_order=opt.order, - approximant=opt.approximant) - - power_chisq = vetoes.SingleDetSkyMaxPowerChisq(num_bins=opt.chisq_bins, - snr_threshold=opt.chisq_snr_threshold) - - autochisq = vetoes.SingleDetSkyMaxAutoChisq(opt.autochi_stride, - opt.autochi_number_points, - onesided=opt.autochi_onesided, - twophase=opt.autochi_two_phase, - reverse_template=opt.autochi_reverse_template, - take_maximum_value=opt.autochi_max_valued, - maximal_value_dof=opt.autochi_max_valued_dof) + matched_filter = MatchedFilterSkyMaxControl( + opt.low_frequency_cutoff, None, opt.snr_threshold, tlen, delta_f, complex64 + ) + + bank_chisq = vetoes.SingleDetSkyMaxBankVeto( + opt.bank_veto_bank_file, + flen, + delta_f, + flow, + complex64, + phase_order=opt.order, + approximant=opt.approximant, + ) + + power_chisq = vetoes.SingleDetSkyMaxPowerChisq( + num_bins=opt.chisq_bins, snr_threshold=opt.chisq_snr_threshold + ) + + autochisq = vetoes.SingleDetSkyMaxAutoChisq( + opt.autochi_stride, + opt.autochi_number_points, + onesided=opt.autochi_onesided, + twophase=opt.autochi_two_phase, + reverse_template=opt.autochi_reverse_template, + take_maximum_value=opt.autochi_max_valued, + maximal_value_dof=opt.autochi_max_valued_dof, + ) logging.info("Overwhitening frequency-domain data segments") for seg in segments: seg /= seg.psd - event_mgr = events.EventManager(opt, names, - [out_types[n] for n in names], psd=segments[0].psd, - gating_info=gwstrain.gating_info) + event_mgr = events.EventManager( + opt, + names, + [out_types[n] for n in names], + psd=segments[0].psd, + gating_info=gwstrain.gating_info, + ) logging.info("Read in template bank") lfc = None if opt.enable_bank_start_frequency else flow - bank = waveform.FilterBankSkyMax(opt.bank_file, flen, - delta_f, complex64, phase_order=opt.order, - taper=opt.taper_template, approximant=opt.approximant, - out_plus=zeros(tlen, dtype=complex64), - out_cross=zeros(tlen, dtype=complex64), - max_template_length=opt.max_template_length, - low_frequency_cutoff=lfc) + bank = waveform.FilterBankSkyMax( + opt.bank_file, + flen, + delta_f, + complex64, + phase_order=opt.order, + taper=opt.taper_template, + approximant=opt.approximant, + out_plus=zeros(tlen, dtype=complex64), + out_cross=zeros(tlen, dtype=complex64), + max_template_length=opt.max_template_length, + low_frequency_cutoff=lfc, + ) sg_chisq = SingleDetSGChisq.from_cli(opt, bank, opt.chisq_bins) @@ -289,70 +420,100 @@ with ctx: # quantities like effective distance may be more complicated! # These may depend on the [complex] overlap between plus and # cross. - event_mgr.new_template(tmplt=hplus.params, - sigmasq_plus=hplus.sigmasq(segments[0].psd), - sigmasq_cross=hcross.sigmasq(segments[0].psd)) + event_mgr.new_template( + tmplt=hplus.params, + sigmasq_plus=hplus.sigmasq(segments[0].psd), + sigmasq_cross=hcross.sigmasq(segments[0].psd), + ) if opt.cluster_method == "window": cluster_window = int(opt.cluster_window * gwstrain.sample_rate) elif opt.cluster_method == "template": cluster_window = int(template.chirp_length * gwstrain.sample_rate) for s_num, stilde in enumerate(segments): - logging.info("Filtering template %d/%d segment %d/%d" % \ - (t_num + 1, len(bank), s_num + 1, len(segments))) + logging.info( + "Filtering template %d/%d segment %d/%d" + % (t_num + 1, len(bank), s_num + 1, len(segments)) + ) # FIXME: Why this takes more inputs than for non-spin? - snr, corr_plus, corr_cross, idx, snrv, u_vals, coa_phase,\ - hplus_cross_corr, hpnorm, hcnorm =\ - matched_filter.matched_filter_and_cluster(hplus, hcross, - hplus.sigmasq(stilde.psd), - hcross.sigmasq(stilde.psd), - stilde.psd, stilde, cluster_window) + ( + snr, + corr_plus, + corr_cross, + idx, + snrv, + u_vals, + coa_phase, + hplus_cross_corr, + hpnorm, + hcnorm, + ) = matched_filter.matched_filter_and_cluster( + hplus, + hcross, + hplus.sigmasq(stilde.psd), + hcross.sigmasq(stilde.psd), + stilde.psd, + stilde, + cluster_window, + ) if not len(idx): continue - out_vals['u_vals'] = u_vals - out_vals['coa_phase'] = coa_phase - out_vals['hplus_cross_corr'] = numpy.repeat(hplus_cross_corr, - len(u_vals)) + out_vals["u_vals"] = u_vals + out_vals["coa_phase"] = coa_phase + out_vals["hplus_cross_corr"] = numpy.repeat(hplus_cross_corr, len(u_vals)) # This hasn't been implemented yet, but the stub is still here. - out_vals['bank_chisq'], out_vals['bank_chisq_dof'] = \ - bank_chisq.values() - - out_vals['chisq'], out_vals['chisq_dof'] = power_chisq.values(\ - corr_plus, corr_cross, snrv, stilde.psd, - idx+stilde.analyze.start, hplus, hcross, u_vals, - hplus_cross_corr, hpnorm, hcnorm) - - out_vals['sg_chisq'] = sg_chisq.values(stilde, hplus, stilde.psd, - snrv, 1., - out_vals['chisq'], - out_vals['chisq_dof'], - idx+stilde.analyze.start) + out_vals["bank_chisq"], out_vals["bank_chisq_dof"] = bank_chisq.values() + + out_vals["chisq"], out_vals["chisq_dof"] = power_chisq.values( + corr_plus, + corr_cross, + snrv, + stilde.psd, + idx + stilde.analyze.start, + hplus, + hcross, + u_vals, + hplus_cross_corr, + hpnorm, + hcnorm, + ) + + out_vals["sg_chisq"] = sg_chisq.values( + stilde, + hplus, + stilde.psd, + snrv, + 1.0, + out_vals["chisq"], + out_vals["chisq_dof"], + idx + stilde.analyze.start, + ) # This hasn't been implemented yet, but the stub is still here. - out_vals['cont_chisq'] = \ - autochisq.values() + out_vals["cont_chisq"] = autochisq.values() idx += stilde.cumulative_index - out_vals['time_index'] = idx - out_vals['snr'] = snrv + out_vals["time_index"] = idx + out_vals["snr"] = snrv if opt.psdvar_short_segment is not None: - out_vals['psd_var_val'] = \ - pycbc.psd.find_trigger_value(psd_var, - out_vals['time_index'], - opt.gps_start_time, opt.sample_rate) + out_vals["psd_var_val"] = pycbc.psd.find_trigger_value( + psd_var, out_vals["time_index"], opt.gps_start_time, opt.sample_rate + ) event_mgr.add_template_events(names, [out_vals[n] for n in names]) event_mgr.cluster_template_events("time_index", "snr", cluster_window) event_mgr.finalize_template_events() - if opt.finalize_events_template_rate is not None and \ - not (t_num+1) % opt.finalize_events_template_rate: + if ( + opt.finalize_events_template_rate is not None + and not (t_num + 1) % opt.finalize_events_template_rate + ): event_mgr.consolidate_events(opt, gwstrain=gwstrain) event_mgr.consolidate_events(opt, gwstrain=gwstrain) @@ -369,4 +530,3 @@ if opt.fftw_output_double_wisdom_file: fft.fftw.export_double_wisdom_to_filename(opt.fftw_output_double_wisdom_file) logging.info("Finished") - diff --git a/bin/pycbc_live b/bin/pycbc_live index 00454ae541d..b820fc9d5ec 100755 --- a/bin/pycbc_live +++ b/bin/pycbc_live @@ -1,48 +1,65 @@ #!/usr/bin/env python -"""Detect compact-binary-merger gravitational-wave signals in low latency. This +""" +Detect compact-binary-merger gravitational-wave signals in low latency. This program takes a fixed template bank and performs matched filtering on a continuous stream of strain data in fixed short blocks of time. It is capable of reading from low latency data directories on the LDG clusters. Work is parallelized through the use of MPI, so this script should typically be launched with mpirun. -See https://arxiv.org/abs/1805.11174 for an overview.""" +See https://arxiv.org/abs/1805.11174 for an overview. +""" -import sys -import argparse, numpy, pycbc, logging, cProfile, h5py, json -import os.path +import argparse +import cProfile import itertools +import json +import logging +import os.path import platform import subprocess +import sys from multiprocessing.dummy import threading + +import h5py +import numpy from matplotlib import use -use('agg') + +import pycbc + +use("agg") from shutil import which + from mpi4py import MPI as mpi -from pycbc.pool import BroadcastPool -from pycbc import fft, version, waveform, scheme, makedir -from pycbc.types import MultiDetOptionAction, MultiDetMultiColonOptionAction -from pycbc.filter import LiveBatchMatchedFilter, compute_followup_snr_series -from pycbc.filter import followup_event_significance -from pycbc.strain import StrainBuffer -from pycbc.events.ranking import newsnr + +import pycbc.waveform.bank +from pycbc import conversions as conv +from pycbc import fft, makedir, mchirp_area, scheme, version, waveform +from pycbc.detector import ppdets from pycbc.events.coinc import LiveCoincTimeslideBackgroundEstimator as Coincer +from pycbc.events.ranking import newsnr from pycbc.events.single import LiveSingle +from pycbc.filter import ( + LiveBatchMatchedFilter, + compute_followup_snr_series, + followup_event_significance, + resample, +) from pycbc.io.gracedb import CandidateForGraceDB from pycbc.io.hdf import recursively_save_dict_contents_to_group -from pycbc.types import positive_int -import pycbc.waveform.bank -from pycbc.vetoes.sgchisq import SingleDetSGChisq -from pycbc.waveform.waveform import props +from pycbc.pool import BroadcastPool from pycbc.population import live_pastro_utils as livepau -from pycbc import mchirp_area -from pycbc.detector import ppdets -from pycbc.filter import resample -from pycbc.psd import estimate -from pycbc.psd import variation -from pycbc import conversions as conv +from pycbc.psd import estimate, variation +from pycbc.strain import StrainBuffer from pycbc.time import gps_to_utc_datetime +from pycbc.types import ( + MultiDetMultiColonOptionAction, + MultiDetOptionAction, + positive_int, +) +from pycbc.vetoes.sgchisq import SingleDetSGChisq +from pycbc.waveform.waveform import props # Use cached class-based FFTs in the resample and estimate module resample.USE_CACHING_FOR_LFILTER = True @@ -52,32 +69,33 @@ estimate.USE_CACHING_FOR_INV_SPEC_TRUNC = True try: from setproctitle import setproctitle except ImportError: + def setproctitle(title): pass def ptof(p, ft): - """Convert p-value to FAR via foreground time `ft`. - """ + """Convert p-value to FAR via foreground time `ft`.""" return numpy.log1p(-p) / -ft + def ftop(f, ft): - """Convert FAR to p-value via foreground time `ft`. - """ + """Convert FAR to p-value via foreground time `ft`.""" return 1 - numpy.exp(-ft * f) + def combine_ifar_pvalue(ifar, pvalue, livetime): - """Convert original IFAR to p-value and combine with followup p-value. - """ + """Convert original IFAR to p-value and combine with followup p-value.""" from scipy.stats import combine_pvalues + # NB units of ifar and livetime must be the same - _, pnew = combine_pvalues([ftop(1. / ifar, livetime), pvalue]) - nifar = 1. / ptof(pnew, livetime) + _, pnew = combine_pvalues([ftop(1.0 / ifar, livetime), pvalue]) + nifar = 1.0 / ptof(pnew, livetime) # take max of original IFAR and combined IFAR & apply trials factor - return numpy.maximum(ifar, nifar) / 2. + return numpy.maximum(ifar, nifar) / 2.0 -class LiveEventManager(object): +class LiveEventManager: def __init__(self, args, bank): self.low_frequency_cutoff = args.low_frequency_cutoff self.bank = bank @@ -125,15 +143,15 @@ class LiveEventManager(object): # to followup processes without slowing down the main search bg_cores = len(tuple(itertools.combinations(self.trigg_ifos, 2))) analysis_cores = 1 + bg_cores - if platform.system() != 'Darwin': + if platform.system() != "Darwin": available_cores = len(os.sched_getaffinity(0)) self.fu_cores = available_cores - analysis_cores if self.fu_cores <= 0: logging.warning( - 'Insufficient number of CPU cores (%d) to ' - 'run search and trigger followups. Uploaded ' - 'triggers will momentarily increase the lag', - available_cores + "Insufficient number of CPU cores (%d) to " + "run search and trigger followups. Uploaded " + "triggers will momentarily increase the lag", + available_cores, ) self.fu_cores = 1 else: @@ -141,12 +159,12 @@ class LiveEventManager(object): self.fu_cores = 1 if args.enable_embright_has_massgap: - if args.embright_massgap_max < self.mc_area_args['mass_bdary']['ns_max']: - parser.error('MAX_GAP value cannot be lower than MAX_NS limit') - self.mc_area_args['embright_mg_max'] = args.embright_massgap_max + if args.embright_massgap_max < self.mc_area_args["mass_bdary"]["ns_max"]: + parser.error("MAX_GAP value cannot be lower than MAX_NS limit") + self.mc_area_args["embright_mg_max"] = args.embright_massgap_max def commit_results(self, results): - logging.info('Committing triggers') + logging.info("Committing triggers") self.comm.gather(results, root=0) def barrier(self): @@ -156,14 +174,14 @@ class LiveEventManager(object): return self.comm.allreduce(status, op=mpi.LAND) def gather_results(self): - """ Collect results from the mpi subprocesses and collate them into + """ + Collect results from the mpi subprocesses and collate them into contiguous sets of arrays. """ - if self.rank != 0: - raise RuntimeError('Not root process') + raise RuntimeError("Not root process") - logging.info('Gathering triggers') + logging.info("Gathering triggers") all_results = self.comm.gather(None, root=0) data_ends = [a[1] for a in all_results if a is not None] results = [a[0] for a in all_results if a is not None] @@ -183,20 +201,20 @@ class LiveEventManager(object): return combined, data_ends[0] - def compute_followup_data(self, ifos, triggers, - followup_ifos=None, recalculate_ifar=False): - """Figure out which of the followup detectors are usable, and compute + def compute_followup_data( + self, ifos, triggers, followup_ifos=None, recalculate_ifar=False + ): + """ + Figure out which of the followup detectors are usable, and compute SNR time series for all the available detectors. """ out = {} followup_ifos = [] if followup_ifos is None else followup_ifos - template_id = triggers[f'foreground/{ifos[0]}/template_id'] + template_id = triggers[f"foreground/{ifos[0]}/template_id"] htilde = self.bank[template_id] - coinc_times = { - ifo: triggers[f'foreground/{ifo}/end_time'] for ifo in ifos - } + coinc_times = {ifo: triggers[f"foreground/{ifo}/end_time"] for ifo in ifos} # Get the SNR series for the ifos that made the initial coinc for ifo in ifos: @@ -204,14 +222,11 @@ class LiveEventManager(object): # IFOs producing the coincidence are assumed to also # produce valid SNR series. snr_series = compute_followup_snr_series( - self.data_readers[ifo], - htilde, - coinc_times[ifo], - check_state=False + self.data_readers[ifo], htilde, coinc_times[ifo], check_state=False ) if snr_series is not None: - out[ifo] = {'snr_series': snr_series} + out[ifo] = {"snr_series": snr_series} # Determine if the other ifos can contribute to the coincident event, # if so then update the info in `triggers` appropriately @@ -222,83 +237,82 @@ class LiveEventManager(object): self.bank, template_id, coinc_times, - lookback=self.pvalue_lookback_time + lookback=self.pvalue_lookback_time, ) if pvalue_info is None: continue - out[ifo] = {'snr_series': pvalue_info['snr_series']} + out[ifo] = {"snr_series": pvalue_info["snr_series"]} self.get_followup_info( ifos[0], ifo, triggers, pvalue_info, - recalculate_ifar=recalculate_ifar and ifo not in self.skymap_only_ifos + recalculate_ifar=recalculate_ifar and ifo not in self.skymap_only_ifos, ) # the SNR time series sample rate can vary slightly due to # rounding errors, so force all of them to be identical fix_delta_t = None for ifo in out: - if 'snr_series' not in out[ifo]: + if "snr_series" not in out[ifo]: continue if fix_delta_t is None: - fix_delta_t = out[ifo]['snr_series']._delta_t + fix_delta_t = out[ifo]["snr_series"]._delta_t else: - out[ifo]['snr_series']._delta_t = fix_delta_t + out[ifo]["snr_series"]._delta_t = fix_delta_t return out def get_followup_info( - self, - coinc_ifo, - ifo, - triggers, - pvalue_info, - recalculate_ifar=False + self, coinc_ifo, ifo, triggers, pvalue_info, recalculate_ifar=False ): - peak_time = pvalue_info['peak_time'] - pv = pvalue_info['pvalue'] - pv_sat = pvalue_info['pvalue_saturated'] + peak_time = pvalue_info["peak_time"] + pv = pvalue_info["pvalue"] + pv_sat = pvalue_info["pvalue_saturated"] # Copy the common fields from the other detector; # ignore fields that contain detector-specific data - fields_to_ignore = set([ - 'end_time', 'snr', 'stat', 'coa_phase', - 'chisq', 'chisq_dof', 'sg_chisq', 'sigmasq' - ]) + fields_to_ignore = set( + [ + "end_time", + "snr", + "stat", + "coa_phase", + "chisq", + "chisq_dof", + "sg_chisq", + "sigmasq", + ] + ) for key in set(triggers): - if f'foreground/{coinc_ifo}/' not in key: + if f"foreground/{coinc_ifo}/" not in key: continue - _, _, name = key.split('/') + _, _, name = key.split("/") if name in fields_to_ignore: continue - triggers[f'foreground/{ifo}/{name}'] = triggers[key] + triggers[f"foreground/{ifo}/{name}"] = triggers[key] # Set the detector-specific fields for which we have data - snr_series_peak = pvalue_info['snr_series'].at_time( - peak_time, - nearest_sample=True + snr_series_peak = pvalue_info["snr_series"].at_time( + peak_time, nearest_sample=True ) - base = f'foreground/{ifo}/' - triggers[base + 'end_time'] = float(peak_time) - triggers[base + 'snr'] = triggers[base + 'stat'] = abs(snr_series_peak) - triggers[base + 'coa_phase'] = numpy.angle(snr_series_peak) - triggers[base + 'sigmasq'] = pvalue_info['sigma2'] + base = f"foreground/{ifo}/" + triggers[base + "end_time"] = float(peak_time) + triggers[base + "snr"] = triggers[base + "stat"] = abs(snr_series_peak) + triggers[base + "coa_phase"] = numpy.angle(snr_series_peak) + triggers[base + "sigmasq"] = pvalue_info["sigma2"] if recalculate_ifar: # Calculate new ifar - triggers['foreground/ifar'] = combine_ifar_pvalue( - triggers['foreground/ifar'], - pv, - self.pvalue_livetime + triggers["foreground/ifar"] = combine_ifar_pvalue( + triggers["foreground/ifar"], pv, self.pvalue_livetime ) - triggers['foreground/ifar_saturated'] |= pv_sat - triggers[f'foreground/pvalue_{ifo}'] = pv - triggers[f'foreground/pvalue_{ifo}_saturated'] = pv_sat + triggers["foreground/ifar_saturated"] |= pv_sat + triggers[f"foreground/pvalue_{ifo}"] = pv + triggers[f"foreground/pvalue_{ifo}_saturated"] = pv_sat - def setup_optimize_snr( - self, results, live_ifos, triggering_ifos, fname - ): - """Setup the network SNR optimization process for a + def setup_optimize_snr(self, results, live_ifos, triggering_ifos, fname): + """ + Setup the network SNR optimization process for a candidate event. See arXiv:2008.07494 for details. Parameters @@ -315,97 +329,93 @@ class LiveEventManager(object): File name where the candidate information has been saved. Used to name other files produced as part of setting up the SNR optimization. + """ - template_id = results[f'foreground/{live_ifos[0]}/template_id'] + template_id = results[f"foreground/{live_ifos[0]}/template_id"] p = props(self.bank.table[template_id]) - p.pop('approximant') + p.pop("approximant") apr = self.bank.approximant(template_id) min_buffer = self.bank.minimum_buffer + 0.5 - buff_size = \ - pycbc.waveform.get_waveform_filter_length_in_time(apr, **p) + buff_size = pycbc.waveform.get_waveform_filter_length_in_time(apr, **p) - tlen = self.bank.round_up((buff_size + min_buffer) \ - * self.bank.sample_rate) + tlen = self.bank.round_up((buff_size + min_buffer) * self.bank.sample_rate) flen = int(tlen / 2 + 1) delta_f = self.bank.sample_rate / float(tlen) - cmd = f'timeout {args.snr_opt_timeout} ' - exepath = which('pycbc_optimize_snr') - cmd += exepath + ' ' + cmd = f"timeout {args.snr_opt_timeout} " + exepath = which("pycbc_optimize_snr") + cmd += exepath + " " # Add SNR optimization options to the command if self.snr_opt_options is not None: cmd += self.snr_opt_options # Add data files according to the event we are optimizing - data_fils_str = '--data-files ' - psd_fils_str = '--psd-files ' + data_fils_str = "--data-files " + psd_fils_str = "--psd-files " for ifo in live_ifos: - curr_fname = fname.replace( - '.xml.gz', f'_{ifo}_data_overwhitened.hdf' - ) + curr_fname = fname.replace(".xml.gz", f"_{ifo}_data_overwhitened.hdf") curr_data = self.data_readers[ifo].overwhitened_data(delta_f) curr_data.save(curr_fname) - data_fils_str += f'{ifo}:{curr_fname} ' - curr_fname = fname.replace('.xml.gz', f'_{ifo}_psd.hdf') + data_fils_str += f"{ifo}:{curr_fname} " + curr_fname = fname.replace(".xml.gz", f"_{ifo}_psd.hdf") curr_psd = curr_data.psd curr_psd.save(curr_fname) - psd_fils_str += f'{ifo}:{curr_fname} ' + psd_fils_str += f"{ifo}:{curr_fname} " cmd += data_fils_str cmd += psd_fils_str - curr_fname = fname.replace('.xml.gz', '_attributes.hdf') - with h5py.File(curr_fname, 'w') as hdfp: + curr_fname = fname.replace(".xml.gz", "_attributes.hdf") + with h5py.File(curr_fname, "w") as hdfp: for ifo in triggering_ifos: - curr_time = results[f'foreground/{ifo}/end_time'] - hdfp[f'coinc_times/{ifo}'] = curr_time + curr_time = results[f"foreground/{ifo}/end_time"] + hdfp[f"coinc_times/{ifo}"] = curr_time f_end = self.bank.end_frequency(template_id) if f_end is None or f_end >= (flen * delta_f): - f_end = (flen-1) * delta_f - hdfp['flen'] = flen - hdfp['delta_f'] = delta_f - hdfp['f_end'] = f_end - hdfp['sample_rate'] = self.bank.sample_rate - hdfp['flow'] = self.bank.table[template_id].f_lower - hdfp['mass1'] = self.bank.table[template_id].mass1 - hdfp['mass2'] = self.bank.table[template_id].mass2 - hdfp['spin1z'] = self.bank.table[template_id].spin1z - hdfp['spin2z'] = self.bank.table[template_id].spin2z - hdfp['template_duration'] = \ - self.bank.table[template_id].template_duration - hdfp['ifar'] = results['foreground/ifar'] - if 'p_terr' in results: - hdfp['p_terr'] = results['p_terr'] + f_end = (flen - 1) * delta_f + hdfp["flen"] = flen + hdfp["delta_f"] = delta_f + hdfp["f_end"] = f_end + hdfp["sample_rate"] = self.bank.sample_rate + hdfp["flow"] = self.bank.table[template_id].f_lower + hdfp["mass1"] = self.bank.table[template_id].mass1 + hdfp["mass2"] = self.bank.table[template_id].mass2 + hdfp["spin1z"] = self.bank.table[template_id].spin1z + hdfp["spin2z"] = self.bank.table[template_id].spin2z + hdfp["template_duration"] = self.bank.table[template_id].template_duration + hdfp["ifar"] = results["foreground/ifar"] + if "p_terr" in results: + hdfp["p_terr"] = results["p_terr"] for ifo in args.channel_name: - hdfp[f'channel_names/{ifo}'] = args.channel_name[ifo] + hdfp[f"channel_names/{ifo}"] = args.channel_name[ifo] - recursively_save_dict_contents_to_group(hdfp, - 'mc_area_args/', - self.mc_area_args) + recursively_save_dict_contents_to_group( + hdfp, "mc_area_args/", self.mc_area_args + ) - cmd += f'--params-file {curr_fname} ' - cmd += f'--approximant {apr} ' - cmd += f'--gracedb-server {self.gracedb_server} ' - cmd += f'--gracedb-search {self.gracedb_search} ' + cmd += f"--params-file {curr_fname} " + cmd += f"--approximant {apr} " + cmd += f"--gracedb-server {self.gracedb_server} " + cmd += f"--gracedb-search {self.gracedb_search} " labels = self.snr_opt_label - labels += ' '.join(self.gracedb_labels or []) - cmd += f'--gracedb-labels {labels} ' + labels += " ".join(self.gracedb_labels or []) + cmd += f"--gracedb-labels {labels} " if not self.gracedb_testing: - cmd += '--production ' - cmd += '--verbose ' + cmd += "--production " + cmd += "--verbose " # set up a place for storing the SNR optimization results and log - out_dir_path = os.path.join(os.path.dirname(fname), 'optimize_snr') + out_dir_path = os.path.join(os.path.dirname(fname), "optimize_snr") makedir(out_dir_path) - cmd += f'--output-path {out_dir_path} ' + cmd += f"--output-path {out_dir_path} " if self.enable_gracedb_upload: - cmd += '--enable-gracedb-upload ' + cmd += "--enable-gracedb-upload " if self.fu_cores: - cmd += f'--cores {self.fu_cores} ' + cmd += f"--cores {self.fu_cores} " if args.processing_scheme: # we will use the cores for multiple workers of the @@ -416,12 +426,12 @@ class LiveEventManager(object): # is expected to be in waveform generation, which is # unlikely to benefit from a processing scheme with more # than 1 thread anyway. - opt_scheme = args.processing_scheme.split(':')[0] - cmd += f'--processing-scheme {opt_scheme}:1 ' + opt_scheme = args.processing_scheme.split(":")[0] + cmd += f"--processing-scheme {opt_scheme}:1 " # Save the command which would be used: - snroc_fname = os.path.join(out_dir_path, 'snr_optimize_command.txt') - with open(snroc_fname, 'w') as snroc_file: + snroc_fname = os.path.join(out_dir_path, "snr_optimize_command.txt") + with open(snroc_fname, "w") as snroc_file: snroc_file.write(cmd) return cmd, out_dir_path @@ -441,17 +451,16 @@ class LiveEventManager(object): The filename of the attributes of the candidate gid: str GraceDB ID of the candidate + """ if gid is not None: - with h5py.File(attribute_fname, 'a') as hdfp: - hdfp['gid'] = gid - log_fname = os.path.join(out_dir_path, 'optimize_snr.log') + with h5py.File(attribute_fname, "a") as hdfp: + hdfp["gid"] = gid + log_fname = os.path.join(out_dir_path, "optimize_snr.log") - logging.info('running %s with log to %s', cmd, log_fname) + logging.info("running %s with log to %s", cmd, log_fname) with open(log_fname, "w") as logfile: - subprocess.Popen( - cmd, shell=True, stdout=logfile, stderr=logfile - ) + subprocess.Popen(cmd, shell=True, stdout=logfile, stderr=logfile) def create_gdb(self): """ @@ -460,19 +469,27 @@ class LiveEventManager(object): """ from ligo.gracedb.rest import GraceDb - logging.info('Creating GraceDB session') + logging.info("Creating GraceDB session") # Set up dict for args that reload the certificate if it will expire in the # reload_buffer time. These args are documented here: # https://ligo-gracedb.readthedocs.io/en/latest/api.html#ligo.gracedb.rest.GraceDb # Because we do not change any of the request session values when running the # code, it should remain thread safe. - gdbargs = {'reload_certificate': True, 'reload_buffer': 300} + gdbargs = {"reload_certificate": True, "reload_buffer": 300} if self.gracedb_server: - gdbargs['service_url'] = self.gracedb_server + gdbargs["service_url"] = self.gracedb_server self.gracedb = GraceDb(**gdbargs) - def upload_in_thread(self, event, fname, comment, cmd, out_dir_path, - upload_checks, optimize_snr_checks): + def upload_in_thread( + self, + event, + fname, + comment, + cmd, + out_dir_path, + upload_checks, + optimize_snr_checks, + ): gid = None if upload_checks: gid = event.upload( @@ -481,45 +498,41 @@ class LiveEventManager(object): testing=self.gracedb_testing, extra_strings=[comment], search=self.gracedb_search, - labels=self.gracedb_labels + labels=self.gracedb_labels, ) if optimize_snr_checks: - logging.info('Optimizing SNR for event above threshold ..') + logging.info("Optimizing SNR for event above threshold ..") self.run_optimize_snr( - cmd, - out_dir_path, - fname.replace('.xml.gz', '_attributes.hdf'), - gid + cmd, out_dir_path, fname.replace(".xml.gz", "_attributes.hdf"), gid ) def check_coincs(self, ifos, coinc_results, psds): - """ Perform any followup and save zerolag triggers to a coinc xml file - """ + """Perform any followup and save zerolag triggers to a coinc xml file""" # Check for hardware injection for ifo in ifos: if self.data_readers[ifo].near_hwinj(): - coinc_results['HWINJ'] = True + coinc_results["HWINJ"] = True break - if 'foreground/ifar' not in coinc_results: + if "foreground/ifar" not in coinc_results: return - logging.info('computing followup data for coinc') - coinc_ifos = coinc_results['foreground/type'].split('-') + logging.info("computing followup data for coinc") + coinc_ifos = coinc_results["foreground/type"].split("-") followup_ifos = set(ifos) - set(coinc_ifos) followup_ifos = list(followup_ifos | self.skymap_only_ifos) - double_ifar = coinc_results['foreground/ifar'] + double_ifar = coinc_results["foreground/ifar"] if double_ifar < args.ifar_double_followup_threshold: - coinc_results['foreground/NO_FOLLOWUP'] = True + coinc_results["foreground/NO_FOLLOWUP"] = True return sld = self.compute_followup_data( coinc_ifos, coinc_results, followup_ifos=followup_ifos, - recalculate_ifar=True + recalculate_ifar=True, ) event = CandidateForGraceDB( @@ -538,32 +551,34 @@ class LiveEventManager(object): fname = os.path.join( self.get_candidate_path(event.merger_time), - f'coinc-{event.merger_time:.3f}.xml.gz' + f"coinc-{event.merger_time:.3f}.xml.gz", ) - logging.info('Coincident candidate! Saving as %s', fname) + logging.info("Coincident candidate! Saving as %s", fname) # Which IFOs were active? - live_ifos = [ifo for ifo in sld if 'snr_series' in sld[ifo]] + live_ifos = [ifo for ifo in sld if "snr_series" in sld[ifo]] # Verbally explain some details not obvious from the other info comment = ( - 'Trigger produced as a {} coincidence.
    ' - 'Two-detector ranking statistic: {}
    ' - 'Detectors used for FAR calculation: {}.
    ' - 'Detectors used for localization: {}.
    ' - 'Detectors used only for localization: {}.' + "Trigger produced as a {} coincidence.
    " + "Two-detector ranking statistic: {}
    " + "Detectors used for FAR calculation: {}.
    " + "Detectors used for localization: {}.
    " + "Detectors used only for localization: {}." ) comment = comment.format( ppdets(coinc_ifos), args.ranking_statistic, ppdets(set(ifos) - self.skymap_only_ifos), ppdets(live_ifos), - ppdets(set(live_ifos) & self.skymap_only_ifos) + ppdets(set(live_ifos) & self.skymap_only_ifos), ) - ifar = coinc_results['foreground/ifar'] + ifar = coinc_results["foreground/ifar"] upload_checks = self.enable_gracedb_upload and self.ifar_upload_threshold < ifar - optimize_snr_checks = self.run_snr_optimization and self.ifar_upload_threshold < ifar + optimize_snr_checks = ( + self.run_snr_optimization and self.ifar_upload_threshold < ifar + ) # Keep track of the last few coincs uploaded in order to # prevent singles being uploaded as well for coinc events @@ -580,8 +595,8 @@ class LiveEventManager(object): # data buffers move on in a possible interim period # Tell SNR optimized event about p_terr - if hasattr(event, 'p_terr') and event.p_terr is not None: - coinc_results['p_terr'] = event.p_terr + if hasattr(event, "p_terr") and event.p_terr is not None: + coinc_results["p_terr"] = event.p_terr # Do the optimizer setup cmd, out_dir_path = self.setup_optimize_snr( @@ -594,38 +609,38 @@ class LiveEventManager(object): # Now start the thread to run the SNR optimizer if upload_checks or optimize_snr_checks: thread_args = ( - event, fname, comment, cmd, out_dir_path, - upload_checks, optimize_snr_checks + event, + fname, + comment, + cmd, + out_dir_path, + upload_checks, + optimize_snr_checks, + ) + gdb_upload_thread = threading.Thread( + target=self.upload_in_thread, args=thread_args ) - gdb_upload_thread = threading.Thread(target=self.upload_in_thread, - args=thread_args) gdb_upload_thread.start() def check_singles(self, results, psds): active = [k for k in results if results[k] is not None] for ifo in active: - single = sngl_estimator[ifo].check( - results[ifo], - self.data_readers[ifo] - ) + single = sngl_estimator[ifo].check(results[ifo], self.data_readers[ifo]) if single is None: continue - sifar = single['foreground/ifar'] - logging.info(f'Found {ifo} single with ifar {sifar}') + sifar = single["foreground/ifar"] + logging.info(f"Found {ifo} single with ifar {sifar}") followup_ifos = [i for i in active if i is not ifo] followup_ifos = list(set(followup_ifos) | self.skymap_only_ifos) # Don't recompute ifar considering other ifos sld = self.compute_followup_data( - [ifo], - single, - followup_ifos=followup_ifos, - recalculate_ifar=False + [ifo], single, followup_ifos=followup_ifos, recalculate_ifar=False ) # Apply a trials factor of the number of active detectors - single['foreground/ifar'] = sifar / len(active) + single["foreground/ifar"] = sifar / len(active) event = CandidateForGraceDB( [ifo], @@ -643,47 +658,49 @@ class LiveEventManager(object): fname = os.path.join( self.get_candidate_path(event.merger_time), - f'single-{ifo}-{event.merger_time:.3f}.xml.gz' + f"single-{ifo}-{event.merger_time:.3f}.xml.gz", ) - logging.info('Single-detector candidate! Saving as %s', fname) + logging.info("Single-detector candidate! Saving as %s", fname) # Which IFOs were active? - live_ifos = [ifo for ifo in sld if 'snr_series' in sld[ifo]] + live_ifos = [ifo for ifo in sld if "snr_series" in sld[ifo]] # Verbally explain some details not obvious from the other info comment = ( - 'Trigger produced as a {0} single.
    ' - 'Detectors used for FAR calculation: {0}.
    ' - 'Detectors used for localization: {1}.
    ' - 'Detectors used only for localization: {2}.' + "Trigger produced as a {0} single.
    " + "Detectors used for FAR calculation: {0}.
    " + "Detectors used for localization: {1}.
    " + "Detectors used only for localization: {2}." ) comment = comment.format( - ifo, - ppdets(live_ifos), - ppdets(set(live_ifos) & self.skymap_only_ifos) + ifo, ppdets(live_ifos), ppdets(set(live_ifos) & self.skymap_only_ifos) ) # Has a coinc event at this time been uploaded recently? # If so, skip upload - Note that this means that we _always_ # prefer an uploaded coincident event over a single - coinc_tdiffs = abs(event.merger_time - - numpy.array(self.last_few_coincs_uploaded)) + coinc_tdiffs = abs( + event.merger_time - numpy.array(self.last_few_coincs_uploaded) + ) nearby_coincs = coinc_tdiffs < 0.1 if any(nearby_coincs): logging.info( "Single-detector candidate at time %.3f was already " "reported as coinc, skipping upload", - event.merger_time + event.merger_time, ) - upload_checks = args.enable_single_detector_upload and \ - self.ifar_upload_threshold < single['foreground/ifar'] and \ - not any(nearby_coincs) - optimize_snr_checks = \ - self.ifar_upload_threshold < single['foreground/ifar'] and \ - not any(nearby_coincs) and \ - self.run_snr_optimization + upload_checks = ( + args.enable_single_detector_upload + and self.ifar_upload_threshold < single["foreground/ifar"] + and not any(nearby_coincs) + ) + optimize_snr_checks = ( + self.ifar_upload_threshold < single["foreground/ifar"] + and not any(nearby_coincs) + and self.run_snr_optimization + ) # Save the event if not upload_checks: @@ -699,8 +716,8 @@ class LiveEventManager(object): continue # Tell SNR optimized event about p_terr - if hasattr(event, 'p_terr') and event.p_terr is not None: - single['p_terr'] = event.p_terr + if hasattr(event, "p_terr") and event.p_terr is not None: + single["p_terr"] = event.p_terr # Do the optimizer setup cmd, out_dir_path = self.setup_optimize_snr( @@ -712,15 +729,22 @@ class LiveEventManager(object): if upload_checks or optimize_snr_checks: thread_args = ( - event, fname, comment, cmd, out_dir_path, - upload_checks, optimize_snr_checks + event, + fname, + comment, + cmd, + out_dir_path, + upload_checks, + optimize_snr_checks, + ) + gdb_upload_thread = threading.Thread( + target=self.upload_in_thread, args=thread_args ) - gdb_upload_thread = threading.Thread(target=self.upload_in_thread, - args=thread_args) gdb_upload_thread.start() def get_out_dir_path(self, time_gps): - """Compose and return the path to a directory to store files associated + """ + Compose and return the path to a directory to store files associated with times in a given day (specified via a GPS time). Create the directory if it does not exist. """ @@ -728,57 +752,61 @@ class LiveEventManager(object): if self.use_date_prefix: time_dt = gps_to_utc_datetime(time_gps) path = os.path.join( - path, - f'{time_dt.year:04d}_{time_dt.month:02d}_{time_dt.day:02d}' + path, f"{time_dt.year:04d}_{time_dt.month:02d}_{time_dt.day:02d}" ) makedir(path) return path def get_candidate_path(self, merger_time): - """Compose and return the path to a directory to store files associated + """ + Compose and return the path to a directory to store files associated with a particular candidate event. Create the directory if it does not exist. """ path = os.path.join( self.get_out_dir_path(merger_time), - f'candidate_{merger_time:.3f}_{pycbc.random_string(4)}' + f"candidate_{merger_time:.3f}_{pycbc.random_string(4)}", ) makedir(path) return path - def dump(self, results, name, store_psd=False, time_index=None, - store_loudest_index=False, raw_results=None, gates=None): - """Save the results from this time block to an hdf output file. - """ - fname = os.path.join( - self.get_out_dir_path(time_index), - name + '.hdf' - ) + def dump( + self, + results, + name, + store_psd=False, + time_index=None, + store_loudest_index=False, + raw_results=None, + gates=None, + ): + """Save the results from this time block to an hdf output file.""" + fname = os.path.join(self.get_out_dir_path(time_index), name + ".hdf") - with h5py.File(fname, 'w') as f: - f.attrs['pycbc_version'] = version.git_verbose_msg - f.attrs['command_line'] = sys.argv - f.attrs['num_live_detectors'] = len(self.live_detectors) + with h5py.File(fname, "w") as f: + f.attrs["pycbc_version"] = version.git_verbose_msg + f.attrs["command_line"] = sys.argv + f.attrs["num_live_detectors"] = len(self.live_detectors) def h5py_unicode_workaround(stuff): # workaround for the fact that h5py cannot handle Unicode # numpy arrays that show up with Python 3 - if hasattr(stuff, 'dtype') and stuff.dtype.kind == 'U': + if hasattr(stuff, "dtype") and stuff.dtype.kind == "U": return [s.encode() for s in stuff] return stuff for ifo in results: for k in results[ifo]: - f[f'{ifo}/{k}'] = h5py_unicode_workaround(results[ifo][k]) + f[f"{ifo}/{k}"] = h5py_unicode_workaround(results[ifo][k]) for key in raw_results: f[key] = h5py_unicode_workaround(raw_results[key]) if store_loudest_index: for ifo in results: - if 'snr' in results[ifo]: - s = numpy.array(results[ifo]['snr'], ndmin=1) - c = numpy.array(results[ifo]['chisq'], ndmin=1) + if "snr" in results[ifo]: + s = numpy.array(results[ifo]["snr"], ndmin=1) + c = numpy.array(results[ifo]["chisq"], ndmin=1) nsnr = numpy.array(newsnr(s, c), ndmin=1) if len(s) > 0 else [] # loudest by newsnr @@ -786,21 +814,24 @@ class LiveEventManager(object): # loudest by snr sloudest = numpy.argsort(s)[::-1][0:store_loudest_index] - f[ifo]['loudest'] = numpy.union1d(nloudest, sloudest) + f[ifo]["loudest"] = numpy.union1d(nloudest, sloudest) - for ifo in (gates or {}): - gate_dtype = [('center_time', float), - ('zero_half_width', float), - ('taper_width', float)] - f[f'{ifo}/gates'] = numpy.array(gates[ifo], dtype=gate_dtype) + for ifo in gates or {}: + gate_dtype = [ + ("center_time", float), + ("zero_half_width", float), + ("taper_width", float), + ] + f[f"{ifo}/gates"] = numpy.array(gates[ifo], dtype=gate_dtype) - for ifo in (store_psd or {}): + for ifo in store_psd or {}: if store_psd[ifo] is not None: - store_psd[ifo].save(fname, group=f'{ifo}/psd') + store_psd[ifo].save(fname, group=f"{ifo}/psd") def check_max_length(args, waveforms): - """Check that the `--max-length` option is sufficient to accomodate the + """ + Check that the `--max-length` option is sufficient to accomodate the longest template in the bank and the PSD estimation options. """ lengths = numpy.array([1.0 / wf.delta_f for wf in waveforms]) @@ -808,245 +839,465 @@ def check_max_length(args, waveforms): max_length = max(lengths.max() + args.pvalue_lookback_time, psd_len) if max_length > args.max_length: raise ValueError( - '--max-length is too short for this template bank. ' - f'Use at least {max_length}.' + "--max-length is too short for this template bank. " + f"Use at least {max_length}." ) parser = argparse.ArgumentParser(description=__doc__) pycbc.waveform.bank.add_approximant_arg(parser) -parser.add_argument('--verbose', action='count') -parser.add_argument('--version', action='version', version=version.git_verbose_msg) -parser.add_argument('--bank-file', required=True, - help="Template bank file in XML or HDF format") -parser.add_argument('--low-frequency-cutoff', help="low frequency cutoff", type=int) -parser.add_argument('--sample-rate', help="output sample rate", type=int) -parser.add_argument('--chisq-bins', help="Number of chisq bins") -parser.add_argument('--analysis-chunk', type=int, required=True, - help="Amount of data to produce triggers in a block") - -parser.add_argument('--snr-threshold', type=float, - help='SNR threshold for generating a trigger') -parser.add_argument('--snr-abort-threshold', type=float) - -parser.add_argument('--channel-name', action=MultiDetMultiColonOptionAction, - required=True) -parser.add_argument('--state-channel', action=MultiDetMultiColonOptionAction, - help="Channel containing frame status information. Used " - "to determine when to analyze the hoft data. This somewhat" - " corresponds to CAT1 information") -parser.add_argument('--analyze-flags', action=MultiDetOptionAction, nargs='+', - help='The flags that must be in the "good" state to analyze data') -parser.add_argument('--idq-channel', action=MultiDetMultiColonOptionAction, - help="Channel to read iDQ timeseries from. iDQ timeseries are " - "required by certain ranking statistics.") -parser.add_argument('--idq-state-channel', action=MultiDetMultiColonOptionAction, - help="Channel containing information about the state of idq." - "Used to determine when idq is usable for statistics. ") -parser.add_argument('--idq-threshold', type=float, - help='Threshold used to veto triggers at times of ' - 'low iDQ False Alarm Probability') -parser.add_argument('--idq-reweighting', action='store_true',default=False, - help='Reweight triggers based on iDQ False Alarm Probability') -parser.add_argument('--data-quality-channel', - action=MultiDetMultiColonOptionAction, - help="Channel containing data quality information. Used " - "to determine when hoft may be suspect and may be used to " - "veto triggers or not analyze a segment of data. This " - "roughly corresponds to CAT2 information") -parser.add_argument('--data-quality-flags', action=MultiDetOptionAction, nargs='+', - help='Flags used to determine when to throw triggers away. ' - 'For each detector, give a comma-separated list of flags.') -parser.add_argument('--data-quality-padding', type=float, default=0, - help='Time in seconds around a bad dq time to additionally remove ' - 'triggers') -parser.add_argument('--frame-src', action=MultiDetOptionAction, nargs='+') -parser.add_argument('--frame-type', action=MultiDetOptionAction, nargs='+') -parser.add_argument('--force-update-cache', action='store_true') -parser.add_argument('--highpass-frequency', type=float, - help="Frequency to apply highpass filtering") -parser.add_argument('--highpass-reduction', type=float, - help="DB to reduce low frequencies") -parser.add_argument('--highpass-bandwidth', type=float, - help="Width of the highpass turnover region in Hz") -parser.add_argument('--psd-recalculate-difference', type=float, default=.01) -parser.add_argument('--psd-abort-difference', type=float, default=.20) -parser.add_argument('--psd-samples', type=int, required=True, - help="Number of PSD segments to use in the rolling estimate") -parser.add_argument('--psd-segment-length', type=int, required=True, - help="Length in seconds of each PSD segment") -parser.add_argument('--psd-inverse-length', type=float, - help="Length in time for the equivalent FIR filter") -parser.add_argument('--psd-recompute-length', type=positive_int, default=1, - help="If given only recompute the PSD after this number of " - "analysis chunks, whose length is set by the option " - "--analysis-chunk.") -parser.add_argument('--trim-padding', type=float, default=0.25, - help="Padding around the overwhitened analysis block") -parser.add_argument("--enable-bank-start-frequency", action='store_true', - help="Read the starting frequency of template waveforms" - " from the template bank") - -parser.add_argument('--autogating-threshold', type=float, metavar='SIGMA', - help='If given, find and gate glitches ' - 'producing a deviation larger than ' - 'SIGMA in the whitened strain time ' - 'series.') -parser.add_argument('--autogating-pad', type=float, default=0.5, - metavar='SECONDS', - help='Ignore the given length of whitened ' - 'strain at the ends of a segment, to ' - 'avoid filters ringing.') -parser.add_argument('--autogating-cluster', type=float, default=1, - metavar='SECONDS', - help='Length of clustering window for ' - 'detecting glitches for autogating.') -parser.add_argument('--autogating-width', type=float, default=0.25, - metavar='SECONDS', help='Half-width of the gating window.') -parser.add_argument('--autogating-taper', type=float, metavar='SECONDS', - default=0.25, - help='Taper the strain before and after ' - 'each gating window over a duration ' - 'of SECONDS.') -parser.add_argument('--autogating-duration', type=float, default=16, - metavar='SECONDS', - help='Amount of data in seconds to apply autogating on.') -parser.add_argument('--autogating-psd-segment-length', type=float, default=2, - metavar='SECONDS', - help='Length in seconds of each segment used to estimate the PSD ' - 'with Welchs method and median averaging.') -parser.add_argument('--autogating-psd-stride', type=float, default=1, - metavar='SECONDS', - help='Length in seconds of the overlap between each segment used ' - 'to estimate the PSD with Welchs method and median averaging.') - -parser.add_argument('--sync', action='store_true', - help="Imposes an MPI synchronization at each transfer of" - " single-detector triggers. Can help with debugging" - " and avoiding memory issues when running offline" - " analyses.") -parser.add_argument('--increment-update-cache', action=MultiDetOptionAction, nargs='+') -parser.add_argument('--frame-read-timeout', type=float, default=30) -parser.add_argument('--increment', type=int, default=8, metavar="T", - help="When generating the template waveforms, their" - " durations are binned such that T is their greatest" - " common divisor.") - -parser.add_argument('--start-time', type=int, default=None, - help='Start the analysis at the given GPS time') -parser.add_argument('--end-time', type=int, default=numpy.inf, - help='Stop the analysis at the given GPS time') - -parser.add_argument('--output-path', required=True, - help='Path to a directory to store results in') -parser.add_argument('--output-status', type=str, metavar='PATH', - help='If given, PyCBC Live will periodically write JSON ' - 'status info to PATH, so the analysis can be ' - 'monitored via Nagios') -parser.add_argument('--day-hour-output-prefix', action='store_true') -parser.add_argument('--store-psd', action='store_true') -parser.add_argument('--output-background', type=str, nargs='+', - help='Takes a period in seconds and a file path and dumps ' - 'the coinc backgrounds to that path with that period') -parser.add_argument('--output-background-n-loudest', type=int, default=10000, - help="If given an integer (assumed positive), it stores loudest n triggers" - "(not sorted) for each of the coinc background. If 0, all bkg will be dumped.") - -parser.add_argument('--newsnr-threshold', type=float, default=0) -parser.add_argument('--max-batch-size', type=int, default=2**27) -parser.add_argument('--store-loudest-index', type=int, default=0) -parser.add_argument('--max-psd-abort-distance', type=float, default=numpy.inf, - help="Safety BNS horizon distance (in Mpc) above which a " - "detector's data is discarded.") -parser.add_argument('--min-psd-abort-distance', type=float, default=-numpy.inf, - help="Safety BNS horizon distance (in Mpc) below which a " - "detector's data is discarded.") -parser.add_argument('--max-triggers-in-batch', type=int, metavar="N", - help="Tells each matched-filtering process to only report " - "the loudest N single-detector triggers by SNR.") -parser.add_argument('--max-length', type=float, - help='Maximum duration of templates, used to set the data buffer size') - -parser.add_argument('--enable-profiling', type=int, metavar='RANK', - help='Dump out profiling information from the MPI process' - ' with given rank at the end of program execution') - -parser.add_argument('--enable-background-estimation', default=False, action='store_true') -parser.add_argument('--ifar-double-followup-threshold', type=float, required=True, - help='Inverse-FAR threshold to followup double coincs with' - 'additional detectors') -parser.add_argument('--pvalue-lookback-time', type=float, default=150, - metavar='SECONDS', - help='Lookback time for the calculation of the p-value in ' - 'followup detectors.') -parser.add_argument('--pvalue-combination-livetime', type=float, required=True, - help="Livetime used for p-value combination with followup " - "detectors, in years") -parser.add_argument('--enable-gracedb-upload', action='store_true', default=False, - help='Upload triggers to GraceDB') -parser.add_argument('--enable-production-gracedb-upload', action='store_true', default=False, - help='Do not mark triggers uploaded to GraceDB as test ' - 'events. This option should *only* be enabled in ' - 'production analyses!') -parser.add_argument('--enable-single-detector-upload', action='store_true', default=False, - help='Upload single ifo events to GraceDB') -parser.add_argument('--gracedb-server', metavar='URL', - help='URL of GraceDB server API for uploading events. ' - 'If not provided, the default URL is used.') -parser.add_argument('--gracedb-search', type=str, default='AllSky', - help='String going into the "search" field of the GraceDB ' - 'events') -parser.add_argument('--gracedb-labels', metavar='LABEL', nargs='+', - help='Apply the given list of labels to events uploaded ' - 'to GraceDB.') -parser.add_argument('--ifar-upload-threshold', type=float, required=True, - help='Inverse-FAR threshold for uploading candidate ' - 'triggers to GraceDB, in years.') -parser.add_argument('--coinc-window-pad', type=float, - help="Amount of time allowed to form a coincidence in " - "addition to the time of flight in seconds.", - default=0.002) -parser.add_argument('--file-prefix', default='Live') - -parser.add_argument('--round-start-time', type=int, metavar='X', - help="Round up the start time to the nearest multiple of X" - " seconds. This is useful for forcing agreement " - " with frame file production.") -parser.add_argument('--size-override', type=int, metavar='N', - help="Override the internal MPI size layout. " - " Useful for debugging and running a portion of a bank") -parser.add_argument('--fftw-planning-limit', type=float, - help="Time in seconds to allow for a plan to be created") -parser.add_argument('--run-snr-optimization', action='store_true', - help='Run spawned followup processes to maximize SNR for ' - 'any trigger uploaded to GraceDB') -parser.add_argument('--snr-opt-timeout', type=int, default=400, metavar='SECONDS', - help='Maximum allowed duration of followup process to maximize SNR') -parser.add_argument('--snr-opt-label', default='SNR_OPTIMIZED', - help='Label to apply to snr-optimized GraceDB uploads') -parser.add_argument('--snr-opt-extra-opts', - help='Extra options to pass to the optimizer subprocess. Example: ' - '--snr-opt-extra-opts "--snr-opt-method differential_evolution ' - '--snr-opt-di-maxiter 50 --snr-opt-di-popsize 100 ' - '--snr-opt-seed 42 --snr-opt-include-candidate "') - -parser.add_argument('--enable-embright-has-massgap', action='store_true', default=False, - help='Estimate HasMassGap probability for EMBright info. Lower limit ' - 'of the mass gap is equal to the maximum NS mass used for ' - 'the source classification.') -parser.add_argument('--embright-massgap-max', type=float, default=5.0, metavar='SOLAR MASSES', - help='Upper limit of the mass gap, used for estimating ' - 'HasMassGap probability.') -parser.add_argument('--skymap-only-ifos', nargs='+', - help="Detectors that only contribute in sky localization") -parser.add_argument('--psd-variation', action='store_true', - help="Run the psd variation code to produce psd variation " - "values for each single detector triggers found by " - "the search. Required when using a single detector " - "ranking statistic that includes psd variation.") -parser.add_argument("--statistic-refresh-rate", type=float, - help="How often to refresh the statistic object, " - "in seconds. If omitted, no refreshing is done.") +parser.add_argument("--verbose", action="count") +parser.add_argument("--version", action="version", version=version.git_verbose_msg) +parser.add_argument( + "--bank-file", required=True, help="Template bank file in XML or HDF format" +) +parser.add_argument("--low-frequency-cutoff", help="low frequency cutoff", type=int) +parser.add_argument("--sample-rate", help="output sample rate", type=int) +parser.add_argument("--chisq-bins", help="Number of chisq bins") +parser.add_argument( + "--analysis-chunk", + type=int, + required=True, + help="Amount of data to produce triggers in a block", +) + +parser.add_argument( + "--snr-threshold", type=float, help="SNR threshold for generating a trigger" +) +parser.add_argument("--snr-abort-threshold", type=float) + +parser.add_argument( + "--channel-name", action=MultiDetMultiColonOptionAction, required=True +) +parser.add_argument( + "--state-channel", + action=MultiDetMultiColonOptionAction, + help="Channel containing frame status information. Used " + "to determine when to analyze the hoft data. This somewhat" + " corresponds to CAT1 information", +) +parser.add_argument( + "--analyze-flags", + action=MultiDetOptionAction, + nargs="+", + help='The flags that must be in the "good" state to analyze data', +) +parser.add_argument( + "--idq-channel", + action=MultiDetMultiColonOptionAction, + help="Channel to read iDQ timeseries from. iDQ timeseries are " + "required by certain ranking statistics.", +) +parser.add_argument( + "--idq-state-channel", + action=MultiDetMultiColonOptionAction, + help="Channel containing information about the state of idq." + "Used to determine when idq is usable for statistics. ", +) +parser.add_argument( + "--idq-threshold", + type=float, + help="Threshold used to veto triggers at times of low iDQ False Alarm Probability", +) +parser.add_argument( + "--idq-reweighting", + action="store_true", + default=False, + help="Reweight triggers based on iDQ False Alarm Probability", +) +parser.add_argument( + "--data-quality-channel", + action=MultiDetMultiColonOptionAction, + help="Channel containing data quality information. Used " + "to determine when hoft may be suspect and may be used to " + "veto triggers or not analyze a segment of data. This " + "roughly corresponds to CAT2 information", +) +parser.add_argument( + "--data-quality-flags", + action=MultiDetOptionAction, + nargs="+", + help="Flags used to determine when to throw triggers away. " + "For each detector, give a comma-separated list of flags.", +) +parser.add_argument( + "--data-quality-padding", + type=float, + default=0, + help="Time in seconds around a bad dq time to additionally remove triggers", +) +parser.add_argument("--frame-src", action=MultiDetOptionAction, nargs="+") +parser.add_argument("--frame-type", action=MultiDetOptionAction, nargs="+") +parser.add_argument("--force-update-cache", action="store_true") +parser.add_argument( + "--highpass-frequency", type=float, help="Frequency to apply highpass filtering" +) +parser.add_argument( + "--highpass-reduction", type=float, help="DB to reduce low frequencies" +) +parser.add_argument( + "--highpass-bandwidth", + type=float, + help="Width of the highpass turnover region in Hz", +) +parser.add_argument("--psd-recalculate-difference", type=float, default=0.01) +parser.add_argument("--psd-abort-difference", type=float, default=0.20) +parser.add_argument( + "--psd-samples", + type=int, + required=True, + help="Number of PSD segments to use in the rolling estimate", +) +parser.add_argument( + "--psd-segment-length", + type=int, + required=True, + help="Length in seconds of each PSD segment", +) +parser.add_argument( + "--psd-inverse-length", + type=float, + help="Length in time for the equivalent FIR filter", +) +parser.add_argument( + "--psd-recompute-length", + type=positive_int, + default=1, + help="If given only recompute the PSD after this number of " + "analysis chunks, whose length is set by the option " + "--analysis-chunk.", +) +parser.add_argument( + "--trim-padding", + type=float, + default=0.25, + help="Padding around the overwhitened analysis block", +) +parser.add_argument( + "--enable-bank-start-frequency", + action="store_true", + help="Read the starting frequency of template waveforms from the template bank", +) + +parser.add_argument( + "--autogating-threshold", + type=float, + metavar="SIGMA", + help="If given, find and gate glitches " + "producing a deviation larger than " + "SIGMA in the whitened strain time " + "series.", +) +parser.add_argument( + "--autogating-pad", + type=float, + default=0.5, + metavar="SECONDS", + help="Ignore the given length of whitened " + "strain at the ends of a segment, to " + "avoid filters ringing.", +) +parser.add_argument( + "--autogating-cluster", + type=float, + default=1, + metavar="SECONDS", + help="Length of clustering window for detecting glitches for autogating.", +) +parser.add_argument( + "--autogating-width", + type=float, + default=0.25, + metavar="SECONDS", + help="Half-width of the gating window.", +) +parser.add_argument( + "--autogating-taper", + type=float, + metavar="SECONDS", + default=0.25, + help="Taper the strain before and after " + "each gating window over a duration " + "of SECONDS.", +) +parser.add_argument( + "--autogating-duration", + type=float, + default=16, + metavar="SECONDS", + help="Amount of data in seconds to apply autogating on.", +) +parser.add_argument( + "--autogating-psd-segment-length", + type=float, + default=2, + metavar="SECONDS", + help="Length in seconds of each segment used to estimate the PSD " + "with Welchs method and median averaging.", +) +parser.add_argument( + "--autogating-psd-stride", + type=float, + default=1, + metavar="SECONDS", + help="Length in seconds of the overlap between each segment used " + "to estimate the PSD with Welchs method and median averaging.", +) + +parser.add_argument( + "--sync", + action="store_true", + help="Imposes an MPI synchronization at each transfer of" + " single-detector triggers. Can help with debugging" + " and avoiding memory issues when running offline" + " analyses.", +) +parser.add_argument("--increment-update-cache", action=MultiDetOptionAction, nargs="+") +parser.add_argument("--frame-read-timeout", type=float, default=30) +parser.add_argument( + "--increment", + type=int, + default=8, + metavar="T", + help="When generating the template waveforms, their" + " durations are binned such that T is their greatest" + " common divisor.", +) + +parser.add_argument( + "--start-time", + type=int, + default=None, + help="Start the analysis at the given GPS time", +) +parser.add_argument( + "--end-time", + type=int, + default=numpy.inf, + help="Stop the analysis at the given GPS time", +) + +parser.add_argument( + "--output-path", required=True, help="Path to a directory to store results in" +) +parser.add_argument( + "--output-status", + type=str, + metavar="PATH", + help="If given, PyCBC Live will periodically write JSON " + "status info to PATH, so the analysis can be " + "monitored via Nagios", +) +parser.add_argument("--day-hour-output-prefix", action="store_true") +parser.add_argument("--store-psd", action="store_true") +parser.add_argument( + "--output-background", + type=str, + nargs="+", + help="Takes a period in seconds and a file path and dumps " + "the coinc backgrounds to that path with that period", +) +parser.add_argument( + "--output-background-n-loudest", + type=int, + default=10000, + help="If given an integer (assumed positive), it stores loudest n triggers" + "(not sorted) for each of the coinc background. If 0, all bkg will be dumped.", +) + +parser.add_argument("--newsnr-threshold", type=float, default=0) +parser.add_argument("--max-batch-size", type=int, default=2**27) +parser.add_argument("--store-loudest-index", type=int, default=0) +parser.add_argument( + "--max-psd-abort-distance", + type=float, + default=numpy.inf, + help="Safety BNS horizon distance (in Mpc) above which a " + "detector's data is discarded.", +) +parser.add_argument( + "--min-psd-abort-distance", + type=float, + default=-numpy.inf, + help="Safety BNS horizon distance (in Mpc) below which a " + "detector's data is discarded.", +) +parser.add_argument( + "--max-triggers-in-batch", + type=int, + metavar="N", + help="Tells each matched-filtering process to only report " + "the loudest N single-detector triggers by SNR.", +) +parser.add_argument( + "--max-length", + type=float, + help="Maximum duration of templates, used to set the data buffer size", +) + +parser.add_argument( + "--enable-profiling", + type=int, + metavar="RANK", + help="Dump out profiling information from the MPI process" + " with given rank at the end of program execution", +) + +parser.add_argument( + "--enable-background-estimation", default=False, action="store_true" +) +parser.add_argument( + "--ifar-double-followup-threshold", + type=float, + required=True, + help="Inverse-FAR threshold to followup double coincs withadditional detectors", +) +parser.add_argument( + "--pvalue-lookback-time", + type=float, + default=150, + metavar="SECONDS", + help="Lookback time for the calculation of the p-value in followup detectors.", +) +parser.add_argument( + "--pvalue-combination-livetime", + type=float, + required=True, + help="Livetime used for p-value combination with followup detectors, in years", +) +parser.add_argument( + "--enable-gracedb-upload", + action="store_true", + default=False, + help="Upload triggers to GraceDB", +) +parser.add_argument( + "--enable-production-gracedb-upload", + action="store_true", + default=False, + help="Do not mark triggers uploaded to GraceDB as test " + "events. This option should *only* be enabled in " + "production analyses!", +) +parser.add_argument( + "--enable-single-detector-upload", + action="store_true", + default=False, + help="Upload single ifo events to GraceDB", +) +parser.add_argument( + "--gracedb-server", + metavar="URL", + help="URL of GraceDB server API for uploading events. " + "If not provided, the default URL is used.", +) +parser.add_argument( + "--gracedb-search", + type=str, + default="AllSky", + help='String going into the "search" field of the GraceDB events', +) +parser.add_argument( + "--gracedb-labels", + metavar="LABEL", + nargs="+", + help="Apply the given list of labels to events uploaded to GraceDB.", +) +parser.add_argument( + "--ifar-upload-threshold", + type=float, + required=True, + help="Inverse-FAR threshold for uploading candidate triggers to GraceDB, in years.", +) +parser.add_argument( + "--coinc-window-pad", + type=float, + help="Amount of time allowed to form a coincidence in " + "addition to the time of flight in seconds.", + default=0.002, +) +parser.add_argument("--file-prefix", default="Live") + +parser.add_argument( + "--round-start-time", + type=int, + metavar="X", + help="Round up the start time to the nearest multiple of X" + " seconds. This is useful for forcing agreement " + " with frame file production.", +) +parser.add_argument( + "--size-override", + type=int, + metavar="N", + help="Override the internal MPI size layout. " + " Useful for debugging and running a portion of a bank", +) +parser.add_argument( + "--fftw-planning-limit", + type=float, + help="Time in seconds to allow for a plan to be created", +) +parser.add_argument( + "--run-snr-optimization", + action="store_true", + help="Run spawned followup processes to maximize SNR for " + "any trigger uploaded to GraceDB", +) +parser.add_argument( + "--snr-opt-timeout", + type=int, + default=400, + metavar="SECONDS", + help="Maximum allowed duration of followup process to maximize SNR", +) +parser.add_argument( + "--snr-opt-label", + default="SNR_OPTIMIZED", + help="Label to apply to snr-optimized GraceDB uploads", +) +parser.add_argument( + "--snr-opt-extra-opts", + help="Extra options to pass to the optimizer subprocess. Example: " + '--snr-opt-extra-opts "--snr-opt-method differential_evolution ' + "--snr-opt-di-maxiter 50 --snr-opt-di-popsize 100 " + '--snr-opt-seed 42 --snr-opt-include-candidate "', +) + +parser.add_argument( + "--enable-embright-has-massgap", + action="store_true", + default=False, + help="Estimate HasMassGap probability for EMBright info. Lower limit " + "of the mass gap is equal to the maximum NS mass used for " + "the source classification.", +) +parser.add_argument( + "--embright-massgap-max", + type=float, + default=5.0, + metavar="SOLAR MASSES", + help="Upper limit of the mass gap, used for estimating HasMassGap probability.", +) +parser.add_argument( + "--skymap-only-ifos", + nargs="+", + help="Detectors that only contribute in sky localization", +) +parser.add_argument( + "--psd-variation", + action="store_true", + help="Run the psd variation code to produce psd variation " + "values for each single detector triggers found by " + "the search. Required when using a single detector " + "ranking statistic that includes psd variation.", +) +parser.add_argument( + "--statistic-refresh-rate", + type=float, + help="How often to refresh the statistic object, " + "in seconds. If omitted, no refreshing is done.", +) scheme.insert_processing_option_group(parser) LiveSingle.insert_args(parser) @@ -1063,18 +1314,18 @@ fft.verify_fft_options(args, parser) Coincer.verify_args(args, parser) if args.output_background is not None and len(args.output_background) != 2: - parser.error('--output-background takes two parameters: period and path') + parser.error("--output-background takes two parameters: period and path") if not args.enable_gracedb_upload and args.enable_single_detector_upload: - parser.error('You are not allowed to enable single ifo upload without the ' - '--enable-gracedb-upload option!') + parser.error( + "You are not allowed to enable single ifo upload without the " + "--enable-gracedb-upload option!" + ) # Configure the log messages so that they are prefixed by the timestamp, the # hostname of the originating node and the MPI rank of the originating process pycbc.init_logging( args.verbose, - format='%(asctime)s {} {} %(message)s'.format( - platform.node(), mpi.COMM_WORLD.Get_rank() - ) + format=f"%(asctime)s {platform.node()} {mpi.COMM_WORLD.Get_rank()} %(message)s", ) ctx = scheme.from_cli(args) @@ -1090,23 +1341,25 @@ bank = waveform.LiveFilterBank( total_pad, low_frequency_cutoff=lfc, approximant=args.approximant, - increment=args.increment + increment=args.increment, ) if bank.min_f_lower < args.low_frequency_cutoff: - parser.error('--low-frequency-cutoff ({} Hz) must not be larger than the ' - 'minimum f_lower across all templates ' - '({} Hz)'.format(args.low_frequency_cutoff, bank.min_f_lower)) + parser.error( + f"--low-frequency-cutoff ({args.low_frequency_cutoff} Hz) must not be larger than the " + "minimum f_lower across all templates " + f"({bank.min_f_lower} Hz)" + ) evnt = LiveEventManager(args, bank) -logging.info('Analyzing data from detectors %s', ppdets(evnt.ifos)) -logging.info('Using %s for localization only', ppdets(evnt.skymap_only_ifos)) +logging.info("Analyzing data from detectors %s", ppdets(evnt.ifos)) +logging.info("Using %s for localization only", ppdets(evnt.skymap_only_ifos)) analyze_singles = LiveSingle.verify_args(args, parser, evnt.trigg_ifos) # include MPI rank and functional description into proctitle -task_name = 'root' if evnt.rank == 0 else 'filtering' -setproctitle(f'PyCBC Live rank {evnt.rank:d} [{task_name}]') +task_name = "root" if evnt.rank == 0 else "filtering" +setproctitle(f"PyCBC Live rank {evnt.rank:d} [{task_name}]") sg_chisq = SingleDetSGChisq.from_cli(args, bank, args.chisq_bins) @@ -1117,8 +1370,12 @@ if args.size_override: if evnt.rank == 0 and args.enable_gracedb_upload: evnt.create_gdb() response = evnt.gracedb.ping() - logging.info('GraceDB ping response: %s %s time elapsed: %s', - response.status, response.reason, response.elapsed.total_seconds()) + logging.info( + "GraceDB ping response: %s %s time elapsed: %s", + response.status, + response.reason, + response.elapsed.total_seconds(), + ) # I'm not the root, so do some actual filtering. with ctx: @@ -1129,10 +1386,14 @@ with ctx: # Read specified user-provided wisdom files if args.fftw_input_float_wisdom_file is not None: - fft.fftw.import_single_wisdom_from_filename(args.fftw_input_float_wisdom_file) + fft.fftw.import_single_wisdom_from_filename( + args.fftw_input_float_wisdom_file + ) if args.fftw_input_double_wisdom_file is not None: - fft.fftw.import_double_wisdom_from_filename(args.fftw_input_double_wisdom_file) + fft.fftw.import_double_wisdom_from_filename( + args.fftw_input_double_wisdom_file + ) if args.fftw_planning_limit: fft.fftw.set_planning_limit(args.fftw_planning_limit) @@ -1141,8 +1402,8 @@ with ctx: exit() if evnt.rank > 0: - bank.table.sort(order='mchirp') - waveforms = list(bank[evnt.rank-1::evnt.size-1]) + bank.table.sort(order="mchirp") + waveforms = list(bank[evnt.rank - 1 :: evnt.size - 1]) check_max_length(args, waveforms) mf = LiveBatchMatchedFilter( waveforms, @@ -1152,7 +1413,7 @@ with ctx: snr_abort_threshold=args.snr_abort_threshold, newsnr_threshold=args.newsnr_threshold, max_triggers_in_batch=args.max_triggers_in_batch, - maxelements=args.max_batch_size + maxelements=args.max_batch_size, ) # Synchronize start time if not provided on the command line @@ -1162,9 +1423,10 @@ with ctx: args.start_time = evnt.comm.bcast(tnow, root=0) if args.round_start_time: - args.start_time = int(args.start_time / args.round_start_time + 1) *\ - args.round_start_time - logging.info('Starting from: %s', args.start_time) + args.start_time = ( + int(args.start_time / args.round_start_time + 1) * args.round_start_time + ) + logging.info("Starting from: %s", args.start_time) # Initialize the data readers for all detectors. For rank 0, we need data # from all detectors, including the localization-only ones. For higher @@ -1177,8 +1439,9 @@ with ctx: # create single-detector background "estimators" if analyze_singles and evnt.rank == 0: - sngl_estimator = {ifo: LiveSingle.from_cli(args, ifo) - for ifo in evnt.trigg_ifos} + sngl_estimator = { + ifo: LiveSingle.from_cli(args, ifo) for ifo in evnt.trigg_ifos + } for estim in sngl_estimator.values(): estim.start_refresh_thread() @@ -1188,19 +1451,18 @@ with ctx: ifo_combos = itertools.combinations(evnt.trigg_ifos, 2) estimators = [] for combo in ifo_combos: - logging.info('Will calculate %s background', ppdets(combo, "-")) - estimators.append(Coincer.from_cli( - args, len(bank), args.analysis_chunk, list(combo) - )) + logging.info("Will calculate %s background", ppdets(combo, "-")) + estimators.append( + Coincer.from_cli(args, len(bank), args.analysis_chunk, list(combo)) + ) my_coinc_id = 999999 + def set_coinc_id(i): global my_coinc_id my_coinc_id = i c = estimators[my_coinc_id] - setproctitle( - 'PyCBC Live {} bg estimator'.format(ppdets(c.ifos, '-')) - ) + setproctitle("PyCBC Live {} bg estimator".format(ppdets(c.ifos, "-"))) def estimator_refresh_threads(_): c = estimators[my_coinc_id] @@ -1209,8 +1471,13 @@ with ctx: def get_coinc(results): c = estimators[my_coinc_id] r = c.add_singles(results) - logging.info('Coincs %i: %s-%s: %s in cbuffer', my_coinc_id, - c.ifos[0], c.ifos[1], c.coincs.index) + logging.info( + "Coincs %i: %s-%s: %s in cbuffer", + my_coinc_id, + c.ifos[0], + c.ifos[1], + c.coincs.index, + ) return r def output_background(_): @@ -1222,7 +1489,7 @@ with ctx: coinc_pool.allmap(set_coinc_id, range(len(estimators))) coinc_pool.broadcast(estimator_refresh_threads, None) - logging.info('Starting') + logging.info("Starting") if args.enable_profiling is not None and evnt.rank == args.enable_profiling: pr = cProfile.Profile() @@ -1231,27 +1498,24 @@ with ctx: # main analysis loop data_end = lambda: data_reader[tuple(data_reader.keys())[0]].end_time last_bg_dump_time = int(data_end()) - psd_count = {ifo:0 for ifo in evnt.ifos} + psd_count = dict.fromkeys(evnt.ifos, 0) # Create dicts to track whether the psd has been recalculated and to hold # psd variation filters - psd_recalculated = { - ifo: True for ifo in (evnt.ifos if evnt.rank == 0 else evnt.trigg_ifos) - } - psd_var_filts = {ifo: None for ifo in evnt.trigg_ifos} + psd_recalculated = dict.fromkeys(evnt.ifos if evnt.rank == 0 else evnt.trigg_ifos, True) + psd_var_filts = dict.fromkeys(evnt.trigg_ifos) while data_end() < args.end_time: t1 = pycbc.gps_now() - logging.info('Analyzing from %s', data_end()) + logging.info("Analyzing from %s", data_end()) results = {} evnt.live_detectors = set() - for ifo in (evnt.ifos if evnt.rank == 0 else evnt.trigg_ifos): + for ifo in evnt.ifos if evnt.rank == 0 else evnt.trigg_ifos: results[ifo] = False status = data_reader[ifo].advance( - valid_pad, - timeout=args.frame_read_timeout + valid_pad, timeout=args.frame_read_timeout ) if status and psd_count[ifo] == 0: status = data_reader[ifo].recalculate_psd() @@ -1265,24 +1529,26 @@ with ctx: psd_count[ifo] -= 1 status &= data_reader[ifo].check_psd_dist( - args.min_psd_abort_distance, - args.max_psd_abort_distance + args.min_psd_abort_distance, args.max_psd_abort_distance ) if status is True: if ifo not in evnt.skymap_only_ifos: evnt.live_detectors.add(ifo) if evnt.rank > 0: - logging.info('Filtering %s', ifo) + logging.info("Filtering %s", ifo) results[ifo] = mf.process_data(data_reader[ifo]) else: - logging.info('Insufficient data for %s analysis', ifo) + logging.info("Insufficient data for %s analysis", ifo) if evnt.rank > 0: evnt.commit_results((results, data_end())) else: - psds = {ifo: data_reader[ifo].psd for ifo in data_reader - if data_reader[ifo].psd is not None} + psds = { + ifo: data_reader[ifo].psd + for ifo in data_reader + if data_reader[ifo].psd is not None + } # Collect together the single detector triggers if evnt.size > 1: @@ -1301,45 +1567,47 @@ with ctx: if data_reader[ifo].dq is not None: logging.info("Checking %s's DQ vector", ifo) start = data_reader[ifo].start_time - times = results[ifo]['end_time'] + times = results[ifo]["end_time"] idx = data_reader[ifo].dq.indices_of_flag( - start, valid_pad, times, - padding=data_reader[ifo].dq_padding) - logging.info('Keeping %d/%d %s triggers after DQ flags', - len(idx), len(times), ifo) + start, valid_pad, times, padding=data_reader[ifo].dq_padding + ) + logging.info( + "Keeping %d/%d %s triggers after DQ flags", + len(idx), + len(times), + ifo, + ) for key in results[ifo]: if len(results[ifo][key]): results[ifo][key] = results[ifo][key][idx] if data_reader[ifo].idq is not None: logging.info("Reading %s's iDQ information", ifo) start = data_reader[ifo].start_time - times = results[ifo]['end_time'] + times = results[ifo]["end_time"] flag_active = data_reader[ifo].idq.flag_at_times( - start, valid_pad, times, - padding=data_reader[ifo].dq_padding - ) + start, valid_pad, times, padding=data_reader[ifo].dq_padding + ) if args.idq_reweighting: logging.info( - 'iDQ flagged %d/%d %s triggers', + "iDQ flagged %d/%d %s triggers", numpy.sum(flag_active), len(times), - ifo + ifo, ) - results[ifo]['dq_state'] = flag_active.astype(int) + results[ifo]["dq_state"] = flag_active.astype(int) else: # use idq as a veto keep = numpy.logical_not(flag_active) logging.info( - 'Keeping %d/%d %s triggers after iDQ', + "Keeping %d/%d %s triggers after iDQ", numpy.sum(keep), len(times), - ifo + ifo, ) for key in results[ifo]: if len(results[ifo][key]): - results[ifo][key] = \ - results[ifo][key][keep] + results[ifo][key] = results[ifo][key][keep] # Calculate and add the psd variation for the results if args.psd_variation: @@ -1351,21 +1619,19 @@ with ctx: psd_var_filts[ifo] = variation.live_create_filter( data_reader[ifo].psd, args.psd_segment_length, - int(args.sample_rate) + int(args.sample_rate), ) psd_recalculated[ifo] = False psd_var_ts = variation.live_calc_psd_variation( - data_reader[ifo].strain, - psd_var_filts[ifo], - args.increment + data_reader[ifo].strain, psd_var_filts[ifo], args.increment ) psd_var_vals = variation.live_find_var_value( results[ifo], psd_var_ts ) - results[ifo]['psd_var_val'] = psd_var_vals + results[ifo]["psd_var_val"] = psd_var_vals # Look for coincident triggers and do background estimation if args.enable_background_estimation: @@ -1383,80 +1649,111 @@ with ctx: gates = {ifo: data_reader[ifo].gate_params for ifo in data_reader} # map the results file to an hdf file - prefix = '{}-{}-{}-{}'.format(''.join(sorted(evnt.ifos)), - args.file_prefix, - data_end() - args.analysis_chunk, - valid_pad) + prefix = "{}-{}-{}-{}".format( + "".join(sorted(evnt.ifos)), + args.file_prefix, + data_end() - args.analysis_chunk, + valid_pad, + ) - evnt.dump(results, prefix, time_index=data_end(), - store_psd=(psds if args.store_psd else False), - store_loudest_index=args.store_loudest_index, - raw_results=best_coinc, gates=gates) + evnt.dump( + results, + prefix, + time_index=data_end(), + store_psd=(psds if args.store_psd else False), + store_loudest_index=args.store_loudest_index, + raw_results=best_coinc, + gates=gates, + ) # dump the background if needed - if args.output_background and \ - data_end() - last_bg_dump_time > float(args.output_background[0]): + if args.output_background and data_end() - last_bg_dump_time > float( + args.output_background[0] + ): last_bg_dump_time = int(data_end()) bg_dists = coinc_pool.broadcast(output_background, None) - bg_fn = '{}-LIVE_BACKGROUND-{}.hdf'.format( - ''.join(sorted(evnt.trigg_ifos)), last_bg_dump_time + bg_fn = "{}-LIVE_BACKGROUND-{}.hdf".format( + "".join(sorted(evnt.trigg_ifos)), last_bg_dump_time ) bg_fn = os.path.join(args.output_background[1], bg_fn) - with h5py.File(bg_fn, 'w') as bgf: + with h5py.File(bg_fn, "w") as bgf: for bg_ifos, bg_data, bg_time in bg_dists: - if args.output_background_n_loudest and \ - args.output_background_n_loudest < (len(bg_data) - 1): + if ( + args.output_background_n_loudest + and args.output_background_n_loudest < (len(bg_data) - 1) + ): n_loudest = args.output_background_n_loudest - assert (n_loudest > 0), \ + assert n_loudest > 0, ( "We can only store positive int loudest triggers." + ) ds = bgf.create_dataset( - ','.join(sorted(bg_ifos)), + ",".join(sorted(bg_ifos)), data=-numpy.partition(-bg_data, n_loudest)[:n_loudest], - compression='gzip' + compression="gzip", ) else: ds = bgf.create_dataset( - ','.join(sorted(bg_ifos)), + ",".join(sorted(bg_ifos)), data=bg_data, - compression='gzip' + compression="gzip", ) - ds.attrs['background_time'] = bg_time - bgf.attrs['gps_time'] = last_bg_dump_time + ds.attrs["background_time"] = bg_time + bgf.attrs["gps_time"] = last_bg_dump_time - logging.info('Finished analyzing up to %s', data_end()) + logging.info("Finished analyzing up to %s", data_end()) if args.sync: evnt.barrier() tdiff = pycbc.gps_now() - t1 lag = float(pycbc.gps_now() - data_end()) - logging.info('Took %1.2f, duty factor of %.2f, lag %.2f s, %d live detectors', - tdiff, tdiff / valid_pad, lag, len(evnt.live_detectors) + logging.info( + "Took %1.2f, duty factor of %.2f, lag %.2f s, %d live detectors", + tdiff, + tdiff / valid_pad, + lag, + len(evnt.live_detectors), ) if args.output_status is not None and evnt.rank == 0: if lag > 120: - status_intervals = [{'num_status': 2, - 'txt_status': 'CRITICAL: lag greater than 2 min', - 'start_sec': 0}] + status_intervals = [ + { + "num_status": 2, + "txt_status": "CRITICAL: lag greater than 2 min", + "start_sec": 0, + } + ] else: - status_intervals = [{'num_status': 0, - 'txt_status': 'OK: No reported problems', - 'start_sec': 0}, - {'num_status': 1, - 'txt_status': 'WARNING: last report between 2 and 4 min ago', - 'start_sec': 120}] - status_intervals.append({'num_status': 2, - 'txt_status': 'CRITICAL: last report more than 4 min ago', - 'start_sec': 240}) - status = {'author': 'Tito Dal Canton', - 'email': 'tito.canton@ligo.org', - 'created_gps': int(pycbc.gps_now()), - 'status_intervals': status_intervals} + status_intervals = [ + { + "num_status": 0, + "txt_status": "OK: No reported problems", + "start_sec": 0, + }, + { + "num_status": 1, + "txt_status": "WARNING: last report between 2 and 4 min ago", + "start_sec": 120, + }, + ] + status_intervals.append( + { + "num_status": 2, + "txt_status": "CRITICAL: last report more than 4 min ago", + "start_sec": 240, + } + ) + status = { + "author": "Tito Dal Canton", + "email": "tito.canton@ligo.org", + "created_gps": int(pycbc.gps_now()), + "status_intervals": status_intervals, + } try: - with open(args.output_status, 'w') as status_fp: + with open(args.output_status, "w") as status_fp: json.dump(status, status_fp) - except IOError: - logging.error('I/O error writing status JSON file!') + except OSError: + logging.exception("I/O error writing status JSON file!") if evnt.rank == 1: if args.fftw_output_float_wisdom_file: @@ -1466,6 +1763,6 @@ if evnt.rank == 1: fft.fftw.export_double_wisdom_to_filename(args.fftw_output_double_wisdom_file) if args.enable_profiling is not None and evnt.rank == args.enable_profiling: - pr.dump_stats(f'profiling_rank_{evnt.rank:03d}') + pr.dump_stats(f"profiling_rank_{evnt.rank:03d}") logging.info("Exiting as the end time has been reached") diff --git a/bin/pycbc_live_nagios_monitor b/bin/pycbc_live_nagios_monitor index c88c5f255ff..562be13201f 100644 --- a/bin/pycbc_live_nagios_monitor +++ b/bin/pycbc_live_nagios_monitor @@ -5,62 +5,63 @@ # Actually check log file for errors and node specific problems # Monitor data transfer to sites as well -import json -import lal import argparse -import time +import json import os.path +import time + +import lal parser = argparse.ArgumentParser( description="This scripts monitors the log file of the " "PyCBC Live process. This is used to generate a json file that can be " - "picked up by nagios to determine if the PyCBC Live process has died.") -parser.add_argument('--log-file', - help="The pycbc live log file") -parser.add_argument('--output-file', - help="The JSON nagios status file") -parser.add_argument('--check-interval', type=int, - help="Time in seconds to wait before rechecking status") + "picked up by nagios to determine if the PyCBC Live process has died." +) +parser.add_argument("--log-file", help="The pycbc live log file") +parser.add_argument("--output-file", help="The JSON nagios status file") +parser.add_argument( + "--check-interval", + type=int, + help="Time in seconds to wait before rechecking status", +) args = parser.parse_args() while 1: everything_ok = True status = {} - status['author'] = "Alexander Harvey Nitz" - status['email'] = "alex.nitz@ligo.org" - status['created_gps'] = int(lal.GPSTimeNow()) + status["author"] = "Alexander Harvey Nitz" + status["email"] = "alex.nitz@ligo.org" + status["created_gps"] = int(lal.GPSTimeNow()) try: tdiff = time.time() - os.path.getmtime(args.log_file) # Check that the pycbc live logfile has been updated recently. - if tdiff >= 60: everything_ok = False - except: - everything_ok = False + if tdiff >= 60: + everything_ok = False + except: + everything_ok = False if everything_ok: - status['status_intervals'] = \ - [ - { - "num_status": 0, - "txt_status": "OK: No reported problems", - "start_sec": 0 - }, - { - "num_status": 1, - "txt_status": "WARNING: The process is slow to report.", - "start_sec": 120 - }, - { - "num_status": 3, - "txt_status": "UNKNOWN: It has been 4 minutes. Has it died?", - "start_sec": 240 - } - ] + status["status_intervals"] = [ + {"num_status": 0, "txt_status": "OK: No reported problems", "start_sec": 0}, + { + "num_status": 1, + "txt_status": "WARNING: The process is slow to report.", + "start_sec": 120, + }, + { + "num_status": 3, + "txt_status": "UNKNOWN: It has been 4 minutes. Has it died?", + "start_sec": 240, + }, + ] else: - status['status_intervals'] = [{"num_status": 2, - "txt_status": "PyCBC Live appears to be down!", - }] - open(args.output_file, 'w').write(json.dumps(status)) + status["status_intervals"] = [ + { + "num_status": 2, + "txt_status": "PyCBC Live appears to be down!", + } + ] + open(args.output_file, "w").write(json.dumps(status)) time.sleep(args.check_interval) - diff --git a/bin/pycbc_make_banksim b/bin/pycbc_make_banksim index 3e4eb98e445..bef673fe50a 100644 --- a/bin/pycbc_make_banksim +++ b/bin/pycbc_make_banksim @@ -1,65 +1,79 @@ #! /usr/bin/env python +import configparser as ConfigParser +import glob import logging import os import shutil -import configparser as ConfigParser import subprocess -import glob import tempfile from argparse import ArgumentParser -from glue.pipeline import CondorDAGJob, CondorDAGNode, CondorDAG, CondorJob -from pycbc import init_logging, add_common_pycbc_options +from glue.pipeline import CondorDAG, CondorDAGJob, CondorDAGNode, CondorJob + +from pycbc import add_common_pycbc_options, init_logging + class BaseJob(CondorDAGJob, CondorJob): - def __init__(self, log_dir, executable, cp, section, gpu=False, - accounting_group=None): + def __init__( + self, log_dir, executable, cp, section, gpu=False, accounting_group=None + ): CondorDAGJob.__init__(self, "vanilla", executable) if gpu: CondorJob.__init__(self, "vanilla", executable, 2) # These are all python jobs so need to pull in the env - self.add_condor_cmd('getenv', 'True') + self.add_condor_cmd("getenv", "True") log_base = os.path.join( - log_dir, os.path.basename(executable) + '-$(cluster)-$(process)') - self.set_stderr_file(log_base + '.err') - self.set_stdout_file(log_base + '.out') - self.set_sub_file(os.path.basename(executable) + '.sub') + log_dir, os.path.basename(executable) + "-$(cluster)-$(process)" + ) + self.set_stderr_file(log_base + ".err") + self.set_stdout_file(log_base + ".out") + self.set_sub_file(os.path.basename(executable) + ".sub") if cp is not None: self.add_ini_opts(cp, section) if accounting_group: - self.add_condor_cmd('accounting_group', accounting_group) + self.add_condor_cmd("accounting_group", accounting_group) + + self.add_condor_cmd("request_disk", 1024) - self.add_condor_cmd('request_disk', 1024) class BanksimNode(CondorDAGNode): - def __init__(self, job, inj_file, tmplt_file, match_file, gpu=True, - gpu_postscript=False, inj_per_job=None): + def __init__( + self, + job, + inj_file, + tmplt_file, + match_file, + gpu=True, + gpu_postscript=False, + inj_per_job=None, + ): CondorDAGNode.__init__(self, job) self.add_file_opt("signal-file", inj_file) self.add_file_opt("template-file", tmplt_file) if gpu: - self.add_var_opt("processing-scheme", 'cuda') + self.add_var_opt("processing-scheme", "cuda") if gpu and gpu_postscript: self.set_retry(5) - mf = match_file+".$(Process)" - mf1 = match_file+".0" - mf2 = match_file+".1" - self.add_file_opt("match-file", match_file+".$(Process)", - file_is_output_file=True) + mf = match_file + ".$(Process)" + mf1 = match_file + ".0" + mf2 = match_file + ".1" + self.add_file_opt( + "match-file", match_file + ".$(Process)", file_is_output_file=True + ) self.job().__queue = 2 # Needed to satisfy the requirements for both running on atlas and spice - job.add_condor_cmd('+WantsGPU', 'true') - job.add_condor_cmd('+WantGPU', 'true') + job.add_condor_cmd("+WantsGPU", "true") + job.add_condor_cmd("+WantGPU", "true") job.add_condor_cmd( - 'Requirements', - '(GPU_PRESENT =?= true) || (HasGPU =?= "gtx580")') + "Requirements", '(GPU_PRESENT =?= true) || (HasGPU =?= "gtx580")' + ) self.set_post_script(gpu_postscript) self.add_post_script_arg(mf1) @@ -70,6 +84,7 @@ class BanksimNode(CondorDAGNode): else: self.add_file_opt("match-file", match_file, file_is_output_file=True) + class CombineNode(CondorDAGNode): def __init__(self, job, inj_num): CondorDAGNode.__init__(self, job) @@ -80,6 +95,7 @@ class CombineNode(CondorDAGNode): self.add_file_opt("output-file", outf) + def get_ini_opts(confs, section): op_str = "" for opt in confs.options(section): @@ -87,17 +103,20 @@ def get_ini_opts(confs, section): op_str += "--" + opt + " " + val + " \\" + "\n" return op_str + def mkdir(dir_name): - try : + try: os.mkdir(dir_name) except OSError: pass + def mc_min_max_from_sorted_file(fname): - from igwn_ligolw.utils import load_filename from igwn_ligolw.ligolw import Table - from pycbc.pnutils import mass1_mass2_to_mchirp_eta + from igwn_ligolw.utils import load_filename + from pycbc.io.ligolw import LIGOLWContentHandler + from pycbc.pnutils import mass1_mass2_to_mchirp_eta doc = load_filename(fname, False, contenthandler=LIGOLWContentHandler) try: @@ -111,6 +130,8 @@ def mc_min_max_from_sorted_file(fname): bf_mchirps = {} sf_mchirps = {} + + def check_outside_mchirp(bf, sf, w): if bf not in bf_mchirps: bf_mchirps[bf] = mc_min_max_from_sorted_file(bf) @@ -118,16 +139,16 @@ def check_outside_mchirp(bf, sf, w): sf_mchirps[sf] = mc_min_max_from_sorted_file(sf) mc_min, mc_max = bf_mchirps[bf] - mc2_min, mc2_max = sf_mchirps[sf] + mc2_min, mc2_max = sf_mchirps[sf] - if (mc_min <= mc2_max * (1+w) ) and (mc_max * (1+w) >= mc2_min): + if (mc_min <= mc2_max * (1 + w)) and (mc_max * (1 + w) >= mc2_min): return False - else: - return True + return True + parser = ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--config', type=str, required=True) +parser.add_argument("--config", type=str, required=True) options = parser.parse_args() # Default logging level is info: --verbose adds to this @@ -141,18 +162,20 @@ bank_file = confs.get("workflow", "bank-file") injections_per_job = confs.get("workflow", "injections-per-job") templates_per_job = confs.get("workflow", "templates-per-job") -log_path = confs.get("workflow", 'log-path') +log_path = confs.get("workflow", "log-path") tempfile.tempdir = log_path -tempfile.template='banksim.dag.log.' +tempfile.template = "banksim.dag.log." logfile = tempfile.mktemp() mchirp_window = None if confs.has_option("banksim", "mchirp-window"): - if ',' in confs.get("banksim", "mchirp-window"): - mchirp_window = max([float(x) for x in confs.get("banksim", "mchirp-window").split(",")]) + if "," in confs.get("banksim", "mchirp-window"): + mchirp_window = max( + [float(x) for x in confs.get("banksim", "mchirp-window").split(",")] + ) else: - mchirp_window = float(confs.get("banksim", "mchirp-window")) + mchirp_window = float(confs.get("banksim", "mchirp-window")) gpu = False try: @@ -163,26 +186,26 @@ except: pass try: - accounting_group = confs.get('workflow', 'accounting-group') + accounting_group = confs.get("workflow", "accounting-group") except: accounting_group = None logging.warning( - 'Warning: accounting-group not specified, LDG clusters may ' - 'reject this workflow!' + "Warning: accounting-group not specified, LDG clusters may " + "reject this workflow!" ) logging.info("Making workspace directories") -mkdir('scripts') -mkdir('bank') -mkdir('match') -mkdir('injection') -mkdir('match-part') -mkdir('log') -mkdir('plots') +mkdir("scripts") +mkdir("bank") +mkdir("match") +mkdir("injection") +mkdir("match-part") +mkdir("log") +mkdir("plots") logging.info("Copying scripts") -shutil.copy(banksim_prog, 'scripts/pycbc_banksim') -os.chmod('scripts/pycbc_banksim', 0o0777) +shutil.copy(banksim_prog, "scripts/pycbc_banksim") +os.chmod("scripts/pycbc_banksim", 0o0777) logging.info("Creating injection file") if confs.has_section("inspinj"): @@ -190,28 +213,46 @@ if confs.has_section("inspinj"): inj_str = "lalapps_inspinj " + get_ini_opts(confs, "inspinj") + "--output inj.xml" os.system(inj_str) elif confs.has_section("external_injection"): - logging.info("Using external injection file. Please ensure the file is in sim_inspiral table (.xml) format.") + logging.info( + "Using external injection file. Please ensure the file is in sim_inspiral table (.xml) format." + ) inj_file_path = confs.get("external_injection", "inj-file") if inj_file_path == "inj.xml": pass else: - os.system("cp {} inj.xml".format(inj_file_path)) + os.system(f"cp {inj_file_path} inj.xml") else: - raise ValueError("Need to specify the injection method. Either provide [inspinj] section or [external_injection]") + raise ValueError( + "Need to specify the injection method. Either provide [inspinj] section or [external_injection]" + ) logging.info("Splitting template bank") -subprocess.call(['pycbc_splitbank', - '--templates-per-bank', str(templates_per_job), - '-t', bank_file, - '-o', 'bank/bank', - '--sort-mchirp']) +subprocess.call( + [ + "pycbc_splitbank", + "--templates-per-bank", + str(templates_per_job), + "-t", + bank_file, + "-o", + "bank/bank", + "--sort-mchirp", + ] +) logging.info("Splitting injection file") -subprocess.call(['pycbc_splitbank', - '--templates-per-bank', str(injections_per_job), - '-t', "inj.xml", - '-o', 'injection/injection', - '--sort-mchirp']) +subprocess.call( + [ + "pycbc_splitbank", + "--templates-per-bank", + str(injections_per_job), + "-t", + "inj.xml", + "-o", + "injection/injection", + "--sort-mchirp", + ] +) num_banks = len(glob.glob("bank/bank*")) num_injs = len(glob.glob("injection/injection*")) @@ -225,14 +266,31 @@ skip_count = 0 dag = CondorDAG(logfile) dag.set_dag_file("banksim") -bsjob = BaseJob("log", "scripts/pycbc_banksim", confs, "banksim", gpu=gpu, - accounting_group=accounting_group) -cjob = BaseJob("log", "scripts/pycbc_banksim_match_combine", None, None, - accounting_group=accounting_group) -rjob = BaseJob("log", "scripts/pycbc_banksim_collect_results", None, None, - accounting_group=accounting_group) -pjob = BaseJob("log", "scripts/pycbc_banksim_plots", None, None, - accounting_group=accounting_group) +bsjob = BaseJob( + "log", + "scripts/pycbc_banksim", + confs, + "banksim", + gpu=gpu, + accounting_group=accounting_group, +) +cjob = BaseJob( + "log", + "scripts/pycbc_banksim_match_combine", + None, + None, + accounting_group=accounting_group, +) +rjob = BaseJob( + "log", + "scripts/pycbc_banksim_collect_results", + None, + None, + accounting_group=accounting_group, +) +pjob = BaseJob( + "log", "scripts/pycbc_banksim_plots", None, None, accounting_group=accounting_group +) rnode = CondorDAGNode(rjob) pnode = CondorDAGNode(pjob) @@ -243,19 +301,25 @@ for inj_num in range(num_injs): for bank_num in range(num_banks): if mchirp_window is not None: bank_part = "bank/bank" + str(bank_num) + ".xml.gz" - sim_part = "injection/injection" + str(inj_num) + ".xml.gz" + sim_part = "injection/injection" + str(inj_num) + ".xml.gz" if check_outside_mchirp(bank_part, sim_part, mchirp_window): skip_count += 1 continue else: do_count += 1 part_num = str(bank_num) - mfn = 'match-part/match' + num +'part' + part_num + '.dat' - sn = 'injection/injection' + num + '.xml.gz' - bn = 'bank/bank' + part_num + '.xml.gz' - bsnode = BanksimNode(bsjob, sn, bn, mfn, gpu=gpu, - gpu_postscript="scripts/diff_match.sh", - inj_per_job=injections_per_job) + mfn = "match-part/match" + num + "part" + part_num + ".dat" + sn = "injection/injection" + num + ".xml.gz" + bn = "bank/bank" + part_num + ".xml.gz" + bsnode = BanksimNode( + bsjob, + sn, + bn, + mfn, + gpu=gpu, + gpu_postscript="scripts/diff_match.sh", + inj_per_job=injections_per_job, + ) cnode.add_parent(bsnode) dag.add_node(bsnode) combine_has_jobs = True @@ -266,7 +330,7 @@ dag.add_node(rnode) pnode.add_parent(rnode) dag.add_node(pnode) -logging.info("DO : %d SKIP %d" %(do_count, skip_count)) +logging.info("DO : %d SKIP %d" % (do_count, skip_count)) f.close() f = open("scripts/pycbc_banksim_match_combine", "w") @@ -297,7 +361,7 @@ for i, j in enumerate(indices): maxmatch=array(maxmatch, dtype =dtypef) savetxt(options.output_file, maxmatch,fmt=('%5.5f', '%s', '%i', '%s', '%i', '%5.5f'), delimiter=' ') """) -os.chmod('scripts/pycbc_banksim_match_combine', 0o0777) +os.chmod("scripts/pycbc_banksim_match_combine", 0o0777) f = open("scripts/pycbc_banksim_collect_results", "w") f.write("""#!/usr/bin/env python @@ -371,7 +435,7 @@ for row in res: f.write(outstr) """) -os.chmod('scripts/pycbc_banksim_collect_results', 0o0777) +os.chmod("scripts/pycbc_banksim_collect_results", 0o0777) if gpu: f = open("cconfig", "w") @@ -429,10 +493,10 @@ if gpu: exit 0 """) - os.chmod('scripts/diff_match.sh', 0o0777) + os.chmod("scripts/diff_match.sh", 0o0777) logging.info("Creating submit script") -f = open("submit.sh","w") +f = open("submit.sh", "w") if gpu: f.write("""#!/bin/bash condor_submit_dag -config cconfig banksim.dag @@ -441,13 +505,13 @@ else: f.write("""#!/bin/bash condor_submit_dag banksim.dag """) -os.chmod('submit.sh', 0o0777) +os.chmod("submit.sh", 0o0777) f = open("partial_results.sh", "w") f.write("""#!/bin/bash scripts/pycbc_banksim_collect_results """) -os.chmod('partial_results.sh', 0o0777) +os.chmod("partial_results.sh", 0o0777) dag.write_sub_files() dag.write_script() diff --git a/bin/pycbc_make_faithsim b/bin/pycbc_make_faithsim index 58e17d17b9f..bd9914887bc 100644 --- a/bin/pycbc_make_faithsim +++ b/bin/pycbc_make_faithsim @@ -1,42 +1,47 @@ #! /usr/bin/env python -import os +import configparser as ConfigParser +import glob import logging +import os import shutil -import configparser as ConfigParser import subprocess -import glob import tempfile from argparse import ArgumentParser -from glue.pipeline import CondorDAGJob, CondorDAGNode, CondorDAG, CondorJob -from pycbc import init_logging, add_common_pycbc_options +from glue.pipeline import CondorDAG, CondorDAGJob, CondorDAGNode, CondorJob + +from pycbc import add_common_pycbc_options, init_logging + class BaseJob(CondorDAGJob, CondorJob): def __init__(self, log_dir, executable, cp, section, accounting_group=None): CondorDAGJob.__init__(self, "vanilla", executable) # These are all python jobs so need to pull in the env - self.add_condor_cmd('getenv', 'True') - self.add_condor_cmd('+Allow_OGrid', 'True') + self.add_condor_cmd("getenv", "True") + self.add_condor_cmd("+Allow_OGrid", "True") log_base = os.path.join( - log_dir, os.path.basename(executable) + '-$(cluster)-$(process)') - self.set_stderr_file(log_base + '.err') - self.set_stdout_file(log_base + '.out') - self.set_sub_file(section + ".sub") + log_dir, os.path.basename(executable) + "-$(cluster)-$(process)" + ) + self.set_stderr_file(log_base + ".err") + self.set_stdout_file(log_base + ".out") + self.set_sub_file(section + ".sub") - if cp is not None: + if cp is not None: self.add_ini_opts(cp, section) if accounting_group: - self.add_condor_cmd('accounting_group', accounting_group) + self.add_condor_cmd("accounting_group", accounting_group) + + self.add_condor_cmd("request_disk", 1024) - self.add_condor_cmd('request_disk', 1024) class FaithsimNode(CondorDAGNode): def __init__(self, job, tmplt_file, match_file, inj_per_job=None): - CondorDAGNode.__init__(self, job) + CondorDAGNode.__init__(self, job) self.add_file_opt("param-file", tmplt_file) - self.add_file_opt("match-file", match_file, file_is_output_file=True) + self.add_file_opt("match-file", match_file, file_is_output_file=True) + def get_ini_opts(confs, section): op_str = "" @@ -44,13 +49,15 @@ def get_ini_opts(confs, section): val = confs.get(section, opt) op_str += "--" + opt + " " + val + " \\" + "\n" return op_str - + + def mkdir(dir_name): - try : + try: os.mkdir(dir_name) except OSError: pass + def matches_in_list(slist, match): matches = [] for st in slist: @@ -58,10 +65,11 @@ def matches_in_list(slist, match): matches.append(st) return matches + parser = ArgumentParser() add_common_pycbc_options(parser) -parser.add_argument('--config', type=str, required=True) -options = parser.parse_args() +parser.add_argument("--config", type=str, required=True) +options = parser.parse_args() # Default logging level is info: --verbose adds to this init_logging(options.verbose, default_level=1) @@ -73,43 +81,50 @@ banksim_prog = confs.get("executables", "faithsim") templates_per_job = confs.get("workflow", "templates-per-job") try: - log_path = confs.get("workflow", 'log-path') + log_path = confs.get("workflow", "log-path") except: - log_path = './' + log_path = "./" tempfile.tempdir = log_path -tempfile.template='faithsim.dag.log.' +tempfile.template = "faithsim.dag.log." logfile = tempfile.mktemp() try: - accounting_group = confs.get('workflow', 'accounting-group') + accounting_group = confs.get("workflow", "accounting-group") except: accounting_group = None logging.warning( - 'Warning: accounting-group not specified, LDG clusters may ' - 'reject this workflow!' + "Warning: accounting-group not specified, LDG clusters may " + "reject this workflow!" ) logging.info("Making workspace directories") -mkdir('scripts') -mkdir('match') -mkdir('bank') -mkdir('log') -mkdir('plots') +mkdir("scripts") +mkdir("match") +mkdir("bank") +mkdir("log") +mkdir("plots") logging.info("Copying scripts") -shutil.copy(banksim_prog, 'scripts/pycbc_faithsim') -os.chmod('scripts/pycbc_faithsim', 0o0777) +shutil.copy(banksim_prog, "scripts/pycbc_faithsim") +os.chmod("scripts/pycbc_faithsim", 0o0777) logging.info("Creating injection file") inj_str = "lalapps_inspinj " + get_ini_opts(confs, "inspinj") + "--output inj.xml" os.system(inj_str) logging.info("Splitting template bank") -subprocess.call(['pycbc_splitbank', - '--templates-per-bank', str(templates_per_job), - '-t', 'inj.xml', - '-o', 'bank/bank']) +subprocess.call( + [ + "pycbc_splitbank", + "--templates-per-bank", + str(templates_per_job), + "-t", + "inj.xml", + "-o", + "bank/bank", + ] +) num_banks = len(glob.glob("bank/bank*")) @@ -118,28 +133,45 @@ logging.info("Creating DAG") dag = CondorDAG(logfile) dag.set_dag_file("faithsim") -fs_secs = matches_in_list(confs.sections(), 'faithsim') +fs_secs = matches_in_list(confs.sections(), "faithsim") fsjobs = [] for sec in fs_secs: - fsjobs.append(BaseJob("log", "scripts/pycbc_faithsim", confs, sec, - accounting_group=accounting_group)) + fsjobs.append( + BaseJob( + "log", + "scripts/pycbc_faithsim", + confs, + sec, + accounting_group=accounting_group, + ) + ) -rjob = BaseJob("log", "scripts/pycbc_faithsim_collect_results", None, - 'collect_results', accounting_group=accounting_group) +rjob = BaseJob( + "log", + "scripts/pycbc_faithsim_collect_results", + None, + "collect_results", + accounting_group=accounting_group, +) rnode = CondorDAGNode(rjob) -pjob = BaseJob("log", "scripts/pycbc_faithsim_plots", None, 'faithsim_plots', - accounting_group=accounting_group) +pjob = BaseJob( + "log", + "scripts/pycbc_faithsim_plots", + None, + "faithsim_plots", + accounting_group=accounting_group, +) pnode = CondorDAGNode(pjob) for inj_num in range(num_banks): - bn = 'bank/bank' + str(inj_num) + '.xml.gz' - + bn = "bank/bank" + str(inj_num) + ".xml.gz" + for fsjob, sec in zip(fsjobs, fs_secs): - sec_sub = str(sec[len('faithsim'):]) - mf = 'match/match' + sec_sub + '-' + str(inj_num) + '.dat' + sec_sub = str(sec[len("faithsim") :]) + mf = "match/match" + sec_sub + "-" + str(inj_num) + ".dat" fsnode = FaithsimNode(fsjob, bn, mf, inj_per_job=templates_per_job) dag.add_node(fsnode) - rnode.add_parent(fsnode) + rnode.add_parent(fsnode) dag.add_node(rnode) pnode.add_parent(rnode) dag.add_node(pnode) @@ -198,20 +230,20 @@ if __name__ == "__main__": data = np.append(data, pdata) savetxt('result-' + tag + '.dat', data) """) -os.chmod('scripts/pycbc_faithsim_collect_results', 0o0777) - +os.chmod("scripts/pycbc_faithsim_collect_results", 0o0777) + logging.info("Creating submit script") -f = open("submit.sh", 'w') +f = open("submit.sh", "w") f.write("""#!/bin/bash condor_submit_dag faithsim.dag """) -os.chmod('submit.sh', 0o0777) +os.chmod("submit.sh", 0o0777) dag.write_sub_files() dag.write_script() dag.write_concrete_dag() -f = open('scripts/pycbc_faithsim_plots', "w") +f = open("scripts/pycbc_faithsim_plots", "w") f.write("""#!/usr/bin/env python import matplotlib matplotlib.use('Agg') @@ -397,6 +429,6 @@ for fil in fils: name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3 basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3, vmin=None, vmax=None) """) -os.chmod('scripts/pycbc_faithsim_plots', 0o0777) +os.chmod("scripts/pycbc_faithsim_plots", 0o0777) -logging.info('Done') +logging.info("Done") diff --git a/bin/pycbc_make_html_page b/bin/pycbc_make_html_page index 424a272b8df..a2a4c900638 100644 --- a/bin/pycbc_make_html_page +++ b/bin/pycbc_make_html_page @@ -17,27 +17,31 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import argparse +import codecs import os -import stat import shutil +import stat import zipfile -import codecs + from jinja2 import Environment, FileSystemLoader import pycbc.results -from pycbc.results.render import get_embedded_config, render_workflow_html_template, setup_template_render +from pycbc.results.render import ( + get_embedded_config, + render_workflow_html_template, + setup_template_render, +) + def examine_dir(cwd): """ Looks in a directory and returns all subdirs and files. If there is a zipped file, it will unzip it, and return the extracted files as well. """ - # list everything in this directory and loop over them names = os.listdir(cwd) dirs, nondirs = [], [] for name in names: - # if it is a directory append the dir list if os.path.isdir(os.path.join(cwd, name)): dirs.append(name) @@ -49,13 +53,14 @@ def examine_dir(cwd): # loop over extracted files for extracted_name in zip.namelist(): - # check if extracted a dir if os.path.isdir(os.path.join(cwd, extracted_name)): dirs.append(os.path.dirname(extracted_name)) # check if extracted a file into this current working dir - elif os.path.exists(os.path.join(cwd, os.path.basename(extracted_name))): + elif os.path.exists( + os.path.join(cwd, os.path.basename(extracted_name)) + ): nondirs.append(os.path.basename(extracted_name)) else: @@ -63,14 +68,16 @@ def examine_dir(cwd): return cwd, dirs, nondirs + def render_sitemap_page(output_path, dirs): """ Creates sitemap page. """ dirs.sort(key=lambda x: x.path) - render_workflow_html_template(output_path, 'sitemap.html', dirs) + render_workflow_html_template(output_path, "sitemap.html", dirs) -class Directory(): + +class Directory: """ Class used to relate all sub-directories and files in a directory. """ @@ -84,26 +91,29 @@ class Directory(): cwd, subdirs, filenames = examine_dir(path) # save subdirs - self.path = path.replace(plots_dir, '') - self.subdirs = [Directory(path+'/'+subdir, plots_dir) for subdir in subdirs] + self.path = path.replace(plots_dir, "") + self.subdirs = [Directory(path + "/" + subdir, plots_dir) for subdir in subdirs] # loop over all filenames self.files = [] for filename in filenames: - # check if this is a configuration file for a file - extension = filename.split('.')[-1] - config_filename = filename.replace(extension, 'file.ini') - if filename.endswith('file.ini'): + extension = filename.split(".")[-1] + config_filename = filename.replace(extension, "file.ini") + if filename.endswith("file.ini"): continue # check if configuration file exists for file - elif not os.path.exists(config_filename): + if not os.path.exists(config_filename): self.config_filename = None # append file to directory - self.files.append(File(plots_dir+'/'+self.path+'/'+filename, - plots_dir+'/'+self.path+'/'+config_filename)) + self.files.append( + File( + plots_dir + "/" + self.path + "/" + filename, + plots_dir + "/" + self.path + "/" + config_filename, + ) + ) # append to class list self.instances.append(self) @@ -112,25 +122,23 @@ class Directory(): """ Returns the name of the directory. """ - - return self.path.split('/')[-1] + return self.path.split("/")[-1] def title(self): """ Returns a string of the directory name with underscores as spaces and capslock. """ - - return self.path.split('/')[-1].replace('_', ' ').title() + return self.path.split("/")[-1].replace("_", " ").title() def level(self): """ Counts how far into the output filesystem this directory is. """ + return self.path.count("/") - return self.path.count('/') -class File(): +class File: """ Class used to keep track of files. """ @@ -138,48 +146,74 @@ class File(): def __init__(self, path, config_path): # save paths - self.path = path + self.path = path self.config_path = config_path def filename(self): """ Returns filename of File as a string. """ - - return self.path.split('/')[-1] + return self.path.split("/")[-1] def render(self): """ Renders a template for the File using the configuration file if present. """ - return setup_template_render(self.path, self.config_path) -default_logo_location = "https://raw.githubusercontent.com/gwastro/" + \ - "pycbc-logo/master/pycbc_logo_name.png" + +default_logo_location = ( + "https://raw.githubusercontent.com/gwastro/" + "pycbc-logo/master/pycbc_logo_name.png" +) # parse command line -parser = argparse.ArgumentParser(usage='pycbc_make_html_page \ -[--options]', - description="Create static html pages of a filesystem's content.") +parser = argparse.ArgumentParser( + usage="pycbc_make_html_page \ +[--options]", + description="Create static html pages of a filesystem's content.", +) pycbc.add_common_pycbc_options(parser) -parser.add_argument('-f', '--template-file', type=str, - help='Template file to use for skeleton html page.') -parser.add_argument('-b', '--output-path', type=str, - help='Path on web server for main html page.') -parser.add_argument('-p', '--plots-dir', type=str, - help='Path to the directory that contains plots.') -parser.add_argument('-t', '--analysis-title', type=str, - help='Title to include at the top of each page.', - default="") -parser.add_argument('-s', '--analysis-subtitle', type=str, - help='Subtitle to include at the top of each page.', - default="") -parser.add_argument('-l', '--logo', type=str, - default=default_logo_location, - help='File location of logo to include at top of page') -parser.add_argument('-P', '--protect-results', action='store_true', - help='Make the output web pages read only.', default=False) +parser.add_argument( + "-f", + "--template-file", + type=str, + help="Template file to use for skeleton html page.", +) +parser.add_argument( + "-b", "--output-path", type=str, help="Path on web server for main html page." +) +parser.add_argument( + "-p", "--plots-dir", type=str, help="Path to the directory that contains plots." +) +parser.add_argument( + "-t", + "--analysis-title", + type=str, + help="Title to include at the top of each page.", + default="", +) +parser.add_argument( + "-s", + "--analysis-subtitle", + type=str, + help="Subtitle to include at the top of each page.", + default="", +) +parser.add_argument( + "-l", + "--logo", + type=str, + default=default_logo_location, + help="File location of logo to include at top of page", +) +parser.add_argument( + "-P", + "--protect-results", + action="store_true", + help="Make the output web pages read only.", + default=False, +) opts = parser.parse_args() pycbc.init_logging(opts.verbose) @@ -188,12 +222,12 @@ pycbc.init_logging(opts.verbose) analysis_title = opts.analysis_title.strip('"').rstrip('"') analysis_subtitle = opts.analysis_subtitle.strip('"').rstrip('"') -if opts.template_file[0] != '/': - full_template_path = pycbc.results.__path__[0] + '/' + opts.template_file +if opts.template_file[0] != "/": + full_template_path = pycbc.results.__path__[0] + "/" + opts.template_file else: full_template_path = opts.template_file -input_template = full_template_path.split('/')[-1] +input_template = full_template_path.split("/")[-1] input_path = full_template_path.rstrip(input_template) # setup template @@ -206,8 +240,8 @@ template = env.get_template(input_template) # find all subdirs and the top-level subdirs Directory(opts.plots_dir, opts.plots_dir) -dirs = [cwd for cwd in Directory.instances] -dirs_0 = [cwd for cwd in Directory.instances if cwd.path.count('/') == 1] +dirs = [cwd for cwd in Directory.instances] +dirs_0 = [cwd for cwd in Directory.instances if cwd.path.count("/") == 1] # sort alphanumerically # FIXME: could move this into Directory when subdirs and files are appended @@ -224,67 +258,73 @@ if local_logo: # loop over all directories for cwd in dirs: - # render template - context = {'analysis_title' : analysis_title, - 'analysis_subtitle' : analysis_subtitle, - 'dirs_0' : dirs_0, - 'dir' : cwd, - 'dot_dot_str' : cwd.level() * '../', - 'plots_dir' : opts.plots_dir} + context = { + "analysis_title": analysis_title, + "analysis_subtitle": analysis_subtitle, + "dirs_0": dirs_0, + "dir": cwd, + "dot_dot_str": cwd.level() * "../", + "plots_dir": opts.plots_dir, + } if local_logo: - context['logo'] = context['dot_dot_str'] + logo_filename + context["logo"] = context["dot_dot_str"] + logo_filename else: - context['logo'] = opts.logo + context["logo"] = opts.logo output = template.render(context) # if directory does not exist make it and copy directory permissions - if not os.path.exists(opts.output_path+cwd.path): - os.makedirs(opts.output_path+cwd.path) - shutil.copymode(opts.plots_dir+cwd.path, opts.output_path+'/'+cwd.path) + if not os.path.exists(opts.output_path + cwd.path): + os.makedirs(opts.output_path + cwd.path) + shutil.copymode(opts.plots_dir + cwd.path, opts.output_path + "/" + cwd.path) # save html page - with codecs.open(opts.output_path+cwd.path+'/index.html', "w", - encoding='utf-8') as fp: + with codecs.open( + opts.output_path + cwd.path + "/index.html", "w", encoding="utf-8" + ) as fp: fp.write(output) # copy all files to html directory for cwd in dirs: for file in cwd.files: try: - shutil.copy2(file.path, opts.output_path+'/'+cwd.path+'/'+file.filename()) - except IOError: + shutil.copy2( + file.path, opts.output_path + "/" + cwd.path + "/" + file.filename() + ) + except OSError: pass # make sitemap page -sitemap_dir = '/sitemap' -if not os.path.exists(opts.plots_dir+sitemap_dir): - os.makedirs(opts.plots_dir+sitemap_dir) -render_sitemap_page(opts.plots_dir+sitemap_dir+'/well.html', dirs) -cwd = Directory(opts.plots_dir+sitemap_dir, opts.plots_dir) -context = {'analysis_title' : analysis_title, - 'analysis_subtitle' : analysis_subtitle, - 'dirs_0' : dirs_0, - 'dir' : cwd, - 'dot_dot_str' : cwd.level() * '../', - 'plots_dir' : opts.plots_dir} +sitemap_dir = "/sitemap" +if not os.path.exists(opts.plots_dir + sitemap_dir): + os.makedirs(opts.plots_dir + sitemap_dir) +render_sitemap_page(opts.plots_dir + sitemap_dir + "/well.html", dirs) +cwd = Directory(opts.plots_dir + sitemap_dir, opts.plots_dir) +context = { + "analysis_title": analysis_title, + "analysis_subtitle": analysis_subtitle, + "dirs_0": dirs_0, + "dir": cwd, + "dot_dot_str": cwd.level() * "../", + "plots_dir": opts.plots_dir, +} if local_logo: - context['logo'] = context['dot_dot_str'] + logo_filename + context["logo"] = context["dot_dot_str"] + logo_filename else: - context['logo'] = opts.logo + context["logo"] = opts.logo output = template.render(context) -if not os.path.exists(opts.output_path+cwd.path): - os.makedirs(opts.output_path+cwd.path) -with open(opts.output_path+cwd.path+'/index.html', "w") as fp: +if not os.path.exists(opts.output_path + cwd.path): + os.makedirs(opts.output_path + cwd.path) +with open(opts.output_path + cwd.path + "/index.html", "w") as fp: fp.write(output) # copy css, js, and font files to html directory -cssDir = pycbc.results.__path__[0] + '/static/css/' -jsDir = pycbc.results.__path__[0] + '/static/js/' -fontsDir = pycbc.results.__path__[0] + '/static/fonts/' -cssOutputDir = opts.output_path + '/static/css/' -jsOutputDir = opts.output_path + '/static/js/' -fontsOutputDir = opts.output_path + '/static/fonts/' +cssDir = pycbc.results.__path__[0] + "/static/css/" +jsDir = pycbc.results.__path__[0] + "/static/js/" +fontsDir = pycbc.results.__path__[0] + "/static/fonts/" +cssOutputDir = opts.output_path + "/static/css/" +jsOutputDir = opts.output_path + "/static/js/" +fontsOutputDir = opts.output_path + "/static/fonts/" if not os.path.exists(cssOutputDir): shutil.copytree(cssDir, cssOutputDir) if not os.path.exists(jsOutputDir): @@ -297,8 +337,11 @@ if not os.path.exists(fontsOutputDir): # created with restrictions that are too restrictive. So, go through and make # sure that directories are readable+executable and files are readable. -for dirpath, dirnames, filenames in os.walk(opts.output_path + '/static'): - os.chmod(dirpath, stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH) +for dirpath, dirnames, filenames in os.walk(opts.output_path + "/static"): + os.chmod( + dirpath, + stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH, + ) for filename in filenames: path = os.path.join(dirpath, filename) @@ -307,7 +350,6 @@ for dirpath, dirnames, filenames in os.walk(opts.output_path + '/static'): # Protect the results if the requested if opts.protect_results: for dirpath, dirnames, filenames in os.walk(opts.output_path): - # Do not open the box dir_permission = stat.S_IREAD | stat.S_IEXEC existing_perm = os.stat(dirpath) diff --git a/bin/pycbc_make_sky_grid b/bin/pycbc_make_sky_grid index e79e1009cb4..9e30f199e3b 100644 --- a/bin/pycbc_make_sky_grid +++ b/bin/pycbc_make_sky_grid @@ -1,6 +1,7 @@ #!/usr/bin/env python -"""Generate a grid of points in the sky to be used by `pycbc_multi_inspiral` to +""" +Generate a grid of points in the sky to be used by `pycbc_multi_inspiral` to find multi-detector gravitational wave triggers and calculate the coherent SNRs and related statistics. The grid is constructed using a stochastic sampling algorithm. The region of the sky covered by the grid can be optionally @@ -14,16 +15,17 @@ https://arxiv.org/abs/1410.6042. Please refer to the documentation of for angle arguments. """ -import logging -import numpy as np import argparse import itertools +import logging + +import numpy as np import pycbc import pycbc.distributions from pycbc.detector import Detector -from pycbc.types import angle_as_radians from pycbc.tmpltbank.sky_grid import SkyGrid +from pycbc.types import angle_as_radians def spher_to_cart(sky_points): @@ -38,21 +40,23 @@ def spher_to_cart(sky_points): def angular_distance(test_point, grid): grid_cart = spher_to_cart(grid) test_cart = spher_to_cart(np.array([test_point])) - dot = np.tensordot(grid_cart, test_cart, ((1,),(1,))).ravel() + dot = np.tensordot(grid_cart, test_cart, ((1,), (1,))).ravel() dists = np.arccos(dot) return min(dists) def dist_from_str(s): - """Instantiate and return a sky-location distribution from the given + """ + Instantiate and return a sky-location distribution from the given string argument. """ - assert s.endswith(')') - return eval('pycbc.distributions.' + s) + assert s.endswith(")") + return eval("pycbc.distributions." + s) def make_single_det_grid(args): - """Construct a sky grid for a single-detector analysis. + """ + Construct a sky grid for a single-detector analysis. Since a one-detector network has no sky localization capability, we just pick a single point. """ @@ -60,21 +64,20 @@ def make_single_det_grid(args): # --ra/--dec arguments ra, dec = args.ra, args.dec args.input_dist = f'UniformDiskSky(mean_ra="{ra}", mean_dec="{dec}", radius="{args.sky_error}")' - elif args.input_dist == 'UniformSky()': + elif args.input_dist == "UniformSky()": # Covering the whole sky, just pick a reference point - ra, dec = 0., 0. + ra, dec = 0.0, 0.0 else: # Restricted non-uniform input distribution input_dist = dist_from_str(args.input_dist) mpp = input_dist.get_max_prob_point() ra, dec = mpp[0], mpp[1] - return SkyGrid( - ra, dec, args.instruments, args.trigger_time - ) + return SkyGrid(ra, dec, args.instruments, args.trigger_time) def make_multi_det_grid(args): - """Construct a sky grid for a network of multiple detectors, + """ + Construct a sky grid for a network of multiple detectors, using the time delay among detectors to determine the required spacing between points. """ @@ -88,7 +91,7 @@ def make_multi_det_grid(args): args.input_dist = f'UniformDiskSky(mean_ra="{args.ra}", mean_dec="{args.dec}", radius="{args.sky_error}")' sky_dist = dist_from_str(args.input_dist) grid = np.vstack((np.empty((0, 2)), (args.ra, args.dec))) - elif args.input_dist == 'UniformSky()': + elif args.input_dist == "UniformSky()": # Covering the whole sky sky_dist = dist_from_str(args.input_dist) grid = np.vstack((np.empty((0, 2)), (0.0, 0.0))) @@ -96,37 +99,33 @@ def make_multi_det_grid(args): # Restricted non-uniform input distribution sky_dist = dist_from_str(args.input_dist) grid = np.vstack( - ( - np.empty((0, 2)), - np.reshape(sky_dist.get_max_prob_point(), (1, 2)) - ) + (np.empty((0, 2)), np.reshape(sky_dist.get_max_prob_point(), (1, 2))) ) sky_dist = sky_dist.to_uniform_patch(args.coverage) # Calculate the max possible light travel times between detector pairs - max_light_travel_times = np.array([ - a.light_travel_time_to_detector(b) - for a, b in detector_pairs - ]) + max_light_travel_times = np.array( + [a.light_travel_time_to_detector(b) for a, b in detector_pairs] + ) while True: prev_size = grid.shape[0] sky_dist_samples = sky_dist.rvs(size=10000) - sky_ra = sky_dist_samples['ra'] - sky_dec = sky_dist_samples['dec'] + sky_ra = sky_dist_samples["ra"] + sky_dec = sky_dist_samples["dec"] for prop_ra, prop_dec in zip(sky_ra, sky_dec): # Calculate the light travel times between detector pairs # for the given proposed sky positions - light_travel_times = np.array([ - a.time_delay_from_detector( - b, prop_ra, prop_dec, args.trigger_time - ) - for a, b in detector_pairs - ]) + light_travel_times = np.array( + [ + a.time_delay_from_detector(b, prop_ra, prop_dec, args.trigger_time) + for a, b in detector_pairs + ] + ) # Calculate the required angular spacing between the sky points ang_spacings = (2 * args.timing_uncertainty) / np.sqrt( - max_light_travel_times ** 2 - light_travel_times ** 2 + max_light_travel_times**2 - light_travel_times**2 ) angular_spacing = np.min(ang_spacings) @@ -138,43 +137,39 @@ def make_multi_det_grid(args): # far enough from other points, accept grid = np.vstack((grid, (prop_ra, prop_dec))) num_new_accepted = grid.shape[0] - prev_size - logging.info( - '%d points accepted, %d total', num_new_accepted, grid.shape[0] - ) + logging.info("%d points accepted, %d total", num_new_accepted, grid.shape[0]) if num_new_accepted == 0: # We reached our convergence criterion break - sky_grid = SkyGrid( - grid[:, 0], grid[:, 1], args.instruments, args.trigger_time - ) + sky_grid = SkyGrid(grid[:, 0], grid[:, 1], args.instruments, args.trigger_time) return sky_grid parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) parser.add_argument( - '--ra', + "--ra", type=angle_as_radians, help="Right ascension of the center of the external trigger " "error box. Use the rad or deg suffix to specify units, " "otherwise radians are assumed.", ) parser.add_argument( - '--dec', + "--dec", type=angle_as_radians, help="Declination of the center of the external trigger " "error box. Use the rad or deg suffix to specify units, " "otherwise radians are assumed.", ) parser.add_argument( - '--instruments', + "--instruments", nargs="+", type=str, required=True, help="List of instruments to analyze.", ) parser.add_argument( - '--sky-error', + "--sky-error", type=angle_as_radians, help="3-sigma confidence radius of the external trigger error " "box. Use the rad or deg suffix to specify units, otherwise " @@ -182,34 +177,32 @@ parser.add_argument( ) parser.add_argument( "--input-dist", - default='UniformSky()', + default="UniformSky()", help="Input distribution of the sky map that you have, e.g. " "HealpixSky or FisherSky. See the sky location documentation for " "more details. If not specified, a UniformSky() distribution is used.", ) parser.add_argument( - '--coverage', + "--coverage", type=float, help="Fraction of probability in the input distribution that you want " "to cover, when using --input-dist with a distribution that is not " "UniformSky.", ) parser.add_argument( - '--trigger-time', + "--trigger-time", type=int, required=True, help="Time (in s) of the external trigger", ) parser.add_argument( - '--timing-uncertainty', + "--timing-uncertainty", type=float, default=0.0005, help="Timing uncertainty (in s) we are willing to accept to determine " "the density of points.", ) -parser.add_argument( - '--output', required=True, help="Name of the output sky grid file." -) +parser.add_argument("--output", required=True, help="Name of the output sky grid file.") args = parser.parse_args() @@ -218,13 +211,16 @@ args.instruments.sort() pycbc.init_logging(args.verbose) -# Coverage argument is needed only for non uniform input distribution -if args.input_dist.startswith('Uniform'): +# Coverage argument is needed only for non uniform input distribution +if args.input_dist.startswith("Uniform"): if args.coverage is not None: - parser.error("please do not specify the coverage when using an uniform distribution") -else: - if args.coverage is None: - parser.error("please specify the coverage when using a non-uniform distribution") + parser.error( + "please do not specify the coverage when using an uniform distribution" + ) +elif args.coverage is None: + parser.error( + "please specify the coverage when using a non-uniform distribution" + ) if len(args.instruments) == 1: sky_grid = make_single_det_grid(args) @@ -233,16 +229,16 @@ else: # Record the input arguments as extra attributes extra_attributes = { - 'input_distribution': args.input_dist, - 'timing_uncertainty': args.timing_uncertainty, + "input_distribution": args.input_dist, + "timing_uncertainty": args.timing_uncertainty, } if "Uniform" not in args.input_dist and args.coverage is not None: - extra_attributes['coverage'] = args.coverage, + extra_attributes["coverage"] = (args.coverage,) elif (args.ra and args.dec and args.sky_error) is not None: - extra_attributes['trigger_ra'] = args.ra - extra_attributes['trigger_dec'] = args.dec - extra_attributes['sky_error'] = args.sky_error + extra_attributes["trigger_ra"] = args.ra + extra_attributes["trigger_dec"] = args.dec + extra_attributes["sky_error"] = args.sky_error sky_grid.write_to_file(args.output, extra_attributes) diff --git a/bin/pycbc_make_skymap b/bin/pycbc_make_skymap index df33bbc9743..87f4cc5d5ba 100755 --- a/bin/pycbc_make_skymap +++ b/bin/pycbc_make_skymap @@ -1,117 +1,145 @@ #!/usr/bin/env python -"""Fast and simple sky localization for a specific compact binary merger event. +""" +Fast and simple sky localization for a specific compact binary merger event. Runs a single-template matched filter on strain data from a number of detectors and calls BAYESTAR to produce a sky localization from the resulting set of SNR -time series.""" +time series. +""" -import os import argparse import logging +import os +import shutil import subprocess import tempfile -import shutil # we will make plots on a likely headless machine, so make sure matplotlib's # backend is set appropriately from matplotlib import use as mpl_use_backend -mpl_use_backend('agg') -import numpy as np +mpl_use_backend("agg") +import numpy as np from ligo.gracedb.rest import GraceDb import pycbc +from pycbc import dq, frame, waveform from pycbc.filter import sigmasq -from pycbc.io import gracedb as pycbc_gracedb, WaveformArray -from pycbc.types import (load_timeseries, - load_frequencyseries, MultiDetMultiColonOptionAction, - MultiDetOptionAction, MultiDetOptionAppendAction) +from pycbc.io import WaveformArray +from pycbc.io import gracedb as pycbc_gracedb from pycbc.pnutils import nearest_larger_binary_number -from pycbc.waveform.spa_tmplt import spa_length_in_time -from pycbc import frame, waveform, dq from pycbc.psd import interpolate +from pycbc.types import ( + MultiDetMultiColonOptionAction, + MultiDetOptionAction, + MultiDetOptionAppendAction, + load_frequencyseries, + load_timeseries, +) +from pycbc.waveform.spa_tmplt import spa_length_in_time def default_frame_type(time, ifo): - """Sensible defaults for frame types based on interferometer and time. - """ + """Sensible defaults for frame types based on interferometer and time.""" if time < 1137254517: # O1 - if ifo in ['H1', 'L1']: - return ifo + '_HOFT_C02' + if ifo in ["H1", "L1"]: + return ifo + "_HOFT_C02" elif time >= 1164556717 and time < 1235433618: # O2 - if ifo == 'V1': - return 'V1O2Repro2A' - elif ifo in ['H1', 'L1']: - return ifo + '_CLEANED_HOFT_C02' + if ifo == "V1": + return "V1O2Repro2A" + if ifo in ["H1", "L1"]: + return ifo + "_CLEANED_HOFT_C02" elif time >= 1235433618 and time < 1368975618: # O3 - if ifo == 'V1': - return 'V1Online' - elif ifo in ['H1', 'L1']: - return ifo + '_HOFT_CLEAN_SUB60HZ_C01' + if ifo == "V1": + return "V1Online" + if ifo in ["H1", "L1"]: + return ifo + "_HOFT_CLEAN_SUB60HZ_C01" elif time >= 1368975618: # O4 - if ifo == 'V1': - return 'HoftOnline' - elif ifo in ['H1', 'L1']: - return ifo + '_HOFT_C00_AR' - raise ValueError('Detector {} not supported at time {}'.format(ifo, time)) + if ifo == "V1": + return "HoftOnline" + if ifo in ["H1", "L1"]: + return ifo + "_HOFT_C00_AR" + raise ValueError(f"Detector {ifo} not supported at time {time}") + def default_channel_name(time, ifo): - """Sensible defaults for channel name based on interferometer and time. - """ + """Sensible defaults for channel name based on interferometer and time.""" if time < 1137254517: # O1 - if ifo in ['H1', 'L1']: - return ifo + ':DCS-CALIB_STRAIN_C02' + if ifo in ["H1", "L1"]: + return ifo + ":DCS-CALIB_STRAIN_C02" elif time > 1164556717 and time < 1235433618: # O2 - if ifo == 'V1': - return ifo + ':Hrec_hoft_V1O2Repro2A_16384Hz' - elif ifo in ['H1', 'L1']: - return ifo + ':DCH-CLEAN_STRAIN_C02' + if ifo == "V1": + return ifo + ":Hrec_hoft_V1O2Repro2A_16384Hz" + if ifo in ["H1", "L1"]: + return ifo + ":DCH-CLEAN_STRAIN_C02" elif time >= 1235433618 and time < 1368975618: # O3 - if ifo == 'V1': - return ifo + ':Hrec_hoft_16384Hz' - elif ifo in ['H1', 'L1']: - return ifo + ':DCS-CALIB_STRAIN_CLEAN_SUB60HZ_C01' + if ifo == "V1": + return ifo + ":Hrec_hoft_16384Hz" + if ifo in ["H1", "L1"]: + return ifo + ":DCS-CALIB_STRAIN_CLEAN_SUB60HZ_C01" elif time > 1368975618: # O4 - if ifo == 'V1': - return ifo + ':Hrec_hoft_16384Hz' - elif ifo in ['H1', 'L1']: - return ifo + ':GDS-CALIB_STRAIN_CLEAN_AR' - raise ValueError('Detector {} not supported at time {}'.format(ifo, time)) - -def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, - ifar, ifos, thresh_SNR, ligolw_skymap_output='.', - ligolw_event_output=None, snr_series_duration=0.1465, - frame_types=None, channel_names=None, - segment_source=None, segment_server=None, - gracedb_server=None, test_event=True, - custom_frame_files=None, approximant=None, detector_state=None, - veto_definer=None, injection_file=None, fake_strain=None, fake_strain_from_file=None, - fake_strain_seed=None, rescale_loglikelihood=None): + if ifo == "V1": + return ifo + ":Hrec_hoft_16384Hz" + if ifo in ["H1", "L1"]: + return ifo + ":GDS-CALIB_STRAIN_CLEAN_AR" + raise ValueError(f"Detector {ifo} not supported at time {time}") + + +def main( + trig_time, + mass1, + mass2, + spin1z, + spin2z, + f_low, + f_upper, + sample_rate, + ifar, + ifos, + thresh_SNR, + ligolw_skymap_output=".", + ligolw_event_output=None, + snr_series_duration=0.1465, + frame_types=None, + channel_names=None, + segment_source=None, + segment_server=None, + gracedb_server=None, + test_event=True, + custom_frame_files=None, + approximant=None, + detector_state=None, + veto_definer=None, + injection_file=None, + fake_strain=None, + fake_strain_from_file=None, + fake_strain_seed=None, + rescale_loglikelihood=None, +): if not test_event and not gracedb_server: - raise RuntimeError('a GraceDB URL must be specified if not a test event.') + raise RuntimeError("a GraceDB URL must be specified if not a test event.") tmpdir = tempfile.mkdtemp() - if len(trig_time) == 1 and ':' not in trig_time[0]: + if len(trig_time) == 1 and ":" not in trig_time[0]: # single approximate time given mean_trig_time = float(trig_time[0]) - trig_time_mode = 'approximate' + trig_time_mode = "approximate" else: # precise per-detector times given - trig_time = {kv.split(':')[0]: float(kv.split(':')[1]) - for kv in trig_time} + trig_time = {kv.split(":")[0]: float(kv.split(":")[1]) for kv in trig_time} mean_trig_time = np.mean([trig_time[ifo] for ifo in trig_time]) - trig_time_mode = 'exact' + trig_time_mode = "exact" if frame_types is None: frame_types = {} @@ -119,7 +147,7 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, channel_names = {} # resulting files will be tagged with this string - file_name_tag = '{:.0f}'.format(mean_trig_time) + file_name_tag = f"{mean_trig_time:.0f}" # parameters to fit a single-template inspiral job nicely # around the trigger time, without requiring too much data @@ -128,13 +156,13 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, # Padding set by 16 * 2 for psd and buffer for other filtering pad = 40 - template_duration = spa_length_in_time(mass1=mass1, mass2=mass2, - f_lower=f_low, phase_order=-1) + template_duration = spa_length_in_time( + mass1=mass1, mass2=mass2, f_lower=f_low, phase_order=-1 + ) segment_length = int(nearest_larger_binary_number(template_duration + pad)) # set minimum so there is enough for a psd estimate - if segment_length < 128: - segment_length = 128 - logging.info('Using segment length: %s', segment_length) + segment_length = max(segment_length, 128) + logging.info("Using segment length: %s", segment_length) gps_end_time = int(mean_trig_time + pad / 2) gps_start_time = gps_end_time - segment_length @@ -147,19 +175,23 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, for ifo in ifos: if detector_state is None or ifo not in detector_state: continue - on_segs = dq.query_str(ifo, detector_state[ifo], - gps_start_time - pad_data, - gps_end_time + pad_data, - veto_definer=veto_definer, - source=segment_source, - server=segment_server) + on_segs = dq.query_str( + ifo, + detector_state[ifo], + gps_start_time - pad_data, + gps_end_time + pad_data, + veto_definer=veto_definer, + source=segment_source, + server=segment_server, + ) if abs(on_segs) < required_duration: - logging.info('Excluding %s due to missing or vetoed data', ifo) + logging.info("Excluding %s due to missing or vetoed data", ifo) unavailable_ifos.add(ifo) ifos = sorted(set(ifos) - unavailable_ifos) if not ifos: - raise RuntimeError('All detectors have been excluded due to ' - 'missing or vetoed data') + raise RuntimeError( + "All detectors have been excluded due to missing or vetoed data" + ) highpass_frequency = int(f_low * 0.7) logging.info("Setting highpass: %s Hz", highpass_frequency) @@ -169,52 +201,86 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, st_out_paths = {} for ifo in ifos: if ifo not in channel_names: - if (fake_strain[ifo] is not None or fake_strain_from_file[ifo] is not None): - channel_names[ifo] = ifo + ':FAKE_DATA' + if fake_strain[ifo] is not None or fake_strain_from_file[ifo] is not None: + channel_names[ifo] = ifo + ":FAKE_DATA" else: channel_names[ifo] = default_channel_name(mean_trig_time, ifo) # compose the command line for the single-template process st_psd_paths[ifo] = os.path.join( - tmpdir, 'PSD_{}_{}.txt'.format(file_name_tag, ifo)) + tmpdir, f"PSD_{file_name_tag}_{ifo}.txt" + ) st_out_paths[ifo] = os.path.join( - tmpdir, 'SNRTS_{}_{}.hdf'.format(file_name_tag, ifo)) - - command = ["pycbc_single_template", - "--verbose", - "--segment-length", str(segment_length), - "--segment-start-pad", "0", - "--segment-end-pad", "0", - "--psd-estimation", "median", - "--psd-segment-length", "16", - "--psd-segment-stride", "8", - "--psd-inverse-length", "16", - "--order", "-1", - "--taper-data", "1", - "--allow-zero-padding", - "--autogating-threshold", "50", - "--autogating-cluster", "0.1", - "--autogating-width", "0.25", - "--autogating-taper", "0.25", - "--autogating-pad", "16", - "--autogating-max-iterations", "5", - "--strain-high-pass", str(highpass_frequency), - "--pad-data", str(pad_data), - "--chisq-bins", '0.72*get_freq("fSEOBNRv4Peak",params.mass1,params.mass2,params.spin1z,params.spin2z)**0.7', - "--sample-rate", str(sample_rate), - "--mass1", str(mass1), - "--mass2", str(mass2), - "--spin1z", str(spin1z), - "--spin2z", str(spin2z), - "--low-frequency-cutoff", str(f_low), - "--gps-start-time", str(gps_start_time), - "--gps-end-time", str(gps_end_time), - "--trigger-time", '{:.6f}'.format(trig_time[ifo] if ifo in trig_time - else mean_trig_time), - "--window", '1', - "--channel-name", channel_names[ifo], - "--psd-output", st_psd_paths[ifo], - "--output-file", st_out_paths[ifo]] + tmpdir, f"SNRTS_{file_name_tag}_{ifo}.hdf" + ) + + command = [ + "pycbc_single_template", + "--verbose", + "--segment-length", + str(segment_length), + "--segment-start-pad", + "0", + "--segment-end-pad", + "0", + "--psd-estimation", + "median", + "--psd-segment-length", + "16", + "--psd-segment-stride", + "8", + "--psd-inverse-length", + "16", + "--order", + "-1", + "--taper-data", + "1", + "--allow-zero-padding", + "--autogating-threshold", + "50", + "--autogating-cluster", + "0.1", + "--autogating-width", + "0.25", + "--autogating-taper", + "0.25", + "--autogating-pad", + "16", + "--autogating-max-iterations", + "5", + "--strain-high-pass", + str(highpass_frequency), + "--pad-data", + str(pad_data), + "--chisq-bins", + '0.72*get_freq("fSEOBNRv4Peak",params.mass1,params.mass2,params.spin1z,params.spin2z)**0.7', + "--sample-rate", + str(sample_rate), + "--mass1", + str(mass1), + "--mass2", + str(mass2), + "--spin1z", + str(spin1z), + "--spin2z", + str(spin2z), + "--low-frequency-cutoff", + str(f_low), + "--gps-start-time", + str(gps_start_time), + "--gps-end-time", + str(gps_end_time), + "--trigger-time", + f"{trig_time[ifo] if ifo in trig_time else mean_trig_time:.6f}", + "--window", + "1", + "--channel-name", + channel_names[ifo], + "--psd-output", + st_psd_paths[ifo], + "--output-file", + st_out_paths[ifo], + ] command.append("--approximant") for apx in approximant: @@ -256,22 +322,23 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, try: frame_data = frame.read_frame(custom_frame, channel_names[ifo]) except RuntimeError: - msg = 'Channel name in {} is not {}'.format( - custom_frame, channel_names[ifo]) + msg = f"Channel name in {custom_frame} is not {channel_names[ifo]}" raise RuntimeError(msg) fr_start_times.append(frame_data.start_time) fr_end_times.append(frame_data.end_time) if gps_start_time < min(fr_start_times): - msg = 'Start time of {} must be before the required start time {}' + msg = "Start time of {} must be before the required start time {}" msg = msg.format(min(fr_start_times), gps_start_time) raise RuntimeError(msg) if max(fr_end_times) < gps_end_time: - msg = 'End time of {} must be after the required end time {}' + msg = "End time of {} must be after the required end time {}" msg = msg.format(max(fr_end_times), gps_end_time) raise RuntimeError(msg) - if mean_trig_time < min(fr_start_times) \ - or max(fr_end_times) < mean_trig_time: - msg = 'Trigger time must be within your frame file(s)' + if ( + mean_trig_time < min(fr_start_times) + or max(fr_end_times) < mean_trig_time + ): + msg = "Trigger time must be within your frame file(s)" raise RuntimeError(msg) command.append("--frame-files") @@ -281,59 +348,59 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, # create and open a file to record the # single-template process' stdout and stderr log_path = os.path.join( - tmpdir, - 'pycbc_single_template_{}_{}_log.txt'.format(file_name_tag, ifo)) - log_file = open(log_path, 'w') - log_file.write(' '.join(command) + '\n\n') + tmpdir, f"pycbc_single_template_{file_name_tag}_{ifo}_log.txt" + ) + log_file = open(log_path, "w") + log_file.write(" ".join(command) + "\n\n") log_file.flush() # start the single-template process proc = subprocess.Popen(command, stdout=log_file, stderr=log_file) procs.append((ifo, proc, log_path, log_file)) - logging.info('Waiting for pycbc_single_template to complete') + logging.info("Waiting for pycbc_single_template to complete") snr_errors = False for ifo, proc, log_path, log_file in procs: proc.wait() if proc.returncode != 0: - logging.error('%s pycbc_single_template failed, see %s', - ifo, log_path) + logging.error("%s pycbc_single_template failed, see %s", ifo, log_path) snr_errors = True log_file.close() if snr_errors: - raise RuntimeError('one or more pycbc_single_template failed, ' - 'please see the messages above for details.') + raise RuntimeError( + "one or more pycbc_single_template failed, " + "please see the messages above for details." + ) - logging.info('Gathering info from pycbc_single_template results') + logging.info("Gathering info from pycbc_single_template results") coinc_results = {} subthreshold_ifos = [] - network_snr2 = 0. + network_snr2 = 0.0 for ifo in ifos: - snr_series = load_timeseries(st_out_paths[ifo], group='snr') - chisq_series = load_timeseries(st_out_paths[ifo], group='chisq') + snr_series = load_timeseries(st_out_paths[ifo], group="snr") + chisq_series = load_timeseries(st_out_paths[ifo], group="chisq") psd_series = load_frequencyseries(st_psd_paths[ifo]) - psd_series *= pycbc.DYN_RANGE_FAC ** 2.0 + psd_series *= pycbc.DYN_RANGE_FAC**2.0 psd_series = psd_series.astype(np.float32) - template_series = load_frequencyseries(st_out_paths[ifo], - group='template') - - key = 'foreground/{}/'.format(ifo) - coinc_results[key + 'mass1'] = mass1 - coinc_results[key + 'mass2'] = mass2 - coinc_results[key + 'spin1z'] = spin1z - coinc_results[key + 'spin2z'] = spin2z - coinc_results[key + 'f_lower'] = f_low + template_series = load_frequencyseries(st_out_paths[ifo], group="template") + + key = f"foreground/{ifo}/" + coinc_results[key + "mass1"] = mass1 + coinc_results[key + "mass2"] = mass2 + coinc_results[key + "spin1z"] = spin1z + coinc_results[key + "spin2z"] = spin2z + coinc_results[key + "f_lower"] = f_low if f_upper: - coinc_results[key + 'f_final'] = f_upper + coinc_results[key + "f_final"] = f_upper # required by CandidateForGraceDB - coinc_results[key + 'template_id'] = -1 - coinc_results[key + 'snr_series'] = snr_series - coinc_results[key + 'psd_series'] = psd_series - coinc_results[key + 'sigmasq'] = sigmasq(template_series, psd_series) + coinc_results[key + "template_id"] = -1 + coinc_results[key + "snr_series"] = snr_series + coinc_results[key + "psd_series"] = psd_series + coinc_results[key + "sigmasq"] = sigmasq(template_series, psd_series) - if trig_time_mode == 'approximate': + if trig_time_mode == "approximate": # find the absolute max SNR peak in the entire SNR time # series; it will be the trigger time in this detector # if above threshold @@ -342,58 +409,61 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, # Any good ideas for how to enforce this in a simple way? snr_series_peak = np.argmax(abs(snr_series)) snr = abs(snr_series[snr_series_peak]) - logging.info('%s peak SNR: %.2f', ifo, snr) + logging.info("%s peak SNR: %.2f", ifo, snr) if snr < thresh_SNR: subthreshold_ifos.append(ifo) continue - elif trig_time_mode == 'exact': + elif trig_time_mode == "exact": if ifo not in trig_time: # treat this detector as subthreshold subthreshold_ifos.append(ifo) continue - snr_series_peak = int((trig_time[ifo] - snr_series.start_time) \ - / snr_series.delta_t) + snr_series_peak = int( + (trig_time[ifo] - snr_series.start_time) / snr_series.delta_t + ) assert snr_series_peak >= 0 assert snr_series_peak < len(snr_series) snr = abs(snr_series[snr_series_peak]) - logging.info('%s SNR at given time: %.2f', ifo, snr) + logging.info("%s SNR at given time: %.2f", ifo, snr) - coinc_results[key + 'end_time'] = \ - float(snr_series.sample_times[snr_series_peak]) - coinc_results[key + 'snr'] = snr - coinc_results[key + 'coa_phase'] = np.angle(snr_series[snr_series_peak]) - coinc_results[key + 'chisq'] = chisq_series[snr_series_peak] + coinc_results[key + "end_time"] = float( + snr_series.sample_times[snr_series_peak] + ) + coinc_results[key + "snr"] = snr + coinc_results[key + "coa_phase"] = np.angle(snr_series[snr_series_peak]) + coinc_results[key + "chisq"] = chisq_series[snr_series_peak] - network_snr2 += snr ** 2 + network_snr2 += snr**2 trig_ifos = sorted(set(ifos) - set(subthreshold_ifos)) if not trig_ifos: - raise RuntimeError(f'All interferometers have SNR below threshold ' - '{thresh_SNR}. Is this really a candidate event?') + raise RuntimeError( + "All interferometers have SNR below threshold " + "{thresh_SNR}. Is this really a candidate event?" + ) - coinc_results['foreground/stat'] = network_snr2 ** 0.5 - coinc_results['foreground/ifar'] = ifar + coinc_results["foreground/stat"] = network_snr2**0.5 + coinc_results["foreground/ifar"] = ifar # BAYESTAR's convention for end time of subthreshold detectors - subthreshold_sngl_time = np.mean([ - coinc_results['foreground/%s/end_time' % ifo] for ifo in trig_ifos]) + subthreshold_sngl_time = np.mean( + [coinc_results["foreground/%s/end_time" % ifo] for ifo in trig_ifos] + ) window_bins = int(snr_series_duration * sample_rate) skyloc_data = {} for ifo in ifos: - key = 'foreground/{}/'.format(ifo) + key = f"foreground/{ifo}/" if ifo in subthreshold_ifos: - coinc_results[key + 'end_time'] = subthreshold_sngl_time + coinc_results[key + "end_time"] = subthreshold_sngl_time - skyloc_data[ifo] = { - 'psd': interpolate(coinc_results[key + 'psd_series'], 0.25) - } + skyloc_data[ifo] = {"psd": interpolate(coinc_results[key + "psd_series"], 0.25)} # Go back to the SNR time series and select a piece centered on the # end time - time = coinc_results[key + 'end_time'] - snr_series = coinc_results[key + 'snr_series'] + time = coinc_results[key + "end_time"] + snr_series = coinc_results[key + "snr_series"] peak_bin = int((time - snr_series.start_time) / snr_series.delta_t) max_bin = peak_bin + window_bins + 1 if max_bin > len(snr_series): @@ -407,201 +477,294 @@ def main(trig_time, mass1, mass2, spin1z, spin2z, f_low, f_upper, sample_rate, max_bin = peak_bin + window_bins + 1 min_bin = peak_bin - window_bins - skyloc_data[ifo]['snr_series'] = \ - snr_series[min_bin:max_bin].astype(np.complex64) + skyloc_data[ifo]["snr_series"] = snr_series[min_bin:max_bin].astype( + np.complex64 + ) - kwargs = {'psds': {ifo: skyloc_data[ifo]['psd'] for ifo in ifos}, - 'low_frequency_cutoff': f_low, - 'high_frequency_cutoff': f_upper, - 'skyloc_data': skyloc_data, - 'channel_names': channel_names} + kwargs = { + "psds": {ifo: skyloc_data[ifo]["psd"] for ifo in ifos}, + "low_frequency_cutoff": f_low, + "high_frequency_cutoff": f_upper, + "skyloc_data": skyloc_data, + "channel_names": channel_names, + } doc = pycbc_gracedb.CandidateForGraceDB( - trig_ifos, - trig_ifos, - coinc_results, - **kwargs + trig_ifos, trig_ifos, coinc_results, **kwargs ) - ligolw_file_path = os.path.join(tmpdir, file_name_tag + '.xml') + ligolw_file_path = os.path.join(tmpdir, file_name_tag + ".xml") if gracedb_server: - comments = ['Manual followup from PyCBC'] - gid = doc.upload(ligolw_file_path, - gracedb_server=gracedb_server, - testing=test_event, - extra_strings=comments) + comments = ["Manual followup from PyCBC"] + gid = doc.upload( + ligolw_file_path, + gracedb_server=gracedb_server, + testing=test_event, + extra_strings=comments, + ) gracedb = GraceDb(gracedb_server) else: doc.save(ligolw_file_path) # Defining the approximant to feed into BAYESTAR row = WaveformArray.from_kwargs( - mass1=mass1, - mass2=mass2, - spin1z=spin1z, - spin2z=spin2z) + mass1=mass1, mass2=mass2, spin1z=spin1z, spin2z=spin2z + ) bs_approximant = waveform.bank.parse_approximant_arg(approximant, row)[0] # BAYESTAR uses TaylorF2 instead of SPAtmplt - if bs_approximant == 'SPAtmplt': - bs_approximant = 'TaylorF2' + if bs_approximant == "SPAtmplt": + bs_approximant = "TaylorF2" # run BAYESTAR to generate the skymap - cmd = ['bayestar-localize-coincs', - ligolw_file_path, - ligolw_file_path, - '--waveform', str(bs_approximant), - '--f-low', str(f_low), - '-o', tmpdir] + cmd = [ + "bayestar-localize-coincs", + ligolw_file_path, + ligolw_file_path, + "--waveform", + str(bs_approximant), + "--f-low", + str(f_low), + "-o", + tmpdir, + ] if rescale_loglikelihood is not None: - cmd += ['--rescale-loglikelihood', rescale_loglikelihood] + cmd += ["--rescale-loglikelihood", rescale_loglikelihood] subprocess.call(cmd) - skymap_fits_name = os.path.join(tmpdir, '0.fits') + skymap_fits_name = os.path.join(tmpdir, "0.fits") # plot the skymap - skymap_plot_path = os.path.join(ligolw_skymap_output, - file_name_tag + '_skymap.png') - cmd = ['ligo-skymap-plot', - skymap_fits_name, - '-o', skymap_plot_path, - '--contour', '50', '90', - '--annotate'] + skymap_plot_path = os.path.join(ligolw_skymap_output, file_name_tag + "_skymap.png") + cmd = [ + "ligo-skymap-plot", + skymap_fits_name, + "-o", + skymap_plot_path, + "--contour", + "50", + "90", + "--annotate", + ] subprocess.call(cmd) - final_fits_path = os.path.join(ligolw_skymap_output, - file_name_tag + '.fits') + final_fits_path = os.path.join(ligolw_skymap_output, file_name_tag + ".fits") shutil.move(skymap_fits_name, final_fits_path) if gracedb_server: gracedb.write_log( gid, - 'Bayestar skymap FITS file upload', + "Bayestar skymap FITS file upload", filename=final_fits_path, - tag_name=['sky_loc'], - displayName=['Bayestar FITS skymap'] + tag_name=["sky_loc"], + displayName=["Bayestar FITS skymap"], ) gracedb.write_log( gid, - 'Bayestar skymap plot upload', + "Bayestar skymap plot upload", filename=skymap_plot_path, - tag_name=['sky_loc'], - displayName=['Bayestar skymap plot'] + tag_name=["sky_loc"], + displayName=["Bayestar skymap plot"], ) if ligolw_event_output: shutil.move(ligolw_file_path, ligolw_event_output) shutil.rmtree(tmpdir) -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) # note that I am not using a MultiDetOptionAction for --trig-time as I # explicitly want to handle cases like `--trig-time 1234` and # `--trig-time H1:1234 L1:1234` in different ways - parser.add_argument('--trig-time', required=True, nargs='+', - help='GPS time of trigger. Can be a single value, or a list of per-detector times' - 'using the format H1:123456 L1:123457. In the first case, the given time is only ' - 'an approximate guess of the trigger time, and ' - 'the code will try to find the maximum-SNR peaks ' - 'at each detector on its own. The given time should be within 1 s ' - 'of the peak. This procedure becomes unreliable ' - 'for low-SNR triggers. When times are given at specific detectors, instead, they ' - 'literally indicate which sample of the ' - 'SNR time series should be taken as peak SNR at ' - 'the given detector. Unspecified detectors will ' - 'be considered subthreshold and their time will ' - 'be chosen as the mean of the given times.') - parser.add_argument('--mass1', type=float, required=True) - parser.add_argument('--mass2', type=float, required=True) - parser.add_argument('--spin1z', type=float, required=True) - parser.add_argument('--spin2z', type=float, required=True) - parser.add_argument('--approximant', type=str, nargs='+', - default=["SPAtmplt:mtotal<4", "SEOBNRv4_ROM:else"]) - parser.add_argument('--f-low', type=float, - help="lower frequency cut-off (float)", default=20.0) - parser.add_argument('--f-upper', type=float, - help="upper frequency cut-off (float)") - parser.add_argument('--sample-rate', type=int, default=2048, - help='sample rate of the data') - parser.add_argument('--ifar', type=float, help="false alarm rate (float)", - default=1) - parser.add_argument('--thresh-SNR', type=float, default=4.5, - help='When `--trig-time` is given a single (approximate) time, only ' - 'detectors with SNR above this threshold are used ' - 'to find the precise trigger times. The others are still ' - 'used for sky localization, but do not determine ' - 'the trigger time.') - parser.add_argument('--ifos', type=str, required=True, nargs='+', - help='List of interferometer names, e.g. H1 L1') - parser.add_argument('--frame-type', type=str, nargs='+') - parser.add_argument('--channel-name', type=str, nargs='+') - parser.add_argument('--detector-state', type=str, nargs='+', - action=MultiDetMultiColonOptionAction, - help='Use the segment database to exclude detectors ' - 'with missing or vetoed data') - parser.add_argument('--veto-definer', type=str, - help='Optional path to a veto definer file to resolve ' - 'macro flags in --detector-state') - parser.add_argument('--segment-source', choices=['any', 'GWOSC', 'dqsegdb'], - default='any', - help='When using `--detector-state`, this controls ' - 'whether to query the GWOSC or dqsegdb server') - parser.add_argument('--segment-server', metavar='URL', - help='URL of segment server to query when using ' - '`--detector-state`', - default='https://segments.ligo.org') - parser.add_argument('--ligolw-skymap-output', type=str, default='.', - help='Option to output sky map files to directory') - parser.add_argument('--ligolw-event-output', type=str, - help='Option to keep coinc file under given name') - parser.add_argument('--enable-production-gracedb-upload', - action='store_true', default=False, - help='Do not mark triggers uploaded to GraceDB as test ' - 'events. This option should *only* be enabled in ' - 'production analyses!') - parser.add_argument('--gracedb-server', metavar='URL', - help='URL of GraceDB server API for uploading events.') - parser.add_argument('--custom-frame-file', type=str, nargs='+', - action=MultiDetOptionAppendAction, - help='Lists of local frame files, e.g., ' - 'H1:/path/to/frame/file L1:/path/to/frame/file') - parser.add_argument('--injection-file', type=str, - help='Optional path to an injection file.') - parser.add_argument('--fake-strain', type=str, - action=MultiDetOptionAction, metavar="IFO:FAKE_STRAIN", - help="The set of fake-strains for each detector") - parser.add_argument('--fake-strain-from-file', type=str, - action=MultiDetOptionAction, metavar="IFO:FAKE_STRAIN_FROM_FILE", - help="The set of fake-strains for each detector from files") - parser.add_argument('--fake-strain-seed', type=int, - action=MultiDetOptionAction, metavar="IFO:FAKE_STRAIN_SEED", - help="The set of fake-strains-seeds for each detector") - parser.add_argument('--rescale-loglikelihood', - help="SNR rescaling factor used in BAYESTAR") + parser.add_argument( + "--trig-time", + required=True, + nargs="+", + help="GPS time of trigger. Can be a single value, or a list of per-detector times" + "using the format H1:123456 L1:123457. In the first case, the given time is only " + "an approximate guess of the trigger time, and " + "the code will try to find the maximum-SNR peaks " + "at each detector on its own. The given time should be within 1 s " + "of the peak. This procedure becomes unreliable " + "for low-SNR triggers. When times are given at specific detectors, instead, they " + "literally indicate which sample of the " + "SNR time series should be taken as peak SNR at " + "the given detector. Unspecified detectors will " + "be considered subthreshold and their time will " + "be chosen as the mean of the given times.", + ) + parser.add_argument("--mass1", type=float, required=True) + parser.add_argument("--mass2", type=float, required=True) + parser.add_argument("--spin1z", type=float, required=True) + parser.add_argument("--spin2z", type=float, required=True) + parser.add_argument( + "--approximant", + type=str, + nargs="+", + default=["SPAtmplt:mtotal<4", "SEOBNRv4_ROM:else"], + ) + parser.add_argument( + "--f-low", type=float, help="lower frequency cut-off (float)", default=20.0 + ) + parser.add_argument("--f-upper", type=float, help="upper frequency cut-off (float)") + parser.add_argument( + "--sample-rate", type=int, default=2048, help="sample rate of the data" + ) + parser.add_argument( + "--ifar", type=float, help="false alarm rate (float)", default=1 + ) + parser.add_argument( + "--thresh-SNR", + type=float, + default=4.5, + help="When `--trig-time` is given a single (approximate) time, only " + "detectors with SNR above this threshold are used " + "to find the precise trigger times. The others are still " + "used for sky localization, but do not determine " + "the trigger time.", + ) + parser.add_argument( + "--ifos", + type=str, + required=True, + nargs="+", + help="List of interferometer names, e.g. H1 L1", + ) + parser.add_argument("--frame-type", type=str, nargs="+") + parser.add_argument("--channel-name", type=str, nargs="+") + parser.add_argument( + "--detector-state", + type=str, + nargs="+", + action=MultiDetMultiColonOptionAction, + help="Use the segment database to exclude detectors " + "with missing or vetoed data", + ) + parser.add_argument( + "--veto-definer", + type=str, + help="Optional path to a veto definer file to resolve " + "macro flags in --detector-state", + ) + parser.add_argument( + "--segment-source", + choices=["any", "GWOSC", "dqsegdb"], + default="any", + help="When using `--detector-state`, this controls " + "whether to query the GWOSC or dqsegdb server", + ) + parser.add_argument( + "--segment-server", + metavar="URL", + help="URL of segment server to query when using `--detector-state`", + default="https://segments.ligo.org", + ) + parser.add_argument( + "--ligolw-skymap-output", + type=str, + default=".", + help="Option to output sky map files to directory", + ) + parser.add_argument( + "--ligolw-event-output", + type=str, + help="Option to keep coinc file under given name", + ) + parser.add_argument( + "--enable-production-gracedb-upload", + action="store_true", + default=False, + help="Do not mark triggers uploaded to GraceDB as test " + "events. This option should *only* be enabled in " + "production analyses!", + ) + parser.add_argument( + "--gracedb-server", + metavar="URL", + help="URL of GraceDB server API for uploading events.", + ) + parser.add_argument( + "--custom-frame-file", + type=str, + nargs="+", + action=MultiDetOptionAppendAction, + help="Lists of local frame files, e.g., " + "H1:/path/to/frame/file L1:/path/to/frame/file", + ) + parser.add_argument( + "--injection-file", type=str, help="Optional path to an injection file." + ) + parser.add_argument( + "--fake-strain", + type=str, + action=MultiDetOptionAction, + metavar="IFO:FAKE_STRAIN", + help="The set of fake-strains for each detector", + ) + parser.add_argument( + "--fake-strain-from-file", + type=str, + action=MultiDetOptionAction, + metavar="IFO:FAKE_STRAIN_FROM_FILE", + help="The set of fake-strains for each detector from files", + ) + parser.add_argument( + "--fake-strain-seed", + type=int, + action=MultiDetOptionAction, + metavar="IFO:FAKE_STRAIN_SEED", + help="The set of fake-strains-seeds for each detector", + ) + parser.add_argument( + "--rescale-loglikelihood", help="SNR rescaling factor used in BAYESTAR" + ) opt = parser.parse_args() pycbc.init_logging(opt.verbose) - frame_type_dict = {f.split(':')[0]: f.split(':')[1] for f in opt.frame_type} \ - if opt.frame_type is not None else None - chan_name_dict = {f.split(':')[0]: f for f in opt.channel_name} \ - if opt.channel_name is not None else None - - main(opt.trig_time, opt.mass1, opt.mass2, - opt.spin1z, opt.spin2z, opt.f_low, opt.f_upper, opt.sample_rate, - opt.ifar, opt.ifos, opt.thresh_SNR, opt.ligolw_skymap_output, - opt.ligolw_event_output, - frame_types=frame_type_dict, channel_names=chan_name_dict, - gracedb_server=opt.gracedb_server, - test_event=not opt.enable_production_gracedb_upload, - custom_frame_files=opt.custom_frame_file, - approximant=opt.approximant, detector_state=opt.detector_state, - segment_source=opt.segment_source, segment_server=opt.segment_server, - veto_definer=opt.veto_definer, injection_file=opt.injection_file, - fake_strain=opt.fake_strain, fake_strain_from_file=opt.fake_strain_from_file, fake_strain_seed=opt.fake_strain_seed, - rescale_loglikelihood=opt.rescale_loglikelihood) + frame_type_dict = ( + {f.split(":")[0]: f.split(":")[1] for f in opt.frame_type} + if opt.frame_type is not None + else None + ) + chan_name_dict = ( + {f.split(":")[0]: f for f in opt.channel_name} + if opt.channel_name is not None + else None + ) + + main( + opt.trig_time, + opt.mass1, + opt.mass2, + opt.spin1z, + opt.spin2z, + opt.f_low, + opt.f_upper, + opt.sample_rate, + opt.ifar, + opt.ifos, + opt.thresh_SNR, + opt.ligolw_skymap_output, + opt.ligolw_event_output, + frame_types=frame_type_dict, + channel_names=chan_name_dict, + gracedb_server=opt.gracedb_server, + test_event=not opt.enable_production_gracedb_upload, + custom_frame_files=opt.custom_frame_file, + approximant=opt.approximant, + detector_state=opt.detector_state, + segment_source=opt.segment_source, + segment_server=opt.segment_server, + veto_definer=opt.veto_definer, + injection_file=opt.injection_file, + fake_strain=opt.fake_strain, + fake_strain_from_file=opt.fake_strain_from_file, + fake_strain_seed=opt.fake_strain_seed, + rescale_loglikelihood=opt.rescale_loglikelihood, + ) diff --git a/bin/pycbc_merge_inj_hdf b/bin/pycbc_merge_inj_hdf index 0d47e3dc627..eb056474b8e 100755 --- a/bin/pycbc_merge_inj_hdf +++ b/bin/pycbc_merge_inj_hdf @@ -20,8 +20,9 @@ Merge hdf injection files """ -import logging import argparse +import logging + import numpy as np import pycbc @@ -29,7 +30,8 @@ import pycbc.inject def get_gc_end_time(injection): - """Return the geocenter end time of an injection. Required for seamless + """ + Return the geocenter end time of an injection. Required for seamless compatibility with LIGOLW and HDF injection objects, which use different names. Copied from pycbc_optimal_snr. """ @@ -40,14 +42,20 @@ def get_gc_end_time(injection): return injection.tc -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) - parser.add_argument('--injection-files', '-i', dest='injection_file', - required=True, nargs='+', - help='Input HDF5 files defining injections') - parser.add_argument('--output-file', '-o', dest='out_file', required=True, - help='Output HDF5 file') + parser.add_argument( + "--injection-files", + "-i", + dest="injection_file", + required=True, + nargs="+", + help="Input HDF5 files defining injections", + ) + parser.add_argument( + "--output-file", "-o", dest="out_file", required=True, help="Output HDF5 file" + ) opts = parser.parse_args() pycbc.init_logging(opts.verbose) @@ -65,33 +73,34 @@ if __name__ == '__main__': inj_table = inj_table[np.sort(list(inj_table.dtype.fields.keys()))] injection_tables.append(injections.table) - inj_dtypes = [inj_table.dtype.fields[k] for k in \ - inj_table.dtype.fields.keys()] + inj_dtypes = [inj_table.dtype.fields[k] for k in inj_table.dtype.fields.keys()] injection_dtypes.append(inj_dtypes) - + # Check that all files contained the same fields - if not np.all([t.dtype.fields.keys() == \ - injection_tables[0].dtype.fields.keys() for t in injection_tables]): + if not np.all( + [ + t.dtype.fields.keys() == injection_tables[0].dtype.fields.keys() + for t in injection_tables + ] + ): raise TypeError("All injection files must contain the same fields!") # Check that all of the dtypes are the same if not np.all([dt == injection_dtypes[0] for dt in injection_dtypes]): - raise TypeError("All injection files must contain the same " + - "data types!") + raise TypeError("All injection files must contain the same " + "data types!") new_inj_table = [] for inj_set in injection_tables: - for inj in inj_set: + for inj in inj_set: new_inj_table.append(inj) # Store injections sorted by coalescence time new_inj_table.sort(key=get_gc_end_time) inj_dtype = injection_tables[0].dtype - new_inj_table = pycbc.io.FieldArray.from_records( - new_inj_table, dtype=inj_dtype) + new_inj_table = pycbc.io.FieldArray.from_records(new_inj_table, dtype=inj_dtype) - logging.info('Writing output') + logging.info("Writing output") pycbc.inject.InjectionSet.write(opts.out_file, new_inj_table) - logging.info('Done') + logging.info("Done") diff --git a/bin/pycbc_multi_inspiral b/bin/pycbc_multi_inspiral index 6fa3c415589..90b9966e0e3 100755 --- a/bin/pycbc_multi_inspiral +++ b/bin/pycbc_multi_inspiral @@ -25,12 +25,15 @@ take a look at: https://github.com/gwastro/pycbc/blob/master/examples/multi_inspiral/run.sh """ +import argparse import logging import time -import argparse + import numpy as np from pycbc import ( + DYN_RANGE_FAC, + add_common_pycbc_options, fft, init_logging, inject, @@ -40,18 +43,17 @@ from pycbc import ( strain, vetoes, waveform, - DYN_RANGE_FAC, - add_common_pycbc_options, ) -from pycbc.events import ranking, coherent as coh, EventManagerCoherent +from pycbc.events import EventManagerCoherent, ranking +from pycbc.events import coherent as coh from pycbc.filter import MatchedFilterControl -from pycbc.types import zeros, float32, complex64, angle_as_radians -from pycbc.vetoes import sgchisq from pycbc.tmpltbank.sky_grid import SkyGrid +from pycbc.types import angle_as_radians, complex64, float32, zeros +from pycbc.vetoes import sgchisq def slide_limiter(n_ifos, slide_shift, segment_length): - ''' + """ This function computes the number of shortslides used by the coherent matched filter statistic to obtain as many background triggers as possible. @@ -72,17 +74,18 @@ def slide_limiter(n_ifos, slide_shift, segment_length): ------- num_slides: int The number of independent timeslides - ''' + + """ low, upp = 1, segment_length if n_ifos == 1: return 1 - stride_dur = segment_length/2 - num_slides = 1 + int(stride_dur / - (slide_shift*(n_ifos-1))) - assert (num_slides >= low) and (num_slides <= upp), \ - "the combination (slideshift, segment_dur)"\ - f" = ({slide_shift:.2f},{stride_dur*2:.2f})"\ - f" goes over the allowed upper bound {upp}" + stride_dur = segment_length / 2 + num_slides = 1 + int(stride_dur / (slide_shift * (n_ifos - 1))) + assert (num_slides >= low) and (num_slides <= upp), ( + "the combination (slideshift, segment_dur)" + f" = ({slide_shift:.2f},{stride_dur * 2:.2f})" + f" goes over the allowed upper bound {upp}" + ) return num_slides @@ -93,10 +96,7 @@ time_init = time.time() parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) parser.add_argument( - "--output", - type=str, - required=True, - help="Path to the output file." + "--output", type=str, required=True, help="Path to the output file." ) parser.add_argument( "--instruments", @@ -106,10 +106,7 @@ parser.add_argument( help="List of instruments to analyze.", ) parser.add_argument( - "--bank-file", - type=str, - required=True, - help="Path to the template bank file." + "--bank-file", type=str, required=True, help="Path to the template bank file." ) parser.add_argument( "--low-frequency-cutoff", @@ -149,52 +146,75 @@ parser.add_argument( parser.add_argument( "--bank-veto-bank-file", type=str, - help="Path to the " - "bank file used to compute the the bank chi-square veto.", + help="Path to the bank file used to compute the the bank chi-square veto.", ) parser.add_argument("--chisq-bins", default=0) # Commenting out unused options: remove if they remain unused # parser.add_argument("--chisq-threshold", type=float, default=0) # parser.add_argument("--chisq-delta", type=float, default=0) -parser.add_argument("--autochi-number-points", type=int, default=0, - help="The number of points to use, in both directions if" - "doing a two-sided auto-chisq, to calculate the" - "auto-chisq statistic.") -parser.add_argument("--autochi-stride", type=int, default=0, - help="The gap, in sample points, between the points at" - "which to calculate auto-chisq.") -parser.add_argument("--autochi-two-phase", action="store_true", - default=False, - help="If given auto-chisq will be calculated by testing " - "against both phases of the SNR time-series. " - "If not given, only the phase matching the trigger " - "will be used.") -parser.add_argument("--autochi-onesided", action='store', - choices=['left', 'right'], - help="Decide whether to calculate auto-chisq using" - "points on both sides of the trigger or only on one" - "side. If not given points on both sides will be" - "used. If given, with either 'left' or 'right'," - "only points on that side (right = forward in time," - "left = back in time) will be used.") -parser.add_argument("--autochi-reverse-template", action="store_true", - default=False, - help="If given, time-reverse the template before" - "calculating the auto-chisq statistic. This will" - "come at additional computational cost as the SNR" - "time-series will need recomputing for the time-" - "reversed template.") -parser.add_argument("--autochi-max-valued", action="store_true", - default=False, - help="If given, store only the maximum value of the auto-" - "chisq over all points tested. A disadvantage of " - "this is that the mean value will not be known " - "analytically.") -parser.add_argument("--autochi-max-valued-dof", action="store", metavar="INT", - type=int, - help="If using --autochi-max-valued this value denotes " - "the pre-calculated mean value that will be stored " - "as the auto-chisq degrees-of-freedom value.") +parser.add_argument( + "--autochi-number-points", + type=int, + default=0, + help="The number of points to use, in both directions if" + "doing a two-sided auto-chisq, to calculate the" + "auto-chisq statistic.", +) +parser.add_argument( + "--autochi-stride", + type=int, + default=0, + help="The gap, in sample points, between the points at" + "which to calculate auto-chisq.", +) +parser.add_argument( + "--autochi-two-phase", + action="store_true", + default=False, + help="If given auto-chisq will be calculated by testing " + "against both phases of the SNR time-series. " + "If not given, only the phase matching the trigger " + "will be used.", +) +parser.add_argument( + "--autochi-onesided", + action="store", + choices=["left", "right"], + help="Decide whether to calculate auto-chisq using" + "points on both sides of the trigger or only on one" + "side. If not given points on both sides will be" + "used. If given, with either 'left' or 'right'," + "only points on that side (right = forward in time," + "left = back in time) will be used.", +) +parser.add_argument( + "--autochi-reverse-template", + action="store_true", + default=False, + help="If given, time-reverse the template before" + "calculating the auto-chisq statistic. This will" + "come at additional computational cost as the SNR" + "time-series will need recomputing for the time-" + "reversed template.", +) +parser.add_argument( + "--autochi-max-valued", + action="store_true", + default=False, + help="If given, store only the maximum value of the auto-" + "chisq over all points tested. A disadvantage of " + "this is that the mean value will not be known " + "analytically.", +) +parser.add_argument( + "--autochi-max-valued-dof", + action="store", + metavar="INT", + type=int, + help="If using --autochi-max-valued this value denotes " + "the pre-calculated mean value that will be stored " + "as the auto-chisq degrees-of-freedom value.", +) parser.add_argument( "--downsample-factor", type=int, @@ -207,13 +227,12 @@ parser.add_argument( parser.add_argument( "--upsample-threshold", type=float, - help="The fraction of the SNR threshold to check the " - "sparse SNR sample.", + help="The fraction of the SNR threshold to check the sparse SNR sample.", ) parser.add_argument( "--upsample-method", choices=["pruned_fft"], - default='pruned_fft', + default="pruned_fft", help="The method to find the SNR points between the sparse SNR sample.", ) parser.add_argument( @@ -228,13 +247,13 @@ parser.add_argument( "--ra", type=angle_as_radians, help="Right ascension of a single sky point to search. Use the rad or deg " - "suffix to specify units, otherwise radians are assumed." + "suffix to specify units, otherwise radians are assumed.", ) parser.add_argument( "--dec", type=angle_as_radians, help="Declination of a single sky point to search. Use the rad or deg " - "suffix to specify units, otherwise radians are assumed." + "suffix to specify units, otherwise radians are assumed.", ) parser.add_argument( "--sky-grid", @@ -246,8 +265,7 @@ parser.add_argument( "--coinc-threshold", type=float, default=0.0, - help="Triggers with coincident/coherent SNR below this " - "value will be discarded.", + help="Triggers with coincident/coherent SNR below this value will be discarded.", ) parser.add_argument( "--nifo-sngl-snr-threshold", @@ -256,14 +274,14 @@ parser.add_argument( default=2, help="How many detectors must pass the --sngl-snr-threshold cut in order " "to generate a trigger; the cut is only applied when more than 1 detector " - "is used in the analysis (default 2)." + "is used in the analysis (default 2).", ) parser.add_argument( "--sngl-snr-threshold", action="store", type=float, default=4.0, - metavar='THRESHOLD', + metavar="THRESHOLD", help="Single detector SNR threshold for trigger generation required in " "at least --nifo-sngl-snr-threshold detectors (default: 4).", ) @@ -285,7 +303,7 @@ parser.add_argument( ) parser.add_argument( "--do-null-cut", - action='store_true', + action="store_true", help="Apply a cut based on null SNR: retained triggers have null SNR " "smaller than null-min and coherent SNR smaller than null-step, or null " "SNR smaller than (null-grad * coherent SNR + null_min) and coherent SNR " @@ -333,10 +351,7 @@ parser.add_argument( default=1.0, help="Size of each time slide shift.", ) -parser.add_argument( - "--do-shortslides", - action="store_true" -) +parser.add_argument("--do-shortslides", action="store_true") # Add options groups strain.insert_strain_option_group_multi_ifo(parser) strain.StrainSegments.insert_segment_option_group_multi_ifo(parser) @@ -353,27 +368,25 @@ nifo = len(args.instruments) num_slides = 1 if args.do_shortslides: num_slides = slide_limiter( - nifo, - args.slide_shift, - args.segment_length[args.instruments[0]] + nifo, args.slide_shift, args.segment_length[args.instruments[0]] ) # Arrange detectors alphabetically so they are always called in the same order args.instruments.sort() # Check that the number of IFOs requested to have single SNR above threshold # is compatible with the number of detectors in the analysis if nifo > 1 and nifo < args.nifo_sngl_snr_threshold: - parser.error("--nifo-sngl-snr-threshold " - f"({args.nifo_sngl_snr_threshold}) cannot exceed the " - "number of detectors in the analysis " - f"({nifo}) because this would " - "never yield any trigger.") + parser.error( + "--nifo-sngl-snr-threshold " + f"({args.nifo_sngl_snr_threshold}) cannot exceed the " + "number of detectors in the analysis " + f"({nifo}) because this would " + "never yield any trigger." + ) # Use class verification methods to check whether input CLI options provided # by parser to pycbc.strain, pycbc.strain.StrainSegments, pycbc.psd, # pycbc.scheme, and pycbc.fft modules are sane. strain.verify_strain_options_multi_ifo(args, parser, args.instruments) -strain.StrainSegments.verify_segment_options_multi_ifo( - args, parser, args.instruments -) +strain.StrainSegments.verify_segment_options_multi_ifo(args, parser, args.instruments) psd.verify_psd_options_multi_ifo(args, parser, args.instruments) scheme.verify_processing_options(args, parser) fft.verify_fft_options(args, parser) @@ -424,15 +437,12 @@ with ctx: vname = "Time length" assert tlen == strain_segments_dict[ifo].time_len, vname + err_msg vname = "delta_f (=segment length inverse)" - assert delta_f == strain_segments_dict[ifo].delta_f, ( - vname + err_msg - ) + assert delta_f == strain_segments_dict[ifo].delta_f, vname + err_msg # segments is a dictionary of frequency domain objects, each one of which # is the Fourier transform of the segments in strain_segments_dict logging.info("Making frequency-domain data segments") segments = { - ifo: strain_segments_dict[ifo].fourier_segments() - for ifo in args.instruments + ifo: strain_segments_dict[ifo].fourier_segments() for ifo in args.instruments } # Memory cleaning del strain_segments_dict @@ -447,7 +457,7 @@ with ctx: flow, args.instruments, dyn_range_factor=DYN_RANGE_FAC, - precision='single', + precision="single", ) logging.info("Determining time slide shifts and time delays") @@ -465,10 +475,7 @@ with ctx: slide: { position_index: { ifo: round( - ( - time_delays_zerolag[ifo][position_index] - + time_slides[ifo][slide] - ) + (time_delays_zerolag[ifo][position_index] + time_slides[ifo][slide]) * sample_rate ) for ifo in args.instruments @@ -533,7 +540,7 @@ with ctx: twophase=args.autochi_two_phase, reverse_template=args.autochi_reverse_template, take_maximum_value=args.autochi_max_valued, - maximal_value_dof=args.autochi_max_valued_dof + maximal_value_dof=args.autochi_max_valued_dof, ) # Overwhiten all frequency-domain segments by dividing by the PSD estimate @@ -545,51 +552,51 @@ with ctx: logging.info("Setting up event manager") # But first build dictionaries to initialize and feed the event manager ifo_out_types = { - 'time_index': int, - 'ifo': int, # IFO is stored as an int internally! - 'snr': complex64, - 'chisq': float32, - 'chisq_dof': int, - 'bank_chisq': float32, - 'bank_chisq_dof': int, - 'auto_chisq': float32, - 'auto_chisq_dof': int, - 'slide_id': int, + "time_index": int, + "ifo": int, # IFO is stored as an int internally! + "snr": complex64, + "chisq": float32, + "chisq_dof": int, + "bank_chisq": float32, + "bank_chisq_dof": int, + "auto_chisq": float32, + "auto_chisq_dof": int, + "slide_id": int, } ifo_out_vals = { - 'time_index': None, - 'ifo': None, - 'snr': None, - 'chisq': None, - 'chisq_dof': None, - 'bank_chisq': None, - 'bank_chisq_dof': None, - 'auto_chisq': None, - 'auto_chisq_dof': None, - 'slide_id': None, + "time_index": None, + "ifo": None, + "snr": None, + "chisq": None, + "chisq_dof": None, + "bank_chisq": None, + "bank_chisq_dof": None, + "auto_chisq": None, + "auto_chisq_dof": None, + "slide_id": None, } ifo_names = sorted(ifo_out_vals.keys()) network_out_types = { - 'dec': float32, - 'ra': float32, - 'time_index': int, - 'coherent_snr': float32, - 'null_snr': float32, - 'nifo': int, - 'my_network_chisq': float32, - 'reweighted_snr': float32, - 'slide_id': int, + "dec": float32, + "ra": float32, + "time_index": int, + "coherent_snr": float32, + "null_snr": float32, + "nifo": int, + "my_network_chisq": float32, + "reweighted_snr": float32, + "slide_id": int, } network_out_vals = { - 'dec': None, - 'ra': None, - 'time_index': None, - 'coherent_snr': None, - 'null_snr': None, - 'nifo': None, - 'my_network_chisq': None, - 'reweighted_snr': None, - 'slide_id': None, + "dec": None, + "ra": None, + "time_index": None, + "coherent_snr": None, + "null_snr": None, + "nifo": None, + "my_network_chisq": None, + "reweighted_snr": None, + "slide_id": None, } network_names = network_out_vals.keys() event_mgr = EventManagerCoherent( @@ -675,7 +682,7 @@ with ctx: t_num + 1, n_bank, s_num + 1, - len(segments[ifo]) + len(segments[ifo]), ) # The following dicts with IFOs as keys are created to store # copies of the matched filtering results computed below. @@ -714,13 +721,11 @@ with ctx: ind = ind.astype(np.int32) if inj_filter_rejector[ifo].enabled: trig_times = ( - (ind + stilde[ifo].cumulative_index) / - snr_ts.sample_rate + - args.gps_start_time[ifo] - ) - inj_idxs = inj_filter_rejector[ifo].find_indices_in_injection_intervals( - trig_times - ) + ind + stilde[ifo].cumulative_index + ) / snr_ts.sample_rate + args.gps_start_time[ifo] + inj_idxs = inj_filter_rejector[ + ifo + ].find_indices_in_injection_intervals(trig_times) ind = ind[inj_idxs] snrv = snrv[inj_idxs] @@ -728,9 +733,7 @@ with ctx: snr_ts[matched_filter[ifo].segments[s_num].analyze] * norm ) wraparound_dict[ifo] = len(snr_dict[ifo]) - assert ( - wraparound_dict[ifo] > 0 - ), f'SNR time series for {ifo} is empty' + assert wraparound_dict[ifo] > 0, f"SNR time series for {ifo} is empty" norm_dict[ifo] = norm corr_dict[ifo] = corr.copy() idx[ifo] = ind @@ -777,8 +780,7 @@ with ctx: idx_dict = { ifo: idx[ifo][ np.logical_and( - idx[ifo] - > time_delay_idx[0][position_index][ifo], + idx[ifo] > time_delay_idx[0][position_index][ifo], idx[ifo] < time_delay_idx[0][position_index][ifo] + wraparound_dict[ifo], @@ -801,20 +803,17 @@ with ctx: # as well as the dictionary to enable the wrap around at # each detector prior to searching for coincidences. coinc_idx = coh.get_coinc_indexes( - idx_dict, time_delay_idx[slide][position_index], - args.nifo_sngl_snr_threshold, wraparound_dict - ) - logging.debug( - "Found %d coincident triggers", len(coinc_idx) + idx_dict, + time_delay_idx[slide][position_index], + args.nifo_sngl_snr_threshold, + wraparound_dict, ) + logging.debug("Found %d coincident triggers", len(coinc_idx)) # Time delay is applied to (coincident) indices at the # geocenter to have them at the IFOs (hence the + which # undoes the delay). The wrap around is performed. coinc_idx_det_frame = { - ifo: ( - coinc_idx - + time_delay_idx[slide][position_index][ifo] - ) + ifo: (coinc_idx + time_delay_idx[slide][position_index][ifo]) % wraparound_dict[ifo] for ifo in args.instruments } @@ -869,10 +868,10 @@ with ctx: } # The cut is applied after the left vs right # comparison to ensure the arrays have equal lengths - if args.projection == 'left+right': + if args.projection == "left+right": # Left polarized coherent SNR project_l = coh.get_projection_matrix( - fp, fc, sigma, projection='left' + fp, fc, sigma, projection="left" ) ( rho_coh_l, @@ -888,7 +887,7 @@ with ctx: ) # Right polarized coherent SNR project_r = coh.get_projection_matrix( - fp, fc, sigma, projection='right' + fp, fc, sigma, projection="right" ) ( rho_coh_r, @@ -912,24 +911,16 @@ with ctx: coinc_idx_l = coinc_idx_l[lr_above] coinc_idx_r = coinc_idx_r[lr_above] for ifo in coinc_triggers_l.keys(): - coinc_triggers_l[ifo] = coinc_triggers_l[ifo][ - lr_above - ] + coinc_triggers_l[ifo] = coinc_triggers_l[ifo][lr_above] for ifo in coinc_triggers_r.keys(): - coinc_triggers_r[ifo] = coinc_triggers_r[ifo][ - lr_above - ] + coinc_triggers_r[ifo] = coinc_triggers_r[ifo][lr_above] rho_coinc_l = rho_coinc_l[lr_above] rho_coinc_r = rho_coinc_r[lr_above] # Point by point, track the larger of the two # and store its information max_idx = np.argmax([rho_coh_l, rho_coh_r], axis=0) - rho_coh = np.where( - max_idx == 0, rho_coh_l, rho_coh_r - ) - coinc_idx = np.where( - max_idx == 0, coinc_idx_l, coinc_idx_r - ) + rho_coh = np.where(max_idx == 0, rho_coh_l, rho_coh_r) + coinc_idx = np.where(max_idx == 0, coinc_idx_l, coinc_idx_r) coinc_triggers = { ifo: np.where( max_idx == 0, @@ -938,9 +929,7 @@ with ctx: ) for ifo in coinc_triggers_l } - rho_coinc = np.where( - max_idx == 0, rho_coinc_l, rho_coinc_r - ) + rho_coinc = np.where(max_idx == 0, rho_coinc_l, rho_coinc_r) else: project = coh.get_projection_matrix( fp, fc, sigma, projection=args.projection @@ -962,9 +951,7 @@ with ctx: len(rho_coh), ) if len(coinc_idx) != 0: - logging.debug( - "With max coherent SNR = %.2f", max(rho_coh) - ) + logging.debug("With max coherent SNR = %.2f", max(rho_coh)) # Calculate the null SNR and apply the null SNR cut ( null, @@ -982,13 +969,9 @@ with ctx: snrv=coinc_triggers, index=coinc_idx, ) - logging.debug( - "%d triggers above null threshold", len(null) - ) + logging.debug("%d triggers above null threshold", len(null)) if len(coinc_idx) != 0: - logging.debug( - "With max null SNR = %.2f", max(null) - ) + logging.debug("With max null SNR = %.2f", max(null)) # Now calculate the individual detector chi2 values # and the SNR reweighted by chi2 and by null SNR # (no cut on reweighted SNR is applied). @@ -1002,8 +985,7 @@ with ctx: # around is performed. coinc_idx_det_frame = { ifo: ( - coinc_idx - + time_delay_idx[slide][position_index][ifo] + coinc_idx + time_delay_idx[slide][position_index][ifo] ) % wraparound_dict[ifo] for ifo in args.instruments @@ -1037,7 +1019,9 @@ with ctx: template, ) power_chisq_arrays[ifo][new_needed_idx] = new_chisq - power_chisq_dof_arrays[ifo][new_needed_idx] = new_chisq_dof + power_chisq_dof_arrays[ifo][new_needed_idx] = ( + new_chisq_dof + ) del new_chisq, new_chisq_dof, new_needed_idx chisq[ifo] = power_chisq_arrays[ifo][needed_idx] chisq_dof[ifo] = power_chisq_dof_arrays[ifo][needed_idx] @@ -1082,20 +1066,19 @@ with ctx: # remaining triggers for ifo in args.instruments: ( - ifo_out_vals['bank_chisq'], - ifo_out_vals['bank_chisq_dof'], + ifo_out_vals["bank_chisq"], + ifo_out_vals["bank_chisq_dof"], ) = bank_chisq.values( template, stilde[ifo].psd, stilde[ifo], coherent_ifo_trigs[ifo], 1, # coherent_ifo_trigs already normalized - coinc_idx_det_frame[ifo] - + stilde[ifo].analyze.start, + coinc_idx_det_frame[ifo] + stilde[ifo].analyze.start, ) ( - ifo_out_vals['auto_chisq'], - ifo_out_vals['auto_chisq_dof'], + ifo_out_vals["auto_chisq"], + ifo_out_vals["auto_chisq_dof"], ) = autochisq.values( snr_dict[ifo], coinc_idx_det_frame[ifo], @@ -1105,51 +1088,47 @@ with ctx: stilde=stilde[ifo], low_frequency_cutoff=flow, ) - ifo_out_vals['chisq'] = chisq[ifo] - ifo_out_vals['chisq_dof'] = chisq_dof[ifo] - ifo_out_vals['time_index'] = ( - coinc_idx_det_frame[ifo] - + stilde[ifo].cumulative_index + ifo_out_vals["chisq"] = chisq[ifo] + ifo_out_vals["chisq_dof"] = chisq_dof[ifo] + ifo_out_vals["time_index"] = ( + coinc_idx_det_frame[ifo] + stilde[ifo].cumulative_index ) - ifo_out_vals['snr'] = coherent_ifo_trigs[ifo] + ifo_out_vals["snr"] = coherent_ifo_trigs[ifo] # IFO is stored as an int - ifo_out_vals['ifo'] = [ - event_mgr.ifo_dict[ifo] - ] * num_events + ifo_out_vals["ifo"] = [event_mgr.ifo_dict[ifo]] * num_events # Time slide ID - ifo_out_vals['slide_id'] = [slide] * num_events + ifo_out_vals["slide_id"] = [slide] * num_events event_mgr.add_template_events_to_ifo( ifo, ifo_names, [ifo_out_vals[n] for n in ifo_names], ) if nifo > 1: - network_out_vals['coherent_snr'] = rho_coh - network_out_vals['null_snr'] = null + network_out_vals["coherent_snr"] = rho_coh + network_out_vals["null_snr"] = null else: - network_out_vals['coherent_snr'] = abs( + network_out_vals["coherent_snr"] = abs( snr_dict[args.instruments[0]][ coinc_idx_det_frame[args.instruments[0]] ] ) - network_out_vals['reweighted_snr'] = reweighted_snr - network_out_vals['my_network_chisq'] = np.real( + network_out_vals["reweighted_snr"] = reweighted_snr + network_out_vals["my_network_chisq"] = np.real( network_chisq_dict ) # coinc_idx has the geocenter indices of triggers. # The wrap around happened in get_coinc_indices. - network_out_vals['time_index'] = ( - coinc_idx - + stilde[args.instruments[0]].cumulative_index + network_out_vals["time_index"] = ( + coinc_idx + stilde[args.instruments[0]].cumulative_index ) - network_out_vals['nifo'] = [nifo] * num_events - network_out_vals['dec'] = [ + network_out_vals["nifo"] = [nifo] * num_events + network_out_vals["dec"] = [ sky_positions.decs[position_index] ] * num_events - network_out_vals['ra'] = [ + network_out_vals["ra"] = [ sky_positions.ras[position_index] ] * num_events - network_out_vals['slide_id'] = [slide] * num_events + network_out_vals["slide_id"] = [slide] * num_events event_mgr.add_template_events_to_network( network_names, [network_out_vals[n] for n in network_names], @@ -1165,7 +1144,7 @@ with ctx: for slide in range(num_slides): logging.debug("Clustering slide %d", slide) event_mgr.cluster_template_network_events( - 'time_index', 'reweighted_snr', cluster_window, slide=slide + "time_index", "reweighted_snr", cluster_window, slide=slide ) # Left loop over segments event_mgr.finalize_template_events() @@ -1176,11 +1155,7 @@ with ctx: time_stop = time.time() run_time = time_stop - time_init event_mgr.save_performance( - num_cpu_cores, - len(segments[args.instruments[0]]), - n_bank, - run_time, - time_setup + num_cpu_cores, len(segments[args.instruments[0]]), n_bank, run_time, time_setup ) logging.info("Writing output") diff --git a/bin/pycbc_optimal_snr b/bin/pycbc_optimal_snr index 39b7a257847..7d0a36daa21 100644 --- a/bin/pycbc_optimal_snr +++ b/bin/pycbc_optimal_snr @@ -17,45 +17,52 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. """ -Compute the optimal SNRs for every injection in an input HDF file or +Compute the optimal SNRs for every injection in an input HDF file or sim_inspiral table and store the result in HDF file datasets or in the same ligolw table. """ -import logging import argparse +import logging import multiprocessing -import numpy as np -from igwn_ligolw import utils as ligolw_utils +import numpy as np from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils import pycbc import pycbc.inject import pycbc.psd -from pycbc.pool import BroadcastPool as Pool -from pycbc.filter import sigma, make_frequency_series -from pycbc.types import TimeSeries, FrequencySeries, zeros, float32, \ - MultiDetOptionAction, load_frequencyseries -from pycbc.io.ligolw import get_table_columns +from pycbc.filter import make_frequency_series, sigma from pycbc.io.hdf import HFile +from pycbc.io.ligolw import get_table_columns +from pycbc.pool import BroadcastPool as Pool +from pycbc.types import ( + FrequencySeries, + MultiDetOptionAction, + TimeSeries, + float32, + load_frequencyseries, + zeros, +) -class TimeIndependentPSD(object): +class TimeIndependentPSD: def __init__(self, psd_series): self.psd_series = psd_series def __call__(self, time=None): return self.psd_series -class TimeVaryingPSD(object): + +class TimeVaryingPSD: def __init__(self, file_name, length=None, delta_f=None, f_low=None): - with HFile(file_name, 'r') as f: + with HFile(file_name, "r") as f: self.file_name = file_name detector = tuple(f.keys())[0] - self.start_times = f[detector + '/start_time'][:] - self.end_times = f[detector + '/end_time'][:] - self.file_f_low = f.attrs['low_frequency_cutoff'] + self.start_times = f[detector + "/start_time"][:] + self.end_times = f[detector + "/end_time"][:] + self.file_f_low = f.attrs["low_frequency_cutoff"] self._curr_psd = {} self._curr_psd_index = {} self.detector = detector @@ -64,11 +71,10 @@ class TimeVaryingPSD(object): self.f_low = f_low def __call__(self, time=None): - mask = np.logical_and(self.start_times <= time, - self.end_times > time) + mask = np.logical_and(self.start_times <= time, self.end_times > time) if not mask.any(): return None - center_times = (self.start_times[mask] + self.end_times[mask]) / 2. + center_times = (self.start_times[mask] + self.end_times[mask]) / 2.0 closest_idx = np.argmin(abs(center_times - time)) return self.get_psd(np.flatnonzero(mask)[closest_idx]) @@ -77,18 +83,19 @@ class TimeVaryingPSD(object): if curr_pid not in self._curr_psd_index.keys(): self._curr_psd_index[curr_pid] = -1 if not index == self._curr_psd_index[curr_pid]: - group = self.detector + '/psds/' + str(index) + group = self.detector + "/psds/" + str(index) psd = load_frequencyseries(self.file_name, group=group) if delta_f is not None and psd.delta_f != delta_f: psd = pycbc.psd.interpolate(psd, delta_f) if self.length is not None and self.length != len(psd): - psd2 = FrequencySeries(zeros(self.length, dtype=psd.dtype), - delta_f=psd.delta_f) + psd2 = FrequencySeries( + zeros(self.length, dtype=psd.dtype), delta_f=psd.delta_f + ) if self.length > len(psd): psd2[:] = np.inf - psd2[0:len(psd)] = psd + psd2[0 : len(psd)] = psd else: - psd2[:] = psd[0:self.length] + psd2[:] = psd[0 : self.length] psd = psd2 if self.f_low is not None and self.f_low < self.file_f_low: # avoid using the PSD below the f_low given in the file @@ -98,15 +105,18 @@ class TimeVaryingPSD(object): self._curr_psd_index[curr_pid] = index return self._curr_psd[curr_pid] + def parse_injection_range(num_inj, rangestr): - part = int(rangestr.split('/')[0]) - pieces = int(rangestr.split('/')[1]) - tmin = num_inj * part // pieces - tmax = num_inj * (part + 1) // pieces + part = int(rangestr.split("/")[0]) + pieces = int(rangestr.split("/")[1]) + tmin = num_inj * part // pieces + tmax = num_inj * (part + 1) // pieces return tmin, tmax + def get_gc_end_time(injection): - """Return the geocenter end time of an injection. Required for seamless + """ + Return the geocenter end time of an injection. Required for seamless compatibility with LIGOLW and HDF injection objects, which use different names. """ @@ -117,67 +127,110 @@ def get_gc_end_time(injection): return injection.tc -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) - parser.add_argument('--input-file', '-i', dest='injection_file', - required=True, - help='Input LIGOLW file defining injections') - parser.add_argument('--injection-f-ref', type=float, - help='Reference frequency in Hz for ' - 'creating CBC injections from an XML ' - 'file.') - parser.add_argument('--injection-f-final', type=float, - help='Override the f_final field of a CBC ' - 'XML injection file.') - parser.add_argument('--input-sort', type=str, choices=['random', 'time'], - default='random', - help='Sort injections by the given criterion before ' - 'processing them. Sorting by time may speed up ' - 'the calculation when the approximant is very ' - 'fast and injections are tightly spaced in time') - parser.add_argument('--output-file', '-o', dest='out_file', required=True, - help='Output LIGOLW file') - parser.add_argument('--f-low', type=float, default=30., - help='Start frequency of matched-filter integration ' - 'in Hz (default %(default)s)') - parser.add_argument('--seg-length', type=float, default=256, - help='Segment duration in seconds (default %(default)s)') - parser.add_argument('--sample-rate', type=float, default=16384, - help='Data sample rate in Hz (default %(default)s)') - parser.add_argument('--ifos', nargs='+', - help='Specify ifos for default HDF dataset output') - parser.add_argument('--snr-columns', nargs='+', action=MultiDetOptionAction, - metavar='DETECTOR:COLUMN', - help='For sim_inspiral table output, specify columns' - ' to store the optimal SNR for each detector. COLUMN' - ' should be an existing sim_inspiral column containing' - ' no useful data, alpha1, alpha2 etc. are good' - ' candidates. For HDF output the --ifos option should be' - ' used, datasets will be named eg "optimal_snr_H1"') - parser.add_argument('--cores', default=1, type=int, - help='Parallelize the computation over the given ' - 'number of cores') - parser.add_argument('--ignore-waveform-errors', action='store_true', - help='Ignore errors in waveform generation and keep ' - 'the corresponding column unchanged') - parser.add_argument('--progress', action='store_true', - help='Show a progress bar (requires tqdm)') + parser.add_argument( + "--input-file", + "-i", + dest="injection_file", + required=True, + help="Input LIGOLW file defining injections", + ) + parser.add_argument( + "--injection-f-ref", + type=float, + help="Reference frequency in Hz for creating CBC injections from an XML file.", + ) + parser.add_argument( + "--injection-f-final", + type=float, + help="Override the f_final field of a CBC XML injection file.", + ) + parser.add_argument( + "--input-sort", + type=str, + choices=["random", "time"], + default="random", + help="Sort injections by the given criterion before " + "processing them. Sorting by time may speed up " + "the calculation when the approximant is very " + "fast and injections are tightly spaced in time", + ) + parser.add_argument( + "--output-file", "-o", dest="out_file", required=True, help="Output LIGOLW file" + ) + parser.add_argument( + "--f-low", + type=float, + default=30.0, + help="Start frequency of matched-filter integration " + "in Hz (default %(default)s)", + ) + parser.add_argument( + "--seg-length", + type=float, + default=256, + help="Segment duration in seconds (default %(default)s)", + ) + parser.add_argument( + "--sample-rate", + type=float, + default=16384, + help="Data sample rate in Hz (default %(default)s)", + ) + parser.add_argument( + "--ifos", nargs="+", help="Specify ifos for default HDF dataset output" + ) + parser.add_argument( + "--snr-columns", + nargs="+", + action=MultiDetOptionAction, + metavar="DETECTOR:COLUMN", + help="For sim_inspiral table output, specify columns" + " to store the optimal SNR for each detector. COLUMN" + " should be an existing sim_inspiral column containing" + " no useful data, alpha1, alpha2 etc. are good" + " candidates. For HDF output the --ifos option should be" + ' used, datasets will be named eg "optimal_snr_H1"', + ) + parser.add_argument( + "--cores", + default=1, + type=int, + help="Parallelize the computation over the given number of cores", + ) + parser.add_argument( + "--ignore-waveform-errors", + action="store_true", + help="Ignore errors in waveform generation and keep " + "the corresponding column unchanged", + ) + parser.add_argument( + "--progress", action="store_true", help="Show a progress bar (requires tqdm)" + ) psd_group = pycbc.psd.insert_psd_option_group_multi_ifo(parser) - psd_group.add_argument('--time-varying-psds', nargs='*', metavar='FILE', - help='Instead of time-independent PSDs, use time-varying ' - 'PSDs from the given HDF5 files and pick the appropriate ' - 'PSD for each injection') - parser.add_argument('--injection-fraction-range', default='0/1', - help='Optional, analyze only a certain range of the ' - 'injections. Format PART/NUM_PARTS') + psd_group.add_argument( + "--time-varying-psds", + nargs="*", + metavar="FILE", + help="Instead of time-independent PSDs, use time-varying " + "PSDs from the given HDF5 files and pick the appropriate " + "PSD for each injection", + ) + parser.add_argument( + "--injection-fraction-range", + default="0/1", + help="Optional, analyze only a certain range of the " + "injections. Format PART/NUM_PARTS", + ) opts = parser.parse_args() if opts.ifos is not None: detectors = opts.ifos if opts.snr_columns: parser.error("Can't use both --ifos and --snr-columns !") - opts.snr_columns = {i: 'optimal_snr_' + i for i in opts.ifos} + opts.snr_columns = {i: "optimal_snr_" + i for i in opts.ifos} else: detectors = opts.snr_columns.keys() @@ -188,8 +241,8 @@ if __name__ == '__main__': seg_len = opts.seg_length sample_rate = opts.sample_rate - delta_t = 1. / sample_rate - delta_f = 1. / seg_len + delta_t = 1.0 / sample_rate + delta_f = 1.0 / seg_len tlen = int(seg_len * sample_rate) flen = tlen // 2 + 1 f_low = opts.f_low @@ -201,8 +254,10 @@ if __name__ == '__main__': tvpsd = TimeVaryingPSD(tvpsd_file, flen, delta_f, f_low) psds[tvpsd.detector] = tvpsd if set(detectors) != set(psds.keys()): - parser.error('Inconsistent detector list in time-varying PSD ' \ - 'specification (%s vs %s)' % (detectors, psds.keys())) + parser.error( + "Inconsistent detector list in time-varying PSD " + "specification (%s vs %s)" % (detectors, psds.keys()) + ) else: psds = pycbc.psd.from_cli_multi_ifos( opts, @@ -211,19 +266,27 @@ if __name__ == '__main__': low_frequency_cutoff_dict=dict((det, f_low) for det in detectors), ifos=detectors, strain_dict=dict((det, None) for det in detectors), - dyn_range_factor=pycbc.DYN_RANGE_FAC) + dyn_range_factor=pycbc.DYN_RANGE_FAC, + ) for det in detectors: psds[det] = TimeIndependentPSD(psds[det].astype(float32)) def get_injection(injections, det, injection_time, simulation_id): - """ Do an injection from the injection XML file, specified by - IFO and end time""" + """ + Do an injection from the injection XML file, specified by + IFO and end time + """ # leave 4 s of padding at the end for possible ringdown - start_time = int(injection_time + 4. - seg_len) - strain = TimeSeries(zeros(tlen, dtype=float32), delta_t=delta_t, - epoch=start_time) - injections.apply(strain, det, distance_scale=1./pycbc.DYN_RANGE_FAC, - simulation_ids=[simulation_id]) + start_time = int(injection_time + 4.0 - seg_len) + strain = TimeSeries( + zeros(tlen, dtype=float32), delta_t=delta_t, epoch=start_time + ) + injections.apply( + strain, + det, + distance_scale=1.0 / pycbc.DYN_RANGE_FAC, + simulation_ids=[simulation_id], + ) return make_frequency_series(strain) def compute_optimal_snr(inj): @@ -234,27 +297,25 @@ if __name__ == '__main__': psd = psds[det](injection_time) if psd is None: continue - logging.debug('Trying injection %s at %s', inj.simulation_id, det) + logging.debug("Trying injection %s at %s", inj.simulation_id, det) try: - wave = get_injection(injections, det, injection_time, - simulation_id=inj.simulation_id) + wave = get_injection( + injections, det, injection_time, simulation_id=inj.simulation_id + ) except Exception as e: if opts.ignore_waveform_errors: logging.debug( - '%s: waveform generation failed, skipping (%s)', + "%s: waveform generation failed, skipping (%s)", inj.simulation_id, - e + e, ) continue - else: - logging.error('%s: waveform generation failed with the ' - 'following exception', inj.simulation_id) - raise - logging.debug( - 'Injection %s at %s completed', - inj.simulation_id, - det - ) + logging.exception( + "%s: waveform generation failed with the following exception", + inj.simulation_id, + ) + raise + logging.debug("Injection %s at %s completed", inj.simulation_id, det) sval = sigma(wave, psd=psd, low_frequency_cutoff=f_low) if ligolw: setattr(inj, column, sval) @@ -271,29 +332,28 @@ if __name__ == '__main__': # to preserve the entire content of the document and only modify the # particular columns of the sim_inspiral table. If HDF injections are # involved, we do not care. - ligolw_suffixes = ('.xml', '.xml.gz') - ligolw = opts.injection_file.endswith(ligolw_suffixes) \ - and opts.out_file.endswith(ligolw_suffixes) + ligolw_suffixes = (".xml", ".xml.gz") + ligolw = opts.injection_file.endswith(ligolw_suffixes) and opts.out_file.endswith( + ligolw_suffixes + ) if not ligolw: # create placeholder fields for FieldArray injections for det, column in opts.snr_columns.items(): if column in inj_table: continue - inj_table = inj_table.add_fields(np.zeros(len(inj_table)), - column) + inj_table = inj_table.add_fields(np.zeros(len(inj_table)), column) # make sure we have simulation IDs - if 'simulation_id' not in inj_table: - inj_table = inj_table.add_fields(np.arange(len(inj_table)), - 'simulation_id') + if "simulation_id" not in inj_table: + inj_table = inj_table.add_fields(np.arange(len(inj_table)), "simulation_id") inj_dtype = inj_table.dtype - if opts.input_sort == 'random': + if opts.input_sort == "random": np.random.seed(100) sort_func = lambda x: np.random.random() - elif opts.input_sort == 'time': + elif opts.input_sort == "time": sort_func = get_gc_end_time inj_table = sorted(inj_table, key=sort_func) @@ -306,12 +366,11 @@ if __name__ == '__main__': # Cut inj_table down to the range defined by opts.injection_fraction_range num_injections = len(inj_table) - imin, imax = parse_injection_range(num_injections, - opts.injection_fraction_range) + imin, imax = parse_injection_range(num_injections, opts.injection_fraction_range) inj_table = inj_table[imin:imax] if opts.cores > 1: - logging.info('Starting workers') + logging.info("Starting workers") pool = Pool(processes=opts.cores) iterator = pool.imap_unordered(compute_optimal_snr, inj_table) else: @@ -324,8 +383,7 @@ if __name__ == '__main__': iterator = tqdm(iterator, total=len(inj_table)) except ImportError: - logging.warning('cannot import tqdm; not showing progress bar') - pass + logging.warning("cannot import tqdm; not showing progress bar") for inj in iterator: new_inj_table.append(inj) @@ -334,18 +392,17 @@ if __name__ == '__main__': new_inj_table.sort(key=get_gc_end_time) if not ligolw: - new_inj_table = pycbc.io.FieldArray.from_records( - new_inj_table, dtype=inj_dtype) + new_inj_table = pycbc.io.FieldArray.from_records(new_inj_table, dtype=inj_dtype) - logging.info('Writing output') + logging.info("Writing output") if ligolw: llw_doc = injections.indoc llw_root = llw_doc.childNodes[0] llw_root.removeChild(injections.table) llw_root.appendChild(new_inj_table) - ligolw_utils.write_filename(llw_doc, opts.out_file, compress='auto') + ligolw_utils.write_filename(llw_doc, opts.out_file, compress="auto") else: pycbc.inject.InjectionSet.write(opts.out_file, new_inj_table) - logging.info('Done') + logging.info("Done") diff --git a/bin/pycbc_optimize_snr b/bin/pycbc_optimize_snr index 6ce3abb9ba0..2c5657a5781 100755 --- a/bin/pycbc_optimize_snr +++ b/bin/pycbc_optimize_snr @@ -2,70 +2,103 @@ """Followup utility to optimize the SNR of a PyCBC Live trigger.""" -import os import argparse import logging +import os + import numpy # we will make plots on a likely headless machine, so make sure matplotlib's # backend is set appropriately from matplotlib import use as mpl_use_backend -mpl_use_backend('agg') + +mpl_use_backend("agg") import pycbc -from pycbc import ( - fft, scheme -) -from pycbc.types import MultiDetOptionAction, load_frequencyseries import pycbc.conversions as cv -from pycbc.io.gracedb import CandidateForGraceDB -from pycbc.io.hdf import load_hdf5_to_dict, HFile +from pycbc import fft, scheme from pycbc.detector import Detector -from pycbc.psd import interpolate +from pycbc.io.gracedb import CandidateForGraceDB +from pycbc.io.hdf import HFile, load_hdf5_to_dict from pycbc.live import snr_optimizer - +from pycbc.psd import interpolate +from pycbc.types import MultiDetOptionAction, load_frequencyseries parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--params-file', required=True, - help='Location of the attributes file created by PyCBC ' - 'Live') -parser.add_argument('--data-files', type=str, nargs='+', - action=MultiDetOptionAction, - metavar='IFO:DATA_FILE', - help='Locations of the overwhitened data files produced ' - 'by PyCBC Live.') -parser.add_argument('--psd-files', type=str, nargs='+', - action=MultiDetOptionAction, - metavar='IFO:PSD_FILE', - help='Locations of the PSD files produced ' - 'by PyCBC Live.') -parser.add_argument('--approximant', required=True, - help='Waveform approximant string.') -parser.add_argument('--snr-threshold', type=float, default=4.0, - help='If the SNR in ifo X is below this threshold do not ' - 'consider it part of the coincidence. Not implemented') -parser.add_argument('--chirp-time-f-lower', type=float, default=20., - help='Starting frequency for chirp time window (Hz).') -parser.add_argument('--chirp-time-window', type=float, default=2., - help='Chirp time window (s).') -parser.add_argument('--gracedb-server', metavar='URL', - help='URL of GraceDB server API for uploading events. ' - 'If not provided, the default URL is used.') -parser.add_argument('--gracedb-search', type=str, default='AllSky', - help='String going into the "search" field of the GraceDB ' - 'events') -parser.add_argument('--gracedb-labels', metavar='LABEL', nargs='+', - help='Apply the given list of labels to events uploaded ' - 'to GraceDB.') -parser.add_argument('--production', action='store_true', - help='Upload a production event rather than a test event') -parser.add_argument('--enable-gracedb-upload', action='store_true', default=False, - help='Upload triggers to GraceDB') -parser.add_argument('--output-path', required=True, - help='Path to a directory to store results in') -parser.add_argument('--cores', type=int, - help='Restrict calculation to given number of CPU cores') +parser.add_argument( + "--params-file", + required=True, + help="Location of the attributes file created by PyCBC Live", +) +parser.add_argument( + "--data-files", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:DATA_FILE", + help="Locations of the overwhitened data files produced by PyCBC Live.", +) +parser.add_argument( + "--psd-files", + type=str, + nargs="+", + action=MultiDetOptionAction, + metavar="IFO:PSD_FILE", + help="Locations of the PSD files produced by PyCBC Live.", +) +parser.add_argument("--approximant", required=True, help="Waveform approximant string.") +parser.add_argument( + "--snr-threshold", + type=float, + default=4.0, + help="If the SNR in ifo X is below this threshold do not " + "consider it part of the coincidence. Not implemented", +) +parser.add_argument( + "--chirp-time-f-lower", + type=float, + default=20.0, + help="Starting frequency for chirp time window (Hz).", +) +parser.add_argument( + "--chirp-time-window", type=float, default=2.0, help="Chirp time window (s)." +) +parser.add_argument( + "--gracedb-server", + metavar="URL", + help="URL of GraceDB server API for uploading events. " + "If not provided, the default URL is used.", +) +parser.add_argument( + "--gracedb-search", + type=str, + default="AllSky", + help='String going into the "search" field of the GraceDB events', +) +parser.add_argument( + "--gracedb-labels", + metavar="LABEL", + nargs="+", + help="Apply the given list of labels to events uploaded to GraceDB.", +) +parser.add_argument( + "--production", + action="store_true", + help="Upload a production event rather than a test event", +) +parser.add_argument( + "--enable-gracedb-upload", + action="store_true", + default=False, + help="Upload triggers to GraceDB", +) +parser.add_argument( + "--output-path", required=True, help="Path to a directory to store results in" +) +parser.add_argument( + "--cores", type=int, help="Restrict calculation to given number of CPU cores" +) snr_optimizer.insert_snr_optimizer_options(parser) scheme.insert_processing_option_group(parser) fft.insert_fft_option_group(parser) @@ -74,8 +107,8 @@ args = parser.parse_args() pycbc.init_logging(args.verbose) -if args.snr_opt_seed is not None and args.snr_opt_seed != 'random': - logging.info('Setting snr optimizer random seed.') +if args.snr_opt_seed is not None and args.snr_opt_seed != "random": + logging.info("Setting snr optimizer random seed.") numpy.random.seed(int(args.snr_opt_seed)) # Input checking @@ -97,7 +130,7 @@ MIN_CPT_MASS = snr_optimizer.MIN_CPT_MASS MAX_ETA = 0.24999999999 MIN_ETA = 0.01 -logging.info('Starting optimize SNR') +logging.info("Starting optimize SNR") data = {} ifos = list(args.data_files.keys()) @@ -105,47 +138,56 @@ for ifo in ifos: data[ifo] = load_frequencyseries(args.data_files[ifo]) data[ifo].psd = load_frequencyseries(args.psd_files[ifo]) -fp = HFile(args.params_file, 'r') +fp = HFile(args.params_file, "r") original_gid = None -if 'gid' in fp: - original_gid = fp['gid'].asstr()[()] +if "gid" in fp: + original_gid = fp["gid"].asstr()[()] if args.enable_gracedb_upload and not original_gid: - raise RuntimeError('Params must include original gracedb ID in order ' - 'to upload followup!') + raise RuntimeError( + "Params must include original gracedb ID in order to upload followup!" + ) if original_gid: - logging.info('Following up GID %s', original_gid) + logging.info("Following up GID %s", original_gid) coinc_times = {} for ifo in ifos: try: - coinc_times[ifo] = fp['coinc_times'][ifo][()] + coinc_times[ifo] = fp["coinc_times"][ifo][()] except KeyError: pass coinc_ifos = list(coinc_times.keys()) -flen = fp['flen'][()] +flen = fp["flen"][()] approximant = args.approximant -flow = float(fp['flow'][()]) -f_end = float(fp['f_end'][()]) -delta_f = fp['delta_f'][()] -sample_rate = fp['sample_rate'][()] - -extra_args = [data, coinc_times, coinc_ifos, flen, - approximant, flow, f_end, delta_f, sample_rate] +flow = float(fp["flow"][()]) +f_end = float(fp["f_end"][()]) +delta_f = fp["delta_f"][()] +sample_rate = fp["sample_rate"][()] + +extra_args = [ + data, + coinc_times, + coinc_ifos, + flen, + approximant, + flow, + f_end, + delta_f, + sample_rate, +] # Determine chirp mass bounds from constant chirp time window tau0flow = args.chirp_time_f_lower -tau0 = cv.tau0_from_mass1_mass2(fp['mass1'][()], fp['mass2'][()], tau0flow) +tau0 = cv.tau0_from_mass1_mass2(fp["mass1"][()], fp["mass2"][()], tau0flow) # Arbitrarily set max mchirp to 200 Msun if tau0 - window is too small -mintau0 = max(tau0 - args.chirp_time_window, - cv.tau0_from_mchirp(200., tau0flow)) +mintau0 = max(tau0 - args.chirp_time_window, cv.tau0_from_mchirp(200.0, tau0flow)) maxtau0 = tau0 + args.chirp_time_window minchirp = cv.mchirp_from_tau0(maxtau0, tau0flow) maxchirp = cv.mchirp_from_tau0(mintau0, tau0flow) # Check basic sanity assert minchirp > 0.5 assert minchirp < maxchirp -assert maxchirp < 250. +assert maxchirp < 250.0 # Establish minimum eta: find the most asymmetric mass point # nb function name is 'mass2' but it doesn't enforce m2 < m1! @@ -164,35 +206,37 @@ maxspin2z = 0.4 if maxm2 < 3 else 0.9 # Boundary of the optimization space bounds = { - 'mchirp': (minchirp, maxchirp), - 'eta': (mineta, MAX_ETA), - 'spin1z': (minspin1z, maxspin1z), - 'spin2z': (minspin2z, maxspin2z) + "mchirp": (minchirp, maxchirp), + "eta": (mineta, MAX_ETA), + "spin1z": (minspin1z, maxspin1z), + "spin2z": (minspin2z, maxspin2z), } if args.snr_opt_include_candidate: # Initial point from found candidate - mchirp_init = cv.mchirp_from_mass1_mass2(fp['mass1'][()], fp['mass2'][()]) - eta_init = cv.eta_from_mass1_mass2(fp['mass1'][()], fp['mass2'][()]) - spin1z_init = fp['spin1z'][()] - spin2z_init = fp['spin2z'][()] - - initial_point = numpy.array([ - mchirp_init, - eta_init, - spin1z_init, - spin2z_init, - ])[numpy.newaxis] + mchirp_init = cv.mchirp_from_mass1_mass2(fp["mass1"][()], fp["mass2"][()]) + eta_init = cv.eta_from_mass1_mass2(fp["mass1"][()], fp["mass2"][()]) + spin1z_init = fp["spin1z"][()] + spin2z_init = fp["spin2z"][()] + + initial_point = numpy.array( + [ + mchirp_init, + eta_init, + spin1z_init, + spin2z_init, + ] + )[numpy.newaxis] else: initial_point = None with scheme_context: - logging.info('Starting optimization') + logging.info("Starting optimization") optimize_func = snr_optimizer.optimize_funcs[args.snr_opt_method] opt_params = optimize_func(bounds, args, extra_args, initial_point) - logging.info('Optimization complete') + logging.info("Optimization complete") fup_ifos = set(ifos) - set(coinc_ifos) for ifo in fup_ifos: @@ -226,16 +270,18 @@ loudest_snr_idxs = {} for ifo in ifos: duration = 0.095 half_dur_samples = int(sample_rate * duration / 2) - onsource_idx = float(coinc_times[ifo] - snr_series_dict[ifo].start_time) \ + onsource_idx = ( + float(coinc_times[ifo] - snr_series_dict[ifo].start_time) * snr_series_dict[ifo].sample_rate + ) onsource_idx = int(round(onsource_idx)) - onsource_slice = slice(onsource_idx - half_dur_samples, - onsource_idx + half_dur_samples + 1) + onsource_slice = slice( + onsource_idx - half_dur_samples, onsource_idx + half_dur_samples + 1 + ) max_snr_idx = numpy.argmax(abs(snr_series_dict[ifo][onsource_slice])) max_snr_idx = max_snr_idx + onsource_idx - half_dur_samples max_snr = snr_series_dict[ifo][max_snr_idx] - max_snr_time = snr_series_dict[ifo].start_time \ - + max_snr_idx*(1./sample_rate) + max_snr_time = snr_series_dict[ifo].start_time + max_snr_idx * (1.0 / sample_rate) loudest_snr_idxs[ifo] = max_snr_idx loudest_snr_times[ifo] = float(max_snr_time) loudest_snrs[ifo] = max_snr @@ -253,36 +299,37 @@ for ifo in ifos: # Add small buffer lttbd += 0.005 - time_window = [loudest_snr_time_allifos-lttbd, - loudest_snr_time_allifos+lttbd] + time_window = [loudest_snr_time_allifos - lttbd, loudest_snr_time_allifos + lttbd] snr_slice = snr_series_dict[ifo].time_slice(time_window[0], time_window[1]) max_snr_idx = numpy.argmax(abs(snr_slice)) idx_offset = snr_slice.start_time - snr_series_dict[ifo].start_time idx_offset = int(float(idx_offset) * sample_rate + 0.5) max_snr_idx = max_snr_idx + idx_offset - new_snr_slice = slice(max_snr_idx - int(sample_rate/10), - max_snr_idx + int(sample_rate/10)+1) + new_snr_slice = slice( + max_snr_idx - int(sample_rate / 10), max_snr_idx + int(sample_rate / 10) + 1 + ) snr_series_dict[ifo] = snr_series_dict[ifo][new_snr_slice] netsnr = 0 for idx, ifo in enumerate(ifos): - coinc_results['foreground/'+ifo+'/mass1'] = mass1 - coinc_results['foreground/'+ifo+'/mass2'] = mass2 - coinc_results['foreground/'+ifo+'/spin1z'] = spin1z - coinc_results['foreground/'+ifo+'/spin2z'] = spin2z - coinc_results['foreground/'+ifo+'/f_lower'] = flow - coinc_results['foreground/'+ifo+'/window'] = 0.1 - coinc_results['foreground/'+ifo+'/sample_rate'] = int(sample_rate) - coinc_results['foreground/'+ifo+'/template_id'] = 0 + coinc_results["foreground/" + ifo + "/mass1"] = mass1 + coinc_results["foreground/" + ifo + "/mass2"] = mass2 + coinc_results["foreground/" + ifo + "/spin1z"] = spin1z + coinc_results["foreground/" + ifo + "/spin2z"] = spin2z + coinc_results["foreground/" + ifo + "/f_lower"] = flow + coinc_results["foreground/" + ifo + "/window"] = 0.1 + coinc_results["foreground/" + ifo + "/sample_rate"] = int(sample_rate) + coinc_results["foreground/" + ifo + "/template_id"] = 0 # Apparently GraceDB gets upset if chi-squared is not set - coinc_results['foreground/'+ifo+'/chisq'] = 1. - coinc_results['foreground/'+ifo+'/chisq_dof'] = 1 - coinc_results['foreground/'+ifo+'/template_duration'] = \ - fp['template_duration'][()] + coinc_results["foreground/" + ifo + "/chisq"] = 1.0 + coinc_results["foreground/" + ifo + "/chisq_dof"] = 1 + coinc_results["foreground/" + ifo + "/template_duration"] = fp["template_duration"][ + () + ] skyloc_data[ifo] = {} - skyloc_data[ifo]['psd'] = interpolate(data[ifo].psd, 0.25) - skyloc_data[ifo]['snr_series'] = snr_series_dict[ifo] + skyloc_data[ifo]["psd"] = interpolate(data[ifo].psd, 0.25) + skyloc_data[ifo]["snr_series"] = snr_series_dict[ifo] # Find loudest SNR within coincidence window of the largest peak snr_peak = abs(loudest_snrs[ifo]) @@ -291,89 +338,88 @@ for idx, ifo in enumerate(ifos): # Add small buffer lttbd += 0.005 - time_window = [loudest_snr_time_allifos - lttbd, - loudest_snr_time_allifos + lttbd] + time_window = [loudest_snr_time_allifos - lttbd, loudest_snr_time_allifos + lttbd] snr_slice = snr_series_dict[ifo].time_slice(time_window[0], time_window[1]) max_snr_idx = numpy.argmax(abs(snr_slice)) loudest_snr_time = snr_slice.sample_times[max_snr_idx] loudest_snr = snr_slice[max_snr_idx] - coinc_results['foreground/'+ifo+'/end_time'] = loudest_snr_time - coinc_results['foreground/'+ifo+'/snr_series'] = snr_series_dict[ifo] - coinc_results['foreground/'+ifo+'/psd_series'] = data[ifo].psd - coinc_results['foreground/'+ifo+'/delta_f'] = delta_f - coinc_results['foreground/'+ifo+'/snr'] = abs(loudest_snr) - netsnr += abs(loudest_snr)**2 - coinc_results['foreground/'+ifo+'/sigmasq'] \ - = snr_series_dict['sigmasq_' + ifo] - coinc_results['foreground/'+ifo+'/coa_phase'] \ - = numpy.angle(loudest_snr) + coinc_results["foreground/" + ifo + "/end_time"] = loudest_snr_time + coinc_results["foreground/" + ifo + "/snr_series"] = snr_series_dict[ifo] + coinc_results["foreground/" + ifo + "/psd_series"] = data[ifo].psd + coinc_results["foreground/" + ifo + "/delta_f"] = delta_f + coinc_results["foreground/" + ifo + "/snr"] = abs(loudest_snr) + netsnr += abs(loudest_snr) ** 2 + coinc_results["foreground/" + ifo + "/sigmasq"] = snr_series_dict["sigmasq_" + ifo] + coinc_results["foreground/" + ifo + "/coa_phase"] = numpy.angle(loudest_snr) -coinc_results['foreground/stat'] = numpy.sqrt(netsnr) -coinc_results['foreground/ifar'] = fp['ifar'][()] +coinc_results["foreground/stat"] = numpy.sqrt(netsnr) +coinc_results["foreground/ifar"] = fp["ifar"][()] channel_names = {} -for ifo in fp['channel_names'].keys(): - channel_names[ifo] = fp['channel_names'][ifo].asstr()[()] +for ifo in fp["channel_names"].keys(): + channel_names[ifo] = fp["channel_names"][ifo].asstr()[()] -mc_area_args = load_hdf5_to_dict(fp, 'mc_area_args/') +mc_area_args = load_hdf5_to_dict(fp, "mc_area_args/") -kwargs = {'psds': {ifo: skyloc_data[ifo]['psd'] for ifo in ifos}, - 'low_frequency_cutoff': flow, - 'skyloc_data': skyloc_data, - 'channel_names': channel_names, - 'mc_area_args': mc_area_args} +kwargs = { + "psds": {ifo: skyloc_data[ifo]["psd"] for ifo in ifos}, + "low_frequency_cutoff": flow, + "skyloc_data": skyloc_data, + "channel_names": channel_names, + "mc_area_args": mc_area_args, +} # Do not recalculate p_terr/p_astro, reuse previous value if available # (source probabilities may be recalculated) -if 'p_terr' in fp: - kwargs['p_terr'] = fp['p_terr'][()] +if "p_terr" in fp: + kwargs["p_terr"] = fp["p_terr"][()] # Treat all ifos as having triggers -doc = CandidateForGraceDB( - ifos, - ifos, - coinc_results, - upload_snr_series=True, - **kwargs -) +doc = CandidateForGraceDB(ifos, ifos, coinc_results, upload_snr_series=True, **kwargs) xml_path = os.path.join( - args.output_path, - f'coinc-{loudest_snr_time_allifos:.3f}.xml.gz' + args.output_path, f"coinc-{loudest_snr_time_allifos:.3f}.xml.gz" ) if args.enable_gracedb_upload: - logging.info('Uploading optimized candidate to GraceDB') - - comment = ('Automatic PyCBC followup of trigger ' - '{0} to find the template ' - 'parameters that maximize the SNR. The FAR of this trigger is ' - 'copied from {0} and does not reflect this trigger\'s template ' - 'parameters.') + logging.info("Uploading optimized candidate to GraceDB") + + comment = ( + "Automatic PyCBC followup of trigger " + '{0} to find the template ' + "parameters that maximize the SNR. The FAR of this trigger is " + "copied from {0} and does not reflect this trigger's template " + "parameters." + ) comment = comment.format(original_gid) - gid = doc.upload(xml_path, gracedb_server=args.gracedb_server, - testing=(not args.production), extra_strings=[comment], - search=args.gracedb_search, labels=args.gracedb_labels) + gid = doc.upload( + xml_path, + gracedb_server=args.gracedb_server, + testing=(not args.production), + extra_strings=[comment], + search=args.gracedb_search, + labels=args.gracedb_labels, + ) if gid is not None: - logging.info('Event uploaded as %s', gid) + logging.info("Event uploaded as %s", gid) # add a note to the original G event pointing to the optimized one from ligo.gracedb.rest import GraceDb - gracedb = GraceDb(args.gracedb_server) \ - if args.gracedb_server is not None else GraceDb() - comment = ('Result of SNR maximization uploaded as ' - '{0}').format(gid) - gracedb.write_log( - original_gid, - comment, - tag_name=['analyst_comments'] + gracedb = ( + GraceDb(args.gracedb_server) + if args.gracedb_server is not None + else GraceDb() + ) + comment = ( + f'Result of SNR maximization uploaded as {gid}' ) + gracedb.write_log(original_gid, comment, tag_name=["analyst_comments"]) else: - logging.info('Saving optimized candidate') + logging.info("Saving optimized candidate") doc.save(xml_path) -logging.info('Done') +logging.info("Done") diff --git a/bin/pycbc_process_sngls b/bin/pycbc_process_sngls index 29b41f4c0ff..24600f7f9b4 100644 --- a/bin/pycbc_process_sngls +++ b/bin/pycbc_process_sngls @@ -4,41 +4,57 @@ import argparse import logging -import numpy + import h5py +import numpy import pycbc -from pycbc.io import SingleDetTriggers, HFile -from pycbc.events import stat, coinc, veto - +from pycbc.events import coinc, stat, veto +from pycbc.io import HFile, SingleDetTriggers parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--single-trig-file', required=True, - help='Path to file containing single-detector triggers in ' - 'HDF5 format. Required') -parser.add_argument('--detector', required=True, - help='Detector. Required') -parser.add_argument('--bank-file', required=True, - help='Path to file containing template bank in HDF5 format' - '. Required') -parser.add_argument('--veto-file', - help='Optional path to file containing veto segments') -parser.add_argument('--segment-name', default=None, - help='Optional, name of segment list to use for vetoes') -parser.add_argument('--filter-string', default=None, - help='Optional, boolean expression for filtering triggers ' - 'e.g. "self.mchirp>5."') -parser.add_argument('--min-snr', default=0., type=float, - help='Only keep triggers above the given SNR') -parser.add_argument('--cluster-window', type=float, - help='If supplied, cluster singles by symmetrical time ' - 'window method, specify window extent from maximum' - 'in seconds') -parser.add_argument('--store-bank-values', default=False, action='store_true', - help='If given also add the template bank parameters into ' - 'the output file.') -parser.add_argument('--output-file', required=True) +parser.add_argument( + "--single-trig-file", + required=True, + help="Path to file containing single-detector triggers in HDF5 format. Required", +) +parser.add_argument("--detector", required=True, help="Detector. Required") +parser.add_argument( + "--bank-file", + required=True, + help="Path to file containing template bank in HDF5 format. Required", +) +parser.add_argument( + "--veto-file", help="Optional path to file containing veto segments" +) +parser.add_argument( + "--segment-name", + default=None, + help="Optional, name of segment list to use for vetoes", +) +parser.add_argument( + "--filter-string", + default=None, + help='Optional, boolean expression for filtering triggers e.g. "self.mchirp>5."', +) +parser.add_argument( + "--min-snr", default=0.0, type=float, help="Only keep triggers above the given SNR" +) +parser.add_argument( + "--cluster-window", + type=float, + help="If supplied, cluster singles by symmetrical time " + "window method, specify window extent from maximum" + "in seconds", +) +parser.add_argument( + "--store-bank-values", + default=False, + action="store_true", + help="If given also add the template bank parameters into the output file.", +) +parser.add_argument("--output-file", required=True) stat.insert_statistic_option_group(parser) args = parser.parse_args() @@ -46,18 +62,18 @@ args = parser.parse_args() pycbc.init_logging(args.verbose) # Munge together SNR cut and any other filter specified -snr_filter = 'self.snr>%f' % (args.min_snr) if args.min_snr > 0. else None +snr_filter = "self.snr>%f" % (args.min_snr) if args.min_snr > 0.0 else None filts = [f for f in [snr_filter, args.filter_string] if f is not None] if len(filts) == 2: # both an explicit filter and a min-snr # io.hdf uses numpy imported as np - filter_func = 'np.logical_and(%s, %s)' % (filts[0], filts[1]) + filter_func = "np.logical_and(%s, %s)" % (filts[0], filts[1]) elif len(filts) == 1: filter_func = filts[0] else: filter_func = None if filter_func is not None: - logging.info('Will filter trigs using %s', filter_func) + logging.info("Will filter trigs using %s", filter_func) # Filter will be stored as self.mask attribute of sngls instance sngls = SingleDetTriggers( args.single_trig_file, @@ -68,7 +84,7 @@ sngls = SingleDetTriggers( filter_func=filter_func, ) -logging.info('Calculating stat') +logging.info("Calculating stat") rank_method = stat.get_statistic_from_opts(args, [args.detector]) # NOTE: inefficient, as we are calculating the stat on all # triggers. Might need to do something complicated to fix this. @@ -76,29 +92,31 @@ rank_method = stat.get_statistic_from_opts(args, [args.detector]) sngl_info = ([args.detector], sngls.trigs) stat = rank_method.rank_stat_single(sngl_info)[sngls.mask] -logging.info('%i stat values found', len(stat)) +logging.info("%i stat values found", len(stat)) -outfile = HFile(args.output_file, 'w') +outfile = HFile(args.output_file, "w") outgroup = outfile.create_group(args.detector) if args.cluster_window is not None: - logging.info('Clustering events over %s s window', args.cluster_window) - out_idx = coinc.cluster_over_time(stat, sngls.end_time, - window=args.cluster_window) - logging.info('%d triggers remaining', len(out_idx)) - outgroup['cluster_window'] = args.cluster_window + logging.info("Clustering events over %s s window", args.cluster_window) + out_idx = coinc.cluster_over_time(stat, sngls.end_time, window=args.cluster_window) + logging.info("%d triggers remaining", len(out_idx)) + outgroup["cluster_window"] = args.cluster_window else: out_idx = numpy.arange(len(sngls.end_time)) -logging.info('Writing %i triggers', len(out_idx)) +logging.info("Writing %i triggers", len(out_idx)) # get the columns to copy over -with HFile(args.single_trig_file, 'r') as trigfile: +with HFile(args.single_trig_file, "r") as trigfile: cnames = [] # only keep datasets parallel to the original trigger list for n, col in trigfile[args.detector].items(): - if n.endswith('_template') or isinstance(col, h5py.Group) \ - or n == u'template_boundaries': + if ( + n.endswith("_template") + or isinstance(col, h5py.Group) + or n == "template_boundaries" + ): continue cnames.append(n) for n in cnames: @@ -106,13 +124,13 @@ for n in cnames: if args.store_bank_values: for n in sngls.bank: - if n == 'template_hash': + if n == "template_hash": continue if not hasattr(sngls, n): logging.warning( "Bank's %s dataset or group is not supported " "by SingleDetTriggers, ignoring it", - n + n, ) continue # don't repeat things that already came from the trigger file @@ -124,29 +142,31 @@ if args.store_bank_values: # copy the live time segments to enable the calculation of trigger rates. # If a veto file has been used, subtract the vetoed time from the live time -live_segs = veto.start_end_to_segments(sngls.trigs['search/start_time'][:], - sngls.trigs['search/end_time'][:]) +live_segs = veto.start_end_to_segments( + sngls.trigs["search/start_time"][:], sngls.trigs["search/end_time"][:] +) live_segs.coalesce() if args.veto_file is not None: - veto_segs = veto.select_segments_by_definer(args.veto_file, - args.segment_name, - args.detector) + veto_segs = veto.select_segments_by_definer( + args.veto_file, args.segment_name, args.detector + ) veto_segs.coalesce() live_segs -= veto_segs -outgroup['search/start_time'], outgroup['search/end_time'] = \ - veto.segments_to_start_end(live_segs) -outgroup['search'].attrs['live_time'] = abs(live_segs) +outgroup["search/start_time"], outgroup["search/end_time"] = veto.segments_to_start_end( + live_segs +) +outgroup["search"].attrs["live_time"] = abs(live_segs) # cannot store None in a h5py attr -outgroup.attrs['filter'] = filter_func or 'None' -outgroup.attrs['cluster_window'] = args.cluster_window or 'None' +outgroup.attrs["filter"] = filter_func or "None" +outgroup.attrs["cluster_window"] = args.cluster_window or "None" -outgroup['stat'] = stat[out_idx] -outgroup.attrs['ranking_statistic'] = args.ranking_statistic -outgroup.attrs['sngl_ranking'] = args.sngl_ranking -outgroup.attrs['statistic_files'] = args.statistic_files +outgroup["stat"] = stat[out_idx] +outgroup.attrs["ranking_statistic"] = args.ranking_statistic +outgroup.attrs["sngl_ranking"] = args.sngl_ranking +outgroup.attrs["statistic_files"] = args.statistic_files outfile.close() -logging.info('Done!') +logging.info("Done!") diff --git a/bin/pycbc_randomize_inj_dist_by_optsnr b/bin/pycbc_randomize_inj_dist_by_optsnr index d5f0fbd71f5..8d1b4a8ee69 100644 --- a/bin/pycbc_randomize_inj_dist_by_optsnr +++ b/bin/pycbc_randomize_inj_dist_by_optsnr @@ -1,97 +1,133 @@ #! /usr/bin/env python -__prog__ = 'pycbc_randr_by_snr' -__author__ = 'Collin Capano , Miriam Cabero Mueller ' -__description__ = 'Resets the distance distribution in a sim_inspiral table based on desired SNR. The SNRs are chosen randomly from a given range.' +__prog__ = "pycbc_randr_by_snr" +__author__ = "Collin Capano , Miriam Cabero Mueller " +__description__ = "Resets the distance distribution in a sim_inspiral table based on desired SNR. The SNRs are chosen randomly from a given range." -import numpy import sys import time -from scipy import stats, special from argparse import ArgumentParser +import numpy from igwn_ligolw import lsctables from igwn_ligolw import utils as ligolw_utils +from scipy import special, stats import pycbc from pycbc.io.ligolw import LIGOLWContentHandler def r_uniform_in_volume(r1, r2, N): - xi = numpy.random.uniform(0., 1., size=N) - return (xi*r2**3. + (1.-xi)*r1**3.)**(1./3) + xi = numpy.random.uniform(0.0, 1.0, size=N) + return (xi * r2**3.0 + (1.0 - xi) * r1**3.0) ** (1.0 / 3) + def r_uniform_in_distance(r1, r2, N): return numpy.random.uniform(r1, r2, size=N) + def r_uniform_in_logdist(r1, r2, N): - return r1*(r2/r1)**numpy.random.uniform(0., 1., size=N) + return r1 * (r2 / r1) ** numpy.random.uniform(0.0, 1.0, size=N) + def r_beta_distribution(rmax, alpha, beta, N=1): - return rmax*stats.beta.rvs(alpha, beta, size=N) + return rmax * stats.beta.rvs(alpha, beta, size=N) + def beta_weight(r, rmax, alpha, beta): - x = r/rmax - return 4.*numpy.pi * rmax**3. * x**(3.-alpha) * (1.-x)**(1.-beta) * \ - special.beta(alpha, beta) - + x = r / rmax + return ( + 4.0 + * numpy.pi + * rmax**3.0 + * x ** (3.0 - alpha) + * (1.0 - x) ** (1.0 - beta) + * special.beta(alpha, beta) + ) + + def get_weights(distribution, r, r1, r2): if distribution == "volume": - min_vol = (4./3)*numpy.pi*r1**3. - weight = (4./3)*numpy.pi*(r2**3. - r1**3.) + min_vol = (4.0 / 3) * numpy.pi * r1**3.0 + weight = (4.0 / 3) * numpy.pi * (r2**3.0 - r1**3.0) elif distribution == "distance": - min_vol = (4./3)*numpy.pi*r1**3. - weight = 4.*numpy.pi*(r2-r1) * r**2. + min_vol = (4.0 / 3) * numpy.pi * r1**3.0 + weight = 4.0 * numpy.pi * (r2 - r1) * r**2.0 elif distribution == "logdist": - min_vol = (4./3)*numpy.pi*r1**3. - weight = 4.*numpy.pi * r**3. * numpy.log(r2/r1) + min_vol = (4.0 / 3) * numpy.pi * r1**3.0 + weight = 4.0 * numpy.pi * r**3.0 * numpy.log(r2 / r1) else: - raise ValueError("unrecognized distribution %s" %(distribution)) + raise ValueError("unrecognized distribution %s" % (distribution)) return min_vol, weight -parser = ArgumentParser(description = __doc__) + +parser = ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--xml-file', required = True, - help = 'Input LIGOLW file defining injections.') -parser.add_argument('--output-file', required=True, - help='Output LIGOLW file.') -parser.add_argument('--snr-columns', nargs='+', required=True, - help='Defines the columns of the sim_inspiral table that ' \ - 'store the optimal SNR for each detector. If the optimal ' \ - 'SNR of an injection is 0 for a detector, that detector ' \ - 'is assumed to be down during the time of the injection. ' \ - 'If all the optimal SNRs are 0 for an injection, ' \ - 'that injection is skipped.') -parser.add_argument('--min-snr', type = float, - help = 'Set the minimum single-detector SNR to use. ' \ - 'This is multipled by N**(1/2), where N is the number ' \ - 'of non-zero snr-columns for each injection') -parser.add_argument('--max-snr', type = float, - help = 'Set the maximum SNR to use; must be larger than min-snr. ' \ - 'To use an exact snr, set to the same value as min-snr. ' \ - 'This is multipled by N**(1/2), where N is the number ' \ - 'of non-zero snr-columns for each injection.') -parser.add_argument('--fixed-distance-range', metavar='MIN_D,MAX_D', - help='Instead of using the SNR to determine the limits ' \ - 'from which to draw the distance, use the given range for all injections.') -parser.add_argument('--distribution', default='distance', metavar="DISTRIBUTION[:PARAMETERS]", - help='What distribution to make the injections uniform in. ' \ - 'Options are "distance", "volume", "logdist", or "beta:${a},${b}". ' \ - 'If "distance", the distance distribution of the injections ' \ - 'will be uniform in distance. If "volume", uniform in volume. ' \ - 'If "logdist", uniform in the log of the distance. If "beta", ' \ - 'distances will be drawn from the beta distribution ' \ - 'p(r; a,b,rmax) = 1/(rmax*B(a,b)) * (r/rmax)**(1-a) * (1-r/rmax)**(b-1), ' \ - 'where B is the beta function. Must provide an a and a b; ' \ - 'rmax is taken --max-snr, or the upper bound of the fixed ' \ - 'distance range. Default is %(default)s.') -parser.add_argument('--scale-by-chirp-dist', action='store_true', - help='Scale the distance by the chirp distance relative to ' \ - 'a 1.4/1.4 binary. Only can use if fixed-distance-range is on.') -parser.add_argument('--seed', type = int, default = int((time.time()*100)% 1e6), - help = 'Set the seed to use for the random number generator. ' \ - 'If none specified, will use the current time.') +parser.add_argument( + "--xml-file", required=True, help="Input LIGOLW file defining injections." +) +parser.add_argument("--output-file", required=True, help="Output LIGOLW file.") +parser.add_argument( + "--snr-columns", + nargs="+", + required=True, + help="Defines the columns of the sim_inspiral table that " + "store the optimal SNR for each detector. If the optimal " + "SNR of an injection is 0 for a detector, that detector " + "is assumed to be down during the time of the injection. " + "If all the optimal SNRs are 0 for an injection, " + "that injection is skipped.", +) +parser.add_argument( + "--min-snr", + type=float, + help="Set the minimum single-detector SNR to use. " + "This is multipled by N**(1/2), where N is the number " + "of non-zero snr-columns for each injection", +) +parser.add_argument( + "--max-snr", + type=float, + help="Set the maximum SNR to use; must be larger than min-snr. " + "To use an exact snr, set to the same value as min-snr. " + "This is multipled by N**(1/2), where N is the number " + "of non-zero snr-columns for each injection.", +) +parser.add_argument( + "--fixed-distance-range", + metavar="MIN_D,MAX_D", + help="Instead of using the SNR to determine the limits " + "from which to draw the distance, use the given range for all injections.", +) +parser.add_argument( + "--distribution", + default="distance", + metavar="DISTRIBUTION[:PARAMETERS]", + help="What distribution to make the injections uniform in. " + 'Options are "distance", "volume", "logdist", or "beta:${a},${b}". ' + 'If "distance", the distance distribution of the injections ' + 'will be uniform in distance. If "volume", uniform in volume. ' + 'If "logdist", uniform in the log of the distance. If "beta", ' + "distances will be drawn from the beta distribution " + "p(r; a,b,rmax) = 1/(rmax*B(a,b)) * (r/rmax)**(1-a) * (1-r/rmax)**(b-1), " + "where B is the beta function. Must provide an a and a b; " + "rmax is taken --max-snr, or the upper bound of the fixed " + "distance range. Default is %(default)s.", +) +parser.add_argument( + "--scale-by-chirp-dist", + action="store_true", + help="Scale the distance by the chirp distance relative to " + "a 1.4/1.4 binary. Only can use if fixed-distance-range is on.", +) +parser.add_argument( + "--seed", + type=int, + default=int((time.time() * 100) % 1e6), + help="Set the seed to use for the random number generator. " + "If none specified, will use the current time.", +) opts = parser.parse_args() @@ -99,43 +135,53 @@ pycbc.init_logging(opts.verbose) # parse the distributions # beta is the only special one -if opts.distribution in ['volume', 'distance', 'logdist']: +if opts.distribution in ["volume", "distance", "logdist"]: distribution = opts.distribution -elif opts.distribution.startswith('beta'): +elif opts.distribution.startswith("beta"): try: - distribution, params = opts.distribution.split(':') + distribution, params = opts.distribution.split(":") except ValueError: - raise ValueError("beta distribution requires an a and b to be " + - "provided; e.g., 'beta:3,2'; see help") + raise ValueError( + "beta distribution requires an a and b to be " + "provided; e.g., 'beta:3,2'; see help" + ) try: - alpha, beta = map(float, params.split(',')) + alpha, beta = map(float, params.split(",")) except ValueError: raise ValueError("beta distribution not formatted correctly; see help") else: raise ValueError("unrecognized distribution; see help") # parse ranges -if opts.fixed_distance_range is not None and (opts.min_snr is not None or \ - opts.max_snr is not None): - raise ValueError("Please specify a fixed-distance-range *or* a min/max " + - "snr, not both.") +if opts.fixed_distance_range is not None and ( + opts.min_snr is not None or opts.max_snr is not None +): + raise ValueError( + "Please specify a fixed-distance-range *or* a min/max " + "snr, not both." + ) if opts.fixed_distance_range is not None: - min_dist, max_dist = map(float, opts.fixed_distance_range.split(',')) - if distribution == 'beta' and min_dist != 0: - raise ValueError("the beta distribution requires the minimum "+ - "distance to be set to 0.") + min_dist, max_dist = map(float, opts.fixed_distance_range.split(",")) + if distribution == "beta" and min_dist != 0: + raise ValueError( + "the beta distribution requires the minimum " + "distance to be set to 0." + ) elif opts.scale_by_chirp_dist: - raise ValueError("cannot use scale-by-chirp-dist if not using a " +\ - "fixed-distance-range (scaling by chirp distance doesn't make " +\ - "sense when using min/max SNR") + raise ValueError( + "cannot use scale-by-chirp-dist if not using a " + "fixed-distance-range (scaling by chirp distance doesn't make " + "sense when using min/max SNR" + ) elif opts.min_snr is None or opts.max_snr is None: if opts.min_snr is None or opts.max_snr is None: if opts.min_snr is None: - raise ValueError("must provide a min SNR if not specifying a " + - "fixed-distance-range") - if opts.max_snr is None and distribution != 'beta': - raise ValueError("must provide a maximum SNR if not specifying a " + - "fixed-distance-range and not using a beta distribution") + raise ValueError( + "must provide a min SNR if not specifying a " + "fixed-distance-range" + ) + if opts.max_snr is None and distribution != "beta": + raise ValueError( + "must provide a maximum SNR if not specifying a " + "fixed-distance-range and not using a beta distribution" + ) numpy.random.seed(opts.seed) @@ -146,31 +192,32 @@ if opts.verbose: # Read in input file xmldoc = ligolw_utils.load_filename( - opts.xml_file, verbose=True, contenthandler=LIGOLWContentHandler) + opts.xml_file, verbose=True, contenthandler=LIGOLWContentHandler +) tabletype = lsctables.SimInspiralTable injections = tabletype.get_table(xmldoc) for ii, inj in enumerate(injections): if opts.verbose: - print("Injection %i\r" %(ii+1), end=' ', file=sys.stdout) + print("Injection %i\r" % (ii + 1), end=" ", file=sys.stdout) sys.stdout.flush() - + # calculate the sigmas in each ifo - sigmas = numpy.array([getattr(inj,col) for col in opts.snr_columns]) * inj.distance + sigmas = numpy.array([getattr(inj, col) for col in opts.snr_columns]) * inj.distance nzidx = sigmas.nonzero()[0] if nzidx.size == 0: continue sigmas = sigmas[nzidx] # calculate the sigmas as the quadrature sum of the sigmas in all # of the ifos - sigma = numpy.sqrt((sigmas**2.).sum()) + sigma = numpy.sqrt((sigmas**2.0).sum()) # get the distance ranges based on the optimal SNR if opts.fixed_distance_range is None: if opts.max_snr is not None: - min_dist = sigma / (numpy.sqrt(sigmas.size)*opts.max_snr) - max_dist = sigma / (numpy.sqrt(sigmas.size)*opts.min_snr) + min_dist = sigma / (numpy.sqrt(sigmas.size) * opts.max_snr) + max_dist = sigma / (numpy.sqrt(sigmas.size) * opts.min_snr) # get a distance to use, and the weights if distribution == "volume": @@ -182,14 +229,15 @@ for ii, inj in enumerate(injections): elif distribution == "beta": distance = r_beta_distribution(max_dist, alpha, beta, 1)[0] else: - raise ValueError("unrecognized distribution %s; " %( - distribution), "see --help for options") + raise ValueError( + "unrecognized distribution %s; " % (distribution), "see --help for options" + ) # scale by chirp distance if desired if opts.scale_by_chirp_dist: - scale_fac = (inj.mchirp/(2.8 * 0.25**0.6))**(5./6) + scale_fac = (inj.mchirp / (2.8 * 0.25**0.6)) ** (5.0 / 6) else: - scale_fac = 1. + scale_fac = 1.0 # save inj.distance = distance * scale_fac @@ -197,14 +245,13 @@ for ii, inj in enumerate(injections): inj.alpha6 = max_dist * scale_fac inj.numrel_data = opts.distribution for kk in range(nzidx.size): - setattr(inj, opts.snr_columns[nzidx[kk]], sigmas[kk]/inj.distance) + setattr(inj, opts.snr_columns[nzidx[kk]], sigmas[kk] / inj.distance) if opts.verbose: - print("", file=sys.stdout) + print(file=sys.stdout) sys.stdout.flush() -ligolw_utils.write_filename(xmldoc, opts.output_file, compress='auto') +ligolw_utils.write_filename(xmldoc, opts.output_file, compress="auto") if opts.verbose: print("Finished!", file=sys.stdout) sys.exit(0) - diff --git a/bin/pycbc_single_template b/bin/pycbc_single_template index 45c6623ef59..86a3cb39f30 100755 --- a/bin/pycbc_single_template +++ b/bin/pycbc_single_template @@ -14,23 +14,25 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Calculate the SNR and CHISQ timeseries for either a chosen template, or +""" +Calculate the SNR and CHISQ timeseries for either a chosen template, or a specific Nth loudest coincident event. """ -import sys -import logging + import argparse +import logging +import sys + import numpy import pycbc -from pycbc import vetoes, psd, waveform, strain, scheme, fft, filter -from pycbc.io import WaveformArray, HFile -from pycbc import events -from pycbc.filter import resample_to_delta_t -from pycbc.types import zeros, complex64 -from pycbc.types import complex_same_precision_as -from pycbc.detector import Detector import pycbc.waveform.utils +from pycbc import events, fft, filter, psd, scheme, strain, vetoes, waveform +from pycbc.detector import Detector +from pycbc.filter import resample_to_delta_t +from pycbc.io import HFile, WaveformArray +from pycbc.types import complex64, complex_same_precision_as, zeros + def subtract_template(stilde, template, snr, trigger_time, flow): idx = int((trigger_time - snr.start_time) / snr.delta_t) @@ -41,6 +43,7 @@ def subtract_template(stilde, template, snr, trigger_time, flow): stilde -= pycbc.waveform.utils.apply_fseries_time_shift(inverse, dt) return stilde + def select_segments(fname, anal_name, data_name, ifo, time, pad_data): anal_segs = events.select_segments_by_definer(fname, anal_name, ifo) data_segs = events.select_segments_by_definer(fname, data_name, ifo) @@ -48,7 +51,7 @@ def select_segments(fname, anal_name, data_name, ifo, time, pad_data): # Anal segs should be disjoint, so first find the seg containing time s = numpy.array([t[0] for t in anal_segs]) e = numpy.array([t[1] for t in anal_segs]) - #ensure sorted + # ensure sorted sorting = s.argsort() s = s[sorting] e = e[sorting] @@ -66,7 +69,7 @@ def select_segments(fname, anal_name, data_name, ifo, time, pad_data): s2 = s2[lgc] e2 = e2[lgc] if len(s2) == 0: - err_msg = "Cannot find a data segment within %s" %(str(time)) + err_msg = "Cannot find a data segment within %s" % (str(time)) raise ValueError(err_msg) if len(s2) == 1: data_time = (s2[0], e2[0]) @@ -88,98 +91,182 @@ def select_segments(fname, anal_name, data_name, ifo, time, pad_data): return anal_time, data_time -parser = argparse.ArgumentParser(usage='', - description="Single template gravitational-wave followup") + +parser = argparse.ArgumentParser( + usage="", description="Single template gravitational-wave followup" +) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--output-file', required=True) -parser.add_argument('--subtract-template', action='store_true') -parser.add_argument("--low-frequency-cutoff", type=float, - help="The low frequency cutoff to use for filtering (Hz)") -parser.add_argument("--high-frequency-cutoff", type=float, - help="The high frequency cutoff to use for filtering (Hz)") -parser.add_argument("--chisq-bins", default="0", type=str, help= - "Number of frequency bins to use for power chisq.") -parser.add_argument("--minimum-chisq-bins", default=0, type=int, help= - "If the chisq bin formula fails, default to this number.") -parser.add_argument("--trigger-time", type=float, default=None, - help="Time used as centre point for the time series " - "calculated. Required with --window option") -parser.add_argument("--use-params-of-closest-injection", action="store_true", - default=False, - help="If given, use the injection with end_time closest to " - "--trigger-time as the waveform for filtering. If " - "using this do not supply mass and spin options. " - "Using this requires trigger-time and injection-file " - "to be given.") +parser.add_argument("--output-file", required=True) +parser.add_argument("--subtract-template", action="store_true") +parser.add_argument( + "--low-frequency-cutoff", + type=float, + help="The low frequency cutoff to use for filtering (Hz)", +) +parser.add_argument( + "--high-frequency-cutoff", + type=float, + help="The high frequency cutoff to use for filtering (Hz)", +) +parser.add_argument( + "--chisq-bins", + default="0", + type=str, + help="Number of frequency bins to use for power chisq.", +) +parser.add_argument( + "--minimum-chisq-bins", + default=0, + type=int, + help="If the chisq bin formula fails, default to this number.", +) +parser.add_argument( + "--trigger-time", + type=float, + default=None, + help="Time used as centre point for the time series " + "calculated. Required with --window option", +) +parser.add_argument( + "--use-params-of-closest-injection", + action="store_true", + default=False, + help="If given, use the injection with end_time closest to " + "--trigger-time as the waveform for filtering. If " + "using this do not supply mass and spin options. " + "Using this requires trigger-time and injection-file " + "to be given.", +) # add approximant arg -waveform.bank.add_approximant_arg(parser, - help="The name of the approximant to use for filtering. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--mass1", type=float, - help="The mass of the first component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--mass2", type=float, - help="The mass of the second component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--spin1z", type=float, default=0, - help="The aligned spin of the first component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--spin2z", type=float, default=0, - help="The aligned pin of the second component object. " - "Do not use if using --use-params-of-closest-injection.") -parser.add_argument("--template-start-frequency", type=float, default=None, - help="If given, use this as a start frequency for " - "generating the template. If not given the " - "--low-frequency-cutoff is used.") +waveform.bank.add_approximant_arg( + parser, + help="The name of the approximant to use for filtering. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--mass1", + type=float, + help="The mass of the first component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--mass2", + type=float, + help="The mass of the second component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--spin1z", + type=float, + default=0, + help="The aligned spin of the first component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--spin2z", + type=float, + default=0, + help="The aligned pin of the second component object. " + "Do not use if using --use-params-of-closest-injection.", +) +parser.add_argument( + "--template-start-frequency", + type=float, + default=None, + help="If given, use this as a start frequency for " + "generating the template. If not given the " + "--low-frequency-cutoff is used.", +) # Optional arguments for precessing templates -parser.add_argument("--spin1x", type=float, default=0, - help="The non-aligned spin of the first component object. " - "Default = 0") -parser.add_argument("--spin2x", type=float, default=0, - help="The non-aligned spin of the second component object. " - "Default = 0") -parser.add_argument("--spin1y", type=float, default=0, - help="The non-aligned spin of the first component object. " - "Default = 0") -parser.add_argument("--spin2y", type=float, default=0, - help="The non-aligned spin of the second component object. " - "Default = 0") -parser.add_argument("--inclination", type=float, default=0, - help="The inclination of the source w.r.t the observer. " - "Default = 0") -parser.add_argument("--coa-phase", type=float, default=0, - help="The orbital azimuth of the source w.r.t the observer. " - "Default = 0") -parser.add_argument("--u-val", type=float, default=None, - help="The ratio between hplus and hcross to use in the " - "template, according to h(t) = hplus * u_val + hcross. " - "If not given only hplus is used.") +parser.add_argument( + "--spin1x", + type=float, + default=0, + help="The non-aligned spin of the first component object. Default = 0", +) +parser.add_argument( + "--spin2x", + type=float, + default=0, + help="The non-aligned spin of the second component object. Default = 0", +) +parser.add_argument( + "--spin1y", + type=float, + default=0, + help="The non-aligned spin of the first component object. Default = 0", +) +parser.add_argument( + "--spin2y", + type=float, + default=0, + help="The non-aligned spin of the second component object. Default = 0", +) +parser.add_argument( + "--inclination", + type=float, + default=0, + help="The inclination of the source w.r.t the observer. Default = 0", +) +parser.add_argument( + "--coa-phase", + type=float, + default=0, + help="The orbital azimuth of the source w.r.t the observer. Default = 0", +) +parser.add_argument( + "--u-val", + type=float, + default=None, + help="The ratio between hplus and hcross to use in the " + "template, according to h(t) = hplus * u_val + hcross. " + "If not given only hplus is used.", +) # Optional arguments for eccentric templates -parser.add_argument("--eccentricity", type=float, default=0, - help="The orbital eccentricity. Default = 0") -parser.add_argument("--rel-anomaly", type=float, default=0, - help="The relativistic anomaly. Default = 0") -parser.add_argument("--window", type=float, - help="Time to save on each side of the given trigger time") -parser.add_argument("--order", type=int, - help="The integer half-PN order at which to generate" - " the approximant. Default is -1 which indicates to use" - " approximant defined default.", default=-1, - choices = numpy.arange(-1, 9, 1)) -parser.add_argument("--taper-template", choices=["start","end","startend"], - help="For time-domain approximants, taper the start and/or" - " end of the waveform before FFTing.") +parser.add_argument( + "--eccentricity", + type=float, + default=0, + help="The orbital eccentricity. Default = 0", +) +parser.add_argument( + "--rel-anomaly", type=float, default=0, help="The relativistic anomaly. Default = 0" +) +parser.add_argument( + "--window", type=float, help="Time to save on each side of the given trigger time" +) +parser.add_argument( + "--order", + type=int, + help="The integer half-PN order at which to generate" + " the approximant. Default is -1 which indicates to use" + " approximant defined default.", + default=-1, + choices=numpy.arange(-1, 9, 1), +) +parser.add_argument( + "--taper-template", + choices=["start", "end", "startend"], + help="For time-domain approximants, taper the start and/or" + " end of the waveform before FFTing.", +) # These options can be used to identify start/end times -parser.add_argument("--inspiral-segments", - help="XML file containing the inspiral analysis segments. " - "Only used with the --statmap-file option") -parser.add_argument("--data-read-name", - help="name of the segmentlist containing the data read in by each job " - "from the inspiral segment file") -parser.add_argument("--data-analyzed-name", - help="name of the segmentlist containing the data analysed by each job " - "from the inspiral segment file") +parser.add_argument( + "--inspiral-segments", + help="XML file containing the inspiral analysis segments. " + "Only used with the --statmap-file option", +) +parser.add_argument( + "--data-read-name", + help="name of the segmentlist containing the data read in by each job " + "from the inspiral segment file", +) +parser.add_argument( + "--data-analyzed-name", + help="name of the segmentlist containing the data analysed by each job " + "from the inspiral segment file", +) # Add options groups psd.insert_psd_option_group(parser) @@ -195,7 +282,7 @@ pycbc.init_logging(opt.verbose) if opt.window and opt.trigger_time is None: raise RuntimeError("Can't use --window option without a valid trigger time!") -f = HFile(opt.output_file, 'w') +f = HFile(opt.output_file, "w") ifo = opt.channel_name[0:2] # If we are choosing start/end times from XML file ############################ @@ -203,9 +290,14 @@ if opt.inspiral_segments: # Important for trig-start/end and data-start/end to match inspiral jobs. # Zero-padding also will not zero-pad unless explicitly told to in the # trig start/end times - anal_seg, data_seg = select_segments(opt.inspiral_segments, - opt.data_analyzed_name, opt.data_read_name, - ifo, opt.trigger_time, opt.pad_data) + anal_seg, data_seg = select_segments( + opt.inspiral_segments, + opt.data_analyzed_name, + opt.data_read_name, + ifo, + opt.trigger_time, + opt.pad_data, + ) opt.trig_start_time = anal_seg[0] opt.trig_end_time = anal_seg[1] opt.gps_start_time = data_seg[0] + opt.pad_data @@ -218,7 +310,7 @@ psd.verify_psd_options(opt, parser) strain.verify_strain_options(opt, parser) strain.StrainSegments.verify_segment_options(opt, parser) scheme.verify_processing_options(opt, parser) -fft.verify_fft_options(opt,parser) +fft.verify_fft_options(opt, parser) ctx = scheme.from_cli(opt) gwstrain = strain.from_cli(opt, pycbc.DYN_RANGE_FAC) @@ -226,16 +318,17 @@ strain_segments = strain.StrainSegments.from_cli(opt, gwstrain) if not opt.use_params_of_closest_injection: row = WaveformArray.from_kwargs( - mass1=opt.mass1, - mass2=opt.mass2, - spin1x=opt.spin1x, - spin1y=opt.spin1y, - spin1z=opt.spin1z, - spin2x=opt.spin2x, - spin2y=opt.spin2y, - spin2z=opt.spin2z, - eccentricity=opt.eccentricity, - rel_anomaly=opt.rel_anomaly) + mass1=opt.mass1, + mass2=opt.mass2, + spin1x=opt.spin1x, + spin1y=opt.spin1y, + spin1z=opt.spin1z, + spin2x=opt.spin2x, + spin2y=opt.spin2y, + spin2z=opt.spin2z, + eccentricity=opt.eccentricity, + rel_anomaly=opt.rel_anomaly, + ) with ctx: fft.from_cli(opt) @@ -249,13 +342,20 @@ with ctx: segments = strain_segments.fourier_segments() logging.info("Calculating the PSDs") - psd.associate_psds_to_segments(opt, segments, gwstrain, flen, delta_f, - flow, dyn_range_factor=pycbc.DYN_RANGE_FAC, - precision='single') + psd.associate_psds_to_segments( + opt, + segments, + gwstrain, + flen, + delta_f, + flow, + dyn_range_factor=pycbc.DYN_RANGE_FAC, + precision="single", + ) logging.info("Making template: %s", opt.approximant) if opt.use_params_of_closest_injection: - if not hasattr(gwstrain, 'injections') or not opt.trigger_time: + if not hasattr(gwstrain, "injections") or not opt.trigger_time: err_msg = "To use --use-params-of-closest-injection you must " err_msg += "be making injections with an injection file and using " err_msg += "the --trigger-time option." @@ -269,21 +369,25 @@ with ctx: row = WaveformArray.from_kwargs( mass1=inj.mass1, mass2=inj.mass2, - spin1x=inj.spin1x if hasattr(inj, 'spin1x') else 0, - spin1y=inj.spin1y if hasattr(inj, 'spin1y') else 0, - spin1z=inj.spin1z if hasattr(inj, 'spin1z') else 0, - spin2x=inj.spin2x if hasattr(inj, 'spin2x') else 0, - spin2y=inj.spin2y if hasattr(inj, 'spin2y') else 0, - spin2z=inj.spin2z if hasattr(inj, 'spin2z') else 0) + spin1x=inj.spin1x if hasattr(inj, "spin1x") else 0, + spin1y=inj.spin1y if hasattr(inj, "spin1y") else 0, + spin1z=inj.spin1z if hasattr(inj, "spin1z") else 0, + spin2x=inj.spin2x if hasattr(inj, "spin2x") else 0, + spin2y=inj.spin2y if hasattr(inj, "spin2y") else 0, + spin2z=inj.spin2z if hasattr(inj, "spin2z") else 0, + ) # FIXME: Don't like hardcoded 16384 here # NOTE: f_lower is set in strain.py as inj.f_lower, but this is a # little unclear to see and caused me some problems! Stating # this clearly here so no-one makes the same mistake as me. - td_template = inj_set.make_strain_from_inj_object(inj, 1./16384., - opt.channel_name[0:2], f_lower=inj.f_lower, - distance_scale=opt.injection_scale_factor) - td_template = resample_to_delta_t(td_template, gwstrain.delta_t, - method='ldas') + td_template = inj_set.make_strain_from_inj_object( + inj, + 1.0 / 16384.0, + opt.channel_name[0:2], + f_lower=inj.f_lower, + distance_scale=opt.injection_scale_factor, + ) + td_template = resample_to_delta_t(td_template, gwstrain.delta_t, method="ldas") td_template = td_template * pycbc.DYN_RANGE_FAC # apply a time shift so that the merger is at time zero, # the usual convention for templates @@ -295,13 +399,13 @@ with ctx: ltt = detector.time_delay_from_earth_center(ra, dec, time) td_template._epoch -= time + ltt # Check if waveform is too long - tlen = (flen-1) * 2 + tlen = (flen - 1) * 2 # FIXME: Hardcoded 7./8. factor here, but I'm not too bothered by this. if len(td_template) > (7 * tlen // 8): new_start_idx = len(td_template) - (7 * tlen // 8) td_template = td_template[new_start_idx:] - td_template = td_template.taper_timeseries(location='TAPER_START') - if hasattr(inj, 'waveform'): + td_template = td_template.taper_timeseries(location="TAPER_START") + if hasattr(inj, "waveform"): approximant = inj.waveform else: approximant = inj.approximant @@ -312,25 +416,30 @@ with ctx: logging.info("Making template: %s", opt.approximant) if opt.u_val is None: template = waveform.get_waveform_filter( - zeros(flen, dtype=complex64), - approximant=approximant, - template=row[0], - inclination=opt.inclination, - coa_phase=opt.coa_phase, - taper=opt.taper_template, - f_lower=template_flow, delta_f=delta_f, - delta_t=gwstrain.delta_t, - distance = 1.0/pycbc.DYN_RANGE_FAC) + zeros(flen, dtype=complex64), + approximant=approximant, + template=row[0], + inclination=opt.inclination, + coa_phase=opt.coa_phase, + taper=opt.taper_template, + f_lower=template_flow, + delta_f=delta_f, + delta_t=gwstrain.delta_t, + distance=1.0 / pycbc.DYN_RANGE_FAC, + ) else: tp, tc = waveform.get_two_pol_waveform_filter( - zeros(flen, dtype=complex64), - zeros(flen, dtype=complex64), row[0], - approximant=approximant, - inclination=opt.inclination, - coa_phase=opt.coa_phase, - taper=opt.taper_template, - f_lower=template_flow, delta_f=delta_f, - delta_t=gwstrain.delta_t) + zeros(flen, dtype=complex64), + zeros(flen, dtype=complex64), + row[0], + approximant=approximant, + inclination=opt.inclination, + coa_phase=opt.coa_phase, + taper=opt.taper_template, + f_lower=template_flow, + delta_f=delta_f, + delta_t=gwstrain.delta_t, + ) template = tc.multiply_and_add(tp, opt.u_val) if opt.high_frequency_cutoff: @@ -340,30 +449,34 @@ with ctx: # FIXME: should probably switch to something like what is done for parsing # the approximant for chisq bins at some point - class t(object): + class t: pass + parse_row = t() parse_row.params = t() for param in row.fieldnames: setattr(parse_row.params, param, row[param][0]) - chisq_bins_float = vetoes.SingleDetPowerChisq.parse_option(parse_row, - opt.chisq_bins) - if numpy.isnan(chisq_bins_float) or \ - (int(chisq_bins_float) < opt.minimum_chisq_bins): + chisq_bins_float = vetoes.SingleDetPowerChisq.parse_option( + parse_row, opt.chisq_bins + ) + if numpy.isnan(chisq_bins_float) or ( + int(chisq_bins_float) < opt.minimum_chisq_bins + ): if opt.minimum_chisq_bins: chisq_bins = opt.minimum_chisq_bins - logging.warning("Number of chisq bins is less than requested " - "minimum or is NaN. Using %d bins.", - opt.minimum_chisq_bins) + logging.warning( + "Number of chisq bins is less than requested " + "minimum or is NaN. Using %d bins.", + opt.minimum_chisq_bins, + ) else: - raise ValueError( - "Chisq bins is NaN or negative and no minimum is set.") + raise ValueError("Chisq bins is NaN or negative and no minimum is set.") else: chisq_bins = int(chisq_bins_float) - f['template'] = template.numpy() - f['template'].attrs['delta_f'] = template.delta_f + f["template"] = template.numpy() + f["template"].attrs["delta_f"] = template.delta_f snrs, chisqs = [], [] raw_bins = [[] for b in range(chisq_bins)] @@ -376,11 +489,9 @@ with ctx: if opt.window: start_time_wind = opt.trigger_time - opt.window - if start_time_wind > start_time: - start_time = start_time_wind + start_time = max(start_time, start_time_wind) end_time_wind = opt.trigger_time + opt.window - if end_time_wind < end_time: - end_time = end_time_wind + end_time = min(end_time, end_time_wind) for s_num, stilde in enumerate(segments): start = stilde.epoch + stilde.analyze.start / float(opt.sample_rate) @@ -393,22 +504,26 @@ with ctx: break logging.info("Filtering segment %s", s_num) - snr, corr, norm = filter.matched_filter_core(template, stilde, - psd=stilde.psd, - low_frequency_cutoff=flow) + snr, corr, norm = filter.matched_filter_core( + template, stilde, psd=stilde.psd, low_frequency_cutoff=flow + ) snr *= norm if opt.subtract_template: - stilde = subtract_template(stilde, template, - snr, opt.trigger_time, flow) - snr, corr, norm = filter.matched_filter_core(template, stilde, - psd=stilde.psd, - low_frequency_cutoff=flow) + stilde = subtract_template(stilde, template, snr, opt.trigger_time, flow) + snr, corr, norm = filter.matched_filter_core( + template, stilde, psd=stilde.psd, low_frequency_cutoff=flow + ) logging.info("calculating chisq") - chisq, raw_bin = vetoes.power_chisq(template, stilde, chisq_bins, stilde.psd, - low_frequency_cutoff=flow, - return_bins=True) + chisq, raw_bin = vetoes.power_chisq( + template, + stilde, + chisq_bins, + stilde.psd, + low_frequency_cutoff=flow, + return_bins=True, + ) chisq /= chisq_bins * 2 - 2 snrs.append(snr[stilde.analyze]) @@ -427,25 +542,31 @@ with ctx: raise ValueError(err_msg) for i in range(chisq_bins): - key = 'chisq_bins/%s' % i + key = "chisq_bins/%s" % i f[key] = numpy.concatenate(raw_bins[i])[sidx:eidx] - f[key].attrs['start_time'] = start_time - f[key].attrs['delta_t'] = snr.delta_t - - f['chisq_boundaries'] = numpy.array(vetoes.power_chisq_bins(template, chisq_bins, stilde.psd, - low_frequency_cutoff=flow)) * template.delta_f - - f['snr'] = numpy.concatenate([snr.numpy() for snr in snrs])[sidx:eidx] - f['snr'].attrs['start_time'] = start_time - f['snr'].attrs['delta_t'] = snr.delta_t - - f['chisq'] = numpy.concatenate([chisq.numpy() for chisq in chisqs])[sidx:eidx] - f['chisq'].attrs['start_time'] = start_time - f['chisq'].attrs['delta_t'] = snr.delta_t + f[key].attrs["start_time"] = start_time + f[key].attrs["delta_t"] = snr.delta_t + + f["chisq_boundaries"] = ( + numpy.array( + vetoes.power_chisq_bins( + template, chisq_bins, stilde.psd, low_frequency_cutoff=flow + ) + ) + * template.delta_f + ) + + f["snr"] = numpy.concatenate([snr.numpy() for snr in snrs])[sidx:eidx] + f["snr"].attrs["start_time"] = start_time + f["snr"].attrs["delta_t"] = snr.delta_t + + f["chisq"] = numpy.concatenate([chisq.numpy() for chisq in chisqs])[sidx:eidx] + f["chisq"].attrs["start_time"] = start_time + f["chisq"].attrs["delta_t"] = snr.delta_t if opt.trigger_time is not None: - f.attrs['event_time'] = opt.trigger_time - f.attrs['approximant'] = approximant.encode() - f.attrs['ifo'] = ifo.encode() - f.attrs['command_line'] = (' '.join(sys.argv)).encode() + f.attrs["event_time"] = opt.trigger_time + f.attrs["approximant"] = approximant.encode() + f.attrs["ifo"] = ifo.encode() + f.attrs["command_line"] = (" ".join(sys.argv)).encode() logging.info("Finished") diff --git a/bin/pycbc_source_probability_offline b/bin/pycbc_source_probability_offline index e79614ce167..f3cdf21afa5 100755 --- a/bin/pycbc_source_probability_offline +++ b/bin/pycbc_source_probability_offline @@ -3,29 +3,37 @@ Compute source probabilities using mchirp estimation method for all events in a chunk with an IFAR above certain threshold. """ -import json -import tqdm + import argparse +import json import logging + import numpy as np +import tqdm import pycbc +from pycbc import mchirp_area from pycbc.io import hdf from pycbc.pnutils import mass1_mass2_to_mchirp_eta -from pycbc import mchirp_area parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trigger-file', required=True) -parser.add_argument('--bank-file', required=True) -parser.add_argument('--single-detector-triggers', nargs='+', required=True) -parser.add_argument('--search-tag', required=True, - help='String to add to the output file names ' - 'identifying the search. eg: PYCBC_AllSky, ' - 'PYCBC_HighMass') -parser.add_argument('--ifar-threshold', type=float, default=None, - help='Select only candidate events with IFAR ' - 'above threshold.') +parser.add_argument("--trigger-file", required=True) +parser.add_argument("--bank-file", required=True) +parser.add_argument("--single-detector-triggers", nargs="+", required=True) +parser.add_argument( + "--search-tag", + required=True, + help="String to add to the output file names " + "identifying the search. eg: PYCBC_AllSky, " + "PYCBC_HighMass", +) +parser.add_argument( + "--ifar-threshold", + type=float, + default=None, + help="Select only candidate events with IFAR above threshold.", +) mchirp_area.insert_args(parser) args = parser.parse_args() @@ -33,40 +41,47 @@ mc_area_args = mchirp_area.from_cli(args, parser) pycbc.init_logging(args.verbose) -TRIGGER_FILE = args.trigger_file.split('/')[-1] -GPS_START_TIME = TRIGGER_FILE.split('-')[2] -GPS_START_TIME_NS = TRIGGER_FILE.split('-')[3].split('.')[0] +TRIGGER_FILE = args.trigger_file.split("/")[-1] +GPS_START_TIME = TRIGGER_FILE.split("-")[2] +GPS_START_TIME_NS = TRIGGER_FILE.split("-")[3].split(".")[0] -logging.info('Using files: %s, %s, %s', TRIGGER_FILE, - args.bank_file.split('/')[-1], - (',').join([name.split('/')[-1] for name in - args.single_detector_triggers])) +logging.info( + "Using files: %s, %s, %s", + TRIGGER_FILE, + args.bank_file.split("/")[-1], + (",").join([name.split("/")[-1] for name in args.single_detector_triggers]), +) -dir_path = 'source_probability_results/CHUNK_%s-%s/' % (GPS_START_TIME, - GPS_START_TIME_NS) +dir_path = "source_probability_results/CHUNK_%s-%s/" % ( + GPS_START_TIME, + GPS_START_TIME_NS, +) pycbc.makedir(dir_path) -logging.info('Saving results in %s', dir_path) +logging.info("Saving results in %s", dir_path) -fortrigs = hdf.ForegroundTriggers(args.trigger_file, args.bank_file, - sngl_files=args.single_detector_triggers) +fortrigs = hdf.ForegroundTriggers( + args.trigger_file, args.bank_file, sngl_files=args.single_detector_triggers +) -ifar = fortrigs.get_coincfile_array('ifar') +ifar = fortrigs.get_coincfile_array("ifar") N_original = len(ifar) if args.ifar_threshold: idx = ifar > args.ifar_threshold ifar = ifar[idx] - logging.info('%i triggers out of %i with IFAR > %s' % - (len(ifar), N_original, str(args.ifar_threshold))) + logging.info( + "%i triggers out of %i with IFAR > %s" + % (len(ifar), N_original, str(args.ifar_threshold)) + ) else: idx = np.full(N_original, True) -mass1 = fortrigs.get_bankfile_array('mass1')[idx] -mass2 = fortrigs.get_bankfile_array('mass2')[idx] -mchirp,_ = mass1_mass2_to_mchirp_eta(mass1, mass2) +mass1 = fortrigs.get_bankfile_array("mass1")[idx] +mass2 = fortrigs.get_bankfile_array("mass2")[idx] +mchirp, _ = mass1_mass2_to_mchirp_eta(mass1, mass2) end_time = fortrigs.get_end_time()[idx] -sngl_snr = fortrigs.get_snglfile_array_dict('snr') -sngl_sigmasq = fortrigs.get_snglfile_array_dict('sigmasq') +sngl_snr = fortrigs.get_snglfile_array_dict("snr") +sngl_sigmasq = fortrigs.get_snglfile_array_dict("sigmasq") for event in tqdm.trange(len(ifar)): # sngl_snr[ifo] contains 2 arrays: first contains SNR values and @@ -75,15 +90,23 @@ for event in tqdm.trange(len(ifar)): ifos_event = [ifo for ifo in sngl_snr.keys() if sngl_snr[ifo][1][event]] snrs_event = [sngl_snr[ifo][0][event] for ifo in ifos_event] coinc_snr = sum([snr**2 for snr in snrs_event]) ** 0.5 - min_eff_dist = min([(sngl_sigmasq[ifo][0][event])**0.5 - / sngl_snr[ifo][0][event] for ifo in ifos_event]) - probs = mchirp_area.calc_probabilities(mchirp[event], coinc_snr, - min_eff_dist, mc_area_args) + min_eff_dist = min( + [ + (sngl_sigmasq[ifo][0][event]) ** 0.5 / sngl_snr[ifo][0][event] + for ifo in ifos_event + ] + ) + probs = mchirp_area.calc_probabilities( + mchirp[event], coinc_snr, min_eff_dist, mc_area_args + ) # We do not expect a mass gap entry, but just in case if "Mass Gap" in probs: probs.pop("Mass Gap") - ifo_names = ''.join(sorted(ifos_event)) - out_name = dir_path + '%s-%s-%i-1.json' % (ifo_names, args.search_tag, - int(end_time[event])) - with open(out_name, 'w') as outfile: + ifo_names = "".join(sorted(ifos_event)) + out_name = dir_path + "%s-%s-%i-1.json" % ( + ifo_names, + args.search_tag, + int(end_time[event]), + ) + with open(out_name, "w") as outfile: json.dump(probs, outfile) diff --git a/bin/pycbc_split_inspinj b/bin/pycbc_split_inspinj index 36cfefc63e9..002465a092e 100644 --- a/bin/pycbc_split_inspinj +++ b/bin/pycbc_split_inspinj @@ -2,10 +2,11 @@ # Copyright (C) 2014 Andrew Lundgren import argparse -from igwn_ligolw import utils as ligolw_utils -from igwn_ligolw import lsctables from itertools import cycle +from igwn_ligolw import lsctables +from igwn_ligolw import utils as ligolw_utils + import pycbc from pycbc.io.ligolw import LIGOLWContentHandler, get_table_columns @@ -13,14 +14,15 @@ from pycbc.io.ligolw import LIGOLWContentHandler, get_table_columns parser = argparse.ArgumentParser() pycbc.add_common_pycbc_options(parser) group = parser.add_mutually_exclusive_group(required=True) -group.add_argument("-n", "--num-splits", type=int, - help="Number of files to be generated") -group.add_argument("-f", "--output-files", nargs='*', default=None, - help="Names of output files") +group.add_argument( + "-n", "--num-splits", type=int, help="Number of files to be generated" +) +group.add_argument( + "-f", "--output-files", nargs="*", default=None, help="Names of output files" +) parser.add_argument("-i", "--input-file", help="Injection file to be split") -parser.add_argument("-o", "--output-dir", default=None, - help="Location of output files") +parser.add_argument("-o", "--output-dir", default=None, help="Location of output files") args = parser.parse_args() @@ -29,9 +31,7 @@ if args.output_files and args.output_dir: # Read in input file xmldoc = ligolw_utils.load_filename( - args.input_file, - verbose=args.verbose, - contenthandler=LIGOLWContentHandler + args.input_file, verbose=args.verbose, contenthandler=LIGOLWContentHandler ) tabletype = lsctables.SimInspiralTable allinjs = tabletype.get_table(xmldoc) @@ -44,8 +44,7 @@ xmlroot.removeChild(allinjs) num_splits = args.num_splits or len(args.output_files) new_inj_tables = [ - tabletype.new(columns=get_table_columns(allinjs)) - for _ in range(num_splits) + tabletype.new(columns=get_table_columns(allinjs)) for _ in range(num_splits) ] table_cycle = cycle(new_inj_tables) @@ -53,15 +52,15 @@ for inj in sorted(allinjs, key=lambda x: x.time_geocent): next(table_cycle).append(inj) if not args.output_files: - temp = args.input_file.split('-') - temp[1] += '_%.4u' - filename_pattern = '-'.join(temp) + temp = args.input_file.split("-") + temp[1] += "_%.4u" + filename_pattern = "-".join(temp) for idx, simtable in enumerate(new_inj_tables): xmlroot.appendChild(simtable) if not args.output_files: - out_path = args.output_dir + '/' + filename_pattern % idx + out_path = args.output_dir + "/" + filename_pattern % idx else: out_path = args.output_files[idx] - ligolw_utils.write_filename(xmldoc, out_path, compress='auto') + ligolw_utils.write_filename(xmldoc, out_path, compress="auto") xmlroot.removeChild(simtable) diff --git a/bin/pycbc_splitbank b/bin/pycbc_splitbank index 41f81d17e40..ec3d6bd98f0 100644 --- a/bin/pycbc_splitbank +++ b/bin/pycbc_splitbank @@ -28,21 +28,21 @@ """Splits a table in an xml file into multiple pieces.""" import argparse -from numpy import random, ceil -from igwn_ligolw import ligolw -from igwn_ligolw import lsctables +from igwn_ligolw import ligolw, lsctables from igwn_ligolw import utils as ligolw_utils +from numpy import ceil, random import pycbc +from pycbc.conversions import mchirp_from_mass1_mass2 from pycbc.io.ligolw import ( - LIGOLWContentHandler, create_process_table, get_table_columns + LIGOLWContentHandler, + create_process_table, + get_table_columns, ) -from pycbc.conversions import mchirp_from_mass1_mass2 from pycbc.pnutils import frequency_cutoff_from_name - -__author__ = "Alex Nitz " +__author__ = "Alex Nitz " __program__ = "pycbc_splitbank" @@ -50,46 +50,81 @@ __program__ = "pycbc_splitbank" parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) group = parser.add_mutually_exclusive_group(required=True) -group.add_argument('--templates-per-bank', metavar='SAMPLES', - help='number of templates in the output banks', type=int) -group.add_argument('-n', '--number-of-banks', metavar='N', - help='Split template bank into N files', type=int) -group.add_argument("-O", "--output-filenames", nargs='*', default=None, - action="store", - metavar="OUTPUT_FILENAME", help="""Directly specify the +group.add_argument( + "--templates-per-bank", + metavar="SAMPLES", + help="number of templates in the output banks", + type=int, +) +group.add_argument( + "-n", + "--number-of-banks", + metavar="N", + help="Split template bank into N files", + type=int, +) +group.add_argument( + "-O", + "--output-filenames", + nargs="*", + default=None, + action="store", + metavar="OUTPUT_FILENAME", + help="""Directly specify the names of the output files. The number of files specified here will dictate how to split the bank. It will be split - equally between all specified files.""") - -parser.add_argument("-o", "--output-prefix", default=None, - help="Prefix to add to the template bank name (name becomes output#.xml[.gz])" ) - -parser.add_argument("-t", "--bank-file", metavar='INPUT_FILE', - help='Template bank to split', required=True) -parser.add_argument("--sort-frequency-cutoff", - help="Frequency cutoff to use for sorting the sub banks") -parser.add_argument("--sort-mchirp", action="store_true", default=False, - help='Sort templates by chirp mass before splitting') -parser.add_argument("--random-sort", action="store_true", default=False, - help='Sort templates randomly before splitting') -parser.add_argument("--random-seed", type=int, - help='Random seed to use when sorting randomly') + equally between all specified files.""", +) + +parser.add_argument( + "-o", + "--output-prefix", + default=None, + help="Prefix to add to the template bank name (name becomes output#.xml[.gz])", +) + +parser.add_argument( + "-t", + "--bank-file", + metavar="INPUT_FILE", + help="Template bank to split", + required=True, +) +parser.add_argument( + "--sort-frequency-cutoff", help="Frequency cutoff to use for sorting the sub banks" +) +parser.add_argument( + "--sort-mchirp", + action="store_true", + default=False, + help="Sort templates by chirp mass before splitting", +) +parser.add_argument( + "--random-sort", + action="store_true", + default=False, + help="Sort templates randomly before splitting", +) +parser.add_argument( + "--random-seed", type=int, help="Random seed to use when sorting randomly" +) args = parser.parse_args() if args.output_filenames and args.output_prefix: - errMsg="Cannot supply --output-filenames with --output-prefix." + errMsg = "Cannot supply --output-filenames with --output-prefix." parser.error(errMsg) if args.sort_mchirp and args.random_sort: - errMsg="You can't sort by Mchirp *and* randomly, dumbass!" + errMsg = "You can't sort by Mchirp *and* randomly, dumbass!" parser.error(errMsg) if args.output_filenames: args.number_of_banks = len(args.output_filenames) -indoc = ligolw_utils.load_filename(args.bank_file, verbose=args.verbose, - contenthandler=LIGOLWContentHandler) +indoc = ligolw_utils.load_filename( + args.bank_file, verbose=args.verbose, contenthandler=LIGOLWContentHandler +) try: template_bank_table = lsctables.SnglInspiralTable.get_table(indoc) @@ -106,7 +141,8 @@ tt = template_bank_table if args.sort_frequency_cutoff: sort_key = lambda x: frequency_cutoff_from_name( - args.sort_frequency_cutoff, x.mass1, x.mass2, x.spin1z, x.spin2z) + args.sort_frequency_cutoff, x.mass1, x.mass2, x.spin1z, x.spin2z + ) tt = sorted(template_bank_table, key=sort_key) if args.sort_mchirp: @@ -129,26 +165,24 @@ elif args.templates_per_bank: num_per_file = args.templates_per_bank num_files = int(ceil(float(length) / num_per_file)) -index_list = [int(round(num_per_file*idx)) for idx in range(num_files)] +index_list = [int(round(num_per_file * idx)) for idx in range(num_files)] index_list.append(length) -assert(index_list[0] == 0) +assert index_list[0] == 0 for num, (idx1, idx2) in enumerate(zip(index_list[:-1], index_list[1:])): - assert(idx2 > idx1) + assert idx2 > idx1 # create a blank xml document and add the process id outdoc = ligolw.Document() outdoc.appendChild(ligolw.LIGO_LW()) process = create_process_table( - outdoc, - program_name=__program__, - options=args.__dict__ + outdoc, program_name=__program__, options=args.__dict__ ) sngl_inspiral_table = tabletype.new(columns=used_columns) outdoc.childNodes[0].appendChild(sngl_inspiral_table) - for i in range(idx2-idx1): + for i in range(idx2 - idx1): row = tt.pop() row.process_id = process.process_id sngl_inspiral_table.append(row) @@ -160,7 +194,7 @@ for num, (idx1, idx2) in enumerate(zip(index_list[:-1], index_list[1:])): if args.output_filenames: outname = args.output_filenames[num] elif args.output_prefix: - outname = args.output_prefix + str(num) + '.xml.gz' + outname = args.output_prefix + str(num) + ".xml.gz" else: errMsg = "Cannot figure out how to set output file names." raise ValueError(errMsg) diff --git a/bin/pycbc_upload_xml_to_gracedb b/bin/pycbc_upload_xml_to_gracedb index 3b0649015e2..1375da6f22a 100755 --- a/bin/pycbc_upload_xml_to_gracedb +++ b/bin/pycbc_upload_xml_to_gracedb @@ -20,28 +20,30 @@ Take a coinc xml file containing multiple events and upload to gracedb. """ -import os import argparse import logging -import numpy as np +import os + import matplotlib -matplotlib.use('agg') +import numpy as np + +matplotlib.use("agg") -from ligo.gracedb.rest import GraceDb import lal import lal.series from igwn_ligolw import lsctables from igwn_ligolw import utils as ligolw_utils from igwn_segments import segment, segmentlist +from ligo.gracedb.rest import GraceDb import pycbc +from pycbc import conversions as conv from pycbc.io.gracedb import gracedb_tag_with_version -from pycbc.io.ligolw import LIGOLWContentHandler from pycbc.io.hdf import HFile +from pycbc.io.ligolw import LIGOLWContentHandler from pycbc.psd import interpolate -from pycbc.types import FrequencySeries from pycbc.results import generate_asd_plot -from pycbc import conversions as conv +from pycbc.types import FrequencySeries def check_gracedb_for_event(gdb_handle, query, far): @@ -50,63 +52,109 @@ def check_gracedb_for_event(gdb_handle, query, far): which matches the FAR given """ gdb_events_match_query = list(gdb_handle.events(query=query)) - ifar = conv.sec_to_year(1. / far) + ifar = conv.sec_to_year(1.0 / far) for gdb_event in gdb_events_match_query: # Test each gracedb event to see if the FAR matches this event - if np.abs(gdb_event['far'] - far) < 1e-16: + if np.abs(gdb_event["far"] - far) < 1e-16: # If an event has been found which matches the FAR - logging.info('Event already exists in GraceDb server with ' - 'time %.3f and IFAR %.3e: %s', - gdb_event['gpstime'], ifar, - gdb_event['graceid']) + logging.info( + "Event already exists in GraceDb server with " + "time %.3f and IFAR %.3e: %s", + gdb_event["gpstime"], + ifar, + gdb_event["graceid"], + ) return True # If no event has been found, log this, and return False - logging.info('No event found in GraceDb with IFAR %.3e when using ' - 'query: "%s"', ifar, query) + logging.info( + 'No event found in GraceDb with IFAR %.3e when using query: "%s"', ifar, query + ) return False parser = argparse.ArgumentParser(description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--psd-files", nargs='+', required=True, - help='HDF file(s) containing the PSDs to upload') -parser.add_argument('--input-file', required=True, type=str, - help='Input LIGOLW XML file of coincidences.') -parser.add_argument('--log-message', type=str, metavar='MESSAGE', - help='Add a log entry to each upload with the given message') -parser.add_argument('--testing', action="store_true", default=False, - help="Upload event to the TEST group of gracedb.") -parser.add_argument('--min-ifar', type=float, metavar='YEARS', - help='Only upload events more significant than given IFAR') -parser.add_argument('--production-server', action="store_true", default=False, - help="Upload event to production graceDB. If not given " - "events will be uploaded to playground server.") -parser.add_argument('--force-overwrite', action='store_true', default=False, - help="GraceDb instance will be checked for if an event " - "with the same event time and FAR already exist. " - "If so, event will not be uploaded") -parser.add_argument('--query-string', default='pycbc', - help="If not using --force-overwrite, add a string to " - "gracedb query to further filter events. " - "Default='pycbc'") -parser.add_argument('--search-id-string', default='AllSky', - help="Using T050017-v1 naming convention for " - "output XML filename. This is a string for " - "search identifier in filename, e.g. 'AllSky' " - "would give H1L1V1-PYCBC_AllSky-1234567890-1.xml. " - "See https://dcc.ligo.org/LIGO-T050017/public. " - "Default: 'AllSky'") -parser.add_argument('--output-directory', default=os.getcwd(), - help="Output directory for locally stored XML and PSD " - "files. Default: current directory") -parser.add_argument('--no-upload', action='store_true', - help="Flag used to indicate that we are not uploading to " - "GraceDb.") -parser.add_argument('--generate-plots', action='store_true', - help="Flag used to indicate that we want to make " - "plots. Uploaded to GraceDB if --no-upload is " - "not given.") +parser.add_argument( + "--psd-files", + nargs="+", + required=True, + help="HDF file(s) containing the PSDs to upload", +) +parser.add_argument( + "--input-file", + required=True, + type=str, + help="Input LIGOLW XML file of coincidences.", +) +parser.add_argument( + "--log-message", + type=str, + metavar="MESSAGE", + help="Add a log entry to each upload with the given message", +) +parser.add_argument( + "--testing", + action="store_true", + default=False, + help="Upload event to the TEST group of gracedb.", +) +parser.add_argument( + "--min-ifar", + type=float, + metavar="YEARS", + help="Only upload events more significant than given IFAR", +) +parser.add_argument( + "--production-server", + action="store_true", + default=False, + help="Upload event to production graceDB. If not given " + "events will be uploaded to playground server.", +) +parser.add_argument( + "--force-overwrite", + action="store_true", + default=False, + help="GraceDb instance will be checked for if an event " + "with the same event time and FAR already exist. " + "If so, event will not be uploaded", +) +parser.add_argument( + "--query-string", + default="pycbc", + help="If not using --force-overwrite, add a string to " + "gracedb query to further filter events. " + "Default='pycbc'", +) +parser.add_argument( + "--search-id-string", + default="AllSky", + help="Using T050017-v1 naming convention for " + "output XML filename. This is a string for " + "search identifier in filename, e.g. 'AllSky' " + "would give H1L1V1-PYCBC_AllSky-1234567890-1.xml. " + "See https://dcc.ligo.org/LIGO-T050017/public. " + "Default: 'AllSky'", +) +parser.add_argument( + "--output-directory", + default=os.getcwd(), + help="Output directory for locally stored XML and PSD " + "files. Default: current directory", +) +parser.add_argument( + "--no-upload", + action="store_true", + help="Flag used to indicate that we are not uploading to GraceDb.", +) +parser.add_argument( + "--generate-plots", + action="store_true", + help="Flag used to indicate that we want to make " + "plots. Uploaded to GraceDB if --no-upload is " + "not given.", +) args = parser.parse_args() @@ -115,23 +163,29 @@ pycbc.init_logging(args.verbose, default_level=1) if args.production_server: gracedb = GraceDb() else: - gracedb = GraceDb(service_url='https://gracedb-playground.ligo.org/api/') + gracedb = GraceDb(service_url="https://gracedb-playground.ligo.org/api/") + +xmldoc = ligolw_utils.load_filename( + args.input_file, contenthandler=LIGOLWContentHandler +) -xmldoc = ligolw_utils.load_filename(args.input_file, - contenthandler=LIGOLWContentHandler) class psd_segment(segment): def __new__(cls, psd, *args): return segment.__new__(cls, *args) + def __init__(self, psd, *args): self.psd = psd + psds = {} for psd_file in args.psd_files: - (ifo, group), = HFile(psd_file, "r").items() + ((ifo, group),) = HFile(psd_file, "r").items() psd = [group["psds"][str(i)] for i in range(len(group["psds"].keys()))] - psds[ifo] = segmentlist(psd_segment(*segargs) for segargs in zip( - psd, group["start_time"], group["end_time"])) + psds[ifo] = segmentlist( + psd_segment(*segargs) + for segargs in zip(psd, group["start_time"], group["end_time"]) + ) coinc_table = lsctables.CoincTable.get_table(xmldoc) coinc_inspiral_table = lsctables.CoincInspiralTable.get_table(xmldoc) @@ -155,14 +209,17 @@ for event in coinc_table: if coinc_insp.coinc_event_id == event.coinc_event_id: coinc_inspiral_table_curr.append(coinc_insp) - if args.min_ifar is not None and \ - conv.sec_to_year(1. / coinc_inspiral_table_curr[0].combined_far) < args.min_ifar: # event IFAR is smaller than the minimum + if ( + args.min_ifar is not None + and conv.sec_to_year(1.0 / coinc_inspiral_table_curr[0].combined_far) + < args.min_ifar + ): # event IFAR is smaller than the minimum continue time = coinc_inspiral_table_curr[0].end_time if not args.force_overwrite and not args.no_upload: far = coinc_inspiral_table_curr[0].combined_far - query = args.query_string + ' %.3f .. %.3f' % (time - 1, time + 1) + query = args.query_string + " %.3f .. %.3f" % (time - 1, time + 1) if check_gracedb_for_event(gracedb, query, far): continue @@ -184,7 +241,9 @@ for event in coinc_table: except KeyError: parser.error( "--psd-files {0}: no PSDs found for detector {1}".format( - " ".join(args.psd_files), sngl.ifo)) + " ".join(args.psd_files), sngl.ifo + ) + ) try: psd = psd[psd.find(sngl.end)].psd @@ -192,22 +251,30 @@ for event in coinc_table: parser.error( "--psd-files {0}: no PSD found for detector {1} " "at GPS time {2}".format( - " ".join(args.psd_files), sngl.ifo, sngl.end)) + " ".join(args.psd_files), sngl.ifo, sngl.end + ) + ) # Resample the psd so it can be uploaded to GDB df = 0.25 - psd_fs = FrequencySeries(psd, delta_f=psd.attrs["delta_f"], - dtype=np.float64) + psd_fs = FrequencySeries( + psd, delta_f=psd.attrs["delta_f"], dtype=np.float64 + ) psd_fs = interpolate(psd_fs, df) psddict[sngl.ifo] = psd_fs - flow = psd.file.attrs['low_frequency_cutoff'] + flow = psd.file.attrs["low_frequency_cutoff"] kmin = int(flow / df) fseries = lal.CreateREAL8FrequencySeries( - "psd", lal.LIGOTimeGPS(int(psd.attrs["epoch"])), kmin * df, df, - lal.StrainUnit**2 / lal.HertzUnit, len(psd_fs) - kmin) + "psd", + lal.LIGOTimeGPS(int(psd.attrs["epoch"])), + kmin * df, + df, + lal.StrainUnit**2 / lal.HertzUnit, + len(psd_fs) - kmin, + ) fseries.data.data = psd_fs[kmin:] / np.square(pycbc.DYN_RANGE_FAC) lal_psddict[sngl.ifo] = fseries @@ -218,9 +285,9 @@ for event in coinc_table: lal.series.make_psd_xmldoc(lal_psddict, xmldoc.childNodes[-1]) ifos = sorted([sngl.ifo for sngl in sngl_inspiral_table_curr]) - ifos_str = ''.join(ifos) + ifos_str = "".join(ifos) id_str = args.search_id_string - filename_xml = "{}-PYCBC_{}-{:d}-1.xml".format(ifos_str, id_str, time) + filename_xml = f"{ifos_str}-PYCBC_{id_str}-{time:d}-1.xml" fullpath_xml = os.path.join(args.output_directory, filename_xml) ligolw_utils.write_filename(xmldoc, fullpath_xml) @@ -228,36 +295,33 @@ for event in coinc_table: if args.no_upload: logging.info("Not uploading event") else: - group_tag = 'Test' if args.testing else 'CBC' + group_tag = "Test" if args.testing else "CBC" r = gracedb.create_event( group_tag, - 'pycbc', + "pycbc", filename_xml, filecontents=open(fullpath_xml, "rb").read(), search=id_str, - offline=True + offline=True, ).json() logging.info("Uploaded event %s.", r["graceid"]) # add info for tracking code version - gracedb_tag_with_version(gracedb, r['graceid']) + gracedb_tag_with_version(gracedb, r["graceid"]) # document the absolute path to the input file - input_file_str = 'Candidate uploaded from ' \ - + os.path.abspath(args.input_file) - gracedb.write_log(r['graceid'], input_file_str) + input_file_str = "Candidate uploaded from " + os.path.abspath(args.input_file) + gracedb.write_log(r["graceid"], input_file_str) # add the custom log message, if provided if args.log_message is not None: gracedb.write_log( - r['graceid'], - args.log_message, - tag_name=['analyst_comments'] + r["graceid"], args.log_message, tag_name=["analyst_comments"] ) if args.generate_plots: - asd_png_filename = f'{ifos_str}-PYCBC_{id_str}_ASD-{time:d}-1.png' + asd_png_filename = f"{ifos_str}-PYCBC_{id_str}_ASD-{time:d}-1.png" fullpath_asd = os.path.join(args.output_directory, asd_png_filename) generate_asd_plot(psddict, fullpath_asd) logging.info("Saved ASD plot %s", asd_png_filename) @@ -267,10 +331,11 @@ for event in coinc_table: "PyCBC ASD estimate from the time of event", filename=fullpath_psd, tag_name=["psd"], - displayName=['ASDs'] + displayName=["ASDs"], + ) + logging.info( + "Uploaded file %s to event %s.", psd_png_filename, r["graceid"] ) - logging.info("Uploaded file %s to event %s.", psd_png_filename, - r["graceid"]) xmldoc.childNodes[-1].removeChild(coinc_event_table_curr) xmldoc.childNodes[-1].removeChild(coinc_inspiral_table_curr) @@ -280,5 +345,5 @@ for event in coinc_table: # through the tag name. This may not work in future if other things use # the LIGO_LW tag which we want to keep for cn in xmldoc.childNodes[-1].childNodes: - if cn.tagName == 'LIGO_LW': + if cn.tagName == "LIGO_LW": xmldoc.childNodes[-1].removeChild(cn) diff --git a/bin/pygrb/pycbc_grb_inj_finder b/bin/pygrb/pycbc_grb_inj_finder index 9d8d6cd6fd2..2f66913154e 100644 --- a/bin/pygrb/pycbc_grb_inj_finder +++ b/bin/pygrb/pycbc_grb_inj_finder @@ -1,5 +1,4 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- # # Copyright (C) 2019 Duncan Macleod, Francesco Pannarale, Erin Vincent # @@ -17,25 +16,20 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Cluster triggers from a pyGRB run -""" +"""Cluster triggers from a pyGRB run""" import argparse +import logging import operator import os -import logging import re from collections import defaultdict from functools import reduce -import tqdm - -import numpy - import h5py - +import numpy +import tqdm from gwdatafind.utils import filename_metadata - from igwn_segments import segmentlist from igwn_segments.utils import fromsegwizard @@ -53,9 +47,11 @@ NETWORK_IFO_EVENT_ID_REGEX = re.compile( ) EVENT_ID_REGEX = re.compile(r"event_id\Z") -TQDM_BAR_FORMAT = ("{desc}: |{bar}| " - "{n_fmt}/{total_fmt} {unit} ({percentage:3.0f}%) " - "[{elapsed} | ETA {remaining}]{postfix}") +TQDM_BAR_FORMAT = ( + "{desc}: |{bar}| " + "{n_fmt}/{total_fmt} {unit} ({percentage:3.0f}%) " + "[{elapsed} | ETA {remaining}]{postfix}" +) TQDM_KW = { "ascii": " -=#", "bar_format": TQDM_BAR_FORMAT, @@ -66,60 +62,50 @@ TQDM_KW = { # -- utilities ---------------------------------- def find_split_files(filelist, split): for fn in filelist: - if (split is None and 'SPLIT' not in fn) or split in fn: + if (split is None and "SPLIT" not in fn) or split in fn: yield fn def read_segment_files(segfiles): def _read(name): - with open(name, "r") as f: + with open(name) as f: return fromsegwizard(f) - return segmentlist(reduce( - operator.or_, - map(_read, segfiles), - segmentlist())) + return segmentlist(reduce(operator.or_, map(_read, segfiles), segmentlist())) def read_hdf5_triggers(inputfiles, verbose=False): - """Merge several HDF5 files into a single file + """ + Merge several HDF5 files into a single file Parameters ---------- inputfiles : `list` of `str` the paths of the input HDF5 files to merge + """ datasets = {} def _scan_dataset(name, obj): - if not isinstance(obj, h5py.Dataset): - return - elif name.startswith("network") and \ - NETWORK_IFO_EVENT_ID_REGEX.match(name): - return - elif "template_id" in name: + if not isinstance(obj, h5py.Dataset) or (name.startswith("network") and NETWORK_IFO_EVENT_ID_REGEX.match(name)): return - elif "search" in name: - return - elif "gating" in name: + if "template_id" in name or "search" in name or "gating" in name: return + shape = obj.shape + dtype = obj.dtype + try: + shape = numpy.sum(datasets[name][0] + shape, keepdims=True) + except KeyError: + pass else: - shape = obj.shape - dtype = obj.dtype - try: - shape = numpy.sum(datasets[name][0] + shape, keepdims=True) - except KeyError: - pass - else: - assert dtype == datasets[name][1], ( - "Cannot merge {0}/{1}, does not match dtype".format( - obj.file.filename, name, - )) - datasets[name] = (shape, dtype) + assert dtype == datasets[name][1], ( + f"Cannot merge {obj.file.filename}/{name}, does not match dtype" + ) + datasets[name] = (shape, dtype) # get list of datasets for filename in inputfiles: - with HFile(filename, 'r') as h5f: + with HFile(filename, "r") as h5f: h5f.visititems(_scan_dataset) position = defaultdict(int) @@ -132,17 +118,17 @@ def read_hdf5_triggers(inputfiles, verbose=False): # copy dataset contents for filename in inputfiles: - with HFile(filename, 'r') as h5in: + with HFile(filename, "r") as h5in: # Skipping empty files - if len(h5in['network']['end_time_gc'][:]): + if len(h5in["network"]["end_time_gc"][:]): for dset in datasets: data = h5in[dset][:] size = data.shape[0] pos = position[dset] if EVENT_ID_REGEX.search(dset): - out_dict[dset][pos:pos+size] = data + pos + out_dict[dset][pos : pos + size] = data + pos else: - out_dict[dset][pos:pos+size] = data + out_dict[dset][pos : pos + size] = data position[dset] += size return out_dict @@ -190,7 +176,7 @@ parser.add_argument( "-W", "--time-window", type=float, - default=0., + default=0.0, help="the found time window (default: %(default)s)", ) parser.add_argument( @@ -205,8 +191,8 @@ parser.add_argument( action="append", default=[], help="ignore injections in segments found within " - "these files, e.g. buffer segments (may be given " - "more than once", + "these files, e.g. buffer segments (may be given " + "more than once", ) args = parser.parse_args() @@ -232,14 +218,18 @@ for key in firstfile.keys(): allinjections = None -with tqdm.tqdm(args.inj_files, desc="Finding injections", - disable=not args.verbose, unit="files", - postfix=dict(found=0, missed=0, excluded=nexcluded), - **TQDM_KW) as bar: +with tqdm.tqdm( + args.inj_files, + desc="Finding injections", + disable=not args.verbose, + unit="files", + postfix=dict(found=0, missed=0, excluded=nexcluded), + **TQDM_KW, +) as bar: for injf in bar: # get list of trigger files try: - split, = SPLIT_FILENAME.search(injf).groups() + (split,) = SPLIT_FILENAME.search(injf).groups() except AttributeError: # no regex match split = None trigfiles = list(find_split_files(args.input_files, split)) @@ -254,41 +244,40 @@ with tqdm.tqdm(args.inj_files, desc="Finding injections", # read triggers triggers = read_hdf5_triggers(trigfiles) - time = triggers['network/end_time_gc'] + time = triggers["network/end_time_gc"] time_sorting = time.argsort() - snr = triggers['network/'+args.rank_column][time_sorting] - event_id = triggers['network/event_id'][time_sorting] + snr = triggers["network/" + args.rank_column][time_sorting] + event_id = triggers["network/event_id"][time_sorting] # determine found or missed _left = numpy.searchsorted( time[time_sorting], injtime - args.time_window, - side='left', + side="left", ) _right = numpy.searchsorted( time[time_sorting], injtime + args.time_window, - side='right', + side="right", ) # get indices of found/missed injections and finding trigger for i, (l, r) in enumerate( - zip(_left, _right), - start=len(allinjections if allinjections is not None else []), + zip(_left, _right), + start=len(allinjections if allinjections is not None else []), ): if not r - l: missed.append(i) else: found.append(i) # the needed event_id is unique within its split_inj file - eid = event_id[l] if r - l == 1 \ - else event_id[l + snr[l:r].argmax()] + eid = event_id[l] if r - l == 1 else event_id[l + snr[l:r].argmax()] for x in triggers: - ij = x.split('/') + ij = x.split("/") # assign a new, overall unique event_id to a found # injection by simply counting found injections - if ij[1] == 'event_id': - out_trigs[ij[0]][ij[1]].append(len(found)-1) + if ij[1] == "event_id": + out_trigs[ij[0]][ij[1]].append(len(found) - 1) else: out_trigs[ij[0]][ij[1]].append(triggers[x][eid]) @@ -310,12 +299,7 @@ ifotag, desc, segment = filename_metadata(args.inj_files[0]) desc = SPLIT_FILENAME.split(desc)[0] outfilename = os.path.join( args.output_dir, - "{}-{}_FOUNDMISSED-{}-{}.h5".format( - args.ifo_tag or ifotag, - desc, - segment[0], - abs(segment), - ), + f"{args.ifo_tag or ifotag}-{desc}_FOUNDMISSED-{segment[0]}-{abs(segment)}.h5", ) with HFile(outfilename, "w") as h5out: @@ -325,7 +309,7 @@ with HFile(outfilename, "w") as h5out: injs = allinjections[injlist] for field in injs.fieldnames: # h5py does not support unicode strings - if 'U' not in str(injs[field].dtype): + if "U" not in str(injs[field].dtype): grp[field] = injs[field] else: grp[field] = [str(s) for s in injs[field]] @@ -334,7 +318,7 @@ with HFile(outfilename, "w") as h5out: for x in out_trigs: xgroup = h5out.create_group(x) for col in out_trigs[x]: - if "template_id" == col: + if col == "template_id": continue xgroup.create_dataset( col, diff --git a/bin/pygrb/pycbc_grb_trig_cluster b/bin/pygrb/pycbc_grb_trig_cluster index ed2d1c99ce7..fd1f22ea9cc 100644 --- a/bin/pygrb/pycbc_grb_trig_cluster +++ b/bin/pygrb/pycbc_grb_trig_cluster @@ -1,5 +1,4 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- # # Copyright (C) 2019 Duncan Macleod # @@ -17,31 +16,29 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Cluster triggers from a pyGRB run -""" +"""Cluster triggers from a pyGRB run""" import argparse +import logging import os import shutil import sys -import logging - -import tqdm - -import numpy import h5py - +import numpy +import tqdm from gwdatafind.utils import filename_metadata -from pycbc import init_logging, add_common_pycbc_options +from pycbc import add_common_pycbc_options, init_logging from pycbc.io.hdf import HFile __author__ = "Duncan Macleod " -TQDM_BAR_FORMAT = ("{desc}: |{bar}| " - "{n_fmt}/{total_fmt} {unit} ({percentage:3.0f}%) " - "[{elapsed} | ETA {remaining}]{postfix}") +TQDM_BAR_FORMAT = ( + "{desc}: |{bar}| " + "{n_fmt}/{total_fmt} {unit} ({percentage:3.0f}%) " + "[{elapsed} | ETA {remaining}]{postfix}" +) TQDM_KW = { "ascii": " -=#", "bar_format": TQDM_BAR_FORMAT, @@ -51,8 +48,10 @@ TQDM_KW = { # -- utilities ---------------------------------- + def slice_hdf5(inputfile, outfile, include, verbose=False): - """Create a new HDF5 file containing a slice of the network events + """ + Create a new HDF5 file containing a slice of the network events Here ``include`` should be an index array """ @@ -68,17 +67,20 @@ def slice_hdf5(inputfile, outfile, include, verbose=False): # find which single-ifo events to keep ifo_index = { ifo: numpy.unique( - h5in["network/{}_event_id".format(ifo)][:][include], - ) for ifo in ifos + h5in[f"network/{ifo}_event_id"][:][include], + ) + for ifo in ifos } - nsets = sum(isinstance(item, h5py.Dataset) - or isinstance(item, h5py.Group) for - group in h5in.values() for - item in group.values()) - msg = "Slicing {} network events into new file".format(nevents) - bar = tqdm.tqdm(total=nsets, desc=msg, disable=not verbose, - unit="datasets", **TQDM_KW) + nsets = sum( + isinstance(item, h5py.Dataset) or isinstance(item, h5py.Group) + for group in h5in.values() + for item in group.values() + ) + msg = f"Slicing {nevents} network events into new file" + bar = tqdm.tqdm( + total=nsets, desc=msg, disable=not verbose, unit="datasets", **TQDM_KW + ) with HFile(outfile, "w") as h5out: for old in h5in["network"].values(): if isinstance(old, h5py.Dataset): @@ -157,12 +159,7 @@ ifotag, filetag, segment = filename_metadata(args.trig_file) start, end = segment outfile = os.path.join( args.output_dir, - "{}-{}_CLUSTERED-{}-{}.h5".format( - ifotag, - filetag, - start, - end - start, - ), + f"{ifotag}-{filetag}_CLUSTERED-{start}-{end - start}.h5", ) # this list contains the indexing of clusters from all slides @@ -179,7 +176,7 @@ with HFile(args.trig_file, "r") as h5f: if not all_times.size: shutil.copyfile(args.trig_file, outfile) msg = "trigger file is empty\n" - msg += "copied input file to {}".format(outfile) + msg += f"copied input file to {outfile}" logging.info(msg) sys.exit(0) @@ -188,9 +185,7 @@ if not all_times.size: unique_slide_ids = numpy.unique(slide_ids) max_slide_id = max(unique_slide_ids) logging.info( - 'Clustering %d triggers from %d slides', - len(slide_ids), - len(unique_slide_ids) + "Clustering %d triggers from %d slides", len(slide_ids), len(unique_slide_ids) ) for slide_id in unique_slide_ids: @@ -209,12 +204,14 @@ for slide_id in unique_slide_ids: clusters = [] # find loudest trigger in each bin, for the current slide - for i in tqdm.tqdm(range(time.size), - desc="Initialising bins", - disable=not args.verbose, - total=time.size, - unit='triggers', - **TQDM_KW): + for i in tqdm.tqdm( + range(time.size), + desc="Initialising bins", + disable=not args.verbose, + total=time.size, + unit="triggers", + **TQDM_KW, + ): t, s = time[i], snr[i] idx = int(float(t - start) // win) bins[idx].append(i) @@ -230,13 +227,15 @@ for slide_id in unique_slide_ids: nclusters = 0 # cluster - bar = tqdm.tqdm(bins, - desc="Clustering bins", - disable=not args.verbose, - total=nbins, - unit='bins', - postfix=dict(nclusters=0), - **TQDM_KW) + bar = tqdm.tqdm( + bins, + desc="Clustering bins", + disable=not args.verbose, + total=nbins, + unit="bins", + postfix=dict(nclusters=0), + **TQDM_KW, + ) for i, bin_ in enumerate(bar): if not bin_: # empty continue @@ -295,14 +294,14 @@ for slide_id in unique_slide_ids: # slide_id_pos which is built at each slide_id all_clusters += list(slide_id_pos[clusters]) logging.info( - 'Slide %s/%s has %d triggers that were clustered to %d', + "Slide %s/%s has %d triggers that were clustered to %d", slide_id, max_slide_id, len(slide_id_pos), - len(clusters) + len(clusters), ) -logging.info('Total clustered triggers: %d', len(all_clusters)) +logging.info("Total clustered triggers: %d", len(all_clusters)) # -- write output -------------------------------- diff --git a/bin/pygrb/pycbc_grb_trig_combiner b/bin/pygrb/pycbc_grb_trig_combiner index cd30e55d7b4..f4ed0dfff6b 100644 --- a/bin/pygrb/pycbc_grb_trig_combiner +++ b/bin/pygrb/pycbc_grb_trig_combiner @@ -1,5 +1,4 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- # # Copyright (C) 2019 Duncan Macleod # @@ -17,33 +16,30 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Combine triggers from a splitbank GRB run -""" +"""Combine triggers from a splitbank GRB run""" import argparse -import os import logging - -import numpy - -import tqdm +import os import h5py - -from gwdatafind.utils import (file_segment, filename_metadata) - import igwn_segments as segments +import numpy +import tqdm +from gwdatafind.utils import file_segment, filename_metadata from igwn_segments.utils import fromsegwizard from pycbc import add_common_pycbc_options, init_logging -from pycbc.results.pygrb_postprocessing_utils import template_hash_to_id from pycbc.io.hdf import HFile +from pycbc.results.pygrb_postprocessing_utils import template_hash_to_id __author__ = "Duncan Macleod " -TQDM_BAR_FORMAT = ("{desc}: |{bar}| " - "{n_fmt}/{total_fmt} {unit} ({percentage:3.0f}%) " - "[{elapsed} | ETA {remaining}]") +TQDM_BAR_FORMAT = ( + "{desc}: |{bar}| " + "{n_fmt}/{total_fmt} {unit} ({percentage:3.0f}%) " + "[{elapsed} | ETA {remaining}]" +) TQDM_KW = { "ascii": " -=#", "bar_format": TQDM_BAR_FORMAT, @@ -53,8 +49,10 @@ TQDM_KW = { # -- utilities ----------------------------------- + def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): - """Merge several HDF5 files into a single file + """ + Merge several HDF5 files into a single file Parameters ---------- @@ -63,6 +61,7 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): outputfile : `str` the path of the output HDF5 file to write + """ attributes = {} datasets = {} @@ -84,22 +83,26 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): pass else: assert dtype == datasets[name][1], ( - "Cannot merge {0}/{1}, does not match dtype".format( - obj.file.filename, name, - )) + f"Cannot merge {obj.file.filename}/{name}, does not match dtype" + ) datasets[name] = (shape, dtype) # use default compression options from this file - for copt in ('compression', 'compression_opts'): + for copt in ("compression", "compression_opts"): compression_kw.setdefault(copt, getattr(obj, copt)) # get list of datasets attributes = {} datasets = {} - for filename in tqdm.tqdm(inputfiles, desc="Scanning trigger files", - disable=not verbose, total=nfiles, unit="files", - **TQDM_KW): - with HFile(filename, 'r') as h5f: + for filename in tqdm.tqdm( + inputfiles, + desc="Scanning trigger files", + disable=not verbose, + total=nfiles, + unit="files", + **TQDM_KW, + ): + with HFile(filename, "r") as h5f: # note: attributes are only recorded from the last file, # since we presume all files have the same attributes attributes = dict(h5f.attrs) @@ -113,7 +116,8 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): ) def _network_first(name): - """Key function to sort datasets in the network group first + """ + Key function to sort datasets in the network group first This is required so that when we merge the network/X1_event_id columns, we can reference the X1/event_id event count from the @@ -127,12 +131,11 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): dataset_names = sorted(datasets, key=_network_first) # get list of event_id columns - sngl_event_id_names = [ - x for x in dataset_names if x.endswith("/event_id") - ] + sngl_event_id_names = [x for x in dataset_names if x.endswith("/event_id")] network_event_id_names = [ - "network/{}".format(x.replace('/', '_')) for - x in sngl_event_id_names if not x.startswith("network") + "network/{}".format(x.replace("/", "_")) + for x in sngl_event_id_names + if not x.startswith("network") ] all_event_id_names = sngl_event_id_names + network_event_id_names @@ -142,7 +145,7 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): gating_datasets = set(filter(lambda x: "/gating/" in x, dataset_names)) once_only_datasets = search_datasets.union(gating_datasets) - with HFile(outputfile, 'w') as h5out: + with HFile(outputfile, "w") as h5out: h5out.attrs.update(attributes) # copy dataset contents @@ -154,19 +157,23 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): logging.info("Merging %s dataset", dset) data = [] - for filename in tqdm.tqdm(inputfiles, desc="Merging trigger files", - disable=not verbose, total=nfiles, - unit="files", **TQDM_KW): - with HFile(filename, 'r') as h5in: - if len(h5in['network']['event_id']) == 0: + for filename in tqdm.tqdm( + inputfiles, + desc="Merging trigger files", + disable=not verbose, + total=nfiles, + unit="files", + **TQDM_KW, + ): + with HFile(filename, "r") as h5in: + if len(h5in["network"]["event_id"]) == 0: continue data += [h5in[dset][:]] # read the search datasets and gating_datasets only once if dset in once_only_datasets: break - h5out.create_dataset(dset, data=numpy.concatenate(data), - **compression_kw) + h5out.create_dataset(dset, data=numpy.concatenate(data), **compression_kw) del data # END OF DATASET LOOP @@ -175,16 +182,21 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): n_events = 0 logging.info("Merging %s dataset", dset) data = [] - for filename in tqdm.tqdm(inputfiles, desc="Merging trigger files", - disable=not verbose, total=nfiles, - unit="files", **TQDM_KW): - with HFile(filename, 'r') as h5in: - if len(h5in['network']['event_id']) == 0: + for filename in tqdm.tqdm( + inputfiles, + desc="Merging trigger files", + disable=not verbose, + total=nfiles, + unit="files", + **TQDM_KW, + ): + with HFile(filename, "r") as h5in: + if len(h5in["network"]["event_id"]) == 0: continue # shift the event ids by the number of events gathered so far n_events_h5in = len(h5in[dset][:]) - data += [h5in[dset][:] + [n_events]*n_events_h5in] + data += [h5in[dset][:] + [n_events] * n_events_h5in] n_events += n_events_h5in data = numpy.concatenate(data) for dset in all_event_id_names: @@ -194,23 +206,23 @@ def merge_hdf5_files(inputfiles, outputfile, verbose=False, **compression_kw): # Handling template_id separately logging.info("Preparing template IDs") - template_ids = template_hash_to_id( - trigger_file=h5out, bank_path=args.bank_file - ) + template_ids = template_hash_to_id(trigger_file=h5out, bank_path=args.bank_file) logging.info("Storing template IDs") - h5out.create_dataset("network/template_id", - data=template_ids, - compression="gzip", - compression_opts=9) + h5out.create_dataset( + "network/template_id", + data=template_ids, + compression="gzip", + compression_opts=9, + ) # END OF WITH H5OUT logging.info("Merged triggers written to %s", outputfile) -def bin_events(inputfile, bins, outdir, filetag, - column="network/end_time_gc", verbose=False): - """Separate events in the inputfile into bins - """ +def bin_events( + inputfile, bins, outdir, filetag, column="network/end_time_gc", verbose=False +): + """Separate events in the inputfile into bins""" ifotag, _, seg = filename_metadata(inputfile) with HFile(inputfile, "r") as h5in: @@ -220,9 +232,11 @@ def bin_events(inputfile, bins, outdir, filetag, for bin_, segl in bins.items(): # find which network triggers to keep if isinstance(segl, list): + def _in_bin(t): return t in segl else: + def _in_bin(t): return segl[0] <= t < segl[1] @@ -232,27 +246,26 @@ def bin_events(inputfile, bins, outdir, filetag, # find which single-ifo events to keep ifo_index = { ifo: numpy.unique( - h5in["network/{}_event_id".format(ifo)][()][include], - ) for ifo in ifos + h5in[f"network/{ifo}_event_id"][()][include], + ) + for ifo in ifos } # generate output file outf = os.path.join( outdir, - '{}-{}_{}-{}-{}.h5'.format( - ifotag, filetag, bin_, seg[0], abs(seg), - ), + f"{ifotag}-{filetag}_{bin_}-{seg[0]}-{abs(seg)}.h5", ) - nsets = sum(isinstance(item, h5py.Dataset) for - group in h5in.values() for - item in group.values()) - msg = "Slicing {} network events for {}".format( - include.sum(), - bin_, + nsets = sum( + isinstance(item, h5py.Dataset) + for group in h5in.values() + for item in group.values() + ) + msg = f"Slicing {include.sum()} network events for {bin_}" + bar = tqdm.tqdm( + total=nsets, desc=msg, disable=not verbose, unit="datasets", **TQDM_KW ) - bar = tqdm.tqdm(total=nsets, desc=msg, disable=not verbose, - unit="datasets", **TQDM_KW) with HFile(outf, "w") as h5out: for old in h5in["network"].values(): if isinstance(old, h5py.Dataset): @@ -290,13 +303,13 @@ def bin_events(inputfile, bins, outdir, filetag, def read_segment_files(segdir): segs = {} for name, filename in { - "buffer": "bufferSeg.txt", - "off": "offSourceSeg.txt", - "on": "onSourceSeg.txt", + "buffer": "bufferSeg.txt", + "off": "offSourceSeg.txt", + "on": "onSourceSeg.txt", }.items(): try: - with open(os.path.join(segdir, filename), "r") as f: - segs[name], = fromsegwizard(f) + with open(os.path.join(segdir, filename)) as f: + (segs[name],) = fromsegwizard(f) except ValueError as exc: exc.args = ("more than one segment, an error has occurred",) raise @@ -328,7 +341,7 @@ parser.add_argument( "--job-tag", type=str.upper, help="the job tag, for use when more than one trig_combiner " - "job is included in a workflow", + "job is included in a workflow", ) parser.add_argument( "-S", @@ -339,7 +352,8 @@ parser.add_argument( # run parameters parser.add_argument( - "-n", "--grb-name", + "-n", + "--grb-name", type=str.upper, help="GRB event name, e.g. 010203", ) @@ -418,9 +432,9 @@ init_logging(args.verbose) logging.info("Welcome to the PyGRB trigger combiner") if args.grb_name: - args.user_tag += "_GRB{}".format(args.grb_name) + args.user_tag += f"_GRB{args.grb_name}" if args.job_tag: - args.user_tag += "_{}".format(args.job_tag) + args.user_tag += f"_{args.job_tag}" analysis = segments.segmentlist([file_segment(args.input_files[0])]) start, end = analysis[0] @@ -447,7 +461,7 @@ logging.info(" on-source segment : %s", bins["ONSOURCE"]) logging.info(" buffer segment : %s", segs["buffer"]) logging.info( " off-source segments : [%s]", - ", \n ".join(map(str, bins["OFFSOURCE"])) + ", \n ".join(map(str, bins["OFFSOURCE"])), ) # The onsource should never be the first or the last bin: @@ -475,10 +489,11 @@ for i, j in enumerate(offtrials): j -= offsource_ntrials[0] _ts += j * trialtime _te = _ts + trialtime - bins["OFFTRIAL_{}".format(i+1)] = seg = segments.segment(_ts, _te) + bins[f"OFFTRIAL_{i + 1}"] = seg = segments.segment(_ts, _te) if seg not in bins["OFFSOURCE"]: - raise ValueError(f"off-trial {i+1} not in off-source segments\n" - f"off-trial {i+1} : {seg}\n") + raise ValueError( + f"off-trial {i + 1} not in off-source segments\noff-trial {i + 1} : {seg}\n" + ) logging.info(" off-trial %d : %s", i + 1, seg) # -- read triggers ------------------------------ @@ -488,13 +503,9 @@ logging.info("Merging events") if args.short_slides and args.long_slides: raise NotImplementedError -outfilename = "{}-{}_ALL_TIMES-{}-{}.h5".format( - args.ifo_tag, args.user_tag, start, end-start, -) +outfilename = f"{args.ifo_tag}-{args.user_tag}_ALL_TIMES-{start}-{end - start}.h5" outfile = os.path.join(args.output_dir, outfilename) -merge_hdf5_files( - args.input_files, outfile, verbose=args.verbose, **compression_kw -) +merge_hdf5_files(args.input_files, outfile, verbose=args.verbose, **compression_kw) logging.info("Binning events") bin_events(outfile, bins, args.output_dir, args.user_tag, verbose=args.verbose) diff --git a/bin/pygrb/pycbc_make_offline_grb_workflow b/bin/pygrb/pycbc_make_offline_grb_workflow index 3331791ae0c..f17f76f8eb4 100644 --- a/bin/pygrb/pycbc_make_offline_grb_workflow +++ b/bin/pygrb/pycbc_make_offline_grb_workflow @@ -20,21 +20,22 @@ Make workflow for the archival, targeted, coherent inspiral pipeline. """ -import sys -import os import argparse import logging +import os +import sys + import matplotlib -matplotlib.use('agg') + +matplotlib.use("agg") from igwn_segments import segment, segmentlist, segmentlistdict -from pycbc import init_logging, add_common_pycbc_options -from pycbc.events.veto import select_segments_by_definer import pycbc.workflow as _workflow -from pycbc.workflow.core import configparser_value_to_file, FileList, SegFile +from pycbc import add_common_pycbc_options, init_logging +from pycbc.events.veto import select_segments_by_definer from pycbc.results.pygrb_plotting_utils import make_grb_segments_plot - +from pycbc.workflow.core import FileList, SegFile, configparser_value_to_file from pycbc.workflow.injection import compute_inj_optimal_snr from pycbc.workflow.psd import make_psd_file @@ -65,22 +66,17 @@ os.chdir(output_dir) # SEGMENTS triggertime = int(wflow.cp.get("workflow", "trigger-time")) -start = triggertime - int(wflow.cp.get("workflow-exttrig_segments", - "max-duration")) -end = triggertime + int(wflow.cp.get("workflow-exttrig_segments", - "max-duration")) +start = triggertime - int(wflow.cp.get("workflow-exttrig_segments", "max-duration")) +end = triggertime + int(wflow.cp.get("workflow-exttrig_segments", "max-duration")) wflow.cp = _workflow.set_grb_start_end(wflow.cp, start, end) # Retrieve science segments curr_dir = os.getcwd() seg_dir = os.path.join(curr_dir, "segments") -sciSegsFile = _workflow.get_segments_file(wflow, - 'science', - 'segments-science', - seg_dir) +sciSegsFile = _workflow.get_segments_file(wflow, "science", "segments-science", seg_dir) sciSegs = {} for ifo in wflow.ifos: - sciSegs[ifo] = sciSegsFile.segment_dict[ifo+':science'] + sciSegs[ifo] = sciSegsFile.segment_dict[ifo + ":science"] # This block of code and PyCBCMultiInspiralExecutable.get_valid_times() # must be consistent @@ -88,18 +84,16 @@ if wflow.cp.has_option("inspiral", "segment-start-pad"): pad_data = int(wflow.cp.get("inspiral", "pad-data")) start_pad = int(wflow.cp.get("inspiral", "segment-start-pad")) end_pad = int(wflow.cp.get("inspiral", "segment-end-pad")) - wflow.cp.set("workflow-exttrig_segments", "min-before", - str(start_pad+pad_data)) - wflow.cp.set("workflow-exttrig_segments", "min-after", - str(end_pad+pad_data)) + wflow.cp.set("workflow-exttrig_segments", "min-before", str(start_pad + pad_data)) + wflow.cp.set("workflow-exttrig_segments", "min-after", str(end_pad + pad_data)) elif wflow.cp.has_option("inspiral", "analyse-segment-end"): safety = 1 deadtime = int(wflow.cp.get("inspiral", "segment-length")) / 2 spec_len = int(wflow.cp.get("inspiral", "inverse-spec-length")) / 2 - wflow.cp.set("workflow-exttrig_segments", "min-before", - str(deadtime - spec_len - safety)) - wflow.cp.set("workflow-exttrig_segments", "min-after", - str(spec_len + safety)) + wflow.cp.set( + "workflow-exttrig_segments", "min-before", str(deadtime - spec_len - safety) + ) + wflow.cp.set("workflow-exttrig_segments", "min-after", str(spec_len + safety)) else: deadtime = int(wflow.cp.get("inspiral", "segment-length")) / 4 wflow.cp.set("workflow-exttrig_segments", "min-before", str(deadtime)) @@ -108,13 +102,13 @@ else: # Do checks for no/single IFO case single_ifo = wflow.cp.has_option("workflow", "allow-single-ifo-search") if len(sciSegs.keys()) == 0: - make_grb_segments_plot(wflow, segmentlistdict(), triggertime, - triggername, seg_dir) + make_grb_segments_plot(wflow, segmentlistdict(), triggertime, triggername, seg_dir) logging.error("No science segments available.") sys.exit() elif len(sciSegs.keys()) < 2 and not single_ifo: - make_grb_segments_plot(wflow, segmentlistdict(sciSegs), - triggertime, triggername, seg_dir) + make_grb_segments_plot( + wflow, segmentlistdict(sciSegs), triggertime, triggername, seg_dir + ) msg = "Science segments exist only for %s. " % tuple(sciSegs.keys())[0] msg += "If you wish to enable single IFO running add the option " msg += "'allow-single-ifo-search' to the [workflow] section of your " @@ -122,9 +116,7 @@ elif len(sciSegs.keys()) < 2 and not single_ifo: logging.error(msg) sys.exit() -onSrc, offSrc, bufferSeg = _workflow.generate_triggered_segment( - wflow, seg_dir, sciSegs -) +onSrc, offSrc, bufferSeg = _workflow.generate_triggered_segment(wflow, seg_dir, sciSegs) sciSegs = segmentlistdict(sciSegs) if onSrc is None: @@ -132,15 +124,19 @@ if onSrc is None: # `make_grb_segments_plot()` wants a single segment for fail_criterion. fail_criterion = offSrc[wflow.ifos[0]][0] make_grb_segments_plot( - wflow, sciSegs, triggertime, triggername, - seg_dir, fail_criterion=fail_criterion) + wflow, sciSegs, triggertime, triggername, seg_dir, fail_criterion=fail_criterion + ) sys.exit() plot_met = make_grb_segments_plot( - wflow, sciSegs, triggertime, triggername, seg_dir, - coherent_seg=offSrc[tuple(offSrc.keys())[0]][0]) -segs_plot = _workflow.File(plot_met[0], plot_met[1], plot_met[2], - file_url=plot_met[3]) + wflow, + sciSegs, + triggertime, + triggername, + seg_dir, + coherent_seg=offSrc[tuple(offSrc.keys())[0]][0], +) +segs_plot = _workflow.File(plot_met[0], plot_met[1], plot_met[2], file_url=plot_met[3]) segs_plot.add_pfn(segs_plot.cache_entry.path, site="local") sciSegs = offSrc all_files.append(segs_plot) @@ -153,8 +149,9 @@ elif len(sciSegs) > 1: # Update analysis time after coherent segment calculation ifo = tuple(sciSegs.keys())[0] -wflow.cp = _workflow.set_grb_start_end(wflow.cp, int(sciSegs[ifo][0][0]), - int(sciSegs[ifo][0][1])) +wflow.cp = _workflow.set_grb_start_end( + wflow.cp, int(sciSegs[ifo][0][0]), int(sciSegs[ifo][0][1]) +) padding = int(wflow.cp.get("inspiral", "pad-data")) if wflow.cp.has_option("workflow-condition_strain", "do-gating"): @@ -164,32 +161,33 @@ if wflow.cp.has_option("workflow-condition_strain", "do-gating"): if wflow.cp.has_option("inspiral", "segment-start-pad"): start_pad = int(wflow.cp.get("inspiral", "segment-start-pad")) end_pad = int(wflow.cp.get("inspiral", "segment-end-pad")) - wflow.analysis_time = segment(int(sciSegs[ifo][0][0]) + - start_pad + padding, - int(sciSegs[ifo][0][1]) - - padding - end_pad) + wflow.analysis_time = segment( + int(sciSegs[ifo][0][0]) + start_pad + padding, + int(sciSegs[ifo][0][1]) - padding - end_pad, + ) elif wflow.cp.has_option("inspiral", "analyse-segment-end"): - wflow.analysis_time = segment(int(sciSegs[ifo][0][0]) + deadtime - - spec_len + padding - safety, - int(sciSegs[ifo][0][1]) - spec_len - - padding - safety) + wflow.analysis_time = segment( + int(sciSegs[ifo][0][0]) + deadtime - spec_len + padding - safety, + int(sciSegs[ifo][0][1]) - spec_len - padding - safety, + ) else: - wflow.analysis_time = segment(int(sciSegs[ifo][0][0]) + deadtime + padding, - int(sciSegs[ifo][0][1]) - deadtime - padding) + wflow.analysis_time = segment( + int(sciSegs[ifo][0][0]) + deadtime + padding, + int(sciSegs[ifo][0][1]) - deadtime - padding, + ) # DATAFIND df_dir = os.path.join(curr_dir, "datafind") -datafind_files, _, sciSegs, _ = _workflow.setup_datafind_workflow(wflow, - sciSegs, - df_dir, - sciSegsFile) +datafind_files, _, sciSegs, _ = _workflow.setup_datafind_workflow( + wflow, sciSegs, df_dir, sciSegsFile +) if wflow.cp.has_option("workflow-condition_strain", "do-gating"): - new_seg = segment(sciSegs[ifo][0][0] + gate_pad, - sciSegs[ifo][0][1] - gate_pad) + new_seg = segment(sciSegs[ifo][0][0] + gate_pad, sciSegs[ifo][0][1] - gate_pad) for iifo in sciSegs: sciSegs[iifo][0] = new_seg - wflow.cp = _workflow.set_grb_start_end(wflow.cp, int(sciSegs[ifo][0][0]), - int(sciSegs[ifo][0][1])) + wflow.cp = _workflow.set_grb_start_end( + wflow.cp, int(sciSegs[ifo][0][0]), int(sciSegs[ifo][0][1]) + ) ifos = sorted(sciSegs.keys()) wflow.ifos = ifos @@ -208,13 +206,13 @@ ifo = ifos[0] # GATING if wflow.cp.has_option("workflow-condition_strain", "do-gating"): logging.info("Creating gating jobs.") - wflow.cp = _workflow.set_grb_start_end(wflow.cp, int(sciSegs[ifo][0][0]), - int(sciSegs[ifo][0][1])) - gating_nodes, gated_files = _workflow.make_gating_node(wflow, - datafind_files, - outdir=df_dir) - gating_method = wflow.cp.get("workflow-condition_strain", - "gating-method") + wflow.cp = _workflow.set_grb_start_end( + wflow.cp, int(sciSegs[ifo][0][0]), int(sciSegs[ifo][0][1]) + ) + gating_nodes, gated_files = _workflow.make_gating_node( + wflow, datafind_files, outdir=df_dir + ) + gating_method = wflow.cp.get("workflow-condition_strain", "gating-method") for gating_node in gating_nodes: if gating_method == "IN_WORKFLOW": wflow.add_node(gating_node) @@ -229,29 +227,31 @@ if wflow.cp.has_option("workflow-condition_strain", "do-gating"): sys.exit() datafind_files = _workflow.FileList([]) for ifo in ifos: - gated_frames = _workflow.FileList([gated_frame for gated_frame in - gated_files - if gated_frame.ifo == ifo]) + gated_frames = _workflow.FileList( + [gated_frame for gated_frame in gated_files if gated_frame.ifo == ifo] + ) # TODO: Remove .lcf cache here gated_cache = _workflow.File( - ifo, "gated", - segment(int(wflow.cp.get("workflow", "start-time")), - int(wflow.cp.get("workflow", "end-time"))), - extension="lcf", directory=df_dir) + ifo, + "gated", + segment( + int(wflow.cp.get("workflow", "start-time")), + int(wflow.cp.get("workflow", "end-time")), + ), + extension="lcf", + directory=df_dir, + ) gated_cache.add_pfn(gated_cache.cache_entry.path, site="local") - gated_frames.convert_to_lal_cache().tofile( - open(gated_cache.storage_path, "w")) + gated_frames.convert_to_lal_cache().tofile(open(gated_cache.storage_path, "w")) datafind_files.append(gated_cache) input_files.extend(datafind_files) # Retrieve vetoes veto_file = None if wflow.cp.has_option("workflow-segments", "segments-vetoes"): - veto_file = _workflow.get_segments_file(wflow, - 'vetoes', - 'segments-vetoes', - seg_dir, - tags=['veto']) + veto_file = _workflow.get_segments_file( + wflow, "vetoes", "segments-vetoes", seg_dir, tags=["veto"] + ) input_files.append(veto_file) # Check that the onsource is free from vetoes for ifo in ifos: @@ -266,25 +266,28 @@ for ifo in ifos: "The onsource %s contains the time %s that is vetoed in %s.", onSrc[ifo][0][:], intersection, - ifo + ifo, ) sys.exit() # Generate sky grid if needed skygrid_file = None -if wflow.cp.has_option("workflow", "sky-error") or wflow.cp.has_option("workflow", "input-dist"): +if wflow.cp.has_option("workflow", "sky-error") or wflow.cp.has_option( + "workflow", "input-dist" +): logging.info("Generating sky-grid file.") - skygrid_file = _workflow.make_skygrid_node(wflow, df_dir, tags=['SEARCH']) + skygrid_file = _workflow.make_skygrid_node(wflow, df_dir, tags=["SEARCH"]) input_files.extend(skygrid_file) # Config file consistency check for IPN GRBs -if wflow.cp.has_option("workflow-inspiral", "ipn-search-points") \ - and wflow.cp.has_option("workflow-injections", "ipn-sim-points"): - wflow.cp.set("injections", "ipn-gps-time", - wflow.cp.get("workflow", "trigger-time")) +if wflow.cp.has_option( + "workflow-inspiral", "ipn-search-points" +) and wflow.cp.has_option("workflow-injections", "ipn-sim-points"): + wflow.cp.set("injections", "ipn-gps-time", wflow.cp.get("workflow", "trigger-time")) IPN = True -elif wflow.cp.has_option("workflow-inspiral", "ipn-search-points") \ - or wflow.cp.has_option("workflow-injections", "ipn-sim-points"): +elif wflow.cp.has_option( + "workflow-inspiral", "ipn-search-points" +) or wflow.cp.has_option("workflow-injections", "ipn-sim-points"): msg = "You have provided only one of 'ipn-search-points' under " msg += "[workflow-inspiral] and 'ipn-sim-points' under " msg += "[workflow-injections] in your configuration files. If this is an " @@ -295,40 +298,38 @@ else: IPN = False # Get bank_veto_bank.xml if running bank veto -if wflow.cp.has_option('workflow-inspiral', 'bank-veto-bank-file'): - bank_veto_file = configparser_value_to_file(wflow.cp, 'workflow-inspiral', - 'bank-veto-bank-file') - bank_veto_file.description += '_BANK_VETO_BANK' +if wflow.cp.has_option("workflow-inspiral", "bank-veto-bank-file"): + bank_veto_file = configparser_value_to_file( + wflow.cp, "workflow-inspiral", "bank-veto-bank-file" + ) + bank_veto_file.description += "_BANK_VETO_BANK" bank_veto_file = _workflow.FileList([bank_veto_file]) input_files.extend(bank_veto_file) if IPN: file_attrs = { - 'ifos': wflow.ifos, - 'segs': wflow.analysis_time, - 'exe_name': "IPN_SKY_POINTS", - 'tags': ["SEARCH"] + "ifos": wflow.ifos, + "segs": wflow.analysis_time, + "exe_name": "IPN_SKY_POINTS", + "tags": ["SEARCH"], } - search_pts_file = configparser_value_to_file(wflow.cp, - 'workflow-inspiral', - 'ipn-search-points', - file_attrs=file_attrs) + search_pts_file = configparser_value_to_file( + wflow.cp, "workflow-inspiral", "ipn-search-points", file_attrs=file_attrs + ) input_files.append(search_pts_file) all_files.extend(input_files) # TEMPLATE BANK AND SPLIT BANK -bank_files = _workflow.setup_tmpltbank_workflow(wflow, sciSegs, - datafind_files, df_dir) +bank_files = _workflow.setup_tmpltbank_workflow(wflow, sciSegs, datafind_files, df_dir) # Check there are not multiple banks if len(bank_files) > 1: raise NotImplementedError("Multiple banks not supported") full_bank_file = bank_files[0] # Note: setup_splittable_workflow requires a FileList as input -splitbank_files = _workflow.setup_splittable_workflow(wflow, - bank_files, - df_dir, - tags=["inspiral"]) +splitbank_files = _workflow.setup_splittable_workflow( + wflow, bank_files, df_dir, tags=["inspiral"] +) all_files.append(full_bank_file) all_files.extend(splitbank_files) @@ -337,12 +338,12 @@ all_files.extend(splitbank_files) if wflow.cp.has_section("calculate_psd"): psd_dir = os.path.join(output_dir, "psds") os.makedirs(psd_dir, exist_ok=True) - + tv_psd_files = FileList([]) - + for ifo in wflow.ifos: logging.info("Creating time-varying PSD job for %s", ifo) - + # Segment file containing the FULL analysis segment psd_segfile = SegFile.from_segment_list( description=f"psd_{ifo}", @@ -351,19 +352,19 @@ if wflow.cp.has_section("calculate_psd"): ifo=ifo, seg_summ_list=segmentlist([wflow.analysis_time]), directory=psd_dir, - extension="xml" + extension="xml", ) - + psd_file = make_psd_file( workflow=wflow, - frame_files=None, # ok: datafind cache already in workflow + frame_files=None, # ok: datafind cache already in workflow segment_file=psd_segfile, segment_name="psd", out_dir=psd_dir, ) - + tv_psd_files.append(psd_file) - + logging.info("Created %d time-varying PSDs", len(tv_psd_files)) # INJECTIONS @@ -381,10 +382,10 @@ if wflow.cp.has_section("workflow-injections"): # Given the stretch of time this workflow will analyse, and the onsource # window with its buffer, generate the configuration file with the prior # for the injections times and add it to the config parser - inj_method = wflow.cp.get("workflow-injections", - "injections-method") - if inj_method == "IN_WORKFLOW" and \ - wflow.cp.has_option("workflow-injections", "tc-prior-at-runtime"): + inj_method = wflow.cp.get("workflow-injections", "injections-method") + if inj_method == "IN_WORKFLOW" and wflow.cp.has_option( + "workflow-injections", "tc-prior-at-runtime" + ): tc_path = os.path.join(output_dir, "tc_prior.ini") _workflow.generate_tc_prior(wflow, tc_path, bufferSeg) @@ -396,15 +397,14 @@ if wflow.cp.has_section("workflow-injections"): # setup_injection_workflow or pycbc_create_injections handle it # directly via configparser_value_to_file file_attrs = { - 'ifos': wflow.ifos, - 'segs': wflow.analysis_time, - 'exe_name': "IPN_SKY_POINTS", - 'tags': ["SIM"] + "ifos": wflow.ifos, + "segs": wflow.analysis_time, + "exe_name": "IPN_SKY_POINTS", + "tags": ["SIM"], } - sim_pts_file = configparser_value_to_file(wflow.cp, - 'workflow-inspiral', - 'ipn-sim-points', - file_attrs=file_attrs) + sim_pts_file = configparser_value_to_file( + wflow.cp, "workflow-inspiral", "ipn-sim-points", file_attrs=file_attrs + ) all_files.append(sim_pts_file) inj_files, inj_tags = _workflow.setup_injection_workflow(wflow, inj_dir) all_files.extend(inj_files) @@ -412,23 +412,21 @@ if wflow.cp.has_section("workflow-injections"): # Here we compute the optimal SNR for the injections before splitting them if wflow.cp.has_section("optimal_snr"): - if not wflow.cp.has_section('calculate_psd'): - raise ValueError('Optimal SNR calculation requires PSD estimation') - + if not wflow.cp.has_section("calculate_psd"): + raise ValueError("Optimal SNR calculation requires PSD estimation") + snr_dir = os.path.join(inj_dir, "optimal_snr") - os.makedirs(snr_dir, exist_ok = True) - + os.makedirs(snr_dir, exist_ok=True) + optimal_snr_files = _workflow.FileList([]) for inj, inj_tag in zip(injs, inj_tags): snr_file = compute_inj_optimal_snr( - wflow, - inj, - tv_psd_files, - snr_dir, - tags=[inj_tag] + wflow, inj, tv_psd_files, snr_dir, tags=[inj_tag] ) optimal_snr_files.append(snr_file) - logging.info("Optimal SNR workflow node created for injection tag %s", inj_tag) + logging.info( + "Optimal SNR workflow node created for injection tag %s", inj_tag + ) all_files.extend(optimal_snr_files) injs = optimal_snr_files inj_files = optimal_snr_files @@ -437,14 +435,19 @@ if wflow.cp.has_section("workflow-injections"): # as for standard matched filter jobs if wflow.cp.has_section("workflow-splittable-injections"): inj_splitbank_files = _workflow.setup_splittable_workflow( - wflow, bank_files, inj_dir, tags=["injections"]) + wflow, bank_files, inj_dir, tags=["injections"] + ) for inj_split in inj_splitbank_files: - split_str = [s for s in inj_split.tagged_description.split("_") - if ("BANK" in s and s[-1].isdigit())] + split_str = [ + s + for s in inj_split.tagged_description.split("_") + if ("BANK" in s and s[-1].isdigit()) + ] if len(split_str) != 0: inj_split.tagged_description += "%s_%d" % ( - inj_split.tag_str, - int(split_str[0].replace("BANK", ""))) + inj_split.tag_str, + int(split_str[0].replace("BANK", "")), + ) all_files.extend(inj_splitbank_files) else: inj_splitbank_files = _workflow.FileList([]) @@ -456,15 +459,18 @@ if wflow.cp.has_section("workflow-injections"): for inj_file, inj_tag in zip(inj_files, inj_tags): file = _workflow.FileList([inj_file]) inj_splits = _workflow.setup_splittable_workflow( - wflow, file, inj_dir, tags=["split_inspinj", inj_tag]) + wflow, file, inj_dir, tags=["split_inspinj", inj_tag] + ) for inj_split in inj_splits: - split_str = [s for s in - inj_split.tagged_description.split("_") - if ("SPLIT" in s and s[-1].isdigit())] + split_str = [ + s + for s in inj_split.tagged_description.split("_") + if ("SPLIT" in s and s[-1].isdigit()) + ] if len(split_str) != 0: new = inj_split.tagged_description.replace( - split_str[0], - "SPLIT_%s" % split_str[0].replace("SPLIT", "")) + split_str[0], "SPLIT_%s" % split_str[0].replace("SPLIT", "") + ) inj_split.tagged_description = new inj_split_files.extend(inj_splits) all_files.extend(inj_split_files) @@ -472,32 +478,51 @@ if wflow.cp.has_section("workflow-injections"): # Generate injection matched filter workflow inj_insp_files = _workflow.setup_matchedfltr_workflow( - wflow, sciSegs, input_files, inj_splitbank_files, - inj_dir, injs, tags=[mf_tag + "_injections"]) + wflow, + sciSegs, + input_files, + inj_splitbank_files, + inj_dir, + injs, + tags=[mf_tag + "_injections"], + ) for inj_insp_file in inj_insp_files: - split_str = [s for s in inj_insp_file.name.split("_") - if ("SPLIT" in s and s[-1].isdigit())] + split_str = [ + s + for s in inj_insp_file.name.split("_") + if ("SPLIT" in s and s[-1].isdigit()) + ] if len(split_str) != 0: num = split_str[0].replace("SPLIT", "_") inj_insp_file.tagged_description += num # Make cache files (needed for post-processing) for inj_tag in inj_tags: - files = _workflow.FileList([file for file in injs - if inj_tag in file.tag_str]) - inj_cache = _workflow.File(ifos, "injections", sciSegs[ifo][0], - extension="lcf", directory=inj_dir, - tags=[inj_tag]) + files = _workflow.FileList([file for file in injs if inj_tag in file.tag_str]) + inj_cache = _workflow.File( + ifos, + "injections", + sciSegs[ifo][0], + extension="lcf", + directory=inj_dir, + tags=[inj_tag], + ) inj_cache.add_pfn(inj_cache.cache_entry.path, site="local") inj_caches.append(inj_cache) inj_cache_entries = files.convert_to_lal_cache() inj_cache_entries.tofile(open(inj_cache.storage_path, "w")) - files = _workflow.FileList([file for file in inj_insp_files - if inj_tag in file.tag_str]) - inj_insp_cache = _workflow.File(ifos, "inspiral_injections", - sciSegs[ifo][0], extension="lcf", - directory=inj_dir, tags=[inj_tag]) + files = _workflow.FileList( + [file for file in inj_insp_files if inj_tag in file.tag_str] + ) + inj_insp_cache = _workflow.File( + ifos, + "inspiral_injections", + sciSegs[ifo][0], + extension="lcf", + directory=inj_dir, + tags=[inj_tag], + ) inj_insp_cache.add_pfn(inj_insp_cache.cache_entry.path, site="local") inj_insp_caches.append(inj_insp_cache) inj_insp_cache_entries = files.convert_to_lal_cache() @@ -510,13 +535,18 @@ if wflow.cp.has_section("workflow-injections"): # MAIN MATCHED FILTERING insp_dir = os.path.join(curr_dir, "inspiral") inspiral_files = _workflow.setup_matchedfltr_workflow( - wflow, sciSegs, - input_files, splitbank_files, insp_dir, - tags=[mf_tag + "_no_injections"]) + wflow, + sciSegs, + input_files, + splitbank_files, + insp_dir, + tags=[mf_tag + "_no_injections"], +) all_files.extend(inspiral_files) # TODO: Remove .lcf caches here? -inspiral_cache = _workflow.File(ifos, "inspiral", sciSegs[ifo][0], - extension="lcf", directory=insp_dir) +inspiral_cache = _workflow.File( + ifos, "inspiral", sciSegs[ifo][0], extension="lcf", directory=insp_dir +) inspiral_cache.add_pfn(inspiral_cache.cache_entry.path, site="local") all_files.append(inspiral_cache) inspiral_cache_entries = inspiral_files.convert_to_lal_cache() @@ -527,31 +557,41 @@ inspiral_cache_entries.tofile(open(inspiral_cache.storage_path, "w")) # POST-PROCESSING pp_dir = os.path.join(curr_dir, "post_processing") os.makedirs(pp_dir) -post_proc_method = wflow.cp.get_opt_tags("workflow-postproc", - "postproc-method", tags) +post_proc_method = wflow.cp.get_opt_tags("workflow-postproc", "postproc-method", tags) pp_files = _workflow.FileList([]) results_files = _workflow.FileList([]) if post_proc_method == "PYGRB_OFFLINE": - trig_comb_files, clustered_files, inj_find_files =\ - _workflow.setup_pygrb_pp_workflow(wflow, pp_dir, seg_dir, - sciSegs[ifo][0], full_bank_file, - inspiral_files, injs, - inj_insp_files, inj_tags) - sec_name = 'workflow-pygrb_results_workflow' + trig_comb_files, clustered_files, inj_find_files = ( + _workflow.setup_pygrb_pp_workflow( + wflow, + pp_dir, + seg_dir, + sciSegs[ifo][0], + full_bank_file, + inspiral_files, + injs, + inj_insp_files, + inj_tags, + ) + ) + sec_name = "workflow-pygrb_results_workflow" if not wflow.cp.has_section(sec_name): - msg = 'No {0} section found in configuration file.'.format(sec_name) + msg = f"No {sec_name} section found in configuration file." logging.info(msg) else: - logging.info('Entering results module') - results_files = _workflow.setup_pygrb_results_workflow(wflow, pp_dir, - clustered_files, - inj_files, - inj_find_files, - full_bank_file, - seg_dir, - skygrid_file[0], - veto_file=veto_file) - logging.info('Leaving results module') + logging.info("Entering results module") + results_files = _workflow.setup_pygrb_results_workflow( + wflow, + pp_dir, + clustered_files, + inj_files, + inj_find_files, + full_bank_file, + seg_dir, + skygrid_file[0], + veto_file=veto_file, + ) + logging.info("Leaving results module") all_files.extend(pp_files) all_files.extend(results_files) diff --git a/bin/pygrb/pycbc_pygrb_efficiency b/bin/pygrb/pycbc_pygrb_efficiency index 7dee81b4002..4f900958a7f 100755 --- a/bin/pygrb/pycbc_pygrb_efficiency +++ b/bin/pygrb/pycbc_pygrb_efficiency @@ -21,25 +21,26 @@ # ============================================================================= # Preamble # ============================================================================= -import sys -import os -import logging import json +import logging +import os +import sys + import matplotlib.pyplot as plt -from matplotlib import rc import numpy as np +import pycbc.version import scipy +from matplotlib import rc from scipy import stats -import pycbc.version from pycbc import init_logging -from pycbc.detector import Detector -from pycbc.results import save_fig_with_metadata -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.conversions import mchirp_from_mass1_mass2 +from pycbc.detector import Detector from pycbc.io.hdf import HFile +from pycbc.results import pygrb_postprocessing_utils as ppu +from pycbc.results import save_fig_with_metadata -plt.switch_backend('Agg') +plt.switch_backend("Agg") rc("image") __author__ = "Francesco Pannarale " @@ -49,9 +50,10 @@ __program__ = "pycbc_pygrb_efficiency" def efficiency_with_errs(found_bestnr, num_injections, num_mc_injs=0): - """Function to calculate the fraction of recovered injections and its - error bars (used for efficiency/sensitive distance plots).""" - + """ + Function to calculate the fraction of recovered injections and its + error bars (used for efficiency/sensitive distance plots). + """ if not isinstance(num_mc_injs, int): err_msg = "The parameter num_mc_injs is the number of Monte-Carlo " err_msg += "injections. It must be an integer." @@ -68,22 +70,21 @@ def efficiency_with_errs(found_bestnr, num_injections, num_mc_injs=0): err_common = all_injs * (2 * only_found_injs + 1) err_denom = 2 * all_injs * (all_injs + 1) - err_vary = 4 * all_injs * only_found_injs * (all_injs - only_found_injs) \ - + all_injs**2 + err_vary = ( + 4 * all_injs * only_found_injs * (all_injs - only_found_injs) + all_injs**2 + ) err_vary = err_vary**0.5 - err_low = (err_common - err_vary)/err_denom + err_low = (err_common - err_vary) / err_denom err_low_mc = fraction - err_low - err_high = (err_common + err_vary)/err_denom + err_high = (err_common + err_vary) / err_denom err_high_mc = err_high - fraction # Check for cases where error bars are negative and set them to zero if (err_low_mc < 0).any(): - logging.warning("Negative lower error bar(s) detected." - " Setting to zero.") + logging.warning("Negative lower error bar(s) detected. Setting to zero.") err_low_mc[err_low_mc < 0] = 0 if (err_high_mc < 0).any(): - logging.warning("Negative upper error bar(s) detected." - " Setting to zero.") + logging.warning("Negative upper error bar(s) detected. Setting to zero.") err_high_mc[err_high_mc < 0] = 0 return err_low_mc, err_high_mc, fraction @@ -93,36 +94,74 @@ def efficiency_with_errs(found_bestnr, num_injections, num_mc_injs=0): # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("-F", "--trig-file", action="store", required=True, - help="Location of off-source trigger file.") -parser.add_argument("--onsource-file", action="store", - help="Location of on-source trigger file (or a " + - "background trigger file to be treated as such).") -parser.add_argument("--background-output-file", - help="Detection efficiency output file.") -parser.add_argument("--onsource-output-file", - help="Exclusion distance output file.") -parser.add_argument("--exclusion-dist-output-file", - help="JSON file containing exclusion distances and efficiency curves.") -parser.add_argument("-g", "--glitch-check-factor", action="store", - type=float, default=1.0, help="When deciding " + - "exclusion efficiencies this value is multiplied " + - "to the offsource around the injection trigger to " + - "determine if it is just a loud glitch.") -parser.add_argument("--found-missed-file", action="store", type=str, - required=True, help="Location of found/missed " + - "injections and trigger file") -parser.add_argument("--injection-set-name", action="store", type=str, - default="", help="Name of the injection set to be " + - "used in the plot title.") -parser.add_argument("--trial-name", action="store", type=str, - required=True, help="Name of trial used " + - "for this run (i.e. ONSOURCE, OFFTRIAL)") -parser.add_argument("-C", "--cluster-window", action="store", type=float, - default=0.1, help="The cluster window used " + - "to cluster triggers in time.") -parser.add_argument("--bank-file", action="store", type=str, required=True, - help="Location of the full template bank used.") +parser.add_argument( + "-F", + "--trig-file", + action="store", + required=True, + help="Location of off-source trigger file.", +) +parser.add_argument( + "--onsource-file", + action="store", + help="Location of on-source trigger file (or a " + "background trigger file to be treated as such).", +) +parser.add_argument( + "--background-output-file", help="Detection efficiency output file." +) +parser.add_argument("--onsource-output-file", help="Exclusion distance output file.") +parser.add_argument( + "--exclusion-dist-output-file", + help="JSON file containing exclusion distances and efficiency curves.", +) +parser.add_argument( + "-g", + "--glitch-check-factor", + action="store", + type=float, + default=1.0, + help="When deciding " + "exclusion efficiencies this value is multiplied " + "to the offsource around the injection trigger to " + "determine if it is just a loud glitch.", +) +parser.add_argument( + "--found-missed-file", + action="store", + type=str, + required=True, + help="Location of found/missed " + "injections and trigger file", +) +parser.add_argument( + "--injection-set-name", + action="store", + type=str, + default="", + help="Name of the injection set to be " + "used in the plot title.", +) +parser.add_argument( + "--trial-name", + action="store", + type=str, + required=True, + help="Name of trial used " + "for this run (i.e. ONSOURCE, OFFTRIAL)", +) +parser.add_argument( + "-C", + "--cluster-window", + action="store", + type=float, + default=0.1, + help="The cluster window used " + "to cluster triggers in time.", +) +parser.add_argument( + "--bank-file", + action="store", + type=str, + required=True, + help="Location of the full template bank used.", +) ppu.pygrb_add_injmc_opts(parser) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_slide_opts(parser) @@ -132,15 +171,16 @@ ppu.slide_opts_helper(opts) init_logging(opts.verbose, format="%(asctime)s: %(levelname)s: %(message)s") # Load bank file -bank_file = HFile(opts.bank_file, 'r') +bank_file = HFile(opts.bank_file, "r") # Check options -if opts.exclusion_dist_output_file is not None or \ - opts.onsource_output_file is not None: +if opts.exclusion_dist_output_file is not None or opts.onsource_output_file is not None: if opts.onsource_file is None: - logging.error("Requesting the --exclusion-dist-output-file or " + - "the --onsource-output-file requires the " + - "--onsource-file as nput.") + logging.error( + "Requesting the --exclusion-dist-output-file or " + "the --onsource-output-file requires the " + "--onsource-file as nput." + ) # Store options used multiple times in local variables trig_file = opts.trig_file @@ -150,17 +190,17 @@ veto_file = opts.veto_file inj_set_name = opts.injection_set_name wf_err = opts.waveform_error cal_errs = {} -cal_errs['G1'] = opts.g1_cal_error -cal_errs['H1'] = opts.h1_cal_error -cal_errs['K1'] = opts.k1_cal_error -cal_errs['L1'] = opts.l1_cal_error -cal_errs['V1'] = opts.v1_cal_error +cal_errs["G1"] = opts.g1_cal_error +cal_errs["H1"] = opts.h1_cal_error +cal_errs["K1"] = opts.k1_cal_error +cal_errs["L1"] = opts.l1_cal_error +cal_errs["V1"] = opts.v1_cal_error cal_dc_errs = {} -cal_dc_errs['G1'] = opts.g1_dc_cal_error -cal_dc_errs['H1'] = opts.h1_dc_cal_error -cal_dc_errs['K1'] = opts.k1_dc_cal_error -cal_dc_errs['L1'] = opts.l1_dc_cal_error -cal_dc_errs['V1'] = opts.v1_dc_cal_error +cal_dc_errs["G1"] = opts.g1_dc_cal_error +cal_dc_errs["H1"] = opts.h1_dc_cal_error +cal_dc_errs["K1"] = opts.k1_dc_cal_error +cal_dc_errs["L1"] = opts.l1_dc_cal_error +cal_dc_errs["V1"] = opts.v1_dc_cal_error # pycbc_multi_inspiral already applies sngl, coinc and null SNR cuts new_snr_thresh = opts.newsnr_threshold upper_dist = opts.upper_inj_dist @@ -176,8 +216,11 @@ logging.info("Setting random seed to %d.", opts.seed) # Set output directories outdir = None -for output_file in [opts.exclusion_dist_output_file, - opts.background_output_file, opts.onsource_output_file]: +for output_file in [ + opts.exclusion_dist_output_file, + opts.background_output_file, + opts.onsource_output_file, +]: if output_file is not None: outdir = os.path.split(os.path.abspath(output_file))[0] if not os.path.isdir(outdir): @@ -195,27 +238,23 @@ segment_dict = ppu.load_segment_dict(trig_file) # Construct trials removing vetoed times trial_dict, total_trials = ppu.construct_trials( - opts.seg_files, - segment_dict, - ifos, - slide_dict, - veto_file + opts.seg_files, segment_dict, ifos, slide_dict, veto_file ) # Load triggers (apply reweighted SNR cut, not vetoes) -all_off_trigs = ppu.load_data(trig_file, ifos, data_tag='offsource', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=opts.slide_id) +all_off_trigs = ppu.load_data( + trig_file, + ifos, + data_tag="offsource", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=opts.slide_id, +) # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors -keys = ['network/end_time_gc', 'network/reweighted_snr'] +keys = ["network/end_time_gc", "network/reweighted_snr"] trig_data = ppu.extract_trig_properties( - trial_dict, - all_off_trigs, - slide_dict, - segment_dict, - keys + trial_dict, all_off_trigs, slide_dict, segment_dict, keys ) # Max BestNR values in each trial: these are stored in a dictionary keyed @@ -234,9 +273,9 @@ for slide_id in slide_dict: # Max and median values of reweighted SNR, # and sorted (loudest in trial) reweighted SNR values -max_bestnr, median_bestnr, sorted_bkgd =\ - ppu.max_median_stat(slide_dict, background, trig_data[keys[1]], - total_trials) +max_bestnr, median_bestnr, sorted_bkgd = ppu.max_median_stat( + slide_dict, background, trig_data[keys[1]], total_trials +) assert total_trials == len(sorted_bkgd) logging.info("Background bestNR calculated.") @@ -244,12 +283,7 @@ logging.info("Background bestNR calculated.") # Output details of loudest offsouce triggers: only triggers compatible # with the trial_dict are considered offsource_trigs = [] -sorted_off_trigs = ppu.sort_trigs( - trial_dict, - all_off_trigs, - slide_dict, - segment_dict -) +sorted_off_trigs = ppu.sort_trigs(trial_dict, all_off_trigs, slide_dict, segment_dict) for slide_id in slide_dict: offsource_trigs.extend( zip(trig_data[keys[1]][slide_id], sorted_off_trigs[slide_id]) @@ -259,35 +293,37 @@ offsource_trigs.reverse() # Calculate chirp masses of templates in bank -logging.info('Reading template chirp masses') -with HFile(opts.bank_file, 'r') as bank_file: +logging.info("Reading template chirp masses") +with HFile(opts.bank_file, "r") as bank_file: template_mchirps = mchirp_from_mass1_mass2( - bank_file['mass1'][:], - bank_file['mass2'][:] + bank_file["mass1"][:], bank_file["mass2"][:] ) # ======================= # Load on source triggers # ======================= if onsource_file: - logging.info("Processing onsource.") # Load onsoource triggers (apply reweighted SNR cut, not vetoes) - on_trigs = ppu.load_data(onsource_file, ifos, data_tag=None, - rw_snr_threshold=opts.newsnr_threshold, - slide_id=0) + on_trigs = ppu.load_data( + onsource_file, + ifos, + data_tag=None, + rw_snr_threshold=opts.newsnr_threshold, + slide_id=0, + ) # Calculate chirp mass values - on_mchirp = template_mchirps[on_trigs['network/template_id']] + on_mchirp = template_mchirps[on_trigs["network/template_id"]] # Set loudest event arrays loud_on_bestnr = 0 # Retrieve BestNRs and record loudest trig by BestNR. # Get indices of loudest triggers and pick the loudest. - if on_trigs and on_trigs['network/reweighted_snr'].size > 0: - loud_on_bestnr_idx = np.argmax(on_trigs['network/reweighted_snr']) - loud_on_bestnr = np.max(on_trigs['network/reweighted_snr']) + if on_trigs and on_trigs["network/reweighted_snr"].size > 0: + loud_on_bestnr_idx = np.argmax(on_trigs["network/reweighted_snr"]) + loud_on_bestnr = np.max(on_trigs["network/reweighted_snr"]) # Convert to float for output loud_on_bestnr = float(loud_on_bestnr) @@ -308,57 +344,55 @@ if onsource_file: # ==================== # injs contains found/missed injections AND triggers they generated # The reweighted SNR cut is applied, vetoes are not -injs = ppu.load_data(found_missed_file, ifos, data_tag='injs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=0) +injs = ppu.load_data( + found_missed_file, + ifos, + data_tag="injs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=0, +) # Gather injections that were not missed found_inj = {} for k in injs.keys(): - if 'missed' not in k: + if "missed" not in k: found_inj[k] = injs[k] # Separate them in found surviving vetoes and found but vetoed found_after_vetoes, vetoed, *_ = ppu.apply_vetoes_to_found_injs( - found_missed_file, - found_inj, - ifos, - veto_file=veto_file + found_missed_file, found_inj, ifos, veto_file=veto_file ) # =================================================================== # Post-process injections: skip this if there are no found injections # =================================================================== -if len(found_after_vetoes['network/template_id']): +if len(found_after_vetoes["network/template_id"]): # Calculate quantities not included in trigger files, such as chirp mass - found_trig_mchirp = \ - template_mchirps[found_after_vetoes['network/template_id']] + found_trig_mchirp = template_mchirps[found_after_vetoes["network/template_id"]] # Construct conditions for injection: # 1) found (surviving vetoes) louder than background, - zero_fap = \ - np.zeros(len(found_after_vetoes['network/end_time_gc'])).astype(bool) - zero_fap_cut = found_after_vetoes['network/reweighted_snr'] > max_bestnr + zero_fap = np.zeros(len(found_after_vetoes["network/end_time_gc"])).astype(bool) + zero_fap_cut = found_after_vetoes["network/reweighted_snr"] > max_bestnr zero_fap = zero_fap | (zero_fap_cut) # 2) found (bestnr>0, and surviving vetoes) but not louder than background - nonzero_fap = ~zero_fap & (found_after_vetoes['network/reweighted_snr'] - != 0) + nonzero_fap = ~zero_fap & (found_after_vetoes["network/reweighted_snr"] != 0) # 3) missed after being recovered (i.e., vetoed) are in vetoed # Non-zero FAP triggers (g_ifar) g_ifar = {} - g_ifar['bestnr'] = \ - found_after_vetoes['network/reweighted_snr'][nonzero_fap] - g_ifar['stat'] = np.zeros([len(g_ifar['bestnr'])]) - for ix, (mc, bestnr) in \ - enumerate(zip(found_trig_mchirp[nonzero_fap], g_ifar['bestnr'])): - g_ifar['stat'][ix] = (sorted_bkgd > bestnr).sum() - g_ifar['stat'] = g_ifar['stat'] / total_trials + g_ifar["bestnr"] = found_after_vetoes["network/reweighted_snr"][nonzero_fap] + g_ifar["stat"] = np.zeros([len(g_ifar["bestnr"])]) + for ix, (mc, bestnr) in enumerate( + zip(found_trig_mchirp[nonzero_fap], g_ifar["bestnr"]) + ): + g_ifar["stat"][ix] = (sorted_bkgd > bestnr).sum() + g_ifar["stat"] = g_ifar["stat"] / total_trials # Set the sigma values - inj_sigma = {ifo: found_after_vetoes[f'{ifo}/sigmasq'][:] for ifo in ifos} + inj_sigma = {ifo: found_after_vetoes[f"{ifo}/sigmasq"][:] for ifo in ifos} # If the sigmasqs are not populated, we can still do calibration errors, # but only in the 1-detector case for ifo in ifos: @@ -371,19 +405,21 @@ if len(found_after_vetoes['network/template_id']): msg += "set to unity for all triggers in order to build the " msg += "calibration errors." logging.info(msg) - inj_sigma[ifo][:] = 1. + inj_sigma[ifo][:] = 1.0 f_resp = {} for ifo in ifos: antenna = Detector(ifo) f_resp[ifo] = ppu.get_antenna_responses( antenna, - found_after_vetoes['found/ra'][:], - found_after_vetoes['found/dec'][:], - found_after_vetoes['found/tc'][:]) + found_after_vetoes["found/ra"][:], + found_after_vetoes["found/dec"][:], + found_after_vetoes["found/tc"][:], + ) - inj_sigma_mult = (np.asarray(list(inj_sigma.values())) * - np.asarray(list(f_resp.values()))) + inj_sigma_mult = np.asarray(list(inj_sigma.values())) * np.asarray( + list(f_resp.values()) + ) inj_sigma_tot = inj_sigma_mult[0, :] for i in range(1, len(ifos)): @@ -391,8 +427,7 @@ if len(found_after_vetoes['network/template_id']): inj_sigma_mean = {} for ifo in ifos: - inj_sigma_mean[ifo] = \ - ((inj_sigma[ifo]*f_resp[ifo])/inj_sigma_tot).mean() + inj_sigma_mean[ifo] = ((inj_sigma[ifo] * f_resp[ifo]) / inj_sigma_tot).mean() msg = f"{len(found_after_vetoes['found/tc'])} injections found and " msg += f"surviving vetoes and {len(injs['missed/tc'])} missed injections " @@ -401,39 +436,48 @@ if len(found_after_vetoes['network/template_id']): # Create new set of injections for efficiency calculations: # these are as many as the original injections - total_injs = len(injs['found/distance']) + len(injs['missed/distance']) + total_injs = len(injs["found/distance"]) + len(injs["missed/distance"]) long_inj = {} - long_inj['dist'] = stats.uniform.rvs(size=total_injs) * \ - (upper_dist-lower_dist) + upper_dist + long_inj["dist"] = ( + stats.uniform.rvs(size=total_injs) * (upper_dist - lower_dist) + upper_dist + ) logging.info("%d long distance injections created.", total_injs) # Set distance bins and data arrays - dist_bins = zip(np.arange(lower_dist, upper_dist + (upper_dist-lower_dist), - (upper_dist-lower_dist)/num_bins), - np.arange(lower_dist, upper_dist + (upper_dist-lower_dist), - (upper_dist-lower_dist)/num_bins) + - (upper_dist-lower_dist)/num_bins) + dist_bins = zip( + np.arange( + lower_dist, + upper_dist + (upper_dist - lower_dist), + (upper_dist - lower_dist) / num_bins, + ), + np.arange( + lower_dist, + upper_dist + (upper_dist - lower_dist), + (upper_dist - lower_dist) / num_bins, + ) + + (upper_dist - lower_dist) / num_bins, + ) dist_bins = list(dist_bins) num_dist_bins_plus_one = len(dist_bins) + 1 num_injections = {} found_max_bestnr = {} found_on_bestnr = {} - for key in ['mc', 'no_mc']: + for key in ["mc", "no_mc"]: num_injections[key] = np.zeros(num_dist_bins_plus_one) found_max_bestnr[key] = np.zeros(num_dist_bins_plus_one) found_on_bestnr[key] = np.zeros(num_dist_bins_plus_one) # Construct FAP list for all found injections - inj_fap = np.zeros(len(found_after_vetoes['found/distance'])) - inj_fap[nonzero_fap] = g_ifar['stat'] + inj_fap = np.zeros(len(found_after_vetoes["found/distance"])) + inj_fap[nonzero_fap] = g_ifar["stat"] # Calculate the amplitude error # Begin by calculating the components from each detector cal_error = 0 for ifo in ifos: - cal_error += cal_errs[ifo]**2 * inj_sigma_mean[ifo]**2 + cal_error += cal_errs[ifo] ** 2 * inj_sigma_mean[ifo] ** 2 cal_error = cal_error**0.5 max_dc_cal_error = max(cal_dc_errs.values()) @@ -452,85 +496,92 @@ if len(found_after_vetoes['network/template_id']): # (this is an MC, order of operations matters!) found_inj_dist_mc = ppu.mc_cal_wf_errs( num_mc_injs, - found_after_vetoes['found/distance'], + found_after_vetoes["found/distance"], cal_error, wav_err, - max_dc_cal_error + max_dc_cal_error, ) missed_inj_dist_mc = ppu.mc_cal_wf_errs( num_mc_injs, - np.concatenate((vetoed['found/distance'], injs['missed/distance'])), + np.concatenate((vetoed["found/distance"], injs["missed/distance"])), cal_error, wav_err, - max_dc_cal_error + max_dc_cal_error, + ) + long_inj["dist_mc"] = ppu.mc_cal_wf_errs( + num_mc_injs, long_inj["dist"], cal_error, wav_err, max_dc_cal_error ) - long_inj['dist_mc'] = ppu.mc_cal_wf_errs(num_mc_injs, - long_inj['dist'], - cal_error, - wav_err, - max_dc_cal_error) - logging.info("MC injection set distributed with %d iterations.", - num_mc_injs) + logging.info("MC injection set distributed with %d iterations.", num_mc_injs) # Check injections against on source if onsource_file: - more_sig_than_onsource = (inj_fap <= loud_on_fap) + more_sig_than_onsource = inj_fap <= loud_on_fap else: - more_sig_than_onsource = (inj_fap <= 0.5) + more_sig_than_onsource = inj_fap <= 0.5 distance_count = np.zeros(len(dist_bins)) - found_trig_max_bestnr = \ - np.empty(len(found_after_vetoes['network/event_id'])) + found_trig_max_bestnr = np.empty(len(found_after_vetoes["network/event_id"])) found_trig_max_bestnr.fill(max_bestnr) - max_bestnr_cut = \ - (found_after_vetoes['network/reweighted_snr'] > found_trig_max_bestnr) + max_bestnr_cut = ( + found_after_vetoes["network/reweighted_snr"] > found_trig_max_bestnr + ) # Check louder than on source - found_trig_loud_on_bestnr = \ - np.empty(len(found_after_vetoes['network/event_id'])) + found_trig_loud_on_bestnr = np.empty(len(found_after_vetoes["network/event_id"])) if onsource_file: found_trig_loud_on_bestnr.fill(loud_on_bestnr) else: found_trig_loud_on_bestnr.fill(median_bestnr) - on_bestnr_cut = found_after_vetoes['network/reweighted_snr'] > \ - found_trig_loud_on_bestnr + on_bestnr_cut = ( + found_after_vetoes["network/reweighted_snr"] > found_trig_loud_on_bestnr + ) # Check whether injection is found for the purposes of exclusion # distance calculation. # Found: if louder than all on source # Missed: if not louder than loudest on source - found_excl = on_bestnr_cut & (more_sig_than_onsource) & \ - (found_after_vetoes['network/reweighted_snr'] != 0) + found_excl = ( + on_bestnr_cut + & (more_sig_than_onsource) + & (found_after_vetoes["network/reweighted_snr"] != 0) + ) # If not missed, double check bestnr against nearby triggers near_test = np.zeros((found_excl).sum()).astype(bool) - for j, (t, bestnr) in enumerate(zip( - found_after_vetoes['found/tc'][found_excl], - found_after_vetoes['network/reweighted_snr'][found_excl])): + for j, (t, bestnr) in enumerate( + zip( + found_after_vetoes["found/tc"][found_excl], + found_after_vetoes["network/reweighted_snr"][found_excl], + ) + ): # 0 is the zero-lag timeslide - near_bestnr = trig_data[keys[1]][0][np.abs(trig_data[keys[0]][0]-t) < - cluster_window] + near_bestnr = trig_data[keys[1]][0][ + np.abs(trig_data[keys[0]][0] - t) < cluster_window + ] near_test[j] = ~((near_bestnr * glitch_check_fac > bestnr).any()) # Apply the local test c = 0 for z, b in enumerate(found_excl): - if found_excl[z]: + if b: found_excl[z] = near_test[c] c += 1 # Loop over each random instance of the injection set - for k in range(num_mc_injs+1): + for k in range(num_mc_injs + 1): # Loop over the distance bins for j, dist_bin in enumerate(dist_bins): # Construct distance cut - found_dist_cut = (dist_bin[0] <= found_inj_dist_mc[k, :]) &\ - (found_inj_dist_mc[k, :] < dist_bin[1]) - missed_dist_cut = (dist_bin[0] <= missed_inj_dist_mc[k, :]) &\ - (missed_inj_dist_mc[k, :] < dist_bin[1]) - long_dist_cut = (dist_bin[0] <= long_inj['dist_mc'][k, :]) &\ - (long_inj['dist_mc'][k, :] < dist_bin[1]) + found_dist_cut = (dist_bin[0] <= found_inj_dist_mc[k, :]) & ( + found_inj_dist_mc[k, :] < dist_bin[1] + ) + missed_dist_cut = (dist_bin[0] <= missed_inj_dist_mc[k, :]) & ( + missed_inj_dist_mc[k, :] < dist_bin[1] + ) + long_dist_cut = (dist_bin[0] <= long_inj["dist_mc"][k, :]) & ( + long_inj["dist_mc"][k, :] < dist_bin[1] + ) # Count all injections in this distance bin num_found_pass = (found_dist_cut).sum() @@ -544,9 +595,9 @@ if len(found_after_vetoes['network/template_id']): # Record number of injections, number found for exclusion # and number of zero FAR if k == 0: - key = 'no_mc' + key = "no_mc" else: - key = 'mc' + key = "mc" num_pass = num_found_pass + num_missed_pass + num_long_pass num_injections[key][j] += num_pass num_injections[key][-1] += num_pass @@ -566,9 +617,8 @@ logging.info("Plotting.") # Calculate distances (as means) for the horizontal axis, only if there are # found injections -if len(found_after_vetoes['network/template_id']): - dist_plot_vals = [np.asarray(dist_bin).mean() - for dist_bin in dist_bins] +if len(found_after_vetoes["network/template_id"]): + dist_plot_vals = [np.asarray(dist_bin).mean() for dist_bin in dist_bins] # Plot efficiency using loudest background if opts.background_output_file: @@ -576,76 +626,78 @@ if opts.background_output_file: ax = fig.gca() # Without found injections, do not determine these quantities to plot - if len(found_after_vetoes['network/template_id']): + if len(found_after_vetoes["network/template_id"]): # Calculate error bars for efficiency/distance plots and datafiles # using max BestNR of background yerr_low_mc, yerr_high_mc, fraction_mc = efficiency_with_errs( - found_max_bestnr['mc'], - num_injections['mc'], - num_mc_injs=num_mc_injs) + found_max_bestnr["mc"], num_injections["mc"], num_mc_injs=num_mc_injs + ) yerr_low_no_mc, yerr_high_no_mc, fraction_no_mc = efficiency_with_errs( - found_max_bestnr['no_mc'], - num_injections['no_mc']) - - ax.plot(dist_plot_vals, (fraction_no_mc), 'g-', - label='No marginalisation') - ax.errorbar(dist_plot_vals, (fraction_no_mc), - yerr=[yerr_low_no_mc, yerr_high_no_mc], c='green') + found_max_bestnr["no_mc"], num_injections["no_mc"] + ) + + ax.plot(dist_plot_vals, (fraction_no_mc), "g-", label="No marginalisation") + ax.errorbar( + dist_plot_vals, + (fraction_no_mc), + yerr=[yerr_low_no_mc, yerr_high_no_mc], + c="green", + ) if np.nansum(fraction_mc) > 0: - ax.plot(dist_plot_vals, fraction_mc, 'r-', label='Marginalised') - ax.errorbar(dist_plot_vals, - fraction_mc, - yerr=[yerr_low_mc, yerr_high_mc], - c='red') + ax.plot(dist_plot_vals, fraction_mc, "r-", label="Marginalised") + ax.errorbar( + dist_plot_vals, fraction_mc, yerr=[yerr_low_mc, yerr_high_mc], c="red" + ) ax.legend() ax.grid() ax.set_ylim([0, 1]) - ax.set_xlim(0, 1.3*upper_dist) - ax.set_ylabel("Fraction of injections found louder than " + - "loudest background") + ax.set_xlim(0, 1.3 * upper_dist) + ax.set_ylabel("Fraction of injections found louder than " + "loudest background") ax.set_xlabel("Distance (Mpc)") - plot_title = "Detection efficiency - "+inj_set_name + plot_title = "Detection efficiency - " + inj_set_name plot_caption = "Injection recovery efficiency using " plot_caption += "BestNR as detection statistic. " plot_caption += "Injections louder than loudest background trigger." fig_path = opts.background_output_file - save_fig_with_metadata(fig, fig_path, cmd=' '.join(sys.argv), - title=plot_title, caption=plot_caption) + save_fig_with_metadata( + fig, fig_path, cmd=" ".join(sys.argv), title=plot_title, caption=plot_caption + ) plt.close() # Calculate and save to disk 50% and 90% exclusion distances # excl_dist dictionary contains 50% and 90% exclusion distances # It contains NaNs unless there are found injections -excl_dist = {f"{percentile}%": 'NaN' for percentile in [50, 90]} -if len(found_after_vetoes['network/template_id']): +excl_dist = {f"{percentile}%": "NaN" for percentile in [50, 90]} +if len(found_after_vetoes["network/template_id"]): # Calculate error bars for efficiency/distance plots and datafiles # using max BestNR of foreground yerr_low_no_mc, yerr_high_no_mc, fraction_no_mc = efficiency_with_errs( - found_on_bestnr['no_mc'], - num_injections['no_mc']) - yerr_low, yerr_high, fraction_mc = \ - efficiency_with_errs(found_on_bestnr['mc'], num_injections['mc'], - num_mc_injs=num_mc_injs) + found_on_bestnr["no_mc"], num_injections["no_mc"] + ) + yerr_low, yerr_high, fraction_mc = efficiency_with_errs( + found_on_bestnr["mc"], num_injections["mc"], num_mc_injs=num_mc_injs + ) # Marginalized efficiency (isf = inverse survival function) red_efficiency = (fraction_mc) - (yerr_low) * scipy.stats.norm.isf(0.1) for percentile in [50, 90]: - eff_idx = np.where(red_efficiency < (percentile / 100.))[0] + eff_idx = np.where(red_efficiency < (percentile / 100.0))[0] if eff_idx.size == 0: - green_efficiency = (fraction_no_mc) + green_efficiency = fraction_no_mc excl_efficiency = green_efficiency - eff_idx = np.where(green_efficiency < (percentile / 100.))[0] + eff_idx = np.where(green_efficiency < (percentile / 100.0))[0] else: excl_efficiency = red_efficiency if eff_idx.size and eff_idx[0] != 0: i = eff_idx[0] d = dist_plot_vals[i] - d_low = dist_plot_vals[i-1] + d_low = dist_plot_vals[i - 1] e = excl_efficiency[i] - e_low = excl_efficiency[i-1] - excl_dist[f"{percentile}%"] = \ - d + (e - (percentile / 100.)) * (d - d_low) / (e_low - e) + e_low = excl_efficiency[i - 1] + excl_dist[f"{percentile}%"] = d + (e - (percentile / 100.0)) * ( + d - d_low + ) / (e_low - e) else: # Warn the user if the exclusion distance cannot be established, # but let the workflow continue: the user will see the plot(s) and @@ -657,17 +709,17 @@ if len(found_after_vetoes['network/template_id']): # Write 50% and 90% exclusion distances to JSON file # Also include injection set name and trial name if opts.exclusion_dist_output_file: - excl_dist['inj_set'] = inj_set_name - excl_dist['trial_name'] = opts.trial_name - excl_dist['inj_dist'] = dist_plot_vals - excl_dist['eff_no_marg'] = list(fraction_no_mc) - excl_dist['eff_no_marg_err_low'] = list(yerr_low_no_mc) - excl_dist['eff_no_marg_err_high'] = list(yerr_high_no_mc) - excl_dist['marg_eff'] = list(fraction_mc) - excl_dist['marg_eff_err_low'] = list(yerr_low) - excl_dist['marg_eff_err_high'] = list(yerr_high) - excl_dist['eff_inc_count_err'] = list(red_efficiency) - with open(opts.exclusion_dist_output_file, 'w') as excl_dist_file: + excl_dist["inj_set"] = inj_set_name + excl_dist["trial_name"] = opts.trial_name + excl_dist["inj_dist"] = dist_plot_vals + excl_dist["eff_no_marg"] = list(fraction_no_mc) + excl_dist["eff_no_marg_err_low"] = list(yerr_low_no_mc) + excl_dist["eff_no_marg_err_high"] = list(yerr_high_no_mc) + excl_dist["marg_eff"] = list(fraction_mc) + excl_dist["marg_eff_err_low"] = list(yerr_low) + excl_dist["marg_eff_err_high"] = list(yerr_high) + excl_dist["eff_inc_count_err"] = list(red_efficiency) + with open(opts.exclusion_dist_output_file, "w") as excl_dist_file: json.dump(excl_dist, excl_dist_file) # Plot efficiency using loudest foreground @@ -675,40 +727,43 @@ if opts.onsource_output_file: fig = plt.figure() ax = fig.gca() ax.grid() - if len(found_after_vetoes['network/template_id']): - ax.plot(dist_plot_vals, (fraction_no_mc), 'g-', - label='No marginalisation') - ax.errorbar(dist_plot_vals, (fraction_no_mc), - yerr=[yerr_low_no_mc, yerr_high_no_mc], c='green') + if len(found_after_vetoes["network/template_id"]): + ax.plot(dist_plot_vals, (fraction_no_mc), "g-", label="No marginalisation") + ax.errorbar( + dist_plot_vals, + (fraction_no_mc), + yerr=[yerr_low_no_mc, yerr_high_no_mc], + c="green", + ) if np.nansum(fraction_mc) > 0: - ax.plot(dist_plot_vals, fraction_mc, 'r-', label='Marginalised') - ax.errorbar(dist_plot_vals, fraction_mc, yerr=[yerr_low, yerr_high], - c='red') + ax.plot(dist_plot_vals, fraction_mc, "r-", label="Marginalised") + ax.errorbar( + dist_plot_vals, fraction_mc, yerr=[yerr_low, yerr_high], c="red" + ) if np.nansum(red_efficiency) > 0: - ax.plot(dist_plot_vals, red_efficiency, 'm-', - label='Inc. counting errors') + ax.plot(dist_plot_vals, red_efficiency, "m-", label="Inc. counting errors") ax.set_ylim([0, 1]) ax.grid() ax.legend() ax.get_legend().get_frame().set_alpha(0.5) ax.grid() ax.set_ylim([0, 1]) - ax.set_xlim(0, 1.3*upper_dist) - ax.plot([excl_dist["90%"]], [0.9], 'gx') - ax.set_ylabel("Fraction of injections found louder than " + - "loudest foreground") + ax.set_xlim(0, 1.3 * upper_dist) + ax.plot([excl_dist["90%"]], [0.9], "gx") + ax.set_ylabel("Fraction of injections found louder than " + "loudest foreground") ax.set_xlabel("Distance (Mpc)") ax.set_ylim([0, 1]) - ax.set_xlim(0, 1.3*upper_dist) - plot_title = "Exclusion distance - "+inj_set_name + ax.set_xlim(0, 1.3 * upper_dist) + plot_title = "Exclusion distance - " + inj_set_name plot_caption = "Injection recovery efficiency using " plot_caption += "BestNR as detection statistic. " plot_caption += "Injections louder than loudest foreground trigger.\n" plot_caption += f" 90%% exclusion distance: {excl_dist['90%']} Mpc\n" plot_caption += f" 50%% sensitive distance: {excl_dist['50%']} Mpc" fig_path = opts.onsource_output_file - save_fig_with_metadata(fig, fig_path, cmd=' '.join(sys.argv), - title=plot_title, caption=plot_caption) + save_fig_with_metadata( + fig, fig_path, cmd=" ".join(sys.argv), title=plot_title, caption=plot_caption + ) plt.close() logging.info("Done.") diff --git a/bin/pygrb/pycbc_pygrb_exclusion_dist_table b/bin/pygrb/pycbc_pygrb_exclusion_dist_table index 873a44c9943..95561b6d8a4 100644 --- a/bin/pygrb/pycbc_pygrb_exclusion_dist_table +++ b/bin/pygrb/pycbc_pygrb_exclusion_dist_table @@ -18,28 +18,36 @@ """Create table of exclusion distances.""" -import sys import argparse import json +import sys + import pycbc.version -import pycbc.results +import pycbc.results __author__ = "Jacob Buchanan " __version__ = pycbc.version.git_verbose_msg __date__ = pycbc.version.date __program__ = "pycbc_pygrb_exclusion_dist_table" -parser = argparse.ArgumentParser(description=__doc__, formatter_class= - argparse.ArgumentDefaultsHelpFormatter) +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-files", nargs="+", required=True, - help="List of JSON input files" + - " output by pycbc_pygrb_efficiency" + - " containing exclusion distances.") -parser.add_argument("--output-file", required=True, - help="HTML output file containing table" + - " of exclusion distances.") +parser.add_argument( + "--input-files", + nargs="+", + required=True, + help="List of JSON input files" + " output by pycbc_pygrb_efficiency" + " containing exclusion distances.", +) +parser.add_argument( + "--output-file", + required=True, + help="HTML output file containing table" + " of exclusion distances.", +) opts = parser.parse_args() pycbc.init_logging(opts.verbose) @@ -47,7 +55,7 @@ pycbc.init_logging(opts.verbose) # Load JSON files as a list of dictionaries file_contents = [] for file_name in opts.input_files: - with open(file_name, "r") as file: + with open(file_name) as file: file_contents.append(json.load(file)) # Get list of trials (i.e. OFFTRIAL_i, ONSOURCE, etc.) @@ -87,13 +95,13 @@ for fc in file_contents: # Prepare dictionary for each trial + injection set results[fc["trial_name"]][fc["inj_set"]] = {} # Add exclusion distances to dictionary - for percent in ('50%', '90%'): + for percent in ("50%", "90%"): results[fc["trial_name"]][fc["inj_set"]][percent] = fc[percent] # Set up rows for table data = [] -for percent in ('50%', '90%'): +for percent in ("50%", "90%"): for trial in trials: row = [f"{trial} ({percent})"] for injection_set in injection_sets: @@ -106,6 +114,6 @@ html = str(pycbc.results.static_table(data, headers)) # Write as figure title = "Exclusion Distances" caption = "Table of exclusion distances for each trial and injection set." -pycbc.results.save_fig_with_metadata(html, opts.output_file, - cmd=' '.join(sys.argv), - title=title, caption=caption) +pycbc.results.save_fig_with_metadata( + html, opts.output_file, cmd=" ".join(sys.argv), title=title, caption=caption +) diff --git a/bin/pygrb/pycbc_pygrb_grb_info_table b/bin/pygrb/pycbc_pygrb_grb_info_table index 97316de672b..87fb84e7ed3 100644 --- a/bin/pygrb/pycbc_pygrb_grb_info_table +++ b/bin/pygrb/pycbc_pygrb_grb_info_table @@ -16,7 +16,8 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Create GRB info table. +""" +Create GRB info table. Please refer to help(pycbc.types.angle_as_radians) for the recommended configuration file syntax for angle arguments. @@ -25,19 +26,19 @@ configuration file syntax for angle arguments. # ============================================================================= # Preamble # ============================================================================= -import sys import argparse import math +import sys -from pycbc import add_common_pycbc_options, init_logging import pycbc.version -import pycbc.results + import pycbc.distributions +import pycbc.results +from pycbc import add_common_pycbc_options, init_logging from pycbc.detector import Detector, ppdets from pycbc.results.pygrb_postprocessing_utils import get_antenna_dist_factor -from pycbc.types import angle_as_radians from pycbc.time import gps_to_utc_str - +from pycbc.types import angle_as_radians __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -47,32 +48,54 @@ __program__ = "pycbc_pygrb_grb_info_table" # ============================================================================= # Main script starts here # ============================================================================= -parser = argparse.ArgumentParser(description=__doc__, formatter_class= - argparse.ArgumentDefaultsHelpFormatter) +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter +) add_common_pycbc_options(parser) -parser.add_argument("--trigger-time", type=int, - required=True, - help="GPS time of the GRB.") -parser.add_argument("--input-dist", type=str, - help="Input Distribution of the GRB. Use distribution " - "functions in pycbc.distributions.sky_location code.") -parser.add_argument("--ra", type=angle_as_radians, - help="Right ascension of the GRB. Use the rad or deg " - "suffix to specify units, otherwise radians are assumed.") -parser.add_argument("--dec", type=angle_as_radians, - help="Declination of the GRB. Use the rad or deg suffix " - "to specify units, otherwise radians are assumed.") -parser.add_argument("--sky-error", type=angle_as_radians, - default=0, - help="Sky-localisation error of the GRB. Use the rad or " - "deg suffix to specify units, otherwise radians are " - "assumed.") -parser.add_argument("--ifos", action="store", nargs='+', - default=None, required=True, - help="List containing the active IFOs.") -parser.add_argument("--output-file", action="store", - default=None, required=True, - help="The output file to write tha table to.") +parser.add_argument( + "--trigger-time", type=int, required=True, help="GPS time of the GRB." +) +parser.add_argument( + "--input-dist", + type=str, + help="Input Distribution of the GRB. Use distribution " + "functions in pycbc.distributions.sky_location code.", +) +parser.add_argument( + "--ra", + type=angle_as_radians, + help="Right ascension of the GRB. Use the rad or deg " + "suffix to specify units, otherwise radians are assumed.", +) +parser.add_argument( + "--dec", + type=angle_as_radians, + help="Declination of the GRB. Use the rad or deg suffix " + "to specify units, otherwise radians are assumed.", +) +parser.add_argument( + "--sky-error", + type=angle_as_radians, + default=0, + help="Sky-localisation error of the GRB. Use the rad or " + "deg suffix to specify units, otherwise radians are " + "assumed.", +) +parser.add_argument( + "--ifos", + action="store", + nargs="+", + default=None, + required=True, + help="List containing the active IFOs.", +) +parser.add_argument( + "--output-file", + action="store", + default=None, + required=True, + help="The output file to write tha table to.", +) opts = parser.parse_args() @@ -82,18 +105,18 @@ headers = [] data = [[]] data[0].append(str(opts.trigger_time)) -headers.append('GPS Time') +headers.append("GPS Time") utc_time_str = gps_to_utc_str(float(opts.trigger_time)) data[0].append(utc_time_str) -headers.append('UTC Time') +headers.append("UTC Time") if opts.input_dist is not None: if opts.ra is not None or opts.dec is not None: parser.error( "You can't use input-dist argument and (ra,dec) at the same time, please choose one" - ) - input_dist = eval('pycbc.distributions.' + opts.input_dist) + ) + input_dist = eval("pycbc.distributions." + opts.input_dist) ra, dec = input_dist.get_max_prob_point() sky_error = None @@ -101,43 +124,37 @@ else: ra = opts.ra dec = opts.dec sky_error = opts.sky_error - -data[0].append(f'{math.degrees(ra):.3f}') -headers.append('R.A. (deg)') -data[0].append(f'{math.degrees(dec):.3f}') -headers.append('Dec (deg)') +data[0].append(f"{math.degrees(ra):.3f}") +headers.append("R.A. (deg)") + +data[0].append(f"{math.degrees(dec):.3f}") +headers.append("Dec (deg)") if sky_error is not None: - data[0].append(f'{math.degrees(sky_error):.3f}') - headers.append('Sky Error (deg)') + data[0].append(f"{math.degrees(sky_error):.3f}") + headers.append("Sky Error (deg)") -data[0].append(ppdets(opts.ifos, '')) -headers.append('IFOs') +data[0].append(ppdets(opts.ifos, "")) +headers.append("IFOs") for ifo in opts.ifos: antenna = Detector(ifo) - factor = get_antenna_dist_factor( - antenna, ra, dec, float(opts.trigger_time) - ) - data[0].append(f'{factor:.3f}') - headers.append(ifo + ' Antenna Factor') + factor = get_antenna_dist_factor(antenna, ra, dec, float(opts.trigger_time)) + data[0].append(f"{factor:.3f}") + headers.append(ifo + " Antenna Factor") -html = pycbc.results.dq.redirect_javascript + \ - str(pycbc.results.static_table(data, headers)) +html = pycbc.results.dq.redirect_javascript + str( + pycbc.results.static_table(data, headers) +) -title = 'External Trigger Summary Information' +title = "External Trigger Summary Information" caption = ( - 'Parameters of the external trigger. The reported antenna factors are the ' - 'dist / eff distance as defined by Eq (4.3) in ' - 'https://arxiv.org/abs/0705.1514.' + "Parameters of the external trigger. The reported antenna factors are the " + "dist / eff distance as defined by Eq (4.3) in " + "https://arxiv.org/abs/0705.1514." ) pycbc.results.save_fig_with_metadata( - html, - opts.output_file, - {}, - cmd=' '.join(sys.argv), - title=title, - caption=caption + html, opts.output_file, {}, cmd=" ".join(sys.argv), title=title, caption=caption ) diff --git a/bin/pygrb/pycbc_pygrb_minifollowups b/bin/pygrb/pycbc_pygrb_minifollowups index dc1320f6183..f95b4029d5e 100644 --- a/bin/pygrb/pycbc_pygrb_minifollowups +++ b/bin/pygrb/pycbc_pygrb_minifollowups @@ -23,20 +23,21 @@ Set up qscans and SNR timeseries plots of loudest triggers/missed injections # ============================================================================= # Preamble # ============================================================================= -import os import argparse import logging +import os -from igwn_segments import segmentlist -import pycbc.workflow as wf -from pycbc.workflow.core import FileList, resolve_url_to_file -import pycbc.workflow.minifollowups as mini import pycbc.version +from igwn_segments import segmentlist + import pycbc.events import pycbc.results.pygrb_postprocessing_utils as ppu +import pycbc.workflow as wf +import pycbc.workflow.minifollowups as mini +from pycbc.io.hdf import HFile from pycbc.results import layout +from pycbc.workflow.core import FileList, resolve_url_to_file from pycbc.workflow.plotting import PlotExecutable -from pycbc.io.hdf import HFile __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -51,61 +52,70 @@ def add_wiki_row(outfile, cols): """ Adds a wiki-formatted row to an output file from a list or a numpy array. """ - with open(outfile, 'a') as f: - f.write('||%s||\n' % '||'.join(map(str, cols))) - - -def make_timeseries_plot(workflow, trig_file, snr_type, central_time, - out_dir, ifo=None, seg_files=None, - veto_file=None, tags=None): + with open(outfile, "a") as f: + f.write("||%s||\n" % "||".join(map(str, cols))) + + +def make_timeseries_plot( + workflow, + trig_file, + snr_type, + central_time, + out_dir, + ifo=None, + seg_files=None, + veto_file=None, + tags=None, +): """Adds a node for a timeseries of PyGRB results to the workflow""" - tags = [] if tags is None else tags # Ensure that zero-lag data is used in these follow-up plots and store # the slide-id option originally provided orig_slide_id = None - if workflow.cp.has_option('pygrb_plot_snr_timeseries', 'slide-id'): - orig_slide_id = workflow.cp.get('pygrb_plot_snr_timeseries', - 'slide-id') + if workflow.cp.has_option("pygrb_plot_snr_timeseries", "slide-id"): + orig_slide_id = workflow.cp.get("pygrb_plot_snr_timeseries", "slide-id") workflow.cp.add_options_to_section( - 'pygrb_plot_snr_timeseries', - [('slide-id', '0')], - True + "pygrb_plot_snr_timeseries", [("slide-id", "0")], True ) # Initialize job node with its tags - grb_name = workflow.cp.get('workflow', 'trigger-name') - extra_tags = ['GRB'+grb_name] + grb_name = workflow.cp.get("workflow", "trigger-name") + extra_tags = ["GRB" + grb_name] extra_tags += [snr_type] if ifo is not None: extra_tags += [ifo] - node = PlotExecutable(workflow.cp, 'pygrb_plot_snr_timeseries', - ifos=workflow.ifos, out_dir=out_dir, - tags=tags+extra_tags).create_node() - node.add_input_opt('--trig-file', trig_file) + node = PlotExecutable( + workflow.cp, + "pygrb_plot_snr_timeseries", + ifos=workflow.ifos, + out_dir=out_dir, + tags=tags + extra_tags, + ).create_node() + node.add_input_opt("--trig-file", trig_file) # Include the onsource trial if this is a follow up on the onsource - if 'loudest_onsource_event' in tags: - node.add_opt('--onsource') + if "loudest_onsource_event" in tags: + node.add_opt("--onsource") # Pass the segments files and veto file if seg_files: - node.add_input_list_opt('--seg-files', seg_files) + node.add_input_list_opt("--seg-files", seg_files) if veto_file: - node.add_input_opt('--veto-file', veto_file) - node.new_output_file_opt(workflow.analysis_time, '.png', - '--output-file', tags=extra_tags) + node.add_input_opt("--veto-file", veto_file) + node.new_output_file_opt( + workflow.analysis_time, ".png", "--output-file", tags=extra_tags + ) # Quantity to be displayed on the y-axis of the plot - node.add_opt('--y-variable', snr_type) + node.add_opt("--y-variable", snr_type) if ifo is not None: - node.add_opt('--ifo', ifo) - node.add_opt('--x-lims=-5.,5.') + node.add_opt("--ifo", ifo) + node.add_opt("--x-lims=-5.,5.") # Plot title if ifo is not None: title_str = f"'{ifo} SNR at {central_time:.3f} (s)'" else: title_str = f"'{snr_type.capitalize()} SNR at {central_time:.3f} (s)'" - node.add_opt('--trigger-time', central_time) - node.add_opt('--plot-title', title_str) + node.add_opt("--trigger-time", central_time) + node.add_opt("--plot-title", title_str) # Add job node to workflow workflow += node @@ -113,12 +123,10 @@ def make_timeseries_plot(workflow, trig_file, snr_type, central_time, # Revert the config parser back to how it was given if orig_slide_id: workflow.cp.add_options_to_section( - 'pygrb_plot_snr_timeseries', - [('slide-id', orig_slide_id)], - True + "pygrb_plot_snr_timeseries", [("slide-id", orig_slide_id)], True ) else: - workflow.cp.remove_option('pygrb_plot_snr_timeseries', 'slide-id') + workflow.cp.remove_option("pygrb_plot_snr_timeseries", "slide-id") return node.output_files @@ -128,24 +136,28 @@ def make_timeseries_plot(workflow, trig_file, snr_type, central_time, # ============================================================================= parser = argparse.ArgumentParser(description=__doc__[1:]) pycbc.add_common_pycbc_options(parser) -parser.add_argument('--trig-file', - help="HDF file with the triggers found by PyGRB") -parser.add_argument('--followups-file', - help="HDF file with the triggers/injections to follow up") -parser.add_argument('--wiki-file', - help="Name of file to save wiki-formatted table in") -parser.add_argument("-a", "--seg-files", nargs="+", action="store", - default=[], help="The location of the buffer, " + - "onsource and offsource txt segment files.") -parser.add_argument("-V", "--veto-file", action="store", - help="The location of the xml veto file.") +parser.add_argument("--trig-file", help="HDF file with the triggers found by PyGRB") +parser.add_argument( + "--followups-file", help="HDF file with the triggers/injections to follow up" +) +parser.add_argument("--wiki-file", help="Name of file to save wiki-formatted table in") +parser.add_argument( + "-a", + "--seg-files", + nargs="+", + action="store", + default=[], + help="The location of the buffer, " + "onsource and offsource txt segment files.", +) +parser.add_argument( + "-V", "--veto-file", action="store", help="The location of the xml veto file." +) wf.add_workflow_command_line_group(parser) wf.add_workflow_settings_cli(parser, include_subdax_opts=True) ppu.pygrb_add_bestnr_cut_opt(parser) args = parser.parse_args() -pycbc.init_logging(args.verbose, - format="%(asctime)s: %(levelname)s: %(message)s") +pycbc.init_logging(args.verbose, format="%(asctime)s: %(levelname)s: %(message)s") workflow = wf.Workflow(args) @@ -155,7 +167,7 @@ wf.makedir(args.output_dir) layouts = [] # Read the file with the triggers/injections to follow up -logging.info('Reading list of triggers/injections to followup') +logging.info("Reading list of triggers/injections to followup") fp = HFile(args.followups_file, "r") # Initialize a wiki table and add the column headers @@ -164,9 +176,8 @@ if args.wiki_file: add_wiki_row(wiki_file, fp.keys()) # Establish the number of follow-ups to perform -num_events = int(workflow.cp.get_opt_tags('workflow-minifollowups', - 'num-events', '')) -num_events = min(num_events, len(fp['BestNR'][:])) +num_events = int(workflow.cp.get_opt_tags("workflow-minifollowups", "num-events", "")) +num_events = min(num_events, len(fp["BestNR"][:])) # File instance of the input trigger file trig_file = resolve_url_to_file(os.path.abspath(args.trig_file)) @@ -185,20 +196,18 @@ if veto_file: seg_file_abs_paths = [os.path.join(start_rundir, f) for f in args.seg_files] # Convert the segments files to a FileList -seg_files = wf.FileList( - [wf.resolve_url_to_file(f) for f in seg_file_abs_paths] -) +seg_files = wf.FileList([wf.resolve_url_to_file(f) for f in seg_file_abs_paths]) # The Q-scan jobs need to be told the valid science segments to deal with # triggers near the boundaries correctly. We use the offsource segment as a # proxy for the valid science segments, and ignore vetoed time inside it. -valid_segs = segmentlist([ppu._read_seg_files(seg_file_abs_paths)['off']]) +valid_segs = segmentlist([ppu._read_seg_files(seg_file_abs_paths)["off"]]) # Determine if we are following up injections or loudest off/on-source triggers. # If the time shift dataset exists, then assume we are doing the latter. is_injection_followup = True try: - time_shift = fp[ifos[0]+' time shift (s)'][0] + time_shift = fp[ifos[0] + " time shift (s)"][0] is_injection_followup = False except KeyError: pass @@ -206,8 +215,8 @@ except KeyError: # Loop over triggers/injections to be followed up for num_event in range(num_events): files = FileList([]) - logging.info('Processing event: %s', num_event+1) - tags = args.tags + [str(num_event+1)] + logging.info("Processing event: %s", num_event + 1) + tags = args.tags + [str(num_event + 1)] if wiki_file: row = [] for key in fp.keys(): @@ -216,36 +225,54 @@ for num_event in range(num_events): # Handle injections # They are on unslid data so geocenter time is used if is_injection_followup: - gps_time = fp['GC GPS time (s)'][num_event] + gps_time = fp["GC GPS time (s)"][num_event] gps_time = gps_time.astype(float) - for snr_type in ['reweighted', 'coherent']: - files += make_timeseries_plot(workflow, trig_file, - snr_type, gps_time, - args.output_dir, ifo=None, - seg_files=seg_files, - veto_file=veto_file, - tags=tags) + for snr_type in ["reweighted", "coherent"]: + files += make_timeseries_plot( + workflow, + trig_file, + snr_type, + gps_time, + args.output_dir, + ifo=None, + seg_files=seg_files, + veto_file=veto_file, + tags=tags, + ) for ifo in ifos: - files += mini.make_qscan_plot(workflow, ifo, gps_time, - args.output_dir, - tags=tags, - data_segments=valid_segs) + files += mini.make_qscan_plot( + workflow, + ifo, + gps_time, + args.output_dir, + tags=tags, + data_segments=valid_segs, + ) # Handle off/on-source loudest triggers follow-up # They are data that may be slid so detector time is used else: for i_ifo, ifo in enumerate(ifos): - ifo_time = fp[ifo+' GPS time (s)'][num_event] + ifo_time = fp[ifo + " GPS time (s)"][num_event] ifo_time = ifo_time.astype(float) - files += make_timeseries_plot(workflow, trig_file, - 'single', ifo_time, - args.output_dir, ifo=ifo, - seg_files=seg_files, - veto_file=veto_file, - tags=tags) - files += mini.make_qscan_plot(workflow, ifo, ifo_time, - args.output_dir, - tags=tags, - data_segments=valid_segs) + files += make_timeseries_plot( + workflow, + trig_file, + "single", + ifo_time, + args.output_dir, + ifo=ifo, + seg_files=seg_files, + veto_file=veto_file, + tags=tags, + ) + files += mini.make_qscan_plot( + workflow, + ifo, + ifo_time, + args.output_dir, + tags=tags, + data_segments=valid_segs, + ) layouts += list(layout.grouper(files, 2)) diff --git a/bin/pygrb/pycbc_pygrb_page_tables b/bin/pygrb/pycbc_pygrb_page_tables index 83ffa507c91..6b96c0ffedd 100755 --- a/bin/pygrb/pycbc_pygrb_page_tables +++ b/bin/pygrb/pycbc_pygrb_page_tables @@ -23,19 +23,20 @@ Processes PyGRB triggers and injections to create html results tables. # ============================================================================= # Preamble # ============================================================================= -import sys -import os import logging -import numpy as np +import os +import sys +import numpy as np import pycbc.version + +import pycbc.results +from pycbc import init_logging from pycbc.conversions import mchirp_from_mass1_mass2 from pycbc.detector import Detector from pycbc.events.coherent import reweightedsnr_cut -from pycbc import init_logging -import pycbc.results -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.io.hdf import HFile +from pycbc.results import pygrb_postprocessing_utils as ppu __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -48,36 +49,36 @@ __program__ = "pycbc_pygrb_page_tables" # ============================================================================= def additional_injection_data(data, ifos): """Provides data with chirp masses and effective distances""" - - data['mchirp'] = mchirp_from_mass1_mass2(data['mass1'], - data['mass2']) + data["mchirp"] = mchirp_from_mass1_mass2(data["mass1"], data["mass2"]) eff_dist = 0 for ifo in ifos: antenna = Detector(ifo) - data['eff_dist_'+ifo] = antenna.effective_distance( - data['distance'], - data['ra'], - data['dec'], - data['polarization'], - data['tc'], - data['inclination'] - ) - eff_dist += 1.0 / data['eff_dist_'+ifo] - data['eff_dist'] = 1.0 / eff_dist + data["eff_dist_" + ifo] = antenna.effective_distance( + data["distance"], + data["ra"], + data["dec"], + data["polarization"], + data["tc"], + data["inclination"], + ) + eff_dist += 1.0 / data["eff_dist_" + ifo] + data["eff_dist"] = 1.0 / eff_dist net_opt_snr_sq = 0 for ifo in ifos: - if 'optimal_snr_'+ifo not in data: - data['optimal_snr_'+ifo] = np.full(len(data['mass1']), np.nan) - net_opt_snr_sq += data['optimal_snr_'+ifo]**2 - data['net_opt_snr'] = np.sqrt(net_opt_snr_sq) + if "optimal_snr_" + ifo not in data: + data["optimal_snr_" + ifo] = np.full(len(data["mass1"]), np.nan) + net_opt_snr_sq += data["optimal_snr_" + ifo] ** 2 + data["net_opt_snr"] = np.sqrt(net_opt_snr_sq) return data -def load_missed_found_injections(hdf_file, ifos, bank_file, snr_threshold=None, - background_bestnrs=None): - """Loads found and missed injections from an hdf file as two dictionaries +def load_missed_found_injections( + hdf_file, ifos, bank_file, snr_threshold=None, background_bestnrs=None +): + """ + Loads found and missed injections from an hdf file as two dictionaries Parameters ---------- @@ -95,16 +96,29 @@ def load_missed_found_injections(hdf_file, ifos, bank_file, snr_threshold=None, data: tuple of dictionaries Found, missed, and missed after the cut in reweighted SNR injection parameter dictionaries. - """ - logging.info('Loading injections...') - inj_data = HFile(hdf_file, 'r') - inj_params = ['mass1', 'mass2', 'distance', 'inclination', 'ra', 'dec', - 'polarization', 'spin1x', 'spin1y', 'spin1z', 'spin2x', - 'spin2y', 'spin2z', 'tc'] + """ + logging.info("Loading injections...") + inj_data = HFile(hdf_file, "r") + inj_params = [ + "mass1", + "mass2", + "distance", + "inclination", + "ra", + "dec", + "polarization", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + "tc", + ] for ifo in ifos: - param = 'optimal_snr_'+ifo - if param in inj_data['missed'] and param in inj_data['found']: + param = "optimal_snr_" + ifo + if param in inj_data["missed"] and param in inj_data["found"]: inj_params.append(param) found_data = {} @@ -113,8 +127,8 @@ def load_missed_found_injections(hdf_file, ifos, bank_file, snr_threshold=None, # Load injections parameters for param in inj_params: - missed_data[param] = inj_data['missed/'+param][...] - found_data[param] = inj_data['found/'+param][...] + missed_data[param] = inj_data["missed/" + param][...] + found_data[param] = inj_data["found/" + param][...] # Calculate effective distance for the ifos found_data = additional_injection_data(found_data, ifos) @@ -122,143 +136,185 @@ def load_missed_found_injections(hdf_file, ifos, bank_file, snr_threshold=None, # Get recovered parameters and statistic values for the found injections # Recovered parameters - for param in ['mass1', 'mass2', 'spin1z', 'spin2z']: - found_data['rec_'+param] = \ - np.array(bank_file[param])[inj_data['network/template_id']] + for param in ["mass1", "mass2", "spin1z", "spin2z"]: + found_data["rec_" + param] = np.array(bank_file[param])[ + inj_data["network/template_id"] + ] # If there are no found injections, simply move on # (there are no injections missed after the cut to return) - if 'network/end_time_gc' not in inj_data.keys(): + if "network/end_time_gc" not in inj_data.keys(): return found_data, missed_data, {} # Otherwise carry on getting the recovered parameters and statistic values # of the found injections - found_data['time_diff'] = \ - found_data['tc'] - inj_data['network/end_time_gc'][...] - found_data['rec_mchirp'] = mchirp_from_mass1_mass2( - found_data['rec_mass1'], - found_data['rec_mass2']) + found_data["time_diff"] = found_data["tc"] - inj_data["network/end_time_gc"][...] + found_data["rec_mchirp"] = mchirp_from_mass1_mass2( + found_data["rec_mass1"], found_data["rec_mass2"] + ) # Recovered RA and Dec - found_data['rec_ra'] = inj_data['network/ra'][...] - found_data['rec_dec'] = inj_data['network/dec'][...] + found_data["rec_ra"] = inj_data["network/ra"][...] + found_data["rec_dec"] = inj_data["network/dec"][...] # Statistics values - for param in ['coherent_snr', 'reweighted_snr', 'null_snr']: - found_data[param] = inj_data['network/'+param][...] - found_data['chisq'] = inj_data['network/my_network_chisq'][...] - found_data['nifos'] = inj_data['network/nifo'][...].astype(int) + for param in ["coherent_snr", "reweighted_snr", "null_snr"]: + found_data[param] = inj_data["network/" + param][...] + found_data["chisq"] = inj_data["network/my_network_chisq"][...] + found_data["nifos"] = inj_data["network/nifo"][...].astype(int) for ifo in ifos: - if np.all(inj_data['network/event_id'][...] == - inj_data[ifo+'/event_id'][...]): - found_data['sigmasq_'+ifo] = inj_data[ifo+'/sigmasq'][...] - found_data['snr_'+ifo] = inj_data[ifo+'/snr'][...] - found_data[ifo+'/end_time'] = inj_data[ifo+'/end_time'][...] + if np.all( + inj_data["network/event_id"][...] == inj_data[ifo + "/event_id"][...] + ): + found_data["sigmasq_" + ifo] = inj_data[ifo + "/sigmasq"][...] + found_data["snr_" + ifo] = inj_data[ifo + "/snr"][...] + found_data[ifo + "/end_time"] = inj_data[ifo + "/end_time"][...] else: # Sort the ifo event_id with respect to the network event_id - ifo_sorted_indices = np.argsort(inj_data['network/event_id'][...][ - np.argsort(inj_data['network/event_id'])].searchsorted( - inj_data[ifo+'/event_id'][...])) - found_data['sigmasq_'+ifo] = \ - inj_data[ifo+'/sigmasq'][...][ifo_sorted_indices] - found_data['snr_'+ifo] = \ - inj_data[ifo+'/snr'][...][ifo_sorted_indices] + ifo_sorted_indices = np.argsort( + inj_data["network/event_id"][...][ + np.argsort(inj_data["network/event_id"]) + ].searchsorted(inj_data[ifo + "/event_id"][...]) + ) + found_data["sigmasq_" + ifo] = inj_data[ifo + "/sigmasq"][...][ + ifo_sorted_indices + ] + found_data["snr_" + ifo] = inj_data[ifo + "/snr"][...][ifo_sorted_indices] # BestNRs - found_data['bestnr'] = reweightedsnr_cut(found_data['reweighted_snr'][...], - snr_threshold) + found_data["bestnr"] = reweightedsnr_cut( + found_data["reweighted_snr"][...], snr_threshold + ) # Apply reweighted SNR cut to the found injections cut_data = {} if snr_threshold: - logging.info("%d found injections loaded.", - len(found_data[inj_params[0]])) - logging.info("%d missed injections loaded.", - len(missed_data[inj_params[0]])) + logging.info("%d found injections loaded.", len(found_data[inj_params[0]])) + logging.info("%d missed injections loaded.", len(missed_data[inj_params[0]])) logging.info("Applying reweighted SNR cut at %s.", snr_threshold) - rw_snr_cut = found_data['reweighted_snr'] < snr_threshold + rw_snr_cut = found_data["reweighted_snr"] < snr_threshold for k, v in found_data.items(): cut_data[k] = v[rw_snr_cut] found_data[k] = v[~rw_snr_cut] - del found_data['reweighted_snr'] - del cut_data['reweighted_snr'] + del found_data["reweighted_snr"] + del cut_data["reweighted_snr"] if background_bestnrs is not None: - found_data['fap'] = np.array( - [sum(background_bestnrs > bestnr) for bestnr in - found_data['bestnr']], - dtype=float) / len(background_bestnrs) + found_data["fap"] = np.array( + [sum(background_bestnrs > bestnr) for bestnr in found_data["bestnr"]], + dtype=float, + ) / len(background_bestnrs) # Antenna responses f_resp = {} for ifo in ifos: - if sum(found_data['sigmasq_'+ifo] == 0): + if sum(found_data["sigmasq_" + ifo] == 0): logging.info("%s: sigmasq not set for at least one trigger.", ifo) - if sum(found_data['sigmasq_'+ifo] != 0) == 0: + if sum(found_data["sigmasq_" + ifo] != 0) == 0: logging.info("%s: sigmasq not set for any trigger.", ifo) if len(ifos) == 1: msg = "This is a single ifo analysis. " msg += "Setting sigmasq to unity for all triggers." logging.info(msg) - found_data['sigmasq_'+ifo][:] = 1.0 + found_data["sigmasq_" + ifo][:] = 1.0 antenna = Detector(ifo) - f_resp[ifo] = ppu.get_antenna_responses(antenna, found_data['ra'], - found_data['dec'], - found_data['tc']) + f_resp[ifo] = ppu.get_antenna_responses( + antenna, found_data["ra"], found_data["dec"], found_data["tc"] + ) - inj_sigma_mult = \ - np.asarray([f_resp[ifo] * - found_data['sigmasq_'+ifo] for ifo in ifos]) + inj_sigma_mult = np.asarray( + [f_resp[ifo] * found_data["sigmasq_" + ifo] for ifo in ifos] + ) inj_sigma_tot = np.sum(inj_sigma_mult, axis=0) for ifo in ifos: - found_data['inj_sigma_mean_'+ifo] = np.mean( - found_data['sigmasq_'+ifo] * f_resp[ifo] / inj_sigma_tot) + found_data["inj_sigma_mean_" + ifo] = np.mean( + found_data["sigmasq_" + ifo] * f_resp[ifo] / inj_sigma_tot + ) # Close the hdf file inj_data.close() - logging.info("%d found injections.", len(found_data['mchirp'])) - logging.info("%d missed injections.", len(missed_data['mchirp'])) - logging.info("%d injections cut.", len(cut_data['mchirp'])) + logging.info("%d found injections.", len(found_data["mchirp"])) + logging.info("%d missed injections.", len(missed_data["mchirp"])) + logging.info("%d injections cut.", len(cut_data["mchirp"])) return found_data, missed_data, cut_data def format_pvalue_str(pvalue, n_trials): """Format p-value as a string.""" - return f'< {(1./n_trials):.3g}' if pvalue == 0 else f'{pvalue:.3g}' + return f"< {(1.0 / n_trials):.3g}" if pvalue == 0 else f"{pvalue:.3g}" # ============================================================================= # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("-F", "--offsource-file", action="store", required=True, - help="Location of off-source trigger file") -parser.add_argument("--onsource-file", action="store", - help="Location of on-source trigger file.") -parser.add_argument("--found-missed-file", action="store", - help="HDF format file with injections to output " + - "details about.") -parser.add_argument("--num-loudest-off-trigs", action="store", - type=int, default=30, help="Number of loudest " + - "offsouce triggers to output details about.") -parser.add_argument("--bank-file", action="store", type=str, required=True, - help="Location of the full template bank used.") -parser.add_argument("--quiet-found-injs-output-file", - help="Quiet-found injections html output file.") -parser.add_argument("--missed-found-injs-output-file", - help="Missed-found injections html output file.") -parser.add_argument("--quiet-found-injs-h5-output-file", - help="Quiet-found injections h5 output file.") -parser.add_argument("--loudest-offsource-trigs-output-file", - help="Loudest offsource triggers html output file.") -parser.add_argument("--loudest-offsource-trigs-h5-output-file", - help="Loudest offsource triggers h5 output file.") -parser.add_argument("--loudest-onsource-trig-output-file", - help="Loudest onsource trigger html output file.") -parser.add_argument("--loudest-onsource-trig-h5-output-file", - help="Loudest onsource trigger h5 output file.") -parser.add_argument("-g", "--glitch-check-factor", action="store", - type=float, default=1.0, help="When deciding " + - "exclusion efficiencies this value is multiplied " + - "to the offsource around the injection trigger to " + - "determine if it is just a loud glitch.") -parser.add_argument("-C", "--cluster-window", action="store", type=float, - default=0.1, help="The cluster window used " + - "to cluster triggers in time.") +parser.add_argument( + "-F", + "--offsource-file", + action="store", + required=True, + help="Location of off-source trigger file", +) +parser.add_argument( + "--onsource-file", action="store", help="Location of on-source trigger file." +) +parser.add_argument( + "--found-missed-file", + action="store", + help="HDF format file with injections to output " + "details about.", +) +parser.add_argument( + "--num-loudest-off-trigs", + action="store", + type=int, + default=30, + help="Number of loudest " + "offsouce triggers to output details about.", +) +parser.add_argument( + "--bank-file", + action="store", + type=str, + required=True, + help="Location of the full template bank used.", +) +parser.add_argument( + "--quiet-found-injs-output-file", help="Quiet-found injections html output file." +) +parser.add_argument( + "--missed-found-injs-output-file", help="Missed-found injections html output file." +) +parser.add_argument( + "--quiet-found-injs-h5-output-file", help="Quiet-found injections h5 output file." +) +parser.add_argument( + "--loudest-offsource-trigs-output-file", + help="Loudest offsource triggers html output file.", +) +parser.add_argument( + "--loudest-offsource-trigs-h5-output-file", + help="Loudest offsource triggers h5 output file.", +) +parser.add_argument( + "--loudest-onsource-trig-output-file", + help="Loudest onsource trigger html output file.", +) +parser.add_argument( + "--loudest-onsource-trig-h5-output-file", + help="Loudest onsource trigger h5 output file.", +) +parser.add_argument( + "-g", + "--glitch-check-factor", + action="store", + type=float, + default=1.0, + help="When deciding " + "exclusion efficiencies this value is multiplied " + "to the offsource around the injection trigger to " + "determine if it is just a loud glitch.", +) +parser.add_argument( + "-C", + "--cluster-window", + action="store", + type=float, + default=0.1, + help="The cluster window used " + "to cluster triggers in time.", +) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_slide_opts(parser) opts = parser.parse_args() @@ -283,28 +339,36 @@ output_files = [] # Check for correct input if [found_missed_file, onsource_file].count(None) == 0: - parser.error('Please provide --found-missed-file to process injections, ' + - '--onsource-file to process the on-source, or neither of ' + - 'them to process the off-source triggers.') + parser.error( + "Please provide --found-missed-file to process injections, " + "--onsource-file to process the on-source, or neither of " + "them to process the off-source triggers." + ) # The user may process injections... elif found_missed_file is not None: output_files = [qf_outfile, mf_outfile, qf_h5_outfile] if None in output_files: - parser.error('Please provide all 3 injections output files when ' + - 'using --found-missed-file') + parser.error( + "Please provide all 3 injections output files when " + "using --found-missed-file" + ) # ...or triggers in the onsource... elif onsource_file is not None: output_files = [lont_outfile, lont_h5_outfile] if None in output_files: - parser.error('Please provide both on-source output files ' + - 'when using --onsource-file.') + parser.error( + "Please provide both on-source output files " + "when using --onsource-file." + ) # ...or triggers in the offsource # (both onsource_file and found_missed_file are None) else: output_files = [lofft_outfile, lofft_h5_outfile] if None in output_files: - parser.error('Please provide both off-source output files ' + - 'when using --offsource-file.') + parser.error( + "Please provide both off-source output files " + "when using --offsource-file." + ) logging.info("Setting output directory.") for output_file in output_files: if output_file: @@ -322,26 +386,25 @@ slide_dict = ppu.load_time_slides(offsource_file) segment_dict = ppu.load_segment_dict(offsource_file) # Construct trials removing vetoed times -trial_dict, total_trials = ppu.construct_trials(opts.seg_files, segment_dict, - ifos, slide_dict, - opts.veto_file) +trial_dict, total_trials = ppu.construct_trials( + opts.seg_files, segment_dict, ifos, slide_dict, opts.veto_file +) # Load triggers (apply reweighted SNR cut, not vetoes) -trig_data = ppu.load_data(offsource_file, ifos, data_tag='offsource', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=opts.slide_id) +trig_data = ppu.load_data( + offsource_file, + ifos, + data_tag="offsource", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=opts.slide_id, +) # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors # _av stands for after vetoes -keys = ['network/end_time_gc', 'network/coherent_snr', - 'network/reweighted_snr'] +keys = ["network/end_time_gc", "network/coherent_snr", "network/reweighted_snr"] trig_data_av = ppu.extract_trig_properties( - trial_dict, - trig_data, - slide_dict, - segment_dict, - keys + trial_dict, trig_data, slide_dict, segment_dict, keys ) # Max SNR and BestNR values in each trial: these are stored in dictionaries @@ -356,22 +419,21 @@ for slide_id in slide_dict: if not trial_cut.any(): continue # Max SNR - background_snr[slide_id][j] = \ - max(trig_data_av[keys[1]][slide_id][trial_cut]) + background_snr[slide_id][j] = max(trig_data_av[keys[1]][slide_id][trial_cut]) # Max BestNR - background[slide_id][j] = \ - max(trig_data_av[keys[2]][slide_id][trial_cut]) + background[slide_id][j] = max(trig_data_av[keys[2]][slide_id][trial_cut]) # Max and median values of reweighted SNR, # and sorted (loudest in trial) reweighted SNR values -max_bestnr, median_bestnr, sorted_bkgd =\ - ppu.max_median_stat(slide_dict, background, - trig_data_av[keys[2]], total_trials) +max_bestnr, median_bestnr, sorted_bkgd = ppu.max_median_stat( + slide_dict, background, trig_data_av[keys[2]], total_trials +) assert total_trials == len(sorted_bkgd) # Median value of SNR -_, median_snr, _ = ppu.max_median_stat(slide_dict, background_snr, - trig_data_av[keys[1]], total_trials) +_, median_snr, _ = ppu.max_median_stat( + slide_dict, background_snr, trig_data_av[keys[1]], total_trials +) logging.info("Background SNR and bestNR of trials calculated.") @@ -380,19 +442,16 @@ logging.info("Background SNR and bestNR of trials calculated.") offsource_trigs = [] sorted_trigs = ppu.sort_trigs(trial_dict, trig_data, slide_dict, segment_dict) for slide_id in slide_dict: - offsource_trigs.extend( - zip(trig_data_av[keys[2]][slide_id], sorted_trigs[slide_id]) - ) + offsource_trigs.extend(zip(trig_data_av[keys[2]][slide_id], sorted_trigs[slide_id])) offsource_trigs.sort(key=lambda element: element[0]) offsource_trigs.reverse() # Calculate chirp masses of templates -logging.info('Loading triggers template masses') -bank_data = HFile(opts.bank_file, 'r') +logging.info("Loading triggers template masses") +bank_data = HFile(opts.bank_file, "r") template_mchirps = mchirp_from_mass1_mass2( - bank_data['mass1'][...], - bank_data['mass2'][...] - ) + bank_data["mass1"][...], bank_data["mass2"][...] +) # ========================================= # Output of loudest offsource triggers data @@ -405,55 +464,68 @@ if lofft_outfile: for i in range(min(len(offsource_trigs), opts.num_loudest_off_trigs)): bestnr = offsource_trigs[i][0] trig_id = offsource_trigs[i][1] - trig_index = \ - np.where(trig_data['network/event_id'] == trig_id)[0][0] + trig_index = np.where(trig_data["network/event_id"] == trig_id)[0][0] ifo_trig_index = { - ifo: np.where(trig_data[ifo+'/event_id'] == trig_id)[0][0] - for ifo in ifos + ifo: np.where(trig_data[ifo + "/event_id"] == trig_id)[0][0] for ifo in ifos } - trig_slide_id = int(trig_data['network/slide_id'][trig_index]) + trig_slide_id = int(trig_data["network/slide_id"][trig_index]) # Get trial of trigger, triggers with 'No trial' should have # already been removed! for j, trial in enumerate(trial_dict[trig_slide_id]): - if trig_data['network/end_time_gc'][trig_index] in trial: + if trig_data["network/end_time_gc"][trig_index] in trial: chunk_num = j break else: - chunk_num = 'No trial' + chunk_num = "No trial" # Get FAP of trigger pval = sum(sorted_bkgd > bestnr) / total_trials pval = format_pvalue_str(pval, total_trials) - d = [chunk_num, trig_slide_id, pval, - trig_data['network/end_time_gc'][trig_index], - bank_data['mass1'][trig_data['network/template_id'][trig_index]], - bank_data['mass2'][trig_data['network/template_id'][trig_index]], - template_mchirps[trig_data['network/template_id'][trig_index]], - bank_data['spin1z'][trig_data['network/template_id'][trig_index]], - bank_data['spin2z'][trig_data['network/template_id'][trig_index]], - trig_data['network/ra'][trig_index], - trig_data['network/dec'][trig_index], - trig_data['network/coherent_snr'][trig_index], - trig_data['network/my_network_chisq'][trig_index], - trig_data['network/null_snr'][trig_index]] - d.extend([trig_data[ifo+'/snr'][ifo_trig_index[ifo]] - for ifo in ifos]) + d = [ + chunk_num, + trig_slide_id, + pval, + trig_data["network/end_time_gc"][trig_index], + bank_data["mass1"][trig_data["network/template_id"][trig_index]], + bank_data["mass2"][trig_data["network/template_id"][trig_index]], + template_mchirps[trig_data["network/template_id"][trig_index]], + bank_data["spin1z"][trig_data["network/template_id"][trig_index]], + bank_data["spin2z"][trig_data["network/template_id"][trig_index]], + trig_data["network/ra"][trig_index], + trig_data["network/dec"][trig_index], + trig_data["network/coherent_snr"][trig_index], + trig_data["network/my_network_chisq"][trig_index], + trig_data["network/null_snr"][trig_index], + ] + d.extend([trig_data[ifo + "/snr"][ifo_trig_index[ifo]] for ifo in ifos]) d.extend([slide_dict[trig_slide_id][ifo] for ifo in ifos]) - d.extend([trig_data[ifo+'/end_time'][ifo_trig_index[ifo]] - for ifo in ifos]) + d.extend([trig_data[ifo + "/end_time"][ifo_trig_index[ifo]] for ifo in ifos]) d.append(bestnr) td.append(d) # th: table header # Check against pycbc_pygrb_minifollowups prior to changing any of these - th = ['Trial', 'Slide Num', 'p-value', 'GC GPS time (s)', - 'Rec. m1', 'Rec. m2', 'Rec. Mc', 'Rec. spin1z', 'Rec. spin2z', - 'Rec. RA', 'Rec. Dec', 'SNR', 'Chi^2', 'Null SNR'] - th.extend([ifo+' SNR' for ifo in ifos]) - th.extend([ifo+' time shift (s)' for ifo in ifos]) - th.extend([ifo+' GPS time (s)' for ifo in ifos]) - th.append('BestNR') + th = [ + "Trial", + "Slide Num", + "p-value", + "GC GPS time (s)", + "Rec. m1", + "Rec. m2", + "Rec. Mc", + "Rec. spin1z", + "Rec. spin2z", + "Rec. RA", + "Rec. Dec", + "SNR", + "Chi^2", + "Null SNR", + ] + th.extend([ifo + " SNR" for ifo in ifos]) + th.extend([ifo + " time shift (s)" for ifo in ifos]) + th.extend([ifo + " GPS time (s)" for ifo in ifos]) + th.append("BestNR") # When len(offsource_trigs) == 0, the loop above leaves td = [] unchanged # and this case needs to be handled adequately prior to moving on @@ -465,9 +537,8 @@ if lofft_outfile: td = list(zip(*td)) # Write to h5 file - logging.info("Writing %d loudest offsource triggers to h5 file.", - len(td[0])) - lofft_h5_fp = HFile(lofft_h5_outfile, 'w') + logging.info("Writing %d loudest offsource triggers to h5 file.", len(td[0])) + lofft_h5_fp = HFile(lofft_h5_outfile, "w") for i, key in enumerate(th): lofft_h5_fp.create_dataset(key, data=td[i]) lofft_h5_fp.close() @@ -482,28 +553,41 @@ if lofft_outfile: td = [np.asarray(d) for d in td] # Format of table data - format_strings = ['##.##', '##.##', None, '##.###', - '##.##', '##.##', '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', '##.##', '##.##'] - format_strings.extend(['##.##' for _ in ifos]) - format_strings.extend(['##.##' for _ in ifos]) - format_strings.extend(['##.###' for _ in ifos]) - format_strings.extend(['##.##']) - html_table = pycbc.results.html_table(td, th, - format_strings=format_strings, - page_size=30) - kwds = {'title': "Parameters of loudest offsource triggers", - 'caption': "Parameters of the " + - str(min(len(offsource_trigs), - opts.num_loudest_off_trigs)) + - " loudest offsource triggers. " + - "The median reweighted SNR value is " + - str(median_bestnr) + - ". The median SNR value is " + - str(median_snr), - 'cmd': ' '.join(sys.argv), } - pycbc.results.save_fig_with_metadata(str(html_table), - lofft_outfile, **kwds) + format_strings = [ + "##.##", + "##.##", + None, + "##.###", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + ] + format_strings.extend(["##.##" for _ in ifos]) + format_strings.extend(["##.##" for _ in ifos]) + format_strings.extend(["##.###" for _ in ifos]) + format_strings.extend(["##.##"]) + html_table = pycbc.results.html_table( + td, th, format_strings=format_strings, page_size=30 + ) + kwds = { + "title": "Parameters of loudest offsource triggers", + "caption": "Parameters of the " + + str(min(len(offsource_trigs), opts.num_loudest_off_trigs)) + + " loudest offsource triggers. " + + "The median reweighted SNR value is " + + str(median_bestnr) + + ". The median SNR value is " + + str(median_snr), + "cmd": " ".join(sys.argv), + } + pycbc.results.save_fig_with_metadata(str(html_table), lofft_outfile, **kwds) # Store BestNR and FAP values: for collective FAP value studies at the # end of an observing run collectively @@ -516,21 +600,24 @@ if lofft_outfile: # Load on source triggers # ======================= if onsource_file: - # Get trigs - on_trigs = ppu.load_data(onsource_file, ifos, data_tag=None, - rw_snr_threshold=opts.newsnr_threshold, - slide_id=0) + on_trigs = ppu.load_data( + onsource_file, + ifos, + data_tag=None, + rw_snr_threshold=opts.newsnr_threshold, + slide_id=0, + ) # Record loudest trig by BestNR loud_on_bestnr = 0 if on_trigs: - on_trigs_bestnrs = on_trigs['network/reweighted_snr'][...] + on_trigs_bestnrs = on_trigs["network/reweighted_snr"][...] # Gather bestNR index if on_trigs_bestnrs.size > 0: bestNR_event = np.argmax(on_trigs_bestnrs) - loud_on_bestnr_trigs = on_trigs['network/event_id'][bestNR_event] + loud_on_bestnr_trigs = on_trigs["network/event_id"][bestNR_event] loud_on_bestnr = on_trigs_bestnrs[bestNR_event] # If the loudest event has bestnr = 0, there is no event at all! if loud_on_bestnr == 0: @@ -544,37 +631,55 @@ if onsource_file: # Gather data if loud_on_bestnr_trigs: trig_id = loud_on_bestnr_trigs - trig_index = np.where(on_trigs['network/event_id'] == trig_id)[0][0] + trig_index = np.where(on_trigs["network/event_id"] == trig_id)[0][0] ifo_trig_index = { - ifo: np.where(on_trigs[ifo+'/event_id'] == trig_id)[0][0] - for ifo in ifos + ifo: np.where(on_trigs[ifo + "/event_id"] == trig_id)[0][0] for ifo in ifos } num_trials_louder = 0 - pval = sum(sorted_bkgd > loud_on_bestnr)/total_trials + pval = sum(sorted_bkgd > loud_on_bestnr) / total_trials pval = format_pvalue_str(pval, total_trials) - d = [pval, - on_trigs['network/end_time_gc'][trig_index], - bank_data['mass1'][on_trigs['network/template_id'][trig_index]], - bank_data['mass2'][on_trigs['network/template_id'][trig_index]], - template_mchirps[on_trigs['network/template_id'][trig_index]], - bank_data['spin1z'][on_trigs['network/template_id'][trig_index]], - bank_data['spin2z'][on_trigs['network/template_id'][trig_index]], - on_trigs['network/ra'][trig_index], - on_trigs['network/dec'][trig_index], - on_trigs['network/coherent_snr'][trig_index], - on_trigs['network/my_network_chisq'][trig_index], - on_trigs['network/null_snr'][trig_index]] + \ - [on_trigs[ifo+'/snr'][ifo_trig_index[ifo]] for ifo in ifos] + \ - [on_trigs[ifo+'/end_time'][ifo_trig_index[ifo]] - for ifo in ifos] + [loud_on_bestnr] + d = ( + [ + pval, + on_trigs["network/end_time_gc"][trig_index], + bank_data["mass1"][on_trigs["network/template_id"][trig_index]], + bank_data["mass2"][on_trigs["network/template_id"][trig_index]], + template_mchirps[on_trigs["network/template_id"][trig_index]], + bank_data["spin1z"][on_trigs["network/template_id"][trig_index]], + bank_data["spin2z"][on_trigs["network/template_id"][trig_index]], + on_trigs["network/ra"][trig_index], + on_trigs["network/dec"][trig_index], + on_trigs["network/coherent_snr"][trig_index], + on_trigs["network/my_network_chisq"][trig_index], + on_trigs["network/null_snr"][trig_index], + ] + + [on_trigs[ifo + "/snr"][ifo_trig_index[ifo]] for ifo in ifos] + + [on_trigs[ifo + "/end_time"][ifo_trig_index[ifo]] for ifo in ifos] + + [loud_on_bestnr] + ) td.append(d) # Table header # Check against pycbc_pygrb_minifollowups prior to changing any of these - th = ['p-value', 'GC GPS time (s)', 'Rec. m1', 'Rec. m2', 'Rec. Mc', - 'Rec. spin1z', 'Rec. spin2z', 'Rec. RA', 'Rec. Dec', 'SNR', 'Chi^2', - 'Null SNR'] + [ifo+' SNR' for ifo in ifos] + \ - [ifo+' GPS time (s)' for ifo in ifos] + ['BestNR'] + th = ( + [ + "p-value", + "GC GPS time (s)", + "Rec. m1", + "Rec. m2", + "Rec. Mc", + "Rec. spin1z", + "Rec. spin2z", + "Rec. RA", + "Rec. Dec", + "SNR", + "Chi^2", + "Null SNR", + ] + + [ifo + " SNR" for ifo in ifos] + + [ifo + " GPS time (s)" for ifo in ifos] + + ["BestNR"] + ) td = list(zip(*td)) @@ -584,7 +689,7 @@ if onsource_file: # Write to h5 file logging.info("Writing loudest onsource trigger to h5 file.") - with HFile(lont_h5_outfile, 'w') as lont_h5_fp: + with HFile(lont_h5_outfile, "w") as lont_h5_fp: for i, key in enumerate(th): lont_h5_fp.create_dataset(key, data=td[i]) @@ -592,26 +697,39 @@ if onsource_file: logging.info("Writing loudest onsource trigger to html file.") # Format of table data - format_strings = [None, '##.###', '##.##', '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', '##.##', '##.##', '##.##'] - format_strings.extend(['##.##' for _ in ifos]) - format_strings.extend(['##.###' for _ in ifos]) - format_strings.extend(['##.##']) + format_strings = [ + None, + "##.###", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + ] + format_strings.extend(["##.##" for _ in ifos]) + format_strings.extend(["##.###" for _ in ifos]) + format_strings.extend(["##.##"]) # Table data: assemble human readable message when no trigger is recovered if not loud_on_bestnr_trigs: td = [["-"] for _ in format_strings] td[0][0] = "There are no events" td = [np.asarray(d) for d in td] - html_table = pycbc.results.html_table(td, th, - format_strings=format_strings, - page_size=1) - kwds = {'title': "Loudest event", - 'caption': "Recovered parameters and statistic values of the \ + html_table = pycbc.results.html_table( + td, th, format_strings=format_strings, page_size=1 + ) + kwds = { + "title": "Loudest event", + "caption": "Recovered parameters and statistic values of the \ loudest trigger.", - 'cmd': ' '.join(sys.argv), } - pycbc.results.save_fig_with_metadata(str(html_table), lont_outfile, - **kwds) + "cmd": " ".join(sys.argv), + } + pycbc.results.save_fig_with_metadata(str(html_table), lont_outfile, **kwds) # ======================= # Post-process injections @@ -619,28 +737,27 @@ if onsource_file: if found_missed_file is not None: # Load injections applying reweighted SNR cut found_injs, missed_injs, cut_injs = load_missed_found_injections( - found_missed_file, ifos, bank_data, + found_missed_file, + ifos, + bank_data, snr_threshold=opts.newsnr_threshold, - background_bestnrs=sorted_bkgd + background_bestnrs=sorted_bkgd, ) # Split in injections found surviving vetoes and ones found but vetoed found_after_vetoes, vetoed, *_ = ppu.apply_vetoes_to_found_injs( - found_missed_file, - found_injs, - ifos, - veto_file=opts.veto_file + found_missed_file, found_injs, ifos, veto_file=opts.veto_file ) - if 'bestnr' not in found_after_vetoes: - found_after_vetoes['bestnr'] = np.array([]) + if "bestnr" not in found_after_vetoes: + found_after_vetoes["bestnr"] = np.array([]) # Construct conditions for injection: # 1) found louder than background, - zero_fap = found_after_vetoes['bestnr'] > max_bestnr + zero_fap = found_after_vetoes["bestnr"] > max_bestnr # 2) found (bestnr > 0) but not louder than background (non-zero FAP) - nonzero_fap = ~zero_fap & (found_after_vetoes['bestnr'] != 0) + nonzero_fap = ~zero_fap & (found_after_vetoes["bestnr"] != 0) # 3) missed after being recovered: vetoed (these have bestnr = 0) @@ -652,53 +769,125 @@ if found_missed_file is not None: # Table header # Check against pycbc_pygrb_minifollowups prior to changing any of these # Injections are in zero-lag so no timing information at the IFOs is needed - th = ['Dist'] + ['Eff. Dist. '+site for site in sites] +\ - ['GC GPS time (s)', 'GPS time - Rec. Time'] +\ - ['Inj. m1', 'Inj. m2', 'Inj. Mc', 'Rec. m1', 'Rec. m2', 'Rec. Mc', - 'Inj. inc', 'Inj. RA', 'Inj. Dec', 'Rec. RA', 'Rec. Dec', 'SNR', - 'Chi^2', 'Null SNR'] +\ - ['SNR '+ifo for ifo in ifos] +\ - ['Net Opt SNR'] +\ - ['Opt SNR '+ifo for ifo in ifos] +\ - ['BestNR', 'p-value', - 'Inj S1x', 'Inj S1y', 'Inj S1z', - 'Inj S2x', 'Inj S2y', 'Inj S2z', - 'Rec S1z', 'Rec S2z'] + th = ( + ["Dist"] + + ["Eff. Dist. " + site for site in sites] + + ["GC GPS time (s)", "GPS time - Rec. Time"] + + [ + "Inj. m1", + "Inj. m2", + "Inj. Mc", + "Rec. m1", + "Rec. m2", + "Rec. Mc", + "Inj. inc", + "Inj. RA", + "Inj. Dec", + "Rec. RA", + "Rec. Dec", + "SNR", + "Chi^2", + "Null SNR", + ] + + ["SNR " + ifo for ifo in ifos] + + ["Net Opt SNR"] + + ["Opt SNR " + ifo for ifo in ifos] + + [ + "BestNR", + "p-value", + "Inj S1x", + "Inj S1y", + "Inj S1z", + "Inj S2x", + "Inj S2y", + "Inj S2z", + "Rec S1z", + "Rec S2z", + ] + ) # Format of table data - format_strings = ['##.##'] - format_strings.extend(['##.##' for _ in ifos]) - format_strings.extend(['##.#####', '##.#####', - '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', - '##.##', '##.##']) - format_strings.extend(['##.##' for _ in ifos]) + format_strings = ["##.##"] + format_strings.extend(["##.##" for _ in ifos]) + format_strings.extend( + [ + "##.#####", + "##.#####", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + ] + ) + format_strings.extend(["##.##" for _ in ifos]) # The +1 takes in account the additional network optimal SNR column - format_strings.extend(['##.##' for _ in range(len(ifos) + 1)]) - format_strings.extend(['##.##', '#.######', - '##.##', '##.##', '##.##', - '##.##', '##.##', '##.##', - '##.##', '##.##']) - sngl_snr_keys = ['snr_'+ifo for ifo in ifos] - keys = ['distance'] - keys += ['eff_dist_'+ifo for ifo in ifos] - keys += ['tc', 'time_diff', 'mass1', 'mass2', 'mchirp', 'rec_mass1', - 'rec_mass2', 'rec_mchirp', 'inclination', 'ra', 'dec', 'rec_ra', - 'rec_dec', 'coherent_snr', 'chisq', 'null_snr'] + format_strings.extend(["##.##" for _ in range(len(ifos) + 1)]) + format_strings.extend( + [ + "##.##", + "#.######", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + "##.##", + ] + ) + sngl_snr_keys = ["snr_" + ifo for ifo in ifos] + keys = ["distance"] + keys += ["eff_dist_" + ifo for ifo in ifos] + keys += [ + "tc", + "time_diff", + "mass1", + "mass2", + "mchirp", + "rec_mass1", + "rec_mass2", + "rec_mchirp", + "inclination", + "ra", + "dec", + "rec_ra", + "rec_dec", + "coherent_snr", + "chisq", + "null_snr", + ] keys += sngl_snr_keys - keys += ['net_opt_snr'] - keys += ['optimal_snr_' + ifo for ifo in ifos] - keys += ['bestnr', 'fap', 'spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y', - 'spin2z', 'rec_spin1z', 'rec_spin2z'] + keys += ["net_opt_snr"] + keys += ["optimal_snr_" + ifo for ifo in ifos] + keys += [ + "bestnr", + "fap", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", + "rec_spin1z", + "rec_spin2z", + ] for key in keys: if key not in found_after_vetoes: found_after_vetoes[key] = np.array([]) td = [found_after_vetoes[key][nonzero_fap] for key in keys] td = list(zip(*td)) td.sort(key=lambda elem: elem[0]) - logging.info("Writing %d quiet-found injections to h5 and html files.", - len(td)) + logging.info("Writing %d quiet-found injections to h5 and html files.", len(td)) td = list(zip(*td)) # Handle the case in which there is no data to be placed in the table @@ -706,59 +895,65 @@ if found_missed_file is not None: td = [[]] * len(th) # Write to h5 file - with HFile(qf_h5_outfile, 'w') as qf_h5_fp: + with HFile(qf_h5_outfile, "w") as qf_h5_fp: for i, key in enumerate(th): qf_h5_fp.create_dataset(key, data=td[i]) # Write to html file td = [np.asarray(d) for d in td] - html_table = pycbc.results.html_table(td, th, - format_strings=format_strings, - page_size=20) - kwds = {'title': "Quiet found injections", - 'caption': "Recovered parameters and statistic values of \ + html_table = pycbc.results.html_table( + td, th, format_strings=format_strings, page_size=20 + ) + kwds = { + "title": "Quiet found injections", + "caption": "Recovered parameters and statistic values of \ injections that are recovered, but not louder than \ - background.", 'cmd': ' '.join(sys.argv), } - pycbc.results.save_fig_with_metadata(str(html_table), qf_outfile, - **kwds) + background.", + "cmd": " ".join(sys.argv), + } + pycbc.results.save_fig_with_metadata(str(html_table), qf_outfile, **kwds) # Write quiet triggers to html file if len(cut_injs) == 0: cut_injs = dict.fromkeys(keys, np.array([])) for key in keys: if key not in vetoed: - vetoed[key] = np.full(len(vetoed['mass1']), -1) + vetoed[key] = np.full(len(vetoed["mass1"]), -1) if key not in cut_injs: - cut_injs[key] = np.full(len(cut_injs['mass1']), -1) + cut_injs[key] = np.full(len(cut_injs["mass1"]), -1) if key not in missed_injs: - missed_injs[key] = np.full(len(missed_injs['mass1']), -1) + missed_injs[key] = np.full(len(missed_injs["mass1"]), -1) # Update the data to include a label stating whether the injection was # vetoed, found and then removed due to the reweighted SNR cut, or missed - vetoed['category'] = np.full(len(vetoed['mass1']), 'Vetoed') - cut_injs['category'] = np.full(len(cut_injs['mass1']), 'Cut') - missed_injs['category'] = np.full(len(missed_injs['mass1']), 'Missed') + vetoed["category"] = np.full(len(vetoed["mass1"]), "Vetoed") + cut_injs["category"] = np.full(len(cut_injs["mass1"]), "Cut") + missed_injs["category"] = np.full(len(missed_injs["mass1"]), "Missed") # Update the header and table formatter - th += ['Category'] - format_strings.extend(['##']) - t_missed = [np.concatenate((missed_injs[key], vetoed[key], cut_injs[key])) - for key in keys+['category']] + th += ["Category"] + format_strings.extend(["##"]) + t_missed = [ + np.concatenate((missed_injs[key], vetoed[key], cut_injs[key])) + for key in keys + ["category"] + ] t_missed = list(zip(*t_missed)) t_missed.sort(key=lambda elem: elem[0]) - logging.info("Writing %d missed, vetoed, and cut injections to html file.", - len(t_missed)) + logging.info( + "Writing %d missed, vetoed, and cut injections to html file.", len(t_missed) + ) t_missed = zip(*t_missed) t_missed = [np.asarray(d) for d in t_missed] - html_table = pycbc.results.html_table(t_missed, th, - format_strings=format_strings, - page_size=20) - kwds = {'title': "Missed, vetoed, and cut injections", - 'caption': "Parameters of missed injections and \ + html_table = pycbc.results.html_table( + t_missed, th, format_strings=format_strings, page_size=20 + ) + kwds = { + "title": "Missed, vetoed, and cut injections", + "caption": "Parameters of missed injections and \ recovered parameters and statistic values of \ injections that are recovered but then vetoed or cut because \ their reweighted SNR is below threshold.", - 'cmd': ' '.join(sys.argv), } - pycbc.results.save_fig_with_metadata(str(html_table), mf_outfile, - **kwds) + "cmd": " ".join(sys.argv), + } + pycbc.results.save_fig_with_metadata(str(html_table), mf_outfile, **kwds) # Close the bank file bank_data.close() diff --git a/bin/pygrb/pycbc_pygrb_plot_chisq_veto b/bin/pygrb/pycbc_pygrb_plot_chisq_veto index 91410472ff8..9c8f619cf82 100644 --- a/bin/pygrb/pycbc_pygrb_plot_chisq_veto +++ b/bin/pygrb/pycbc_pygrb_plot_chisq_veto @@ -25,19 +25,21 @@ single detector chi-square vs single IFO/coherent/reweighted/null/coinc SNR. # ============================================================================= # Preamble # ============================================================================= -import sys -import os import logging +import os +import sys + import numpy +import pycbc.version from matplotlib import pyplot as plt from matplotlib import rc -import pycbc.version + from pycbc import init_logging -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.results import pygrb_plotting_utils as plu +from pycbc.results import pygrb_postprocessing_utils as ppu -plt.switch_backend('Agg') -rc('font', size=14) +plt.switch_backend("Agg") +rc("font", size=14) __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -51,18 +53,16 @@ __program__ = "pycbc_pygrb_plot_chisq_veto" # Function to calculate chi-square weight for the reweighted SNR def new_snr_chisq(snr, new_snr, chisq_index=4.0, chisq_nhigh=3.0): """Returns the chi-square value needed to weight SNR into new SNR""" - - chisqnorm = (snr/new_snr)**chisq_index + chisqnorm = (snr / new_snr) ** chisq_index if chisqnorm <= 1: - return 1E-20 + return 1e-20 - return (2*chisqnorm - 1)**(chisq_nhigh/chisq_index) + return (2 * chisqnorm - 1) ** (chisq_nhigh / chisq_index) # Function that produces the contours to be plotted def calculate_contours(opts, new_snrs=None): """Generate the contours for the veto plots""" - # Add the new SNR threshold contour to the list if necessary # and keep track of where it is if new_snrs is None: @@ -79,20 +79,20 @@ def calculate_contours(opts, new_snrs=None): snr_vals = numpy.asarray(list(snr_low_vals) + list(snr_high_vals)) # Initialise contours - contours = numpy.zeros([len(new_snrs), len(snr_vals)], - dtype=numpy.float64) + contours = numpy.zeros([len(new_snrs), len(snr_vals)], dtype=numpy.float64) # Loop over SNR values and calculate chisq variable needed for j, snr in enumerate(snr_vals): for i, new_snr in enumerate(new_snrs): - contours[i][j] = new_snr_chisq(snr, new_snr, - opts.chisq_index, - opts.chisq_nhigh) + contours[i][j] = new_snr_chisq( + snr, new_snr, opts.chisq_index, opts.chisq_nhigh + ) # Colors and styles of the contours - colors = ["k-" if snr == opts.newsnr_threshold else - "y-" if snr == int(snr) else - "y--" for snr in new_snrs] + colors = [ + "k-" if snr == opts.newsnr_threshold else "y-" if snr == int(snr) else "y--" + for snr in new_snrs + ] return contours, snr_vals, cont_value, colors @@ -101,19 +101,37 @@ def calculate_contours(opts, new_snrs=None): # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("-t", "--trig-file", action="store", - default=None, required=True, - help="The location of the trigger file") -parser.add_argument("--found-missed-file", - help="The hdf injection results file", required=False) -parser.add_argument("-z", "--zoom-in", default=False, action="store_true", - help="Output file a zoomed in version of the plot.") -parser.add_argument("-y", "--y-variable", required=True, - choices=['network', 'bank', 'auto', 'power'], - help="Quantity to plot on the vertical axis.") -parser.add_argument("--snr-type", default='coherent', - choices=['coherent', 'coincident', 'null', 'reweighted', - 'single'], help="SNR value to plot on x-axis.") +parser.add_argument( + "-t", + "--trig-file", + action="store", + default=None, + required=True, + help="The location of the trigger file", +) +parser.add_argument( + "--found-missed-file", help="The hdf injection results file", required=False +) +parser.add_argument( + "-z", + "--zoom-in", + default=False, + action="store_true", + help="Output file a zoomed in version of the plot.", +) +parser.add_argument( + "-y", + "--y-variable", + required=True, + choices=["network", "bank", "auto", "power"], + help="Quantity to plot on the vertical axis.", +) +parser.add_argument( + "--snr-type", + default="coherent", + choices=["coherent", "coincident", "null", "reweighted", "single"], + help="SNR value to plot on x-axis.", +) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_bestnr_opts(parser) ppu.pygrb_add_slide_opts(parser) @@ -124,15 +142,16 @@ init_logging(opts.verbose, format="%(asctime)s: %(levelname)s: %(message)s") # Check options trig_file = os.path.abspath(opts.trig_file) -found_missed_file = os.path.abspath(opts.found_missed_file) \ - if opts.found_missed_file else None +found_missed_file = ( + os.path.abspath(opts.found_missed_file) if opts.found_missed_file else None +) zoom_in = opts.zoom_in veto_type = opts.y_variable ifo = opts.ifo snr_type = opts.snr_type # If this is false, coherent SNR is used on the horizontal axis # otherwise the single IFO SNR is used -if snr_type == 'single': +if snr_type == "single": if ifo is None: err_msg = "--ifo must be given to plot single IFO SNR veto" parser.error(err_msg) @@ -141,27 +160,31 @@ if snr_type == 'single': # TODO: fix vetoes # Prepare plot title and caption -veto_labels = {'network': "Network Power", - 'bank': "Bank", - 'auto': "Auto", - 'power': "Power"} +veto_labels = { + "network": "Network Power", + "bank": "Bank", + "auto": "Auto", + "power": "Power", +} if opts.plot_title is None: opts.plot_title = veto_labels[veto_type] + " Chi Square" - if veto_type != 'network': + if veto_type != "network": opts.plot_title = ifo + opts.plot_title - if snr_type == 'single': + if snr_type == "single": opts.plot_title += f" vs {ifo} SNR" else: opts.plot_title += f" vs {snr_type.capitalize()} SNR" if opts.plot_caption is None: - opts.plot_caption = ("Blue crosses: background triggers. ") + opts.plot_caption = "Blue crosses: background triggers. " if found_missed_file: opts.plot_caption += "Red crosses: injections triggers. " - if veto_type == 'network': - opts.plot_caption += ("Gray shaded region: area cut by the " + - "reweighted SNR threshold. " + - "Black line: reweighted SNR threshold. Yellow " + - "lines: contours of constant reweighted SNR.") + if veto_type == "network": + opts.plot_caption += ( + "Gray shaded region: area cut by the " + "reweighted SNR threshold. " + "Black line: reweighted SNR threshold. Yellow " + "lines: contours of constant reweighted SNR." + ) logging.info("Imported and ready to go.") @@ -181,69 +204,65 @@ segment_dict = ppu.load_segment_dict(trig_file) # Construct trials removing vetoed times trial_dict, total_trials = ppu.construct_trials( - opts.seg_files, - segment_dict, - ifos, - slide_dict, - opts.veto_file + opts.seg_files, segment_dict, ifos, slide_dict, opts.veto_file ) # Load trigger and injections data: ensure that newtwork power chi-square plots # show all the data to see the impact of the reweighted SNR cut, otherwise # remove points with reweighted SNR below threshold -rw_snr_threshold = None if veto_type == 'network' else opts.newsnr_threshold -trig_data = ppu.load_data(trig_file, ifos, data_tag='trigs', - rw_snr_threshold=rw_snr_threshold, - slide_id=opts.slide_id) -inj_data = ppu.load_data(found_missed_file, ifos, data_tag='injs', - rw_snr_threshold=rw_snr_threshold, - slide_id=0) +rw_snr_threshold = None if veto_type == "network" else opts.newsnr_threshold +trig_data = ppu.load_data( + trig_file, + ifos, + data_tag="trigs", + rw_snr_threshold=rw_snr_threshold, + slide_id=opts.slide_id, +) +inj_data = ppu.load_data( + found_missed_file, + ifos, + data_tag="injs", + rw_snr_threshold=rw_snr_threshold, + slide_id=0, +) # Dataset name for the horizontal direction -if snr_type == 'single': - x_key = ifo + '/snr' +if snr_type == "single": + x_key = ifo + "/snr" else: - x_key = 'network/' + snr_type + '_snr' + x_key = "network/" + snr_type + "_snr" # Dataset name for the vertical direction and for normalization -if veto_type == 'power': - y_key = opts.ifo + '/chisq' -elif veto_type in ['bank', 'auto']: - y_key = opts.ifo + '/' + veto_type + '_chisq' +if veto_type == "power": + y_key = opts.ifo + "/chisq" +elif veto_type in ["bank", "auto"]: + y_key = opts.ifo + "/" + veto_type + "_chisq" else: - y_key = 'network/my_network_chisq' + y_key = "network/my_network_chisq" keys = [x_key, y_key] # The network chi-square is already normalized so it does not require a key # for the number of degrees of freedom -if veto_type != 'network': - keys += [y_key + '_dof'] +if veto_type != "network": + keys += [y_key + "_dof"] # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors found_trigs_slides = ppu.extract_trig_properties( - trial_dict, - trig_data, - slide_dict, - segment_dict, - keys + trial_dict, trig_data, slide_dict, segment_dict, keys ) found_trigs = {} for key in keys: found_trigs[key] = numpy.concatenate( - [found_trigs_slides[key][slide_id][:] for slide_id in slide_dict] + [found_trigs_slides[key][slide_id][:] for slide_id in slide_dict] ) # Gather injections found surviving vetoes found_injs, *_ = ppu.apply_vetoes_to_found_injs( - opts.found_missed_file, - inj_data, - ifos, - veto_file=opts.veto_file, - keys=keys + opts.found_missed_file, inj_data, ifos, veto_file=opts.veto_file, keys=keys ) # Sanity checks -for test in zip(keys[0:2], ['x', 'y']): +for test in zip(keys[0:2], ["x", "y"]): if found_trigs[test[0]] is None and found_injs[test[0]] is None: err_msg = "No data to be plotted on the " + test[1] + "-axis was found" raise RuntimeError(err_msg) @@ -262,24 +281,31 @@ logging.info("Plotting...") # Determine x-axis values of triggers and injections # Default is coherent SNR -x_label = ifo if snr_type == 'single' else snr_type.capitalize() +x_label = ifo if snr_type == "single" else snr_type.capitalize() x_label += " SNR" # Determine the minumum and maximum SNR value we are dealing with -x_min = 0.9*plu.axis_min_value(found_trigs[x_key], found_injs[x_key], - found_missed_file) -x_max = 1.1*plu.axis_max_value(found_trigs[x_key], found_injs[x_key], - found_missed_file) +x_min = 0.9 * plu.axis_min_value( + found_trigs[x_key], found_injs[x_key], found_missed_file +) +x_max = 1.1 * plu.axis_max_value( + found_trigs[x_key], found_injs[x_key], found_missed_file +) # Determine the minimum and maximum chi-square value we are dealing with -y_min = 0.9*plu.axis_min_value(found_trigs[y_key], found_injs[y_key], - found_missed_file) -y_max = 1.1*plu.axis_max_value(found_trigs[y_key], found_injs[y_key], - found_missed_file) +y_min = 0.9 * plu.axis_min_value( + found_trigs[y_key], found_injs[y_key], found_missed_file +) +y_max = 1.1 * plu.axis_max_value( + found_trigs[y_key], found_injs[y_key], found_missed_file +) # Determine y-axis label -y_label = "Network power chi-square" if veto_type == 'network' \ +y_label = ( + "Network power chi-square" + if veto_type == "network" else f"{ifo} Single {veto_labels[veto_type].lower()} chi-square" +) # Determine contours for plots conts = None @@ -287,24 +313,30 @@ snr_vals = None cont_value = None colors = None # Enable countours of constant reweighted SNR as a function of coherent SNR -if snr_type == 'coherent': - conts, snr_vals, cont_value, colors = calculate_contours(opts, - new_snrs=None) +if snr_type == "coherent": + conts, snr_vals, cont_value, colors = calculate_contours(opts, new_snrs=None) # The cut in reweighted SNR involves only the network power chi-square -if veto_type != 'network': +if veto_type != "network": cont_value = None # Produce the veto vs. SNR plot if not opts.x_lims: if zoom_in: - opts.x_lims = str(x_min)+',50' - opts.y_lims = str(y_min)+',20000' + opts.x_lims = str(x_min) + ",50" + opts.y_lims = str(y_min) + ",20000" else: - opts.x_lims = str(x_min)+','+str(x_max) - opts.y_lims = str(y_min)+','+str(10*y_max) -plu.pygrb_plotter([found_trigs[x_key], found_trigs[y_key]], - [found_injs[x_key], found_injs[y_key]], - x_label, y_label, opts, - snr_vals=snr_vals, conts=conts, colors=colors, - shade_cont_value=cont_value, vert_spike=True, - cmd=' '.join(sys.argv)) + opts.x_lims = str(x_min) + "," + str(x_max) + opts.y_lims = str(y_min) + "," + str(10 * y_max) +plu.pygrb_plotter( + [found_trigs[x_key], found_trigs[y_key]], + [found_injs[x_key], found_injs[y_key]], + x_label, + y_label, + opts, + snr_vals=snr_vals, + conts=conts, + colors=colors, + shade_cont_value=cont_value, + vert_spike=True, + cmd=" ".join(sys.argv), +) diff --git a/bin/pygrb/pycbc_pygrb_plot_coh_ifosnr b/bin/pygrb/pycbc_pygrb_plot_coh_ifosnr index 7fd881a2932..77bc5006c8f 100644 --- a/bin/pygrb/pycbc_pygrb_plot_coh_ifosnr +++ b/bin/pygrb/pycbc_pygrb_plot_coh_ifosnr @@ -24,25 +24,26 @@ Plot single IFO SNR vs coherent SNR for a PyGRB run. # ============================================================================= # Preamble # ============================================================================= -import sys -import os -import logging import collections +import logging import operator -from matplotlib import pyplot as plt -from matplotlib import rc +import os +import sys + import numpy +import pycbc.version import scipy +from matplotlib import pyplot as plt +from matplotlib import rc -import pycbc.version from pycbc import init_logging from pycbc.detector import Detector -from pycbc.results import save_fig_with_metadata -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.results import pygrb_plotting_utils as plu +from pycbc.results import pygrb_postprocessing_utils as ppu +from pycbc.results import save_fig_with_metadata -plt.switch_backend('Agg') -rc('font', size=14) +plt.switch_backend("Agg") +rc("font", size=14) __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -56,7 +57,6 @@ __program__ = "pycbc_pygrb_plot_coh_ifosnr" # Plot lines representing deviations based on non-central chi-square def plot_deviation(percentile, snr_grid, y, ax, style): """Plot deviations based on non-central chi-square""" - # ncx2: non-central chi-squared; ppf: percent point function # ax.plot(snr_grid, scipy.stats.ncx2.ppf(percentile, 2, y*y)**0.5, style) @@ -64,7 +64,7 @@ def plot_deviation(percentile, snr_grid, y, ax, style): # original code line (commented out above) y_vals = scipy.stats.ncx2.ppf(percentile, 2, y * y) ** 0.5 y_vals = numpy.unique(y_vals) - x_vals = snr_grid[0:len(y_vals)] + x_vals = snr_grid[0 : len(y_vals)] n_vals = int(len(y_vals) / 2) f = scipy.interpolate.interp1d( x_vals[0:n_vals], @@ -152,46 +152,46 @@ segment_dict = ppu.load_segment_dict(trig_file) # Construct trials removing vetoed times trial_dict, total_trials = ppu.construct_trials( - opts.seg_files, - segment_dict, - ifos, - slide_dict, - opts.veto_file + opts.seg_files, segment_dict, ifos, slide_dict, opts.veto_file ) # Load triggers/injections (apply reweighted SNR cut, not vetoes) -trig_data = ppu.load_data(trig_file, ifos, data_tag='trigs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=opts.slide_id) -inj_data = ppu.load_data(found_missed_file, ifos, data_tag='injs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=0) +trig_data = ppu.load_data( + trig_file, + ifos, + data_tag="trigs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=opts.slide_id, +) +inj_data = ppu.load_data( + found_missed_file, + ifos, + data_tag="injs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=0, +) # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors # Coherent SNR is always used -x_key = 'network/coherent_snr' +x_key = "network/coherent_snr" keys = [x_key] # Get parameters necessary for antenna responses -keys += ['network/ra', 'network/dec', 'network/end_time_gc'] +keys += ["network/ra", "network/dec", "network/end_time_gc"] # Get event_ids -keys += [ifo+'/event_id' for ifo in ifos] -keys += ['network/event_id'] +keys += [ifo + "/event_id" for ifo in ifos] +keys += ["network/event_id"] # Get single ifo SNR data -keys += [ifo+'/snr' for ifo in ifos] +keys += [ifo + "/snr" for ifo in ifos] # Get sigma for each ifo -keys += [ifo+'/sigmasq' for ifo in ifos] +keys += [ifo + "/sigmasq" for ifo in ifos] found_trigs_slides = ppu.extract_trig_properties( - trial_dict, - trig_data, - slide_dict, - segment_dict, - keys + trial_dict, trig_data, slide_dict, segment_dict, keys ) found_trigs = {} for key in keys: found_trigs[key] = numpy.concatenate( - [found_trigs_slides[key][slide_id][:] for slide_id in slide_dict] + [found_trigs_slides[key][slide_id][:] for slide_id in slide_dict] ) # Complete the dictionary found_trigs @@ -199,43 +199,38 @@ for key in keys: for ifo in ifos: sorted_ifo_ids = numpy.array( [ - numpy.nonzero(found_trigs[ifo + '/event_id'] == idx)[0][0] - for idx in found_trigs['network/event_id'] + numpy.nonzero(found_trigs[ifo + "/event_id"] == idx)[0][0] + for idx in found_trigs["network/event_id"] ] ) - for key in [ifo+'/snr', ifo+'/sigmasq']: + for key in [ifo + "/snr", ifo + "/sigmasq"]: found_trigs[key] = found_trigs[key][sorted_ifo_ids] # 2) Get antenna response based parameters -found_trigs['sigma_tot'] = numpy.zeros(len(found_trigs[x_key])) +found_trigs["sigma_tot"] = numpy.zeros(len(found_trigs[x_key])) for ifo in ifos: antenna = Detector(ifo) ifo_f_resp = ppu.get_antenna_responses( antenna, - found_trigs['network/ra'], - found_trigs['network/dec'], - found_trigs['network/end_time_gc'] + found_trigs["network/ra"], + found_trigs["network/dec"], + found_trigs["network/end_time_gc"], ) # Get the average for f_resp_mean and calculate sigma_tot - found_trigs[ifo+'/f_resp_mean'] = ifo_f_resp.mean() - found_trigs['sigma_tot'] += found_trigs[ifo+'/sigmasq'] * ifo_f_resp + found_trigs[ifo + "/f_resp_mean"] = ifo_f_resp.mean() + found_trigs["sigma_tot"] += found_trigs[ifo + "/sigmasq"] * ifo_f_resp # 3) Calculate the mean, max, and min sigmas for ifo in ifos: - sigma_norm = found_trigs[ifo+'/sigmasq'] / found_trigs['sigma_tot'] - found_trigs[ifo+'/sigma_mean'] = sigma_norm.mean() \ - if len(sigma_norm) else 0 + sigma_norm = found_trigs[ifo + "/sigmasq"] / found_trigs["sigma_tot"] + found_trigs[ifo + "/sigma_mean"] = sigma_norm.mean() if len(sigma_norm) else 0 if ifo == opts.ifo: - found_trigs['sigma_max'] = sigma_norm.max() if len(sigma_norm) else 0 - found_trigs['sigma_min'] = sigma_norm.min() if len(sigma_norm) else 0 + found_trigs["sigma_max"] = sigma_norm.max() if len(sigma_norm) else 0 + found_trigs["sigma_min"] = sigma_norm.min() if len(sigma_norm) else 0 # Gather injections found surviving vetoes found_injs, *_ = ppu.apply_vetoes_to_found_injs( - found_missed_file, - inj_data, - ifos, - veto_file=opts.veto_file, - keys=keys + found_missed_file, inj_data, ifos, veto_file=opts.veto_file, keys=keys ) # Generate plots @@ -243,27 +238,23 @@ logging.info("Plotting...") # Order the IFOs by sensitivity ifo_sensitivity = { - ifo: found_trigs[ifo+'/f_resp_mean'] * found_trigs[ifo+'/sigma_mean'] + ifo: found_trigs[ifo + "/f_resp_mean"] * found_trigs[ifo + "/sigma_mean"] for ifo in ifos } ifo_sensitivity = collections.OrderedDict( sorted(ifo_sensitivity.items(), key=operator.itemgetter(1), reverse=True) ) -loudness_labels = ['first', 'second', 'third'] +loudness_labels = ["first", "second", "third"] # Determine the maximum coherent SNR value we are dealing with -x_max = plu.axis_max_value( - found_trigs[x_key], found_injs[x_key], found_missed_file -) +x_max = plu.axis_max_value(found_trigs[x_key], found_injs[x_key], found_missed_file) max_snr = x_max if x_max < 50.0: max_snr = 50.0 # Determine the maximum auto veto value we are dealing with -y_key = opts.ifo+'/snr' -y_max = plu.axis_max_value( - found_trigs[y_key], found_injs[y_key], found_missed_file -) +y_key = opts.ifo + "/snr" +y_max = plu.axis_max_value(found_trigs[y_key], found_injs[y_key], found_missed_file) # Setup the plots x_label = "Coherent SNR" @@ -271,34 +262,33 @@ y_label = opts.ifo + " SNR" fig = plt.figure() ax = fig.gca() # Plot trigger data -ax.plot(found_trigs[x_key], found_trigs[y_key], 'bx') +ax.plot(found_trigs[x_key], found_trigs[y_key], "bx") ax.grid() # Plot injection data if found_missed_file: - ax.plot(found_injs[x_key], found_injs[y_key], 'r+') + ax.plot(found_injs[x_key], found_injs[y_key], "r+") # Sigma-mean, min, max y_data = { - 'mean': found_trigs[opts.ifo+'/sigma_mean'], - 'min': found_trigs['sigma_min'], - 'max': found_trigs['sigma_max'], + "mean": found_trigs[opts.ifo + "/sigma_mean"], + "min": found_trigs["sigma_min"], + "max": found_trigs["sigma_max"], } # Calculate: zoom-snr * sqrt(response * sigma-mean, min, max) snr_grid = numpy.arange(0.01, max_snr, 1) y_data = dict( - (key, - snr_grid * (found_trigs[opts.ifo+'/f_resp_mean'] * val) ** 0.5) + (key, snr_grid * (found_trigs[opts.ifo + "/f_resp_mean"] * val) ** 0.5) for key, val in y_data.items() ) for key in y_data: - ax.plot(snr_grid, y_data[key], 'g-') + ax.plot(snr_grid, y_data[key], "g-") # 2 sigma (0.9545) -plot_deviation(0.02275, snr_grid, y_data['min'], ax, 'm-') -plot_deviation(1 - 0.02275, snr_grid, y_data['max'], ax, 'm-') +plot_deviation(0.02275, snr_grid, y_data["min"], ax, "m-") +plot_deviation(1 - 0.02275, snr_grid, y_data["max"], ax, "m-") # 3 sigma (0.9973) -plot_deviation(0.00135, snr_grid, y_data['min'], ax, 'c-') -plot_deviation(1 - 0.00135, snr_grid, y_data['max'], ax, 'c-') +plot_deviation(0.00135, snr_grid, y_data["min"], ax, "c-") +plot_deviation(1 - 0.00135, snr_grid, y_data["max"], ax, "c-") # Non-zoomed plot -ax.plot([0, max_snr], [4, 4], 'k-') +ax.plot([0, max_snr], [4, 4], "k-") ax.set_xlabel(x_label) ax.set_ylabel(y_label) ax.set_xlim([0, 1.01 * x_max]) @@ -311,10 +301,8 @@ if loudness_index < 2: polyy = [4, 4] polyx.extend([max_snr, 0]) polyy.extend([limy, limy]) - ax.fill(polyx, polyy, color='#dddddd') -opts.plot_title += ( - f" ({loudness_labels[loudness_index]} average sensitivity)" -) + ax.fill(polyx, polyy, color="#dddddd") +opts.plot_title += f" ({loudness_labels[loudness_index]} average sensitivity)" # Zoom in if asked to do so if opts.zoom_in: ax.set_xlim([6, 50]) @@ -323,7 +311,7 @@ if opts.zoom_in: save_fig_with_metadata( fig, opts.output_file, - cmd=' '.join(sys.argv), + cmd=" ".join(sys.argv), title=opts.plot_title, caption=opts.plot_caption, ) diff --git a/bin/pygrb/pycbc_pygrb_plot_injs_results b/bin/pygrb/pycbc_pygrb_plot_injs_results index 4e73cded613..0d75b349c61 100644 --- a/bin/pygrb/pycbc_pygrb_plot_injs_results +++ b/bin/pygrb/pycbc_pygrb_plot_injs_results @@ -24,18 +24,19 @@ Plot found/missed injection properties for the triggered search (PyGRB). import logging import os.path import sys -import matplotlib.pyplot as plt + import matplotlib +import matplotlib.pyplot as plt import numpy as np +import pycbc.version import pycbc.conversions import pycbc.results -import pycbc.version from pycbc import init_logging -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.detector import Detector +from pycbc.results import pygrb_postprocessing_utils as ppu -plt.switch_backend('Agg') +plt.switch_backend("Agg") matplotlib.rc("image") __author__ = "Francesco Pannarale " @@ -49,37 +50,34 @@ __program__ = "pycbc_pygrb_plot_injs_results" # ============================================================================= def process_var_strings(qty): """Add underscores to match HDF column name conventions""" - - qty = qty.replace('skyerror', 'sky_error') - qty = qty.replace('cos', 'cos_') - qty = qty.replace('abs', 'abs_') - qty = qty.replace('endtime', 'end_time') - qty = qty.replace('spin1a', 'spin1_a') - qty = qty.replace('spin2a', 'spin2_a') - qty = qty.replace('chip', 'chi_p') - qty = qty.replace('netoptsnr', 'net_opt_snr') + qty = qty.replace("skyerror", "sky_error") + qty = qty.replace("cos", "cos_") + qty = qty.replace("abs", "abs_") + qty = qty.replace("endtime", "end_time") + qty = qty.replace("spin1a", "spin1_a") + qty = qty.replace("spin2a", "spin2_a") + qty = qty.replace("chip", "chi_p") + qty = qty.replace("netoptsnr", "net_opt_snr") return qty def complete_incl_data(injs, key, tag): """Extract data related to inclination from raw injection data""" - local_dict = {} # Whether the user requests incl, |incl|, cos(incl), or cos(|incl|) # the following information is needed # TODO: make sure it is clear that inclination is interpreted as theta_jn - local_dict['incl'] = injs[tag+'/thetajn'] + local_dict["incl"] = injs[tag + "/thetajn"] # Requesting |incl| or cos(|incl|) - if 'abs_' in key: - local_dict['abs_incl'] = 0.5*np.pi - \ - abs(local_dict['incl'] - 0.5*np.pi) + if "abs_" in key: + local_dict["abs_incl"] = 0.5 * np.pi - abs(local_dict["incl"] - 0.5 * np.pi) # Requesting cos(incl) or cos(|incl|): take cosine - if 'cos_' in key: - angle = key.replace('cos_', '') + if "cos_" in key: + angle = key.replace("cos_", "") angle_data = local_dict[angle] data = np.cos(angle_data) # Requesting incl or abs_incl: convert to degrees @@ -90,32 +88,32 @@ def complete_incl_data(injs, key, tag): def complete_mass_data(injs, key, tag): - """Extract data related to mass ratio, chirp mass or total mass from raw - injection data""" - - mass1 = injs[tag+'/mass1'] - mass2 = injs[tag+'/mass2'] + """ + Extract data related to mass ratio, chirp mass or total mass from raw + injection data + """ + mass1 = injs[tag + "/mass1"] + mass2 = injs[tag + "/mass2"] - if key == 'mtotal': + if key == "mtotal": data = mass1 + mass2 - elif key == 'mchirp': + elif key == "mchirp": data = pycbc.conversions.mchirp_from_mass1_mass2(mass1, mass2) else: data = mass2 / mass1 - data = np.where(data > 1, 1./data, data) + data = np.where(data > 1, 1.0 / data, data) return data -def complete_spin_data(injs, tag): +def complete_spin_data(injs, tag): """Extract data related to effective precession spin from raw injection data""" - - mass1 = injs[tag+'/mass1'] - mass2 = injs[tag+'/mass2'] - spin1x = injs[tag+'/spin1x'] - spin1y = injs[tag+'/spin1y'] - spin2x = injs[tag+'/spin2x'] - spin2y = injs[tag+'/spin2y'] + mass1 = injs[tag + "/mass1"] + mass2 = injs[tag + "/mass2"] + spin1x = injs[tag + "/spin1x"] + spin1y = injs[tag + "/spin1y"] + spin2x = injs[tag + "/spin2x"] + spin2y = injs[tag + "/spin2y"] data = pycbc.conversions.chi_p(mass1, mass2, spin1x, spin1y, spin2x, spin2y) return data @@ -123,91 +121,112 @@ def complete_spin_data(injs, tag): def complete_sky_error_data(injs, tag): """Extract data related to sky_error from raw injection and trigger data""" - # Missed injections are assigned null values - data = np.full(len(injs[tag+'/mass1']), None) - if tag == 'found': + data = np.full(len(injs[tag + "/mass1"]), None) + if tag == "found": inj = {} - inj['ra'] = injs[tag+'/ra'] - inj['dec'] = injs[tag+'/dec'] + inj["ra"] = injs[tag + "/ra"] + inj["dec"] = injs[tag + "/dec"] trig = {} - trig['ra'] = injs['network/ra'] - trig['dec'] = injs['network/dec'] - data = np.arccos(np.cos(inj['dec'] - trig['dec']) - - np.cos(inj['dec']) * np.cos(trig['dec']) * - (1 - np.cos(inj['ra'] - trig['ra']))) + trig["ra"] = injs["network/ra"] + trig["dec"] = injs["network/dec"] + data = np.arccos( + np.cos(inj["dec"] - trig["dec"]) + - np.cos(inj["dec"]) + * np.cos(trig["dec"]) + * (1 - np.cos(inj["ra"] - trig["ra"])) + ) return data -def complete_opt_snr_data (injs, tag, ifos): - optimal_snr = np.empty((injs[tag + '/tc'].size, len(ifos))) +def complete_opt_snr_data(injs, tag, ifos): + optimal_snr = np.empty((injs[tag + "/tc"].size, len(ifos))) for i, ifo in enumerate(ifos): - optimal_snr[:, i] = injs[tag + '/optimal_snr_' + ifo] - - data = np.sqrt((optimal_snr ** 2).sum(axis=1)) + optimal_snr[:, i] = injs[tag + "/optimal_snr_" + ifo] + + data = np.sqrt((optimal_snr**2).sum(axis=1)) return data + # These are keys in the found-missed-file under 'found' and 'missed' # TODO: also in the found-missed injs file is 'inclination' -easy_keys = ['distance', 'mass1', 'mass2', 'polarization', - 'spin1_a', 'spin1x', 'spin1y', 'spin1z', - 'spin2_a', 'spin2x', 'spin2y', 'spin2z', - 'spin1_azimuthal', 'spin1_polar', - 'spin2_azimuthal', 'spin2_polar', - 'dec', 'ra', 'phi_ref', 'optimal_snr_H1', 'optimal_snr_L1', - 'optimal_snr_V1', 'optimal_snr_K1'] +easy_keys = [ + "distance", + "mass1", + "mass2", + "polarization", + "spin1_a", + "spin1x", + "spin1y", + "spin1z", + "spin2_a", + "spin2x", + "spin2y", + "spin2z", + "spin1_azimuthal", + "spin1_polar", + "spin2_azimuthal", + "spin2_polar", + "dec", + "ra", + "phi_ref", + "optimal_snr_H1", + "optimal_snr_L1", + "optimal_snr_V1", + "optimal_snr_K1", +] def complete_inj_data(injs, keys, tag, ifos=[]): - """Create a dictionary containing the data specified by the - list of keys extracted from an injection file""" - + """ + Create a dictionary containing the data specified by the + list of keys extracted from an injection file + """ data_dict = {} for key in keys: data_dict[key] = np.array([]) try: - if key == 'end_time': - data_dict[key] = injs[tag+'/tc'] + if key == "end_time": + data_dict[key] = injs[tag + "/tc"] # data_dict[key] -= grb_time - elif key in ['mchirp', 'mtotal', 'q']: + elif key in ["mchirp", "mtotal", "q"]: data_dict[key] = complete_mass_data(injs, key, tag) - elif key == 'chi_p': + elif key == "chi_p": data_dict[key] = complete_spin_data(injs, tag) - elif 'incl' in key: + elif "incl" in key: data_dict[key] = complete_incl_data(injs, key, tag) - elif key == 'sky_error': + elif key == "sky_error": data_dict[key] = complete_sky_error_data(injs, tag) - elif key == 'net_opt_snr': + elif key == "net_opt_snr": data_dict[key] = complete_opt_snr_data(injs, tag, ifos) - elif key == 'eff_dist': - eff_dists = np.empty((injs[tag + '/tc'].size, len(ifos))) + elif key == "eff_dist": + eff_dists = np.empty((injs[tag + "/tc"].size, len(ifos))) for i, ifo in enumerate(ifos): eff_dists[:, i] = Detector(ifo).effective_distance( - injs[tag + '/distance'], - injs[tag + '/ra'], - injs[tag + '/dec'], - injs[tag + '/polarization'], - injs[tag + '/tc'], - injs[tag + '/inclination'], + injs[tag + "/distance"], + injs[tag + "/ra"], + injs[tag + "/dec"], + injs[tag + "/polarization"], + injs[tag + "/tc"], + injs[tag + "/inclination"], ) data_dict[key] = 1 / (1 / eff_dists).sum(axis=1) else: - data_dict[key] = injs[tag+'/'+key] + data_dict[key] = injs[tag + "/" + key] except KeyError: # raise NotImplemented(key+' not allowed: returning empty entry') logging.warning("%s/%s not allowed yet", tag, key) # Include information needed to apply vetoes and reweighted SNR cut - if tag == 'found': + if tag == "found": for ifo in ifos: - key = ifo+'/end_time' + key = ifo + "/end_time" data_dict[key] = injs[key] if key in injs else np.array([]) - key = 'network/reweighted_snr' - data_dict['reweighted_snr'] = \ - injs[key] if key in injs else np.array([]) + key = "network/reweighted_snr" + data_dict["reweighted_snr"] = injs[key] if key in injs else np.array([]) logging.info("%d %s injections analysed.", len(data_dict[keys[0]]), tag) @@ -219,47 +238,86 @@ def complete_inj_data(injs, keys, tag, ifos=[]): # ============================================================================= # FIXME: not all of these work -admitted_vars = easy_keys + ['mtotal', 'q', 'mchirp', - 'spin1a', 'spin2a', - 'incl', 'cos_incl', 'abs_incl', 'cos_abs_incl', - 'cosincl', 'absincl', 'cosabsincl', - 'sky_error', 'skyerror', 'end_time', 'endtime', - 'dec', 'ra', 'coaphase', 'coa_phase', - 'eff_site_dist', 'eff_dist', - 'effsitedist', 'effdist', 'chip', 'chi_p', - 'netoptsnr', 'net_opt_snr'] +admitted_vars = easy_keys + [ + "mtotal", + "q", + "mchirp", + "spin1a", + "spin2a", + "incl", + "cos_incl", + "abs_incl", + "cos_abs_incl", + "cosincl", + "absincl", + "cosabsincl", + "sky_error", + "skyerror", + "end_time", + "endtime", + "dec", + "ra", + "coaphase", + "coa_phase", + "eff_site_dist", + "eff_dist", + "effsitedist", + "effdist", + "chip", + "chi_p", + "netoptsnr", + "net_opt_snr", +] admitted_vars = sorted(set(admitted_vars)) parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("--found-missed-file", - help="The hdf injection results file", required=True) -parser.add_argument("--trig-file", - help="The hdf offsource trigger file", required=True) -parser.add_argument("--trigger-time", - help="The GPS time of the external trigger", required=True) -parser.add_argument("-x", "--x-variable", default=None, required=True, - choices=admitted_vars, - help="Quantity to plot on the horizontal axis. " + - "(Underscores may be omitted in specifying this option).") -parser.add_argument("--x-log", action="store_true", - help="Use log horizontal axis") -parser.add_argument("-y", "--y-variable", default=None, required=True, - choices=admitted_vars, - help="Quantity to plot on the vertical axis. " + - "(Underscores may be omitted in specifying this option).") -parser.add_argument("--y-log", action="store_true", - help="Use log vertical axis") -parser.add_argument("--colormap", default='cividis_r', - help="Type of colormap to be used for the plots.") -parser.add_argument("--linear-colormap", action='store_true', - help="Use a linear instead of log colormap.") -parser.add_argument("--far-type", choices=('inclusive', 'exclusive'), - default='inclusive', - help="Type of far to plot for the color. Choices are " - "'inclusive' or 'exclusive'. Default = 'inclusive'") -parser.add_argument("--missed-on-top", action='store_true', - help="Plot missed injections on top of found ones and " - "high FAR on top of low FAR") +parser.add_argument( + "--found-missed-file", help="The hdf injection results file", required=True +) +parser.add_argument("--trig-file", help="The hdf offsource trigger file", required=True) +parser.add_argument( + "--trigger-time", help="The GPS time of the external trigger", required=True +) +parser.add_argument( + "-x", + "--x-variable", + default=None, + required=True, + choices=admitted_vars, + help="Quantity to plot on the horizontal axis. " + "(Underscores may be omitted in specifying this option).", +) +parser.add_argument("--x-log", action="store_true", help="Use log horizontal axis") +parser.add_argument( + "-y", + "--y-variable", + default=None, + required=True, + choices=admitted_vars, + help="Quantity to plot on the vertical axis. " + "(Underscores may be omitted in specifying this option).", +) +parser.add_argument("--y-log", action="store_true", help="Use log vertical axis") +parser.add_argument( + "--colormap", default="cividis_r", help="Type of colormap to be used for the plots." +) +parser.add_argument( + "--linear-colormap", + action="store_true", + help="Use a linear instead of log colormap.", +) +parser.add_argument( + "--far-type", + choices=("inclusive", "exclusive"), + default="inclusive", + help="Type of far to plot for the color. Choices are " + "'inclusive' or 'exclusive'. Default = 'inclusive'", +) +parser.add_argument( + "--missed-on-top", + action="store_true", + help="Plot missed injections on top of found ones and high FAR on top of low FAR", +) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_slide_opts(parser) opts = parser.parse_args() @@ -270,7 +328,7 @@ init_logging(opts.verbose, format="%(asctime)s: %(levelname)s: %(message)s") x_qty = process_var_strings(opts.x_variable) y_qty = process_var_strings(opts.y_variable) -if 'eff_site_dist' in [x_qty, y_qty] and opts.ifo is None: +if "eff_site_dist" in [x_qty, y_qty] and opts.ifo is None: parser.error( "A value for --ifo must be provided for site-specific effective distance" ) @@ -297,27 +355,23 @@ segment_dict = ppu.load_segment_dict(trig_file) # Construct trials removing vetoed times trial_dict, total_trials = ppu.construct_trials( - opts.seg_files, - segment_dict, - ifos, - slide_dict, - opts.veto_file + opts.seg_files, segment_dict, ifos, slide_dict, opts.veto_file ) # Load triggers (apply reweighted SNR cut, not vetoes) -all_off_trigs = ppu.load_data(trig_file, ifos, data_tag='trigs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=opts.slide_id) +all_off_trigs = ppu.load_data( + trig_file, + ifos, + data_tag="trigs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=opts.slide_id, +) # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors -keys = ['network/end_time_gc', 'network/reweighted_snr'] +keys = ["network/end_time_gc", "network/reweighted_snr"] trig_data = ppu.extract_trig_properties( - trial_dict, - all_off_trigs, - slide_dict, - segment_dict, - keys + trial_dict, all_off_trigs, slide_dict, segment_dict, keys ) # Max BestNR values in each trial: these are stored in a dictionary keyed @@ -344,25 +398,23 @@ logging.info("Background bestNR calculated.") # Post-process injections # ======================= # Load injection data without applying the reweighted SNR cut, nor vetoes -inj_data = ppu.load_data(opts.found_missed_file, ifos, data_tag='injs', - slide_id=0) +inj_data = ppu.load_data(opts.found_missed_file, ifos, data_tag="injs", slide_id=0) # Extract the necessary data from the missed injections for the plot -missed_inj = complete_inj_data(inj_data, [x_qty, y_qty], 'missed', ifos=ifos) +missed_inj = complete_inj_data(inj_data, [x_qty, y_qty], "missed", ifos=ifos) # Extract the necessary data from the found injections for the plot -found_inj = complete_inj_data(inj_data, [x_qty, y_qty], 'found', ifos=ifos) +found_inj = complete_inj_data(inj_data, [x_qty, y_qty], "found", ifos=ifos) # Apply reweighted SNR cut if opts.newsnr_threshold: - rw_snr_cut = found_inj['reweighted_snr'] < opts.newsnr_threshold + rw_snr_cut = found_inj["reweighted_snr"] < opts.newsnr_threshold for key in [x_qty, y_qty]: - missed_inj[key] = np.concatenate((missed_inj[key], - found_inj[key][rw_snr_cut])) + missed_inj[key] = np.concatenate((missed_inj[key], found_inj[key][rw_snr_cut])) found_inj[key] = found_inj[key][~rw_snr_cut] for ifo in ifos: - found_inj[ifo+'/end_time'] = found_inj[ifo+'/end_time'][~rw_snr_cut] - found_inj['reweighted_snr'] = found_inj['reweighted_snr'][~rw_snr_cut] + found_inj[ifo + "/end_time"] = found_inj[ifo + "/end_time"][~rw_snr_cut] + found_inj["reweighted_snr"] = found_inj["reweighted_snr"][~rw_snr_cut] msg = f"After applying reweighted SNR cut at {opts.newsnr_threshold}: " msg += f"{len(found_inj[x_qty])} found injections and " msg += f"{len(missed_inj[x_qty])} missed injections" @@ -370,19 +422,17 @@ if opts.newsnr_threshold: # Split in injections found surviving vetoes and injections found but vetoed found_after_vetoes, vetoed, found_idx, _ = ppu.apply_vetoes_to_found_injs( - opts.found_missed_file, - found_inj, - ifos, - veto_file=opts.veto_file + opts.found_missed_file, found_inj, ifos, veto_file=opts.veto_file ) # Extract the detection statistic of injections found after vetoes # Injections found but not louder than background are populated further down -if len(list(found_after_vetoes['reweighted_snr'])) > 0: - found_after_vetoes_stat = \ - found_inj['reweighted_snr'][found_idx] \ - if 'found_after_vetoes/stat' not in found_after_vetoes.keys() else \ - found_after_vetoes['found_after_vetoes/stat'] +if len(list(found_after_vetoes["reweighted_snr"])) > 0: + found_after_vetoes_stat = ( + found_inj["reweighted_snr"][found_idx] + if "found_after_vetoes/stat" not in found_after_vetoes.keys() + else found_after_vetoes["found_after_vetoes/stat"] + ) # Separate triggers into: # 1) Found louder than background @@ -390,16 +440,17 @@ if len(list(found_after_vetoes['reweighted_snr'])) > 0: found_louder = {} for key in found_after_vetoes.keys(): found_louder[key] = found_after_vetoes[key][louder_mask] - found_louder['reweighted_snr'] = found_after_vetoes_stat[louder_mask] + found_louder["reweighted_snr"] = found_after_vetoes_stat[louder_mask] # 2) Found quieter than background: injections found (bestnr > 0) # but not louder than background (non-zero FAP) - quieter_mask = (found_after_vetoes_stat <= max_bkgd_reweighted_snr) \ - & (found_after_vetoes_stat != 0) + quieter_mask = (found_after_vetoes_stat <= max_bkgd_reweighted_snr) & ( + found_after_vetoes_stat != 0 + ) found_quieter = {} for key in found_after_vetoes.keys(): found_quieter[key] = found_after_vetoes[key][quieter_mask] - found_quieter['reweighted_snr'] = found_after_vetoes_stat[quieter_mask] + found_quieter["reweighted_snr"] = found_after_vetoes_stat[quieter_mask] # TODO: ifar still missing """ @@ -409,9 +460,13 @@ if len(list(found_after_vetoes['reweighted_snr'])) > 0: else 'found_after_vetoes/ifar_exc' found_quieter['ifar'] = f[ifar_string][quieter_mask] """ - found_quieter['fap'] = np.array([sum(background > bestnr) for bestnr in - found_quieter['reweighted_snr']], - dtype=float) / total_trials + found_quieter["fap"] = ( + np.array( + [sum(background > bestnr) for bestnr in found_quieter["reweighted_snr"]], + dtype=float, + ) + / total_trials + ) # 3) Missed due to vetoes # TODO: needs function to cherry-pick a subset of inj_data specified by @@ -424,7 +479,7 @@ else: # TODO: ifar still missing if found_quieter[x_qty].size: # Statistics: found on top (found-missed) - FM = np.argsort(found_quieter['fap']) + FM = np.argsort(found_quieter["fap"]) # Statistics: missed on top (missed-found) MF = FM[::-1] @@ -435,34 +490,36 @@ if found_quieter[x_qty].size: # ========== # Take care of axes labels -axis_labels_dict = {'mchirp': "Chirp Mass (solar masses)", - 'mtotal': "Total mass (solar masses)", - 'q': "Mass ratio", - 'distance': "Distance (Mpc)", - 'eff_site_dist': f"{opts.ifo} effective distance (Mpc)", - 'eff_dist': "Inverse sum of inverse effective distances (Mpc)", - 'end_time': "GPS time (s)", - 'sky_error': "Rec. sky error (rad)", - 'coa_phase': "Phase of complex SNR (rad)", - 'dec': "Declination (rad)", - 'ra': "Right ascension (rad)", - 'incl': "Inclination iota (deg)", - 'abs_incl': 'Magnitude of inclination |iota| (deg)', - 'cos_incl': "cos(iota)", - 'cos_abs_incl': "cos(|iota|)", - 'mass1': "Mass of 1st binary component (solar masses)", - 'mass2': "Mass of 2nd binary component (solar masses)", - 'polarization': "Polarization phase (rad)", - 'spin1_a': "Spin on 1st binary component", - 'spin1x': "Spin x-component of 1st binary component", - 'spin1y': "Spin y-component of 1st binary component", - 'spin1z': "Spin z-component of 1st binary component", - 'spin2_a': "Spin on 2nd binary component", - 'spin2x': "Spin x-component of 2nd binary component", - 'spin2y': "Spin y-component of 2nd binary component", - 'spin2z': "Spin z-component of 2nd binary component", - 'chi_p': "Effective precession spin", - 'net_opt_snr': "Network optimal SNR"} +axis_labels_dict = { + "mchirp": "Chirp Mass (solar masses)", + "mtotal": "Total mass (solar masses)", + "q": "Mass ratio", + "distance": "Distance (Mpc)", + "eff_site_dist": f"{opts.ifo} effective distance (Mpc)", + "eff_dist": "Inverse sum of inverse effective distances (Mpc)", + "end_time": "GPS time (s)", + "sky_error": "Rec. sky error (rad)", + "coa_phase": "Phase of complex SNR (rad)", + "dec": "Declination (rad)", + "ra": "Right ascension (rad)", + "incl": "Inclination iota (deg)", + "abs_incl": "Magnitude of inclination |iota| (deg)", + "cos_incl": "cos(iota)", + "cos_abs_incl": "cos(|iota|)", + "mass1": "Mass of 1st binary component (solar masses)", + "mass2": "Mass of 2nd binary component (solar masses)", + "polarization": "Polarization phase (rad)", + "spin1_a": "Spin on 1st binary component", + "spin1x": "Spin x-component of 1st binary component", + "spin1y": "Spin y-component of 1st binary component", + "spin1z": "Spin z-component of 1st binary component", + "spin2_a": "Spin on 2nd binary component", + "spin2x": "Spin x-component of 2nd binary component", + "spin2y": "Spin y-component of 2nd binary component", + "spin2z": "Spin z-component of 2nd binary component", + "chi_p": "Effective precession spin", + "net_opt_snr": "Network optimal SNR", +} fig = plt.figure() ax = fig.gca() @@ -480,10 +537,10 @@ cmap = plt.get_cmap(opts.colormap) fnd_col = cmap(0) fnd_col = np.array([fnd_col]) if opts.missed_on_top: - z_orders = ['found louder', 'found quieter', 'vetoed', 'missed'] + z_orders = ["found louder", "found quieter", "vetoed", "missed"] marker_sizes = [15, 40, 40, 40] else: - z_orders = ['missed', 'vetoed', 'found louder', 'found quieter'] + z_orders = ["missed", "vetoed", "found louder", "found quieter"] marker_sizes = [10, 10, 30, 40] if missed_inj[x_qty].size and missed_inj[y_qty].size: @@ -493,8 +550,8 @@ if missed_inj[x_qty].size and missed_inj[y_qty].size: c="black", marker="x", s=marker_sizes[0], - zorder=z_orders.index('missed') / len(z_orders) + 1, - label="Missed" + zorder=z_orders.index("missed") / len(z_orders) + 1, + label="Missed", ) if vetoed[x_qty].size: ax.scatter( @@ -503,8 +560,8 @@ if vetoed[x_qty].size: c="red", marker="x", s=marker_sizes[1], - zorder=z_orders.index('vetoed') / len(z_orders) + 1, - label="Found but vetoed" + zorder=z_orders.index("vetoed") / len(z_orders) + 1, + label="Found but vetoed", ) if found_louder[x_qty].size: ax.scatter( @@ -513,27 +570,25 @@ if found_louder[x_qty].size: c=fnd_col, marker="+", s=marker_sizes[2], - zorder=z_orders.index('found louder') / len(z_orders) + 1, - label="Found louder than offsource" + zorder=z_orders.index("found louder") / len(z_orders) + 1, + label="Found louder than offsource", ) if found_quieter[x_qty].size: if opts.linear_colormap: fq_norm = matplotlib.colors.Normalize(vmin=0, vmax=1) else: - fq_norm = matplotlib.colors.LogNorm( - vmin=min(found_quieter['fap'][FM]), vmax=1 - ) + fq_norm = matplotlib.colors.LogNorm(vmin=min(found_quieter["fap"][FM]), vmax=1) p = ax.scatter( found_quieter[x_qty][FM], found_quieter[y_qty][FM], - c=found_quieter['fap'][FM], + c=found_quieter["fap"][FM], cmap=cmap, norm=fq_norm, s=marker_sizes[3], edgecolor="w", linewidths=1, - zorder=z_orders.index('found quieter') / len(z_orders) + 1, - label="Found within offsource" + zorder=z_orders.index("found quieter") / len(z_orders) + 1, + label="Found within offsource", ) cb = plt.colorbar(p, label="p-value") @@ -541,48 +596,47 @@ ax.grid() ax.legend(loc=(0, 1), ncols=4, fontsize="xx-small", frameon=False) # Handle axis limits when plotting spins -if "spin" in x_qty and missed_inj['spin1_a'].size: - max_missed_inj = missed_inj['spin1_a'].max() +if "spin" in x_qty and missed_inj["spin1_a"].size: + max_missed_inj = missed_inj["spin1_a"].max() ax.set_xlim([0, np.ceil(10 * max_missed_inj) / 10]) if found_inj[x_qty].size: - ax.set_xlim([0, np.ceil(10 * max(max_missed_inj, - found_inj[x_qty].max())) / 10]) -if "spin" in y_qty and missed_inj['spin2_a'].size: - max_missed_inj = missed_inj['spin2_a'].max() + ax.set_xlim([0, np.ceil(10 * max(max_missed_inj, found_inj[x_qty].max())) / 10]) +if "spin" in y_qty and missed_inj["spin2_a"].size: + max_missed_inj = missed_inj["spin2_a"].max() ax.set_ylim([0, np.ceil(10 * max_missed_inj) / 10]) if found_inj[y_qty].size: - ax.set_ylim([0, np.ceil(10 * max(max_missed_inj, - found_inj[y_qty].max())) / 10]) + ax.set_ylim([0, np.ceil(10 * max(max_missed_inj, found_inj[y_qty].max())) / 10]) # Handle axis limits when plotting chi_p -if x_qty == 'chi_p': +if x_qty == "chi_p": ax.set_xlim(0, 1) -if y_qty == 'chi_p': +if y_qty == "chi_p": ax.set_ylim(0, 1) - + # Handle axis limits when plotting inclination if "incl" in x_qty or "incl" in y_qty: max_inc = np.pi # max_inc = max(np.concatenate((g_found[qty], g_ifar[qty], \ # g_missed2[qty], missed_inj[qty]))) max_inc_deg = np.rad2deg(max_inc) - max_inc_deg = np.ceil(max_inc_deg/10.0)*10 + max_inc_deg = np.ceil(max_inc_deg / 10.0) * 10 max_inc = np.deg2rad(max_inc_deg) if x_qty == "incl": ax.set_xlim(0, max_inc_deg) elif x_qty == "abs_incl": - ax.set_xlim(0, max_inc_deg*0.5) + ax.set_xlim(0, max_inc_deg * 0.5) if y_qty == "incl": ax.set_ylim(0, max_inc_deg) elif y_qty == "abs_incl": - ax.set_ylim(0, max_inc_deg*0.5) + ax.set_ylim(0, max_inc_deg * 0.5) # if "cos_incl" in [x_qty, y_qty]: if "cos_" in [x_qty, y_qty]: # tt = np.arange(0, max_inc_deg + 10, 10) - tt = np.asarray([0, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, - 130, 140, 150, 180]) + tt = np.asarray( + [0, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 180] + ) tks = np.cos(np.deg2rad(tt)) - tk_labs = [f'cos({tk} deg)' for tk in tt] + tk_labs = [f"cos({tk} deg)" for tk in tt] # if x_qty == "cos_incl": if "cos_" in x_qty: plt.xticks(tks, tk_labs, fontsize=10) @@ -616,33 +670,35 @@ if plot_caption is None: # Take care of title plot_title = opts.plot_title if plot_title is None: - title_dict = {'mchirp': "chirp mass", - 'mtotal': "total mass", - 'q': "mass ratio", - 'distance': "distance (Mpc)", - 'eff_dist': "inverse sum of inverse effective distances", - 'eff_site_dist': "site specific effective distance", - 'end_time': "time", - 'coa_phase': "phase of complex SNR", - 'dec': "declination", - 'ra': "right ascension", - 'incl': "inclination", - 'cos_incl': "inclination", - 'abs_incl': "inclination", - 'cos_abs_incl': "inclination", - 'mass1': "mass", - 'mass2': "mass", - 'polarization': "polarization", - 'spin1_a': "spin", - 'spin1x': "spin x-component", - 'spin1y': "spin y-component", - 'spin1z': "spin z-component", - 'spin2_a': "spin", - 'spin2x': "spin x-component", - 'spin2y': "spin y-component", - 'spin2z': "spin z-component", - 'chi_p': "Effective precession spin", - 'net_opt_snr': "Network optimal SNR"} + title_dict = { + "mchirp": "chirp mass", + "mtotal": "total mass", + "q": "mass ratio", + "distance": "distance (Mpc)", + "eff_dist": "inverse sum of inverse effective distances", + "eff_site_dist": "site specific effective distance", + "end_time": "time", + "coa_phase": "phase of complex SNR", + "dec": "declination", + "ra": "right ascension", + "incl": "inclination", + "cos_incl": "inclination", + "abs_incl": "inclination", + "cos_abs_incl": "inclination", + "mass1": "mass", + "mass2": "mass", + "polarization": "polarization", + "spin1_a": "spin", + "spin1x": "spin x-component", + "spin1y": "spin y-component", + "spin1z": "spin z-component", + "spin2_a": "spin", + "spin2x": "spin x-component", + "spin2y": "spin y-component", + "spin2z": "spin z-component", + "chi_p": "Effective precession spin", + "net_opt_snr": "Network optimal SNR", + } if "sky_error" in [x_qty, y_qty]: plot_title = "Sky error of recovered injections" @@ -656,16 +712,16 @@ if plot_title is None: plt.tight_layout() fig_kwds = {} -if outfile.endswith('.png'): - fig_kwds['dpi'] = 200 +if outfile.endswith(".png"): + fig_kwds["dpi"] = 200 pycbc.results.save_fig_with_metadata( fig, outfile, - cmd=' '.join(sys.argv), + cmd=" ".join(sys.argv), title=plot_title, caption=plot_caption, - fig_kwds=fig_kwds + fig_kwds=fig_kwds, ) plt.close() logging.info("Plot complete.") diff --git a/bin/pygrb/pycbc_pygrb_plot_null_stats b/bin/pygrb/pycbc_pygrb_plot_null_stats index 3e902bddcc3..61c8f376faa 100644 --- a/bin/pygrb/pycbc_pygrb_plot_null_stats +++ b/bin/pygrb/pycbc_pygrb_plot_null_stats @@ -24,18 +24,20 @@ Plot null statistic or coincident SNR vs coherent SNR for a PyGRB run. # ============================================================================= # Preamble # ============================================================================= -import sys -import os import logging -import numpy -from matplotlib import rc +import os +import sys + import matplotlib.pyplot as plt +import numpy import pycbc.version +from matplotlib import rc + from pycbc import init_logging -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.results import pygrb_plotting_utils as plu +from pycbc.results import pygrb_postprocessing_utils as ppu -plt.switch_backend('Agg') +plt.switch_backend("Agg") __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -49,7 +51,6 @@ __program__ = "pycbc_pygrb_plot_null_stats" # Function that produces the contrours to be plotted def calculate_contours(opts, new_snrs=None): """Generate the contours to plot""" - # Add the new SNR threshold contour to the list if necessary if new_snrs is None: new_snrs = [5.5, 6, 6.5, 7, 8, 9, 10, 11] @@ -69,7 +70,7 @@ def calculate_contours(opts, new_snrs=None): null_grad_val = opts.null_grad_val for snr in snr_vals: if snr > null_grad_snr: - null_cont.append(null_thresh + (snr-null_grad_snr)*null_grad_val) + null_cont.append(null_thresh + (snr - null_grad_snr) * null_grad_val) else: null_cont.append(null_thresh) null_cont = numpy.asarray(null_cont) @@ -81,16 +82,31 @@ def calculate_contours(opts, new_snrs=None): # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("-t", "--trig-file", action="store", - default=None, required=True, - help="The location of the trigger file") -parser.add_argument("--found-missed-file", - help="The hdf injection results file", required=False) -parser.add_argument("-z", "--zoom-in", default=False, action="store_true", - help="Output file a zoomed in version of the plot.") -parser.add_argument("-y", "--y-variable", default=None, - choices=['coincident', 'null'], # TODO: overwhitened? - help="Quantity to plot on the vertical axis.") +parser.add_argument( + "-t", + "--trig-file", + action="store", + default=None, + required=True, + help="The location of the trigger file", +) +parser.add_argument( + "--found-missed-file", help="The hdf injection results file", required=False +) +parser.add_argument( + "-z", + "--zoom-in", + default=False, + action="store_true", + help="Output file a zoomed in version of the plot.", +) +parser.add_argument( + "-y", + "--y-variable", + default=None, + choices=["coincident", "null"], # TODO: overwhitened? + help="Quantity to plot on the vertical axis.", +) ppu.pygrb_add_null_snr_opts(parser) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_slide_opts(parser) @@ -101,28 +117,29 @@ init_logging(opts.verbose, format="%(asctime)s: %(levelname)s: %(message)s") # Check options trig_file = os.path.abspath(opts.trig_file) -found_missed_file = os.path.abspath(opts.found_missed_file)\ - if opts.found_missed_file else None +found_missed_file = ( + os.path.abspath(opts.found_missed_file) if opts.found_missed_file else None +) zoom_in = opts.zoom_in # Prepare plot title and caption -y_labels = {'null': "Null SNR", - 'coincident': "Coincident SNR"} # TODO: overwhitened +y_labels = {"null": "Null SNR", "coincident": "Coincident SNR"} # TODO: overwhitened if opts.plot_title is None: opts.plot_title = y_labels[opts.y_variable] + " vs Coherent SNR" if opts.plot_caption is None: - opts.plot_caption = ("Blue crosses: background triggers. ") + opts.plot_caption = "Blue crosses: background triggers. " if found_missed_file: - opts.plot_caption = opts.plot_caption +\ - ("Red crosses: injections triggers. ") + opts.plot_caption = opts.plot_caption + ("Red crosses: injections triggers. ") - if opts.y_variable == 'coincident': - opts.plot_caption += ("Green line: coincident SNR = coherent SNR.") + if opts.y_variable == "coincident": + opts.plot_caption += "Green line: coincident SNR = coherent SNR." else: - opts.plot_caption = opts.plot_caption +\ - "Black line: veto line. " +\ - "Magenta line: above this triggers have " +\ - "reduced detection statistic." + opts.plot_caption = ( + opts.plot_caption + + "Black line: veto line. " + + "Magenta line: above this triggers have " + + "reduced detection statistic." + ) logging.info("Imported and ready to go.") @@ -148,65 +165,60 @@ segment_dict = ppu.load_segment_dict(trig_file) # Construct trials removing vetoed times trial_dict, total_trials = ppu.construct_trials( - opts.seg_files, - segment_dict, - ifos, - slide_dict, - opts.veto_file + opts.seg_files, segment_dict, ifos, slide_dict, opts.veto_file ) # Load triggers/injections (apply reweighted SNR cut, not vetoes) -trig_data = ppu.load_data(trig_file, ifos, data_tag='trigs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=opts.slide_id) -inj_data = ppu.load_data(found_missed_file, ifos, data_tag='injs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=0) +trig_data = ppu.load_data( + trig_file, + ifos, + data_tag="trigs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=opts.slide_id, +) +inj_data = ppu.load_data( + found_missed_file, + ifos, + data_tag="injs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=0, +) # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors # Coherent SNR is always used -x_key = 'network/coherent_snr' +x_key = "network/coherent_snr" # The other SNR may vary -y_key = 'network/' + opts.y_variable + '_snr' +y_key = "network/" + opts.y_variable + "_snr" found_trigs_slides = ppu.extract_trig_properties( - trial_dict, - trig_data, - slide_dict, - segment_dict, - [x_key, y_key] + trial_dict, trig_data, slide_dict, segment_dict, [x_key, y_key] ) found_trigs = {} for key in [x_key, y_key]: found_trigs[key] = numpy.concatenate( - [found_trigs_slides[key][slide_id][:] for slide_id in slide_dict] + [found_trigs_slides[key][slide_id][:] for slide_id in slide_dict] ) # Gather injections found surviving vetoes found_injs, *_ = ppu.apply_vetoes_to_found_injs( - found_missed_file, - inj_data, - ifos, - veto_file=opts.veto_file, - keys=[x_key, y_key] + found_missed_file, inj_data, ifos, veto_file=opts.veto_file, keys=[x_key, y_key] ) # Generate plots logging.info("Plotting...") # Contours -cont_colors = ['g-'] +cont_colors = ["g-"] snr_vals = None shade_cont_value = None # Coincident SNR plot case: we want a coinc=coh diagonal line on the plot -if y_key == 'network/coincident_snr': - x_max = plu.axis_max_value(found_trigs[x_key], found_injs[x_key], - found_missed_file) +if y_key == "network/coincident_snr": + x_max = plu.axis_max_value(found_trigs[x_key], found_injs[x_key], found_missed_file) snr_vals = [4, x_max] null_stat_conts = [[4, x_max]] # Overwhitened null stat (null SNR) and null stat cases: newSNR contours else: - cont_colors = ['k-', 'm-'] + cont_colors = ["k-", "m-"] null_cont, snr_vals = calculate_contours(opts, new_snrs=None) null_stat_conts = [null_cont] if zoom_in: @@ -217,15 +229,21 @@ else: # Overwhitened null stat (null SNR), null stat or coincident SNR vs # Coherent SNR plot if not opts.x_lims and zoom_in: - opts.x_lims = '6,30' + opts.x_lims = "6,30" if not opts.y_lims and zoom_in: - opts.y_lims = '0,30' + opts.y_lims = "0,30" # Get rcParams -rc('font', size=14) -plu.pygrb_plotter([found_trigs[x_key], found_trigs[y_key]], - [found_injs[x_key], found_injs[y_key]], - "Coherent SNR", y_labels[opts.y_variable], opts, - snr_vals=snr_vals, conts=null_stat_conts, - shade_cont_value=shade_cont_value, - colors=cont_colors, vert_spike=True, - cmd=' '.join(sys.argv)) +rc("font", size=14) +plu.pygrb_plotter( + [found_trigs[x_key], found_trigs[y_key]], + [found_injs[x_key], found_injs[y_key]], + "Coherent SNR", + y_labels[opts.y_variable], + opts, + snr_vals=snr_vals, + conts=null_stat_conts, + shade_cont_value=shade_cont_value, + colors=cont_colors, + vert_spike=True, + cmd=" ".join(sys.argv), +) diff --git a/bin/pygrb/pycbc_pygrb_plot_skygrid b/bin/pygrb/pycbc_pygrb_plot_skygrid index 2f398fc9794..035d4b1dfd3 100644 --- a/bin/pygrb/pycbc_pygrb_plot_skygrid +++ b/bin/pygrb/pycbc_pygrb_plot_skygrid @@ -22,24 +22,25 @@ # Preamble # ============================================================================= -import sys -import os import logging -import numpy +import os +import sys + import h5py +import numpy +import pycbc.version +from matplotlib import colors, rc from matplotlib import pyplot as plt -import matplotlib.colors as colors -from matplotlib import rc from matplotlib.ticker import MaxNLocator -import pycbc.version + +import pycbc.distributions from pycbc import init_logging -from pycbc.results import save_fig_with_metadata -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.detector import Detector -import pycbc.distributions +from pycbc.results import pygrb_postprocessing_utils as ppu +from pycbc.results import save_fig_with_metadata -plt.switch_backend('Agg') -rc('font', size=14) +plt.switch_backend("Agg") +rc("font", size=14) __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -49,24 +50,35 @@ __program__ = "pycbc_pygrb_plot_skygrid" def define_rows_cols_subplot(nplots): cols = int(numpy.ceil(numpy.sqrt(nplots))) - rows = int(numpy.ceil(nplots/cols)) + rows = int(numpy.ceil(nplots / cols)) return cols, rows + def ra_to_ra_mollweide(ra): - mollweide_ra = numpy.remainder(ra + 2*numpy.pi, 2*numpy.pi) - mollweide_ra[mollweide_ra > numpy.pi] -= 2*numpy.pi + mollweide_ra = numpy.remainder(ra + 2 * numpy.pi, 2 * numpy.pi) + mollweide_ra[mollweide_ra > numpy.pi] -= 2 * numpy.pi return mollweide_ra + # ============================================================================= # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("--sky-grid", required=True, - help="The location of the sky grid file") -parser.add_argument("--num-density-bins", default=1200, type=int, - help="Bins for the input distribution plotting") -parser.add_argument("--num-density-samples", default=2000000, type=int, - help="Samples for plotting the input distribution") +parser.add_argument( + "--sky-grid", required=True, help="The location of the sky grid file" +) +parser.add_argument( + "--num-density-bins", + default=1200, + type=int, + help="Bins for the input distribution plotting", +) +parser.add_argument( + "--num-density-samples", + default=2000000, + type=int, + help="Samples for plotting the input distribution", +) opts = parser.parse_args() init_logging(opts.verbose, format="%(asctime)s:%(levelname)s : %(message)s") @@ -74,7 +86,7 @@ init_logging(opts.verbose, format="%(asctime)s:%(levelname)s : %(message)s") sky_grid = os.path.abspath(opts.sky_grid) outfile = opts.output_file if opts.plot_title is None: - opts.plot_title = 'PyGRB sky grid' + opts.plot_title = "PyGRB sky grid" logging.info("Imported and ready to go.") @@ -84,23 +96,23 @@ for outdir in outdirs: if not os.path.isdir(outdir): os.makedirs(outdir) -#Extract all informations from sky grid +# Extract all informations from sky grid with h5py.File(sky_grid, "r") as f: - ra, dec = f['ra'][:], f['dec'][:] - dist = f.attrs['input_distribution'] - input_dist = eval("pycbc.distributions."+dist) + ra, dec = f["ra"][:], f["dec"][:] + dist = f.attrs["input_distribution"] + input_dist = eval("pycbc.distributions." + dist) samples = input_dist.rvs(opts.num_density_samples) - input_ra, input_dec = samples['ra'], samples['dec'] + input_ra, input_dec = samples["ra"], samples["dec"] ifos = f.attrs["detectors"] detectors = [Detector(d) for d in ifos] - gps_time = f.attrs['ref_gps_time'] + gps_time = f.attrs["ref_gps_time"] xlabel = "Right ascension [deg]" ylabel = "Declination [deg]" uni_points = pycbc.distributions.UniformSky().rvs(1000000) -#Convert ra from [0,2*pi] to [-pi,pi] for the mollweide plot +# Convert ra from [0,2*pi] to [-pi,pi] for the mollweide plot uni_ra = ra_to_ra_mollweide(uni_points["ra"]) uni_dec = uni_points["dec"] @@ -110,76 +122,102 @@ grb_ra = ra_to_ra_mollweide(ra) num_points = len(grb_ra) -#Generation of the input density for the mollweide plot +# Generation of the input density for the mollweide plot input_density, ra_edge, dec_edge = numpy.histogram2d( - in_ra,input_dec, bins=opts.num_density_bins, - range=[[-numpy.pi, numpy.pi], [-numpy.pi/2, numpy.pi/2]] + in_ra, + input_dec, + bins=opts.num_density_bins, + range=[[-numpy.pi, numpy.pi], [-numpy.pi / 2, numpy.pi / 2]], ) -cols , rows = define_rows_cols_subplot(len(detectors)+2) -fig, ax = plt.subplots(nrows=rows, ncols=cols,subplot_kw=dict(projection="mollweide"), figsize=(20,20)) +cols, rows = define_rows_cols_subplot(len(detectors) + 2) +fig, ax = plt.subplots( + nrows=rows, ncols=cols, subplot_kw=dict(projection="mollweide"), figsize=(20, 20) +) -cmap = 'Oranges' -skygrid_color = 'black' +cmap = "Oranges" +skygrid_color = "black" -#Sky grid over input distribution plot +# Sky grid over input distribution plot levels = MaxNLocator(nbins=100).tick_values(0, input_density.max()) -ax[0,0].contourf(ra_edge[:-1],dec_edge[:-1],input_density.T, levels=levels,cmap=cmap) -ax[0,0].plot(grb_ra, dec, 'x', c=skygrid_color) +ax[0, 0].contourf( + ra_edge[:-1], dec_edge[:-1], input_density.T, levels=levels, cmap=cmap +) +ax[0, 0].plot(grb_ra, dec, "x", c=skygrid_color) cb = fig.colorbar( - plt.cm.ScalarMappable(colors.Normalize(vmin=0, vmax=input_density.max()), cmap=cmap), - location="bottom", ax=ax[0,0], label="Probability density [a.u.]" + plt.cm.ScalarMappable( + colors.Normalize(vmin=0, vmax=input_density.max()), cmap=cmap + ), + location="bottom", + ax=ax[0, 0], + label="Probability density [a.u.]", +) +ax[0, 0].set_xlabel(xlabel) +ax[0, 0].set_ylabel(ylabel) +ax[0, 0].set_title(f"Sky grid ({num_points} points) over input distribution") +ax[0, 0].grid(True) + +# Hide the second "plot" +ax[0, 1].axis("Off") + +# Zoomed in plot +ax[0, 1] = plt.subplot(rows, cols, 2, projection="rectilinear") +ax[0, 1].set_xlim(numpy.degrees(grb_ra.min()) - 5, numpy.degrees(grb_ra.max()) + 5) +ax[0, 1].set_ylim(numpy.degrees(dec.min()) - 5, numpy.degrees(dec.max()) + 5) +ax[0, 1].contourf( + numpy.degrees(ra_edge[:-1]), + numpy.degrees(dec_edge[:-1]), + input_density.T, + levels=levels, + cmap=cmap, ) -ax[0,0].set_xlabel(xlabel) -ax[0,0].set_ylabel(ylabel) -ax[0,0].set_title(f"Sky grid ({num_points} points) over input distribution") -ax[0,0].grid(True) - -#Hide the second "plot" -ax[0,1].axis("Off") - -#Zoomed in plot -ax[0,1] = plt.subplot(rows,cols,2, projection="rectilinear") -ax[0,1].set_xlim(numpy.degrees(grb_ra.min())-5, numpy.degrees(grb_ra.max())+5) -ax[0,1].set_ylim(numpy.degrees(dec.min())-5, numpy.degrees(dec.max())+5) -ax[0,1].contourf(numpy.degrees(ra_edge[:-1]),numpy.degrees(dec_edge[:-1]),input_density.T, levels=levels, cmap=cmap) -ax[0,1].plot(numpy.degrees(grb_ra), numpy.degrees(dec), 'x', c=skygrid_color) -ax[0,1].set_xlabel(xlabel) -ax[0,1].set_ylabel(ylabel) -ax[0,1].set_title(f"Sky grid ({num_points} points) over input distribution (zoom)") -ax[0,1].grid(True) +ax[0, 1].plot(numpy.degrees(grb_ra), numpy.degrees(dec), "x", c=skygrid_color) +ax[0, 1].set_xlabel(xlabel) +ax[0, 1].set_ylabel(ylabel) +ax[0, 1].set_title(f"Sky grid ({num_points} points) over input distribution (zoom)") +ax[0, 1].grid(True) if len(detectors) <= 2: - idx_col, idx_row = 0, 1 + idx_col, idx_row = 0, 1 else: idx_col, idx_row = 2, 0 -#Sky grid over Antenna pattern plots +# Sky grid over Antenna pattern plots for det in detectors: if idx_col >= cols: idx_col = 0 idx_row += 1 - ant_pat = det.antenna_pattern(uni_points["ra"], uni_points["dec"], 0, t_gps=gps_time) - quad_ant = numpy.sqrt(ant_pat[0]**2 + ant_pat[1]**2) - logging.info('Plotting %s', os.path.basename(outfile)) - sc = ax[idx_row,idx_col].scatter(uni_ra, uni_dec, c=quad_ant, marker='.', cmap=cmap) - ax[idx_row,idx_col].scatter((grb_ra), (dec),c=skygrid_color, marker='x') - ax[idx_row,idx_col].set_title(f"Sky grid over {det.name} antenna pattern") - ax[idx_row,idx_col].set_xlabel(xlabel) - ax[idx_row,idx_col].set_ylabel(ylabel) - plt.colorbar(sc, ax=ax[idx_row,idx_col], label=r"$\sqrt{F_+^2 + F_\times^2}$", location="bottom") - ax[idx_row,idx_col].grid(True) + ant_pat = det.antenna_pattern( + uni_points["ra"], uni_points["dec"], 0, t_gps=gps_time + ) + quad_ant = numpy.sqrt(ant_pat[0] ** 2 + ant_pat[1] ** 2) + logging.info("Plotting %s", os.path.basename(outfile)) + sc = ax[idx_row, idx_col].scatter( + uni_ra, uni_dec, c=quad_ant, marker=".", cmap=cmap + ) + ax[idx_row, idx_col].scatter((grb_ra), (dec), c=skygrid_color, marker="x") + ax[idx_row, idx_col].set_title(f"Sky grid over {det.name} antenna pattern") + ax[idx_row, idx_col].set_xlabel(xlabel) + ax[idx_row, idx_col].set_ylabel(ylabel) + plt.colorbar( + sc, + ax=ax[idx_row, idx_col], + label=r"$\sqrt{F_+^2 + F_\times^2}$", + location="bottom", + ) + ax[idx_row, idx_col].grid(True) idx_col += 1 -#Hiding axis on which there are no plot -if len(detectors)+2 < rows*cols: - ax[-1,-1].axis("Off") - if len(detectors)+2 < rows*cols - 1: - ax[-1,-2].axis("Off") +# Hiding axis on which there are no plot +if len(detectors) + 2 < rows * cols: + ax[-1, -1].axis("Off") + if len(detectors) + 2 < rows * cols - 1: + ax[-1, -2].axis("Off") # Wrap up -plot_caption = f'First panel: search sky grid points (in black) associated to the input external trigger skymap (color scale). Second panel: zoom on the external trigger sky region. Other panels: sky grid points over the corresponding antenna pattern of the interferometers used for the search. Number of sky grid points: {num_points}' +plot_caption = f"First panel: search sky grid points (in black) associated to the input external trigger skymap (color scale). Second panel: zoom on the external trigger sky region. Other panels: sky grid points over the corresponding antenna pattern of the interferometers used for the search. Number of sky grid points: {num_points}" -save_fig_with_metadata(fig, outfile, cmd=' '.join(sys.argv), - title=opts.plot_title, caption=plot_caption) +save_fig_with_metadata( + fig, outfile, cmd=" ".join(sys.argv), title=opts.plot_title, caption=plot_caption +) plt.close() diff --git a/bin/pygrb/pycbc_pygrb_plot_snr_timeseries b/bin/pygrb/pycbc_pygrb_plot_snr_timeseries index 9a0f710c030..67aa1299e08 100644 --- a/bin/pygrb/pycbc_pygrb_plot_snr_timeseries +++ b/bin/pygrb/pycbc_pygrb_plot_snr_timeseries @@ -24,20 +24,22 @@ Plot single IFO/coherent/reweighted/null SNR timeseries for a PyGRB run. # ============================================================================= # Preamble # ============================================================================= -import sys -import os import logging -import numpy +import os +import sys + import matplotlib.pyplot as plt -from matplotlib import rc +import numpy import pycbc.version +from matplotlib import rc + from pycbc import init_logging from pycbc.events.coherent import reweightedsnr_cut -from pycbc.results import pygrb_postprocessing_utils as ppu from pycbc.results import pygrb_plotting_utils as plu +from pycbc.results import pygrb_postprocessing_utils as ppu -plt.switch_backend('Agg') -rc('font', size=14) +plt.switch_backend("Agg") +rc("font", size=14) __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -51,12 +53,11 @@ __program__ = "pycbc_pygrb_plot_snr_timeseries" # Find start and end times of trigger/injecton data relative to a given time def get_start_end_times(data_time, central_time): """Determine padded start and end times of data relative to central_time""" - start = int(min(data_time)) - central_time end = int(max(data_time)) - central_time duration = end - start - start -= duration*0.05 - end += duration*0.05 + start -= duration * 0.05 + end += duration * 0.05 return start, end @@ -64,8 +65,7 @@ def get_start_end_times(data_time, central_time): # Reset times so that t=0 is corresponds to the given trigger time def reset_times(data_time, trig_time): """Reset times so that t=0 corresponds to the trigger time provided""" - - data_time = [t-trig_time for t in data_time] + data_time = [t - trig_time for t in data_time] return data_time @@ -74,18 +74,36 @@ def reset_times(data_time, trig_time): # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("-t", "--trig-file", action="store", - default=None, required=True, - help="The location of the trigger file") -parser.add_argument("--found-missed-file", - help="The hdf injection results file", required=False) -parser.add_argument("--trigger-time", type=float, default=0, - help="External GPS time. Used to center the plot.") -parser.add_argument("-y", "--y-variable", default=None, - choices=['coherent', 'single', 'reweighted', 'null'], - help="Quantity to plot on the vertical axis.") -parser.add_argument("--onsource", default=False, action="store_true", - help="Include onsource data in the plot (CAUTION!)") +parser.add_argument( + "-t", + "--trig-file", + action="store", + default=None, + required=True, + help="The location of the trigger file", +) +parser.add_argument( + "--found-missed-file", help="The hdf injection results file", required=False +) +parser.add_argument( + "--trigger-time", + type=float, + default=0, + help="External GPS time. Used to center the plot.", +) +parser.add_argument( + "-y", + "--y-variable", + default=None, + choices=["coherent", "single", "reweighted", "null"], + help="Quantity to plot on the vertical axis.", +) +parser.add_argument( + "--onsource", + default=False, + action="store_true", + help="Include onsource data in the plot (CAUTION!)", +) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_slide_opts(parser) opts = parser.parse_args() @@ -95,11 +113,10 @@ init_logging(opts.verbose, format="%(asctime)s: %(levelname)s: %(message)s") # Check options trig_file = os.path.abspath(opts.trig_file) -inj_file = \ - os.path.abspath(opts.found_missed_file) if opts.found_missed_file else None +inj_file = os.path.abspath(opts.found_missed_file) if opts.found_missed_file else None snr_type = opts.y_variable ifo = opts.ifo -if snr_type == 'single' and ifo is None: +if snr_type == "single" and ifo is None: err_msg = "Please specify an interferometer for a single IFO plot" parser.error(err_msg) @@ -126,41 +143,41 @@ trial_dict, total_trials = ppu.construct_trials( ifos, slide_dict, opts.veto_file, - hide_onsource=(not opts.onsource) + hide_onsource=(not opts.onsource), ) # Load trigger and injections data: when plotting reweighted SNR, keep all # points to show the impact of the cut, otherwise remove points with # reweighted SNR below threshold -rw_snr_threshold = None if snr_type == 'reweighted' else opts.newsnr_threshold +rw_snr_threshold = None if snr_type == "reweighted" else opts.newsnr_threshold # When including the onsource, avoid printing to screen via logging the number # of triggers found -data_tag = 'trigs' if opts.onsource else None -trig_data = ppu.load_data(trig_file, ifos, data_tag=data_tag, - rw_snr_threshold=rw_snr_threshold, - slide_id=opts.slide_id) -inj_data = ppu.load_data(inj_file, ifos, data_tag='injs', - rw_snr_threshold=rw_snr_threshold, - slide_id=0) +data_tag = "trigs" if opts.onsource else None +trig_data = ppu.load_data( + trig_file, + ifos, + data_tag=data_tag, + rw_snr_threshold=rw_snr_threshold, + slide_id=opts.slide_id, +) +inj_data = ppu.load_data( + inj_file, ifos, data_tag="injs", rw_snr_threshold=rw_snr_threshold, slide_id=0 +) # Specify HDF file keys for x quantity (time) and y quantity (SNR) -x_key = 'network/end_time_gc' -y_key = 'network/' + opts.y_variable + '_snr' -if opts.y_variable == 'single': - y_key = opts.ifo + '/snr' +x_key = "network/end_time_gc" +y_key = "network/" + opts.y_variable + "_snr" +if opts.y_variable == "single": + y_key = opts.ifo + "/snr" # When looking at a single slide in a detector use the time at the detector # Otherwise use the time at the geocenter to look at all slides at once - if opts.slide_id != 'all': - x_key = opts.ifo + '/end_time' + if opts.slide_id != "all": + x_key = opts.ifo + "/end_time" # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors found_trigs = ppu.extract_trig_properties( - trial_dict, - trig_data, - slide_dict, - segment_dict, - [x_key, y_key] + trial_dict, trig_data, slide_dict, segment_dict, [x_key, y_key] ) # Gather injections found surviving vetoes @@ -169,28 +186,24 @@ found_after_vetoes, *_ = ppu.apply_vetoes_to_found_injs( inj_data, ifos, veto_file=opts.veto_file, - keys=[x_key, y_key] + keys=[x_key, y_key], ) # Obtain times and SNRs -trig_data_time = numpy.concatenate([found_trigs[x_key][slide_id][:] - for slide_id in slide_dict]) -trig_data_snr = numpy.concatenate([found_trigs[y_key][slide_id][:] - for slide_id in slide_dict]) +trig_data_time = numpy.concatenate( + [found_trigs[x_key][slide_id][:] for slide_id in slide_dict] +) +trig_data_snr = numpy.concatenate( + [found_trigs[y_key][slide_id][:] for slide_id in slide_dict] +) inj_data_time = found_after_vetoes[x_key][:] if inj_file else None inj_data_snr = found_after_vetoes[y_key][:] if inj_file else None # Apply reweighted SNR threshold keeping downweighted points -if snr_type == 'reweighted': - trig_data_snr = reweightedsnr_cut( - trig_data_snr, - opts.newsnr_threshold - ) +if snr_type == "reweighted": + trig_data_snr = reweightedsnr_cut(trig_data_snr, opts.newsnr_threshold) if inj_file: - inj_data_snr = reweightedsnr_cut( - inj_data_snr, - opts.newsnr_threshold - ) + inj_data_snr = reweightedsnr_cut(inj_data_snr, opts.newsnr_threshold) # Determine the central time (t=0) for the plot central_time = opts.trigger_time @@ -207,28 +220,34 @@ if inj_file: logging.info("Plotting...") # Prepare the horizontal axis label -x_label = "Geocenter" if 'network' in x_key else opts.ifo +x_label = "Geocenter" if "network" in x_key else opts.ifo x_label += f" time since {central_time:.3f} (s)" # Determine what goes on the vertical axis -y_labels = {'coherent': "Coherent SNR", - 'single': f"{ifo} SNR", - 'null': "Null SNR", - 'reweighted': "Reweighted SNR"} +y_labels = { + "coherent": "Coherent SNR", + "single": f"{ifo} SNR", + "null": "Null SNR", + "reweighted": "Reweighted SNR", +} y_label = y_labels[snr_type] # Determine title and caption if opts.plot_title is None: opts.plot_title = y_label + " vs Time" if opts.plot_caption is None: - opts.plot_caption = ("Blue crosses: background triggers. ") + opts.plot_caption = "Blue crosses: background triggers. " if inj_file: - opts.plot_caption += ("Red crosses: injections triggers.") + opts.plot_caption += "Red crosses: injections triggers." # Single IFO SNR versus time plots if not opts.x_lims: opts.x_lims = f"{start},{end}" -plu.pygrb_plotter([trig_data_time, trig_data_snr], - [inj_data_time, inj_data_snr], - x_label, y_label, - opts, cmd=' '.join(sys.argv)) +plu.pygrb_plotter( + [trig_data_time, trig_data_snr], + [inj_data_time, inj_data_snr], + x_label, + y_label, + opts, + cmd=" ".join(sys.argv), +) diff --git a/bin/pygrb/pycbc_pygrb_plot_stats_distribution b/bin/pygrb/pycbc_pygrb_plot_stats_distribution index 048314fed0f..845274e87d4 100644 --- a/bin/pygrb/pycbc_pygrb_plot_stats_distribution +++ b/bin/pygrb/pycbc_pygrb_plot_stats_distribution @@ -24,18 +24,20 @@ Plot distribution of BestNR/SNR/SNR after cuts of triggers in a PyGRB run. # ============================================================================= # Preamble # ============================================================================= -import os import logging +import os import sys + import matplotlib.pyplot as plt -from matplotlib import rc import numpy as np import pycbc.version +from matplotlib import rc + from pycbc import init_logging -from pycbc.results import save_fig_with_metadata from pycbc.results import pygrb_postprocessing_utils as ppu +from pycbc.results import save_fig_with_metadata -plt.switch_backend('Agg') +plt.switch_backend("Agg") rc("image", cmap="cividis_r") __author__ = "Francesco Pannarale " @@ -48,11 +50,20 @@ __program__ = "pycbc_pygrb_plot_stats_distribution" # Main script starts here # ============================================================================= parser = ppu.pygrb_initialize_plot_parser(description=__doc__) -parser.add_argument("-F", "--trig-file", action="store", required=True, - help="Location of off-source trigger file") -parser.add_argument("-x", "--x-variable", required=True, - choices=["bestnr", "snr", "snraftercuts"], - help="Quantity to plot on the horizontal axis.") +parser.add_argument( + "-F", + "--trig-file", + action="store", + required=True, + help="Location of off-source trigger file", +) +parser.add_argument( + "-x", + "--x-variable", + required=True, + choices=["bestnr", "snr", "snraftercuts"], + help="Quantity to plot on the horizontal axis.", +) ppu.pygrb_add_bestnr_cut_opt(parser) ppu.pygrb_add_slide_opts(parser) opts = parser.parse_args() @@ -87,28 +98,23 @@ segment_dict = ppu.load_segment_dict(trig_file) # Construct trials removing vetoed times trial_dict, total_trials = ppu.construct_trials( - opts.seg_files, - segment_dict, - ifos, - slide_dict, - opts.veto_file + opts.seg_files, segment_dict, ifos, slide_dict, opts.veto_file ) # Load triggers (apply reweighted SNR cut, not vetoes) -all_off_trigs = ppu.load_data(trig_file, ifos, data_tag='trigs', - rw_snr_threshold=opts.newsnr_threshold, - slide_id=opts.slide_id) +all_off_trigs = ppu.load_data( + trig_file, + ifos, + data_tag="trigs", + rw_snr_threshold=opts.newsnr_threshold, + slide_id=opts.slide_id, +) # Extract needed trigger properties and store them as dictionaries # Based on trial_dict: if vetoes were applied, trig_* are the veto survivors -keys = \ - ['network/end_time_gc', 'network/coherent_snr', 'network/reweighted_snr'] +keys = ["network/end_time_gc", "network/coherent_snr", "network/reweighted_snr"] trig_data = ppu.extract_trig_properties( - trial_dict, - all_off_trigs, - slide_dict, - segment_dict, - keys + trial_dict, all_off_trigs, slide_dict, segment_dict, keys ) # Calculate SNR and BestNR values and maxima @@ -124,37 +130,37 @@ for slide_id in slide_dict: for slide_id in slide_dict: for j, trial in enumerate(trial_dict[slide_id]): - trial_cut = (trial[0] <= trig_data[keys[0]][slide_id][:])\ - & (trig_data[keys[0]][slide_id][:] < trial[1]) + trial_cut = (trial[0] <= trig_data[keys[0]][slide_id][:]) & ( + trig_data[keys[0]][slide_id][:] < trial[1] + ) if not trial_cut.any(): continue # Max SNR - time_veto_max_snr[slide_id][j] = \ - max(trig_data[keys[1]][slide_id][trial_cut]) + time_veto_max_snr[slide_id][j] = max(trig_data[keys[1]][slide_id][trial_cut]) # Max BestNR - time_veto_max_bestnr[slide_id][j] = \ - max(trig_data[keys[2]][slide_id][trial_cut]) + time_veto_max_bestnr[slide_id][j] = max(trig_data[keys[2]][slide_id][trial_cut]) # Max SNR for triggers passing the cut on reweighted SNR sbv_cut = trig_data[keys[2]][slide_id][:] >= opts.newsnr_threshold if not (trial_cut & sbv_cut).any(): continue - time_veto_max_snr_aftercuts[slide_id][j] =\ - max(trig_data[keys[1]][slide_id][trial_cut & sbv_cut]) + time_veto_max_snr_aftercuts[slide_id][j] = max( + trig_data[keys[1]][slide_id][trial_cut & sbv_cut] + ) # This is the data that will be plotted full_time_veto_max_snr = ppu.sort_stat(time_veto_max_snr) full_time_veto_max_snr_aftercuts = ppu.sort_stat(time_veto_max_snr_aftercuts) -_, _, full_time_veto_max_bestnr = \ - ppu.max_median_stat(slide_dict, time_veto_max_bestnr, trig_data[keys[2]], - total_trials) +_, _, full_time_veto_max_bestnr = ppu.max_median_stat( + slide_dict, time_veto_max_bestnr, trig_data[keys[2]], total_trials +) # The 0.'s here force the histograms to start at (0, 1) if no trial # returned a no-event (i.e., BestNR = 0) -if full_time_veto_max_bestnr[0] != 0.: - full_time_veto_max_snr = np.concatenate(([0.], full_time_veto_max_snr)) - full_time_veto_max_snr_aftercuts = \ - np.concatenate(([0.], full_time_veto_max_snr_aftercuts)) - full_time_veto_max_bestnr = \ - np.concatenate(([0.], full_time_veto_max_bestnr)) +if full_time_veto_max_bestnr[0] != 0.0: + full_time_veto_max_snr = np.concatenate(([0.0], full_time_veto_max_snr)) + full_time_veto_max_snr_aftercuts = np.concatenate( + ([0.0], full_time_veto_max_snr_aftercuts) + ) + full_time_veto_max_bestnr = np.concatenate(([0.0], full_time_veto_max_bestnr)) logging.info("SNR and bestNR maxima calculated.") @@ -162,12 +168,16 @@ logging.info("SNR and bestNR maxima calculated.") # ========= # Make plot # ========= -x_label_dict = {"bestnr": "BestNR", - "snr": "SNR", - "snraftercuts": "SNR after signal based vetoes"} -data_dict = {"bestnr": full_time_veto_max_bestnr, - "snr": full_time_veto_max_snr, - "snraftercuts": full_time_veto_max_snr_aftercuts} +x_label_dict = { + "bestnr": "BestNR", + "snr": "SNR", + "snraftercuts": "SNR after signal based vetoes", +} +data_dict = { + "bestnr": full_time_veto_max_bestnr, + "snr": full_time_veto_max_snr, + "snraftercuts": full_time_veto_max_snr_aftercuts, +} fig = plt.figure() ax = fig.gca() ax.grid(True) @@ -176,14 +186,14 @@ ax.set_xlabel(x_label_dict[stat]) ax.set_ylabel("False alarm probability") # Some standard plot settings ymax = 1.2 -normalization = 1. +normalization = 1.0 epsilon = 1e-8 num_bins = 50 # Figure out binning min_stat = min(data_dict[stat]) max_stat = max(data_dict[stat]) bins = np.linspace(min_stat, max_stat, num_bins + 1, endpoint=True) -bins_bg = np.append(bins, float('Inf')) +bins_bg = np.append(bins, float("Inf")) dx = bins[1] - bins[0] # Histogram each background instance separately and take stats hist_sum = np.zeros(len(bins), dtype=float) @@ -194,37 +204,37 @@ for instance in data_dict[stat]: x = np.delete(x, -1) y = y[::-1].cumsum()[::-1] hist_sum += y - sq_hist_sum += y*y + sq_hist_sum += y * y # Get statistics N = len(data_dict[stat]) means = hist_sum / N -stds = np.sqrt((sq_hist_sum - hist_sum*means) / (N - 1)) +stds = np.sqrt((sq_hist_sum - hist_sum * means) / (N - 1)) means[means <= epsilon] = epsilon upper = means + stds lower = means - stds lower[lower <= epsilon] = epsilon # Normalize -means = means*normalization -upper = upper*normalization -lower = lower*normalization +means = means * normalization +upper = upper * normalization +lower = lower * normalization # Plot mean -ax.plot(x + dx/2, means, 'r+', markersize=12) +ax.plot(x + dx / 2, means, "r+", markersize=12) # 3 lines for aesthetic reasons: ensure that the width of the # bar on the far right right is the same as all the others -x = np.append(x, x[-1]+dx/2) +x = np.append(x, x[-1] + dx / 2) upper = np.append(upper, upper[-1]) lower = np.append(lower, lower[-1]) # Shade in the 1-standard deviation area -ax.fill_between(x + dx/2, lower, upper, alpha=0.3, facecolor='y', step='mid') +ax.fill_between(x + dx / 2, lower, upper, alpha=0.3, facecolor="y", step="mid") # Wrap up ax.set_xlim((0.9 * min_stat, 1.1 * max_stat)) -ax.set_ylim((0.6/N, ymax)) +ax.set_ylim((0.6 / N, ymax)) plot_title = "Cumulative distribution of background triggers" -plot_caption = \ - f"Background cumulative distribution of the {x_label_dict[stat]}" +plot_caption = f"Background cumulative distribution of the {x_label_dict[stat]}" fig_path = opts.output_file -save_fig_with_metadata(fig, fig_path, cmd=' '.join(sys.argv), - title=plot_title, caption=plot_caption) +save_fig_with_metadata( + fig, fig_path, cmd=" ".join(sys.argv), title=plot_title, caption=plot_caption +) plt.close() logging.info("Plots complete.") diff --git a/bin/pygrb/pycbc_pygrb_results_workflow b/bin/pygrb/pycbc_pygrb_results_workflow index 1fd10276087..14c7a4452ac 100644 --- a/bin/pygrb/pycbc_pygrb_results_workflow +++ b/bin/pygrb/pycbc_pygrb_results_workflow @@ -24,18 +24,18 @@ Workflow generator to run PyGRB offline post-processing. # ============================================================================= # Preamble # ============================================================================= -import sys -import socket -import logging import argparse +import logging import os +import socket +import sys -import pycbc import pycbc.version + +import pycbc import pycbc.workflow as _workflow -from pycbc.results import layout +from pycbc.results import layout, save_fig_with_metadata from pycbc.results.pygrb_postprocessing_utils import extract_ifos -from pycbc.results import save_fig_with_metadata __author__ = "Francesco Pannarale " __version__ = pycbc.version.git_verbose_msg @@ -46,7 +46,8 @@ __program__ = "pycbc_pygrb_results_workflow" # Function that converts a list of file paths to a FileList in which each File # is provided with a label def labels_in_files_metadata(labels, rundir, file_paths): - """Return a FileList based on the list of file paths. Each File in it + """ + Return a FileList based on the list of file paths. Each File in it is provided with its tag from labels: the correctess of the tag is based on the file name. """ @@ -56,8 +57,7 @@ def labels_in_files_metadata(labels, rundir, file_paths): # Sort the file paths according to labels up_labels = [l.upper() for l in labels] sorted_file_paths = [ - list(filter(lambda x: l in x.upper(), file_paths))[0] - for l in up_labels + list(filter(lambda x: l in x.upper(), file_paths))[0] for l in up_labels ] # Convert sorted_file_paths to a FileList where each File has its label @@ -65,9 +65,7 @@ def labels_in_files_metadata(labels, rundir, file_paths): labels_paths = zip(up_labels, sorted_file_paths) out = _workflow.FileList( [ - _workflow.resolve_url_to_file( - os.path.join(rundir, p), attrs={'tags': [l]} - ) + _workflow.resolve_url_to_file(os.path.join(rundir, p), attrs={"tags": [l]}) for (l, p) in labels_paths ] ) @@ -79,7 +77,8 @@ def labels_in_files_metadata(labels, rundir, file_paths): # that ensures its copy is saved in the output directory before returning to # the original working directory. Appropriate meta data is added to the plot. def display_seg_plot(output_dir, segment_dir): - """Return the File of the segments plot (which was already produced during + """ + Return the File of the segments plot (which was already produced during the pre-processing) after adding the appropriate metadata to it. """ from PIL import Image, PngImagePlugin @@ -87,7 +86,7 @@ def display_seg_plot(output_dir, segment_dir): curr_dir = os.getcwd() os.chdir(output_dir) segments_plot_name = ( - 'GRB' + wflow.cp.get('workflow', 'trigger-name') + '_segments.png' + "GRB" + wflow.cp.get("workflow", "trigger-name") + "_segments.png" ) segments_plot = _workflow.resolve_url_to_file( os.path.join(segment_dir, segments_plot_name) @@ -95,18 +94,14 @@ def display_seg_plot(output_dir, segment_dir): segments_plot_path = segments_plot.storage_path im = Image.open(segments_plot_path) meta = PngImagePlugin.PngInfo() - meta.add_text('title', str('Segments and analysis time')) + meta.add_text("title", "Segments and analysis time") meta.add_text( - 'caption', - str( - 'Segments (horizontal bands) available around the trigger time (orange vertical line) and analysis time used (orange box).' - ), + "caption", + "Segments (horizontal bands) available around the trigger time (orange vertical line) and analysis time used (orange box).", ) meta.add_text( - 'cmd', - str( - 'This plot is generated by the make_grb_segments_plot function during the pre-processing stage of the analysis.' - ), + "cmd", + "This plot is generated by the make_grb_segments_plot function during the pre-processing stage of the analysis.", ) im.save(segments_plot_path, "png", pnginfo=meta) os.chdir(curr_dir) @@ -134,7 +129,7 @@ parser.add_argument( "--sky-grid", required=True, action="store", - help="The location of the sky grid file" + help="The location of the sky grid file", ) parser.add_argument( "-i", @@ -169,9 +164,7 @@ _workflow.add_workflow_command_line_group(parser) _workflow.add_workflow_settings_cli(parser, include_subdax_opts=True) args = parser.parse_args() -pycbc.init_logging( - args.verbose, format="%(asctime)s: %(levelname)s: %(message)s" -) +pycbc.init_logging(args.verbose, format="%(asctime)s: %(levelname)s: %(message)s") # Store starting run directory start_rundir = os.getcwd() @@ -187,19 +180,19 @@ os.chdir(args.output_dir) # Setup results directory rdir = layout.SectionNumber( - 'webpage', + "webpage", [ - 'offsource_triggers_vs_time', - 'signal_consistency', - 'injections', - 'loudest_offsource_events', - 'exclusion_distances', - 'open_box', - 'workflow', + "offsource_triggers_vs_time", + "signal_consistency", + "injections", + "loudest_offsource_events", + "exclusion_distances", + "open_box", + "workflow", ], ) _workflow.makedir(rdir.base) -_workflow.makedir(rdir['workflow']) +_workflow.makedir(rdir["workflow"]) # File instances of all input trigger files # Expected structure of args.trig_files: @@ -224,8 +217,8 @@ bank_file = os.path.join(start_rundir, args.bank_file) bank_file = _workflow.resolve_url_to_file(bank_file) # Sanity check and File instances of the input injection files -inj_sets = wflow.cp.get_subsections('injections') -eff_secs = wflow.cp.get_subsections('pygrb_efficiency') +inj_sets = wflow.cp.get_subsections("injections") +eff_secs = wflow.cp.get_subsections("pygrb_efficiency") if not set(inj_sets).issubset(eff_secs): err_msg = "Each [injections] subsection requires at " err_msg += "least one dedicated [pygrb_efficiency] subsection. " @@ -238,7 +231,7 @@ full_inj_files = labels_in_files_metadata(inj_sets, start_rundir, args.full_inj_ # File instance of the veto file veto_file = args.veto_file -if veto_file: +if veto_file: veto_file = os.path.join(start_rundir, args.veto_file) veto_file = _workflow.resolve_url_to_file(veto_file) @@ -256,16 +249,16 @@ seg_files = _workflow.build_segment_filelist(args.segment_dir) # Logfile of this workflow wf_log_file = _workflow.File( wflow.ifos, - 'workflow-log', + "workflow-log", wflow.analysis_time, - extension='.txt', - directory=rdir['workflow'], + extension=".txt", + directory=rdir["workflow"], ) -logfile = logging.FileHandler(filename=wf_log_file.storage_path, mode='w') +logfile = logging.FileHandler(filename=wf_log_file.storage_path, mode="w") logfile.setLevel(logging.INFO) -formatter = logging.Formatter('%(asctime)s: %(levelname)s: %(message)s') +formatter = logging.Formatter("%(asctime)s: %(levelname)s: %(message)s") logfile.setFormatter(formatter) -logging.getLogger('').addHandler(logfile) +logging.getLogger("").addHandler(logfile) logging.info("Created log file %s", wf_log_file.storage_path) # TODO: Pick up inifile, segments plot, GRB time and location, and @@ -282,34 +275,32 @@ _workflow.makedir(out_dir) # # Create the information table about the GRB trigger info_table_node, grb_info_table = _workflow.make_pygrb_info_table( - wflow, 'pygrb_grb_info_table', out_dir + wflow, "pygrb_grb_info_table", out_dir ) html_nodes.append(info_table_node) # Plot the search grid plot_node, skygrid_plot = _workflow.make_pygrb_plot( - wflow, 'pygrb_plot_skygrid', out_dir, - sky_grid_file=sky_grid_file + wflow, "pygrb_plot_skygrid", out_dir, sky_grid_file=sky_grid_file ) plotting_nodes.append(plot_node) # Plot the template bank plot_node, bank_plot = _workflow.make_pygrb_plot( - wflow, 'pycbc_plot_bank_corner', out_dir, bank_file=bank_file + wflow, "pycbc_plot_bank_corner", out_dir, bank_file=bank_file ) plotting_nodes.append(plot_node) # Retrieve the segments plot (produced in the preprocessing stage) seg_plot = display_seg_plot(out_dir, args.segment_dir) -summary_layout = [(grb_info_table[0],), (seg_plot, skygrid_plot[0]), - (bank_plot[0],)] +summary_layout = [(grb_info_table[0],), (seg_plot, skygrid_plot[0]), (bank_plot[0],)] layout.two_column_layout(out_dir, summary_layout) # # Plot SNR timeseries # -out_dir = rdir['offsource_triggers_vs_time'] +out_dir = rdir["offsource_triggers_vs_time"] _workflow.makedir(out_dir) # Retrieve the name of the preferred injection set for these plots -tuning_inj_set = wflow.cp.get('workflow-pygrb_results_workflow', 'tuning-inj-set') +tuning_inj_set = wflow.cp.get("workflow-pygrb_results_workflow", "tuning-inj-set") tuning_inj_set = tuning_inj_set.upper() logging.info("The tuning injections set is %s", tuning_inj_set) @@ -317,24 +308,22 @@ logging.info("The tuning injections set is %s", tuning_inj_set) tuning_inj_file = inj_files.find_output_with_tag( tuning_inj_set, fail_if_not_single_file=True ) -logging.info( - "The tuning injections results file is %s", tuning_inj_file.storage_path -) +logging.info("The tuning injections results file is %s", tuning_inj_file.storage_path) # Loop over timeseries request by the user -timeseries = wflow.cp.get_subsections('pygrb_plot_snr_timeseries') +timeseries = wflow.cp.get_subsections("pygrb_plot_snr_timeseries") out_dirs_dict = { - 'coherent': 'offsource_triggers_vs_time/coh_snr_timeseries', - 'reweighted': 'offsource_triggers_vs_time/reweighted_snr_timeseries', - 'single': 'offsource_triggers_vs_time/single_ifo_snr_timeseries', - 'null': 'offsource_triggers_vs_time/null_snr_timeseries', + "coherent": "offsource_triggers_vs_time/coh_snr_timeseries", + "reweighted": "offsource_triggers_vs_time/reweighted_snr_timeseries", + "single": "offsource_triggers_vs_time/single_ifo_snr_timeseries", + "null": "offsource_triggers_vs_time/null_snr_timeseries", } for snr_type in timeseries: out_dir = rdir[out_dirs_dict[snr_type]] _workflow.makedir(out_dir) files = _workflow.FileList([]) # Only single SNR timeseries requires looping over IFOs - ifos_to_loop = ifos if snr_type == 'single' else [None] + ifos_to_loop = ifos if snr_type == "single" else [None] for ifo in ifos_to_loop: timeseries_plots = _workflow.FileList([]) # Plots without and with injections @@ -345,7 +334,7 @@ for snr_type in timeseries: tags = [snr_type] + inj_tags plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_snr_timeseries', + "pygrb_plot_snr_timeseries", out_dir, trig_file=offsource_file, inj_file=inj_file, @@ -365,17 +354,17 @@ for snr_type in timeseries: # # Signal consistency plots # -out_dir = rdir['signal_consistency'] +out_dir = rdir["signal_consistency"] _workflow.makedir(out_dir) # Chisq veto vs SNR plots -out_dir = rdir['signal_consistency/chi_squared_tests'] +out_dir = rdir["signal_consistency/chi_squared_tests"] _workflow.makedir(out_dir) files = _workflow.FileList([]) # Loop over vetoes request by the user -vetoes = wflow.cp.get_subsections('pygrb_plot_chisq_veto') +vetoes = wflow.cp.get_subsections("pygrb_plot_chisq_veto") for veto in vetoes: # Loop over IFOs only in the case of single detector chi-squares - ifo_loop = [ifos[0]] if veto == 'network' else ifos + ifo_loop = [ifos[0]] if veto == "network" else ifos for ifo in ifo_loop: # Plot with and without injections for inj_file in set([tuning_inj_file, None]): @@ -383,10 +372,10 @@ for veto in vetoes: # tags in assembling the plot name inj_tags = inj_file.tags if inj_file else [] tags = [veto] + inj_tags - ifo_arg = None if veto == 'network' else ifo + ifo_arg = None if veto == "network" else ifo plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_chisq_veto', + "pygrb_plot_chisq_veto", out_dir, trig_file=offsource_file, inj_file=inj_file, @@ -403,17 +392,17 @@ chisq_layout = list(layout.grouper(files, 2)) layout.two_column_layout(out_dir, chisq_layout) # Single detector chi-square plots: zoomed in and zoomed out -out_dir = rdir['signal_consistency/individual_detector_snrs'] +out_dir = rdir["signal_consistency/individual_detector_snrs"] _workflow.makedir(out_dir) files = _workflow.FileList([]) # Single IFO SNR vs Coherent SNR plots: zoomed in and zoomed out # Requires looping over IFOs -if wflow.cp.has_section('pygrb_plot_coh_ifosnr'): +if wflow.cp.has_section("pygrb_plot_coh_ifosnr"): for ifo in ifos: # Plots with and without injections for inj_file in set([tuning_inj_file, None]): sngl_snr_plots = _workflow.FileList([]) - for zoom_tag in [['zoomin'], []]: + for zoom_tag in [["zoomin"], []]: # If the tuning injection file is used for the plot, include # its tags in assembling the plot name inj_tags = inj_file.tags if inj_file else [] @@ -421,7 +410,7 @@ if wflow.cp.has_section('pygrb_plot_coh_ifosnr'): # Single IFO SNR vs Coherent SNR plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_coh_ifosnr', + "pygrb_plot_coh_ifosnr", out_dir, trig_file=offsource_file, inj_file=inj_file, @@ -437,38 +426,38 @@ if wflow.cp.has_section('pygrb_plot_coh_ifosnr'): files.append(sngl_snr_plots) layout.two_column_layout(out_dir, files) else: - msg = 'No pygrb_plot_coh_ifosnr section found in the configuration file. ' - msg += 'No coherent vs single detector SNR plots will be generated.' + msg = "No pygrb_plot_coh_ifosnr section found in the configuration file. " + msg += "No coherent vs single detector SNR plots will be generated." logging.info(msg) # Null SNR/Overwhitened null stat vs Coherent SNR plots -null_snr_out_dir = rdir['signal_consistency/null_snrs'] +null_snr_out_dir = rdir["signal_consistency/null_snrs"] _workflow.makedir(null_snr_out_dir) null_snr_files = _workflow.FileList([]) # Coincident SNR vs Coherent SNR plots -coinc_out_dir = rdir['signal_consistency/coincident_snr'] +coinc_out_dir = rdir["signal_consistency/coincident_snr"] _workflow.makedir(coinc_out_dir) coinc_files = _workflow.FileList([]) # Loop over null statistics requested by the user (including coincident SNR) -nstats = wflow.cp.get_subsections('pygrb_plot_null_stats') +nstats = wflow.cp.get_subsections("pygrb_plot_null_stats") for nstat in nstats: # Plots with and without injections, zoomed in and zoomed out for inj_file in set([tuning_inj_file, None]): - if nstat == 'coincident': + if nstat == "coincident": out_dir = coinc_out_dir files = coinc_files else: out_dir = null_snr_out_dir files = null_snr_files null_stats_plots = _workflow.FileList([]) - for zoom_tag in [['zoomin'], []]: + for zoom_tag in [["zoomin"], []]: # If the tuning injection file is used for the plot, include its # tags in assembling the plot name inj_tags = inj_file.tags if inj_file else [] tags = [nstat] + zoom_tag + inj_tags plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_null_stats', + "pygrb_plot_null_stats", out_dir, trig_file=offsource_file, inj_file=inj_file, @@ -490,7 +479,7 @@ layout.two_column_layout(coinc_out_dir, coinc_files) # # Found/missed injections plots and tables # -out_dir = rdir['injections'] +out_dir = rdir["injections"] _workflow.makedir(out_dir) ## Create corner plot inj_corner_plot_files = _workflow.FileList([]) @@ -502,40 +491,38 @@ for inj_set in inj_sets: # Generate a corner bank plot for each injection set tags = [inj_set] plot_node, output_files = _workflow.make_pygrb_plot( - wflow, - 'pycbc_plot_bank_corner', - out_dir, - inj_file=inj_file, - tags=tags, - ) + wflow, + "pycbc_plot_bank_corner", + out_dir, + inj_file=inj_file, + tags=tags, + ) plotting_nodes.append(plot_node) inj_corner_plot_files.append(output_files[0]) layout.single_layout(out_dir, inj_corner_plot_files) # Loop over injection plots requested by the user -inj_plots = wflow.cp.get_subsections('pygrb_plot_injs_results') +inj_plots = wflow.cp.get_subsections("pygrb_plot_injs_results") # The command above also picks up the injection set names so we remove them # from the set of requested injection plot types inj_plots = [inj_plot for inj_plot in inj_plots if inj_plot not in inj_sets] for inj_set in inj_sets: # Retrieve the injections result File corresponding to the inj_set label - inj_file = inj_files.find_output_with_tag( - inj_set, fail_if_not_single_file=True - ) - out_dir = rdir['injections/' + inj_set] + inj_file = inj_files.find_output_with_tag(inj_set, fail_if_not_single_file=True) + out_dir = rdir["injections/" + inj_set] _workflow.makedir(out_dir) files = _workflow.FileList([]) # Generate plots: loop over inj_plots and found-missed/missed-found for inj_plot in inj_plots: - y_qty, x_qty = inj_plot.split('_') + y_qty, x_qty = inj_plot.split("_") ifos_to_loop = [None] - if y_qty == 'effsitedist': + if y_qty == "effsitedist": ifos_to_loop = ifos for ifo in ifos_to_loop: - for fm_or_mf in ['missed-on-top', 'found-on-top']: + for fm_or_mf in ["missed-on-top", "found-on-top"]: tags = [y_qty, x_qty, fm_or_mf, inj_set] plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_injs_results', + "pygrb_plot_injs_results", out_dir, trig_file=offsource_file, ifo=ifo, @@ -567,38 +554,38 @@ for inj_set in inj_sets: # Follow up of loudest N quiet/missed-found injections: # qf_h5 includes quiet and missed found injections - out_dir = rdir['injections/' + inj_set + 'loudest_quiet_found_followups'] - if not wflow.cp.has_section('workflow-minifollowups'): - msg = 'No [workflow-minifollowups] section found in ' - msg += 'configuration file' + out_dir = rdir["injections/" + inj_set + "loudest_quiet_found_followups"] + if not wflow.cp.has_section("workflow-minifollowups"): + msg = "No [workflow-minifollowups] section found in " + msg += "configuration file" logging.info(msg) else: - logging.info('Entering minifollowups module for injections') + logging.info("Entering minifollowups module for injections") mfu_node = _workflow.setup_pygrb_minifollowups( wflow, qf_h5, offsource_file, - 'daxes', + "daxes", out_dir, seg_files=seg_files, veto_file=veto_file, - tags=inj_file.tags + ['loudest_quiet_found_injs'], + tags=inj_file.tags + ["loudest_quiet_found_injs"], ) mfu_nodes.append(mfu_node) - logging.info('Leaving minifollowups') + logging.info("Leaving minifollowups") # # FAP distributions # -out_dir = rdir['loudest_offsource_events'] +out_dir = rdir["loudest_offsource_events"] _workflow.makedir(out_dir) files = [] # Loop over statistics requested by the user -stats = wflow.cp.get_subsections('pygrb_plot_stats_distribution') +stats = wflow.cp.get_subsections("pygrb_plot_stats_distribution") for stat in stats: plot_node, output_file = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_stats_distribution', + "pygrb_plot_stats_distribution", out_dir, trig_file=offsource_file, inj_file=inj_file, @@ -621,50 +608,48 @@ lofft_layout = list(layout.grouper(files, 2)) + [(lofft_table,)] layout.two_column_layout(out_dir, lofft_layout) # Follow up N loudest offsource triggers (parent-child) -out_dir = rdir['loudest_offsource_events/followups'] -if not wflow.cp.has_section('workflow-minifollowups'): - msg = 'No [workflow-minifollowups] section found in ' - msg += 'configuration file' +out_dir = rdir["loudest_offsource_events/followups"] +if not wflow.cp.has_section("workflow-minifollowups"): + msg = "No [workflow-minifollowups] section found in " + msg += "configuration file" logging.info(msg) else: - logging.info('Entering minifollowups module for offsource') + logging.info("Entering minifollowups module for offsource") mfu_node = _workflow.setup_pygrb_minifollowups( wflow, lofft_h5, offsource_file, - 'daxes', + "daxes", out_dir, seg_files=seg_files, veto_file=veto_file, - tags=['loudest_offsource_events'], + tags=["loudest_offsource_events"], ) mfu_nodes.append(mfu_node) - logging.info('Leaving minifollowups') + logging.info("Leaving minifollowups") # # Exclusion distance and efficiency plots based on offtrials # -out_dir = rdir['exclusion_distances'] +out_dir = rdir["exclusion_distances"] _workflow.makedir(out_dir) # Offtrials and injection sets requested by the user -num_trials = int(wflow.cp.get('trig_combiner', 'num-trials')) -offtrials = [f"offtrial_{i+1}" for i in range(num_trials)] +num_trials = int(wflow.cp.get("trig_combiner", "num-trials")) +offtrials = [f"offtrial_{i + 1}" for i in range(num_trials)] # Sensitivity plots of each injection set -out_dir = rdir['exclusion_distances'] +out_dir = rdir["exclusion_distances"] _workflow.makedir(out_dir) bkgd_plots = [] for inj_set in inj_sets: # Retrieve the injections result File for the inj_set label - inj_file = inj_files.find_output_with_tag( - inj_set, fail_if_not_single_file=True - ) + inj_file = inj_files.find_output_with_tag(inj_set, fail_if_not_single_file=True) tags = [offtrials[0]] + inj_file.tags plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_efficiency', + "pygrb_efficiency", out_dir, trig_file=offsource_file, inj_file=inj_file, @@ -681,18 +666,16 @@ layout.two_column_layout(out_dir, bkgd_plots) # Exclusion distances of each injection set for i, offtrial in enumerate(offtrials): - out_dir = rdir[f'exclusion_distances/{offtrial}'] + out_dir = rdir[f"exclusion_distances/{offtrial}"] _workflow.makedir(out_dir) eff_layout = [] for inj_set in inj_sets: # Retrieve the injections result File for the inj_set label - inj_file = inj_files.find_output_with_tag( - inj_set, fail_if_not_single_file=True - ) + inj_file = inj_files.find_output_with_tag(inj_set, fail_if_not_single_file=True) tags = [offtrial] + inj_file.tags plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_efficiency', + "pygrb_efficiency", out_dir, trig_file=offsource_file, onsource_file=offtrial_files[i], @@ -707,36 +690,34 @@ for i, offtrial in enumerate(offtrials): eff_plot = output_files[0] json_input = output_files[1] # Create information table about exclusion distances - excl_dist_table_node, excl_dist_table = ( - _workflow.make_pygrb_info_table( - wflow, - 'pygrb_exclusion_dist_table', - out_dir, - in_files=json_input, - tags=tags, - ) + excl_dist_table_node, excl_dist_table = _workflow.make_pygrb_info_table( + wflow, + "pygrb_exclusion_dist_table", + out_dir, + in_files=json_input, + tags=tags, ) html_nodes.append(excl_dist_table_node) eff_layout += [(eff_plot, excl_dist_table[0])] layout.two_column_layout(out_dir, eff_layout, offtrial) # Make room for throughput histograms (TODO) -base = rdir['workflow/throughput'] +base = rdir["workflow/throughput"] _workflow.makedir(base) # Save global config file -base = rdir['workflow/configuration'] +base = rdir["workflow/configuration"] _workflow.makedir(base) -ini_file_path = os.path.join(base, args.workflow_name + '.ini') -with open(ini_file_path, 'w') as ini_fh: +ini_file_path = os.path.join(base, args.workflow_name + ".ini") +with open(ini_file_path, "w") as ini_fh: wflow.cp.write(ini_fh) ini_file = _workflow.FileList( [ _workflow.File( wflow.ifos, - '', + "", wflow.analysis_time, - file_url='file://' + ini_file_path, + file_url="file://" + ini_file_path, ) ] ) @@ -746,42 +727,42 @@ layout.single_layout(base, ini_file) _workflow.make_versioning_page( wflow, wflow.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) # Create the final log file log_file_html = _workflow.File( wflow.ifos, - 'WORKFLOW-LOG', + "WORKFLOW-LOG", wflow.analysis_time, - extension='.html', - directory=rdir['workflow'], + extension=".html", + directory=rdir["workflow"], ) # Create a page to contain a dashboard link dashboard_file = _workflow.File( wflow.ifos, - 'DASHBOARD', + "DASHBOARD", wflow.analysis_time, - extension='.html', - directory=rdir['workflow'], + extension=".html", + directory=rdir["workflow"], ) dashboard_str = """

    Pegasus Dashboard Page

    """ kwds = { - 'title': 'Pegasus Dashboard', - 'caption': "Link to Pegasus Dashboard", - 'cmd': "PYCBC_SUBMIT_DAX_ARGV", + "title": "Pegasus Dashboard", + "caption": "Link to Pegasus Dashboard", + "cmd": "PYCBC_SUBMIT_DAX_ARGV", } save_fig_with_metadata(dashboard_str, dashboard_file.storage_path, **kwds) # Create pages for the submission script to write data -_workflow.makedir(rdir['workflow/dax']) -_workflow.makedir(rdir['workflow/input_map']) -_workflow.makedir(rdir['workflow/output_map']) -_workflow.makedir(rdir['workflow/planning']) +_workflow.makedir(rdir["workflow/dax"]) +_workflow.makedir(rdir["workflow/input_map"]) +_workflow.makedir(rdir["workflow/output_map"]) +_workflow.makedir(rdir["workflow/planning"]) # Protect the open box results folder -out_dir = rdir['open_box'] +out_dir = rdir["open_box"] _workflow.makedir(out_dir) os.chmod(out_dir, 0o0700) @@ -794,15 +775,15 @@ html_nodes.append(html_node) # Reweighted and coherent SNR timeseries for the full data stretch timeseries_plots = _workflow.FileList([]) -for snr_type in ['reweighted', 'coherent']: +for snr_type in ["reweighted", "coherent"]: plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_plot_snr_timeseries', + "pygrb_plot_snr_timeseries", out_dir, trig_file=all_times_file, seg_files=seg_files, veto_file=veto_file, - tags=[snr_type, 'alltimes'], + tags=[snr_type, "alltimes"], ) plotting_nodes.append(plot_node) timeseries_plots.extend(output_files) @@ -810,10 +791,10 @@ openbox_layout = [(lont_table,)] + list(layout.grouper(timeseries_plots, 2)) layout.two_column_layout(out_dir, openbox_layout) # Loudest on-source trigger follow up -out_dir = rdir['open_box/loudest_event_followup'] -if not wflow.cp.has_section('workflow-minifollowups'): - msg = 'No [workflow-minifollowups] section found in ' - msg += 'configuration file' +out_dir = rdir["open_box/loudest_event_followup"] +if not wflow.cp.has_section("workflow-minifollowups"): + msg = "No [workflow-minifollowups] section found in " + msg += "configuration file" logging.info(msg) else: logging.info("Entering minifollowups module for loudest onsource") @@ -821,27 +802,25 @@ else: wflow, lont_h5, onsource_file, - 'daxes', + "daxes", out_dir, seg_files=seg_files, veto_file=veto_file, - tags=['loudest_onsource_event'], + tags=["loudest_onsource_event"], ) mfu_nodes.append(mfu_node) logging.info("Leaving onsource minifollowups") # Exclusion distances and efficiency plots based on the on-source -out_dir = rdir['open_box/exclusion_distances'] +out_dir = rdir["open_box/exclusion_distances"] eff_layout = [] for inj_set in inj_sets: # Retrieve the injections result File corresponding to the inj_set label - inj_file = inj_files.find_output_with_tag( - inj_set, fail_if_not_single_file=True - ) - tags = ['onsource'] + inj_file.tags + inj_file = inj_files.find_output_with_tag(inj_set, fail_if_not_single_file=True) + tags = ["onsource"] + inj_file.tags plot_node, output_files = _workflow.make_pygrb_plot( wflow, - 'pygrb_efficiency', + "pygrb_efficiency", out_dir, trig_file=offsource_file, onsource_file=onsource_file, @@ -858,7 +837,7 @@ for inj_set in inj_sets: # Create information table about exclusion distances excl_dist_table_node, excl_dist_table = _workflow.make_pygrb_info_table( wflow, - 'pygrb_exclusion_dist_table', + "pygrb_exclusion_dist_table", out_dir, in_files=json_input, tags=tags, @@ -868,19 +847,17 @@ for inj_set in inj_sets: layout.two_column_layout(out_dir, eff_layout) # Job to gather all html material in the webpage -logging.info( - "Path for make_results_web_page: %s", os.path.join(os.getcwd(), rdir.base) -) +logging.info("Path for make_results_web_page: %s", os.path.join(os.getcwd(), rdir.base)) _workflow.make_results_web_page( wflow, os.path.join(os.getcwd(), rdir.base), - template='red', + template="red", explicit_dependencies=plotting_nodes + html_nodes + mfu_nodes, ) # Close the log and flush to the html file logging.shutdown() -with open(wf_log_file.storage_path, "r") as logfile: +with open(wf_log_file.storage_path) as logfile: logdata = logfile.read() log_str = """

    Workflow generation script created workflow in output directory: %s

    @@ -894,12 +871,12 @@ log_str = """ logdata, ) kwds = { - 'title': 'Workflow Generation Log', - 'caption': f"Log of the workflow script {sys.argv[0]}", - 'cmd': ' '.join(sys.argv), + "title": "Workflow Generation Log", + "caption": f"Log of the workflow script {sys.argv[0]}", + "cmd": " ".join(sys.argv), } save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds) -layout.single_layout(rdir['workflow'], ([dashboard_file, log_file_html])) +layout.single_layout(rdir["workflow"], ([dashboard_file, log_file_html])) # Go back to the run directory and save the workflow os.chdir(start_rundir) diff --git a/bin/workflow_comparisons/offline_search/pycbc_combine_injection_comparisons b/bin/workflow_comparisons/offline_search/pycbc_combine_injection_comparisons index 4624603bfaf..36789749312 100755 --- a/bin/workflow_comparisons/offline_search/pycbc_combine_injection_comparisons +++ b/bin/workflow_comparisons/offline_search/pycbc_combine_injection_comparisons @@ -16,95 +16,96 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import argparse + import numpy as np from h5py import File -from pycbc.pnutils import mass1_mass2_to_mchirp_eta import pycbc - +from pycbc.pnutils import mass1_mass2_to_mchirp_eta # Globals -ratio_keys = ['ifar_exc', 'ifar', 'stat', 'minimum_single_detector_statistic'] -other_keys = ['injection_index', 'template_hash'] +ratio_keys = ["ifar_exc", "ifar", "stat", "minimum_single_detector_statistic"] +other_keys = ["injection_index", "template_hash"] # Functions def update_found_one_only(filep, groupp, fileno, data_dict): - keys = ratio_keys+other_keys + keys = ratio_keys + other_keys for key in keys: data_dict[key].append(groupp[key][:]) - idx = groupp['injection_index'][:] - injs = filep['injections'] - cm_arr = mass1_mass2_to_mchirp_eta(injs['mass1'][:][idx], - injs['mass2'][:][idx])[0] - data_dict['chirp_mass'].append(cm_arr) - data_dict['file_number'].append(fileno * np.ones(len(idx), dtype=np.int32)) - return + idx = groupp["injection_index"][:] + injs = filep["injections"] + cm_arr = mass1_mass2_to_mchirp_eta(injs["mass1"][:][idx], injs["mass2"][:][idx])[0] + data_dict["chirp_mass"].append(cm_arr) + data_dict["file_number"].append(fileno * np.ones(len(idx), dtype=np.int32)) + def update_missed(filep, groupp, fileno, data_dict): - keys = ['injection_index', 'loudest_rank'] + keys = ["injection_index", "loudest_rank"] for key in keys: data_dict[key].append(groupp[key][:]) - idx = groupp['injection_index'][:] - injs = filep['injections'] - cm_arr = mass1_mass2_to_mchirp_eta(injs['mass1'][:][idx], - injs['mass2'][:][idx])[0] - data_dict['chirp_mass'].append(cm_arr) - data_dict['file_number'].append(fileno * np.ones(len(idx), dtype=np.int32)) - return + idx = groupp["injection_index"][:] + injs = filep["injections"] + cm_arr = mass1_mass2_to_mchirp_eta(injs["mass1"][:][idx], injs["mass2"][:][idx])[0] + data_dict["chirp_mass"].append(cm_arr) + data_dict["file_number"].append(fileno * np.ones(len(idx), dtype=np.int32)) + def update_found_both(filep, groupp, fileno, data_dict): for key in ratio_keys: - carr = groupp['comparison'][key][:] - rarr = groupp['reference'][key][:] - darr = carr/rarr - data_dict[key]['comparison'].append(carr) - data_dict[key]['reference'].append(rarr) - data_dict[key]['ratio'].append(darr) - for run in ['comparison', 'reference']: - data_dict['template_hash'][run].append(groupp[run]['template_hash'][:]) - idx = groupp[run]['injection_index'][:] - data_dict['injection_index'][run].append(idx) - injs = filep['injections'] - cm_arr = mass1_mass2_to_mchirp_eta(injs['mass1'][:][idx], - injs['mass2'][:][idx])[0] - data_dict['chirp_mass'][run].append(cm_arr) - data_dict['file_number'].append(fileno * np.ones(len(idx), dtype=np.int32)) - return + carr = groupp["comparison"][key][:] + rarr = groupp["reference"][key][:] + darr = carr / rarr + data_dict[key]["comparison"].append(carr) + data_dict[key]["reference"].append(rarr) + data_dict[key]["ratio"].append(darr) + for run in ["comparison", "reference"]: + data_dict["template_hash"][run].append(groupp[run]["template_hash"][:]) + idx = groupp[run]["injection_index"][:] + data_dict["injection_index"][run].append(idx) + injs = filep["injections"] + cm_arr = mass1_mass2_to_mchirp_eta( + injs["mass1"][:][idx], injs["mass2"][:][idx] + )[0] + data_dict["chirp_mass"][run].append(cm_arr) + data_dict["file_number"].append(fileno * np.ones(len(idx), dtype=np.int32)) + def concat_one_only(data_dict): for key in data_dict.keys(): arr = np.concatenate(data_dict[key]) data_dict[key] = arr - return + def concat_both(data_dict): - for key in ratio_keys+other_keys+['chirp_mass']: + for key in ratio_keys + other_keys + ["chirp_mass"]: d = data_dict[key] for k in d.keys(): arr = np.concatenate(d[k]) d[k] = arr - arr = np.concatenate(data_dict['file_number']) - data_dict['file_number'] = arr + arr = np.concatenate(data_dict["file_number"]) + data_dict["file_number"] = arr + def write_combined_results(outgrp, both_dict, ref_dict, com_dict): - grp = outgrp.create_group('found_in_both') - for key in ratio_keys+other_keys+['chirp_mass']: + grp = outgrp.create_group("found_in_both") + for key in ratio_keys + other_keys + ["chirp_mass"]: g = grp.create_group(key) for k in both_dict[key].keys(): - g.create_dataset(k, data = both_dict[key][k]) - grp.create_dataset('file_number', data = both_dict['file_number']) - for nm, ddict in zip(['found_reference_only', 'found_comparison_only'], - [ref_dict, com_dict]): + g.create_dataset(k, data=both_dict[key][k]) + grp.create_dataset("file_number", data=both_dict["file_number"]) + for nm, ddict in zip( + ["found_reference_only", "found_comparison_only"], [ref_dict, com_dict] + ): grp = outgrp.create_group(nm) for key in ddict.keys(): grp.create_dataset(key, data=ddict[key]) + def write_combined_missed(outgrp, ref_dict, com_dict): - rgrp = outgrp.create_group('missed_only_reference') - cgrp = outgrp.create_group('missed_only_comparison') - for key in ['injection_index', 'loudest_rank', - 'chirp_mass', 'file_number']: + rgrp = outgrp.create_group("missed_only_reference") + cgrp = outgrp.create_group("missed_only_comparison") + for key in ["injection_index", "loudest_rank", "chirp_mass", "file_number"]: rgrp.create_dataset(key, data=ref_dict[key]) cgrp.create_dataset(key, data=com_dict[key]) @@ -138,18 +139,31 @@ found. """ formatter = argparse.RawDescriptionHelpFormatter parser = argparse.ArgumentParser() -parser = argparse.ArgumentParser(formatter_class=formatter, - description=long_description) +parser = argparse.ArgumentParser( + formatter_class=formatter, description=long_description +) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--input-files", nargs='+', required=True, - help="List of comparison files created by running" - " 'pycbc_injection_set_comparison' on several injection" - " sets") -parser.add_argument("--output-file", type=str, required=True, - help="Name of HDF output file in which to store results") -parser.add_argument("--found-type", type=str, required=True, - choices=['found', 'found_after_vetoes'], - help="Which class of found injections to collate") +parser.add_argument( + "--input-files", + nargs="+", + required=True, + help="List of comparison files created by running" + " 'pycbc_injection_set_comparison' on several injection" + " sets", +) +parser.add_argument( + "--output-file", + type=str, + required=True, + help="Name of HDF output file in which to store results", +) +parser.add_argument( + "--found-type", + type=str, + required=True, + choices=["found", "found_after_vetoes"], + help="Which class of found injections to collate", +) args = parser.parse_args() pycbc.init_logging(args.verbose) @@ -157,51 +171,54 @@ pycbc.init_logging(args.verbose) outfp = File(args.output_file, "w") both_dict = { - 'ifar_exc' : {'comparison' : [], 'reference' : [], 'ratio' : []}, - 'ifar' : {'comparison' : [], 'reference' : [], 'ratio' : []}, - 'stat' : {'comparison' : [], 'reference' : [], 'ratio' : []}, - 'minimum_single_detector_statistic' : {'comparison' : [], 'reference' : [], - 'ratio' : []}, - 'template_hash' : {'comparison' : [], 'reference' : []}, - 'injection_index' : {'comparison' : [], 'reference' : []}, - 'chirp_mass' : {'comparison' : [], 'reference' : []}, - 'file_number' : [] - } + "ifar_exc": {"comparison": [], "reference": [], "ratio": []}, + "ifar": {"comparison": [], "reference": [], "ratio": []}, + "stat": {"comparison": [], "reference": [], "ratio": []}, + "minimum_single_detector_statistic": { + "comparison": [], + "reference": [], + "ratio": [], + }, + "template_hash": {"comparison": [], "reference": []}, + "injection_index": {"comparison": [], "reference": []}, + "chirp_mass": {"comparison": [], "reference": []}, + "file_number": [], +} ref_dict = { - 'ifar_exc' : [], - 'ifar' : [], - 'stat' : [], - 'minimum_single_detector_statistic' : [], - 'injection_index' : [], - 'chirp_mass' : [], - 'file_number' : [], - 'template_hash' : [] - } + "ifar_exc": [], + "ifar": [], + "stat": [], + "minimum_single_detector_statistic": [], + "injection_index": [], + "chirp_mass": [], + "file_number": [], + "template_hash": [], +} com_dict = { - 'ifar_exc' : [], - 'ifar' : [], - 'stat' : [], - 'minimum_single_detector_statistic' : [], - 'injection_index' : [], - 'chirp_mass' : [], - 'file_number' : [], - 'template_hash' : [] - } + "ifar_exc": [], + "ifar": [], + "stat": [], + "minimum_single_detector_statistic": [], + "injection_index": [], + "chirp_mass": [], + "file_number": [], + "template_hash": [], +} missed_ref = { - 'injection_index' : [], - 'loudest_rank' : [], - 'file_number': [], - 'chirp_mass' : [] + "injection_index": [], + "loudest_rank": [], + "file_number": [], + "chirp_mass": [], } missed_com = { - 'injection_index' : [], - 'loudest_rank' : [], - 'file_number': [], - 'chirp_mass' : [] + "injection_index": [], + "loudest_rank": [], + "file_number": [], + "chirp_mass": [], } # Loop over the files, parsing each one into the appropriate dictionaries @@ -210,26 +227,31 @@ i = 0 read_first = False same_dict = { - 'detectors' : None, - 'single_detector_statistic' : None, - 'reference_dir' : None, - 'comparison_dir' : None + "detectors": None, + "single_detector_statistic": None, + "reference_dir": None, + "comparison_dir": None, } curr_dict = { - 'detectors' : None, - 'single_detector_statistic' : None, - 'reference_dir' : None, - 'comparison_dir' : None + "detectors": None, + "single_detector_statistic": None, + "reference_dir": None, + "comparison_dir": None, } nfiles = len(args.input_files) -outfp.attrs['nfiles'] = nfiles +outfp.attrs["nfiles"] = nfiles for f in args.input_files: fp = File(f, "r") - curr_dict['detectors'] = [fp.attrs['detector_1'], fp.attrs['detector_2']] - for key in ['single_detector_statistic', 'reference_dir', 'comparison_dir', - 'number_missed', 'ifar_threshold']: + curr_dict["detectors"] = [fp.attrs["detector_1"], fp.attrs["detector_2"]] + for key in [ + "single_detector_statistic", + "reference_dir", + "comparison_dir", + "number_missed", + "ifar_threshold", + ]: curr_dict[key] = fp.attrs[key] # The next if/else block checks that as we read each new # comparison file, we are only combining them if they @@ -239,36 +261,38 @@ for f in args.input_files: # path of the comparison directories do not change is inadequate. if not read_first: read_first = True - for key in curr_dict.keys(): + for key in curr_dict: same_dict[key] = curr_dict[key] else: - for key in curr_dict.keys(): + for key in curr_dict: cval = curr_dict[key] sval = same_dict[key] if isinstance(cval, list): - is_same = (set(cval) == set(sval)) + is_same = set(cval) == set(sval) else: - is_same = (cval == sval) + is_same = cval == sval if not is_same: - raise RuntimeError("Incompatible injection sets: file" - " {0}.attrs[{1}] does not match" - " previous injection sets".format(f, key)) - file_key = "injfile_{0}".format(i) - outfp.attrs[file_key] = fp.attrs['injection_label'] + raise RuntimeError( + "Incompatible injection sets: file" + f" {f}.attrs[{key}] does not match" + " previous injection sets" + ) + file_key = f"injfile_{i}" + outfp.attrs[file_key] = fp.attrs["injection_label"] found_group = fp[args.found_type] - update_found_both(fp, found_group['found_in_both'], i, both_dict) - update_found_one_only(fp, found_group['found_reference_only'], i, ref_dict) - update_found_one_only(fp, found_group['found_comparison_only'], i, com_dict) - if args.found_type == 'found_after_vetoes': - missed_group = fp['missed_after_vetoes'] - update_missed(fp, missed_group['missed_only_reference'], i, missed_ref) - update_missed(fp, missed_group['missed_only_comparison'], i, missed_com) + update_found_both(fp, found_group["found_in_both"], i, both_dict) + update_found_one_only(fp, found_group["found_reference_only"], i, ref_dict) + update_found_one_only(fp, found_group["found_comparison_only"], i, com_dict) + if args.found_type == "found_after_vetoes": + missed_group = fp["missed_after_vetoes"] + update_missed(fp, missed_group["missed_only_reference"], i, missed_ref) + update_missed(fp, missed_group["missed_only_comparison"], i, missed_com) i += 1 fp.close() -outfp.attrs['detector_1'] = same_dict['detectors'][0] -outfp.attrs['detector_2'] = same_dict['detectors'][1] -for key in ['single_detector_statistic', 'reference_dir', 'comparison_dir']: +outfp.attrs["detector_1"] = same_dict["detectors"][0] +outfp.attrs["detector_2"] = same_dict["detectors"][1] +for key in ["single_detector_statistic", "reference_dir", "comparison_dir"]: outfp.attrs[key] = same_dict[key] # Now collapse the lists of arrays in these dictionaries into @@ -276,38 +300,40 @@ for key in ['single_detector_statistic', 'reference_dir', 'comparison_dir']: concat_both(both_dict) concat_one_only(ref_dict) concat_one_only(com_dict) -if args.found_type == 'found_after_vetoes': +if args.found_type == "found_after_vetoes": concat_one_only(missed_ref) concat_one_only(missed_com) -cth = both_dict['template_hash']['comparison'] -rth = both_dict['template_hash']['reference'] +cth = both_dict["template_hash"]["comparison"] +rth = both_dict["template_hash"]["reference"] same = (cth == rth).sum() nboth = len(cth) -same_pct = 100.0*float(same)/float(nboth) -nref = len(ref_dict['template_hash']) -ncom = len(com_dict['template_hash']) -ref_pct = 100.0*float(nref)/float(nboth) -com_pct = 100.0*float(ncom)/float(nboth) +same_pct = 100.0 * float(same) / float(nboth) +nref = len(ref_dict["template_hash"]) +ncom = len(com_dict["template_hash"]) +ref_pct = 100.0 * float(nref) / float(nboth) +com_pct = 100.0 * float(ncom) / float(nboth) # Now we print out some summary information that is useful -print("Reference run directory path: {0}".format(same_dict['reference_dir'])) -print("Comparison run directory path: {0}".format(same_dict['comparison_dir'])) -print("Comparing {0} distinct injection sets between" - " the two searches\n".format(nfiles)) -print("Number of triggers found in both: {0}".format(len(rth))) -print("Number from same template: {0} ({1:.2f} % of" - " found in both)\n".format(same, same_pct)) -print("Number found in reference only: {0} ({1:.2f} % of" - " found in both)".format(nref, ref_pct)) -print("Number found in comparison only: {0} ({1:.2f} % of" - " found in both)".format(ncom, com_pct)) +print("Reference run directory path: {0}".format(same_dict["reference_dir"])) +print("Comparison run directory path: {0}".format(same_dict["comparison_dir"])) +print(f"Comparing {nfiles} distinct injection sets between the two searches\n") +print(f"Number of triggers found in both: {len(rth)}") +print( + f"Number from same template: {same} ({same_pct:.2f} % of found in both)\n" +) +print( + f"Number found in reference only: {nref} ({ref_pct:.2f} % of found in both)" +) +print( + f"Number found in comparison only: {ncom} ({com_pct:.2f} % of found in both)" +) # Write everything and finish outgrp = outfp.create_group(args.found_type) write_combined_results(outgrp, both_dict, ref_dict, com_dict) -if args.found_type == 'found_after_vetoes': - outgrp = outfp.create_group('missed_after_vetoes') +if args.found_type == "found_after_vetoes": + outgrp = outfp.create_group("missed_after_vetoes") write_combined_missed(outgrp, missed_ref, missed_com) outfp.close() diff --git a/bin/workflow_comparisons/offline_search/pycbc_injection_set_comparison b/bin/workflow_comparisons/offline_search/pycbc_injection_set_comparison index 0b21dfe3ca6..f54fb71cdaa 100755 --- a/bin/workflow_comparisons/offline_search/pycbc_injection_set_comparison +++ b/bin/workflow_comparisons/offline_search/pycbc_injection_set_comparison @@ -18,15 +18,16 @@ Detailed comparison of a specified injection set between two PyCBC runs """ +import argparse import logging from glob import glob from os import path -import argparse + import numpy as np from h5py import File -from pycbc.events import ranking import pycbc +from pycbc.events import ranking def parse_injection_path(injname, basedir): @@ -34,103 +35,127 @@ def parse_injection_path(injname, basedir): dirpath = path.expanduser(dirpath) dirpath = path.normpath(dirpath) if not path.exists(dirpath): - raise RuntimeError("Directory {0} does not exist".format(basedir)) + raise RuntimeError(f"Directory {basedir} does not exist") if not path.isdir(dirpath): - raise RuntimeError("Path {0} does not specify a" - " directory".format(basedir)) + raise RuntimeError(f"Path {basedir} does not specify a directory") bankpath = path.join(dirpath, "bank") if not path.exists(bankpath): - raise RuntimeError("There is no bank sub-directory of" - " {0}".format(basedir)) + raise RuntimeError(f"There is no bank sub-directory of {basedir}") if not path.isdir(bankpath): - raise RuntimeError("Path {0} does not specify a" - " directory".format(bankpath)) + raise RuntimeError(f"Path {bankpath} does not specify a directory") # The following is really not that robust - bank_files = glob(bankpath+"/*-BANK2HDF-*.hdf") + bank_files = glob(bankpath + "/*-BANK2HDF-*.hdf") if len(bank_files) != 1: - raise RuntimeError("There is not exactly one complete HDF bank file in" - " the path {0}".format(bankpath)) + raise RuntimeError( + f"There is not exactly one complete HDF bank file in the path {bankpath}" + ) bank_file = bank_files[0] - dirpath = path.join(dirpath, "{0}_INJ_coinc".format(injname)) + dirpath = path.join(dirpath, f"{injname}_INJ_coinc") if not path.exists(dirpath): - raise RuntimeError("There is no sub-directory {0}_INJ_coinc in" - " {1}".format(injname, basedir)) + raise RuntimeError( + f"There is no sub-directory {injname}_INJ_coinc in {basedir}" + ) if not path.isdir(dirpath): - raise RuntimeError("Path {0}/{1}_INJ_coinc is not a" - " directory".format(basedir, injname)) - injfile = glob(dirpath+"/*HDFINJFIND*") + raise RuntimeError( + f"Path {basedir}/{injname}_INJ_coinc is not a directory" + ) + injfile = glob(dirpath + "/*HDFINJFIND*") if len(injfile) == 0: - raise RuntimeError("No found-injections file in directory" - " {0}".format(dirpath)) + raise RuntimeError(f"No found-injections file in directory {dirpath}") if len(injfile) > 1: - raise RuntimeError("More than one found-injections file in directory" - " {0}".format(dirpath)) + raise RuntimeError( + f"More than one found-injections file in directory {dirpath}" + ) injfile = injfile[0] return dirpath, injfile, bank_file + # Below are the parameters we use in hashing---for injections especially, we may # need to watch out for additional meaningful parameters. Note, it does not work # to simply take all of the keys present in the injection data group of an # HDFINJFIND file---some of them can contain tiny numerical differences that # spoil their use in a hash. -hash_inj_params = ['coa_phase', 'distance', 'end_time', 'inclination', - 'latitude', 'longitude', 'mass1', 'mass2', - 'polarization', 'spin1x', 'spin1y', 'spin1z', - 'spin2x', 'spin2y', 'spin2z'] +hash_inj_params = [ + "coa_phase", + "distance", + "end_time", + "inclination", + "latitude", + "longitude", + "mass1", + "mass2", + "polarization", + "spin1x", + "spin1y", + "spin1z", + "spin2x", + "spin2y", + "spin2z", +] -class Injections(object): + +class Injections: def __init__(self, inj_group): self.injgroup = inj_group self.inj_params = self.injgroup.keys() self.ninj = len(self.injgroup[self.inj_params[0]][:]) - self.inj_hashes = np.array([hash(v) for v in - zip(*[self.injgroup[p] for p \ - in hash_inj_params])]) + self.inj_hashes = np.array( + [hash(v) for v in zip(*[self.injgroup[p] for p in hash_inj_params])] + ) if not np.all(self.inj_hashes): raise RuntimeError("Not all injection hashes were finite") def __eq__(self, other): if not isinstance(other, Injections): - raise ValueError("{0} cannot be compared to Injections instance; is" - " not itself an instance".format(other)) + raise ValueError( + f"{other} cannot be compared to Injections instance; is" + " not itself an instance" + ) if self.ninj == other.ninj: if set(self.inj_params) == set(other.inj_params): return (self.inj_hashes == other.inj_hashes).all() - else: - return False - else: return False + return False + -class RunInjectionResults(object): - def __init__(self, injname, basedir, - single_detector_statistic='newsnr', - nmissed=10, ifar_threshold=None): +class RunInjectionResults: + def __init__( + self, + injname, + basedir, + single_detector_statistic="newsnr", + nmissed=10, + ifar_threshold=None, + ): self.injname = injname # First, get the necessary files and load them up dirpath, injfile, bankfile = parse_injection_path(injname, basedir) self.dirpath = dirpath self.nmissed = nmissed self.ifar_threshold = ifar_threshold - logging.info("Reading injection directory {0}".format(self.dirpath)) + logging.info(f"Reading injection directory {self.dirpath}") self.injfile = File(injfile, "r") self.bankfile = File(bankfile, "r") - self.inj_dets = [self.injfile.attrs['detector_1'], - self.injfile.attrs['detector_2']] + self.inj_dets = [ + self.injfile.attrs["detector_1"], + self.injfile.attrs["detector_2"], + ] self.trigger_merge_dict = {} for ifo in self.inj_dets: self.trigger_merge_dict.update({ifo: {}}) - merge_file_list = glob(self.dirpath+ - "/{0}-HDF_TRIGGER_MERGE*".format(ifo)) + merge_file_list = glob(self.dirpath + f"/{ifo}-HDF_TRIGGER_MERGE*") if len(merge_file_list) != 1: - raise RuntimeError("There is not exactly one trigger file for" - " IFO {0} in directory {1}".format( - ifo, self.dirpath)) - self.trigger_merge_dict[ifo].update({'fileptr': - File(merge_file_list[0], "r")}) - dgroup = self.trigger_merge_dict[ifo]['fileptr'][ifo] - self.trigger_merge_dict[ifo].update({'data': dgroup}) + raise RuntimeError( + "There is not exactly one trigger file for" + f" IFO {ifo} in directory {self.dirpath}" + ) + self.trigger_merge_dict[ifo].update( + {"fileptr": File(merge_file_list[0], "r")} + ) + dgroup = self.trigger_merge_dict[ifo]["fileptr"][ifo] + self.trigger_merge_dict[ifo].update({"data": dgroup}) # Next, load the information about the injections themselves. - self.injgroup = self.injfile['injections'] + self.injgroup = self.injfile["injections"] self.injections = Injections(self.injgroup) # Now, read the information from the found injection file. The trigger # merge files are typically quite large; we do not automatically load @@ -141,92 +166,104 @@ class RunInjectionResults(object): for ifo in self.inj_dets: self.found.update({ifo: {}}) self.found_after_vetoes.update({ifo: {}}) - for param in ['fap', 'fap_exc', 'ifar', 'ifar_exc', - 'injection_index', 'stat']: - self.found.update({param: self.injfile['found'][param][:]}) + for param in ["fap", "fap_exc", "ifar", "ifar_exc", "injection_index", "stat"]: + self.found.update({param: self.injfile["found"][param][:]}) self.found_after_vetoes.update( - {param: self.injfile['found_after_vetoes'][param][:]}) - template_hashes = self.bankfile['template_hash'][:] - for found_class in ['found', 'found_after_vetoes']: + {param: self.injfile["found_after_vetoes"][param][:]} + ) + template_hashes = self.bankfile["template_hash"][:] + for found_class in ["found", "found_after_vetoes"]: fdict = getattr(self, found_class) - idx = self.injfile[found_class]['template_id'][:] - fdict.update({'template_hash': template_hashes[idx]}) + idx = self.injfile[found_class]["template_id"][:] + fdict.update({"template_hash": template_hashes[idx]}) # The ordering of detectors is arbitrary in the found # injection file, so we disentangle it and only record # via the IFO name - d1 = self.injfile.attrs['detector_1'] - d2 = self.injfile.attrs['detector_2'] - fdict[d1].update({'trigger_time': - self.injfile[found_class]['time1'][:]}) - fdict[d1].update({'trigger_id': - self.injfile[found_class]['trigger_id1'][:]}) - fdict[d2].update({'trigger_time': - self.injfile[found_class]['time2'][:]}) - fdict[d2].update({'trigger_id': - self.injfile[found_class]['trigger_id2'][:]}) - for param in ['after_vetoes', 'all', 'within_analysis']: - self.missed.update({param: self.injfile['missed'][param][:]}) + d1 = self.injfile.attrs["detector_1"] + d2 = self.injfile.attrs["detector_2"] + fdict[d1].update({"trigger_time": self.injfile[found_class]["time1"][:]}) + fdict[d1].update( + {"trigger_id": self.injfile[found_class]["trigger_id1"][:]} + ) + fdict[d2].update({"trigger_time": self.injfile[found_class]["time2"][:]}) + fdict[d2].update( + {"trigger_id": self.injfile[found_class]["trigger_id2"][:]} + ) + for param in ["after_vetoes", "all", "within_analysis"]: + self.missed.update({param: self.injfile["missed"][param][:]}) # Record what the single detector statistic is self.single_detector_statistic = single_detector_statistic # Next, read information from the two trigger merge files - logging.info("Reading single detector triggers from {0}; finding" - "minimum of {1}".format( - self.dirpath, self.single_detector_statistic)) + logging.info( + f"Reading single detector triggers from {self.dirpath}; findingminimum of {self.single_detector_statistic}" + ) for fdict in [self.found, self.found_after_vetoes]: for ifo in self.inj_dets: - for param in ['chisq', 'chisq_dof', 'snr', 'sg_chisq', - 'sigmasq', 'coa_phase', 'end_time']: - if param in self.trigger_merge_dict[ifo]['data'].keys(): - idx = fdict[ifo]['trigger_id'] - vals = \ - self.trigger_merge_dict[ifo]['data'][param][:][idx] + for param in [ + "chisq", + "chisq_dof", + "snr", + "sg_chisq", + "sigmasq", + "coa_phase", + "end_time", + ]: + if param in self.trigger_merge_dict[ifo]["data"].keys(): + idx = fdict[ifo]["trigger_id"] + vals = self.trigger_merge_dict[ifo]["data"][param][:][idx] # A hack: - if param == 'chisq_dof': + if param == "chisq_dof": vals = vals / 2 + 1 fdict[ifo].update({param: vals}) # Initialize statclass with an empty file list curr_rank = ranking.get_sngls_ranking_from_trigs( - fdict[ifo], - single_detector_statistic + fdict[ifo], single_detector_statistic ) - fdict[ifo].update({single_detector_statistic: - curr_rank}) + fdict[ifo].update({single_detector_statistic: curr_rank}) statname = self.single_detector_statistic - fdict.update({'minimum_single_detector_statistic': - np.minimum(fdict[self.inj_dets[0]][statname], - fdict[self.inj_dets[1]][statname])}) + fdict.update( + { + "minimum_single_detector_statistic": np.minimum( + fdict[self.inj_dets[0]][statname], + fdict[self.inj_dets[1]][statname], + ) + } + ) def close(self): self.injfile.close() self.bankfile.close() for ifo in self.inj_dets: - self.trigger_merge_dict[ifo]['fileptr'].close() + self.trigger_merge_dict[ifo]["fileptr"].close() def __str__(self): - return "RunInjectionResults instance for injection set {0} from" \ - " directory {1}".format(self.injname, self.dirpath) + return ( + f"RunInjectionResults instance for injection set {self.injname} from" + f" directory {self.dirpath}" + ) def __eq__(self, other): if not isinstance(other, RunInjectionResults): - raise ValueError("{0} cannot be compared to RunInjectionResults" - " instance; is not itself an instance".format( - other)) - return ((set(self.inj_dets) == set(other.inj_dets)) and - (self.injections == other.injections)) + raise ValueError( + f"{other} cannot be compared to RunInjectionResults" + " instance; is not itself an instance" + ) + return (set(self.inj_dets) == set(other.inj_dets)) and ( + self.injections == other.injections + ) def write_to_hdf(self, outpath): outfp = File(outpath, "w") - outfp.attrs['injection_set'] = self.injname - outfp.attrs['original_directory'] = self.dirpath - outfp.attrs['single_detector_statistic'] = \ - self.single_detector_statistic - outfp.attrs['detector_1'] = self.inj_dets[0] - outfp.attrs['detector_2'] = self.inj_dets[1] - self.injfile.copy('injections', outfp) - missed_group = outfp.create_group('missed') + outfp.attrs["injection_set"] = self.injname + outfp.attrs["original_directory"] = self.dirpath + outfp.attrs["single_detector_statistic"] = self.single_detector_statistic + outfp.attrs["detector_1"] = self.inj_dets[0] + outfp.attrs["detector_2"] = self.inj_dets[1] + self.injfile.copy("injections", outfp) + missed_group = outfp.create_group("missed") for col in self.missed.keys(): missed_group.create_dataset(col, data=self.missed[col]) - for fname in ['found', 'found_after_vetoes']: + for fname in ["found", "found_after_vetoes"]: fgroup = outfp.create_group(fname) fdict = getattr(self, fname) keys_no_ifos = fdict.keys() @@ -252,33 +289,35 @@ def populate_injection_group(fgroup, fidx, fdict, ifos): for key in fdict[ifo].keys(): igroup.create_dataset(key, data=fdict[ifo][key][fidx]) + def compare_found_injs(reference_run, comparison_run, outfp): ifos = reference_run.inj_dets - for found_class in ['found', 'found_after_vetoes']: - logging.info("Comparing found injections for class" - " '{0}'".format(found_class)) + for found_class in ["found", "found_after_vetoes"]: + logging.info(f"Comparing found injections for class '{found_class}'") fgroup = outfp.create_group(found_class) ref_found = getattr(reference_run, found_class) com_found = getattr(comparison_run, found_class) - ref_injs = ref_found['injection_index'] - com_injs = com_found['injection_index'] + ref_injs = ref_found["injection_index"] + com_injs = com_found["injection_index"] # Note that intersect1d returns sorted, unique elements in both arrays # Because we set 'return_indices' to True, we also get the index into # each array of these common elements - logging.info("Calculating indices of found/missed injections between" - " reference and comparison runs") - both, ref_both_idx, com_both_idx = np.intersect1d(ref_injs, - com_injs, - return_indices=True) + logging.info( + "Calculating indices of found/missed injections between" + " reference and comparison runs" + ) + both, ref_both_idx, com_both_idx = np.intersect1d( + ref_injs, com_injs, return_indices=True + ) # Now we want to find the indices of those injections found in one # run but not the other ref_only_idx = np.isin(ref_injs, com_injs, invert=True) com_only_idx = np.isin(com_injs, ref_injs, invert=True) # We now have all of our indices, so create the hierarchy of data groups - ref_both_group = fgroup.create_group('found_in_both/reference') - com_both_group = fgroup.create_group('found_in_both/comparison') - ref_only_group = fgroup.create_group('found_reference_only') - com_only_group = fgroup.create_group('found_comparison_only') + ref_both_group = fgroup.create_group("found_in_both/reference") + com_both_group = fgroup.create_group("found_in_both/comparison") + ref_only_group = fgroup.create_group("found_reference_only") + com_only_group = fgroup.create_group("found_comparison_only") # Now fill these groups with the various data they should contain logging.info("Writing injections found in both, reference run") populate_injection_group(ref_both_group, ref_both_idx, ref_found, ifos) @@ -289,139 +328,173 @@ def compare_found_injs(reference_run, comparison_run, outfp): logging.info("Writing injections found only in comparison run") populate_injection_group(com_only_group, com_only_idx, com_found, ifos) - return + def compare_missed_injs(reference_run, comparison_run, outfp): logging.info("Comparing missed injections after vetoes") ifos = reference_run.inj_dets - fgroup = outfp.create_group('missed_after_vetoes') - ref_missed = reference_run.missed['after_vetoes'] - com_missed = comparison_run.missed['after_vetoes'] - ref_found_injs = reference_run.found_after_vetoes['injection_index'] - com_found_injs = comparison_run.found_after_vetoes['injection_index'] + fgroup = outfp.create_group("missed_after_vetoes") + ref_missed = reference_run.missed["after_vetoes"] + com_missed = comparison_run.missed["after_vetoes"] + ref_found_injs = reference_run.found_after_vetoes["injection_index"] + com_found_injs = comparison_run.found_after_vetoes["injection_index"] if reference_run.ifar_threshold is not None: - ifars = reference_run.found_after_vetoes['ifar'] + ifars = reference_run.found_after_vetoes["ifar"] idx = ifars < reference_run.ifar_threshold ref_missed = np.append(ref_missed, ref_found_injs[idx]) - ref_above = reference_run.found_after_vetoes['injection_index'][~idx] + ref_above = reference_run.found_after_vetoes["injection_index"][~idx] else: - ref_above = reference_run.found_after_vetoes['injection_index'] + ref_above = reference_run.found_after_vetoes["injection_index"] if comparison_run.ifar_threshold is not None: - ifars = comparison_run.found_after_vetoes['ifar'] + ifars = comparison_run.found_after_vetoes["ifar"] idx = ifars < comparison_run.ifar_threshold com_missed = np.append(com_missed, com_found_injs[idx]) - com_above = comparison_run.found_after_vetoes['injection_index'][~idx] + com_above = comparison_run.found_after_vetoes["injection_index"][~idx] else: - com_above = comparison_run.found_after_vetoes['injection_index'] + com_above = comparison_run.found_after_vetoes["injection_index"] # Now find the indices of the required loudest missed injections # Note that we are assuming that the injection table contains # optimal SNR columns, which may not always be true. - logging.info("Finding {0} loudest missed injections in" - " reference run".format(reference_run.nmissed)) - o1 = reference_run.injgroup['optimal_snr_1'][:][ref_missed] - o2 = reference_run.injgroup['optimal_snr_2'][:][ref_missed] + logging.info( + f"Finding {reference_run.nmissed} loudest missed injections in reference run" + ) + o1 = reference_run.injgroup["optimal_snr_1"][:][ref_missed] + o2 = reference_run.injgroup["optimal_snr_2"][:][ref_missed] ref_dec_snr = np.minimum(o1, o2) ref_sort = ref_dec_snr.argsort() ref_sort = ref_sort[::-1] - ref_missed_n = ref_missed[ref_sort][0:reference_run.nmissed] - logging.info("Finding {0} loudest missed injections in" - " comparison run".format(comparison_run.nmissed)) - o1 = comparison_run.injgroup['optimal_snr_1'][:][com_missed] - o2 = comparison_run.injgroup['optimal_snr_2'][:][com_missed] + ref_missed_n = ref_missed[ref_sort][0 : reference_run.nmissed] + logging.info( + f"Finding {comparison_run.nmissed} loudest missed injections in comparison run" + ) + o1 = comparison_run.injgroup["optimal_snr_1"][:][com_missed] + o2 = comparison_run.injgroup["optimal_snr_2"][:][com_missed] com_dec_snr = np.minimum(o1, o2) com_sort = com_dec_snr.argsort() com_sort = com_sort[::-1] - com_missed_n = com_missed[com_sort][0:comparison_run.nmissed] + com_missed_n = com_missed[com_sort][0 : comparison_run.nmissed] # Now see which missed injections in one run were found in the other logging.info("Finding loud injections missed only in reference run") ref_missed_only_idx = np.isin(ref_missed_n, com_above) ref_missed_only = ref_missed_n[ref_missed_only_idx] ref_missed_only_locs = np.where(ref_missed_only_idx)[0] - com_ii = comparison_run.found_after_vetoes['injection_index'] + com_ii = comparison_run.found_after_vetoes["injection_index"] ref_missed_com_idx = np.isin(com_ii, ref_missed_only) logging.info("Finding loud injections missed only in comparison run") com_missed_only_idx = np.isin(com_missed_n, ref_above) com_missed_only = com_missed_n[com_missed_only_idx] com_missed_only_locs = np.where(com_missed_only_idx)[0] - ref_ii = reference_run.found_after_vetoes['injection_index'] + ref_ii = reference_run.found_after_vetoes["injection_index"] com_missed_ref_idx = np.isin(ref_ii, com_missed_only) # We now have all of our indices, so create the hierarchy of data groups - ref_only_group = outfp.create_group( - "missed_after_vetoes/missed_only_reference") - com_only_group = outfp.create_group( - "missed_after_vetoes/missed_only_comparison") + ref_only_group = outfp.create_group("missed_after_vetoes/missed_only_reference") + com_only_group = outfp.create_group("missed_after_vetoes/missed_only_comparison") # Now fill these groups with the various data they should contain logging.info("Writing loud injections missed only in reference run") ref_only_group.create_dataset("injection_index", data=ref_missed_only) ref_only_group.create_dataset("loudest_rank", data=ref_missed_only_locs) ref_missed_com_group = ref_only_group.create_group("comparison") - populate_injection_group(ref_missed_com_group, ref_missed_com_idx, - comparison_run.found_after_vetoes, ifos) + populate_injection_group( + ref_missed_com_group, + ref_missed_com_idx, + comparison_run.found_after_vetoes, + ifos, + ) logging.info("Writing loud injections missed only in comparison run") com_only_group.create_dataset("injection_index", data=com_missed_only) com_only_group.create_dataset("loudest_rank", data=com_missed_only_locs) com_missed_ref_group = com_only_group.create_group("reference") - populate_injection_group(com_missed_ref_group, com_missed_ref_idx, - reference_run.found_after_vetoes, ifos) + populate_injection_group( + com_missed_ref_group, com_missed_ref_idx, reference_run.found_after_vetoes, ifos + ) - return parser = argparse.ArgumentParser(usage="", description=__doc__) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--injection-label", type=str, required=True, - help="Label of injection set") -parser.add_argument("--reference-dir", type=str, required=True, - help="Directory containing reference run of this injection" - " set. This should contain a 'bank/' sub-directory, as well" - " as a sub-directory '_INJ_coinc'") -parser.add_argument("--comparison-dir", type=str, required=True, - help="Directory containing comparison run of this" - " injection set. This should contain a 'bank/' sub-" - "directory, as well as a sub-directory" - " '_INJ_coinc'") -parser.add_argument("--output-file", type=str, required=True, - help="Name of HDF output file in which to store results") -parser.add_argument("--number-missed", type=int, - default=10, required=True, - help="Number of the loudest missed injections to compare" - " between runs") -parser.add_argument('--ifar-threshold', type=float, default=None, - help="If given, also followup injections with ifar smaller " - "than this threshold.") -parser.add_argument("--single-detector-statistic", type=str, default='newsnr', - choices=ranking.sngls_ranking_function_dict.keys(), - help="Which single-detector statistic to calculate for" - " found injections") +parser.add_argument( + "--injection-label", type=str, required=True, help="Label of injection set" +) +parser.add_argument( + "--reference-dir", + type=str, + required=True, + help="Directory containing reference run of this injection" + " set. This should contain a 'bank/' sub-directory, as well" + " as a sub-directory '_INJ_coinc'", +) +parser.add_argument( + "--comparison-dir", + type=str, + required=True, + help="Directory containing comparison run of this" + " injection set. This should contain a 'bank/' sub-" + "directory, as well as a sub-directory" + " '_INJ_coinc'", +) +parser.add_argument( + "--output-file", + type=str, + required=True, + help="Name of HDF output file in which to store results", +) +parser.add_argument( + "--number-missed", + type=int, + default=10, + required=True, + help="Number of the loudest missed injections to compare between runs", +) +parser.add_argument( + "--ifar-threshold", + type=float, + default=None, + help="If given, also followup injections with ifar smaller than this threshold.", +) +parser.add_argument( + "--single-detector-statistic", + type=str, + default="newsnr", + choices=ranking.sngls_ranking_function_dict.keys(), + help="Which single-detector statistic to calculate for found injections", +) args = parser.parse_args() pycbc.init_logging(args.verbose) -ref_found_injs = RunInjectionResults(args.injection_label, args.reference_dir, - args.single_detector_statistic, - args.number_missed, args.ifar_threshold) -com_found_injs = RunInjectionResults(args.injection_label, args.comparison_dir, - args.single_detector_statistic, - args.number_missed, args.ifar_threshold) +ref_found_injs = RunInjectionResults( + args.injection_label, + args.reference_dir, + args.single_detector_statistic, + args.number_missed, + args.ifar_threshold, +) +com_found_injs = RunInjectionResults( + args.injection_label, + args.comparison_dir, + args.single_detector_statistic, + args.number_missed, + args.ifar_threshold, +) -logging.info("Comparing injection parameters between reference and comparison" - " runs") -same_inj = (ref_found_injs == com_found_injs) +logging.info("Comparing injection parameters between reference and comparison runs") +same_inj = ref_found_injs == com_found_injs if not same_inj: - raise RuntimeError("Reference and comparison runs did not perform the same" - " injections") + raise RuntimeError( + "Reference and comparison runs did not perform the same injections" + ) else: outfp = File(args.output_file, "w") - ref_found_injs.injfile.copy('injections', outfp) - -outfp.attrs['injection_label'] = args.injection_label -outfp.attrs['reference_dir'] = args.reference_dir -outfp.attrs['comparison_dir'] = args.comparison_dir -outfp.attrs['single_detector_statistic'] = args.single_detector_statistic -outfp.attrs['detector_1'] = ref_found_injs.inj_dets[0] -outfp.attrs['detector_2'] = ref_found_injs.inj_dets[1] -outfp.attrs['number_missed'] = args.number_missed -outfp.attrs['ifar_threshold'] = args.ifar_threshold + ref_found_injs.injfile.copy("injections", outfp) + +outfp.attrs["injection_label"] = args.injection_label +outfp.attrs["reference_dir"] = args.reference_dir +outfp.attrs["comparison_dir"] = args.comparison_dir +outfp.attrs["single_detector_statistic"] = args.single_detector_statistic +outfp.attrs["detector_1"] = ref_found_injs.inj_dets[0] +outfp.attrs["detector_2"] = ref_found_injs.inj_dets[1] +outfp.attrs["number_missed"] = args.number_missed +outfp.attrs["ifar_threshold"] = args.ifar_threshold compare_found_injs(ref_found_injs, com_found_injs, outfp) compare_missed_injs(ref_found_injs, com_found_injs, outfp) diff --git a/bin/workflow_comparisons/offline_search/pycbc_plot_injections_found_both_workflows b/bin/workflow_comparisons/offline_search/pycbc_plot_injections_found_both_workflows index de9f2817691..dd4c264c924 100755 --- a/bin/workflow_comparisons/offline_search/pycbc_plot_injections_found_both_workflows +++ b/bin/workflow_comparisons/offline_search/pycbc_plot_injections_found_both_workflows @@ -20,30 +20,46 @@ between two comparable runs """ import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") +import argparse + import matplotlib.pyplot as plt import numpy as np -import argparse from pycbc import add_common_pycbc_options, init_logging from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--combined-comparison-file', required=True, - help="HDF file holding output from" - " 'pycbc_combine_injection_comparisons'") -parser.add_argument('--outfile', type=str, required=True, - help='Output file to save to') -parser.add_argument('--plot-title', type=str, required=True, - help='(Possibly) quoted string to be title of plot') -parser.add_argument('--found-category', type=str, required=True, - choices=['found', 'found_after_vetoes'], - help='Which class of found injections to compare') -parser.add_argument('--nbins', type=int, default=10, - help='Number of bins to use for template duration (x-axis)') -parser.add_argument('--log-y', action='store_true', default=False, - help='Use logarithmic y-axis') +parser.add_argument( + "--combined-comparison-file", + required=True, + help="HDF file holding output from 'pycbc_combine_injection_comparisons'", +) +parser.add_argument("--outfile", type=str, required=True, help="Output file to save to") +parser.add_argument( + "--plot-title", + type=str, + required=True, + help="(Possibly) quoted string to be title of plot", +) +parser.add_argument( + "--found-category", + type=str, + required=True, + choices=["found", "found_after_vetoes"], + help="Which class of found injections to compare", +) +parser.add_argument( + "--nbins", + type=int, + default=10, + help="Number of bins to use for template duration (x-axis)", +) +parser.add_argument( + "--log-y", action="store_true", default=False, help="Use logarithmic y-axis" +) args = parser.parse_args() init_logging(args.verbose) @@ -51,8 +67,8 @@ init_logging(args.verbose) # Load in the two datasets f = HFile(args.combined_comparison_file) -ifar_ratio = f[args.found_category]['found_in_both']['ifar']['ratio'][:] -stat_ratio = f[args.found_category]['found_in_both']['stat']['ratio'][:] +ifar_ratio = f[args.found_category]["found_in_both"]["ifar"]["ratio"][:] +stat_ratio = f[args.found_category]["found_in_both"]["stat"]["ratio"][:] nbins = args.nbins @@ -62,17 +78,17 @@ stitle = fig.suptitle(args.plot_title) _, bins = np.histogram(np.log10(ifar_ratio), bins=nbins) ax_ifar.hist(ifar_ratio, bins=10.0**bins) -ax_ifar.set_xscale('log') -ax_ifar.set(xlabel='IFAR Ratio') +ax_ifar.set_xscale("log") +ax_ifar.set(xlabel="IFAR Ratio") ax_stat.hist(stat_ratio, bins=nbins) -ax_stat.set(xlabel='Ranking Statistic Ratio') +ax_stat.set(xlabel="Ranking Statistic Ratio") -ax_ifar.set(ylabel='Count') +ax_ifar.set(ylabel="Count") for ax in [ax_ifar, ax_stat]: if args.log_y: - ax.set_yscale('log') + ax.set_yscale("log") fig.savefig(args.outfile) plt.close() diff --git a/bin/workflow_comparisons/offline_search/pycbc_plot_injections_missed_one_workflow b/bin/workflow_comparisons/offline_search/pycbc_plot_injections_missed_one_workflow index 84ffc4ccb0a..c6b72edba55 100755 --- a/bin/workflow_comparisons/offline_search/pycbc_plot_injections_missed_one_workflow +++ b/bin/workflow_comparisons/offline_search/pycbc_plot_injections_missed_one_workflow @@ -14,38 +14,59 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Plot histograms of IFAR and ranking statistic of injections missed +""" +Plot histograms of IFAR and ranking statistic of injections missed in only one of two comparable runs """ import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") +import argparse + import matplotlib.pyplot as plt import numpy as np -import argparse from pycbc import add_common_pycbc_options, init_logging from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--combined-comparison-file', required=True, - help="HDF file holding output of" - " 'pycbc_combine_injection_comparisons'") -parser.add_argument('--outfile', type=str, required=True, - help='Output file to save to') -parser.add_argument('--plot-title', type=str, required=True, - help='(Possibly) quoted string to be title of plot') -parser.add_argument('--found-category', type=str, required=True, - choices=['found', 'found_after_vetoes'], - help='Which class of found injections to plot') -parser.add_argument('--missed-run', type=str, required=True, - choices=['reference', 'comparison'], - help='Which run missed the injections to plot') -parser.add_argument('--nbins', type=int, default=10, - help='Number of bins to use for template duration (x-axis)') -parser.add_argument('--log-y', action='store_true', default=False, - help='Use logarithmic y-axis') +parser.add_argument( + "--combined-comparison-file", + required=True, + help="HDF file holding output of 'pycbc_combine_injection_comparisons'", +) +parser.add_argument("--outfile", type=str, required=True, help="Output file to save to") +parser.add_argument( + "--plot-title", + type=str, + required=True, + help="(Possibly) quoted string to be title of plot", +) +parser.add_argument( + "--found-category", + type=str, + required=True, + choices=["found", "found_after_vetoes"], + help="Which class of found injections to plot", +) +parser.add_argument( + "--missed-run", + type=str, + required=True, + choices=["reference", "comparison"], + help="Which run missed the injections to plot", +) +parser.add_argument( + "--nbins", + type=int, + default=10, + help="Number of bins to use for template duration (x-axis)", +) +parser.add_argument( + "--log-y", action="store_true", default=False, help="Use logarithmic y-axis" +) args = parser.parse_args() init_logging(args.verbose) @@ -53,11 +74,13 @@ init_logging(args.verbose) # Load in the two datasets f = HFile(args.combined_comparison_file) -conversion_dict = { 'reference' : 'found_comparison_only', - 'comparison' : 'found_reference_only'} +conversion_dict = { + "reference": "found_comparison_only", + "comparison": "found_reference_only", +} -ifar = f[args.found_category][conversion_dict[args.missed_run]]['ifar'][:] -stat = f[args.found_category][conversion_dict[args.missed_run]]['stat'][:] +ifar = f[args.found_category][conversion_dict[args.missed_run]]["ifar"][:] +stat = f[args.found_category][conversion_dict[args.missed_run]]["stat"][:] nbins = args.nbins @@ -67,17 +90,17 @@ stitle = fig.suptitle(args.plot_title) _, bins = np.histogram(np.log10(ifar), bins=nbins) ax_ifar.hist(ifar, bins=10.0**bins) -ax_ifar.set_xscale('log') -ax_ifar.set(xlabel='IFAR in found run') +ax_ifar.set_xscale("log") +ax_ifar.set(xlabel="IFAR in found run") ax_stat.hist(stat, bins=nbins) -ax_stat.set(xlabel='Ranking Statistic in found run') +ax_stat.set(xlabel="Ranking Statistic in found run") -ax_ifar.set(ylabel='Count') +ax_ifar.set(ylabel="Count") for ax in [ax_ifar, ax_stat]: if args.log_y: - ax.set_yscale('log') + ax.set_yscale("log") fig.savefig(args.outfile) plt.close() diff --git a/bin/workflow_comparisons/offline_search/pycbc_plot_vt_ratio_vs_ifar b/bin/workflow_comparisons/offline_search/pycbc_plot_vt_ratio_vs_ifar index e9a03096b8b..adbd437a196 100755 --- a/bin/workflow_comparisons/offline_search/pycbc_plot_vt_ratio_vs_ifar +++ b/bin/workflow_comparisons/offline_search/pycbc_plot_vt_ratio_vs_ifar @@ -21,32 +21,43 @@ Plot ratios of VTs calculated at various IFARs using pycbc_page_sensitivity """ import matplotlib -matplotlib.use('Agg') + +matplotlib.use("Agg") +import argparse +from math import ceil + import matplotlib.pyplot as plt import numpy as np -import argparse from matplotlib.pyplot import cm -from math import ceil from pycbc import add_common_pycbc_options, init_logging from pycbc.io.hdf import HFile parser = argparse.ArgumentParser(description=__doc__) add_common_pycbc_options(parser) -parser.add_argument('--vt-file-one', required=True, - help="HDF file containing VT curves, first set of data" - " for comparison") -parser.add_argument('--vt-file-two', required=True, - help="HDF file containing VT curves, second set of data" - " for comparison") -parser.add_argument('--outfile', type=str, required=True, - help='Output file to save to') -parser.add_argument('--plot-title', type=str, required=True, - help='(Possibly) quoted string to be title of plot') -parser.add_argument('--log-x', action='store_true', default=False, - help='Use logarithmic x-axis') -parser.add_argument('--log-y', action='store_true', default=False, - help='Use logarithmic y-axis') +parser.add_argument( + "--vt-file-one", + required=True, + help="HDF file containing VT curves, first set of data for comparison", +) +parser.add_argument( + "--vt-file-two", + required=True, + help="HDF file containing VT curves, second set of data for comparison", +) +parser.add_argument("--outfile", type=str, required=True, help="Output file to save to") +parser.add_argument( + "--plot-title", + type=str, + required=True, + help="(Possibly) quoted string to be title of plot", +) +parser.add_argument( + "--log-x", action="store_true", default=False, help="Use logarithmic x-axis" +) +parser.add_argument( + "--log-y", action="store_true", default=False, help="Use logarithmic y-axis" +) args = parser.parse_args() init_logging(args.verbose) @@ -55,23 +66,23 @@ init_logging(args.verbose) f1 = HFile(args.vt_file_one) f2 = HFile(args.vt_file_two) -x1 = f1['xvals'][:] -x2 = f2['xvals'][:] +x1 = f1["xvals"][:] +x2 = f2["xvals"][:] if not np.array_equal(x1, x2): - raise RuntimeError("IFAR values are not the same between the two files") + raise RuntimeError("IFAR values are not the same between the two files") xvals = x1 -keys = f1['data'].keys() +keys = f1["data"].keys() # sanitise the input so that the files have the same binning parameter and bins -assert keys == f2['data'].keys(), "keys do not match for the two files" +assert keys == f2["data"].keys(), "keys do not match for the two files" nkeys = len(keys) -#if nkeys != 6: +# if nkeys != 6: # raise RuntimeError("Only prepared for number of chirp mass bins to be six") -#fig, axs = plt.subplots(3, 2, sharex=True, sharey=True) -nrows = int(ceil(nkeys/2.0)) +# fig, axs = plt.subplots(3, 2, sharex=True, sharey=True) +nrows = int(ceil(nkeys / 2.0)) fig, axs = plt.subplots(nrows, 2, sharex=True, sharey=True) stitle = fig.suptitle(args.plot_title) @@ -85,50 +96,63 @@ lines = [] for n, key in zip(range(nkeys), keys): c = next(color) - data1 = f1['data'][key][:] - data2 = f2['data'][key][:] + data1 = f1["data"][key][:] + data2 = f2["data"][key][:] - errhigh1 = f1['errorhigh'][key][:] - errlow1 = f1['errorlow'][key][:] + errhigh1 = f1["errorhigh"][key][:] + errlow1 = f1["errorlow"][key][:] - errhigh2 = f2['errorhigh'][key][:] - errlow2 = f2['errorlow'][key][:] + errhigh2 = f2["errorhigh"][key][:] + errlow2 = f2["errorlow"][key][:] ys = np.divide(data1, data2) - yerr_errlow = np.multiply(np.sqrt(np.divide(errlow1, data1)**2 + - np.divide(errlow2, data2)**2), ys) - yerr_errhigh = np.multiply(np.sqrt(np.divide(errhigh1, data1)**2 + - np.divide(errhigh2, data2)**2), ys) - ax = axs[n%nrows, n//nrows] - ax.axhline(y=1.0, alpha=0.5, ls='--', color='k') - ax.plot(xvals, ys, c='k') - l, = ax.plot(xvals, ys, c=c) + yerr_errlow = np.multiply( + np.sqrt(np.divide(errlow1, data1) ** 2 + np.divide(errlow2, data2) ** 2), ys + ) + yerr_errhigh = np.multiply( + np.sqrt(np.divide(errhigh1, data1) ** 2 + np.divide(errhigh2, data2) ** 2), ys + ) + ax = axs[n % nrows, n // nrows] + ax.axhline(y=1.0, alpha=0.5, ls="--", color="k") + ax.plot(xvals, ys, c="k") + (l,) = ax.plot(xvals, ys, c=c) lines.append(l) - ax.fill_between(xvals, ys-yerr_errlow, ys+yerr_errhigh, - facecolor=c, edgecolor=c, alpha=alpha) + ax.fill_between( + xvals, + ys - yerr_errlow, + ys + yerr_errhigh, + facecolor=c, + edgecolor=c, + alpha=alpha, + ) if args.log_x: - ax.set_xscale('log') + ax.set_xscale("log") if args.log_y: - ax.set_yscale('log') + ax.set_yscale("log") - print("Bin {0}: minimum ratio {1}, maximum ratio {2}".format(key, - ys.min(), ys.max())) + print( + f"Bin {key}: minimum ratio {ys.min()}, maximum ratio {ys.max()}" + ) for ax in axs.flat: - ax.set(xlabel='Inverse False Alarm Rate (years)', - ylabel='VT ratio') + ax.set(xlabel="Inverse False Alarm Rate (years)", ylabel="VT ratio") for ax in axs.flat: - ax.label_outer() + ax.label_outer() # Positioning the legend seems to be a nightmare # First, grab the center, rightmost of the axes: ax = axs[1, 1] -lgd = ax.legend(handles=lines, labels=keys, loc="center left", - borderaxespad=0.1, title="Chirp mass bins", - bbox_to_anchor=(1.05,0.5)) - -fig.savefig(args.outfile, bbox_extra_artists=(lgd, stitle), bbox_inches='tight') +lgd = ax.legend( + handles=lines, + labels=keys, + loc="center left", + borderaxespad=0.1, + title="Chirp mass bins", + bbox_to_anchor=(1.05, 0.5), +) + +fig.savefig(args.outfile, bbox_extra_artists=(lgd, stitle), bbox_inches="tight") plt.close() diff --git a/bin/workflows/pycbc_make_bank_compression_workflow b/bin/workflows/pycbc_make_bank_compression_workflow index 8d9edcd7fc0..138380ace3b 100644 --- a/bin/workflows/pycbc_make_bank_compression_workflow +++ b/bin/workflows/pycbc_make_bank_compression_workflow @@ -2,38 +2,53 @@ """ Create a workflow for adding compressed waveforms to a template bank """ -import pycbc -import sys -import os import argparse import logging +import os import socket +import sys +import pycbc import pycbc.workflow as wf from pycbc.results import layout, save_fig_with_metadata + def finalize(container, workflow, finalize_workflow): # Create the final log file - log_file_html = wf.File(workflow.ifos, 'WORKFLOW-LOG', workflow.analysis_time, - extension='.html', directory=rdir['workflow']) + log_file_html = wf.File( + workflow.ifos, + "WORKFLOW-LOG", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], + ) - gen_file_html = wf.File(workflow.ifos, 'WORKFLOW-GEN', workflow.analysis_time, - extension='.html', directory=rdir['workflow']) + gen_file_html = wf.File( + workflow.ifos, + "WORKFLOW-GEN", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], + ) # Create a page to contain a dashboard link - dashboard_file = wf.File(workflow.ifos, 'DASHBOARD', workflow.analysis_time, - extension='.html', directory=rdir['workflow']) + dashboard_file = wf.File( + workflow.ifos, + "DASHBOARD", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], + ) dashboard_str = """

    Pegasus Dashboard Page

    """ - kwds = {'title': "Pegasus Dashboard", - 'caption': "Link to Pegasus Dashboard", - 'cmd': "PYCBC_SUBMIT_DAX_ARGV", } + kwds = { + "title": "Pegasus Dashboard", + "caption": "Link to Pegasus Dashboard", + "cmd": "PYCBC_SUBMIT_DAX_ARGV", + } save_fig_with_metadata(dashboard_str, dashboard_file.storage_path, **kwds) - wf.make_results_web_page( - finalize_workflow, - os.path.join(os.getcwd(), rdir.base) - ) + wf.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(), rdir.base)) container += workflow container += finalize_workflow @@ -46,7 +61,7 @@ def finalize(container, workflow, finalize_workflow): # Close the log and flush to the html file logging.shutdown() - with open(wf_log_file.storage_path, "r") as logfile: + with open(wf_log_file.storage_path) as logfile: logdata = logfile.read() log_str = """

    Workflow generation script created workflow in output directory: %s

    @@ -54,27 +69,33 @@ def finalize(container, workflow, finalize_workflow):

    Workflow generation script run on host: %s

    %s
    """ % (os.getcwd(), args.workflow_name, socket.gethostname(), logdata) - kwds = {'title': 'Workflow Generation Log', - 'caption': "Log of the workflow script %s" % sys.argv[0], - 'cmd': ' '.join(sys.argv), } + kwds = { + "title": "Workflow Generation Log", + "caption": "Log of the workflow script %s" % sys.argv[0], + "cmd": " ".join(sys.argv), + } save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds) # Add the command line used to a specific file args_to_output = [sys.argv[0]] for arg in sys.argv[1:]: - if arg.startswith('--'): + if arg.startswith("--"): # This is an option, add tab - args_to_output.append(' ' + arg) + args_to_output.append(" " + arg) else: # This is a parameter, add two tabs - args_to_output.append(' ' + arg) - - gen_str = '
    ' + ' \\\n'.join(args_to_output) + '
    ' - kwds = {'title': 'Workflow Generation Command', - 'caption': "Command used to generate the workflow.", - 'cmd': ' '.join(sys.argv), } + args_to_output.append(" " + arg) + + gen_str = "
    " + " \\\n".join(args_to_output) + "
    " + kwds = { + "title": "Workflow Generation Command", + "caption": "Command used to generate the workflow.", + "cmd": " ".join(sys.argv), + } save_fig_with_metadata(gen_str, gen_file_html.storage_path, **kwds) - layout.single_layout(rdir['workflow'], ([dashboard_file, gen_file_html, log_file_html])) + layout.single_layout( + rdir["workflow"], ([dashboard_file, gen_file_html, log_file_html]) + ) sys.exit(0) @@ -88,27 +109,32 @@ args = parser.parse_args() pycbc.init_logging(args.verbose, default_level=1) container = wf.Workflow(args, args.workflow_name) -workflow = wf.Workflow(args, args.workflow_name + '-main') -finalize_workflow = wf.Workflow(args, args.workflow_name + '-finalization') +workflow = wf.Workflow(args, args.workflow_name + "-main") +finalize_workflow = wf.Workflow(args, args.workflow_name + "-finalization") wf.makedir(args.output_dir) os.chdir(args.output_dir) rdir = layout.SectionNumber( - 'results', - ['workflow',], + "results", + [ + "workflow", + ], ) wf.makedir(rdir.base) -wf.makedir(rdir['workflow']) +wf.makedir(rdir["workflow"]) # We are _also_ logging to a file -wf_log_file = wf.File(workflow.ifos, 'workflow-log', workflow.analysis_time, - extension='.txt', - directory=rdir['workflow']) +wf_log_file = wf.File( + workflow.ifos, + "workflow-log", + workflow.analysis_time, + extension=".txt", + directory=rdir["workflow"], +) -pycbc.init_logging(args.verbose, default_level=1, - to_file=wf_log_file.storage_path) +pycbc.init_logging(args.verbose, default_level=1, to_file=wf_log_file.storage_path) logging.info("Created log file %s" % wf_log_file.storage_path) # Setup the workflow with the bank input file, @@ -117,14 +143,14 @@ hdfbank = wf.setup_tmpltbank_pregenerated( workflow, ) -assert( len(hdfbank) == 1 ) +assert len(hdfbank) == 1 hdfbank = hdfbank[0] -# Split the bank so that we can parallelise the +# Split the bank so that we can parallelise the splitbank_files = wf.setup_splittable_dax_generated( workflow, [hdfbank], - out_dir='split_bank', + out_dir="split_bank", tags=None, ) @@ -132,7 +158,7 @@ splitbank_files = wf.setup_splittable_dax_generated( compressed_files = wf.make_compress_split_banks( workflow, splitbank_files, - out_dir='compress_bank', + out_dir="compress_bank", tags=None, ) @@ -141,7 +167,7 @@ compressed_files = wf.make_compress_split_banks( combine_banks = wf.make_combine_split_banks( workflow, compressed_files, - out_dir='combine_bank', + out_dir="combine_bank", tags=None, ) @@ -159,21 +185,28 @@ layout.single_layout( ) # Save global config file to results directory -base = rdir['workflow/configuration'] +base = rdir["workflow/configuration"] wf.makedir(base) -ini_file_path = os.path.join(base, 'configuration.ini') -with open(ini_file_path, 'w') as ini_fh: +ini_file_path = os.path.join(base, "configuration.ini") +with open(ini_file_path, "w") as ini_fh: container.cp.write(ini_fh) -ini_file = wf.FileList([wf.File(workflow.ifos, '', workflow.analysis_time, - file_url='file://' + ini_file_path)]) +ini_file = wf.FileList( + [ + wf.File( + workflow.ifos, + "", + workflow.analysis_time, + file_url="file://" + ini_file_path, + ) + ] +) layout.single_layout(base, ini_file) # Create versioning information wf.make_versioning_page( workflow, workflow.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) finalize(container, workflow, finalize_workflow) - diff --git a/bin/workflows/pycbc_make_bank_verifier_workflow b/bin/workflows/pycbc_make_bank_verifier_workflow index ae47e893116..93b5507d99c 100644 --- a/bin/workflows/pycbc_make_bank_verifier_workflow +++ b/bin/workflows/pycbc_make_bank_verifier_workflow @@ -22,33 +22,36 @@ Workflow generator to create diagnosis plots and figures of merit for an input template bank. """ -#imports -import os +# imports import argparse +import os import shutil import igwn_segments as segments - -from pycbc import add_common_pycbc_options, init_logging import pycbc.version + import pycbc.workflow as wf +from pycbc import add_common_pycbc_options, init_logging from pycbc.results import layout -from pycbc.workflow.jobsetup import (select_generic_executable, - int_gps_time_to_str, - PycbcCreateInjectionsExecutable, - LalappsInspinjExecutable) from pycbc.workflow import setup_splittable_dax_generated +from pycbc.workflow.jobsetup import ( + LalappsInspinjExecutable, + PycbcCreateInjectionsExecutable, + int_gps_time_to_str, + select_generic_executable, +) # Boiler-plate stuff -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_make_bank_verifier_workflow" + # Some new executable classes. These can be moved into modules if needed class BanksimExecutable(wf.Executable): - """Class for running pycbc_banksim - """ + """Class for running pycbc_banksim""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.ALL_TRIGGERS @@ -56,15 +59,17 @@ class BanksimExecutable(wf.Executable): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) - node.add_input_opt('--signal-file', inj_file) - node.add_input_opt('--template-file', bank_file) - node.new_output_file_opt(analysis_time, '.dat', '--match-file', - tags=self.tags + extra_tags) + node.add_input_opt("--signal-file", inj_file) + node.add_input_opt("--template-file", bank_file) + node.new_output_file_opt( + analysis_time, ".dat", "--match-file", tags=self.tags + extra_tags + ) return node + class BanksimBankCombineExecutable(wf.Executable): - """Class for running pycbc_banksim_combine_banks - """ + """Class for running pycbc_banksim_combine_banks""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.ALL_TRIGGERS @@ -72,37 +77,41 @@ class BanksimBankCombineExecutable(wf.Executable): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) - node.add_input_list_opt('--input-files', inp_files) - node.new_output_file_opt(analysis_time, '.dat', '--output-file', - tags=self.tags + extra_tags) + node.add_input_list_opt("--input-files", inp_files) + node.new_output_file_opt( + analysis_time, ".dat", "--output-file", tags=self.tags + extra_tags + ) return node + class BanksimMatchCombineExecutable(wf.Executable): - """Class for running pycbc_banksim_match_combine - """ + """Class for running pycbc_banksim_match_combine""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.FINAL_RESULT - file_input_options = wf.Executable.file_input_options + \ - ['--filter-func-file'] + file_input_options = wf.Executable.file_input_options + ["--filter-func-file"] - def create_node(self, analysis_time, match_files, inj_files, bank_files, - extra_tags=None): + def create_node( + self, analysis_time, match_files, inj_files, bank_files, extra_tags=None + ): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) - node.add_input_list_opt('--match-files', match_files) + node.add_input_list_opt("--match-files", match_files) for curr_file in inj_files: node._add_input(curr_file) for curr_file in bank_files: node._add_input(curr_file) - node.new_output_file_opt(analysis_time, '.h5', '--output-file', - tags=self.tags + extra_tags) + node.new_output_file_opt( + analysis_time, ".h5", "--output-file", tags=self.tags + extra_tags + ) return node + class BanksimPlotFittingFactorsExecutable(wf.Executable): - """Class for running pycbc_banksim_plot_fitting_factors - """ + """Class for running pycbc_banksim_plot_fitting_factors""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.FINAL_RESULT @@ -110,14 +119,16 @@ class BanksimPlotFittingFactorsExecutable(wf.Executable): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) - node.add_input_opt('--input-file', input_file) - node.new_output_file_opt(analysis_time, '.png', '--output-file', - tags=self.tags + extra_tags) + node.add_input_opt("--input-file", input_file) + node.new_output_file_opt( + analysis_time, ".png", "--output-file", tags=self.tags + extra_tags + ) return node + class BanksimPlotEffFittingFactorsExecutable(wf.Executable): - """Class for running pycbc_banksim_plot_eff_fitting_factor - """ + """Class for running pycbc_banksim_plot_eff_fitting_factor""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.FINAL_RESULT @@ -125,32 +136,34 @@ class BanksimPlotEffFittingFactorsExecutable(wf.Executable): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) - node.add_input_list_opt('--input-files', input_files) - node.new_output_file_opt(analysis_time, '.png', '--output-file', - tags=self.tags + extra_tags) + node.add_input_list_opt("--input-files", input_files) + node.new_output_file_opt( + analysis_time, ".png", "--output-file", tags=self.tags + extra_tags + ) return node + class BanksimTablePointInjsExecutable(wf.Executable): - """Class for running pycbc_banksim_table_point_injs - """ + """Class for running pycbc_banksim_table_point_injs""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.FINAL_RESULT - def create_node(self, analysis_time, input_files, relative_dirs, - extra_tags=None): + def create_node(self, analysis_time, input_files, relative_dirs, extra_tags=None): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) - node.add_input_list_opt('--input-files', input_files) - node.add_list_opt('--directory-links', relative_dirs) - node.new_output_file_opt(analysis_time, '.html', '--output-file', - tags=self.tags + extra_tags) + node.add_input_list_opt("--input-files", input_files) + node.add_list_opt("--directory-links", relative_dirs) + node.new_output_file_opt( + analysis_time, ".html", "--output-file", tags=self.tags + extra_tags + ) return node # Argument parsing and setup of workflow -# Use the standard workflow command-line parsing routines. Things like a +# Use the standard workflow command-line parsing routines. Things like a # configuration file are specified within the "workflow command line group" # so run this with --help to see what options are added. _desc = __doc__[1:] @@ -169,50 +182,52 @@ workflow = wf.Workflow(args) wf.makedir(args.output_dir) os.chdir(args.output_dir) -args.output_dir = '.' +args.output_dir = "." -rdir = layout.SectionNumber('results', ['point_injection_sets', - 'broad_injection_sets', - 'workflow']) +rdir = layout.SectionNumber( + "results", ["point_injection_sets", "broad_injection_sets", "workflow"] +) wf.makedir(rdir.base) -wf.makedir(rdir['workflow']) +wf.makedir(rdir["workflow"]) # Save config file to results directory -conf_dir = rdir['workflow/configuration'] +conf_dir = rdir["workflow/configuration"] wf.makedir(conf_dir) -conf_path = os.path.join(conf_dir, 'configuration.ini') -with open(conf_path, 'w') as conf_fh: +conf_path = os.path.join(conf_dir, "configuration.ini") +with open(conf_path, "w") as conf_fh: workflow.cp.write(conf_fh) -conf_file = wf.FileList([wf.File(workflow.ifos, '', workflow.analysis_time, - file_url='file://' + conf_path)]) +conf_file = wf.FileList( + [wf.File(workflow.ifos, "", workflow.analysis_time, file_url="file://" + conf_path)] +) # Input bank file file_attrs = { - 'segs': workflow.analysis_time, - 'tags': [], - 'ifo_list': workflow.ifos, - 'description': 'TEMPLATEBANK' + "segs": workflow.analysis_time, + "tags": [], + "ifo_list": workflow.ifos, + "description": "TEMPLATEBANK", } -inp_bank = workflow.cp.get('workflow', 'input-bank') +inp_bank = workflow.cp.get("workflow", "input-bank") inp_bank = wf.resolve_url_to_file(inp_bank, attrs=file_attrs) # Inspinj Executable -inspinj_exe = select_generic_executable(workflow, 'injection') +inspinj_exe = select_generic_executable(workflow, "injection") # The output must be in xml format for pycbc_split_inspinj to work, # while h5 will work when splitting with pycbc_hdf5_splitbank -splitter = workflow.cp.get('executables', 'splitinj') -inspinj_exe.extension = ".xml" if 'pycbc_split_inspinj' in splitter else ".h5" +splitter = workflow.cp.get("executables", "splitinj") +inspinj_exe.extension = ".xml" if "pycbc_split_inspinj" in splitter else ".h5" # The following line is redundant for lalapps_inspinj, but not # for pycbc_create_injections inspinj_exe.current_retention_level = wf.Executable.FINAL_RESULT # Inspinj job -inspinj_job = inspinj_exe(workflow.cp, 'injection', out_dir='.', - ifos=workflow.ifos, tags=[]) +inspinj_job = inspinj_exe( + workflow.cp, "injection", out_dir=".", ifos=workflow.ifos, tags=[] +) + def add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks): - """Add a group of jobs that does a complete banksim. - """ - t_seg = segments.segment([1000000000, 1000000000+int(num_injs)]) + """Add a group of jobs that does a complete banksim.""" + t_seg = segments.segment([1000000000, 1000000000 + int(num_injs)]) inspinj_job.update_current_tags(curr_tags) if inspinj_exe is LalappsInspinjExecutable: node = inspinj_job.create_node(t_seg) @@ -220,7 +235,7 @@ def add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks): elif inspinj_exe is PycbcCreateInjectionsExecutable: # Ensure pycbc_create_injections dedicated configuration files # are copied over to the results directory - shutil.copy2(inspinj_job.get_opt('config-files'), conf_dir) + shutil.copy2(inspinj_job.get_opt("config-files"), conf_dir) node, inj_file = inspinj_job.create_node() node.add_opt("--gps-start-time", int_gps_time_to_str(t_seg[0])) node.add_opt("--gps-end-time", int_gps_time_to_str(t_seg[1])) @@ -228,88 +243,114 @@ def add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks): raise NotImplementedError workflow += node # Here we apply the em-bright criterion - if workflow.cp.has_option('workflow-injections', 'em-bright-only'): + if workflow.cp.has_option("workflow-injections", "em-bright-only"): # Job to carry on with em-bright injections only - em_filter_exe = select_generic_executable(workflow, 'em_bright_filter') - em_filter_job = em_filter_exe(workflow.cp, 'em_bright_filter', - out_dir='.', ifos=workflow.ifos, - tags=curr_tags) + em_filter_exe = select_generic_executable(workflow, "em_bright_filter") + em_filter_job = em_filter_exe( + workflow.cp, + "em_bright_filter", + out_dir=".", + ifos=workflow.ifos, + tags=curr_tags, + ) node = em_filter_job.create_node(inj_file, t_seg, curr_tags) workflow += node inj_file = node.output_files[0] - split_injs = setup_splittable_dax_generated(workflow, [inj_file], - 'splitinjfiles', curr_tags) + split_injs = setup_splittable_dax_generated( + workflow, [inj_file], "splitinjfiles", curr_tags + ) # Banksim job - banksim_job = BanksimExecutable(workflow.cp, 'banksim', - out_dir=file_tag+'match', - ifos=workflow.ifos, tags=[file_tag]) - bscombine_job = \ - BanksimBankCombineExecutable(workflow.cp, 'banksim_bank_combine', - out_dir=file_tag+'match', - ifos=workflow.ifos, tags=[file_tag]) - mcombine_job = \ - BanksimMatchCombineExecutable(workflow.cp, 'banksim_match_combine', - out_dir=file_tag+'match', - ifos=workflow.ifos, tags=[file_tag]) + banksim_job = BanksimExecutable( + workflow.cp, + "banksim", + out_dir=file_tag + "match", + ifos=workflow.ifos, + tags=[file_tag], + ) + bscombine_job = BanksimBankCombineExecutable( + workflow.cp, + "banksim_bank_combine", + out_dir=file_tag + "match", + ifos=workflow.ifos, + tags=[file_tag], + ) + mcombine_job = BanksimMatchCombineExecutable( + workflow.cp, + "banksim_match_combine", + out_dir=file_tag + "match", + ifos=workflow.ifos, + tags=[file_tag], + ) banksim_files = wf.FileList([]) for inj_idx, split_inj in enumerate(split_injs): - inj_tag = 'INJ{}'.format(inj_idx) + inj_tag = f"INJ{inj_idx}" currinj_banksim_files = wf.FileList([]) for bank_idx, split_bank in enumerate(split_banks): - bank_tag = 'BANK{}'.format(bank_idx) - inj_tag = 'INJ{}'.format(inj_idx) - node = banksim_job.create_node(workflow.analysis_time, split_inj, - split_bank, - extra_tags=[bank_tag,inj_tag]) - workflow+=node + bank_tag = f"BANK{bank_idx}" + inj_tag = f"INJ{inj_idx}" + node = banksim_job.create_node( + workflow.analysis_time, + split_inj, + split_bank, + extra_tags=[bank_tag, inj_tag], + ) + workflow += node currinj_banksim_files.append(node.output_file) - curr_node = bscombine_job.create_node(workflow.analysis_time, - currinj_banksim_files, - extra_tags=[inj_tag]) + curr_node = bscombine_job.create_node( + workflow.analysis_time, currinj_banksim_files, extra_tags=[inj_tag] + ) workflow += curr_node banksim_files.append(curr_node.output_file) - curr_node = mcombine_job.create_node(workflow.analysis_time, banksim_files, - split_injs, split_banks) + curr_node = mcombine_job.create_node( + workflow.analysis_time, banksim_files, split_injs, split_banks + ) workflow += curr_node return curr_node.output_file + # Set up the actual banksims -curr_tags = ['shortinjbanksplit'] -split_banks = setup_splittable_dax_generated(workflow, [inp_bank], - 'splitbankfiles', curr_tags) +curr_tags = ["shortinjbanksplit"] +split_banks = setup_splittable_dax_generated( + workflow, [inp_bank], "splitbankfiles", curr_tags +) output_pointinjs = {} -for file_tag, num_injs in workflow.cp.items('workflow-pointinjs'): - curr_tags = ['shortinjs', file_tag] - curr_file = add_banksim_set(workflow, file_tag, num_injs, curr_tags, - split_banks) +for file_tag, num_injs in workflow.cp.items("workflow-pointinjs"): + curr_tags = ["shortinjs", file_tag] + curr_file = add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks) output_pointinjs[file_tag] = curr_file -curr_tags = ['broadinjbanksplit'] -split_banks = setup_splittable_dax_generated(workflow, [inp_bank], - 'splitbankfiles', curr_tags) +curr_tags = ["broadinjbanksplit"] +split_banks = setup_splittable_dax_generated( + workflow, [inp_bank], "splitbankfiles", curr_tags +) output_broadinjs = {} -for file_tag, num_injs in workflow.cp.items('workflow-broadinjs'): - curr_tags = ['broadinjs', file_tag] - curr_file = add_banksim_set(workflow, file_tag, num_injs, curr_tags, - split_banks) +for file_tag, num_injs in workflow.cp.items("workflow-broadinjs"): + curr_tags = ["broadinjs", file_tag] + curr_file = add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks) output_broadinjs[file_tag] = curr_file plotting_nodes = [] out_dir = rdir.base -point_injs_table_exe = BanksimTablePointInjsExecutable\ - (workflow.cp, 'banksim_table_point_injs', - out_dir=rdir['point_injection_sets'], ifos=workflow.ifos) -eff_fitting_facs_exe = BanksimPlotEffFittingFactorsExecutable\ - (workflow.cp, 'banksim_plot_eff_fitting_fac', out_dir=out_dir, - ifos=workflow.ifos) -plot_fitting_facs_exe = BanksimPlotFittingFactorsExecutable\ - (workflow.cp, 'banksim_plot_fitting_factors', - out_dir=rdir['point_injection_sets'], ifos=workflow.ifos) +point_injs_table_exe = BanksimTablePointInjsExecutable( + workflow.cp, + "banksim_table_point_injs", + out_dir=rdir["point_injection_sets"], + ifos=workflow.ifos, +) +eff_fitting_facs_exe = BanksimPlotEffFittingFactorsExecutable( + workflow.cp, "banksim_plot_eff_fitting_fac", out_dir=out_dir, ifos=workflow.ifos +) +plot_fitting_facs_exe = BanksimPlotFittingFactorsExecutable( + workflow.cp, + "banksim_plot_fitting_factors", + out_dir=rdir["point_injection_sets"], + ifos=workflow.ifos, +) summary_page_files = [] # Add files to point_inj_summ_files in pairs. A tuple of one entry will span @@ -320,66 +361,67 @@ broad_inj_summ_files = [] # Set the point injection names for f in sorted(output_pointinjs): - rdir['point_injection_sets/{}'.format(f)] - -curr_node = point_injs_table_exe.create_node\ - (workflow.analysis_time, - [output_pointinjs[f] for f in sorted(output_pointinjs)], - ['../' + rdir.name['point_injection_sets/{}'.format(f)] - for f in sorted(output_pointinjs)]) + rdir[f"point_injection_sets/{f}"] + +curr_node = point_injs_table_exe.create_node( + workflow.analysis_time, + [output_pointinjs[f] for f in sorted(output_pointinjs)], + [ + "../" + rdir.name[f"point_injection_sets/{f}"] + for f in sorted(output_pointinjs) + ], +) workflow += curr_node plotting_nodes.append(curr_node) point_inj_summ_files.append((curr_node.output_file,)) -secs = workflow.cp.get_subsections('banksim_plot_eff_fitting_fac') +secs = workflow.cp.get_subsections("banksim_plot_eff_fitting_fac") for tag in secs: eff_fitting_facs_exe.update_current_tags([tag]) - curr_node = eff_fitting_facs_exe.create_node\ - (workflow.analysis_time, - [output_pointinjs[f] for f in output_pointinjs]) + curr_node = eff_fitting_facs_exe.create_node( + workflow.analysis_time, [output_pointinjs[f] for f in output_pointinjs] + ) workflow += curr_node plotting_nodes.append(curr_node) summary_page_files.append(curr_node.output_file) # Set up layouts layout.group_layout(rdir.base, summary_page_files) -layout.two_column_layout(rdir['point_injection_sets'], point_inj_summ_files) +layout.two_column_layout(rdir["point_injection_sets"], point_inj_summ_files) # Also add broad_injs when ready -secs = workflow.cp.get_subsections('banksim_plot_fitting_factors') +secs = workflow.cp.get_subsections("banksim_plot_fitting_factors") # Note a sorted(dict) returns a list of sorted *keys*. Works in python 2.4+ # and python 3 (dict.keys is removed in python 3) for tag in sorted(output_broadinjs): curr_outs = [] curr_file = output_broadinjs[tag] - plot_fitting_facs_exe.update_output_directory\ - (rdir['broad_injection_sets/{}'.format(tag)]) + plot_fitting_facs_exe.update_output_directory( + rdir[f"broad_injection_sets/{tag}"] + ) for tag2 in secs: - plot_fitting_facs_exe.update_current_tags([tag,tag2]) - curr_node = plot_fitting_facs_exe.create_node\ - (workflow.analysis_time, curr_file) + plot_fitting_facs_exe.update_current_tags([tag, tag2]) + curr_node = plot_fitting_facs_exe.create_node(workflow.analysis_time, curr_file) workflow += curr_node plotting_nodes.append(curr_node) curr_outs.append((curr_node.output_file,)) # Other outputs could go here, before running the layout - layout.two_column_layout(rdir['broad_injection_sets/{}'.format(tag)], - curr_outs) + layout.two_column_layout(rdir[f"broad_injection_sets/{tag}"], curr_outs) for tag in sorted(output_pointinjs): curr_outs = [] curr_file = output_pointinjs[tag] - plot_fitting_facs_exe.update_output_directory\ - (rdir['point_injection_sets/{}'.format(tag)]) + plot_fitting_facs_exe.update_output_directory( + rdir[f"point_injection_sets/{tag}"] + ) for tag2 in secs: - plot_fitting_facs_exe.update_current_tags([tag,tag2]) - curr_node = plot_fitting_facs_exe.create_node\ - (workflow.analysis_time, curr_file) + plot_fitting_facs_exe.update_current_tags([tag, tag2]) + curr_node = plot_fitting_facs_exe.create_node(workflow.analysis_time, curr_file) workflow += curr_node plotting_nodes.append(curr_node) curr_outs.append((curr_node.output_file,)) # Other outputs could go here, before running the layout - layout.two_column_layout(rdir['point_injection_sets/{}'.format(tag)], - curr_outs) + layout.two_column_layout(rdir[f"point_injection_sets/{tag}"], curr_outs) # Save config file(s) to results directory layout.single_layout(conf_dir, conf_file) @@ -388,11 +430,12 @@ layout.single_layout(conf_dir, conf_file) wf.make_versioning_page( workflow, workflow.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) -wf.make_results_web_page(workflow, os.path.join(os.getcwd(), rdir.base), - explicit_dependencies=plotting_nodes) +wf.make_results_web_page( + workflow, os.path.join(os.getcwd(), rdir.base), explicit_dependencies=plotting_nodes +) workflow.save() diff --git a/bin/workflows/pycbc_make_faithsim_workflow b/bin/workflows/pycbc_make_faithsim_workflow index d515366583f..a1aeb40b25a 100755 --- a/bin/workflows/pycbc_make_faithsim_workflow +++ b/bin/workflows/pycbc_make_faithsim_workflow @@ -5,12 +5,13 @@ Program for running a faithfulness comparisons workflow analysis between two app and generate files containing the match between them and plots. """ -import pycbc.workflow as wf import argparse +import pycbc.workflow as wf from pycbc import add_common_pycbc_options, init_logging from pycbc.workflow.plotting import PlotExecutable + def make_faithsim_plot(workflow, analysis_time, input_file, out_dir, tags=None): tags = [] if tags is None else tags secs = workflow.cp.get_subsections("pycbc_faithsim_plots") @@ -114,21 +115,24 @@ workflow += inj_node inj = inj_node.output_files[0] split_exe = wf.PycbcSplitBankXmlExecutable( - workflow.cp, "pycbc_splitbank", num_banks=num_banks, out_dir=workflow.output_dir +"bank" + workflow.cp, + "pycbc_splitbank", + num_banks=num_banks, + out_dir=workflow.output_dir + "bank", ) splitbank_node = split_exe.create_node(inj) workflow += spltbank_node faithsim_exe = FaithsimExecutable( - workflow.cp, "pycbc_faithsim", ifos=["X1"], out_dir=workflow.output_dir +"match" + workflow.cp, "pycbc_faithsim", ifos=["X1"], out_dir=workflow.output_dir + "match" ) collect_exe = CollectResultsExecutable( workflow.cp, "pycbc_faithsim_collect_results", ifos=["X1"], - out_dir=workflow.output_dir +"collect_results", + out_dir=workflow.output_dir + "collect_results", ) faithsim_files = wf.FileList([]) @@ -155,7 +159,7 @@ make_faithsim_plot( workflow, workflow.analysis_time, collect_results, - out_dir=workflow.output_dir +"plots", + out_dir=workflow.output_dir + "plots", tags=None, ) diff --git a/bin/workflows/pycbc_make_geom_aligned_bank_workflow b/bin/workflows/pycbc_make_geom_aligned_bank_workflow index f6233065352..3f681f13168 100755 --- a/bin/workflows/pycbc_make_geom_aligned_bank_workflow +++ b/bin/workflows/pycbc_make_geom_aligned_bank_workflow @@ -1,8 +1,9 @@ #!/usr/bin/env python -import os import argparse import logging +import os + import pycbc import pycbc.workflow as wf @@ -38,9 +39,7 @@ class GeomAligned2DStackExecutable(wf.Executable): current_retention_level = wf.Executable.MERGED_TRIGGERS file_input_options = ["--asd-file", "--psd-file"] - def create_node( - self, input_file, split_bank_num, analysis_time, extra_tags=None - ): + def create_node(self, input_file, split_bank_num, analysis_time, extra_tags=None): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) @@ -77,9 +76,7 @@ class AlignedBankCatExecutable(wf.Executable): out_file = wf.File.from_path(output_file_path) if self.retain_files: if not os.path.isabs(output_file_path): - out_file.storage_path = os.path.join( - self.out_dir, output_file_path - ) + out_file.storage_path = os.path.join(self.out_dir, output_file_path) else: out_file.storage_path = output_file_path node.add_output_opt("--output-file", out_file) @@ -93,9 +90,7 @@ class TmpltbankToChiParams(wf.Executable): def create_node(self, input_bank, analysis_time): node = wf.Executable.create_node(self) node.add_input_opt("--input-bank", input_bank) - node.new_output_file_opt( - analysis_time, ".dat", "--output-file", tags=self.tags - ) + node.new_output_file_opt(analysis_time, ".dat", "--output-file", tags=self.tags) return node @@ -130,9 +125,7 @@ pycbc.init_logging(args.verbose) workflow = wf.Workflow(args) wf.makedir(args.output_dir) -num_split_jobs = int( - workflow.cp.get_opt_tags("workflow", "num-split-jobs", args.tags) -) +num_split_jobs = int(workflow.cp.get_opt_tags("workflow", "num-split-jobs", args.tags)) setup_output_dir = os.path.join(args.output_dir, "geom_aligned_bank_setup") setup_exe = GeomAlignedBankSetupExecutable( @@ -140,12 +133,11 @@ setup_exe = GeomAlignedBankSetupExecutable( "geom_aligned_bank_setup", ifos=workflow.ifos, out_dir=setup_output_dir, - tags=args.tags + tags=args.tags, ) setup_node, intermediate_file, metadata_file = setup_exe.create_node( - workflow.analysis_time, - num_split_jobs + workflow.analysis_time, num_split_jobs ) workflow += setup_node @@ -162,10 +154,7 @@ all_outs = wf.FileList([]) for idx in range(num_split_jobs): stackid = "job%05d" % (idx) stack_node = stack_exe.create_node( - intermediate_file, - idx, - workflow.analysis_time, - extra_tags=[stackid] + intermediate_file, idx, workflow.analysis_time, extra_tags=[stackid] ) workflow += stack_node assert len(stack_node.output_files) == 1 diff --git a/bin/workflows/pycbc_make_inference_inj_workflow b/bin/workflows/pycbc_make_inference_inj_workflow index 7a443ba3cd3..c6ee830b5d2 100644 --- a/bin/workflows/pycbc_make_inference_inj_workflow +++ b/bin/workflows/pycbc_make_inference_inj_workflow @@ -15,28 +15,27 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -""" Creates a DAX for a parameter estimation injection study. -""" +"""Creates a DAX for a parameter estimation injection study.""" import argparse import logging import os import shlex -import numpy import socket import sys -from pycbc import results, init_logging, add_common_pycbc_options -from pycbc.results import layout -from pycbc.workflow import configuration -from pycbc.workflow import core -from pycbc.workflow.jobsetup import (PycbcCreateInjectionsExecutable, - PycbcInferenceExecutable) -from pycbc.workflow import inference_followups as inffu -from pycbc.workflow import plotting -from pycbc.workflow import versioning +import numpy + +from pycbc import add_common_pycbc_options, init_logging, results from pycbc.inject import InjectionSet from pycbc.io import FieldArray +from pycbc.results import layout +from pycbc.workflow import configuration, core, plotting, versioning +from pycbc.workflow import inference_followups as inffu +from pycbc.workflow.jobsetup import ( + PycbcCreateInjectionsExecutable, + PycbcInferenceExecutable, +) def config_from_config(cp, section, skip_opts=None): @@ -52,16 +51,15 @@ def config_from_config(cp, section, skip_opts=None): def read_inference_settings_from_config(cp, section): """Loads the config parser and gets the number of inference runs to do.""" - if cp.has_option(section, 'nruns'): - nruns = int(cp.get(section, 'nruns')) + if cp.has_option(section, "nruns"): + nruns = int(cp.get(section, "nruns")) else: nruns = 1 - return config_from_config(cp, section, skip_opts=['nruns']), nruns - + return config_from_config(cp, section, skip_opts=["nruns"]), nruns + def symlink_path(f, path): - """ Symlinks a path. - """ + """Symlinks a path.""" if f is None: return try: @@ -76,18 +74,26 @@ add_common_pycbc_options(parser) # injection options: either specify a number to create, or use the given file group = parser.add_mutually_exclusive_group(required=True) -group.add_argument("--num-injections", type=int, - help="The number of injections to create.") -group.add_argument("--injection-file", type=str, nargs="+", - help="Analyze injections in the given file(s) instead of " - "creating them.") +group.add_argument( + "--num-injections", type=int, help="The number of injections to create." +) +group.add_argument( + "--injection-file", + type=str, + nargs="+", + help="Analyze injections in the given file(s) instead of creating them.", +) # add option groups configuration.add_workflow_command_line_group(parser) # add workflow group core.add_workflow_settings_cli(parser, include_subdax_opts=True) -parser.add_argument("--seed", type=int, default=0, - help="Starting to seed to use. This will be incremented " - "one for each injection analyzed. Default is 0.") +parser.add_argument( + "--seed", + type=int, + default=0, + help="Starting to seed to use. This will be incremented " + "one for each injection analyzed. Default is 0.", +) # parser command line opts = parser.parse_args() @@ -95,22 +101,22 @@ opts = parser.parse_args() init_logging(opts.verbose, default_level=1) # configuration files -config_file_tmplt = 'inference-{}.ini' -config_file_dir = 'config_files' +config_file_tmplt = "inference-{}.ini" +config_file_dir = "config_files" # the directory we'll store samples files to -samples_file_dir = 'samples_files' +samples_file_dir = "samples_files" # the directory we'll store posterior files to -posterior_file_dir = 'posterior_files' +posterior_file_dir = "posterior_files" # the directory we'll store the injection files to -injection_file_dir = 'injection_files' +injection_file_dir = "injection_files" # make data output directory if opts.output_dir is None: - opts.output_dir = opts.workflow_name + '_output' + opts.output_dir = opts.workflow_name + "_output" core.makedir(opts.output_dir) -core.makedir('{}/{}'.format(opts.output_dir, config_file_dir)) -core.makedir('{}/{}'.format(opts.output_dir, posterior_file_dir)) -core.makedir('{}/{}'.format(opts.output_dir, injection_file_dir)) +core.makedir(f"{opts.output_dir}/{config_file_dir}") +core.makedir(f"{opts.output_dir}/{posterior_file_dir}") +core.makedir(f"{opts.output_dir}/{injection_file_dir}") # create workflow and sub-workflows workflow = core.Workflow(opts, name=opts.workflow_name) @@ -118,7 +124,8 @@ finalize_workflow = core.Workflow(opts, name="finalization") # load the inference config file inference_cp, nruns = read_inference_settings_from_config( - workflow.cp, 'workflow-inference') + workflow.cp, "workflow-inference" +) # change working directory to the output origdir = os.path.abspath(os.curdir) @@ -130,21 +137,28 @@ if opts.injection_file: injection_files = core.FileList([]) n_injections = 0 for injection_file in opts.injection_file: - injections = InjectionSet('{}/{}'.format(origdir, injection_file)) + injections = InjectionSet(f"{origdir}/{injection_file}") local_numinj = len(injections.table) - st_args = {k: injections.table[k][0] - for k in injections._injhandler.static_args} + st_args = { + k: injections.table[k][0] for k in injections._injhandler.static_args + } for ii in range(local_numinj): - samples = {f: injections.table[ii][f] - for f in injections.table.fieldnames - if f not in injections._injhandler.static_args} - outfile = '{}/injection{}.hdf'.format(injection_file_dir, - n_injections) + samples = { + f: injections.table[ii][f] + for f in injections.table.fieldnames + if f not in injections._injhandler.static_args + } + outfile = f"{injection_file_dir}/injection{n_injections}.hdf" injections.write( - outfile, FieldArray.from_kwargs(**samples), - write_params=[k for k in injections.table.fieldnames - if k not in injections._injhandler.static_args], - static_args=st_args) + outfile, + FieldArray.from_kwargs(**samples), + write_params=[ + k + for k in injections.table.fieldnames + if k not in injections._injhandler.static_args + ], + static_args=st_args, + ) injection_files.append(core.resolve_url_to_file(outfile)) n_injections += 1 else: @@ -154,28 +168,31 @@ else: diagnostics = inffu.get_diagnostic_plots(workflow) # figure out if we're doing pp tests -do_pp_test = workflow.cp.has_option('executables', 'pp_table_summary') +do_pp_test = workflow.cp.has_option("executables", "pp_table_summary") if do_pp_test: # get the parameters to do the test - pp_params = ["'"+s+"'" for s in - shlex.split(workflow.cp.get('workflow-pp_test', 'pp-params'))] + pp_params = [ + "'" + s + "'" + for s in shlex.split(workflow.cp.get("workflow-pp_test", "pp-params")) + ] # see if there is an injection map provided - if workflow.cp.has_option('workflow-pp_test', 'injection-samples-map'): - inj_samples_map = workflow.cp.get('workflow-pp_test', - 'injection-samples-map') + if workflow.cp.has_option("workflow-pp_test", "injection-samples-map"): + inj_samples_map = workflow.cp.get("workflow-pp_test", "injection-samples-map") else: inj_samples_map = None - pp_section = 'percentile-percentile_test' + pp_section = "percentile-percentile_test" pp_dir = [pp_section] else: pp_dir = [] # sections for output HTML pages -rdir = layout.SectionNumber("results", - pp_dir + - ["detector_sensitivity", "priors", "posteriors"] + - diagnostics + - ["config_files", "workflow"]) +rdir = layout.SectionNumber( + "results", + pp_dir + + ["detector_sensitivity", "priors", "posteriors"] + + diagnostics + + ["config_files", "workflow"], +) # make results directories core.makedir(rdir.base) @@ -183,12 +200,20 @@ core.makedir(rdir["workflow"]) core.makedir(rdir["config_files"]) # create files for workflow log -log_file_txt = core.File(workflow.ifos, "workflow-log", - workflow.analysis_time, - extension=".txt", directory=rdir["workflow"]) -log_file_html = core.File(workflow.ifos, "WORKFLOW-LOG", - workflow.analysis_time, - extension=".html", directory=rdir["workflow"]) +log_file_txt = core.File( + workflow.ifos, + "workflow-log", + workflow.analysis_time, + extension=".txt", + directory=rdir["workflow"], +) +log_file_html = core.File( + workflow.ifos, + "WORKFLOW-LOG", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], +) # Save log to file as well init_logging(opts.verbose, default_level=1, to_file=log_file_txt.storage_path) @@ -201,8 +226,8 @@ seed = opts.seed sub_workflows = [] for num_inj in range(n_injections): zpad = int(numpy.ceil(numpy.log10(n_injections))) - label = 'Injection {}'.format(str(num_inj+1).zfill(zpad)) - event = label.lower().replace(' ', '_') + label = f"Injection {str(num_inj + 1).zfill(zpad)}" + event = label.lower().replace(" ", "_") # create a sub workflow for this event # we need to go back to the original directory to do this for all the file @@ -214,23 +239,26 @@ for num_inj in range(n_injections): # if the config file has a fake-strain-seed set, increment to get # independent noise realizations - if inference_cp.has_option('data', 'fake-strain-seed'): + if inference_cp.has_option("data", "fake-strain-seed"): # get the detectors - detectors = shlex.split(inference_cp.get('data', 'instruments')) + detectors = shlex.split(inference_cp.get("data", "instruments")) fake_strain_seeds = {} for det in detectors: fake_strain_seeds[det] = seed seed = seed + 1 - inference_cp.set('data', 'fake-strain-seed', - ' '.join(['{}:{}'.format(d, s) - for d, s in fake_strain_seeds.items()])) + inference_cp.set( + "data", + "fake-strain-seed", + " ".join([f"{d}:{s}" for d, s in fake_strain_seeds.items()]), + ) # write the configuration file to the config files directory - config_file = sub_workflow.save_config(config_file_tmplt.format(event), - config_file_dir, inference_cp)[0] + config_file = sub_workflow.save_config( + config_file_tmplt.format(event), config_file_dir, inference_cp + )[0] # create sym links to config file for results page - base = "config_files/{}".format(event) + base = f"config_files/{event}" layout.single_layout(rdir[base], [config_file]) symlink_path(config_file, rdir[base]) @@ -240,32 +268,36 @@ for num_inj in range(n_injections): else: # construct Executable for creating injections create_injections_exe = PycbcCreateInjectionsExecutable( - sub_workflow.cp, "create_injections", ifos=sub_workflow.ifos, - out_dir=injection_file_dir) + sub_workflow.cp, + "create_injections", + ifos=sub_workflow.ifos, + out_dir=injection_file_dir, + ) node, injection_file = create_injections_exe.create_node( - [config_file], seed=seed, tags=opts.tags+[event]) + [config_file], seed=seed, tags=opts.tags + [event] + ) node.add_opt("--ninjections", 1) sub_workflow += node seed = seed + 1 # add the injection file to the inference config file cp = configuration.WorkflowConfigParser([config_file.storage_path]) - cp.set('data', 'injection-file', injection_file.name) + cp.set("data", "injection-file", injection_file.name) with open(config_file.storage_path, "w") as fp: cp.write(fp) # make node(s) for running sampler samples_files = [] - inference_exe = PycbcInferenceExecutable(sub_workflow.cp, "inference", - ifos=sub_workflow.ifos, - out_dir=samples_file_dir) + inference_exe = PycbcInferenceExecutable( + sub_workflow.cp, "inference", ifos=sub_workflow.ifos, out_dir=samples_file_dir + ) for nn in range(nruns): tags = opts.tags + [event] if nruns > 1: tags.append(str(nn)) node, samples_file = inference_exe.create_node( - config_file, seed=seed, tags=tags, - analysis_time=sub_workflow.analysis_time) + config_file, seed=seed, tags=tags, analysis_time=sub_workflow.analysis_time + ) # declare the injection file as a needed input file node.add_input(injection_file) # add node to workflow @@ -276,31 +308,49 @@ for num_inj in range(n_injections): # create the posterior file and plots posterior_file, summary_files, _, _ = inffu.make_posterior_workflow( - sub_workflow, samples_files, config_file, event, rdir, - posterior_file_dir=posterior_file_dir, tags=opts.tags) + sub_workflow, + samples_files, + config_file, + event, + rdir, + posterior_file_dir=posterior_file_dir, + tags=opts.tags, + ) posterior_files.append(posterior_file) # create the diagnostic plots - _ = inffu.make_diagnostic_plots(sub_workflow, diagnostics, samples_files, - event, rdir, tags=opts.tags) + _ = inffu.make_diagnostic_plots( + sub_workflow, diagnostics, samples_files, event, rdir, tags=opts.tags + ) # files for detector_sensitivity summary subsection base = "detector_sensitivity" psd_plot = plotting.make_spectrum_plot( - sub_workflow, [samples_files[0]], rdir[base], - tags=opts.tags+[event], - hdf_group="data") + sub_workflow, + [samples_files[0]], + rdir[base], + tags=opts.tags + [event], + hdf_group="data", + ) # build the summary page zpad = int(numpy.ceil(numpy.log10(len(samples_files)))) - layout.two_column_layout(rdir.base, summary_files, - unique=str(num_inj).zfill(zpad), - title=label, collapse=True) + layout.two_column_layout( + rdir.base, + summary_files, + unique=str(num_inj).zfill(zpad), + title=label, + collapse=True, + ) # build the psd page - layout.single_layout(rdir['detector_sensitivity'], [psd_plot], - unique=str(num_inj).zfill(zpad), - title=label, collapse=True) + layout.single_layout( + rdir["detector_sensitivity"], + [psd_plot], + unique=str(num_inj).zfill(zpad), + title=label, + collapse=True, + ) # add the sub workflow to the main workflow workflow += sub_workflow @@ -312,30 +362,45 @@ if do_pp_test: # we need to go back to the original directory to do this for all the file # references to work correctly os.chdir(origdir) - pp_workflow = core.Workflow(opts, name='pp_test') + pp_workflow = core.Workflow(opts, name="pp_test") # now go back to the output os.chdir(opts.output_dir) # create the pp summary table pp_table = inffu.make_inference_pp_table( - pp_workflow, posterior_files, rdir[pp_section], - parameters=pp_params, injection_samples_map=inj_samples_map, - analysis_seg=workflow.analysis_time, tags=opts.tags)[0] + pp_workflow, + posterior_files, + rdir[pp_section], + parameters=pp_params, + injection_samples_map=inj_samples_map, + analysis_seg=workflow.analysis_time, + tags=opts.tags, + )[0] # now the pp plots and injection recovery pp_plots = [] inj_recovery_plots = [] for pi, param in enumerate(pp_params): - tags = opts.tags + ['PARAM_{}'.format(pi)] + tags = opts.tags + [f"PARAM_{pi}"] # the pp plot pp_plot = inffu.make_inference_pp_plot( - pp_workflow, posterior_files, rdir[pp_section], - parameters=param, injection_samples_map=inj_samples_map, - analysis_seg=workflow.analysis_time, tags=tags)[0] + pp_workflow, + posterior_files, + rdir[pp_section], + parameters=param, + injection_samples_map=inj_samples_map, + analysis_seg=workflow.analysis_time, + tags=tags, + )[0] pp_plots.append(pp_plot) # the injection recovery plot injrec_plot = inffu.make_inference_inj_recovery_plot( - pp_workflow, posterior_files, rdir[pp_section], param, + pp_workflow, + posterior_files, + rdir[pp_section], + param, injection_samples_map=inj_samples_map, - analysis_seg=workflow.analysis_time, tags=tags)[0] + analysis_seg=workflow.analysis_time, + tags=tags, + )[0] inj_recovery_plots.append(injrec_plot) # add to the results page pp_files = [(pp_table,)] + list(zip(pp_plots, inj_recovery_plots)) @@ -347,12 +412,11 @@ if do_pp_test: versioning.make_versioning_page( workflow, workflow.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) # create node for making HTML pages -plotting.make_results_web_page(finalize_workflow, - os.path.join(os.getcwd(), rdir.base)) +plotting.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(), rdir.base)) # add finalize workflow to workflow and make it depend on the others workflow += finalize_workflow @@ -372,7 +436,7 @@ layout.single_layout(base, wf_ini) # close the log and flush to the html file logging.shutdown() -with open (log_file_txt.storage_path, "r") as log_file: +with open(log_file_txt.storage_path) as log_file: log_data = log_file.read() log_str = """

    Workflow generation script created workflow in output directory: %s

    @@ -380,8 +444,10 @@ log_str = """

    Workflow generation script run on host: %s

    %s
    """ % (os.getcwd(), opts.workflow_name, socket.gethostname(), log_data) -kwds = {"title" : "Workflow Generation Log", - "caption" : "Log of the workflow script %s" % sys.argv[0], - "cmd" : " ".join(sys.argv)} +kwds = { + "title": "Workflow Generation Log", + "caption": "Log of the workflow script %s" % sys.argv[0], + "cmd": " ".join(sys.argv), +} results.save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds) layout.single_layout(rdir["workflow"], ([log_file_html])) diff --git a/bin/workflows/pycbc_make_inference_plots_workflow b/bin/workflows/pycbc_make_inference_plots_workflow index 0a97d7206a7..ea3ad4dc547 100644 --- a/bin/workflows/pycbc_make_inference_plots_workflow +++ b/bin/workflows/pycbc_make_inference_plots_workflow @@ -15,29 +15,29 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Creates a DAX that generates a posterior file and plots from one or more +""" +Creates a DAX that generates a posterior file and plots from one or more inference samples files. """ import argparse import logging import os +import shlex import socket import sys -import shlex + import numpy -from pycbc import results, init_logging, add_common_pycbc_options -from pycbc.results import layout -from pycbc.workflow import configuration -from pycbc.workflow import core -from pycbc.workflow import plotting -from pycbc.workflow import versioning import pycbc.workflow.inference_followups as inffu +from pycbc import add_common_pycbc_options, init_logging, results +from pycbc.results import layout +from pycbc.workflow import configuration, core, plotting, versioning def read_events_from_config(cp): - """Gets events to load from a config file. + """ + Gets events to load from a config file. Each event should have its own section with header ``[event-{{NAME}}]``, where ``NAME`` is a unique identifier. The section must have a @@ -77,9 +77,10 @@ def read_events_from_config(cp): List of the configuration file(s) for each event. samples-files : lsit of lists List of the samples file(s) for each event. + """ # get the events - events = cp.get_subsections('event') + events = cp.get_subsections("event") # create a dummy command-line parser for getting the config files and # options cfparser = argparse.ArgumentParser() @@ -89,13 +90,13 @@ def read_events_from_config(cp): config_files = [] samples_files = [] for event in events: - section = '-'.join(['event', event]) - if cp.has_option(section, 'label'): - label = cp.get(section, 'label') + section = "-".join(["event", event]) + if cp.has_option(section, "label"): + label = cp.get(section, "label") else: label = event - cf = shlex.split(cp.get(section, 'config-files')) - sf = shlex.split(cp.get(section, 'samples-files')) + cf = shlex.split(cp.get(section, "config-files")) + sf = shlex.split(cp.get(section, "samples-files")) labels.append(label) config_files.append(list(map(os.path.abspath, cf))) samples_files.append(list(map(os.path.abspath, sf))) @@ -104,12 +105,11 @@ def read_events_from_config(cp): def event_slug(label): """Slugifies an event label.""" - return label.replace(' ', '_').replace(':', '_').replace('+', '_') + return label.replace(" ", "_").replace(":", "_").replace("+", "_") def symlink_path(f, path): - """ Symlinks a path. - """ + """Symlinks a path.""" if f is None: return try: @@ -127,28 +127,27 @@ configuration.add_workflow_command_line_group(parser) core.add_workflow_settings_cli(parser, include_subdax_opts=True) opts = parser.parse_args() -posterior_file_dir = 'posterior_files' -config_file_dir = 'config_files' -config_file_tmplt = 'inference-{}.ini' +posterior_file_dir = "posterior_files" +config_file_dir = "config_files" +config_file_tmplt = "inference-{}.ini" # make data output directory if opts.output_dir is None: - opts.output_dir = opts.workflow_name + '_output' + opts.output_dir = opts.workflow_name + "_output" core.makedir(opts.output_dir) -core.makedir('{}/{}'.format(opts.output_dir, config_file_dir)) -core.makedir('{}/{}'.format(opts.output_dir, posterior_file_dir)) +core.makedir(f"{opts.output_dir}/{config_file_dir}") +core.makedir(f"{opts.output_dir}/{posterior_file_dir}") # log to terminal until we know where the path to log output file init_logging(opts.verbose, default_level=1) # create workflow and sub-workflows container = core.Workflow(opts, opts.workflow_name) -workflow = core.Workflow(opts, 'main') +workflow = core.Workflow(opts, "main") finalize_workflow = core.Workflow(opts, "finalization") # get the events -events, labels, infconfig_files, samples_files = \ - read_events_from_config(workflow.cp) +events, labels, infconfig_files, samples_files = read_events_from_config(workflow.cp) # change working directory to the output origdir = os.path.abspath(os.curdir) @@ -158,10 +157,12 @@ os.chdir(opts.output_dir) diagnostics = inffu.get_diagnostic_plots(workflow) # sections for output HTML pages -rdir = layout.SectionNumber("results", - ["detector_sensitivity", "priors", "posteriors"] + - diagnostics + - ["config_files", "workflow"]) +rdir = layout.SectionNumber( + "results", + ["detector_sensitivity", "priors", "posteriors"] + + diagnostics + + ["config_files", "workflow"], +) # make results directories core.makedir(rdir.base) @@ -169,10 +170,20 @@ core.makedir(rdir["workflow"]) core.makedir(rdir["config_files"]) # create files for workflow log -log_file_txt = core.File(workflow.ifos, "workflow-log", workflow.analysis_time, - extension=".txt", directory=rdir["workflow"]) -log_file_html = core.File(workflow.ifos, "WORKFLOW-LOG", workflow.analysis_time, - extension=".html", directory=rdir["workflow"]) +log_file_txt = core.File( + workflow.ifos, + "workflow-log", + workflow.analysis_time, + extension=".txt", + directory=rdir["workflow"], +) +log_file_html = core.File( + workflow.ifos, + "WORKFLOW-LOG", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], +) # Save log to file as well init_logging(opts.verbose, default_level=1, to_file=log_file_txt.storage_path) @@ -189,10 +200,11 @@ for num_event, event in enumerate(events): # write the configuration file to the config files directory cp = configuration.WorkflowConfigParser(config_fnames) - config_file = workflow.save_config(config_file_tmplt.format(event), - config_file_dir, cp)[0] + config_file = workflow.save_config( + config_file_tmplt.format(event), config_file_dir, cp + )[0] # create sym links to config file for results page - base = "config_files/{}".format(event) + base = f"config_files/{event}" layout.single_layout(rdir[base], [config_file]) symlink_path(config_file, rdir[base]) @@ -207,42 +219,59 @@ for num_event, event in enumerate(events): # create the posterior file and plots posterior_file, summary_files, _, _ = inffu.make_posterior_workflow( - workflow, samples_filelist, config_file, event, rdir, - posterior_file_dir=posterior_file_dir, tags=opts.tags) + workflow, + samples_filelist, + config_file, + event, + rdir, + posterior_file_dir=posterior_file_dir, + tags=opts.tags, + ) # create the diagnostic plots - _ = inffu.make_diagnostic_plots(workflow, diagnostics, samples_filelist, - event, rdir, tags=opts.tags) + _ = inffu.make_diagnostic_plots( + workflow, diagnostics, samples_filelist, event, rdir, tags=opts.tags + ) # files for detector_sensitivity summary subsection base = "detector_sensitivity" # we'll just use the first file, and assume the rest are the same psd_plot = plotting.make_spectrum_plot( - workflow, [samples_filelist[0]], rdir[base], - tags=opts.tags+[event], - hdf_group="data") + workflow, + [samples_filelist[0]], + rdir[base], + tags=opts.tags + [event], + hdf_group="data", + ) # build the summary page zpad = int(numpy.ceil(numpy.log10(len(samples_files)))) - layout.two_column_layout(rdir.base, summary_files, - unique=str(num_event).zfill(zpad), - title=label, collapse=True) + layout.two_column_layout( + rdir.base, + summary_files, + unique=str(num_event).zfill(zpad), + title=label, + collapse=True, + ) # build the psd page - layout.single_layout(rdir['detector_sensitivity'], [psd_plot], - unique=str(num_event).zfill(zpad), - title=label, collapse=True) - + layout.single_layout( + rdir["detector_sensitivity"], + [psd_plot], + unique=str(num_event).zfill(zpad), + title=label, + collapse=True, + ) + # Create versioning information versioning.make_versioning_page( workflow, container.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) # create node for making HTML pages -plotting.make_results_web_page(finalize_workflow, - os.path.join(os.getcwd(), rdir.base)) +plotting.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(), rdir.base)) # add sub-workflows to workflow container += workflow @@ -262,7 +291,7 @@ layout.single_layout(base, wf_ini) # close the log and flush to the html file logging.shutdown() -with open (log_file_txt.storage_path, "r") as log_file: +with open(log_file_txt.storage_path) as log_file: log_data = log_file.read() log_str = """

    Workflow generation script created workflow in output directory: %s

    @@ -270,8 +299,10 @@ log_str = """

    Workflow generation script run on host: %s

    %s
    """ % (os.getcwd(), opts.workflow_name, socket.gethostname(), log_data) -kwds = {"title" : "Workflow Generation Log", - "caption" : "Log of the workflow script %s" % sys.argv[0], - "cmd" : " ".join(sys.argv)} +kwds = { + "title": "Workflow Generation Log", + "caption": "Log of the workflow script %s" % sys.argv[0], + "cmd": " ".join(sys.argv), +} results.save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds) layout.single_layout(rdir["workflow"], ([log_file_html])) diff --git a/bin/workflows/pycbc_make_inference_workflow b/bin/workflows/pycbc_make_inference_workflow index 9f820350069..04a93af28e2 100644 --- a/bin/workflows/pycbc_make_inference_workflow +++ b/bin/workflows/pycbc_make_inference_workflow @@ -15,8 +15,7 @@ # You should have received a copy of the GNU General Public License along # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Creates a DAX for a parameter estimation workflow. -""" +"""Creates a DAX for a parameter estimation workflow.""" import argparse import logging @@ -24,21 +23,20 @@ import os import shlex import socket import sys + import numpy import pycbc +import pycbc.workflow.inference_followups as inffu from pycbc import results from pycbc.results import layout -from pycbc.workflow import configuration -from pycbc.workflow import core -from pycbc.workflow import plotting -from pycbc.workflow import versioning -import pycbc.workflow.inference_followups as inffu +from pycbc.workflow import configuration, core, plotting, versioning from pycbc.workflow.jobsetup import PycbcInferenceExecutable def read_events_from_config(cp): - """Gets events to load from a config file. + """ + Gets events to load from a config file. Each event should have its own section with header ``[event-{{NAME}}]``, where ``NAME`` is a unique identifier. The section must have a @@ -94,9 +92,10 @@ def read_events_from_config(cp): argparse.ArgumentParser options. nruns : list of int The number of inference jobs to create for the event. + """ # get the events - events = cp.get_subsections('event') + events = cp.get_subsections("event") # create a dummy command-line parser for getting the config files and # options cfparser = argparse.ArgumentParser() @@ -106,35 +105,34 @@ def read_events_from_config(cp): cpopts = [] nruns = [] for event in events: - section = '-'.join(['event', event]) + section = "-".join(["event", event]) # get the label - if cp.has_option(section, 'label'): - label = cp.get(section, 'label') + if cp.has_option(section, "label"): + label = cp.get(section, "label") else: label = event labels.append(label) # convert the config-file options to a command line string - cli = cp.section_to_cli(section, skip_opts=['label', 'nruns']) + cli = cp.section_to_cli(section, skip_opts=["label", "nruns"]) cpopts.append(cfparser.parse_args(shlex.split(cli))) # get the number of times to run the event - if cp.has_option(section, 'nruns'): - nrun = int(cp.get(section, 'nruns')) + if cp.has_option(section, "nruns"): + nrun = int(cp.get(section, "nruns")) else: nrun = 1 if nrun < 1: - raise ValueError('nruns must be >= 1') + raise ValueError("nruns must be >= 1") nruns.append(nrun) return events, labels, cpopts, nruns def event_slug(label): """Slugifies an event label.""" - return label.replace(' ', '_').replace(':', '_').replace('+', '_') + return label.replace(" ", "_").replace(":", "_").replace("+", "_") def symlink_path(f, path): - """ Symlinks a path. - """ + """Symlinks a path.""" if f is None: return try: @@ -150,10 +148,14 @@ pycbc.add_common_pycbc_options(parser) configuration.add_workflow_command_line_group(parser) # workflow options core.add_workflow_settings_cli(parser, include_subdax_opts=True) -parser.add_argument("--seed", type=int, default=0, - help="Seed to use for inference job(s). If multiple " - "events are analyzed, the seed will be incremented " - "by one for each event.") +parser.add_argument( + "--seed", + type=int, + default=0, + help="Seed to use for inference job(s). If multiple " + "events are analyzed, the seed will be incremented " + "by one for each event.", +) # parser command line @@ -163,23 +165,23 @@ opts = parser.parse_args() pycbc.init_logging(opts.verbose, default_level=1) # configuration files -config_file_tmplt = 'inference-{}.ini' -config_file_dir = 'config_files' +config_file_tmplt = "inference-{}.ini" +config_file_dir = "config_files" # the directory we'll store samples files to -samples_file_dir = 'samples_files' +samples_file_dir = "samples_files" # the directory we'll store posterior files to -posterior_file_dir = 'posterior_files' +posterior_file_dir = "posterior_files" # make data output directory if opts.output_dir is None: - opts.output_dir = opts.workflow_name + '_output' + opts.output_dir = opts.workflow_name + "_output" core.makedir(opts.output_dir) -core.makedir('{}/{}'.format(opts.output_dir, config_file_dir)) -core.makedir('{}/{}'.format(opts.output_dir, posterior_file_dir)) +core.makedir(f"{opts.output_dir}/{config_file_dir}") +core.makedir(f"{opts.output_dir}/{posterior_file_dir}") # create workflow and sub-workflows container = core.Workflow(opts, opts.workflow_name) -workflow = core.Workflow(opts, 'main') +workflow = core.Workflow(opts, "main") finalize_workflow = core.Workflow(opts, "finalization") # read the events to analyze @@ -193,10 +195,12 @@ os.chdir(opts.output_dir) diagnostics = inffu.get_diagnostic_plots(workflow) # sections for output HTML pages -rdir = layout.SectionNumber("results", - ["detector_sensitivity", "priors", "posteriors"] + - diagnostics + - ["config_files", "workflow"]) +rdir = layout.SectionNumber( + "results", + ["detector_sensitivity", "priors", "posteriors"] + + diagnostics + + ["config_files", "workflow"], +) # make results directories core.makedir(rdir.base) @@ -204,11 +208,20 @@ core.makedir(rdir["workflow"]) core.makedir(rdir["config_files"]) # create files for workflow log -log_file_txt = core.File(workflow.ifos, "workflow-log", workflow.analysis_time, - extension=".txt", directory=rdir["workflow"]) -log_file_html = core.File(workflow.ifos, "WORKFLOW-LOG", - workflow.analysis_time, - extension=".html", directory=rdir["workflow"]) +log_file_txt = core.File( + workflow.ifos, + "workflow-log", + workflow.analysis_time, + extension=".txt", + directory=rdir["workflow"], +) +log_file_html = core.File( + workflow.ifos, + "WORKFLOW-LOG", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], +) # Save log to file pycbc.init_logging(opts.verbose, default_level=1, to_file=log_file_txt.storage_path) @@ -236,26 +249,27 @@ for num_event, event in enumerate(events): os.chdir(opts.output_dir) # write the configuration file to the config files directory - config_file = sub_workflow.save_config(config_file_tmplt.format(event), - config_file_dir, cp)[0] + config_file = sub_workflow.save_config( + config_file_tmplt.format(event), config_file_dir, cp + )[0] # create sym links to config file for results page - base = "config_files/{}".format(event) + base = f"config_files/{event}" layout.single_layout(rdir[base], [config_file]) symlink_path(config_file, rdir[base]) # make node(s) for running sampler samples_files = [] - inference_exe = PycbcInferenceExecutable(sub_workflow.cp, "inference", - ifos=sub_workflow.ifos, - out_dir=samples_file_dir) + inference_exe = PycbcInferenceExecutable( + sub_workflow.cp, "inference", ifos=sub_workflow.ifos, out_dir=samples_file_dir + ) for nn in range(nrun): tags = opts.tags + [event] if nrun > 1: tags.append(str(nn)) node, samples_file = inference_exe.create_node( - config_file, seed=seed, tags=tags, - analysis_time=sub_workflow.analysis_time) + config_file, seed=seed, tags=tags, analysis_time=sub_workflow.analysis_time + ) samples_files.append(samples_file) # add node to workflow sub_workflow += node @@ -264,30 +278,48 @@ for num_event, event in enumerate(events): # create the posterior file and plots posterior_file, summary_files, _, _ = inffu.make_posterior_workflow( - sub_workflow, samples_files, config_file, event, rdir, - posterior_file_dir=posterior_file_dir, tags=opts.tags) + sub_workflow, + samples_files, + config_file, + event, + rdir, + posterior_file_dir=posterior_file_dir, + tags=opts.tags, + ) # create the diagnostic plots - _ = inffu.make_diagnostic_plots(sub_workflow, diagnostics, samples_files, - event, rdir, tags=opts.tags) + _ = inffu.make_diagnostic_plots( + sub_workflow, diagnostics, samples_files, event, rdir, tags=opts.tags + ) # files for detector_sensitivity summary subsection base = "detector_sensitivity" psd_plot = plotting.make_spectrum_plot( - sub_workflow, [samples_files[0]], rdir[base], - tags=opts.tags+[event], - hdf_group="data") + sub_workflow, + [samples_files[0]], + rdir[base], + tags=opts.tags + [event], + hdf_group="data", + ) # build the summary page zpad = int(numpy.ceil(numpy.log10(len(samples_files)))) - layout.two_column_layout(rdir.base, summary_files, - unique=str(num_event).zfill(zpad), - title=label, collapse=True) + layout.two_column_layout( + rdir.base, + summary_files, + unique=str(num_event).zfill(zpad), + title=label, + collapse=True, + ) # build the psd page - layout.single_layout(rdir['detector_sensitivity'], [psd_plot], - unique=str(num_event).zfill(zpad), - title=label, collapse=True) + layout.single_layout( + rdir["detector_sensitivity"], + [psd_plot], + unique=str(num_event).zfill(zpad), + title=label, + collapse=True, + ) # add the sub workflow to the main workflow workflow += sub_workflow @@ -297,12 +329,11 @@ for num_event, event in enumerate(events): versioning.make_versioning_page( workflow, container.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) # create node for making HTML pages -plotting.make_results_web_page(finalize_workflow, - os.path.join(os.getcwd(), rdir.base)) +plotting.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(), rdir.base)) # add sub-workflows to workflow container += workflow @@ -322,7 +353,7 @@ layout.single_layout(base, wf_ini) # close the log and flush to the html file logging.shutdown() -with open (log_file_txt.storage_path, "r") as log_file: +with open(log_file_txt.storage_path) as log_file: log_data = log_file.read() log_str = """

    Workflow generation script created workflow in output directory: %s

    @@ -330,8 +361,10 @@ log_str = """

    Workflow generation script run on host: %s

    %s
    """ % (os.getcwd(), opts.workflow_name, socket.gethostname(), log_data) -kwds = {"title" : "Workflow Generation Log", - "caption" : "Log of the workflow script %s" % sys.argv[0], - "cmd" : " ".join(sys.argv)} +kwds = { + "title": "Workflow Generation Log", + "caption": "Log of the workflow script %s" % sys.argv[0], + "cmd": " ".join(sys.argv), +} results.save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds) layout.single_layout(rdir["workflow"], ([log_file_html])) diff --git a/bin/workflows/pycbc_make_offline_search_workflow b/bin/workflows/pycbc_make_offline_search_workflow index 7e251cb55e8..57862403e01 100755 --- a/bin/workflows/pycbc_make_offline_search_workflow +++ b/bin/workflows/pycbc_make_offline_search_workflow @@ -21,17 +21,17 @@ Program for running offline analysis through event finding and ranking then generate post-processing and plots. """ -import pycbc -import sys -import socket -import os import argparse -import logging import itertools +import logging +import os +import socket +import sys import igwn_segments as segments +import pycbc import pycbc.events import pycbc.workflow as wf from pycbc.results import layout, save_fig_with_metadata @@ -53,7 +53,7 @@ def symlink_result(f, rdir_path): # Generator for producing ifo combinations def ifo_combos(ifos): - for i in range(2, len(ifos)+1): + for i in range(2, len(ifos) + 1): combinations = itertools.combinations(ifos, i) for ifocomb in combinations: yield ifocomb @@ -61,23 +61,39 @@ def ifo_combos(ifos): def finalize(container, workflow, finalize_workflow): # Create the final log file - log_file_html = wf.File(workflow.ifos, 'WORKFLOW-LOG', workflow.analysis_time, - extension='.html', directory=rdir['workflow']) + log_file_html = wf.File( + workflow.ifos, + "WORKFLOW-LOG", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], + ) - gen_file_html = wf.File(workflow.ifos, 'WORKFLOW-GEN', workflow.analysis_time, - extension='.html', directory=rdir['workflow']) + gen_file_html = wf.File( + workflow.ifos, + "WORKFLOW-GEN", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], + ) # Create a page to contain a dashboard link - dashboard_file = wf.File(workflow.ifos, 'DASHBOARD', workflow.analysis_time, - extension='.html', directory=rdir['workflow']) + dashboard_file = wf.File( + workflow.ifos, + "DASHBOARD", + workflow.analysis_time, + extension=".html", + directory=rdir["workflow"], + ) dashboard_str = """

    Pegasus Dashboard Page

    """ - kwds = {'title': "Pegasus Dashboard", - 'caption': "Link to Pegasus Dashboard", - 'cmd': "PYCBC_SUBMIT_DAX_ARGV", } + kwds = { + "title": "Pegasus Dashboard", + "caption": "Link to Pegasus Dashboard", + "cmd": "PYCBC_SUBMIT_DAX_ARGV", + } save_fig_with_metadata(dashboard_str, dashboard_file.storage_path, **kwds) - wf.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(), - rdir.base)) + wf.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(), rdir.base)) container += workflow container += finalize_workflow @@ -86,7 +102,7 @@ def finalize(container, workflow, finalize_workflow): container.save() - open_box_result_path = rdir['open_box_result'] + open_box_result_path = rdir["open_box_result"] if os.path.exists(open_box_result_path): os.chmod(open_box_result_path, 0o0700) else: @@ -96,7 +112,7 @@ def finalize(container, workflow, finalize_workflow): # Close the log and flush to the html file logging.shutdown() - with open(wf_log_file.storage_path, "r") as logfile: + with open(wf_log_file.storage_path) as logfile: logdata = logfile.read() log_str = """

    Workflow generation script created workflow in output directory: %s

    @@ -104,78 +120,92 @@ def finalize(container, workflow, finalize_workflow):

    Workflow generation script run on host: %s

    %s
    """ % (os.getcwd(), args.workflow_name, socket.gethostname(), logdata) - kwds = {'title': 'Workflow Generation Log', - 'caption': "Log of the workflow script %s" % sys.argv[0], - 'cmd': ' '.join(sys.argv), } + kwds = { + "title": "Workflow Generation Log", + "caption": "Log of the workflow script %s" % sys.argv[0], + "cmd": " ".join(sys.argv), + } save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds) # Add the command line used to a specific file args_to_output = [sys.argv[0]] for arg in sys.argv[1:]: - if arg.startswith('--'): + if arg.startswith("--"): # This is an option, add tab - args_to_output.append(' ' + arg) + args_to_output.append(" " + arg) else: # This is a parameter, add two tabs - args_to_output.append(' ' + arg) - - gen_str = '
    ' + ' \\\n'.join(args_to_output) + '
    ' - kwds = {'title': 'Workflow Generation Command', - 'caption': "Command used to generate the workflow.", - 'cmd': ' '.join(sys.argv), } + args_to_output.append(" " + arg) + + gen_str = "
    " + " \\\n".join(args_to_output) + "
    " + kwds = { + "title": "Workflow Generation Command", + "caption": "Command used to generate the workflow.", + "cmd": " ".join(sys.argv), + } save_fig_with_metadata(gen_str, gen_file_html.storage_path, **kwds) - layout.single_layout(rdir['workflow'], ([dashboard_file, gen_file_html, log_file_html])) + layout.single_layout( + rdir["workflow"], ([dashboard_file, gen_file_html, log_file_html]) + ) sys.exit(0) def check_stop(job_name, container, workflow, finalize_workflow): - #This function will finalize the workflow and stop it at the job specified. - #Under [workflow] supply the option stop-after = and the job name you want the workflow to stop at. Current options are [inspiral, hdf_trigger_merge, statmap] - if workflow.cp.has_option('workflow', 'stop-after'): - stop_after = workflow.cp.get('workflow', 'stop-after') + # This function will finalize the workflow and stop it at the job specified. + # Under [workflow] supply the option stop-after = and the job name you want the workflow to stop at. Current options are [inspiral, hdf_trigger_merge, statmap] + if workflow.cp.has_option("workflow", "stop-after"): + stop_after = workflow.cp.get("workflow", "stop-after") if stop_after == job_name: logging.info("Search worflow will be stopped after " + str(job_name)) finalize(container, workflow, finalize_workflow) - else: + else: pass + parser = argparse.ArgumentParser(description=__doc__[1:]) pycbc.add_common_pycbc_options(parser) wf.add_workflow_command_line_group(parser) wf.add_workflow_settings_cli(parser) args = parser.parse_args() -# Default logging level is info: --verbose adds to this +# Default logging level is info: --verbose adds to this pycbc.init_logging(args.verbose, default_level=1) container = wf.Workflow(args, args.workflow_name) -workflow = wf.Workflow(args, args.workflow_name + '-main') -finalize_workflow = wf.Workflow(args, args.workflow_name + '-finalization') +workflow = wf.Workflow(args, args.workflow_name + "-main") +finalize_workflow = wf.Workflow(args, args.workflow_name + "-finalization") wf.makedir(args.output_dir) os.chdir(args.output_dir) -rdir = layout.SectionNumber('results', ['analysis_time', - 'detector_sensitivity', - 'data_quality', - 'single_triggers', - 'background_triggers', - 'injections', - 'search_sensitivity', - 'open_box_result', - 'workflow', - ]) +rdir = layout.SectionNumber( + "results", + [ + "analysis_time", + "detector_sensitivity", + "data_quality", + "single_triggers", + "background_triggers", + "injections", + "search_sensitivity", + "open_box_result", + "workflow", + ], +) wf.makedir(rdir.base) -wf.makedir(rdir['workflow']) +wf.makedir(rdir["workflow"]) # We are _also_ logging to a file -wf_log_file = wf.File(workflow.ifos, 'workflow-log', workflow.analysis_time, - extension='.txt', - directory=rdir['workflow']) +wf_log_file = wf.File( + workflow.ifos, + "workflow-log", + workflow.analysis_time, + extension=".txt", + directory=rdir["workflow"], +) -pycbc.init_logging(args.verbose, default_level=1, - to_file=wf_log_file.storage_path) +pycbc.init_logging(args.verbose, default_level=1, to_file=wf_log_file.storage_path) logging.info("Created log file %s" % wf_log_file.storage_path) # put start / end time at top of summary page @@ -183,105 +213,137 @@ time = workflow.analysis_time s, e = int(time[0]), int(time[1]) s_utc = gps_to_utc_str(s) e_utc = gps_to_utc_str(e) -time_str = '

    GPS Interval [%s,%s). ' %(s,e) -time_str += 'UTC Interval %s - %s. ' %(s_utc, e_utc) -time_str += 'Interval duration = %.3f days.

    '\ - %(float(e-s)/86400.0,) -time_file = wf.File(workflow.ifos, 'time', workflow.analysis_time, - extension='.html', - directory=rdir.base) -kwds = { 'title' : 'Search Workflow Duration (Wall Clock Time)', - 'caption' : "Wall clock start and end times for this invocation of " - "the workflow. The command line button shows the " - "arguments used to invoke the workflow creation script.", - 'cmd' :' '.join(sys.argv), } +time_str = "

    GPS Interval [%s,%s). " % (s, e) +time_str += "UTC Interval %s - %s. " % (s_utc, e_utc) +time_str += "Interval duration = %.3f days.

    " % ( + float(e - s) / 86400.0, +) +time_file = wf.File( + workflow.ifos, + "time", + workflow.analysis_time, + extension=".html", + directory=rdir.base, +) +kwds = { + "title": "Search Workflow Duration (Wall Clock Time)", + "caption": "Wall clock start and end times for this invocation of " + "the workflow. The command line button shows the " + "arguments used to invoke the workflow creation script.", + "cmd": " ".join(sys.argv), +} save_fig_with_metadata(time_str, time_file.storage_path, **kwds) # Get segments and find the data locations -sci_seg_name = 'science' -science_seg_file = wf.get_segments_file(workflow, sci_seg_name, 'segments-science', - rdir['analysis_time/segment_data'], - tags=['science']) +sci_seg_name = "science" +science_seg_file = wf.get_segments_file( + workflow, + sci_seg_name, + "segments-science", + rdir["analysis_time/segment_data"], + tags=["science"], +) ssegs = {} for ifo in workflow.ifos: ssegs[ifo] = science_seg_file.segment_dict["%s:science" % ifo] -hoft_tags=[] -if 'hoft' in workflow.cp.get_subsections('workflow-datafind'): - hoft_tags=['hoft'] +hoft_tags = [] +if "hoft" in workflow.cp.get_subsections("workflow-datafind"): + hoft_tags = ["hoft"] -datafind_files, analyzable_file, analyzable_segs, analyzable_name = \ - wf.setup_datafind_workflow(workflow, - ssegs, "datafind", - seg_file=science_seg_file, tags=hoft_tags) +datafind_files, analyzable_file, analyzable_segs, analyzable_name = ( + wf.setup_datafind_workflow( + workflow, ssegs, "datafind", seg_file=science_seg_file, tags=hoft_tags + ) +) -final_veto_name = 'vetoes' -final_veto_file = wf.get_segments_file(workflow, final_veto_name, - 'segments-vetoes', - rdir['analysis_time/segment_data'], - tags=['veto']) +final_veto_name = "vetoes" +final_veto_file = wf.get_segments_file( + workflow, + final_veto_name, + "segments-vetoes", + rdir["analysis_time/segment_data"], + tags=["veto"], +) # Get dq segments from veto definer and calculate data quality timeseries -dq_flag_name = 'dq_flag' -dq_segment_file = wf.get_flag_segments_file(workflow, dq_flag_name, - 'segments-dq', - rdir['analysis_time/segment_data'], - tags=['dq']) +dq_flag_name = "dq_flag" +dq_segment_file = wf.get_flag_segments_file( + workflow, + dq_flag_name, + "segments-dq", + rdir["analysis_time/segment_data"], + tags=["dq"], +) # Template bank stuff -hdfbank = wf.setup_tmpltbank_workflow(workflow, analyzable_segs, - datafind_files, output_dir="bank", - return_format='hdf') -assert( len(hdfbank) == 1 ) +hdfbank = wf.setup_tmpltbank_workflow( + workflow, analyzable_segs, datafind_files, output_dir="bank", return_format="hdf" +) +assert len(hdfbank) == 1 hdfbank = hdfbank[0] -splitbank_files_fd = wf.setup_splittable_workflow(workflow, [hdfbank], - out_dir="bank", - tags=['full_data']) +splitbank_files_fd = wf.setup_splittable_workflow( + workflow, [hdfbank], out_dir="bank", tags=["full_data"] +) bank_tags = [] -if 'mass1_mass2' in workflow.cp.get_subsections('plot_bank'): - bank_tags=['mass1_mass2'] +if "mass1_mass2" in workflow.cp.get_subsections("plot_bank"): + bank_tags = ["mass1_mass2"] -bank_plot = wf.make_template_plot(workflow, hdfbank, - rdir['background_triggers'], - tags=bank_tags) +bank_plot = wf.make_template_plot( + workflow, hdfbank, rdir["background_triggers"], tags=bank_tags +) ######################## Setup the FULL DATA run ############################## output_dir = "full_data" # setup the matchedfilter jobs -ind_insps = insps = wf.setup_matchedfltr_workflow(workflow, analyzable_segs, - datafind_files, splitbank_files_fd, - output_dir, tags=['full_data']) +ind_insps = insps = wf.setup_matchedfltr_workflow( + workflow, + analyzable_segs, + datafind_files, + splitbank_files_fd, + output_dir, + tags=["full_data"], +) -#check to see if workflow should stop at inspiral jobs. -check_stop('inspiral', container, workflow, finalize_workflow) +# check to see if workflow should stop at inspiral jobs. +check_stop("inspiral", container, workflow, finalize_workflow) -insps = wf.merge_single_detector_hdf_files(workflow, hdfbank, - insps, output_dir, - tags=['full_data']) +insps = wf.merge_single_detector_hdf_files( + workflow, hdfbank, insps, output_dir, tags=["full_data"] +) -check_stop('hdf_trigger_merge', container, workflow, finalize_workflow) +check_stop("hdf_trigger_merge", container, workflow, finalize_workflow) # setup sngl trigger distribution fitting jobs # 'statfiles' is list of files used in calculating statistic # 'dqfiles' is the subset of files containing data quality information statfiles = [] -dqfiles, dqfile_labels = wf.setup_dq_reranking(workflow, insps, - hdfbank, analyzable_file, - analyzable_name, - dq_segment_file, - output_dir='dq', - tags=['full_data']) +dqfiles, dqfile_labels = wf.setup_dq_reranking( + workflow, + insps, + hdfbank, + analyzable_file, + analyzable_name, + dq_segment_file, + output_dir="dq", + tags=["full_data"], +) statfiles += dqfiles -statfiles += wf.setup_trigger_fitting(workflow, insps, hdfbank, - final_veto_file, final_veto_name, - output_dir=output_dir, - tags=['full_data']) +statfiles += wf.setup_trigger_fitting( + workflow, + insps, + hdfbank, + final_veto_file, + final_veto_name, + output_dir=output_dir, + tags=["full_data"], +) # Set up the multi-ifo coinc jobs # final_bg_files contains coinc results using vetoes final_veto_files @@ -291,62 +353,80 @@ no_fg_exc_files = {} ifo_ids = {} # Get the ifo precedence values -ifo_precedence_list = workflow.cp.get_opt_tags('workflow-coincidence', 'timeslide-precedence', ['full_data']) -for ifo, _ in zip(*insps.categorize_by_attr('ifo')): +ifo_precedence_list = workflow.cp.get_opt_tags( + "workflow-coincidence", "timeslide-precedence", ["full_data"] +) +for ifo, _ in zip(*insps.categorize_by_attr("ifo")): ifo_ids[ifo] = ifo_precedence_list.index(ifo) # Generate the possible detector combinations -if workflow.cp.has_option_tags('workflow-data_quality', 'no-coinc-veto', - tags=None): - logging.info("no-coinc-veto option enabled, " + - "no longer passing veto segments to coinc jobs.") +if workflow.cp.has_option_tags("workflow-data_quality", "no-coinc-veto", tags=None): + logging.info( + "no-coinc-veto option enabled, " + "no longer passing veto segments to coinc jobs." + ) coinc_veto_file = None else: coinc_veto_file = final_veto_file for ifocomb in ifo_combos(ifo_ids.keys()): inspcomb = wf.select_files_by_ifo_combination(ifocomb, insps) - pivot_ifo, fixed_ifo, ordered_ifo_list = wf.get_ordered_ifo_list(ifocomb, - ifo_ids) + pivot_ifo, fixed_ifo, ordered_ifo_list = wf.get_ordered_ifo_list(ifocomb, ifo_ids) # Create coinc tag, and set up the coinc job for the combination - coinctag = '{}det'.format(len(ifocomb)) - ctagcomb = ['full_data', coinctag] + coinctag = f"{len(ifocomb)}det" + ctagcomb = ["full_data", coinctag] bg_file = wf.setup_interval_coinc( - workflow, hdfbank, inspcomb, statfiles, coinc_veto_file, - final_veto_name, output_dir, pivot_ifo, fixed_ifo, tags=ctagcomb) + workflow, + hdfbank, + inspcomb, + statfiles, + coinc_veto_file, + final_veto_name, + output_dir, + pivot_ifo, + fixed_ifo, + tags=ctagcomb, + ) # Optionally perform follow-up on triggers and rerank the candidates # Returns the input file if not enabled. no_fg_exc_files[ordered_ifo_list] = wf.rerank_coinc_followup( - workflow, bg_file, hdfbank, output_dir, - tags=ctagcomb) + workflow, bg_file, hdfbank, output_dir, tags=ctagcomb + ) # Are we analysing single-detector candidates? -analyze_singles = workflow.cp.has_section('workflow-singles') \ - and workflow.cp.has_option_tags('workflow-singles', - 'analyze', tags=None) +analyze_singles = workflow.cp.has_section( + "workflow-singles" +) and workflow.cp.has_option_tags("workflow-singles", "analyze", tags=None) # The single-detector findtrigs and statmap jobs work differently # - set these up here -for ifo in ifo_ids.keys(): +for ifo in ifo_ids: if not analyze_singles: continue inspcomb = wf.select_files_by_ifo_combination([ifo], insps) # Create coinc tag, and set up the findtrigs job for the combination - ctagsngl = ['full_data', '1det'] + ctagsngl = ["full_data", "1det"] no_fg_exc_files[ifo] = wf.setup_sngls( - workflow, hdfbank, inspcomb, statfiles, final_veto_file, - final_veto_name, output_dir, tags=ctagsngl) + workflow, + hdfbank, + inspcomb, + statfiles, + final_veto_file, + final_veto_name, + output_dir, + tags=ctagsngl, + ) -#check to see if workflow should stop at statmap jobs. -check_stop('statmap', container, workflow, finalize_workflow) +# check to see if workflow should stop at statmap jobs. +check_stop("statmap", container, workflow, finalize_workflow) ifo_sets = list(ifo_combos(ifo_ids.keys())) if analyze_singles: - ifo_sets += [(ifo,) for ifo in ifo_ids.keys()] + ifo_sets += [(ifo,) for ifo in ifo_ids] final_bg_files = {} # set up exclude-zerolag jobs for each ifo combination @@ -354,10 +434,10 @@ for ifocomb in ifo_sets: if len(ifocomb) > 1: _, _, ordered_ifo_list = wf.get_ordered_ifo_list(ifocomb, ifo_ids) # Create coinc tag - coinctag = '{}det'.format(len(ifocomb)) + coinctag = f"{len(ifocomb)}det" else: ordered_ifo_list = ifocomb[0] - coinctag= '1det' + coinctag = "1det" if len(ifo_sets) == 1: # Just one event type - pass it through @@ -366,180 +446,225 @@ for ifocomb in ifo_sets: # Combining the statmap files, multiple possible event types other_ifo_keys = list(no_fg_exc_files.keys()) other_ifo_keys.remove(ordered_ifo_list) - ctagcomb = ['full_data', coinctag] - other_bg_files = {ctype: no_fg_exc_files[ctype] - for ctype in other_ifo_keys} + ctagcomb = ["full_data", coinctag] + other_bg_files = {ctype: no_fg_exc_files[ctype] for ctype in other_ifo_keys} final_bg_files[ordered_ifo_list] = wf.setup_exclude_zerolag( workflow, no_fg_exc_files[ordered_ifo_list], wf.FileList(other_bg_files.values()), - output_dir, ordered_ifo_list, - tags=ctagcomb + output_dir, + ordered_ifo_list, + tags=ctagcomb, ) # Sort the IFO combinations: longest (most detectors) first, then alphabetical final_bg_files = { - k: v for k, v in sorted( - final_bg_files.items(), - key=(lambda kv: (-len(kv[0]), kv[0])) - ) + k: v + for k, v in sorted(final_bg_files.items(), key=(lambda kv: (-len(kv[0]), kv[0]))) } if len(ifo_sets) == 1: combined_bg_file = no_fg_exc_files[ordered_ifo_list] else: combined_bg_file = wf.setup_combine_statmap( - workflow, - wf.FileList(final_bg_files.values()), - wf.FileList([]), - output_dir, - tags=['full_data']) + workflow, + wf.FileList(final_bg_files.values()), + wf.FileList([]), + output_dir, + tags=["full_data"], + ) -censored_veto_name = 'closed_box' -censored_veto = wf.make_foreground_censored_veto(workflow, - combined_bg_file, final_veto_file, final_veto_name, - censored_veto_name, 'segments') +censored_veto_name = "closed_box" +censored_veto = wf.make_foreground_censored_veto( + workflow, + combined_bg_file, + final_veto_file, + final_veto_name, + censored_veto_name, + "segments", +) # Calculate the inspiral psds psd_files = [] -trig_generated_name = 'TRIGGERS_GENERATED' +trig_generated_name = "TRIGGERS_GENERATED" trig_generated_segs = {} -data_analysed_name = 'DATA_ANALYSED' +data_analysed_name = "DATA_ANALYSED" data_analysed_segs = {} insp_files_seg_dict = segments.segmentlistdict() -for ifo, files in zip(*ind_insps.categorize_by_attr('ifo')): +for ifo, files in zip(*ind_insps.categorize_by_attr("ifo")): trig_generated_segs[ifo] = segments.segmentlist([f.segment for f in files]) - data_analysed_segs[ifo] = \ - segments.segmentlist([f.metadata['data_seg'] for f in files]) + data_analysed_segs[ifo] = segments.segmentlist( + [f.metadata["data_seg"] for f in files] + ) # Remove duplicates from splitbank - trig_generated_segs[ifo] = \ - segments.segmentlist(set(trig_generated_segs[ifo])) - data_analysed_segs[ifo] = \ - segments.segmentlist(set(data_analysed_segs[ifo])) + trig_generated_segs[ifo] = segments.segmentlist(set(trig_generated_segs[ifo])) + data_analysed_segs[ifo] = segments.segmentlist(set(data_analysed_segs[ifo])) - insp_files_seg_dict[ifo + ":" + trig_generated_name] = \ - trig_generated_segs[ifo] - insp_files_seg_dict[ifo + ":" + data_analysed_name] = \ - data_analysed_segs[ifo] + insp_files_seg_dict[ifo + ":" + trig_generated_name] = trig_generated_segs[ifo] + insp_files_seg_dict[ifo + ":" + data_analysed_name] = data_analysed_segs[ifo] if datafind_files: frame_files = datafind_files.find_output_with_ifo(ifo) else: frame_files = None - psd_files += [wf.setup_psd_calculate(workflow, frame_files, ifo, - data_analysed_segs[ifo], data_analysed_name, 'psds')] - -insp_files_seg_file = wf.SegFile.from_segment_list_dict('INSP_SEGMENTS', - insp_files_seg_dict, valid_segment=workflow.analysis_time, - extension='xml', directory=rdir['analysis_time/segment_data']) + psd_files += [ + wf.setup_psd_calculate( + workflow, + frame_files, + ifo, + data_analysed_segs[ifo], + data_analysed_name, + "psds", + ) + ] + +insp_files_seg_file = wf.SegFile.from_segment_list_dict( + "INSP_SEGMENTS", + insp_files_seg_dict, + valid_segment=workflow.analysis_time, + extension="xml", + directory=rdir["analysis_time/segment_data"], +) ################### Range, spectrum and segments plots ####################### -s = wf.make_spectrum_plot(workflow, psd_files, rdir['detector_sensitivity']) -r = wf.make_range_plot(workflow, psd_files, rdir['detector_sensitivity'], - require='summ') -r2 = wf.make_range_plot(workflow, psd_files, rdir['detector_sensitivity'], - exclude='summ') +s = wf.make_spectrum_plot(workflow, psd_files, rdir["detector_sensitivity"]) +r = wf.make_range_plot( + workflow, psd_files, rdir["detector_sensitivity"], require="summ" +) +r2 = wf.make_range_plot( + workflow, psd_files, rdir["detector_sensitivity"], exclude="summ" +) det_summ = [(s, r[0] if len(r) != 0 else None)] -layout.two_column_layout(rdir['detector_sensitivity'], - det_summ + list(layout.grouper(r2, 2))) +layout.two_column_layout( + rdir["detector_sensitivity"], det_summ + list(layout.grouper(r2, 2)) +) # do plotting of segments / veto times -wf.make_segments_plot(workflow, [insp_files_seg_file], - rdir['analysis_time/segments'], - tags=[trig_generated_name]) -wf.make_gating_plot(workflow, full_insps, rdir['analysis_time/gating'], - tags=['full_data']) +wf.make_segments_plot( + workflow, + [insp_files_seg_file], + rdir["analysis_time/segments"], + tags=[trig_generated_name], +) +wf.make_gating_plot( + workflow, full_insps, rdir["analysis_time/gating"], tags=["full_data"] +) # make segment table and plot for summary page -curr_files = [science_seg_file, analyzable_file, - insp_files_seg_file] -curr_names = [sci_seg_name, analyzable_name, - trig_generated_name] -seg_summ_table = wf.make_seg_table\ - (workflow, curr_files, curr_names, rdir['analysis_time/segments'], - ['SUMMARY'], title_text='Input and output', - description='This shows the total amount of input data, analyzable data, ' - 'and the time for which triggers are produced.') -seg_summ_plot = wf.make_seg_plot(workflow, curr_files, - rdir['analysis_time/segments'], - curr_names, ['SUMMARY']) +curr_files = [science_seg_file, analyzable_file, insp_files_seg_file] +curr_names = [sci_seg_name, analyzable_name, trig_generated_name] +seg_summ_table = wf.make_seg_table( + workflow, + curr_files, + curr_names, + rdir["analysis_time/segments"], + ["SUMMARY"], + title_text="Input and output", + description="This shows the total amount of input data, analyzable data, " + "and the time for which triggers are produced.", +) +seg_summ_plot = wf.make_seg_plot( + workflow, curr_files, rdir["analysis_time/segments"], curr_names, ["SUMMARY"] +) curr_files = [insp_files_seg_file] + [final_veto_file] # Add in singular veto files curr_files = curr_files + [science_seg_file] -curr_names = [trig_generated_name + '&' + final_veto_name] +curr_names = [trig_generated_name + "&" + final_veto_name] # And SCIENCE - CAT 1 vetoes explicitly. -curr_names += [sci_seg_name + '&' + 'VETO_CAT1'] +curr_names += [sci_seg_name + "&" + "VETO_CAT1"] -veto_summ_table = wf.make_seg_table\ - (workflow, curr_files, curr_names, rdir['analysis_time/segments'], - ['VETO_SUMMARY'], title_text='Time removed by vetoes', - description='This shows the time removed from the output time by the ' - 'vetoes applied to the triggers.') +veto_summ_table = wf.make_seg_table( + workflow, + curr_files, + curr_names, + rdir["analysis_time/segments"], + ["VETO_SUMMARY"], + title_text="Time removed by vetoes", + description="This shows the time removed from the output time by the " + "vetoes applied to the triggers.", +) # make veto definer table -vetodef_table = wf.make_veto_table(workflow, rdir['analysis_time/veto_definer']) +vetodef_table = wf.make_veto_table(workflow, rdir["analysis_time/veto_definer"]) if vetodef_table is not None: - layout.single_layout(rdir['analysis_time/veto_definer'], ([vetodef_table])) + layout.single_layout(rdir["analysis_time/veto_definer"], ([vetodef_table])) #################### Plotting on FULL_DATA results ########################## ##################### SINGLES plots first ################################### -snrchi = wf.make_snrchi_plot(workflow, insps, censored_veto, - 'closed_box', rdir['single_triggers'], - tags=['full_data']) -layout.group_layout(rdir['single_triggers'], snrchi) +snrchi = wf.make_snrchi_plot( + workflow, + insps, + censored_veto, + "closed_box", + rdir["single_triggers"], + tags=["full_data"], +) +layout.group_layout(rdir["single_triggers"], snrchi) hist_summ = [] for insp in full_insps: - outdir = rdir['single_triggers/%s_binned_triggers' % insp.ifo] + outdir = rdir["single_triggers/%s_binned_triggers" % insp.ifo] singles_plots = wf.make_singles_plot( workflow, [insp], hdfbank, censored_veto, - 'closed_box', + "closed_box", outdir, - tags=['full_data'] + tags=["full_data"], ) layout.group_layout(outdir, singles_plots) # make non-summary hists using the bank file # currently, none of these are made - outdir = rdir['single_triggers/%s_trigger_histograms' % insp.ifo] - wf.make_single_hist(workflow, insp, censored_veto, 'closed_box', outdir, - bank_file=hdfbank, exclude='summ', - tags=['full_data']) + outdir = rdir["single_triggers/%s_trigger_histograms" % insp.ifo] + wf.make_single_hist( + workflow, + insp, + censored_veto, + "closed_box", + outdir, + bank_file=hdfbank, + exclude="summ", + tags=["full_data"], + ) # make summary hists for all templates together # currently, 2 per ifo: snr and newsnr - outdir = rdir['single_triggers/%s_trigger_histograms' % insp.ifo] - allhists = wf.make_single_hist(workflow, insp, censored_veto, 'closed_box', - outdir, require='summ', tags=['full_data']) + outdir = rdir["single_triggers/%s_trigger_histograms" % insp.ifo] + allhists = wf.make_single_hist( + workflow, + insp, + censored_veto, + "closed_box", + outdir, + require="summ", + tags=["full_data"], + ) layout.group_layout(outdir, allhists) # make hists of newsnr split up by parameter # currently, 1 per ifo split by template duration - outdir = rdir['single_triggers/%s_binned_histograms' % insp.ifo] - binhists = wf.make_binned_hist(workflow, insp, censored_veto, - 'closed_box', outdir, hdfbank, - tags=['full_data']) + outdir = rdir["single_triggers/%s_binned_histograms" % insp.ifo] + binhists = wf.make_binned_hist( + workflow, insp, censored_veto, "closed_box", outdir, hdfbank, tags=["full_data"] + ) layout.group_layout(outdir, binhists) # put raw SNR and binned newsnr hist in summary hist_summ += list(layout.grouper([allhists[0], binhists[0]], 2)) -if workflow.cp.has_option_tags('workflow-matchedfilter', - 'plot-throughput', tags=['full_data']): +if workflow.cp.has_option_tags( + "workflow-matchedfilter", "plot-throughput", tags=["full_data"] +): throughput_plot = wf.make_throughput_plot( - workflow, - full_insps, - rdir['workflow/throughput'], - tags=['full_data'] + workflow, full_insps, rdir["workflow/throughput"], tags=["full_data"] ) - layout.single_layout(rdir['workflow/throughput'], [throughput_plot]) + layout.single_layout(rdir["workflow/throughput"], [throughput_plot]) # Run minifollowups on loudest sngl detector events excl_subsecs = set([]) @@ -550,80 +675,134 @@ for insp_file in full_insps: for insp_file in full_insps: curr_ifo = insp_file.ifo - for subsec in workflow.cp.get_subsections('workflow-sngl_minifollowups'): + for subsec in workflow.cp.get_subsections("workflow-sngl_minifollowups"): if subsec in excl_subsecs: continue - sec_name = 'workflow-sngl_minifollowups-{}'.format(subsec) - dir_str = workflow.cp.get(sec_name, 'section-header') - currdir = rdir['single_triggers/{}_{}'.format(curr_ifo, dir_str)] - wf.setup_single_det_minifollowups\ - (workflow, insp_file, hdfbank, insp_files_seg_file, - data_analysed_name, trig_generated_name, 'daxes', currdir, - statfiles=wf.FileList(statfiles), - fg_file=censored_veto, fg_name='closed_box', - tags=insp_file.tags + [subsec]) + sec_name = f"workflow-sngl_minifollowups-{subsec}" + dir_str = workflow.cp.get(sec_name, "section-header") + currdir = rdir[f"single_triggers/{curr_ifo}_{dir_str}"] + wf.setup_single_det_minifollowups( + workflow, + insp_file, + hdfbank, + insp_files_seg_file, + data_analysed_name, + trig_generated_name, + "daxes", + currdir, + statfiles=wf.FileList(statfiles), + fg_file=censored_veto, + fg_name="closed_box", + tags=insp_file.tags + [subsec], + ) ##################### COINC FULL_DATA plots ################################### # Main results with combined file (we mix open and closed box here, but # separate them in the result page) -ifar_ob = wf.make_ifar_plot(workflow, combined_bg_file, - rdir['open_box_result'], - tags=combined_bg_file.tags + ['open_box'], - executable='page_ifar_catalog') +ifar_ob = wf.make_ifar_plot( + workflow, + combined_bg_file, + rdir["open_box_result"], + tags=combined_bg_file.tags + ["open_box"], + executable="page_ifar_catalog", +) -table = wf.make_foreground_table(workflow, combined_bg_file, - hdfbank, rdir['open_box_result'], - singles=insps, extension='.html', - tags=combined_bg_file.tags) +table = wf.make_foreground_table( + workflow, + combined_bg_file, + hdfbank, + rdir["open_box_result"], + singles=insps, + extension=".html", + tags=combined_bg_file.tags, +) -fore_xmlall = wf.make_foreground_table(workflow, combined_bg_file, - hdfbank, rdir['open_box_result'], singles=insps, - extension='.xml', tags=["xmlall"]) +fore_xmlall = wf.make_foreground_table( + workflow, + combined_bg_file, + hdfbank, + rdir["open_box_result"], + singles=insps, + extension=".xml", + tags=["xmlall"], +) -if workflow.cp.has_option_tags('workflow-minifollowups', - 'prepare-gracedb-uploads', ''): +if workflow.cp.has_option_tags("workflow-minifollowups", "prepare-gracedb-uploads", ""): # Need to get absolute path here - upload_path = os.path.join(workflow.out_dir, 'upload_data') - wf.setup_upload_prep_minifollowups(workflow, combined_bg_file, fore_xmlall, - full_insps, psd_files, hdfbank, insp_files_seg_file, - data_analysed_name, trig_generated_name, - 'daxes', upload_path, - tags=combined_bg_file.tags) + upload_path = os.path.join(workflow.out_dir, "upload_data") + wf.setup_upload_prep_minifollowups( + workflow, + combined_bg_file, + fore_xmlall, + full_insps, + psd_files, + hdfbank, + insp_files_seg_file, + data_analysed_name, + trig_generated_name, + "daxes", + upload_path, + tags=combined_bg_file.tags, + ) -fore_xmlloudest = wf.make_foreground_table(workflow, combined_bg_file, - hdfbank, rdir['open_box_result'], singles=insps, - extension='.xml', tags=["xmlloudest"]) +fore_xmlloudest = wf.make_foreground_table( + workflow, + combined_bg_file, + hdfbank, + rdir["open_box_result"], + singles=insps, + extension=".xml", + tags=["xmlloudest"], +) -symlink_result(table, 'open_box_result/significance') +symlink_result(table, "open_box_result/significance") # Set html pages -main_page = [(ifar_ob,), (table, )] -layout.two_column_layout(rdir['open_box_result'], main_page) +main_page = [(ifar_ob,), (table,)] +layout.two_column_layout(rdir["open_box_result"], main_page) # run minifollowups on the output of the loudest events -mfup_dir_fg = rdir['open_box_result/loudest_events_followup'] -mfup_dir_bg = rdir['background_triggers/loudest_background_followup'] -wf.setup_foreground_minifollowups(workflow, combined_bg_file, - full_insps, hdfbank, insp_files_seg_file, - data_analysed_name, trig_generated_name, - 'daxes', mfup_dir_fg, - tags=combined_bg_file.tags + ['foreground']) - -wf.setup_foreground_minifollowups(workflow, combined_bg_file, - full_insps, hdfbank, insp_files_seg_file, - data_analysed_name, trig_generated_name, - 'daxes', mfup_dir_bg, - tags=combined_bg_file.tags + ['background']) +mfup_dir_fg = rdir["open_box_result/loudest_events_followup"] +mfup_dir_bg = rdir["background_triggers/loudest_background_followup"] +wf.setup_foreground_minifollowups( + workflow, + combined_bg_file, + full_insps, + hdfbank, + insp_files_seg_file, + data_analysed_name, + trig_generated_name, + "daxes", + mfup_dir_fg, + tags=combined_bg_file.tags + ["foreground"], +) + +wf.setup_foreground_minifollowups( + workflow, + combined_bg_file, + full_insps, + hdfbank, + insp_files_seg_file, + data_analysed_name, + trig_generated_name, + "daxes", + mfup_dir_bg, + tags=combined_bg_file.tags + ["background"], +) # Far vs stat for all ifo combinations in a single plot -farstat = wf.make_farstat_plot(workflow, final_bg_files, rdir['background_triggers'], - require='summ', - tags=bg_file.tags + ['closed']) +farstat = wf.make_farstat_plot( + workflow, + final_bg_files, + rdir["background_triggers"], + require="summ", + tags=bg_file.tags + ["closed"], +) closed_page = [(bank_plot, farstat)] -layout.two_column_layout(rdir['background_triggers'], closed_page) +layout.two_column_layout(rdir["background_triggers"], closed_page) # Add farstat plot to summary page snrifar_summ = [(farstat,)] @@ -631,33 +810,39 @@ snrifar_summ = [(farstat,)] # Sub-pages for each ifo combination for key in final_bg_files: bg_file = final_bg_files[key] - open_dir = rdir['open_box_result/{}_candidates'.format(key)] - closed_dir = rdir['background_triggers/{}_background'.format(key)] - snrifar = wf.make_snrifar_plot(workflow, bg_file, open_dir, - tags=bg_file.tags) - snrifar_cb = wf.make_snrifar_plot(workflow, bg_file, closed_dir, - closed_box=True, - tags=bg_file.tags + ['closed']) - ratehist = wf.make_snrratehist_plot(workflow, bg_file, open_dir, - tags=bg_file.tags) - snrifar_ifar = wf.make_snrifar_plot(workflow, bg_file, open_dir, - cumulative=False, - tags=bg_file.tags + ['ifar']) - ifar_ob = wf.make_ifar_plot(workflow, bg_file, open_dir, - tags=bg_file.tags + ['open_box']) + open_dir = rdir[f"open_box_result/{key}_candidates"] + closed_dir = rdir[f"background_triggers/{key}_background"] + snrifar = wf.make_snrifar_plot(workflow, bg_file, open_dir, tags=bg_file.tags) + snrifar_cb = wf.make_snrifar_plot( + workflow, bg_file, closed_dir, closed_box=True, tags=bg_file.tags + ["closed"] + ) + ratehist = wf.make_snrratehist_plot(workflow, bg_file, open_dir, tags=bg_file.tags) + snrifar_ifar = wf.make_snrifar_plot( + workflow, bg_file, open_dir, cumulative=False, tags=bg_file.tags + ["ifar"] + ) + ifar_ob = wf.make_ifar_plot( + workflow, bg_file, open_dir, tags=bg_file.tags + ["open_box"] + ) if len(key) > 2: # don't do the background plot for single-detector stuff, # as it is just blank - ifar_cb = wf.make_ifar_plot(workflow, bg_file, closed_dir, - tags=bg_file.tags + ['closed_box']) + ifar_cb = wf.make_ifar_plot( + workflow, bg_file, closed_dir, tags=bg_file.tags + ["closed_box"] + ) closed_page = [(snrifar_cb, ifar_cb)] else: closed_page = [(snrifar_cb,)] if len(ifo_sets) > 1: # If there is only one ifo_set, then this table has already been made - table = wf.make_foreground_table(workflow, bg_file, hdfbank, open_dir, - singles=insps, extension='.html', - tags=bg_file.tags) + table = wf.make_foreground_table( + workflow, + bg_file, + hdfbank, + open_dir, + singles=insps, + extension=".html", + tags=bg_file.tags, + ) detailed_page = [(snrifar, ratehist), (snrifar_ifar, ifar_ob), (table,)] layout.two_column_layout(open_dir, detailed_page) @@ -668,20 +853,13 @@ for key in final_bg_files: # DQ log likelihood plots -for dqf, dql in sorted( - zip(dqfiles, dqfile_labels), - key=lambda dqfl: dqfl[0].ifo -): +for dqf, dql in sorted(zip(dqfiles, dqfile_labels), key=lambda dqfl: dqfl[0].ifo): ifo = dqf.ifo - dq_dir = rdir[f'data_quality/{dqf.ifo}_DQ_results'] + dq_dir = rdir[f"data_quality/{dqf.ifo}_DQ_results"] # plot rates when flag was on trig_rate_plot = wf.make_dq_flag_trigger_rate_plot( - workflow, - dqf, - dql, - dq_dir, - tags=[dql] + workflow, dqf, dql, dq_dir, tags=[dql] ) # make table of dq segment info @@ -689,13 +867,14 @@ for dqf, dql in sorted( # plot background bins background_bins = workflow.cp.get_opt_tags( - 'bin_templates', 'background-bins', tags=[ifo]) + "bin_templates", "background-bins", tags=[ifo] + ) bbin_plot = wf.make_template_plot( workflow, hdfbank, dq_dir, bins=background_bins, - tags=[ifo, 'dq_bins'] + bank_tags + tags=[ifo, "dq_bins"] + bank_tags, ) # make table of template bin info @@ -705,32 +884,33 @@ for dqf, dql in sorted( ############################## Setup the injection runs ####################### -splitbank_files_inj = wf.setup_splittable_workflow(workflow, [hdfbank], - out_dir="bank", - tags=['injections']) +splitbank_files_inj = wf.setup_splittable_workflow( + workflow, [hdfbank], out_dir="bank", tags=["injections"] +) # setup the injection files -inj_files_base, inj_tags = wf.setup_injection_workflow(workflow, - output_dir="inj_files") +inj_files_base, inj_tags = wf.setup_injection_workflow(workflow, output_dir="inj_files") inj_files = [] for inj_file, tag in zip(inj_files_base, inj_tags): - inj_files.append(wf.inj_to_hdf(workflow, inj_file, 'inj_files', [tag])) + inj_files.append(wf.inj_to_hdf(workflow, inj_file, "inj_files", [tag])) inj_coincs = wf.FileList() -found_inj_dict ={} +found_inj_dict = {} insps_dict = {} files_for_combined_injfind = [] for inj_file, tag in zip(inj_files, inj_tags): - ctags = [tag, 'injections'] - output_dir = '%s_coinc' % tag + ctags = [tag, "injections"] + output_dir = "%s_coinc" % tag - if workflow.cp.has_option_tags('workflow-injections', - 'compute-optimal-snr', tags=ctags): + if workflow.cp.has_option_tags( + "workflow-injections", "compute-optimal-snr", tags=ctags + ): optimal_snr_file = wf.compute_inj_optimal_snr( - workflow, inj_file, psd_files, 'inj_files', tags=ctags) + workflow, inj_file, psd_files, "inj_files", tags=ctags + ) file_for_injfind = optimal_snr_file else: file_for_injfind = inj_file @@ -738,23 +918,30 @@ for inj_file, tag in zip(inj_files, inj_tags): files_for_combined_injfind.append(file_for_injfind) # setup the matchedfilter jobs - insps = wf.setup_matchedfltr_workflow(workflow, analyzable_segs, - datafind_files, splitbank_files_inj, - output_dir, tags=ctags, - injection_file=inj_file) + insps = wf.setup_matchedfltr_workflow( + workflow, + analyzable_segs, + datafind_files, + splitbank_files_inj, + output_dir, + tags=ctags, + injection_file=inj_file, + ) - insps = wf.merge_single_detector_hdf_files(workflow, hdfbank, - insps, output_dir, tags=ctags) + insps = wf.merge_single_detector_hdf_files( + workflow, hdfbank, insps, output_dir, tags=ctags + ) # coincidence for injections inj_coinc = {} for ifocomb in ifo_combos(ifo_ids.keys()): inspcomb = wf.select_files_by_ifo_combination(ifocomb, insps) - pivot_ifo, fixed_ifo, ordered_ifo_list = \ - wf.get_ordered_ifo_list(ifocomb, ifo_ids) + pivot_ifo, fixed_ifo, ordered_ifo_list = wf.get_ordered_ifo_list( + ifocomb, ifo_ids + ) # Create coinc tag, and set up the coinc job for the combination - coinctag = '{}det'.format(len(ifocomb)) - ctagcomb = [tag, 'injections', coinctag] + coinctag = f"{len(ifocomb)}det" + ctagcomb = [tag, "injections", coinctag] curr_out = wf.setup_interval_coinc_inj( workflow, hdfbank, @@ -766,7 +953,7 @@ for inj_file, tag in zip(inj_files, inj_tags): output_dir, pivot_ifo, fixed_ifo, - tags=ctagcomb + tags=ctagcomb, ) # Rerank events @@ -777,18 +964,18 @@ for inj_file, tag in zip(inj_files, inj_tags): output_dir, tags=ctagcomb, injection_file=inj_file, - ranking_file=final_bg_files[ordered_ifo_list] + ranking_file=final_bg_files[ordered_ifo_list], ) inj_coinc[ordered_ifo_list] = curr_out # get sngls for injections - for ifo in ifo_ids.keys(): + for ifo in ifo_ids: if not analyze_singles: continue inspcomb = wf.select_files_by_ifo_combination([ifo], insps) # Create sngls tag, and set up the findtrigs job for the combination - ctagsngl = [tag, 'injections', '1det'] + ctagsngl = [tag, "injections", "1det"] inj_coinc[ifo] = wf.setup_sngls_inj( workflow, hdfbank, @@ -798,10 +985,10 @@ for inj_file, tag in zip(inj_files, inj_tags): final_veto_file, final_veto_name, output_dir, - tags=ctagsngl + tags=ctagsngl, ) - combctags = [tag, 'injections'] + combctags = [tag, "injections"] final_inj_bg_file_list = wf.FileList(inj_coinc.values()) combined_inj_bg_file = wf.setup_combine_statmap( @@ -809,7 +996,7 @@ for inj_file, tag in zip(inj_files, inj_tags): final_inj_bg_file_list, wf.FileList(final_bg_files.values()), output_dir, - tags=combctags + tags=combctags, ) found_inj = wf.find_injections_in_hdf_coinc( @@ -819,7 +1006,8 @@ for inj_file, tag in zip(inj_files, inj_tags): censored_veto, censored_veto_name, output_dir, - tags=combctags) + tags=combctags, + ) inj_coincs += [combined_inj_bg_file] @@ -835,8 +1023,8 @@ if len(files_for_combined_injfind) > 0: files_for_combined_injfind, censored_veto, censored_veto_name, - 'allinj', - tags=['all','injections'] + "allinj", + tags=["all", "injections"], ) ############################ Injection plots ################################# @@ -844,21 +1032,23 @@ if len(files_for_combined_injfind) > 0: for inj_file, tag in zip(inj_files, inj_tags): found_inj = found_inj_dict[tag] insps = insps_dict[tag] - injdir = rdir['injections/%s' % tag] - sensdir = rdir['search_sensitivity/%s' % tag] - - #foundmissed/sensitivity plots - s = wf.make_sensitivity_plot(workflow, found_inj, sensdir, - exclude=['all', 'summ'], require='sub', - tags=[tag]) - f = wf.make_foundmissed_plot(workflow, found_inj, injdir, - exclude=['all', 'summ'], require='sub', - tags=[tag]) - - found_table = wf.make_inj_table(workflow, found_inj, injdir, - singles=insps, tags=[tag, 'found']) - missed_table = wf.make_inj_table(workflow, found_inj, injdir, - missed=True, tags=[tag, 'missed']) + injdir = rdir["injections/%s" % tag] + sensdir = rdir["search_sensitivity/%s" % tag] + + # foundmissed/sensitivity plots + s = wf.make_sensitivity_plot( + workflow, found_inj, sensdir, exclude=["all", "summ"], require="sub", tags=[tag] + ) + f = wf.make_foundmissed_plot( + workflow, found_inj, injdir, exclude=["all", "summ"], require="sub", tags=[tag] + ) + + found_table = wf.make_inj_table( + workflow, found_inj, injdir, singles=insps, tags=[tag, "found"] + ) + missed_table = wf.make_inj_table( + workflow, found_inj, injdir, missed=True, tags=[tag, "missed"] + ) for ifocomb in ifo_combos(ifo_ids.keys()): inspcomb = wf.select_files_by_ifo_combination(ifocomb, insps) @@ -868,79 +1058,109 @@ for inj_file, tag in zip(inj_files, inj_tags): for inj_insp, trig_insp in zip(inspcomb, fdinspcmb): final_bg_file = final_bg_files[ordered_ifo_list] curr_tags = [tag, ordered_ifo_list] - f += wf.make_coinc_snrchi_plot(workflow, found_inj, inj_insp, - final_bg_file, trig_insp, - injdir, tags=curr_tags) + f += wf.make_coinc_snrchi_plot( + workflow, + found_inj, + inj_insp, + final_bg_file, + trig_insp, + injdir, + tags=curr_tags, + ) # Make pages from plots inj_layout = list(layout.grouper(f, 2)) + [(found_table,), (missed_table,)] - layout.two_column_layout(rdir['injections/%s' % tag], inj_layout) + layout.two_column_layout(rdir["injections/%s" % tag], inj_layout) if len(s) > 0: - layout.group_layout(rdir['search_sensitivity/%s' % tag], s) + layout.group_layout(rdir["search_sensitivity/%s" % tag], s) # Minifollowups - curr_dir_nam = 'injections/followup_of_{}'.format(tag) - if workflow.cp.has_option_tags('workflow-injection_minifollowups', - 'subsection-suffix', tags=[tag]): - suf_str = workflow.cp.get_opt_tags('workflow-injection_minifollowups', - 'subsection-suffix', tags=[tag]) - curr_dir_nam += '_' + suf_str + curr_dir_nam = f"injections/followup_of_{tag}" + if workflow.cp.has_option_tags( + "workflow-injection_minifollowups", "subsection-suffix", tags=[tag] + ): + suf_str = workflow.cp.get_opt_tags( + "workflow-injection_minifollowups", "subsection-suffix", tags=[tag] + ) + curr_dir_nam += "_" + suf_str currdir = rdir[curr_dir_nam] - wf.setup_injection_minifollowups(workflow, found_inj, inj_file, - insps, hdfbank, insp_files_seg_file, - data_analysed_name, trig_generated_name, - 'daxes', currdir, tags=[tag]) + wf.setup_injection_minifollowups( + workflow, + found_inj, + inj_file, + insps, + hdfbank, + insp_files_seg_file, + data_analysed_name, + trig_generated_name, + "daxes", + currdir, + tags=[tag], + ) # If option given, make throughput plots - if workflow.cp.has_option_tags('workflow-matchedfilter', - 'plot-throughput', tags=[tag]): + if workflow.cp.has_option_tags( + "workflow-matchedfilter", "plot-throughput", tags=[tag] + ): # These will show up under the "expand all" button of the page wf.make_throughput_plot( - workflow, - insps, - rdir['workflow/throughput'], - tags=[tag] + workflow, insps, rdir["workflow/throughput"], tags=[tag] ) ######################## Make combined injection plots ########################## if len(files_for_combined_injfind) > 0: - sen_all = wf.make_sensitivity_plot(workflow, found_inj_comb, - rdir['search_sensitivity'], - require='all', - exclude="summ") - inj_all = wf.make_foundmissed_plot(workflow, found_inj_comb, - rdir['injections'], - require='all', - exclude="summ") + sen_all = wf.make_sensitivity_plot( + workflow, + found_inj_comb, + rdir["search_sensitivity"], + require="all", + exclude="summ", + ) + inj_all = wf.make_foundmissed_plot( + workflow, found_inj_comb, rdir["injections"], require="all", exclude="summ" + ) # Make summary page foundmissed and sensitivity plot - sen_s = wf.make_sensitivity_plot(workflow, found_inj_comb, - rdir['search_sensitivity'], - require='summ') - inj_s = wf.make_foundmissed_plot(workflow, found_inj_comb, - rdir['injections'], - require='summ') + sen_s = wf.make_sensitivity_plot( + workflow, found_inj_comb, rdir["search_sensitivity"], require="summ" + ) + inj_s = wf.make_foundmissed_plot( + workflow, found_inj_comb, rdir["injections"], require="summ" + ) inj_summ = list(layout.grouper(inj_s + sen_s, 2)) - layout.group_layout(rdir['injections'], inj_s + inj_all) - layout.group_layout(rdir['search_sensitivity'], sen_s + sen_all) + layout.group_layout(rdir["injections"], inj_s + inj_all) + layout.group_layout(rdir["search_sensitivity"], sen_s + sen_all) # Make analysis time summary analysis_time_summ = [time_file, seg_summ_plot] for f in analysis_time_summ: - symlink_result(f, 'analysis_time') -layout.single_layout(rdir['analysis_time'], (analysis_time_summ)) + symlink_result(f, "analysis_time") +layout.single_layout(rdir["analysis_time"], (analysis_time_summ)) ########################## Make full summary #################################### if len(files_for_combined_injfind) > 0: - summ = ([(time_file,)] + [(seg_summ_plot,)] + - [(seg_summ_table, veto_summ_table)] + det_summ + hist_summ + - [(bank_plot,)] + inj_summ + snrifar_summ) + summ = ( + [(time_file,)] + + [(seg_summ_plot,)] + + [(seg_summ_table, veto_summ_table)] + + det_summ + + hist_summ + + [(bank_plot,)] + + inj_summ + + snrifar_summ + ) else: - summ = ([(time_file,)] + [(seg_summ_plot,)] + - [(seg_summ_table, veto_summ_table)] + det_summ + hist_summ + - [(bank_plot,)] + snrifar_summ) + summ = ( + [(time_file,)] + + [(seg_summ_plot,)] + + [(seg_summ_table, veto_summ_table)] + + det_summ + + hist_summ + + [(bank_plot,)] + + snrifar_summ + ) for row in summ: for f in row: @@ -948,22 +1168,30 @@ for row in summ: layout.two_column_layout(rdir.base, summ) # Save global config file to results directory -base = rdir['workflow/configuration'] +base = rdir["workflow/configuration"] wf.makedir(base) -ini_file_path = os.path.join(base, 'configuration.ini') -with open(ini_file_path, 'w') as ini_fh: +ini_file_path = os.path.join(base, "configuration.ini") +with open(ini_file_path, "w") as ini_fh: container.cp.write(ini_fh) -ini_file = wf.FileList([wf.File(workflow.ifos, '', workflow.analysis_time, - file_url='file://' + ini_file_path)]) +ini_file = wf.FileList( + [ + wf.File( + workflow.ifos, + "", + workflow.analysis_time, + file_url="file://" + ini_file_path, + ) + ] +) layout.single_layout(base, ini_file) # Create versioning information _, version_page = wf.make_versioning_page( workflow, container.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) -layout.single_layout(rdir['workflow/version'], version_page) +layout.single_layout(rdir["workflow/version"], version_page) ############################ Finalization #################################### finalize(container, workflow, finalize_workflow) diff --git a/bin/workflows/pycbc_make_psd_estimation_workflow b/bin/workflows/pycbc_make_psd_estimation_workflow index b09132069a1..b6f4070adb2 100644 --- a/bin/workflows/pycbc_make_psd_estimation_workflow +++ b/bin/workflows/pycbc_make_psd_estimation_workflow @@ -16,12 +16,14 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. -"""Program for setting up a workflow which estimates the average PSD of a given -portion of strain data.""" +""" +Program for setting up a workflow which estimates the average PSD of a given +portion of strain data. +""" -import os.path import argparse import logging +import os.path import sys import igwn_segments as _segments @@ -29,7 +31,6 @@ import igwn_segments as _segments import pycbc import pycbc.workflow from pycbc.results import save_fig_with_metadata, two_column_layout -import pycbc.workflow from pycbc.time import gps_to_utc_str parser = argparse.ArgumentParser(description=__doc__) @@ -43,11 +44,10 @@ pycbc.init_logging(args.verbose, default_level=1) # FIXME: opts.tags is currently unused here. container = pycbc.workflow.Workflow(args) -workflow = pycbc.workflow.Workflow(args, subworkflow_name='main', - is_subworkflow=True) -finalize_workflow = pycbc.workflow.Workflow(args, - subworkflow_name='finalization', - is_subworkflow=True) +workflow = pycbc.workflow.Workflow(args, subworkflow_name="main", is_subworkflow=True) +finalize_workflow = pycbc.workflow.Workflow( + args, subworkflow_name="finalization", is_subworkflow=True +) pycbc.workflow.makedir(args.output_dir) os.chdir(args.output_dir) @@ -60,40 +60,50 @@ time = workflow.analysis_time s, e = int(time[0]), int(time[1]) s_utc = gps_to_utc_str(s) e_utc = gps_to_utc_str(e) -time_str = '

    GPS Interval [%s,%s). UTC Interval %s - %s. Interval duration = %.3f days.

    ' % (s, e, s_utc, e_utc, float(e-s)/86400.0) -time_file = pycbc.workflow.File(workflow.ifos, 'time', workflow.analysis_time, - extension='.html', - directory='plots') -pycbc.workflow.makedir('plots') -kwds = { 'title' : 'Search Workflow Duration (Wall Clock Time)', - 'caption' : "Wall clock start and end times for this invocation of the workflow. " - " The command line button shows the arguments used to invoke the workflow " - " creation script.", - 'cmd' :' '.join(sys.argv), } +time_str = ( + "

    GPS Interval [%s,%s). UTC Interval %s - %s. Interval duration = %.3f days.

    " + % (s, e, s_utc, e_utc, float(e - s) / 86400.0) +) +time_file = pycbc.workflow.File( + workflow.ifos, "time", workflow.analysis_time, extension=".html", directory="plots" +) +pycbc.workflow.makedir("plots") +kwds = { + "title": "Search Workflow Duration (Wall Clock Time)", + "caption": "Wall clock start and end times for this invocation of the workflow. " + " The command line button shows the arguments used to invoke the workflow " + " creation script.", + "cmd": " ".join(sys.argv), +} save_fig_with_metadata(time_str, time_file.storage_path, **kwds) result_plots += [(time_file,)] # Get segments and find where the data is -seg_dir = 'segments' +seg_dir = "segments" veto_cat_files = pycbc.workflow.get_files_for_vetoes( - workflow, seg_dir, runtime_names=['segments-science-veto']) + workflow, seg_dir, runtime_names=["segments-science-veto"] +) -science_seg_file, science_segs, _ = \ - pycbc.workflow.get_science_segments(workflow, seg_dir) +science_seg_file, science_segs, _ = pycbc.workflow.get_science_segments( + workflow, seg_dir +) -sci_ok_seg_file, science_ok_segs, _ = \ - pycbc.workflow.get_analyzable_segments(workflow, science_segs, - veto_cat_files, seg_dir) +sci_ok_seg_file, science_ok_segs, _ = pycbc.workflow.get_analyzable_segments( + workflow, science_segs, veto_cat_files, seg_dir +) -datafind_files, analyzable_file, analyzable_segs, analyzable_name = \ - pycbc.workflow.setup_datafind_workflow(workflow, science_ok_segs, - 'datafind', - seg_file=science_seg_file) +datafind_files, analyzable_file, analyzable_segs, analyzable_name = ( + pycbc.workflow.setup_datafind_workflow( + workflow, science_ok_segs, "datafind", seg_file=science_seg_file + ) +) -psd_job_length = int(workflow.cp.get('workflow-matchedfilter', 'analysis-length')) -pad_data = int(workflow.cp.get('calculate_psd', 'pad-data')) -max_segs_per_job = int(workflow.cp.get('workflow-matchedfilter', 'max-segments-per-job')) +psd_job_length = int(workflow.cp.get("workflow-matchedfilter", "analysis-length")) +pad_data = int(workflow.cp.get("calculate_psd", "pad-data")) +max_segs_per_job = int( + workflow.cp.get("workflow-matchedfilter", "max-segments-per-job") +) # calculate noise PSDs over analyzable segments psd_files = {} @@ -121,85 +131,110 @@ for ifo, segments in analyzable_segs.items(): # each batch goes into a PSD estimation job for job_index, job_segs in enumerate(batches): - tag = 'PART%.4d' % job_index - name = 'ANALYZABLE_SPLIT_' + tag + tag = "PART%.4d" % job_index + name = "ANALYZABLE_SPLIT_" + tag job_segs_file = pycbc.workflow.SegFile.from_segment_list( - name, job_segs, name, ifo, - valid_segment=workflow.analysis_time, extension='xml', - directory=seg_dir) + name, + job_segs, + name, + ifo, + valid_segment=workflow.analysis_time, + extension="xml", + directory=seg_dir, + ) ifo_psd_file = pycbc.workflow.make_psd_file( - workflow, datafind_files.find_output_with_ifo(ifo), - job_segs_file, name, 'psds', tags=[tag]) + workflow, + datafind_files.find_output_with_ifo(ifo), + job_segs_file, + name, + "psds", + tags=[tag], + ) if ifo not in psd_files: psd_files[ifo] = [] psd_files[ifo].append(ifo_psd_file) - flat_split_segments = \ - _segments.segmentlist([s for batch in batches for s in batch]) - logging.info('%.1f s of analyzable %s data reduced to %.1f s after ' - 'segmentation', abs(segments), ifo, abs(flat_split_segments)) + flat_split_segments = _segments.segmentlist([s for batch in batches for s in batch]) + logging.info( + "%.1f s of analyzable %s data reduced to %.1f s after segmentation", + abs(segments), + ifo, + abs(flat_split_segments), + ) # merge all PSDs for each detector merged_psd_files = [] for ifo, ifo_psd_files in psd_files.items(): - merged_psd_file = pycbc.workflow.merge_psds(workflow, ifo_psd_files, ifo, - 'psds') + merged_psd_file = pycbc.workflow.merge_psds(workflow, ifo_psd_files, ifo, "psds") merged_psd_files.append(merged_psd_file) # average noise PSDs and save to .txt and .xml -pycbc.workflow.make_average_psd(workflow, merged_psd_files, 'psds', - output_fmt='.txt') -pycbc.workflow.make_average_psd(workflow, merged_psd_files, 'psds', - output_fmt='.xml.gz') +pycbc.workflow.make_average_psd(workflow, merged_psd_files, "psds", output_fmt=".txt") +pycbc.workflow.make_average_psd( + workflow, merged_psd_files, "psds", output_fmt=".xml.gz" +) -s = pycbc.workflow.make_spectrum_plot(workflow, merged_psd_files, 'plots') +s = pycbc.workflow.make_spectrum_plot(workflow, merged_psd_files, "plots") result_plots += [(s,)] -pycbc.workflow.make_segments_plot(workflow, - pycbc.workflow.FileList([science_seg_file]), - 'plots', tags=['SCIENCE_MINUS_CAT1']) +pycbc.workflow.make_segments_plot( + workflow, + pycbc.workflow.FileList([science_seg_file]), + "plots", + tags=["SCIENCE_MINUS_CAT1"], +) # get data segments to write to segment summary XML file -seg_summ_names = ['DATA', 'SCIENCE_OK', 'ANALYZABLE_DATA'] +seg_summ_names = ["DATA", "SCIENCE_OK", "ANALYZABLE_DATA"] seg_summ_seglists = [science_segs, science_ok_segs, analyzable_segs] # write segment summary XML file seg_list = [] names = [] ifos = [] -for segment_list,segment_name in zip(seg_summ_seglists, seg_summ_names): +for segment_list, segment_name in zip(seg_summ_seglists, seg_summ_names): for ifo in workflow.ifos: seg_list.append(segment_list[ifo]) names.append(segment_name) ifos.append(ifo) seg_summ_file = pycbc.workflow.SegFile.from_multi_segment_list( - 'WORKFLOW_SEGMENT_SUMMARY', seg_list, names, ifos, - valid_segment=workflow.analysis_time, extension='xml', - directory=seg_dir) + "WORKFLOW_SEGMENT_SUMMARY", + seg_list, + names, + ifos, + valid_segment=workflow.analysis_time, + extension="xml", + directory=seg_dir, +) # make segment table for summary page -seg_summ_table = pycbc.workflow.make_seg_table(workflow, [seg_summ_file], - seg_summ_names, 'plots', ['SUMMARY'], - title_text='Input and output time', - description='This shows the total amount of input data, analyzable data, and the time for which PSDs are produced.') +seg_summ_table = pycbc.workflow.make_seg_table( + workflow, + [seg_summ_file], + seg_summ_names, + "plots", + ["SUMMARY"], + title_text="Input and output time", + description="This shows the total amount of input data, analyzable data, and the time for which PSDs are produced.", +) result_plots += [(seg_summ_table,)] -two_column_layout('plots', result_plots) +two_column_layout("plots", result_plots) # Create versioning information pycbc.workflow.make_versioning_page( workflow, container.cp, - rdir['workflow/version'], + rdir["workflow/version"], ) -pycbc.workflow.make_results_web_page(finalize_workflow, - os.path.join(os.getcwd(), - 'plots')) +pycbc.workflow.make_results_web_page( + finalize_workflow, os.path.join(os.getcwd(), "plots") +) container += workflow container += finalize_workflow diff --git a/bin/workflows/pycbc_make_sbank_workflow b/bin/workflows/pycbc_make_sbank_workflow index 0045b015f9c..99d386946e1 100644 --- a/bin/workflows/pycbc_make_sbank_workflow +++ b/bin/workflows/pycbc_make_sbank_workflow @@ -23,29 +23,36 @@ SbankExecutable class in the pycbc.workflow module, to give an illustration of how a simple workflow is constructed with pycbc.workflow. """ -#imports -import os +# imports import argparse +import os import pycbc import pycbc.workflow as wf # We define classes for all executables used in the workflow + class SbankExecutable(wf.Executable): - """ Class for running lalapps_cbc_sbank - """ + """Class for running lalapps_cbc_sbank""" + # This can be altered if you don't always want to store output files current_retention_level = wf.Executable.FINAL_RESULT # This tells us that reference-psd is a file option - file_input_options = wf.Executable.file_input_options + ['--reference-psd'] + file_input_options = wf.Executable.file_input_options + ["--reference-psd"] sbank_job_seed = 0 - def create_node(self, analysis_time, seed_bank=None, trial_bank=None, - mchirp_boundaries_file=None, mchirp_boundary_idx=None, - extra_tags=None): + def create_node( + self, + analysis_time, + seed_bank=None, + trial_bank=None, + mchirp_boundaries_file=None, + mchirp_boundary_idx=None, + extra_tags=None, + ): if extra_tags is None: extra_tags = [] node = wf.Executable.create_node(self) @@ -54,81 +61,85 @@ class SbankExecutable(wf.Executable): # the create_node function. *DO NOT* specify these in the config file. # The seed must be unique for each job and reproducible - node.add_opt('--seed', str(self.sbank_job_seed)) + node.add_opt("--seed", str(self.sbank_job_seed)) SbankExecutable.sbank_job_seed += 1 # These input files are optional. If given, add them if seed_bank is not None: - node.add_input_opt('--bank-seed', seed_bank) + node.add_input_opt("--bank-seed", seed_bank) if trial_bank is not None: - node.add_input_opt('--trial-waveforms', trial_bank) + node.add_input_opt("--trial-waveforms", trial_bank) if mchirp_boundaries_file is not None: - node.add_input_opt('--mchirp-boundaries-file', - mchirp_boundaries_file) + node.add_input_opt("--mchirp-boundaries-file", mchirp_boundaries_file) # The boundaries file option also requires the boundary idx - assert(mchirp_boundary_idx is not None) - node.add_opt('--mchirp-boundaries-index', mchirp_boundary_idx) + assert mchirp_boundary_idx is not None + node.add_opt("--mchirp-boundaries-index", mchirp_boundary_idx) # Here we add the output file, but we are letting pycbc.workflow # handle how to name the file - node.new_output_file_opt(analysis_time, '.h5', - '--output-filename', tags=self.tags + extra_tags) + node.new_output_file_opt( + analysis_time, ".h5", "--output-filename", tags=self.tags + extra_tags + ) return node + class SbankChooseMchirpBinsExecutable(wf.Executable): - """ Class for running lalapps_cbc_sbank_choose_mchirp_boundaries - """ + """Class for running lalapps_cbc_sbank_choose_mchirp_boundaries""" + current_retention_level = wf.Executable.ALL_TRIGGERS def create_node(self, analysis_time, input_file, nbanks): node = wf.Executable.create_node(self) # Here we add the output file - node.new_output_file_opt(analysis_time, '.txt', - '--output-file', tags=self.tags) + node.new_output_file_opt(analysis_time, ".txt", "--output-file", tags=self.tags) # And the input file, which is an argument, not an option node.add_input_arg(input_file) # nbanks is just a normal option, but as it affects the workflow # structure, it is supplied here and not directly in the config file - node.add_opt('--nbanks', nbanks) + node.add_opt("--nbanks", nbanks) return node + # There is already a ligolw_add executable (wf.LigolwAddExecutable), this needs # a minor change because we are potentially dealing with sub-daxes here. class LigolwAddExecutable(wf.LigolwAddExecutable): - - def create_node(self, jobSegment, input_files, output=None, - use_tmp_subdirs=True, tags=None): + def create_node( + self, jobSegment, input_files, output=None, use_tmp_subdirs=True, tags=None + ): if output is not None: # Convert path to file out_file = wf.File.from_path(output) if self.retain_files: if not os.path.isabs(output): - out_file.storage_path = os.path.join(self.out_dir, - output) + out_file.storage_path = os.path.join(self.out_dir, output) else: out_file.storage_path = output else: out_file = output - return super(LigolwAddExecutable, self).create_node\ - (jobSegment, input_files, output=out_file, - use_tmp_subdirs=use_tmp_subdirs, tags=tags) + return super().create_node( + jobSegment, + input_files, + output=out_file, + use_tmp_subdirs=use_tmp_subdirs, + tags=tags, + ) + class CombineHDFBanksExecutable(wf.Executable): - """ Class for running a combination of hdf banks - """ + """Class for running a combination of hdf banks""" + current_retention_level = wf.Executable.ALL_TRIGGERS - def create_node(self, analysis_time, input_file_list, output=None, - tags=None): + def create_node(self, analysis_time, input_file_list, output=None, tags=None): node = wf.Executable.create_node(self) # Here we add the input files - node.add_input_list_opt('--input-filenames', input_file_list) + node.add_input_list_opt("--input-filenames", input_file_list) curr_tags = self.tags if tags is not None: @@ -140,14 +151,14 @@ class CombineHDFBanksExecutable(wf.Executable): out_file = wf.File.from_path(output) if self.retain_files: if not os.path.isabs(output): - out_file.storage_path = os.path.join(self.out_dir, - output) + out_file.storage_path = os.path.join(self.out_dir, output) else: out_file.storage_path = output - node.add_output_opt('--output-file', out_file) + node.add_output_opt("--output-file", out_file) else: - node.new_output_file_opt(analysis_time, '.h5', - '--output-file', tags=curr_tags) + node.new_output_file_opt( + analysis_time, ".h5", "--output-file", tags=curr_tags + ) return node @@ -163,10 +174,14 @@ class CombineHDFBanksExecutable(wf.Executable): _desc = __doc__[1:] parser = argparse.ArgumentParser(description=_desc) pycbc.add_common_pycbc_options(parser) -parser.add_argument("--output-file", type=str, default=None, - help="Specify the output file name. Either a name can be " - "provided or a full path to file. Is this is not " - "given a filename and location is chosen ") +parser.add_argument( + "--output-file", + type=str, + default=None, + help="Specify the output file name. Either a name can be " + "provided or a full path to file. Is this is not " + "given a filename and location is chosen ", +) wf.add_workflow_command_line_group(parser) @@ -186,10 +201,10 @@ wf.makedir(args.output_dir) seed_file = None bins_inp_file = None -if workflow.cp.has_option_tags('workflow', 'seed-bank', args.tags): +if workflow.cp.has_option_tags("workflow", "seed-bank", args.tags): # If a seed bank is provided register it as a File object - seed_banks = workflow.cp.get_opt_tags('workflow', 'seed-bank', args.tags) - seed_banks = seed_banks.split(' ') + seed_banks = workflow.cp.get_opt_tags("workflow", "seed-bank", args.tags) + seed_banks = seed_banks.split(" ") if len(seed_banks) == 0: raise ValueError("No seed bank actually provided!") seed_files = [] @@ -203,52 +218,57 @@ if workflow.cp.has_option_tags('workflow', 'seed-bank', args.tags): seed_file = seed_files[0] else: # Combine with h5add - out_dir = os.path.join(args.output_dir, 'input_combine') - h5add_exe = CombineHDFBanksExecutable(workflow.cp, 'h5add', - ifos=['H1L1V1'], - out_dir=out_dir, - tags=['INPUT'] + args.tags) + out_dir = os.path.join(args.output_dir, "input_combine") + h5add_exe = CombineHDFBanksExecutable( + workflow.cp, + "h5add", + ifos=["H1L1V1"], + out_dir=out_dir, + tags=["INPUT"] + args.tags, + ) h5add_exe.update_current_retention_level(wf.Executable.ALL_TRIGGERS) - h5add_node = h5add_exe.create_node(workflow.analysis_time, - seed_files) + h5add_node = h5add_exe.create_node(workflow.analysis_time, seed_files) workflow += h5add_node - assert(len(h5add_node.output_files) == 1) + assert len(h5add_node.output_files) == 1 seed_file = h5add_node.output_files[0] # bins_inp_file will go to the mchirp_bins generator. seed file will go to # the first set of sbank jobs if not None. bins_inp_file = seed_file -if not workflow.cp.has_option_tags('workflow', 'use-seed-bank-for-chirp-bins', - args.tags): - out_dir = os.path.join(args.output_dir, 'coarse') +if not workflow.cp.has_option_tags( + "workflow", "use-seed-bank-for-chirp-bins", args.tags +): + out_dir = os.path.join(args.output_dir, "coarse") # Generate Executable class (similar to Job in the old terminology) # The tags=coarse option is used to ensure that options in the # ['sbank-coarse']section of the ini file are sent to this job, and *only* # this job - coarse_sbank_exe = SbankExecutable(workflow.cp, 'sbank', - ifos=workflow.ifos, - out_dir=out_dir, - tags=['coarse']+args.tags) - coarse_sbank_exe.update_current_retention_level\ - (wf.Executable.MERGED_TRIGGERS) + coarse_sbank_exe = SbankExecutable( + workflow.cp, + "sbank", + ifos=workflow.ifos, + out_dir=out_dir, + tags=["coarse"] + args.tags, + ) + coarse_sbank_exe.update_current_retention_level(wf.Executable.MERGED_TRIGGERS) # Then make a specific node coarse_node = coarse_sbank_exe.create_node(workflow.analysis_time) # Add to workflow workflow += coarse_node # And record output file, as it will be needed later - assert(len(coarse_node.output_files) == 1) + assert len(coarse_node.output_files) == 1 bins_inp_file = coarse_node.output_files[0] if seed_file is None: - if not workflow.cp.has_option_tags('workflow', - 'do-not-use-coarse-job-as-seed', - args.tags): + if not workflow.cp.has_option_tags( + "workflow", "do-not-use-coarse-job-as-seed", args.tags + ): seed_file = bins_inp_file if bins_inp_file is None: # Only get here if weird options are given - err_msg = 'You have not given a seed bank but have asked to use the seed ' - err_msg += 'bank for generating the chirp mass bins. This is not possible.' + err_msg = "You have not given a seed bank but have asked to use the seed " + err_msg += "bank for generating the chirp mass bins. This is not possible." raise ValueError(err_msg) ############################################################################## @@ -258,35 +278,38 @@ if bins_inp_file is None: # How many repetitions to try? Get this from config-parser. Special # config-parser options like this go in the [workflow] section -num_cycles = int(workflow.cp.get_opt_tags('workflow', 'num-cycles', args.tags)) +num_cycles = int(workflow.cp.get_opt_tags("workflow", "num-cycles", args.tags)) # Create executables up front to make plots nicer in the dashboard -bins_exe = SbankChooseMchirpBinsExecutable(workflow.cp, 'sbank_mchirp_bins', - ifos=workflow.ifos, tags=args.tags) -main_sbank_exe = SbankExecutable(workflow.cp, 'sbank', ifos=workflow.ifos, - tags=['parallel'] + args.tags) -h5add_first_exe = CombineHDFBanksExecutable(workflow.cp, 'h5add', - ifos=['H1L1V1'], - tags=['FIRST'] + args.tags) -readder_sbank_exe = SbankExecutable(workflow.cp, 'sbank', - ifos=workflow.ifos, - tags=['readder'] + args.tags) -h5add_final_exe = CombineHDFBanksExecutable(workflow.cp, 'h5add', - ifos=['H1L1V1'], - tags=['FINAL'] + args.tags) +bins_exe = SbankChooseMchirpBinsExecutable( + workflow.cp, "sbank_mchirp_bins", ifos=workflow.ifos, tags=args.tags +) +main_sbank_exe = SbankExecutable( + workflow.cp, "sbank", ifos=workflow.ifos, tags=["parallel"] + args.tags +) +h5add_first_exe = CombineHDFBanksExecutable( + workflow.cp, "h5add", ifos=["H1L1V1"], tags=["FIRST"] + args.tags +) +readder_sbank_exe = SbankExecutable( + workflow.cp, "sbank", ifos=workflow.ifos, tags=["readder"] + args.tags +) +h5add_final_exe = CombineHDFBanksExecutable( + workflow.cp, "h5add", ifos=["H1L1V1"], tags=["FINAL"] + args.tags +) for cycle_idx in range(num_cycles): ######### # SETUP # ######### - cycle_tag = 'cycle%d' %(cycle_idx) + cycle_tag = "cycle%d" % (cycle_idx) out_dir = os.path.join(args.output_dir, cycle_tag) # How many banks to use? This can vary cycle to cycle, or be the same for # all. Either supply it once in [workflow], or in [workflow-cycleN] for # N in range(num_cycles) - nbanks = workflow.cp.get_opt_tags('workflow', 'nbanks', - tags=[cycle_tag]+args.tags) + nbanks = workflow.cp.get_opt_tags( + "workflow", "nbanks", tags=[cycle_tag] + args.tags + ) nbanks = int(nbanks) ############# @@ -294,59 +317,58 @@ for cycle_idx in range(num_cycles): ############# bins_exe.update_current_tags([cycle_tag] + args.tags) bins_exe.update_output_directory(out_dir=out_dir) - bins_node = bins_exe.create_node(workflow.analysis_time, - bins_inp_file, nbanks) + bins_node = bins_exe.create_node(workflow.analysis_time, bins_inp_file, nbanks) workflow += bins_node - assert(len(bins_node.output_files) == 1) + assert len(bins_node.output_files) == 1 bins_out_file = bins_node.output_files[0] ####################### # PARALELLIZED SBANKS # ####################### - main_sbank_exe.update_current_tags(['parallel', cycle_tag] + args.tags) + main_sbank_exe.update_current_tags(["parallel", cycle_tag] + args.tags) main_sbank_exe.update_output_directory(out_dir=out_dir) # These jobs we don't always want to store - main_sbank_exe.update_current_retention_level\ - (wf.Executable.INTERMEDIATE_PRODUCT) + main_sbank_exe.update_current_retention_level(wf.Executable.INTERMEDIATE_PRODUCT) main_sbank_files = wf.FileList([]) for nbank_idx in range(nbanks): - nbank_tag = 'nbank%d' %(nbank_idx) - main_sbank_node = (main_sbank_exe.create_node\ - (workflow.analysis_time, seed_bank=seed_file, - mchirp_boundaries_file=bins_out_file, - mchirp_boundary_idx=nbank_idx, - extra_tags=[nbank_tag])) + nbank_tag = "nbank%d" % (nbank_idx) + main_sbank_node = main_sbank_exe.create_node( + workflow.analysis_time, + seed_bank=seed_file, + mchirp_boundaries_file=bins_out_file, + mchirp_boundary_idx=nbank_idx, + extra_tags=[nbank_tag], + ) workflow += main_sbank_node - assert(len(main_sbank_node.output_files) == 1) + assert len(main_sbank_node.output_files) == 1 main_sbank_files += main_sbank_node.output_files ############ # COMBINER # ############ - h5add_first_exe.update_current_tags([cycle_tag, 'FIRST'] + args.tags) + h5add_first_exe.update_current_tags([cycle_tag, "FIRST"] + args.tags) h5add_first_exe.update_output_directory(out_dir=out_dir) h5add_first_exe.update_current_retention_level(wf.Executable.ALL_TRIGGERS) - h5add_node = h5add_first_exe.create_node(workflow.analysis_time, - main_sbank_files) + h5add_node = h5add_first_exe.create_node(workflow.analysis_time, main_sbank_files) workflow += h5add_node - assert(len(h5add_node.output_files) == 1) + assert len(h5add_node.output_files) == 1 h5add_out = h5add_node.output_files[0] ########### # READDER # ########### - readder_sbank_exe.update_current_tags([cycle_tag, 'readder'] + args.tags) + readder_sbank_exe.update_current_tags([cycle_tag, "readder"] + args.tags) readder_sbank_exe.update_output_directory(out_dir=out_dir) - readder_sbank_exe.update_current_retention_level\ - (wf.Executable.ALL_TRIGGERS) - readder_sbank_node = readder_sbank_exe.create_node(workflow.analysis_time, - trial_bank=h5add_out) + readder_sbank_exe.update_current_retention_level(wf.Executable.ALL_TRIGGERS) + readder_sbank_node = readder_sbank_exe.create_node( + workflow.analysis_time, trial_bank=h5add_out + ) workflow += readder_sbank_node - assert(len(readder_sbank_node.output_files) == 1) + assert len(readder_sbank_node.output_files) == 1 readder_out = readder_sbank_node.output_files[0] ################# @@ -362,7 +384,7 @@ for cycle_idx in range(num_cycles): crl = wf.Executable.MERGED_TRIGGERS output_path = None - h5add_final_exe.update_current_tags([cycle_tag, 'FINAL'] + args.tags) + h5add_final_exe.update_current_tags([cycle_tag, "FINAL"] + args.tags) h5add_final_exe.update_output_directory(out_dir) h5add_final_exe.update_current_retention_level(crl) @@ -371,10 +393,11 @@ for cycle_idx in range(num_cycles): else: inputs = [seed_file, readder_out] - h5add_node = h5add_final_exe.create_node(workflow.analysis_time, - inputs, output=output_path) + h5add_node = h5add_final_exe.create_node( + workflow.analysis_time, inputs, output=output_path + ) workflow += h5add_node - assert(len(h5add_node.output_files) == 1) + assert len(h5add_node.output_files) == 1 # This becomes the input file for the next loop if going again seed_file = h5add_node.output_files[0] bins_inp_file = seed_file diff --git a/bin/workflows/pycbc_make_uberbank_workflow b/bin/workflows/pycbc_make_uberbank_workflow index cde03c2e140..e3f97e09d44 100644 --- a/bin/workflows/pycbc_make_uberbank_workflow +++ b/bin/workflows/pycbc_make_uberbank_workflow @@ -31,27 +31,29 @@ one described, so hopefully the script is clear enough to allow someone to do that. """ -#imports -import os +# imports import argparse import logging +import os -import pycbc import pycbc.version + +import pycbc import pycbc.workflow as wf from pycbc.workflow.pegasus_workflow import SubWorkflow # Boiler-plate stuff -__author__ = "Ian Harry " +__author__ = "Ian Harry " __version__ = pycbc.version.git_verbose_msg -__date__ = pycbc.version.date +__date__ = pycbc.version.date __program__ = "pycbc_make_uberbank_workflow" # Define classes for executables + class GeomBankExecutable(wf.Executable): - """ Class for running pycbc_geom_aligned_bank. - """ + """Class for running pycbc_geom_aligned_bank.""" + # This outputs a dax file. current_retention_level = wf.Executable.FINAL_RESULT @@ -61,41 +63,42 @@ class GeomBankExecutable(wf.Executable): config_file, out_storage_path, output_file, - workflow_name='geom_bank', - sub_wf_tags=None + workflow_name="geom_bank", + sub_wf_tags=None, ): - """ Create a sbank dax generator node. - """ + """Create a sbank dax generator node.""" node = wf.Executable.create_node(self) # Output file: This one is done weirdly because this job doesn't # actually produce this file, the resulting dax will do that. - node.add_opt('--output-file', output_file.storage_path) + node.add_opt("--output-file", output_file.storage_path) - - node.add_opt('--workflow-name', workflow_name) - node.add_opt('--output-dir', out_storage_path) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", workflow_name) + node.add_opt("--output-dir", out_storage_path) + node.add_opt("--dax-file-directory", ".") if sub_wf_tags is not None: - node.add_opt('--tags', ' '.join(sub_wf_tags)) + node.add_opt("--tags", " ".join(sub_wf_tags)) else: sub_wf_tags = [] - node.add_input_opt('--config-files', config_file) - - node.new_output_file_opt(analysis_time, '.dax', - '--dax-file', - tags=self.tags + sub_wf_tags + ['DAX']) - node.new_output_file_opt(analysis_time, '.map', - '--output-map', - tags=self.tags + sub_wf_tags + ['MAP']) + node.add_input_opt("--config-files", config_file) + + node.new_output_file_opt( + analysis_time, ".dax", "--dax-file", tags=self.tags + sub_wf_tags + ["DAX"] + ) + node.new_output_file_opt( + analysis_time, + ".map", + "--output-map", + tags=self.tags + sub_wf_tags + ["MAP"], + ) return node class SbankDaxGenerator(wf.Executable): - """ Class for running pycbc_make_sbank_workflow. - """ + """Class for running pycbc_make_sbank_workflow.""" + # This outputs a dax file. current_retention_level = wf.Executable.FINAL_RESULT @@ -106,28 +109,26 @@ class SbankDaxGenerator(wf.Executable): out_storage_path, output_file, seed_files=None, - workflow_name='sbank_workflow', - sub_wf_tags=None + workflow_name="sbank_workflow", + sub_wf_tags=None, ): - """ Create a sbank dax generator node. - """ + """Create a sbank dax generator node.""" node = wf.Executable.create_node(self) # Output file: This one is done weirdly because this job doesn't # actually produce this file, the resulting dax will do that. - node.add_opt('--output-file', output_file.storage_path) + node.add_opt("--output-file", output_file.storage_path) - - node.add_opt('--workflow-name', workflow_name) - node.add_opt('--output-dir', out_storage_path) - node.add_opt('--dax-file-directory', '.') + node.add_opt("--workflow-name", workflow_name) + node.add_opt("--output-dir", out_storage_path) + node.add_opt("--dax-file-directory", ".") if sub_wf_tags is not None: - node.add_opt('--tags', ' '.join(sub_wf_tags)) + node.add_opt("--tags", " ".join(sub_wf_tags)) else: sub_wf_tags = [] - node.add_input_opt('--config-files', config_file) + node.add_input_opt("--config-files", config_file) config_overrides = [] @@ -135,20 +136,23 @@ class SbankDaxGenerator(wf.Executable): sfs = [] for seed_file in seed_files: sfs.append(seed_file.name) - config_overrides.append('workflow:seed-bank:"%s"' % ' '.join(sfs)) - - if len(config_overrides): - node.add_opt('--config-overrides', ' '.join(config_overrides)) - - - node.new_output_file_opt(analysis_time, '.dax', - '--dax-file', - tags=self.tags + sub_wf_tags + ['DAX']) - node.new_output_file_opt(analysis_time, '.map', - '--output-map', - tags=self.tags + sub_wf_tags + ['MAP']) + config_overrides.append('workflow:seed-bank:"%s"' % " ".join(sfs)) + + if config_overrides: + node.add_opt("--config-overrides", " ".join(config_overrides)) + + node.new_output_file_opt( + analysis_time, ".dax", "--dax-file", tags=self.tags + sub_wf_tags + ["DAX"] + ) + node.new_output_file_opt( + analysis_time, + ".map", + "--output-map", + tags=self.tags + sub_wf_tags + ["MAP"], + ) return node + ############################################################################## # Argument parsing and setup of workflow # ############################################################################## @@ -171,55 +175,52 @@ workflow = wf.Workflow(args) wf.makedir(args.output_dir) os.chdir(args.output_dir) -wf.makedir('daxes') +wf.makedir("daxes") ############################################## # Run the geom_bank workflow the first time # ############################################## logging.info("Setting up the round of geometric bank jobs.") -config_path = os.path.abspath('daxes' + '/' + 'geom_workflow.ini') -workflow.cp.write(open(config_path, 'w')) +config_path = os.path.abspath("daxes" + "/" + "geom_workflow.ini") +workflow.cp.write(open(config_path, "w")) config_file = wf.resolve_url_to_file(config_path) geom_exe = GeomBankExecutable( - workflow.cp, - 'geom_aligned_bank_wflow', - ifos=workflow.ifos, - out_dir='daxes' + workflow.cp, "geom_aligned_bank_wflow", ifos=workflow.ifos, out_dir="daxes" ) # This one's a bit weird as the dag generator names the final output file, # but that job doesn't *generate* it, so we create this file first. -geom_out_dir = os.path.abspath('initial_geom_bank') +geom_out_dir = os.path.abspath("initial_geom_bank") geom_out_file = wf.File( geom_exe.ifo_list, geom_exe.name, workflow.analysis_time, - extension='.h5', + extension=".h5", store_file=True, directory=geom_out_dir, tags=[], - use_tmp_subdirs=False + use_tmp_subdirs=False, ) geom_node = geom_exe.create_node( workflow.analysis_time, config_file, geom_out_dir, geom_out_file, - workflow_name='initial_geom_bank' + workflow_name="initial_geom_bank", ) workflow += geom_node # Verify files are the correct ones, they have unique extensions -dax_file = [x for x in geom_node.output_files if x.name.endswith('.dax')][0] -map_file = [x for x in geom_node.output_files if x.name.endswith('.map')][0] +dax_file = [x for x in geom_node.output_files if x.name.endswith(".dax")][0] +map_file = [x for x in geom_node.output_files if x.name.endswith(".map")][0] # NOTE: One can use workflow += dax_job, but this seems to miss a bunch of # the stuff done below, and does not consider file management -dax_job = SubWorkflow(dax_file, is_planned=False, _id='geom') -dax_job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) +dax_job = SubWorkflow(dax_file, is_planned=False, _id="geom") +dax_job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file +) # And this output is used again so declare it. dax_job.add_outputs(geom_out_file, stage_out=True) @@ -232,11 +233,12 @@ geom_dax_job = dax_job # Run the sbank workflow for the first time # ############################################### -sbank_exe = SbankDaxGenerator(workflow.cp, 'sbank_workflow', - ifos=workflow.ifos, out_dir='daxes') +sbank_exe = SbankDaxGenerator( + workflow.cp, "sbank_workflow", ifos=workflow.ifos, out_dir="daxes" +) -config_path = os.path.abspath('daxes' + '/' + 'sbank_workflow.ini') -workflow.cp.write(open(config_path, 'w')) +config_path = os.path.abspath("daxes" + "/" + "sbank_workflow.ini") +workflow.cp.write(open(config_path, "w")) config_file = wf.resolve_url_to_file(config_path) # This first round of sbank is used for BBHs for the standard uberbank; @@ -244,26 +246,34 @@ config_file = wf.resolve_url_to_file(config_path) # skips this step. if not workflow.cp.has_option("workflow-bank_structure", "skip-coarse-bank"): logging.info("Setting up the (optional) first round of sbank jobs.") - sbank_out_file_bbh = wf.File(sbank_exe.ifo_list, sbank_exe.name, - workflow.analysis_time, extension='.h5', - store_file=True, - directory=os.path.abspath('sbank_bbh'), - tags=['sbank_bbh'], use_tmp_subdirs=False) - sbank_node = sbank_exe.create_node(workflow.analysis_time, config_file, - os.path.abspath('sbank_bbh'), - sbank_out_file_bbh, - workflow_name='sbank_bbh', - sub_wf_tags=['bbh']) + sbank_out_file_bbh = wf.File( + sbank_exe.ifo_list, + sbank_exe.name, + workflow.analysis_time, + extension=".h5", + store_file=True, + directory=os.path.abspath("sbank_bbh"), + tags=["sbank_bbh"], + use_tmp_subdirs=False, + ) + sbank_node = sbank_exe.create_node( + workflow.analysis_time, + config_file, + os.path.abspath("sbank_bbh"), + sbank_out_file_bbh, + workflow_name="sbank_bbh", + sub_wf_tags=["bbh"], + ) workflow += sbank_node ofiles = sbank_node.output_files - dax_file = [x for x in ofiles if x.name.endswith('.dax')][0] - map_file = [x for x in ofiles if x.name.endswith('.map')][0] - dax_job = SubWorkflow(dax_file, is_planned=False, _id='sbank_bbh') - dax_job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) + dax_file = [x for x in ofiles if x.name.endswith(".dax")][0] + map_file = [x for x in ofiles if x.name.endswith(".map")][0] + dax_job = SubWorkflow(dax_file, is_planned=False, _id="sbank_bbh") + dax_job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file + ) dax_job.add_outputs(sbank_out_file_bbh, stage_out=True) workflow._outputs.append(sbank_out_file_bbh) @@ -275,11 +285,16 @@ if not workflow.cp.has_option("workflow-bank_structure", "skip-coarse-bank"): ############################################### logging.info("Setting up the final round of sbank jobs.") -sbank_out_file_final = wf.File(sbank_exe.ifo_list, sbank_exe.name, - workflow.analysis_time, extension='.h5', - store_file=True, - directory=os.path.abspath('sbank_final'), - tags=['sbank_final'], use_tmp_subdirs=False) +sbank_out_file_final = wf.File( + sbank_exe.ifo_list, + sbank_exe.name, + workflow.analysis_time, + extension=".h5", + store_file=True, + directory=os.path.abspath("sbank_final"), + tags=["sbank_final"], + use_tmp_subdirs=False, +) seed_files = [geom_out_file] if not workflow.cp.has_option("workflow-bank_structure", "skip-coarse-bank"): @@ -288,22 +303,25 @@ if not workflow.cp.has_option("workflow-bank_structure", "skip-coarse-bank"): # NOTE: The call to generate this workflow actually doesn't depend on the input # seed banks being generated. It will only pass the name of these files # along, and not do anything with them. -sbank_node = sbank_exe.create_node(workflow.analysis_time, config_file, - os.path.abspath('sbank_final'), - sbank_out_file_final, - seed_files=seed_files, - workflow_name='sbank_final', - sub_wf_tags=['final']) +sbank_node = sbank_exe.create_node( + workflow.analysis_time, + config_file, + os.path.abspath("sbank_final"), + sbank_out_file_final, + seed_files=seed_files, + workflow_name="sbank_final", + sub_wf_tags=["final"], +) workflow += sbank_node ofiles = sbank_node.output_files -dax_file = [x for x in ofiles if x.name.endswith('.dax')][0] -map_file = [x for x in ofiles if x.name.endswith('.map')][0] -dax_job = SubWorkflow(dax_file, is_planned=False, _id='sbank_final') -dax_job.set_subworkflow_properties(map_file, - staging_site=workflow.staging_site, - cache_file=workflow.cache_file) +dax_file = [x for x in ofiles if x.name.endswith(".dax")][0] +map_file = [x for x in ofiles if x.name.endswith(".map")][0] +dax_job = SubWorkflow(dax_file, is_planned=False, _id="sbank_final") +dax_job.set_subworkflow_properties( + map_file, staging_site=workflow.staging_site, cache_file=workflow.cache_file +) dax_job.add_inputs(*seed_files) dax_job.add_into_workflow(workflow) From 18460a675cec907f136f376a0104288ed3638592 Mon Sep 17 00:00:00 2001 From: Ian Harry Date: Wed, 29 Jul 2026 13:13:10 +0100 Subject: [PATCH 3/5] Still need monkeypatch here, add NOQA to it --- pycbc/inference/sampler/emcee_pt.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pycbc/inference/sampler/emcee_pt.py b/pycbc/inference/sampler/emcee_pt.py index 054c3bbb85f..a954edf4ac5 100644 --- a/pycbc/inference/sampler/emcee_pt.py +++ b/pycbc/inference/sampler/emcee_pt.py @@ -44,7 +44,10 @@ # This is a hack that will allow us to continue using emcee's abandoned # PTSampler, which relied on `numpy.float`, until the end of time. -float = float +# NOTE: ruff's NPY001 fixer will try to "fix" this to `float = float` since +# it doesn't distinguish an assignment target from a read of the deprecated +# alias -- that turns this into a no-op and breaks the actual monkey-patch. +numpy.float = float # noqa: NPY001 if emcee.__version__ >= "3.0.0": raise ImportError From ba2f9ad74b584d01b6e8e2dc52d6a2a24be16eda Mon Sep 17 00:00:00 2001 From: Ian Harry Date: Wed, 29 Jul 2026 13:32:54 +0100 Subject: [PATCH 4/5] Avoid circular import --- pycbc/results/__init__.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/pycbc/results/__init__.py b/pycbc/results/__init__.py index ed3b54e1503..14fec29bb94 100644 --- a/pycbc/results/__init__.py +++ b/pycbc/results/__init__.py @@ -1,12 +1,21 @@ +# ruff: noqa: I001 +# NOTE: pygrb_plotting_utils / pygrb_postprocessing_utils MUST be imported +# last. They pull in pycbc.io.gracedb, which does +# `from pycbc.results import generate_asd_plot, generate_snr_plot, +# source_color` -- a circular import back into this partially-initialized +# module. That only works if color/psd/snr (which define those names) have +# already been imported above. A "safe" isort/ruff I001 autofix alphabetized +# these imports once already, moving the pygrb_* imports before snr and +# breaking this exact thing -- do not let that happen again. from pycbc.results.color import * from pycbc.results.dq import * from pycbc.results.layout import * from pycbc.results.metadata import * from pycbc.results.plot import * from pycbc.results.psd import * -from pycbc.results.pygrb_plotting_utils import * -from pycbc.results.pygrb_postprocessing_utils import * from pycbc.results.snr import * from pycbc.results.str_utils import * from pycbc.results.table_utils import * from pycbc.results.versioning import * +from pycbc.results.pygrb_plotting_utils import * +from pycbc.results.pygrb_postprocessing_utils import * From 5c325662a5f14639f0638acec2d0ce8c34dd6723 Mon Sep 17 00:00:00 2001 From: Ian Harry Date: Wed, 29 Jul 2026 14:45:02 +0100 Subject: [PATCH 5/5] Fix pre-existing typo that wasn't noticed before --- pycbc/waveform/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pycbc/waveform/utils.py b/pycbc/waveform/utils.py index feedbc92d5e..4427ecaa3b7 100644 --- a/pycbc/waveform/utils.py +++ b/pycbc/waveform/utils.py @@ -245,7 +245,6 @@ def phase_from_polarizations(h_plus, h_cross, remove_start_phase=True): Examples -------- - --------s >>> from pycbc.waveform import get_td_waveform, phase_from_polarizations >>> hp, hc = get_td_waveform(approximant="TaylorT4", mass1=10, mass2=10, f_lower=30, delta_t=1.0/4096)