algorithms/reco/UndoAfterBurner.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Alex Jentsch, Jihee Kim, Brian Page
0003 //
0004
0005 #include "UndoAfterBurner.h"
0006
0007 #include <Math/GenVector/Boost.h>
0008 #include <Math/GenVector/Cartesian3D.h>
0009 #include <Math/GenVector/LorentzVector.h>
0010 #include <Math/GenVector/PxPyPzE4D.h>
0011 #include <Math/GenVector/RotationX.h>
0012 #include <Math/GenVector/RotationY.h>
0013 #include <Math/Vector4Dfwd.h>
0014 #include <TMath.h>
0015 #include <edm4hep/Vector3d.h>
0016 #include <podio/ObjectID.h>
0017 #include <podio/RelationRange.h>
0018 #include <cstddef>
0019 #include <tuple>
0020 #include <unordered_map>
0021 #include <utility>
0022 #include <vector>
0023
0024 #include "algorithms/reco/Beam.h"
0025 #include "algorithms/reco/UndoAfterBurnerConfig.h"
0026
0027 void eicrecon::UndoAfterBurner::init() {}
0028
0029 void eicrecon::UndoAfterBurner::process(const UndoAfterBurner::Input& input,
0030                                         const UndoAfterBurner::Output& output) const {
0031
0032   const auto [mcparts]   = input;
0033   auto [outputParticles] = output;
0034
0035   bool pidAssumePionMass = m_cfg.m_pid_assume_pion_mass;
0036   double crossingAngle   = m_cfg.m_crossing_angle;
0037   bool correctBeamFX     = m_cfg.m_correct_beam_FX;
0038   bool pidUseMCTruth     = m_cfg.m_pid_use_MC_truth;
0039
0040   bool hasBeamHadron = true;
0041   bool hasBeamLepton = true;
0042
0043   //read MCParticles information and "postburn" to remove the afterburner effects.
0044   //The output is then the original MC input produced by the generator.
0045
0046   ROOT::Math::PxPyPzEVector e_beam(0., 0., 0., 0.);
0047   ROOT::Math::PxPyPzEVector h_beam(0., 0., 0., 0.);
0048
0049   auto incoming_lepton = find_first_beam_electron(mcparts);
0050   if (incoming_lepton.empty()) {
0051     debug("No beam electron found -- particleGun input");
0052     hasBeamLepton = false;
0053   }
0054
0055   auto incoming_hadron = find_first_beam_hadron(mcparts);
0056   if (incoming_hadron.empty()) {
0057     debug("No beam hadron found -- particleGun input");
0058     hasBeamHadron = false;
0059   }
0060
0061   if ((hasBeamHadron && !hasBeamLepton) || (!hasBeamHadron && hasBeamLepton)) {
0062     debug("Only one beam defined! Not a possible configuration!");
0063     return;
0064   }
0065
0066   // Handling for FF particle gun input!!
0067   if (!hasBeamHadron || !hasBeamLepton) {
0068     for (const auto& p : *mcparts) {
0069       if ((p.getPDG() == 2212 || p.getPDG() == 2112)) { //look for "gun" proton/neutron
0070         hasBeamHadron = true;
0071         h_beam.SetPxPyPzE(crossingAngle * p.getEnergy(), 0.0, p.getEnergy(), p.getEnergy());
0072         if (p.getEnergy() > 270.0 && p.getEnergy() < 280.0) {
0073           hasBeamLepton = true;
0074           e_beam.SetPxPyPzE(0.0, 0.0, -18.0, 18.0);
0075         }
0076       }
0077     }
0078
0079   } else {
0080
0081     if (correctBeamFX) {
0082
0083       h_beam.SetPxPyPzE(incoming_hadron[0].getMomentum().x, incoming_hadron[0].getMomentum().y,
0084                         incoming_hadron[0].getMomentum().z, incoming_hadron[0].getEnergy());
0085       e_beam.SetPxPyPzE(incoming_lepton[0].getMomentum().x, incoming_lepton[0].getMomentum().y,
0086                         incoming_lepton[0].getMomentum().z, incoming_lepton[0].getEnergy());
0087
0088     } else {
0089
0090       h_beam.SetPxPyPzE(crossingAngle * incoming_hadron[0].getEnergy(), 0.0,
0091                         incoming_hadron[0].getEnergy(), incoming_hadron[0].getEnergy());
0092       e_beam.SetPxPyPzE(0.0, 0.0, -incoming_lepton[0].getEnergy(), incoming_lepton[0].getEnergy());
0093     }
0094   }
0095
0096   // Bail out if either beam is missing, since this leads to division by zero or unphysical boosts
0097   if (!hasBeamHadron || !hasBeamLepton) {
0098     return;
0099   }
0100
0101   // Calculate boost vectors and rotations here
0102   ROOT::Math::PxPyPzEVector cm_frame_boost = e_beam + h_beam;
0103   ROOT::Math::Cartesian3D beta(-cm_frame_boost.Px() / cm_frame_boost.E(),
0104                                -cm_frame_boost.Py() / cm_frame_boost.E(),
0105                                -cm_frame_boost.Pz() / cm_frame_boost.E());
0106   ROOT::Math::Boost boostVector(beta);
0107
0108   // Boost to CM frame
0109   e_beam = boostVector(e_beam);
0110   h_beam = boostVector(h_beam);
0111
0112   double rotationAngleY = -1.0 * TMath::ATan2(h_beam.Px(), h_beam.Pz());
0113   double rotationAngleX = 1.0 * TMath::ATan2(h_beam.Py(), h_beam.Pz());
0114
0115   ROOT::Math::RotationY rotationAboutY(rotationAngleY);
0116   ROOT::Math::RotationX rotationAboutX(rotationAngleX);
0117
0118   // Boost back to proper head-on frame
0119   ROOT::Math::PxPyPzEVector head_on_frame_boost(0., 0., cm_frame_boost.Pz(), cm_frame_boost.E());
0120   ROOT::Math::Boost headOnBoostVector(head_on_frame_boost.Px() / head_on_frame_boost.E(),
0121                                       head_on_frame_boost.Py() / head_on_frame_boost.E(),
0122                                       head_on_frame_boost.Pz() / head_on_frame_boost.E());
0123
0124   // Now, loop through events and apply operations to the MCparticles
0125   const int maxGenStatus = m_cfg.m_max_gen_status;
0126
0127   // Helper lambda to check if particle should be processed
0128   auto shouldProcessParticle = [maxGenStatus](const edm4hep::MCParticle& p) {
0129     if (p.isCreatedInSimulation()) {
0130       return false;
0131     }
0132     // Filter by generator status to exclude background particles and conserve memory
0133     if (maxGenStatus >= 0 && p.getGeneratorStatus() > maxGenStatus) {
0134       return false;
0135     }
0136     return true;
0137   };
0138
0139   // Map from input MCParticle ObjectID to output MCParticle index
0140   std::unordered_map<podio::ObjectID, size_t> inputToOutputMap;
0141
0142   // First pass: create all output particles
0143   for (const auto& p : *mcparts) {
0144     if (!shouldProcessParticle(p)) {
0145       continue;
0146     }
0147
0148     ROOT::Math::PxPyPzEVector mc(p.getMomentum().x, p.getMomentum().y, p.getMomentum().z,
0149                                  p.getEnergy());
0150
0151     mc = boostVector(mc);
0152     mc = rotationAboutY(mc);
0153     mc = rotationAboutX(mc);
0154     mc = headOnBoostVector(mc);
0155
0156     decltype(edm4hep::MCParticleData::momentum) mcMom(mc.Px(), mc.Py(), mc.Pz());
0157
0158     // Create new particle without cloning relationships (which would point to input collection)
0159     // We manually copy only the fields we need, and will add relationships in the second pass
0160     auto MCTrack = outputParticles->create();
0161     MCTrack.setPDG(pidUseMCTruth ? p.getPDG() : (pidAssumePionMass ? 211 : p.getPDG()));
0162     MCTrack.setGeneratorStatus(p.getGeneratorStatus());
0163     MCTrack.setSimulatorStatus(p.getSimulatorStatus());
0164     MCTrack.setCharge(p.getCharge());
0165     MCTrack.setTime(p.getTime());
0166     MCTrack.setMass(pidUseMCTruth ? p.getMass() : (pidAssumePionMass ? 0.13957 : p.getMass()));
0167     MCTrack.setVertex(p.getVertex());
0168     MCTrack.setEndpoint(p.getEndpoint());
0169     MCTrack.setMomentum(mcMom);
0170     MCTrack.setMomentumAtEndpoint(p.getMomentumAtEndpoint());
0171     MCTrack.setHelicity(p.getHelicity());
0172     // Store mapping from input particle ObjectID to output particle index
0173     inputToOutputMap[p.getObjectID()] = outputParticles->size() - 1;
0174   }
0175
0176   // Second pass: establish parent-daughter relationships
0177   for (const auto& p : *mcparts) {
0178     if (!shouldProcessParticle(p)) {
0179       continue;
0180     }
0181
0182     // Get the output particle corresponding to this input particle
0183     auto outputIter = inputToOutputMap.find(p.getObjectID());
0184     if (outputIter == inputToOutputMap.end()) {
0185       continue;
0186     }
0187     auto outputParticle = outputParticles->at(outputIter->second);
0188
0189     // Add parent relationships
0190     for (const auto& parent : p.getParents()) {
0191       auto parentIter = inputToOutputMap.find(parent.getObjectID());
0192       if (parentIter != inputToOutputMap.end()) {
0193         auto outputParent = outputParticles->at(parentIter->second);
0194         outputParticle.addToParents(outputParent);
0195       }
0196     }
0197
0198     // Add daughter relationships
0199     for (const auto& daughter : p.getDaughters()) {
0200       auto daughterIter = inputToOutputMap.find(daughter.getObjectID());
0201       if (daughterIter != inputToOutputMap.end()) {
0202         auto outputDaughter = outputParticles->at(daughterIter->second);
0203         outputParticle.addToDaughters(outputDaughter);
0204       }
0205     }
0206   }
0207 }