benchmarks/beamline/shared_functions.h

0001 #include "DD4hep/Detector.h"
0002 #include "DDRec/CellIDPositionConverter.h"
0003 #include "edm4hep/MCParticleCollection.h"
0004 #include "edm4hep/SimTrackerHitCollection.h"
0005 #include "edm4hep/SimCalorimeterHitCollection.h"
0006 #include "DD4hep/VolumeManager.h"
0007 #include "TFile.h"
0008
0009 using RVecHits = ROOT::VecOps::RVec<edm4hep::SimTrackerHitData>;
0010
0011 //-----------------------------------------------------------------------------------------
0012 // Grab Component functor
0013 //-----------------------------------------------------------------------------------------
0014 struct getSubID{
0015   getSubID(std::string cname, dd4hep::Detector& det, std::string rname = "BackwardsBeamlineHits") : componentName(cname), detector(det), readoutName(rname){}
0016
0017   ROOT::VecOps::RVec<int> operator()(const RVecHits& evt) {
0018     auto decoder = detector.readout(readoutName).idSpec().decoder();
0019     auto indexID = decoder->index(componentName);
0020     ROOT::VecOps::RVec<int> result;
0021     for(const auto& h: evt) {
0022       result.push_back(decoder->get(h.cellID,indexID));
0023     }
0024     return result;
0025   };
0026
0027   void SetComponent(std::string cname){
0028     componentName = cname;
0029   }
0030   void SetReadout(std::string rname){
0031     readoutName = rname;
0032   }
0033
0034   private:
0035   std::string componentName;
0036   dd4hep::Detector& detector;
0037   std::string readoutName;
0038 };
0039
0040 //-----------------------------------------------------------------------------------------
0041 // Transform global x,y,z position and momentum into local coordinates
0042 //-----------------------------------------------------------------------------------------
0043 struct globalToLocal{
0044   globalToLocal(dd4hep::Detector& det) : detector(det){
0045      volumeManager = dd4hep::VolumeManager::getVolumeManager(det);
0046   }
0047
0048   ROOT::VecOps::RVec<ROOT::VecOps::RVec<double>> operator()(const RVecHits& evt) {
0049
0050     ROOT::VecOps::RVec<ROOT::VecOps::RVec<double>> result;
0051     ROOT::VecOps::RVec<double> xvec;
0052     ROOT::VecOps::RVec<double> yvec;
0053     ROOT::VecOps::RVec<double> zvec;
0054     ROOT::VecOps::RVec<double> pxvec;
0055     ROOT::VecOps::RVec<double> pyvec;
0056     ROOT::VecOps::RVec<double> pzvec;
0057
0058     for(const auto& h: evt) {
0059       auto cellID = h.cellID;
0060       // dd4hep::DetElement detelement = volumeManager.lookupDetElement(cellID);
0061       // auto detelement = volumeManager.lookupVolumePlacement(cellID);
0062       auto detelement = volumeManager.lookupDetElement(cellID);
0063       const TGeoMatrix& transform = detelement.nominal().worldTransformation();
0064       // transform.Print();
0065       //auto transform = volumeManager.worldTransformation(cellID);
0066       //Convert position to local coordinates
0067       auto pos = h.position;
0068       Double_t globalPos[3] = {pos.x/10, pos.y/10, pos.z/10};
0069       Double_t localPos[3];
0070       transform.MasterToLocal(globalPos, localPos);
0071       // std::cout << "Global: " << globalPos[0] << " " << globalPos[1] << " " << globalPos[2] << std::endl;
0072       // std::cout << "Local: " << localPos[0] << " " << localPos[1] << " " << localPos[2] << std::endl;
0073       //Transform global momentum to local coordinates
0074       auto mom = h.momentum;
0075       Double_t globalMom[3] = {mom.x, mom.y, mom.z};
0076       Double_t localMom[3];
0077       transform.MasterToLocalVect(globalMom, localMom);
0078       // std::cout << "Global: " << globalMom[0] << " " << globalMom[1] << " " << globalMom[2] << std::endl;
0079       // std::cout << "Local: " << localMom[0] << " " << localMom[1] << " " << localMom[2] << std::endl;
0080
0081       xvec.push_back(localPos[0]);
0082       yvec.push_back(localPos[1]);
0083       zvec.push_back(localPos[2]);
0084       pxvec.push_back(localMom[0]);
0085       pyvec.push_back(localMom[1]);
0086       pzvec.push_back(localMom[2]);
0087
0088     }
0089
0090     result.push_back(xvec);
0091     result.push_back(yvec);
0092     result.push_back(zvec);
0093     result.push_back(pxvec);
0094     result.push_back(pyvec);
0095     result.push_back(pzvec);
0096
0097     return result;
0098   };
0099
0100   private:
0101   dd4hep::Detector& detector;
0102   dd4hep::VolumeManager volumeManager;
0103 };
0104
0105 struct volParams{
0106   double radius;
0107   double xPos;
0108   double yPos;
0109   double zPos;
0110   double rotation;
0111 };
0112
0113 // Functor to get the volume element parameters from the CellID
0114 struct getVolumeParametersFromCellID {
0115   getVolumeParametersFromCellID(dd4hep::Detector& det) : detector(det) {
0116     volumeManager = dd4hep::VolumeManager::getVolumeManager(det);
0117   }
0118
0119   ROOT::VecOps::RVec<volParams> operator()(const RVecHits& evt) {
0120     ROOT::VecOps::RVec<volParams> result;
0121     // Look up the detector element using the CellID
0122     for(const auto& h : evt) {
0123       auto  cellID = h.cellID;
0124       auto  detelement = volumeManager.lookupDetElement(cellID);
0125       const TGeoMatrix& transform      = detelement.nominal().worldTransformation();
0126       const Double_t*   translation    = transform.GetTranslation();
0127       const Double_t*   rotationMatrix = transform.GetRotationMatrix(); // Compute rotation angle around the Y-axis
0128       double rotationAngleY = std::atan2(rotationMatrix[2], rotationMatrix[8]); // atan2(r13, r33)
0129       auto volume = detelement.volume();
0130       dd4hep::ConeSegment cone = volume.solid();
0131       volParams params{
0132         cone.rMax1(),
0133         translation[0],
0134         translation[1],
0135         translation[2],
0136         rotationAngleY
0137       };
0138       result.push_back(params);
0139     }
0140     return result;
0141   }
0142
0143 private:
0144   dd4hep::Detector& detector;
0145   dd4hep::VolumeManager volumeManager;
0146 };
0147
0148 TH1F* CreateFittedHistogram(const std::string& histName,
0149   const std::string& histTitle,
0150   const std::map<TString, double>& valueMap,
0151   const std::map<TString, double>& errorMap,
0152   const std::string& xAxisLabel) {
0153   // Number of bins corresponds to the number of entries in the value map
0154   int nPoints = valueMap.size();
0155   TH1F* hist = new TH1F(histName.c_str(), (";" + xAxisLabel + ";" + histTitle).c_str(), nPoints, 0.5, nPoints + 0.5);
0156
0157   // Fill the histogram with values and errors, and set custom bin labels
0158   int binIndex = 1; // Start from bin 1
0159   for (const auto& [name, value] : valueMap) {
0160     hist->SetBinContent(binIndex, value); // Set the bin content
0161     if(errorMap.size()==valueMap.size()){
0162       hist->SetBinError(binIndex, errorMap.at(name)); // Set the bin error
0163     }
0164     else{
0165       hist->SetBinError(binIndex, 0); // Set the bin error to 0 if not provided
0166     }
0167     hist->GetXaxis()->SetBinLabel(binIndex, name); // Set the bin label
0168     binIndex++;
0169   }
0170
0171   // Customize the histogram
0172   hist->SetMarkerStyle(20);
0173   hist->SetMarkerSize(1.0);
0174   hist->SetLineColor(kBlue);
0175   hist->SetMarkerColor(kRed);
0176
0177   return hist;
0178 }