Add calculation of source type probabilities to PyCBC Live

bin/pycbc_live

@@ -31,7 +31,7 @@ from pycbc.io.live import SingleCoincForGraceDB
 import pycbc.waveform.bank
 from pycbc.vetoes.sgchisq import SingleDetSGChisq
 from pycbc.waveform.waveform import props
-
+from pycbc import mchirp_area
 
 def ptof(p, ft):
     """Convert p-value to FAR via foreground time `ft`.
@@ -192,7 +192,7 @@ class LiveEventManager(object):
 
 
     def check_coincs(self, ifos, coinc_results, psds, f_low,
-                     data_readers, bank):
+                     data_readers, bank, mc_area_args):
         """ Perform any followup and save zerolag triggers to a coinc xml file
         """
         #check for hardware injection
@@ -227,7 +227,8 @@ class LiveEventManager(object):
             event = SingleCoincForGraceDB(live_ifos, coinc_results, bank=bank,
                                           psds=psds, followup_data=fud,
                                           low_frequency_cutoff=f_low,
-                                          channel_names=args.channel_name)
+                                          channel_names=args.channel_name,
+                                          mc_area_args=mc_area_args)
 
             end_time = int(coinc_results['foreground/%s/end_time'
                                          % coinc_ifos[0]])
@@ -583,6 +584,7 @@ LiveSingle.insert_args(parser)
 fft.insert_fft_option_group(parser)
 Coincer.insert_args(parser)
 SingleDetSGChisq.insert_option_group(parser)
+mchirp_area.insert_args(parser)
 args = parser.parse_args()
 scheme.verify_processing_options(args, parser)
 fft.verify_fft_options(args, parser)
@@ -595,6 +597,7 @@ pycbc.init_logging(args.verbose, format=log_format)
 
 ctx = scheme.from_cli(args)
 fft.from_cli(args)
+mc_area_args = mchirp_area.from_cli(args)
 
 # Approximant guess of the total padding
 valid_pad = args.analysis_chunk
@@ -799,7 +802,7 @@ with ctx:
 
                 evnt.check_coincs(list(results.keys()), best_coinc,
                                   psds, args.low_frequency_cutoff,
-                                  data_reader, bank)
+                                  data_reader, bank, mc_area_args)
 
             # Check for singles
             if args.enable_single_detector_background:

pycbc/io/live.py

@@ -3,6 +3,7 @@
 import pycbc
 import numpy
 import lal
+import json
 from six import u as unicode
 from glue.ligolw import ligolw
 from glue.ligolw import lsctables
@@ -13,7 +14,7 @@
 from pycbc import pnutils
 from pycbc.tmpltbank import return_empty_sngl
 from pycbc.results import ifo_color
-
+from pycbc.mchirp_area import calc_probabilities
 
 #FIXME Legacy build PSD xml helpers, delete me when we move away entirely from
 # xml formats
@@ -104,6 +105,9 @@ def __init__(self, ifos, coinc_results, **kwargs):
         channel_names: dict of strings, optional
             Strain channel names for each detector.
             Will be recorded in the sngl_inspiral table.
+        mc_area_args: dict of dicts, optional
+            Dictionary providing arguments to be used in source probability
+            estimation with pycbc/mchirp_area.py
         """
         self.template_id = coinc_results['foreground/%s/template_id' % ifos[0]]
         self.coinc_results = coinc_results
@@ -251,6 +255,17 @@ def __init__(self, ifos, coinc_results, **kwargs):
             psds_lal[ifo] = fseries
         make_psd_xmldoc(psds_lal, outdoc)
 
+        # source probabilities estimation
+        if 'mc_area_args' in kwargs:
+            eff_distances = [sngl.eff_distance for sngl in sngl_inspiral_table]
+            probabilities = calc_probabilities(coinc_inspiral_row.mchirp,
+                                               coinc_inspiral_row.snr,
+                                               min(eff_distances),
+                                               kwargs['mc_area_args'])
+            self.probabilities = probabilities
+        else:
+            self.probabilities = None
+
         self.outdoc = outdoc
         self.time = sngl_populated.get_end()
 
@@ -265,6 +280,13 @@ def save(self, filename):
         gz = filename.endswith('.gz')
         ligolw_utils.write_filename(self.outdoc, filename, gz=gz)
 
+        # save source probabilities in a json file
+        if self.probabilities is not None:
+            prob_fname = filename.replace('.xml.gz', '_probs.json')
+            with open(prob_fname, 'w') as prob_outfile:
+                json.dump(self.probabilities, prob_outfile)
+            logging.info('Source probabilities file saved as %s', prob_fname)
+
     def upload(self, fname, gracedb_server=None, testing=True,
                extra_strings=None):
         """Upload this trigger to gracedb
@@ -281,8 +303,6 @@ def upload(self, fname, gracedb_server=None, testing=True,
             test trigger (True) or a production trigger (False).
         """
         from ligo.gracedb.rest import GraceDb
-        import matplotlib
-        matplotlib.use('Agg')
         import pylab
 
         # first of all, make sure the event is saved on disk
@@ -332,6 +352,21 @@ def upload(self, fname, gracedb_server=None, testing=True,
             pylab.xlabel('Frequency (Hz)')
             pylab.ylabel('ASD')
             pylab.savefig(psd_series_plot_fname)
+            pylab.close()
+
+        if self.probabilities is not None:
+            prob_fname = fname.replace('.xml.gz', '_probs.json')
+            prob_plot_fname = prob_fname.replace('.json', '.png')
+
+            prob_plot = {k: v for (k, v) in self.probabilities.items()
+                         if v != 0.0}
+            labels, sizes = zip(*prob_plot.items())
+            fig, ax = pylab.subplots()
+            ax.pie(sizes, labels=labels, autopct='%1.1f%%',
+                   textprops={'fontsize': 15})
+            ax.axis('equal')
+            fig.savefig(prob_plot_fname)
+            pylab.close()
 
         gid = None
         try:
@@ -378,6 +413,16 @@ def upload(self, fname, gracedb_server=None, testing=True,
                                  filename=psd_series_plot_fname,
                                  tag_name=['psd'], displayName=['PSDs'])
 
+            # upload source probabilities in json format and plot
+            if self.probabilities is not None:
+                gracedb.writeLog(gid, 'source probabilities JSON file upload',
+                                 filename=prob_fname)
+                logging.info('Uploaded source probabilities for event %s', gid)
+                gracedb.writeLog(gid, 'source probabilities plot upload',
+                                 filename=prob_plot_fname)
+                logging.info('Uploaded source probabilities pie chart for '
+                             'event %s', gid)
+
         except Exception as exc:
             logging.error('Something failed during the upload/annotation of '
                           'event %s on GraceDB. The event may not have been '

pycbc/mchirp_area.py

@@ -8,8 +8,55 @@
 detector frame mass and redshift.
 """
 
+import math
 from pycbc.conversions import mass2_from_mchirp_mass1 as m2mcm1
 from scipy.integrate import quad
+from pycbc.cosmology import _redshift
+
+
+def insert_args(parser):
+    mchirp_group = parser.add_argument_group("Arguments for computing the "
+                                             "areas of the CBC regions using "
+                                             "mchirp and effective distance.")
+    mchirp_group.add_argument('--max-m1', type=float, default=45.0,
+                              help="Maximum value for m1")
+    mchirp_group.add_argument('--min-m2', type=float, default=1.0,
+                              help="Minimum value for m2")
+    mchirp_group.add_argument('--ns-max', type=float, default=3.0,
+                              help="Maximum neutron star mass")
+    mchirp_group.add_argument('--gap-max', type=float, default=5.0,
+                              help="Minimum black hole mass")
+    mchirp_group.add_argument('--mchirp-to-delta-coeff', type=float,
+                              metavar='m0', default=0.01,
+                              help='Coefficient to estimate the value of the '
+                                   'mchirp uncertainty by mchirp_delta = '
+                                   'm0 * mchirp.')
+    mchirp_group.add_argument('--eff-to-lum-distance-coeff', type=float,
+                              metavar='a0', default=0.759,
+                              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('--lum-distance-to-delta-coeff', type=float,
+                              nargs=2, metavar=('b0', 'b1'),
+                              default=[-0.449, -0.342],
+                              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('--mass-gap-separate', action='store_true',
+                              help='Gives separate probabilities for each kind'
+                                   ' of mass gap CBC sources: gns, gg, bhg.')
+
+
+def from_cli(args):
+    return {'mass_limits': {'max_m1': args.max_m1, 'min_m2': args.min_m2},
+            'mass_bdary': {'ns_max': args.ns_max, 'gap_max': args.gap_max},
+            'estimation_coeff': {'a0': args.eff_to_lum_distance_coeff,
+                                 'b0': args.lum_distance_to_delta_coeff[0],
+                                 'b1': args.lum_distance_to_delta_coeff[1],
+                                 'm0': args.mchirp_to_delta_coeff},
+            'mass_gap': args.mass_gap_separate}
 
 
 def src_mass_from_z_det_mass(z, del_z, mdet, del_mdet):
@@ -22,23 +69,15 @@ def src_mass_from_z_det_mass(z, del_z, mdet, del_mdet):
     return (msrc, del_msrc)
 
 
-# Integration function
-def mchange(x, mc):
-    """Returns a component mass as a function of mchirp and the other
-    component mass.
-    """
-    return m2mcm1(mc, x)
-
-
 def intmc(mc, x_min, x_max):
     """Returns the integral of mchange between the minimum and maximum values
     of a component mass taking mchirp as an argument.
     """
-    integral = quad(mchange, x_min, x_max, args=mc)
+    integral = quad(lambda x, mc: m2mcm1(mc, x), x_min, x_max, args=mc)
     return integral[0]
 
 
-def calc_areas(trig_mc_det, mass_limits, mass_bdary, z):
+def calc_areas(trig_mc_det, mass_limits, mass_bdary, z, mass_gap):
     """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
@@ -227,12 +266,50 @@ def calc_areas(trig_mc_det, mass_limits, mass_bdary, z):
         int_inf_nsbh = ints_nsbh + intline_inf_nsbh
 
         ansbh = int_sup_nsbh - int_inf_nsbh
-
+    if mass_gap is not False:
+        return {
+            "BNS": abns,
+            "GNS": agns,
+            "NSBH": ansbh,
+            "GG": agg,
+            "BHG": abhg,
+            "BBH": abbh
+            }
     return {
-        "bns": abns,
-        "gns": agns,
-        "nsbh": ansbh,
-        "gg": agg,
-        "bhg": abhg,
-        "bbh": abbh
+        "BNS": abns,
+        "NSBH": ansbh,
+        "BBH": abbh,
+        "Mass Gap": agns + agg + abhg
         }
+
+
+def calc_probabilities(mchirp, snr, eff_distance, src_args):
+    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_estimation = _redshift(dist_estimation)
+    z_est_max = _redshift(dist_estimation + dist_std_estimation)
+    z_est_min = _redshift(dist_estimation - dist_std_estimation)
+    z_std_estimation = 0.5 * (z_est_max - z_est_min)
+    z = {'central': z_estimation, 'delta': z_std_estimation}
+    mass_gap = src_args['mass_gap']
+
+# If the mchirp is greater than a mchirp corresponding to two masses
+    # equal to the maximum mass, the probability for BBH is 100%
+    mc_max = mass_limits['max_m1'] / (2 ** 0.2)
+    if trig_mc_det['central'] > mc_max * (1 + z['central']):
+        if mass_gap is not False:
+            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}
+    else:
+        areas = calc_areas(trig_mc_det, mass_limits, mass_bdary, z, mass_gap)
+        total_area = sum(areas.values())
+        probabilities = {key: areas[key]/total_area for key in areas}
+    return probabilities

pycbc/results/color.py

@@ -10,6 +10,19 @@
     '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
+}
+
 
 def ifo_color(ifo):
     return _ifo_color_map[ifo]
+
+def source_color(source):
+    return _source_color_map[source]