pycbc_create_injections

#! /usr/bin/env python

# Copyright (C) 2017 Collin Capano
#
# 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 3 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.

"""Generates injections drawn from a distribution read from a config file.

Config file syntax
------------------
The configuration file should have a [variable_params], a [static_params],
and one or more [prior] sections. Multiple [constraint] sections and
one or more [waveform_transforms] sections may also be provided. An example:

    [variable_params]
    mass1 =
    mass2 =
    spin1_a =
    spin1_azimuthal =
    spin1_polar =

    [static_params]
    approximant = IMRPhenomPv2
    f_lower = 19

    [prior-mass1]
    name = uniform
    min-mass1 = 3
    max-mass1 = 12

    [prior-mass2]
    name = uniform
    min-mass2 = 1
    max-mass2 = 3

    [constraint-1]
    name = custom
    constraint_arg = q_from_mass1_mass2(mass1, mass2) <= 4

    [prior-spin1_a]
    name = uniform
    min-spin1_a = 0.0
    max-spin1_a = 0.9

    [prior-spin1_polar+spin1_azimuthal]
    name = uniform_solidangle
    polar-angle = spin1_polar
    azimuthal-angle = spin1_azimuthal

    [waveform_transforms-spin1x+spin1y+spin1z]
    name = spherical_spin_1_to_cartesian_spin_1

This config file would generate injections uniform in mass1 and mass2, with a
cut on mass ratio (q) <= 4. The spin of the larger body is also varied,
uniform in magnitude < 0.9, and isotropic in orientation (spin1_polar,
spin1_azimuthal). Because the waveform module expects spins to be provided in
cartesian coordinates, the spherical spin coordinates are transformed prior to
being written out to file.

The [variable_params] section gives the list of waveform parameters to
randomize.  The [static_params] section gives parameters that are fixed for all
of the injections.

The [prior-{tag}] sections provide the arguments needed to initialize
the prior for each parameter. Every parameter specified in
[variable_params] must have a prior section; the name of the parameter(s)
used by that prior must be in the 'tag' part of the section header (the
bit after the dash). If a prior covers multiple parameters, the
parameters should be separated by a '+'.  Any prior in the distributions
module may be used.  The rest of the options should provide the necessary
arguments to initialize that prior; see the distributions module for
details.

The variable_params need not be parameters understood by the waveform module.  In
that case, a [waveform_transforms] sections must be provided that map
variable_params not understood by the waveform module to parameters that are
understood. In the above example, spin1_a, spin1_azimuthal, and spin1_polar are
converted to spin1x, spin1y, and spin1z before being written out. Any transform
in the transforms module may be used to do this; see that module for details.
No attempt is made to check that provided parameter names are sensible. It is
up to the user to ensure that written parameters are understood by the waveform
approximant that will be used.

One or more constraints may be applied to the priors; these are
specified by the [constraint] section(s). Additional constraints may be
supplied by adding more [constraint-{tag}] sections. Any tag may be used; the
only requirement is that they be unique. If multiple constaint sections are
provided, the union of all constraints are applied. Alternatively, multiple
constraints may be joined in a single argument using numpy's logical operators.

The parameter that constraints are applied to may be any parameter in
variable_params or any output parameter of the transforms. Functions may be
applied on these parameters to obtain constraints on derived parameters. Any
function in the conversions, coordinates, or cosmology module may be used,
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 numpy
from numpy.random import uniform

import pycbc
from pycbc.inject import InjectionSet
from pycbc import distributions
from pycbc import transforms
from pycbc.distributions import JointDistribution
from pycbc.workflow import configuration
from pycbc.workflow import WorkflowConfigParser

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.")
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.")

numpy.random.seed(opts.seed)

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,
    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')
else:
    waveform_transforms = None
    write_args = variable_params

# get prior distribution for each variable parameter
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})

if opts.ninjections:
    draw_size = opts.ninjections
else:
    draw_size = 4000 # Just a default so it's not super slow drawing large sets
    old_samples = None

while True:
    logging.info("Drawing samples")
    samples = randomsampler.rvs(size=draw_size)

    if waveform_transforms is not None:
        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])
        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 old_samples is None:
            tstart = opts.gps_start_time
        else:
            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'])
        
        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))
            break

    # We got as many samples as we needed so we can stop
    if opts.ninjections and len(samples) >= opts.ninjections:
        break

# 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))