Scripts/Python/gnn_module_map_odd.py

0001 #!/usr/bin/env python3
0002
0003 """
0004 Example: GNN track finding with module maps on ODD.
0005
0006 Full simulation pipeline using module map (geometry-based) graph construction on the
0007 OpenDataDetector. Demonstrates the complete workflow from event generation to track finding.
0008
0009 All parameters are hardcoded except file paths. Users should copy and modify this script
0010 for their specific needs.
0011 """
0012
0013 from pathlib import Path
0014 import os
0015
0016 import acts
0017 from acts import UnitConstants as u
0018 from acts.examples import Sequencer
0019 from acts.examples.odd import getOpenDataDetector
0020 from acts.examples.simulation import (
0021     addParticleGun,
0022     EtaConfig,
0023     PhiConfig,
0024     MomentumConfig,
0025     ParticleConfig,
0026     addFatras,
0027     addDigitization,
0028     addGenParticleSelection,
0029     ParticleSelectorConfig,
0030 )
0031 from acts.examples.reconstruction import addGnn, addSpacePointsMaking
0032 from acts.examples.gnn import (
0033     ModuleMapCuda,
0034     CudaTrackBuilding,
0035     NodeFeature,
0036 )
0037
0038
0039 def runGnnModuleMap(
0040     trackingGeometry,
0041     field,
0042     geometrySelection,
0043     stripGeometrySelection,
0044     digiConfigFile,
0045     moduleMapPath,
0046     gnnModel,
0047     outputDir,
0048     events=100,
0049     s=None,
0050 ):
0051     """
0052     Run GNN tracking with module maps on ODD.
0053
0054     This example shows the full simulation chain with module map-based GNN.
0055     All GNN parameters are hardcoded for reproducibility.
0056
0057     Args:
0058         trackingGeometry: Tracking geometry
0059         field: Magnetic field
0060         geometrySelection: Geometry selection JSON file path
0061         digiConfigFile: Digitization config file path
0062         moduleMapPath: Path prefix for module map files
0063                       (will load .doublets.root and .triplets.root)
0064         gnnModel: Path to trained model (.pt, .onnx, or .engine)
0065         outputDir: Output directory for performance files
0066         events: Number of events to process
0067         s: Optional sequencer (creates new one if None)
0068     """
0069     # Validate inputs
0070     assert Path(
0071         moduleMapPath + ".doublets.root"
0072     ).exists(), f"Module map not found: {moduleMapPath}.doublets.root"
0073     assert Path(
0074         moduleMapPath + ".triplets.root"
0075     ).exists(), f"Module map not found: {moduleMapPath}.triplets.root"
0076     assert Path(gnnModel).exists(), f"Model file not found: {gnnModel}"
0077
0078     s = s or Sequencer(events=events, numThreads=1)
0079
0080     # Random number generator
0081     rnd = acts.examples.RandomNumbers(seed=42)
0082
0083     addParticleGun(
0084         s,
0085         MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0086         EtaConfig(-3.0, 3.0, uniform=True),
0087         PhiConfig(0.0, 360.0 * u.degree),
0088         ParticleConfig(10, acts.PdgParticle.eMuon, randomizeCharge=True),
0089         vtxGen=acts.examples.GaussianVertexGenerator(
0090             mean=acts.Vector4(0, 0, 0, 0),
0091             stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns),
0092         ),
0093         multiplicity=50,
0094         rnd=rnd,
0095     )
0096
0097     addGenParticleSelection(
0098         s,
0099         ParticleSelectorConfig(
0100             rho=(0.0, 24 * u.mm),
0101             absZ=(0.0, 1.0 * u.m),
0102             eta=(-3.0, 3.0),
0103             pt=(150 * u.MeV, None),
0104         ),
0105     )
0106
0107     # FATRAS simulation
0108     addFatras(
0109         s,
0110         trackingGeometry,
0111         field,
0112         rnd=rnd,
0113         enableInteractions=True,
0114         outputDirRoot=None,
0115         outputDirCsv=None,
0116         outputDirObj=None,
0117         logLevel=acts.logging.INFO,
0118     )
0119
0120     # Digitization
0121     addDigitization(
0122         s,
0123         trackingGeometry,
0124         field,
0125         digiConfigFile=digiConfigFile,
0126         rnd=rnd,
0127         logLevel=acts.logging.INFO,
0128     )
0129
0130     addSpacePointsMaking(
0131         s,
0132         trackingGeometry,
0133         geoSelectionConfigFile=geometrySelection,
0134         stripGeoSelectionConfigFile=stripGeometrySelection,
0135         logLevel=acts.logging.INFO,
0136     )
0137
0138     moduleMapConfig = {
0139         "level": acts.logging.INFO,
0140         "moduleMapPath": moduleMapPath,
0141         "rScale": 1000.0,
0142         "phiScale": 3.141592654,
0143         "zScale": 1000.0,
0144         "etaScale": 1.0,
0145         "gpuDevice": 0,
0146         "gpuBlocks": 512,
0147         "moreParallel": True,
0148     }
0149     graphConstructor = ModuleMapCuda(**moduleMapConfig)
0150
0151     gnnModel = Path(gnnModel)
0152     edgeClassifierConfig = {
0153         "level": acts.logging.INFO,
0154         "modelPath": str(gnnModel),
0155         "cut": 0.5,
0156     }
0157
0158     if gnnModel.suffix == ".pt":
0159         edgeClassifierConfig["useEdgeFeatures"] = True
0160         from acts.examples.gnn import TorchEdgeClassifier
0161
0162         edgeClassifiers = [TorchEdgeClassifier(**edgeClassifierConfig)]
0163     elif gnnModel.suffix == ".onnx":
0164         from acts.examples.gnn import OnnxEdgeClassifier
0165
0166         edgeClassifiers = [OnnxEdgeClassifier(**edgeClassifierConfig)]
0167     elif gnnModel.suffix == ".engine":
0168         from acts.examples.gnn import TensorRTEdgeClassifier
0169
0170         edgeClassifiers = [TensorRTEdgeClassifier(**edgeClassifierConfig)]
0171     else:
0172         raise ValueError(f"Unsupported model format: {gnnModel.suffix}")
0173
0174     trackBuilderConfig = {
0175         "level": acts.logging.INFO,
0176         "useOneBlockImplementation": False,
0177         "doJunctionRemoval": True,
0178     }
0179     trackBuilder = CudaTrackBuilding(**trackBuilderConfig)
0180
0181     e = NodeFeature
0182     nodeFeatures = [
0183         e.R,
0184         e.Phi,
0185         e.Z,
0186         e.Eta,
0187         e.Cluster1R,
0188         e.Cluster1Phi,
0189         e.Cluster1Z,
0190         e.Cluster1Eta,
0191         e.Cluster2R,
0192         e.Cluster2Phi,
0193         e.Cluster2Z,
0194         e.Cluster2Eta,
0195     ]
0196     featureScales = [1000.0, 3.141592654, 1000.0, 1.0] * 3
0197
0198     addGnn(
0199         s,
0200         graphConstructor=graphConstructor,
0201         edgeClassifiers=edgeClassifiers,
0202         trackBuilder=trackBuilder,
0203         nodeFeatures=nodeFeatures,
0204         featureScales=featureScales,
0205         inputClusters="clusters",
0206         outputDirRoot=str(outputDir),
0207         logLevel=acts.logging.INFO,
0208     )
0209
0210     s.run()
0211     return s
0212
0213
0214 if __name__ == "__main__":
0215     # Setup detector (ODD)
0216     detector = getOpenDataDetector()
0217     trackingGeometry = detector.trackingGeometry()
0218     decorators = detector.contextDecorators()
0219
0220     # Magnetic field
0221     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0222
0223     # Configuration files
0224     srcdir = Path(__file__).resolve().parent.parent.parent.parent
0225     geometrySelection = srcdir / "Examples/Configs/odd-seeding-config.json"
0226     assert geometrySelection.exists(), f"File not found: {geometrySelection}"
0227
0228     stripGeometrySelection = (
0229         srcdir / "Examples/Configs/odd-strip-spacepoint-selection.json"
0230     )
0231     assert srcdir.exists(), f"File not found: {stripGeometrySelection}"
0232
0233     digiConfigFile = srcdir / "Examples/Configs/odd-digi-smearing-config.json"
0234     assert digiConfigFile.exists(), f"File not found: {digiConfigFile}"
0235
0236     # Model paths from MODEL_STORAGE environment variable
0237     model_storage = os.environ.get("MODEL_STORAGE")
0238     assert model_storage is not None, "MODEL_STORAGE environment variable is not set"
0239     ci_models_odd = Path(model_storage)
0240     moduleMapPath = str(ci_models_odd / "module_map_odd_2k_events.1e-03.float")
0241     gnnModel = str(ci_models_odd / "gnn_odd_module_map.pt")
0242     outputDir = Path.cwd()
0243     events = 100
0244
0245     # Run the workflow
0246     runGnnModuleMap(
0247         trackingGeometry=trackingGeometry,
0248         field=field,
0249         geometrySelection=geometrySelection,
0250         stripGeometrySelection=stripGeometrySelection,
0251         digiConfigFile=digiConfigFile,
0252         moduleMapPath=moduleMapPath,
0253         gnnModel=gnnModel,
0254         outputDir=outputDir,
0255         events=events,
0256     )