Skip to content
Snippets Groups Projects
DD4hepSimulation.py 35.1 KiB
Newer Older
"""

DD4hep simulation with some argument parsing
Based on M. Frank and F. Gaede runSim.py
   @author  A.Sailer
   @version 0.1

"""
import argparse
import io
import logging
Marko Petric's avatar
Marko Petric committed
import os
from urllib.parse import urlparse
Marko Petric's avatar
Marko Petric committed
from DDSim.Helper.Meta import Meta
from DDSim.Helper.LCIO import LCIO
Marko Petric's avatar
Marko Petric committed
from DDSim.Helper.GuineaPig import GuineaPig
from DDSim.Helper.Physics import Physics
from DDSim.Helper.Filter import Filter
from DDSim.Helper.Geometry import Geometry
Marko Petric's avatar
Marko Petric committed
from DDSim.Helper.Random import Random
from DDSim.Helper.Action import Action
from DDSim.Helper.Output import Output, outputLevel, outputLevelType
from DDSim.Helper.OutputConfig import OutputConfig, defaultOutputFile
from DDSim.Helper.InputConfig import InputConfig
from DDSim.Helper.ConfigHelper import ConfigHelper
Marko Petric's avatar
Marko Petric committed
from DDSim.Helper.MagneticField import MagneticField
from DDSim.Helper.ParticleHandler import ParticleHandler
from DDSim.Helper.Gun import Gun
Marko Petric's avatar
Marko Petric committed
logger = logging.getLogger('DDSim')
try:
  import argcomplete
Marko Petric's avatar
Marko Petric committed
  ARGCOMPLETEENABLED = True
except ImportError:
Marko Petric's avatar
Marko Petric committed
  ARGCOMPLETEENABLED = False
HEPMC3_SUPPORTED_EXTENSIONS = [
    ".hepmc.gz", ".hepmc.xz", ".hepmc.bz2",
Dmitry Kalinkin's avatar
Dmitry Kalinkin committed
    ".hepmc3", ".hepmc3.gz", ".hepmc3.xz", ".hepmc3.bz2",
    ".hepmc3.tree.root",
EDM4HEP_INPUT_EXTENSIONS = [
    ".root",
    ".sio",
    ]
POSSIBLEINPUTFILES = [
    ".stdhep", ".slcio", ".HEPEvt", ".hepevt",
Dmitry Kalinkin's avatar
Dmitry Kalinkin committed
    ".pairs",
    ]
POSSIBLEINPUTFILES += HEPMC3_SUPPORTED_EXTENSIONS
POSSIBLEINPUTFILES += EDM4HEP_INPUT_EXTENSIONS
Marko Petric's avatar
Marko Petric committed

class DD4hepSimulation(object):
  """Class to hold all the parameters and functions to run simulation"""

  def __init__(self):
    self.steeringFile = None
    self.compactFile = []
    self.outputFile = defaultOutputFile()
    self.printLevel = 3
Marko Petric's avatar
Marko Petric committed
    self.physicsList = None  # deprecated use physics.list
    self.crossingAngleBoost = 0.0
    self.macroFile = ''
    self.enableGun = False
    self.enableG4GPS = False
    self.enableG4Gun = False
    self._g4gun = None
    self._g4gps = None
    self.vertexSigma = [0.0, 0.0, 0.0, 0.0]
    self.vertexOffset = [0.0, 0.0, 0.0, 0.0]
    self.enableDetailedShowerMode = False
    self._errorMessages = []
    self._dumpParameter = False
    self._dumpSteeringFile = False
Marko Petric's avatar
Marko Petric committed
    # objects for extended configuration option
    self.output = Output()
    self.random = Random()
    self.gun = Gun()
    self.part = ParticleHandler()
    self.field = MagneticField()
    self.outputConfig = OutputConfig()
    self.inputConfig = InputConfig()
Jan Strube's avatar
Jan Strube committed
    self.lcio = LCIO()
    self.geometry = Geometry()
  def readSteeringFile(self):
    """Reads a steering file and sets the parameters to that of the
    DD4hepSimulation object present in the steering file.
    """
    globs = {}
Marko Petric's avatar
Marko Petric committed
    locs = {}
    if not self.steeringFile:
    sFileTemp = self.steeringFile
    exec(compile(io.open(self.steeringFile).read(), self.steeringFile, 'exec'), globs, locs)
    for _name, obj in locs.items():
      if isinstance(obj, DD4hepSimulation):
        self.__dict__ = obj.__dict__
    self.steeringFile = os.path.abspath(sFileTemp)
    parser = argparse.ArgumentParser("Running DD4hep Simulations:",
                                     formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument("--steeringFile", "-S", action="store", default=self.steeringFile,
                        help="Steering file to change default behaviour")

Marko Petric's avatar
Marko Petric committed
    # first we parse just the steering file, but only if we don't want to see the help message
    if not any(opt in self._argv for opt in ('-h', '--help')):
      parsed, _unknown = parser.parse_known_args()
      self.steeringFile = parsed.steeringFile
      self.readSteeringFile()
Marko Petric's avatar
Marko Petric committed
    # readSteeringFile will set self._argv to None if there is a steering file
    if self._argv is None:
      self._argv = list(argv) if argv else list(sys.argv)

    parser.add_argument("--compactFile", nargs='+', action="store",
                        default=ConfigHelper.makeList(self.compactFile), type=str,
                        help="The compact XML file, or multiple compact files, if the last one is the closer.")
    parser.add_argument("--runType", action="store", choices=("batch", "vis", "run", "shell", "qt"),
                        default=self.runType,
Marko Petric's avatar
Marko Petric committed
                        help="The type of action to do in this invocation"  # Note: implicit string concatenation
                        "\nbatch: just simulate some events, needs numberOfEvents, and input file or gun"
                        "\nvis: enable visualisation, run the macroFile if it is set"
                        "\nqt: enable visualisation in Qt shell, run the macroFile if it is set"
                        "\nrun: run the macroFile and exit"
                        "\nshell: enable interactive session")

    parser.add_argument("--inputFiles", "-I", nargs='+', action="store", default=self.inputFiles,
Marko Petric's avatar
Marko Petric committed
                        help="InputFiles for simulation %s files are supported" % ", ".join(POSSIBLEINPUTFILES))
Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--outputFile", "-O", action="store", default=self.outputFile,
                        help="Outputfile from the simulation: .slcio, edm4hep.root and .root"
                        " output files are supported")

    parser.add_argument("-v", "--printLevel", action="store", default=self.printLevel, dest="printLevel",
Marko Petric's avatar
Marko Petric committed
                        choices=(1, 2, 3, 4, 5, 6, 7, 'VERBOSE', 'DEBUG',
                                 'INFO', 'WARNING', 'ERROR', 'FATAL', 'ALWAYS'),
                        help="Verbosity use integers from 1(most) to 7(least) verbose"
                        "\nor strings: VERBOSE, DEBUG, INFO, WARNING, ERROR, FATAL, ALWAYS")

    parser.add_argument("--numberOfEvents", "-N", action="store", dest="numberOfEvents", default=self.numberOfEvents,
                        type=int, help="number of events to simulate, used in batch mode")

    parser.add_argument("--skipNEvents", action="store", dest="skipNEvents", default=self.skipNEvents, type=int,
                        help="Skip first N events when reading a file")

    parser.add_argument("--physicsList", action="store", dest="physicsList", default=self.physicsList,
                        help="Physics list to use in simulation")

Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--crossingAngleBoost", action="store", dest="crossingAngleBoost",
                        default=self.crossingAngleBoost,
                        type=float, help="Lorentz boost for the crossing angle, in radian!")

Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--vertexSigma", nargs=4, action="store", dest="vertexSigma",
                        default=self.vertexSigma, metavar=('X', 'Y', 'Z', 'T'),
                        type=float, help="FourVector of the Sigma for the Smearing of the Vertex position: x y z t")

Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--vertexOffset", nargs=4, action="store", dest="vertexOffset",
                        default=self.vertexOffset, metavar=('X', 'Y', 'Z', 'T'),
                        type=float, help="FourVector of translation for the Smearing of the Vertex position: x y z t")

    parser.add_argument("--macroFile", "-M", action="store", dest="macroFile", default=self.macroFile,
                        help="Macro file to execute for runType 'run' or 'vis'")

    parser.add_argument("--enableGun", "-G", action="store_true", dest="enableGun", default=self.enableGun,
    parser.add_argument("--enableG4GPS", action="store_true", dest="enableG4GPS", default=self.enableG4GPS,
Marko Petric's avatar
Marko Petric committed
                        help="enable the Geant4 GeneralParticleSource. Needs a macroFile (runType run)"
                        "or use it with the shell (runType shell)")

    parser.add_argument("--enableG4Gun", action="store_true", dest="enableG4Gun", default=self.enableG4Gun,
Marko Petric's avatar
Marko Petric committed
                        help="enable the Geant4 particle gun. Needs a macroFile (runType run)"
                        " or use it with the shell (runType shell)")
Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--dumpParameter", "--dump", action="store_true", dest="dumpParameter",
                        default=self._dumpParameter, help="Print all configuration Parameters and exit")
Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--enableDetailedShowerMode", action="store_true", dest="enableDetailedShowerMode",
                        default=self.enableDetailedShowerMode,
    parser.add_argument("--disableSignalHandler", action="store_true", dest="disableSignalHandler",
                        default=self.disableSignalHandler,
                        help="disable the Signal Handler of DD4hep")

Marko Petric's avatar
Marko Petric committed
    parser.add_argument("--dumpSteeringFile", action="store_true", dest="dumpSteeringFile",
                        default=self._dumpSteeringFile, help="print an example steering file to stdout")
Marko Petric's avatar
Marko Petric committed
    # output, or do something smarter with fullHelp only for example
Marko Petric's avatar
Marko Petric committed
    # now parse everything. The default values are now taken from the
    # steeringFile if they were set so that the steering file parameters can be
    # overwritten from the command line
    if ARGCOMPLETEENABLED:
      argcomplete.autocomplete(parser)
    self._dumpParameter = parsed.dumpParameter
    self._dumpSteeringFile = parsed.dumpSteeringFile
    self.compactFile = ConfigHelper.makeList(parsed.compactFile)
    self.__checkFilesExist(self.compactFile, fileType='compact')
Marko Petric's avatar
Marko Petric committed
    self.inputFiles = self.__checkFileFormat(self.inputFiles, POSSIBLEINPUTFILES)
    self.__checkFilesExist(self.inputFiles, fileType='input')
Marko Petric's avatar
Marko Petric committed
    self.__checkFileFormat(self.outputFile, ('.root', '.slcio'))
    self.printLevel = self.__checkOutputLevel(parsed.printLevel)

    self.numberOfEvents = parsed.numberOfEvents
    self.skipNEvents = parsed.skipNEvents
    self.physicsList = parsed.physicsList
    self.crossingAngleBoost = parsed.crossingAngleBoost
    self.macroFile = parsed.macroFile
    self.enableGun = parsed.enableGun
    self.enableG4Gun = parsed.enableG4Gun
    self.enableG4GPS = parsed.enableG4GPS
    self.enableDetailedShowerMode = parsed.enableDetailedShowerMode
    self.vertexOffset = parsed.vertexOffset
    self.vertexSigma = parsed.vertexSigma

    self._consistencyChecks()
    if self.printLevel <= 2:  # VERBOSE or DEBUG
      logger.setLevel(logging.DEBUG)

Marko Petric's avatar
Marko Petric committed
    # self.__treatUnknownArgs( parsed, unknown )
Marko Petric's avatar
Marko Petric committed
    self.__parseAllHelper(parsed)
    if self._errorMessages and not (self._dumpParameter or self._dumpSteeringFile):
      parser.epilog = "\n".join(self._errorMessages)
      parser.print_help()
      exit(1)

    if self._dumpParameter:
      from pprint import pprint
Marko Petric's avatar
Marko Petric committed
      logger.info("=" * 80)
Marko Petric's avatar
Marko Petric committed
      logger.info("=" * 80)
Marko Petric's avatar
Marko Petric committed
      self.__printSteeringFile(parser)
    ''' get lists of trackers and calorimeters that are defined in detectorDescription (the compact xml file)'''
    trackers, calos, unknown = [], [], []
    for i in detectorDescription.detectors():
Andre Sailer's avatar
Andre Sailer committed
      det = DDG4.DetElement(i.second.ptr())
Marko Petric's avatar
Marko Petric committed
      sd = detectorDescription.sensitiveDetector(name)
        logger.info('getDetectorLists - found active detector %s type: %s', name, detType)
        if any(pat.lower() in detType.lower() for pat in self.action.trackerSDTypes):
          logger.info('getDetectorLists - Identified %s as a tracker', name)
Marko Petric's avatar
Marko Petric committed
          trackers.append(det.name())
        elif any(pat.lower() in detType.lower() for pat in self.action.calorimeterSDTypes):
          logger.info('getDetectorLists - Identified %s as a calorimeter', name)
Marko Petric's avatar
Marko Petric committed
          calos.append(det.name())
          logger.warning('getDetectorLists - Unknown sensitive detector type: %s', detType)
Marko Petric's avatar
Marko Petric committed
# ==================================================================================

  def run(self):
    """setup the geometry and dd4hep and geant4 and do what was asked to be done"""
    import ROOT
    ROOT.PyConfig.IgnoreCommandLineOptions = True

Marko Petric's avatar
Marko Petric committed
    import DDG4
    import dd4hep

    self.printLevel = getOutputLevel(self.printLevel)
    dd4hep.setPrintLevel(self.printLevel)
    for compactFile in self.compactFile:
      kernel.loadGeometry(str("file:" + os.path.abspath(compactFile)))
    detectorDescription = kernel.detectorDescription()
Marko Petric's avatar
Marko Petric committed
    DDG4.importConstants(detectorDescription)
Marko Petric's avatar
Marko Petric committed
  # ----------------------------------------------------------------------------------
Marko Petric's avatar
Marko Petric committed
    # simple = DDG4.Geant4( kernel, tracker='Geant4TrackerAction',calo='Geant4CalorimeterAction')
    # geant4 = DDG4.Geant4( kernel, tracker='Geant4TrackerCombineAction',calo='Geant4ScintillatorCalorimeterAction')
    geant4 = DDG4.Geant4(kernel, tracker=self.action.tracker, calo=self.action.calo)
    if not self.disableSignalHandler:
      geant4.registerInterruptHandler()
    geant4.printDetectors()
Marko Petric's avatar
Marko Petric committed
    if self.runType == "vis":
      uiaction = geant4.setupUI(typ="tcsh", vis=True, macro=self.macroFile)
    elif self.runType == "qt":
      uiaction = geant4.setupUI(typ="qt", vis=True, macro=self.macroFile)
Marko Petric's avatar
Marko Petric committed
    elif self.runType == "run":
      uiaction = geant4.setupUI(typ="tcsh", vis=False, macro=self.macroFile, ui=False)
Marko Petric's avatar
Marko Petric committed
    elif self.runType == "shell":
      uiaction = geant4.setupUI(typ="tcsh", vis=False, macro=None, ui=True)
      uiaction = geant4.setupUI(typ="tcsh", vis=False, macro=None, ui=False)
Marko Petric's avatar
Marko Petric committed
      logger.error("unknown runType")
    # User Configuration for the Geant4Phases
    uiaction.ConfigureCommands = self.ui._commandsConfigure
    uiaction.InitializeCommands = self.ui._commandsInitialize
    uiaction.PostRunCommands = self.ui._commandsPostRun
    uiaction.PreRunCommands = self.ui._commandsPreRun
    uiaction.TerminateCommands = self.ui._commandsTerminate
Marko Petric's avatar
Marko Petric committed
    kernel.NumEvents = self.numberOfEvents
Marko Petric's avatar
Marko Petric committed
    # -----------------------------------------------------------------------------------
    self.__setMagneticFieldOptions(geant4)
    # configure geometry creation
    self.geometry.constructGeometry(kernel, geant4, self.output.geometry)
    # ----------------------------------------------------------------------------------
    # Configure run, event, track, step, and stack actions, if present
    for action_list, DDG4_Action, kernel_Action in \
        [(self.action.run, DDG4.RunAction, kernel.runAction),
         (self.action.event, DDG4.EventAction, kernel.eventAction),
         (self.action.track, DDG4.TrackingAction, kernel.trackingAction),
         (self.action.step, DDG4.SteppingAction, kernel.steppingAction),
         (self.action.stack, DDG4.StackingAction, kernel.stackingAction)]:
      for action_dict in action_list:
        action = DDG4_Action(kernel, action_dict["name"])
        for parameter, value in action_dict.get('parameter', {}).items():
          setattr(action, parameter, value)

    # ----------------------------------------------------------------------------------
Marko Petric's avatar
Marko Petric committed
    run1 = DDG4.RunAction(kernel, 'Geant4TestRunAction/RunInit')
    kernel.registerGlobalAction(run1)
    kernel.runAction().add(run1)

    # Configure the random seed, do it before the I/O because we might change the seed!
Marko Petric's avatar
Marko Petric committed
    self.random.initialize(DDG4, kernel, self.output.random)
    # Configure the output file format and plugin
    self.outputConfig.initialize(dd4hepsimulation=self, geant4=geant4)
    if self.enableGun:
Marko Petric's avatar
Marko Petric committed
      gun = DDG4.GeneratorAction(kernel, "Geant4ParticleGun/" + "Gun")
      self.gun.setOptions(gun)
      gun.Standalone = False
      gun.Mask = 1
      actionList.append(gun)
      self.__applyBoostOrSmear(kernel, actionList, 1)
      logger.info("++++ Adding DD4hep Particle Gun ++++")
Marko Petric's avatar
Marko Petric committed
      # GPS Create something
      self._g4gun = DDG4.GeneratorAction(kernel, "Geant4GeneratorWrapper/Gun")
      self._g4gun.Uses = 'G4ParticleGun'
      self._g4gun.Mask = 2
      logger.info("++++ Adding Geant4 Particle Gun ++++")
      actionList.append(self._g4gun)

    if self.enableG4GPS:
Marko Petric's avatar
Marko Petric committed
      # GPS Create something
      self._g4gps = DDG4.GeneratorAction(kernel, "Geant4GeneratorWrapper/GPS")
      self._g4gps.Uses = 'G4GeneralParticleSource'
      self._g4gps.Mask = 3
      logger.info("++++ Adding Geant4 General Particle Source ++++")
    start = 4
    for index, plugin in enumerate(self.inputConfig.userInputPlugin, start=start):
      gen = plugin(self)
      gen.Mask = index
      start = index + 1
      self.__applyBoostOrSmear(kernel, actionList, index)
      logger.info("++++ Adding User Plugin %s ++++", gen.Name)
    for index, inputFile in enumerate(self.inputFiles, start=start):
Marko Petric's avatar
Marko Petric committed
        gen = DDG4.GeneratorAction(kernel, "LCIOInputAction/LCIO%d" % index)
Jan Strube's avatar
Jan Strube committed
        gen.Parameters = self.lcio.getParameters()
Marko Petric's avatar
Marko Petric committed
        gen.Input = "LCIOFileReader|" + inputFile
Marko Petric's avatar
Marko Petric committed
        gen = DDG4.GeneratorAction(kernel, "LCIOInputAction/STDHEP%d" % index)
        gen.Input = "LCIOStdHepReader|" + inputFile
Marko Petric's avatar
Marko Petric committed
        gen = DDG4.GeneratorAction(kernel, "Geant4InputAction/HEPEvt%d" % index)
        gen.Input = "Geant4EventReaderHepEvtShort|" + inputFile
Marko Petric's avatar
Marko Petric committed
        gen = DDG4.GeneratorAction(kernel, "Geant4InputAction/hepevt%d" % index)
        gen.Input = "Geant4EventReaderHepEvtLong|" + inputFile
      elif inputFile.endswith(tuple([".hepmc"] + HEPMC3_SUPPORTED_EXTENSIONS)):
        if self.hepmc3.useHepMC3:
          gen = DDG4.GeneratorAction(kernel, "Geant4InputAction/hepmc%d" % index)
          gen.Parameters = self.hepmc3.getParameters()
          gen.Input = "HEPMC3FileReader|" + inputFile
        else:
          gen = DDG4.GeneratorAction(kernel, "Geant4InputAction/hepmc%d" % index)
          gen.Input = "Geant4EventReaderHepMC|" + inputFile
      elif inputFile.endswith(".pairs"):
Marko Petric's avatar
Marko Petric committed
        gen = DDG4.GeneratorAction(kernel, "Geant4InputAction/GuineaPig%d" % index)
        gen.Input = "Geant4EventReaderGuineaPig|" + inputFile
        gen.Parameters = self.guineapig.getParameters()
      elif inputFile.endswith(tuple(EDM4HEP_INPUT_EXTENSIONS)):
        # EDM4HEP must come after HEPMC3 because of .root also part of hepmc3 extensions
        gen = DDG4.GeneratorAction(kernel, "Geant4InputAction/EDM4hep%d" % index)
        gen.Input = "EDM4hepFileReader|" + inputFile
Marko Petric's avatar
Marko Petric committed
        # this should never happen because we already check at the top, but in case of some LogicError...
        raise RuntimeError("Unknown input file type: %s" % inputFile)
      gen.AlternativeDecayStatuses = self.physics.alternativeDecayStatuses
      gen.Sync = self.skipNEvents
      gen.Mask = index
      actionList.append(gen)
      self.__applyBoostOrSmear(kernel, actionList, index)

      generationInit = self._buildInputStage(geant4, actionList, output_level=self.output.inputStage,
                                             have_mctruth=self._enablePrimaryHandler())
Marko Petric's avatar
Marko Petric committed
    # ================================================================================================
Marko Petric's avatar
Marko Petric committed
    part = DDG4.GeneratorAction(kernel, "Geant4ParticleHandler/ParticleHandler")
Marko Petric's avatar
Marko Petric committed
    # part.SaveProcesses = ['conv','Decay']
Marko Petric's avatar
Marko Petric committed
    part.SaveProcesses = self.part.saveProcesses
    part.MinimalKineticEnergy = self.part.minimalKineticEnergy
Marko Petric's avatar
Marko Petric committed
    part.KeepAllParticles = self.part.keepAllParticles
    part.PrintEndTracking = self.part.printEndTracking
    part.PrintStartTracking = self.part.printStartTracking
    part.MinDistToParentVertex = self.part.minDistToParentVertex
    part.OutputLevel = self.output.part
    if self.part.enableDetailedHitsAndParticleInfo:
Marko Petric's avatar
Marko Petric committed
      self.part.setDumpDetailedParticleInfo(kernel, DDG4)
    self.part.setupUserParticleHandler(part, kernel, DDG4)
Marko Petric's avatar
Marko Petric committed
    # =================================================================================

    # Setup global filters for use in sensitive detectors
Marko Petric's avatar
Marko Petric committed
      self.filter.setupFilters(kernel)
    except RuntimeError as e:
Marko Petric's avatar
Marko Petric committed
    # =================================================================================
    # get lists of trackers and calorimeters in detectorDescription
    trk, cal, unk = self.getDetectorLists(detectorDescription)
    for detectors, function, defFilter, defAction, abort in \
        [(trk, geant4.setupTracker, self.filter.tracker, self.action.tracker, False),
         (cal, geant4.setupCalorimeter, self.filter.calo, self.action.calo, False),
         (unk, geant4.setupDetector, None, "No Default", True),
         ]:
        self.__setupSensitiveDetectors(detectors, function, defFilter, defAction, abort)
      except Exception as e:
        logger.error("Failed setting up sensitive detector %s", e)
        raise
Marko Petric's avatar
Marko Petric committed
  # =================================================================================
Marko Petric's avatar
Marko Petric committed
    _phys = self.physics.setupPhysics(kernel, name=self.physicsList)
    _phys.verbosity = self.output.physics
Marko Petric's avatar
Marko Petric committed
    # add the G4StepLimiterPhysics to activate the max step limits in volumes
    ph = DDG4.PhysicsList(kernel, 'Geant4PhysicsList/Myphysics')
    ph.addPhysicsConstructor(str('G4StepLimiterPhysics'))
    _phys.add(ph)
    dd4hep.setPrintLevel(self.printLevel)
Marko Petric's avatar
Marko Petric committed
    # GPS
    if self._g4gun is not None:
      self._g4gun.generator()
    if self._g4gps is not None:
      self._g4gps.generator()

    startUpTime, _sysTime, _cuTime, _csTime, _elapsedTime = os.times()

    exitCode = 0
    if not kernel.run():
      logger.error("Simulation failed!")
      exitCode += 1
    if not kernel.terminate():
      exitCode += 1
      logger.error("Termination failed!")
    totalTimeUser, totalTimeSys, _cuTime, _csTime, _elapsedTime = os.times()
    processedEvents = self.numberOfEvents
    if generationInit:
      processedEvents = int(generationInit.numberOfEvents)
      if self.numberOfEvents < 0:
        processedEvents -= 1
        logger.debug(f"Correcting number of events to: {processedEvents}")

      logger.info("Total Time:   %3.2f s (User), %3.2f s (System)" %
Marko Petric's avatar
Marko Petric committed
                  (totalTimeUser, totalTimeSys))
        eventTime = totalTimeUser - startUpTime
        perEventTime = eventTime / processedEvents
        logger.info("StartUp Time: %3.2f s, Processing and Init: %3.2f s (~%3.2f s/Event) "
Marko Petric's avatar
Marko Petric committed
                    % (startUpTime, eventTime, perEventTime))
  def __setMagneticFieldOptions(self, geant4):
    """ create and configure the magnetic tracking setup """
    field = geant4.addConfig('Geant4FieldTrackingSetupAction/MagFieldTrackingSetup')
Marko Petric's avatar
Marko Petric committed
    field.stepper = self.field.stepper
    field.equation = self.field.equation
    field.eps_min = self.field.eps_min
    field.eps_max = self.field.eps_max
    field.min_chord_step = self.field.min_chord_step
    field.delta_chord = self.field.delta_chord
    field.delta_intersection = self.field.delta_intersection
Marko Petric's avatar
Marko Petric committed
    field.delta_one_step = self.field.delta_one_step
    field.largest_step = self.field.largest_step
  def __checkFilesExist(self, fileNames, fileType=''):
    """Make sure all files in the given list exist, add to errorMessage otherwise.


    :param list fileNames: list of files to check for existence
    :param str fileType: type if file, for nicer error message
    """
    if isinstance(fileNames, str):
      fileNames = [fileNames]
    for fileName in fileNames:
      if not os.path.exists(fileName) and not urlparse(fileName).scheme:
        self._errorMessages.append(f"ERROR: The {fileType}file '{fileName}' does not exist")

  def __checkFileFormat(self, fileNames, extensions):
    """check if the fileName is allowed, note that the filenames are case
    sensitive, and in case of hepevt we depend on this to identify short and long versions of the content
    """
    if not all(fileName.endswith(tuple(extensions)) for fileName in fileNames):
      self._errorMessages.append(f"ERROR: Unknown fileformat for file(s): {','.join(fileNames)}")
    is_hepmc3_extension = any(fileName.endswith(tuple(HEPMC3_SUPPORTED_EXTENSIONS)) for fileName in fileNames)
    if not self.hepmc3.useHepMC3 and is_hepmc3_extension:
      self._errorMessages.append("ERROR: HepMC3 files or compressed HepMC2 require the use of HepMC3 library")
Marko Petric's avatar
Marko Petric committed
  def __applyBoostOrSmear(self, kernel, actionList, mask):
    """apply boost or smearing for given mask index"""
    if self.crossingAngleBoost:
      lbo = DDG4.GeneratorAction(kernel, "Geant4InteractionVertexBoost")
      lbo.Angle = self.crossingAngleBoost
      lbo.Mask = mask
      actionList.append(lbo)

    if any(self.vertexSigma) or any(self.vertexOffset):
      vSmear = DDG4.GeneratorAction(kernel, "Geant4InteractionVertexSmear")
      vSmear.Offset = self.vertexOffset
      vSmear.Sigma = self.vertexSigma
      vSmear.Mask = mask
      actionList.append(vSmear)
Marko Petric's avatar
Marko Petric committed
  def __parseAllHelper(self, parsed):
    """ parse all the options for the helper """
    parsedDict = vars(parsed)
    for name, obj in vars(self).items():
Marko Petric's avatar
Marko Petric committed
      if isinstance(obj, ConfigHelper):
        for var in obj.getOptions():
Marko Petric's avatar
Marko Petric committed
          key = "%s.%s" % (name, var)
          if key in parsedDict:
Marko Petric's avatar
Marko Petric committed
              obj.setOption(var, parsedDict[key])
Marko Petric's avatar
Marko Petric committed
              self._errorMessages.append("ERROR: %s " % e)
              if logger.level <= logging.DEBUG:
                self._errorMessages.append(traceback.format_exc())

  def __checkOutputLevel(self, level):
    """return outputlevel as int so we don't have to import anything for faster startup"""
    try:
      return outputLevel(level)
    except ValueError:
      self._errorMessages.append("ERROR: printLevel is neither integer nor string")
      return -1
    except KeyError:
Marko Petric's avatar
Marko Petric committed
      self._errorMessages.append("ERROR: printLevel '%s' unknown" % level)
Andre Sailer's avatar
Andre Sailer committed
  def __setupSensitiveDetectors(self, detectors, setupFunction, defaultFilter=None,
                                defaultAction=None, abortForMissingAction=False,
                                ):
    """Attach sensitive detector actions for all subdetectors.

    Can be steered with the `Action` ConfigHelpers

    :param detectors: list of detectors
    :param setupFunction: function used to register the sensitive detector
    :param defaultFilter: default filter to apply for given types
    :param abortForMissingAction: if true end program if there is no action found
      logger.info('Setting up SD for %s with %s', det, defaultAction)
      action = None
      for pattern in self.action.mapActions:
        if pattern.lower() in det.lower():
          action = self.action.mapActions[pattern]
Marko Petric's avatar
Marko Petric committed
          logger.info('       replace default action with : %s', action)
      if abortForMissingAction and action is None:
        logger.error('Cannot find Action for detector %s. You have to extend "action.mapAction"', det)
        raise RuntimeError("Cannot find Action")
Andre Sailer's avatar
Andre Sailer committed
      seq, act = setupFunction(det, action)
Marko Petric's avatar
Marko Petric committed
      self.filter.applyFilters(seq, det, defaultFilter)
Marko Petric's avatar
Marko Petric committed
      # set detailed hit creation mode for this
      if self.enableDetailedShowerMode:
        if isinstance(act, list):
          for a in act:
            a.HitCreationMode = 2
        else:
          act.HitCreationMode = 2
Marko Petric's avatar
Marko Petric committed
  def __printSteeringFile(self, parser):
    """print the parameters formated as a steering file"""

    steeringFileBase = textwrap.dedent("""\
        from DDSim.DD4hepSimulation import DD4hepSimulation
        from g4units import mm, GeV, MeV
        SIM = DD4hepSimulation()
        """)
    steeringFileBase += "\n"
    optionDict = parser._option_string_actions
    parameters = vars(self)
    for parName, parameter in sorted(list(parameters.items()), key=sortParameters):
      if parName.startswith("_"):
        continue
Marko Petric's avatar
Marko Petric committed
      if isinstance(parameter, ConfigHelper):
        steeringFileBase += "\n\n"
        steeringFileBase += "################################################################################\n"
Marko Petric's avatar
Marko Petric committed
        steeringFileBase += "## %s \n" % "\n## ".join(parameter.__doc__.splitlines())
        steeringFileBase += "################################################################################\n"
        options = parameter.getOptions()
        for opt, optionsDict in sorted(options.items(), key=sortParameters):
          if isinstance(optionsDict.get('help'), str):
            steeringFileBase += "\n## %s\n" % "\n## ".join(optionsDict.get('help').splitlines())
Marko Petric's avatar
Marko Petric committed
          # add quotes if it is a string
Marko Petric's avatar
Marko Petric committed
            steeringFileBase += "SIM.%s.%s = \"%s\"\n" % (parName, opt, parValue)
Marko Petric's avatar
Marko Petric committed
            steeringFileBase += "SIM.%s.%s = %s\n" % (parName, opt, parValue)
Marko Petric's avatar
Marko Petric committed
        # get the docstring from the command line parameter
        optionObj = optionDict.get("--" + parName, None)
        if isinstance(optionObj, argparse._StoreAction):
          steeringFileBase += "## %s\n" % "\n## ".join(optionObj.help.splitlines())
Marko Petric's avatar
Marko Petric committed
        # add quotes if it is a string
Marko Petric's avatar
Marko Petric committed
          steeringFileBase += "SIM.%s = \"%s\"" % (parName, str(parameter))
Marko Petric's avatar
Marko Petric committed
          steeringFileBase += "SIM.%s = %s" % (parName, str(parameter))
    for line in steeringFileBase.splitlines():
      print(line)
Marko Petric's avatar
Marko Petric committed
  def _consistencyChecks(self):
    """Check if the requested setup makes sense, or if there is something preventing it from working correctly

    Appends error messages to self._errorMessages

    :returns: None
    """

    if not self.compactFile:
      self._errorMessages.append("ERROR: No geometry compact file provided")

    if self.runType == "batch":
      if not self.numberOfEvents:
        self._errorMessages.append("ERROR: Batch mode requested, but did not set number of events")
      if not (self.inputFiles or self.enableGun or self.inputConfig.userInputPlugin):
        self._errorMessages.append("ERROR: Batch mode requested, but did not set inputFile(s), gun, or userInputPlugin")
Marko Petric's avatar
Marko Petric committed
    if self.inputFiles and (self.enableG4Gun or self.enableG4GPS):
      self._errorMessages.append("ERROR: Cannot use both inputFiles and Geant4Gun or GeneralParticleSource")

Marko Petric's avatar
Marko Petric committed
    if self.enableGun and (self.enableG4Gun or self.enableG4GPS):
      self._errorMessages.append("ERROR: Cannot use both DD4hepGun and Geant4 Gun or GeneralParticleSource")

    if self.inputConfig.userInputPlugin and (self.enableG4Gun or self.enableG4GPS):
      self._errorMessages.append("ERROR: Cannot use both userInputPlugin and Geant4 Gun or GeneralParticleSource")

    if self.numberOfEvents < 0 and not self.inputFiles:
      self._errorMessages.append("ERROR: Negative number of events only sensible for inputFiles")
Marko Petric's avatar
Marko Petric committed
  def _enablePrimaryHandler(self):
    """ the geant4 Gun or GeneralParticleSource cannot be used together with the PrimaryHandler.
        Particles would be simulated multiple times

    :returns: True or False
    """
    enablePrimaryHandler = not (self.enableG4Gun or self.enableG4GPS)
    if enablePrimaryHandler:
    return enablePrimaryHandler
  def _buildInputStage(self, geant4, generator_input_modules, output_level=None, have_mctruth=True):
    """
    Generic build of the input stage with multiple input modules.
    Actions executed are:
    1) Register Generation initialization action
    2) Append all modules to build the complete input record
    These modules are readers/particle sources, boosters and/or smearing actions.
    3) Merge all existing interaction records
    4) Add the MC truth handler
    """
    from DDG4 import GeneratorAction
    ga = geant4.kernel().generatorAction()

    # Register Generation initialization action
    gen = GeneratorAction(geant4.kernel(), "Geant4GeneratorActionInit/GenerationInit")
    if output_level is not None:
      gen.OutputLevel = output_level
    ga.adopt(gen)

    # Now append all modules to build the complete input record
    # These modules are readers/particle sources, boosters and/or smearing actions
    for gen in generator_input_modules:
      gen.enableUI()
      if output_level is not None:
        gen.OutputLevel = output_level
      ga.adopt(gen)

    # Merge all existing interaction records
    gen = GeneratorAction(geant4.kernel(), "Geant4InteractionMerger/InteractionMerger")
    gen.enableUI()
    if output_level is not None:
      gen.OutputLevel = output_level
    ga.adopt(gen)

    # Finally generate Geant4 primaries
    if have_mctruth:
      gen = GeneratorAction(geant4.kernel(), "Geant4PrimaryHandler/PrimaryHandler")
Marko Petric's avatar
Marko Petric committed
      gen.RejectPDGs = ConfigHelper.makeString(self.physics.rejectPDGs)
      gen.ZeroTimePDGs = ConfigHelper.makeString(self.physics.zeroTimePDGs)
      gen.enableUI()
      if output_level is not None:
        gen.OutputLevel = output_level
      ga.adopt(gen)
    # Puuuhh! All done.
################################################################################
Marko Petric's avatar
Marko Petric committed
# MODULE FUNCTIONS GO HERE
################################################################################


def sortParameters(key):
  from functools import cmp_to_key

  def _sortParameters(parA, parB):
    """sort the parameters by name: first normal parameters, then set of
    parameters based on ConfigHelper objects
    """
    parTypeA = parA[1]
    parTypeB = parB[1]
    if isinstance(parTypeA, ConfigHelper) and isinstance(parTypeB, ConfigHelper):
      return 1 if str(parA[0]) > str(parB[0]) else -1
    elif isinstance(parTypeA, ConfigHelper):
      return 1
    elif isinstance(parTypeB, ConfigHelper):
      return -1
    else:
      return 1 if str(parA[0]) > str(parB[0]) else -1

  return cmp_to_key(_sortParameters)(key)
def getOutputLevel(level):
  """return output.LEVEL"""
  from DDG4 import OutputLevel
Marko Petric's avatar
Marko Petric committed
  levels = {1: OutputLevel.VERBOSE,
            2: OutputLevel.DEBUG,
            3: OutputLevel.INFO,
            4: OutputLevel.WARNING,
            5: OutputLevel.ERROR,
            6: OutputLevel.FATAL,
            7: OutputLevel.ALWAYS}
  return levels[level]