python/statistics/counter.py

0001  # Copyright (C) 2014 Colin Bernet
0002 # https://github.com/cbernet/heppy/blob/master/LICENSE
0003
0004 from builtins import range
0005 import pickle
0006 from PhysicsTools.HeppyCore.utils.diclist import diclist
0007
0008 class Counter(diclist):
0009
0010     def __init__(self, name):
0011         self.name = name
0012         super(Counter, self).__init__()
0013
0014     def register(self, level):
0015         self.add( level, [level, 0] )
0016
0017     def inc(self, level, nentries=1):
0018         '''increment an existing level
0019         '''
0020         if level not in self.dico:
0021             raise ValueError('level', level, 'has not been registered')
0022         else:
0023             self[level][1] += nentries
0024
0025     def __add__(self, other):
0026         '''Add two counters (+).'''
0027         size = max( len(self), len(other))
0028         for i in range(0, size):
0029             if i>=len(other):
0030                 # this line exists only in this counter, leave it as is
0031                 continue
0032             elif i>=len(self):
0033                 self.register( other[i][0])
0034                 self.inc( other[i][0], other[i][1] )
0035             else:
0036                 if self[i][0] != other[i][0]:
0037                     err = ['cannot add these counters:', str(self), str(other)]
0038                     raise ValueError('\n'.join(err))
0039                 else:
0040                     self.inc( other[i][0], other[i][1] )
0041         return self
0042
0043     def __iadd__(self, other):
0044         '''Add two counters (+=).'''
0045         return self.__add__(other)
0046
0047     def write(self, dirname):
0048         '''Dump the counter to a pickle file and to a text file in dirname.'''
0049         pckfname = '{d}/{f}.pck'.format(d=dirname, f=self.name)
0050         pckfname = pckfname.replace('*','STAR')
0051         pckfile = open( pckfname, 'w' )
0052         pickle.dump(self, pckfile)
0053         txtfile = open( pckfname.replace('.pck', '.txt'), 'w')
0054         txtfile.write( str(self) )
0055         txtfile.write( '\n' )
0056         txtfile.close()
0057
0058     def __str__(self):
0059         retstr = 'Counter %s :\n' % self.name
0060         prev = None
0061         init = None
0062         for level, count in self:
0063             if prev == None:
0064                 prev = count
0065                 init = count
0066             if prev == 0:
0067                 eff1 = -1.
0068             else:
0069                 eff1 = float(count)/prev
0070             if init == 0:
0071                 eff2 = -1.
0072             else:
0073                 eff2 = float(count)/init
0074             retstr += '\t {level:<40} {count:>9} \t {eff1:4.2f} \t {eff2:6.4f}\n'.format(
0075                 level=level,
0076                 count=count,
0077                 eff1=eff1,
0078                 eff2=eff2 )
0079             prev = count
0080         return retstr
0081
0082
0083
0084
0085 class Counters(object):
0086     '''
0087     TODO: could be a diclist?
0088     '''
0089
0090     def __init__( self ):
0091         self.counters = []
0092         self.ranks = {}
0093
0094     def addCounter(self, name):
0095         self.ranks[ name ] = len( self.counters )
0096         self.counters.append( Counter(name) )
0097
0098     def counter(self, name):
0099         return self.counters[ self.ranks[name] ]
0100
0101     def write(self, dirname):
0102         map( lambda x: x.write(dirname), self.counters)
0103
0104     def __str__(self):
0105         prints = map( str, self.counters )
0106         return '\n'.join(prints)
0107
0108     def __getitem__(self, name):
0109         return self.counter(name)
0110