algorithms/reco/HadronicFinalState.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Tyler Kutz
0003
0004 #include <Math/GenVector/LorentzVector.h>
0005 #include <Math/GenVector/PxPyPzE4D.h>
0006 #include <Math/Vector4Dfwd.h>
0007 #include <edm4eic/HadronicFinalStateCollection.h>
0008 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0009 #include <edm4eic/ReconstructedParticleCollection.h>
0010 #include <edm4hep/MCParticleCollection.h>
0011 #include <edm4hep/Vector3f.h>
0012 #include <podio/ObjectID.h>
0013 #include <algorithm>
0014 #include <cmath>
0015 #include <tuple>
0016
0017 #include "Beam.h"
0018 #include "Boost.h"
0019 #include "HadronicFinalState.h"
0020
0021 using ROOT::Math::PxPyPzEVector;
0022
0023 namespace eicrecon {
0024
0025 void HadronicFinalState::init() {}
0026
0027 void HadronicFinalState::process(const HadronicFinalState::Input& input,
0028                                  const HadronicFinalState::Output& output) const {
0029
0030   const auto [mc_beam_electrons, mc_beam_protons, mcparts, rcparts, rcassoc] = input;
0031   auto [hadronicfinalstate]                                                  = output;
0032
0033   // Get first (should be only) beam electron
0034   if (mc_beam_electrons->empty()) {
0035     debug("No beam electron found");
0036     return;
0037   }
0038   const auto& ei_particle = (*mc_beam_electrons)[0];
0039   const PxPyPzEVector ei(round_beam_four_momentum(ei_particle.getMomentum(),
0040                                                   m_particleSvc.particle(ei_particle.getPDG()).mass,
0041                                                   electron_beam_pz_set, 0.0));
0042
0043   // Get first (should be only) beam proton
0044   if (mc_beam_protons->empty()) {
0045     debug("No beam hadron found");
0046     return;
0047   }
0048   const auto& pi_particle = (*mc_beam_protons)[0];
0049   const PxPyPzEVector pi(round_beam_four_momentum(pi_particle.getMomentum(),
0050                                                   m_particleSvc.particle(pi_particle.getPDG()).mass,
0051                                                   hadron_beam_pz_set, m_crossingAngle));
0052
0053   // Get first scattered electron from full MCParticles collection
0054   if (mcparts == nullptr) {
0055     debug("No MCParticles collection available");
0056     return;
0057   }
0058   const auto ef_coll = find_first_scattered_electron(mcparts);
0059   if (ef_coll.empty()) {
0060     debug("No truth scattered electron found");
0061     return;
0062   }
0063
0064   // Check if associations are available
0065   if (!rcassoc) {
0066     debug("No associations available");
0067     return;
0068   }
0069
0070   // Associate first scattered electron with reconstructed electrons
0071   const auto ef_assoc = std::find_if(rcassoc->begin(), rcassoc->end(), [&ef_coll](const auto& a) {
0072     return a.getSim().getObjectID() == ef_coll[0].getObjectID();
0073   });
0074
0075   if (ef_assoc == rcassoc->end()) {
0076     debug("Truth scattered electron not in reconstructed particles");
0077     return;
0078   }
0079   const auto ef_rc{(*ef_assoc).getRec()};
0080   const auto ef_rc_id{ef_rc.getObjectID().index};
0081
0082   // Sums in colinear frame
0083   double pxsum = 0;
0084   double pysum = 0;
0085   double pzsum = 0;
0086   double Esum  = 0;
0087
0088   // Get boost to colinear frame
0089   auto boost = determine_boost(ei, pi);
0090
0091   auto hfs = hadronicfinalstate->create(0., 0., 0.);
0092
0093   for (const auto& p : *rcparts) {
0094     // Check if it's the scattered electron
0095     if (p.getObjectID().index != ef_rc_id) {
0096       // Lorentz vector in lab frame
0097       PxPyPzEVector hf_lab(p.getMomentum().x, p.getMomentum().y, p.getMomentum().z, p.getEnergy());
0098       // Boost to colinear frame
0099       PxPyPzEVector hf_boosted = apply_boost(boost, hf_lab);
0100
0101       pxsum += hf_boosted.Px();
0102       pysum += hf_boosted.Py();
0103       pzsum += hf_boosted.Pz();
0104       Esum += hf_boosted.E();
0105
0106       hfs.addToHadrons(p);
0107     }
0108   }
0109
0110   // Hadronic final state calculations
0111   auto sigma = Esum - pzsum;
0112   auto pT    = sqrt(pxsum * pxsum + pysum * pysum);
0113   auto gamma = acos((pT * pT - sigma * sigma) / (pT * pT + sigma * sigma));
0114
0115   hfs.setSigma(sigma);
0116   hfs.setPT(pT);
0117   hfs.setGamma(gamma);
0118
0119   // Sigma zero or negative
0120   if (sigma <= 0) {
0121     debug("Sigma zero or negative");
0122     return;
0123   }
0124
0125   debug("sigma_h, pT_h, gamma_h = {},{},{}", hfs.getSigma(), hfs.getPT(), hfs.getGamma());
0126 }
0127
0128 } // namespace eicrecon