Skip to content
Snippets Groups Projects
track_reconstruction.py 8.21 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 "JUGGLER_DETECTOR" in os.environ :
      detector_name = str(os.environ["JUGGLER_DETECTOR"])
    
    detector_path = ""
    if "DETECTOR_PATH" in os.environ :
      detector_path = str(os.environ["DETECTOR_PATH"])
    
    # todo add checks
    input_sim_file  = str(os.environ["JUGGLER_SIM_FILE"])
    output_rec_file = str(os.environ["JUGGLER_REC_FILE"])
    n_events = str(os.environ["JUGGLER_N_EVENTS"])
    
    geo_service  = GeoSvc("GeoSvc", detectors=["{}/{}.xml".format(detector_path,detector_name)], OutputLevel=WARNING)
    podioevent   = EICDataSvc("EventDataSvc", inputs=[input_sim_file], OutputLevel=WARNING)
    
    from Configurables import PodioInput
    from Configurables import Jug__Base__InputCopier_dd4pod__Geant4ParticleCollection_dd4pod__Geant4ParticleCollection_ as MCCopier
    from Configurables import Jug__Base__InputCopier_dd4pod__CalorimeterHitCollection_dd4pod__CalorimeterHitCollection_ as CalCopier
    from Configurables import Jug__Base__InputCopier_dd4pod__TrackerHitCollection_dd4pod__TrackerHitCollection_ as TrkCopier
    
    from Configurables import Jug__Digi__ExampleCaloDigi as ExampleCaloDigi
    from Configurables import Jug__Digi__UFSDTrackerDigi as TrackerDigi
    from Configurables import Jug__Digi__EMCalorimeterDigi as EMCalorimeterDigi
    
    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__TrackerSourcesLinker as TrackerSourcesLinker
    #from Configurables import Jug__Reco__TrackingHitsSourceLinker as TrackingHitsSourceLinker
    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__ConformalXYPeakProtoTracks as ConformalXYPeakProtoTracks
    
    from Configurables import Jug__Reco__TrackFindingAlgorithm as TrackFindingAlgorithm
    from Configurables import Jug__Reco__ParticlesFromTrackFit as ParticlesFromTrackFit
    from Configurables import Jug__Reco__EMCalReconstruction as EMCalReconstruction
    
    from Configurables import Jug__Reco__SimpleClustering as SimpleClustering
    
    
    algorithms = [ ]
    
    podioinput = PodioInput("PodioReader", 
    
                            collections=["mcparticles","TrackerEndcapHits","TrackerBarrelHits","VertexBarrelHits","VertexEndcapHits","GEMTrackerEndcapHits"])
    
    algorithms.append( podioinput )
    
    ## copiers to get around input --> output copy bug. Note the "2" appended to the output collection.
    copier = MCCopier("MCCopier", 
            inputCollection="mcparticles", 
            outputCollection="mcparticles2") 
    algorithms.append( copier )
    
    trkcopier = TrkCopier("TrkCopier", 
            inputCollection="TrackerBarrelHits", 
            outputCollection="TrackerBarrelHits2") 
    algorithms.append( trkcopier )
    
    trk_b_digi = TrackerDigi("trk_b_digi", 
            inputHitCollection="TrackerBarrelHits",
            outputHitCollection="TrackerBarrelRawHits",
            timeResolution=8)
    algorithms.append( trk_b_digi )
    trk_ec_digi = TrackerDigi("trk_ec_digi", 
            inputHitCollection="TrackerEndcapHits",
            outputHitCollection="TrackerEndcapRawHits",
            timeResolution=8)
    algorithms.append( trk_ec_digi )
    
    vtx_b_digi = TrackerDigi("vtx_b_digi", 
            inputHitCollection="VertexBarrelHits",
            outputHitCollection="VertexBarrelRawHits",
            timeResolution=8)
    algorithms.append( vtx_b_digi )
    
    vtx_ec_digi = TrackerDigi("vtx_ec_digi", 
            inputHitCollection="VertexEndcapHits",
            outputHitCollection="VertexEndcapRawHits",
            timeResolution=8)
    algorithms.append( vtx_ec_digi )
    
    gem_ec_digi = TrackerDigi("gem_ec_digi",
            inputHitCollection="GEMTrackerEndcapHits",
            outputHitCollection="GEMTrackerEndcapRawHits",
            timeResolution=10)
    algorithms.append(gem_ec_digi)
    
    # Tracker and vertex reconstruction
    trk_b_reco = TrackerHitReconstruction("trk_b_reco",
            inputHitCollection = trk_b_digi.outputHitCollection,
            outputHitCollection="TrackerBarrelRecHits")
    algorithms.append( trk_b_reco )
    
    trk_ec_reco = TrackerHitReconstruction("trk_ec_reco",
            inputHitCollection = trk_ec_digi.outputHitCollection,
            outputHitCollection="TrackerEndcapRecHits")
    algorithms.append( trk_ec_reco )
    
    vtx_b_reco = TrackerHitReconstruction("vtx_b_reco",
            inputHitCollection = vtx_b_digi.outputHitCollection,
            outputHitCollection="VertexBarrelRecHits")
    algorithms.append( vtx_b_reco )
    
    vtx_ec_reco = TrackerHitReconstruction("vtx_ec_reco",
            inputHitCollection = vtx_ec_digi.outputHitCollection,
            outputHitCollection="VertexEndcapRecHits")
    algorithms.append( vtx_ec_reco )
    
    gem_ec_reco = TrackerHitReconstruction("gem_ec_reco",
            inputHitCollection=gem_ec_digi.outputHitCollection,
            outputHitCollection="GEMTrackerEndcapRecHits")
    algorithms.append(gem_ec_reco)
    
    trk_hit_col = TrackingHitsCollector("trk_hit_col",
            inputTrackingHits=[
                str(trk_b_reco.outputHitCollection),
                str(trk_ec_reco.outputHitCollection),
                str(vtx_b_reco.outputHitCollection),
                str(vtx_ec_reco.outputHitCollection),
                str(gem_ec_reco.outputHitCollection) ],
            trackingHits="trackingHits",
            OutputLevel=DEBUG)
    algorithms.append( trk_hit_col )
    
    
    conformal_find = ConformalXYPeakProtoTracks("conformal_find",
            inputTrackerHits=trk_hit_col.trackingHits,
            outputProtoTracks="outputProtoTracks",
            nProtoTracks="nProtoTracks",
            OutputLevel=DEBUG)
    algorithms.append( conformal_find )
    
    
    # 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")
            #OutputLevel=DEBUG)
    algorithms.append( truth_trk_init )
    
    # Tracking algorithms
    trk_find_alg = TrackFindingAlgorithm("trk_find_alg",
            inputSourceLinks = sourcelinker.outputSourceLinks,
            inputMeasurements = sourcelinker.outputMeasurements,
            inputInitialTrackParameters= "InitTrackParams",#"InitTrackParamsFromClusters", 
            outputTrajectories="trajectories")
            #OutputLevel=DEBUG)
    algorithms.append( trk_find_alg )
    
    parts_from_fit = ParticlesFromTrackFit("parts_from_fit",
            inputTrajectories="trajectories",
            outputParticles="ReconstructedParticles",
            outputTrackParameters="outputTrackParameters")
            #OutputLevel=DEBUG)
    algorithms.append( parts_from_fit )
    
    #types = []
    ## this printout is useful to check that the type information is passed to python correctly
    #print("---------------------------------------\n")
    #print("---\n# List of input and output types by class")
    #for configurable in sorted([ PodioInput, EICDataSvc, PodioOutput,
    #                             TrackerHitReconstruction,ExampleCaloDigi, 
    #                             UFSDTrackerDigi, TrackerSourceLinker,
    #                             PodioOutput],
    #                           key=lambda c: c.getType()):
    #    print("\"{}\":".format(configurable.getType()))
    #    props = configurable.getDefaultProperties()
    #    for propname, prop in sorted(props.items()):
    #        print(" prop name: {}".format(propname))
    #        if isinstance(prop, DataHandleBase):
    #            types.append(prop.type())
    #            print("  {}: \"{}\"".format(propname, prop.type()))
    #print("---")
    
    out = PodioOutput("out", filename=output_rec_file)
    out.outputCommands = ["keep *", 
            "drop BarrelTrackSourceLinks", 
            "drop InitTrackParams",
            "drop trajectories",
            "drop outputSourceLinks",
            "drop outputInitialTrackParameters",
            "drop mcparticles"
            ]
    algorithms.append(out)
    
    ApplicationMgr(
        TopAlg = algorithms,
        EvtSel = 'NONE',
        EvtMax   = n_events,
        ExtSvc = [podioevent,geo_service],
        OutputLevel=WARNING
     )