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