Use ndf instead of probability for vertexes in EDM4hep (#71)

* Use ndf instead of probability for vertexes in EDM4hep * Add a comment about the find_ndf function * Reduce a few steps when above the initial guess for the upper bound * Add a comment about the initial upper bound * Move find_ndf to the header and .cpp file * Fix pre-commit * Add a check for invalid arguments --------- Co-authored-by: jmcarcell <[email protected]>
key4hep · Jun 28, 2024 · 74efd62 · 74efd62
1 parent 6f0c538
commit 74efd62
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 5 deletions.
diff --git a/k4EDM4hep2LcioConv/CMakeLists.txt b/k4EDM4hep2LcioConv/CMakeLists.txt
@@ -1,3 +1,5 @@
+find_package(ROOT REQUIRED COMPONENTS MathCore)
+
 add_library(k4EDM4hep2LcioConv SHARED
   src/k4EDM4hep2LcioConv.cpp
   src/k4Lcio2EDM4hepConv.cpp
@@ -12,6 +14,7 @@ target_link_libraries(k4EDM4hep2LcioConv PUBLIC
   LCIO::lcio
   EDM4HEP::edm4hep
   EDM4HEP::utils
+  ROOT::MathCore
 )
 
 set(public_headers

diff --git a/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4EDM4hep2LcioConv.ipp b/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4EDM4hep2LcioConv.ipp
@@ -2,6 +2,8 @@
 
 #include <cassert>
 
+#include "TMath.h"
+
 namespace EDM4hep2LCIOConv {
 
 template <typename TrackMapT>
@@ -420,7 +422,7 @@ std::unique_ptr<lcio::LCCollectionVec> convertVertices(const edm4hep::VertexColl
       lcio_vertex->setPrimary(edm_vertex.getPrimary());
       lcio_vertex->setAlgorithmType(std::to_string(edm_vertex.getAlgorithmType()));
       lcio_vertex->setChi2(edm_vertex.getChi2());
-      lcio_vertex->setProbability(edm_vertex.getProbability());
+      lcio_vertex->setProbability(TMath::Prob(edm_vertex.getChi2(), edm_vertex.getNdf()));
       lcio_vertex->setPosition(edm_vertex.getPosition()[0], edm_vertex.getPosition()[1], edm_vertex.getPosition()[2]);
       lcio_vertex->setCovMatrix(edm_vertex.getCovMatrix().data());
 

diff --git a/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h b/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h
@@ -397,6 +397,11 @@ void resolveRelationsTracks(TrackMapT& tracksMap, const TrackHitMapT& trackerHit
 template <typename VertexMapT, typename RecoParticleMapT>
 void resolveRelationsVertices(VertexMapT& vertexMap, const RecoParticleMapT& recoparticleMap);
 
+/**
+ * Go from chi^2 and probability (1 - CDF(chi^2, ndf)) to ndf by a binary search
+ */
+int find_ndf(double chi2, double prob);
+
 } // namespace LCIO2EDM4hepConv
 
 #include "k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.ipp"

diff --git a/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.ipp b/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.ipp
@@ -184,7 +184,7 @@ std::unique_ptr<edm4hep::VertexCollection> convertVertices(const std::string& na
 
     lval.setPrimary(rval->isPrimary() ? 1 : 0); // 1 for primary and 0 for not primary
     lval.setChi2(rval->getChi2());
-    lval.setProbability(rval->getProbability());
+    lval.setNdf(find_ndf(rval->getChi2(), rval->getProbability()));
     lval.setPosition(rval->getPosition());
     auto& m = rval->getCovMatrix(); // 6 parameters
     lval.setCovMatrix({m[0], m[1], m[2], m[3], m[4], m[5]});

diff --git a/k4EDM4hep2LcioConv/src/k4Lcio2EDM4hepConv.cpp b/k4EDM4hep2LcioConv/src/k4Lcio2EDM4hepConv.cpp
@@ -2,7 +2,7 @@
 
 #include <UTIL/PIDHandler.h>
 
-#include <iostream>
+#include "TMath.h"
 
 namespace LCIO2EDM4hepConv {
 
@@ -152,4 +152,30 @@ podio::Frame convertRunHeader(EVENT::LCRunHeader* rheader) {
   return runHeaderFrame;
 }
 
+int find_ndf(double chi2, double prob) {
+  if (chi2 < 0 || prob < 0 || prob > 1) {
+    throw std::invalid_argument("Invalid input for find_ndf. Either chi2 < 0, prob < 0 or prob > 1 in a LCIO Vertex.");
+  }
+  int lower = 0;
+  // Initial guess for the upper bound. If it's not enough, it will be increased
+  int upper = 100;
+  while (TMath::Prob(chi2, upper) < prob) {
+    lower = upper;
+    upper *= 2;
+  }
+  while (lower < upper - 1) {
+    int mid = (lower + upper) / 2;
+    if (TMath::Prob(chi2, mid) < prob) {
+      lower = mid;
+    } else {
+      upper = mid;
+    }
+  }
+  if (std::abs(TMath::Prob(chi2, lower) - prob) < std::abs(TMath::Prob(chi2, upper) - prob)) {
+    return lower;
+  } else {
+    return upper;
+  }
+}
+
 } // namespace LCIO2EDM4hepConv
diff --git a/tests/src/CompareEDM4hepLCIO.cc b/tests/src/CompareEDM4hepLCIO.cc
@@ -5,6 +5,8 @@
 
 #include <cstdint>
 
+#include "TMath.h"
+
 /**
  * The basic implementation of the functionality has been generated via modified
  * podio templates, employing some handwritten macros to facilitate the task.
@@ -409,7 +411,8 @@ bool compare(const EVENT::Vertex* lcioElem, const edm4hep::Vertex& edm4hepElem,
   // LCIO has isPrimary (bool), EDM4hep has getPrimary (int32_t)
   ASSERT_COMPARE_VALS(lcioElem->isPrimary(), edm4hepElem.getPrimary(), "primary in Vertex");
   ASSERT_COMPARE(lcioElem, edm4hepElem, getChi2, "chi2 in Vertex");
-  ASSERT_COMPARE(lcioElem, edm4hepElem, getProbability, "probability in Vertex");
+  ASSERT_COMPARE_VALS_FLOAT(lcioElem->getProbability(), TMath::Prob(edm4hepElem.getChi2(), edm4hepElem.getNdf()), 1e-6,
+                            "probability in Vertex");
   ASSERT_COMPARE(lcioElem, edm4hepElem, getPosition, "position in Vertex");
   ASSERT_COMPARE(lcioElem, edm4hepElem, getCovMatrix, "covMatrix in Vertex");
   ASSERT_COMPARE(lcioElem, edm4hepElem, getParameters, "parameters in Vertex");

diff --git a/tests/src/ComparisonUtils.h b/tests/src/ComparisonUtils.h
@@ -69,7 +69,7 @@ template <typename T, typename... Ts>
 constexpr bool isAnyOf = (std::is_same_v<T, Ts> || ...);
 
 // Helper function for comparing values and vectors of values element by
-// element, ignoring cases where both values aer nan
+// element, ignoring cases where both values are nan
 template <typename LCIO, typename EDM4hepT>
 bool compareValuesNanSafe(LCIO lcioV, EDM4hepT edm4hepV, const std::string& msg) {
   constexpr auto isVectorLike =
@@ -118,6 +118,12 @@ bool compareValuesNanSafe(LCIO lcioV, EDM4hepT edm4hepV, const std::string& msg)
     return false;                                                                                                      \
   }
 
+#define ASSERT_COMPARE_VALS_FLOAT(lcioV, edm4hepV, tol, msg)                                                           \
+  if (std::abs(lcioV - edm4hepV) > tol) {                                                                              \
+    std::cerr << msg << " (LCIO: " << (lcioV) << ", EDM4hep: " << (edm4hepV) << ")" << std::endl;                      \
+    return false;                                                                                                      \
+  }
+
 #define ASSERT_COMPARE(lcioE, edm4hepE, func, msg)                                                                     \
   {                                                                                                                    \
     const auto lcioV = lcioE->func();                                                                                  \