Scripts/Python/vertex_fitting.py

0001 #!/usr/bin/env python3
0002 from pathlib import Path
0003 from typing import Optional
0004
0005 import acts
0006 from acts.examples import (
0007     Sequencer,
0008     ParticleSelector,
0009     TrackParameterSmearing,
0010     TrackParameterSelector,
0011 )
0012 from acts.examples.simulation import addPythia8
0013 from acts.examples.reconstruction import (
0014     addVertexFitting,
0015     VertexFinder,
0016 )
0017
0018 u = acts.UnitConstants
0019
0020
0021 def runVertexFitting(
0022     field,
0023     outputDir: Path,
0024     outputRoot: bool = True,
0025     inputParticlePath: Optional[Path] = None,
0026     inputTrackSummary: Path = None,
0027     vertexFinder: VertexFinder = VertexFinder.Truth,
0028     s=None,
0029 ):
0030     s = s or Sequencer(events=100, numThreads=-1)
0031
0032     logger = acts.logging.getLogger("VertexFittingExample")
0033
0034     rnd = acts.examples.RandomNumbers(seed=42)
0035
0036     inputParticles = "particles_input"
0037     if inputParticlePath is None:
0038         logger.info("Generating particles using Pythia8")
0039         addPythia8(s, rnd)
0040     else:
0041         logger.info("Reading particles from %s", inputParticlePath.resolve())
0042         assert inputParticlePath.exists()
0043         s.addReader(
0044             acts.examples.RootParticleReader(
0045                 level=acts.logging.INFO,
0046                 filePath=str(inputParticlePath.resolve()),
0047                 outputParticles=inputParticles,
0048             )
0049         )
0050
0051     selectedParticles = "particles_selected"
0052     ptclSelector = ParticleSelector(
0053         level=acts.logging.INFO,
0054         inputParticles=inputParticles,
0055         outputParticles=selectedParticles,
0056         removeNeutral=True,
0057         absEtaMax=2.5,
0058         rhoMax=4.0 * u.mm,
0059         ptMin=500 * u.MeV,
0060     )
0061     s.addAlgorithm(ptclSelector)
0062
0063     trackParameters = "fittedTrackParameters"
0064
0065     if inputTrackSummary is None or inputParticlePath is None:
0066         logger.info("Using smeared particles")
0067
0068         trkParamExtractor = acts.examples.ParticleTrackParamExtractor(
0069             level=acts.logging.WARNING,
0070             inputParticles=selectedParticles,
0071             outputTrackParameters="params_particles_input",
0072         )
0073         s.addAlgorithm(trkParamExtractor)
0074
0075         ptclSmearing = TrackParameterSmearing(
0076             level=acts.logging.INFO,
0077             inputTrackParameters="params_particles_input",
0078             outputTrackParameters=trackParameters,
0079             randomNumbers=rnd,
0080         )
0081         s.addAlgorithm(ptclSmearing)
0082         associatedParticles = selectedParticles
0083     else:
0084         logger.info("Reading track summary from %s", inputTrackSummary.resolve())
0085         assert inputTrackSummary.exists()
0086         associatedParticles = "associatedTruthParticles"
0087         trackSummaryReader = acts.examples.RootTrackSummaryReader(
0088             level=acts.logging.VERBOSE,
0089             outputTracks=trackParameters,
0090             outputParticles=associatedParticles,
0091             filePath=str(inputTrackSummary.resolve()),
0092         )
0093         s.addReader(trackSummaryReader)
0094
0095         trackParamSelector = TrackParameterSelector(
0096             level=acts.logging.INFO,
0097             inputTrackParameters=trackSummaryReader.config.outputTracks,
0098             outputTrackParameters="selectedTrackParameters",
0099             absEtaMax=2.5,
0100             loc0Max=4.0 * u.mm,  # rho max
0101             ptMin=500 * u.MeV,
0102         )
0103         s.addAlgorithm(trackParamSelector)
0104         trackParameters = trackParamSelector.config.outputTrackParameters
0105
0106     logger.info("Using vertex finder: %s", vertexFinder.name)
0107
0108     addVertexFitting(
0109         s,
0110         field,
0111         trackParameters=trackParameters,
0112         associatedParticles=associatedParticles,
0113         trajectories=None,
0114         vertexFinder=vertexFinder,
0115         outputDirRoot=outputDir if outputRoot else None,
0116     )
0117
0118     return s
0119
0120
0121 if "__main__" == __name__:
0122     detector = acts.examples.GenericDetector()
0123     trackingGeometry = detector.trackingGeometry()
0124
0125     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0126
0127     inputParticlePath = Path("particles.root")
0128     if not inputParticlePath.exists():
0129         inputParticlePath = None
0130
0131     inputTrackSummary = None
0132     for p in ("tracksummary_fitter.root", "tracksummary_ckf.root"):
0133         p = Path(p)
0134         if p.exists():
0135             inputTrackSummary = p
0136             break
0137
0138     runVertexFitting(
0139         field,
0140         vertexFinder=VertexFinder.Truth,
0141         inputParticlePath=inputParticlePath,
0142         inputTrackSummary=inputTrackSummary,
0143         outputDir=Path.cwd(),
0144     ).run()