ci: Reenable some ExaTrkX tests (#3395)

These tests should work even though cugraph is not available anymore...  ## Summary by CodeRabbit - **New Features** - Introduced new test jobs for enhanced testing capabilities. - Added a new algorithm for converting prototracks to tracks, improving data processing. - **Bug Fixes** - Updated input parameters for existing algorithms to ensure consistency and accuracy in data flow. - **Chores** - Improved cache management in the CI pipeline for better build efficiency. - Updated hash values for test files to reflect recent changes.
acts-project · Dec 6, 2024 · 12ea68c · 12ea68c
1 parent 2862173
commit 12ea68c
Show file tree

Hide file tree

Showing 5 changed files with 102 additions and 63 deletions.
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
@@ -142,33 +142,46 @@ build_exatrkx:
     - cmake --build build -- -j6
     - ccache -s
 
-# test_exatrkx_unittests:
-#   stage: test
-#   needs:
-#     - build_exatrkx
-#   image: ghcr.io/acts-project/ubuntu2204_exatrkx:63
-#   tags:
-#     - docker-gpu-nvidia
-#   script:
-#     - ctest --test-dir build -R ExaTrkX
-#
-# test_exatrkx_python:
-#   stage: test
-#   needs:
-#     - build_exatrkx
-#   image: ghcr.io/acts-project/ubuntu2204_exatrkx:63
-#   tags:
-#     - docker-gpu-nvidia
-#   script:
-#     - apt-get update -y
-#     - apt-get install -y python3 libxxhash0
-#     - source build/this_acts_withdeps.sh
-#     - git clone $CLONE_URL src
-#     - cd src
-#     - git checkout $HEAD_SHA
-#     - pip3 install -r Examples/Python/tests/requirements.txt
-#     - nvidia-smi
-#     - pytest -rFsv -k test_exatrkx
+test_exatrkx_unittests:
+  stage: test
+  needs:
+    - build_exatrkx
+  image: ghcr.io/acts-project/ubuntu2204_exatrkx:63
+  variables:
+    DEPENDENCY_URL: https://acts.web.cern.ch/ACTS/ci/ubuntu-22.04/deps.$DEPENDENCY_TAG.tar.zst
+  tags:
+    - docker-gpu-nvidia
+  script:
+
+    - apt-get update -y
+    - git clone $CLONE_URL src
+    - source src/CI/dependencies.sh
+    - ctest --test-dir build -R ExaTrkX
+
+test_exatrkx_python:
+  stage: test
+  needs:
+    - build_exatrkx
+  image: ghcr.io/acts-project/ubuntu2204_exatrkx:63
+  variables:
+    DEPENDENCY_URL: https://acts.web.cern.ch/ACTS/ci/ubuntu-22.04/deps.$DEPENDENCY_TAG.tar.zst
+  tags:
+    - docker-gpu-nvidia
+  script:
+    - apt-get update -y
+    - git clone $CLONE_URL src
+    - cd src
+    - git checkout $HEAD_SHA
+    - nvidia-smi
+    - source CI/dependencies.sh
+    - source ../build/this_acts_withdeps.sh
+    - python3 -m pip install -r Examples/Python/tests/requirements.txt
+    - echo $PYTHONPATH
+    - which python3
+    - python3 --version
+    - python3 -c "import acts"
+    - pytest -rFsv -k torch --collect-only
+    - pytest -rFsv -k gpu-torch # For now only test torch GPU pipeline
 
 build_linux_ubuntu:
   stage: build

diff --git a/Examples/Algorithms/TrackFindingExaTrkX/CMakeLists.txt b/Examples/Algorithms/TrackFindingExaTrkX/CMakeLists.txt
@@ -5,6 +5,7 @@ add_library(
     src/PrototracksToParameters.cpp
     src/TrackFindingFromPrototrackAlgorithm.cpp
     src/TruthGraphBuilder.cpp
+    src/createFeatures.cpp
 )
 
 target_include_directories(

diff --git a/Examples/Python/python/acts/examples/reconstruction.py b/Examples/Python/python/acts/examples/reconstruction.py
@@ -1786,13 +1786,12 @@ def addExaTrkX(
 
     if backend == ExaTrkXBackend.Torch:
         metricLearningConfig["modelPath"] = str(modelDir / "embed.pt")
-        metricLearningConfig["numFeatures"] = 3
+        metricLearningConfig["selectedFeatures"] = [0, 1, 2]
         filterConfig["modelPath"] = str(modelDir / "filter.pt")
-        filterConfig["nChunks"] = 10
-        filterConfig["numFeatures"] = 3
+        filterConfig["selectedFeatures"] = [0, 1, 2]
         gnnConfig["modelPath"] = str(modelDir / "gnn.pt")
         gnnConfig["undirected"] = True
-        gnnConfig["numFeatures"] = 3
+        gnnConfig["selectedFeatures"] = [0, 1, 2]
 
         graphConstructor = acts.examples.TorchMetricLearning(**metricLearningConfig)
         edgeClassifiers = [
@@ -1824,11 +1823,18 @@ def addExaTrkX(
     s.addAlgorithm(findingAlg)
     s.addWhiteboardAlias("prototracks", findingAlg.config.outputProtoTracks)
 
-    # TODO convert prototracks to tracks
+    s.addAlgorithm(
+        acts.examples.PrototracksToTracks(
+            level=customLogLevel(),
+            inputProtoTracks="prototracks",
+            inputMeasurements="measurements",
+            outputTracks="tracks",
+        )
+    )
 
     matchAlg = acts.examples.TrackTruthMatcher(
         level=customLogLevel(),
-        inputProtoTracks=findingAlg.config.outputProtoTracks,
+        inputTracks="tracks",
         inputParticles="particles",
         inputMeasurementParticlesMap="measurement_particles_map",
         outputTrackParticleMatching="exatrkx_track_particle_matching",
@@ -1843,14 +1849,12 @@ def addExaTrkX(
         "particle_track_matching", matchAlg.config.outputParticleTrackMatching
     )
 
-    # Write truth track finding / seeding performance
     if outputDirRoot is not None:
         s.addWriter(
             acts.examples.TrackFinderNTupleWriter(
                 level=customLogLevel(),
-                inputProtoTracks=findingAlg.config.outputProtoTracks,
-                # the original selected particles after digitization
-                inputParticles="particles_initial",
+                inputTracks="tracks",
+                inputParticles="particles",
                 inputParticleMeasurementsMap="particle_measurements_map",
                 inputTrackParticleMatching=matchAlg.config.outputTrackParticleMatching,
                 filePath=str(Path(outputDirRoot) / "performance_track_finding.root"),

diff --git a/Examples/Python/tests/root_file_hashes.txt b/Examples/Python/tests/root_file_hashes.txt
@@ -72,7 +72,7 @@ test_root_clusters_writer[configKwConstructor]__clusters.root: e842df4fe04eefff3
 test_root_clusters_writer[kwargsConstructor]__clusters.root: e842df4fe04eefff3df5f32cd1026e93286be62b8040dc700a2aff557c56dec8
 test_exatrkx[cpu-torch]__performance_track_finding.root: 36b3045589c4c17c038dbc87943366f4af4440f7eea6887afb763871ac149b05
 test_exatrkx[gpu-onnx]__performance_track_finding.root: 9090de10ffb1489d3f1993e2a3081a3038227e3e5c453e98a9a4f33ea3d6d817
-test_exatrkx[gpu-torch]__performance_track_finding.root: 36b3045589c4c17c038dbc87943366f4af4440f7eea6887afb763871ac149b05
+test_exatrkx[gpu-torch]__performance_track_finding.root: 6b658fa22c7532e082eaab7aa4b71b852f1c324adcc59d1156aff45124b222d9
 test_ML_Ambiguity_Solver__performance_finding_ambiML.root: 166dd8bb189097c4957b7b02c04c41267868d72d9a08c4bb892985b06849cb76
 test_refitting[odd]__trackstates_gsf_refit.root: e297749dc1e7eda3b8dea13defa0499986c584740d93e723a901b498b8e90c71
 test_refitting[odd]__tracksummary_gsf_refit.root: d5085882e45a0b699194dff9f40a36e9291227bf65f9aaaf9087f9242ef5ae22

diff --git a/Examples/Scripts/Python/exatrkx.py b/Examples/Scripts/Python/exatrkx.py
@@ -1,16 +1,52 @@
 #!/usr/bin/env python3
 from pathlib import Path
+import os
+import sys
 
 import acts.examples
 import acts
+from acts.examples.reconstruction import addExaTrkX, ExaTrkXBackend
 from acts import UnitConstants as u
 
+from digitization import runDigitization
+
+
+def runGNNTrackFinding(
+    trackingGeometry,
+    field,
+    outputDir,
+    digiConfigFile,
+    geometrySelection,
+    backend,
+    modelDir,
+    outputRoot=False,
+    outputCsv=False,
+    s=None,
+):
+    s = runDigitization(
+        trackingGeometry,
+        field,
+        outputDir,
+        digiConfigFile=digiConfigFile,
+        particlesInput=None,
+        outputRoot=outputRoot,
+        outputCsv=outputCsv,
+        s=s,
+    )
+
+    addExaTrkX(
+        s,
+        trackingGeometry,
+        geometrySelection,
+        modelDir,
+        backend=backend,
+        outputDirRoot=outputDir if outputRoot else None,
+    )
+
+    s.run()
+
 
 if "__main__" == __name__:
-    import os
-    import sys
-    from digitization import runDigitization
-    from acts.examples.reconstruction import addExaTrkX, ExaTrkXBackend
 
     backend = ExaTrkXBackend.Torch
 
@@ -19,16 +55,10 @@
     if "torch" in sys.argv:
         backend = ExaTrkXBackend.Torch
 
-    srcdir = Path(__file__).resolve().parent.parent.parent.parent
-
     detector, trackingGeometry, decorators = acts.examples.GenericDetector.create()
 
     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
 
-    inputParticlePath = Path("particles.root")
-    if not inputParticlePath.exists():
-        inputParticlePath = None
-
     srcdir = Path(__file__).resolve().parent.parent.parent.parent
 
     geometrySelection = (
@@ -60,24 +90,15 @@
     rnd = acts.examples.RandomNumbers()
     outputDir = Path(os.getcwd())
 
-    s = runDigitization(
+    runGNNTrackFinding(
         trackingGeometry,
         field,
         outputDir,
-        digiConfigFile=digiConfigFile,
-        particlesInput=inputParticlePath,
-        outputRoot=True,
-        outputCsv=True,
-        s=s,
-    )
-
-    addExaTrkX(
-        s,
-        trackingGeometry,
+        digiConfigFile,
         geometrySelection,
+        backend,
         modelDir,
-        outputDir,
-        backend=backend,
+        outputRoot=True,
+        outputCsv=False,
+        s=s,
     )
-
-    s.run()