Scripts/Python/track_finding_python_only.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 import os
0011 from pathlib import Path
0012
0013 os.environ["ACTS_SEQUENCER_DISABLE_FPEMON"] = "1"
0014
0015 import acts
0016 import acts.examples
0017 from acts import UnitConstants as u
0018
0019
0020 def runTrackFindingPythonOnly(
0021     trackingGeometry,
0022     field,
0023     digiConfigFile,
0024     geoSelectionConfigFile,
0025     outputDir,
0026     decorators=[],
0027     s=None,
0028 ):
0029     from acts.examples.simulation import (
0030         addParticleGun,
0031         MomentumConfig,
0032         EtaConfig,
0033         PhiConfig,
0034         ParticleConfig,
0035         addFatras,
0036         addDigitization,
0037     )
0038
0039     s = s or acts.examples.Sequencer(events=1, numThreads=1, logLevel=acts.logging.INFO)
0040     outputDir = Path(outputDir)
0041     rnd = acts.examples.RandomNumbers(seed=42)
0042
0043     for d in decorators:
0044         s.addContextDecorator(d)
0045
0046     addParticleGun(
0047         s,
0048         MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0049         EtaConfig(-2.0, 2.0, uniform=True),
0050         PhiConfig(0.0, 360.0 * u.degree),
0051         ParticleConfig(1, acts.PdgParticle.eMuon, randomizeCharge=True),
0052         rnd=rnd,
0053     )
0054
0055     addFatras(
0056         s,
0057         trackingGeometry,
0058         field,
0059         rnd=rnd,
0060     )
0061
0062     addDigitization(
0063         s,
0064         trackingGeometry,
0065         field,
0066         digiConfigFile=digiConfigFile,
0067         rnd=rnd,
0068     )
0069
0070     s.addAlgorithm(
0071         acts.examples.SpacePointMaker(
0072             level=acts.logging.INFO,
0073             trackingGeometry=trackingGeometry,
0074             inputMeasurements="measurements",
0075             outputSpacePoints="spacepoints",
0076             geometrySelection=acts.examples.json.readJsonGeometryList(
0077                 str(geoSelectionConfigFile)
0078             ),
0079         )
0080     )
0081
0082     class PythonTrackFinder(acts.examples.IAlgorithm):
0083         def __init__(self, name, level):
0084             acts.examples.IAlgorithm.__init__(self, name, level)
0085
0086             self.spacepoints = acts.examples.ReadDataHandle(
0087                 self, acts.SpacePointContainer2, "Spacepoints"
0088             )
0089             self.spacepoints.initialize("spacepoints")
0090
0091             self.prototracks = acts.examples.WriteDataHandle(
0092                 self, acts.examples.ProtoTrackContainer, "Prototracks"
0093             )
0094             self.prototracks.initialize("prototracks")
0095
0096         def execute(self, context):
0097             spacepoints = self.spacepoints(context.eventStore)
0098
0099             track = acts.examples.ProtoTrack()
0100             for sp in sorted(spacepoints, key=lambda sp: sp.r):
0101                 for sl in sp.sourceLinks:
0102                     isl = acts.examples.IndexSourceLink.FromSourceLink(sl)
0103                     track.append(isl.index())
0104
0105             prototracks = acts.examples.ProtoTrackContainer()
0106             prototracks.append(track)
0107
0108             self.prototracks(context, prototracks)
0109             return acts.examples.ProcessCode.SUCCESS
0110
0111     s.addAlgorithm(PythonTrackFinder("PythonTrackFinder", acts.logging.INFO))
0112
0113     class PythonTrackFitter(acts.examples.IAlgorithm):
0114         def __init__(self, name, level):
0115             acts.examples.IAlgorithm.__init__(self, name, level)
0116
0117             self.prototracks = acts.examples.ReadDataHandle(
0118                 self, acts.examples.ProtoTrackContainer, "Prototracks"
0119             )
0120             self.prototracks.initialize("prototracks")
0121
0122             self.tracks = acts.examples.WriteDataHandle(
0123                 self, acts.examples.ConstTrackContainer, "Tracks"
0124             )
0125             self.tracks.initialize("fitted_tracks")
0126
0127         def execute(self, context):
0128             prototracks = self.prototracks(context.eventStore)
0129
0130             container = acts.examples.TrackContainer()
0131             for prototrack in prototracks:
0132                 track = container.makeTrack()
0133                 track.parameters = acts.BoundVector(1.0, 1.0, 1.0, 1.0, 1.0, 1.0)
0134                 track.nMeasurements = len(prototrack)
0135
0136             self.tracks(context, container.makeConst())
0137             return acts.examples.ProcessCode.SUCCESS
0138
0139     s.addAlgorithm(PythonTrackFitter("PythonTrackFitter", acts.logging.INFO))
0140
0141     return s
0142
0143
0144 if __name__ == "__main__":
0145     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0146
0147     detector = acts.examples.GenericDetector(acts.examples.GenericDetector.Config())
0148     trackingGeometry = detector.trackingGeometry()
0149     decorators = detector.contextDecorators()
0150
0151     field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0152
0153     digiConfigFile = srcdir / "Examples/Configs/generic-digi-smearing-config.json"
0154     geoSelectionConfigFile = srcdir / "Examples/Configs/generic-seeding-config.json"
0155
0156     outputDir = Path.cwd() / "output_track_finding_python_only"
0157     outputDir.mkdir(exist_ok=True)
0158
0159     runTrackFindingPythonOnly(
0160         trackingGeometry=trackingGeometry,
0161         field=field,
0162         digiConfigFile=digiConfigFile,
0163         geoSelectionConfigFile=geoSelectionConfigFile,
0164         outputDir=outputDir,
0165         decorators=decorators,
0166     ).run()