refactor: Improve GNN module map script

Examples/Scripts/Python/gnn_module_map_odd.py

@@ -45,7 +45,6 @@ def runGnnModuleMap(
     moduleMapPath,
     gnnModel,
     outputDir,
-    events=100,
     s=None,
 ):
     """
@@ -75,7 +74,7 @@ def runGnnModuleMap(
     ).exists(), f"Module map not found: {moduleMapPath}.triplets.root"
     assert Path(gnnModel).exists(), f"Model file not found: {gnnModel}"
 
-    s = s or Sequencer(events=events, numThreads=1)
+    s = s or Sequencer(events=100, numThreads=1)
 
     # Random number generator
     rnd = acts.examples.RandomNumbers(seed=42)
@@ -240,7 +239,6 @@ def runGnnModuleMap(
     moduleMapPath = str(ci_models_odd / "module_map_odd_2k_events.1e-03.float")
     gnnModel = str(ci_models_odd / "gnn_odd_module_map.pt")
     outputDir = Path.cwd()
-    events = 100
 
     # Run the workflow
     runGnnModuleMap(
@@ -252,5 +250,4 @@ def runGnnModuleMap(
         moduleMapPath=moduleMapPath,
         gnnModel=gnnModel,
         outputDir=outputDir,
-        events=events,
     )

Python/Examples/python/reconstruction.py

@@ -1817,6 +1817,7 @@ def addTrackWriters(
     writeStates: bool = False,
     writeFitterPerformance: bool = False,
     writeFinderPerformance: bool = False,
+    writeFinderNTuple: bool = False,
     logLevel: Optional[acts.logging.Level] = None,
     writeCovMat=False,
 ):
@@ -1877,6 +1878,17 @@ def addTrackWriters(
             )
             s.addWriter(trackFinderPerfWriter)
 
+        if writeFinderNTuple:
+            trackFinderNTupleWriter = acts.examples.root.RootTrackFinderNTupleWriter(
+                level=customLogLevel(),
+                inputTracks=tracks,
+                inputParticles="particles_selected",
+                inputParticleMeasurementsMap="particle_measurements_map",
+                inputTrackParticleMatching="track_particle_matching",
+                filePath=str(Path(outputDirRoot) / f"track_finding_ntuple_{name}.root"),
+            )
+            s.addWriter(trackFinderNTupleWriter)
+
     if outputDirCsv is not None:
         outputDirCsv = Path(outputDirCsv)
         if not outputDirCsv.exists():
@@ -2024,21 +2036,15 @@ def addGnn(
         "particle_track_matching", matchAlg.config.outputParticleTrackMatching
     )
 
-    # Optional performance writer
-    if outputDirRoot is not None:
-        assert (
-            ACTS_EXAMPLES_ROOT_AVAILABLE
-        ), "ROOT output requested but ROOT is not available"
-        s.addWriter(
-            RootTrackFinderNTupleWriter(
-                level=customLogLevel(),
-                inputTracks="tracks",
-                inputParticles="particles",
-                inputParticleMeasurementsMap="particle_measurements_map",
-                inputTrackParticleMatching=matchAlg.config.outputTrackParticleMatching,
-                filePath=str(Path(outputDirRoot) / "performance_track_finding.root"),
-            )
-        )
+    addTrackWriters(
+        s,
+        name="gnn",
+        tracks="tracks",
+        outputDirRoot=outputDirRoot,
+        writeFinderPerformance=True,
+        writeFinderNTuple=True,
+        logLevel=logLevel,
+    )
 
     return s
 

Python/Examples/tests/test_examples.py

@@ -1325,9 +1325,10 @@ def test_gnn_module_map(tmp_path, assert_root_hash, backend, hardware):
         )
 
     # Verify output
-    output_file = tmp_path / "performance_track_finding.root"
-    assert output_file.exists()
-    assert_root_hash("performance_track_finding.root", output_file)
+    for f in ["performance_finding_gnn.root", "track_finding_ntuple_gnn.root"]:
+        output_file = tmp_path / f
+        assert output_file.exists()
+        assert_root_hash(f, output_file)
 
 
 @pytest.mark.odd