generators.py

Go to the documentation of this file.

"""! @package generators
Event generation tools."""
 
import os
import shutil
import glob
import gzip
import logging
 
from hpsmc.component import Component
 
 
class EventGenerator(Component):
    """! Event generator base class."""
 
    def __init__(self, name, command=None, **kwargs):
        Component.__init__(self, name, command=command, description='', **kwargs)
 
    def required_parameters(self):
        return ['nevents']
 
    def get_install_dir(self):
        if os.getenv("HPSMC_DIR") is None:
            raise Exception("HPSMC_DIR is not set!")
        return "{}/share/generators".format(os.getenv("HPSMC_DIR", None))
 
 
class EGS5(EventGenerator):
    """!
    Run the EGS5 event generator to produce a StdHep file.
 
    Required parameters are **seed**, **target_thickness**, **beam_energy**, **num_electrons** \n
    Optional parameters are: **bunches**
    """
 
    def __init__(self, name='', **kwargs):
        
        self.bunches = 5e5
        
        self.target_thickness = None
        
        self.egs5_dir = None
        
        self.beam_energy = None
        
        self.num_electrons = None
        EventGenerator.__init__(self, name, "egs5_" + name, **kwargs)
 
    def get_install_dir(self):
        """! Get installation directory."""
        return EventGenerator.get_install_dir(self) + "/egs5"
 
    def setup(self):
        """! Setup of egs5 event generator."""
        EventGenerator.setup(self)
 
        if self.egs5_dir is None:
            self.egs5_dir = self.get_install_dirget_install_dir()
            self.logger.debug("Using EGS5 from install dir: " + self.egs5_dir)
 
        
        self.egs5_data_dir = os.path.join(self.egs5_dir, "data")
        
        self.egs5_config_dir = os.path.join(self.egs5_dir, "config")
 
        self.logger.debug("egs5_data_dir=%s" % self.egs5_data_dir)
        self.logger.debug("egs5_config_dir=%s" % self.egs5_config_dir)
 
        if os.path.exists("data"):
            os.unlink("data")
        os.symlink(self.egs5_data_dir, "data")
 
        if os.path.exists("pgs5job.pegs5inp"):
            os.unlink("pgs5job.pegs5inp")
        os.symlink(self.egs5_config_dir + "/src/esa.inp", "pgs5job.pegs5inp")
 
        # Set target thickness from job parameter or use the default from run parameters
        if self.target_thickness is not None:
            self.target_dz = self.target_thickness
            self.logger.debug("Target thickness set from job param: {}".format(self.target_dz))
        else:
            raise Exception("Target thickness not set!")
 
        ebeam = self.beam_energy
        electrons = int(self.num_electrons * self.bunches)
 
        # seed_data = "%d %f %f %d" % (self.seed, self.target_dz, ebeam, electrons)
        seed_data = " ".join(str(item) for item in [self.seed, self.target_dz, ebeam, electrons])
        self.logger.debug("Seed data (seed, target_dz, ebeam, electrons): {}".format(seed_data))
        seed_file = open("seed.dat", 'w')
        seed_file.write(seed_data)
        seed_file.close()
 
    def output_files(self):
        """! Generate output file name.
        @return moller.stdhep if name of generator contains moller; beam.stdhep else"""
        # Output file for Moller generation
        if 'moller' in self.name:
            return ['moller.stdhep']
        # Output file for beam generation
        return ['beam.stdhep']
 
    def execute(self, log_out, log_err):
        """! Execute event generator.
        @param log_out  name of log file for output
        @param log_err  name of log file for error
        """
        EventGenerator.execute(self, log_out, log_err)
        if 'moller' not in self.name:
            src = os.path.join(self.rundir, 'brems.stdhep')
            dest = os.path.join(self.rundir, self.output_filesoutput_files()[0])
            self.logger.debug("Copying '%s' to '%s'" % (src, dest))
            shutil.copy(src, dest)
 
    def required_parameters(self):
        """!
        Return required parameters.
 
        Required parameters are **seed**, **target_thickness**, **beam_energy**, **num_electrons**
        """
        return ['seed', 'target_thickness', 'beam_energy', 'num_electrons']
 
    def optional_parameters(self):
        """!
        Return optional parameters.
 
        Optional parameters are: **bunches**
        """
        return ['bunches']
 
    # def required_config(self):
    #    return ['egs5_dir']
 
 
class StdHepConverter(EGS5):
    """! Target processing and conversion of LHE files to StdHep using EGS5."""
 
    def __init__(self, name="lhe_v1", **kwargs):
        EGS5.__init__(self, name, **kwargs)
 
        if self.namename not in ["lhe_v1", "lhe_rad", "lhe_prompt", "lhe_uniform", "lhe_exponential"]:
            raise Exception("The name '%s' is not valid for StdHepConverter tools." % self.namename)
 
    def config(self, parser):
        EGS5.config(self, parser)
 
    def setup(self):
        """! Setup egs5 generator.
        Throws exception if input LHE file is missing."""
        EGS5.setup(self)
        if not len(self.inputsinputs):
            raise Exception("Missing required input LHE file.")
 
    def execute(self, log_out, log_err):
        """! Execute converter.
        Calls egs5 generator.
 
        @param log_out  name of log file for output
        @param log_err  name of log file for error
        @return egs5 error code
        """
        input_file = self.inputsinputs[0]
        base, ext = os.path.splitext(input_file)
        if ext == ".lhe":
            os.symlink(input_file, "egs5job.inp")
        elif ext == ".gz":
            with open("egs5job.inp", 'wb') as outfile:
                with gzip.open(self.inputsinputs[0], 'rb') as infile:
                    outfile.write(infile.read())
        else:
            raise Exception('Input file has an unknown extension: %s' % input_file)
        return EGS5.execute(self, log_out, log_err)
 
    def output_files(self):
        """! Converts *.lhe.gz and *.lhe to *.stdhep files."""
        return [self.input_files()[0].replace(".lhe.gz", ".stdhep").replace(".lhe", ".stdhep")]
 
 
class MG(EventGenerator):
    """!
    Abstract class for MadGraph generators.
 
    Required parameters are: **nevents**, **run_params** \n
    Optional parameters are: **seed**, **param_card**, **apmass**, **map**, **mpid**, **mrhod**
    """
 
    def __init__(self, name, **kwargs):
 
        
        self.madgraph_dir = None
 
        
        self.param_card = "param_card.dat"
 
        # Extra parameters for param card:
        # map = mass of the A-prime
        # mpid = mass of the dark pion
        # mrhod = mass of the dark rho
        
        self.apmass = None
        
        self.map = None
        
        self.Map = None
        
        self.mpid = None
        
        self.mrhod = None
        
        self.mchi = None
        
        self.dmchi = None
 
        if 'event_types' in kwargs:
            
            self.event_types = kwargs['event_types']
        else:
            self.event_types = ['unweighted', 'weighted']
 
        EventGenerator.__init__(self, name, **kwargs)
 
    def output_files(self):
        """! Generate output file name."""
        o = []
        if 'unweighted' in self.event_types:
            o.append(self.namename + "_unweighted_events.lhe.gz")
        if 'weighted' in self.event_types:
            o.append(self.namename + "_events.lhe.gz")
        return o
 
    def set_parameters(self, params):
        """! Set parameters."""
        Component.set_parameters(self, params)
        self.run_card = "run_card_" + self.run_params + ".dat"
        self.logger.debug("Set run card to '%s'" % self.run_card)
 
    def required_parameters(self):
        """!
        Return required parameters.
 
        Required parameters are: **nevents**, **run_params**
        """
        return ['nevents', 'run_params']
 
    def optional_parameters(self):
        """!
        Return optional parameters.
 
        Optional parameters are: **seed**, **param_card**, **apmass**, **map**, **mpid**, **mrhod**, **Map**, **mchi**, **dmchi**
        """
        return ['seed', 'param_card', 'apmass', 'map', 'mpid', 'mrhod', 'Map', 'mchi', 'dmchi']
 
    # def required_config(self):
    #    return ['madgraph_dir']
 
    def make_run_card(self, run_card):
        """! Make run card."""
        
        self.logger.info("Making run card '%s' with nevents=%d, seed=%d" % (run_card, self.neventsnevents, self.seedseed))
 
        with open(run_card, 'r') as cardin:
            data = cardin.readlines()
 
        for i in range(0, len(data)):
            if "= nevents" in data[i]:
                self.logger.debug("Setting nevents=%d in run card" % self.neventsnevents)
                data[i] = " " + str(self.neventsnevents) + " = nevents ! Number of unweighted events requested" + '\n'
            if "= iseed" in data[i]:
                self.logger.debug("Setting seed=%d in run card" % self.seedseed)
                data[i] = " " + str(self.seedseed) + " = iseed   ! rnd seed (0=assigned automatically=default))" + '\n'
 
        with open(run_card, 'w') as cardout:
            cardout.writelines(data)
 
    def make_param_card(self, param_card):
        """! Make parameter card."""
        
        self.logger.debug("Making param card '%s'" % param_card)
 
        with open(param_card, 'r') as paramin:
            data = paramin.readlines()
 
        for i in range(0, len(data)):
            if "APMASS" in data[i] and self.apmass is not None:
                data[i] = "       622     %.7fe-03   # APMASS" % (self.apmass) + '\n'
                self.logger.debug("APMASS in param card set to %d" % self.apmass)
            if "MAp" in data[i] and "622" in data[i] and self.apmass is not None:
                data[i] = "  622 %.7fe-03 # MAp" % (self.apmass) + '\n'
            if "map" in data[i] and self.map is not None:
                data[i] = "      622 %.7fe-03 # map" % (self.map) + '\n'
            if "mpid" in data[i] and self.mpid is not None:
                data[i] = "      624 %.7fe-03 # mpid" % (self.mpid) + '\n'
            if "mrhod" in data[i] and self.mrhod is not None:
                data[i] = "      625 %.7fe-03 # mrhod" % (self.mrhod) + '\n'
            if "Map" in data[i] and "ap :" not in data[i] and self.Map is not None:
                data[i] = "    1 %.7e # Map\n" % (self.Map/1000)
            if "Mchi" in data[i] and "dMchi" not in data[i] and self.mchi is not None:
                data[i] = "    1 %.7e # Mchi\n" % (self.mchi/1000)
            if "dMchi" in data[i] and "Mchi/2" not in data[i] and self.dmchi is not None:
                data[i] = "    2 %.7e # dMchi" % (self.dmchi/1000)
 
        with open(param_card, 'w') as paramout:
            paramout.writelines(data)
 
    def cmd_args(self):
        """! Return command arguments."""
        return ["0", self.namename]
 
    def execute(self, log_out, log_err):
        """! Execute MadGraph generator.
        @param log_out  name of log file for output
        @param log_err  name of log file for error
        @return  error code
        """
        os.chdir(os.path.dirname(self.commandcommand))
        # need python3.7 for MG5, but s3df only has python3.6 for now; python2 works though
        # if "WAB" in self.name:
        #     self.command = "python2 " + self.command
        self.logger.debug("Executing '%s' from '%s'" % (self.namename, os.getcwd()))
        return Component.execute(self, log_out, log_err)
 
    def setup(self):
        """! Setup event generator."""
        EventGenerator.setup(self)
 
        if self.madgraph_dir is None:
            self.madgraph_dir = self.get_install_dir()
            self.logger.debug("Using Madgraph from install dir: " + self.madgraph_dir)
 
        if self.namename == 'ap' and self.apmass is None:
            raise Exception("Missing apmass param for AP generation.")
 
 
class MG4(MG):
    """! Run the MadGraph 4 event generator."""
 
    
    dir_map = {"BH": "BH/MG_mini_BH/apBH",
               "RAD": "RAD/MG_mini_Rad/apRad",
               "TM": "TM/MG_mini/ap",
               "ap": "ap/MG_mini/ap",
               "trigg": "trigg/MG_mini_Trigg/apTri",
               "tritrig": "tritrig/MG_mini_Tri_W/apTri",
               "wab": "wab/MG_mini_WAB/AP_6W_XSec2_HallB"}
 
    def __init__(self, name='ap', **kwargs):
 
        MG.__init__(self, name, **kwargs)
 
        if self.namenamename not in MG4.dir_map:
            raise Exception("The name '%s' is not valid for MG4." % self.namenamename)
 
    def get_install_dir(self):
        """! Get installation directory of MadGraph4."""
        return EventGenerator.get_install_dir(self) + "/madgraph4/src"
 
    def setup(self):
        """! Setup MadGraph4 generator."""
        MG.setup(self)
 
        proc_dirs = MG4.dir_map[self.namenamename].split(os.sep)
        src = os.path.join(self.madgraph_dir, proc_dirs[0], proc_dirs[1])
        dest = proc_dirs[1]
        self.logger.debug("Copying '%s' to '%s'" % (src, dest))
        shutil.copytree(src, dest, symlinks=True)
 
        self.event_dir = os.path.join(self.rundir, proc_dirs[1], proc_dirs[2], "Events")
        if not os.path.isdir(self.event_dir):
            os.makedirs(self.event_dir)
 
        self.commandcommandcommand = os.path.join(os.getcwd(), proc_dirs[1], proc_dirs[2], "bin", "generate_events")
        self.logger.debug("Command set to '%s'" % self.commandcommandcommand)
 
        run_card_src = os.path.join(self.madgraph_dir, proc_dirs[0], self.run_card)
        run_card_dest = os.path.join(self.rundir, proc_dirs[1], proc_dirs[2], "Cards", "run_card.dat")
        self.logger.debug("Copying run card from '%s' to '%s'" % (run_card_src, run_card_dest))
        shutil.copyfile(run_card_src, run_card_dest)
 
        self.make_run_card(run_card_dest)
 
        param_card_src = os.path.join(self.madgraph_dir, proc_dirs[0], self.param_card)
        param_card_dest = os.path.join(self.rundir, proc_dirs[1], proc_dirs[2], "Cards", "param_card.dat")
        self.logger.debug("Copying param card from '%s' to '%s'" % (param_card_src, param_card_dest))
        shutil.copyfile(param_card_src, param_card_dest)
 
        self.make_param_card(param_card_dest)
 
    def execute(self, log_out, log_err):
        """! Execute MadGraph4 generator.
        @param log_out  name of log file for output
        @param log_err  name of log file for error
        @return  error code
        """
        returncode = MG.execute(self, log_out, log_err)
        lhe_files = glob.glob(os.path.join(self.event_dir, "*.lhe.gz"))
        for f in lhe_files:
            dest = os.path.join(self.rundir, os.path.basename(f))
            self.logger.debug("Copying '%s' to '%s'" % (f, dest))
            shutil.copy(f, dest)
        os.chdir(self.rundir)
        return returncode
 
 
class MG5(MG):
    """! Run the MadGraph 5 event generator."""
 
    
    dir_map = {"BH": "BH",
               "RAD": "RAD",
               "tritrig": "tritrig",
               "simp": "simp",
               "simp-3body": "simp-3body",
               "idm": "idm",
               "WAB": "WAB",
               "ap": "ap"
               }
 
    def __init__(self, name='tritrig', **kwargs):
 
        MG.__init__(self, name, **kwargs)
 
        if self.namenamename not in MG5.dir_map:
            raise Exception("The name '%s' is not valid for MG5." % self.namenamename)
 
    def get_install_dir(self):
        """! Get installation directory of MadGraph5."""
        return EventGenerator.get_install_dir(self) + "/madgraph5/src"
 
    def setup(self):
        """! Setup MadGraph5 generator."""
        MG.setup(self)
 
        self.proc_dir = MG5.dir_map[self.namenamename]
        self.event_dir = os.path.join(self.rundir, self.proc_dir, "Events", self.proc_dir)
 
        src = os.path.join(self.madgraph_dir, self.proc_dir)
        dest = os.path.join(self.rundir, self.proc_dir)
        self.logger.debug("Copying '%s' to '%s'" % (src, dest))
        shutil.copytree(src, dest, symlinks=True)
 
        
        """
        input_dir = os.path.join(self.madgraph_dir, "input")
        dest_input_dir = os.path.join(self.rundir, self.proc_dir, "input")
        self.logger.info("Copying '%s' to '%s'" % (input_dir, dest_input_dir))
        shutil.copytree(input_dir, dest_input_dir)
        """
 
        self.commandcommandcommand = os.path.join(dest, "bin", "generate_events")
        self.logger.debug("Command set to '%s'" % self.commandcommandcommand)
 
        run_card_src = os.path.join(src, "Cards", self.run_card)
        run_card_dest = os.path.join(dest, "Cards", "run_card.dat")
        self.logger.debug("Copying run card from '%s' to '%s'" % (run_card_src, run_card_dest))
        shutil.copyfile(run_card_src, run_card_dest)
 
        self.make_run_card(run_card_dest)
 
        param_card_src = os.path.join(src, "Cards", self.param_card)
        param_card_dest = os.path.join(dest, "Cards", "param_card.dat")
        self.logger.debug("Copying param card from '%s' to '%s'" % (param_card_src, param_card_dest))
        shutil.copyfile(param_card_src, param_card_dest)
 
        self.make_param_card(param_card_dest)
 
    def execute(self, log_out, log_err):
        """! Execute MadGraph5 generator.
        @param log_out  name of log file for output
        @param log_err  name of log file for error
        @return  error code
        """
        os.chdir(os.path.dirname(self.commandcommandcommand))
        # need python3.7 for MG5, but s3df only has python3.6 for now; python2 works though
        if "WAB" or "ap" in self.namenamename:
            self.commandcommandcommand = "python2 " + self.commandcommandcommand
        self.logger.debug("Executing '%s' from '%s'" % (self.namenamename, os.getcwd()))
        returncode = Component.execute(self, log_out, log_err)
 
        # returncode = MG.execute(self, log_out, log_err)
        lhe_files = glob.glob(os.path.join(self.event_dir, "*.lhe.gz"))
        for f in lhe_files:
            dest = os.path.join(self.rundir, '%s_%s' % (self.namenamename, os.path.basename(f)))
            self.logger.debug("Copying '%s' to '%s'" % (f, dest))
            shutil.copy(f, dest)
        os.chdir(self.rundir)
        return returncode