tools.py

Go to the documentation of this file.

"""! Tools that can be used in HPSMC jobs."""
 
import os
import gzip
import shutil
import subprocess
import tarfile
 
from subprocess import PIPE
 
from hpsmc.component import Component
import hpsmc.func as func
 
 
class SLIC(Component):
    """!
    Run the SLIC Geant4 simulation.
 
    Optional parameters are: **nevents**, **macros**, **run_number**, **disable_particle_table** \n
    Required parameters are: **detector** \n
    Required configurations are: **slic_dir**, **detector_dir**
    """
 
    def __init__(self, **kwargs):
        
        self.macrosmacros = []
        
        self.run_numberrun_number = None
        
        self.detector_dirdetector_dir = None
        
        self.disable_particle_tabledisable_particle_table = False
 
        Component.__init__(self,
                           name='slic',
                           command='slic',
                           output_ext='.slcio',
                           **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        if not len(self.input_filesinput_files()):
            raise Exception("No inputs given for SLIC.")
 
        args = ["-g", self.__detector_file__detector_file(),
                # "-i", self.input_files()[0],
                "-o", self.output_filesoutput_files()[0],
                "-d%s" % str(self.seedseed)]
 
        if self.neventsnevents is not None:
            args.extend(["-r", str(self.neventsnevents)])
 
        if self.run_numberrun_number is not None:
            args.extend(["-m", "run_number.mac"])
 
        if not self.disable_particle_tabledisable_particle_table:
            tbl = self.__particle_tbl__particle_tbl()
            if os.path.exists(tbl):
                args.extend(["-P", tbl])
            else:
                raise Exception('SLIC particle.tbl does not exist: %s' % tbl)
 
        if len(self.macrosmacros):
            # args = []
            for macro in self.macrosmacros:
                if macro == "run_number.mac":
                    raise Exception("Macro name '%s' is not allowed." % macro)
                if not os.path.isabs(macro):
                    raise Exception("Macro '%s' is not an absolute path." % macro)
                args.extend(["-m", macro])
        else:
            args.extend(["-i", self.input_filesinput_files()[0]])
 
        return args
 
    def __detector_file(self):
        """! Return path to detector file."""
        return os.path.join(self.detector_dirdetector_dir, self.detector, self.detector + ".lcdd")
 
    def __particle_tbl(self):
        """! Return path to particle table."""
        return os.path.join(self.slic_dir, "share", "particle.tbl")
 
    def config(self, parser):
        """! Configure SLIC component."""
        super().config(parser)
 
        if self.detector_dirdetector_dir is None:
            self.detector_dirdetector_dir = "{}/share/detectors".format(self.hpsmc_dirhpsmc_dir)
            if not os.path.isdir(self.detector_dirdetector_dir):
                raise Exception('Failed to find valid detector_dir')
            self.loggerlogger.debug("Using detector_dir from install: {}".format(self.detector_dirdetector_dir))
 
    def setup(self):
        """! Setup SLIC component."""
        if not os.path.exists(self.slic_dir):
            raise Exception("slic_dir does not exist: %s" % self.slic_dir)
 
        self.env_scriptenv_script = self.slic_dir + os.sep + "bin" + os.sep + "slic-env.sh"
        if not os.path.exists(self.env_scriptenv_script):
            raise Exception('SLIC setup script does not exist: %s' % self.namename)
 
        if self.run_numberrun_number is not None:
            run_number_cmd = "/lcio/runNumber %d" % self.run_numberrun_number
            run_number_mac = open("run_number.mac", 'w')
            run_number_mac.write(run_number_cmd)
            run_number_mac.close()
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **nevents**, **macros**, **run_number**
        @return  list of optional parameters
        """
        return ['nevents', 'macros', 'run_number', 'disable_particle_table']
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: **detector**
        @return  list of required parameters
        """
        return ['detector']
 
    def required_config(self):
        """!
        Return list of required configurations.
 
        Required configurations are: **slic_dir**, **detector_dir**
        @return  list of required configurations
        """
        return ['slic_dir', 'detector_dir']
 
    def execute(self, log_out, log_err):
        """!
        Execute SLIC component.
 
        Component is executed by creating command line input
        from command and command arguments.
        @return  return code of process
        """
        # SLIC needs to be run inside bash as the Geant4 setup script is a piece of #@$@#$.
        cl = 'bash -c ". %s && %s %s"' % (self.env_scriptenv_script, self.commandcommand, ' '.join(self.cmd_argscmd_argscmd_args()))
 
        # self.logger.info("Executing '%s' with command: %s" % (self.name, cl))
        proc = subprocess.Popen(cl, shell=True, stdout=log_out, stderr=log_err)
        proc.communicate()
        proc.wait()
 
        return proc.returncode
 
 
class JobManager(Component):
    """!
    Run the hps-java JobManager class.
 
    Input files have slcio format.
 
    Required parameters are: **steering_files** \n
    Optional parameters are: **detector**, **run_number**, **defs**
    """
 
    def __init__(self, steering=None, **kwargs):
        
        self.run_numberrun_number = None
        
        self.neventsneventsnevents = None
        
        self.detectordetector = None
        
        self.event_print_intervalevent_print_interval = None
        
        self.defsdefs = None
        
        self.java_argsjava_args = None
        
        self.logging_config_filelogging_config_file = None
        
        self.lcsim_cache_dirlcsim_cache_dir = None
        
        self.conditions_userconditions_user = None
        
        self.conditions_passwordconditions_password = None
        
        self.conditions_urlconditions_url = None
        
        self.steeringsteering = steering
        
        self.hps_java_bin_jarhps_java_bin_jar = None
 
        if 'overlay_file' in kwargs:
            self.overlay_fileoverlay_file = kwargs['overlay_file']
        else:
            self.overlay_fileoverlay_file = None
 
        Component.__init__(self,
                           name='job_manager',
                           command='java',
                           description='HPS Java Job Manager',
                           output_ext='.slcio',
                           **kwargs)
 
        # Automatically append steering file key to output file name
        if self.append_tokappend_tokappend_tok is None:
            self.append_tokappend_tokappend_tok = self.steeringsteering
            self.loggerlogger.debug("Append token for '%s' automatically set to '%s' from steering key." % (self.namename, self.append_tokappend_tokappend_tok))
 
    def config(self, parser):
        """! Configure JobManager component."""
        super().config(parser)
        # if installed these are set in the environment script...
        if self.hps_java_bin_jarhps_java_bin_jar is None:
            if os.getenv('HPS_JAVA_BIN_JAR', None) is not None:
                self.hps_java_bin_jarhps_java_bin_jar = os.getenv('HPS_JAVA_BIN_JAR', None)
                self.loggerlogger.debug('Set HPS_JAVA_BIN_JAR from environment: {}'.format(self.hps_java_bin_jarhps_java_bin_jar))
            else:
                raise Exception('hps_java_bin_jar not set in environment or config file!')
        if self.conditions_urlconditions_url is None:
            if os.getenv("CONDITIONS_URL", None) is not None:
                self.conditions_urlconditions_url = os.getenv("CONDITIONS_URL", None)
                self.loggerlogger.debug('Set CONDITIONS_URL from environment: {}'.format(self.hps_java_bin_jarhps_java_bin_jar))
 
    def required_config(self):
        """!
        Return list of required configurations.
 
        Required configurations are: **hps_java_bin_jar**
        @retun list of required configurations.
        """
        return ['hps_java_bin_jar']
 
    def setup(self):
        """! Setup JobManager component."""
        if not len(self.input_filesinput_files()):
            raise Exception("No inputs provided to hps-java.")
 
        if self.steeringsteering not in self.steering_files:
            raise Exception("Steering '%s' not found in: %s" % (self.steeringsteering, self.steering_files))
        self.steering_filesteering_file = self.steering_files[self.steeringsteering]
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = []
 
        if self.java_argsjava_args is not None:
            self.loggerlogger.debug('Setting java_args from config: %s' % self.java_argsjava_args)
            args.append(self.java_argsjava_args)
 
        if self.logging_config_filelogging_config_file is not None:
            self.loggerlogger.debug('Setting logging_config_file from config: %s' % self.logging_config_filelogging_config_file)
            args.append('-Djava.util.logging.config.file=%s' % self.logging_config_filelogging_config_file)
 
        if self.lcsim_cache_dirlcsim_cache_dir is not None:
            self.loggerlogger.debug('Setting lcsim_cache_dir from config: %s' % self.lcsim_cache_dirlcsim_cache_dir)
            args.append('-Dorg.lcsim.cacheDir=%s' % self.lcsim_cache_dirlcsim_cache_dir)
 
        if self.conditions_userconditions_user is not None:
            self.loggerlogger.debug('Setting conditions_user from config: %s' % self.conditions_userconditions_user)
            args.append('-Dorg.hps.conditions.user=%s' % self.conditions_userconditions_user)
        if self.conditions_passwordconditions_password is not None:
            self.loggerlogger.debug('Setting conditions_password from config (not shown)')
            args.append('-Dorg.hps.conditions.password=%s' % self.conditions_passwordconditions_password)
        if self.conditions_urlconditions_url is not None:
            self.loggerlogger.debug('Setting conditions_url from config: %s' % self.conditions_urlconditions_url)
            args.append('-Dorg.hps.conditions.url=%s' % self.conditions_urlconditions_url)
 
        args.append("-jar")
        args.append(self.hps_java_bin_jarhps_java_bin_jar)
 
        
        if self.event_print_intervalevent_print_interval is not None:
            args.append("-e")
            args.append(str(self.event_print_intervalevent_print_interval))
 
        if self.run_numberrun_number is not None:
            args.append("-R")
            args.append(str(self.run_numberrun_number))
 
        if self.detectordetector is not None:
            args.append("-d")
            args.append(self.detectordetector)
 
        if len(self.output_filesoutput_files()):
            args.append("-D")
            args.append("outputFile=" + os.path.splitext(self.output_filesoutput_files()[0])[0])
 
        if self.defsdefs:
            for k, v in self.defsdefs.items():
                args.append("-D")
                args.append(k + "=" + str(v))
 
        if not os.path.isfile(self.steering_filesteering_file):
            args.append("-r")
            self.loggerlogger.debug("Steering does not exist at '%s' so assuming it is a resource." % self.steering_filesteering_file)
        else:
            if not os.path.isabs(self.steering_filesteering_file):
                raise Exception('Steering looks like a file but is not an abs path: %s' % self.steering_filesteering_file)
        args.append(self.steering_filesteering_file)
 
        if self.neventsneventsnevents is not None:
            args.append("-n")
            args.append(str(self.neventsneventsnevents))
 
        for input_file in self.input_filesinput_files():
            args.append("-i")
            args.append(input_file)
 
        if self.overlay_fileoverlay_file is not None:
            args.append("-D")
            args.append("overlayFile=" + os.path.splitext(self.overlay_fileoverlay_file)[0])
 
        return args
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: **steering_files**
        @return  list of required parameters
        """
        return ['steering_files']
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **detector**, **run_number**, **defs**
        @return list of optional parameters
        """
        return ['detector', 'run_number', 'defs', 'nevents']
 
 
class HPSTR(Component):
    """!
    Run the hpstr analysis tool.
 
    Required parameters are: **config_files** \n
    Optional parameters are: **year**, **is_data**, **nevents** \n
    Required configs are: **hpstr_install_dir**, **hpstr_base**
    """
 
    def __init__(self, cfg=None, is_data=0, year=None, tracking=None, **kwargs):
        
        self.cfgcfg = cfg
        
        self.is_datais_data = is_data
        
        self.yearyear = year
        
        self.trackingtracking = tracking
 
        self.hpstr_install_dirhpstr_install_dir = None
        self.hpstr_basehpstr_base = None
 
        Component.__init__(self,
                           name='hpstr',
                           command='hpstr',
                           **kwargs)
 
    def setup(self):
        """! Setup HPSTR component."""
        if not os.path.exists(self.hpstr_install_dirhpstr_install_dir):
            raise Exception('hpstr_install_dir does not exist: %s' % self.hpstr_install_dirhpstr_install_dir)
        self.env_scriptenv_script = self.hpstr_install_dirhpstr_install_dir + os.sep + "bin" + os.sep + "hpstr-env.sh"
 
        # The config file to use is read from a dict in the JSON parameters.
        if self.cfgcfg not in self.config_files:
            raise Exception("Config '%s' was not found in: %s" % (self.cfgcfg, self.config_files))
        config_file = self.config_files[self.cfgcfg]
        if len(os.path.dirname(config_file)):
            # If there is a directory name then we expect an absolute path not in the hpstr dir.
            if os.path.isabs(config_file):
                self.cfg_pathcfg_path = config_file
            else:
                # The config must be an abs path.
                raise Exception('The config has a directory but is not an abs path: %s' % self.cfgcfg)
        else:
            # Assume the cfg file is within the hpstr base dir.
            self.cfg_pathcfg_path = os.path.join(self.hpstr_basehpstr_base, "processors", "config", config_file)
        self.loggerlogger.debug('Set config path: %s' % self.cfg_pathcfg_path)
 
        # For ROOT output, automatically append the cfg key from the job params.
        if os.path.splitext(self.input_filesinput_files()[0])[1] == '.root':
            self.append_tokappend_tokappend_tok = self.cfgcfg
            self.loggerlogger.debug('Automatically appending token to output file: %s' % self.append_tokappend_tokappend_tok)
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: **config_files**
        @return  list of required parameters
        """
        return ['config_files']
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **year**, **is_data**, **nevents**
        @return  list of optional parameters
        """
        return ['year', 'is_data', 'nevents', 'tracking']
 
    def required_config(self):
        """!
        Return list of required configs.
 
        Required configs are: **hpstr_install_dir**, **hpstr_base**
        @return  list of required configs
        """
        return ['hpstr_install_dir', 'hpstr_base']
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = [self.cfg_pathcfg_path,
                "-t", str(self.is_datais_data),
                "-i", self.input_filesinput_files()[0],
                "-o", self.output_filesoutput_filesoutput_files()[0]]
        if self.neventsnevents is not None:
            args.extend(["-n", str(self.neventsnevents)])
        if self.yearyear is not None:
            args.extend(["-y", str(self.yearyear)])
        if self.trackingtracking is not None:
            args.extend(["-w", str(self.trackingtracking)])
        return args
 
    def output_files(self):
        """! Adjust names of output files."""
        f, ext = os.path.splitext(self.input_filesinput_files()[0])
        if '.slcio' in ext:
            return ['%s.root' % f]
        else:
            if not self.append_tokappend_tokappend_tok:
                self.append_tokappend_tokappend_tok = self.cfgcfg
            return ['%s_%s.root' % (f, self.append_tokappend_tokappend_tok)]
 
    def execute(self, log_out, log_err):
        """! Execute HPSTR component."""
        args = self.cmd_argscmd_argscmd_args()
        cl = 'bash -c ". %s && %s %s"' % (self.env_scriptenv_script, self.commandcommand,
                                          ' '.join(self.cmd_argscmd_argscmd_args()))
 
        self.loggerlogger.debug("Executing '%s' with command: %s" % (self.namename, cl))
        proc = subprocess.Popen(cl, shell=True, stdout=log_out, stderr=log_err)
        proc.communicate()
        proc.wait()
 
        return proc.returncode
 
 
 
 
class StdHepTool(Component):
    """!
    Generic class for StdHep tools.
    """
 
    
    seed_names = ['beam_coords',
                  'beam_coords_old',
                  'lhe_tridents',
                  'lhe_tridents_displacetime',
                  'lhe_tridents_displaceuni',
                  'merge_poisson',
                  'mix_signal',
                  'random_sample']
 
    def __init__(self, name=None, **kwargs):
 
        Component.__init__(self,
                           name=name,
                           command="stdhep_" + name,
                           **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = []
 
        if self.namename in StdHepTool.seed_names:
            args.extend(["-s", str(self.seedseed)])
 
        if len(self.output_filesoutput_files()) == 1:
            args.insert(0, self.output_filesoutput_files()[0])
        elif len(self.output_filesoutput_files()) > 1:
            raise Exception("Too many outputs specified for StdHepTool.")
        else:
            raise Exception("No outputs specified for StdHepTool.")
 
        if len(self.input_filesinput_files()):
            for i in self.inputsinputs[::-1]:
                args.insert(0, i)
        else:
            raise Exception("No inputs specified for StdHepTool.")
 
        return args
 
 
class BeamCoords(StdHepTool):
    """!
    Transform StdHep events into beam coordinates.
 
    Optional parameters are: **beam_sigma_x**, **beam_sigma_y**, **beam_rot_x**,
    **beam_rot_y**, **beam_rot_z**, **target_x**, **target_y**, **target_z**
    """
 
    def __init__(self, **kwargs):
        
        self.beam_sigma_xbeam_sigma_x = None
        
        self.beam_sigma_ybeam_sigma_y = None
        
        self.target_xtarget_x = None
        
        self.target_ytarget_y = None
        
        self.target_ztarget_z = None
        
        self.beam_rot_xbeam_rot_x = None
        
        self.beam_rot_ybeam_rot_y = None
        
        self.beam_rot_zbeam_rot_z = None
 
        StdHepTool.__init__(self,
                            name='beam_coords',
                            append_tok='rot',
                            **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = StdHepTool.cmd_args(self)
 
        if self.beam_sigma_xbeam_sigma_x is not None:
            args.extend(['-x', str(self.beam_sigma_xbeam_sigma_x)])
        if self.beam_sigma_ybeam_sigma_y is not None:
            args.extend(['-y', str(self.beam_sigma_ybeam_sigma_y)])
 
        if self.beam_rot_xbeam_rot_x is not None:
            args.extend(['-u', str(self.beam_rot_xbeam_rot_x)])
        if self.beam_rot_ybeam_rot_y is not None:
            args.extend(['-v', str(self.beam_rot_ybeam_rot_y)])
        if self.beam_rot_zbeam_rot_z is not None:
            args.extend(['-w', str(self.beam_rot_zbeam_rot_z)])
 
        if self.target_xtarget_x is not None:
            args.extend(['-X', str(self.target_xtarget_x)])
        if self.target_ytarget_y is not None:
            args.extend(['-Y', str(self.target_ytarget_y)])
        if self.target_ztarget_z is not None:
            args.extend(['-Z', str(self.target_ztarget_z)])
 
        return args
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **beam_sigma_x**, **beam_sigma_y**, **beam_rot_x**,
        **beam_rot_y**, **beam_rot_z**, **target_x**, **target_y**, **target_z**
        @return list of optional parameters
        """
        return ['beam_sigma_x', 'beam_sigma_y', 'beam_rot_x',
                'beam_rot_y', 'beam_rot_z',
                'target_x', 'target_y', 'target_z']
 
 
class RandomSample(StdHepTool):
    """!
    Randomly sample StdHep events into a new file.
 
    Optional parameters are: **nevents**, **mu**
    """
 
    def __init__(self, **kwargs):
        StdHepTool.__init__(self,
                            name='random_sample',
                            append_tok='sampled',
                            **kwargs)
        
        self.mumu = None
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = []
 
        if self.namename in StdHepTool.seed_names:
            args.extend(["-s", str(self.seedseed)])
 
        args.extend(["-N", str(1)])
 
        if self.neventsnevents is not None:
            args.extend(["-n", str(self.neventsnevents)])
 
        if self.mumu is not None:
            args.extend(["-m", str(self.mumu)])
 
        if len(self.output_filesoutput_files()) == 1:
            # only use file name, not extension because extension is added by tool
            args.insert(0, os.path.splitext(self.output_filesoutput_files()[0])[0])
        elif len(self.output_filesoutput_files()) > 1:
            raise Exception("Too many outputs specified for RandomSample.")
        else:
            raise Exception("No outputs specified for RandomSample.")
 
        if len(self.input_filesinput_files()):
            for i in self.inputsinputs[::-1]:
                args.insert(0, i)
        else:
            raise Exception("No inputs were provided.")
 
        return args
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **nevents**, **mu**
        @return list of optional parameters
        """
        return ['nevents', 'mu']
 
    def execute(self, log_out, log_err):
        """! Execute RandomSample component"""
        returncode = Component.execute(self, log_out, log_err)
 
        # Move file to proper output file location.
        src = '%s_1.stdhep' % os.path.splitext(self.output_filesoutput_files()[0])[0]
        dest = '%s.stdhep' % os.path.splitext(self.output_filesoutput_files()[0])[0]
        self.loggerlogger.debug("Moving '%s' to '%s'" % (src, dest))
        shutil.move(src, dest)
 
        return returncode
 
 
class DisplaceTime(StdHepTool):
    """!
    Convert LHE files to StdHep, displacing the time by given ctau.
 
    Optional parameters are: **ctau**
    """
 
    def __init__(self, **kwargs):
        
        self.ctauctau = None
        StdHepTool.__init__(self,
                            name='lhe_tridents_displacetime',
                            output_ext='.stdhep',
                            **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = StdHepTool.cmd_args(self)
        if self.ctauctau is not None:
            args.extend(["-l", str(self.ctauctau)])
        return args
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **ctau**
        @return list of optional parameters
        """
        return ['ctau']
 
 
class DisplaceUni(StdHepTool):
    """!
    Convert LHE files to StdHep, displacing the time by given ctau.
 
    Optional parameters are: **ctau**
    """
 
    def __init__(self, **kwargs):
        
        self.ctauctau = None
        StdHepTool.__init__(self,
                            name='lhe_tridents_displaceuni',
                            output_ext='.stdhep',
                            **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = StdHepTool.cmd_args(self)
        if self.ctauctau is not None:
            args.extend(["-l", str(self.ctauctau)])
        return args
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **ctau**
        @return list of optional parameters
        """
        return ['ctau']
 
 
class AddMother(StdHepTool):
    """!
    Add mother particles for physics samples.
    """
 
    def __init__(self, **kwargs):
        StdHepTool.__init__(self,
                            name='add_mother',
                            append_tok='mom',
                            **kwargs)
 
 
class AddMotherFullTruth(StdHepTool):
    """! Add full truth mother particles for physics samples"""
 
    def __init__(self, **kwargs):
        StdHepTool.__init__(self,
                            'add_mother_full_truth',
                            append_tok='mom_full_truth',
                            **kwargs)
        if len(self.inputsinputs) != 2:
            raise Exception("Must have 2 input files: a stdhep file and a lhe file in order")
        self.input_file_1input_file_1 = self.inputsinputs[0]
        base, ext = os.path.splitext(self.input_file_1input_file_1)
        if ext != '.stdhep':
            raise Exception("The first input file must be a stdhep file")
        self.input_file_2input_file_2 = self.inputsinputs[1]
        base, ext = os.path.splitext(self.input_file_2input_file_2)
        if ext != '.lhe':
            raise Exception("The second input file must be a lhe file")
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        return super().cmd_args()
 
 
class MergePoisson(StdHepTool):
    """!
    Merge StdHep files, applying poisson sampling.
 
    Required parameters are: **target_thickness**, **num_electrons**
    """
 
    def __init__(self, xsec=0, **kwargs):
        
        self.xsecxsec = xsec
        
        self.target_thicknesstarget_thickness = None
        
        self.num_electronsnum_electrons = None
 
        StdHepTool.__init__(self,
                            name='merge_poisson',
                            append_tok='sampled',
                            **kwargs)
 
    def setup(self):
        """! Setup MergePoisson component."""
        if self.xsecxsec > 0:
            self.mumu = func.lint(self.target_thicknesstarget_thickness, self.num_electronsnum_electrons) * self.xsecxsec
        else:
            raise Exception("Cross section is missing.")
        self.loggerlogger.info("mu is %f", self.mumu)
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: **target_thickness**, **num_electrons**
        @return list of required parameters
        """
        return ['target_thickness', 'num_electrons']
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = []
        if self.namename in StdHepTool.seed_names:
            args.extend(["-s", str(self.seedseed)])
 
        args.extend(["-m", str(self.mumu), "-N", str(1), "-n", str(self.neventsnevents)])
 
        if len(self.output_filesoutput_files()) == 1:
            # only use file name, not extension because extension is added by tool
            args.insert(0, os.path.splitext(self.output_filesoutput_files()[0])[0])
        elif len(self.output_filesoutput_files()) > 1:
            raise Exception("Too many outputs specified for MergePoisson.")
        else:
            raise Exception("No outputs specified for MergePoisson.")
 
        if len(self.input_filesinput_files()):
            for i in self.inputsinputs[::-1]:
                args.insert(0, i)
        else:
            raise Exception("No inputs were provided.")
 
        return args
 
    def execute(self, log_out, log_err):
        """! Execute MergePoisson component."""
        returncode = Component.execute(self, log_out, log_err)
 
        # Move file from tool to proper output file location.
        src = '%s_1.stdhep' % os.path.splitext(self.output_filesoutput_files()[0])[0]
        dest = '%s.stdhep' % os.path.splitext(self.output_filesoutput_files()[0])[0]
        self.loggerlogger.debug("Moving '%s' to '%s'" % (src, dest))
        shutil.move(src, dest)
 
        return returncode
 
 
class MergeFiles(StdHepTool):
    """!
    Merge StdHep files.
 
    Optional parameters are: none \n
    Required parameters are: none
    """
 
    def __init__(self, **kwargs):
        StdHepTool.__init__(self,
                            name='merge_files',
                            **kwargs)
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: none
        @return list of optional parameters
        """
        return []
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: none
        @return list of required parameters
        """
        return []
 
 
class StdHepCount(Component):
    """!
    Count number of events in a StdHep file.
    """
 
    def __init__(self, **kwargs):
        Component.__init__(self,
                           name='stdhep_count',
                           command='stdhep_count.sh',
                           **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        
        return [self.input_filesinput_files()[0]]
 
    def execute(self, log_out, log_err):
        """! Execute StdHepCount component."""
        cl = [self.commandcommand]
        cl.extend(self.cmd_argscmd_argscmd_args())
        proc = subprocess.Popen(cl, stdout=PIPE)
        (output, err) = proc.communicate()
 
        nevents = int(output.split()[1])
        print("StdHep file '%s' has %d events." % (self.input_filesinput_files()[0], nevents))
 
        return proc.returncode
 
 
class JavaTool(Component):
    """!
    Generic base class for Java based tools.
    """
 
    def __init__(self, name, java_class, **kwargs):
        
        self.java_classjava_class = java_class
        
        self.java_argsjava_args = None
        
        self.conditions_urlconditions_url = None
        Component.__init__(self,
                           name,
                           "java",
                           **kwargs)
 
    def required_config(self):
        """!
        Return list of required config.
 
        Required config are: **hps_java_bin_jar**
        @return list of required config
        """
        return ['hps_java_bin_jar']
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = []
        if self.java_argsjava_args is not None:
            self.loggerlogger.debug("Setting java_args from config: %s" + self.java_argsjava_args)
            args.append(self.java_argsjava_args)
        if self.conditions_urlconditions_url is not None:
            self.loggerlogger.debug('Setting conditions_url from config: %s' % self.conditions_urlconditions_url)
            args.append('-Dorg.hps.conditions.url=%s' % self.conditions_urlconditions_url)
        args.append("-cp")
        args.append(self.hps_java_bin_jar)
        args.append(self.java_classjava_class)
        return args
 
    def config(self, parser):
        super().config(parser)
 
 
class EvioToLcio(JavaTool):
    """!
    Convert EVIO events to LCIO using the hps-java EvioToLcio command line tool.
 
    Input files have evio format (format used by DAQ system).
 
    Required parameters are: **detector**, **steering_files** \n
    Optional parameters are: **run_number**, **skip_events**, **nevents**, **event_print_interval**
    """
 
    def __init__(self, steering=None, **kwargs):
        
        self.detectordetector = None
        
        self.run_numberrun_number = None
        
        self.skip_eventsskip_events = None
        
        self.event_print_intervalevent_print_interval = None
        
        self.steeringsteering = steering
 
        JavaTool.__init__(self,
                          name='evio_to_lcio',
                          java_class='org.hps.evio.EvioToLcio',
                          output_ext='.slcio',
                          **kwargs)
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: **detector**, **steering_files**
        @return list of required parameters
        """
        return ['detector', 'steering_files']
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **run_number**, **skip_events**, **nevents**, **event_print_interval**
        @return list of optional parameters
        """
        return ['run_number', 'skip_events', 'nevents', 'event_print_interval']
 
    def setup(self):
        """! Setup EvioToLcio component."""
        super().setup()
        if self.steeringsteering not in self.steering_files:
            raise Exception("Steering '%s' not found in: %s" % (self.steeringsteering, self.steering_files))
        self.steering_filesteering_file = self.steering_files[self.steeringsteering]
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = JavaTool.cmd_args(self)
        if not len(self.output_filesoutput_files()):
            raise Exception('No output files were provided.')
        output_file = self.output_filesoutput_files()[0]
        args.append('-DoutputFile=%s' % os.path.splitext(output_file)[0])
        args.extend(['-d', self.detectordetector])
        if self.run_numberrun_number is not None:
            args.extend(['-R', str(self.run_numberrun_number)])
        if self.skip_eventsskip_events is not None:
            args.extend(['-s', str(self.skip_eventsskip_events)])
 
        if not os.path.isfile(self.steering_filesteering_file):
            args.append('-r')
            self.loggerlogger.debug("Steering does not exist at '%s' so assuming it is a resource." % self.steering_filesteering_file)
        else:
            if not os.path.isabs(self.steering_filesteering_file):
                raise Exception("Steering looks like a file but is not an abs path: %s" % self.steering_filesteering_file)
        args.extend(['-x', self.steering_filesteering_file])
 
        if self.neventsnevents is not None:
            args.extend(['-n', str(self.neventsnevents)])
 
        args.append('-b')
 
        for inputfile in self.input_filesinput_files():
            args.append(inputfile)
 
        if self.event_print_intervalevent_print_interval is not None:
            args.extend(['-e', str(self.event_print_intervalevent_print_interval)])
 
        return args
 
 
class FilterBunches(JavaTool):
    """!
    Space MC events and apply energy filters to process before readout.
 
    Optional parameters are: **filter_ecal_hit_ecut**, **filter_event_interval**,
    **filter_nevents_read**, **filter_nevents_write**, **filter_no_cuts** \n
    Required config are: **hps_java_bin_jar**
    """
 
    def __init__(self, **kwargs):
        if 'filter_no_cuts' in kwargs:
            self.filter_no_cutsfilter_no_cuts = kwargs['filter_no_cuts']
        else:
            
            self.filter_no_cutsfilter_no_cuts = False
 
        if 'filter_ecal_pairs' in kwargs:
            self.filter_ecal_pairsfilter_ecal_pairs = kwargs['filter_ecal_pairs']
        else:
            self.filter_ecal_pairsfilter_ecal_pairs = False
 
        if 'filter_ecal_hit_ecut' in kwargs:
            self.filter_ecal_hit_ecutfilter_ecal_hit_ecut = kwargs['filter_ecal_hit_ecut']
        else:
            
            self.filter_ecal_hit_ecutfilter_ecal_hit_ecut = -1.0
            # self.filter_ecal_hit_ecut = 0.05
 
        if 'filter_event_interval' in kwargs:
            self.filter_event_intervalfilter_event_interval = kwargs['filter_event_interval']
        else:
            
            self.filter_event_intervalfilter_event_interval = 250
 
        if 'filter_nevents_read' in kwargs:
            self.filter_nevents_readfilter_nevents_read = kwargs['filter_nevents_read']
        else:
            
            self.filter_nevents_readfilter_nevents_read = -1
 
        if 'filter_nevents_write' in kwargs:
            self.filter_nevents_writefilter_nevents_write = kwargs['filter_nevents_write']
        else:
            
            self.filter_nevents_writefilter_nevents_write = -1
 
        self.hps_java_bin_jarhps_java_bin_jar = None
 
        JavaTool.__init__(self,
                          name='filter_bunches',
                          java_class='org.hps.util.FilterMCBunches',
                          append_tok='filt',
                          **kwargs)
 
    def config(self, parser):
        """! Configure FilterBunches component."""
        super().config(parser)
        if self.hps_java_bin_jarhps_java_bin_jar is None:
            if os.getenv('HPS_JAVA_BIN_JAR', None) is not None:
                self.hps_java_bin_jarhps_java_bin_jar = os.getenv('HPS_JAVA_BIN_JAR', None)
                self.loggerlogger.debug('Set HPS_JAVA_BIN_JAR from environment: {}'.format(self.hps_java_bin_jarhps_java_bin_jar))
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = JavaTool.cmd_args(self)
        args.append("-e")
        args.append(str(self.filter_event_intervalfilter_event_interval))
        for i in self.input_filesinput_files():
            args.append(i)
        args.append(self.output_filesoutput_files()[0])
        if self.filter_ecal_pairsfilter_ecal_pairs:
            args.append("-d")
        if self.filter_ecal_hit_ecutfilter_ecal_hit_ecut > 0:
            args.append("-E")
            args.append(str(self.filter_ecal_hit_ecutfilter_ecal_hit_ecut))
        if self.filter_nevents_readfilter_nevents_read > 0:
            args.append('-n')
            args.append(str(self.filter_nevents_readfilter_nevents_read))
        if self.filter_nevents_writefilter_nevents_write > 0:
            args.append('-w')
            args.append(str(self.filter_nevents_writefilter_nevents_write))
        if self.filter_no_cutsfilter_no_cuts:
            args.append('-a')
        return args
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **filter_ecal_hit_ecut**, **filter_event_interval**,
        **filter_nevents_read**, **filter_nevents_write**, **filter_no_cuts** \n
        @return list of optional parameters
        """
        return ['filter_ecal_hit_ecut',
                'filter_event_interval',
                'filter_nevents_read',
                'filter_nevents_write',
                'filter_no_cuts']
 
    def required_config(self):
        """!
        Return list of required config.
 
        Required config are: **hps_java_bin_jar**
        @return list of required config
        """
        return ['hps_java_bin_jar']
 
 
class ExtractEventsWithHitAtHodoEcal(JavaTool):
    """!
    Apply hodo-hit filter and space MC events to process before readout.
 
    The nevents parameter is not settable from JSON in this class. It should
    be supplied as an init argument in the job script if it needs to be
    customized (the default nevents and event_interval used to apply spacing
    should usually not need to be changed by the user). \n
 
    Optional parameters are: **num_hodo_hits**, **event_interval**
    """
 
    def __init__(self, **kwargs):
        if "num_hodo_hits" in kwargs:
            self.num_hodo_hitsnum_hodo_hits = kwargs['num_hodo_hits']
        else:
            self.num_hodo_hitsnum_hodo_hits = 0
 
        if "event_interval" in kwargs:
            self.event_intervalevent_interval = kwargs['event_interval']
        else:
            self.event_intervalevent_interval = 250
 
        JavaTool.__init__(self,
                          name='filter_events',
                          java_class='org.hps.util.ExtractEventsWithHitAtHodoEcal',
                          append_tok='filt',
                          **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = JavaTool.cmd_args(self)
        args.append("-e")
        args.append(str(self.event_intervalevent_interval))
        for i in self.input_filesinput_files():
            args.append(i)
        args.append(self.output_filesoutput_files()[0])
        if self.num_hodo_hitsnum_hodo_hits > 0:
            args.append("-M")
            args.append(str(self.num_hodo_hitsnum_hodo_hits))
        if self.neventsnevents:
            args.append("-w")
            args.append(str(self.neventsnevents))
        return args
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: **num_hodo_hits**, **event_interval**
        @return list of optional parameters
        """
        return ['num_hodo_hits', 'event_interval']
 
 
class Unzip(Component):
    """!
    Unzip the input files to outputs.
    """
 
    def __init__(self, **kwargs):
        Component.__init__(self,
                           name='unzip',
                           command='gunzip',
                           **kwargs)
 
    def output_files(self):
        """! Return list of output files."""
        if self.outputsoutputs:
            return self.outputsoutputs
        return [os.path.splitext(i)[0] for i in self.input_filesinput_files()]
 
    def execute(self, log_out, log_err):
        """! Execute Unzip component."""
        for i in range(0, len(self.input_filesinput_files())):
            inputfile = self.input_filesinput_files()[i]
            outputfile = self.output_filesoutput_filesoutput_files()[i]
            with gzip.open(inputfile, 'rb') as in_file, open(outputfile, 'wb') as out_file:
                shutil.copyfileobj(in_file, out_file)
                self.loggerlogger.debug("Unzipped '%s' to '%s'" % (inputfile, outputfile))
        return 0
 
 
class LCIODumpEvent(Component):
    """!
    Dump LCIO event information.
 
    Required parameters are: none \n
    Required config are: **lcio_dir**
    """
 
    def __init__(self, **kwargs):
        
        self.lcio_dirlcio_dir = None
        Component.__init__(self,
                           name='lcio_dump_event',
                           command='dumpevent',
                           **kwargs)
 
        if "event_num" in kwargs:
            self.event_numevent_num = kwargs["event_num"]
        else:
            self.event_numevent_num = 1
 
    def config(self, parser):
        """! Configure LCIODumpEvent component."""
        super().config(parser)
        if self.lcio_dirlcio_dir is None:
            self.lcio_dirlcio_dir = self.hpsmc_dirhpsmc_dir
 
    def setup(self):
        """! Setup LCIODumpEvent component."""
        self.commandcommandcommand = self.lcio_dirlcio_dir + "/bin/dumpevent"
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        if not len(self.input_filesinput_files()):
            raise Exception("Missing required inputs for LCIODumpEvent.")
        args = []
        args.append(self.input_filesinput_files()[0])
        args.append(str(self.event_numevent_num))
        return args
 
    def required_config(self):
        """!
        Return list of required config.
 
        Required config are: **lcio_dir**
        @return list of required config
        """
        return ['lcio_dir']
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: none
        @return list of required parameters
        """
        return []
 
 
class LHECount(Component):
    """!
    Count events in an LHE file.
    """
 
    def __init__(self, minevents=0, fail_on_underflow=False, **kwargs):
        self.mineventsminevents = minevents
        Component.__init__(self,
                           name="lhe_count",
                           **kwargs)
 
    def setup(self):
        """! Setup LHECount component."""
        if not len(self.input_filesinput_files()):
            raise Exception("Missing at least one input file.")
 
    def cmd_exists(self):
        """!
        Check if command exists.
        @return True if command exists
        """
        return True
 
    def execute(self, log_out, log_err):
        """! Execute LHECount component."""
        for i in self.inputsinputs:
            with gzip.open(i, 'rb') as in_file:
                lines = in_file.readlines()
 
            nevents = 0
            for line in lines:
                if "<event>" in line:
                    nevents += 1
 
            print("LHE file '%s' has %d events." % (i, nevents))
 
            if nevents < self.mineventsminevents:
                msg = "LHE file '%s' does not contain the minimum %d events." % (i, nevents)
                if self.fail_on_underflow:
                    raise Exception(msg)
                else:
                    self.loggerlogger.warning(msg)
        return 0
 
 
class TarFiles(Component):
    """!
    Tar files into an archive.
    """
 
    def __init__(self, **kwargs):
        Component.__init__(self,
                           name='tar_files',
                           **kwargs)
 
    def cmd_exists(self):
        """!
        Check if command exists.
        @return True if command exists
        """
        return True
 
    def execute(self, log_out, log_err):
        """! Execute TarFiles component."""
        self.loggerlogger.debug("Opening '%s' for writing ..." % self.outputsoutputs[0])
        tar = tarfile.open(self.outputsoutputs[0], "w")
        for i in self.inputsinputs:
            self.loggerlogger.debug("Adding '%s' to archive" % i)
            tar.add(i)
        tar.close()
        self.loggerlogger.info("Wrote archive '%s'" % self.outputsoutputs[0])
        return 0
 
 
class MoveFiles(Component):
    """!
    Move input files to new locations.
    """
 
    def __init__(self, **kwargs):
        Component.__init__(self,
                           name='move_files',
                           **kwargs)
 
    def cmd_exists(self):
        """!
        Check if command exists.
        @return True if command exists
        """
        return True
 
    def execute(self, log_out, log_err):
        """! Execute TarFiles component."""
        if len(self.inputsinputs) != len(self.outputsoutputs):
            raise Exception("Input and output lists are not the same length!")
        for io in zip(self.inputsinputs, self.outputsoutputs):
            src = io[0]
            dest = io[1]
            self.loggerlogger.info("Moving %s -> %s" % (src, dest))
            shutil.move(src, dest)
        return 0
 
 
class LCIOTool(Component):
    """!
    Generic component for LCIO tools.
 
    Required parameters are: none \n
    Required config are: **lcio_bin_jar**
    """
 
    def __init__(self, name=None, **kwargs):
        
        self.lcio_bin_jarlcio_bin_jar = None
        Component.__init__(self,
                           name,
                           command='java',
                           **kwargs)
 
    def config(self, parser):
        """! Configure LCIOTool component."""
        super().config(parser)
        if self.lcio_bin_jarlcio_bin_jar is None:
            self.config_from_environconfig_from_environ()
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        if not self.namename:
            raise Exception("Name required to write cmd args for LCIOTool.")
        return ['-jar', self.lcio_bin_jarlcio_bin_jar, self.namename]
 
    def required_config(self):
        """!
        Return list of required config.
 
        Required config are: **lcio_bin_jar**
        @return list of required config
        """
        return ['lcio_bin_jar']
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: none
        @return list of required parameters
        """
        return []
 
 
class LCIOConcat(LCIOTool):
    """!
    Concatenate LCIO files together.
    """
 
    def __init__(self, **kwargs):
        LCIOTool.__init__(self,
                          name='concat',
                          **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = LCIOTool.cmd_args(self)
        if not len(self.input_filesinput_files()):
            raise Exception("Missing at least one input file.")
        if not len(self.output_filesoutput_files()):
            raise Exception("Missing an output file.")
        for i in self.input_filesinput_files():
            args.extend(["-f", i])
        args.extend(["-o", self.outputsoutputs[0]])
        return args
 
 
class LCIOCount(LCIOTool):
    """!
    Count events in LCIO files.
 
    Required parameters are: none \n
    Optional parameters are: none
    """
 
    def __init__(self, **kwargs):
        LCIOTool.__init__(self,
                          name='count',
                          **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = LCIOTool.cmd_args(self)
        if not len(self.inputsinputs):
            raise Exception("Missing an input file.")
        args.extend(["-f", self.inputsinputs[0]])
        return args
 
    def required_parameters(self):
        """!
        Return list of required parameters.
 
        Required parameters are: none
        @return list of required parameters
        """
        return []
 
    def optional_parameters(self):
        """!
        Return list of optional parameters.
 
        Optional parameters are: none
        @return list of optional parameters
        """
        return []
 
 
class LCIOMerge(LCIOTool):
    """!
    Merge LCIO files.
    """
 
    def __init__(self, **kwargs):
        LCIOTool.__init__(self,
                          name='merge',
                          **kwargs)
 
    def cmd_args(self):
        """!
        Setup command arguments.
        @return  list of arguments
        """
        args = LCIOTool.cmd_args(self)
        if not len(self.input_filesinput_files()):
            raise Exception("Missing at least one input file.")
        if not len(self.output_filesoutput_files()):
            raise Exception("Missing an output file.")
        for i in self.input_filesinput_files():
            args.extend(["-f", i])
        args.extend(["-o", self.outputsoutputs[0]])
        if self.neventsnevents is not None:
            args.extend(['-n', str(self.neventsnevents)])
        return args