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