From 7b9c90003cb83926267c5cdf9b8c77e66057621d Mon Sep 17 00:00:00 2001
From: seley <sinead.mary.eley@cern.ch>
Date: Fri, 23 Aug 2024 10:26:00 +0100
Subject: [PATCH] Added truth_st1_x, truth_st1_px etc. for truth positions and
 momenta extrapolated to each station for first 10 MC truth particles

---
 .../NtupleDumper/src/NtupleDumperAlg.cxx      | 51 +++++++++++++++++++
 .../NtupleDumper/src/NtupleDumperAlg.h        |  7 +++
 2 files changed, 58 insertions(+)

diff --git a/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.cxx b/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.cxx
index 5b789e0e..937b2e7e 100644
--- a/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.cxx
+++ b/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.cxx
@@ -377,6 +377,32 @@ StatusCode NtupleDumperAlg::initialize()
   m_tree->Branch("truth_dec_z", &m_truth_dec_z);
   m_tree->Branch("truth_isFiducial", &m_truth_isFiducial);
 
+
+  m_tree->Branch("truth_st0_x", &m_truth_st_x[0]);
+  m_tree->Branch("truth_st0_y", &m_truth_st_y[0]);
+  m_tree->Branch("truth_st0_z", &m_truth_st_z[0]);
+  m_tree->Branch("truth_st1_x", &m_truth_st_x[1]);
+  m_tree->Branch("truth_st1_y", &m_truth_st_y[1]);
+  m_tree->Branch("truth_st1_z", &m_truth_st_z[1]);
+  m_tree->Branch("truth_st2_x", &m_truth_st_x[2]);
+  m_tree->Branch("truth_st2_y", &m_truth_st_y[2]);
+  m_tree->Branch("truth_st2_z", &m_truth_st_z[2]);
+  m_tree->Branch("truth_st3_x", &m_truth_st_x[3]);
+  m_tree->Branch("truth_st3_y", &m_truth_st_y[3]);
+  m_tree->Branch("truth_st3_z", &m_truth_st_z[3]);
+  m_tree->Branch("truth_st0_px", &m_truth_st_px[0]);
+  m_tree->Branch("truth_st0_py", &m_truth_st_py[0]);
+  m_tree->Branch("truth_st0_pz", &m_truth_st_pz[0]);
+  m_tree->Branch("truth_st1_px", &m_truth_st_px[1]);
+  m_tree->Branch("truth_st1_py", &m_truth_st_py[1]);
+  m_tree->Branch("truth_st1_pz", &m_truth_st_pz[1]);
+  m_tree->Branch("truth_st2_px", &m_truth_st_px[2]);
+  m_tree->Branch("truth_st2_py", &m_truth_st_py[2]);
+  m_tree->Branch("truth_st2_pz", &m_truth_st_pz[2]);
+  m_tree->Branch("truth_st3_px", &m_truth_st_px[3]);
+  m_tree->Branch("truth_st3_py", &m_truth_st_py[3]);
+  m_tree->Branch("truth_st3_pz", &m_truth_st_pz[3]);
+
   // for mother + daughter particle truth infomation 
 
   m_tree->Branch("truthM_P", &m_truthM_P);
@@ -764,6 +790,20 @@ StatusCode NtupleDumperAlg::execute(const EventContext &ctx) const
         m_truth_pdg.push_back(particle->pdgId());
         m_truth_isFiducial.push_back(m_fiducialParticleTool->isFiducial(particle->barcode()));
 
+        auto positions = m_fiducialParticleTool->getTruthPositions(particle->barcode()); // Add truth position information at each station
+        for (int station = 0; station < 4; ++station) {
+          m_truth_st_x[station].push_back(positions[station].x());
+          m_truth_st_y[station].push_back(positions[station].y());
+          m_truth_st_z[station].push_back(positions[station].z());
+        }
+
+        auto momenta = m_fiducialParticleTool->getTruthMomenta(particle->barcode()); //  Add truth momentum information at each station
+        for (int station = 0; station < 4; ++station) {
+          m_truth_st_px[station].push_back(momenta[station].x());
+          m_truth_st_py[station].push_back(momenta[station].y());
+          m_truth_st_pz[station].push_back(momenta[station].z());
+        }
+
         if ( particle->hasProdVtx()) {
           m_truth_prod_x.push_back(particle->prodVtx()->x());
           m_truth_prod_y.push_back(particle->prodVtx()->y());
@@ -1604,6 +1644,17 @@ NtupleDumperAlg::clearTree() const
   m_truth_dec_z.clear();
   m_truth_isFiducial.clear();
 
+  for (int station = 0; station < 4; ++station) {
+    m_truth_st_x[station].clear();
+    m_truth_st_y[station].clear();
+    m_truth_st_z[station].clear();
+  }
+  for (int station = 0; station < 4; ++station) {
+    m_truth_st_px[station].clear();
+    m_truth_st_py[station].clear();
+    m_truth_st_pz[station].clear();
+  }
+
   m_truthM_P.clear();
   m_truthM_px.clear();
   m_truthM_py.clear();
diff --git a/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.h b/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.h
index 0786017d..c57c803c 100644
--- a/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.h
+++ b/PhysicsAnalysis/NtupleDumper/src/NtupleDumperAlg.h
@@ -380,6 +380,13 @@ private:
   mutable std::vector<int> m_truth_pdg; // pdg of first 10 truth particles 
   mutable std::vector<bool> m_truth_isFiducial; // Boolean for if first 10 truth particles are fiducial
 
+  mutable std::array<std::vector<double>, 4> m_truth_st_x; // vector of the x components of the simulated hits of the truth particle for each station
+  mutable std::array<std::vector<double>, 4> m_truth_st_y; // vector of the y components of the simulated hits of the truth particle for each station
+  mutable std::array<std::vector<double>, 4> m_truth_st_z; // vector of the z components of the simulated hits of the truth particle for each station
+  mutable std::array<std::vector<double>, 4> m_truth_st_px; // x components of the true momentum at each station
+  mutable std::array<std::vector<double>, 4> m_truth_st_py; // y components of the true momentum at each station
+  mutable std::array<std::vector<double>, 4> m_truth_st_pz; // z components of the true momentum at each station
+
   mutable double m_truthLeptonMomentum; 
   mutable int    m_truthBarcode;
   mutable int    m_truthPdg;
-- 
GitLab