Scripts/Python/telescope_track_params_lookup_generation.py

0001 #!/usr/bin/env python3
0002
0003 import argparse
0004
0005 import acts
0006 import acts.examples
0007 from acts.examples.simulation import (
0008     addParticleGun,
0009     addFatras,
0010     MomentumConfig,
0011     EtaConfig,
0012     PhiConfig,
0013     ParticleConfig,
0014 )
0015
0016 u = acts.UnitConstants
0017
0018
0019 def estimateLookup(trackingGeometry, numEvents, outputPath):
0020
0021     # Set up the dipole magnetic field
0022     field = acts.ConstantBField(acts.Vector3(50 * u.T, 0, 0))
0023
0024     # Fatras simulation of muons
0025     rnd = acts.examples.RandomNumbers(seed=42)
0026
0027     s = acts.examples.Sequencer(
0028         events=numEvents, numThreads=1, logLevel=acts.logging.INFO
0029     )
0030
0031     vertexGen = acts.examples.GaussianVertexGenerator(
0032         stddev=acts.Vector4(0, 0, 0, 0), mean=acts.Vector4(0, 9, 0, 0)
0033     )
0034
0035     addParticleGun(
0036         s=s,
0037         etaConfig=EtaConfig(10.0, 10.0),
0038         phiConfig=PhiConfig(0, 0),
0039         momentumConfig=MomentumConfig(0.5 * u.GeV, 10 * u.GeV),
0040         particleConfig=ParticleConfig(1, acts.PdgParticle.eMuon, False),
0041         multiplicity=1,
0042         rnd=rnd,
0043         vtxGen=vertexGen,
0044     )
0045
0046     addFatras(
0047         s,
0048         trackingGeometry,
0049         field,
0050         inputParticles="particles_generated",
0051         outputSimHits="sim_hits",
0052         rnd=rnd,
0053     )
0054
0055     # Set up the track lookup grid writer
0056     jsonWriterConfig = acts.examples.JsonTrackParamsLookupWriter.Config(path=outputPath)
0057     jsonWriter = acts.examples.JsonTrackParamsLookupWriter(jsonWriterConfig)
0058
0059     # Set up the track estimation algorithm
0060     surfaces = list(trackingGeometry.geoIdSurfaceMap().values())
0061     refSurface = surfaces[0]
0062     refGeometryId = refSurface.geometryId
0063
0064     trackEstConfig = acts.examples.TrackParamsLookupEstimation.Config(
0065         refLayers={refGeometryId: refSurface},
0066         bins=(1, 1000),
0067         inputHits="sim_hits",
0068         inputParticles="particles_generated",
0069         trackLookupGridWriters=[jsonWriter],
0070     )
0071     trackEstAlg = acts.examples.TrackParamsLookupEstimation(
0072         trackEstConfig, acts.logging.INFO
0073     )
0074
0075     s.addAlgorithm(trackEstAlg)
0076
0077     s.run()
0078
0079
0080 if __name__ == "__main__":
0081     p = argparse.ArgumentParser()
0082
0083     p.add_argument(
0084         "-n",
0085         "--events",
0086         type=int,
0087         default=100000,
0088         help="Number of events for lookup estimation",
0089     )
0090     p.add_argument(
0091         "-o",
0092         "--output",
0093         type=str,
0094         default="lookup.json",
0095         help="Output lookup file name",
0096     )
0097
0098     args = p.parse_args()
0099
0100     # Initialize the geometry
0101     detector = acts.examples.TelescopeDetector(
0102         bounds=[4, 10],
0103         positions=[30, 60, 90],
0104         stereos=[0, 0, 0],
0105         binValue=2,
0106         surfaceType=0,
0107     )
0108     trackingGeometry = detector.trackingGeometry()
0109
0110     # Estimate the lookup
0111     estimateLookup(trackingGeometry, args.events, args.output)