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