fitting_example_curvefit.py 1.1 KB

12345678910111213141516171819202122232425262728293031323334
  1. import numpy as np
  2. from scipy.optimize import curve_fit
  3. import matplotlib.pyplot as plt
  4. from atlas_fit_function import atlas_invMass_mumu_core
  5. # Load data:
  6. inFileName = "DATA/original_histograms/mass_mm_higgs_Background.npz"
  7. with np.load(inFileName) as data:
  8. bin_edges = data['bin_edges']
  9. bin_centers = data['bin_centers']
  10. bin_values = data['bin_values']
  11. bin_errors = data['bin_errors']
  12. #############################################################################################################################
  13. def poly3(x, a, b, c, d):
  14. """Fit poly3"""
  15. return a * x**3 + b * x**2 + c * x + d
  16. fitfun = poly3
  17. popt, pcov = curve_fit(fitfun, bin_centers, bin_values, sigma=bin_errors, p0=[1, 1, 1, 1])
  18. perr = np.sqrt(np.diag(pcov))
  19. a, b, c, d = popt
  20. my_fit = np.array(fitfun(bin_centers, a, b, c, d))
  21. xerrs = 0.5 * (bin_edges[1:] - bin_edges[:-1])
  22. plt.figure()
  23. plt.errorbar(bin_centers, bin_values, bin_errors, xerrs, fmt="none", color='b', ecolor='b', label='Original histogram')
  24. plt.plot(bin_centers, my_fit, 'g-', label='fit poly3')
  25. plt.legend()