ppfx_plot_utils.h

Go to the documentation of this file.

#pragma once

#include "TH1.h"
#include "TH2.h"
#include "TLegend.h"
#include "TLine.h"
#include "TText.h"
#include "TLatex.h"
#include "TLegend.h"
#include "TPad.h"

#include <fstream>
#include <iostream>
#include <cmath>
#include <algorithm>

void set_labels_1D(std::vector<std::string> labels, TH1D* hist)
{
  int nlabels = labels.size();
  
  double left = gPad->GetLeftMargin();
  double right = 1-gPad->GetRightMargin();
  double top = 1-gPad->GetTopMargin();
  double bottom = gPad->GetBottomMargin();
  
  TLatex *tLabel = new TLatex();
  if(nlabels == 8)
    tLabel->SetTextSize(0.05);
  if(nlabels == 16)
    tLabel->SetTextSize(0.0125);

  tLabel->SetNDC();
  
  double start = 0;
  int nbins = hist->GetNbinsX();

  double hxmin = hist->GetXaxis()->GetBinLowEdge(1);
  double hxmax = hist->GetXaxis()->GetBinLowEdge(nbins+1);
  start = left + 0.25*(right-left)*nbins/(nlabels*(hxmax-hxmin));
  
  std::for_each(labels.begin(), labels.end(), 
                 [&](std::string label){
                    tLabel->DrawLatex(start, 0.25*bottom, label.c_str()); 
                    start += (right-left)*nbins/(nlabels*(hxmax-hxmin));
               });

  gPad->Update();
}  

void set_labels_2D(std::vector<std::string> labels_x, std::vector<std::string> labels_y, TH2D* hist)
{
  int nlabels_x = labels_x.size();
  int nlabels_y = labels_y.size();
  
  double left = gPad->GetLeftMargin();
  double right = 1-gPad->GetRightMargin();
  double top = 1-gPad->GetTopMargin();
  double bottom = gPad->GetBottomMargin();
  
  TLatex *tLabelX = new TLatex();
  if(nlabels_x == 8)
    tLabelX->SetTextSize(0.02);
  if(nlabels_x == 16)
    tLabelX->SetTextSize(0.0125);
  tLabelX->SetNDC();
  
  TLatex *tLabelY = new TLatex();
  if(nlabels_y == 8)
    tLabelY->SetTextSize(0.0125);
  if(nlabels_y == 16)
    tLabelY->SetTextSize(0.01);

  tLabelY->SetNDC();
  tLabelY->SetTextAngle(90);
  
  double start_x = 0;
  double start_y = 0;
  int nbins_x = hist->GetNbinsX();
  int nbins_y = hist->GetNbinsY();

  double hxmin = hist->GetXaxis()->GetBinLowEdge(1);
  double hxmax = hist->GetXaxis()->GetBinLowEdge(nbins_x+1);
  
  double hymin = hist->GetYaxis()->GetBinLowEdge(1);
  double hymax = hist->GetYaxis()->GetBinLowEdge(nbins_y+1);
  
  start_x = left + 0.25*(right-left)*nbins_x/(nlabels_y*(hxmax-hxmin));
  start_y = bottom + 0.25*(top-bottom)*nbins_y/(nlabels_y*(hymax-hymin));
  
  std::for_each(labels_x.begin(), labels_x.end(), 
                  [&](std::string label){
                    tLabelX->DrawLatex(start_x, 0.25*bottom, label.c_str()); 
                    start_x += (right-left)*nbins_x/(nlabels_x*(hxmax-hxmin));
                });
  
  std::for_each(labels_y.begin(), labels_y.end(), 
                  [&](std::string label){
                    tLabelY->DrawLatex(0.25*left, start_y, label.c_str());
                    start_y += (top-bottom)*nbins_y/(nlabels_y*(hymax-hymin));
                });
  
  gPad->Update();
}  

void draw_vlines(int n, double xmax, double y0, double y1, Color_t color = kBlack)
{
  TLine *line = new TLine();
  line->SetLineWidth(2);
  line->SetLineColor(color);
  double x = 0;
  for(int i = 0; i < n; i++) {
    line->DrawLine(x, y0, x, y1);
    x += (double)xmax/n;
  }
  gPad->Update();
}

struct plot_h
{
  TH1* hist;
  Option_t* drawOption;
  Color_t color;
  bool drawLegend;
  std::string legendLabel;
  Option_t* legendOption;
};

TLegend* auto_legend(std::vector <TH1*> hists, double dx, double dy, double hymin, double hymax)
{
  //get pad coordinates first
  double left = gPad->GetLeftMargin();
  double right = 1-gPad->GetRightMargin();
  double top = 1-gPad->GetTopMargin();
  double bottom = gPad->GetBottomMargin();

  // start at centre
  double x = 0.5*(right + left) - 0.5*dx;
  double y = 0.5*(top + bottom) - 0.5*dy;

  for(int i = 0; i < 1000; i++){
    double overlap = 0.;
    double doverlap_x = 0.;
    double doverlap_y = 0.;

    for(int histIdx = 0; histIdx < (int)hists.size(); histIdx++){
      double hxmin = hists[histIdx]->GetXaxis()->GetBinLowEdge(1);
      double hxmax = hists[histIdx]->GetXaxis()->GetBinLowEdge(hists[histIdx]->GetNbinsX()+1);

      for(int binIdx = 1; binIdx < hists[histIdx]->GetNbinsX()+1; ++binIdx){

        // convert histogram coordinates to pad coordinates
        double bx = (hists[histIdx]->GetXaxis()->GetBinWidth(binIdx))*(right - left)/(hxmax - hxmin);
        double hx = ((hists[histIdx]->GetXaxis()->GetBinLowEdge(binIdx) - hxmin)*(right - left)/(hxmax - hxmin)) + left;
        double hy = (hists[histIdx]->GetBinContent(binIdx) - hymin)*(top - bottom)/(hymax - hymin) + bottom;
        double hy_prev = (hists[histIdx]->GetBinContent(binIdx-1))*(top - bottom)/(hymax - hymin) + bottom;
        double by = std::abs(hy - hy_prev);
        double low_hy = hy;

        // calculate overlap
        double overlap_x = (dx + bx) - (std::max(x+dx, hx+bx) - std::min(x, hx));
        double overlap_y = 0.;
        if (hy > hy_prev){
          overlap_y = (dy + by) - (std::max(y+dy, hy) - std::min(y, hy_prev));
          low_hy = hy_prev;
        }
        else {
          overlap_y = (dy + by) - (std::max(y+dy, hy_prev) - std::min(y, hy));
        }
        
        // negative overlap is as good as 0
        if (overlap_x < 0.) overlap_x = 0.;
        if (overlap_y < 0.) overlap_y = 0.;

        // gradient
        if(overlap_x == bx || overlap_x == dx || overlap_x == 0.)
          doverlap_x += 0.;
        else if(x+dx >= hx + bx)
          doverlap_x += -1.*overlap_y;
        else
          doverlap_x += overlap_y;

        if(overlap_y == by || overlap_y == dy || overlap_y == 0.)
          doverlap_y += 0.;
        else if(y+dy >= low_hy + by)
          doverlap_y += -1.*overlap_x;
        else
          doverlap_y += overlap_x;
        
        overlap += overlap_x*overlap_y;
      }
    }

    // descend along gradient
    double xprev = x;
    double yprev = y;
    x -= 0.02*doverlap_x;
    y -= 0.02*doverlap_y;

    // if it hits an edge, throw it to the other end with offset to prevent bouncing back
    if(x+dx >= right){
      x = right - x + 0.1;
      continue;
    }
    if(x <= left) {
      x = right - x - dx - 0.1;
      continue;
    }
    if(y+dy >= top) {
      y = top - y + 0.1;
      continue;
    }
    if(y <= bottom) {
      y = top - y - dy - 0.1;
      continue;
    }
    // if it has found a space or converges
    if(overlap == 0. || (std::abs(x - xprev) < 0.00005 && std::abs(y - yprev) < 0.00005)) break;

  }
  return new TLegend(x, y, x+dx, y+dy);

}

void plot_hists(std::vector <plot_h> hists, std::string title="", 
                     bool turnOnlegend = true, bool turnOnMaxRange = true, 
                     double leg_w = 0.2, double leg_h = 0.3,
                     double offset = 1.5)
{      

  double yMax = 0;
  std::vector <TH1*> h_leg;

  for(int histIdx = 0; histIdx<(int)hists.size(); ++histIdx){

    hists[histIdx].hist->SetLineColor(hists[histIdx].color);
    double tmp = hists[histIdx].hist->GetBinContent(hists[histIdx].hist->GetMaximumBin()) 
                 + std::sqrt(hists[histIdx].hist->GetBinContent(hists[histIdx].hist->GetMaximumBin()));
    tmp *= offset;
    if(yMax < tmp) yMax = tmp;

    h_leg.push_back(hists[histIdx].hist);
  }
  
  hists[0].hist->SetTitle(title.c_str());

  if(turnOnMaxRange) 
    hists[0].hist->GetYaxis()->SetRangeUser(0, yMax);
  
  
  for(int histIdx = 0; histIdx<(int)hists.size(); ++histIdx)
    hists[histIdx].hist->Draw(hists[histIdx].drawOption);

  if(turnOnlegend) { 
    TLegend* legnd = auto_legend(h_leg, leg_w, leg_h, 0., yMax);
    for(int histIdx = 0; histIdx<(int)hists.size(); ++histIdx){
      if(hists[histIdx].drawLegend) 
        legnd->AddEntry(hists[histIdx].hist,hists[histIdx].legendLabel.c_str(),hists[histIdx].legendOption);
    }
    legnd->Draw(); 
  }

  gPad->Update();     
      
}

TH1* plot_ratios(std::vector <std::pair<plot_h, TH1*>> hists, std::string title = "", 
                 bool turnOnlegend = false, double low_y = -1, double high_y = -1)
{
  TH1* ret = 0;
  bool turnOnMaxRange = false;
  std::vector <plot_h> ratios;
  if(low_y == high_y)  turnOnMaxRange = true;
  for(int histIdx = 0; histIdx < (int) hists.size(); histIdx++){

    TH1F* hRatio = (TH1F*)hists[histIdx].first.hist->Clone("");
    hRatio->Sumw2();
    hRatio->Divide(hists[histIdx].second);
    hRatio->GetYaxis()->SetTitle("Ratio");
    if(!turnOnMaxRange) hRatio->GetYaxis()->SetRangeUser(low_y, high_y);

    ratios.push_back({hRatio, hists[histIdx].first.drawOption, hists[histIdx].first.color, 
                      hists[histIdx].first.drawLegend, hists[histIdx].first.legendLabel, hists[histIdx].first.legendOption});
  
  }
  
  plot_hists(ratios, title, turnOnlegend, turnOnMaxRange);
  
  // return histogram object to make modifications later
  ret = ratios[0].hist;
  return ret;
}

void plot_label(double x, double y, std::string label)
{
  TLatex *tLabel = new TLatex(x, y, label.c_str());
  tLabel->SetTextSize(0.035);
  tLabel->SetNDC();
  tLabel->Draw();

  gPad->Update();
}

void draw_h_line(double x0, double x1, double y, Color_t color)
{
  TLine *line1 = new TLine(x0,y,x1, y);
  line1->SetLineWidth(2);
  line1->SetLineStyle(kDashed);
  line1->SetLineColor(color);
  line1->Draw("same");

  gPad->Update();
}

void draw_v_line(double x, double y0, double y1, Color_t color)
{
  TLine *line1 = new TLine(x,y0,x, y1);
  line1->SetLineWidth(2);
  line1->SetLineStyle(kDashed);
  line1->SetLineColor(color);
  line1->Draw("same");

  gPad->Update();
}