Scripts/Python/gbts_layer_connection_training.py

0001 #!/usr/bin/env python3
0002 import pathlib, acts, acts.root, acts.examples, acts.examples.itk
0003 from acts.examples.simulation import (
0004     addParticleGun,
0005     MomentumConfig,
0006     EtaConfig,
0007     ParticleConfig,
0008     addPythia8,
0009     ParticleSelectorConfig,
0010     addGenParticleSelection,
0011     addFatras,
0012     addDigitization,
0013     addDigiParticleSelection,
0014 )
0015 from acts.examples.reconstruction import addGbtsTraining
0016
0017 ttbar_pu200 = False
0018 u = acts.UnitConstants
0019 geo_dir = pathlib.Path("acts-itk")
0020 outputDir = pathlib.Path.cwd() / "itk_output"
0021 # acts.examples.dump_args_calls()  # show acts.examples python binding calls
0022
0023 detector = acts.examples.itk.buildITkGeometry(geo_dir)
0024 trackingGeometry = detector.trackingGeometry()
0025 field = acts.root.MagneticFieldMapXyz(str(geo_dir / "bfield/ATLAS-BField-xyz.root"))
0026 rnd = acts.examples.RandomNumbers(seed=42)
0027
0028 s = acts.examples.Sequencer(events=1, numThreads=96, outputDir=str(outputDir))
0029
0030 if not ttbar_pu200:
0031     addParticleGun(
0032         s,
0033         MomentumConfig(1.0 * u.GeV, 1.0 * u.GeV, transverse=True),
0034         EtaConfig(-4.0, 4.0, uniform=True),
0035         ParticleConfig(2, acts.PdgParticle.eMuon, randomizeCharge=True),
0036         rnd=rnd,
0037     )
0038 else:
0039     addPythia8(
0040         s,
0041         hardProcess=["Top:qqbar2ttbar=on"],
0042         npileup=200,
0043         vtxGen=acts.examples.GaussianVertexGenerator(
0044             stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 5.0 * u.ns),
0045             mean=acts.Vector4(0, 0, 0, 0),
0046         ),
0047         rnd=rnd,
0048         outputDirRoot=outputDir,
0049     )
0050
0051     addGenParticleSelection(
0052         s,
0053         ParticleSelectorConfig(
0054             rho=(0.0 * u.mm, 28.0 * u.mm),
0055             absZ=(0.0 * u.mm, 1.0 * u.m),
0056             eta=(-4.0, 4.0),
0057             pt=(150 * u.MeV, None),
0058         ),
0059     )
0060
0061 addFatras(
0062     s,
0063     trackingGeometry,
0064     field,
0065     rnd=rnd,
0066     outputDirRoot=outputDir,
0067 )
0068
0069 addDigitization(
0070     s,
0071     trackingGeometry,
0072     field,
0073     digiConfigFile=geo_dir / "itk-hgtd/itk-smearing-config.json",
0074     outputDirRoot=outputDir,
0075     rnd=rnd,
0076 )
0077
0078 addDigiParticleSelection(
0079     s,
0080     ParticleSelectorConfig(
0081         pt=(1.0 * u.GeV, None),
0082         eta=(-4.0, 4.0),
0083         measurements=(9, None),
0084         removeNeutral=True,
0085     ),
0086 )
0087
0088 addGbtsTraining(
0089     s,
0090     selectedParticles="particles_selected",
0091     geometryFile=geo_dir / "itk-hgtd/gbts_layer_geometry.txt",
0092     outputConnectionTable=outputDir / "new_implementation_symm_1000k_singleMu.txt",
0093     probThreshold=-1.0,
0094     zMinTol=0.2340,
0095     zMaxTol=0.2340,
0096     rMinTol=2.5337,
0097     rMaxTol=2.5337,
0098     doSymmetrization=True,
0099     useOldFormatting=True,
0100     logLevel=acts.logging.INFO,
0101 )
0102
0103 s.run()