eic-smear/eicpy/binning.py

0001 import array
0002 import math
0003
0004 import ROOT
0005
0006 class Bins(object):
0007     '''
0008     Base class for generating bin boundaries in some range.
0009     '''
0010     def __init__(self, n, lower, upper):
0011         '''
0012         Initialise with the number of ranges for which
0013         to generate bin edges, in the range [lower, upper].
0014         '''
0015         self.n = n
0016         self.lower = lower
0017         self.upper = upper
0018
0019
0020 class BinsLog10(Bins):
0021     '''
0022     Generator class for bins with equal intervals in log10.
0023     '''
0024     def __init__(self, n, lower, upper):
0025         super(BinsLog10, self).__init__(n, lower, upper)
0026         self.width = math.log10(self.upper / self.lower) / self.n
0027         # Redefine lower and upper as log10 equivalents
0028         self.lower = math.log10(self.lower)
0029         self.upper = math.log10(self.upper)
0030
0031     def edges(self):
0032         '''
0033         Generator function for (n+1) bin edges.
0034         '''
0035         for i in range(0, self.n + 1):
0036             yield math.pow(10., self.lower + i * self.width)
0037
0038     def lower_edges(self):
0039         '''
0040         Generator function for n bin lower edges.
0041         '''
0042         for i in range(0, self.n):
0043             yield math.pow(10., self.lower + i * self.width)
0044
0045     def upper_edges(self):
0046         '''
0047         Generator function for n bin upper edges.
0048         '''
0049         for i in range(1, self.n + 1):
0050             yield math.pow(10., self.lower + i * self.width)
0051
0052
0053 def bin_log10(obj):
0054     '''
0055     Rebins an axis such that bins span equal intervals in log_10(x).
0056     '''
0057     if isinstance(obj, ROOT.TAxis):
0058         binner = BinsLog10(obj.GetNbins(), obj.GetXmin(), obj.GetXmax())
0059         # Need array('d') to pass to TAxis.Set() in place of Double_t*
0060         edgearray = array.array('d', [x for x in binner.edges()])
0061         obj.Set(obj.GetNbins(), edgearray)
0062     elif isinstance(obj, ROOT.TH1):
0063         axes = [obj.GetXaxis(), obj.GetYaxis(), obj.GetZaxis()]
0064         for i in range(0, obj.GetDimension()):
0065             bin_log10(axes[i])