Examples/tests/test_colliderml_reader.py

0001 from pathlib import Path
0002
0003 import pytest
0004
0005 import acts
0006 import acts.examples
0007 from acts import UnitConstants as u
0008
0009 from helpers import arrowEnabled, dd4hepEnabled
0010
0011 pytestmark = [
0012     pytest.mark.skipif(not arrowEnabled, reason="Arrow/Parquet bindings not built"),
0013     pytest.mark.skipif(not dd4hepEnabled, reason="DD4hep not set up"),
0014 ]
0015
0016 _srcdir = Path(__file__).resolve().parent.parent.parent.parent
0017
0018 # Number of events generated by the FATRAS fixture and processed by the test.
0019 _N_EVENTS = 3
0020
0021
0022 @pytest.fixture(scope="session")
0023 def colliderml_fatras_sample(tmp_path_factory):
0024     """Generate a small ColliderML-compatible parquet dataset using FATRAS + ODD.
0025
0026     Runs a minimal sequencer: ParametricParticleGenerator → FATRAS → Digitization
0027     → ArrowParticleOutputConverter + ArrowSimHitOutputConverter → ParquetWriter.
0028
0029     Returns (base_dir, sample_name).
0030     """
0031     from acts.arrow import particleSchema, simHitSchema
0032     from acts.examples.arrow import (
0033         ArrowParticleOutputConverter,
0034         ArrowSimHitOutputConverter,
0035         ParquetWriter,
0036     )
0037     from acts.examples.json import readDigiConfigFromJson
0038     from acts.examples.odd import getOpenDataDetector
0039
0040     sample = "fatras_test"
0041     tmp = tmp_path_factory.mktemp("colliderml_fatras")
0042
0043     particles_dir = tmp / f"{sample}_particles" / "data" / f"{sample}_particles"
0044     hits_dir = tmp / f"{sample}_tracker_hits" / "data" / f"{sample}_tracker_hits"
0045     particles_dir.mkdir(parents=True)
0046     hits_dir.mkdir(parents=True)
0047
0048     rng = acts.examples.RandomNumbers(seed=42)
0049     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0050
0051     with getOpenDataDetector() as detector:
0052         tgeo = detector.trackingGeometry()
0053
0054         s = acts.examples.Sequencer(
0055             events=_N_EVENTS,
0056             numThreads=1,
0057             logLevel=acts.logging.WARNING,
0058             outputDir=str(tmp),
0059             failOnUnmaskedFpe=False,
0060         )
0061
0062         # --- Particle gun: 10 muons per event, pT=2 GeV, full eta/phi ---
0063         evGen = acts.examples.EventGenerator(
0064             level=acts.logging.WARNING,
0065             generators=[
0066                 acts.examples.EventGenerator.Generator(
0067                     multiplicity=acts.examples.FixedMultiplicityGenerator(n=1),
0068                     vertex=acts.examples.GaussianVertexGenerator(
0069                         stddev=acts.Vector4(0, 0, 0, 0),
0070                         mean=acts.Vector4(0, 0, 0, 0),
0071                     ),
0072                     particles=acts.examples.ParametricParticleGenerator(
0073                         p=(2 * u.GeV, 2 * u.GeV),
0074                         eta=(-2, 2),
0075                         phi=(0, 360 * u.degree),
0076                         pdg=acts.PdgParticle.eMuon,
0077                         randomizeCharge=True,
0078                         numParticles=10,
0079                     ),
0080                 )
0081             ],
0082             outputEvent="gun_event",
0083             randomNumbers=rng,
0084         )
0085         s.addReader(evGen)
0086
0087         hepmc3Conv = acts.examples.hepmc3.HepMC3InputConverter(
0088             level=acts.logging.WARNING,
0089             inputEvent=evGen.config.outputEvent,
0090             outputParticles="particles_gen",
0091             outputVertices="vertices_gen",
0092         )
0093         s.addAlgorithm(hepmc3Conv)
0094
0095         # --- FATRAS simulation ---
0096         fatrasAlg = acts.examples.FatrasSimulation(
0097             level=acts.logging.WARNING,
0098             inputParticles="particles_gen",
0099             outputParticles="particles_sim",
0100             outputSimHits="simhits",
0101             randomNumbers=rng,
0102             trackingGeometry=tgeo,
0103             magneticField=field,
0104             generateHitsOnSensitive=True,
0105             emScattering=False,
0106             emEnergyLossIonisation=False,
0107             emEnergyLossRadiation=False,
0108             emPhotonConversion=False,
0109         )
0110         s.addAlgorithm(fatrasAlg)
0111
0112         # --- Digitization (needed for valid x/y/z cluster positions) ---
0113         digiCfg = acts.examples.DigitizationAlgorithm.Config(
0114             digitizationConfigs=readDigiConfigFromJson(
0115                 str(_srcdir / "Examples/Configs/odd-digi-smearing-config-notime.json")
0116             ),
0117             surfaceByIdentifier=tgeo.geoIdSurfaceMap(),
0118             randomNumbers=rng,
0119             inputSimHits="simhits",
0120         )
0121         s.addAlgorithm(
0122             acts.examples.DigitizationAlgorithm(digiCfg, acts.logging.WARNING)
0123         )
0124
0125         # --- Arrow output converters ---
0126         s.addAlgorithm(
0127             ArrowParticleOutputConverter(
0128                 level=acts.logging.WARNING,
0129                 inputParticles="particles_sim",
0130                 outputTable="cml_particles_arrow",
0131             )
0132         )
0133         s.addAlgorithm(
0134             ArrowSimHitOutputConverter(
0135                 level=acts.logging.WARNING,
0136                 inputSimHits="simhits",
0137                 inputParticles="particles_sim",
0138                 inputClusters="clusters",
0139                 inputSimHitMeasurementsMap="simhit_measurements_map",
0140                 outputTable="cml_hits_arrow",
0141             )
0142         )
0143
0144         # --- ParquetWriter: absolute paths so files land in the nested layout ---
0145         s.addWriter(
0146             ParquetWriter(
0147                 level=acts.logging.WARNING,
0148                 outputDir=str(tmp),
0149                 collections={
0150                     "cml_particles_arrow": str(particles_dir),
0151                     "cml_hits_arrow": str(hits_dir),
0152                 },
0153                 expectedSchemas={
0154                     "cml_particles_arrow": particleSchema(),
0155                     "cml_hits_arrow": simHitSchema(),
0156                 },
0157                 eventsPerShard=_N_EVENTS,
0158             )
0159         )
0160
0161         s.run()
0162
0163     return tmp, sample
0164
0165
0166 @pytest.mark.parametrize("reco_geo", ["gen1", "gen3"])
0167 def test_colliderml_truth_tracking(
0168     tmp_path, assert_root_hash, colliderml_fatras_sample, reco_geo
0169 ):
0170     """Read FATRAS-generated ColliderML-format data and run truth-seeded KF tracking.
0171
0172     Parametrized on the reconstruction geometry:
0173       gen1 — same Gen1 ODD used for simulation; no geo-id map needed.
0174       gen3 — Gen3 ODD; geo-id map and digi config generated on the fly.
0175
0176     Verifies that the ColliderML reader pipeline (ParquetReader +
0177     ColliderMLRelease1InputConverter + TruthEstimated seeding + KF) runs without
0178     error and produces ROOT output files with the expected content hashes.
0179     """
0180     import sys
0181
0182     sys.path.insert(0, str(_srcdir / "Examples/Scripts/Python"))
0183     from acts.examples.odd import getOpenDataDetector
0184     from colliderml_truth_tracking import runColliderMLTruthTracking
0185     from generate_geoid_map import generate_geoid_map
0186
0187     inputDir, sample = colliderml_fatras_sample
0188     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0189
0190     if reco_geo == "gen3":
0191         # Build both geometries to generate the mapping on the fly
0192         detector_gen1 = getOpenDataDetector()
0193         geo_gen1 = detector_gen1.trackingGeometry()
0194
0195         detector_gen3 = getOpenDataDetector(gen3=True)
0196         trackingGeometry = detector_gen3.trackingGeometry()
0197         decorators = detector_gen3.contextDecorators()
0198
0199         geoid_map_path = tmp_path / "geoid_map.csv"
0200         generate_geoid_map(
0201             geo_gen1,
0202             trackingGeometry,
0203             output_path=geoid_map_path,
0204             prefix_a="gen1",
0205             prefix_b="gen3",
0206         )
0207
0208         ctx = None
0209     else:
0210         odd_dir = acts.examples.odd.getOpenDataDetectorDirectory()
0211         matDeco = acts.IMaterialDecorator.fromFile(
0212             odd_dir / "data/odd-material-maps.root", level=acts.logging.WARNING
0213         )
0214         detector = getOpenDataDetector(matDeco)
0215         trackingGeometry = detector.trackingGeometry()
0216         decorators = detector.contextDecorators()
0217         ctx = detector
0218
0219     def _run():
0220         s, _perf_proto, _perf_kf = runColliderMLTruthTracking(
0221             trackingGeometry=trackingGeometry,
0222             field=field,
0223             outputDir=tmp_path,
0224             inputDir=inputDir,
0225             geoIdMapPath=geoid_map_path if reco_geo == "gen3" else None,
0226             decorators=decorators,
0227             events=_N_EVENTS,
0228             numThreads=1,
0229             sample=sample,
0230         )
0231         s.run()
0232
0233     if ctx is not None:
0234         with ctx:
0235             _run()
0236     else:
0237         _run()
0238
0239     root_files = [
0240         "trackstates_kf.root",
0241         "tracksummary_kf.root",
0242         "performance_kf.root",
0243     ]
0244
0245     for fn in root_files:
0246         fp = tmp_path / fn
0247         assert fp.exists(), f"{fn} was not produced"
0248         assert fp.stat().st_size > 1024, f"{fn} is suspiciously small"
0249         assert_root_hash(fn, fp)