ci: Reenable some ExaTrkX tests

.gitlab-ci.yml

@@ -142,33 +142,42 @@ 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:
+    - source 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
+    - apt-get install -y python3 libxxhash0
+    - git clone $CLONE_URL src
+    - cd src
+    - git checkout $HEAD_SHA
+    - pip3 install -r Examples/Python/tests/requirements.txt
+    - nvidia-smi
+    - source CI/dependencies.sh
+    - source build/this_acts_withdeps.sh
+    - source build/python/setup.sh
+    - echo "PYTHONPATH: $PYTHONPATH"
+    - pytest -rFsv -k torch --collect-only
+    - pytest -rFsv -k gpu-torch # For now only test torch GPU pipeline
 
 build_linux_ubuntu:
   stage: build

Examples/Python/python/acts/examples/reconstruction.py

@@ -1792,13 +1792,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 = [
@@ -1830,11 +1829,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",
@@ -1854,7 +1860,7 @@ def addExaTrkX(
         s.addWriter(
             acts.examples.TrackFinderNTupleWriter(
                 level=customLogLevel(),
-                inputProtoTracks=findingAlg.config.outputProtoTracks,
+                inputTracks="tracks",
                 # the original selected particles after digitization
                 inputParticles="particles_initial",
                 inputMeasurementParticlesMap="measurement_particles_map",

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()