Scripts/Python/gbts_layer_connection_training_odd.py

0001 #!/usr/bin/env python3
0002
0003 import argparse
0004 import pathlib
0005
0006 import acts
0007 import acts.examples
0008 from acts.examples.simulation import (
0009     MomentumConfig,
0010     EtaConfig,
0011     ParticleConfig,
0012     ParticleSelectorConfig,
0013     addParticleGun,
0014     addFatras,
0015     addDigitization,
0016     addDigiParticleSelection,
0017 )
0018 from acts.examples.reconstruction import addGbtsTraining
0019 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0020
0021 u = acts.UnitConstants
0022
0023 parser = argparse.ArgumentParser(
0024     description="Layer connection training tool for the GBTS on ODD"
0025 )
0026 parser.add_argument(
0027     "--output",
0028     "-o",
0029     help="Output directory",
0030     type=pathlib.Path,
0031     default=pathlib.Path.cwd() / "odd_gbts_training",
0032 )
0033 parser.add_argument("--events", "-n", help="Number of events", type=int, default=400)
0034 parser.add_argument("--particles", help="Muons per event", type=int, default=50)
0035 parser.add_argument(
0036     "--jobs",
0037     "-j",
0038     help="Number of worker threads (-1 uses all cores)",
0039     type=int,
0040     default=-1,
0041 )
0042 parser.add_argument(
0043     "--prob-threshold",
0044     help="Drop transitions below this probability (-1 keeps all)",
0045     type=float,
0046     default=-1.0,
0047 )
0048 args = parser.parse_args()
0049
0050 actsDir = pathlib.Path(__file__).parent.parent.parent.parent
0051 configDir = actsDir / "Examples/Configs"
0052 outputDir = args.output
0053 outputDir.mkdir(parents=True, exist_ok=True)
0054
0055 oddDir = getOpenDataDetectorDirectory()
0056 oddMaterialDeco = acts.IMaterialDecorator.fromFile(
0057     oddDir / "data/odd-material-maps.root"
0058 )
0059
0060 detector = getOpenDataDetector(materialDecorator=oddMaterialDeco)
0061 trackingGeometry = detector.trackingGeometry()
0062 field = acts.ConstantBField(acts.Vector3(0.0, 0.0, 2.0 * u.T))
0063 rnd = acts.examples.RandomNumbers(seed=42)
0064
0065 s = acts.examples.Sequencer(
0066     events=args.events, numThreads=args.jobs, outputDir=str(outputDir)
0067 )
0068
0069 addParticleGun(
0070     s,
0071     MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0072     EtaConfig(-3.0, 3.0, uniform=True),
0073     ParticleConfig(args.particles, acts.PdgParticle.eMuon, randomizeCharge=True),
0074     vtxGen=acts.examples.GaussianVertexGenerator(
0075         mean=acts.Vector4(0, 0, 0, 0),
0076         stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns),
0077     ),
0078     rnd=rnd,
0079 )
0080
0081 addFatras(
0082     s,
0083     trackingGeometry,
0084     field,
0085     rnd=rnd,
0086 )
0087
0088 addDigitization(
0089     s,
0090     trackingGeometry,
0091     field,
0092     digiConfigFile=configDir / "odd-digi-smearing-config.json",
0093     rnd=rnd,
0094 )
0095
0096 addDigiParticleSelection(
0097     s,
0098     ParticleSelectorConfig(
0099         pt=(1.0 * u.GeV, None),
0100         eta=(-3.0, 3.0),
0101         measurements=(3, None),
0102         removeNeutral=True,
0103     ),
0104 )
0105
0106 addGbtsTraining(
0107     s,
0108     trackingGeometry,
0109     selectedParticles="particles_selected",
0110     geometryFile=configDir / "odd-gbts-layer-map.json",
0111     outputConnectionTable=outputDir / "odd-gbts-connection-table.json",
0112     probThreshold=args.prob_threshold,
0113     doSymmetrization=True,
0114     logLevel=acts.logging.INFO,
0115 )
0116
0117 s.run()