gstlal_inspiral_iir_bank_pipe

#!/usr/bin/env python
#
# Copyright (C) 2011-2012  Shaun Hooper, Chad Hanna
#
# 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 2 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.

## @file gstlal_inspiral_iir_bank_pipe
# A program to make an HTCondor DAG to create banks of IIR filters for an online CBC analysis
#
# ### Command line interface
#   + `--flow` [Hz] (float): Set the template low-frequency cut-off (default = 40.0).
#   + `--sampleRate` [Hz] (float): Set the sample rate of the IIR template bank (optional).
#   + `--padding` [pad] (float): Fractional amount to pad time slices (default = 1.1).
#   + `--epsilon` [pad] (float): Second order correction factor (default = 0.02).
#   + `--reference-psd` [file]: Set the name of the reference psd file, required.
#   + `--bank-cache` [file]: Set the name of the bank cache, required.
#   + `--output-name`: Set the base name of the output, required.
#   + `--instrument`: Set the name of the instrument, required.
#   + `--verbose`: Be verbose.
#   + `--downsample`: Choose if you want to downsample IIR bank (recommended).

"""
This program makes a dag to generate iir banks
"""

__author__ = 'Chad Hanna <channa@caltech.edu>, Shaun Hooper <shaun.hooper@uwa.edu.au>'

##############################################################################
# import standard modules and append the lalapps prefix to the python path
import sys, os, copy, math
import subprocess, socket, tempfile

##############################################################################
# import the modules we need to build the pipeline
from glue import iterutils
from glue import pipeline
from glue import lal
from glue.ligolw import lsctables
from glue import segments
from optparse import OptionParser

def which(prog):
    which = subprocess.Popen(['which',prog], stdout=subprocess.PIPE)
    out = which.stdout.read().strip()
    if not out: 
        print >>sys.stderr, "ERROR: could not find %s in your path, have you built the proper software and source the proper env. scripts?" % (prog,prog)
        raise ValueError 
    return out

def log_path():
    host = socket.getfqdn()
    #FIXME add more hosts as you need them
    if 'caltech.edu' in host: return '/usr1/' + os.environ['USER']
    if 'phys.uwm.edu' in host: return '/localscratch/' + os.environ['USER']
    if 'aei.uni-hannover.de' in host: return '/local/user/' + os.environ['USER']
    if 'phy.syr.edu' in host: return '/usr1/' + os.environ['USER']


class bank_DAG(pipeline.CondorDAG):

    def __init__(self, name, logpath = log_path()):
        self.basename = name
        tempfile.tempdir = logpath
        tempfile.template = self.basename + '.dag.log.'
        logfile = tempfile.mktemp()
        fh = open( logfile, "w" )
        fh.close()
        pipeline.CondorDAG.__init__(self,logfile)
        self.set_dag_file(self.basename)
        self.jobsDict = {}
        self.node_id = 0
        self.output_cache = []

    def add_node(self, node):
        node.set_retry(3)
        self.node_id += 1
        node.add_macro("macroid", self.node_id)
        pipeline.CondorDAG.add_node(self, node)

    def write_cache(self):
        out = self.basename + ".cache"
        f = open(out,"w")
        for c in self.output_cache:
            f.write(str(c)+"\n")
        f.close()

class gstlal_iir_bank_job(pipeline.CondorDAGJob):
    """
    A gstlal_iir_bank job
    """
    def __init__(self, executable=which('gstlal_iir_bank'), tag_base='gstlal_iir_bank'):
        """
        """
        self.__prog__ = 'gstlal_iir_bank'
        self.__executable = executable
        self.__universe = 'vanilla'
        pipeline.CondorDAGJob.__init__(self,self.__universe,self.__executable)
        self.add_condor_cmd('getenv','True')
        self.add_condor_cmd('requirements', 'Memory > 1999') #FIXME is this enough?
        self.tag_base = tag_base
        self.add_condor_cmd('environment',"KMP_LIBRARY=serial;MKL_SERIAL=yes")
        self.set_sub_file(tag_base+'.sub')
        self.set_stdout_file('logs/'+tag_base+'-$(macroid)-$(process).out')
        self.set_stderr_file('logs/'+tag_base+'-$(macroid)-$(process).err')


class gstlal_iir_bank_node(pipeline.CondorDAGNode):
    """
    """
    def __init__(self, job, dag, template_bank, ifo, flow, reference_psd, sampleRate, padding, downsample, epsilon=0.02, p_node=[]):

        pipeline.CondorDAGNode.__init__(self,job)
        self.add_var_opt("template-bank", template_bank)
        self.add_var_opt("instrument", ifo)
        self.add_var_opt("flow", flow)
        self.add_var_opt("reference-psd", reference_psd)
        self.add_var_opt("sampleRate",sampleRate)
        self.add_var_opt("padding",padding)
        if downsample: self.add_var_opt("downsample","")
        self.add_var_opt("epsilon", epsilon)
        iir_bank_name_path = os.path.split(template_bank)
        iir_bank_name = iir_bank_name_path[0] + "/iir_" + iir_bank_name_path[1]
        self.add_var_opt("output", iir_bank_name)
        dag.output_cache.append(lal.CacheEntry(ifo, "-", segments.segment(0, 999999999), "file://localhost%s" % (iir_bank_name,)))
        for p in p_node:
            self.add_parent(p)
        dag.add_node(self)

def parse_command_line():
    parser = OptionParser()
    parser.add_option("--flow", metavar = "Hz", type = "float", default = 40.0, help = "Set the template low-frequency cut-off (default = 40.0).")
    parser.add_option("--sampleRate", metavar = "Hz", type = "float", help = "Set the sample rate of the IIR template bank (optional).")
    parser.add_option("--padding", metavar = "pad", type = "float", default = 1.1, help = "Fractional amount to pad time slices.")
    parser.add_option("--epsilon", metavar = "pad", type = "float", default = 0.02, help = "Second order correction factor.")
    parser.add_option("--reference-psd", metavar = "file", help = "Set the name of the reference psd file, required")
    parser.add_option("--bank-cache", metavar = "file", help = "Set the name of the bank cache, required")
    parser.add_option("--output-name", help = "set the base name of the output, required")
    parser.add_option("--instrument", help = "set the name of the instrument, required")
    parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
    parser.add_option("--downsample", action = "store_true", help = "Choose if you want to downsample IIR bank (recommended)")
    options, filenames = parser.parse_args()

    return options, filenames

options, filenames = parse_command_line()

# get input arguments
ifo = options.instrument
ref_psd = options.reference_psd
input_cache = options.bank_cache
psdname = ref_psd.split('.')[0]

try: os.mkdir("logs")
except: pass
dag = bank_DAG(options.output_name)

iirJob = gstlal_iir_bank_job(tag_base=psdname + "_gstlal_iir_bank")
iirNode = {}


for f in [lal.CacheEntry(line).path for line in open(input_cache)]:
    print f
    iirNode[f] = gstlal_iir_bank_node(iirJob, dag, f, ifo, options.flow, ref_psd, options.sampleRate, options.padding, options.downsample, options.epsilon)

dag.write_sub_files()
dag.write_dag()
dag.write_script()
dag.write_cache()