Convert dQ/dx information in both directions for Tracks (#69)

* Convert dE/dx information into separate collection

* Attach dE/dx information during conversion from EDM4hep to LCIO

* Remove unused function

* Use more appropriate function name

---------

Co-authored-by: tmadlener <[email protected]>

k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4EDM4hep2LcioConv.h

@@ -17,6 +17,7 @@
 #include <edm4hep/ParticleIDCollection.h>
 #include <edm4hep/RawCalorimeterHitCollection.h>
 #include <edm4hep/RawTimeSeriesCollection.h>
+#include <edm4hep/RecDqdxCollection.h>
 #include <edm4hep/RecoParticleVertexAssociationCollection.h>
 #include <edm4hep/ReconstructedParticleCollection.h>
 #include <edm4hep/SimCalorimeterHitCollection.h>
@@ -106,6 +107,11 @@ struct ParticleIDConvData {
   std::optional<edm4hep::utils::ParticleIDMeta> metadata;
 };
 
+struct TrackDqdxConvData {
+  std::string name;
+  const edm4hep::RecDqdxCollection* coll;
+};
+
 /// Sort the ParticleIDs according to their algorithmType.
 ///
 /// This sorting allows for a fully consistent roundtrip conversion under the
@@ -399,6 +405,20 @@ void resolveRelationsClusters(ClusterMapT& clustersMap, const CaloHitMapT& caloH
 template <typename PidMapT, typename RecoParticleMapT>
 void resolveRelationsParticleIDs(PidMapT& pidMap, const RecoParticleMapT& recoMap);
 
+/// Attach the dE/dx information that is stored in the RecDqdxCollections to the
+/// corresponding tracks
+///
+/// @note: This assumes that all tracks have been converted already
+template <typename TrackMapT>
+void attachDqdxInfo(TrackMapT& trackMap, const std::vector<TrackDqdxConvData>& dQdxCollections);
+
+/// Attach the dE/dx information that is stored in the RecDqdxCollection to the
+/// corresponding tracks
+///
+/// @note: This assumes that all tracks have been converted already
+template <typename TrackMapT>
+void attachDqdxInfo(TrackMapT& trackMap, const TrackDqdxConvData& dQdxCollection);
+
 /**
  * Resolve all relations in all converted objects that are held in the map.
  * Dispatch to the correpsonding implementation for all the types that have

k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4EDM4hep2LcioConv.ipp

@@ -31,8 +31,6 @@ std::unique_ptr<lcio::LCCollectionVec> convertTracks(const edm4hep::TrackCollect
       }
       lcio_tr->setChi2(edm_tr.getChi2());
       lcio_tr->setNdf(edm_tr.getNdf());
-      lcio_tr->setdEdx(edm_tr.getDEdx());
-      lcio_tr->setdEdxError(edm_tr.getDEdxError());
       lcio_tr->setRadiusOfInnermostHit(getRadiusOfStateAtFirstHit(edm_tr).value_or(-1.0));
 
       // Loop over the hit Numbers in the track
@@ -729,6 +727,29 @@ void resolveRelationsParticleIDs(PidMapT& pidMap, const RecoParticleMapT& recoMa
   }
 }
 
+template <typename TrackMapT>
+void attachDqdxInfo(TrackMapT& trackMap, const std::vector<TrackDqdxConvData>& dQdxCollections) {
+  for (const auto& coll : dQdxCollections) {
+    attachDqdxInfo(trackMap, coll);
+  }
+}
+
+template <typename TrackMapT>
+void attachDqdxInfo(TrackMapT& trackMap, const TrackDqdxConvData& dQdxCollection) {
+  const auto& [name, coll] = dQdxCollection;
+
+  for (const auto& elem : *coll) {
+    const auto dQdxTrack = elem.getTrack();
+    const auto lcioTrack = k4EDM4hep2LcioConv::detail::mapLookupFrom(dQdxTrack, trackMap);
+    if (lcioTrack) {
+      lcioTrack.value()->setdEdx(elem.getDQdx().value);
+      lcioTrack.value()->setdEdxError(elem.getDQdx().error);
+    } else {
+      std::cerr << "Cannot find a track to attach dQ/dx information to for collection: " << name << std::endl;
+    }
+  }
+}
+
 template <typename ObjectMappingT>
 void resolveRelations(ObjectMappingT& collection_pairs) {
   resolveRelations(collection_pairs, collection_pairs);

k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h

@@ -275,8 +275,8 @@ convertTrackerHitPlanes(const std::string& name, EVENT::LCCollection* LCCollecti
  * Simultaneously populates the mapping from LCIO to EDM4hep objects.
  */
 template <typename TrackMapT>
-std::unique_ptr<edm4hep::TrackCollection> convertTracks(const std::string& name, EVENT::LCCollection* LCCollection,
-                                                        TrackMapT& TrackMap);
+std::vector<CollNamePair> convertTracks(const std::string& name, EVENT::LCCollection* LCCollection,
+                                        TrackMapT& TrackMap);
 
 /**
  * Convert a SimCalorimeterHit collection and return the resulting collection.

k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.ipp

@@ -1,7 +1,9 @@
 #include "k4EDM4hep2LcioConv/MappingUtils.h"
 
 #include <UTIL/PIDHandler.h>
+
 #include <edm4hep/ParticleIDCollection.h>
+#include <edm4hep/RecDqdxCollection.h>
 
 namespace LCIO2EDM4hepConv {
 template <typename LCIOType>
@@ -335,19 +337,24 @@ convertTrackerHitPlanes(const std::string& name, EVENT::LCCollection* LCCollecti
 }
 
 template <typename TrackMapT>
-std::unique_ptr<edm4hep::TrackCollection> convertTracks(const std::string& name, EVENT::LCCollection* LCCollection,
-                                                        TrackMapT& TrackMap) {
+std::vector<CollNamePair> convertTracks(const std::string& name, EVENT::LCCollection* LCCollection,
+                                        TrackMapT& TrackMap) {
   auto dest = std::make_unique<edm4hep::TrackCollection>();
+  auto dQdxColl = std::make_unique<edm4hep::RecDqdxCollection>();
 
   for (unsigned i = 0, N = LCCollection->getNumberOfElements(); i < N; ++i) {
     auto* rval = static_cast<EVENT::Track*>(LCCollection->getElementAt(i));
     auto lval = dest->create();
+    auto trackDqdx = dQdxColl->create();
+    trackDqdx.setTrack(lval);
 
     lval.setType(rval->getType());
     lval.setChi2(rval->getChi2());
     lval.setNdf(rval->getNdf());
-    lval.setDEdx(rval->getdEdx());
-    lval.setDEdxError(rval->getdEdxError());
+
+    auto& dqdx = trackDqdx.getDQdx();
+    dqdx.value = rval->getdEdx();
+    dqdx.error = rval->getdEdxError();
 
     auto subdetectorHitNum = rval->getSubdetectorHitNumbers();
     for (auto hitNum : subdetectorHitNum) {
@@ -357,10 +364,7 @@ std::unique_ptr<edm4hep::TrackCollection> convertTracks(const std::string& name,
     for (auto& trackState : trackStates) {
       lval.addToTrackStates(convertTrackState(trackState));
     }
-    auto quantities = edm4hep::Quantity{};
-    quantities.value = rval->getdEdx();
-    quantities.error = rval->getdEdxError();
-    lval.addToDxQuantities(quantities);
+
     const auto [iterator, inserted] = k4EDM4hep2LcioConv::detail::mapInsert(rval, lval, TrackMap);
     if (!inserted) {
       auto existing = k4EDM4hep2LcioConv::detail::getMapped(iterator);
@@ -370,7 +374,10 @@ std::unique_ptr<edm4hep::TrackCollection> convertTracks(const std::string& name,
     }
   }
 
-  return dest;
+  std::vector<CollNamePair> results{};
+  results.emplace_back(name, std::move(dest));
+  results.emplace_back(name + "_dQdx", std::move(dQdxColl));
+  return results;
 }
 
 template <typename HitMapT>
@@ -497,7 +504,7 @@ std::vector<CollNamePair> convertCollection(const std::string& name, EVENT::LCCo
   } else if (type == "Vertex") {
     retColls.emplace_back(name, convertVertices(name, LCCollection, typeMapping.vertices));
   } else if (type == "Track") {
-    retColls.emplace_back(name, convertTracks(name, LCCollection, typeMapping.tracks));
+    return convertTracks(name, LCCollection, typeMapping.tracks);
   } else if (type == "Cluster") {
     retColls.emplace_back(name, convertClusters(name, LCCollection, typeMapping.clusters));
   } else if (type == "SimCalorimeterHit") {

k4EDM4hep2LcioConv/src/k4EDM4hep2LcioConv.cpp

@@ -76,6 +76,7 @@ std::unique_ptr<lcio::LCEventImpl> convertEvent(const podio::Frame& edmEvent, co
   // We convert these at the very end, once all the necessary information is
   // available
   std::vector<std::tuple<std::string, const podio::CollectionBase*>> associations{};
+  std::vector<TrackDqdxConvData> dQdxCollections{};
 
   const auto& collections = edmEvent.getAvailableCollections();
   for (const auto& name : collections) {
@@ -124,6 +125,8 @@ std::unique_ptr<lcio::LCEventImpl> convertEvent(const podio::Frame& edmEvent, co
       convertEventHeader(coll, lcioEvent.get());
     } else if (auto coll = dynamic_cast<const edm4hep::ParticleIDCollection*>(edmCollection)) {
       pidCollections.emplace_back(name, coll, edm4hep::utils::PIDHandler::getAlgoInfo(metadata, name));
+    } else if (auto coll = dynamic_cast<const edm4hep::RecDqdxCollection*>(edmCollection)) {
+      dQdxCollections.emplace_back(name, coll);
     } else if (dynamic_cast<const edm4hep::CaloHitContributionCollection*>(edmCollection)) {
       // "converted" during relation resolving later
       continue;
@@ -148,6 +151,8 @@ std::unique_ptr<lcio::LCEventImpl> convertEvent(const podio::Frame& edmEvent, co
     convertParticleIDs(pidCollMeta.coll, objectMappings.particleIDs, algoId);
   }
 
+  attachDqdxInfo(objectMappings.tracks, dQdxCollections);
+
   resolveRelations(objectMappings);
 
   for (auto& [name, coll] : createLCRelationCollections(associations, objectMappings)) {

tests/edm4hep_roundtrip.cpp

@@ -32,6 +32,7 @@ int main() {
   ASSERT_SAME_OR_ABORT(edm4hep::ParticleIDCollection, "ParticleID_coll_3");
   ASSERT_SAME_OR_ABORT(edm4hep::MCRecoParticleAssociationCollection, "mcRecoAssocs");
   ASSERT_SAME_OR_ABORT(edm4hep::MCRecoCaloAssociationCollection, "mcCaloHitsAssocs");
+  ASSERT_SAME_OR_ABORT(edm4hep::RecDqdxCollection, "tracks_dQdx")
 
   return 0;
 }

tests/edm4hep_to_lcio.cpp

@@ -21,7 +21,7 @@ int main() {
     const auto edmColl = edmEvent.get(name);
     const auto typeName = edmColl->getValueTypeName();
     if (typeName == "edm4hep::CaloHitContribution" || typeName == "edm4hep::ParticleID" ||
-        typeName == "edm4hep::EventHeader") {
+        typeName == "edm4hep::EventHeader" || typeName == "edm4hep::RecDqdx") {
       continue;
     }
     try {
@@ -42,7 +42,7 @@ int main() {
   for (const auto& name : edmEvent.getAvailableCollections()) {
     const auto type = edmEvent.get(name)->getTypeName();
     if (type == "edm4hep::CaloHitContributionCollection" || type == "edm4hep::ParticleIDCollection" ||
-        type == "edm4hep::EventHeaderCollection") {
+        type == "edm4hep::EventHeaderCollection" || type == "edm4hep::RecDqdxCollection") {
       continue;
     }
     const auto* lcioColl = lcioEvent->getCollection(name);

tests/src/CompareEDM4hepEDM4hep.cc

@@ -1,4 +1,28 @@
 #include "CompareEDM4hepEDM4hep.h"
+#include "ComparisonUtils.h"
+
+#include "edm4hep/CalorimeterHitCollection.h"
+#include "edm4hep/ClusterCollection.h"
+#include "edm4hep/MCParticleCollection.h"
+#include "edm4hep/ParticleIDCollection.h"
+#include "edm4hep/RecDqdxCollection.h"
+#include "edm4hep/ReconstructedParticleCollection.h"
+#include "edm4hep/SimCalorimeterHitCollection.h"
+#include "edm4hep/TrackCollection.h"
+#include "edm4hep/TrackerHitPlaneCollection.h"
+
+#include <edm4hep/TrackState.h>
+#include <iostream>
+#include <podio/RelationRange.h>
+
+#define REQUIRE_SAME(expected, actual, msg)                                                                            \
+  {                                                                                                                    \
+    if (!((expected) == (actual))) {                                                                                   \
+      std::cerr << msg << " are not the same (expected: " << (expected) << ", actual: " << (actual) << ")"             \
+                << std::endl;                                                                                          \
+      return false;                                                                                                    \
+    }                                                                                                                  \
+  }
 
 bool compare(const edm4hep::CalorimeterHitCollection& origColl,
              const edm4hep::CalorimeterHitCollection& roundtripColl) {
@@ -123,8 +147,6 @@ bool compare(const edm4hep::TrackCollection& origColl, const edm4hep::TrackColle
 
     REQUIRE_SAME(origTrack.getChi2(), track.getChi2(), "chi2 in track " << i);
     REQUIRE_SAME(origTrack.getNdf(), track.getNdf(), "chi2 in track " << i);
-    REQUIRE_SAME(origTrack.getDEdx(), track.getDEdx(), "dEdx in track " << i);
-    REQUIRE_SAME(origTrack.getDEdxError(), track.getDEdxError(), "dEdxError in track " << i);
 
     const auto origSubDetHitNumbers = origTrack.getSubdetectorHitNumbers();
     const auto subDetHitNumbers = track.getSubdetectorHitNumbers();
@@ -314,5 +336,25 @@ bool compare(const edm4hep::RecoParticleVertexAssociationCollection& origColl,
     REQUIRE_SAME(origAssoc.getVertex().id(), assoc.getVertex().id(), "vertex in association " << i);
     REQUIRE_SAME(origAssoc.getRec().id(), assoc.getRec().id(), "reco particle in association " << i);
   }
+
+  return true;
+}
+
+bool compare(const edm4hep::RecDqdxCollection& origColl, const edm4hep::RecDqdxCollection& roundtripColl) {
+  REQUIRE_SAME(origColl.size(), roundtripColl.size(), "collection sizes");
+  for (size_t i = 0; i < origColl.size(); ++i) {
+    const auto origDqdx = origColl[i];
+    const auto dQdx = roundtripColl[i];
+
+    const auto origQuant = origDqdx.getDQdx();
+    const auto quant = dQdx.getDQdx();
+
+    REQUIRE_SAME(origQuant.value, quant.value, "value of quantity in dQ/dx " << i);
+    REQUIRE_SAME(origQuant.error, quant.error, "error of quantity in dQ/dx " << i);
+    REQUIRE_SAME(origQuant.type, quant.value, "value of quantity in dQ/dx " << i);
+
+    REQUIRE_SAME(origDqdx.getTrack().id(), dQdx.getTrack().id(), "related track in dQ/dx " << i);
+  }
+
   return true;
 }

tests/src/CompareEDM4hepEDM4hep.h

@@ -53,4 +53,6 @@ bool compare(const AssociationCollT& origColl, const AssociationCollT& roundtrip
   return true;
 }
 
+bool compare(const edm4hep::RecDqdxCollection& origColl, const edm4hep::RecDqdxCollection& roundtripColl);
+
 #endif // K4EDM4HEP2LCIOCONV_TEST_COMPAREEDM4HEPEDM4HEP_H

tests/src/CompareEDM4hepLCIO.cc

@@ -303,18 +303,16 @@ bool compare(const EVENT::Track* lcioElem, const edm4hep::Track& edm4hepElem, co
   ASSERT_COMPARE(lcioElem, edm4hepElem, getType, "type in Track");
   ASSERT_COMPARE(lcioElem, edm4hepElem, getChi2, "chi2 in Track");
   ASSERT_COMPARE(lcioElem, edm4hepElem, getNdf, "ndf in Track");
-  // LCIO has getdEdx instead of getDEdx
-  ASSERT_COMPARE_VALS(lcioElem->getdEdx(), edm4hepElem.getDEdx(), "dEdx in Track");
-  ASSERT_COMPARE_VALS(lcioElem->getdEdxError(), edm4hepElem.getDEdxError(), "dEdxError in Track");
-  // Also check whether these have been corretly put into the dQQuantities
-  const auto dxQuantities = edm4hepElem.getDxQuantities();
-  if (dxQuantities.size() != 1) {
-    std::cerr << "DxQuantities have not been filled correctly, expected exactly 1, got " << dxQuantities.size()
-              << " in Track" << std::endl;
+
+  // EDM4hep does not have the dEdx information inside the track, but rather
+  // inside a RecDqdx object
+  const auto& dQdxInfos = objectMaps.trackPidHandler.getDqdxValues(edm4hepElem);
+  if (dQdxInfos.size() != 1) {
+    std::cerr << "Could not find a unique RecDqdx object for track" << std::endl;
     return false;
   }
-  ASSERT_COMPARE_VALS(lcioElem->getdEdx(), dxQuantities[0].value, "dEdx in DxQuantities in Track");
-  ASSERT_COMPARE_VALS(lcioElem->getdEdxError(), dxQuantities[0].error, "dEdxError in DxQuantities in Track");
+  ASSERT_COMPARE_VALS(lcioElem->getdEdx(), dQdxInfos[0].getDQdx().value, "dEdx in Track");
+  ASSERT_COMPARE_VALS(lcioElem->getdEdxError(), dQdxInfos[0].getDQdx().error, "dEdxError in Track");
 
   double radius = EDM4hep2LCIOConv::getRadiusOfStateAtFirstHit(edm4hepElem).value_or(-1.0);
   double radius3D = EDM4hep2LCIOConv::getRadiusOfStateAtFirstHit(edm4hepElem, true).value_or(-1.0);

tests/src/EDM4hep2LCIOUtilities.cc

@@ -179,11 +179,6 @@ edm4hep::TrackCollection createTracks(const int num_elements, const int subdetec
     elem.setChi2(i * 10.f);
     elem.setNdf(i * 12);
 
-    elem.setDEdx(i);
-    elem.setDEdxError(i / std::sqrt(i + 1));
-    // Also add a DxQuantity since the comparison expects that
-    elem.addToDxQuantities({0, i * 1.f, i / std::sqrt(i + 1.f)});
-
     for (int j = 0; j < subdetectorhitnumbers; ++j) {
       elem.addToSubdetectorHitNumbers(i + 10 * j);
     }
@@ -217,6 +212,21 @@ edm4hep::TrackCollection createTracks(const int num_elements, const int subdetec
   return track_coll;
 }
 
+edm4hep::RecDqdxCollection createDqDxColl(const edm4hep::TrackCollection& tracks) {
+  auto dqdxColl = edm4hep::RecDqdxCollection();
+
+  int i = 0;
+  for (const auto& track : tracks) {
+    auto elem = dqdxColl.create();
+    elem.setTrack(track);
+    auto& dqdx = elem.getDQdx();
+    dqdx.value = i;
+    dqdx.error = i / std::sqrt(i + 1);
+  }
+
+  return dqdxColl;
+}
+
 std::pair<edm4hep::SimCalorimeterHitCollection, edm4hep::CaloHitContributionCollection>
 createSimCalorimeterHits(const int num_elements, const int num_contributions,
                          const edm4hep::MCParticleCollection& mcParticles,
@@ -379,6 +389,7 @@ std::tuple<podio::Frame, podio::Frame> createExampleEvent() {
                                               test_config::nTrackStates, trackerHits, trackerHitPlanes,
                                               test_config::trackTrackerHitIdcs, test_config::trackTrackIdcs),
                                  "tracks");
+  event.put(createDqDxColl(tracks), "tracks_dQdx");
 
   const auto& clusters = event.put(createClusters(test_config::nClusters, caloHits, test_config::nSubdetectorEnergies,
                                                   test_config::clusterHitIdcs, test_config::clusterClusterIdcs),

tests/src/EDM4hep2LCIOUtilities.h

@@ -18,6 +18,7 @@
 #include <edm4hep/MCRecoTrackerAssociationCollection.h>
 #include <edm4hep/ParticleIDCollection.h>
 #include <edm4hep/RawTimeSeriesCollection.h>
+#include <edm4hep/RecDqdxCollection.h>
 #include <edm4hep/RecoParticleVertexAssociationCollection.h>
 #include <edm4hep/TrackerHitPlaneCollection.h>
 #if __has_include("edm4hep/TrackerHit3DCollection.h")

tests/src/ObjectMapping.cc

@@ -19,6 +19,7 @@
 #include "UTIL/PIDHandler.h"
 
 #include "edm4hep/TrackCollection.h"
+#include <edm4hep/RecDqdxCollection.h>
 #if __has_include("edm4hep/TrackerHit3DCollection.h")
 #include "edm4hep/TrackerHit3DCollection.h"
 #else
@@ -89,14 +90,6 @@ void fillRecoPIDMaps(ObjectMappings::Map<const EVENT::ReconstructedParticle*>& r
   }
 }
 
-std::string getRecoName(const std::string& pidName) {
-  const auto pos = pidName.find("_PID_");
-  if (pos != std::string::npos) {
-    return pidName.substr(0, pos);
-  }
-  return "";
-}
-
 ObjectMappings ObjectMappings::fromEvent(EVENT::LCEvent* lcEvt, const podio::Frame& edmEvt) {
   ObjectMappings mapping{};
   for (const auto& name : *(lcEvt->getCollectionNames())) {
@@ -133,6 +126,9 @@ ObjectMappings ObjectMappings::fromEvent(EVENT::LCEvent* lcEvt, const podio::Fra
     if (type == "ReconstructedParticle") {
       fillRecoPIDMaps(mapping.recoParticles, mapping.particleIDs, name, lcEvt, edmEvt);
     }
+    if (type == "Track") {
+      mapping.trackPidHandler.addColl(edmEvt.get<edm4hep::RecDqdxCollection>(name + "_dQdx"));
+    }
   }
 
   return mapping;