Bank inj update

bin/workflows/pycbc_make_bank_verifier_workflow

@@ -26,6 +26,7 @@ template bank.
 from __future__ import division
 import os
 import argparse
+import logging
 
 from glue import segments
 
@@ -34,10 +35,12 @@ import pycbc.workflow as wf
 import pycbc.workflow.pegasus_workflow as pwf
 from pycbc.results import create_versioning_page, static_table, layout
 from pycbc.workflow import LalappsInspinjExecutable
+from pycbc.workflow import PycbcCreateInjectionsExecutable
 from pycbc.workflow import setup_splittable_dax_generated
 
 # Boiler-plate stuff
-__author__  = "Ian Harry <[email protected]>"
+__author__ = "Ian Harry <[email protected]> and "
+__author__ += "Sebastian Khan <[email protected]>"
 __version__ = pycbc.version.git_verbose_msg
 __date__    = pycbc.version.date
 __program__ = "pycbc_make_bank_verifier_workflow"
@@ -148,7 +151,7 @@ class BanksimTablePointInjsExecutable(wf.Executable):
 
 # Argument parsing and setup of workflow
 
-# Use the standard workflow command-line parsing routines. Things like a 
+# Use the standard workflow command-line parsing routines. Things like a
 # configuration file are specified within the "workflow command line group"
 # so run this with --help to see what options are added.
 _desc = __doc__[1:]
@@ -158,9 +161,15 @@ parser.add_argument("--workflow-name", type=str, default='bank_verifier',
                     help="Descriptive name of the analysis.")
 parser.add_argument("-d", "--output-dir", default=None,
                     help="Path to output directory.")
+
+parser.add_argument("--verbose", action="store_true", default=False,
+                    help="Print logging messages.")
+
 wf.add_workflow_command_line_group(parser)
 args = parser.parse_args()
 
+pycbc.init_logging(args.verbose)
+
 # Create the workflow object
 workflow = wf.Workflow(args, args.workflow_name)
 
@@ -180,20 +189,57 @@ inp_bank = wf.File.from_path(inp_bank)
 inp_bank.tags = []
 inp_bank.description='TEMPLATEBANK'
 inp_bank.ifo_list=(['H1','L1','V1'])
+# seems strange to hardcode this to 3 IFO
+# maybe derive this from [workflow-ifos]
 inp_bank.segment = workflow.analysis_time
 
-# Inspinj job
-inspinj_job = LalappsInspinjExecutable(workflow.cp, 'injection', out_dir='.',
-                                       ifos='HL', tags=[])
+# Injection Job
+injection_exe = workflow.cp.get('executables', 'injection')
+injection_exe = os.path.basename(injection_exe)
+logging.info("Using {} for injection generation.".format(injection_exe))
+
+if injection_exe == "lalapps_inspinj":
+    injection_job = LalappsInspinjExecutable(workflow.cp,
+        'injection', out_dir='.', ifos='HL', tags=[])
+elif injection_exe == "pycbc_create_injections":
+    injection_job = PycbcCreateInjectionsExecutable(
+        workflow.cp, 'injection', out_dir=".", ifo="HL", tags=[])
+else:
+    err_msg = "Injection executable {} is not supported.".format(injection_exe)
+    err_msg += " Please use, lalapps_inspinj or pycbc_create_injections."
+    raise ValueError(err_msg)
+
 
 def add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks):
     """Add a group of jobs that does a complete banksim.
     """
-    inspinj_job.update_current_tags(curr_tags)
-    t_seg = segments.segment([1000000000, 1000000000+int(num_injs)])
-    node = inspinj_job.create_node(t_seg)
+    # curr_tags are like [section_heading, injection_set_name]
+    injection_job.update_current_tags(curr_tags)
+
+    start_time = 1000000000
+    end_time = 1000000000+int(num_injs)
+
+    if injection_exe == "lalapps_inspinj":
+        t_seg = segments.segment([start_time, end_time])
+        node = injection_job.create_node(t_seg)
+        inj_file = node.output_file
+
+    elif injection_exe == "pycbc_create_injections":
+        # Set a default value. When .hdf is supported this should be .hdf
+        inj_file_ext = ".xml"
+        # Allow overrides from config file
+        if workflow.cp.has_option_tags("workflow-injection",
+                                       "output-file-format", curr_tags):
+            inj_file_ext = workflow.cp.get_option_tags("workflow-injection",
+                                                       "output-file-format",
+                                                       curr_tags)
+        node = injection_job.create_node(ext=inj_file_ext,
+                                         start_time=start_time,
+                                         end_time=end_time,
+                                         num_injs=num_injs)
+        inj_file = node.output_file
+
     workflow += node
-    inj_file = node.output_file
     split_injs = setup_splittable_dax_generated(workflow, [inj_file],
                                                 'splitinjfiles', curr_tags)
     # Banksim job
@@ -215,7 +261,6 @@ def add_banksim_set(workflow, file_tag, num_injs, curr_tags, split_banks):
         currinj_banksim_files = wf.FileList([])
         for bank_idx, split_bank in enumerate(split_banks):
             bank_tag = 'BANK{}'.format(bank_idx)
-            inj_tag = 'INJ{}'.format(inj_idx)
             node = banksim_job.create_node(workflow.analysis_time, split_inj,
                                            split_bank,
                                            extra_tags=[bank_tag,inj_tag])
@@ -239,17 +284,19 @@ split_banks = setup_splittable_dax_generated(workflow, [inp_bank],
 output_pointinjs = {}
 for file_tag, num_injs in workflow.cp.items('workflow-pointinjs'):
     curr_tags = ['shortinjs', file_tag]
+    logging.info("working curr_tags = {}".format(curr_tags))
     curr_file = add_banksim_set(workflow, file_tag, num_injs, curr_tags,
                                 split_banks)
     output_pointinjs[file_tag] = curr_file
 
 curr_tags = ['broadinjbanksplit']
-split_banks = setup_splittable_dax_generated(workflow, [inp_bank], 
+split_banks = setup_splittable_dax_generated(workflow, [inp_bank],
                                              'splitbankfiles', curr_tags)
 
 output_broadinjs = {}
 for file_tag, num_injs in workflow.cp.items('workflow-broadinjs'):
     curr_tags = ['broadinjs', file_tag]
+    logging.info("working curr_tags = {}".format(curr_tags))
     curr_file = add_banksim_set(workflow, file_tag, num_injs, curr_tags,
                                 split_banks)
     output_broadinjs[file_tag] = curr_file

pycbc/workflow/jobsetup.py

@@ -1622,23 +1622,39 @@ class PycbcCreateInjectionsExecutable(Executable):
     """
 
     current_retention_level = Executable.ALL_TRIGGERS
+    file_input_options = ["--config-files"]
     def __init__(self, cp, exe_name, ifo=None, out_dir=None,
                  universe=None, tags=None):
         super(PycbcCreateInjectionsExecutable, self).__init__(
                                cp, exe_name, universe, ifo, out_dir, tags)
 
-    def create_node(self, config_file=None, seed=None, tags=None):
+    def create_node(self, seed=None, tags=None, ext=".hdf",
+                    start_time=None, end_time=None, num_injs=None):
         """ Set up a CondorDagmanNode class to run ``pycbc_create_injections``.
 
         Parameters
         ----------
-        config_file : pycbc.workflow.core.File
-            A ``pycbc.workflow.core.File`` for inference configuration file
-            to be used with ``--config-files`` option.
         seed : int
             Seed to use for generating injections.
         tags : list
             A list of tags to include in filenames.
+        ext : str
+            Output file extension. Use '.hdf' or '.xml' for sim_inspiral table
+            ( Default = '.hdf' )
+        start_time : int
+            Use this to override the value of the 'start-time' given in
+            the 'config_file' [workflow] section
+            ( Default = None )
+        end_time : int
+            Use this to override the value of the 'end-time' given in
+            the 'config_file' [workflow] section
+            ( Default = None )
+        num_injs : int
+            Use this to specify the number of injections (the --ninjections
+            command line argument). If not provided (or set to None), this
+            command line argument is not set and it is assumed this will be
+            specified in the configuration file.
+            ( Default = None )
 
         Returns
         --------
@@ -1650,20 +1666,24 @@ def create_node(self, config_file=None, seed=None, tags=None):
         tags = [] if tags is None else tags
 
         # get analysis start and end time
-        start_time = self.cp.get("workflow", "start-time")
-        end_time = self.cp.get("workflow", "end-time")
+        if start_time is None:
+            start_time = self.cp.get("workflow", "start-time")
+        if end_time is None:
+            end_time = self.cp.get("workflow", "end-time")
         analysis_time = segments.segment(int(start_time), int(end_time))
 
+        if num_injs is not None:
+            node.add_opt("--ninjections", num_injs)
+
         # make node for running executable
         node = Node(self)
-        node.add_input_opt("--config-file", config_file)
         if seed:
             node.add_opt("--seed", seed)
         injection_file = node.new_output_file_opt(analysis_time,
-                                                  ".hdf", "--output-file",
+                                                  ext, "--output-file",
                                                   tags=tags)
 
-        return node, injection_file
+        return node
 
 class PycbcInferenceExecutable(Executable):
     """ The class responsible for creating jobs for ``pycbc_inference``.