Skip to content
Snippets Groups Projects
track_reconstruction_with_overlay.py 7.82 KiB
Newer Older
  • Learn to ignore specific revisions
  • 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
     )