https://github.com/gwastro/pycbc
Revision b574af5aeb8e5b22c040f0e6cecd6796e6a5cb69 authored by Duncan Brown on 27 November 2018, 23:45:28 UTC, committed by GitHub on 27 November 2018, 23:45:28 UTC
1 parent b1a1759
Tip revision: b574af5aeb8e5b22c040f0e6cecd6796e6a5cb69 authored by Duncan Brown on 27 November 2018, 23:45:28 UTC
Set for 1.13.2 release
Set for 1.13.2 release
Tip revision: b574af5
pycbc_split_inspinj
#!/usr/bin/env /usr/bin/python
# Copyright (C) 2014 Andrew Lundgren
import sys
import argparse
from glue.ligolw import utils as ligolw_utils
from glue.ligolw import ligolw, lsctables
from itertools import cycle
import pycbc.version
# The following is required for reasons
class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
pass
# Parse command line
parser = argparse.ArgumentParser()
parser.add_argument("--version", action="version",
version=pycbc.version.git_verbose_msg)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument("-n", "--num-splits", type=int,
help="Number of files to be generated")
group.add_argument("-f", "--output-files", nargs='*', default=None,
help="Names of output files")
parser.add_argument("-i", "--input-file", help="Injection file to be split")
parser.add_argument("-o", "--output-dir", default=None,
help="Location of output files")
args = parser.parse_args()
if args.output_files and args.output_dir:
parser.error("Provide only one of --output-dir or --output-files")
# The XML reading will hang if you forget this line
lsctables.use_in(LIGOLWContentHandler)
# Read in input file
xmldoc = ligolw_utils.load_filename\
(args.input_file, verbose=True, contenthandler=LIGOLWContentHandler)
tabletype = lsctables.SimInspiralTable
allinjs = tabletype.get_table(xmldoc)
# The sim_inspiral table is a grandchild of the document, I guess
xmlroot = xmldoc.childNodes[0]
xmlroot.removeChild(allinjs)
if args.num_splits:
num_splits = args.num_splits
else:
num_splits = len(args.output_files)
new_inj_tables = [lsctables.New(tabletype, columns=allinjs.columnnames) \
for idx in xrange(num_splits)]
table_cycle = cycle(new_inj_tables)
for inj in sorted(allinjs, key=lambda x: x.get_time_geocent()):
table_cycle.next().append(inj)
if not args.output_files:
temp = args.input_file.split('-')
temp[1] += '_%.4u'
filename_pattern = '-'.join(temp)
for idx, simtable in enumerate(new_inj_tables):
xmlroot.appendChild(simtable)
if not args.output_files:
out_path = args.output_dir + '/' + filename_pattern % idx
else:
out_path = args.output_files[idx]
ligolw_utils.write_filename(xmldoc, out_path,
gz=out_path.endswith('gz'))
xmlroot.removeChild(simtable)
Computing file changes ...