Skip to content
Snippets Groups Projects
DDG4.py 24.3 KiB
Newer Older
Marko Petric's avatar
Marko Petric committed
# ==========================================================================
Marko Petric's avatar
Marko Petric committed
#  AIDA Detector description implementation
Marko Petric's avatar
Marko Petric committed
# --------------------------------------------------------------------------
Markus Frank's avatar
Markus Frank committed
# Copyright (C) Organisation europeenne pour la Recherche nucleaire (CERN)
# All rights reserved.
#
# For the licensing terms see $DD4hepINSTALL/LICENSE.
# For the list of contributors see $DD4hepINSTALL/doc/CREDITS.
#
Marko Petric's avatar
Marko Petric committed
# ==========================================================================
from __future__ import absolute_import, unicode_literals
Marko Petric's avatar
Marko Petric committed
from dd4hep_base import *  # noqa: F403
Marko Petric's avatar
Marko Petric committed

Markus Frank's avatar
Markus Frank committed
def loadDDG4():
Marko Petric's avatar
Marko Petric committed
  import ROOT
Markus Frank's avatar
Markus Frank committed
  from ROOT import gSystem
Marko Petric's avatar
Marko Petric committed
  # Try to load libglapi to avoid issues with TLS Static
  # Turn off all errors from ROOT about the library missing
  if('libglapi' not in gSystem.GetLibraries()):
    orgLevel = ROOT.gErrorIgnoreLevel
    ROOT.gErrorIgnoreLevel = 6000
    gSystem.Load("libglapi")
    ROOT.gErrorIgnoreLevel = orgLevel
Marko Petric's avatar
Marko Petric committed
  if platform.system() == "Darwin":
    gSystem.SetDynamicPath(os.environ['DD4HEP_LIBRARY_PATH'])
    os.environ['DYLD_LIBRARY_PATH'] = os.pathsep.join([os.environ['DD4HEP_LIBRARY_PATH'],
                                                       os.environ.get('DYLD_LIBRARY_PATH', '')]).strip(os.pathsep)
  result = gSystem.Load("libDDG4Plugins")
Marko Petric's avatar
Marko Petric committed
    raise Exception('DDG4.py: Failed to load the DDG4 library libDDG4Plugins: ' + gSystem.GetErrorStr())
Markus Frank's avatar
Markus Frank committed
  from ROOT import dd4hep as module
Markus Frank's avatar
Markus Frank committed
  return module

# We are nearly there ....
current = __import__(__name__)
Marko Petric's avatar
Marko Petric committed


def _import_class(ns, nam):
  scope = getattr(current, ns)
  setattr(current, nam, getattr(scope, nam))

Markus Frank's avatar
Markus Frank committed
try:
Marko Petric's avatar
Marko Petric committed
  dd4hep = loadDDG4()
Marko Petric's avatar
Marko Petric committed
  logger.error('+--%-100s--+', 100 * '-')
  logger.error('|  %-100s  |', 'Failed to load DDG4 library:')
  logger.error('|  %-100s  |', str(X))
  logger.error('+--%-100s--+', 100 * '-')
Marko Petric's avatar
Marko Petric committed
from ROOT import CLHEP as CLHEP  # noqa
Marko Petric's avatar
Marko Petric committed
Core = dd4hep
Sim = dd4hep.sim
Markus Frank's avatar
Markus Frank committed
Simulation = dd4hep.sim
Marko Petric's avatar
Marko Petric committed
Kernel = Sim.KernelHandle
Interface = Sim.Geant4ActionCreation
Detector = Core.Detector


def _constant(self, name):
  return self.constantAsString(name)

Marko Petric's avatar
Marko Petric committed

Markus Frank's avatar
Markus Frank committed
Detector.globalVal = _constant
Marko Petric's avatar
Marko Petric committed

def importConstants(description, namespace=None, debug=False):
Marko Petric's avatar
Marko Petric committed
  """
  Import the Detector constants into the DDG4 namespace
  """
  ns = current
Marko Petric's avatar
Marko Petric committed
  if namespace is not None and not hasattr(current, namespace):
    import imp
Marko Petric's avatar
Marko Petric committed
    m = imp.new_module('DDG4.' + namespace)
    setattr(current, namespace, m)
    ns = m
Markus Frank's avatar
Markus Frank committed
  evaluator = dd4hep.g4Evaluator()
Markus Frank's avatar
Markus Frank committed
  for c in description.constants():
    if c.second.dataType == 'string':
      strings[c.first] = c.second.GetTitle()
    else:
Marko Petric's avatar
Marko Petric committed
      todo[c.first] = c.second.GetTitle().replace('(int)', '')
  while len(todo) and cnt < 100:
Marko Petric's avatar
Marko Petric committed
      logger.error('%s %d out of %d %s "%s": [%s]\n+++ %s',
                   '+++ FAILED to import',
                   len(todo), len(todo) + num,
                   'global values into namespace',
                   ns.__name__, 'Try to continue anyway', 100 * '=')
      for k, v in todo.items():
        if not hasattr(ns, k):
          logger.error('+++ FAILED to import: "' + k + '" = "' + str(v) + '"')
      logger.info('+++ %s', 100 * '=')

    for k, v in list(todo.items()):
      if not hasattr(ns, k):
        val = evaluator.evaluate(v)
        status = evaluator.status()
        if status == 0:
Marko Petric's avatar
Marko Petric committed
          evaluator.setVariable(k, val)
          setattr(ns, k, val)
          if debug:
            logger.info('Imported global value: "' + k + '" = "' + str(val) + '" into namespace' + ns.__name__)
Marko Petric's avatar
Marko Petric committed
  if cnt < 100:
    logger.info('+++ Imported %d global values to namespace:%s', num, ns.__name__,)

Marko Petric's avatar
Marko Petric committed
def _registerGlobalAction(self, action):
  self.get().registerGlobalAction(Interface.toAction(action))
Marko Petric's avatar
Marko Petric committed


def _registerGlobalFilter(self, filter):
  self.get().registerGlobalFilter(Interface.toAction(filter))
Markus Frank's avatar
Markus Frank committed
def _getKernelProperty(self, name):
Marko Petric's avatar
Marko Petric committed
  ret = Interface.getPropertyKernel(self.get(), name)
Markus Frank's avatar
Markus Frank committed
  if ret.status > 0:
    return ret.data
Marko Petric's avatar
Marko Petric committed
  elif hasattr(self.get(), name):
    return getattr(self.get(), name)
  elif hasattr(self, name):
    return getattr(self, name)
  msg = 'Geant4Kernel::GetProperty [Unhandled]: Cannot access Kernel.' + name
Ercan Pilicer's avatar
Ercan Pilicer committed
  raise KeyError(msg)
Markus Frank's avatar
Markus Frank committed
def _setKernelProperty(self, name, value):
  if Interface.setPropertyKernel(self.get(), str(name), str(value)):
Markus Frank's avatar
Markus Frank committed
    return
Marko Petric's avatar
Marko Petric committed
  msg = 'Geant4Kernel::SetProperty [Unhandled]: Cannot set Kernel.' + name + ' = ' + str(value)
Ercan Pilicer's avatar
Ercan Pilicer committed
  raise KeyError(msg)
Marko Petric's avatar
Marko Petric committed
def _kernel_phase(self, name):
  return self.addSimplePhase(str(name), False)
Marko Petric's avatar
Marko Petric committed
def _kernel_worker(self):
  return Kernel(self.get().createWorker())
Marko Petric's avatar
Marko Petric committed
def _kernel_terminate(self):
  return self.get().terminate()
Kernel.registerGlobalAction = _registerGlobalAction
Kernel.registerGlobalFilter = _registerGlobalFilter
Markus Frank's avatar
Markus Frank committed
Kernel.__getattr__ = _getKernelProperty
Kernel.__setattr__ = _setKernelProperty
Kernel.terminate = _kernel_terminate
Marko Petric's avatar
Marko Petric committed

Marko Petric's avatar
Marko Petric committed


def SensitiveAction(kernel, nam, det, shared=False):
  return Interface.createSensitive(kernel, str(nam), str(det), shared)


def Action(kernel, nam, shared=False):
  return Interface.createAction(kernel, str(nam), shared)


def Filter(kernel, nam, shared=False):
  return Interface.createFilter(kernel, str(nam), shared)


def PhaseAction(kernel, nam, shared=False):
  return Interface.createPhaseAction(kernel, str(nam), shared)


def RunAction(kernel, nam, shared=False):
  return Interface.createRunAction(kernel, str(nam), shared)


def EventAction(kernel, nam, shared=False):
  return Interface.createEventAction(kernel, str(nam), shared)


def GeneratorAction(kernel, nam, shared=False):
  return Interface.createGeneratorAction(kernel, str(nam), shared)


def TrackingAction(kernel, nam, shared=False):
  return Interface.createTrackingAction(kernel, str(nam), shared)


def SteppingAction(kernel, nam, shared=False):
  return Interface.createSteppingAction(kernel, str(nam), shared)


def StackingAction(kernel, nam, shared=False):
  return Interface.createStackingAction(kernel, str(nam), shared)


def DetectorConstruction(kernel, nam):
  return Interface.createDetectorConstruction(kernel, str(nam))


def PhysicsList(kernel, nam):
  return Interface.createPhysicsList(kernel, str(nam))


Marko Petric's avatar
Marko Petric committed
  return Interface.createUserInitialization(kernel, str(nam))


Marko Petric's avatar
Marko Petric committed
  return Interface.createSensDetSequence(kernel, str(nam))


def _setup(obj):
Marko Petric's avatar
Marko Petric committed
  def _adopt(self, action):
    self.__adopt(action.get())
Marko Petric's avatar
Marko Petric committed
  _import_class('Sim', obj)
  o = getattr(current, obj)
  setattr(o, '__adopt', getattr(o, 'adopt'))
  setattr(o, 'adopt', _adopt)
  setattr(o, 'add', _adopt)

Marko Petric's avatar
Marko Petric committed
  def _adopt(self, action):
    self.__adopt(action.get(), action.callback())
Marko Petric's avatar
Marko Petric committed
  _import_class('Sim', obj)
  o = getattr(current, obj)
  setattr(o, '__adopt', getattr(o, 'add'))
  setattr(o, 'add', _adopt)

_setup('Geant4RunActionSequence')
_setup('Geant4EventActionSequence')
_setup('Geant4GeneratorActionSequence')
_setup('Geant4TrackingActionSequence')
_setup('Geant4SteppingActionSequence')
_setup('Geant4StackingActionSequence')
_setup('Geant4PhysicsListActionSequence')
_setup('Geant4SensDetActionSequence')
_setup('Geant4DetectorConstructionSequence')
_setup('Geant4UserInitializationSequence')
_setup('Geant4Sensitive')
_setup('Geant4ParticleHandler')
Marko Petric's avatar
Marko Petric committed
_import_class('Sim', 'Geant4Vertex')
_import_class('Sim', 'Geant4Particle')
_import_class('Sim', 'Geant4VertexVector')
_import_class('Sim', 'Geant4ParticleVector')
_import_class('Sim', 'Geant4Action')
_import_class('Sim', 'Geant4Filter')
_import_class('Sim', 'Geant4RunAction')
_import_class('Sim', 'Geant4TrackingAction')
_import_class('Sim', 'Geant4StackingAction')
_import_class('Sim', 'Geant4PhaseAction')
_import_class('Sim', 'Geant4UserParticleHandler')
_import_class('Sim', 'Geant4UserInitialization')
_import_class('Sim', 'Geant4DetectorConstruction')
_import_class('Sim', 'Geant4GeneratorWrapper')
_import_class('Sim', 'Geant4Random')
_import_class('CLHEP', 'HepRandom')
_import_class('CLHEP', 'HepRandomEngine')
def _get(self, name):
  a = Interface.toAction(self)
Marko Petric's avatar
Marko Petric committed
  ret = Interface.getProperty(a, name)
  if ret.status > 0:
    return ret.data
Marko Petric's avatar
Marko Petric committed
  elif hasattr(self.action, name):
    return getattr(self.action, name)
  elif hasattr(a, name):
    return getattr(a, name)
  msg = 'Geant4Action::GetProperty [Unhandled]: Cannot access property ' + a.name() + '.' + name
Ercan Pilicer's avatar
Ercan Pilicer committed
  raise KeyError(msg)
def _set(self, name, value):
  a = Interface.toAction(self)
Marko Petric's avatar
Marko Petric committed
  if isinstance(value, list):
    value = [str(x) for x in value]
  if isinstance(value, dict):
    tempDict = {}
    for key, val in value.items():
      if isinstance(val, six.string_types):
        val = str(val)
      tempDict[str(key)] = val
    value = tempDict
  if Interface.setProperty(a, str(name), str(value)):
Marko Petric's avatar
Marko Petric committed
  msg = 'Geant4Action::SetProperty [Unhandled]: Cannot set ' + a.name() + '.' + name + ' = ' + str(value)
Ercan Pilicer's avatar
Ercan Pilicer committed
  raise KeyError(msg)
def _props(obj):
Marko Petric's avatar
Marko Petric committed
  _import_class('Sim', obj)
  cl = getattr(current, obj)
  cl.__getattr__ = _get
  cl.__setattr__ = _set

_props('FilterHandle')
_props('ActionHandle')
_props('RunActionHandle')
_props('EventActionHandle')
_props('GeneratorActionHandle')
_props('PhysicsListHandle')
_props('TrackingActionHandle')
_props('SteppingActionHandle')
_props('StackingActionHandle')
_props('SensitiveHandle')
_props('Geant4ParticleHandler')
_props('Geant4UserParticleHandler')

_props('GeneratorActionSequenceHandle')
_props('RunActionSequenceHandle')
_props('EventActionSequenceHandle')
_props('TrackingActionSequenceHandle')
_props('SteppingActionSequenceHandle')
_props('StackingActionSequenceHandle')
_props('DetectorConstructionSequenceHandle')
_props('PhysicsListActionSequenceHandle')
_props('SensDetActionSequenceHandle')
_props('UserInitializationSequenceHandle')

_props('Geant4PhysicsListActionSequence')
Marko Petric's avatar
Marko Petric committed
class Geant4:
  """
Marko Petric's avatar
Marko Petric committed
  Helper object to perform stuff, which occurs very often.
  I am sick of typing the same over and over again.
  Hence, I grouped often used python fragments to this small
  class to re-usage.
Marko Petric's avatar
Marko Petric committed
  \author  M.Frank
  \version 1.0
Marko Petric's avatar
Marko Petric committed
  """
  def __init__(self, kernel=None,
               calo='Geant4CalorimeterAction',
               tracker='Geant4SimpleTrackerAction'):
    kernel.UI = "UI"
    kernel.printProperties()
Markus Frank's avatar
Markus Frank committed
    if kernel is None:
Markus Frank's avatar
Markus Frank committed
    self.description = self._kernel.detectorDescription()
    self.sensitive_types['tracker'] = tracker
    self.sensitive_types['calorimeter'] = calo
    self.sensitive_types['escape_counter'] = 'Geant4EscapeCounter'
Marko Petric's avatar
Marko Petric committed
    """
    Access the worker kernel object.
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
    """
    return self._kernel.worker()
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Access the master kernel object.
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
    return self._kernel
Marko Petric's avatar
Marko Petric committed

  def setupUI(self, typ='csh', vis=False, ui=True, macro=None):
Marko Petric's avatar
Marko Petric committed
    """
    Configure the Geant4 command executive

    \author  M.Frank
    """
Marko Petric's avatar
Marko Petric committed
    ui_action = Action(self.master(), "Geant4UIManager/UI")
    if vis:
      ui_action.HaveVIS = True
    else:
      ui_action.HaveVIS = False
    if ui:
      ui_action.HaveUI = True
    else:
      ui_action.HaveUI = False
    ui_action.SessionType = typ
    if macro:
      ui_action.SetupUI = macro
    self.master().registerGlobalAction(ui_action)
Marko Petric's avatar
Marko Petric committed
  def setupCshUI(self, typ='csh', vis=False, ui=True, macro=None):
Marko Petric's avatar
Marko Petric committed
    """
    Configure the Geant4 command executive with a csh like command prompt
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
    """
    return self.setupUI(typ='csh', vis=vis, ui=ui, macro=macro)
  def addUserInitialization(self, worker, worker_args=None, master=None, master_args=None):
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Configure Geant4 user initialization for optionasl multi-threading mode
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
    init_seq = self.master().userInitialization(True)
Marko Petric's avatar
Marko Petric committed
    init_action = UserInitialization(self.master(), 'Geant4PythonInitialization/PyG4Init')
    #
    if worker:
      init_action.setWorkerSetup(worker, worker_args)
    else:
Ercan Pilicer's avatar
Ercan Pilicer committed
      raise RuntimeError('Invalid argument for Geant4 worker initialization')
Marko Petric's avatar
Marko Petric committed
      init_action.setMasterSetup(master, master_args)
Marko Petric's avatar
Marko Petric committed
    return init_seq, init_action
  def detectorConstruction(self):
    seq = self.master().detectorConstruction(True)
    return seq
Marko Petric's avatar
Marko Petric committed

  def addDetectorConstruction(self, name_type,
                              field=None, field_args=None,
                              geometry=None, geometry_args=None,
                              sensitives=None, sensitives_args=None,
                              allow_threads=False):
Marko Petric's avatar
Marko Petric committed
    """
    Configure Geant4 user initialization for optionasl multi-threading mode

    \author  M.Frank
    """
    init_seq = self.master().detectorConstruction(True)
Marko Petric's avatar
Marko Petric committed
    init_action = DetectorConstruction(self.master(), name_type)
Marko Petric's avatar
Marko Petric committed
      init_action.setConstructGeo(geometry, geometry_args)
Marko Petric's avatar
Marko Petric committed
      init_action.setConstructField(field, field_args)
Marko Petric's avatar
Marko Petric committed
      init_action.setConstructSensitives(sensitives, sensitives_args)
Marko Petric's avatar
Marko Petric committed
      last_action = DetectorConstruction(self.master(), "Geant4PythonDetectorConstructionLast/LastDetectorAction")
Marko Petric's avatar
Marko Petric committed
    return init_seq, init_action
Marko Petric's avatar
Marko Petric committed
  def addPhaseAction(self, phase_name, factory_specification, ui=True, instance=None):
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Add a new phase action to an arbitrary step.
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
    if instance is None:
      instance = self.kernel()
Marko Petric's avatar
Marko Petric committed
    action = PhaseAction(instance, factory_specification)
Marko Petric's avatar
Marko Petric committed
    if ui:
      action.enableUI()
    return action

  def addConfig(self, factory_specification):
Marko Petric's avatar
Marko Petric committed
    """
    Add a new phase action to the 'configure' step.
    Called at the beginning of Geant4Exec::configure.
    The factory specification is the typical string "<factory_name>/<instance name>".
    If no instance name is specified it defaults to the factory name.

    \author  M.Frank
    """
Marko Petric's avatar
Marko Petric committed
    return self.addPhaseAction('configure', factory_specification, instance=self.master())
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Add a new phase action to the 'initialize' step.
    Called at the beginning of Geant4Exec::initialize.
    The factory specification is the typical string "<factory_name>/<instance name>".
    If no instance name is specified it defaults to the factory name.
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    return self.addPhaseAction('initialize', factory_specification)
Marko Petric's avatar
Marko Petric committed
    """
    Add a new phase action to the 'start' step.
    Called at the beginning of Geant4Exec::run.
    The factory specification is the typical string "<factory_name>/<instance name>".
    If no instance name is specified it defaults to the factory name.

    \author  M.Frank
    """
Marko Petric's avatar
Marko Petric committed
    return self.addPhaseAction('start', factory_specification)
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Add a new phase action to the 'stop' step.
    Called at the end of Geant4Exec::run.
    The factory specification is the typical string "<factory_name>/<instance name>".
    If no instance name is specified it defaults to the factory name.
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    return self.addPhaseAction('stop', factory_specification)
Markus Frank's avatar
Markus Frank committed
  def execute(self):
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Execute the Geant 4 program with all steps.
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
    self.kernel().configure()
    self.kernel().initialize()
    self.kernel().run()
    self.kernel().terminate()
Markus Frank's avatar
Markus Frank committed
    return self

  def printDetectors(self):
    logger.info('+++  List of sensitive detectors:')
Markus Frank's avatar
Markus Frank committed
    for i in self.description.detectors():
Marko Petric's avatar
Marko Petric committed
      o = DetElement(i.second.ptr())  # noqa: F405
      sd = self.description.sensitiveDetector(str(o.name()))
      if sd.isValid():
Markus Frank's avatar
Markus Frank committed
        sdtyp = 'Unknown'
        if typ in self.sensitive_types:
Markus Frank's avatar
Markus Frank committed
          sdtyp = self.sensitive_types[typ]
Marko Petric's avatar
Marko Petric committed
        logger.info('+++  %-32s type:%-12s  --> Sensitive type: %s', o.name(), typ, sdtyp)
Marko Petric's avatar
Marko Petric committed
  def setupDetector(self, name, action, collections=None):
    # fg: allow the action to be a tuple with parameter dictionary
Marko Petric's avatar
Marko Petric committed
    if isinstance(action, tuple) or isinstance(action, list):
      sensitive_type = action[0]
      parameterDict = action[1]
    else:
      sensitive_type = action
Marko Petric's avatar
Marko Petric committed
    seq = SensitiveSequence(self.kernel(), 'Geant4SensDetActionSequence/' + name)
    acts = []
    if collections is None:
      sd = self.description.sensitiveDetector(str(name))
      ro = sd.readout()
      collections = ro.collectionNames()
Marko Petric's avatar
Marko Petric committed
      if len(collections) == 0:
        act = SensitiveAction(self.kernel(), sensitive_type + '/' + name + 'Handler', name)
        for parameter, value in six.iteritems(parameterDict):
Marko Petric's avatar
Marko Petric committed
          setattr(act, parameter, value)
        acts.append(act)

    # Work down the collections if present
    if collections is not None:
      for coll in collections:
        params = {}
Marko Petric's avatar
Marko Petric committed
        if isinstance(coll, tuple) or isinstance(coll, list):
          if len(coll) > 2:
            coll_nam = coll[0]
            sensitive_type = coll[1]
            params = coll[2]
Marko Petric's avatar
Marko Petric committed
          elif len(coll) > 1:
            coll_nam = coll[0]
            sensitive_type = coll[1]
          else:
            coll_nam = coll[0]
        else:
          coll_nam = coll
Marko Petric's avatar
Marko Petric committed
        act = SensitiveAction(self.kernel(), sensitive_type + '/' + coll_nam + 'Handler', name)
        act.CollectionName = coll_nam
        for parameter, value in six.iteritems(params):
Marko Petric's avatar
Marko Petric committed
          setattr(act, parameter, value)
        acts.append(act)

    for act in acts:
      act.enableUI()
      seq.add(act)
Marko Petric's avatar
Marko Petric committed
    if len(acts) > 1:
      return (seq, acts)
    return (seq, acts[0])
Marko Petric's avatar
Marko Petric committed
  def setupCalorimeter(self, name, type=None, collections=None):
Marko Petric's avatar
Marko Petric committed
    self.description.sensitiveDetector(str(name))
Marko Petric's avatar
Marko Petric committed
    # sd.setType('calorimeter')
    if type is None:
      type = self.sensitive_types['calorimeter']
    return self.setupDetector(name, type, collections)
Marko Petric's avatar
Marko Petric committed
  def setupTracker(self, name, type=None, collections=None):
Marko Petric's avatar
Marko Petric committed
    self.description.sensitiveDetector(str(name))
Marko Petric's avatar
Marko Petric committed
    # sd.setType('tracker')
    if type is None:
      type = self.sensitive_types['tracker']
    return self.setupDetector(name, type, collections)
  def _private_setupField(self, field, stepper, equation, prt):
Marko Petric's avatar
Marko Petric committed
    field.stepper = stepper
    field.equation = equation
    field.eps_min = 5e-05 * g4units.mm
    field.eps_max = 0.001 * g4units.mm
    field.min_chord_step = 0.01 * g4units.mm
    field.delta_chord = 0.25 * g4units.mm
    field.delta_intersection = 0.001 * g4units.mm
    field.delta_one_step = 0.01 * g4units.mm
    field.largest_step = 1000 * g4units.m
    if prt:
Marko Petric's avatar
Marko Petric committed
      logger.info('+++++> %s %s %s %s ', field.name, '-> stepper  = ', str(field.stepper), '')
      logger.info('+++++> %s %s %s %s ', field.name, '-> equation = ', str(field.equation), '')
      logger.info('+++++> %s %s %s %s ', field.name, '-> eps_min  = ', str(field.eps_min), '[mm]')
      logger.info('+++++> %s %s %s %s ', field.name, '-> eps_max  = ', str(field.eps_max), '[mm]')
      logger.info('+++++> %s %s %s %s ', field.name, '-> delta_chord        = ', str(field.delta_chord), '[mm]')
      logger.info('+++++> %s %s %s %s ', field.name, '-> min_chord_step     = ', str(field.min_chord_step), '[mm]')
      logger.info('+++++> %s %s %s %s ', field.name, '-> delta_one_step     = ', str(field.delta_one_step), '[mm]')
      logger.info('+++++> %s %s %s %s ', field.name, '-> delta_intersection = ', str(field.delta_intersection), '[mm]')
      logger.info('+++++> %s %s %s %s ', field.name, '-> largest_step       = ', str(field.largest_step), '[mm]')
Marko Petric's avatar
Marko Petric committed
  def setupTrackingFieldMT(self, name='MagFieldTrackingSetup',
                           stepper='ClassicalRK4', equation='Mag_UsualEqRhs', prt=False):
Marko Petric's avatar
Marko Petric committed
    seq, fld = self.addDetectorConstruction("Geant4FieldTrackingConstruction/" + name)
    self._private_setupField(fld, stepper, equation, prt)
Marko Petric's avatar
Marko Petric committed
    return (seq, fld)
Marko Petric's avatar
Marko Petric committed
  def setupTrackingField(self, name='MagFieldTrackingSetup',
                         stepper='ClassicalRK4', equation='Mag_UsualEqRhs', prt=False):
Marko Petric's avatar
Marko Petric committed
    field = self.addConfig('Geant4FieldTrackingSetupAction/' + name)
    self._private_setupField(field, stepper, equation, prt)
    return field
  def setupPhysics(self, name):
    phys.extends = name
Marko Petric's avatar
Marko Petric committed
    phys.decays = True
    phys.enableUI()
    phys.dump()
    return phys
  def addPhysics(self, name):
    phys = self.master().physicsList()
Marko Petric's avatar
Marko Petric committed
    opt = PhysicsList(self.master(), name)
    opt.enableUI()
    phys.adopt(opt)
    return opt
Marko Petric's avatar
Marko Petric committed
  def setupGun(self, name, particle, energy, isotrop=True,
               multiplicity=1, position=(0.0, 0.0, 0.0), register=True, **args):
Marko Petric's avatar
Marko Petric committed
    gun = GeneratorAction(self.kernel(), "Geant4ParticleGun/" + name, True)
    for i in args.items():
Marko Petric's avatar
Marko Petric committed
      setattr(gun, i[0], i[1])
    gun.energy = energy
    gun.particle = particle
    gun.multiplicity = multiplicity
Marko Petric's avatar
Marko Petric committed
    gun.position = position
    gun.isotrop = isotrop
    if register:
      self.kernel().generatorAction().add(gun)
    return gun

Marko Petric's avatar
Marko Petric committed
  def setupROOTOutput(self, name, output, mc_truth=True):
Marko Petric's avatar
Marko Petric committed
    """
    Configure ROOT output for the simulated events

    \author  M.Frank
    """
Marko Petric's avatar
Marko Petric committed
    evt_root = EventAction(self.kernel(), 'Geant4Output2ROOT/' + name, True)
Markus Frank's avatar
Markus Frank committed
    evt_root.HandleMCTruth = mc_truth
    evt_root.Control = True
    if not output.endswith('.root'):
      output = output + '.root'
    evt_root.Output = output
    evt_root.enableUI()
    self.kernel().eventAction().add(evt_root)
    return evt_root
Marko Petric's avatar
Marko Petric committed
  def setupLCIOOutput(self, name, output):
Marko Petric's avatar
Marko Petric committed
    """
    Configure LCIO output for the simulated events

    \author  M.Frank
    """
Marko Petric's avatar
Marko Petric committed
    evt_lcio = EventAction(self.kernel(), 'Geant4Output2LCIO/' + name, True)
    evt_lcio.Control = True
Marko Petric's avatar
Marko Petric committed
    evt_lcio.Output = output
    evt_lcio.enableUI()
    self.kernel().eventAction().add(evt_lcio)
    return evt_lcio
  def setupEDM4hepOutput(self, name, output):
    """Configure EDM4hep root output for the simulated events."""
    evt_edm4hep = EventAction(self.kernel(), 'Geant4Output2EDM4hep/' + name, True)
    evt_edm4hep.Control = True
    evt_edm4hep.Output = output
    evt_edm4hep.enableUI()
    self.kernel().eventAction().add(evt_edm4hep)
    return evt_edm4hep

  def buildInputStage(self, generator_input_modules, output_level=None, have_mctruth=True):
Marko Petric's avatar
Marko Petric committed
    """
Marko Petric's avatar
Marko Petric committed
    Generic build of the input stage with multiple input modules.
Marko Petric's avatar
Marko Petric committed
    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
Marko Petric's avatar
Marko Petric committed
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
    # Register Generation initialization action
Marko Petric's avatar
Marko Petric committed
    gen = GeneratorAction(self.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
Marko Petric's avatar
Marko Petric committed
    gen = GeneratorAction(self.kernel(), "Geant4InteractionMerger/InteractionMerger")
    gen.enableUI()
    if output_level is not None:
      gen.OutputLevel = output_level
    ga.adopt(gen)

    # Finally generate Geant4 primaries
Marko Petric's avatar
Marko Petric committed
      gen = GeneratorAction(self.kernel(), "Geant4PrimaryHandler/PrimaryHandler")
      gen.RejectPDGs = "{1,2,3,4,5,6,21,23,24}"
      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
    """
Marko Petric's avatar
Marko Petric committed
    Execute the main Geant4 action
    \author  M.Frank
Marko Petric's avatar
Marko Petric committed
    """
    from ROOT import PyDDG4
    PyDDG4.run(self.master().get())
    return self