Merge pull request #315 from AndrewEdmonds11/trkpid-rooutil

Add trkpid branch to RooUtil

Author: AndrewEdmonds11

doc/branches.md

@@ -35,6 +35,7 @@ If a Kalman fit fails or there are multiple downstream tracks to fit, the number
 | trkcalohit |  Vector branch |   the calorimeter cluster assigned to a track| [see TrkCaloHitInfo.hh](../inc/TrkCaloHitInfo.hh)
 | trkcalohitmc |  Vector branch |   MC-truth infromation for calorimeter clusters| [see CaloClusterInfoMC.hh](../inc/CaloClusterInfoMC.hh)
 | trkqual |  Vector branch |   the output of a multi-variate analysis (MVA)| [see MVAResultInfo.hh](../inc/MVAResultInfo.hh)
+| trkpid |  Vector branch |   the output of a multi-variate analysis (MVA)| [see MVAResultInfo.hh](../inc/MVAResultInfo.hh)
 ## Track segments Branches
 
 These branches contain 4 elements per event corresponding to different Kalman fit hypotheses (see Track branches).

helper/ntuplehelper.py

@@ -3,11 +3,11 @@
 class nthelper:
 
     single_object_branches = ['evtinfo', 'evtinfomc', 'hitcount', 'tcnt', 'crvsummary', 'crvsummarymc']
-    vector_object_branches = ['trk', 'trkmc', 'trkcalohit', 'trkcalohitmc', 'caloclusters', 'calohits', 'calorecodigis', 'calodigis', 'crvcoincs', 'crvcoincsmc', 'crvcoincsmcplane', 'trkqual']
+    vector_object_branches = ['trk', 'trkmc', 'trkcalohit', 'trkcalohitmc', 'caloclusters', 'calohits', 'calorecodigis', 'calodigis', 'crvcoincs', 'crvcoincsmc', 'crvcoincsmcplane', 'trkqual', 'trkpid']
     vector_vector_object_branches = ['trksegs', 'trksegpars_lh', 'trksegpars_ch', 'trksegpars_kl', 'trkmcsim', 'trkhits', 'trkhitsmc', 'trkmats', 'trkhitcalibs', 'trkmcsci', 'trkmcssi', 'trksegsmc' ]
 
     evt_branches = ['evtinfo','evtinfomc','hitcount','tcnt']
-    trk_branches = ['trk', 'trkmc', 'trkcalohit', 'trkcalohitmc', 'trkqual']
+    trk_branches = ['trk', 'trkmc', 'trkcalohit', 'trkcalohitmc', 'trkqual', 'trkpid']
     trksegs_branches = ['trksegs', 'trksegpars_lh', 'trksegpars_ch', 'trksegpars_kl', 'trksegsmc']
     straw_branches = ['trkhits', 'trkmats', 'trkhitsmc', 'trkhitcalibs']
     mc_branches = ['trkmcsim']
@@ -56,6 +56,7 @@ class nthelper:
                            "crvcoincsmc" : "CrvHitInfoMC",
                            "crvcoincsmcplane" : "CrvPlaneInfoMC",
                            "trkqual" : "MVAResultInfo",
+                           "trkpid" : "MVAResultInfo",
                            "helices" : "HelixInfo",
                            "trksegsmc" : "SurfaceStepInfo"
                           }

rooutil/README.md

@@ -59,7 +59,7 @@ Because some vector-object branches in the EventNtuple should be paired and loop
 ### The ```Track``` Class
 The ```Track``` class contains all information related to a single track
 
-* single objects: ```trk```, ```trkmc```, ```trkcalohit```, ```trkqual```
+* single objects: ```trk```, ```trkmc```, ```trkcalohit```, ```trkqual```, ```trkpid```
 * vectors: ```trksegs```, ```trksegmcs```, ```trksegpars_{lh,ch,kl}```, ```trkmcsim```, ```trkhits```, ```trkhitsmc```, ```trkmats```, ```trkhitcalibs```
 
 Example: [PlotTrackNHits_RecoVsTrue.C](./examples/PlotTrackNHits_RecoVsTrue.C)

rooutil/examples/PlotEoverP_TrkPIDCut.C

@@ -0,0 +1,48 @@
+//
+// An example of how to plot the momentum of electrons at the tracker entrance
+// This uses cut functions defined in common_cuts.hh
+//
+
+#include "EventNtuple/rooutil/inc/RooUtil.hh"
+#include "EventNtuple/rooutil/inc/common_cuts.hh"
+
+#include "TH1F.h"
+#include "TCanvas.h"
+
+void PlotEoverP_TrkPIDCut(std::string filename) {
+
+  // Create the histogram you want to fill
+  TH1F* hEoverP_All = new TH1F("hEoverP_All", "E / p", 200,0, 1);
+  TH1F* hEoverP_TrkPIDCut = new TH1F("hEoverP_TrkPIDCut", "E / p", 200,0, 1);
+
+  // Set up RooUtil
+  RooUtil util(filename);
+  //  util.Debug(true);
+  // Loop through the events
+  for (int i_event = 0; i_event < util.GetNEvents(); ++i_event) {
+    // Get the next event
+    auto& event = util.GetEvent(i_event);
+
+    // Get the e_minus tracks from the event
+    auto e_minus_tracks = event.GetTracks(is_e_minus);
+
+    // Loop through the e_minus tracks
+    for (auto& track : e_minus_tracks) {
+
+      double EoverP = track.trkcalohit->edep  / track.trkcalohit->mom.R();
+      hEoverP_All->Fill(EoverP);
+
+      if (passes_trkpid(track, 0.8)) {
+        hEoverP_TrkPIDCut->Fill(EoverP);
+      }
+    }
+  }
+
+  // Draw the histogram
+  TCanvas* c1 = new TCanvas();
+  hEoverP_All->SetLineColor(kBlack);
+  hEoverP_All->Draw("HIST E");
+
+  hEoverP_TrkPIDCut->SetLineColor(kRed);
+  hEoverP_TrkPIDCut->Draw("HIST E SAME");
+}

rooutil/examples/PrintEvents.C

@@ -108,6 +108,13 @@ void PrintEvents(std::string filename) {
       }
     }
 
+    // trkpid branch
+    if (event.trkpid != nullptr) {
+      for (const auto& trkpid : *(event.trkpid)) {
+        std::cout << "trkpid: " << trkpid.valid << ", " << trkpid.result << std::endl;
+      }
+    }
+
     // trksegpars_lh
     if (event.trksegpars_lh != nullptr) { // might not have this branch
       for (const auto& track : *(event.trksegpars_lh)) {

rooutil/inc/Event.hh

@@ -83,6 +83,9 @@ struct Event {
     if (ntuple->GetBranch("crvcoincs")) {
       ntuple->SetBranchAddress("crvcoincs", &this->crvcoincs);
     }
+    if (ntuple->GetBranch("trkpid")) {
+      ntuple->SetBranchAddress("trkpid", &this->trkpid);
+    }
 
     // Check if the MC branches exist
     if (ntuple->GetBranch("evtinfomc")) {
@@ -211,6 +214,10 @@ struct Event {
         if (debug) { std::cout << "Event::Update(): Adding trkqual_alt to Track " << i_track << "... " << std::endl; }
         track.trkqual_alt = &(trkqual_alt->at(i_track));
       }
+      if (trkpid != nullptr) {
+        if (debug) { std::cout << "Event::Update(): Adding trkpid to Track " << i_track << "... " << std::endl; }
+        track.trkpid = &(trkpid->at(i_track));
+      }
 
       if (debug) { std::cout << "Event::Update(): Updating Track " << i_track << "... " << std::endl; }
       track.Update(debug);
@@ -284,6 +291,7 @@ struct Event {
         if (trkcalohit) { trkcalohit->erase(trkcalohit->begin()+trks_to_remove[i_trk]); }
         if (trkqual) { trkqual->erase(trkqual->begin()+trks_to_remove[i_trk]); }
         if (trkqual_alt) { trkqual_alt->erase(trkqual_alt->begin()+trks_to_remove[i_trk]); }
+        if (trkpid) { trkpid->erase(trkpid->begin()+trks_to_remove[i_trk]); }
         if (trksegpars_lh) { trksegpars_lh->erase(trksegpars_lh->begin()+trks_to_remove[i_trk]); }
         if (trksegpars_ch) { trksegpars_ch->erase(trksegpars_ch->begin()+trks_to_remove[i_trk]); }
         if (trksegpars_kl) { trksegpars_kl->erase(trksegpars_kl->begin()+trks_to_remove[i_trk]); }
@@ -382,6 +390,7 @@ struct Event {
   std::vector<mu2e::CaloClusterInfoMC>* trkcalohitmc = nullptr;
   std::vector<mu2e::MVAResultInfo>* trkqual = nullptr;
   std::vector<mu2e::MVAResultInfo>* trkqual_alt = nullptr; // an optional trkqual branch to also use
+  std::vector<mu2e::MVAResultInfo>* trkpid = nullptr;
   std::vector<std::vector<mu2e::TrkSegInfo>>* trksegs = nullptr;
   std::vector<std::vector<mu2e::SurfaceStepInfo>>* trksegsmc = nullptr;
   std::vector<std::vector<mu2e::LoopHelixInfo>>* trksegpars_lh = nullptr;

rooutil/inc/RooUtil.hh

@@ -118,6 +118,7 @@ public:
     if(event->trkcalohitmc) { output_ntuple->Branch("trkcalohitmc", event->trkcalohitmc); }
     if(event->trkqual) { output_ntuple->Branch("trkqual", event->trkqual); }
     if(event->trkqual_alt) { output_ntuple->Branch("trkqual_alt", event->trkqual_alt); }
+    if(event->trkpid) { output_ntuple->Branch("trkpid", event->trkpid); }
     if(event->trksegs) { output_ntuple->Branch("trksegs", event->trksegs); }
     if(event->trksegsmc) { output_ntuple->Branch("trksegsmc", event->trksegsmc); }
     if(event->trksegpars_lh) { output_ntuple->Branch("trksegpars_lh", event->trksegpars_lh); }

rooutil/inc/Track.hh

@@ -177,6 +177,7 @@ struct Track {
   std::vector<mu2e::SimInfo>* trkmcsim = nullptr;
   mu2e::MVAResultInfo* trkqual = nullptr;
   mu2e::MVAResultInfo* trkqual_alt = nullptr; // TODO: is there a way to allow for more than two...
+  mu2e::MVAResultInfo* trkpid = nullptr;
 };
 
 typedef std::function<bool(Track&)> TrackCut;

rooutil/inc/common_cuts.hh

@@ -184,6 +184,11 @@ bool passes_trkqual(Track& track, double cut_val) { // true if trkqual > cut_val
   else { return false; }
 }
 
+bool passes_trkpid(Track& track, double cut_val) { // true if trkpid > cut_val
+  if (track.trkpid->result > cut_val) { return true; }
+  else { return false; }
+}
+
 //+ CrvCoinc Cuts
 bool three_of_four_coinc(CrvCoinc& crv_coinc) { // CRV coincidence has exactly three layers hit
   if (crv_coinc.reco->nLayers == 3) { return true; }