Selection/src/applySelection.cc

//
// System headers
//
#include <vector>                   // STL vector class
#include <iostream>                 // standard I/O
#include <iomanip>                  // functions to format standard I/O
#include <bitset>  
#include <map>  
using namespace std;

//
// root headers
//
#include <TFile.h>
#include <TRegexp.h>
#include <TTree.h>                  
#include <TChain.h>                 
#include <TBranch.h>                
#include <TClonesArray.h>           


//
// TMVA
//
#include "TMVA/Reader.h"
#include "TMVA/Tools.h"
#include "TMVA/Config.h"
#include "TMVA/MethodBDT.h"

#include "FactorizedJetCorrector.h"
#include "JetCorrectorParameters.h"

//
// ntuple format headers
//
#include "EventHeader.h"
#include "Electron.h"
#include "Muon.h"
#include "Vertex.h"
#include "PFMet.h"
#include "PFJetCol.h"
#include "PFCandidate.h"
#include "PFCandidateCol.h"
#include "PileupInfoCol.h"
#include "PileupEnergyDensity.h"
#include "MCParticle.h"
#include "TriggerMask.h"
#include "TriggerTable.h"
#include "TriggerObject.h"
#include "TriggerObjectRel.h"
#include "Names.h"
#include "BaseMod.h"
#include "TriggerObjectsTable.h"
#include "JetIDMVA.h"

//
// our headers
//
#include "ParseArgs.h"
#include "MuonSelection.h"
#include "ElectronSelection.h"
#include "ElectronMomentumCorrection.h"
#include "JetSelection.h"
#include "IsolationSelection.h"
#include "ReferenceSelection.h"
#include "MuonMomentumCorrection.h"

#include "TriggerUtils.h"
#include "PassHLT.h"
#include "Angles.h"
#include "KinematicsStruct.h"
#include "InfoStruct.h"
#include "GenInfoStruct.h"
#include "WeightStruct.h"
#include "AngleTuple.h"
#include "JetInfoStruct.h"
#include "FOTuple.h"

#include "RunLumiRangeMap.h"
#include "SampleWeight.h"
#include "EfficiencyWeightsInterface.h"
#include "SelectionFuncs.h"

#include "SimpleLepton.h"

#ifndef CMSSW_BASE
#define CMSSW_BASE "../"
#endif

using namespace mithep;

//
// globals
//----------------------------------------------------------------------------
TH1D *hpu_2011, *hpu_2012;
RunLumiRangeMap rlrm;
vector<SimpleLepton> failingLeptons ; // for fake estimation
vector<SimpleLepton> passingLeptons;      // for fake estimation
vector<unsigned> cutvec;
vector<vector<unsigned> > zcutvec;
vector<vector<unsigned> > zzcutvec;
map<unsigned,float>       evtrhoMap;
TrigInfo ti;
vector<string> cutstrs;
bool passes_HLT;
vector<bool>   PFnoPUflag;;
map<TString, map<TString,int>* > counts;
ElectronMomentumCorrection electron_momentum_correction;
MuCorr *muCorr;
void addDummyZ2Lepotons(EventData &ret4l);
void setNtupleVals(ControlFlags &ctrl, EventData &ret4l, InfoStruct &evtinfo, Angles &angles, KinematicsStruct &kine,
                   JetInfoStruct &ji, WeightStruct &weights, EventHeader *info, Array<PFJet> *jetArr, Array<PFMet> *metArr,
                   vector<SimpleLepton> &simpleLeptonJets, FactorizedJetCorrector *fJetCorrector, Array<PileupInfo> *puArr,
                   Array<PileupEnergyDensity> *puDArr, JetIDMVA *fJetIDMVA, Array<Vertex> *vtxArr, bitset<1024> &triggerBits,
                   TriggerTable *hltTable, Array<TriggerObject> *hltObjArr, TriggerObjectsTable *fTrigObjs, AngleTuple *nt);
//----------------------------------------------------------------------------
void initRunLumiRangeMap(ControlFlags &ctrl) 
//----------------------------------------------------------------------------
{
  if(ctrl.jsonFile==""){
    cout << "WARNING: running with empty jsonFile, so setting noJSON to true" << endl;
    ctrl.noJSON = true;
  }
  rlrm.AddJSONFile(string(ctrl.jsonFile));  
}

//
// prototypes
//----------------------------------------------------------------------------

//
// MAIN
//----------------------------------------------------------------------------
int main(int argc, char** argv) 
//----------------------------------------------------------------------------
{
  vector<TString> cuts;
  cuts.push_back("start"); // NOTE: is incremented in ReferenceSelection, so if you turn on JSONS be careful
  cuts.push_back("trigger");
  cuts.push_back("sfOsHiPt");
  cuts.push_back("4<mZ2<120");
  cuts.push_back("pt>20,10");
  cuts.push_back("mll>4");
  cuts.push_back("m4l>70");
  cuts.push_back("mZ2>12");
  cuts.push_back("m4l>100");

  for(unsigned icut=0; icut<cuts.size(); icut++)
    counts[cuts[icut]] = new map<TString,int>;

  map<TString,map<TString,int>* >::iterator cit;
  for(cit=counts.begin(); cit!=counts.end(); cit++) {
    map<TString,int> *cts = (*cit).second;
    (*cts)["all"]  = int(0);
    (*cts)["4e"]   = int(0);
    (*cts)["4m"]   = int(0);
    (*cts)["2e2m"] = int(0);
  }
  
  string cmsswpath(CMSSW_BASE);
  cmsswpath.append("/src");
  bitset<1024> triggerBits;
  vector<vector<string> > inputFiles;
  inputFiles.push_back(vector<string>());
  map<string,unsigned> entrymap; // number of unskimmed entries in each file

  //
  // args
  //--------------------------------------------------------------------------------------------------------------
  ControlFlags ctrl;
  parse_args( argc, argv, ctrl );
  ctrl.dump();
  cout << "\n\nNOTE: now setting era with arguments... double check this doesnt mess something else up\n\n" << endl;
  assert(ctrl.doFSR);

  //
  // File I/O
  //--------------------------------------------------------------------------------------------------------------
  TChain * chain = new TChain("Events");
  TChain * hltchain = new TChain("HLT");
  TChain * runchain = new TChain(Names::gkRunTreeName);

  string fname;
  unsigned total_unskimmed=0;
  if( !ctrl.inputfiles.empty() ) { 
    ifstream f(ctrl.inputfiles.c_str());
    while (f >> fname) {
      if( !(strncmp( fname.c_str(), "#", 1 ) ) ) continue; // skip commented lines
      if(fname.find("/store/") != string::npos) fname = "root://eoscms//"+fname;
      cout << "adding inputfile : " << fname.c_str() << endl;
      entrymap[string(fname.c_str())] = unskimmedEntries(fname.c_str());
      cerr << "unskimmed entries: " << entrymap[string(fname.c_str())] << endl;
      total_unskimmed += entrymap[string(fname.c_str())];
      chain->AddFile(fname.c_str());
      hltchain->AddFile(fname.c_str());
      runchain->AddFile(fname.c_str());
    }
  } else { 
    if(ctrl.inputfile.find("/store/") != string::npos) ctrl.inputfile = "root://eoscms//"+ctrl.inputfile;
    cout << "adding inputfile : " << ctrl.inputfile.c_str() << endl;
    unsigned unsk_ents = unskimmedEntries(ctrl.inputfile.c_str());
    entrymap[string(ctrl.inputfile.c_str())] = unsk_ents;
    cerr << "unskimmed entries: " << unsk_ents << endl;
    total_unskimmed += unsk_ents;
    chain->AddFile(ctrl.inputfile.c_str());
    hltchain->AddFile(ctrl.inputfile.c_str());
    runchain->AddFile(ctrl.inputfile.c_str());
  }
  // write the total number of unskimmed events that went into making this output file to a text file
  writeEntries(ctrl,total_unskimmed);
  ctrl.totalMC = total_unskimmed;

  const char * ofname;
  if( strcmp( ctrl.outputfile.c_str(), "") ) {
    ofname = ctrl.outputfile.c_str();
  } else {
    ofname = "tmp.root";
  }
  // table of cross section values   
  string xspath = (cmsswpath+string("/MitPhysics/data/xs.dat"));
  cout << "xspath: " << xspath.c_str() << endl;
  SimpleTable xstab(xspath.c_str());
  if(ctrl.mc) {
    getWeight(xstab,entrymap,chain); // test *now* whether we've got the xs in here, so it fails before it creates the output
  }

  //
  // Setup
  //--------------------------------------------------------------------------------------------------------------
  Angles angles;
  KinematicsStruct kine;
  InfoStruct evtinfo;
  WeightStruct weights;
  GenInfoStruct geninfo;
  JetInfoStruct ji;
  vector<SimpleLepton> simpleLeptonJets;

  AngleTuple nt( (const char*)ofname, (const char*)"zznt");
  nt.makeAngleBranch(angles);
  nt.makeKinematicsBranch(kine);
  nt.makeInfoBranch(evtinfo);
  nt.makeJetInfoBranch(ji);
  nt.makeJetAngleBranches("DataStruct/data/jetAngleVars.txt");

  FOTuple *foTree=0;
  if(ctrl.fakeScheme != "none") {
    foTree = new FOTuple( nt.getFile(), (const char*)"FOtree");
    foTree->makeInfoBranch(evtinfo);
    foTree->makeFailedLeptonBranch(failingLeptons);
    foTree->makePassedLeptonBranch(passingLeptons);
  }

  int genSampleType(0);
  if(ctrl.mc) { 
    nt.makeWeightBranch(weights);
    if(ctrl.fillGen) {
      genSampleType = getGenSampleType(ofname);
      if(genSampleType == -1) {
        cout << "\ngenSampleType returned as -1, setting fillGen to false" << endl;
        ctrl.fillGen = false;
      } else
        nt.makeGenInfoBranch(geninfo);
    }
    initEfficiencyWeights();
    // initPUWeights();
  } else { 
    if(!(ctrl.noJSON) )
      initRunLumiRangeMap(ctrl);
  }

  initElectronIDMVAV1();
  initTrigger();
  initAnalysisTriggers(ti);
  ti.dump();

  if(ctrl.do_escale) {
    string regPath("EGamma/EGammaAnalysisTools/data/eleEnergyRegWeights_V1.root");
    // if(TString(getenv("HOSTNAME")).Contains(TRegexp("t3[bd][te][cs][hk]")))
      regPath = cmsswpath + "/" + regPath;
    electron_momentum_correction.initialize_regression(regPath);
    if(ctrl.mc) {
      if(ctrl.era == 2011) ctrl.smearing_scaling_dataset = "Fall11";
      else                 ctrl.smearing_scaling_dataset = "Summer12_DR53X_HCP2012";
    } else {
      if(ctrl.era == 2011) ctrl.smearing_scaling_dataset = "Jan16ReReco";
      else                 ctrl.smearing_scaling_dataset = "2012Jul13ReReco";
    }
  }
  if(ctrl.correct_muon_momentum) {
    muCorr = new MuCorr;
    muCorr->corr2011 = new rochcor;
    muCorr->corr2012 = new rochcor2012;
  }

  vector<TString> fJetCorrParsv;
  vector<JetCorrectorParameters> correctionParameters;
  FactorizedJetCorrector *fJetCorrector;
  JetIDMVA *fJetIDMVA;
  initJets(ctrl, cmsswpath, fJetCorrParsv, correctionParameters, fJetCorrector, fJetIDMVA);

  //
  // Setup tree I/O
  //--------------------------------------------------------------------------------------------------------------
  TFile *inputFile=0;
  TTree *eventTree=0;  
  string currentFile("");

  UInt_t fNMaxTriggers = TRIGGER_BIG_NUMBER;
  EventHeader                   *info                    = new EventHeader();
  Array<PFMet>                  *metArr                  = new Array<PFMet>();
  Array<Electron>               *electronArr             = new Array<Electron>();
  Array<Muon>                   *muonArr                 = new Array<Muon>();
  Array<Vertex>                 *vtxArr                  = new Array<Vertex>();
  Array<PFCandidate>            *pfArr                   = new Array<PFCandidate>();
  Array<PFJet>                  *jetArr                  = new Array<PFJet>();
  Array<PileupInfo>             *puArr                   = new Array<PileupInfo>();
  Array<PileupEnergyDensity>    *puDArr                  = new Array<PileupEnergyDensity>();
  Array<Track>                  *trkArr                  = new Array<Track>();
  Array<MCParticle>             *mcArr                   = new Array<MCParticle>();
  MCEventInfo                   *mcEvtInfo               = new MCEventInfo();
  TriggerMask                   *trigMask                = new TriggerMask();
  TriggerTable                  *hltTable                = new TriggerTable(fNMaxTriggers);
  vector<string>                *hltTableStrings         = new vector<string>();
  vector<string>                *hltLabelStrings         = new vector<string>();
  Array<TriggerObject>          *hltObjArr               = new Array<TriggerObject>();
  Array<TriggerObjectRel>       *hltRelsArr              = new Array<TriggerObjectRel>();
  TriggerObjectsTable           *fTrigObjs               = new TriggerObjectsTable(hltTable,fNMaxTriggers);
  RunInfo                       *runInfo                 = new RunInfo();
  
  chain->SetBranchAddress(Names::gkEvtHeaderBrn,                &info);
  chain->SetBranchAddress("PFMet",                              &metArr);
  chain->SetBranchAddress(Names::gkElectronBrn,                 &electronArr);
  chain->SetBranchAddress(Names::gkMuonBrn,                     &muonArr);
  chain->SetBranchAddress(Names::gkPVBrn,                       &vtxArr);
  chain->SetBranchAddress(Names::gkPFCandidatesBrn,             &pfArr);
  chain->SetBranchAddress(Names::gkPFJetBrn,                    &jetArr);
  if( ctrl.mc ) { 
    chain->SetBranchAddress(Names::gkPileupInfoBrn,             &puArr);
    chain->SetBranchAddress(Names::gkMCPartBrn,                 &mcArr);
    chain->SetBranchAddress(Names::gkMCEvtInfoBrn,              &mcEvtInfo);
  }
  chain->SetBranchAddress(Names::gkPileupEnergyDensityBrn,      &puDArr);
  chain->SetBranchAddress(Names::gkTrackBrn,                    &trkArr);
  chain->SetBranchAddress(Names::gkHltBitBrn,                   &trigMask);
  chain->SetBranchAddress(Names::gkHltObjBrn,                   &hltObjArr);
  chain->SetBranchAddress("HLTObjectsRelation",                 &hltRelsArr);
  hltchain->SetBranchAddress(Names::gkHltTableBrn,              &hltTableStrings);
  hltchain->SetBranchAddress(Names::gkHltLabelBrn,              &hltLabelStrings);

  runchain->SetBranchAddress(Names::gkRunInfoBrn, &runInfo);

  Vertex              vtx;          // best primary vertex in the event

  hltchain->GetEntry(0);

  int imax = (ctrl.n_events_to_process < 0) ? chain->GetEntries() : ctrl.n_events_to_process;
  cerr << "processing: " << imax << " ( / " << chain->GetEntries() << " )" << endl;

  if(ctrl.mc && ctrl.debugOffMc) // necessary for batch mode -- we need debug output for, data but size is too large to store MC output
    ctrl.debug = false;

  int lastHltEntry=-2;
  for(UInt_t ientry=0; ientry<imax; ientry++) {
    chain->GetEntry(ientry);
    if(!(ientry%1000)) cerr << "entry: " << ientry << " ( / " << imax << " )" << endl;

    increment(counts,"start");
    if( ctrl.debug ) {
      cout << "-----------------------------------------------------------------" << endl;
      cout << "-----------------------------------------------------------------" << endl;
      cout << "Run: " << info->RunNum() 
           << "\tEvt: " << info->EvtNum()
           << "\tLumi: " << info->LumiSec()
           << endl;
      cerr << "Run: " << info->RunNum() 
           << "\tEvt: " << info->EvtNum()
           << "\tLumi: " << info->LumiSec()
           << endl;
      cout << "-----------------------------------------------------------------" << endl;
    }

    runchain->GetEvent(info->RunEntry());
    hltchain->GetEntry(runInfo->HltEntry());
    hltTable->Clear();
    fillTriggerNames(hltTable, hltTableStrings );
    // if(ctrl.debug && (lastHltEntry != runInfo->HltEntry())) hltTable->Print();
    lastHltEntry = runInfo->HltEntry();
    fillTriggerBits( hltTable, trigMask, triggerBits );
    setHLTObjectRelations( hltObjArr, hltRelsArr, hltTableStrings, hltLabelStrings, fTrigObjs );
    // printTriggerObjs( hltTable, fTrigObjs);
      
    string fname = string(chain->GetFile()->GetEndpointUrl()->GetFile());
    if(ctrl.debug) cout << "era is " << ctrl.era << endl;
    if( ctrl.era == 0  ) {
      setEra( fname, ctrl );
      if(ctrl.debug) cout << "era is now " << ctrl.era << endl;
    }

    // pfNoPU 
    PFnoPUflag.clear();
    int pfnopu_size = makePFnoPUArray( pfArr, PFnoPUflag, vtxArr );
    assert(pfnopu_size == pfArr->GetEntries());

    // trigger
    if( string(chain->GetFile()->GetEndpointUrl()->GetFile()) != currentFile ) {
      currentFile = string(chain->GetFile()->GetEndpointUrl()->GetFile());
    }
    if( ctrl.debug ) cout << "file is : " << currentFile  << endl;
    passes_HLT = passHLT(ctrl, triggerBits, ti, info->RunNum(), hltTable, fTrigObjs); // NOTE: run value is not used if it's MC

    // data/MC
    if(ctrl.mc) {
      if(ctrl.fillGen) fillGenInfo( mcArr, mcEvtInfo, geninfo, genSampleType, ctrl);
      // NOTE: xs weight and npuw are reset in merge.cc
      weights.w = getWeight(xstab,entrymap,chain)/ctrl.totalMC;
      if(ctrl.debug) cout ;
      weights.npu = getNPU(puArr, 0);
      weights.npuw = getPUWeight(ctrl.era, fname, weights.npu);
    } else { 
      // JSON
      if(!(ctrl.noJSON) ) {
        RunLumiRangeMap::RunLumiPairType rl(info->RunNum(), info->LumiSec());      
        if( !(rlrm.HasRunLumi(rl)) )  {
          if( ctrl.debug ) cout << "fails JSON" << endl;
          continue;
        }
      }
    }

    //
    // lepton/kinematic selection ...
    //
    failingLeptons.clear();
    passingLeptons.clear();
    EventData ret4l = apply_HZZ4L_reference_selection(ctrl, info, 
                                                      vtxArr,
                                                      pfArr,
                                                      puDArr,
                                                      electronArr, 
                                                      &electronReferencePreSelection,
                                                      &electronReferenceIDMVASelectionV1, // not used anymore! (Blame the damn regression implementation, not me!)
                                                      &electronReferenceIsoSelection,
                                                      muonArr, 
                                                      &muonReferencePreSelection,
                                                      &muonIDPFSelection,
                                                      &muonReferenceIsoSelection);
      
    if( ctrl.debug ) cout << endl;

    if(ctrl.fakeScheme == "none") { // fill angletuple if event passes full selection
      if( ret4l.status.pass() ) {
        setNtupleVals(ctrl, ret4l, evtinfo, angles, kine, ji, weights, info, jetArr, metArr, simpleLeptonJets, fJetCorrector, puArr, puDArr, fJetIDMVA, vtxArr, triggerBits, hltTable, hltObjArr, fTrigObjs, &nt);
        nt.Fill();
      } else {
        if(ctrl.debug) cout << "failing with code : " << ret4l.status.selectionBits << endl;
      }
    } else { // fill FOtuple and angletuple if event has a good Z and one or two extra leptons
      assert(ctrl.fakeScheme=="ZPlusF" || ctrl.fakeScheme=="ZPlusFF");
      if(ret4l.status.selectionBits.test(PASS_GOODZ1) ) {
        unsigned nleps = passingLeptons.size() + failingLeptons.size();
        if( (nleps>=3 && ctrl.fakeScheme=="ZPlusF" ) || (nleps>=4 && ctrl.fakeScheme=="ZPlusFF" ) ) {
          addDummyZ2Lepotons(ret4l);
          setNtupleVals(ctrl, ret4l, evtinfo, angles, kine, ji, weights, info, jetArr, metArr, simpleLeptonJets, fJetCorrector, puArr, puDArr, fJetIDMVA, vtxArr, triggerBits, hltTable, hltObjArr, fTrigObjs, &nt);
          nt.Fill();
          foTree->Fill();
        }
      } else {
        if(ctrl.debug) cout << "failing PASS_GOODZ1" << endl;
      }
    }
  }

  if(ctrl.fakeScheme != "none") {
    foTree->getFile()->cd();
    foTree->getTree()->Write();
  }
  nt.WriteClose();

  for(unsigned icut=0; icut<cuts.size(); icut++) {
    map<TString,int> thisCount(*(counts[cuts[icut]]));
    cerr << setw(20) << cuts[icut] << "\t" << setw(9) << thisCount["all"] << "  ("
         << setw(9) << thisCount["4e"] << setw(9) << thisCount["4m"] << setw(9) << thisCount["2e2m"] << ")" << endl;
  }
}
//----------------------------------------------------------------------------------------
void addDummyZ2Lepotons(EventData &ret4l)
// for fake ntuples, we don't know which are the Z2 leptons yet, but need to fill something so we don't seg fault
{
  SimpleLepton dum1,dum2;
  dum1.vec.SetPtEtaPhiM(1,2,3,4);
  dum2.vec.SetPtEtaPhiM(5,6,7,8);
  ret4l.Z2leptons.push_back(dum1);
  ret4l.Z2leptons.push_back(dum2);
}
//----------------------------------------------------------------------------------------
void setNtupleVals(ControlFlags &ctrl, EventData &ret4l, InfoStruct &evtinfo, Angles &angles, KinematicsStruct &kine,
                   JetInfoStruct &ji, WeightStruct &weights, EventHeader *info, Array<PFJet> *jetArr, Array<PFMet> *metArr,
                   vector<SimpleLepton> &simpleLeptonJets, FactorizedJetCorrector *fJetCorrector, Array<PileupInfo> *puArr,
                   Array<PileupEnergyDensity> *puDArr, JetIDMVA *fJetIDMVA, Array<Vertex> *vtxArr, bitset<1024> &triggerBits,
                   TriggerTable *hltTable, Array<TriggerObject> *hltObjArr, TriggerObjectsTable *fTrigObjs, AngleTuple *nt)
{
  fillKinematics(ret4l,kine);
  // fillMassErrors(ret4l,muonArr,electronArr,kine);
  TLorentzVector pfmet; pfmet.SetPxPyPzE(metArr->At(0)->Mex(),metArr->At(0)->Mey(),0,0);
  fillEventInfo( info, pfmet, evtinfo, ctrl.mc ? getNPU(puArr) : 0, vtxArr->GetEntries());
  fillJets(ctrl, jetArr, simpleLeptonJets, fJetCorrector, puDArr, fJetIDMVA, vtxArr, ret4l);
  fillJetInfo(simpleLeptonJets, ji, ctrl);
  fillJetAngleBranches(ji, ret4l, nt, ctrl);
  evtinfo.status = 0;
  evtinfo.status = setHltBits(ti, triggerBits);

  if(ctrl.fakeScheme == "none") {
    fillAngles(ret4l,angles);

    if( ctrl.mc) { 
      setEfficiencyWeights(ctrl, ctrl.era, ret4l, triggerBits, hltTable, hltObjArr, fTrigObjs, weights); 
      if(ctrl.debug) 
        cout << "w: " << weights.w << "\t" << "npuw: " << weights.npuw << "\t" << "won: " << weights.won << "\t" << "woff: " << weights.woff << endl;
    }

    if(ctrl.debug) {
      stringstream ss;
      ss << "PASS:: " << "\trun: " << evtinfo.run << "\tevt: " << evtinfo.evt << "\tlumi: " << evtinfo.lumi << "\tchannel: "
         << kine.channel << "\tm4l: " << kine.m4l << "\tmZ1: " << kine.mZ1 << "\tmZ2: " << kine.mZ2 << endl;
      cerr << ss.str();
      cout << ss.str();
    }
  } else {
    // fillAngles(ret4l,angles); // at this point we don't know which four leptons will constitute the ZZ system, so doesn't make sense to fill the angles
  }
}
Revision:	1.62
Committed:	Mon Dec 17 17:12:21 2012 UTC (12 years, 5 months ago) by dkralph
Content type:	text/plain
Branch:	MAIN
CVS Tags:	compiled
Changes since 1.61:	+180 -128 lines
Log Message:	* empty log message *