drich-dev/src/draw_hits.cpp

0001 // draw hits, and make some other related plots
0002 // (cf. event_display.cpp for readout)
0003 #include <cstdlib>
0004 #include <iostream>
0005 #include <fmt/format.h>
0006
0007 // ROOT
0008 #include "TSystem.h"
0009 #include "TStyle.h"
0010 #include "TRegexp.h"
0011 #include "TCanvas.h"
0012 #include "TApplication.h"
0013 #include "TBox.h"
0014 #include "ROOT/RDataFrame.hxx"
0015
0016 // edm4hep
0017 #include "edm4hep/MCParticleCollection.h"
0018 #include "edm4hep/SimTrackerHitCollection.h"
0019
0020 // local
0021 #include "WhichRICH.h"
0022
0023 using namespace ROOT;
0024 using namespace ROOT::VecOps;
0025 using namespace edm4hep;
0026
0027 TCanvas *CreateCanvas(TString name, Bool_t logx=0, Bool_t logy=0, Bool_t logz=0);
0028
0029 int main(int argc, char** argv) {
0030
0031   // args
0032   TString infileN="out/sim.edm4hep.root";
0033   if(argc<=1) {
0034     fmt::print("\nUSAGE: {} [d/p] [simulation_output_file(optional)]\n\n",argv[0]);
0035     fmt::print("    [d/p]: d for dRICH\n");
0036     fmt::print("           p for pfRICH\n");
0037     fmt::print("    [simulation_output_file]: output from `npsim` (`simulate.py`)\n");
0038     fmt::print("                              default: {}\n",infileN);
0039     return 2;
0040   }
0041   std::string zDirectionStr = argv[1];
0042   if(argc>2) infileN = TString(argv[2]);
0043   WhichRICH wr(zDirectionStr);
0044   if(!wr.valid) return 1;
0045
0046   // setup
0047   //TApplication mainApp("mainApp",&argc,argv); // keep canvases open
0048   //EnableImplicitMT();
0049   RDataFrame dfIn("events",infileN.Data());
0050   TString outfileN = infileN;
0051   outfileN(TRegexp("\\.root$"))=".";
0052   TFile *outfile = new TFile(outfileN+"plots.root","RECREATE");
0053   gStyle->SetOptStat(0);
0054   gStyle->SetPalette(55);
0055
0056
0057   /* lambdas
0058    * - most of these transform an `RVec<T1>` to an `RVec<T2>` using `VecOps::Map` or `VecOps::Filter`
0059    * - see NPdet/src/dd4pod/dd4hep.yaml for POD syntax
0060    */
0061   // calculate number of hits
0062   auto numHits = [](RVec<SimTrackerHitData> hits) { return hits.size(); };
0063   // calculate momentum magnitude for each particle (units=GeV)
0064   auto momentum = [](RVec<MCParticleData> parts){
0065     return Map(parts, [](auto p){
0066         auto mom = p.momentum;
0067         return sqrt( mom[0]*mom[0] + mom[1]*mom[1] + mom[2]*mom[2] );
0068         });
0069   };
0070   // filter for thrown particles
0071   auto isThrown = [](RVec<MCParticleData> parts){
0072     return Filter(parts, [](auto p){ return p.generatorStatus==1; } );
0073   };
0074   // get positions for each hit (units=cm)
0075   auto hitPos = [](RVec<SimTrackerHitData> hits){ return Map(hits,[](auto h){ return h.position; }); };
0076   auto hitPosX = [](RVec<Vector3d> v){ return Map(v,[](auto p){ return p.x/10; }); };
0077   auto hitPosY = [](RVec<Vector3d> v){ return Map(v,[](auto p){ return p.y/10; }); };
0078   auto hitPosZ = [](RVec<Vector3d> v){ return Map(v,[](auto p){ return p.z/10; }); };
0079
0080
0081   // transformations
0082   auto df1 = dfIn
0083     .Define("thrownParticles",isThrown,{"MCParticles"})
0084     .Define("thrownP",momentum,{"thrownParticles"})
0085     .Define("numHits",numHits,{wr.readoutName})
0086     .Define("hitPos",hitPos,{wr.readoutName})
0087     .Define("hitX",hitPosX,{"hitPos"})
0088     .Define("hitY",hitPosY,{"hitPos"})
0089     ;
0090   auto dfFinal = df1;
0091
0092
0093   // actions
0094   auto hitPositionHist = dfFinal.Histo2D(
0095       { "hitPositions",TString(wr.xRICH)+" hit positions (units=cm)",
0096       1000,-200,200, 1000,-200,200 },
0097       "hitX","hitY"
0098       );
0099   auto numHitsVsThrownP = dfFinal.Histo2D(
0100       { "numHitsVsThrownP","number of "+TString(wr.xRICH)+" hits vs. thrown momentum",
0101       65,0,65, 100,0,400 },
0102       "thrownP","numHits"
0103       ); // TODO: cut opticalphotons (may not be needed, double check PID)
0104
0105
0106   // execution
0107   TCanvas *canv;
0108   canv = CreateCanvas("hits",0,0,1);
0109   hitPositionHist->Draw("colz");
0110   if(wr.zDirection>0) {
0111     hitPositionHist->GetXaxis()->SetRangeUser(100,200);
0112     hitPositionHist->GetYaxis()->SetRangeUser(-50,50);
0113   } else {
0114     hitPositionHist->GetXaxis()->SetRangeUser(-70,70);
0115     hitPositionHist->GetYaxis()->SetRangeUser(-70,70);
0116   }
0117   canv->Print(outfileN+"hits.png");
0118   canv->Write();
0119   //
0120   canv = CreateCanvas("photon_yield");
0121   numHitsVsThrownP->Draw("box");
0122   TProfile * aveHitsVsP;
0123   aveHitsVsP = numHitsVsThrownP->ProfileX("_pfx",1,-1,"i"); // TODO: maybe not the right errors, see TProfile::BuildOptions `i`
0124   aveHitsVsP->SetLineColor(kBlack);
0125   aveHitsVsP->SetLineWidth(3);
0126   aveHitsVsP->Draw("same");
0127   canv->Print(outfileN+"photon_count.png");
0128   canv->Write();
0129   aveHitsVsP->Write("aveHitsVsP");
0130   outfile->Close();
0131
0132
0133   // exit
0134   //fmt::print("\n\npress ^C to exit.\n\n");
0135   //mainApp.Run(); // keep canvases open
0136   return 0;
0137 }
0138
0139
0140 TCanvas *CreateCanvas(TString name, Bool_t logx, Bool_t logy, Bool_t logz) {
0141   TCanvas *c = new TCanvas("canv_"+name,"canv_"+name,800,600);
0142   c->SetGrid(1,1);
0143   if(logx) c->SetLogx(1);
0144   if(logy) c->SetLogy(1);
0145   if(logz) c->SetLogz(1);
0146   return c;
0147 }
0148