tschum/FWlite_Analysis/PlotTool.cc

#include "PlotTool.h"
//------------------------------------------------------------------
//Constructur: Set some global Variables to define canvas output
PlotTool::PlotTool() {

        this->SetClass("TChain", 100);

        samePad_trees = true;
        samePad_vars = false;
        samePad_cuts = false;
        sameCanv_trees = false;
        sameCanv_vars = false;
        sameCanv_cuts = false;

        showLegend = false;
        logY = true;
        addTrackJets = true;

}

//------------------------------------------------------------------
//Fill object PlotTool with Chains constructed from files from given source 

int PlotTool::init(string fileName, string dirPath, string treeName,
                string fileLabel) {
        this->New(this->GetEntries() );
        int currChain = this->GetEntries() - 1;

        if (currChain < 0)
                return currChain;

        fileNames.clear();

        //check if alternative label is different from default(searchString)
        if (fileLabel=="") {
                fileLabel=fileName;
        }
        ((TChain*) this->At(currChain))->SetName(fileLabel.c_str());
        TSystemDirectory dir("sourceDir", dirPath.c_str());
        TList *files = dir.GetListOfFiles();
        if (files) {
                TIter next(files);
                TSystemFile *file;
                TString fname;
                string filePath;

                cout<<"Open"<<dirPath.c_str()<<" Search for .root files that contain: "
                                <<fileName.c_str()<<endl;

                while ((file=(TSystemFile*)next())) {
                        fname = file->GetName();
                        if (!fname.EndsWith(".root"))
                                continue;
                        if (!fname.Contains(fileName.c_str()))
                                continue;

                        filePath = dirPath;
                        filePath += fname.Data();
                        //fwlite::ChainEvent to lop over events, jets, etc
                        fileNames.push_back(filePath.c_str());

                        if (((TChain*) this->At(currChain))->AddFile(filePath.c_str(), -1,
                                        treeName.c_str()) )
                                cout<<"Chained "<<((TChain*) this->At(currChain))->GetNtrees()<<" file(s) with "<<((TChain*) this->At(currChain))->GetEntries()<<" events."<<endl;
                        else
                                return -1;

                }
        } else
                return -1;

        for (int i=0; i<((TChain*) this->At(currChain))->GetListOfBranches()->GetEntries(); ++i) {

                string s(((TChain*) this->At(currChain))->GetListOfBranches()->At(i)->GetName());

                string branch_alias = s;
                string branch_name = s;
                if (s.find(".", s.size()-1) != string::npos)
                        branch_name += "obj";

                size_t a = s.find("_");
                if (a != string::npos) {
                  size_t b = s.find("_", a+1);
                  if (b != string::npos) {
                    size_t c = s.find("_", b+1);
                    if (c != string::npos) {
                      string _label    =s.substr(a+1, b-a-1);
                      string _instance =s.substr(b+1, c-b-1);
                      branch_alias = _label;
                      if(_instance.length() > 0 )  {
                        branch_alias += "_";
                        branch_alias += _instance;
                      }
                    }
                  }
                }

                ((TChain*) this->At(currChain))->SetAlias(branch_alias.c_str(),
                                branch_name.c_str());

        }

        // add branch with track  Jets

        if( addTrackJets ) {

          //make file for tree friends (adding additional, computed branches)
          string friendFileName("/scratch/hh/lustre/cms/user/thomsen/temp/");
          friendFileName += fileName;
          friendFileName += ".root";
          TFile *f = new TFile(friendFileName.c_str(),"recreate");
          string friendTreeName("friendTree");
          //char number = currChain;
          //friendTreeName += number;
          //TTree *friendTree = new TTree("friendTree","friendTree");
          TTree *friendTree = new TTree(friendTreeName.c_str(),friendTreeName.c_str());

          int nJetsKT;
          TrackJetKT = new float [100];
          fwlite::ChainEvent ev(fileNames);

          friendTree->Branch("nJetsKT", &nJetsKT, "nJetsKT/I");
          friendTree->Branch("TrackJetKT", TrackJetKT, "TrackJetKT[nJetsKT]/F");

          cout<<"calculating additional variables..."<<endl;
          for (ev.toBegin(); !ev.atEnd(); ++ev) {

            fwlite::Handle<reco::CaloJetCollection> jets;
            jets.getByLabel(ev, "antikt5CaloJets");

            fwlite::Handle<reco::TrackCollection> tracks;
            tracks.getByLabel(ev, "generalTracks");

            if (!jets.isValid())
              continue;
            if (!tracks.isValid())
              continue;
            double trackSum;
            nJetsKT = 0;
            for (reco::CaloJetCollection::const_iterator jet = jets->begin(); jet
                   !=jets->end(); jet++) {
              trackSum = 0;
              for (reco::TrackCollection::const_iterator track = tracks->begin(); track
                     !=tracks->end(); track++) {
                if (deltaR(jet->eta(), jet->phi(), track->eta(), track->phi())
                    < 0.5)
                  trackSum += track->pt();
              }
              TrackJetKT[nJetsKT] = trackSum;
              nJetsKT++;
            }
            friendTree->Fill();
          }
          f->cd();
          friendTree->Write();
          f->Close();
          ((TChain*) this->At(currChain))->AddFriend(friendTreeName.c_str(),
                                                     friendFileName.c_str());
        }
        return this->GetEntries();

}
//------------------------------------------------------------------
//Draw given Chain with given variable and given cuts
int PlotTool::plot(int chainIndex, string histName, string cutName,
                int nEntries, string drwOpt) {

        if (chainIndex >= this->GetEntries() )
                return -1;

        int currN = nEntries;
        if (nEntries < 0)
                currN = ((TChain*) this->At(chainIndex))->GetEntries(); //nEntries<0 : all entries are plotted!

        //++++ Create and name Canvases according to global variables +++++++++++++
        ostringstream currHistName;
        if (samePad_trees)
                currHistName<<((TChain*) this->At(chainIndex))->GetName()<<":";
        if (samePad_vars)
                currHistName<<histName;
        if (samePad_cuts)
                currHistName<<"{"<<cutName<<"}";

        ostringstream currPadName;
        if (!samePad_trees)
                currPadName<<((TChain*) this->At(chainIndex))->GetName()<<":";
        if (!samePad_vars)
                currPadName<<histName;
        if (!samePad_cuts)
                currPadName<<"{"<<cutName<<"}";

        ostringstream currCanvName;
        if (!sameCanv_trees && !samePad_trees)
                currCanvName<<((TChain*) this->At(chainIndex))->GetName()<<":";
        if (!sameCanv_vars && !samePad_vars)
                currCanvName<<histName;
        if (!sameCanv_cuts && !samePad_cuts)
                currCanvName<<"{"<<cutName<<"}";

        if ( (sameCanv_trees || samePad_trees) && (sameCanv_vars || samePad_vars)
                        && (sameCanv_cuts || samePad_cuts))
                currCanvName<<"All";

        string currOpt = drwOpt;
        bool useSubPads = (sameCanv_trees && !samePad_trees) || (sameCanv_vars
                        && !samePad_vars) || (sameCanv_cuts && !samePad_cuts);

        if (useSubPads) {
                if ( !canvases_[currCanvName.str()] || !gROOT->GetListOfCanvases()->FindObject(currCanvName.str().c_str()) ) {
                        canvases_[currCanvName.str()] = new TCanvas(currCanvName.str().c_str(), currCanvName.str().c_str());
                } else if (canvases_[currCanvName.str()]->GetNumber() >= 0) {
                        canvases_[currCanvName.str()]->SetNumber(canvases_[currCanvName.str()]->GetNumber()+1);
                }
        }

        if ( !pads_[currPadName.str()] || !gROOT->GetListOfCanvases()->FindObject(currPadName.str().c_str()) ) {
                pads_[currPadName.str()] = new TCanvas(currPadName.str().c_str(), currPadName.str().c_str());
                if (logY)
                        pads_[currPadName.str()]->SetLogy(1);
        } else {
                pads_[currPadName.str()]->cd();
                currOpt += "SAMES";
                if (useSubPads)
                        canvases_[currCanvName.str()]->SetNumber(canvases_[currCanvName.str()]->GetNumber()-1);
        }
        //++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++


        //Draw histogram:
        int draw_err = ((TChain*) this->At(chainIndex))->Draw(histName.c_str(), cutName.c_str(),
                        currOpt.c_str(), currN);
        if (draw_err < 0)
                return draw_err;

        //++++ Fix for histos with no entries +++++++++++++
        if (draw_err == 0) {
                if (currOpt.find("SAMES") == string::npos) {
                        pads_[currPadName.str()]->Close();
                        pads_.erase(currPadName.str() );
                        if (useSubPads)
                                canvases_[currCanvName.str()]->SetNumber(canvases_[currCanvName.str()]->GetNumber()-1);
                }
                cout<< "Warning: "<<currHistName.str().c_str()<<" in "<<currPadName.str()
                                <<" has no entries and is not drawn!"<<endl;
                return draw_err;
        }
        //++++++++++++++++++++++++++++++++++++++++++++++++++

        ((TH1F*) pads_[currPadName.str()]->GetPrimitive("htemp"))->SetName(currHistName.str().c_str()); // Set Name of histogram


        return draw_err;

}
//------------------------------------------------------------------
//Standard loop to draw all chains and multiple variables and cuts
int PlotTool::loop(vector<string> _histName, vector<string> _cutName,
                int nEntries, string drwOpt, bool correspond) {

        int numProcessed = 0;

        if (correspond == false) {

                for (int i=0; i<this->GetEntries(); ++i) {
                        for (vector<string>::iterator it1 =_histName.begin(); it1
                                        != _histName.end(); ++it1) {
                                for (vector<string>::iterator it2 =_cutName.begin(); it2
                                                != _cutName.end(); ++it2) {

                                        numProcessed += plot(i, *it1, *it2, nEntries, drwOpt);
                                        updatePads();

                                }
                        }
                }
        } else {

                if (_histName.size() != _cutName.size() )
                        return -1;

                for (int i=0; i<this->GetEntries(); ++i) {
                        for (vector<string>::iterator it1 =_histName.begin(), it2 =
                                        _cutName.begin(); it1 != _cutName.end()&&it2
                                        != _cutName.end(); ++it1, ++it2) {

                                numProcessed += plot(i, *it1, *it2, nEntries, drwOpt);

                        }
                }
        }

        //  updatePads();
        fillCanvases();

        return numProcessed;
}

//Helper methods to allow for using simple strings as input
int PlotTool::loop(string histName, string cutName, int nEntries, string drwOpt) {

        vector<string> _histName;
        _histName.push_back(histName);
        vector<string> _cutName;
        _cutName.push_back(cutName);

        return loop(_histName, _cutName, nEntries, drwOpt);

}

int PlotTool::loop(vector<string> _histName, string cutName, int nEntries,
                string drwOpt) {

        vector<string> _cutName;
        _cutName.push_back(cutName);

        return loop(_histName, _cutName, nEntries, drwOpt);

}

int PlotTool::loop(string histName, vector<string> _cutName, int nEntries,
                string drwOpt) {

        vector<string> _histName;
        _histName.push_back(histName);

        return loop(_histName, _cutName, nEntries, drwOpt);
}
//------------------------------------------------------------------
//redraw the canvas to make changes in style visible
int PlotTool::updatePads() {

        for (map< string, TCanvas* >::iterator it=pads_.begin() ; it != pads_.end(); ++it) {
                if (gROOT->GetListOfCanvases()->FindObject((*it).first.c_str() ) ) {
                        (*it).second->Draw();
                        setCanvas((*it).second);
                } else
                        pads_.erase(it);
        }

        return pads_.size();

}

int PlotTool::fillCanvases() {

        for (map< string, TCanvas* >::iterator it=canvases_.begin() ; it
                        != canvases_.end(); ++it) {
                const char* canvName = (*it).first.c_str();
                if (gROOT->GetListOfCanvases()->FindObject(canvName) ) {

                        int numP = (*it).second->GetNumber()+1;

                        if (numP <= 0)
                                continue;

                        int x = int( sqrt(numP) );
                        int y = x;
                        if (x*y < numP)
                                x += 1;
                        if (x*y < numP)
                                x += 1;
                        if (x*y < numP) {
                                x -= 1;
                                y += 1;
                        }

                        (*it).second->Divide(x, y);
                        int padIndex = 1;
                        for (map< string, TCanvas* >::iterator it2=pads_.begin() ; it2
                                        != pads_.end(); ++it2) {
                                string padName = (*it2).first;
                                if (gROOT->GetListOfCanvases()->FindObject(padName.c_str() ) ) {
                                        if ( (padName.find(canvName) != string::npos ) || !strcmp(
                                                        canvName, "All") ) {
                                                (*it).second->cd(padIndex);
                                                (*it2).second->DrawClonePad();
                                                (*it2).second->Close();
                                                pads_.erase(it2);
                                                ++padIndex;
                                        }
                                } else
                                        pads_.erase(it2);
                        }
                        (*it).second->SetNumber(-1);

                } else
                        canvases_.erase(it);
        }

        return canvases_.size();

}
//------------------------------------------------------------------
void PlotTool::setCanvas(TCanvas* thisCanvas) {

        TH1* thisHist = 0;
        TPaveStats* thisStatsBox = 0;
        TPaletteAxis* palette =0;
        TLegend* thisLeg = 0;
        int counter =0;
        double maxEntry=0;

        ((TFrame*) thisCanvas->GetFrame())->Delete();
        thisCanvas->GetFrame()->SetLineWidth(gStyle->GetLineWidth() );

        thisCanvas->SetLeftMargin(0.2);
        thisCanvas->SetRightMargin(0.06);
        thisCanvas->SetBottomMargin(0.2);
        thisCanvas->SetTopMargin(0.1);

        if (logY)
                thisCanvas->SetLogy(1);
        else
                thisCanvas->SetLogy(0);

        if (showLegend)
                thisLeg = new TLegend();

        for (int i = 0; i != thisCanvas->GetListOfPrimitives()->GetSize(); ++i) {

                if ( !thisCanvas->GetListOfPrimitives()->At(i)->InheritsFrom("TH1"))
                        continue;

                thisHist = ((TH1*) thisCanvas->GetListOfPrimitives()->At(i));
                setColor(thisHist, counter);
                setMathLabels(thisHist);

                if (thisHist->GetMaximum() > maxEntry)
                        maxEntry = thisHist->GetMaximum();

                thisStatsBox = (TPaveStats*) thisHist->GetListOfFunctions()->FindObject("stats");
                if (thisStatsBox)
                        setStats(thisCanvas, thisStatsBox, thisHist, counter);

                palette = (TPaletteAxis*) thisHist->GetListOfFunctions()->FindObject("palette");
                if (palette)
                        setPalette(thisCanvas, palette);

                if (thisLeg)
                        thisLeg->AddEntry(thisHist,thisHist->GetName())->SetTextSize(0.04);

                ++counter;

        }

        if (maxEntry != 0)
                setHistMax(thisCanvas, maxEntry);

        thisCanvas->cd();
        thisCanvas->Update();

        if (thisLeg)
                setLegend(thisCanvas, thisLeg, counter);

}
//------------------------------------------------------------------
//private helper classes to set the canvas and hist style
void PlotTool::setStats(TCanvas* thisCanvas, TPaveStats* thisStatsBox,
                TH1* thisHist, int counter) {

        if (thisCanvas->GetRightMargin() < .2)
                thisCanvas->SetRightMargin(.2);

        thisStatsBox->SetLineColor(thisHist->GetLineColor());
        thisStatsBox->SetX2NDC(1.);
        thisStatsBox->SetY2NDC(1-thisCanvas->GetTopMargin()-0.16*counter);
        thisStatsBox->SetX1NDC(1-thisCanvas->GetRightMargin()+0.01);
        thisStatsBox->SetY1NDC(thisStatsBox->GetY2NDC()-.15);

}

void PlotTool::setMathLabels(TH1* thisHist) {

        string t = thisHist ->GetTitle();
        string x = thisHist->GetXaxis()->GetTitle();
        string y = thisHist->GetYaxis()->GetTitle();

        if (t.find(".phi()") != string::npos)
                t.replace(t.find(".phi()"), 6, " #phi");
        if (x.find(".phi()") != string::npos)
                x.replace(x.find(".phi()"), 6, " #phi");
        if (y.find(".phi()") != string::npos)
                y.replace(y.find(".phi()"), 6, " #phi");

        if (t.find(".eta()") != string::npos)
                t.replace(t.find(".eta()"), 6, " #eta");
        if (x.find(".eta()") != string::npos)
                x.replace(x.find(".eta()"), 6, " #eta");
        if (y.find(".eta()") != string::npos)
                y.replace(y.find(".eta()"), 6, " #eta");

        if (t.find(".pt()") != string::npos)
                t.replace(t.find(".pt()"), 5, " p_{T}");
        if (x.find(".pt()") != string::npos)
                x.replace(x.find(".pt()"), 5, " p_{T}");
        if (y.find(".pt()") != string::npos)
                y.replace(y.find(".pt()"), 5, " p_{T}");

        thisHist ->SetTitle(t.c_str());
        thisHist->GetXaxis()->SetTitle(x.c_str());
        thisHist->GetYaxis()->SetTitle(y.c_str());

}

void PlotTool::setColor(TH1* thisHist, int counter) {

        if (counter == 0) {

                thisHist->SetLineColor(kRed+1);
                //if( gStyle->GetHistFillStyle() ) thisHist->SetFillColor  (kRed+1);
                //else 
                thisHist->SetFillStyle(0);
                thisHist->SetMarkerColor(kRed+1);
        } else if (counter == 1) {
                thisHist->SetLineColor(kBlue+1);
                //if( gStyle->GetHistFillStyle() ) thisHist->SetFillColor  (kBlue+1);
                //else 
                thisHist->SetFillStyle(0);
                thisHist->SetMarkerColor(kBlue+1);

        } else if (counter == 2) {
                thisHist->SetLineColor(kGreen+2);
                //if( gStyle->GetHistFillStyle() ) thisHist->SetFillColor  (kGreen+2);
                //else
                thisHist->SetFillStyle(0);
                thisHist->SetMarkerColor(kGreen+2);

        } else if (counter == 3) {
                thisHist->SetLineColor(kMagenta+2);
                //if( gStyle->GetHistFillStyle() ) thisHist->SetFillColor  (kMagenta+2);
                //else
                thisHist->SetFillStyle(0);
                thisHist->SetMarkerColor(kMagenta+2);

        } else if (counter == 4) {
                thisHist->SetLineColor(kCyan+2);
                //if( gStyle->GetHistFillStyle() ) thisHist->SetFillColor  (kCyan+2);
                //else
                thisHist->SetFillStyle(0);
                thisHist->SetMarkerColor(kCyan+2);

        } else {
                thisHist->SetLineColor(kBlack);
                //if( gStyle->GetHistFillStyle() ) thisHist->SetFillColor  (kBlack);
                //else
                thisHist->SetFillStyle(0);
                thisHist->SetMarkerColor(kBlack);

        }

}

void PlotTool::setPalette(TCanvas* thisCanvas, TPaletteAxis* palette) {

        palette->SetLabelSize(0.045);
        if (thisCanvas->GetRightMargin() < .15)
                thisCanvas->SetRightMargin(.15);
        palette->SetX1NDC(1-thisCanvas->GetRightMargin()+0.01);
        palette->SetY1NDC(thisCanvas->GetBottomMargin());
        palette->SetX2NDC(palette->GetX1NDC()+0.05);
        palette->SetY2NDC(1-thisCanvas->GetTopMargin());

}

void PlotTool::setHistMax(TCanvas* thisCanvas, double maxEntry) {

        TH1* thisHist = 0;

        for (int i = 0; i != thisCanvas->GetListOfPrimitives()->GetSize(); ++i) {

                if ( !thisCanvas->GetListOfPrimitives()->At(i)->InheritsFrom("TH1") )
                        continue;
                thisHist = ((TH1*) thisCanvas->GetListOfPrimitives()->At(i));
                if (thisHist->GetMaximum() < maxEntry) {
                        double minEntry=thisHist->GetBinContent(thisHist->GetMinimumBin());
                        if (logY) {
                                int bin = (thisHist->GetMinimumBin()+1);
                                minEntry=thisHist->GetBinContent(bin);
                        }
                        thisHist->GetYaxis()->SetRangeUser(minEntry, maxEntry*1.2);
                }
                break;

        }

}

void PlotTool::setLegend(TCanvas* thisCanvas, TLegend* thisLeg, int counter) {

        thisLeg->SetFillStyle(0);
        thisLeg->SetX1NDC(1-thisCanvas->GetRightMargin()-0.5);
        thisLeg->SetY1NDC(1-thisCanvas->GetTopMargin()-0.05*counter);
        thisLeg->SetX2NDC(1-thisCanvas->GetRightMargin());
        thisLeg->SetY2NDC(1-thisCanvas->GetTopMargin());
        thisLeg->SetEntrySeparation(0.5);
        thisLeg->Draw("NDC");

}

void PlotTool::createColors() {
        const Int_t NRGBs = 5;
        const Int_t NCont = 255;

        Double_t stops[NRGBs] = { 0.00, 0.34, 0.61, 0.84, 1.00 };
        Double_t red[NRGBs] = { 0.00, 0.00, 0.87, 1.00, 0.51 };
        Double_t green[NRGBs] = { 0.00, 0.81, 1.00, 0.20, 0.00 };
        Double_t blue[NRGBs] = { 0.51, 1.00, 0.12, 0.00, 0.00 };
        TColor::CreateGradientColorTable(NRGBs, stops, red, green, blue, NCont);
        gStyle->SetNumberContours(NCont);

}

//------------------------------------------------------------------

int PlotTool::saveCanvases(string type, string path) {

        TSystemDirectory d("", path.c_str());
        if (!d.GetListOfFiles())
                return -1;
        if (type.find(".") == string::npos)
                return -1;

        int savedFiles =0;

        TIter next(gROOT->GetListOfCanvases());
        TCanvas* canv;
        while ((canv=(TCanvas*)next())) {

                string s = "";
                s += path;
                s += canv->GetName();
                s += type;

                canv->SaveAs(s.c_str());
                ++savedFiles;

        }

        return savedFiles;

}
//------------------------------------------------------------------
void PlotTool::showChainInfo()
{


  cout<<endl;
  cout<<"****** Show Chain Information:"<<endl;


  this->ls();

  cout<<endl;
  cout<<"We have "<<this->GetEntries()<<" TChains with following aliases set:"<<endl;
        for (int i=0; i< this->GetEntries(); ++i) {
          cout<<endl;
          cout<<((TChain*) this->At(i))->GetName()<<endl;
          ((TChain*) this->At(i))->GetListOfAliases()->ls();

        }


}
Revision:	1.7
Committed:	Fri Nov 27 17:39:46 2009 UTC (15 years, 5 months ago) by tschum
Content type:	text/plain
Branch:	MAIN
Changes since 1.6:	+102 -62 lines
Log Message:	fixed bug with slash in name and label ambiguity