JetMultSub/src/JetMultSub.cc

0001 /*
0002 JetMultSub SubsysReco module.
0003
0004   Description:
0005     This is a module that takes in reconstructed jets and subtracts the background using the multiplicity method.
0006     Estimates rho using the median of the average pt in a kT jet with radius 0.4.
0007     Corrects the multiplicity using the average truth multiplicity of a a jet with a given pt_reco.
0008     Subtracts the background from the jet pt using the formula:
0009       pt_sub = pt_reco - rho*(N_reco - <N_signal>(pt_reco))
0010
0011   Author: Tanner Mengel
0012
0013   Date: 10/04/2023
0014
0015 */
0016
0017 //____________________________________________________________________________..
0018 // JetMultSub.h
0019 #include "JetMultSub.h"
0020 //____________________________________________________________________________..
0021 // Fun4All includes
0022 #include <fun4all/Fun4AllReturnCodes.h>
0023 #include <fun4all/SubsysReco.h>
0024 //____________________________________________________________________________..
0025 // PHOOL includes
0026 #include <phool/PHCompositeNode.h>
0027 #include <phool/getClass.h>
0028 #include <phool/PHCompositeNode.h>
0029 #include <phool/PHIODataNode.h>
0030 #include <phool/PHNode.h>
0031 #include <phool/PHNodeIterator.h>
0032 #include <phool/PHObject.h>
0033 #include <phool/phool.h>
0034 //____________________________________________________________________________..
0035 // Jet includes
0036 #include <jetbase/Jet.h>
0037 #include <jetbase/JetMap.h>
0038 #include <jetbase/JetMapv1.h>
0039 #include <jetbase/Jetv1.h>
0040 //____________________________________________________________________________..
0041 // standard includes
0042 #include <cmath>
0043 #include <iostream>
0044 #include <map>
0045 #include <vector>
0046 #include <utility>
0047 //____________________________________________________________________________..
0048 // Starting function definitions...
0049
0050
0051 //____________________________________________________________________________..
0052 // constructor. Sets default values
0053 JetMultSub::JetMultSub(const std::string &name):
0054   SubsysReco(name)
0055   , m_reco_input("AntiKt_Tower_r04")
0056   , m_kt_input("Kt_Tower_r04")
0057   , m_subtracted_output("AntiKt_Tower_r04_sub")
0058   , m_etaRange(-1.0, 1.0)
0059   , m_ptRange(5.0, 100.0)
0060 {
0061   std::cout << "JetMultSub::JetMultSub(const std::string &name) Calling ctor" << std::endl;
0062 }
0063
0064 //____________________________________________________________________________..
0065 // destructor. Does nothing
0066 JetMultSub::~JetMultSub()
0067 {
0068   std::cout << "JetMultSub::~JetMultSub() Calling dtor" << std::endl;
0069 }
0070
0071 //____________________________________________________________________________..
0072 // init. Does nothing
0073 int JetMultSub::Init(PHCompositeNode *topNode)
0074 {
0075   std::cout << "JetMultSub::Init(PHCompositeNode *topNode) Initializing" << std::endl;
0076   return Fun4AllReturnCodes::EVENT_OK;
0077 }
0078
0079 //____________________________________________________________________________..
0080 // init run. Very important! Creates output node if it doesn't exist
0081 int JetMultSub::InitRun(PHCompositeNode *topNode)
0082 {
0083   CreateNode(topNode);
0084   return Fun4AllReturnCodes::EVENT_OK;
0085 }
0086
0087 //____________________________________________________________________________..
0088 // process event. This is where the magic happens
0089 int JetMultSub::process_event(PHCompositeNode *topNode)
0090 {
0091   if (Verbosity() > 0){
0092     std::cout << "JetMultSub::process_event(PHCompositeNode *topNode) Processing Event" << std::endl;
0093   }
0094
0095   // get unsubtracted and subtracted jet collections
0096   JetMap *reco_jets = findNode::getClass<JetMap>(topNode, m_reco_input);
0097   JetMap *sub_jets = findNode::getClass<JetMap>(topNode, m_subtracted_output);
0098   if(!reco_jets){
0099     std::cout << "JetMultSub::process_event(PHCompositeNode *topNode) Missing input reco jet collection" << std::endl;
0100     return Fun4AllReturnCodes::ABORTRUN;
0101   }
0102
0103   // loop over reco jets and subtract background
0104   int ijet = 0;
0105   for (JetMap::Iter iter = reco_jets->begin(); iter != reco_jets->end(); ++iter)
0106   {
0107     Jet *this_jet = iter->second;
0108
0109     float this_pt = this_jet->get_pt();
0110     float this_phi = this_jet->get_phi();
0111     float this_eta = this_jet->get_eta();
0112
0113     Jet *new_jet = new Jetv1();
0114
0115     // subtract background
0116     int nConstituents =this_jet->size_comp();
0117     float rho = EstimateRho(topNode);
0118     float multiplicity_correction = GetMultiplicityCorrection(this_pt);
0119     float nConstituents_minus_multiplicity_correction = nConstituents - multiplicity_correction;
0120     float new_pt = this_pt - rho*nConstituents_minus_multiplicity_correction;
0121
0122     // calculate new px, py, pz, e
0123     float new_et = new_pt*cosh(this_eta);
0124     float new_e = sqrt(pow(new_et,2) + pow(this_jet->get_pz(),2));
0125     float new_px = new_pt*cos(this_phi);
0126     float new_py = new_pt*sin(this_phi);
0127     float new_pz = this_jet->get_pz();
0128
0129
0130     // set new jet properties
0131     new_jet->set_px(new_px);
0132     new_jet->set_py(new_py);
0133     new_jet->set_pz(new_pz);
0134     new_jet->set_e(new_e);
0135
0136     // add new jet to subtracted jet collection
0137     sub_jets->insert(new_jet);
0138
0139     // debug info
0140     if (Verbosity() > 1 && this_pt > 5)
0141       {
0142         std::cout << "JetMultSub::process_event: old jet #" << ijet << ", old px / py / pz / e = " << this_jet->get_px() << " / " << this_jet->get_py() << " / " << this_jet->get_pz() << " / " << this_jet->get_e() << std::endl;
0143         std::cout << "JetMultSub::process_event: new jet #" << ijet << ", new px / py / pz / e = " << new_jet->get_px() << " / " << new_jet->get_py() << " / " << new_jet->get_pz() << " / " << new_jet->get_e() << std::endl;
0144       }
0145
0146     ijet++;
0147
0148   }
0149
0150   return Fun4AllReturnCodes::EVENT_OK;
0151 }
0152
0153 //____________________________________________________________________________..
0154 // reset event. Does nothing
0155 int JetMultSub::ResetEvent(PHCompositeNode *topNode)
0156 {
0157   return Fun4AllReturnCodes::EVENT_OK;
0158 }
0159
0160 //____________________________________________________________________________..
0161 // end run. You guessed it! does nothing
0162 int JetMultSub::EndRun(const int runnumber)
0163 {
0164   return Fun4AllReturnCodes::EVENT_OK;
0165 }
0166
0167 //____________________________________________________________________________..'
0168 // this is the end. Does nothing
0169 int JetMultSub::End(PHCompositeNode *topNode)
0170 {
0171   std::cout << "JetMultSub::End(PHCompositeNode *topNode) This is the End..." << std::endl;
0172   return Fun4AllReturnCodes::EVENT_OK;
0173 }
0174
0175 //____________________________________________________________________________..
0176 // reset. Does nothing
0177 int JetMultSub::Reset(PHCompositeNode *topNode)
0178 {
0179   if (Verbosity() > 0)
0180   {
0181     std::cout << "JetMultSub::Reset(PHCompositeNode *topNode) being Reset" << std::endl;
0182   }
0183   return Fun4AllReturnCodes::EVENT_OK;
0184 }
0185
0186 //____________________________________________________________________________..
0187 // print. Prints meaningless message to screen
0188 void JetMultSub::Print(const std::string &what) const
0189 {
0190   if (Verbosity() > 0){
0191     std::cout << "JetMultSub::Print(const std::string &what) const Printing info for " << what << std::endl;
0192     std::cout << "Current status: 50% working" << std::endl;
0193     std::cout << "Errors codes: 49512323, 304, b.ss3, sos" << std::endl;
0194   }
0195 }
0196
0197 //____________________________________________________________________________..
0198 // create node. This is actually important. Creates the node if it doesn't exist
0199 int JetMultSub::CreateNode(PHCompositeNode *topNode)
0200 {
0201   // iterate to the DST node and add the jet collection if necessary
0202   PHNodeIterator iter(topNode);
0203
0204   // Looking for the DST node
0205   PHCompositeNode *dstNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "DST"));
0206   if (!dstNode)
0207   {
0208     std::cout << PHWHERE << "DST Node missing, doing nothing." << std::endl;
0209     return Fun4AllReturnCodes::ABORTRUN;
0210   }
0211
0212   // Looking for the ANTIKT node
0213   PHCompositeNode *antiktNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "ANTIKT"));
0214   if (!antiktNode)
0215   {
0216     std::cout << PHWHERE << "ANTIKT node not found, doing nothing." << std::endl;
0217   }
0218
0219   // store the new jet collection if it doesn't exist already
0220   JetMap *test_jets;
0221   test_jets = findNode::getClass<JetMap>(topNode, m_subtracted_output);
0222   if(!test_jets)
0223   {
0224     JetMap *sub_jets = new JetMapv1();
0225     PHIODataNode<PHObject> *subjetNode ;
0226     if (Verbosity() > 0)
0227     {
0228       std::cout << "JetMultSub::CreateNode : creating " << m_subtracted_output << std::endl;
0229     }
0230     subjetNode = new PHIODataNode<PHObject>(sub_jets, m_subtracted_output, "PHObject");
0231     antiktNode->addNode(subjetNode);
0232   }
0233   else
0234   {
0235     std::cout << "JetMultSub::CreateNode : " << m_subtracted_output << " already exists" << std::endl;
0236   }
0237
0238   // all done!
0239
0240   return Fun4AllReturnCodes::EVENT_OK;
0241
0242 }
0243
0244 //____________________________________________________________________________..
0245 // estimate rho. Gets the event rho value from r=0.4 kT jets
0246 float JetMultSub::EstimateRho(PHCompositeNode *topNode)
0247 {
0248
0249   if (Verbosity() > 0){
0250     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Estimating rho" << std::endl;
0251   }
0252
0253   // get the kT jet collection
0254   JetMap *kt_jets = findNode::getClass<JetMap>(topNode, m_kt_input);
0255   if(!kt_jets){
0256     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Missing input kT jet collection" << std::endl;
0257     return -1.0;
0258   }
0259
0260   // debug info
0261   if (Verbosity() > 0)
0262   {
0263     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Found input kT jet collection" << std::endl;
0264     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Number of kT jets: " << kt_jets->size() << std::endl;
0265   }
0266
0267   // loop over kT jets and find the 2 hardest jets in the eta range
0268   int hardest_kT_jet_indices[2] = {-1,-1};
0269   float hardest_kT_jet_pts[2] = {-1.0,-1.0};
0270   for (JetMap::Iter iter = kt_jets->begin(); iter != kt_jets->end(); ++iter)
0271   {
0272     Jet *jet = iter->second;
0273     float pt = jet->get_pt();
0274     float eta = jet->get_eta();
0275     // float et = jet->get_et();
0276
0277     if (eta > m_etaRange.first && eta < m_etaRange.second)
0278     {
0279       if (pt > hardest_kT_jet_pts[0])
0280       {
0281         hardest_kT_jet_pts[1] = hardest_kT_jet_pts[0];
0282         hardest_kT_jet_indices[1] = hardest_kT_jet_indices[0];
0283         hardest_kT_jet_pts[0] = pt;
0284         hardest_kT_jet_indices[0] = jet->get_id();
0285       }
0286       else if (pt > hardest_kT_jet_pts[1])
0287       {
0288         hardest_kT_jet_pts[1] = pt;
0289         hardest_kT_jet_indices[1] = jet->get_id();
0290       }
0291     }
0292   }
0293
0294
0295   // debug info
0296   if(Verbosity() > 0)
0297   {
0298     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Hardest kT jet indices: " << hardest_kT_jet_indices[0] << ", " << hardest_kT_jet_indices[1] << std::endl;
0299     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Hardest kT jet et: " << hardest_kT_jet_pts[0] << ", " << hardest_kT_jet_pts[1] << std::endl;
0300   }
0301
0302   // get median value of the average et of constituents in all but the 2 hardest jets
0303   std::vector<float> et_over_nConstituents;
0304   for (JetMap::Iter iter = kt_jets->begin(); iter != kt_jets->end(); ++iter)
0305   {
0306     Jet *jet = iter->second;
0307     float pt = jet->get_pt();
0308     float eta = jet->get_eta();
0309     // float et = jet->get_et();
0310     int nConstituents = jet->size_comp();
0311
0312     if (eta > m_etaRange.first && eta < m_etaRange.second)
0313     {
0314       int id = jet->get_id();
0315       if (id != hardest_kT_jet_indices[0] && id != hardest_kT_jet_indices[1])
0316       {
0317         et_over_nConstituents.push_back(pt/nConstituents);
0318       }
0319     }
0320   }
0321
0322   // get median value
0323   float median_et_over_nConstituents = 0.0;
0324   if (et_over_nConstituents.size() > 0)
0325   {
0326     std::sort(et_over_nConstituents.begin(), et_over_nConstituents.end());
0327     median_et_over_nConstituents = et_over_nConstituents[et_over_nConstituents.size()/2];
0328   }
0329
0330   // debug info
0331   if (Verbosity() > 0)
0332   {
0333     std::cout << "JetMultSub::EstimateRho(PHCompositeNode *topNode) Median et over nConstituents: " << median_et_over_nConstituents << std::endl;
0334   }
0335
0336   // return median et over nConstituents
0337   return median_et_over_nConstituents;
0338 }
0339
0340 //____________________________________________________________________________..
0341 // get multiplicity correction. Gets the multiplicity correction for a given pt
0342 float JetMultSub::GetMultiplicityCorrection(float pt)
0343 {
0344   if (Verbosity() > 0)
0345   {
0346     std::cout << "JetMultSub::GetMultiplicityCorrection(float pt) Getting multiplicity correction" << std::endl;
0347   }
0348
0349
0350   // to do! Match jets and get the values for this function!
0351   std::vector<float> pt_reco_bins = {0.0,100.0};
0352   std::vector<float> multiplicity_corrections = {0.0};
0353   // get pt  bin
0354   int pt_bin = -1;
0355   int pt_size = pt_reco_bins.size();
0356   for (int i = 0; i < pt_size; i++){
0357     if (pt > pt_reco_bins[i]){
0358       pt_bin = i;
0359     }
0360   }
0361   float correction;
0362   if(pt_bin == -1){
0363     correction = 0.0;
0364   }
0365   else{
0366     correction = multiplicity_corrections[pt_bin];
0367   }
0368   return correction;
0369 }