gstlal_inspiral_create_prior_diststats

#!/usr/bin/env python
#
# Copyright (C) 2010--2014  Kipp Cannon, Chad Hanna
#
# 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
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# 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.
#
## @file gstlal_inspiral_create_prior_diststats
# A program to create some prior likelihood data to seed an online analysis
#
# ### Command line interface
#   + `--verbose`: Be verbose.
#   + `--synthesize-injection-count` [N] (int): Synthesize an injection distribution with N injections (default is 1000000).
#   + `--write-likelihood` [filename]: Write merged raw likelihood data to this file.
#   + `--instrument`: Append to a list of instruments to create dist stats for.  Must be whatever instruments you intend to analyze.
#   + `--trials-far-thresh` (float): Set the far threshold for the thresh parameter in the trials table.
#   + `--background-prior` [N] (float): Include an exponential background prior with the maximum bin count = N (default is 1).

#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


import sys
from optparse import OptionParser
import numpy


from glue import iterutils
from glue.ligolw import ligolw
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
from glue.text_progress_bar import ProgressBar


from pylal import series as lalseries
from pylal import rate


from gstlal import far
from gstlal import reference_psd


__author__ = "Chad Hanna <chad.hanna@ligo.org>"
__version__ = "git id %s" % ""  # FIXME
__date__ = ""   # FIXME


#
# =============================================================================
#
#                                 Command Line
#
# =============================================================================
#


def parse_command_line():
    parser = OptionParser(
        version = "Name: %%prog\n%s" % "" # FIXME
    )
    parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
    parser.add_option("-s", "--synthesize-injection-count", metavar = "N", default = 100000, type = "int", help = "Synthesize an injection distribution with N injections. default 1000000")
    parser.add_option("--write-likelihood", metavar = "filename", help = "Write merged raw likelihood data to this file.")
    parser.add_option("--instrument", action = "append", help = "Append to a list of instruments to create dist stats for.  List must be whatever instruments you intend to analyze.")
    parser.add_option("--horizon-distances", action = "append", help = "Cache SNR PDFs for these instruments and horizon distances.  Format is \"instrument=distance,instrument=distance,...\", e.g., H1=120,L1=120,V1=48.  Units for distance are irrelevant (PDFs depend only on their ratios).  A PDF will be constructed for every combination of two or more instruments from the set.  It is an error for an instrument to be named here and not in a --instrument option.")
    parser.add_option("--horizon-distance-masses", metavar = "m1,m2", action = "append", default = ["1.4,1.4"], help = "When computing pre-cached SNR PDFs from a collection of PSD files, compute horizon distances for these masses in solar mass units (default = 1.4,1.4).  Can be given multiple times.")
    parser.add_option("--horizon-distance-flow", metavar = "Hz", default = 10., type = "float", help = "When computing pre-cached SNR PDFs from a collection PSD files, start horizon distance integral at this frequency in Hertz (default = 10 Hz).")
    parser.add_option("-p", "--background-prior", metavar = "N", default = 1, type = "float", help = "Include an exponential background prior with the maximum bin count = N, default 1")
    options, filenames = parser.parse_args()

    process_params = dict(options.__dict__)

    options.instrument = set(options.instrument)
    if not options.instrument or len(options.instrument) < 2:
        raise ValueError("must specify at least two distinct --instrument's")

    if options.horizon_distances is None:
        options.horizon_distances = []
    options.horizon_distances = [dict((name, float(dist)) for name, dist in [name_dist.split("=") for name_dist in horizon_distances.strip().split(",")]) for horizon_distances in options.horizon_distances]
    if options.horizon_distances:
        all_horizon_instruments = reduce(lambda a, b: a | b, (set(horizon_distances) for horizon_distances in options.horizon_distances))
        if all_horizon_instruments > options.instrument:
            raise ValueError("missing %s for instruments named in --horizon-distances options" % ", ".join("--instrument %s" % inst for inst in all_horizon_instruments - options.instrument))

    if len(filenames) > 1:
        raise ValueError("At this point only one PSD file is supported on the command line")

    options.horizon_distance_masses = [map(float, s.split(",")) for s in options.horizon_distance_masses]

    return options, process_params, filenames



#
# =============================================================================
#
#                                     Main
#
# =============================================================================
#


#
# command line
#


options, process_params, filenames = parse_command_line()


#
# initialize output document (records process start time)
#


xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process = ligolw_process.register_to_xmldoc(xmldoc, u"gstlal_inspiral_create_prior_diststats", ifos = options.instrument, paramdict = process_params)


#
# create parameter distribution priors
#


coincparamsdistributions = far.ThincaCoincParamsDistributions()

if options.background_prior > 0:
    coincparamsdistributions.add_background_prior(n = options.background_prior, instruments = options.instrument, verbose = options.verbose)

if options.synthesize_injection_count > 0:
    coincparamsdistributions.add_foreground_snrchi_prior(instruments = options.instrument, n = options.synthesize_injection_count, prefactors_range = (0.0, 0.10), df = 40, verbose = options.verbose)


#
# seed SNR PDF cache
#


for n, filename in enumerate(filenames, 1):
    if options.verbose:
        print >>sys.stderr, "%d/%d:" % (n, len(filenames)),
    psd = lalseries.read_psd_xmldoc(ligolw_utils.load_filename(filename, contenthandler = lalseries.LIGOLWContentHandler, verbose = options.verbose))
    for m1, m2 in options.horizon_distance_masses:
        horizon_distances = dict((instrument, (0. if instrument not in psd else reference_psd.horizon_distance(psd[instrument], m1, m2, 8., options.horizon_distance_flow))) for instrument in options.instrument)
        if options.verbose:
            print >>sys.stderr, "\t%s" % ", ".join("%s = %.4g Mpc" % x for x in sorted(horizon_distances.items()))
        options.horizon_distances.append(horizon_distances)

for horizon_distances in options.horizon_distances:
    for n in range(2, len(horizon_distances) + 1):
        for instruments in iterutils.choices(horizon_distances.keys(), n):
            # FIXME get_snr_joint_pdf() should be called in the
            # future, but for now we just make pdfs for the actual
            # horizon distances requested not the quantized ones.
            # coincparamsdistributions.get_snr_joint_pdf(instruments, horizon_distances, verbose = options.verbose)

            # Force the SNR pdf to be generated to be at the actual horizon distance values not the quantized ones
            key = frozenset(instruments), frozenset(coincparamsdistributions.quantize_horizon_distances(horizon_distances).items())
            if options.verbose:
                print >>sys.stderr, "For horizon distances %s" % ", ".join("%s = %.4g Mpc" % item for item in sorted(horizon_distances.items()))
                progressbar = ProgressBar(text = "%s SNR PDF" % ", ".join(sorted(key[0])))
            else:
                progressbar = None
            binnedarray = coincparamsdistributions.joint_pdf_of_snrs(key[0], horizon_distances, progressbar = progressbar)
            del progressbar
            coincparamsdistributions.snr_joint_pdf_cache[key] = None, binnedarray, 0


#
# record results in output file
#


far.gen_likelihood_control_doc(xmldoc, process, coincparamsdistributions, None, {}, comment = u"background and signal priors (no real data)")


#
# done
#


ligolw_process.set_process_end_time(process)
ligolw_utils.write_filename(xmldoc, options.write_likelihood, gz = options.write_likelihood.endswith(".gz"), verbose = options.verbose)