Merge pull request #121 from kjvbrt/ana_thrust

Adding determination of event thrust not by minimization

Author: clementhelsens

.github/workflows/bench.yml

@@ -3,7 +3,7 @@ name: benchmark
 on: [push]
 
 jobs:
-  test:
+  bench:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false

.gitignore

@@ -9,7 +9,10 @@ __pycache__/
 
 # ROOT files
 *.root
+
+# Editors
 *~
+.vimlocal
 
 # Distribution / packaging
 .Python

analyzers/dataframe/Utils.h

@@ -72,7 +72,7 @@ namespace Algorithms{
   };
 
 
-  /// Calculates the thrust axis based on a list of px, py, pz
+  /// Finds the thrust axis based on a list of px, py, pz
   struct minimize_thrust {
     minimize_thrust(std::string arg_minname="Minuit2",
                     std::string arg_algoname="Migrad",
@@ -92,6 +92,22 @@ namespace Algorithms{
     double _variable[3]={1.0,1.0,1.0};
   };
 
+  /// Calculates the thrust axis by looping over all possible combinations
+  struct calculate_thrust {
+    calculate_thrust(){}
+    ROOT::VecOps::RVec<float> operator()(const ROOT::VecOps::RVec<float>& px,
+                                         const ROOT::VecOps::RVec<float>& py,
+                                         const ROOT::VecOps::RVec<float>& pz);
+
+    // Helper functions, to ease manipulation with the elements of internal array
+    inline void mag2(float (&vec)[4]);
+    inline float dot(float vec1[4], float vec2[4]);
+    inline void cross(float (&vec)[4], float vec1[4], float vec2[4]);
+    inline void unit(float (&vec)[4]);
+    inline void plus(float (&vec)[4], float vecIn1[4], float vecIn2[4]);
+    inline void minus(float (&vecOut)[4], float vecIn1[4], float vecIn2[4]);
+    inline void copy(float (&vecOut)[4], float vecIn[4]);
+  };
 
   /// Get the weighted charge in a given hemisphere (defined by it's angle wrt to axis). For definition see eq1 https://arxiv.org/pdf/1209.2421.pdf
   struct getAxisCharge {

analyzers/dataframe/include/FCCAnalyses/Algorithms.h

@@ -0,0 +1,14 @@
+#ifndef  UTILS_ANALYZERS_H
+#define  UTILS_ANALYZERS_H
+
+#include <cmath>
+
+namespace FCCAnalyses {
+  namespace Utils {
+
+    template<typename T> inline auto getsize( T& vec){ return vec.size();};
+
+  }
+}
+
+#endif

analyzers/dataframe/include/FCCAnalyses/Utils.h

@@ -1,4 +1,5 @@
 #include "FCCAnalyses/Algorithms.h"
+#include "FCCAnalyses/Utils.h"
 #include "Math/Minimizer.h"
 #include "Math/IFunction.h"
 #include "Math/Factory.h"
@@ -163,6 +164,150 @@ ROOT::VecOps::RVec<float> minimize_thrust::operator()(const ROOT::VecOps::RVec<f
 }
 
 
+ROOT::VecOps::RVec<float> Algorithms::calculate_thrust::operator()(
+    const ROOT::VecOps::RVec<float>& px,
+    const ROOT::VecOps::RVec<float>& py,
+    const ROOT::VecOps::RVec<float>& pz) {
+
+  if (px.size() != py.size()) {
+    throw std::domain_error("calculate_thrust: Input vector sizes are not equal.");
+  }
+  if (px.size() != pz.size()) {
+    throw std::domain_error("calculate_thrust: Input vector sizes are not equal.");
+  }
+
+  ROOT::VecOps::RVec<float> result;
+
+  size_t nParticles = px.size();
+  if (nParticles < 2) {
+    // std::cout << "calculate_thrust: Number of input particles too small."
+    //           << std::endl;
+
+    result.push_back(-1);
+    result.push_back(-1);
+    result.push_back(-1);
+    result.push_back(-1);
+
+    return result;
+  }
+
+  // std::cout << "calculate_thrust: Calculating sum of all particles in event"
+  //           << std::endl;
+
+  // Array to store x, y, z and magnitude squared of the particles.
+  // 0 -- magnitude squared, 1 -- x, 2 -- y, 3 -- z
+  float pArr[nParticles][4];
+  float pSum = 0.;
+  for (size_t i = 0; i < nParticles; ++i) {
+    pArr[i][1] = px[i];
+    pArr[i][2] = py[i];
+    pArr[i][3] = pz[i];
+    mag2(pArr[i]);
+    pSum += std::sqrt(pArr[i][0]);
+  }
+
+  // Trying all combinations of reference vector orthogonal to two selected
+  // particles.
+  // std::cout << "Trying all combinations..." << std::endl;
+  float pMax[4] = {0., 0., 0., 0.};
+  for (size_t i = 0; i < nParticles - 1; ++i) {
+    for (size_t j = i + 1; j < nParticles; ++j) {
+      float nRef[4];
+      cross(nRef, pArr[i], pArr[j]);
+      mag2(nRef);
+      unit(nRef);
+
+      float pPart[4] = {0., 0., 0., 0.};
+      for (size_t k = 0; k < nParticles; ++k) {
+        if (k == i || k == j) {
+          continue;
+        }
+
+        if (dot(nRef, pArr[k]) > 0.) {
+          plus(pPart, pPart, pArr[k]);
+        } else {
+          minus(pPart, pPart, pArr[k]);
+        }
+      }
+
+      float pFullArr[4][4];
+      // pPart + pArr[i] + pArr[j]
+      plus(pFullArr[0], pPart, pArr[i]);
+      plus(pFullArr[0], pFullArr[0], pArr[j]);
+
+      // pPart + pArr[i] - pArr[j]
+      plus(pFullArr[1], pPart, pArr[i]);
+      minus(pFullArr[1], pFullArr[1], pArr[j]);
+
+      // pPart - pArr[i] + pArr[j]
+      minus(pFullArr[2], pPart, pArr[i]);
+      plus(pFullArr[2], pFullArr[2], pArr[j]);
+
+      // pPart - pArr[i] - pArr[j]
+      minus(pFullArr[3], pPart, pArr[i]);
+      minus(pFullArr[3], pFullArr[3], pArr[j]);
+
+      for (size_t k = 0; k < 4; ++k) {
+        mag2(pFullArr[k]);
+        if (pFullArr[k][0] > pMax[0]) {
+          copy(pMax, pFullArr[k]);
+        }
+      }
+    }
+  }
+
+  float pMaxMag = std::sqrt(pMax[0]);
+  // std::cout << "Thrust value arr: " << pMaxMag / pSum << std::endl;
+  result.push_back(pMaxMag / pSum);
+  // Normalizing the thrust vector
+  result.push_back(pMax[1]/pMaxMag);
+  result.push_back(pMax[2]/pMaxMag);
+  result.push_back(pMax[3]/pMaxMag);
+
+  return result;
+}
+
+inline void Algorithms::calculate_thrust::mag2(float (&vec)[4]) {
+  vec[0] = vec[1]*vec[1] + vec[2]*vec[2] + vec[3]*vec[3];
+}
+
+inline float Algorithms::calculate_thrust::dot(float vec1[4], float vec2[4]) {
+  return vec1[1]*vec2[1] + vec1[2]*vec2[2] + vec1[3]*vec2[3];
+}
+
+inline void Algorithms::calculate_thrust::cross(float (&vec)[4], float vec1[4], float vec2[4]) {
+  vec[1] = vec1[2]*vec2[3] - vec1[3]*vec2[2];
+  vec[2] = vec1[3]*vec2[1] - vec1[1]*vec2[3];
+  vec[3] = vec1[1]*vec2[2] - vec1[2]*vec2[1];
+}
+
+inline void Algorithms::calculate_thrust::unit(float (&vec)[4]) {
+  float mag = std::sqrt(vec[0]);
+  vec[1] = vec[1]/mag;
+  vec[2] = vec[2]/mag;
+  vec[3] = vec[3]/mag;
+}
+
+inline void Algorithms::calculate_thrust::plus(float (&vec)[4], float vecIn1[4], float vecIn2[4]) {
+  vec[1] = vecIn1[1] + vecIn2[1];
+  vec[2] = vecIn1[2] + vecIn2[2];
+  vec[3] = vecIn1[3] + vecIn2[3];
+}
+
+inline void Algorithms::calculate_thrust::minus(float (&vecOut)[4], float vecIn1[4], float vecIn2[4]) {
+  vecOut[1] = vecIn1[1] - vecIn2[1];
+  vecOut[2] = vecIn1[2] - vecIn2[2];
+  vecOut[3] = vecIn1[3] - vecIn2[3];
+}
+
+inline void Algorithms::calculate_thrust::copy(float (&vecOut)[4], float vecIn[4]) {
+  vecOut[0] = vecIn[0];
+  vecOut[1] = vecIn[1];
+  vecOut[2] = vecIn[2];
+  vecOut[3] = vecIn[3];
+}
+
+
 getAxisCharge::getAxisCharge(bool arg_pos,
 	                           float arg_power){
   _pos = arg_pos;
@@ -185,7 +330,6 @@ float  getAxisCharge::operator() (const ROOT::VecOps::RVec<float> & angle,
 }
 
 
-
 getAxisMass::getAxisMass(bool arg_pos){
   _pos=arg_pos;
 }

analyzers/dataframe/src/Algorithms.cc

@@ -1,6 +1,8 @@
 
 add_subdirectory(unittest)
 
+add_subdirectory(benchmark)
+
 function(add_integration_test _testname)
 
   add_test(NAME ${_testname} 

tests/CMakeLists.txt

@@ -0,0 +1,46 @@
+
+if(USE_EXTERNAL_CATCH2)
+  find_package(Catch2 REQUIRED)
+else()
+  message(STATUS "Fetching local copy of Catch2 library for unit-tests...")
+  # Build Catch2 with the default flags, to avoid generating warnings when we
+  # build it
+  set(CXX_FLAGS_CMAKE_USED ${CMAKE_CXX_FLAGS})
+  set(CMAKE_CXX_FLAGS ${CXX_FLAGS_CMAKE_DEFAULTS})
+  Include(FetchContent)
+  FetchContent_Declare(
+    Catch2
+    GIT_REPOSITORY https://github.com/catchorg/Catch2.git
+    GIT_TAG        037ddbc75cc5e58b93cf5a010a94b32333ad824d
+    )
+  FetchContent_MakeAvailable(Catch2)
+  set(CMAKE_MODULE_PATH ${Catch2_SOURCE_DIR}/extras ${CMAKE_MODULE_PATH})
+  # Disable clang-tidy on external contents
+  set_target_properties(Catch2 PROPERTIES CXX_CLANG_TIDY "")
+
+  # Hack around the fact, that the include directories are not declared as
+  # SYSTEM for the targets defined this way. Otherwise warnings can still occur
+  # in Catch2 code when templates are evaluated (which happens quite a bit)
+  get_target_property(CATCH2_IF_INC_DIRS Catch2 INTERFACE_INCLUDE_DIRECTORIES)
+  set_target_properties(Catch2 PROPERTIES INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${CATCH2_IF_INC_DIRS}")
+
+  # Reset the flags
+  set(CMAKE_CXX_FLAGS ${CXX_FLAGS_CMAKE_USED})
+
+endif()
+
+
+# The unittests are a bit better and they are labelled so we can put together a
+# list of labels that we want to ignore
+set(filter_tests "")
+
+add_executable(bench algorithms.cpp myutils.cpp)
+target_link_libraries(bench PUBLIC FCCAnalyses gfortran PRIVATE Catch2::Catch2WithMain)
+target_include_directories(bench PUBLIC ${VDT_INCLUDE_DIR})
+
+include(Catch)
+catch_discover_tests(bench
+  WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}
+  TEST_PREFIX "B_" # make it possible to filter easily with -R ^UT
+  TEST_SPEC ${filter_tests} # discover only tests that are known to not fail
+)

tests/benchmark/CMakeLists.txt

@@ -0,0 +1,15 @@
+#include "FCCAnalyses/Algorithms.h"
+
+#include "catch2/catch_test_macros.hpp"
+#include <catch2/benchmark/catch_benchmark.hpp>
+
+
+TEST_CASE("calculate_thrust", "[algorithms]") {
+  ROOT::VecOps::RVec<float> x {0., 1., 3., 7., 11., 3.};
+  ROOT::VecOps::RVec<float> y {0., -1., 3., -7., -11., .3};
+  ROOT::VecOps::RVec<float> z {5., -3., 1., 4., 2., -4};
+
+  BENCHMARK("calculate_thrust short vector") {
+        return FCCAnalyses::Algorithms::calculate_thrust()(x, y, z);
+  };
+}

tests/benchmark/algorithms.cpp

@@ -0,0 +1,15 @@
+#include "FCCAnalyses/myUtils.h"
+
+#include <catch2/catch_test_macros.hpp>
+#include <catch2/benchmark/catch_benchmark.hpp>
+
+TEST_CASE("isPV", "[basics]") {
+  edm4hep::ReconstructedParticleData p;
+  p.tracks_begin = 0;
+  ROOT::VecOps::RVec<int> index1 = {3, 0, 7};
+  ROOT::VecOps::RVec<int> index2 = {3, 1, 7};
+
+  BENCHMARK("isPV") {
+        return FCCAnalyses::myUtils::isPV(p, index1);
+  };
+}