UtilAlgos/interface/VariableNTupler.h

0001 #ifndef VariableNtupler_NTupler_H
0002 #define VariableNtupler_NTupler_H
0003
0004 #include "PhysicsTools/UtilAlgos/interface/VariableHelper.h"
0005
0006 #include "FWCore/Utilities/interface/InputTag.h"
0007 #include "FWCore/Framework/interface/Frameworkfwd.h"
0008 #include "FWCore/Framework/interface/ProducesCollector.h"
0009
0010 #include "CommonTools/UtilAlgos/interface/TFileService.h"
0011 #include "TTree.h"
0012 #include "TBranch.h"
0013 #include "TFile.h"
0014
0015 #include "PhysicsTools/UtilAlgos/interface/NTupler.h"
0016
0017 #include <algorithm>
0018
0019 class VariableNTupler : public NTupler {
0020 public:
0021   VariableNTupler(const edm::ParameterSet& iConfig) {
0022     ownTheTree_ = false;
0023     edm::ParameterSet variablePSet = iConfig.getParameter<edm::ParameterSet>("variablesPSet");
0024     if (variablePSet.getParameter<bool>("allVariables")) {
0025       VariableHelper::iterator v = edm::Service<VariableHelperService>()->get().begin();
0026       VariableHelper::iterator v_end = edm::Service<VariableHelperService>()->get().end();
0027       for (; v != v_end; ++v) {
0028         leaves_[v->second->name()] = v->second;
0029       }
0030     } else {
0031       std::vector<std::string> leaves = variablePSet.getParameter<std::vector<std::string> >("leaves");
0032       for (uint i = 0; i != leaves.size(); ++i) {
0033         leaves_[leaves[i]] = edm::Service<VariableHelperService>()->get().variable(leaves[i]);
0034       }
0035     }
0036     if (variablePSet.exists("useTFileService"))
0037       useTFileService_ = variablePSet.getParameter<bool>("useTFileService");
0038     else
0039       useTFileService_ = iConfig.getParameter<bool>("useTFileService");
0040
0041     if (useTFileService_) {
0042       if (variablePSet.exists("treeName"))
0043         treeName_ = variablePSet.getParameter<std::string>("treeName");
0044       else
0045         treeName_ = iConfig.getParameter<std::string>("treeName");
0046     }
0047   }
0048
0049   uint registerleaves(edm::ProducesCollector producesCollector) override {
0050     uint nLeaves = 0;
0051     if (useTFileService_) {
0052       //loop the leaves registered
0053       nLeaves = leaves_.size();
0054       // make arrays of pointer to the actual values
0055       dataHolder_ = new double[nLeaves];
0056       iterator i = leaves_.begin();
0057       iterator i_end = leaves_.end();
0058       edm::Service<TFileService> fs;
0059       if (ownTheTree_) {
0060         ownTheTree_ = true;
0061         tree_ = fs->make<TTree>(treeName_.c_str(), "VariableNTupler tree");
0062       } else {
0063         TObject* object = fs->file().Get(treeName_.c_str());
0064         if (!object) {
0065           ownTheTree_ = true;
0066           tree_ = fs->make<TTree>(treeName_.c_str(), "VariableNTupler tree");
0067         } else {
0068           tree_ = dynamic_cast<TTree*>(object);
0069           if (!tree_) {
0070             ownTheTree_ = true;
0071             tree_ = fs->make<TTree>(treeName_.c_str(), "VariableNTupler tree");
0072           } else
0073             ownTheTree_ = false;
0074         }
0075       }
0076       uint iInDataHolder = 0;
0077       for (; i != i_end; ++i, ++iInDataHolder) {
0078         tree_->Branch(i->first.c_str(), &(dataHolder_[iInDataHolder]), (i->first + "/D").c_str());
0079       }
0080     } else {
0081       //loop the leaves registered
0082       iterator i = leaves_.begin();
0083       iterator i_end = leaves_.end();
0084       for (; i != i_end; ++i) {
0085         nLeaves++;
0086         std::string lName(i->first);
0087         std::replace(lName.begin(), lName.end(), '_', '0');
0088         producesCollector.produces<double>(lName).setBranchAlias(i->first);
0089       }
0090     }
0091     return nLeaves;
0092   }
0093
0094   void fill(edm::Event& iEvent) override {
0095     if (useTFileService_) {
0096       //fill the data holder
0097       iterator i = leaves_.begin();
0098       iterator i_end = leaves_.end();
0099       uint iInDataHolder = 0;
0100       for (; i != i_end; ++i, ++iInDataHolder) {
0101         dataHolder_[iInDataHolder] = (*i->second)(iEvent);
0102       }
0103       //fill into root;
0104       if (ownTheTree_) {
0105         tree_->Fill();
0106       }
0107     } else {
0108       //other leaves
0109       iterator i = leaves_.begin();
0110       iterator i_end = leaves_.end();
0111       for (; i != i_end; ++i) {
0112         auto leafValue = std::make_unique<double>((*i->second)(iEvent));
0113         std::string lName(i->first);
0114         std::replace(lName.begin(), lName.end(), '_', '0');
0115         iEvent.put(std::move(leafValue), lName);
0116       }
0117     }
0118   }
0119   void callBack() {}
0120
0121 protected:
0122   typedef std::map<std::string, const CachingVariable*>::iterator iterator;
0123   std::map<std::string, const CachingVariable*> leaves_;
0124
0125   bool ownTheTree_;
0126   std::string treeName_;
0127   double* dataHolder_;
0128 };
0129
0130 #endif