feat: addCKFTracks option to disable some output (#1447)

`addCKFTracks(writeTrajectories=False)` disables writing potentially enormous `trackstates_ckf.root` and `tracksummary_ckf.root` ntuples, while still producing histograms in `performance_ckf.root`.
acts-project · Aug 19, 2022 · 48f2993 · 48f2993
1 parent 9587b8d
commit 48f2993
Showing 1 changed file with 35 additions and 31 deletions.
diff --git a/Examples/Python/python/acts/examples/reconstruction.py b/Examples/Python/python/acts/examples/reconstruction.py
@@ -645,6 +645,7 @@ def addCKFTracks(
     outputDirCsv: Optional[Union[Path, str]] = None,
     outputDirRoot: Optional[Union[Path, str]] = None,
     selectedParticles: str = "truth_seeds_selected",
+    writeTrajectories: bool = True,
 ) -> None:
     """This function steers the seeding
 
@@ -663,6 +664,8 @@ def addCKFTracks(
         the output folder for the Root output, None triggers no output
     selectedParticles : str, "truth_seeds_selected"
         CKFPerformanceWriter truth input
+    writeTrajectories : bool, True
+        write trackstates_ckf.root and tracksummary_ckf.root ntuples? These can be quite large.
     """
 
     if int(s.config.logLevel) <= int(acts.logging.DEBUG):
@@ -692,37 +695,38 @@ def addCKFTracks(
         if not outputDirRoot.exists():
             outputDirRoot.mkdir()
 
-        # write track states from CKF
-        trackStatesWriter = acts.examples.RootTrajectoryStatesWriter(
-            level=s.config.logLevel,
-            inputTrajectories=trackFinder.config.outputTrajectories,
-            # @note The full particles collection is used here to avoid lots of warnings
-            # since the unselected CKF track might have a majority particle not in the
-            # filtered particle collection. This could be avoided when a seperate track
-            # selection algorithm is used.
-            inputParticles="particles_selected",
-            inputSimHits="simhits",
-            inputMeasurementParticlesMap="measurement_particles_map",
-            inputMeasurementSimHitsMap="measurement_simhits_map",
-            filePath=str(outputDirRoot / "trackstates_ckf.root"),
-            treeName="trackstates",
-        )
-        s.addWriter(trackStatesWriter)
-
-        # write track summary from CKF
-        trackSummaryWriter = acts.examples.RootTrajectorySummaryWriter(
-            level=s.config.logLevel,
-            inputTrajectories=trackFinder.config.outputTrajectories,
-            # @note The full particles collection is used here to avoid lots of warnings
-            # since the unselected CKF track might have a majority particle not in the
-            # filtered particle collection. This could be avoided when a seperate track
-            # selection algorithm is used.
-            inputParticles="particles_selected",
-            inputMeasurementParticlesMap="measurement_particles_map",
-            filePath=str(outputDirRoot / "tracksummary_ckf.root"),
-            treeName="tracksummary",
-        )
-        s.addWriter(trackSummaryWriter)
+        if writeTrajectories:
+            # write track states from CKF
+            trackStatesWriter = acts.examples.RootTrajectoryStatesWriter(
+                level=s.config.logLevel,
+                inputTrajectories=trackFinder.config.outputTrajectories,
+                # @note The full particles collection is used here to avoid lots of warnings
+                # since the unselected CKF track might have a majority particle not in the
+                # filtered particle collection. This could be avoided when a seperate track
+                # selection algorithm is used.
+                inputParticles="particles_selected",
+                inputSimHits="simhits",
+                inputMeasurementParticlesMap="measurement_particles_map",
+                inputMeasurementSimHitsMap="measurement_simhits_map",
+                filePath=str(outputDirRoot / "trackstates_ckf.root"),
+                treeName="trackstates",
+            )
+            s.addWriter(trackStatesWriter)
+
+            # write track summary from CKF
+            trackSummaryWriter = acts.examples.RootTrajectorySummaryWriter(
+                level=s.config.logLevel,
+                inputTrajectories=trackFinder.config.outputTrajectories,
+                # @note The full particles collection is used here to avoid lots of warnings
+                # since the unselected CKF track might have a majority particle not in the
+                # filtered particle collection. This could be avoided when a seperate track
+                # selection algorithm is used.
+                inputParticles="particles_selected",
+                inputMeasurementParticlesMap="measurement_particles_map",
+                filePath=str(outputDirRoot / "tracksummary_ckf.root"),
+                treeName="tracksummary",
+            )
+            s.addWriter(trackSummaryWriter)
 
         # Write CKF performance data
         ckfPerfWriter = acts.examples.CKFPerformanceWriter(