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 <algorithm>
0010 #include <format>
0011 #include <set>
0012 #include <stdexcept>
0013 #include <vector>
0014 #include <cmath>
0015
0016 using ROOT::Math::PxPyPzEVector;
0017
0018 namespace eicrecon {
0019
0020 template <class T> auto find_first_with_pdg(const T* parts, const std::set<int32_t>& pdg) {
0021   T c;
0022   c.setSubsetCollection();
0023   const auto it = std::find_if(parts->begin(), parts->end(),
0024                                [&pdg](const auto& p) { return pdg.count(p.getPDG()) > 0; });
0025   if (it != parts->end()) {
0026     c.push_back(*it);
0027   }
0028   return c;
0029 }
0030
0031 template <class T>
0032 auto find_first_with_status_pdg(const T* parts, const std::set<int32_t>& status,
0033                                 const std::set<int32_t>& pdg) {
0034   T c;
0035   c.setSubsetCollection();
0036   const auto it = std::find_if(parts->begin(), parts->end(), [&status, &pdg](const auto& p) {
0037     return status.count(p.getGeneratorStatus()) > 0 && pdg.count(p.getPDG()) > 0;
0038   });
0039   if (it != parts->end()) {
0040     c.push_back(*it);
0041   }
0042   return c;
0043 }
0044
0045 inline auto find_first_beam_electron(const edm4hep::MCParticleCollection* mcparts) {
0046   return find_first_with_status_pdg(mcparts, {4}, {11});
0047 }
0048
0049 inline auto find_first_beam_hadron(const edm4hep::MCParticleCollection* mcparts) {
0050   return find_first_with_status_pdg(mcparts, {4}, {2212, 2112});
0051 }
0052
0053 inline auto find_first_scattered_electron(const edm4hep::MCParticleCollection* mcparts) {
0054   return find_first_with_status_pdg(mcparts, {1}, {11});
0055 }
0056
0057 inline auto find_first_scattered_electron(const edm4eic::ReconstructedParticleCollection* rcparts) {
0058   return find_first_with_pdg(rcparts, {11});
0059 }
0060
0061 // Canonical beam momentum allowlists used by all kinematics algorithms.
0062 // Electron beam: negative pz (beam goes in -z direction).
0063 inline const std::vector<float> electron_beam_pz_set{-5.0, -9.0, -10.0, -18.0};
0064 // Hadron beam: positive pz (beam goes in +z direction).
0065 // Proton entries: 41, 100, 130, 250, 275 GeV.
0066 // He-3 entries (Z=2, A=3): per-nucleon momenta = proton × (Z/A) = × 2/3.
0067 //   He-3 at proton-250-equivalent rigidity: 250 × 2/3 ≈ 166.7 GeV/nucleon (e.g. 9×166 GeV).
0068 //   He-3 at proton-275-equivalent rigidity: 275 × 2/3 ≈ 183.3 GeV/nucleon.
0069 // Add further ion/nucleon momenta here as new beam configurations are commissioned.
0070 inline const std::vector<float> hadron_beam_pz_set{41.0, 100.0, 130.0, 166.7, 183.3, 250.0, 275.0};
0071
0072 template <typename Vector3>
0073 PxPyPzEVector round_beam_four_momentum(const Vector3& p_in, const float mass,
0074                                        const std::vector<float>& pz_set,
0075                                        const float crossing_angle = 0.0) {
0076   // Find the closest pz within 10% relative tolerance
0077   float best_pz    = 0.0F;
0078   float best_err   = 0.1F; // 10% tolerance — entries above this are not accepted
0079   bool found_match = false;
0080   for (const auto& pz : pz_set) {
0081     const float err = std::abs(p_in.z / pz - 1);
0082     if (err < best_err) {
0083       best_err    = err;
0084       best_pz     = pz;
0085       found_match = true;
0086     }
0087   }
0088   if (!found_match) {
0089     throw std::runtime_error(
0090         std::format("round_beam_four_momentum: no match for beam momentum {:.3f} GeV within 10% "
0091                     "of any of the allowed values",
0092                     p_in.z));
0093   }
0094   PxPyPzEVector p_out;
0095   p_out.SetPz(best_pz);
0096   p_out.SetPx(p_out.Pz() * sin(crossing_angle));
0097   p_out.SetPz(p_out.Pz() * cos(crossing_angle));
0098   p_out.SetE(std::hypot(p_out.Px(), p_out.Pz(), mass));
0099   return p_out;
0100 }
0101
0102 } // namespace eicrecon