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>
0016   auto find_first_with_pdg(
0017       const T* parts,
0018       const std::set<int32_t>& pdg) {
0019     T c;
0020     c.setSubsetCollection();
0021     for (const auto& p: *parts) {
0022       if (pdg.count(p.getPDG()) > 0) {
0023         c.push_back(p);
0024         break;
0025       }
0026     }
0027     return c;
0028   }
0029
0030   template<class T>
0031   auto find_first_with_status_pdg(
0032       const T* parts,
0033       const std::set<int32_t>& status,
0034       const std::set<int32_t>& pdg) {
0035     T c;
0036     c.setSubsetCollection();
0037     for (const auto& p: *parts) {
0038       if (status.count(p.getGeneratorStatus()) > 0 &&
0039           pdg.count(p.getPDG()) > 0) {
0040         c.push_back(p);
0041         break;
0042       }
0043     }
0044     return c;
0045   }
0046
0047   inline auto find_first_beam_electron(const edm4hep::MCParticleCollection* mcparts) {
0048     return find_first_with_status_pdg(mcparts, {4}, {11});
0049   }
0050
0051   inline auto find_first_beam_hadron(const edm4hep::MCParticleCollection* mcparts) {
0052     return find_first_with_status_pdg(mcparts, {4}, {2212, 2112});
0053   }
0054
0055   inline auto find_first_scattered_electron(const edm4hep::MCParticleCollection* mcparts) {
0056     return find_first_with_status_pdg(mcparts, {1}, {11});
0057   }
0058
0059   inline auto find_first_scattered_electron(const edm4eic::ReconstructedParticleCollection* rcparts) {
0060     return find_first_with_pdg(rcparts, {11});
0061   }
0062
0063   inline
0064   PxPyPzEVector
0065   round_beam_four_momentum(
0066       const edm4hep::Vector3f& p_in,
0067       const float mass,
0068       const std::vector<float>& pz_set,
0069       const float crossing_angle = 0.0) {
0070     PxPyPzEVector p_out;
0071     for (const auto& pz : pz_set) {
0072       if (fabs(p_in.z / pz - 1) < 0.1) {
0073         p_out.SetPz(pz);
0074         break;
0075       }
0076     }
0077     p_out.SetPx(p_out.Pz() * sin(crossing_angle));
0078     p_out.SetPz(p_out.Pz() * cos(crossing_angle));
0079     p_out.SetE(std::hypot(p_out.Px(), p_out.Pz(), mass));
0080     return p_out;
0081   }
0082
0083 } // namespace eicrecon