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