pycbc_inference

#!/usr/bin/env python

# Copyright (C) 2016 Christopher M. Biwer, 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.
""" Runs a sampler to find the posterior distributions.
"""

import os
import sys
import argparse
import logging
import shutil

import numpy

import pycbc
from pycbc import (distributions, transforms, fft,
                   opt, scheme, pool)
from pycbc.waveform import generator

from pycbc import __version__
from pycbc import inference
from pycbc.inference import (models, burn_in, option_utils)
from pycbc.inference.io import loadfile
from pycbc.workflow import configuration

# command line usage
parser = argparse.ArgumentParser(usage=__file__ + " [--options]",
                                 description=__doc__)
parser.add_argument("--version", action="version", version=__version__,
                    help="Prints version information.")
parser.add_argument("--verbose", action="store_true", default=False,
                    help="Print logging messages.")
# 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.")
# 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.")
# add config options
configuration.add_workflow_command_line_group(parser)
# add module pre-defined options
fft.insert_fft_option_group(parser)
opt.insert_optimization_option_group(parser)
scheme.insert_processing_option_group(parser)

# parse command line
opts = parser.parse_args()

# setup log
# If we're running in MPI mode, only allow the parent to print
use_mpi, size, rank = pycbc.pool.use_mpi(opts.use_mpi, log=False)
if pool.is_main_process():
    pycbc.init_logging(opts.verbose)

# verify options are sane
fft.verify_fft_options(opts, parser)
opt.verify_optimization_options(opts, parser)
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.")


# set seed
numpy.random.seed(opts.seed)
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')

# get scheme
ctx = scheme.from_cli(opts)
fft.from_cli(opts)

with ctx:

    # read configuration file
    cp = configuration.WorkflowConfigParser.from_cli(opts)

    # create an empty checkpoint file, if needed
    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')))
        # create an empty output file to keep condor happy
        open(opts.output_file, 'a').close()

    logging.info("Setting up model")

    # construct class that will return the natural logarithm of likelihood
    model = models.read_from_config(cp)

    logging.info("Setting up sampler")

    # Create sampler that will run.
    # Note: the pool is created at this point. This means that,
    # 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)

    # store the checkpoint file
    if pool.is_main_process():
        for fn in [sampler.checkpoint_file, sampler.backup_file]:
            with loadfile(fn, 'a') as fp:
                fp.write_config_file(cp)

    # Run the sampler
    sampler.run()
    
    if not pool.is_main_process():
        sys.exit(0)

    # Finalize the output
    sampler.finalize()

if condor_ckpt:
   # 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")
os.rename(sampler.checkpoint_file, opts.output_file)
if not opts.save_backup:
    logging.info("Deleting backup file")
    os.remove(sampler.backup_file)

# write the end time
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()

# exit
logging.info("Done")