ClientTests/scripts/DDG4TestSetup.py

0001 # ==========================================================================
0002 #  AIDA Detector description implementation
0003 # --------------------------------------------------------------------------
0004 # Copyright (C) Organisation europeenne pour la Recherche nucleaire (CERN)
0005 # All rights reserved.
0006 #
0007 # For the licensing terms see $DD4hepINSTALL/LICENSE.
0008 # For the list of contributors see $DD4hepINSTALL/doc/CREDITS.
0009 #
0010 # ==========================================================================
0011 from __future__ import absolute_import, unicode_literals
0012 import sys
0013 import logging
0014 import DDG4
0015 from DDG4 import OutputLevel as Output
0016 from g4units import GeV, MeV
0017
0018 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)
0019 logger = logging.getLogger(__name__)
0020
0021 """
0022
0023    dd4hep example setup using the python configuration
0024
0025    \author  M.Frank
0026    \version 1.0
0027
0028 """
0029
0030
0031 class Setup:
0032   def __init__(self, geometry_file, macro=None, vis=None):
0033     self.kernel = DDG4.Kernel()
0034     self.kernel.setOutputLevel(str('Geant4Converter'), Output.DEBUG)
0035     self.kernel.setOutputLevel(str('Gun'), Output.INFO)
0036     self.description = self.kernel.detectorDescription()
0037     self.kernel.loadGeometry(str(geometry_file))
0038     self.geant4 = DDG4.Geant4(self.kernel)
0039     self.geant4.printDetectors()
0040     self.ui = self.geant4.setupCshUI(macro=macro, vis=vis)
0041
0042   def configure(self):
0043     # Configure field
0044     self.field = self.geant4.setupTrackingField(prt=True)
0045     return self
0046
0047   def defineOutput(self, output):
0048     # Configure I/O
0049     evt_write = self.geant4.setupROOTOutput('RootOutput', output, mc_truth=True)
0050     return evt_write
0051
0052   def defineEdm4hepOutput(self, output):
0053     # Configure I/O
0054     evt_write = self.geant4.setupEDM4hepOutput('Edm4hepOutput', output)
0055     return evt_write
0056
0057   def setupGun(self, name="Gun", particle='pi-', energy=100 * GeV, multiplicity=1):
0058     # Setup particle gun
0059     return self.geant4.setupGun(name, particle=particle, energy=energy, multiplicity=multiplicity)
0060
0061   def setupInput(self, spec, mask=1):
0062     input_action = DDG4.GeneratorAction(self.kernel, "Geant4InputAction/Input")
0063     input_action.Input = spec
0064     input_action.MomentumScale = 1.0
0065     input_action.Mask = mask
0066     self.geant4.buildInputStage([input_action])
0067
0068   def setupGenerator(self):
0069     # And handle the simulation particles.
0070     part = DDG4.GeneratorAction(self.kernel, "Geant4ParticleHandler/ParticleHandler")
0071     self.kernel.generatorAction().adopt(part)
0072     part.SaveProcesses = ['conv', 'Decay']
0073     part.MinimalKineticEnergy = 1 * MeV
0074     part.OutputLevel = 5  # generator_output_level
0075     part.enableUI()
0076     return part
0077
0078   def setupPhysics(self, model='QGSP_BERT', dump=False):
0079     # Now build the physics list:
0080     self.phys = self.kernel.physicsList()
0081     self.phys.extends = model
0082     self.phys.decays = True
0083     self.phys.enableUI()
0084     if dump:
0085       self.phys.dump()
0086     return self
0087
0088   def run(self, num_events=None):
0089     # and run
0090     self.geant4.execute(num_events)
0091     return self
0092
0093   # Stop the entire excercise
0094   def terminate(self):
0095     self.kernel.terminate()
0096     logger.info('+++++ All Done....\n\nTEST_PASSED')
0097     sys.exit(0)
0098
0099   # Test the configuration
0100   def test_config(self, have_geo=True):
0101     self.kernel.configure()
0102     if have_geo:
0103       self.kernel.initialize()
0104
0105   # Test runner
0106   def test_run(self, have_geo=True, have_physics=False):
0107     self.test_config(have_geo)
0108     if have_geo:
0109       self.kernel.UI = ''
0110       self.kernel.NumEvents = 0
0111       self.kernel.run()
0112     self.terminate()