include/common_bench/util.h

0001 #ifndef UTIL_H
0002 #define UTIL_H
0003
0004
0005 #include <algorithm>
0006 #include <cmath>
0007 #include <exception>
0008 #include <fmt/core.h>
0009 #include <limits>
0010 #include <string>
0011 #include <string_view>
0012 #include <vector>
0013 #include "TF1.h"
0014 #include "TFitResult.h"
0015 #include "TFitResultPtr.h"
0016
0017 #include <Math/Vector4D.h>
0018
0019 #include "edm4hep/MCParticleCollection.h"
0020 #include "edm4eic/TrackParametersCollection.h"
0021 #include "edm4eic/ReconstructedParticleCollection.h"
0022 #include "edm4eic/ReconstructedParticleData.h"
0023
0024 namespace common_bench {
0025
0026 /** Exception definition for unknown particle errors.
0027  *  \todo  Fixme: A utility exception base class should be included in the
0028  * analysis utility library, so we can skip most of this boilerplate
0029  */
0030 class unknown_particle_error : public std::exception {
0031 public:
0032   unknown_particle_error(std::string_view particle)
0033       : m_particle{particle},
0034         m_msg(fmt::format("Unknown particle type: {}", m_particle)) {}
0035   const char *what() const noexcept override {
0036     return m_msg.c_str();
0037   }
0038   const char *type() const noexcept { return "unknown_particle_error"; }
0039
0040 private:
0041   const std::string m_particle;
0042   const std::string m_msg;
0043 };
0044
0045 /** Simple function for pdg masses.
0046  *  Return the appropriate PDG mass for the particles
0047  *  we care about for this process.
0048  *  \todo FIXME: consider something more robust (maybe based on hepPDT) to the
0049  *         analysis utility library
0050  */
0051 inline double get_pdg_mass(std::string_view part) {
0052   if (part == "electron") {
0053     return 0.0005109989461;
0054   } else if (part == "muon") {
0055     return .1056583745;
0056   } else if (part == "jpsi") {
0057     return 3.0969;
0058   } else if (part == "upsilon") {
0059     return 9.49630;
0060   } else if (part == "proton") {
0061     return 0.938272;
0062   } else {
0063     throw unknown_particle_error{part};
0064   }
0065 }
0066
0067 /** Compute momentum from track parameters.
0068  * Get a vector of 4-momenta from raw tracking info, using an externally
0069  * provided particle mass assumption. //outputTrackParameters
0070  */
0071 inline auto
0072 momenta_from_tracking(const std::vector<edm4eic::TrackParametersData> &tracks,
0073                       const double mass) {
0074   std::vector<ROOT::Math::PxPyPzMVector> momenta{tracks.size()};
0075   // transform our raw tracker info into proper 4-momenta
0076   std::transform(tracks.begin(), tracks.end(), momenta.begin(),
0077                  [mass](const auto &track) {
0078                    // make sure we don't divide by zero
0079                    if (fabs(track.qOverP) < 1e-9) {
0080                      return ROOT::Math::PxPyPzMVector{};
0081                    }
0082                    const double p = fabs(1. / track.qOverP);
0083                    const double px = p * cos(track.phi) * sin(track.theta);
0084                    const double py = p * sin(track.phi) * sin(track.theta);
0085                    const double pz = p * cos(track.theta);
0086                    return ROOT::Math::PxPyPzMVector{px, py, pz, mass};
0087                  });
0088   return momenta;
0089 }
0090
0091 /** Helper to get momentum 4 vector.
0092  */
0093 inline auto
0094 momenta_RC(const std::vector<edm4eic::ReconstructedParticleData> &parts) {
0095   std::vector<ROOT::Math::PxPyPzMVector> momenta{parts.size()};
0096   // transform our raw tracker info into proper 4-momenta
0097   std::transform(parts.begin(), parts.end(), momenta.begin(),
0098                  [](const auto &part) {
0099                    return ROOT::Math::PxPyPzMVector{part.momentum.x, part.momentum.y,
0100                                                     part.momentum.z, part.mass};
0101                  });
0102   return momenta;
0103 }
0104
0105 /** Get a vector of 4-momenta from the simulation data.
0106  *  \todo Add PID selector (maybe using ranges?)
0107  */
0108 inline auto
0109 momenta_from_simulation(const std::vector<edm4hep::MCParticleData> &parts) {
0110   std::vector<ROOT::Math::PxPyPzMVector> momenta{parts.size()};
0111   // transform our simulation particle data into 4-momenta
0112   std::transform(parts.begin(), parts.end(), momenta.begin(),
0113                  [](const auto &part) {
0114                    return ROOT::Math::PxPyPzMVector{part.momentum.x, part.momentum.y,
0115                                                     part.momentum.z, part.mass};
0116                  });
0117   return momenta;
0118 }
0119
0120 /**  find the decay lepton pair.
0121  *  Find the decay pair candidates from a vector of particles (parts),
0122  *  with invariant mass closest to a desired value (pdg_mass)
0123  */
0124 inline std::pair<ROOT::Math::PxPyPzMVector, ROOT::Math::PxPyPzMVector>
0125 find_decay_pair(const std::vector<ROOT::Math::PxPyPzMVector> &parts,
0126                 const double pdg_mass, const double daughter_mass) {
0127   int first = -1;
0128   int second = -1;
0129   double best_mass = -1;
0130
0131   // go through all particle combinatorics, calculate the invariant mass
0132   // for each combination, and remember which combination is the closest
0133   // to the desired pdg_mass
0134   for (size_t i = 0; i < parts.size(); ++i) {
0135     if (fabs(parts[i].mass() - daughter_mass) / daughter_mass > 0.01)
0136       continue;
0137     for (size_t j = i + 1; j < parts.size(); ++j) {
0138       if (fabs(parts[j].mass() - daughter_mass) / daughter_mass > 0.01)
0139         continue;
0140       const double new_mass{(parts[i] + parts[j]).mass()};
0141       if (fabs(new_mass - pdg_mass) < fabs(best_mass - pdg_mass)) {
0142         first = i;
0143         second = j;
0144         best_mass = new_mass;
0145       }
0146     }
0147   }
0148   if (first < 0) {
0149     return {{}, {}};
0150   }
0151   return {parts[first], parts[second]};
0152 }
0153
0154 /**
0155  * Calculate the magnitude of the momentum of a vector of 4-vectors
0156  */
0157 inline auto mom(const std::vector<ROOT::Math::PxPyPzMVector> &momenta) {
0158   std::vector<double> P(momenta.size());
0159   // transform our raw tracker info into proper 4-momenta
0160   std::transform(momenta.begin(), momenta.end(), P.begin(),
0161                  [](const auto &mom) { return mom.P(); });
0162   return P;
0163 }
0164
0165 /**
0166  * Calculate the transverse momentum of a vector of 4-vectors
0167  */
0168 inline auto pt(const std::vector<ROOT::Math::PxPyPzMVector> &momenta) {
0169   std::vector<double> pt(momenta.size());
0170   // transform our raw tracker info into proper 4-momenta
0171   std::transform(momenta.begin(), momenta.end(), pt.begin(),
0172                  [](const auto &mom) { return mom.pt(); });
0173   return pt;
0174 }
0175
0176 /** Calculate the azimuthal angle phi of a vector of 4-vectors.
0177  */
0178 inline auto phi(const std::vector<ROOT::Math::PxPyPzMVector> &momenta) {
0179   std::vector<double> phi(momenta.size());
0180   // transform our raw tracker info into proper 4-momenta
0181   std::transform(momenta.begin(), momenta.end(), phi.begin(),
0182                  [](const auto &mom) { return mom.phi(); });
0183   return phi;
0184 }
0185
0186 /** Calculate the pseudo-rapidity of a vector of particles.
0187  */
0188 inline auto eta(const std::vector<ROOT::Math::PxPyPzMVector> &momenta) {
0189   std::vector<double> eta(momenta.size());
0190   // transform our raw tracker info into proper 4-momenta
0191   std::transform(momenta.begin(), momenta.end(), eta.begin(),
0192                  [](const auto &mom) { return mom.eta(); });
0193   return eta;
0194 }
0195
0196 } // namespace common_bench
0197
0198 #endif