track_reconstruction_with_overlay.py

from Gaudi.Configuration import *

from Configurables import ApplicationMgr, EICDataSvc, PodioOutput, GeoSvc
from GaudiKernel import SystemOfUnits as units

detector_name = "topside"
if "DETECTOR_CONFIG" in os.environ :
    detector_name = str(os.environ["DETECTOR_CONFIG"])

detector_path = ""
if "DETECTOR_PATH" in os.environ :
    detector_path = str(os.environ["DETECTOR_PATH"])

detector_version = 'default'
if "DETECTOR_VERSION" in os.environ:
    env_version = str(os.environ["DETECTOR_VERSION"])

# todo add checks
input_sim_file  = str(os.environ["JUGGLER_SIM_FILE"])
input_bkg_file  = str(os.environ["JUGGLER_BKG_FILE"])
output_rec_file = str(os.environ["JUGGLER_REC_FILE"])
n_events = str(os.environ["JUGGLER_N_EVENTS"])

## note: old version of material map is called material-maps.XXX, new version is materials-map.XXX
##       these names are somewhat inconsistent, and should probably all be renamed to 'material-map.XXX' 
##       FIXME
geo_service = GeoSvc("GeoSvc",
                     detectors=["{}/{}.xml".format(detector_path,detector_name)],
                     materials="calibrations/materials-map.cbor",
                     OutputLevel=WARNING)

podioevent   = EICDataSvc("EventDataSvc", inputs=[input_sim_file], OutputLevel=WARNING)

from Configurables import PodioInput

from Configurables import Jug__Digi__SimTrackerHitsCollector as SimTrackerHitsCollector
from Configurables import Jug__Digi__SiliconTrackerDigi as TrackerDigi

from Configurables import Jug__Reco__TrackerHitReconstruction as TrackerHitReconstruction
from Configurables import Jug__Reco__TrackingHitsCollector2 as TrackingHitsCollector

from Configurables import Jug__Reco__TrackerSourceLinker as TrackerSourceLinker
from Configurables import Jug__Reco__TrackParamTruthInit as TrackParamTruthInit
from Configurables import Jug__Reco__TrackParamClusterInit as TrackParamClusterInit
from Configurables import Jug__Reco__TrackParamVertexClusterInit as TrackParamVertexClusterInit

from Configurables import Jug__Reco__TrackFindingAlgorithm as TrackFindingAlgorithm
from Configurables import Jug__Reco__ParticlesFromTrackFit as ParticlesFromTrackFit
from Configurables import Jug__Reco__TrackProjector as TrackProjector

from Configurables import Jug__Reco__SimpleClustering as SimpleClustering

algorithms = [ ]

tracker_endcap_collections = [
    'TrackerEndcapHits1',
    'TrackerEndcapHits2',
    'TrackerEndcapHits3',
    'TrackerEndcapHits4',
    'TrackerEndcapHits5',
    'TrackerEndcapHits6'
]
tracker_barrel_collections = [
    'TrackerBarrelHits'
]
vertex_barrel_collections = [
    'VertexBarrelHits'
]
gem_endcap_collections = [
    'GEMTrackerEndcapHits1',
    'GEMTrackerEndcapHits2',
    'GEMTrackerEndcapHits3'
]
vertex_endcap_collections = [
    'VertexEndcapHits'
]
mpgd_barrel_collections = [
    'MPGDTrackerBarrelHits1',
    'MPGDTrackerBarrelHits2'
]

input_collections = ['MCParticles'] + tracker_endcap_collections + tracker_barrel_collections + vertex_barrel_collections + gem_endcap_collections

input_collections += mpgd_barrel_collections

podioinput = PodioInput("PodioReader",
                        collections=input_collections)
algorithms.append( podioinput )

# use the pileuptool to specify the number of pileup
from Configurables import ConstPileUp
pileuptool = ConstPileUp("MyPileupTool", numPileUpEvents=10)

# edm data from simulation: hits and positioned hits
from Configurables import PileupHitMergeTool_edm4hep__SimTrackerHitCollection_ as PileupTrackHitMergeTool
trackhitsmergetool = PileupTrackHitMergeTool("MyTrackHitMergeTool")
# branchnames for the pileup
trackhitsmergetool.pileupHitsBranch = "hits"
# branchnames for the signal
trackhitsmergetool.signalHits = "hits"
# branchnames for the output
trackhitsmergetool.mergedHits = "overlaidTrackHits"

# algorithm for the overlay
from Configurables import PileupOverlayAlg
overlay = PileupOverlayAlg()
overlay.pileupFilenames = [input_bkg_file]
overlay.randomizePileup = False
overlay.mergeTools = ["PileupTrackHitMergeTool/MyTrackHitMergeTool"]
overlay.PileUpTool = pileuptool
overlay.OutputLevel = DEBUG

trk_b_coll = SimTrackerHitsCollector("trk_b_coll",
        inputSimTrackerHits = tracker_barrel_collections,
        outputSimTrackerHits = "TrackerBarrelAllHits")
algorithms.append( trk_b_coll )
trk_b_digi = TrackerDigi("trk_b_digi",
        inputHitCollection = trk_b_coll.outputSimTrackerHits,
        outputHitCollection = "TrackerBarrelRawHits",
        timeResolution=8)
algorithms.append( trk_b_digi )

vtx_b_coll = SimTrackerHitsCollector("vtx_b_coll",
        inputSimTrackerHits = vertex_barrel_collections,
        outputSimTrackerHits = "VertexBarrelAllHits")
algorithms.append( vtx_b_coll )
vtx_b_digi = TrackerDigi("vtx_b_digi",
        inputHitCollection = vtx_b_coll.outputSimTrackerHits,
        outputHitCollection = "VertexBarrelRawHits",
        timeResolution=8)
algorithms.append( vtx_b_digi )

mm_b_coll = SimTrackerHitsCollector("mm_b_coll",
        inputSimTrackerHits = mpgd_barrel_collections,
        outputSimTrackerHits = "MPGDTrackerBarrelAllHits")
algorithms.append( mm_b_coll )
mm_b_digi = TrackerDigi("mm_b_digi",
        inputHitCollection = mm_b_coll.outputSimTrackerHits,
        outputHitCollection = "MPGDTrackerBarrelRawHits",
        timeResolution=8)
algorithms.append( mm_b_digi )

# Tracker and vertex reconstruction
trk_b_reco = TrackerHitReconstruction("trk_b_reco",
        inputHitCollection = trk_b_digi.outputHitCollection,
        outputHitCollection="TrackerBarrelRecHits")
algorithms.append( trk_b_reco )

vtx_b_reco = TrackerHitReconstruction("vtx_b_reco",
        inputHitCollection = vtx_b_digi.outputHitCollection,
        outputHitCollection="VertexBarrelRecHits")
algorithms.append( vtx_b_reco )

mm_b_reco = TrackerHitReconstruction("mm_b_reco",
        inputHitCollection = mm_b_digi.outputHitCollection,
        outputHitCollection="MPGDTrackerBarrelRecHits")
algorithms.append( mm_b_reco )

input_tracking_hits = [
    str(trk_b_reco.outputHitCollection),
    str(vtx_b_reco.outputHitCollection),
]
input_tracking_hits.append(str(mm_b_reco.outputHitCollection))

trk_hit_col = TrackingHitsCollector("trk_hit_col",
        inputTrackingHits=input_tracking_hits,
        trackingHits="trackingHits")
algorithms.append( trk_hit_col )

# Hit Source linker 
sourcelinker = TrackerSourceLinker("sourcelinker",
        inputHitCollection=trk_hit_col.trackingHits,
        outputSourceLinks="TrackSourceLinks",
        outputMeasurements="TrackMeasurements")
algorithms.append( sourcelinker )

## Track param init
truth_trk_init = TrackParamTruthInit("truth_trk_init",
        inputMCParticles="MCParticles",
        outputInitialTrackParameters="InitTrackParams")
algorithms.append( truth_trk_init )

# Tracking algorithms
trk_find_alg = TrackFindingAlgorithm("trk_find_alg",
        inputSourceLinks = sourcelinker.outputSourceLinks,
        inputMeasurements = sourcelinker.outputMeasurements,
        inputInitialTrackParameters = "InitTrackParams",
        outputTrajectories = "trajectories")
algorithms.append( trk_find_alg )

parts_from_fit = ParticlesFromTrackFit("parts_from_fit",
        inputTrajectories = "trajectories",
        outputParticles = "ReconstructedParticles",
        outputTrackParameters = "outputTrackParameters")
algorithms.append( parts_from_fit )

trk_proj = TrackProjector("trajs_from_fit",
        inputTrajectories = trk_find_alg.outputTrajectories,
        outputTrackSegments = "outputTrackSegments")
algorithms.append( trk_proj )

out = PodioOutput("out", filename=output_rec_file)
out.outputCommands = ["keep *", 
        "drop BarrelTrackSourceLinks", 
        "drop InitTrackParams",
        "drop trajectories",
        "drop outputSourceLinks",
        "drop outputInitialTrackParameters",
        "keep MCParticles"
        ]
algorithms.append(out)

ApplicationMgr(
    TopAlg = algorithms,
    EvtSel = 'NONE',
    EvtMax   = n_events,
    ExtSvc = [podioevent,geo_service],
    OutputLevel=WARNING
 )