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     ParticleSmearing,
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         ptclSmearing = ParticleSmearing(
0069             level=acts.logging.INFO,
0070             inputParticles=selectedParticles,
0071             outputTrackParameters=trackParameters,
0072             randomNumbers=rnd,
0073         )
0074         s.addAlgorithm(ptclSmearing)
0075         associatedParticles = selectedParticles
0076     else:
0077         logger.info("Reading track summary from %s", inputTrackSummary.resolve())
0078         assert inputTrackSummary.exists()
0079         associatedParticles = "associatedTruthParticles"
0080         trackSummaryReader = acts.examples.RootTrackSummaryReader(
0081             level=acts.logging.VERBOSE,
0082             outputTracks=trackParameters,
0083             outputParticles=associatedParticles,
0084             filePath=str(inputTrackSummary.resolve()),
0085         )
0086         s.addReader(trackSummaryReader)
0087
0088         trackParamSelector = TrackParameterSelector(
0089             level=acts.logging.INFO,
0090             inputTrackParameters=trackSummaryReader.config.outputTracks,
0091             outputTrackParameters="selectedTrackParameters",
0092             absEtaMax=2.5,
0093             loc0Max=4.0 * u.mm,  # rho max
0094             ptMin=500 * u.MeV,
0095         )
0096         s.addAlgorithm(trackParamSelector)
0097         trackParameters = trackParamSelector.config.outputTrackParameters
0098
0099     logger.info("Using vertex finder: %s", vertexFinder.name)
0100
0101     addVertexFitting(
0102         s,
0103         field,
0104         trackParameters=trackParameters,
0105         associatedParticles=associatedParticles,
0106         trajectories=None,
0107         vertexFinder=vertexFinder,
0108         outputDirRoot=outputDir if outputRoot else None,
0109     )
0110
0111     return s
0112
0113
0114 if "__main__" == __name__:
0115     detector, trackingGeometry, decorators = acts.examples.GenericDetector.create()
0116
0117     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0118
0119     inputParticlePath = Path("particles.root")
0120     if not inputParticlePath.exists():
0121         inputParticlePath = None
0122
0123     inputTrackSummary = None
0124     for p in ("tracksummary_fitter.root", "tracksummary_ckf.root"):
0125         p = Path(p)
0126         if p.exists():
0127             inputTrackSummary = p
0128             break
0129
0130     runVertexFitting(
0131         field,
0132         vertexFinder=VertexFinder.Truth,
0133         inputParticlePath=inputParticlePath,
0134         inputTrackSummary=inputTrackSummary,
0135         outputDir=Path.cwd(),
0136     ).run()