#!/usr/bin/env python """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.""" import h5py, pycbc, subprocess, argparse, tempfile, shutil import numpy as np import logging from pycbc.filter import sigmasq from pycbc.io import live from pycbc.types import TimeSeries, FrequencySeries, MultiDetOptionAppendAction from pycbc.pnutils import nearest_larger_binary_number from pycbc.waveform.spa_tmplt import spa_length_in_time from pycbc import frame from ligo.gracedb.rest import GraceDb def default_frame_type(time, ifo): """Sensible defaults for frame types based on interferometer and time. """ if time < 1137254517: # O1 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' elif time >= 1235433618: # O3 if ifo == 'V1': return 'V1Online' elif ifo in ['H1', 'L1']: return ifo + '_HOFT_C00' raise ValueError('Detector {} not supported at time {}'.format(ifo, time)) def default_channel_name(time, ifo): """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' 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' elif time >= 1235433618: # O3 if ifo == 'V1': return ifo + ':Hrec_hoft_16384Hz' elif ifo in ['H1', 'L1']: return ifo + ':GDS-CALIB_STRAIN' 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, window_bins=300, frame_types=None, channel_names=None, gracedb_server=None, test_event=True, custom_frame_files=None, approximant=None): if not test_event and not gracedb_server: raise RuntimeError('a GraceDB URL must be specified if not a test event.') tmpdir = tempfile.mkdtemp() coinc_results = {} follow_up = [] followup_data = {} procs = [] if frame_types is None: frame_types = {} if channel_names is None: channel_names = {} for ifo in ifos: if ifo not in frame_types: frame_types[ifo] = default_frame_type(trig_time, ifo) if ifo not in channel_names: channel_names[ifo] = default_channel_name(trig_time, ifo) rough_time = int(np.round(trig_time)) window = 1. start_time = trig_time - window # parameters to fit a single-template inspiral job nicely # around the trigger time, without requiring too much data # Padding set by 16 * 2 for psd and buffer for other filtering pad = 60 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) gps_end_time = int(rough_time + pad // 2) gps_start_time = gps_end_time - segment_length logging.info("Using data: %s-%s", gps_start_time, gps_end_time) highpass_frequency = int(f_low * 0.7) logging.info("Setting highpass: %s Hz", highpass_frequency) for i, ifo in enumerate(ifos): followup_data[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","100", "--autogating-cluster","0.5", "--autogating-width","0.25", "--autogating-taper","0.25", "--autogating-pad","16", "--strain-high-pass", str(highpass_frequency), "--pad-data", "8", "--chisq-bins",'0.72*get_freq("fSEOBNRv4Peak",params.mass1,params.mass2,params.spin1z,params.spin2z)**0.7'] coinc_results['foreground/'+ifo+'/end_time'] = float(trig_time) 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'] = f_low if f_upper: coinc_results['foreground/'+ifo+'/f_final'] = f_upper coinc_results['foreground/'+ifo+'/window'] = window coinc_results['foreground/'+ifo+'/sample_rate'] = int(sample_rate) coinc_results['foreground/'+ifo+'/template_id'] = i command.append("--approximant") for apx in approximant: command.append(apx) command.append("--sample-rate"),command.append(str(int(sample_rate))) command.append("--mass1"),command.append(str(mass1)) command.append("--mass2"),command.append(str(mass2)) command.append("--spin1z"),command.append(str(spin1z)) command.append("--spin2z"),command.append(str(spin2z)) command.append("--low-frequency-cutoff"),command.append(str(f_low)) if f_upper: command.append("--high-frequency-cutoff"),command.append(str(f_upper)) command.append("--gps-start-time"),command.append(str(gps_start_time)) command.append("--gps-end-time"),command.append(str(gps_end_time)) command.append("--trigger-time"),command.append(str(np.round(trig_time,5))) command.append("--window"),command.append(str(window)) if custom_frame_files is None or ifo not in custom_frame_files: command.append("--frame-type") command.append(frame_types[ifo]) else: logging.info("Check if the segment in the custom frame file is safe") fr_start_times = [] fr_end_times = [] for custom_frame in custom_frame_files[ifo]: 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]) raise RuntimeError(msg) fr_start_times.append(frame_data.start_time) fr_end_times.append(frame_data.end_time) if gps_start_time < np.min(fr_start_times): msg = 'Start time of {} must be before the required start time {}' msg = msg.format(np.min(fr_start_times), gps_start_time) raise RuntimeError(msg) if np.max(fr_end_times) < gps_end_time: msg = 'End time of {} must be after the required end time {}' msg = msg.format(np.max(fr_end_times), gps_end_time) raise RuntimeError(msg) if trig_time < np.min(fr_start_times) or np.max(fr_end_times) < trig_time: msg = 'Trigger time must be within your frame file(s)' raise RuntimeError(msg) command.append("--frame-files") for custom_frame in custom_frame_files[ifo]: command.append(custom_frame) command.append("--channel-name") command.append(channel_names[ifo]) command.append("--psd-output") command.append(tmpdir+"/PSD_"+str(rough_time)+"_"+ifo+".txt") command.append("--output-file") command.append(tmpdir+"/SNRTS_"+str(rough_time)+"_"+ifo+".hdf") log_path = tmpdir + '/pycbc_single_template_{}_{}_log.txt'.format(str(rough_time), ifo) log_file = open(log_path, 'w') log_file.write(' '.join(command) + '\n\n') log_file.flush() 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') 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) 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.') logging.info('Gathering info from pycbc_single_template results') for i, ifo in enumerate(ifos): f = h5py.File(tmpdir+"/SNRTS_"+str(rough_time)+"_"+ifo+".hdf",'r') g = np.loadtxt(tmpdir+"/PSD_"+str(rough_time)+"_"+ifo+".txt") coinc_results['foreground/'+ifo+'/snr_series'] = f['snr'][()] coinc_results['foreground/'+ifo+'/psd_series'] = g[:,1] * pycbc.DYN_RANGE_FAC ** 2.0 coinc_results['foreground/'+ifo+'/delta_f'] = float(g[1][0])-float(g[0][0]) coinc_results['foreground/'+ifo+'/event_id'] = 'sngl_inspiral:event_id:'+str(i) df = float(g[1][0])-float(g[0][0]) dt = 1.0/sample_rate snr_series_peak = np.argmax(np.absolute(coinc_results['foreground/'+ifo+'/snr_series'])) SNR=abs(coinc_results['foreground/'+ifo+'/snr_series'][snr_series_peak]) coa_phase = np.angle(coinc_results['foreground/'+ifo+'/snr_series'][snr_series_peak]) chisq = f['chisq'][:][snr_series_peak] logging.info('%s SNR: %.2f', ifo, SNR) if SNR < thresh_SNR: follow_up.append(ifo) continue coinc_results['foreground/'+ifo+'/snr'] = SNR coinc_results['foreground/'+ifo+'/sigmasq'] = \ sigmasq(FrequencySeries(np.complex128(f['template'][()]), delta_f=df), psd=FrequencySeries(coinc_results['foreground/'+ifo+'/psd_series'], delta_f=df)) coinc_results['foreground/'+ifo+'/coa_phase'] = coa_phase coinc_results['foreground/'+ifo+'/chisq'] = chisq ifos = [x for x in ifos if x not in follow_up] if not ifos: raise RuntimeError('all interferometers have SNR below threshold.' ' Is this really a candidate event?') sumsq_snr = 0.0 for i,ifo in enumerate(ifos): snr_series_peak = np.argmax(np.absolute(coinc_results['foreground/'+ifo+'/snr_series'])) coinc_results['foreground/'+ifo+'/end_time'] = start_time + dt*snr_series_peak sumsq_snr += np.power(coinc_results['foreground/'+ifo+'/snr'],2) coinc_results['foreground/stat'] = np.sqrt(sumsq_snr) coinc_results['foreground/ifar'] = ifar subthreshold_sngl_time = np.mean([coinc_results['foreground/%s/end_time' % ifo] for ifo in ifos]) for ifo in ifos + follow_up: if ifo in follow_up: coinc_results['foreground/'+ifo+'/end_time'] = subthreshold_sngl_time time = coinc_results['foreground/'+ifo+'/end_time'] name='psd' fseries = FrequencySeries(coinc_results['foreground/'+ifo+'/psd_series'], delta_f=df) from pycbc.psd import interpolate fseries = interpolate(fseries, 0.25) followup_data[ifo]['psd'] = fseries name='snr' units='' peak_bin = int(sample_rate * (time - start_time)) max_bin = peak_bin + window_bins + 1 if max_bin > coinc_results['foreground/'+ifo+'/snr_series'].size: window_bins = coinc_results['foreground/'+ifo+'/snr_series'].size - peak_bin max_bin = coinc_results['foreground/'+ifo+'/snr_series'].size min_bin = peak_bin - window_bins + 1 else: min_bin = peak_bin - window_bins if min_bin < 0: window_bins = peak_bin max_bin = peak_bin + window_bins + 1 min_bin = peak_bin - window_bins epoch = time - window_bins*dt series = TimeSeries(coinc_results['foreground/'+ifo+'/snr_series'][min_bin:max_bin].astype(np.complex64), delta_t=dt, epoch=epoch) followup_data[ifo]['snr_series'] = series kwargs = {'psds': {ifo: followup_data[ifo]['psd'] for ifo in ifos + follow_up}, 'low_frequency_cutoff': f_low, 'high_frequency_cutoff': f_upper, 'followup_data': followup_data, 'channel_names': channel_names} comments = ['Manual followup up from PyCBC'] doc = live.SingleCoincForGraceDB(ifos, coinc_results, **kwargs) ligolw_file_path = tmpdir + '/' + str(rough_time) + '.xml' if gracedb_server: 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) # run BAYESTAR to generate the skymap cmd = ['bayestar-localize-coincs', ligolw_file_path, ligolw_file_path, '--f-low', str(f_low), '-o', tmpdir] subprocess.call(cmd) skymap_fits_name = tmpdir + '/0.fits' # plot the skymap skymap_plot_name = tmpdir + '/' + str(rough_time) + '_skymap.png' cmd = ['ligo-skymap-plot', skymap_fits_name, '-o', skymap_plot_name, '--contour', '50', '90', '--annotate'] subprocess.call(cmd) final_fits_dir = ligolw_skymap_output + '/' + str(rough_time) + '.fits' final_png_dir = ligolw_skymap_output + '/' + str(rough_time) + '_skymap.png' shutil.move(skymap_fits_name, final_fits_dir) shutil.move(skymap_plot_name, final_png_dir) if gracedb_server: gracedb.writeLog(gid, 'Bayestar skymap FITS file upload', filename=skymap_fits_name, tag_name=['sky_loc'], displayName=['Bayestar FITS skymap']) gracedb.writeLog(gid, 'Bayestar skymap plot upload', filename=skymap_plot_name, tag_name=['sky_loc'], displayName=['Bayestar skymap plot']) if ligolw_event_output: shutil.move(ligolw_file_path, ligolw_event_output) if __name__ == '__main__': pycbc.init_logging(True) parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--version', action=pycbc.version.Version) parser.add_argument('--trig-time', required=True, type=float, help='GPS time of trigger (float)') 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=float, default=2048.0, 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, help="Threshold SNR (float)", default=4.5) 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('--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') opt = parser.parse_args() 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)