Wrap fit to a function

[bdrum]

I have 4 fits with the same data part, obviously I have to hide it into a function. For drawing part see #47

import matplotlib.pyplot as plt
from numpy import exp, loadtxt, pi, sqrt, random, histogram
from lmfit import Model
from lmfit.models import BreitWignerModel, LinearModel, GaussianModel, ConstantModel, PolynomialModel
from FourTracks.analysis import fit

# %matplotlib inline
fig = plt.figure() #figsize=(10,7))
ax = fig.add_subplot()

b = 30
r = 0.8,2.2
y, x = np.histogram(mass_events(kinematics.GetTracksWithPtCut(four_tracks_zq_after_all_cuts)), bins=b, range=r)
x = x[:-1] + np.round((r[1]-r[0]) / b, 3) / 2

mod_bw = Model(fit.bw)

bcg_y, bcg_x = np.histogram(mass_events(kinematics.GetTracksWithPtCut(data.four_tracks_nzq,maxPt=0.6)), bins=b, range=r)
bcg_x = bcg_x[:-1]
bcg_y = bcg_y*0.2
fit.bckg_y = bcg_y

mod_bckg = PolynomialModel(degree=4)
par_bckg = mod_bckg.guess(bcg_y, x = bcg_x)

mod = Model(fit.bw, prefix='bw1_') + Model(fit.bw, prefix='bw2_') +  Model(fit.polyn, prefix='bckg_')
pars = mod.make_params()

amp1, M1, G1 = 150, 1.40, 0.2
amp2, M2, G2 = 135, 1.55, 0.3

pars['bw1_amp'].set(amp1)
pars['bw1_M'].set(M1)
pars['bw1_G'].set(G1)

pars['bw2_amp'].set(amp2)
pars['bw2_M'].set(M2)
pars['bw2_G'].set(G2)

pars['bckg_amp'].set(0.4)
pars['bckg_c0'].set(r_bckg.best_values['c0'], vary=False)
pars['bckg_c1'].set(r_bckg.best_values['c1'], vary=False)
pars['bckg_c2'].set(r_bckg.best_values['c2'], vary=False)
pars['bckg_c3'].set(r_bckg.best_values['c3'], vary=False)
pars['bckg_c4'].set(r_bckg.best_values['c4'], vary=False)
pars['bckg_c5'].set(0, vary=False)
pars['bckg_c6'].set(0, vary=False)
pars['bckg_c7'].set(0, vary=False)

result = mod.fit(y,pars, x=x, method='lm', weights=1/np.sqrt(y), nan_policy='omit')

print(result.fit_report())

dat, = hep.histplot(np.histogram(mass_after_all_cuts, bins=b, range=r), yerr=True, xerr=0.022, color='black',histtype="errorbar")
bw1, = plt.plot(np.linspace(x.min(), x.max(),1000), fit.bw(x=np.linspace(x.min(), x.max(),1000), M = result.best_values['bw1_M'], G=result.best_values['bw1_G'], amp=result.best_values['bw1_amp'] ), 'g-',label='B-W')
bw2, = plt.plot(np.linspace(x.min(), x.max(),1000), fit.bw(x=np.linspace(x.min(), x.max(),1000), M = result.best_values['bw2_M'], G=result.best_values['bw2_G'], amp=result.best_values['bw2_amp'] ), 'm-',label='B-W')
bkg, = plt.plot(np.linspace(x.min(), x.max(),1000), fit.polyn(np.linspace(x.min(), x.max(),1000), result.best_values['bckg_c0'], result.best_values['bckg_c1'], result.best_values['bckg_c2'], result.best_values['bckg_c3'], result.best_values['bckg_c4'],0,0,0, result.best_values['bckg_amp']), 'b--', label='Bckgr fit')
ft, = plt.plot(np.linspace(x.min(), x.max(),1000), result.eval(x=np.linspace(x.min(), x.max(),1000)), 'r-', label='fit')
# result.plot_fit(datafmt='ko',fitfmt='r-',numpoints=100,data_kws={'xerr':0.022})
ax.set_title("")
_ = ax.set_ylabel(r'Counts per 45 MeV/$c^2$')
_ = ax.set_xlabel(r'M($\pi^+\pi^-\pi^+\pi^-$) (GeV/$c^2$)')
_ = ax.legend([dat,ft,bw1,bw2, bkg],['Data','Fit', 'Breit-Wigner', 'Breit-Wigner','Bkgr'])