Skip to content
Snippets Groups Projects
tracker_reconstruction.py 9.57 KiB
Newer Older
from Gaudi.Configuration import *

from GaudiKernel.DataObjectHandleBase import DataObjectHandleBase
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"])

# 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"])
Whitney Armstrong's avatar
Whitney Armstrong committed

detector_path = detector_name
if "DETECTOR_PATH" in os.environ :
    detector_path = str(os.environ["DETECTOR_PATH"])
Sylvester Joosten's avatar
Sylvester Joosten committed
geo_service  = GeoSvc("GeoSvc",
        detectors=["{}/{}.xml".format(detector_path, detector_name)])
podioevent   = EICDataSvc("EventDataSvc", inputs=[input_sim_file], OutputLevel=DEBUG)

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 UFSDTrackerDigi
from Configurables import Jug__Digi__EMCalorimeterDigi as EMCalorimeterDigi

from Configurables import Jug__Base__MC2DummyParticle as MC2DummyParticle

from Configurables import Jug__Reco__TrackerHitReconstruction as TrackerHitReconstruction

from Configurables import Jug__Reco__TrackerSourceLinker as TrackerSourceLinker
Whitney Armstrong's avatar
Whitney Armstrong committed
from Configurables import Jug__Reco__Tracker2SourceLinker as Tracker2SourceLinker
#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__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



podioinput = PodioInput("PodioReader", 
Whitney Armstrong's avatar
Whitney Armstrong committed
                        collections=["mcparticles","SiTrackerEndcapHits","SiTrackerBarrelHits","EcalBarrelHits"])#, OutputLevel=DEBUG)
#"SiVertexBarrelHits",
dummy = MC2DummyParticle("MC2Dummy",
        inputCollection="mcparticles",
        outputCollection="DummyReconstructedParticles")

## copiers to get around input --> output copy bug. Note the "2" appended to the output collection.
copier = MCCopier("MCCopier", 
        inputCollection="mcparticles", 
        outputCollection="mcparticles2") 
trkcopier = TrkCopier("TrkCopier", 
        inputCollection="SiTrackerBarrelHits", 
        outputCollection="SiTrackerBarrelHits2") 

ecal_digi = EMCalorimeterDigi("ecal_digi", 
        inputHitCollection="EcalBarrelHits", 
        outputHitCollection="RawEcalBarrelHits")

ufsd_digi = UFSDTrackerDigi("ufsd_digi", 
        inputHitCollection="SiTrackerBarrelHits",
        outputHitCollection="SiTrackerBarrelRawHits",
        timeResolution=8)
Whitney Armstrong's avatar
Whitney Armstrong committed
ufsd_digi2 = UFSDTrackerDigi("ufsd_digi2", 
        inputHitCollection="SiTrackerEndcapHits",
        outputHitCollection="SiTrackerEndcapRawHits",
        timeResolution=8)

Whitney Armstrong's avatar
Whitney Armstrong committed
#vtx_digi = UFSDTrackerDigi("vtx_digi", 
#        inputHitCollection="SiVertexBarrelHits",
#        outputHitCollection="SiVertexBarrelRawHits",
#        timeResolution=8)


ecal_reco = EMCalReconstruction("ecal_reco", 
        inputHitCollection="RawEcalBarrelHits", 
        outputHitCollection="RecEcalBarrelHits",
        minModuleEdep=0.0*units.MeV,
        OutputLevel=DEBUG)

simple_cluster = SimpleClustering("simple_cluster", 
        inputHitCollection="RecEcalBarrelHits", 
        outputClusters="SimpleClusters",
        minModuleEdep=1.0*units.MeV,
        maxDistance=50.0*units.cm,
        OutputLevel=DEBUG)

trk_barrel_reco = TrackerHitReconstruction("trk_barrel_reco",
        inputHitCollection="SiTrackerBarrelRawHits",
        outputHitCollection="TrackerBarrelRecHits")

Whitney Armstrong's avatar
Whitney Armstrong committed
trk_endcap_reco = TrackerHitReconstruction("trk_endcap_reco",
        inputHitCollection="SiTrackerEndcapRawHits",
        outputHitCollection="TrackerEndcapRecHits")

#vtx_barrel_reco = TrackerHitReconstruction("vtx_barrel_reco",
#        inputHitCollection = vtx_digi.outputHitCollection,
#        outputHitCollection="VertexBarrelRecHits")

# Source linker 
sourcelinker = TrackerSourceLinker("trk_srclinker",
        inputHitCollection="TrackerBarrelRecHits",
        outputSourceLinks="BarrelTrackSourceLinks",
        OutputLevel=DEBUG)

Whitney Armstrong's avatar
Whitney Armstrong committed
trk_hits_srclnkr = Tracker2SourceLinker("trk_hits_srclnkr",
        TrackerBarrelHits="TrackerBarrelRecHits",
        TrackerEndcapHits="TrackerEndcapRecHits",
        outputMeasurements="lnker2Measurements",
        outputSourceLinks="lnker2Links",
        allTrackerHits="linker2AllHits",
        OutputLevel=DEBUG)

## Track param init
truth_trk_init = TrackParamTruthInit("truth_trk_init",
        inputMCParticles="mcparticles",
        outputInitialTrackParameters="InitTrackParams",
        OutputLevel=DEBUG)

clust_trk_init = TrackParamClusterInit("clust_trk_init",
        inputClusters="SimpleClusters",
        outputInitialTrackParameters="InitTrackParamsFromClusters",
        OutputLevel=DEBUG)

Whitney Armstrong's avatar
Whitney Armstrong committed
#vtxcluster_trk_init = TrackParamVertexClusterInit("vtxcluster_trk_init",
#        inputVertexHits="VertexBarrelRecHits",
#        inputClusters="SimpleClusters",
#        outputInitialTrackParameters="InitTrackParamsFromVtxClusters",
#        maxHitRadius=40.0*units.mm,
#        OutputLevel=DEBUG)

# Tracking algorithms
trk_find_alg = TrackFindingAlgorithm("trk_find_alg",
Whitney Armstrong's avatar
Whitney Armstrong committed
        inputSourceLinks = sourcelinker.outputSourceLinks,
        inputMeasurements = sourcelinker.outputMeasurements,
        inputInitialTrackParameters= "InitTrackParams",#"InitTrackParamsFromClusters", 
        outputTrajectories="trajectories",
        OutputLevel=DEBUG)
parts_from_fit = ParticlesFromTrackFit("parts_from_fit",
        inputTrajectories="trajectories",
        outputParticles="ReconstructedParticles",
        outputTrackParameters="outputTrackParameters",
        OutputLevel=DEBUG)

trk_find_alg1 = TrackFindingAlgorithm("trk_find_alg1",
Whitney Armstrong's avatar
Whitney Armstrong committed
        inputSourceLinks = trk_hits_srclnkr.outputSourceLinks,
        inputMeasurements = trk_hits_srclnkr.outputMeasurements,
        inputInitialTrackParameters= "InitTrackParamsFromClusters", 
        outputTrajectories="trajectories1",
        OutputLevel=DEBUG)
parts_from_fit1 = ParticlesFromTrackFit("parts_from_fit1",
        inputTrajectories="trajectories1",
        outputParticles="ReconstructedParticles1",
        outputTrackParameters="outputTrackParameters1",
        OutputLevel=DEBUG)

trk_find_alg2 = TrackFindingAlgorithm("trk_find_alg2",
Whitney Armstrong's avatar
Whitney Armstrong committed
        inputSourceLinks = trk_hits_srclnkr.outputSourceLinks,
        inputMeasurements = trk_hits_srclnkr.outputMeasurements,
        inputInitialTrackParameters= "InitTrackParams",#"InitTrackParamsFromClusters", 
        #inputInitialTrackParameters= "InitTrackParamsFromVtxClusters", 
        outputTrajectories="trajectories2",
        OutputLevel=DEBUG)
parts_from_fit2 = ParticlesFromTrackFit("parts_from_fit2",
        inputTrajectories="trajectories2",
        outputParticles="ReconstructedParticles2",
        outputTrackParameters="outputTrackParameters2",
        OutputLevel=DEBUG)


#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, DataObjectHandleBase):
#            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"
        ]

ApplicationMgr(
    TopAlg = [podioinput, 
              copier, trkcopier,
Whitney Armstrong's avatar
Whitney Armstrong committed
              ecal_digi, ufsd_digi2,ufsd_digi, #vtx_digi, 
              ecal_reco, 
              simple_cluster,
              trk_barrel_reco, 
Whitney Armstrong's avatar
Whitney Armstrong committed
              trk_endcap_reco, 
              #vtx_barrel_reco, 
              sourcelinker, trk_hits_srclnkr,
              clust_trk_init, 
              truth_trk_init, 
Whitney Armstrong's avatar
Whitney Armstrong committed
              #vtxcluster_trk_init, 
              trk_find_alg, parts_from_fit,
              trk_find_alg1, parts_from_fit1,
              trk_find_alg2, parts_from_fit2,
              out
              ],
    EvtSel = 'NONE',
    EvtMax   = n_events,
    ExtSvc = [podioevent,geo_service],
    OutputLevel=DEBUG
 )