Scripts/Python/gnn.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 import os
0005 import sys
0006
0007 import acts
0008 import acts.examples
0009 import acts.examples.gnn
0010 from acts.examples.reconstruction import addGnn, addSpacePointsMaking
0011 from acts.examples.gnn import (
0012     TorchMetricLearning,
0013     TorchEdgeClassifier,
0014     BoostTrackBuilding,
0015     NodeFeature,
0016 )
0017 from acts import UnitConstants as u
0018
0019 from digitization import runDigitization
0020
0021
0022 def runGnnMetricLearning(
0023     trackingGeometry,
0024     field,
0025     outputDir,
0026     digiConfigFile,
0027     geometrySelection,
0028     embedModelPath,
0029     filterModelPath,
0030     gnnModelPath,
0031     outputRoot=False,
0032     outputCsv=False,
0033     s=None,
0034 ):
0035     s = runDigitization(
0036         trackingGeometry,
0037         field,
0038         outputDir,
0039         digiConfigFile=digiConfigFile,
0040         particlesInput=None,
0041         outputRoot=outputRoot,
0042         outputCsv=outputCsv,
0043         s=s,
0044     )
0045
0046     addSpacePointsMaking(
0047         s,
0048         geoSelectionConfigFile=geometrySelection,
0049         stripGeoSelectionConfigFile=None,
0050         trackingGeometry=trackingGeometry,
0051         logLevel=acts.logging.INFO,
0052     )
0053
0054     graphConstructorConfig = {
0055         "level": acts.logging.INFO,
0056         "modelPath": str(embedModelPath),
0057         "embeddingDim": 8,
0058         "rVal": 1.6,
0059         "knnVal": 100,
0060         "selectedFeatures": [0, 1, 2],  # R, Phi, Z
0061     }
0062     graphConstructor = TorchMetricLearning(**graphConstructorConfig)
0063
0064     filterConfig = {
0065         "level": acts.logging.INFO,
0066         "modelPath": str(filterModelPath),
0067         "cut": 0.01,
0068     }
0069     gnnConfig = {
0070         "level": acts.logging.INFO,
0071         "modelPath": str(gnnModelPath),
0072         "cut": 0.5,
0073     }
0074
0075     edgeClassifiers = []
0076
0077     if filterModelPath.suffix == ".pt":
0078         edgeClassifiers.append(
0079             TorchEdgeClassifier(
0080                 **filterConfig,
0081                 nChunks=5,
0082                 undirected=False,
0083                 selectedFeatures=[0, 1, 2],
0084             )
0085         )
0086     elif filterModelPath.suffix == ".onnx":
0087         from acts.examples.gnn import OnnxEdgeClassifier
0088
0089         edgeClassifiers.append(OnnxEdgeClassifier(**filterConfig))
0090     else:
0091         raise ValueError(f"Unsupported model format: {filterModelPath.suffix}")
0092
0093     if gnnModelPath.suffix == ".pt":
0094         edgeClassifiers.append(
0095             TorchEdgeClassifier(
0096                 **gnnConfig,
0097                 undirected=True,
0098                 selectedFeatures=[0, 1, 2],
0099             )
0100         )
0101     elif gnnModelPath.suffix == ".onnx":
0102         edgeClassifiers.append(
0103             OnnxEdgeClassifier(**gnnConfig),
0104         )
0105     else:
0106         raise ValueError(f"Unsupported model format: {filterModelPath.suffix}")
0107
0108     # Stage 3: CPU track building
0109     trackBuilderConfig = {
0110         "level": acts.logging.INFO,
0111     }
0112     trackBuilder = BoostTrackBuilding(**trackBuilderConfig)
0113
0114     # Node features: Standard 3 features (R, Phi, Z)
0115     nodeFeatures = [
0116         NodeFeature.R,
0117         NodeFeature.Phi,
0118         NodeFeature.Z,
0119     ]
0120     featureScales = [1.0, 1.0, 1.0]
0121
0122     # Add GNN tracking
0123     addGnn(
0124         s,
0125         graphConstructor=graphConstructor,
0126         edgeClassifiers=edgeClassifiers,
0127         trackBuilder=trackBuilder,
0128         nodeFeatures=nodeFeatures,
0129         featureScales=featureScales,
0130         outputDirRoot=outputDir if outputRoot else None,
0131         logLevel=acts.logging.INFO,
0132     )
0133
0134     s.run()
0135
0136
0137 if "__main__" == __name__:
0138     detector = acts.examples.GenericDetector()
0139     trackingGeometry = detector.trackingGeometry()
0140
0141     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0142
0143     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0144
0145     geometrySelection = srcdir / "Examples/Configs/generic-seeding-config.json"
0146     assert geometrySelection.exists()
0147
0148     digiConfigFile = srcdir / "Examples/Configs/generic-digi-smearing-config.json"
0149     assert digiConfigFile.exists()
0150
0151     # Model paths from MODEL_STORAGE environment variable
0152     model_storage = os.environ.get("MODEL_STORAGE")
0153     assert model_storage is not None, "MODEL_STORAGE environment variable is not set"
0154     ci_models = Path(model_storage)
0155     if "onnx" in sys.argv:
0156         embedModelPath = ci_models / "torchscript_models/embed.pt"
0157         filterModelPath = ci_models / "torchscript_models/filter.pt"
0158         gnnModelPath = ci_models / "torchscript_models/gnn.pt"
0159     elif "torch" in sys.argv:
0160         embedModelPath = ci_models / "torchscript_models/embed.pt"
0161         filterModelPath = ci_models / "onnx_models/filtering.onnx"
0162         gnnModelPath = ci_models / "onnx_models/gnn.onnx"
0163     else:
0164         raise ValueError("Please specify backend: 'torch' or 'onnx'")
0165
0166     s = acts.examples.Sequencer(events=2, numThreads=1)
0167     s.config.logLevel = acts.logging.INFO
0168
0169     rnd = acts.examples.RandomNumbers()
0170     outputDir = Path(os.getcwd())
0171
0172     runGnnMetricLearning(
0173         trackingGeometry,
0174         field,
0175         outputDir,
0176         digiConfigFile,
0177         geometrySelection,
0178         embedModelPath,
0179         filterModelPath,
0180         gnnModelPath,
0181         outputRoot=True,
0182         outputCsv=False,
0183         s=s,
0184     )