feat: Add truth tracking python example + test (#1027)

This PR adds a python example for truth tracking and corresponding tests.

I tried to model the functionality of the truth tracking example executables, but I would hope a reviewer could check if I missed anything important.

Examples/Python/tests/root_file_hashes.txt

@@ -40,6 +40,10 @@ test_root_clusters_writer[configKwConstructor]__clusters.root: 7e452af7243d282dd
 test_root_clusters_writer[kwargsConstructor]__clusters.root: 7e452af7243d282dd0a8f5aa2844e150ef44364980bf3641718899068a1a1ecb
 test_root_material_writer__material.root: e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
 
+test_truth_tracking__trackstates_fitter.root: d01b99db16c353bb7224ca51378da4c3a21041054a769ee09a7588d0083bbd50
+test_truth_tracking__tracksummary_fitter.root: 662959c1b8dbdcf8efe9d512bf36ac0332f33ba1ff785676a97b700b0b84495f
+test_truth_tracking__performance_track_finder.root: 76a990d595b6e097da2bed447783bd63044956e5649a5dd6fd7a6a3434786877
+
 # these are currently broken
 test_digitization_example_input__measurements.root: 9e4293264448809926c7011347f66ba5df5d37adbf31e0e7157627cb37d1fef5
 test_digitization_example_input__particles.root: 78a89f365177423d0834ea6f1bd8afe1488e72b12a25066a20bd9050f5407860

Examples/Python/tests/test_examples.py

@@ -279,6 +279,50 @@ def test_event_recording(tmp_path):
     assert alg.events_seen == 1
 
 
+@pytest.mark.skipif(not dd4hepEnabled, reason="DD4hep not set up")
+def test_truth_tracking(tmp_path, assert_root_hash):
+    from truth_tracking import runTruthTracking
+
+    detector, trackingGeometry, _ = getOpenDataDetector()
+
+    field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
+
+    seq = Sequencer(events=10, numThreads=1)
+
+    root_files = [
+        ("trackstates_fitter.root", "trackstates", 19),
+        ("tracksummary_fitter.root", "tracksummary", 10),
+        ("performance_track_finder.root", "track_finder_tracks", 19),
+        ("performance_track_fitter.root", None, -1),
+    ]
+
+    for fn, _, _ in root_files:
+        fp = tmp_path / fn
+        assert not fp.exists()
+
+    runTruthTracking(
+        trackingGeometry,
+        field,
+        digiConfigFile=Path(
+            "thirdparty/OpenDataDetector/config/odd-digi-smearing-config.json",
+        ),
+        outputDir=tmp_path,
+        s=seq,
+    )
+
+    seq.run()
+
+    del seq
+
+    for fn, tn, ee in root_files:
+        fp = tmp_path / fn
+        assert fp.exists()
+        assert fp.stat().st_size > 1024
+        if tn is not None:
+            assert_entries(fp, tn, ee)
+            assert_root_hash(fn, fp)
+
+
 def test_particle_gun(tmp_path, assert_root_hash):
     from particle_gun import runParticleGun
 

Examples/Python/tests/test_writer.py

@@ -458,6 +458,40 @@ def test_json_material_writer(tmp_path, fmt):
     assert out.stat().st_size > 1000
 
 
+@pytest.mark.csv
+def test_csv_multitrajectory_writer(tmp_path):
+    detector, trackingGeometry, decorators = GenericDetector.create()
+    field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
+
+    from truth_tracking import runTruthTracking
+
+    s = Sequencer(numThreads=1, events=10)
+    runTruthTracking(
+        trackingGeometry,
+        field,
+        digiConfigFile=Path(
+            "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json"
+        ),
+        outputDir=tmp_path,
+        s=s,
+    )
+
+    csv_dir = tmp_path / "csv"
+    csv_dir.mkdir()
+    s.addWriter(
+        CsvMultiTrajectoryWriter(
+            level=acts.logging.INFO,
+            inputTrajectories="trajectories",
+            inputMeasurementParticlesMap="measurement_particles_map",
+            outputDir=str(csv_dir),
+        )
+    )
+    s.run()
+    del s
+    assert len([f for f in csv_dir.iterdir() if f.is_file()]) == 10
+    assert all(f.stat().st_size > 20 for f in csv_dir.iterdir())
+
+
 @pytest.fixture(scope="session")
 def hepmc_data_impl(tmp_path_factory):
 

Examples/Scripts/Python/truth_tracking.py

@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+import os
+from pathlib import Path
+
+import acts
+import acts.examples
+
+u = acts.UnitConstants
+
+
+def runTruthTracking(
+    trackingGeometry,
+    field,
+    outputDir: Path,
+    digiConfigFile: Path,
+    directNavigation=False,
+    s: acts.examples.Sequencer = None,
+):
+
+    # Preliminaries
+    rnd = acts.examples.RandomNumbers()
+
+    # Sequencer
+    s = s or acts.examples.Sequencer(
+        events=10, numThreads=-1, logLevel=acts.logging.INFO
+    )
+
+    # Input
+    vtxGen = acts.examples.GaussianVertexGenerator()
+    vtxGen.stddev = acts.Vector4(0, 0, 0, 0)
+
+    ptclGen = acts.examples.ParametricParticleGenerator(
+        p=(1 * u.GeV, 10 * u.GeV), eta=(-2, 2), numParticles=2
+    )
+
+    g = acts.examples.EventGenerator.Generator()
+    g.multiplicity = acts.examples.FixedMultiplicityGenerator()
+    g.vertex = vtxGen
+    g.particles = ptclGen
+
+    evGen = acts.examples.EventGenerator(
+        level=acts.logging.INFO,
+        generators=[g],
+        outputParticles="particles_input",
+        randomNumbers=rnd,
+    )
+    s.addReader(evGen)
+
+    s.addWriter(
+        acts.examples.RootParticleWriter(
+            level=acts.logging.INFO,
+            inputParticles=evGen.config.outputParticles,
+            filePath=str(outputDir / "particles.root"),
+        )
+    )
+
+    # Selector
+    selector = acts.examples.ParticleSelector(
+        level=acts.logging.INFO,
+        inputParticles=evGen.config.outputParticles,
+        outputParticles="particles_selected",
+    )
+    s.addAlgorithm(selector)
+
+    # Simulation
+    simAlg = acts.examples.FatrasSimulation(
+        level=acts.logging.INFO,
+        inputParticles=selector.config.outputParticles,
+        outputParticlesInitial="particles_initial",
+        outputParticlesFinal="particles_final",
+        outputSimHits="simhits",
+        randomNumbers=rnd,
+        trackingGeometry=trackingGeometry,
+        magneticField=field,
+        generateHitsOnSensitive=True,
+    )
+    s.addAlgorithm(simAlg)
+
+    # Digitization
+    digiCfg = acts.examples.DigitizationConfig(
+        acts.examples.readDigiConfigFromJson(
+            str(digiConfigFile),
+        ),
+        trackingGeometry=trackingGeometry,
+        randomNumbers=rnd,
+        inputSimHits=simAlg.config.outputSimHits,
+    )
+    digiAlg = acts.examples.DigitizationAlgorithm(digiCfg, acts.logging.INFO)
+    s.addAlgorithm(digiAlg)
+
+    selAlg = acts.examples.TruthSeedSelector(
+        level=acts.logging.INFO,
+        ptMin=500 * u.MeV,
+        nHitsMin=9,
+        inputParticles=simAlg.config.outputParticlesInitial,
+        inputMeasurementParticlesMap=digiCfg.outputMeasurementParticlesMap,
+        outputParticles="particles_seed_selected",
+    )
+    s.addAlgorithm(selAlg)
+
+    inputParticles = selAlg.config.outputParticles
+
+    smearAlg = acts.examples.ParticleSmearing(
+        level=acts.logging.INFO,
+        inputParticles=inputParticles,
+        outputTrackParameters="smearedparameters",
+        randomNumbers=rnd,
+        # Gaussian sigmas to smear particle parameters
+        # sigmaD0=20 * u.um,
+        # sigmaD0PtA=30 * u.um,
+        # sigmaD0PtB=0.3 / u.GeV,
+        # sigmaZ0=20 * u.um,
+        # sigmaZ0PtA=30 * u.um,
+        # sigmaZ0PtB=0.3 / u.GeV,
+        # sigmaPhi=1 * u.degree,
+        # sigmaTheta=1 * u.degree,
+        # sigmaPRel=0.01,
+        # sigmaT0=1 * u.ns,
+        # initialVarInflation=[1, 1, 1, 1, 1, 1],
+    )
+    s.addAlgorithm(smearAlg)
+
+    truthTrkFndAlg = acts.examples.TruthTrackFinder(
+        level=acts.logging.INFO,
+        inputParticles=inputParticles,
+        inputMeasurementParticlesMap=digiAlg.config.outputMeasurementParticlesMap,
+        outputProtoTracks="prototracks",
+    )
+    s.addAlgorithm(truthTrkFndAlg)
+
+    if directNavigation:
+        srfSortAlg = acts.examples.SurfaceSortingAlgorithm(
+            level=acts.logging.INFO,
+            inputProtoTracks=truthTrkFndAlg.config.outputProtoTracks,
+            inputSimulatedHits=simAlg.config.outputSimHits,
+            inputMeasurementSimHitsMap=digiAlg.config.outputMeasurementSimHitsMap,
+            outputProtoTracks="sortedprototracks",
+        )
+        s.addAlgorithm(srfSortAlg)
+        inputProtoTracks = srfSortAlg.config.outputProtoTracks
+    else:
+        inputProtoTracks = truthTrkFndAlg.config.outputProtoTracks
+
+    fitAlg = acts.examples.TrackFittingAlgorithm(
+        level=acts.logging.INFO,
+        inputMeasurements=digiAlg.config.outputMeasurements,
+        inputSourceLinks=digiAlg.config.outputSourceLinks,
+        inputProtoTracks=inputProtoTracks,
+        inputInitialTrackParameters=smearAlg.config.outputTrackParameters,
+        outputTrajectories="trajectories",
+        directNavigation=directNavigation,
+        multipleScattering=True,
+        energyLoss=True,
+        pickTrack=-1,
+        trackingGeometry=trackingGeometry,
+        dFit=acts.examples.TrackFittingAlgorithm.makeTrackFitterFunction(field),
+        fit=acts.examples.TrackFittingAlgorithm.makeTrackFitterFunction(
+            trackingGeometry, field
+        ),
+    )
+    s.addAlgorithm(fitAlg)
+
+    # Output
+
+    s.addWriter(
+        acts.examples.RootTrajectoryStatesWriter(
+            level=acts.logging.INFO,
+            inputTrajectories=fitAlg.config.outputTrajectories,
+            inputParticles=inputParticles,
+            inputSimHits=simAlg.config.outputSimHits,
+            inputMeasurementParticlesMap=digiAlg.config.outputMeasurementParticlesMap,
+            inputMeasurementSimHitsMap=digiAlg.config.outputMeasurementSimHitsMap,
+            filePath=str(outputDir / "trackstates_fitter.root"),
+        )
+    )
+
+    s.addWriter(
+        acts.examples.RootTrajectorySummaryWriter(
+            level=acts.logging.INFO,
+            inputTrajectories=fitAlg.config.outputTrajectories,
+            inputParticles=inputParticles,
+            inputMeasurementParticlesMap=digiAlg.config.outputMeasurementParticlesMap,
+            filePath=str(outputDir / "tracksummary_fitter.root"),
+        )
+    )
+
+    s.addWriter(
+        acts.examples.TrackFinderPerformanceWriter(
+            level=acts.logging.INFO,
+            inputProtoTracks=truthTrkFndAlg.config.outputProtoTracks,
+            inputParticles=inputParticles,
+            inputMeasurementParticlesMap=digiAlg.config.outputMeasurementParticlesMap,
+            filePath=str(outputDir / "performance_track_finder.root"),
+        )
+    )
+
+    s.addWriter(
+        acts.examples.TrackFitterPerformanceWriter(
+            level=acts.logging.INFO,
+            inputTrajectories=fitAlg.config.outputTrajectories,
+            inputParticles=inputParticles,
+            inputMeasurementParticlesMap=digiAlg.config.outputMeasurementParticlesMap,
+            filePath=str(outputDir / "performance_track_fitter.root"),
+        )
+    )
+
+    return s
+
+
+if "__main__" == __name__:
+
+    # detector, trackingGeometry, _ = getOpenDataDetector()
+    detector, trackingGeometry, decorators = acts.examples.GenericDetector.create()
+
+    field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
+
+    runTruthTracking(
+        trackingGeometry,
+        field,
+        digiConfigFile=Path(
+            # "thirdparty/OpenDataDetector/config/odd-digi-smearing-config.json",
+            "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json"
+        ),
+        outputDir=Path.cwd(),
+    ).run()