histogram_utils.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
  1. '''
  2. histogram_utils.py
  3. The functions in this module use a representation of a histogram that is a
  4. tuple containing an arr of N bin values, an array of N bin errors(symmetric)
  5. and an array of N+1 bin edges(N lower edges + 1 upper edge).
  6. For 2d histograms, It is similar, but the arrays are two dimensional and
  7. there are separate arrays for x-edges and y-edges.
  8. '''
  9. import numpy as np
  10. from scipy.optimize import curve_fit
  11. def hist(th1, rescale_x=1.0, rescale_y=1.0):
  12. nbins = th1.GetNbinsX()
  13. edges = np.zeros(nbins+1, np.float32)
  14. values = np.zeros(nbins, np.float32)
  15. errors = np.zeros(nbins, np.float32)
  16. for i in range(nbins):
  17. edges[i] = th1.GetXaxis().GetBinLowEdge(i+1)
  18. values[i] = th1.GetBinContent(i+1)
  19. errors[i] = th1.GetBinError(i+1)
  20. edges[nbins] = th1.GetXaxis().GetBinUpEdge(nbins)
  21. edges *= rescale_x
  22. values *= rescale_y
  23. errors *= rescale_y
  24. return values, errors, edges
  25. def hist_bin_centers(h):
  26. _, _, edges = h
  27. return (edges[:-1] + edges[1:])/2.0
  28. def hist2d(th2, rescale_x=1.0, rescale_y=1.0, rescale_z=1.0):
  29. """ Converts TH2 object to something amenable to
  30. plotting w/ matplotlab's pcolormesh.
  31. """
  32. nbins_x = th2.GetNbinsX()
  33. nbins_y = th2.GetNbinsY()
  34. xs = np.zeros((nbins_y+1, nbins_x+1), np.float32)
  35. ys = np.zeros((nbins_y+1, nbins_x+1), np.float32)
  36. values = np.zeros((nbins_y, nbins_x), np.float32)
  37. errors = np.zeros((nbins_y, nbins_x), np.float32)
  38. for i in range(nbins_x):
  39. for j in range(nbins_y):
  40. xs[j][i] = th2.GetXaxis().GetBinLowEdge(i+1)
  41. ys[j][i] = th2.GetYaxis().GetBinLowEdge(j+1)
  42. values[j][i] = th2.GetBinContent(i+1, j+1)
  43. errors[j][i] = th2.GetBinError(i+1, j+1)
  44. xs[nbins_y][i] = th2.GetXaxis().GetBinUpEdge(i+1)
  45. ys[nbins_y][i] = th2.GetYaxis().GetBinUpEdge(nbins_y+1)
  46. for j in range(nbins_y+1):
  47. xs[j][nbins_x] = th2.GetXaxis().GetBinUpEdge(nbins_x+1)
  48. ys[j][nbins_x] = th2.GetYaxis().GetBinUpEdge(j+1)
  49. xs *= rescale_x
  50. ys *= rescale_y
  51. values *= rescale_z
  52. errors *= rescale_z
  53. return values, errors, xs, ys
  54. def hist_slice(hist, range_):
  55. values, errors, edges = hist
  56. lim_low, lim_high = range_
  57. slice_ = np.logical_and(edges[:-1] > lim_low, edges[1:] < lim_high)
  58. last = len(slice_) - np.argmax(slice_[::-1])
  59. return (values[slice_],
  60. errors[slice_],
  61. np.concatenate([edges[:-1][slice_], [edges[last]]]))
  62. def hist_add(hist1, hist2, w1=1.0, w2=1.0):
  63. v1, e1, *lim1 = hist1
  64. v2, e2, *lim2 = hist2
  65. # print(hist1)
  66. if v1.shape != v2.shape:
  67. raise ValueError(f'Mismatched histograms to add {v1.shape} != {v2.shape}')
  68. # print(lim1)
  69. # print(lim2)
  70. # nlims_equal = (lim1 != lim2).any()
  71. # print(nlims_equal)
  72. # if nlims_equal:
  73. # raise ValueError(f'Histograms have different limits!')
  74. if len(v1.shape) == 1: # 1D histograms
  75. return ((v1+v2), np.sqrt(e1*e1 + e2*e2), *lim1)
  76. def hist_integral(hist, times_bin_width=True):
  77. values, errors, edges = hist
  78. if times_bin_width:
  79. bin_widths = [abs(x2 - x1) for x1, x2 in zip(edges[:-1], edges[1:])]
  80. return sum(val*width for val, width in zip(values, bin_widths))
  81. else:
  82. return sum(values)
  83. def hist_normalize(hist, norm):
  84. values, errors, edges = hist
  85. scale = norm/np.sum(values)
  86. return values*scale, errors*scale, edges
  87. def hist_mean(hist):
  88. xs = hist_bin_centers(hist)
  89. ys, _, _ = hist
  90. return sum(x*y for x, y in zip(xs, ys)) / sum(ys)
  91. def hist_var(hist):
  92. xs = hist_bin_centers(hist)
  93. ys, _, _ = hist
  94. mean = sum(x*y for x, y in zip(xs, ys)) / sum(ys)
  95. mean2 = sum((x**2)*y for x, y in zip(xs, ys)) / sum(ys)
  96. return mean2 - mean**2
  97. def hist_std(hist):
  98. return np.sqrt(hist_var(hist))
  99. def hist_stats(hist):
  100. return {'int': hist_integral(hist),
  101. 'sum': hist_integral(hist, False),
  102. 'mean': hist_mean(hist),
  103. 'var': hist_var(hist),
  104. 'std': hist_std(hist)}
  105. # def hist_slice2d(h, range_):
  106. # values, errors, xs, ys = h
  107. # last = len(slice_) - np.argmax(slice_[::-1])
  108. # (xlim_low, xlim_high), (ylim_low, ylim_high) = range_
  109. # slice_ = np.logical_and(xs[:-1, :-1] > xlim_low, xs[1:, 1:] < xlim_high,
  110. # ys[:-1, :-1] > ylim_low, ys[1:, 1:] < ylim_high)
  111. # last = len(slice_) - np.argmax(slice_[::-1])
  112. # return (values[slice_],
  113. # errors[slice_],
  114. # np.concatenate([edges[:-1][slice_], [edges[last]]]))
  115. def hist_fit(h, f, p0=None):
  116. values, errors, edges = h
  117. xs = hist_bin_centers(h)
  118. # popt, pcov = curve_fit(f, xs, values, p0=p0, sigma=errors)
  119. popt, pcov = curve_fit(f, xs, values, p0=p0)
  120. return popt, pcov
  121. def hist_rebin(hist, range_, nbins):
  122. raise NotImplementedError()