[checkpoint] Doing great on XS extraction, now focus on MC part

analysis/scripts/adaptive_fit.sharedwidth.py

+#!/usr/env python
+# -*- coding: utf-8 -*-
+# -*- format: python -*-
+# -*- author: G. Henning -*-
+# -*- created: may 2020 -*-
+# -*- copyright: GH/IPHC 2020 -*-
+# -*- file: fit_one_peak.py -*-
+# -*- purpose: -*-
+
+'''
+Fit peaks with a given energy, bg profile and parasitic peaks
+'''
+
+import argparse
+import yaml 
+import traceback
+
+from operator import itemgetter
+
+import numpy as np
+import scipy
+from scipy.optimize import curve_fit
+
+import matplotlib
+matplotlib.use('agg')
+import matplotlib.pyplot as plt
+
+from pyhisto.histogram import Histogram
+from pyhisto.tools import GaussRealFit
+
+
+def adaptive_fit(histo_file: str,
+                 xpeak: float = 0.0, xwide: float = 10.,
+                 guessstdev: float = 1., 
+                 stdevrange: float = 0.25,
+                 xmin: float = None, xmax: float = None,
+                 bg_type: str = 'lin',
+                 parasites: list = [], 
+                 plot: str = None) -> str:
+    '''
+    Fit with parasites
+    '''
+    the_histo: Histogram = None
+    if (xmin is None) and (xmax is None) and (xwide is None):
+        the_histo = Histogram(fromfile=histo_file)
+    else:
+        if xmin is None:
+            xmin = xpeak - xwide
+        if xmax is None:
+            xmax = xpeak + xwide
+        the_histo = Histogram(fromfile=histo_file).slice(xmin, xmax)
+    # if no parasites,
+    if parasites is None or len(parasites) == 0:
+        # fall back to realfit
+        the_fit = GaussRealFit(the_histo)
+        if plot is not None:
+            the_fit.plot(plot)
+        return yaml.dump({'integral': float(the_fit.integral),
+                          'integral_u': float(the_fit.integral_u),
+                          'mean': float(the_fit.mean),
+                          'mean_u': float(the_fit.mean_u),
+                          'stdev': float(the_fit.stdev),
+                          'stdev_u': float(the_fit.stdev_u)
+                         })
+    # if some parasites, now we are talking.
+    functions_to_fit = []
+    variables_guess = []
+    variables_limits = []
+    variables_parameters = []
+    last_variable = 0
+    width_variable = None
+    #   .. Background is either linear or constant
+    if bg_type == 'lin':
+        functions_to_fit.append(lambda x, B, SL: B + SL * (x - xpeak))
+        variables_guess.extend([the_histo[0].count, 0.])
+        variables_limits.extend([(0.0, np.inf),
+                                 (-np.inf, np.inf)])
+        variables_parameters.append((0, 1,))
+        last_variable = 1
+    else:
+        functions_to_fit.append(lambda x, B: B)
+        variables_guess.extend([the_histo[0].count])
+        variables_limits.extend([(0.0, np.inf)])
+        variables_parameters.append((0,))
+        last_variable = 0
+    #   .. Main peak
+    functions_to_fit.append(lambda x, S, M, R: S / np.sqrt(2. * np.pi * R * R) * np.exp(-(x - M) ** 2. / 2. / R / R))
+    variables_guess.extend([max(1, sum(the_histo) - the_histo[0].count * len(the_histo)),
+                            xpeak, 
+                            guessstdev])
+    variables_limits.extend([(0., max(1, sum(the_histo))),
+                             (xpeak - 3 * guessstdev, xpeak + 3 * guessstdev),
+                             ( (1. - stdevrange) * guessstdev, (1. + stdevrange) * guessstdev)])
+    variables_parameters.append((last_variable + 1, 
+                                 last_variable + 2,
+                                 last_variable + 3,))
+    width_variable = last_variable + 3
+    last_variable = last_variable+3
+    #   .. now the parasites:
+    for extra_peak in parasites:
+        functions_to_fit.append(lambda x, S, M, R: S / np.sqrt(2. * np.pi * R * R) * np.exp(-(x - M) ** 2. / 2. / R / R))
+        variables_guess.extend([(max(1, sum(the_histo)  - the_histo[0].count * len(the_histo)) / 2.),
+                                extra_peak])
+        variables_limits.extend([(0., max(1, sum(the_histo))),
+                                 (extra_peak - 1.0 * guessstdev, extra_peak + 1.0 * guessstdev)])
+        variables_parameters.append((last_variable + 1,
+                                     last_variable + 2,
+                                     width_variable,))
+        last_variable = last_variable + 2
+    variables_limits = list(zip(*variables_limits))
+    # Now, defining our final function
+    def the_function(x, *p):
+        S = 0.0
+        for f, s in zip(functions_to_fit, variables_parameters):
+            S += f(x, *itemgetter(*s)(p))
+        return S
+    # And ready to fit!
+    fitrslt, cov = curve_fit(the_function,
+                             *the_histo.xy(),
+                             p0=variables_guess,
+                             bounds=variables_limits,
+                             max_nfev=len(variables_guess)*len(the_histo)*250)
+    #print(variables_parameters)
+    #print(fitrslt)
+    i_integral, i_mean, i_stdev = variables_parameters[1][0], variables_parameters[1][1], variables_parameters[1][2]
+    uncerts = np.sqrt(np.diag(cov))
+    integral = float(fitrslt[i_integral])/the_histo.binwidth
+    integral_u = float(uncerts[i_integral])/the_histo.binwidth
+    mean = float(fitrslt[i_mean])
+    mean_u = float(uncerts[i_mean])
+    stdev = float(fitrslt[i_stdev])
+    stdev_u = float(uncerts[i_stdev])
+    the_result: dict = {'integral': integral,
+                        'integral_u': integral_u,
+                        'mean': mean,
+                        'mean_u': mean_u,
+                        'stdev': stdev,
+                        'stdev_u': stdev_u}
+    # Plotting the result! (if)
+    if plot is None:
+        return yaml.dump(the_result)
+    plt.figure()
+    plt.grid(True)
+    plt.step(*the_histo.xy(), label='original')
+    X_values = list(the_histo.xy())[0]
+    BG_const = float(fitrslt[0])
+    plt.step(X_values,
+             list(map(lambda x: functions_to_fit[0](x, *itemgetter(*variables_parameters[0])(fitrslt)), 
+                      X_values)),
+             label='background')
+    plt.step(X_values,
+             list(map(lambda x: BG_const+functions_to_fit[1](x, *itemgetter(*variables_parameters[1])(fitrslt)),
+                     X_values)),
+             label='Mainpeak')
+    for f, s in zip(functions_to_fit[2:], variables_parameters[2:]):
+        plt.step(X_values,
+                 list(map(lambda x: BG_const+f(x, *itemgetter(*s)(fitrslt)),
+                     X_values)),
+                 label='Parasite ')
+    plt.legend()
+    plt.savefig(plot)
+    
+    return yaml.dump(the_result)
+    # end of fit
+
+# Now, execution part
+
+if __name__ == "__main__":
+    arg_parser = argparse.ArgumentParser(description=__doc__)
+    arg_parser.add_argument('--xpeak', type=float,
+                            default=0.0, nargs='?',
+                            help='guess of energy')
+    arg_parser.add_argument('--xwide', type=float,
+                            default=None, nargs='?',
+                            help='fit window size')
+    arg_parser.add_argument('--guessstdev', type=float,
+                           default=1.0, nargs='?',
+                           help="guess for stdev of peaks")
+    arg_parser.add_argument('--stdevrange', type=float,
+                            default=0.25, nargs='?',
+                            help="relative range around the guess value of stdev (1. = 100%)")
+    arg_parser.add_argument('--xmin', type=float, 
+                           default=None, nargs='?')
+    arg_parser.add_argument('--xmax', type=float, 
+                           default=None, nargs='?')
+    arg_parser.add_argument('--bg_type', type=str,
+                            choices=['const', 'lin'],
+                            default='lin', 
+                            help='type of bg')
+    arg_parser.add_argument('--plot', type=str,
+                             default=None, nargs='?',
+                             help="add file name to save in png")
+    arg_parser.add_argument('--parasites', type=float,
+                           default=None, nargs="*")
+    arg_parser.add_argument('--histo_file', type=str,
+                            nargs='?', help="Faster files to read")
+    the_arguments = arg_parser.parse_args()
+    #print(the_arguments)
+    try:
+        print(adaptive_fit(**vars(the_arguments)))
+    except Exception as e:
+        print(f"""
+integral: 0.0
+integral_u: 0.0
+mean: 0.0
+mean_u: 0.0
+stdev: 0.0
+stdev_u: 0.0
+error: True
+description: {str(e)}
+""")
+
+
+
+# end of file