scripts/src/momentum_scan_juggler_draw.C

0001 // draw various momentum scan plots
0002 R__LOAD_LIBRARY(podioDict)
0003 R__LOAD_LIBRARY(podioRootIO)
0004 R__LOAD_LIBRARY(edm4hep)
0005 R__LOAD_LIBRARY(edm4eic)
0006 R__LOAD_LIBRARY(fmt)
0007
0008 #include "podio/EventStore.h"
0009 #include "podio/ROOTReader.h"
0010 #include "edm4hep/utils/kinematics.h"
0011 #include "fmt/format.h"
0012
0013 // radiators
0014 enum rad_enum {kAgl=0, kGas=1};
0015
0016 //=========================================================
0017
0018 void momentum_scan_juggler_draw(
0019     std::string rec_file_name  = "out/rec.edm4hep.root",
0020     std::string out_file_name  = "out/rec.scan_plots.root",
0021     std::string det_name       = "DRICH", // or PFRICH
0022     unsigned    which_radiator = 0 // see `rad_enum` above
0023     )
0024 {
0025   // FIXME: upgrade to frame reader
0026   fmt::print(stderr,"WARNING: this script still uses podio::ROOTReader, which may be deprecated soon!\n");
0027
0028   // open event store
0029   podio::ROOTReader reader;
0030   podio::EventStore store;
0031   reader.openFile(rec_file_name);
0032   store.setReader(&reader);
0033
0034   // output root file
0035   TFile *out_file = new TFile(TString(out_file_name),"RECREATE");
0036
0037   // radiators
0038   std::string radiator_name;
0039   double mom_max;
0040   switch(which_radiator) {
0041     case kAgl:
0042       radiator_name = "Aerogel";
0043       mom_max       = 22;
0044       break;
0045     case kGas:
0046       radiator_name = "Gas";
0047       mom_max       = 65;
0048       break;
0049     default:
0050       fmt::print(stderr,"ERROR: unknown which_radiator={}\n",which_radiator);
0051       return;
0052   }
0053
0054   // plot limits
0055   const double theta_max = 260.0;
0056   const int    nphot_max = 500;
0057   const int    npe_max   = 0.3*nphot_max;
0058
0059   // define histograms
0060   std::map<std::string,TH2D*> scans;
0061   auto make_scan = [&scans,&mom_max,&radiator_name] (
0062       std::string name,
0063       std::string title,
0064       std::string units,
0065       auto yn, auto yi, auto yf
0066       )
0067   {
0068     if(units!="") units=" ["+units+"]";
0069     std::string for_radiator = (name!="nphot") ? ", for "+radiator_name : "";
0070     auto hist = new TH2D(
0071         TString(name + "_scan"),
0072         TString(title + " vs. Thrown Momentum" + for_radiator +";p [GeV];" + title + units),
0073         (int)mom_max+1, 0, mom_max,
0074         yn, yi, yf
0075         );
0076     hist->StatOverflows(true);
0077     scans.insert({name,hist});
0078   };
0079   make_scan( "nphot", "N_{photons}",        "",     nphot_max, 0, nphot_max );
0080   make_scan( "npe",   "N_{photoelectrons}", "",     npe_max,   0, npe_max   );
0081   make_scan( "theta", "#theta_{Cherenkov}", "mrad", 100,       0, theta_max );
0082
0083   // function to Fill a scan: warns for under/overflows
0084   auto FillScan = [&scans] (std::string name, auto x, auto y) {
0085     auto check = [] (auto val, std::string ax_name, TAxis *ax) {
0086       if(val<ax->GetXmin() || val>ax->GetXmax())
0087         fmt::print(stderr,"WARNING: {} overflow: {}\n",ax_name,val);
0088     };
0089     auto scan = scans.at(name);
0090     check(x, name+" scan x-axis", scan->GetXaxis());
0091     check(y, name+" scan y-axis", scan->GetYaxis());
0092     scan->Fill(x,y);
0093   };
0094
0095   // event loop =================================================================
0096   for(unsigned e=0; e<reader.getEntries(); e++) {
0097
0098     // get next event
0099     // fmt::print("read event #{}\n",e);
0100     if(e>0) {
0101       store.clear();
0102       reader.endOfEvent();
0103     };
0104
0105     // get collections
0106     const auto& mcParts = store.get<edm4hep::MCParticleCollection>("MCParticles");
0107     const auto& hits    = store.get<edm4hep::SimTrackerHitCollection>(det_name+"Hits");
0108     const auto& cpids   = store.get<edm4eic::CherenkovParticleIDCollection>(det_name+"PID");
0109
0110     // thrown momentum
0111     int nthrown = 0;
0112     float mom;
0113     for(const auto& mcPart : mcParts) {
0114       if(mcPart.getGeneratorStatus()==1) {
0115         mom = edm4hep::utils::p(mcPart);
0116         nthrown++;
0117       }
0118     }
0119     if(nthrown == 0) { fmt::print(stderr,"WARNING: no thrown particle found for event #{}\n",e); continue; };
0120     if(nthrown >  1)   fmt::print(stderr,"WARNING: this script does not yet support multi-track events (nthrown={})\n",nthrown);
0121
0122     // number of photons
0123     FillScan("nphot", mom, hits.size());
0124
0125     // NPE, Cherenkov angle
0126     if(cpids.size() != 1) fmt::print(stderr,"WARNING: CherenkovParticleIDCollection size = {} != 1\n",cpids.size());
0127     for(const auto& cpid : cpids) {
0128       for(const auto &meas : cpid.getAngles()) {
0129         auto npe = meas.npe;
0130         auto theta = meas.theta * 1000; // [mrad]
0131         bool pass =
0132           meas.radiator == which_radiator
0133           && npe>0
0134           ;
0135         if(pass) {
0136           FillScan( "npe",   mom, npe   );
0137           FillScan( "theta", mom, theta );
0138         }
0139       }
0140     }
0141
0142   } // end event loop =================================================================
0143   store.clear();
0144   reader.endOfEvent();
0145   reader.closeFile();
0146
0147   // make scan profiles
0148   std::map<std::string,TProfile*> scan_profs;
0149   std::map<std::string,TCanvas*> scan_canvs;
0150   for(const auto& [name,scan] : scans) {
0151     auto prof = scan->ProfileX("_pfx",1,-1,"");
0152     prof->SetLineColor(kBlack);
0153     prof->SetLineWidth(3);
0154     scan_profs.insert({name,prof});
0155     TString canv_name(name+"_canv");
0156     auto canv = new TCanvas(canv_name,canv_name);
0157     canv->SetGrid(1,1);
0158     scan->Draw("box");
0159     prof->Draw("same");
0160     scan_canvs.insert({name,canv});
0161   }
0162
0163   // write output
0164   auto WriteScans = [] (auto vec) { for(const auto& [name,obj] : vec) obj->Write(); };
0165   WriteScans(scans);
0166   WriteScans(scan_profs);
0167   WriteScans(scan_canvs);
0168   out_file->Close();
0169   fmt::print("\n{:=<50}\nwrote: {}\n\n","",out_file_name);
0170 }