algorithms/reco/Beam.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Wouter Deconinck
0003
0004 #pragma once
0005
0006 #include <Math/Vector4D.h>
0007 #include <edm4hep/MCParticleCollection.h>
0008 #include <edm4eic/ReconstructedParticleCollection.h>
0009 #include <set>
0010
0011 using ROOT::Math::PxPyPzEVector;
0012
0013 namespace eicrecon {
0014
0015 template <class T> auto find_first_with_pdg(const T* parts, const std::set<int32_t>& pdg) {
0016   T c;
0017   c.setSubsetCollection();
0018   for (const auto& p : *parts) {
0019     if (pdg.count(p.getPDG()) > 0) {
0020       c.push_back(p);
0021       break;
0022     }
0023   }
0024   return c;
0025 }
0026
0027 template <class T>
0028 auto find_first_with_status_pdg(const T* parts, const std::set<int32_t>& status,
0029                                 const std::set<int32_t>& pdg) {
0030   T c;
0031   c.setSubsetCollection();
0032   for (const auto& p : *parts) {
0033     if (status.count(p.getGeneratorStatus()) > 0 && pdg.count(p.getPDG()) > 0) {
0034       c.push_back(p);
0035       break;
0036     }
0037   }
0038   return c;
0039 }
0040
0041 inline auto find_first_beam_electron(const edm4hep::MCParticleCollection* mcparts) {
0042   return find_first_with_status_pdg(mcparts, {4}, {11});
0043 }
0044
0045 inline auto find_first_beam_hadron(const edm4hep::MCParticleCollection* mcparts) {
0046   return find_first_with_status_pdg(mcparts, {4}, {2212, 2112});
0047 }
0048
0049 inline auto find_first_scattered_electron(const edm4hep::MCParticleCollection* mcparts) {
0050   return find_first_with_status_pdg(mcparts, {1}, {11});
0051 }
0052
0053 inline auto find_first_scattered_electron(const edm4eic::ReconstructedParticleCollection* rcparts) {
0054   return find_first_with_pdg(rcparts, {11});
0055 }
0056
0057 template <typename Vector3>
0058 PxPyPzEVector round_beam_four_momentum(const Vector3& p_in, const float mass,
0059                                        const std::vector<float>& pz_set,
0060                                        const float crossing_angle = 0.0) {
0061   PxPyPzEVector p_out;
0062   for (const auto& pz : pz_set) {
0063     if (std::abs(p_in.z / pz - 1) < 0.1) {
0064       p_out.SetPz(pz);
0065       break;
0066     }
0067   }
0068   p_out.SetPx(p_out.Pz() * sin(crossing_angle));
0069   p_out.SetPz(p_out.Pz() * cos(crossing_angle));
0070   p_out.SetE(std::hypot(p_out.Px(), p_out.Pz(), mass));
0071   return p_out;
0072 }
0073
0074 } // namespace eicrecon