Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-01-18 09:15:24

0001 // dRICH event display
0002 
0003 // test readout segmentation
0004 #include <cstdlib>
0005 #include <iostream>
0006 #include <bitset>
0007 #include <map>
0008 #include <vector>
0009 #include <algorithm>
0010 #include <fmt/format.h>
0011 
0012 // ROOT
0013 #include <TSystem.h>
0014 #include <TStyle.h>
0015 #include <TCanvas.h>
0016 #include <TApplication.h>
0017 #include <TBox.h>
0018 #include <ROOT/RDataFrame.hxx>
0019 #include <ROOT/RDF/HistoModels.hxx>
0020 #include <ROOT/RVec.hxx>
0021 
0022 // DD4Hep
0023 #include <DD4hep/Detector.h>
0024 #include <DDRec/DetectorData.h>
0025 
0026 // local
0027 #include "WhichRICH.h"
0028 
0029 using namespace ROOT;
0030 using namespace dd4hep;
0031 
0032 int main(int argc, char** argv) {
0033 
0034   // arguments
0035   // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
0036   if(argc<=3) {
0037     fmt::print("\nUSAGE: {} [d/p] [s/r] [input_root_file] [i/n] [event_num_min] [event_num_max]\n\n",argv[0]);
0038     fmt::print("    [d/p]: detector\n");
0039     fmt::print("         - d: for dRICH\n");
0040     fmt::print("         - p: for pfRICH\n");
0041     fmt::print("\n");
0042     fmt::print("    [s/r]: file type:\n");
0043     fmt::print("         - s: for simulation file (all photons)\n");
0044     fmt::print("         - r: for reconstructed file (digitized hits)\n");
0045     fmt::print("\n");
0046     fmt::print("    [input_root_file]: output from simulation or reconstruction\n");
0047     fmt::print("\n");
0048     fmt::print("    [i/n]: interactive mode (optional)\n");
0049     fmt::print("         - i: keep TCanvases open for interaction\n");
0050     fmt::print("         - n: do not open TCanvases and just produce images (default)\n");
0051     fmt::print("\n");
0052     fmt::print("    [event_num_min]: minimum event number (optional)\n");
0053     fmt::print("         - if unspecified, draw sum of all events\n");
0054     fmt::print("         - if specified, but without [event_num_max], draw only\n");
0055     fmt::print("           this event\n");
0056     fmt::print("         - if specified with [event_num_max], draw the range\n");
0057     fmt::print("           of events, ONE at a time\n");
0058     fmt::print("\n");
0059     fmt::print("    [event_num_max]: maximum event number (optional)\n");
0060     fmt::print("         - set to 0 if you want the maximum possible\n");
0061     fmt::print("\n");
0062     return 2;
0063   }
0064   std::string zDirectionStr  = argv[1];
0065   std::string fileType       = argv[2];
0066   TString     infileN        = TString(argv[3]);
0067   TString     interactiveOpt = argc>4 ? TString(argv[4]) : "n";
0068   int         evnumMin       = argc>5 ? std::atoi(argv[5]) : -1;
0069   int         evnumMax       = argc>6 ? std::atoi(argv[6]) : -1;
0070   WhichRICH wr(zDirectionStr);
0071   if(!wr.valid) return 1;
0072   bool interactiveOn = interactiveOpt=="i";
0073 
0074   // settings
0075   // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
0076 
0077   // number of pixels along sensor side
0078   const Int_t numPx = 8;
0079 
0080   // expected range of values of `x` and `y` `cellID` bit fields
0081   const Int_t segXmin = 0;
0082   const Int_t segXmax = numPx-1;
0083 
0084   // dilations: for re-scaling module positions and segment positions
0085   // for drawing; if you change `numPx`, consider tuning these parameters
0086   // as well
0087   const Double_t dilation = 4.3;
0088 
0089   // drawing
0090   gStyle->SetPalette(55);
0091   gStyle->SetOptStat(0);
0092   Bool_t singleCanvas = true; // if true, draw all hitmaps on one canvas
0093   if(!interactiveOn)
0094     singleCanvas = true;
0095 
0096   // data collections
0097   std::string inputCollection;
0098   if(fileType=="s")
0099     inputCollection = wr.readoutName;
0100   else if(fileType=="r")
0101     inputCollection = wr.rawHitsName;
0102   else {
0103     fmt::print(stderr,"ERROR: unknown file type '{}'\n",fileType);
0104     return 1;
0105   }
0106   fmt::print("Reading collection '{}'\n",inputCollection);
0107 
0108   // setup
0109   // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
0110 
0111   // define application environment, to keep canvases open, if interactive mode is on
0112   TApplication *mainApp;
0113   if(interactiveOn)
0114     mainApp = new TApplication("mainApp",&argc,argv);
0115 
0116   // main dataframe
0117   RDataFrame dfIn("events",infileN.Data());
0118   auto numEvents = dfIn.Count().GetValue();
0119 
0120   // determine event numbers to read
0121   std::vector<std::pair<int,int>> evnumRanges;
0122   if(evnumMin<0) {
0123     fmt::print("Reading all events\n");
0124     evnumRanges.push_back({ 0, 0 });
0125   }
0126   else if(evnumMax<0) {
0127     fmt::print("Reading only event number {}\n",evnumMin);
0128     evnumRanges.push_back({ evnumMin, evnumMin+1 });
0129   }
0130   else {
0131     if(evnumMax==0)
0132       evnumMax = numEvents-1;
0133     if(evnumMax>=numEvents) {
0134       fmt::print("WARNING: there are only {} events\n",numEvents);
0135       evnumMax = numEvents-1;
0136     }
0137     fmt::print("Reading event numbers {} to {}, one at a time\n",evnumMin,evnumMax);
0138     for(int e=evnumMin; e<=evnumMax; e++)
0139       evnumRanges.push_back({ e, e+1 });
0140     if(interactiveOn) {
0141       fmt::print(stderr,"ERROR: cannot yet run with interactive mode enabled with an event number range (FIXME)\n");
0142       return 1;
0143     }
0144   }
0145 
0146 
0147   // compact file name
0148   std::string DETECTOR_PATH(getenv("DETECTOR_PATH"));
0149   std::string DETECTOR(getenv("DETECTOR"));
0150   if(DETECTOR_PATH.empty() || DETECTOR.empty()) {
0151     fmt::print(stderr,"ERROR: source environ.sh\n");
0152     return 1;
0153   }
0154   std::string compactFile = DETECTOR_PATH + "/" + DETECTOR + ".xml";
0155 
0156   // get detector handle and some constants
0157   const std::string richName = wr.XRICH;
0158   const auto det = &(Detector::getInstance());
0159   det->fromXML(compactFile);
0160   const auto detRich  = det->detector(richName);
0161   const auto nSectors = wr.zDirection>0 ? det->constant<int>(wr.XRICH+"_num_sectors") : 1;
0162 
0163   // cellID decoder
0164   /* - `decodeCellID(fieldName)` returns a "decoder" for the field with name `fieldName`
0165    * - this decoder is a function that maps an `RVecUL` of `cellID`s to an
0166    *   `RVecUL` of correpsonding field element values
0167    */
0168   const auto readoutCoder = det->readout(wr.readoutName).idSpec().decoder();
0169   auto decodeCellID = [&readoutCoder] (std::string fieldName) {
0170     return [&readoutCoder,&fieldName] (RVecUL cellIDvec) {
0171       RVecL result;
0172       bool found=false;
0173       for(auto elem : readoutCoder->fields())
0174         if(fieldName == elem.name()) { found = true; break; }
0175       // if(!found) fmt::print("- skipping missing bit field \"{}\"\n",fieldName);
0176       for(const auto& cellID : cellIDvec) {
0177         if(found) {
0178           auto val = readoutCoder->get(cellID,fieldName); // get BitFieldElement value
0179           result.emplace_back(val);
0180           // fmt::print("decode {}: {:64b} -> {}\n",fieldName,cellID,val);
0181         } else result.emplace_back(0);
0182       }
0183       return result;
0184     };
0185   };
0186 
0187   // build sensor position LUT `imod2hitmapXY`
0188   /* - find the sector 0 sensors, and build a map of their module number `pdu`,`sipm` to
0189    *   X and Y coordinates to use in the hitmap
0190    *   - these hitmap coordinates are from the sensor position X and Y, rescaled
0191    *     by the factor `dilation` and rounded to the nearest integer
0192    *   - also builds a list of `TBox`es, for drawing the sensors on the hitmap
0193    */
0194   std::map<std::pair<Long_t,Long64_t>,std::pair<Long64_t,Long64_t>> imod2hitmapXY;
0195   std::vector<TBox*> boxList;
0196   for(auto const& [de_name, detSensor] : detRich.children()) {
0197     if(de_name.find(wr.sensorNamePattern)!=std::string::npos) {
0198       // convert global position to hitmapX and Y
0199       const auto detSensorPars = detSensor.extension<rec::VariantParameters>(true);
0200       if(detSensorPars==nullptr)
0201         throw std::runtime_error(fmt::format("sensor '{}' does not have VariantParameters", de_name));
0202       auto posSensorX = detSensorPars->get<double>("pos_x");
0203       auto posSensorY = detSensorPars->get<double>("pos_y");
0204       auto hitmapX    = Long64_t(dilation*posSensorX + 0.5);
0205       auto hitmapY    = Long64_t(dilation*posSensorY + 0.5);
0206       // convert unique cellID to module number, using the cellID decoder
0207       auto sensorID = ULong_t(detSensor.id());
0208       auto pdu      = decodeCellID("pdu")(RVecUL({sensorID})).front();
0209       auto sipm     = decodeCellID("sipm")(RVecUL({sensorID})).front();
0210       // add to `imod2hitmapXY` and create sensor `TBox`
0211       imod2hitmapXY.insert({{pdu,sipm}, {hitmapX,hitmapY}});
0212       boxList.push_back(new TBox(
0213             hitmapX, 
0214             hitmapY,
0215             hitmapX + numPx,
0216             hitmapY + numPx
0217             ));
0218       boxList.back()->SetLineColor(kGreen-10);
0219       boxList.back()->SetFillColor(kGreen-10);
0220       boxList.back()->SetFillStyle(1001);
0221     }
0222   }
0223 
0224   // convert vectors of `pdu`, `sipm` to vector of hitmap X or Y
0225   auto imod2hitmapXY_get = [&imod2hitmapXY] (RVecL pduVec, RVecL sipmVec, int c) {
0226     RVecL result;
0227     Long64_t pos;
0228     if(pduVec.size() != sipmVec.size())
0229       throw std::runtime_error("imod2hitmapXY_get input vectors differ in size");
0230     for(unsigned long i=0; i<pduVec.size(); i++) {
0231       auto pdu  = pduVec.at(i);
0232       auto sipm = sipmVec.at(i);
0233       try {
0234         pos = (c==0) ?
0235           imod2hitmapXY.at({pdu,sipm}).first :
0236           imod2hitmapXY.at({pdu,sipm}).second;
0237       }
0238       catch (const std::out_of_range &ex) {
0239         fmt::print(stderr,"ERROR: cannot find sensor module pdu={} sipm={}\n",pdu,sipm);
0240         pos = 0.0;
0241       };
0242       result.emplace_back(pos);
0243     }
0244     return result;
0245   };
0246   auto imod2hitmapX = [&imod2hitmapXY_get] (RVecL pduVec, RVecL sipmVec) { return imod2hitmapXY_get(pduVec,sipmVec,0); };
0247   auto imod2hitmapY = [&imod2hitmapXY_get] (RVecL pduVec, RVecL sipmVec) { return imod2hitmapXY_get(pduVec,sipmVec,1); };
0248 
0249   // convert vector of hitmap X (or Y) + vector of segmentation X (or Y) to vector of pixel X (or Y)
0250   auto pixelCoord = [] (RVecL hitmapXvec, RVecL segXvec) {
0251     RVecL result = hitmapXvec + segXvec;
0252     return result;
0253   };
0254 
0255 
0256   // dataframe transformations
0257   // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
0258 
0259   // decode cellID to bit field element values
0260   auto dfDecoded = dfIn
0261       .Alias("cellID", inputCollection+".cellID")
0262       .Define("system", decodeCellID("system"), {"cellID"})
0263       .Define("sector", decodeCellID("sector"), {"cellID"})
0264       .Define("pdu",    decodeCellID("pdu"),    {"cellID"})
0265       .Define("sipm",   decodeCellID("sipm"),   {"cellID"})
0266       .Define("x",      decodeCellID("x"),      {"cellID"})
0267       .Define("y",      decodeCellID("y"),      {"cellID"})
0268       ;
0269 
0270   // map `(pdu,sipm,x,y)` to pixel hitmap bins
0271   auto dfHitmap = dfDecoded
0272       .Define("hitmapX", imod2hitmapX, {"pdu","sipm"})
0273       .Define("hitmapY", imod2hitmapY, {"pdu","sipm"})
0274       .Define("pixelX", pixelCoord, {"hitmapX","x"})
0275       .Define("pixelY", pixelCoord, {"hitmapY","y"})
0276       ;
0277 
0278   // count how many hits are inside the expected segmentation box
0279   auto countInBox = [&segXmin,&segXmax] (RVecL xVec, RVecL yVec) {
0280     return xVec[ xVec>=segXmin && xVec<=segXmax && yVec>=segXmin && yVec<=segXmax ].size();
0281   };
0282   auto countOutBox = [&segXmin,&segXmax] (RVecL xVec, RVecL yVec) {
0283     return xVec[ xVec<segXmin  || xVec>segXmax  || yVec<segXmin  || yVec>segXmax  ].size();
0284   };
0285   auto numInBox  = dfDecoded.Define( "numInBox",  countInBox,  {"x","y"} ).Sum("numInBox").GetValue();
0286   auto numOutBox = dfDecoded.Define( "numOutBox", countOutBox, {"x","y"} ).Sum("numOutBox").GetValue();
0287   fmt::print("{:=<60}\nNUMBER OF HITS OUTSIDE EXPECTED BOX: {} / {} ({:.4f}%)\n{:=<60}\n",
0288       "", numOutBox, numInBox+numOutBox, 100*Double_t(numOutBox)/Double_t(numInBox+numOutBox), "");
0289 
0290   // loop over event number(s)
0291   for(auto& [evnum,evnum_stop] : evnumRanges) {
0292 
0293     // cut on specified event(s)
0294     auto dfFinal = dfHitmap.Range(evnum,evnum_stop);
0295 
0296     // histograms
0297     // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
0298 
0299     // cellID field histograms
0300     auto fieldHists = std::vector({
0301       dfFinal.Histo1D("system"),
0302       dfFinal.Histo1D("sector"),
0303       dfFinal.Histo1D("pdu"),
0304       dfFinal.Histo1D("sipm"),
0305       dfFinal.Histo1D("x"),
0306       dfFinal.Histo1D("y")
0307     });
0308     const int segXmaxPlot = 10;
0309     auto segXY = dfFinal.Histo2D(
0310         { "segXY", "SiPM pixels;x;y",
0311           2*segXmaxPlot, -segXmaxPlot, segXmaxPlot,
0312           2*segXmaxPlot, -segXmaxPlot, segXmaxPlot },
0313         "x","y"
0314         );
0315 
0316 
0317     // pixel hitmap
0318     Double_t pixelXmin = dilation * wr.plotXmin;
0319     Double_t pixelXmax = dilation * wr.plotXmax;
0320     Double_t pixelYmin = dilation * wr.plotYmin;
0321     Double_t pixelYmax = dilation * wr.plotYmax;
0322     auto pixelHitmapModel = RDF::TH3DModel(
0323         "pixelHitmap", "Pixel Hit Map;x;y;sector",
0324         (Int_t)(pixelXmax-pixelXmin), pixelXmin, pixelXmax,
0325         (Int_t)(pixelYmax-pixelYmin), pixelYmin, pixelYmax,
0326         nSectors,0,double(nSectors)
0327         );
0328     auto pixelHitmap = fileType=="r" ? // weight by ADC counts, if reading digitized hits
0329       dfFinal.Histo3D(pixelHitmapModel, "pixelX", "pixelY", "sector", inputCollection+".charge") :
0330       dfFinal.Histo3D(pixelHitmapModel, "pixelX", "pixelY", "sector");
0331 
0332     // draw
0333     // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
0334 
0335     TCanvas *c;
0336     if(interactiveOn) {
0337       // draw cellID field histograms
0338       c = new TCanvas();
0339       c->Divide(4,2);
0340       int pad=1;
0341       for(auto hist : fieldHists) {
0342         c->GetPad(pad)->SetLogy();
0343         c->cd(pad);
0344         // if(TString(hist->GetName())!="pdu") hist->SetBarWidth(4);
0345         hist->SetLineColor(kBlack);
0346         hist->SetFillColor(kBlack);
0347         hist->Draw("bar");
0348         pad++;
0349       }
0350 
0351       // draw segmentation XY plot, along with expected box
0352       c->cd(pad);
0353       c->GetPad(pad)->SetGrid(1,1);
0354       segXY->Draw("colz");
0355       auto expectedBox = new TBox(segXmin, segXmin, segXmax+1, segXmax+1);
0356       expectedBox->SetFillStyle(0);
0357       expectedBox->SetLineColor(kBlack);
0358       expectedBox->SetLineWidth(4);
0359       expectedBox->Draw("same");
0360     }
0361 
0362     // draw pixel hitmap
0363     if(singleCanvas) { c = new TCanvas("canv","canv",3*700,2*600); c->Divide(3,2); };
0364     int secBin;
0365     TH2D *pixelHitmapSec[nSectors];
0366     for(int sec=0; sec<nSectors; sec++) {
0367       if(singleCanvas) c->cd(sec+1); else c = new TCanvas();
0368       secBin = pixelHitmap->GetZaxis()->FindBin(Double_t(sec));
0369       pixelHitmap->GetZaxis()->SetRange(secBin,secBin);
0370 
0371       pixelHitmapSec[sec] = (TH2D*) pixelHitmap->Project3D("yx");
0372       if(fileType=="r") pixelHitmapSec[sec]->SetMaximum(4096);
0373       pixelHitmapSec[sec]->SetName(Form("pixelHitmap_s%d",sec));
0374 
0375       TString pixelHitmapTitle;
0376       if(fileType=="s")      pixelHitmapTitle = "Photon hits";
0377       else if(fileType=="r") pixelHitmapTitle = "Digitized hits";
0378       pixelHitmapTitle += Form(", sector %d",sec);
0379       if(evnumMin<0) pixelHitmapTitle += ", all events";
0380       else           pixelHitmapTitle += Form(", event %d",evnum);
0381 
0382       pixelHitmapSec[sec]->SetTitle(pixelHitmapTitle);
0383       pixelHitmapSec[sec]->Draw("colz");
0384       for(auto box : boxList) {
0385         box->Draw("same");
0386       };
0387       pixelHitmapSec[sec]->Draw("colz same");
0388       // pixelHitmapSec[sec]->GetXaxis()->SetRangeUser(475, 725);  // zoom in
0389       // pixelHitmapSec[sec]->GetYaxis()->SetRangeUser(-150, 150);
0390     };
0391 
0392     // either hold the TCanvases open, or save them as PNG files
0393     if(interactiveOn) {
0394       fmt::print("\n\npress ^C to exit.\n\n");
0395       mainApp->Run();
0396     } else {
0397       gROOT->ProcessLine(".! mkdir -p out/ev");
0398       TString pngName = Form("out/ev/%s.png", fmt::format("{:08}",evnum).c_str()); 
0399       c->SaveAs(pngName);
0400       fmt::print("Saved image: {}", pngName);
0401       // cleanup and avoid memory leaks # FIXME: refactor this... and there is still a slow leak...
0402       delete c;
0403       for(int sec=0; sec<nSectors; sec++) delete pixelHitmapSec[sec];
0404     }
0405 
0406   } // end evnumRanges loop
0407 
0408   if(!interactiveOn)
0409     fmt::print("\n\nEvent display images written to out/ev/*.png\n\n");
0410 
0411   return 0;
0412 };