canMan.py

Go to the documentation of this file.

from ROOT import *
import sys, os, math


def drange(start, stop, step=1):
  r = float(start)
  while r < stop:
    yield r
    r += step
        

cans = []
canNum = 0 
canPositions = [[10, 10, 600, 450],[615, 10, 600, 450], [10, 480, 600, 450], [615, 480, 600, 450]]




def newCan(title):
        global canNum
        global cans
        global canPositions     

        cans.append(TCanvas("can" + str(canNum), title))
        cans[canNum].SetWindowPosition(canPositions[canNum % 4][0], canPositions[canNum % 4][1])
        cans[canNum].SetWindowSize(canPositions[canNum % 4][2], canPositions[canNum % 4][3])
        canNum += 1
        if canNum % 4 == 0:
                posUp()
        return cans[canNum-1]


def canUp():
        for i in range(0, len(cans)):
                cans[i].RedrawAxis()
                cans[i].Update()


def saveCans(prefix = "", suffix=".pdf"):
        for i in range(0, len(cans)):
                cans[i].SaveAs(prefix + cans[i].GetTitle() + suffix)

def posUp():
        global canPositions
        for i in range(0,len(canPositions)):
                canPositions[i][0] += 30
                canPositions[i][1] += 30

def getTVectorFromList(list):
  vec = TVectorD(len(list))
  for i in range(0, len(list)):
    vec[i] = list[i]
  return vec




overlayNum = 0
legs = []

def drawOverlay(listOfHists, name="", opt="hist", varOpts=[""], colors=[kRed, kBlue -4, kGreen + 2, kMagenta + 1,kOrange+2, kBlack, kViolet - 7], lineStyles=[1], \
legEntries = [], logY = False, axisMin = 0, legPos = "right", legTextSize=0, inputCan=None, legHeight=0.15, legWidth=0.2, supressZero=False,\
padFactor=1.1, lineWidths=[gStyle.GetHistLineWidth()], fillStyles=[0], preserveStyle=False, 
yOffset=1.05):
  global overlayNum
  global legs
  if logY:
    supressZero = True
  
  histMax = 0 
  histMin = sys.float_info.max
  for hist in listOfHists:
    hist.GetYaxis().SetTitleOffset(yOffset)
    if hist.GetBinContent(hist.GetMaximumBin()) > histMax:
      histMax = hist.GetBinContent(hist.GetMaximumBin())
    if supressZero:
      for i in range(1, hist.GetNbinsX()+ 1):
        content = hist.GetBinContent(i)
        if content >0 and content < histMin:
            histMin = content
  if supressZero:
    axisMin = 0.9 * histMin
  if not inputCan:
    if name == "":
      can = newCan("Overlay" + str(overlayNum) )
    else:
      can = newCan(name + "Overlay")
  else:
    can = inputCan
    can.cd()

  counter  = 0
  
  if logY:
    can.SetLogy()
  if not preserveStyle:
    for hist in listOfHists:
      hist.SetLineColor(colors[counter % len(colors)])
      hist.SetLineWidth(lineWidths[counter % len(lineWidths)])
      hist.SetLineStyle(lineStyles[counter % len(lineStyles)])
      if fillStyles[counter % len(fillStyles)]:
        hist.SetFillStyle(fillStyles[counter % len(fillStyles)])
        hist.SetFillColor(colors[counter % len(colors)])
      counter += 1
  counter  = 0
  for hist in listOfHists:
    hist.GetYaxis().SetRangeUser(axisMin, padFactor * histMax)
    hist.Draw(opt + varOpts[counter % len(varOpts)])
    line = TLine(0,1,10,0)
    line.SetLineWidth(3)
    line.SetLineColor(kGray)
    line.SetLineStyle(2)  
   # line.DrawLineNDC(2,pow(listOfHists[0].GetMaximum(),0.58),10,pow(listOfHists[0].GetMaximum(),0.58))
    if hist.GetSize() > 40:
      line.DrawLineNDC(0.863,0.1,0.863,0.9)
    else:
      line.DrawLineNDC(0.72,0.1,0.72,0.9)
    line.Draw("same")

    counter += 1    
    if not "SAME" in opt:
      opt += "SAME"
    
  can.Update()
    
  if legPos == "right":
    legs.append(TLegend(0.65, 0.6, 0.65 + legWidth, 0.6 + legHeight))
  elif legPos == "mid":
    legs.append(TLegend(0.4, 0.6, 0.4 + legWidth, 0.6 + legHeight))
  elif legPos =="auto":
    legs.append(autoPlaceLegend(legWidth, legHeight))
  elif legPos == "left":
    legs.append(TLegend(0.15, 0.6, 0.15 + legWidth, 0.6 + legHeight))
  elif type(legPos) == type([]) and len(legPos) == 2:
    legs.append(TLegend(legPos[0], legPos[1], legPos[0] + legWidth, legPos[1] + legHeight))  
  else:
    print "Legend position", legPos, "is invalid."

  legs[overlayNum].SetFillColor(0)
  if legTextSize:
    legs[overlayNum].SetTextSize(legTextSize)
  counter = 0
    
  for hist in listOfHists:
    legStyle=''
    if hist.GetFillStyle != 0:
      legStyle += 'f'
    if hist.GetLineWidth != 0:
      legStyle += 'l'
    if len(legEntries) == 0:
      legs[overlayNum].AddEntry(hist, hist.GetName(), legStyle)
    else:
      legs[overlayNum].AddEntry(hist, legEntries[counter], legStyle)
    counter += 1


  legs[overlayNum].Draw()
  overlayNum += 1
  listOfHists[0].Draw("sameaxis")
  return can
  
def drawHists(listOfHists, opt= "", logY=False):
  for hist in listOfHists:
    c = newCan(hist.GetName())
    if logY:
      c.SetLogy()
    hist.Draw(opt)
    
def drawHist(hist, opt= "", logY=False):
  c = newCan(hist.GetName())
  if logY:
    c.SetLogy()
  hist.Draw(opt)
  return c

def getFileId(filePath):
  fileName = os.path.split(filePath)[1]
  extPos = fileName.find(".")
  fileId = fileName[0:extPos]
  return fileId
  
class ProgressTracker:
  def __init__(self, tree, increment=10000):
    self.entries = float(tree.GetEntries())
    self.increment = increment
    self.progress = 0
    
    
  def advance(self):
    self.progress+=1
    if not self.progress%self.increment:
      print "Event: ",  self.progress, str(self.progress / self.entries)[:6]


def sqr(x):
  return x*x
def sign(x):
  if x >= 0:
    return 1
  return -1
  
def PointDistanceToBox(x, y, x0, y0, x1, y1):
  # Inside
  if x > x0 and x < x1 and y > y0 and y < y1:
    return 0

  # Corners
  d = sqr(x-x0)+sqr(y-y0);
  d = min(d, sqr(x-x1)+sqr(y-y0))
  d = min(d, sqr(x-x1)+sqr(y-y1))
  d = min(d, sqr(x-x0)+sqr(y-y1))

  # Top and bottom edges
  if x > x0 and x < x1:
    d = min(d, sqr(y-y0))
    d = min(d, sqr(y-y1))
  
  # Left and right
  if y > y0 and y < y1: 
    d = min(d, sqr(x-x0))
    d = min(d, sqr(x-x1))
  

  return d;



def autoPlaceLegend(dx, dy, gridPoints=30):
  gPad.Update()
  # Convert requested width and height into physics coordinates
  dx *= (gPad.GetX2()-gPad.GetX1())
  dy *= (gPad.GetY2()-gPad.GetY1())

  # Range of axes in physics units, if log, these already have the log applied
  x0 = gPad.GetUxmin()
  x1 = gPad.GetUxmax()
  y0 = gPad.GetUymin()
  y1 = gPad.GetUymax()

   
  X = x1-x0
  Y = y1-y0

  bestd = 0
  bestx = 0
  besty = 0
  for fallback in [False, True]:  
    if fallback:
      print "falling back"
    for x in drange(x0+dx/2, x1-dx/2, X/float(gridPoints)):
      for y in drange(y0+dy/2, y1-dy/2, Y/float(gridPoints)):
        d = 999999
          
          
        # Repel from edges
        """                

        dToEdges = [sqr((x-dx/2-x0)/X), sqr((x+dx/2-x1)/X), sqr((y-dy/2-y0)/Y), sqr((y+dy/2-y1)/Y)]
        for dToEdge in dToEdges:
          print dToEdge
          if math.fabs(dToEdge) < 0.05:
            d = 0
            continue
        
        """

        d = min(d, sqr((x-dx/2-x0)/X))
        d = min(d, sqr((x+dx/2-x1)/X))
        d = min(d, sqr((y-dy/2-y0)/Y))
        d = min(d, sqr((y+dy/2-y1)/Y))
        if d < bestd:
          continue

        prims = gPad.GetListOfPrimitives() 
        for obj in prims:

          if not obj.InheritsFrom(TH1.Class()):
            continue
          h = obj

          for n in range(0, h.GetNbinsX() + 1):
            px = h.GetBinCenter(n)
            py = h.GetBinContent(n)
            if gPad.GetLogy():
              if py > 0:
                py = math.log10(py)
              else:
                py = gPad.GetUymin()
            if fallback:
              d = min(d, sqr((px-x)/X)+sqr((py-y)/Y))
            else:
              
              d = min(d, PointDistanceToBox(px/X, py/Y,
                                                 (x-dx/2)/X, (y-dy/2)/Y,
                                                 (x+dx/2)/X, (y+dy/2)/Y))
            
            if d < bestd:
              break;
          
          if d < bestd:
            break;
        

        if d > bestd:
          bestd = d;
          bestx = x;
          besty = y;
        
    
    

    if bestd != 0:
      break # If we always collide, have to do fallback

  # Convert to pad coordinates
  nx = float(bestx-gPad.GetX1())/(gPad.GetX2()-gPad.GetX1());
  ny = float(besty-gPad.GetY1())/(gPad.GetY2()-gPad.GetY1());

  ndx = dx/float(gPad.GetX2()-gPad.GetX1());
  ndy = dy/float(gPad.GetY2()-gPad.GetY1());

  return TLegend(nx-ndx/2, ny-ndy/2, nx+ndx/2, ny+ndy/2);



class Table():
  def __init__(self):
    self.table = []
    
  def append(self, row):
    self.table.append(row)
    
  def getLatex(self):
    str = ""
    table = self.table
    nCol = len(table[0])
    colDef = "| " + (" c |" * nCol)
    str +=   "\\begin{center} \n \\begin{tabular} {" + colDef + "} \n \hline \n"

    col_width = [max(len(x) for x in col) for col in zip(*table)]
    for line in table:
      str += " " + " & ".join("{:{}}".format(x, col_width[i])
                                for i, x in enumerate(line)) + " \\\\ \hline \n"
    str += "  \\end{tabular} \n \\end{center}"
    return str

  def getReadable(self):
    str = ""
    table = self.table
    col_width = [max(len(x) for x in col) for col in zip(*table)]
    for line in table:
        str += "| " + " | ".join("{:{}}".format(x, col_width[i])
                                for i, x in enumerate(line)) + " | \n"
  
                            
    return str
  
  def printLatex(self):
    print self.getLatex()

  def printReadable(self):
    print self.getReadable()

  def writeLatex(self, outputFile):
    outputFile.write(self.getReadable())

  def writeLatex(self, outputFile):
    outputFile.write(self.getLatex())