edm4hep/utils/kinematics.h

0001 #ifndef EDM4HEP_UTILS_KINEMATICS_H
0002 #define EDM4HEP_UTILS_KINEMATICS_H
0003
0004 #include "Math/Vector4D.h"
0005
0006 #include <cmath>
0007
0008 namespace edm4hep {
0009 /**
0010  * A LorentzVector with (px, py, pz) and M
0011  */
0012 using LorentzVectorM = ROOT::Math::PxPyPzMVector;
0013
0014 /**
0015  * A LorentzVector with (px, py, pz) and E
0016  */
0017 using LorentzVectorE = ROOT::Math::PxPyPzEVector;
0018
0019 namespace utils {
0020
0021   /**
0022    * Get the transverse momentum from a Particle
0023    */
0024   template <typename ParticleT>
0025   inline float pT(ParticleT const& p) {
0026     return std::hypot(p.getMomentum()[0], p.getMomentum()[1]);
0027   }
0028
0029   /**
0030    * Get the transverse momentum from a Particle
0031    */
0032   template <typename ParticleT>
0033   inline float pt(ParticleT const& p) {
0034     return pT(p);
0035   }
0036
0037   /**
0038    * Get the total momentum from a Particle
0039    */
0040   template <typename ParticleT>
0041   inline float p(ParticleT const& part) {
0042     const auto mom = part.getMomentum();
0043     return std::hypot(mom[0], mom[1], mom[2]);
0044   }
0045
0046   namespace detail {
0047     /**
0048      * Tag struct for getting 4-momenta with momentum + mass
0049      */
0050     struct UseMassTag {
0051       using type = ::edm4hep::LorentzVectorM;
0052       static constexpr bool has_value = false;
0053     };
0054
0055     /**
0056      * Tag struct for getting 4-momenta with momentum + energy
0057      */
0058     struct UseEnergyTag {
0059       using type = ::edm4hep::LorentzVectorE;
0060       static constexpr bool has_value = false;
0061     };
0062
0063     /**
0064      * Tag struct holding an additional value for getting 4 momenta with arbitrary
0065      * values instead of the one from the Particle.
0066      */
0067     template <typename T, typename TagT>
0068     struct TaggedUserValue {
0069       TaggedUserValue() = default;
0070       TaggedUserValue(T v) : value(v) {}
0071       using type = typename TagT::type;
0072       static constexpr bool has_value = true;
0073       T value{};
0074     };
0075
0076     /**
0077      * Tag-dispatched implementation for getting the 4-momentum from a particle
0078      * according to the desired type.
0079      */
0080     template <typename ParticleT, typename LorentzVectorTypeTag>
0081     inline typename LorentzVectorTypeTag::type p4(ParticleT const& part, [[maybe_unused]] LorentzVectorTypeTag* tag) {
0082       const auto mom = part.getMomentum();
0083       // Either the user wants to set a specific value
0084       if constexpr (LorentzVectorTypeTag::has_value) {
0085         return typename LorentzVectorTypeTag::type{mom[0], mom[1], mom[2], tag->value};
0086       }
0087
0088       // Or we take the one from the underlying particle
0089       if constexpr (std::is_same_v<typename LorentzVectorTypeTag::type, LorentzVectorM>) {
0090         return LorentzVectorM{mom[0], mom[1], mom[2], part.getMass()};
0091       }
0092       if constexpr (std::is_same_v<typename LorentzVectorTypeTag::type, LorentzVectorE>) {
0093         return LorentzVectorE{mom[0], mom[1], mom[2], part.getEnergy()};
0094       }
0095     }
0096   } // namespace detail
0097
0098   using UseMassTag = detail::UseMassTag;
0099   /**
0100    * Static tag to select the mass in 4 momentum vectors
0101    */
0102   constexpr static UseMassTag UseMass;
0103
0104   using UseEnergyTag = detail::UseEnergyTag;
0105   /**
0106    * Static tag to select the energy in 4 momentum vectors
0107    */
0108   constexpr static UseEnergyTag UseEnergy;
0109
0110   /**
0111    * Class to inject a user-defined mass into 4 momentum vectors in the call to
0112    * p4
0113    */
0114   using SetMass = detail::TaggedUserValue<float, UseMassTag>;
0115
0116   /**
0117    * Class to inject a user-defined energy into 4 momentum vectors in the call to
0118    * p4
0119    */
0120   using SetEnergy = detail::TaggedUserValue<float, UseEnergyTag>;
0121
0122   /**
0123    * Get the 4 momentum vector from a Particle using the momentum and the mass.
0124    */
0125   template <typename ParticleT>
0126   inline LorentzVectorM p4(ParticleT const& part) {
0127     return detail::p4(part, &UseMass);
0128   }
0129
0130   /**
0131    * Get the 4 momentum vector from a Particle. Use the second argument to choose
0132    * something other than the mass. E.g. the UseEnergy as second argument to get
0133    * the energy stored within the particle. Additionally it is possible to take
0134    * the momentum from the particle but set a specific mass or energy value by
0135    * using SetMass or SetEnergy as the second argument. The underlying particle
0136    * will not be changed in this case.
0137    *
0138    * NOTE: split into two overloads, in order to work around a short-coming of
0139    * cppyy: See https://github.com/wlav/cppyy/issues/78
0140    */
0141   template <typename ParticleT, typename LorentzVectorTag>
0142   inline typename LorentzVectorTag::type p4(ParticleT const& part, LorentzVectorTag tag) {
0143     return detail::p4(part, &tag);
0144   }
0145
0146 } // namespace utils
0147 } // namespace edm4hep
0148
0149 #endif