Setup mc truth UndoAfterBurner (#1492)

### Briefly, what does this PR introduce?

This PR introduces the postburner to EICrecon, which will remove the
effects of the afterburner from the MCParticles branch, and store the
resulting output (the actual MC Truth information from the generator)
into a new branch, "MCParticlesHeadOnFrameNoBeamFX".

The postburner does not erase any information - the user will simply
have access to both the afterburned MC information via MCParticles (as
before), and the new branch, which handles the correct transformations
to remove the effects.

The code is also staged to be able to handle removing *only* the
crossing angle from reconstructed branches, and this additional
functionality will come with a future PR.

### What kind of change does this PR introduce?
- [ ] Bug fix (issue #__)
- [x] New feature (issue #__)
- [ ] Documentation update
- [ ] Other: __

### Please check if this PR fulfills the following:
- [ x] Tests for the changes have been added
- [ x] Documentation has been added / updated
- [ x] Changes have been communicated to collaborators


### Does this PR introduce breaking changes? What changes might users
need to make to their code?
No.

### Does this PR change default behavior?

[pt_DVCS_protons_MCParticles_POSTBURN_ePIC_eicrecon_6_4_2024_run_0.pdf](https://github.com/user-attachments/files/15570067/pt_DVCS_protons_MCParticles_POSTBURN_ePIC_eicrecon_6_4_2024_run_0.pdf)
[Analysis of ePIC Output With Afterburned
Events.pdf](https://github.com/user-attachments/files/15570110/Analysis.of.ePIC.Output.With.Afterburned.Events.pdf)

It only adds functionality and a new output branch, nothing else is
affected.

---------

Co-authored-by: Alexander Jentsch <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Dmitry Kalinkin <[email protected]>
Co-authored-by: Simon Gardner <[email protected]>
Co-authored-by: Wouter Deconinck <[email protected]>
Co-authored-by: Sebouh Paul <[email protected]>

src/algorithms/reco/UndoAfterBurner.cc

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 Alex Jentsch, Jihee Kim, Brian Page
+//
+
+#include "UndoAfterBurner.h"
+
+#include <Math/GenVector/Boost.h>
+#include <Math/GenVector/Cartesian3D.h>
+#include <Math/GenVector/LorentzVector.h>
+#include <Math/GenVector/PxPyPzE4D.h>
+#include <Math/GenVector/RotationX.h>
+#include <Math/GenVector/RotationY.h>
+#include <Math/Vector4Dfwd.h>
+#include <TMath.h>
+#include <edm4hep/Vector3f.h>
+#include <gsl/pointers>
+
+#include "algorithms/reco/Beam.h"
+#include "algorithms/reco/UndoAfterBurnerConfig.h"
+
+void eicrecon::UndoAfterBurner::init() {
+
+}
+
+void eicrecon::UndoAfterBurner::process(
+    const UndoAfterBurner::Input& input,
+    const UndoAfterBurner::Output& output) const {
+
+    const auto [mcparts] = input;
+    auto [outputParticles] = output;
+
+    bool      pidAssumePionMass = m_cfg.m_pid_assume_pion_mass;
+    double    crossingAngle    = m_cfg.m_crossing_angle;
+    double    pidPurity        = m_cfg.m_pid_purity;
+    bool      correctBeamFX    = m_cfg.m_correct_beam_FX;
+    bool      pidUseMCTruth    = m_cfg.m_pid_use_MC_truth;
+
+    bool      hasBeamHadron    = true;
+    bool      hasBeamLepton    = true;
+
+    //read MCParticles information and "postburn" to remove the afterburner effects.
+    //The output is then the original MC input produced by the generator.
+
+    ROOT::Math::PxPyPzEVector  e_beam(0.,0.,0.,0.);
+    ROOT::Math::PxPyPzEVector  h_beam(0.,0.,0.,0.);
+
+    auto incoming_lepton = find_first_beam_electron(mcparts);
+    if (incoming_lepton.size() == 0) {
+        debug("No beam electron found -- particleGun input");
+        hasBeamLepton = false;
+    }
+
+    auto incoming_hadron = find_first_beam_hadron(mcparts);
+    if (incoming_hadron.size() == 0) {
+        debug("No beam hadron found -- particleGun input");
+        hasBeamHadron = false;
+    }
+
+    if((hasBeamHadron && !hasBeamLepton) || (!hasBeamHadron && hasBeamLepton)){
+        debug("Only one beam defined! Not a possible configuration!");
+        return;
+    }
+
+        //handling for FF particle gun input!!
+    if(!hasBeamHadron || !hasBeamLepton){
+        for (const auto& p: *mcparts) {
+            if((p.getPDG() == 2212 || p.getPDG() == 2112)) { //look for "gun" proton/neutron
+                h_beam.SetPxPyPzE(crossingAngle*p.getEnergy(), 0.0, p.getEnergy(), p.getEnergy());
+                if(p.getEnergy() > 270.0 && p.getEnergy() < 280.0){
+                    e_beam.SetPxPyPzE(0.0, 0.0, -18.0, 18.0);
+                }
+            }
+        }
+    }
+        else{
+
+        if(correctBeamFX){
+
+            h_beam.SetPxPyPzE(incoming_hadron[0].getMomentum().x, incoming_hadron[0].getMomentum().y, incoming_hadron[0].getMomentum().z, incoming_hadron[0].getEnergy());
+            e_beam.SetPxPyPzE(incoming_lepton[0].getMomentum().x, incoming_lepton[0].getMomentum().y, incoming_lepton[0].getMomentum().z, incoming_lepton[0].getEnergy());
+
+        }
+        else{
+
+            h_beam.SetPxPyPzE(crossingAngle*incoming_hadron[0].getEnergy(), 0.0, incoming_hadron[0].getEnergy(), incoming_hadron[0].getEnergy());
+            e_beam.SetPxPyPzE(0.0, 0.0, -incoming_lepton[0].getEnergy(), incoming_lepton[0].getEnergy());
+
+        }
+    }
+
+
+
+    //Calculate boost vectors and rotations here
+
+    ROOT::Math::PxPyPzEVector cm_frame_boost = e_beam + h_beam;
+    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());
+
+    ROOT::Math::Boost boostVector(beta);
+
+    //Boost to CM frame
+    e_beam = boostVector(e_beam);
+    h_beam = boostVector(h_beam);
+
+    double rotationAngleY = -1.0*TMath::ATan2(h_beam.Px(), h_beam.Pz());
+    double rotationAngleX = 1.0*TMath::ATan2(h_beam.Py(), h_beam.Pz());
+
+    ROOT::Math::RotationY rotationAboutY(rotationAngleY);
+    ROOT::Math::RotationX rotationAboutX(rotationAngleX);
+
+    //Boost back to proper head-on frame
+
+    ROOT::Math::PxPyPzEVector head_on_frame_boost(0., 0., cm_frame_boost.Pz(), cm_frame_boost.E());
+    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());
+
+    //Now, loop through events and apply operations to the MCparticles
+    for (const auto& p: *mcparts) {
+
+        ROOT::Math::PxPyPzEVector mc(p.getMomentum().x, p.getMomentum().y, p.getMomentum().z, p.getEnergy());
+
+        mc = boostVector(mc);
+        mc = rotationAboutY(mc);
+        mc = rotationAboutX(mc);
+        mc = headOnBoostVector(mc);
+
+        edm4hep::Vector3f mcMom(mc.Px(), mc.Py(), mc.Pz());
+        edm4hep::MutableMCParticle MCTrack(p.clone());
+        MCTrack.setMomentum(mcMom);
+
+        if(pidUseMCTruth){
+            MCTrack.setPDG(p.getPDG());
+            MCTrack.setMass(p.getMass());
+        }
+        if(!pidUseMCTruth && pidAssumePionMass){
+            MCTrack.setPDG(211);
+            MCTrack.setMass(0.13957);
+        }
+
+        outputParticles->push_back(MCTrack);
+    }
+
+}

src/algorithms/reco/UndoAfterBurner.h

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 Alex Jentsch, Jihee Kim, Brian Page
+//
+
+#include <algorithms/algorithm.h>
+#include <edm4hep/MCParticleCollection.h>
+#include <string>
+#include <string_view>
+
+#include "UndoAfterBurnerConfig.h"
+#include "algorithms/interfaces/WithPodConfig.h"
+
+namespace eicrecon {
+
+  using UndoAfterBurnerAlgorithm = algorithms::Algorithm<
+    algorithms::Input<
+      edm4hep::MCParticleCollection
+    >,
+    algorithms::Output<
+      edm4hep::MCParticleCollection
+    >
+  >;
+
+  class UndoAfterBurner
+  : public UndoAfterBurnerAlgorithm,
+    public WithPodConfig<UndoAfterBurnerConfig> {
+
+  public:
+    UndoAfterBurner(std::string_view name)
+      : UndoAfterBurnerAlgorithm{name,
+                            {"inputMCParticles"},
+                            {"outputMCParticles"},
+                            "Apply boosts and rotations to remove crossing angle and beam effects."} {}
+
+    void init();
+    void process(const Input&, const Output&) const final;
+
+  private:
+
+
+  };
+}

src/algorithms/reco/UndoAfterBurnerConfig.h

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 Alex Jentsch, Jihee Kim, Brian Page
+//
+
+#pragma once
+
+namespace eicrecon {
+
+    struct UndoAfterBurnerConfig {
+
+        bool      m_pid_assume_pion_mass = false;
+        double    m_crossing_angle       = -0.025;
+        double    m_pid_purity           = 0.51;
+        bool      m_correct_beam_FX      = true;
+        bool      m_pid_use_MC_truth     = true;
+
+    };
+
+}

src/factories/reco/UndoAfterBurnerMCParticles_factory.h

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 Alex Jentsch, Jihee Kim, Brian Page
+//
+
+#include "algorithms/reco/UndoAfterBurner.h"
+#include "algorithms/reco/UndoAfterBurnerConfig.h"
+
+// Event Model related classes
+#include <edm4eic/MCRecoParticleAssociationCollection.h>
+#include <edm4eic/ReconstructedParticleCollection.h>
+#include <edm4hep/MCParticleCollection.h>
+
+#include "extensions/jana/JOmniFactory.h"
+
+namespace eicrecon {
+
+class UndoAfterBurnerMCParticles_factory :
+    public JOmniFactory<UndoAfterBurnerMCParticles_factory, UndoAfterBurnerConfig> {
+
+public:
+    using AlgoT = eicrecon::UndoAfterBurner;
+private:
+    std::unique_ptr<AlgoT> m_algo;
+
+    PodioInput<edm4hep::MCParticle> m_mcparts_input {this};
+    PodioOutput<edm4hep::MCParticle> m_postburn_output {this};
+
+    ParameterRef<bool>      m_pid_assume_pion_mass {this, "m_pid_assume_pion_mass", config().m_pid_assume_pion_mass};
+    ParameterRef<double>    m_crossing_angle       {this, "m_crossing_angle", config().m_crossing_angle};
+    ParameterRef<double>    m_pid_purity           {this, "m_pid_purity", config().m_pid_purity};
+    ParameterRef<bool>      m_correct_beam_FX      {this, "m_correct_beam_FX", config().m_correct_beam_FX};
+    ParameterRef<bool>      m_pid_use_MC_truth     {this, "m_pid_use_MC_truth", config().m_pid_use_MC_truth};
+
+
+public:
+    void Configure() {
+        m_algo = std::make_unique<AlgoT>(GetPrefix());
+        m_algo->applyConfig(config());
+    }
+
+    void ChangeRun(int64_t run_number) {
+    }
+
+    void Process(int64_t run_number, uint64_t event_number) {
+        m_algo->process({m_mcparts_input()}, {m_postburn_output().get()});
+    }
+
+};
+
+} // eicrecon

src/global/reco/reco.cc

@@ -1,9 +1,8 @@
-// Copyright 2022, Dmitry Romanov
-// Subject to the terms in the LICENSE file found in the top-level directory.
-//
-//
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2022 - 2024, Dmitry Romanov, Nathan Brei, Tooba Ali, Wouter Deconinck, Dmitry Kalinkin, John Lajoie, Simon Gardner, Tristan Protzman, Daniel Brandenburg, Derek M Anderson, Sebouh Paul, Tyler Kutz, Alex Jentsch, Jihee Kim, Brian Page
 
 
+#include <Evaluator/DD4hepUnits.h>
 #include <JANA/JApplication.h>
 #include <edm4eic/Cluster.h>
 #include <edm4eic/EDM4eicVersion.h>
@@ -40,6 +39,7 @@
 #if EDM4EIC_VERSION_MAJOR >= 6
 #include "factories/reco/HadronicFinalState_factory.h"
 #endif
+#include "factories/reco/UndoAfterBurnerMCParticles_factory.h"
 #include "global/reco/ChargedReconstructedParticleSelector_factory.h"
 #include "global/reco/MC2SmearedParticle_factory.h"
 #include "global/reco/MatchClusters_factory.h"
@@ -351,5 +351,26 @@ void InitPlugin(JApplication *app) {
             app
     ));
 
+
+    //Full correction for MCParticles --> MCParticlesHeadOnFrame
+    app->Add(new JOmniFactoryGeneratorT<UndoAfterBurnerMCParticles_factory>(
+            "MCParticlesHeadOnFrameNoBeamFX",
+            {
+                "MCParticles"
+            },
+            {
+                "MCParticlesHeadOnFrameNoBeamFX"
+            },
+            {
+              .m_pid_assume_pion_mass = false,
+              .m_crossing_angle = -0.025 * dd4hep::rad,
+              .m_pid_purity = 0.51, //dummy value for MC truth information
+              .m_correct_beam_FX = true,
+              .m_pid_use_MC_truth = true,
+            },
+            app
+    ));
+
+
 }
 } // extern "C"

src/services/io/podio/JEventProcessorPODIO.cc

@@ -55,6 +55,7 @@ JEventProcessorPODIO::JEventProcessorPODIO() {
             "MCBeamProtons",
             "MCScatteredElectrons",
             "MCScatteredProtons",
+            "MCParticlesHeadOnFrameNoBeamFX",
 
             // All tracking hits combined
             "CentralTrackingRecHits",