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