Plotting.h

Go to the documentation of this file.

#ifndef Core_Plotting_h
#define Core_Plotting_h
 
#include <map>
#include <string>
#include <iostream>
#include <vector>
#include <cmath>
 
#include "TStyle.h"
#include "TGraph.h"
#include "TGraph2D.h"
#include "TCanvas.h"
#include "TPad.h"
#include "TLatex.h"
#include "TAxis.h"
#include "TList.h"
#include "TH1F.h"
#include "TLegend.h"
#include "TColor.h"
#include "TROOT.h"
#include "TFrame.h"
#include "TTree.h"
#include "TFile.h"
#include "TLine.h"
#include "TMultiGraph.h"
#include "THStack.h"
 
 
 
std::vector<TPad*> OnePad();
 
std::vector<TPad*> TwoPadSplit(double split_point, double gap_low,
                               double gap_high);
 
TH1* CreateAxisHist(TH1* src, double xmin = 0, double xmax = -1);
 
std::vector<TH1*> CreateAxisHists(unsigned n, TH1* src, double xmin = 0,
                                  double xmax = -1);
 
TH1* CreateAxisHist(TGraph* src, bool at_limits);
 
std::vector<TH1*> CreateAxisHists(unsigned n, TGraph* src,  bool at_limits);
 
 
TH1* GetAxisHist(TPad *pad);
TH1* MakeRatioHist(TH1* num, TH1* den, bool num_err, bool den_err);
 
 
TGraph TGraphFromTree(TTree* tree, TString const& xvar, TString const& yvar,
                      TString const& selection = "");
 
TGraph2D TGraph2DFromTree(TTree* tree, TString const& xvar, TString const& yvar,
                          TString const& zvar, TString const& selection = "");
 
void ReZeroTGraph(TGraph *gr);
int CreateTransparentColor(int color, float alpha);
 
void SetupTwoPadSplitAsRatio(std::vector<TPad*> const& pads, TH1* upper,
                             TH1* lower, TString y_title, bool y_centered,
                             float y_min, float y_max);
 
void SetupTwoPadSplitAsRatio(std::vector<TPad*> const& pads, TString y_title,
                             bool y_centered, float y_min, float y_max);
 
void StandardAxes(TAxis* xaxis, TAxis* yaxis, TString var, TString units);
 
void UnitAxes(TAxis* xaxis, TAxis* yaxis, TString var, TString units);
 
TLegend* PositionedLegend(double width, double height, int pos, double offset);
 
TLine* DrawHorizontalLine(TPad *pad, TLine* line, double yval);
 
void DrawTitle(TPad* pad, TString text, int align);
 
void DrawCMSLogo(TPad* pad, TString cmsText, TString extraText, int iPosX,
                 float relPosX, float relPosY, float relExtraDY);
 
void DrawCMSLogo(TPad* pad, TString cmsText, TString extraText, int iPosX);
double GetPadYMax(TPad* pad, double x_min, double x_max);
 
double GetPadYMax(TPad* pad);
 
void FixTopRange(TPad *pad, double fix_y, double fraction);
 
 
void FixBoxPadding(TPad *pad, TBox *box, double frac);
 
void FixOverlay();
//----------------------------------------------------------------------------
// Pad and axis histogram creation
//----------------------------------------------------------------------------
std::vector<TPad*> OnePad() {
  TPad* pad = new TPad("pad", "pad", 0., 0., 1., 1.);
  pad->Draw();
  pad->cd();
  std::vector<TPad*> result;
  result.push_back(pad);
  return result;
}
 
std::vector<TPad*> TwoPadSplit(double split_point, double gap_low,
                               double gap_high) {
  TPad* upper = new TPad("upper", "upper", 0., 0., 1., 1.);
  upper->SetBottomMargin(split_point + gap_high);
  upper->SetFillStyle(4000);
  upper->Draw();
  TPad* lower = new TPad("lower", "lower", 0., 0., 1., 1.);
  lower->SetTopMargin(1 - split_point + gap_low);
  lower->SetFillStyle(4000);
  lower->Draw();
  upper->cd();
  std::vector<TPad*> result;
  result.push_back(upper);
  result.push_back(lower);
  return result;
}
 
TH1* CreateAxisHist(TH1* src, double xmin, double xmax) {
  TH1 *res = reinterpret_cast<TH1F*>(src->Clone());
  res->Reset();
  if (xmax > xmin) {
    res->SetAxisRange(xmin, xmax);
  }
  return res;
}
 
std::vector<TH1*> CreateAxisHists(unsigned n, TH1* src, double xmin,
                                  double xmax) {
  std::vector<TH1*> result;
  for (unsigned i = 0; i < n; ++i) {
    result.push_back(CreateAxisHist(src, xmin, xmax));
  }
  return result;
}
 
TH1* CreateAxisHist(TGraph* src, bool at_limits) {
  TVirtualPad *backup = gPad;
  TCanvas tmp;
  tmp.cd();
  src->Draw("AP");
  TH1 * result = (TH1*)src->GetHistogram()->Clone();
  if (at_limits) {
    double min = 0;
    double max = 0;
    double x = 0;
    double y = 0;
    src->GetPoint(0, x, y);
    min = x;
    max = x;
    for (int i = 1; i < src->GetN(); ++i) {
      src->GetPoint(i, x, y);
      if (x < min) min = x;
      if (x > max) max = x;
    }
    result->GetXaxis()->SetLimits(min, max);
  }
  gPad = backup;
  return result;
}
 
std::vector<TH1*> CreateAxisHists(unsigned n, TGraph* src,  bool at_limits) {
  std::vector<TH1*> result;
  for (unsigned i = 0; i < n; ++i) {
    result.push_back(CreateAxisHist(src, at_limits));
  }
  return result;
}
 
TH1* GetAxisHist(TPad *pad) {
  TList* pad_obs = pad->GetListOfPrimitives();
  if (!pad_obs) return NULL;
  TIter next(pad_obs);
  TObject* obj;
  TH1 *hobj = 0;
  while ((obj = next())) {
    if (obj->InheritsFrom(TH1::Class())) {
      hobj = reinterpret_cast<TH1*>(obj);
      break;
    }
    if (obj->InheritsFrom(TMultiGraph::Class())) {
      TMultiGraph *mg = reinterpret_cast<TMultiGraph*>(obj);
      if (mg) return mg->GetHistogram();
    }
    if (obj->InheritsFrom(TGraph::Class())) {
      TGraph *g = reinterpret_cast<TGraph*>(obj);
      if (g) return g->GetHistogram();
    }
    if (obj->InheritsFrom(THStack::Class())) {
      THStack *hs = reinterpret_cast<THStack*>(obj);
      if (hs) return hs->GetHistogram();
    }
  }
  return hobj;
}
 
 
//----------------------------------------------------------------------------
// Object extraction and manipulation
//----------------------------------------------------------------------------
TH1* MakeRatioHist(TH1* num, TH1* den, bool num_err, bool den_err) {
  TH1* result = reinterpret_cast<TH1*>(num->Clone());
  if (!num_err) {
    for (int i = 1; i <= result->GetNbinsX(); ++i) result->SetBinError(i, 0.);
  }
  TH1* den_fix = reinterpret_cast<TH1*>(den->Clone());
  if (!den_err) {
    for (int i = 1; i <= den_fix->GetNbinsX(); ++i) den_fix->SetBinError(i, 0.);
  }
  result->Divide(den_fix);
  delete den_fix;
  return result;
}
 
TGraph TGraphFromTree(TTree * tree, TString const& xvar, TString const& yvar,
                     TString const& selection) {
  tree->Draw(xvar+":"+yvar, selection, "goff");
  TGraph gr(tree->GetSelectedRows(), tree->GetV1(), tree->GetV2());
  return gr;
}
 
TGraph2D TGraph2DFromTree(TTree* tree, TString const& xvar, TString const& yvar,
                          TString const& zvar, TString const& selection) {
  tree->Draw(xvar+":"+yvar+":"+zvar, selection, "goff");
  TGraph2D gr(tree->GetSelectedRows(), tree->GetV1(), tree->GetV2(), tree->GetV3());
  return gr;
}
 
void ReZeroTGraph(TGraph *gr) {
  unsigned n = gr->GetN();
  double min = 0.;
  double x = 0;
  double y = 0;
  for (unsigned i = 0; i < n; ++i) {
    gr->GetPoint(i, x, y);
    if (y < min) min = y;
  }
  for (unsigned i = 0; i < n; ++i) {
    gr->GetPoint(i, x, y);
    y = y - min;
    gr->SetPoint(i, x, y);
  }
}
 
//----------------------------------------------------------------------------
// Plot decoration (colours, styles, text and legends)
//----------------------------------------------------------------------------
int CreateTransparentColor(int color, float alpha) {
  TColor* adapt = gROOT->GetColor(color);
  int new_idx = gROOT->GetListOfColors()->GetSize() + 1;
  TColor* trans = new TColor(new_idx, adapt->GetRed(), adapt->GetGreen(),
                             adapt->GetBlue(), "", alpha);
  trans->SetName(Form("userColor%i", new_idx));
  return new_idx;
}
 
void SetupTwoPadSplitAsRatio(std::vector<TPad*> const& pads,
                      TString y_title, bool y_centered, float y_min,
                      float y_max) {
  TH1 *upper = GetAxisHist(pads[0]);
  TH1 *lower = GetAxisHist(pads[1]);
  SetupTwoPadSplitAsRatio(pads, upper, lower, y_title, y_centered, y_min,
                          y_max);
}
 
void SetupTwoPadSplitAsRatio(std::vector<TPad*> const& pads, TH1* upper,
                             TH1* lower, TString y_title, bool y_centered,
                             float y_min, float y_max) {
  upper->GetXaxis()->SetTitle("");
  upper->GetXaxis()->SetLabelSize(0);
  double upper_h = 1. - pads[0]->GetTopMargin() - pads[0]->GetBottomMargin();
  double lower_h = 1. - pads[1]->GetTopMargin() - pads[1]->GetBottomMargin();
  lower->GetYaxis()->SetTickLength(gStyle->GetTickLength() *
                                   (upper_h / lower_h));
  pads[1]->SetTickx(1);
  pads[1]->SetTicky(1);
  lower->GetYaxis()->SetTitle(y_title);
  lower->GetYaxis()->CenterTitle(y_centered);
  if (y_max > y_min) {
    lower->SetMinimum(y_min);
    lower->SetMaximum(y_max);
  }
}
 
void StandardAxes(TAxis* xaxis, TAxis* yaxis, TString var, TString units) {
  double width = xaxis->GetBinWidth(1);
  TString w_label = TString::Format("%.1f", width);
  if (units == "") {
    xaxis->SetTitle(var);
    yaxis->SetTitle("Events / " + w_label);
  } else {
    xaxis->SetTitle(var + " (" + units + ")");
    yaxis->SetTitle("Events / " + w_label + " " + units);
  }
}
 
 
void UnitAxes(TAxis* xaxis, TAxis* yaxis, TString var, TString units) {
  xaxis->SetTitle(var + " (" + units + ")");
  yaxis->SetTitle("dN/d"+ var + " (1/" + units + ")");
}
 
TLegend* PositionedLegend(double width, double height, int pos, double offset) {
  double o = offset;
  double w = width;
  double h = height;
  double l = gPad->GetLeftMargin();
  double t = gPad->GetTopMargin();
  double b = gPad->GetBottomMargin();
  double r = gPad->GetRightMargin();
  TLegend* leg = 0;
  if (pos == 1) {
    leg = new TLegend(l + o, 1 - t - o - h, l + o + w, 1 - t - o, "", "NBNDC");
  }
  if (pos == 2) {
    double c = l + 0.5 * (1 - l - r);
    leg = new TLegend(c - 0.5 * w, 1 - t - o - h, c + 0.5 * w, 1 - t - o, "",
                      "NBNDC");
  }
  if (pos == 3) {
    leg = new TLegend(1 - r - o - w, 1 - t - o - h, 1 - r - o, 1 - t - o, "",
                      "NBNDC");
  }
  if (pos == 4) {
    leg = new TLegend(l + o, b + o, l + o + w, b + o + h, "", "NBNDC");
  }
  if (pos == 5) {
    double c = l + 0.5 * (1 - l - r);
    leg = new TLegend(c - 0.5 * w, b + o, c + 0.5 * w, b + o + h, "",
                      "NBNDC");
  }
  if (pos == 6) {
    leg = new TLegend(1 - r - o - w, b + o, 1 - r - o, b + o + h, "",
                      "NBNDC");
  }
  return leg;
}
 
TLine* DrawHorizontalLine(TPad *pad, TLine* line, double yval) {
  TH1 *axis = GetAxisHist(pad);
  std::cout << "pad: " << pad << "\n";
  double xmin = axis->GetXaxis()->GetXmin();
  double xmax = axis->GetXaxis()->GetXmax();
  return line->DrawLine(xmin, yval, xmax, yval);
}
 
void DrawTitle(TPad* pad, TString text, int align) {
  TVirtualPad *pad_backup = gPad;
  pad->cd();
  float t = pad->GetTopMargin();
  float l = pad->GetLeftMargin();
  float r = pad->GetRightMargin();
 
  float pad_ratio = (static_cast<float>(pad->GetWh()) * pad->GetAbsHNDC()) /
                    (static_cast<float>(pad->GetWw()) * pad->GetAbsWNDC());
  if (pad_ratio < 1.) pad_ratio = 1.;
 
  float textSize = 0.6;
  float textOffset = 0.2;
 
  TLatex latex;
  latex.SetNDC();
  latex.SetTextAngle(0);
  latex.SetTextColor(kBlack);
 
  latex.SetTextFont(42);
  latex.SetTextSize(textSize * t * pad_ratio);
  double y_off = 1 - t + textOffset * t;
  if (align == 1) latex.SetTextAlign(11);
  if (align == 1) latex.DrawLatex(l, y_off, text);
  if (align == 2) latex.SetTextAlign(21);
  if (align == 2) latex.DrawLatex(l + (1 - l - r) * 0.5, y_off, text);
  if (align == 3) latex.SetTextAlign(31);
  if (align == 3) latex.DrawLatex(1 - r, y_off, text);
  pad_backup->cd();
}
 
void DrawCMSLogo(TPad* pad, TString cmsText, TString extraText, int iPosX,
                 float relPosX, float relPosY, float relExtraDY) {
  TVirtualPad *pad_backup = gPad;
  pad->cd();
  float cmsTextFont = 61;  // default is helvetic-bold
 
  bool writeExtraText = extraText.Length() > 0;
  float extraTextFont = 52;  // default is helvetica-italics
 
  // text sizes and text offsets with respect to the top frame
  // in unit of the top margin size
  float lumiTextOffset = 0.2;
  float cmsTextSize = 0.8;
  // float cmsTextOffset    = 0.1;  // only used in outOfFrame version
 
  // ratio of "CMS" and extra text size
  float extraOverCmsTextSize = 0.76;
 
  TString lumi_13TeV = "20.1 fb^{-1}";
  TString lumi_8TeV = "19.7 fb^{-1}";
  TString lumi_7TeV = "5.1 fb^{-1}";
 
  bool outOfFrame = false;
  if (iPosX / 10 == 0) {
    outOfFrame = true;
  }
  int alignY_ = 3;
  int alignX_ = 2;
  if (iPosX / 10 == 0) alignX_ = 1;
  if (iPosX == 0) alignX_ = 1;
  if (iPosX == 0) alignY_ = 1;
  if (iPosX / 10 == 1) alignX_ = 1;
  if (iPosX / 10 == 2) alignX_ = 2;
  if (iPosX / 10 == 3) alignX_ = 3;
  if (iPosX == 0) relPosX = 0.14;
  int align_ = 10 * alignX_ + alignY_;
 
  float l = pad->GetLeftMargin();
  float t = pad->GetTopMargin();
  float r = pad->GetRightMargin();
  float b = pad->GetBottomMargin();
 
  TLatex latex;
  latex.SetNDC();
  latex.SetTextAngle(0);
  latex.SetTextColor(kBlack);
 
  float extraTextSize = extraOverCmsTextSize * cmsTextSize;
  float pad_ratio = (static_cast<float>(pad->GetWh()) * pad->GetAbsHNDC()) /
                    (static_cast<float>(pad->GetWw()) * pad->GetAbsWNDC());
  if (pad_ratio < 1.) pad_ratio = 1.;
 
 
  if (outOfFrame) {
    latex.SetTextFont(cmsTextFont);
    latex.SetTextAlign(11);
    latex.SetTextSize(cmsTextSize * t * pad_ratio);
    latex.DrawLatex(l, 1 - t + lumiTextOffset * t, cmsText);
  }
 
  float posX_ = 0;
  if (iPosX % 10 <= 1) {
    posX_ = l + relPosX * (1 - l - r);
  } else if (iPosX % 10 == 2) {
    posX_ = l + 0.5 * (1 - l - r);
  } else if (iPosX % 10 == 3) {
    posX_ = 1 - r - relPosX * (1 - l - r);
  }
  float posY_ = 1 - t - relPosY * (1 - t - b);
  if (!outOfFrame) {
    latex.SetTextFont(cmsTextFont);
    latex.SetTextSize(cmsTextSize * t * pad_ratio);
    latex.SetTextAlign(align_);
    latex.DrawLatex(posX_, posY_, cmsText);
    if (writeExtraText) {
      latex.SetTextFont(extraTextFont);
      latex.SetTextAlign(align_);
      latex.SetTextSize(extraTextSize * t * pad_ratio);
      latex.DrawLatex(posX_, posY_ - relExtraDY * cmsTextSize * t, extraText);
    }
  } else if (writeExtraText) {
    if (iPosX == 0) {
      posX_ = l + relPosX * (1 - l - r);
      posY_ = 1 - t + lumiTextOffset * t;
    }
    latex.SetTextFont(extraTextFont);
    latex.SetTextSize(extraTextSize * t * pad_ratio);
    latex.SetTextAlign(align_);
    latex.DrawLatex(posX_, posY_, extraText);
  }
  pad_backup->cd();
}
 
void DrawCMSLogo(TPad* pad, TString cmsText, TString extraText, int iPosX) {
  DrawCMSLogo(pad, cmsText, extraText, iPosX, 0.045, 0.035, 1.2);
}
 
//----------------------------------------------------------------------------
// Axis adjustment and overlap checks
//----------------------------------------------------------------------------
double GetPadYMax(TPad* pad, double x_min, double x_max) {
  TList *pad_obs = pad->GetListOfPrimitives();
  if (!pad_obs) return 0.;
  TIter next(pad_obs);
  double h_max = -99999;
  TObject *obj;
  while ((obj = next())) {
    if (obj->InheritsFrom(TH1::Class())) {
      TH1 *hobj = reinterpret_cast<TH1*>(obj);
      for (int j = 1; j < hobj->GetNbinsX(); ++j) {
        if (hobj->GetBinLowEdge(j) + hobj->GetBinWidth(j) < x_min ||
            hobj->GetBinLowEdge(j) > x_max)
          continue;
        if (hobj->GetBinContent(j) + hobj->GetBinError(j) > h_max) {
          h_max = hobj->GetBinContent(j) + hobj->GetBinError(j);
        }
      }
    }
    if (obj->InheritsFrom(TGraph::Class())) {
      TGraph *gobj = reinterpret_cast<TGraph*>(obj);
      unsigned n = gobj->GetN();
      double x = 0;
      double y = 0;
      for (unsigned k = 0; k < n; ++k) {
        gobj->GetPoint(k, x, y);
        if (x < x_min || x > x_max) continue;
        if (y > h_max) h_max = y;
      }
    }
  }
  return h_max;
}
 
double GetPadYMax(TPad* pad) {
  TList *pad_obs = pad->GetListOfPrimitives();
  if (!pad_obs) return 0.;
  double xmin = GetAxisHist(pad)->GetXaxis()->GetXmin();
  double xmax = GetAxisHist(pad)->GetXaxis()->GetXmax();
  return GetPadYMax(pad, xmin, xmax);
}
 
void FixTopRange(TPad *pad, double fix_y, double fraction) {
  TH1* hobj = GetAxisHist(pad);
  double ymin = hobj->GetMinimum();
  hobj->SetMaximum((fix_y - fraction * ymin) / (1. - fraction));
  if (gPad->GetLogy()) {
    if (ymin == 0.) {
      std::cout
          << "Can't adjust log-scale y-axis range if the minimum is zero!\n";
      return;
    }
    double max =
        (std::log10(fix_y) - fraction * std::log10(ymin)) / (1 - fraction);
    max = std::pow(10, max);
    hobj->SetMaximum(max);
  }
}
 
void FixBoxPadding(TPad *pad, TBox *box, double frac) {
  // Get the bounds of the box - these are in the normalised
  // Pad co-ordinates.
  double p_x1 = box->GetX1();
  double p_x2 = box->GetX2();
  double p_y1 = box->GetY1();
  // double p_y2 = box->GetY2();
 
  // Convert to normalised co-ordinates in the frame
  double f_x1 = (p_x1 - pad->GetLeftMargin()) /
                (1. - pad->GetLeftMargin() - pad->GetRightMargin());
  double f_x2 = (p_x2 - pad->GetLeftMargin()) /
                (1. - pad->GetLeftMargin() - pad->GetRightMargin());
  double f_y1 = (p_y1 - pad->GetBottomMargin()) /
                (1. - pad->GetTopMargin() - pad->GetBottomMargin());
 
  // Extract histogram providing the frame and axes
  TH1 *hobj = GetAxisHist(pad);
 
  double xmin = hobj->GetBinLowEdge(hobj->GetXaxis()->GetFirst());
  double xmax = hobj->GetBinLowEdge(hobj->GetXaxis()->GetLast()+1);
  double ymin = hobj->GetMinimum();
  double ymax = hobj->GetMaximum();
 
  // Convert box bounds to x-axis values
  double a_x1 = xmin + (xmax - xmin) * f_x1;
  double a_x2 = xmin + (xmax - xmin) * f_x2;
 
  // Get the histogram maximum in this range, given as y-axis value
  double a_max_h = GetPadYMax(pad, a_x1, a_x2);
 
  // Convert this to a normalised frame value
  double f_max_h = (a_max_h - ymin) / (ymax - ymin);
  if (gPad->GetLogy()) {
    f_max_h = (log10(a_max_h) - log10(ymin)) / (log10(ymax) - log10(ymin));
  }
 
  if (f_y1 - f_max_h < frac) {
    double f_target = 1. - (f_y1 - frac);
    FixTopRange(pad, a_max_h, f_target);
  }
}
 
void FixOverlay() { gPad->GetFrame()->Draw(); gPad->RedrawAxis(); }
 
#endif