file_size_ana.py

Go to the documentation of this file.

#!/bin/env python

import argparse
import sys, os
import re
from math import *
import code

parser = argparse.ArgumentParser()

parser.add_argument('-t', '--text', action = 'store_true',
                    help = 'textual output')
parser.add_argument('-r', '--radial', action = 'store_true',
                    help = 'pie-chart-style output')
parser.add_argument('-l', '--linear', action = 'store_true',
                    help = 'bar-chart-style output')
parser.add_argument('-j', '--json', action = 'store_true',
                    help = 'json output')
parser.add_argument('-b', '--batch', action = 'store_true',
                    help = 'don\'t open windows for graphics')

parser.add_argument('-f', '--focus', default='', metavar='BR', help = 'center this branch')

parser.add_argument('filename.root',
                    help = 'the CAF or ART file to analyze')

opts = vars(parser.parse_args())

if not (opts['text'] or opts['radial'] or opts ['linear'] or opts['json']):
    print 'You must specify at least one of --text, --radial, --linear, or --json'
    exit()

# ROOT seems unhappy about having its arguments messed with. Import it too late
# for it to notice
from ROOT import *

class Node:
    def __init__(self, t, s):
        self.title = t
        self.size = s
        self.children = []

    def __str__(self):
        return self.title + ':\t' + str(self.FullSize())

    def __repr__(self):
        return '(' + self.title + ', ' + str(self.FullSize()) + ')'

    def FullSize(self):
        return sum([c.FullSize() for c in self.children], self.size)

##### Parse the file #####

# First, figure out what sort of file it is
se = sys.stderr
sys.stderr = open(os.devnull, 'w') # swallow errors about missing dictionaries
f = TFile(opts['filename.root'],"READ")
sys.stderr = se # put it back

isArt = bool(f.Get('Events'))

if isArt:
    treeNames = ['Events', 'Runs', 'SubRuns',
                 'EventHistory', 'MetaData', 'Parentage',
                 'EventMetaData', 'SubRunMetaData', 'RunMetaData']
else:
    treeNames = ['recTree', 'spillTree']

    if not f.Get('recTree'):
        print 'This doesn\'t appear to be an ART or CAF file. Aborting.'
        exit(1)


nodes = {}

# Recurse through the tree and store the size on disk of every branch
def AddNodes(branch):
    title = branch.GetName()

    # In ART files you get type_label_instance_process.therest
    # Transform to label.instance.type.therest
    m = re.match('(.*)_(.*)_(.*)_(.*)\\.(.*)', title)
    if m:
        # Often instance is blank
        if len(m.group(3)) > 0:
            title = m.group(2)+'.'+m.group(3)+'.'+m.group(1)+'.'+m.group(5)
        else:
            title = m.group(2)+'.'+m.group(1)+'.'+m.group(5)

    label = title.split('.')[-1]
    nodes[title] = Node(label, branch.GetZipBytes())
    for b in branch.GetListOfBranches(): AddNodes(b)


for n in treeNames:
    for b in f.Get(n).GetListOfBranches(): AddNodes(b)


##### Insert implicit nodes #####

progress = True
while progress:
    progress = False
    newnodes = {}

    for key in nodes:
        node = nodes[key]

        parentKey = '.'.join(key.split('.')[:-1])

        if parentKey.endswith('.obj'): parentKey = parentKey[:-4]

        if parentKey not in nodes:
            label = parentKey.split('.')[-1]
            newnodes[parentKey] = Node(label, 0)
            progress = True

    nodes.update(newnodes)


##### Link up the nodes tree #####

root = nodes[opts['focus']]

if opts['focus'] == '':
    if isArt:
        root.title = 'Events'
    else:
        root.title = 'rec'

for key in nodes:
    node = nodes[key]

    parentKey = '.'.join(key.split('.')[:-1])

    if parentKey.endswith('.obj'): parentKey = parentKey[:-4]

    parent = nodes[parentKey]
    if parent is not node: # special case root node
        parent.children.append(node)


##### Display the node tree #####

curNodes = [root]

def prettyPrint(node, total, depth = 100, indent = ''):
    if depth == 0: return
    if node.FullSize() == 0: return
    percent = int((100*node.FullSize())/total+.5)
    if percent > 0:
        print indent, node, '\t('+str(percent)+'%)'
    else:
        print indent, node

    for c in sorted(node.children, key = lambda n: -n.FullSize()):
        prettyPrint(c, node.FullSize(), depth-1, indent+'\t')



if opts['text']:
    # Empirically this is the maximum depth
    for depth in range(1, 7):
        prettyPrint(root, root.FullSize(), depth)
        print
        print



def toJSON(node, indent = ''):
    ret = ''

    if len(node.children) == 0:
        ret += indent+'{"name": "'+node.title+'", "size": '+str(node.FullSize())+'}'
    else:
        ret += indent+'{"name": "'+node.title+'", "children": [\n'

        ret += ',\n'.join([toJSON(c, indent+'    ') for c in sorted(node.children, key = lambda n: -n.FullSize()) if c.FullSize() > 0])
        ret += ' ] }'

    return ret


if opts['json']:
    print toJSON(root)




gs = []
# When ROOT objects get garbage-collected they disappear from the screen
# again. Need to leak them like in C++
def New(cons, *args):
    ret = cons(*args)
    gs.append(ret)
    return ret


def RingChart(node, total, depth = 1, startang = 0):
    span = 360.*node.FullSize()/total
    meanang = (startang+span/2)*3.14159/180

    if span < .5: return

    arc = New(TArc, 200, 200, 20*depth, startang, startang+span)
    arc.SetLineWidth(2)

    for c in sorted(node.children, key = lambda n: -n.FullSize()):
        RingChart(c, total, depth+1, startang)
        startang += (360.*c.FullSize())/total

    arc.Draw()

    meanr = 20*depth-10
    if depth == 1: meanr = 0
    if span > 10:
        label = New(TLatex, 200+meanr*cos(meanang),
                            200+meanr*sin(meanang),
                            node.title)
        textscale = .06*span/len(node.title)
        if textscale > 1: textscale = 1
        label.SetTextSize(label.GetTextSize()*textscale)
        label.SetTextAlign(22)
        textang = meanang*180/3.14159265358979-90
        if textang > +90: textang -= 180
        if textang < -90: textang += 180
        if depth > 1: label.SetTextAngle(textang)
        label.Draw()
    elif span > 1:
        label = New(TLatex, 200+meanr*cos(meanang),
                            200+meanr*sin(meanang),
                            node.title)
        textscale = .06*span
        if textscale > 1: textscale = 1
        label.SetTextSize(label.GetTextSize()*textscale)
        label.SetTextAlign(22)
        textang = meanang*180/3.14159265358979
        while textang > +90: textang -= 180
        while textang < -90: textang += 180
        if depth > 1: label.SetTextAngle(textang)
        label.Draw()



def BarChart(node, total, depth = 0, startpos = 0):
    span = 100.*node.FullSize()/total
    meanpos = startpos+span/2
    meany = -(depth*20+10)

    if span < .3: return

    box = New(TBox, startpos, -depth*20, startpos+span, -(depth*20+20))
    box.SetLineWidth(2)
    box.SetFillStyle(0)

    for c in sorted(node.children, key = lambda n: -n.FullSize()):
        BarChart(c, total, depth+1, startpos)
        startpos += (100.*c.FullSize())/total

    box.Draw()

    if span > 20:
        label = New(TLatex, meanpos, meany, node.title)
        textscale = .4*min(20, span)/len(node.title)
        if textscale > 2: textscale = 2
        label.SetTextSize(label.GetTextSize()*textscale)
        label.SetTextAlign(22)
        label.Draw()
    elif span > 1:
        label = New(TLatex, meanpos, meany, node.title)
        textscale = .4*min(20, span)/len(node.title)
        if textscale > 1: textscale = 1
        label.SetTextSize(label.GetTextSize()*textscale)
        label.SetTextAlign(22)
        label.SetTextAngle(90)
        label.Draw()


if opts['radial']:
    canv = New(TCanvas, 'rings', 'rings', 1000, 1000)

    axes = New(TH2F, "", "", 100, 70, 330, 100, 70, 330)
    axes.Draw()

    axes.GetXaxis().SetTickLength(0)
    axes.GetYaxis().SetTickLength(0)
    gPad.SetLeftMargin(0)
    gPad.SetRightMargin(0)
    gPad.SetBottomMargin(0)
    gPad.SetTopMargin(0)
    gStyle.SetOptStat(0)

    RingChart(root, root.FullSize())

    gPad.Print("rings.png")
    gPad.Print("rings.eps")


if opts['linear']:
    canv = New(TCanvas, 'bars', 'bars', 1000, 1000)

    axes = New(TH2F, "", "", 100, -10, 110, 100, -110, +10)
    axes.Draw()

    axes.GetXaxis().SetTickLength(0)
    axes.GetYaxis().SetTickLength(0)
    gPad.SetLeftMargin(0)
    gPad.SetRightMargin(0)
    gPad.SetBottomMargin(0)
    gPad.SetTopMargin(0)
    gStyle.SetOptStat(0)

    BarChart(root, root.FullSize())

    gPad.Print("bars.png")
    gPad.Print("bars.eps")


if (opts['radial'] or opts['linear']) and not opts['batch']:
    # Just want the console to hang around so the user can look at the plots
    # without them disappearing. Ctrl-D or "exit()" to quit.
    print 'Ctrl-D to quit'
    code.interact(local = locals(), banner='')