DDSim/Helper/OutputConfig.py

0001 """Class for output file configuration"""
0002 import logging
0003
0004 from DDSim.Helper.ConfigHelper import ConfigHelper
0005
0006 logger = logging.getLogger(__name__)
0007
0008 #: True if DD4hep was built with LCIO
0009 DD4HEP_USE_LCIO = "@DD4HEP_USE_LCIO@" != "OFF"
0010 #: True if DD4hep was built with EDM4hep
0011 DD4HEP_USE_EDM4HEP = "@DD4HEP_USE_EDM4HEP@" != "OFF"
0012
0013
0014 def defaultOutputFile():
0015   if DD4HEP_USE_LCIO and not DD4HEP_USE_EDM4HEP:
0016     return "ddsimOutput.slcio"
0017   return "ddsimOutput.root"
0018
0019
0020 class OutputConfig(ConfigHelper):
0021   """Configuration for Output Files."""
0022
0023   def __init__(self):
0024     super(OutputConfig, self).__init__()
0025     self._userPlugin = None
0026     self._forceLCIO = False
0027     self._forceEDM4HEP = False
0028     self._forceDD4HEP = False
0029     # no closeProperties, allow custom ones for userPlugin configuration
0030
0031   def _checkConsistency(self):
0032     """Raise error if more than one force flag is true."""
0033     if self._forceLCIO + self._forceEDM4HEP + self._forceDD4HEP > 1:
0034       raise RuntimeError(f"OutputConfig error: More than one force flag enabled: LCIO({self._forceLCIO}),"
0035                          f" EDM4HEP({self._forceEDM4HEP}), DD4HEP({self._forceDD4HEP})")
0036
0037   @property
0038   def forceLCIO(self):
0039     """Use the LCIO output plugin regardless of outputfilename."""
0040     return self._forceLCIO
0041
0042   @forceLCIO.setter
0043   def forceLCIO(self, val):
0044     self._forceLCIO = self.makeBool(val)
0045     if self._forceLCIO:
0046       if not DD4HEP_USE_LCIO:
0047         raise RuntimeError("OutputConfig error: forceLCIO requested, but LCIO not available!")
0048       self._checkConsistency()
0049
0050   @property
0051   def forceEDM4HEP(self):
0052     """Use the EDM4HEP output plugin regardless of outputfilename."""
0053     return self._forceEDM4HEP
0054
0055   @forceEDM4HEP.setter
0056   def forceEDM4HEP(self, val):
0057     self._forceEDM4HEP = self.makeBool(val)
0058     if self._forceEDM4HEP:
0059       if not DD4HEP_USE_EDM4HEP:
0060         raise RuntimeError("OutputConfig error: forceEDM4HEP requested, but EDM4HEP not available!")
0061       self._checkConsistency()
0062
0063   @property
0064   def forceDD4HEP(self):
0065     """Use the DD4HEP output plugin regardless of outputfilename."""
0066     return self._forceDD4HEP
0067
0068   @forceDD4HEP.setter
0069   def forceDD4HEP(self, val):
0070     self._forceDD4HEP = self.makeBool(val)
0071     if self._forceDD4HEP:
0072       self._checkConsistency()
0073
0074   @property
0075   def userOutputPlugin(self):
0076     """Set a function to configure the outputFile.
0077
0078     The function must take a ``DD4hepSimulation`` object as its only argument and return ``None``.
0079
0080     For example one can add this to the ddsim steering file:
0081
0082       def exampleUserPlugin(dd4hepSimulation):
0083         '''Example code for user created plugin.
0084
0085         :param DD4hepSimulation dd4hepSimulation: The DD4hepSimulation instance, so all parameters can be accessed
0086         :return: None
0087         '''
0088         from DDG4 import EventAction, Kernel
0089         dd = dd4hepSimulation  # just shorter variable name
0090         evt_root = EventAction(Kernel(), 'Geant4Output2ROOT/' + dd.outputFile, True)
0091         evt_root.HandleMCTruth = True or False
0092         evt_root.Control = True
0093         output = dd.outputFile
0094         if not dd.outputFile.endswith(dd.outputConfig.myExtension):
0095           output = dd.outputFile + dd.outputConfig.myExtension
0096         evt_root.Output = output
0097         evt_root.enableUI()
0098         Kernel().eventAction().add(evt_root)
0099         return None
0100
0101       SIM.outputConfig.userOutputPlugin = exampleUserPlugin
0102       # arbitrary options can be created and set via the steering file or command line
0103       SIM.outputConfig.myExtension = '.csv'
0104     """
0105     return self._userPlugin
0106
0107   @userOutputPlugin.setter
0108   def userOutputPlugin(self, userOutputPluginConfig):
0109     if userOutputPluginConfig is None:
0110       return
0111     if not callable(userOutputPluginConfig):
0112       raise RuntimeError("The provided userPlugin is not a callable function.")
0113     self._userPlugin = userOutputPluginConfig
0114
0115   def initialize(self, dd4hepsimulation, geant4):
0116     """Configure the output file and plugin."""
0117     if callable(self._userPlugin):
0118       logger.info("++++ Setting up UserPlugin for Output ++++")
0119       return self._userPlugin(dd4hepsimulation)
0120
0121     if self.forceLCIO:
0122       return self._configureLCIO(dd4hepsimulation, geant4)
0123
0124     if self.forceEDM4HEP:
0125       return self._configureEDM4HEP(dd4hepsimulation, geant4)
0126
0127     if self.forceDD4HEP:
0128       return self._configureDD4HEP(dd4hepsimulation, geant4)
0129
0130     if dd4hepsimulation.outputFile.endswith(".slcio"):
0131       return self._configureLCIO(dd4hepsimulation, geant4)
0132
0133     if dd4hepsimulation.outputFile.endswith(".root") and DD4HEP_USE_EDM4HEP:
0134       return self._configureEDM4HEP(dd4hepsimulation, geant4)
0135
0136     if dd4hepsimulation.outputFile.endswith(".root"):
0137       return self._configureDD4HEP(dd4hepsimulation, geant4)
0138
0139   def _configureLCIO(self, dds, geant4):
0140     if not DD4HEP_USE_LCIO:
0141       raise RuntimeError("DD4HEP was not build wiht LCIO support: please change output format %s" % dds.outputFile)
0142     logger.info("++++ Setting up LCIO Output ++++")
0143     lcOut = geant4.setupLCIOOutput('LcioOutput', dds.outputFile)
0144     lcOut.RunHeader = dds.meta.addParametersToRunHeader(dds)
0145     eventPars = dds.meta.parseEventParameters()
0146     lcOut.EventParametersString, lcOut.EventParametersInt, lcOut.EventParametersFloat = eventPars
0147     lcOut.RunNumberOffset = dds.meta.runNumberOffset if dds.meta.runNumberOffset > 0 else 0
0148     lcOut.EventNumberOffset = dds.meta.eventNumberOffset if dds.meta.eventNumberOffset > 0 else 0
0149     return
0150
0151   def _configureEDM4HEP(self, dds, geant4):
0152     logger.info("++++ Setting up EDM4hep ROOT Output ++++")
0153     e4Out = geant4.setupEDM4hepOutput('EDM4hepOutput', dds.outputFile)
0154     eventPars = dds.meta.parseEventParameters()
0155     e4Out.RunHeader = dds.meta.addParametersToRunHeader(dds)
0156     e4Out.EventParametersString, e4Out.EventParametersInt, e4Out.EventParametersFloat = eventPars
0157     e4Out.RunNumberOffset = dds.meta.runNumberOffset if dds.meta.runNumberOffset > 0 else 0
0158     e4Out.EventNumberOffset = dds.meta.eventNumberOffset if dds.meta.eventNumberOffset > 0 else 0
0159     return
0160
0161   def _configureDD4HEP(self, dds, geant4):
0162     logger.info("++++ Setting up DD4hep's ROOT Output ++++")
0163     geant4.setupROOTOutput('RootOutput', dds.outputFile)
0164     return