Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-12-16 09:20:04

0001 #include "JetCalib.h"
0002 
0003 #include "JetContainerv1.h"
0004 
0005 #include <cdbobjects/CDBTF.h>  // for CDBTF1
0006 
0007 #include <ffamodules/CDBInterface.h>
0008 #include <ffaobjects/EventHeader.h>
0009 
0010 #include <globalvertex/GlobalVertex.h>
0011 #include <globalvertex/GlobalVertexMap.h>
0012 
0013 #include <fun4all/Fun4AllReturnCodes.h>
0014 #include <fun4all/SubsysReco.h>  // for SubsysReco
0015 
0016 #include <phool/PHCompositeNode.h>
0017 #include <phool/PHIODataNode.h>    // for PHIODataNode
0018 #include <phool/PHNode.h>          // for PHNode
0019 #include <phool/PHNodeIterator.h>  // for PHNodeIterator
0020 #include <phool/PHObject.h>        // for PHObject
0021 #include <phool/getClass.h>
0022 #include <phool/phool.h>
0023 #include <phool/recoConsts.h>
0024 
0025 #include <TF1.h>
0026 #include <TSystem.h>
0027 
0028 #include <boost/format.hpp>
0029 
0030 #include <cmath>
0031 #include <cstdlib>    // for exit
0032 #include <exception>  // for exception
0033 #include <iostream>   // for operator<<, basic_ostream
0034 #include <stdexcept>  // for runtime_error
0035 
0036 JetCalib::JetCalib(const std::string &name)
0037   : SubsysReco(name)
0038 {
0039   if (Verbosity() > 0)
0040   {
0041     std::cout << "JetCalib::JetCalib(const std::string &name) Calling ctor." << std::endl;
0042   }
0043 }
0044 
0045 JetCalib::~JetCalib()
0046 {
0047   if (Verbosity() > 0)
0048   {
0049     std::cout << "JetCalib::~JetCalib() : Calling dtor." << std::endl;
0050   }
0051 }
0052 
0053 int JetCalib::InitRun(PHCompositeNode *topNode)
0054 {
0055   // Create calib jet node.
0056   CreateNodeTree(topNode);
0057 
0058   // Load calibration file and function.
0059   if (fetchCalibDir("Default").empty())
0060   {
0061     std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : No default calibration available! Will apply calib factor 1." << std::endl;
0062   }
0063   else
0064   {
0065     m_JetCalibFile = new CDBTF(fetchCalibDir("Default"));
0066     if (m_JetCalibFile)
0067     {
0068       m_JetCalibFile->LoadCalibrations();
0069       std::string JetCalibFunc_name = "JES_Calib_Func_R0" + std::to_string((int)(jet_radius * 10));
0070       if (!ApplyZvrtxDependentCalib && !ApplyEtaDependentCalib)
0071       {
0072         int nZvrtxBins = 1;
0073         int nEtaBins = 1;
0074         m_JetCalibFunc.resize(nZvrtxBins);
0075         for (auto &row : m_JetCalibFunc)
0076         {
0077           row.resize(nEtaBins, nullptr);
0078         }
0079         TF1 *func_temp = m_JetCalibFile->getTF(JetCalibFunc_name + "_Default");
0080         if (!func_temp)
0081         {
0082           std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Could not find calibration function: " << JetCalibFunc_name << "_Default" << std::endl;
0083           exit(1);
0084         }
0085         m_JetCalibFunc[0][0] = func_temp;
0086       }
0087       else if (ApplyZvrtxDependentCalib && !ApplyEtaDependentCalib)
0088       {
0089         int nZvrtxBins = 5;
0090         int nEtaBins = 1;
0091         m_JetCalibFunc.resize(nZvrtxBins);
0092         for (auto &row : m_JetCalibFunc)
0093         {
0094           row.resize(nEtaBins, nullptr);
0095         }
0096         for (int iz = 0; iz < nZvrtxBins; ++iz)
0097         {
0098           TF1 *func_temp = m_JetCalibFile->getTF(JetCalibFunc_name + "_Z" + std::to_string(iz));
0099           if (!func_temp)
0100           {
0101             std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Could not find calibration function: " << JetCalibFunc_name << "_Z" << iz << std::endl;
0102             exit(1);
0103           }
0104           m_JetCalibFunc[iz][0] = func_temp;
0105         }
0106       }
0107       else if (ApplyEtaDependentCalib)
0108       {
0109         if (!ApplyZvrtxDependentCalib)
0110         {
0111           std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Must apply Zvrtx dependent calibration to apply eta dependent calibration. Applying Zvrtx + eta dependent calibration." << std::endl;
0112         }
0113         int nZvrtxBins = 5;
0114         int nEtaBins = 4;
0115         m_JetCalibFunc.resize(nZvrtxBins);
0116         for (auto &row : m_JetCalibFunc)
0117         {
0118           row.resize(nEtaBins, nullptr);
0119         }
0120         for (int iz = 0; iz < nZvrtxBins; ++iz)
0121         {
0122           if (iz < 3 && jet_radius <= 0.4) {
0123             for (int ieta = 0; ieta < nEtaBins; ++ieta)
0124             {
0125               TF1 *func_temp = m_JetCalibFile->getTF(JetCalibFunc_name + "_Z" + std::to_string(iz) + "_Eta" + std::to_string(ieta));
0126               if (!func_temp)
0127               {
0128                 std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Could not find calibration function: " << JetCalibFunc_name << "_Z" << iz << "_Eta" << ieta << std::endl;
0129                 exit(1);
0130               }
0131               m_JetCalibFunc[iz][ieta] = func_temp;
0132             }
0133           }
0134           else if (iz < 3 && jet_radius > 0.4)
0135           {
0136             for (int ieta = 0; ieta < nEtaBins; ++ieta)
0137             {
0138               TF1 *func_temp = m_JetCalibFile->getTF(JetCalibFunc_name + "_Z" + std::to_string(iz));
0139               if (!func_temp)
0140               {
0141                 std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Could not find calibration function: " << JetCalibFunc_name << "_Z" << iz << std::endl;
0142                 exit(1);
0143               }
0144               m_JetCalibFunc[iz][ieta] = func_temp;
0145             }
0146           }
0147           else
0148           {
0149             TF1 *func_temp = m_JetCalibFile->getTF(JetCalibFunc_name + "_Z" + std::to_string(iz));
0150             if (!func_temp)
0151             {
0152               std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Could not find calibration function: " << JetCalibFunc_name << "_Z" << iz << std::endl;
0153               exit(1);
0154             }
0155             m_JetCalibFunc[iz][0] = func_temp;
0156           }
0157         }
0158       }
0159     }
0160     else
0161     {
0162       std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) : Could not open calibration file!" << std::endl;
0163     }
0164   }
0165 
0166   if (Verbosity() > 0)
0167   {
0168     std::cout << "JetCalib::InitRun(PHCompositeNode *topNode) - InitRun with inputNode: " << m_inputNode << " outputNode: " << m_outputNode << std::endl;
0169   }
0170   return Fun4AllReturnCodes::EVENT_OK;
0171 }
0172 
0173 int JetCalib::process_event(PHCompositeNode *topNode)
0174 {
0175   // Get raw and calib jet nodes.
0176   JetContainer *_raw_jets = findNode::getClass<JetContainer>(topNode, m_inputNode);
0177   if (!_raw_jets)
0178   {
0179     std::cout << "JetCalib::process_event(PHCompositeNode *topNode) : Raw jet node not found! Cannot apply calibration." << std::endl;
0180     return Fun4AllReturnCodes::ABORTRUN;
0181   }
0182   JetContainer *_calib_jets = findNode::getClass<JetContainer>(topNode, m_outputNode);
0183   if (!_calib_jets)
0184   {
0185     std::cout << "JetCalib::process_event(PHCompositeNode *topNode) : Calib jet node not found! Cannot apply calibration." << std::endl;
0186     return Fun4AllReturnCodes::ABORTRUN;
0187   }
0188   GlobalVertexMap *vertexmap = nullptr;
0189   if (ApplyZvrtxDependentCalib)
0190   {
0191     vertexmap = findNode::getClass<GlobalVertexMap>(topNode, m_zvrtxNode);
0192     if (!vertexmap)
0193     {
0194       std::cout << "JetCalib::process_event(PHCompositeNode *topNode) : GlobalVertexMap node not found! Cannot apply Z-vertex dependent calibration." << std::endl;
0195       return Fun4AllReturnCodes::ABORTRUN;
0196     }
0197   }
0198   // Get z-vertex if needed.
0199   float zvertex = -999.0;
0200   if (ApplyZvrtxDependentCalib)
0201   {
0202     if (!vertexmap->empty())
0203     {
0204       GlobalVertex *vtx = vertexmap->begin()->second;
0205       auto mbdStartIter = vtx->find_vertexes(GlobalVertex::MBD);
0206       auto mbdEndIter = vtx->end_vertexes();
0207       for (auto iter = mbdStartIter; iter != mbdEndIter; ++iter)
0208       {
0209         const auto &[type, vertexVec] = *iter;
0210         if (type != GlobalVertex::MBD)
0211         {
0212           continue;
0213         }
0214         for (const auto *vertex : vertexVec)
0215         {
0216           if (!vertex)
0217           {
0218             continue;
0219           }
0220           zvertex = vertex->get_z();
0221         }
0222       }
0223     }
0224     else
0225     {
0226       if (Verbosity() > 0)
0227       {
0228         std::cout << "JetCalib::process_event(PHCompositeNode *topNode) : GlobalVertexMap is empty." << std::endl;
0229       }
0230     }
0231   }
0232 
0233   // Apply calibration to jets and add to calib jet node.
0234   int ijet = 0;
0235   for (auto *jet : *_raw_jets)
0236   {
0237     float pt = jet->get_pt();
0238     float eta = jet->get_eta();
0239     float phi = jet->get_phi();
0240 
0241     float calib_pt = doCalibration(m_JetCalibFunc, pt, zvertex, eta);
0242     auto *calib_jet = _calib_jets->add_jet();
0243     calib_jet->set_px(calib_pt * std::cos(phi));
0244     calib_jet->set_py(calib_pt * std::sin(phi));
0245     calib_jet->set_pz(calib_pt * std::sinh(eta));
0246     calib_jet->set_id(ijet);
0247     calib_jet->set_isCalib(1);
0248     ijet++;
0249   }
0250 
0251   if (Verbosity() > 0)
0252   {
0253     std::cout << "JetCalib::process_event(PHCompositeNode *topNode) : nRawJets: " << _raw_jets->size() << " nCalibJets: " << _calib_jets->size() << std::endl;
0254     if (_calib_jets->size() != _raw_jets->size())
0255     {
0256       std::cout << "JetCalib::process_event(PHCompositeNode *topNode) : different number of raw jets vs. calib jets! Something is amiss! " << std::endl;
0257     }
0258   }
0259 
0260   return Fun4AllReturnCodes::EVENT_OK;
0261 }
0262 
0263 int JetCalib::CreateNodeTree(PHCompositeNode *topNode)
0264 {
0265   // Check nodes.
0266   PHNodeIterator iter(topNode);
0267   PHCompositeNode *dstNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "DST"));
0268   if (!dstNode)
0269   {
0270     std::cout << PHWHERE << "JetCalib::CreateNodeTree : DST Node missing, aborting!" << std::endl;
0271     return Fun4AllReturnCodes::ABORTRUN;
0272   }
0273   PHCompositeNode *antiktNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "ANTIKT"));
0274   if (!antiktNode)
0275   {
0276     std::cout << PHWHERE << "JetCalib::CreateNodeTree : ANTIKT node missing, aborting!" << std::endl;
0277     return Fun4AllReturnCodes::ABORTRUN;
0278   }
0279   PHCompositeNode *towerNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "TOWER"));
0280   if (!towerNode)
0281   {
0282     std::cout << PHWHERE << "TOWER node not found, aborting!" << std::endl;
0283     return Fun4AllReturnCodes::ABORTRUN;
0284   }
0285 
0286   // Create calib jet node if it doesn't exist.
0287   JetContainer *test_jets = findNode::getClass<JetContainer>(topNode, m_outputNode);
0288   if (!test_jets)
0289   {
0290     JetContainer *calib_jets = new JetContainerv1();
0291     PHIODataNode<PHObject> *calibjetNode;
0292     if (Verbosity() > 0)
0293     {
0294       std::cout << "JetCalib::CreateNode : creating " << m_outputNode << std::endl;
0295     }
0296     calibjetNode = new PHIODataNode<PHObject>(calib_jets, m_outputNode, "PHObject");
0297     towerNode->addNode(calibjetNode);
0298   }
0299   else
0300   {
0301     std::cout << "JetCalib::CreateNode : " << m_outputNode << " already exists! " << std::endl;
0302   }
0303 
0304   return Fun4AllReturnCodes::EVENT_OK;
0305 }
0306 
0307 std::string JetCalib::fetchCalibDir(const char *calibType)
0308 {
0309   std::string calibName = std::string("JES_Calib_") + calibType;
0310   return CDBInterface::instance()->getUrl(calibName);
0311 }
0312 
0313 int getZvrtxBin(float zvrtx)
0314 {
0315   if (zvrtx >= -60.0 && zvrtx < -30.0)
0316   {
0317     return 1;  // -60 to -30
0318   }
0319   if (zvrtx >= -30.0 && zvrtx < 30.0)
0320   {
0321     return 0;  // -30 to 30
0322   }
0323   if (zvrtx >= 30.0 && zvrtx < 60)
0324   {
0325     return 2;  // 30 to 60
0326   }
0327   if ((zvrtx < -60.0 && zvrtx >= -900) || (zvrtx >= 60.0 && zvrtx < 900))
0328   {
0329     return 3;  // -inf to -60 or 60 to inf
0330   }
0331   if (zvrtx < -900)
0332   {
0333     return 4;  // -999 no zvertex
0334   }
0335 
0336   return 0;  // Default case, should not happen
0337 }
0338 
0339 int getEtaBin(int zvrtxbin, float eta, float jet_radius)
0340 {
0341   float eta_low = 0;
0342   float eta_high = 0;
0343   if (zvrtxbin == 0)  // -30 < zvrtx < 30
0344   {
0345     eta_low = -1.2 + jet_radius;
0346     eta_high = 1.2 - jet_radius;
0347   }
0348   else if (zvrtxbin == 1)  // -60 < zvrtx < -30
0349   {
0350     eta_low = -0.95 + jet_radius;
0351     eta_high = 1.25 - jet_radius;
0352   }
0353   else if (zvrtxbin == 2)  // 30 < zvrtx < inf
0354   {
0355     eta_low = -1.25 + jet_radius;
0356     eta_high = 0.95 - jet_radius;
0357   }
0358   else if (zvrtxbin == 3 || zvrtxbin == 4)
0359   {
0360     return 0;
0361   }
0362 
0363   float threshold1 = eta_low + ((eta_high - eta_low) / 4.0);
0364   float threshold2 = eta_low + ((eta_high - eta_low) / 2.0);
0365   float threshold3 = eta_low + (3.0 * (eta_high - eta_low) / 4.0);
0366 
0367   if (eta < threshold1)
0368   {
0369     return 0;  // -inf to threshold1
0370   }
0371   if (eta < threshold2)
0372   {
0373     return 1;  // threshold1 to threshold2
0374   }
0375   if (eta < threshold3)
0376   {
0377     return 2;  // threshold2 to threshold3
0378   }
0379 
0380   return 3;  // threshold3 to inf
0381 }
0382 
0383 float JetCalib::doCalibration(const std::vector<std::vector<TF1 *>> &JetCalibFunc, float jetPt, float zvrtx, float eta) const
0384 {
0385   float calib = jetPt;
0386   int zvrtxbin = 0;
0387   int etabin = 0;
0388   if (ApplyZvrtxDependentCalib || ApplyEtaDependentCalib)
0389   {
0390     zvrtxbin = getZvrtxBin(zvrtx);
0391     if (ApplyEtaDependentCalib)
0392     {
0393       if (zvrtxbin < 3)
0394       {
0395         etabin = getEtaBin(zvrtxbin, eta, jet_radius);
0396       }
0397       else
0398       {
0399         etabin = 0;
0400       }
0401     }
0402   }
0403   calib = JetCalibFunc[zvrtxbin][etabin]->Eval(jetPt);
0404   return calib;
0405 }