scripts/src/momentum_scan_2D_draw.C

0001 // draw combined theta vs. p plot
0002 R__LOAD_LIBRARY(fmt)
0003 #include "fmt/format.h"
0004
0005 // radiators
0006 enum rad_enum {kAgl, kGas, nRad};
0007
0008 //=========================================================
0009
0010 void momentum_scan_2D_draw(
0011     TString  ana_file_names = "out/momentum_scan.drich/*.rec.agl.ana.root", // input files (will be hadded)
0012     TString  out_base_name  = "out/2D",                                     // output; should be png or pdf
0013     unsigned which_radiator = 0,                                            // see `rad_enum` above
0014     Int_t    n_group_rebin  = 1                                             // ngroup for rebinning momentum bins (1=disable)
0015     )
0016 {
0017
0018   // hadd input files
0019   TString hadd_file_name = out_base_name + ".root";
0020   gROOT->ProcessLine(".! hadd -f " + hadd_file_name + " " + ana_file_names);
0021   auto hadd_file = new TFile(hadd_file_name,"UPDATE");
0022
0023   // radiators
0024   TString radiator_name;
0025   double mom_max, theta_max, rindex_ref;
0026   switch(which_radiator) {
0027     case kAgl:
0028       radiator_name = "Aerogel";
0029       mom_max       = 20;  // for plot zoom
0030       theta_max     = 230; // for plot zoom
0031       rindex_ref    = 1.0190;
0032       break;
0033     case kGas:
0034       radiator_name = "Gas";
0035       mom_max       = 60; // for plot zoom
0036       theta_max     = 50; // for plot zoom
0037       rindex_ref    = 1.00076;
0038       break;
0039     default:
0040       fmt::print(stderr,"ERROR: unknown which_radiator={}\n",which_radiator);
0041       return;
0042   }
0043
0044   // theta curves
0045   std::map<TString,double> particle_masses = {
0046     { "e-",     0.00051 },
0047     { "pi+",    0.13957 },
0048     { "kaon+",  0.49368 },
0049     { "proton", 0.93827 }
0050   };
0051   std::map<TString,TF1*> theta_curves;
0052   for(auto [particle,mass] : particle_masses)
0053     theta_curves.insert({
0054         particle,
0055         new TF1(
0056             "ftn_theta_" + particle,
0057             Form("1000*TMath::ACos(TMath::Sqrt(x^2+%f^2)/(%f*x))", mass, rindex_ref),
0058             mass / TMath::Sqrt(rindex_ref * rindex_ref - 1),
0059             mom_max
0060           )
0061         });
0062
0063   // get plots
0064   std::map<TString,TH2D*> scans;
0065   auto get_scan = [&] (TString dir, TString name, TString suffix="") {
0066     TString hist_name = dir + "/" + name + "_vs_p";
0067     if(suffix!="") hist_name += "_" + suffix;
0068     fmt::print("get histogram '{}'\n",hist_name);
0069     auto hist = hadd_file->Get<TH2D>(hist_name);
0070     hist->SetName(name+"_scan");
0071     scans.insert({name,hist});
0072   };
0073   get_scan( "pid_irt", "theta", radiator_name );
0074
0075   // re-binning
0076   for(const auto& [name,scan] : scans) {
0077     if(n_group_rebin>1) {
0078       fmt::print("REBIN: n_group_rebin={}\n", n_group_rebin);
0079       scan->RebinX(n_group_rebin);
0080     }
0081   }
0082
0083   // draw
0084   gStyle->SetOptStat(0);
0085   std::map<TString,TCanvas*> scan_canvs;
0086   for(const auto& [name,scan] : scans) {
0087     TString canv_name(name+"_canv");
0088     auto canv = new TCanvas(canv_name,canv_name);
0089     canv->SetGrid(1,1);
0090     scan->Draw("COLZ");
0091     scan->GetXaxis()->SetRangeUser(0,mom_max);
0092     scan->GetYaxis()->SetRangeUser(0,theta_max);
0093     for(auto [particle,curve] : theta_curves) {
0094       curve->Draw("SAME");
0095     }
0096     canv->SaveAs(out_base_name + "_" + name + ".png");
0097     scan_canvs.insert({name,canv});
0098   }
0099
0100   // write output
0101   auto WriteScans = [] (auto vec) { for(const auto& [name,obj] : vec) obj->Write(); };
0102   WriteScans(scans);
0103   WriteScans(scan_canvs);
0104   hadd_file->Close();
0105   fmt::print("\n{:=<50}\nwrote: {}\n\n","",hadd_file_name);
0106 }