Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-07-01 07:56:13

0001 // Script to plot the x and y positions and momentum of beamline particles as they pass through the magnets
0002 
0003 #include <iostream>
0004 #include "DD4hep/Detector.h"
0005 #include "DDRec/CellIDPositionConverter.h"
0006 #include "edm4hep/MCParticleCollection.h"
0007 #include "edm4hep/SimTrackerHitCollection.h"
0008 #include "edm4hep/SimCalorimeterHitCollection.h"
0009 #include "edm4hep/Vector3f.h"
0010 #include "edm4hep/Vector3d.h"
0011 #include "ROOT/RDataFrame.hxx"
0012 #include "ROOT/RDF/RInterface.hxx"
0013 #include "ROOT/RVec.hxx"
0014 #include "shared_functions.h"
0015 #include "TCanvas.h"
0016 #include "TStyle.h"
0017 
0018 
0019 using RVecS       = ROOT::VecOps::RVec<string>;
0020 using RNode       = ROOT::RDF::RNode;
0021 
0022 int acceptanceAnalysis( TString inFile             = "/scratch/EIC/G4out/beamline/acceptanceTest.edm4hep.root",
0023                         TString outFile            = "output.root",
0024                         std::string compactName    = "/home/simong/EIC/epic/install/share/epic/epic_ip6_extended.xml",
0025                         TString beampipeCanvasName = "acceptance_in_beampipe.png",
0026                         TString EThetaCanvasName   = "acceptance_energy_theta.png",
0027                         TString EThetaAccCanvasName= "acceptance_energy_theta_acceptance.png",
0028                         TString entryFractionCanvasName = "acceptance_entries.png") {
0029 
0030     //Set ROOT style    
0031     gStyle->SetPadLeftMargin(0.1);  // Set left margin
0032     gStyle->SetPadRightMargin(0.0); // Set right margin
0033     gStyle->SetPadTopMargin(0.0);   // Set top margin
0034     gStyle->SetPadBottomMargin(0.1);// Set bottom margin
0035     gStyle->SetTitleAlign(13);
0036     gStyle->SetTitleX(0.12);          // Place the title on the top right
0037     gStyle->SetTitleY(0.985);         // Place the title on the top right
0038     gStyle->SetTitleSize(0.08, "t");
0039     gStyle->SetTitleXOffset(1.0);    // Adjust y-axis title offset
0040     gStyle->SetTitleYOffset(1.0);    // Adjust y-axis title offset
0041     gStyle->SetOptStat(0);
0042 
0043     int pass = 0;
0044 
0045     //Set implicit multi-threading
0046     // ROOT::EnableImplicitMT();
0047        
0048     //Load the detector config
0049     dd4hep::Detector& detector = dd4hep::Detector::getInstance();
0050     detector.fromCompact(compactName);
0051  
0052     ROOT::RDataFrame d0("events",inFile, {"BackwardsBeamlineHits"});
0053     RNode d1 = d0;
0054     RVecS colNames = d1.GetColumnNames();
0055 
0056     //Get total number of entries
0057     int nEntries = d1.Count().GetValue();
0058     //Set number of entries to process
0059     // d1 = d1.Range(0,1000);
0060     
0061     //Get the collection 
0062     std::string mcParticlesName = "MCParticles";
0063     std::string readoutName = "BackwardsBeamlineHits";  
0064 
0065     std::cout << "Running lazy RDataframe execution" << std::endl;
0066 
0067     if(Any(colNames==mcParticlesName)){
0068 
0069         // Get theta and energy of the particles
0070         d1 = d1.Define("theta",[](const vector<edm4hep::MCParticleData>& mcParticles) {
0071             // Calculate theta from momentum components as angle from negative z axis
0072             double px = mcParticles[0].momentum.x;
0073             double py = mcParticles[0].momentum.y;
0074             double pz = mcParticles[0].momentum.z;
0075             double p = std::sqrt(px*px + py*py + pz*pz);
0076             double theta = M_PI-std::acos(pz / p); // Angle from the z-axis
0077             
0078             return theta;
0079         }, {mcParticlesName})
0080         .Define("energy",[](const vector<edm4hep::MCParticleData>& mcParticles) {
0081            
0082             //Calculate energy from mass and momentum
0083             double mass = mcParticles[0].mass;
0084             double px = mcParticles[0].momentum.x;
0085             double py = mcParticles[0].momentum.y;
0086             double pz = mcParticles[0].momentum.z;
0087             double energy = std::sqrt(px*px + py*py + pz*pz + mass*mass);
0088             
0089             return energy;
0090         }, {mcParticlesName});
0091 
0092     }
0093     else{
0094         std::cout << "Collection " << mcParticlesName << " not found in file" << std::endl;
0095         return 1;
0096     }
0097 
0098     int eBins = 100;
0099     int thetaBins = 100;
0100 
0101     auto totalETheta = d1.Histo2D({"Energy vs Theta","Energy vs Theta",eBins,6,18,thetaBins,0,0.011},"energy","theta");
0102 
0103     if(Any(colNames==readoutName)){
0104 
0105         d1 = d1.Define("pipeID",getSubID("pipe",detector),{readoutName})
0106                .Define("endID",getSubID("end",detector),{readoutName})
0107                .Define("pipeParameters",getVolumeParametersFromCellID(detector),{readoutName})
0108                .Define("pipeRadius",[](const ROOT::VecOps::RVec<volParams>& params) {
0109                 ROOT::VecOps::RVec<double> radii;
0110                 for (const auto& param : params) {
0111                     radii.push_back(param.radius);
0112                 }
0113                 return radii;
0114                 }, {"pipeParameters"});
0115                 
0116 
0117         //global x,y,z position and momentum
0118         d1 = d1 .Define("NHits","BackwardsBeamlineHits.size()");
0119         
0120         d1 = d1.Define("hitPosMom",globalToLocal(detector),{readoutName})
0121                 .Define("xpos","hitPosMom[0]")
0122                 .Define("ypos","hitPosMom[1]")
0123                 .Define("zpos","hitPosMom[2]");        
0124 
0125     }
0126     else{
0127         std::cout << "Collection " << readoutName << " not found in file" << std::endl;
0128         return 1;
0129     }    
0130 
0131     // Calculate the maximum pipe radius for plotting
0132     auto maxPipeRadius = 1.2*d1.Max("pipeRadius").GetValue();
0133 
0134     std::cout << "Executing Analysis and creating histograms" << std::endl;
0135 
0136     //Create array of histogram results
0137     std::map<TString,ROOT::RDF::RResultPtr<TH2D>> hHistsxy;
0138     std::map<TString,ROOT::RDF::RResultPtr<TH2D>> hHistsETheta;
0139     
0140 
0141     std::map<TString,ROOT::RDF::RResultPtr<double>> pipeRadii;
0142     std::map<TString,double> filterEntries;
0143     std::map<TString,double> filterAcceptanceIntegral;
0144     
0145     //Create histograms
0146     for(int i=0; i<=7; i++){
0147 
0148         std::string name = "pipeID";
0149         std::string str_i = std::to_string(i);
0150         name += str_i;
0151         auto filterDF = d1.Define("xposf","xpos[pipeID=="+str_i+"]")
0152                           .Define("yposf","ypos[pipeID=="+str_i+"]")
0153                           .Define("pipeRadiusf","pipeRadius[pipeID=="+str_i+"]");
0154                    
0155 
0156         TString beamspotName = "Beamspot ID"+str_i+";x offset [cm]; y offset [cm]";
0157         TString xyname = name+";x Offset [cm]; y Offset [cm]";
0158         TString xname = name+";x Offset [cm]; x trajectory component";
0159         TString yname = name+";y Offset [cm]; y trajectory component";
0160         hHistsxy[name] = filterDF.Histo2D({beamspotName,beamspotName,400,-maxPipeRadius,maxPipeRadius,400,-maxPipeRadius,maxPipeRadius},"xposf","yposf");
0161 
0162         auto extraFilterDF = filterDF.Filter("All(pipeID<"+std::to_string(i+1)+")&&Any(pipeID=="+std::to_string(i)+")" );
0163         TString EThetaName = "Energy vs Theta ID"+str_i+";Energy [GeV]; Theta [rad]";
0164         TString nameETheta = name+";Energy [GeV]; Theta [rad]";
0165         hHistsETheta[name] = extraFilterDF.Histo2D({EThetaName,EThetaName,eBins,4,18,thetaBins,0,0.011},"energy","theta");
0166 
0167         //Parameters of the pipe
0168         pipeRadii[name]    = filterDF.Max("pipeRadiusf");        
0169         // std::cout << "Pipe ID: " << name << " Radius: " << pipeRadii[name] << " " << filterDF.Min("pipeRadiusf").GetValue() << std::endl;
0170     
0171     }
0172 
0173     // Create histograms of the beamspot
0174     TCanvas *cXY = new TCanvas("beamspot_canvas","beamspot_canvas",3000,1600);
0175     cXY->Divide(4,2);
0176     int i=1;
0177     for(auto [name,h] : hHistsxy){
0178         // Get the pipe radius for this histogram
0179         auto pipeRadius = pipeRadii[name].GetValue();
0180         cXY->cd(i++);
0181 
0182         h->Draw("col");
0183         //Draw cicle
0184         TEllipse *circle = new TEllipse(0,0,pipeRadius);
0185         circle->SetLineColor(kRed);
0186         circle->SetLineWidth(2);
0187         circle->SetFillStyle(0);
0188         circle->Draw("same");
0189 
0190     }
0191 
0192     // Cnavas for energy vs theta
0193     TCanvas *cETheta = new TCanvas("energy_theta_canvas","energy_theta_canvas",3000,1600);
0194     cETheta->Divide(4,2);
0195     i=1;
0196     for(auto [name,h] : hHistsETheta){        
0197         cETheta->cd(i++);
0198         h->Draw("colz");
0199         filterEntries[name] = h->GetEntries()/ nEntries;
0200     }
0201   
0202     // Canvas for energy vs theta acceptance
0203     TCanvas *cEThetaAcc = new TCanvas("energy_theta_acceptance_canvas","energy_theta_acceptance_canvas",3000,1600);
0204     cEThetaAcc->Divide(4,2);
0205     i=1;
0206     for(auto [name,h] : hHistsETheta){
0207         cEThetaAcc->cd(i++);
0208         h->Divide(totalETheta.GetPtr());
0209         h->Draw("colz");
0210         filterAcceptanceIntegral[name] = h->Integral()/ (eBins*thetaBins); // Normalize by the number of bins
0211     }
0212 
0213     TH1F* hPipeEntries = CreateFittedHistogram("hNumberOfHits",
0214         "Number of Hits per Pipe ID",
0215         filterEntries,
0216         {},
0217         "Pipe ID");
0218     TH1F* hPipeAcceptance = CreateFittedHistogram("hPipeAcceptance",
0219         "Acceptance per Pipe ID",
0220         filterAcceptanceIntegral,
0221         {},
0222         "Pipe ID");
0223 
0224     TCanvas *cPipeAcceptance = new TCanvas("cPipeAcceptance", "Pipe Acceptance", 1200, 400);    
0225     cPipeAcceptance->Divide(2, 1);
0226     cPipeAcceptance->cd(1);
0227     hPipeEntries->Draw("");
0228     cPipeAcceptance->cd(2);
0229     hPipeAcceptance->Draw("");
0230     cPipeAcceptance->SetGrid();
0231     cPipeAcceptance->Update();
0232 
0233     // Save 2D canvases as pngs
0234     cXY->SaveAs(beampipeCanvasName);
0235     cETheta->SaveAs(EThetaCanvasName);
0236     cEThetaAcc->SaveAs(EThetaAccCanvasName);
0237     cPipeAcceptance->SaveAs(entryFractionCanvasName);
0238 
0239     TFile *f = new TFile(outFile,"RECREATE");
0240     cXY->Write();
0241     cETheta->Write();
0242     cEThetaAcc->Write();
0243     cPipeAcceptance->Write();
0244 
0245     f->Close();
0246 
0247     std::cout << "Analysis complete. Results saved to " << outFile << std::endl;
0248     return pass;
0249 
0250     // std::cout << "Saving events to file" << std::endl;
0251 
0252     //  ROOT::RDF::RSnapshotOptions opts;
0253     // opts.fMode = "UPDATE";
0254     // d1.Snapshot("events",outFile,{"pipeID","endID","pipeRadius","xpos","ypos","zpos","xmom","ymom","zmom","xpos_global","ypos_global","zpos_global","px_global","py_global","pz_global"},opts);
0255 }