LightFlavorRatios/bco_correction/V0DuplicateReader_mod.h

0001 #ifndef V0_DUPLICATE_READER_H
0002 #define V0_DUPLICATE_READER_H
0003
0004 #include <TTree.h>
0005
0006 #include <cmath>
0007 #include <set>
0008 #include <stdexcept>
0009 #include <string>
0010 #include <vector>
0011
0012 class V0DuplicateReader
0013 {
0014  public:
0015   enum class ParticleType
0016   {
0017     K0s,
0018     Lambda
0019   };
0020
0021   V0DuplicateReader(TTree* tree, ParticleType particleType)
0022   : m_tree(tree)
0023   , m_particleType(particleType)
0024   {
0025     if (!m_tree)
0026     {
0027       throw std::runtime_error("V0DuplicateReader: input tree is null");
0028     }
0029
0030     prefix = (m_particleType == ParticleType::K0s) ? "K_S0" : "Lambda0";
0031
0032     setupBranches();
0033
0034     //setupCommonBranches();
0035     //setupParticleBranches();
0036   }
0037
0038   void enableDeltaBCOCut(long long minDeltaBCO, long long maxDeltaBCO)
0039   {
0040     m_useDeltaBCOCut = true;
0041     m_minDeltaBCO = minDeltaBCO;
0042     m_maxDeltaBCO = maxDeltaBCO;
0043   }
0044
0045   void disableDeltaBCOCut()
0046   {
0047     m_useDeltaBCOCut = false;
0048   }
0049
0050   bool loadEntry(Long64_t entry)
0051   {
0052     if (!m_tree) return false;
0053
0054     m_tree->GetEntry(entry);
0055
0056     Long64_t m_bco = get<Long64_t>("Collision_BCO");
0057     Long64_t m_eventBCO = get<Long64_t>("GL1_BCO");
0058     int m_runNumber = get<int>("runNumber");
0059
0060     m_deltaBCO = (m_bco >= 0 && m_eventBCO >= 0) ? (m_bco - m_eventBCO) : -1;
0061
0062     if (m_useDeltaBCOCut)
0063     {
0064       m_passesDeltaBCO = (m_deltaBCO >= m_minDeltaBCO && m_deltaBCO < m_maxDeltaBCO);
0065     }
0066     else
0067     {
0068       m_passesDeltaBCO = true;
0069     }
0070
0071     if (!m_passesDeltaBCO)
0072     {
0073       m_currentEntryIsUnique = false;
0074       return true;
0075     }
0076
0077     std::vector<std::string> candidateBranches = {
0078       "mass",
0079       "pT",
0080       "pseudorapidity",
0081       "phi",
0082       "decayLength",
0083       "DIRA",
0084       "IP",
0085       "x",
0086       "y",
0087       "z",
0088       "chi2"
0089     };
0090
0091     std::vector<float> candidateVariables;
0092
0093     for(std::string name : candidateBranches)
0094     {
0095       candidateVariables.push_back(get<float>(prefix+"_"+name));
0096     }
0097
0098     const std::string eventKey =
0099       std::to_string(m_runNumber) + "_" + std::to_string(m_bco);
0100
0101     const std::string candidateKey = buildCandidateKey(candidateVariables);
0102     const std::string fullKey = eventKey + "__" + candidateKey;
0103
0104     const auto [it, inserted] = m_seenCandidateKeys.insert(fullKey);
0105
0106     if (inserted)
0107     {
0108       ++m_numberOfUniqueCandidates;
0109       m_currentEntryIsUnique = true;
0110     }
0111     else
0112     {
0113       ++m_numberOfDuplicateCandidates;
0114       m_currentEntryIsUnique = false;
0115     }
0116
0117     return true;
0118   }
0119
0120   bool isCurrentEntryUnique() const
0121   {
0122     return m_currentEntryIsUnique;
0123   }
0124
0125   bool passesDeltaBCOCut() const
0126   {
0127     return m_passesDeltaBCO;
0128   }
0129
0130   long long deltaBCO() const
0131   {
0132     return m_deltaBCO;
0133   }
0134
0135   void reset()
0136   {
0137     m_seenCandidateKeys.clear();
0138     m_numberOfUniqueCandidates = 0;
0139     m_numberOfDuplicateCandidates = 0;
0140     m_currentEntryIsUnique = false;
0141     m_deltaBCO = -1;
0142     m_passesDeltaBCO = true;
0143   }
0144
0145   Long64_t entries() const
0146   {
0147     return m_tree ? m_tree->GetEntries() : 0;
0148   }
0149 /*
0150   int runNumber() const { return m_runNumber; }
0151   Long64_t bco() const { return m_bco; }
0152   Long64_t eventBCO() const { return m_eventBCO; }
0153   int eventNumber() const { return m_eventNumber; }
0154
0155   float mass() const { return m_mass; }
0156   float pt() const { return m_pt; }
0157   float eta() const { return m_eta; }
0158   float phi() const { return m_phi; }
0159   float decayLength() const { return m_decayLength; }
0160   float dira() const { return m_dira; }
0161   float ip() const { return m_ip; }
0162   float x() const { return m_x; }
0163   float y() const { return m_y; }
0164   float z() const { return m_z; }
0165   float chi2() const { return m_chi2; }
0166   float rapidity() const { return m_rapidity; }
0167 */
0168   long long numberOfUniqueCandidates() const
0169   {
0170     return m_numberOfUniqueCandidates;
0171   }
0172
0173   long long numberOfDuplicateCandidates() const
0174   {
0175     return m_numberOfDuplicateCandidates;
0176   }
0177
0178   template<typename T>
0179   T get(const std::string& name)
0180   {
0181     T val = *(T*)leafmap[name.c_str()];
0182     return val;
0183   }
0184
0185  private:
0186   void setupBranches()
0187   {
0188     TObjArray* branches = m_tree->GetListOfBranches();
0189     size_t nbranches = m_tree->GetNbranches();
0190     for(size_t i=0;i<nbranches;i++)
0191     {
0192       std::string branchname = branches->At(i)->GetName();
0193       leafmap[branchname] = m_tree->GetLeaf(branchname.c_str())->GetValuePointer();
0194     }
0195   }
0196
0197 /*
0198   void setupCommonBranches()
0199   {
0200     m_tree->SetBranchAddress("runNumber", &m_runNumber);
0201     m_tree->SetBranchAddress("eventNumber", &m_eventNumber);
0202     m_tree->SetBranchAddress("BCO", &m_bco);
0203     m_tree->SetBranchAddress("event_bco", &m_eventBCO);
0204   }
0205
0206   void setupParticleBranches()
0207   {
0208     const std::string prefix = (m_particleType == ParticleType::K0s) ? "K_S0" : "Lambda0";
0209
0210     m_tree->SetBranchAddress((prefix + "_mass").c_str(), &m_mass);
0211     m_tree->SetBranchAddress((prefix + "_pT").c_str(), &m_pt);
0212     m_tree->SetBranchAddress((prefix + "_pseudorapidity").c_str(), &m_eta);
0213     m_tree->SetBranchAddress((prefix + "_phi").c_str(), &m_phi);
0214     m_tree->SetBranchAddress((prefix + "_decayLength").c_str(), &m_decayLength);
0215     m_tree->SetBranchAddress((prefix + "_DIRA").c_str(), &m_dira);
0216     m_tree->SetBranchAddress((prefix + "_IP").c_str(), &m_ip);
0217     m_tree->SetBranchAddress((prefix + "_x").c_str(), &m_x);
0218     m_tree->SetBranchAddress((prefix + "_y").c_str(), &m_y);
0219     m_tree->SetBranchAddress((prefix + "_z").c_str(), &m_z);
0220     m_tree->SetBranchAddress((prefix + "_chi2").c_str(), &m_chi2);
0221     m_tree->SetBranchAddress((prefix + "_rapidity").c_str(), &m_rapidity);
0222   }
0223 */
0224   static std::string buildCandidateKey(const std::vector<float>& candidateVariables)
0225   {
0226     std::string key;
0227     key.reserve(candidateVariables.size() * 16);
0228
0229     for (float value : candidateVariables)
0230     {
0231       const long long roundedValue = llround(value * 1e6);
0232       key += std::to_string(roundedValue);
0233       key += "_";
0234     }
0235
0236     return key;
0237   }
0238
0239  private:
0240   TTree* m_tree = nullptr;
0241   ParticleType m_particleType;
0242
0243   std::string prefix;
0244
0245   std::map<std::string,void*> leafmap;
0246
0247   std::set<std::string> m_seenCandidateKeys;
0248
0249   long long m_numberOfUniqueCandidates = 0;
0250   long long m_numberOfDuplicateCandidates = 0;
0251   bool m_currentEntryIsUnique = false;
0252
0253   bool m_useDeltaBCOCut = false;
0254   long long m_minDeltaBCO = 0;
0255   long long m_maxDeltaBCO = 350;
0256   long long m_deltaBCO = -1;
0257   bool m_passesDeltaBCO = true;
0258 /*
0259   int m_runNumber = -1;
0260   int m_eventNumber = -1;
0261   Long64_t m_bco = -1;
0262   Long64_t m_eventBCO = -1;
0263
0264   float m_mass = 0.0;
0265   float m_pt = 0.0;
0266   float m_eta = 0.0;
0267   float m_phi = 0.0;
0268   float m_decayLength = 0.0;
0269   float m_dira = 0.0;
0270   float m_ip = 0.0;
0271   float m_x = 0.0;
0272   float m_y = 0.0;
0273   float m_z = 0.0;
0274   float m_chi2 = 0.0;
0275   float m_rapidity = 0.0;
0276 */
0277 };
0278
0279 #endif