[WIP] Simple track merging

Tracking/components/ExtractTrackParams.cpp

@@ -0,0 +1,140 @@
+#include "utils.hpp"
+
+#include <edm4hep/Constants.h>
+#include <edm4hep/TrackCollection.h>
+
+#include <podio/UserDataCollection.h>
+
+#include "k4FWCore/Transformer.h"
+
+#include "Gaudi/Property.h"
+
+#include <optional>
+#include <string>
+#include <tuple>
+
+// Which type of collection we are reading
+using FloatColl = podio::UserDataCollection<float>;
+using TrackColl = edm4hep::TrackCollection;
+using TS = edm4hep::TrackState;
+using TP = edm4hep::TrackParams;
+using edm4hep::utils::detail::to_index;
+
+struct TrackParamExtractor final
+    : k4FWCore::MultiTransformer<std::tuple<FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl,
+                                            FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl,
+                                            FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl>(
+          const TrackColl&, const TrackColl&)> {
+  TrackParamExtractor(const std::string& name, ISvcLocator* svcLoc)
+      : MultiTransformer(name, svcLoc,
+                         {
+                             KeyValues("InputSiTracks", {"SiTracks"}),
+                             KeyValues("InputCluTracks", {"ClupatraTracks"}),
+                         },
+                         {
+                             KeyValues("OutCollSiD0", {"SiTrackD0"}),
+                             KeyValues("OutCollSiPhi", {"SiTrackPhi"}),
+                             KeyValues("OutCollSiOmega", {"SiTrackOmega"}),
+                             KeyValues("OutCollSiZ0", {"SiTrackZ0"}),
+                             KeyValues("OutCollSiTanL", {"SiTrackTanL"}),
+                             KeyValues("OutCollCluD0", {"CluTrackD0"}),
+                             KeyValues("OutCollCluPhi", {"CluTrackPhi"}),
+                             KeyValues("OutCollCluOmega", {"CluTrackOmega"}),
+                             KeyValues("OutCollCluZ0", {"CluTrackZ0"}),
+                             KeyValues("OutCollCluTanL", {"CluTrackTanL"}),
+                             // uncertainties
+                             KeyValues("OutCollSiUncD0", {"SiTrackUncD0"}),
+                             KeyValues("OutCollSiUncPhi", {"SiTrackUncPhi"}),
+                             KeyValues("OutCollSiUncOmega", {"SiTrackUncOmega"}),
+                             KeyValues("OutCollSiUncZ0", {"SiTrackUncZ0"}),
+                             KeyValues("OutCollSiUncTanL", {"SiTrackUncTanL"}),
+                             KeyValues("OutCollCluUncD0", {"CluTrackUncD0"}),
+                             KeyValues("OutCollCluUncPhi", {"CluTrackUncPhi"}),
+                             KeyValues("OutCollCluUncOmega", {"CluTrackUncOmega"}),
+                             KeyValues("OutCollCluUncZ0", {"CluTrackUncZ0"}),
+                             KeyValues("OutCollCluUncTanL", {"CluTrackUncTanL"}),
+                         }) {}
+
+  // This is the function that will be called to transform the data
+  // Note that the function has to be const, as well as the collections
+  // we get from the input
+  std::tuple<FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl,
+             FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl, FloatColl,
+             FloatColl, FloatColl>
+  operator()(const TrackColl& inSiTracks, const TrackColl& inCluTracks) const override {
+
+    printInStars(this, "New Event", n_stars);
+
+    debug() << "Received SiTracks collection with " << inSiTracks.size() << " tracks" << endmsg;
+    debug() << "Received ClupatraTracks collection with " << inCluTracks.size() << " tracks" << endmsg;
+
+    // Check that both collections are of the same size
+    if (inSiTracks.size() != inCluTracks.size()) {
+      fatal() << "Track collections have different sizes: SiTracks (" << inSiTracks.size() << ") and ClupatraTracks ("
+              << inCluTracks.size() << ")." << endmsg;
+    }
+
+    // Process both collections (SiTracks and ClupatraTracks) at the same time
+    std::tuple<FloatColl, FloatColl, FloatColl, FloatColl, FloatColl> siColls;
+    std::tuple<FloatColl, FloatColl, FloatColl, FloatColl, FloatColl> cluColls;
+    std::tuple<FloatColl, FloatColl, FloatColl, FloatColl, FloatColl> siUncColls;
+    std::tuple<FloatColl, FloatColl, FloatColl, FloatColl, FloatColl> cluUncColls;
+
+    for (size_t i = 0; i < inSiTracks.size(); ++i) {
+
+      // Process SiTrack
+      const auto oSiTrackStateIP = getOTrackAtIP(inSiTracks[i], "SiTrack");
+      if (oSiTrackStateIP.has_value()) {
+        // values
+        std::get<to_index(TP::d0)>(siColls).push_back(oSiTrackStateIP->D0);
+        std::get<to_index(TP::phi)>(siColls).push_back(oSiTrackStateIP->phi);
+        std::get<to_index(TP::omega)>(siColls).push_back(oSiTrackStateIP->omega);
+        std::get<to_index(TP::z0)>(siColls).push_back(oSiTrackStateIP->Z0);
+        std::get<to_index(TP::tanLambda)>(siColls).push_back(oSiTrackStateIP->tanLambda);
+        // uncertainties
+        std::get<to_index(TP::d0)>(siUncColls).push_back(getSigmaVar(*oSiTrackStateIP, TP::d0));
+        std::get<to_index(TP::phi)>(siUncColls).push_back(getSigmaVar(*oSiTrackStateIP, TP::phi));
+        std::get<to_index(TP::omega)>(siUncColls).push_back(getSigmaVar(*oSiTrackStateIP, TP::omega));
+        std::get<to_index(TP::z0)>(siUncColls).push_back(getSigmaVar(*oSiTrackStateIP, TP::z0));
+        std::get<to_index(TP::tanLambda)>(siUncColls).push_back(getSigmaVar(*oSiTrackStateIP, TP::tanLambda));
+      }
+
+      // Process CluTrack
+      const auto oCluTrackStateIP = getOTrackAtIP(inCluTracks[i], "CluTrack");
+      if (oCluTrackStateIP.has_value()) {
+        std::get<to_index(TP::d0)>(cluColls).push_back(oCluTrackStateIP->D0);
+        std::get<to_index(TP::phi)>(cluColls).push_back(oCluTrackStateIP->phi);
+        std::get<to_index(TP::omega)>(cluColls).push_back(oCluTrackStateIP->omega);
+        std::get<to_index(TP::z0)>(cluColls).push_back(oCluTrackStateIP->Z0);
+        std::get<to_index(TP::tanLambda)>(cluColls).push_back(oCluTrackStateIP->tanLambda);
+        // uncertainties
+        std::get<to_index(TP::d0)>(cluUncColls).push_back(getSigmaVar(*oCluTrackStateIP, TP::d0));
+        std::get<to_index(TP::phi)>(cluUncColls).push_back(getSigmaVar(*oCluTrackStateIP, TP::phi));
+        std::get<to_index(TP::omega)>(cluUncColls).push_back(getSigmaVar(*oCluTrackStateIP, TP::omega));
+        std::get<to_index(TP::z0)>(cluUncColls).push_back(getSigmaVar(*oCluTrackStateIP, TP::z0));
+        std::get<to_index(TP::tanLambda)>(cluUncColls).push_back(getSigmaVar(*oCluTrackStateIP, TP::tanLambda));
+      }
+    }
+
+    return std::tuple_cat(std::move(siColls), std::move(cluColls), std::move(siUncColls), std::move(cluUncColls));
+  };
+
+private:
+  Gaudi::Property<size_t> n_stars{this, "nStars", 20, "line-width of message in star box"};
+
+  std::optional<edm4hep::TrackState> getOTrackAtIP(const edm4hep::Track& track, const std::string& trackType) const {
+
+    // assuming there is only one Track State at IP
+    auto trackAtIP = std::ranges::find(track.getTrackStates(), TS::AtIP, &TS::location);
+    if (trackAtIP != track.getTrackStates().end()) {
+      verbose() << fmt::format("Track at IP found for {}.", trackType) << endmsg;
+      return *trackAtIP;
+    } else {
+      fatal() << fmt::format("No track at IP found for {}!", trackType) << endmsg;
+      return std::nullopt;
+    }
+  }
+
+  float getSigmaVar(const edm4hep::TrackState& ts, const TP var) const { return std::sqrt(ts.getCovMatrix(var, var)); }
+};
+DECLARE_COMPONENT(TrackParamExtractor)

Tracking/components/TrackD0Printer.cpp

@@ -0,0 +1,92 @@
+#include "utils.hpp"
+
+#include "edm4hep/Track.h"
+#include "edm4hep/TrackCollection.h"
+#include "edm4hep/TrackState.h"
+#include "podio/RelationRange.h"
+
+#include "Gaudi/Property.h"
+#include "GaudiKernel/MsgStream.h"
+
+#include "k4FWCore/Consumer.h"
+
+#include <fmt/core.h>
+
+#include <algorithm>
+#include <cmath>
+#include <iostream>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+
+// Type aliases for improved readability
+using TrackColl = edm4hep::TrackCollection;
+using TP = edm4hep::TrackParams;
+using TS = edm4hep::TrackState;
+
+// Consumer that processes track collections and prints phi values
+struct TrackD0Printer final : k4FWCore::Consumer<void(const TrackColl&, const TrackColl&)> {
+  // Constructor: define the input collections (ClupatraTracks and SiTracks)
+  TrackD0Printer(const std::string& name, ISvcLocator* svcLoc)
+      : Consumer(name, svcLoc,
+                 {
+                     KeyValues("InputSiTracks", {"SiTracks"}),
+                     KeyValues("InputCluTracks", {"ClupatraTracks"}),
+                 }) {}
+
+  // This function will be called to process the data
+  void operator()(const TrackColl& inSiTracks, const TrackColl& inCluTracks) const override {
+
+    printInStars(this, "New Event", n_stars);
+
+    debug() << "Received SiTracks collection with " << inSiTracks.size() << " tracks" << endmsg;
+    debug() << "Received ClupatraTracks collection with " << inCluTracks.size() << " tracks" << endmsg;
+
+    // Check that both collections are of the same size
+    if (inSiTracks.size() != inCluTracks.size()) {
+      fatal() << "Track collections have different sizes: SiTracks (" << inSiTracks.size() << ") and ClupatraTracks ("
+              << inCluTracks.size() << "). Exiting!" << endmsg;
+      return;
+    }
+
+    // Process both collections (SiTracks and ClupatraTracks) at the same time
+    for (size_t i = 0; i < inSiTracks.size(); ++i) {
+
+      // Process SiTrack
+      processTrack(inSiTracks[i], "SiTrack");
+
+      // Process CluTrack
+      processTrack(inCluTracks[i], "CluTrack");
+    }
+  }
+
+private:
+  Gaudi::Property<size_t> n_stars{this, "nStars", 20, "line-width of message in star box"};
+
+  float getSigmaPhi(const edm4hep::TrackState& ts) const { return std::sqrt(ts.getCovMatrix(TP::phi, TP::phi)); }
+
+  std::string printValueUnc(const std::string& strTrType, const std::string& vName, const int maxLabelWidth,
+                            const float varValue) const {
+    return fmt::format("{:<10}{:>{}}{:>10.5f}", strTrType, vName, maxLabelWidth, varValue);
+  }
+
+  // Helper function to process each track (either SiTrack or CluTrack)
+  void processTrack(const edm4hep::Track& track, const std::string& trackType) const {
+
+    std::string varName = "phi";
+    std::string sigmaVarName = "sigma " + varName;
+    int maxVarWidth = std::max(varName.size(), sigmaVarName.size()) + 2;
+
+    // assuming there is only one Track State at IP
+    if (auto trackAtIP = std::ranges::find(track.getTrackStates(), TS::AtIP, &TS::location);
+        trackAtIP != track.getTrackStates().end()) {
+      info() << printValueUnc(trackType, varName, maxVarWidth, trackAtIP->phi) << endmsg;
+      info() << printValueUnc(trackType, sigmaVarName, maxVarWidth, getSigmaPhi(*trackAtIP)) << endmsg;
+    } else {
+      fatal() << fmt::format("No track at IP found for {}!", trackType) << endmsg;
+    }
+  }
+};
+
+// Declare the consumer component
+DECLARE_COMPONENT(TrackD0Printer)

Tracking/include/utils.hpp

@@ -20,6 +20,8 @@
 #include "extension/MutableTrack.h"
 #include "extension/TrackCollection.h"
 
+#include <Gaudi/Algorithm.h>
+
 //=== Others ===
 #include <Objects/Helix.h>
 #include <TMatrixDSym.h>
@@ -189,4 +191,9 @@ int getHypotesisCharge(int pdg);
 TMatrixDSym computeTrackStateCovMatrix(TVectorD stateTrack, TVectorD params, TVector3 referencePoint, double timeError,
                                        TMatrixDSym statecovMatrix);
 
-#endif // UTILS_HPP
+/**
+ * @brief Print the message in a block of stars ('*') at level DEBUG
+ */
+void printInStars(const Gaudi::Algorithm* thisAlg, const std::string& msg, const int lineWidth);
+
+#endif // UTILS_HPP

Tracking/options/ExtractTrackParams.py

@@ -0,0 +1,51 @@
+import os
+from pathlib import Path
+
+from commonArgParsing import add_common_args, detModNames, registry
+from Configurables import TrackParamExtractor
+from Gaudi.Configuration import INFO, VERBOSE
+from k4FWCore import ApplicationMgr, IOSvc
+from k4FWCore.parseArgs import parser
+
+ARGS = add_common_args(parser).parse_known_args()[0]
+assert len(ARGS.detectorModels) == 1, (
+    f"Only provide one detector model! You provided {ARGS.detectorModels}"
+)
+ARGS.detectorModels = ARGS.detectorModels[0]
+
+FILE_SUFFIX = ".edm4hep.root"
+PROCESSOR_NAME = "TrackParamExtractor"
+BASE_PATH = Path(os.getenv("prmDir", Path.home() / "promotion"))
+IN_OUT_BASE_PATH = BASE_PATH / "data" / PROCESSOR_NAME
+CORE_PATH = f"{ARGS.version}_{detModNames[ARGS.detectorModels]}"
+
+# assert that the input path exists
+INPUT_PATH = (IN_OUT_BASE_PATH / "input_data" / f"{CORE_PATH}_REC").with_suffix(FILE_SUFFIX)
+assert INPUT_PATH.exists(), f"ERROR: The input path ({INPUT_PATH}) does not exist!"
+
+iosvc = IOSvc()
+iosvc.Input = str(INPUT_PATH)
+iosvc.Output = str(
+    (IN_OUT_BASE_PATH / "out_track_params" / f"{CORE_PATH}_track_params").with_suffix(FILE_SUFFIX)
+)
+# iosvc.outputCommands = ["drop *", "keep SiTrackPhi"]
+
+printer = TrackParamExtractor(PROCESSOR_NAME, nStars=40)
+printer.OutputLevel = VERBOSE
+
+# the collection name with the SiTracks differs between ILC and FCC models
+if registry.get(ARGS.detectorModels).at_fcc:
+    printer.InputSiTracks = ["SiTracksCT"]
+    SI_TRACK_COLL_NAME = "SiTracksCT"
+else:
+    SI_TRACK_COLL_NAME = "SiTracks"
+iosvc.CollectionNames = ["ClupatraTracks", SI_TRACK_COLL_NAME]
+
+
+ApplicationMgr(
+    TopAlg=[printer],
+    EvtSel="NONE",
+    EvtMax=10,
+    ExtSvc=[iosvc],
+    OutputLevel=INFO,
+)