Skip to content

Commit

Permalink
Setup mc truth UndoAfterBurner (#1492)
Browse files Browse the repository at this point in the history
### 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]>
  • Loading branch information
7 people authored Jun 26, 2024
1 parent 9a032a2 commit 87f64cd
Show file tree
Hide file tree
Showing 6 changed files with 278 additions and 4 deletions.
141 changes: 141 additions & 0 deletions src/algorithms/reco/UndoAfterBurner.cc
Original file line number Diff line number Diff line change
@@ -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);
}

}
42 changes: 42 additions & 0 deletions src/algorithms/reco/UndoAfterBurner.h
Original file line number Diff line number Diff line change
@@ -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:


};
}
19 changes: 19 additions & 0 deletions src/algorithms/reco/UndoAfterBurnerConfig.h
Original file line number Diff line number Diff line change
@@ -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;

};

}
50 changes: 50 additions & 0 deletions src/factories/reco/UndoAfterBurnerMCParticles_factory.h
Original file line number Diff line number Diff line change
@@ -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
29 changes: 25 additions & 4 deletions src/global/reco/reco.cc
Original file line number Diff line number Diff line change
@@ -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>
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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"
1 change: 1 addition & 0 deletions src/services/io/podio/JEventProcessorPODIO.cc
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ JEventProcessorPODIO::JEventProcessorPODIO() {
"MCBeamProtons",
"MCScatteredElectrons",
"MCScatteredProtons",
"MCParticlesHeadOnFrameNoBeamFX",

// All tracking hits combined
"CentralTrackingRecHits",
Expand Down

0 comments on commit 87f64cd

Please sign in to comment.