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_input",
0051         outputSimHits="sim_hits",
0052         rnd=rnd,
0053         preSelectParticles=None,
0054     )
0055
0056     # Set up the track lookup grid writer
0057     jsonWriterConfig = acts.examples.JsonTrackParamsLookupWriter.Config(path=outputPath)
0058     jsonWriter = acts.examples.JsonTrackParamsLookupWriter(jsonWriterConfig)
0059
0060     # Set up the track estimation algorithm
0061     surfaces = list(trackingGeometry.geoIdSurfaceMap().values())
0062     refSurface = surfaces[0]
0063     refGeometryId = refSurface.geometryId()
0064
0065     trackEstConfig = acts.examples.TrackParamsLookupEstimation.Config(
0066         refLayers={refGeometryId: refSurface},
0067         bins=(1, 1000),
0068         inputHits="sim_hits",
0069         inputParticles="particles_input",
0070         trackLookupGridWriters=[jsonWriter],
0071     )
0072     trackEstAlg = acts.examples.TrackParamsLookupEstimation(
0073         trackEstConfig, acts.logging.INFO
0074     )
0075
0076     s.addAlgorithm(trackEstAlg)
0077
0078     s.run()
0079
0080
0081 if __name__ == "__main__":
0082     p = argparse.ArgumentParser()
0083
0084     p.add_argument(
0085         "-n",
0086         "--events",
0087         type=int,
0088         default=100000,
0089         help="Number of events for lookup estimation",
0090     )
0091     p.add_argument(
0092         "-o",
0093         "--output",
0094         type=str,
0095         default="lookup.json",
0096         help="Output lookup file name",
0097     )
0098
0099     args = p.parse_args()
0100
0101     # Initialize the geometry
0102     detector = acts.examples.TelescopeDetector(
0103         bounds=[4, 10],
0104         positions=[30, 60, 90],
0105         stereos=[0, 0, 0],
0106         binValue=2,
0107         surfaceType=0,
0108     )
0109     trackingGeometry = detector.trackingGeometry()
0110
0111     # Estimate the lookup
0112     estimateLookup(trackingGeometry, args.events, args.output)