bwedges.py (Source)

#! /usr/bin/env python3
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("DATFILE", help="unbinned data file to read in")
ap.add_argument("OUTNAME", nargs="?", default="mee",
                help="hist name to write out as .dat and .pdf")
ap.add_argument("--dyn", dest="DYNBIN", action="store_true",
                help="hist name to write out as .dat and .pdf")
args = ap.parse_args()
import numpy as np
vals = np.loadtxt(args.DATFILE)
## Binning
NBINS = 50
RANGE = [70,120]
binedges = np.linspace(*RANGE, NBINS)
if args.DYNBIN:
    ## Dynamic binning, by inversion of the Breit-Wigner CDF:
    ##  https://en.wikipedia.org/wiki/Cauchy_distribution
    ## PDF = 1 / [pi gamma (1 + (x-x0)^2/gamma^2)]
    ##   with x = E2,  x0 = M2, gamma = M Gamma
    ## CDF = arctan( (x - x0) / gamma) / pi + 1/2  
    ##   -> x_samp = gamma tan((rand - 0.5) pi) + x0
    M, Gamma = 91.2, 5.5
    gamma, x0 = M*Gamma, M**2
    qmin, qmax = 0.05, 0.95 #< quantile range to map
    qs = np.linspace(qmin, qmax, NBINS)
    xs = gamma * np.tan( (qs-0.5) * np.pi ) + x0
    binedges = np.sqrt(xs[xs > 0]) #< eliminate any negative E^2s
    
## Plot and save
import matplotlib.pyplot as plt
fig = plt.figure(figsize=(10,7))
fig.patch.set_alpha(0.0)
counts, edges, _ = plt.hist(vals, bins=binedges,
                            density=True, histtype="step")
plt.xlim(RANGE)
plt.xlabel("$e^+ e^-$ pair invariant mass, $m_{ee}$ [GeV]")
plt.ylabel("$\mathrm{d}N / \mathrm{d}m_{ee}$ [count/GeV]")
plt.yscale("log")
for ext in [".pdf", ".png"]:
    plt.savefig(args.OUTNAME+ext, dpi=100) # transparent=True
np.savetxt(args.OUTNAME+"-hist.dat",
           np.stack((edges[:-1], edges[1:], counts)).T)