python/statistics/average.py

0001 # Copyright (C) 2014 Colin Bernet
0002 # https://github.com/cbernet/heppy/blob/master/LICENSE
0003
0004 import math
0005 import pickle
0006 from PhysicsTools.HeppyCore.utils.diclist import diclist
0007
0008 class Average(object):
0009
0010     def __init__(self, name):
0011         self.name = name
0012         self.sumw = 0
0013         self.sumwx = 0
0014         self.sumwx2 = 0
0015
0016     def add(self, value, weight=1.0):
0017         """
0018         Add a new sample to the average.
0019         """
0020         value = float(value) # avoids surprising results with integers
0021         weight = float(weight)
0022         self.sumw += weight
0023         self.sumwx += weight * value
0024         self.sumwx2 += weight * value * value
0025
0026     def variance(self):
0027         return abs( self.sumwx2 / self.sumw - \
0028                     self.sumwx * self.sumwx / (self.sumw*self.sumw) )
0029
0030     def value(self):
0031         """
0032         Mean value
0033         """
0034         if self.sumw:
0035             return self.sumwx / self.sumw
0036         else:
0037             return None
0038
0039     def uncertainty(self):
0040         """
0041         Uncertainty on the mean value
0042         """
0043         if self.sumw:
0044             return math.sqrt( self.variance() ) / math.sqrt( self.sumw )
0045         else:
0046             return None
0047
0048     def average( self ):
0049         """
0050         Returns: mean value, uncertainty on mean value.
0051         """
0052         return self.value(), self.uncertainty()
0053
0054     def __add__(self, other):
0055         '''Add two averages, merging the two samples.'''
0056         self.sumw += other.sumw
0057         self.sumwx += other.sumwx
0058         self.sumwx2 += other.sumwx2
0059         return self
0060
0061     def __iadd__(self, other):
0062         '''Add two averages.'''
0063         return self.__add__(other)
0064
0065     def write(self, dirname):
0066         '''Dump the average to a pickle file and to a text file in dirname.'''
0067         pckfname = '{d}/{f}.pck'.format(d=dirname, f=self.name)
0068         pckfile = open( pckfname, 'w' )
0069         pickle.dump(self, pckfile)
0070         txtfile = open( pckfname.replace('.pck', '.txt'), 'w')
0071         txtfile.write( str(self) )
0072         txtfile.write( '\n' )
0073         txtfile.close()
0074
0075     def __str__(self):
0076         ave, unc = self.average()
0077         tmp = None
0078         if ave is not None:
0079             tmp = 'Average {name:<15}: {average: 8.4f} +- {unc:8.4f}'
0080             # tmp = 'Average {name:<15}: {average: 8.4f}'
0081             tmp = tmp.format( name = self.name,
0082                               average = ave,
0083                               unc = unc
0084                               )
0085         else:
0086             tmp = 'Average {name:<15}: undefined (call Average.add)'\
0087                   .format( name = self.name)
0088         return tmp
0089
0090
0091
0092 class Averages(diclist):
0093     def write(self, dirname):
0094         map( lambda x: x.write(dirname), self)