Scripts/Python/exatrkx.py

0001 #!/usr/bin/env python3
0002 from pathlib import Path
0003
0004 import acts.examples
0005 import acts
0006 from acts import UnitConstants as u
0007
0008
0009 if "__main__" == __name__:
0010     import os
0011     import sys
0012     from digitization import runDigitization
0013     from acts.examples.reconstruction import addExaTrkX, ExaTrkXBackend
0014
0015     backend = ExaTrkXBackend.Torch
0016
0017     if "onnx" in sys.argv:
0018         backend = ExaTrkXBackend.Onnx
0019     if "torch" in sys.argv:
0020         backend = ExaTrkXBackend.Torch
0021
0022     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0023
0024     detector, trackingGeometry, decorators = acts.examples.GenericDetector.create()
0025
0026     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0027
0028     inputParticlePath = Path("particles.root")
0029     if not inputParticlePath.exists():
0030         inputParticlePath = None
0031
0032     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0033
0034     geometrySelection = (
0035         srcdir
0036         / "Examples/Algorithms/TrackFinding/share/geoSelection-genericDetector.json"
0037     )
0038     assert geometrySelection.exists()
0039
0040     digiConfigFile = (
0041         srcdir
0042         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json"
0043     )
0044     assert digiConfigFile.exists()
0045
0046     if backend == ExaTrkXBackend.Torch:
0047         modelDir = Path.cwd() / "torchscript_models"
0048         assert (modelDir / "embed.pt").exists()
0049         assert (modelDir / "filter.pt").exists()
0050         assert (modelDir / "gnn.pt").exists()
0051     else:
0052         modelDir = Path.cwd() / "onnx_models"
0053         assert (modelDir / "embedding.onnx").exists()
0054         assert (modelDir / "filtering.onnx").exists()
0055         assert (modelDir / "gnn.onnx").exists()
0056
0057     s = acts.examples.Sequencer(events=2, numThreads=1)
0058     s.config.logLevel = acts.logging.INFO
0059
0060     rnd = acts.examples.RandomNumbers()
0061     outputDir = Path(os.getcwd())
0062
0063     s = runDigitization(
0064         trackingGeometry,
0065         field,
0066         outputDir,
0067         digiConfigFile=digiConfigFile,
0068         particlesInput=inputParticlePath,
0069         outputRoot=True,
0070         outputCsv=True,
0071         s=s,
0072     )
0073
0074     addExaTrkX(
0075         s,
0076         trackingGeometry,
0077         geometrySelection,
0078         modelDir,
0079         outputDir,
0080         backend=backend,
0081     )
0082
0083     s.run()