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