def gauss_chi2(data,
ax,
bins=None,
range=None,
coords=(0.05, 0.95),
decimals=4):
vals, bincenter, binwidth = hist(data, bins, range)
svals = np.sqrt(vals)
mask = vals > 0
def fitfunc(x, mu, sigma):
normalization = len(data) * binwidth
return normalization / np.sqrt(2 * np.pi) / sigma * np.exp(
-0.5 * (x - mu)**2 / sigma**2)
chi2_obj = Chi2Regression(fitfunc, bincenter[mask], vals[mask],
svals[mask])
minuit = Minuit(chi2_obj,
pedantic=False,
mu=np.mean(data),
sigma=np.std(data))
minuit.migrad()
if not minuit.migrad_ok():
print('minuit.migrad() did not converge!')
if not minuit.matrix_accurate():
print('Hessematrix is not accurate!')
ax.errorbar(bincenter[mask],
vals[mask],
svals[mask],
drawstyle='steps-mid',
capsize=2,
linewidth=1,
color='k',
ecolor='r',
label='data')
ax.plot(bincenter[np.logical_not(mask)],
vals[np.logical_not(mask)],
'gx',
label='Bins with 0')
x = np.linspace(min(bincenter), max(bincenter), 500)
ax.plot(x, fitfunc(x, *minuit.args), 'b', label='$\\chi^2$ fit')
ndof = np.sum(mask) - len(minuit.args)
d = {
'Chi2/ndof:': f"{minuit.fval:.3f}/{ndof:d}",
"p": stats.chi2.sf(minuit.fval, ndof),
"mu": minuit.values['mu'],
"sigma": minuit.values['sigma']
}
add_text_to_ax(*coords,
nice_string_output(d, decimals=decimals),
ax,
fontsize=12)
return minuit
def weighted_mean(x,
errs,
ax=None,
plot_ticks=None,
point_label=None,
coords=(0.1, 0.9),
dec=3):
"""
This function takes as input measurents and errors and returns the weighted mean along with the error.
The weighted mean is calculated by doing a Chi-Square fit with a constant.
if ax is given, the function will plot the fit, data and errors on it.
The function return weighted_mean, err, and a dictionairy with chi-square value and p-value
"""
def constant(x, k):
return k
from AppStatFunctions import Chi2Regression, nice_string_output, add_text_to_ax
ticks = np.arange(len(x))
chi2_object = Chi2Regression(constant, ticks, x, errs)
chi2_min = Minuit(chi2_object, pedantic=False)
chi2_min.migrad()
Chi2 = chi2_min.fval
p_val = chi2.sf(Chi2, len(x) - 1)
k = chi2_min.args[0]
err = chi2_min.errors[0]
if ax:
if type(plot_ticks) != type(None):
ticks = plot_ticks
ax.plot(ticks, x, 'r.', label=point_label)
ax.errorbar(ticks, x, errs, c = 'k', elinewidth = 1, \
capsize = 2, ls = 'none')
ax.hlines(k, min(ticks), max(ticks), ls='--', label="Weighted mean")
d = {"chisquare": Chi2, \
"p:": p_val, \
"mean:": k,\
"mean_err": err}
add_text_to_ax(*coords,
nice_string_output(d, decimals=dec),
ax,
fontsize=10)
ax.legend()
return k, err, {"chi2": Chi2, "p": p_val}
def fit_mass(xs,
vals,
errs,
ax=None,
guesses_bkgr=[0, 0, -10, 2000],
guesses_sig=[498, 6, 17000]):
if not ax:
fig, ax = plt.subplots(figsize=(16, 10), ncols=2)
ax_sig = ax[1]
ax_all = ax[0]
ax_all.plot(xs, vals, 'r.')
ax_all.errorbar(xs,
vals,
errs,
color='k',
elinewidth=1,
capsize=2,
ls='none')
def background_fit(x, a, b, c, d):
return a * (x - 498)**3 + b * (x - 498)**2 + c * (x - 498) + d
# The signal fit Here gauss
def add_signal(x, mean, sig, size):
return size * norm.pdf(x, mean, sig)
# The full fit
def full_fit(x, mean, sig, size, a, b, c, d):
return background_fit(x, a, b, c, d) + add_signal(x, mean, sig, size)
# Background fit under here
vals_b, cov_b = curve_fit(background_fit, xs, vals, p0=guesses_bkgr)
b1, b2, b3, b4 = vals_b
bkgr_chi2 = Chi2Regression(background_fit, xs, vals, errs)
bkgr_min = Minuit(bkgr_chi2, pedantic=False, a=b1, b=b2, c=b3, d=b4)
bkgr_min.migrad()
# Plot result and save guesses
# ax_all.plot(xs, background_fit(xs, *bkgr_min.args),'b--', label = "background_fit")
b1, b2, b3, b4 = bkgr_min.args
s1, s2, s3 = guesses_sig
# Full fit
full_chi2 = Chi2Regression(full_fit, xs, vals, errs)
full_min = Minuit(full_chi2, pedantic = False, a = b1, b = b2, c = b3, d = b4, \
mean = s1, sig = s2, size = s3)
full_min.migrad()
s1, s2, s3, b1, b2, b3, b4 = full_min.args
ax_all.plot(xs, full_fit(xs, *full_min.args), "k-", label="full_fit")
ax_all.plot(xs,
background_fit(xs, *full_min.args[3:]),
'b--',
label="background_fit")
ax_all.legend(loc="upper right")
# Details:
text = {'chi2': full_min.fval, \
'pval': chi2.sf(full_min.fval, len(xs) - len(full_min.args)), \
'mean': f"{full_min.values['mean']:.1f} +/- {full_min.errors['mean']:.1f}",\
'N': f"{full_min.values['size']:.1f} +/- {full_min.errors['size']:.1f}"}
text_output = nice_string_output(text)
add_text_to_ax(0.60, 0.925, text_output, ax_all)
# Plot signal seperately
ax_sig.fill_between(xs,
add_signal(xs, s1, s2, s3),
color='red',
alpha=0.5,
label="sig fit")
vals_sig = vals - background_fit(xs, b1, b2, b3, b4)
ax_sig.plot(xs, vals_sig, 'r.')
ax_sig.errorbar(xs,
vals_sig,
errs,
color='k',
elinewidth=1,
capsize=2,
ls='none')
sig_amount = np.sum(add_signal(xs, s1, s2, s3))
bak_amount = np.sum(background_fit(xs, b1, b2, b3, b4))
text_a = {'sig': np.round(sig_amount), \
'bkgr': np.round(bak_amount), \
's/b': sig_amount / bak_amount}
text_output = nice_string_output(text_a, decimals=2)
add_text_to_ax(0.70, 0.90, text_output, ax_sig)
fig.tight_layout()
bak_func = lambda x: background_fit(x, b1, b2, b3, b4)
sig_func = lambda x: add_signal(x, s1, s2, s3)
return fig, ax, full_min, bak_func, sig_func, [s1, s2, s3, b1, b2, b3, b4]
def double_gauss_fit(mass,
bins=100,
range=(400, 600),
ax=None,
verbose=True,
guesses=None,
max_size=None,
plimit=0,
color="red",
type='ks'):
# Fit third degree polynomium to background
if type == 'ks':
def background_fit(x, a, b, c, d):
res = a * (x - 498)**3 + b * (x - 498)**2 + c * (x - 498) + d
return res
elif type == 'la':
def background_fit(x, a, b, c, d):
res = a * (x - 1116)**3 + b * (x - 1116)**2 + c * (x - 1116) + d
return res
# The double gauss signal
def add_signal(x, mean, sig, size, ratio, sig_ratio):
return size * binwidth * (ratio * norm.pdf(x, mean, sig) + \
(1 - ratio) * norm.pdf(x, mean, sig_ratio * sig))
# The full fit
def full_fit(x, mean, sig, size, ratio, sig_ratio, a, b, c, d):
return background_fit(x, a, b, c, d) + add_signal(
x, mean, sig, size, ratio, sig_ratio)
# Make histogram
vals, edges = np.histogram(mass, bins=bins, range=range)
xs = (edges[1:] + edges[:-1]) / 2
binwidth = xs[1] - xs[0]
mask = vals > 0
vals = vals[mask]
xs = xs[mask]
errs = np.sqrt(vals)
# Get guesses for a background fit
if not guesses:
if type == 'ks':
back_data_mask = abs(xs - xs[np.argmax(vals)]) > 10
if type == 'la':
back_data_mask = abs(xs - 1117) > 5
background_guess = [0, 0, (vals[-1] - vals[0]) / 100, vals.min()]
if len(vals[back_data_mask]) == 0:
return None, None, None, None
try:
vals_b, cov_b = curve_fit(background_fit,
xs[back_data_mask],
vals[back_data_mask],
p0=background_guess)
except:
vals_b = background_guess
b1, b2, b3, b4 = vals_b
bkgr_chi2 = Chi2Regression(background_fit, xs[back_data_mask],
vals[back_data_mask], errs[back_data_mask])
bkgr_min = Minuit(bkgr_chi2, pedantic=False, a=b1, b=b2, c=b3, d=b4)
bkgr_min.migrad()
counter = 0
while not bkgr_min.valid and counter < 50:
bkgr_min.migrad()
counter += 1
if not bkgr_min.valid: print("No background valid minimum found!")
#Save guesses
b1, b2, b3, b4 = bkgr_min.args
if type == 'ks':
guesses_sig = [498, 7, 2000, 0.5, 2]
elif type == 'la':
guesses_sig = [1116, 3, 3700, 0.5, 2]
try:
vals_f, cov_f = curve_fit(full_fit,
xs,
vals,
p0=guesses_sig + [b1, b2, b3, b4])
except:
vals_f = np.hstack([guesses_sig, vals_b])
s1, s2, s3, s4, s5, b1, b2, b3, b4 = vals_f
else:
s1, s2, s3, s4, s5, b1, b2, b3, b4 = guesses
full_chi2 = Chi2Regression(full_fit, xs, vals, errs)
if type == 'ks':
full_min = Minuit(full_chi2, pedantic = False, a = b1, b = b2, c = b3, d = b4, \
mean = s1, sig = s2, size = s3, ratio = s4, sig_ratio = s5, limit_sig_ratio = (1, 4), \
limit_ratio = (0, 1.0), limit_mean = (490, 510), limit_size = (0, max_size), limit_sig = (3, 10))
elif type == 'la':
full_min = Minuit(full_chi2, pedantic = False, a = b1, b = b2, c = b3, d = b4, \
mean = s1, sig = s2, size = s3, ratio = s4, sig_ratio = s5, limit_sig_ratio = (1, 4), \
limit_ratio = (0, 1.0), limit_mean = (1115,1118), limit_size = (0, max_size), limit_sig = (0.1,10))
full_min.migrad()
full_min.migrad()
counter = 0
while not full_min.valid and counter < 200:
full_min.migrad()
counter += 1
if not full_min.valid: print("No valid minimum found!")
# Check fit
chi = full_min.fval
pval = chi2.sf(chi, np.sum(mask) - len(full_min.args))
if verbose:
print(f"Completed fit with Chi2: {chi:.1f}, p-val: {pval:.3f} and the total amount of signal " + \
f"{full_min.values['size']:.0f} +/- {full_min.errors['size']:.0f}, background: {len(mass) - int(full_min.values['size'])}")
if ax:
ax.plot(xs, vals, alpha=1, color=color)
# ax.errorbar(xs, vals, errs, elinewidth = 1, color = 'k', capsize = 2, linestyle = 'none', alpha = 0.25)
# ax.plot(xs, full_fit(xs, *full_min.args), '--', alpha = 0.5)
# s1, s2, s3, s4, s5, b1, b2, b3, b4 = full_min.args
# ax.plot(xs,background_fit(xs,b1,b2,b3,b4),color=color)
# ax.plot(xs,add_signal(xs,s1,s2,s3,s4,s5)+background_fit(xs,b1,b2,b3,b4),color='k',ls='--')
if True: #full_min.errors['size'] < full_min.values['size'] and full_min.valid and pval > plimit:
return full_min.values['size'], len(mass) - full_min.values[
'size'], full_min.errors['size'], full_min.args
else:
return None, None, None, None
def fit_mass2(xs,
vals,
errs,
ax=None,
guesses_bkgr=[0, 0, -10, 2000],
guesses_sig=[497, 6, 17000],
plot=True,
type='ks',
second_axis=None):
guesses_bkgr[-1] = 0.5 * (vals[0] + vals[-1])
guesses_sig[-1] = 20 * max(
vals
) #np.sqrt(2*np.pi*guesses_sig[1]**2)*(max(vals))# - guesses_bkgr[-1])
if not ax and plot:
fig, ax = plt.subplots(figsize=(16, 10), ncols=2)
ax_sig = ax[1]
ax_all = ax[0]
ax_all.plot(xs, vals, 'r.')
ax_all.errorbar(xs,
vals,
errs,
color='k',
elinewidth=1,
capsize=2,
ls='none')
elif plot:
fig = None
ax_sig = ax[1]
ax_all = ax[0]
ax_all.plot(xs, vals, 'r.')
ax_all.errorbar(xs,
vals,
errs,
color='k',
elinewidth=1,
capsize=2,
ls='none')
def background_fit(x, a, b, c, d):
return a * (x - guesses_sig[0])**3 + b * (
x - guesses_sig[0])**2 + c * (x - guesses_sig[0]) + d
# The signal fit Here gauss
def sig1(x, mean, sig, size):
return size * norm.pdf(x, mean, sig)
def sig2(x, mean, sig, size):
return size * norm.pdf(x, mean, sig)
# The full fit
def full_fit(x, mean, sig, size, f, sigmp, a, b, c, d):
return background_fit(x, a, b, c, d) + f * sig1(
x, mean, sig, size) + (1 - f) * sig2(x, mean, sigmp * sig, size)
# Background fit under here
if type == 'ks':
bkgr_mask = (xs < 475) | (xs > 525)
elif type == 'la':
bkgr_mask = abs(xs - 1117) > 5
try:
vals_b, cov_b = curve_fit(background_fit,
xs[bkgr_mask],
vals[bkgr_mask],
p0=guesses_bkgr)
b1, b2, b3, b4 = vals_b
except:
b1, b2, b3, b4 = guesses_bkgr
bkgr_chi2 = Chi2Regression(background_fit, xs[bkgr_mask], vals[bkgr_mask],
errs[bkgr_mask])
bkgr_min = Minuit(bkgr_chi2, pedantic=False, a=b1, b=b2, c=b3, d=b4)
with warnings.catch_warnings():
warnings.simplefilter('ignore')
bkgr_min.migrad()
counter = 0
while not bkgr_min.valid and counter < 50:
bkgr_min.migrad()
counter += 1
if not bkgr_min.valid: print("No background valid minimum found!")
#Save guesses
b1, b2, b3, b4 = bkgr_min.args
s1, s2, s3 = guesses_sig
if type == 'la':
s1 = 1117
# Full fit
full_chi2 = Chi2Regression(full_fit, xs, vals, errs)
full_min = Minuit(full_chi2, pedantic = False, a = b1, b = b2, c = b3, d = b4, \
mean = s1, sig = s2, size = s3, f = 0.5, sigmp = 2, \
limit_mean=(475,525), limit_f=(0,1), limit_size=(0,None))
# full_min = Minuit(full_chi2, pedantic = False, a = b1, b = b2, c = b3, d = b4, \
# mean = s1, sig = s2, size = s3, f = 0.5, sigmp = 2, \
# limit_f=(0,1), limit_size=(0,None),limit_sig=(0.1,10))
with warnings.catch_warnings():
warnings.simplefilter('ignore')
full_min.migrad()
counter = 0
while not full_min.valid and counter < 200:
full_min.migrad()
counter += 1
if not full_min.valid: print("No valid minimum found!")
mean, sig, size, f, sigmp, b1, b2, b3, b4 = full_min.args
sig_amount = np.sum(f * sig1(xs, mean, sig, size) +
(1 - f) * sig2(xs, mean, sigmp * sig, size))
bak_amount = np.sum(background_fit(xs, b1, b2, b3, b4))
neg_bkgr = any(background_fit(xs, b1, b2, b3, b4) < 0)
def signal():
def signal_func(x):
return f * sig1(x, mean, sig, size) + (1 - f) * sig2(
x, mean, sigmp * sig, size)
return signal_func
def background():
def background_func(x):
return background_fit(x, b1, b2, b3, b4)
return background_func
if plot:
ax_all.plot(xs, full_fit(xs, *full_min.args), "k-", label="full_fit")
ax_all.plot(xs,
background_fit(xs, b1, b2, b3, b4),
'b--',
label="background_fit")
ax_all.legend(loc="upper right")
# Details:
text = {'chi2': full_min.fval, \
'pval': chi2.sf(full_min.fval, len(xs) - len(full_min.args)), \
'mean': f"{full_min.values['mean']:.1f} +/- {full_min.errors['mean']:.1f}",\
'N': f"{full_min.values['size']:.1f} +/- {full_min.errors['size']:.1f}"}
text_output = nice_string_output(text)
add_text_to_ax(0.60, 0.925, text_output, ax_all)
# Plot signal seperately
ax_sig.fill_between(xs,
f * sig1(xs, mean, sig, size) +
(1 - f) * sig2(xs, mean, sigmp * sig, size),
color='red',
alpha=0.5,
label="sig fit")
ax_sig.plot(xs,
f * sig1(xs, mean, sig, size),
ls='--',
color='blue',
alpha=0.5,
label="sig fit")
ax_sig.plot(xs, (1 - f) * sig2(xs, mean, sigmp * sig, size),
ls='--',
color='green',
alpha=0.5,
label="sig fit")
vals_sig = vals - background_fit(xs, b1, b2, b3, b4)
ax_sig.plot(xs, vals_sig, 'r.')
ax_sig.errorbar(xs,
vals_sig,
errs,
color='k',
elinewidth=1,
capsize=2,
ls='none')
text_a = {'sig': np.round(sig_amount), \
'bkgr': np.round(bak_amount), \
's/b': sig_amount / bak_amount}
text_output = nice_string_output(text_a, decimals=2)
add_text_to_ax(0.70, 0.90, text_output, ax_sig)
try:
fig.tight_layout()
except:
pass
return {
'fig': fig,
'ax': ax,
'M': full_min,
'sig': sig_amount,
'bkgr': bak_amount,
'neg_bkgr': neg_bkgr,
'sig_func': signal(),
'bkgr_func': background()
}
only_text = True
if only_text and second_axis:
N_sig = full_min.values['size']
N_background = sum(vals) - N_sig
text = {
'mean': f"{full_min.values['mean']:.1f} +/- {full_min.errors['mean']:.1f}",\
'sigma': f"{full_min.values['sig']:.2f} +/- {full_min.errors['sig']:.2f}",\
'sigma-mp': f"{full_min.values['sigmp']:.2f} +/- {full_min.errors['sigmp']:.2f}",\
'f': f"{full_min.values['f']:.2f} +/- {full_min.errors['f']:.2f}",\
'signal': f"{N_sig:.0f} +/- {full_min.errors['size']:.0f}",\
'background': f"{N_background:.0f} +/- {full_min.errors['size']:.0f}", \
's/b': f"{N_sig/N_background:.3f} +/- {N_sig/N_background*full_min.errors['size']*np.sqrt(1/N_sig**2 + 1/N_background**2):.3f}"}
text_output = nice_string_output(text, extra_spacing=1)
add_text_to_ax(0.01, 0.99, text_output, second_axis)
return {
'M': full_min,
'sig': sig_amount,
'bkgr': bak_amount,
'neg_bkgr': neg_bkgr,
'sig_func': signal(),
'bkgr_func': background()
}