Scripts/Python/pypi_finding_fitting_demo.py

0001 #!/usr/bin/env python3
0002 # This file is part of the ACTS project.
0003 #
0004 # Copyright (C) 2016 CERN for the benefit of the ACTS project
0005 #
0006 # This Source Code Form is subject to the terms of the Mozilla Public
0007 # License, v. 2.0. If a copy of the MPL was not distributed with this
0008 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
0009
0010 from pathlib import Path
0011
0012 import acts
0013 import acts.examples
0014 from acts import UnitConstants as u
0015
0016
0017 def runPypiFindingFittingDemo(
0018     trackingGeometry,
0019     field,
0020     digiConfigFile,
0021     geoSelectionConfigFile,
0022     outputDir,
0023     decorators=[],
0024     s=None,
0025 ):
0026     from acts.examples.simulation import (
0027         addParticleGun,
0028         MomentumConfig,
0029         EtaConfig,
0030         PhiConfig,
0031         ParticleConfig,
0032         addFatras,
0033         addDigitization,
0034     )
0035
0036     s = s or acts.examples.Sequencer(events=1, numThreads=1, logLevel=acts.logging.INFO)
0037     outputDir = Path(outputDir)
0038     rnd = acts.examples.RandomNumbers(seed=42)
0039
0040     for d in decorators:
0041         s.addContextDecorator(d)
0042
0043     addParticleGun(
0044         s,
0045         MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0046         EtaConfig(-2.0, 2.0, uniform=True),
0047         PhiConfig(0.0, 360.0 * u.degree),
0048         ParticleConfig(1, acts.PdgParticle.eMuon, randomizeCharge=True),
0049         rnd=rnd,
0050     )
0051
0052     addFatras(
0053         s,
0054         trackingGeometry,
0055         field,
0056         rnd=rnd,
0057     )
0058
0059     addDigitization(
0060         s,
0061         trackingGeometry,
0062         field,
0063         digiConfigFile=digiConfigFile,
0064         rnd=rnd,
0065     )
0066
0067     s.addAlgorithm(
0068         acts.examples.SpacePointMaker(
0069             level=acts.logging.INFO,
0070             trackingGeometry=trackingGeometry,
0071             inputMeasurements="measurement_subset",
0072             outputSpacePoints="spacepoints",
0073             geometrySelection=acts.examples.json.readJsonGeometryList(
0074                 str(geoSelectionConfigFile)
0075             ),
0076         )
0077     )
0078
0079     # Option 1: Implement a track finding algorithm...
0080     class PythonTrackFinder(acts.examples.IAlgorithm):
0081         def __init__(self, name, level):
0082             acts.examples.IAlgorithm.__init__(self, name, level)
0083
0084             self.spacepoints = acts.examples.ReadDataHandle(
0085                 self, acts.SpacePointContainer, "Spacepoints"
0086             )
0087             self.spacepoints.initialize("spacepoints")
0088
0089             self.prototracks = acts.examples.WriteDataHandle(
0090                 self, acts.examples.ProtoTrackContainer, "Prototracks"
0091             )
0092             self.prototracks.initialize("prototracks")
0093
0094         def execute(self, context):
0095             spacepoints = self.spacepoints(context.eventStore)
0096
0097             track = acts.examples.ProtoTrack()
0098             for sp in sorted(spacepoints, key=lambda sp: sp.r):
0099                 for sl in sp.sourceLinks:
0100                     isl = acts.examples.IndexSourceLink.fromSourceLink(sl)
0101                     track.append(isl.index())
0102
0103             prototracks = acts.examples.ProtoTrackContainer()
0104             prototracks.append(track)
0105
0106             self.prototracks(context, prototracks)
0107             return acts.examples.ProcessCode.SUCCESS
0108
0109     s.addAlgorithm(PythonTrackFinder("PythonTrackFinder", acts.logging.INFO))
0110
0111     # ... or option 2: use truth values
0112     truthTrkFndAlg = acts.examples.TruthTrackFinder(
0113         level=acts.logging.INFO,
0114         inputParticles="particles_generated_selected",
0115         inputMeasurements="measurements",
0116         inputParticleMeasurementsMap="particle_measurements_map",
0117         inputSimHits="simhits",
0118         inputMeasurementSimHitsMap="measurement_simhits_map",
0119         outputProtoTracks="prototracks",
0120     )
0121     # s.addAlgorithm(truthTrkFndAlg)
0122
0123     class PythonTrackFitter(acts.examples.IAlgorithm):
0124         def __init__(self, name, level):
0125             acts.examples.IAlgorithm.__init__(self, name, level)
0126
0127             self.prototracks = acts.examples.ReadDataHandle(
0128                 self, acts.examples.ProtoTrackContainer, "Prototracks"
0129             )
0130             self.prototracks.initialize("prototracks")
0131
0132             self.tracks = acts.examples.WriteDataHandle(
0133                 self, acts.examples.ConstTrackContainer, "Tracks"
0134             )
0135             self.tracks.initialize("fitted_tracks")
0136
0137             self.spacepoints = acts.examples.ReadDataHandle(
0138                 self, acts.SpacePointContainer, "Spacepoints"
0139             )
0140             self.spacepoints.initialize("spacepoints")
0141
0142         def execute(self, context):
0143             prototracks = self.prototracks(context.eventStore)
0144             spacepoints = self.spacepoints(context.eventStore)
0145
0146             measurement_to_spacepoint = {}
0147             measurement_to_sourcelink = {}
0148             for sp in spacepoints:
0149                 for sl in sp.sourceLinks:
0150                     isl = acts.examples.IndexSourceLink.fromSourceLink(sl)
0151                     meas_id = isl.index()
0152                     measurement_to_spacepoint[meas_id] = sp
0153                     measurement_to_sourcelink[meas_id] = sl
0154             surface_map = trackingGeometry.geoIdSurfaceMap()
0155
0156             container = acts.examples.TrackContainer()
0157             for prototrack in prototracks:
0158                 track = container.makeTrack()
0159                 track.parameters = acts.BoundVector(1.0, 1.0, 1.0, 1.0, 1.0, 1.0)
0160                 track.nMeasurements = len(prototrack)
0161
0162                 for meas_id in prototrack:
0163                     sp = measurement_to_spacepoint[meas_id]
0164                     sl = measurement_to_sourcelink[meas_id]
0165                     isl = acts.examples.IndexSourceLink.fromSourceLink(sl)
0166                     sf = surface_map[isl.geometryId()]
0167
0168                     trackState = track.appendTrackState()
0169                     trackState.typeFlags.isMeasurement = True
0170                     trackState.uncalibratedSourceLink = sl
0171                     trackState.referenceSurface = sf
0172
0173             self.tracks(context, container.makeConst())
0174             return acts.examples.ProcessCode.SUCCESS
0175
0176     s.addAlgorithm(PythonTrackFitter("PythonTrackFitter", acts.logging.INFO))
0177
0178     s.addAlgorithm(
0179         acts.examples.TrackTruthMatcher(
0180             level=acts.logging.INFO,
0181             inputTracks="fitted_tracks",
0182             inputParticles="particles",
0183             inputMeasurementParticlesMap="measurement_particles_map",
0184             outputTrackParticleMatching="track_particle_matching",
0185             outputParticleTrackMatching="particle_track_matching",
0186             doubleMatching=True,
0187         )
0188     )
0189
0190     # Add track finder performance writer
0191     cfg_finder = acts.examples.PythonPatternRecognitionPerformanceWriter.Config()
0192     cfg_finder.inputTracks = "fitted_tracks"
0193     cfg_finder.inputParticles = "particles"
0194     cfg_finder.inputTrackParticleMatching = "track_particle_matching"
0195     cfg_finder.inputParticleTrackMatching = "particle_track_matching"
0196     cfg_finder.inputParticleMeasurementsMap = "particle_measurements_map"
0197     perfWriterFinder = acts.examples.PythonPatternRecognitionPerformanceWriter(
0198         cfg_finder, acts.logging.INFO
0199     )
0200     s.addWriter(perfWriterFinder)
0201
0202     # Add track fitter performance writer
0203     cfg_fitter = acts.examples.PythonTrackParameterPerformanceWriter.Config()
0204     cfg_fitter.inputTracks = "fitted_tracks"
0205     cfg_fitter.inputParticles = "particles"
0206     cfg_fitter.inputTrackParticleMatching = "track_particle_matching"
0207     perfWriterFitter = acts.examples.PythonTrackParameterPerformanceWriter(
0208         cfg_fitter, acts.logging.INFO
0209     )
0210     s.addWriter(perfWriterFitter)
0211
0212     return s, perfWriterFinder, perfWriterFitter
0213
0214
0215 if __name__ == "__main__":
0216     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0217
0218     detector = acts.examples.GenericDetector(acts.examples.GenericDetector.Config())
0219     trackingGeometry = detector.trackingGeometry()
0220     decorators = detector.contextDecorators()
0221
0222     field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0223
0224     digiConfigFile = srcdir / "Examples/Configs/generic-digi-smearing-config.json"
0225     geoSelectionConfigFile = (
0226         srcdir / "Examples/Configs/generic-pixel-sstrips-lstrips-spacepoints.json"
0227     )
0228
0229     outputDir = Path.cwd() / "output_pypi_finding_fitting_demo"
0230     outputDir.mkdir(exist_ok=True)
0231
0232     s, perfWriterFinder, perfWriterFitter = runPypiFindingFittingDemo(
0233         trackingGeometry=trackingGeometry,
0234         field=field,
0235         digiConfigFile=digiConfigFile,
0236         geoSelectionConfigFile=geoSelectionConfigFile,
0237         outputDir=outputDir,
0238         decorators=decorators,
0239     )
0240     s.run()
0241
0242     histograms_finder = perfWriterFinder.histograms()
0243     print(
0244         f"Retrieved {len(histograms_finder)} finder performance histograms: {list(histograms_finder.keys())}"
0245     )
0246
0247     histograms_fitter = perfWriterFitter.histograms()
0248     print(
0249         f"Retrieved {len(histograms_fitter)} fitter performance histograms: {list(histograms_fitter.keys())}"
0250     )