File indexing completed on 2023-10-25 09:32:04
0001 #include <iostream>
0002 #include <vector>
0003 #include <memory>
0004 #include <utility>
0005 #include <cmath>
0006 #include <string>
0007 #include "TLorentzVector.h"
0008 #include "TTree.h"
0009 #include "TH1F.h"
0010 #include "TMath.h"
0011
0012
0013 #include "FWCore/Framework/interface/Frameworkfwd.h"
0014 #include "FWCore/Framework/interface/one/EDAnalyzer.h"
0015
0016 #include "FWCore/Framework/interface/Event.h"
0017 #include "FWCore/Framework/interface/MakerMacros.h"
0018 #include "FWCore/ServiceRegistry/interface/Service.h"
0019 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0020 #include "CommonTools/UtilAlgos/interface/TFileService.h"
0021 #include "DataFormats/VertexReco/interface/Vertex.h"
0022 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0023 #include "DataFormats/TrackReco/interface/Track.h"
0024
0025 #include "TrackingTools/TrajectoryState/interface/TrajectoryStateOnSurface.h"
0026 #include "TrackingTools/GeomPropagators/interface/AnalyticalImpactPointExtrapolator.h"
0027 #include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
0028 #include "TrackingTools/Records/interface/TransientTrackRecord.h"
0029 #include "RecoVertex/KalmanVertexFit/interface/KalmanVertexFitter.h"
0030 #include "RecoVertex/VertexPrimitives/interface/ConvertToFromReco.h"
0031
0032 class HIPTwoBodyDecayAnalyzer : public edm::one::EDAnalyzer<edm::one::SharedResources> {
0033 public:
0034 explicit HIPTwoBodyDecayAnalyzer(const edm::ParameterSet&);
0035 ~HIPTwoBodyDecayAnalyzer() override;
0036
0037 edm::EDGetTokenT<reco::TrackCollection> alcareco_trackCollToken_;
0038 edm::EDGetTokenT<reco::TrackCollection> refit1_trackCollToken_;
0039 edm::EDGetTokenT<reco::TrackCollection> ctf_trackCollToken_;
0040 edm::EDGetTokenT<reco::TrackCollection> final_trackCollToken_;
0041
0042 static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0043
0044 TTree* tree;
0045 std::vector<std::pair<std::string, float*>> floatBranches;
0046 std::vector<std::pair<std::string, int*>> intBranches;
0047 std::vector<std::pair<std::string, short*>> shortBranches;
0048
0049 private:
0050
0051 const edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> ttkbuilderToken_;
0052
0053 enum BranchType { BranchType_short_t, BranchType_int_t, BranchType_float_t, BranchType_unknown_t };
0054
0055 void beginJob() override;
0056 void analyze(const edm::Event&, const edm::EventSetup&) override;
0057 void endJob() override;
0058
0059 void bookAllBranches();
0060 bool bookBranch(std::string bname, BranchType btype);
0061 BranchType searchArray(std::string branchname, int& position);
0062 template <typename varType>
0063 void setVal(std::string bname, varType value) {
0064 int varposition = -1;
0065 BranchType varbranchtype = searchArray(bname, varposition);
0066 if (varposition == -1)
0067 std::cerr << "HIPTwoBodyDecayAnalyzer::setVal -> Could not find the branch called " << bname << "!" << std::endl;
0068 else if (varbranchtype == BranchType_short_t)
0069 *(shortBranches.at(varposition).second) = value;
0070 else if (varbranchtype == BranchType_int_t)
0071 *(intBranches.at(varposition).second) = value;
0072 else if (varbranchtype == BranchType_float_t)
0073 *(floatBranches.at(varposition).second) = value;
0074 else
0075 std::cerr << "HIPTwoBodyDecayAnalyzer::setVal -> Could not find the type " << varbranchtype
0076 << " for branch called " << bname << "!" << std::endl;
0077 }
0078 void cleanBranches();
0079 void initializeBranches();
0080 bool actuateBranches();
0081
0082 void analyzeTrackCollection(std::string strTrackType,
0083 const TransientTrackBuilder& theTTBuilder,
0084 edm::Handle<reco::TrackCollection>& hTrackColl,
0085 bool verbose = false);
0086 reco::Vertex fitDimuonVertex(const TransientTrackBuilder& theTTBuilder,
0087 edm::Handle<reco::TrackCollection>& hTrackColl,
0088 bool& fitOk);
0089 };
0090
0091 HIPTwoBodyDecayAnalyzer::HIPTwoBodyDecayAnalyzer(const edm::ParameterSet& iConfig)
0092 : ttkbuilderToken_(esConsumes(edm::ESInputTag("", "TransientTrackBuilder"))) {
0093 usesResource(TFileService::kSharedResource);
0094
0095 alcareco_trackCollToken_ = consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("alcarecotracks"));
0096 refit1_trackCollToken_ = consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("refit1tracks"));
0097 ctf_trackCollToken_ = consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("refit2tracks"));
0098 final_trackCollToken_ = consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("finaltracks"));
0099
0100 edm::Service<TFileService> fs;
0101
0102 tree = fs->make<TTree>("TestTree", "");
0103 bookAllBranches();
0104 }
0105
0106 HIPTwoBodyDecayAnalyzer::BranchType HIPTwoBodyDecayAnalyzer::searchArray(std::string branchname, int& position) {
0107 for (unsigned short el = 0; el < shortBranches.size(); el++) {
0108 if (branchname == shortBranches.at(el).first) {
0109 position = el;
0110 return BranchType_short_t;
0111 }
0112 }
0113 for (unsigned int el = 0; el < intBranches.size(); el++) {
0114 if (branchname == intBranches.at(el).first) {
0115 position = el;
0116 return BranchType_int_t;
0117 }
0118 }
0119 for (unsigned int el = 0; el < floatBranches.size(); el++) {
0120 if (branchname == floatBranches.at(el).first) {
0121 position = el;
0122 return BranchType_float_t;
0123 }
0124 }
0125 return BranchType_unknown_t;
0126 }
0127 void HIPTwoBodyDecayAnalyzer::cleanBranches() {
0128 for (unsigned short el = 0; el < shortBranches.size(); el++) {
0129 if (shortBranches.at(el).second != nullptr)
0130 delete shortBranches.at(el).second;
0131 shortBranches.at(el).second = nullptr;
0132 }
0133 shortBranches.clear();
0134 for (unsigned int el = 0; el < intBranches.size(); el++) {
0135 if (intBranches.at(el).second != nullptr)
0136 delete intBranches.at(el).second;
0137 intBranches.at(el).second = nullptr;
0138 }
0139 intBranches.clear();
0140 for (unsigned int el = 0; el < floatBranches.size(); el++) {
0141 if (floatBranches.at(el).second != nullptr)
0142 delete floatBranches.at(el).second;
0143 floatBranches.at(el).second = nullptr;
0144 }
0145 floatBranches.clear();
0146 }
0147 void HIPTwoBodyDecayAnalyzer::initializeBranches() {
0148 for (unsigned short el = 0; el < shortBranches.size(); el++) {
0149 if (shortBranches.at(el).second != nullptr)
0150 *(shortBranches.at(el).second) = 0;
0151 }
0152 for (unsigned int el = 0; el < intBranches.size(); el++) {
0153 if (intBranches.at(el).second != nullptr)
0154 *(intBranches.at(el).second) = 0;
0155 }
0156 for (unsigned int el = 0; el < floatBranches.size(); el++) {
0157 if (floatBranches.at(el).second != nullptr)
0158 *(floatBranches.at(el).second) = 0;
0159 }
0160 }
0161
0162 void HIPTwoBodyDecayAnalyzer::bookAllBranches() {
0163 const int nTrackTypes = 4;
0164 std::vector<std::string> strTrackTypes;
0165 strTrackTypes.reserve(nTrackTypes);
0166 strTrackTypes.push_back("alcareco");
0167 strTrackTypes.push_back("refit1");
0168 strTrackTypes.push_back("refit2");
0169 strTrackTypes.push_back("final");
0170 for (unsigned int it = 0; it < nTrackTypes; it++) {
0171 std::string& strTrackType = strTrackTypes[it];
0172 bookBranch(strTrackType + "_present", BranchType_short_t);
0173 bookBranch(strTrackType + "_ZVtxFitOk", BranchType_short_t);
0174 bookBranch(strTrackType + "_ZMass", BranchType_float_t);
0175 bookBranch(strTrackType + "_ZPt", BranchType_float_t);
0176 bookBranch(strTrackType + "_ZPz", BranchType_float_t);
0177 bookBranch(strTrackType + "_ZPhi", BranchType_float_t);
0178 bookBranch(strTrackType + "_ZVertex_x", BranchType_float_t);
0179 bookBranch(strTrackType + "_ZVertex_y", BranchType_float_t);
0180 bookBranch(strTrackType + "_ZVertex_z", BranchType_float_t);
0181 bookBranch(strTrackType + "_ZVertex_NormChi2", BranchType_float_t);
0182 bookBranch(strTrackType + "_MuPlusVertex_x", BranchType_float_t);
0183 bookBranch(strTrackType + "_MuPlusVertex_y", BranchType_float_t);
0184 bookBranch(strTrackType + "_MuPlusVertex_z", BranchType_float_t);
0185 bookBranch(strTrackType + "_MuMinusPt_AfterZVtxFit", BranchType_float_t);
0186 bookBranch(strTrackType + "_MuMinusPz_AfterZVtxFit", BranchType_float_t);
0187 bookBranch(strTrackType + "_MuMinusPhi_AfterZVtxFit", BranchType_float_t);
0188 bookBranch(strTrackType + "_MuPlusPt_AfterZVtxFit", BranchType_float_t);
0189 bookBranch(strTrackType + "_MuPlusPz_AfterZVtxFit", BranchType_float_t);
0190 bookBranch(strTrackType + "_MuPlusPhi_AfterZVtxFit", BranchType_float_t);
0191 bookBranch(strTrackType + "_ZMass_AfterZVtxFit", BranchType_float_t);
0192 bookBranch(strTrackType + "_ZPt_AfterZVtxFit", BranchType_float_t);
0193 bookBranch(strTrackType + "_ZPz_AfterZVtxFit", BranchType_float_t);
0194 bookBranch(strTrackType + "_ZPhi_AfterZVtxFit", BranchType_float_t);
0195 bookBranch(strTrackType + "_MuMinusPt", BranchType_float_t);
0196 bookBranch(strTrackType + "_MuMinusPz", BranchType_float_t);
0197 bookBranch(strTrackType + "_MuMinusPhi", BranchType_float_t);
0198 bookBranch(strTrackType + "_MuMinusVertex_x", BranchType_float_t);
0199 bookBranch(strTrackType + "_MuMinusVertex_y", BranchType_float_t);
0200 bookBranch(strTrackType + "_MuMinusVertex_z", BranchType_float_t);
0201 bookBranch(strTrackType + "_MuPlusPt", BranchType_float_t);
0202 bookBranch(strTrackType + "_MuPlusPz", BranchType_float_t);
0203 bookBranch(strTrackType + "_MuPlusPhi", BranchType_float_t);
0204 }
0205 actuateBranches();
0206 }
0207 bool HIPTwoBodyDecayAnalyzer::bookBranch(std::string bname, BranchType btype) {
0208 if (btype == BranchType_float_t)
0209 floatBranches.emplace_back(bname, new float);
0210 else if (btype == BranchType_int_t)
0211 intBranches.emplace_back(bname, new int);
0212 else if (btype == BranchType_short_t)
0213 shortBranches.emplace_back(bname, new short);
0214 else {
0215 std::cerr << "HIPTwoBodyDecayAnalyzer::bookBranch: No support for type " << btype << " for the branch " << bname
0216 << " is available." << std::endl;
0217 return false;
0218 }
0219 return true;
0220 }
0221 bool HIPTwoBodyDecayAnalyzer::actuateBranches() {
0222 bool success = true;
0223 std::cout << "Begin HIPTwoBodyDecayAnalyzer::actuateBranches" << std::endl;
0224 std::cout << "Number of short branches: " << shortBranches.size() << std::endl;
0225 std::cout << "Number of int branches: " << intBranches.size() << std::endl;
0226 std::cout << "Number of float branches: " << floatBranches.size() << std::endl;
0227 if (tree != nullptr) {
0228 for (unsigned short el = 0; el < shortBranches.size(); el++) {
0229 std::cout << "Actuating branch " << shortBranches.at(el).first << " at address " << shortBranches.at(el).second
0230 << std::endl;
0231 if (!tree->GetBranchStatus(shortBranches.at(el).first.c_str()))
0232 tree->Branch(shortBranches.at(el).first.c_str(), shortBranches.at(el).second);
0233 else
0234 std::cout << "Failed!" << std::endl;
0235 }
0236 for (unsigned int el = 0; el < intBranches.size(); el++) {
0237 std::cout << "Actuating branch " << intBranches.at(el).first.c_str() << " at address "
0238 << intBranches.at(el).second << std::endl;
0239 if (!tree->GetBranchStatus(intBranches.at(el).first.c_str()))
0240 tree->Branch(intBranches.at(el).first.c_str(), intBranches.at(el).second);
0241 else
0242 std::cout << "Failed!" << std::endl;
0243 }
0244 for (unsigned int el = 0; el < floatBranches.size(); el++) {
0245 std::cout << "Actuating branch " << floatBranches.at(el).first.c_str() << " at address "
0246 << floatBranches.at(el).second << std::endl;
0247 if (!tree->GetBranchStatus(floatBranches.at(el).first.c_str()))
0248 tree->Branch(floatBranches.at(el).first.c_str(), floatBranches.at(el).second);
0249 else
0250 std::cout << "Failed!" << std::endl;
0251 }
0252 } else
0253 success = false;
0254 if (!success)
0255 std::cerr << "HIPTwoBodyDecayAnalyzer::actuateBranch: Failed to actuate the branches!" << std::endl;
0256 return success;
0257 }
0258 HIPTwoBodyDecayAnalyzer::~HIPTwoBodyDecayAnalyzer() { cleanBranches(); }
0259
0260 void HIPTwoBodyDecayAnalyzer::analyze(const edm::Event& iEvent, const edm::EventSetup& iSetup) {
0261 using namespace edm;
0262 using namespace reco;
0263 using reco::TrackCollection;
0264
0265 edm::Handle<reco::TrackCollection> alcarecotracks;
0266 iEvent.getByToken(alcareco_trackCollToken_, alcarecotracks);
0267 edm::Handle<reco::TrackCollection> refit1tracks;
0268 iEvent.getByToken(refit1_trackCollToken_, refit1tracks);
0269 edm::Handle<reco::TrackCollection> ctftracks;
0270 iEvent.getByToken(ctf_trackCollToken_, ctftracks);
0271 edm::Handle<reco::TrackCollection> finaltracks;
0272 iEvent.getByToken(final_trackCollToken_, finaltracks);
0273
0274 const auto& theTTBuilder = iSetup.getData(ttkbuilderToken_);
0275 initializeBranches();
0276
0277 analyzeTrackCollection("alcareco", theTTBuilder, alcarecotracks);
0278 analyzeTrackCollection("refit1", theTTBuilder, refit1tracks);
0279 analyzeTrackCollection("refit2", theTTBuilder, ctftracks);
0280 analyzeTrackCollection("final", theTTBuilder, finaltracks);
0281
0282 tree->Fill();
0283 }
0284
0285 void HIPTwoBodyDecayAnalyzer::beginJob() {}
0286 void HIPTwoBodyDecayAnalyzer::endJob() {}
0287
0288 void HIPTwoBodyDecayAnalyzer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0289 edm::ParameterSetDescription desc;
0290 desc.setUnknown();
0291 descriptions.addDefault(desc);
0292 }
0293
0294 void HIPTwoBodyDecayAnalyzer::analyzeTrackCollection(std::string strTrackType,
0295 const TransientTrackBuilder& theTTBuilder,
0296 edm::Handle<reco::TrackCollection>& hTrackColl,
0297 bool verbose) {
0298 if (verbose)
0299 std::cout << "Starting to process the track collection for " << strTrackType << std::endl;
0300
0301 using namespace edm;
0302 using namespace reco;
0303 using reco::TrackCollection;
0304
0305 if (!hTrackColl.isValid()) {
0306 if (verbose)
0307 std::cout << "Track collection is invalid." << std::endl;
0308 return;
0309 }
0310 if (hTrackColl->size() < 2) {
0311 if (verbose)
0312 std::cout << "Track collection size<2." << std::endl;
0313 return;
0314 }
0315
0316 unsigned int itrk = 0;
0317 unsigned int j = 0;
0318 int totalcharge = 0;
0319 bool isValidPair = true;
0320 bool ZVtxOk = false;
0321 TLorentzVector trackMom[2];
0322 TLorentzVector trackMomAfterZVtxFit[2];
0323 TVector3 trackVtx[2];
0324
0325 for (unsigned int jtrk = 0; jtrk < 2; jtrk++) {
0326 trackMom[jtrk].SetXYZT(0, 0, 0, 0);
0327 trackVtx[jtrk].SetXYZ(0, 0, 0);
0328 }
0329 for (reco::TrackCollection::const_iterator track = hTrackColl->begin(); track != hTrackColl->end(); ++track) {
0330 int charge = track->charge();
0331 totalcharge += charge;
0332 if (j == 0) {
0333 itrk = (charge > 0 ? 1 : 0);
0334 } else
0335 itrk = 1 - itrk;
0336 trackMom[itrk].SetPtEtaPhiM(track->pt(), track->eta(), track->phi(), 0.105);
0337 trackVtx[itrk].SetXYZ(track->vx(), track->vy(), track->vz());
0338 j++;
0339 if (j == 2)
0340 break;
0341 }
0342
0343 isValidPair = (totalcharge == 0 && trackMom[0].P() != 0. && trackMom[1].P() != 0.);
0344 if (verbose && !isValidPair)
0345 std::cout << "Track collection does not contain a valid std::pair." << std::endl;
0346 setVal(strTrackType + "_present", (isValidPair ? (short)1 : (short)0));
0347 if (isValidPair) {
0348 TLorentzVector ZMom = trackMom[0] + trackMom[1];
0349 setVal(strTrackType + "_ZPt", (float)ZMom.Pt());
0350 setVal(strTrackType + "_ZPz", (float)ZMom.Pz());
0351 setVal(strTrackType + "_ZPhi", (float)ZMom.Phi());
0352 setVal(strTrackType + "_ZMass", (float)ZMom.M());
0353
0354 reco::Vertex ZVtx = fitDimuonVertex(theTTBuilder, hTrackColl, ZVtxOk);
0355 if (ZVtxOk) {
0356 setVal(strTrackType + "_ZVertex_x", (float)ZVtx.x());
0357 setVal(strTrackType + "_ZVertex_y", (float)ZVtx.y());
0358 setVal(strTrackType + "_ZVertex_z", (float)ZVtx.z());
0359 setVal(strTrackType + "_ZVertex_NormChi2", (float)ZVtx.normalizedChi2());
0360
0361
0362 j = 0;
0363 for (reco::TrackCollection::const_iterator track = hTrackColl->begin(); track != hTrackColl->end(); ++track) {
0364 TransientTrack t_track = theTTBuilder.build(&(*track));
0365 AnalyticalImpactPointExtrapolator extrapolator(t_track.field());
0366 TrajectoryStateOnSurface closestIn3DSpaceState =
0367 extrapolator.extrapolate(t_track.impactPointState(), RecoVertex::convertPos(ZVtx.position()));
0368 GlobalVector mom = closestIn3DSpaceState.globalMomentum();
0369 int charge = track->charge();
0370 totalcharge += charge;
0371 if (j == 0) {
0372 itrk = (charge > 0 ? 1 : 0);
0373 } else
0374 itrk = 1 - itrk;
0375 trackMomAfterZVtxFit[itrk].SetXYZT(mom.x(), mom.y(), mom.z(), sqrt(pow(0.105, 2) + pow(mom.mag(), 2)));
0376 j++;
0377 if (j == 2)
0378 break;
0379 }
0380 if (totalcharge != 0)
0381 std::cerr
0382 << "HIPTwoBodyDecayAnalyzer::analyzeTrackCollection: Something went wrong! The total charge is no longer 0!"
0383 << std::endl;
0384 for (unsigned int jtrk = 0; jtrk < 2; jtrk++) {
0385 std::string strMuCore = (jtrk == 0 ? "MuMinus" : "MuPlus");
0386 setVal(strTrackType + "_" + strMuCore + "Pt_AfterZVtxFit", (float)trackMomAfterZVtxFit[jtrk].Pt());
0387 setVal(strTrackType + "_" + strMuCore + "Pz_AfterZVtxFit", (float)trackMomAfterZVtxFit[jtrk].Pz());
0388 setVal(strTrackType + "_" + strMuCore + "Phi_AfterZVtxFit", (float)trackMomAfterZVtxFit[jtrk].Phi());
0389 }
0390 TLorentzVector ZMom_AfterZVtxFit = trackMomAfterZVtxFit[0] + trackMomAfterZVtxFit[1];
0391 setVal(strTrackType + "_ZPt_AfterZVtxFit", (float)ZMom_AfterZVtxFit.Pt());
0392 setVal(strTrackType + "_ZPz_AfterZVtxFit", (float)ZMom_AfterZVtxFit.Pz());
0393 setVal(strTrackType + "_ZPhi_AfterZVtxFit", (float)ZMom_AfterZVtxFit.Phi());
0394 setVal(strTrackType + "_ZMass_AfterZVtxFit", (float)ZMom_AfterZVtxFit.M());
0395 } else
0396 std::cerr << "HIPTwoBodyDecayAnalyzer::analyzeTrackCollection: Z vertex fit failed for track collection "
0397 << strTrackType << std::endl;
0398 }
0399 setVal(strTrackType + "_ZVtxFitOk", (ZVtxOk ? (short)1 : (short)0));
0400 for (unsigned int jtrk = 0; jtrk < 2; jtrk++) {
0401 std::string strMuCore = (jtrk == 0 ? "MuMinus" : "MuPlus");
0402 setVal(strTrackType + "_" + strMuCore + "Pt", (float)trackMom[jtrk].Pt());
0403 setVal(strTrackType + "_" + strMuCore + "Pz", (float)trackMom[jtrk].Pz());
0404 setVal(strTrackType + "_" + strMuCore + "Phi", (float)trackMom[jtrk].Phi());
0405 setVal(strTrackType + "_" + strMuCore + "Vertex_x", (float)trackVtx[jtrk].X());
0406 setVal(strTrackType + "_" + strMuCore + "Vertex_y", (float)trackVtx[jtrk].Y());
0407 setVal(strTrackType + "_" + strMuCore + "Vertex_z", (float)trackVtx[jtrk].Z());
0408 }
0409 }
0410
0411 reco::Vertex HIPTwoBodyDecayAnalyzer::fitDimuonVertex(const TransientTrackBuilder& theTTBuilder,
0412 edm::Handle<reco::TrackCollection>& hTrackColl,
0413 bool& fitOk) {
0414 using namespace edm;
0415 using namespace reco;
0416
0417 std::vector<TransientTrack> t_tks;
0418 for (TrackCollection::const_iterator track = hTrackColl->begin(); track != hTrackColl->end(); ++track) {
0419 TransientTrack tt = theTTBuilder.build(&(*track));
0420 t_tks.push_back(tt);
0421 }
0422
0423
0424 KalmanVertexFitter vtxFitter;
0425 TransientVertex stdVertex = vtxFitter.vertex(t_tks);
0426 fitOk = stdVertex.isValid();
0427 if (fitOk) {
0428 reco::Vertex stdRecoVertex(Vertex::Point(stdVertex.position()),
0429 stdVertex.positionError().matrix(),
0430 stdVertex.totalChiSquared(),
0431 stdVertex.degreesOfFreedom(),
0432 0);
0433 return stdRecoVertex;
0434 } else {
0435 reco::Vertex stdRecoVertex;
0436 return stdRecoVertex;
0437 }
0438 }
0439
0440 DEFINE_FWK_MODULE(HIPTwoBodyDecayAnalyzer);
