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/Vector3f.h>
0016 #include <gsl/pointers>
0017
0018 #include "algorithms/reco/Beam.h"
0019 #include "algorithms/reco/UndoAfterBurnerConfig.h"
0020
0021 void eicrecon::UndoAfterBurner::init() {
0022
0023 }
0024
0025 void eicrecon::UndoAfterBurner::process(
0026     const UndoAfterBurner::Input& input,
0027     const UndoAfterBurner::Output& output) const {
0028
0029     const auto [mcparts] = input;
0030     auto [outputParticles] = output;
0031
0032     bool      pidAssumePionMass = m_cfg.m_pid_assume_pion_mass;
0033     double    crossingAngle    = m_cfg.m_crossing_angle;
0034     double    pidPurity        = m_cfg.m_pid_purity;
0035     bool      correctBeamFX    = m_cfg.m_correct_beam_FX;
0036     bool      pidUseMCTruth    = m_cfg.m_pid_use_MC_truth;
0037
0038     bool      hasBeamHadron    = true;
0039     bool      hasBeamLepton    = true;
0040
0041     //read MCParticles information and "postburn" to remove the afterburner effects.
0042     //The output is then the original MC input produced by the generator.
0043
0044     ROOT::Math::PxPyPzEVector  e_beam(0.,0.,0.,0.);
0045     ROOT::Math::PxPyPzEVector  h_beam(0.,0.,0.,0.);
0046
0047     auto incoming_lepton = find_first_beam_electron(mcparts);
0048     if (incoming_lepton.size() == 0) {
0049         debug("No beam electron found -- particleGun input");
0050         hasBeamLepton = false;
0051     }
0052
0053     auto incoming_hadron = find_first_beam_hadron(mcparts);
0054     if (incoming_hadron.size() == 0) {
0055         debug("No beam hadron found -- particleGun input");
0056         hasBeamHadron = false;
0057     }
0058
0059     if((hasBeamHadron && !hasBeamLepton) || (!hasBeamHadron && hasBeamLepton)){
0060         debug("Only one beam defined! Not a possible configuration!");
0061         return;
0062     }
0063
0064     // Handling for FF particle gun input!!
0065     if (!hasBeamHadron || !hasBeamLepton) {
0066         for (const auto& p: *mcparts) {
0067             if((p.getPDG() == 2212 || p.getPDG() == 2112)) { //look for "gun" proton/neutron
0068                 hasBeamHadron = true;
0069                 h_beam.SetPxPyPzE(crossingAngle*p.getEnergy(), 0.0, p.getEnergy(), p.getEnergy());
0070                 if(p.getEnergy() > 270.0 && p.getEnergy() < 280.0){
0071                     hasBeamLepton = true;
0072                     e_beam.SetPxPyPzE(0.0, 0.0, -18.0, 18.0);
0073                 }
0074             }
0075         }
0076
0077     } else {
0078
0079         if (correctBeamFX) {
0080
0081             h_beam.SetPxPyPzE(incoming_hadron[0].getMomentum().x, incoming_hadron[0].getMomentum().y, incoming_hadron[0].getMomentum().z, incoming_hadron[0].getEnergy());
0082             e_beam.SetPxPyPzE(incoming_lepton[0].getMomentum().x, incoming_lepton[0].getMomentum().y, incoming_lepton[0].getMomentum().z, incoming_lepton[0].getEnergy());
0083
0084         } else {
0085
0086             h_beam.SetPxPyPzE(crossingAngle*incoming_hadron[0].getEnergy(), 0.0, incoming_hadron[0].getEnergy(), incoming_hadron[0].getEnergy());
0087             e_beam.SetPxPyPzE(0.0, 0.0, -incoming_lepton[0].getEnergy(), incoming_lepton[0].getEnergy());
0088
0089         }
0090     }
0091
0092     // Bail out if still no beam particles, since this leads to division by zero
0093     if (!hasBeamHadron && !hasBeamLepton) {
0094       return;
0095     }
0096
0097     // Calculate boost vectors and rotations here
0098     ROOT::Math::PxPyPzEVector cm_frame_boost = e_beam + h_beam;
0099     ROOT::Math::Cartesian3D beta(-cm_frame_boost.Px() / cm_frame_boost.E(), -cm_frame_boost.Py() / cm_frame_boost.E(), -cm_frame_boost.Pz() / cm_frame_boost.E());
0100     ROOT::Math::Boost boostVector(beta);
0101
0102     // Boost to CM frame
0103     e_beam = boostVector(e_beam);
0104     h_beam = boostVector(h_beam);
0105
0106     double rotationAngleY = -1.0*TMath::ATan2(h_beam.Px(), h_beam.Pz());
0107     double rotationAngleX = 1.0*TMath::ATan2(h_beam.Py(), h_beam.Pz());
0108
0109     ROOT::Math::RotationY rotationAboutY(rotationAngleY);
0110     ROOT::Math::RotationX rotationAboutX(rotationAngleX);
0111
0112     // Boost back to proper head-on frame
0113     ROOT::Math::PxPyPzEVector head_on_frame_boost(0., 0., cm_frame_boost.Pz(), cm_frame_boost.E());
0114     ROOT::Math::Boost headOnBoostVector(head_on_frame_boost.Px()/head_on_frame_boost.E(), head_on_frame_boost.Py()/head_on_frame_boost.E(), head_on_frame_boost.Pz()/head_on_frame_boost.E());
0115
0116     // Now, loop through events and apply operations to the MCparticles
0117     for (const auto& p: *mcparts) {
0118
0119         ROOT::Math::PxPyPzEVector mc(p.getMomentum().x, p.getMomentum().y, p.getMomentum().z, p.getEnergy());
0120
0121         mc = boostVector(mc);
0122         mc = rotationAboutY(mc);
0123         mc = rotationAboutX(mc);
0124         mc = headOnBoostVector(mc);
0125
0126         edm4hep::Vector3f mcMom(mc.Px(), mc.Py(), mc.Pz());
0127         edm4hep::MutableMCParticle MCTrack(p.clone());
0128         MCTrack.setMomentum(mcMom);
0129
0130         if(pidUseMCTruth){
0131             MCTrack.setPDG(p.getPDG());
0132             MCTrack.setMass(p.getMass());
0133         }
0134         if(!pidUseMCTruth && pidAssumePionMass){
0135             MCTrack.setPDG(211);
0136             MCTrack.setMass(0.13957);
0137         }
0138
0139         outputParticles->push_back(MCTrack);
0140     }
0141
0142 }