Scripts/Python/full_chain_itk.py

0001 #!/usr/bin/env python3
0002 import pathlib, acts, acts.examples, acts.examples.itk
0003 from acts.examples.simulation import (
0004     addParticleGun,
0005     MomentumConfig,
0006     EtaConfig,
0007     ParticleConfig,
0008     ParticleSelectorConfig,
0009     addPythia8,
0010     addFatras,
0011     ParticleSelectorConfig,
0012     addDigitization,
0013 )
0014 from acts.examples.reconstruction import (
0015     addSeeding,
0016     SeedingAlgorithm,
0017     addCKFTracks,
0018     CkfConfig,
0019     TrackSelectorConfig,
0020     addAmbiguityResolution,
0021     AmbiguityResolutionConfig,
0022     addVertexFitting,
0023     VertexFinder,
0024 )
0025
0026 ttbar_pu200 = False
0027 u = acts.UnitConstants
0028 geo_dir = pathlib.Path("acts-itk")
0029 outputDir = pathlib.Path.cwd() / "itk_output"
0030 # acts.examples.dump_args_calls(locals())  # show acts.examples python binding calls
0031
0032 detector, trackingGeometry, decorators = acts.examples.itk.buildITkGeometry(geo_dir)
0033 field = acts.examples.MagneticFieldMapXyz(str(geo_dir / "bfield/ATLAS-BField-xyz.root"))
0034 rnd = acts.examples.RandomNumbers(seed=42)
0035
0036 s = acts.examples.Sequencer(events=100, numThreads=-1, outputDir=str(outputDir))
0037
0038 if not ttbar_pu200:
0039     addParticleGun(
0040         s,
0041         MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0042         EtaConfig(-4.0, 4.0, uniform=True),
0043         ParticleConfig(2, acts.PdgParticle.eMuon, randomizeCharge=True),
0044         rnd=rnd,
0045     )
0046 else:
0047     addPythia8(
0048         s,
0049         hardProcess=["Top:qqbar2ttbar=on"],
0050         npileup=200,
0051         vtxGen=acts.examples.GaussianVertexGenerator(
0052             stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns),
0053             mean=acts.Vector4(0, 0, 0, 0),
0054         ),
0055         rnd=rnd,
0056         outputDirRoot=outputDir,
0057     )
0058
0059 addFatras(
0060     s,
0061     trackingGeometry,
0062     field,
0063     rnd=rnd,
0064     preSelectParticles=(
0065         ParticleSelectorConfig(
0066             rho=(0.0 * u.mm, 28.0 * u.mm),
0067             absZ=(0.0 * u.mm, 1.0 * u.m),
0068             eta=(-4.0, 4.0),
0069             pt=(150 * u.MeV, None),
0070         )
0071         if ttbar_pu200
0072         else ParticleSelectorConfig()
0073     ),
0074     postSelectParticles=ParticleSelectorConfig(
0075         pt=(1.0 * u.GeV, None),
0076         eta=(-4.0, 4.0),
0077         hits=(9, None),
0078         removeNeutral=True,
0079     ),
0080     outputDirRoot=outputDir,
0081 )
0082
0083 addDigitization(
0084     s,
0085     trackingGeometry,
0086     field,
0087     digiConfigFile=geo_dir / "itk-hgtd/itk-smearing-config.json",
0088     outputDirRoot=outputDir,
0089     rnd=rnd,
0090 )
0091
0092 addSeeding(
0093     s,
0094     trackingGeometry,
0095     field,
0096     seedingAlgorithm=SeedingAlgorithm.Default,
0097     *acts.examples.itk.itkSeedingAlgConfig(
0098         acts.examples.itk.InputSpacePointsType.PixelSpacePoints
0099     ),
0100     initialSigmas=[
0101         1 * u.mm,
0102         1 * u.mm,
0103         1 * u.degree,
0104         1 * u.degree,
0105         0.1 * u.e / u.GeV,
0106         1 * u.ns,
0107     ],
0108     initialSigmaPtRel=0.1,
0109     initialVarInflation=[1.0] * 6,
0110     geoSelectionConfigFile=geo_dir / "itk-hgtd/geoSelection-ITk.json",
0111     outputDirRoot=outputDir,
0112 )
0113
0114 addCKFTracks(
0115     s,
0116     trackingGeometry,
0117     field,
0118     trackSelectorConfig=(
0119         # fmt: off
0120         TrackSelectorConfig(absEta=(None, 2.0), pt=(0.9 * u.GeV, None), nMeasurementsMin=9, maxHoles=2, maxOutliers=2, maxSharedHits=2),
0121         TrackSelectorConfig(absEta=(None, 2.6), pt=(0.4 * u.GeV, None), nMeasurementsMin=8, maxHoles=2, maxOutliers=2, maxSharedHits=2),
0122         TrackSelectorConfig(absEta=(None, 4.0), pt=(0.4 * u.GeV, None), nMeasurementsMin=7, maxHoles=2, maxOutliers=2, maxSharedHits=2),
0123         # fmt: on
0124     ),
0125     ckfConfig=CkfConfig(
0126         seedDeduplication=True,
0127         stayOnSeed=True,
0128         # ITk volumes from Noemi's plot
0129         pixelVolumes=[8, 9, 10, 13, 14, 15, 16, 18, 19, 20],
0130         stripVolumes=[22, 23, 24],
0131         maxPixelHoles=1,
0132         maxStripHoles=2,
0133     ),
0134     outputDirRoot=outputDir,
0135 )
0136
0137 addAmbiguityResolution(
0138     s,
0139     AmbiguityResolutionConfig(
0140         maximumSharedHits=3,
0141         maximumIterations=10000,
0142         nMeasurementsMin=6,
0143     ),
0144     outputDirRoot=outputDir,
0145 )
0146
0147 addVertexFitting(
0148     s,
0149     field,
0150     vertexFinder=VertexFinder.AMVF,
0151     outputDirRoot=outputDir,
0152 )
0153
0154 s.run()