Created
October 16, 2015 07:12
-
-
Save blinkseb/9a1c497ec98aafbc57ab to your computer and use it in GitHub Desktop.
Small PyROOT script to print statistics about a TTree
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#! /bin/env python | |
""" | |
Print disk usage of a TTree | |
""" | |
from __future__ import division | |
import sys, array, os | |
import ROOT | |
from math import log | |
unit_list = zip(['bytes', 'kB', 'MB', 'GB', 'TB', 'PB'], [0, 0, 1, 2, 2, 2]) | |
def sizeof_fmt(num): | |
"""Human friendly file size""" | |
if num > 1: | |
exponent = min(int(log(num, 1024)), len(unit_list) - 1) | |
quotient = float(num) / 1024**exponent | |
unit, num_decimals = unit_list[exponent] | |
format_string = '{:.%sf} {}' % (num_decimals) | |
return format_string.format(quotient, unit) | |
if num == 0: | |
return '0 bytes' | |
if num == 1: | |
return '1 byte' | |
def printTreeStat(filesize, d, t): | |
print("Statistics for tree %s/%s" % (d.GetName() if d is not None else "", t.GetName())) | |
s = 0 | |
skey = 0 | |
if t.GetDirectory() is not None: | |
key = t.GetDirectory().GetKey(t.GetName()) | |
if key is not None: | |
skey = key.GetKeylen() | |
s = key.GetNbytes() | |
total = skey | |
if t.GetZipBytes() > 0: | |
total += t.GetTotBytes() | |
b = ROOT.TBufferFile(ROOT.TBuffer.kWrite, 10000) | |
ROOT.TTree.Class().WriteBuffer(b, t) | |
total += b.Length() | |
file = t.GetZipBytes() + s | |
cx = 1 | |
if t.GetZipBytes() > 0: | |
cx = (t.GetTotBytes() + 0.00001) / t.GetZipBytes() | |
percent = file / filesize * 100 | |
print("******************************************************************************") | |
print("*Tree :%-10s: %-54s *" % (t.GetName(), t.GetTitle())) | |
print("*Entries : %8d : Total = %10s File Size = %10s ~ %5.2f %% *" % (t.GetEntries(), sizeof_fmt(total), sizeof_fmt(file), percent)) | |
print("* : : Tree compression factor = %6.2f *" % cx) | |
print("******************************************************************************") | |
nl = t.GetListOfLeaves().GetEntries() | |
br = None | |
leaf = None | |
count = array.array('f', [0] * nl) | |
names = [] | |
keep = 0 | |
for l in xrange(0, nl): | |
leaf = t.GetListOfLeaves().At(l) | |
br = leaf.GetBranch() | |
if '.' in br.GetName() or br.GetMother() != br: | |
count[l] = -1 | |
count[keep] += br.GetZipBytes() | |
else: | |
keep = l | |
count[keep] = br.GetZipBytes() | |
names.append(br.GetName()) | |
percent_sum = 0 | |
for size, name in sorted(zip(count, names), reverse=True): | |
if size < 0: | |
continue | |
percent = size / file * 100 | |
percent_sum += percent | |
print("branch %-60s %9s ~%.2f %%" % (name, sizeof_fmt(size), percent)) | |
print("------------------") | |
print(" %-60s %9s ~%.2f %%" % (" ", " ", percent_sum)) | |
print("\n\n\n") | |
filesize = os.path.getsize(sys.argv[1]) | |
f = ROOT.TFile.Open(sys.argv[1], "read") | |
for key in f.GetListOfKeys(): | |
d = f.Get(key.GetName()) | |
if isinstance(d, ROOT.TTree): | |
printTreeStat(filesize, None, d) | |
elif isinstance(d, ROOT.TDirectory): | |
for treeKey in d.GetListOfKeys(): | |
tree = d.Get(treeKey.GetName()) | |
printTreeStat(filesize, d, tree) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment