#!/usr/bin/python
#
# Copyright (C) 2014 LIGO Scientific Collaboration
#
# 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.
#
# =============================================================================
#
# Preamble
#
# =============================================================================
#
"""
Splits a table in an xml file into multiple pieces
"""
import pycbc, pycbc.version, pycbc.pnutils
__author__ = "Alex Nitz <alex.nitz@ligo.org>"
__version__ = pycbc.version.git_verbose_msg
__date__ = pycbc.version.date
__program__ = "pycbc_splitbank"
import time
import argparse
from pycbc_glue import gpstime
from pycbc_glue.ligolw import ligolw
from pycbc_glue.ligolw import table
from pycbc_glue.ligolw import lsctables
from pycbc_glue.ligolw import utils as ligolw_utils
from pycbc_glue.ligolw.utils import process as ligolw_process
from pycbc_glue import git_version
from numpy import random, ceil
class ContentHandler(ligolw.LIGOLWContentHandler):
pass
lsctables.use_in(ContentHandler)
# Command line parsing
_desc = __doc__[1:]
parser = argparse.ArgumentParser(description=_desc)
parser.add_argument('--version', action='version', version=__version__)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('--templates-per-bank', metavar='SAMPLES',
help='number of templates in the output banks', type=int)
group.add_argument('-n', '--number-of-banks', metavar='N',
help='Split template bank into N files', type=int)
group.add_argument("-O", "--output-filenames", nargs='*', default=None,
action="store",
metavar="OUTPUT_FILENAME", help="""Directly specify the
names of the output files. The number of files specified
here will dictate how to split the bank. It will be split
equally between all specified files.""")
parser.add_argument("-o", "--output-prefix", default=None,
help="Prefix to add to the template bank name (name becomes output#.xml[.gz])" )
parser.add_argument("-z", "--write-compress", action="store_true",
help="Write compressed xml.gz files.")
parser.add_argument("-V", "--verbose", action="store_true",
help="Print extra debugging information", default=False )
parser.add_argument("-t", "--bank-file", metavar='INPUT_FILE',
help='Template bank to split', required=True)
parser.add_argument("--sort-frequency-cutoff",
help="Frequency cutoff to use for sorting the sub banks")
parser.add_argument("--sort-mchirp", action="store_true", default=False,
help='Sort templates by chirp mass before splitting')
parser.add_argument("--random-sort", action="store_true", default=False,
help='Sort templates randomly before splitting')
parser.add_argument("--random-seed", type=int,
help='Random seed to use when sorting randomly')
args = parser.parse_args()
if args.output_filenames and args.output_prefix:
errMsg="Cannot supply --output-filenames with --output-prefix."
parser.error(errMsg)
if args.sort_mchirp and args.random_sort:
errMsg="You can't sort by Mchirp *and* randomly, dumbass!"
parser.error(errMsg)
if args.output_filenames:
args.number_of_banks = len(args.output_filenames)
indoc = ligolw_utils.load_filename(args.bank_file, verbose=args.verbose,
contenthandler=ContentHandler)
try:
template_bank_table = table.get_table(indoc, lsctables.SnglInspiralTable.tableName)
tabletype = lsctables.SnglInspiralTable
except:
template_bank_table = table.get_table(indoc, lsctables.SimInspiralTable.tableName)
tabletype = lsctables.SimInspiralTable
length = len(template_bank_table)
def mchirp_sort(x, y):
mc1, e1 = pycbc.pnutils.mass1_mass2_to_mchirp_eta(x.mass1, x.mass2)
mc2, e2 = pycbc.pnutils.mass1_mass2_to_mchirp_eta(y.mass1, y.mass2)
return cmp(mc1, mc2)
def frequency_cutoff_sort(x, y):
p1 = pycbc.pnutils.frequency_cutoff_from_name(args.sort_frequency_cutoff,
x.mass1, x.mass2,
x.spin1z, x.spin2z)
p2 = pycbc.pnutils.frequency_cutoff_from_name(args.sort_frequency_cutoff,
y.mass1, y.mass2,
y.spin1z, y.spin2z)
return cmp(p1, p2)
tt = template_bank_table
if args.sort_frequency_cutoff:
tt = sorted(template_bank_table, cmp=frequency_cutoff_sort)
if args.sort_mchirp:
tt = sorted(template_bank_table, cmp=mchirp_sort)
if args.random_sort:
if args.random_seed is not None:
random.seed(args.random_seed)
random.shuffle(template_bank_table)
if args.number_of_banks:
# Decide how to split up the templates
# Put approximately the requested number of templates in each file
# But try to make each file very nearly the same size
num_files = args.number_of_banks
num_per_file = length / float(num_files)
elif args.templates_per_bank:
num_per_file = args.templates_per_bank
num_files = int(ceil(float(length) / num_per_file))
index_list = [int(round(num_per_file*idx)) for idx in range(num_files)]
index_list.append(length)
assert(index_list[0] == 0)
for num, (idx1, idx2) in enumerate(zip(index_list[:-1], index_list[1:])):
assert(idx2 > idx1)
# create a blank xml document and add the process id
outdoc = ligolw.Document()
outdoc.appendChild(ligolw.LIGO_LW())
proc_id = ligolw_process.register_to_xmldoc(outdoc,
__program__, args.__dict__, ifos=["G1"],
version=git_version.id, cvs_repository=git_version.branch,
cvs_entry_time=git_version.date).process_id
sngl_inspiral_table = lsctables.New(tabletype,columns=template_bank_table.columnnames)
outdoc.childNodes[0].appendChild(sngl_inspiral_table)
for i in range(idx2-idx1):
row = tt.pop()
sngl_inspiral_table.append(row)
# write the xml doc to disk
proctable = table.get_table(outdoc, lsctables.ProcessTable.tableName)
proctable[0].end_time = gpstime.GpsSecondsFromPyUTC(time.time())
if args.output_filenames:
outname = args.output_filenames[num]
elif args.output_prefix:
outname = args.output_prefix + str(num) + '.xml'
if args.write_compress:
outname +='.gz'
else:
errMsg = "Cannot figure out how to set output file names."
raise ValueError(errMsg)
ligolw_utils.write_filename(outdoc, outname, gz=outname.endswith('.gz'))