Newer
Older
Master_thesis / raremodel.py
import numpy as np
from pdg_const import pdg
import matplotlib
import matplotlib.pyplot as plt
import pickle as pkl
import sys
import time
from helperfunctions import display_time, prepare_plot
import cmath as c
import scipy.integrate as integrate
from scipy.optimize import fminbound
from array import array as arr
import collections
from itertools import compress

mmu = pdg['muon_M']
mb = pdg["bquark_M"]
ms = pdg["squark_M"]
mK = pdg["Ks_M"]
mB = pdg["Bplus_M"]

class model:

    def __init__(self):

        #Initialize constants

        self.mmu = pdg['muon_M']
        self.mb = pdg["bquark_M"]
        self.ms = pdg["squark_M"]
        self.mK = pdg["Ks_M"]
        self.mB = pdg["Bplus_M"]

        self.C7eff = pdg["C7eff"]
        self.C9eff = pdg["C9eff"]
        self.C10eff = pdg["C10eff"]

        self.number_of_decays = pdg["number_of_decays"]

        self.nr_of_part_list = {}

        self.resolution = 7.0

        #self.C1 = pdg["C1"]
        #self.C2 = pdg["C2"]
        #self.C3 = pdg["C3"]
        #self.C4 = pdg["C4"]

        self.GF = pdg["GF"] #Fermi coupling const.
        self.alpha_ew = pdg["alpha_ew"]
        self.Vts = pdg["Vts"]
        self.Vtb = pdg["Vtb"]

        self.param_val = collections.OrderedDict()
        self.param_val_orig = collections.OrderedDict()

        #helperconstants

        self.x_min = 2*self.mmu
        self.x_max = (self.mB - self.mK) - 0.1
        self.total_pdf_list = []
        self.total_pdf_names = []
        self.mode = ""
        self.param_val["total_scale_amp"] = 1.0
        self.param_val["_mean"] = 0.0
        self.reco_steps = 10000

        self.param_val_orig = self.param_val

    def formfactor(self, q2, subscript): #returns real value

        #check if subscript is viable

        if subscript != "0" and subscript != "+" and subscript != "T":
            raise ValueError('Wrong subscript entered, choose either 0, + or T')

        #get constants

        mh = self.mK
        mbstar0 = pdg["mbstar0"]
        mbstar = pdg["mbstar"]
        b0 = pdg["b0"]
        bplus = pdg["bplus"]
        bT = pdg["bT"]

        #N comes from derivation in paper

        N = 3

        #some helperfunctions

        tpos = (self.mB - self.mK)**2
        tzero = (self.mB + self.mK)*(np.sqrt(self.mB)-np.sqrt(self.mK))**2

        z_oben = np.sqrt(tpos - q2) - np.sqrt(tpos - tzero)
        z_unten = np.sqrt(tpos - q2) + np.sqrt(tpos - tzero)
        z = np.divide(z_oben, z_unten)

        #calculate f0

        if subscript == "0":
            prefactor = 1/(1 - q2/(mbstar0**2))
            _sum = 0

            for i in range(N):
                _sum += b0[i]*(np.power(z,i))

            return prefactor * _sum

        #calculate f+ or fT

        else:
            prefactor = 1/(1 - q2/(mbstar**2))
            _sum = 0

            if subscript == "T":
                b = bT
            else:
                b = bplus

            for i in range(N):
                _sum += b[i] * (np.power(z, i) - ((-1)**(i-N)) * (i/N) * np.power(z, N))

            return prefactor * _sum

    def axiv_nonres(self, q2): #returns real value

        #Load constants

        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C10eff = self.C10eff

        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB

        #Some helperfunctions

        beta = np.sqrt(np.abs(1. - 4. * self.mmu**2. / q2))

        kabs = np.sqrt(mB**2. + np.power(q2, 2)/mB**2. + mK**4./mB**2. - 2. * (mB**2. * mK**2. + mK**2. * q2 + mB**2. * q2) / mB**2.)

        #prefactor in front of whole bracket

        prefactor1 = GF**2. *alpha_ew**2. * (np.abs(Vtb*Vts))**2. * kabs * beta / (128. * np.pi**5.)

        #left term in bracket

        bracket_left = 2./3. * kabs**2. * beta**2. * np.abs(C10eff*self.formfactor(q2, "+"))**2.

        #middle term in bracket

        _top = 4. * mmu**2. * (mB**2. - mK**2.) * (mB**2. - mK**2.)

        _under = q2 * mB**2.

        bracket_middle = _top/_under * np.power(np.abs(C10eff * self.formfactor(q2, "0")), 2)

        #Note sqrt(q2) comes from derivation as we use q2 and plot q

        return prefactor1 * (bracket_left + bracket_middle) * 2 * np.sqrt(q2)

    def vec_nonres(self, q2): #returns real value

        #Load constants

        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C7eff = self.C7eff
        C9eff = self.C9eff

        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB

        #Some helperfunctions

        beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2))

        kabs = np.sqrt(mB**2. + np.power(q2, 2)/mB**2. + mK**4./mB**2. - 2 * (mB**2 * mK**2 + mK**2 * q2 + mB**2 * q2) / mB**2)

        #prefactor in front of whole bracket

        prefactor1 = GF**2. *alpha_ew**2. * (np.abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)

        #right term in bracket

        prefactor2 = kabs**2 * (1. - 1./3. * beta**2)

        abs_bracket = np.abs(C9eff * self.formfactor(q2, "+") + 2 * C7eff * (mb + ms)/(mB + mK) * self.formfactor(q2, "T"))**2

        bracket_right = prefactor2 * abs_bracket

        #Note sqrt(q2) comes from derivation as we use q2 and plot q

        return prefactor1 * bracket_right * 2 * np.sqrt(q2)

    def total_nonres(self, q2, parameters = 0, name = "", absolut = False): #returns complex value

        #Load constants

        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C10eff = self.C10eff
        C9eff = self.C9eff
        C7eff = self.C7eff

        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB

        #vector nonresonant part

        vec_nonres_part = self.vec_nonres(q2)

        #axial verctor nonresonant part including C7

        axiv_nonres_part = self.axiv_nonres(q2)

        #Normalized factor

        # for nr in self.nr_of_part_list:
        #     _sum += nr
        #
        # factor = self.nr_of_part_list[name]/_sum

        factor = 1.0

        #Complete term

        if absolut:
            return factor*(vec_nonres_part + axiv_nonres_part)

        else:
            return factor * np.vectorize(complex)(vec_nonres_part + axiv_nonres_part, 0.0)

    def generate_points(self, set_size = 10000, x_min = 2* mmu, x_max = (mB - mK) - 0.1, save = True, verbose = 1, mode = "true_data", resolution = 7.0, nonres_set_size = 44000, nr_of_toys = 1): #returns part_set (=dic)

        #Check if entered mode and verbose are valid

        if mode != "slim_points" and mode != "fast_binned" and mode != "true_data":
            raise ValueError('Invalid mode entered, choose either slim_points, fast_binned or true_data')

        if verbose != 0 and verbose != 1:
            raise ValueError('Invalid verbose entered, choose either 0 (no verbose) or 1 (verbose)')

        #Setup contants and variables

        self.mode = mode

        mB = self.mB
        mK = self.mK
        mmu = self.mmu


        #slim_mode: Generates data of a set size in range(x_min,x_max) according to total_pdf -> Only saves the valid q values (Hit and miss method)

        if mode == "slim_points":

            #Prepare variables

            set_sizes = np.random.normal(set_size, np.sqrt(set_size), nr_of_toys)

            def total_pdf_neg(q2):
                return -1*self.total_pdf(q2)

            maxi = fminbound(total_pdf_neg, x_min**2, x_max**2)

            maxi = self.total_pdf(np.array(maxi))

            for toy in range(nr_of_toys):

                set_size_intermed = 0

                print("Generating set of size {}...".format(int(set_sizes[toy])))

                x_part = []

                _ = 0

                while set_size_intermed < set_sizes[toy]:

                    x_part_raw = np.random.uniform(low = x_min, high = x_max, size = 1000000)

                    # print(x_part_raw[:40])

                    y_part_raw = np.random.uniform(low = 0.0, high = maxi, size = 1000000)

                    # print(y_part_raw[:40])

                    choose = np.where(y_part_raw < self.total_pdf(np.power(x_part_raw, 2)), True, False)

                    print(choose[:40])

                    x_part_intermed = list(compress(x_part_raw, choose))

                    x_part += x_part_intermed

                    set_size_intermed += len(x_part_intermed)

                x_part = x_part[:int(set_sizes[toy])]

                x_part = np.array(x_part)

                print(" Toy {0} of {1} generated!".format(toy, nr_of_toys))

                #Save the set

                if save:

                    part_set = {"x_part" : x_part}

                    pkl.dump( part_set, open("./data/slim_points/slim_points_toy_{0}_range({1}-{2}).pkl".format(toy, int(x_min), int(x_max)) , "wb" ) )

                    print(" Set saved!\n")

                print

            #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)

            return part_set

        elif mode == "true_data":

            #Change Seed here if necessary

            #ROOT.TRandom1().SetSeed(0)

            #Range of the whole nonres spectrum (used to count total hits in the nonres part)

            non_res_sizes = np.random.normal(nonres_set_size, np.sqrt(nonres_set_size), nr_of_toys)

            if x_min < 3096.0 and x_max > 3096.0:
                _max = self.total_pdf(3096.0**2)

            elif self.total_pdf(x_min**2) < self.total_pdf(x_max**2):
                _max = self.total_pdf(x_max**2)
            else:
                _max = self.total_pdf(x_min**2)

            x_min_nr = self.x_min
            x_max_nr = self.x_max

            _start = time.time()
            _end = _start

            tot = []
            for toy_nr in range(nr_of_toys):
                tot.append(0)

            for toy_nr in range(nr_of_toys):

                #Prepare variables

                o = 0
                x_part = []
                _nonres_set_size = int(ROOT.TRandom1().Gaus(nonres_set_size, np.sqrt(nonres_set_size)))

                print("Generating toy set {1} of {2} with {0} total rare nonresonant particles...".format(int(_nonres_set_size), toy_nr, nr_of_toys - 1))

                #Take starting time and other variables for projected time

                start = time.time()

                counter = 0
                counter_x = 0
                counter_x_nr = 0

                #Prepare verbose calls

                if verbose != 0:
                    verbose_calls = []
                    o = 0
                    for j in range(100):
                        verbose_calls.append(int(_nonres_set_size*(j+1)/100))

                #Loop until enough particles in nonresonant part

                while counter_x_nr < _nonres_set_size:

                    #Generate random points

                    x = ROOT.TRandom1().Uniform(x_min_nr, x_max_nr)
                    y = ROOT.TRandom1().Uniform(0, _max)

                    #Add the part to the set if within range(x_min,x_max) and under pdf

                    if x > x_min and x < x_max and y < self.total_pdf(x**2):
                        x_part.append(x)
                        counter_x += 1

                    #Count the hits in the nonres part -> Range of the whole pdf!

                    if y < self.total_nonres(x**2, absolut = True):
                        counter_x_nr += 1

                    #Calculate total estimated time to produce all toys

                    end = time.time()
                    tot[toy_nr] = (end-start)*_nonres_set_size/(counter_x_nr+1)
                    total_estim = nr_of_toys*np.mean(tot[:toy_nr+1])

                    #progress calls
                    if verbose != 0:
                        if o < 100 and counter%10000 == 0:
                            print(" {0}{1} completed of this toy".format(o, "%"))
                            print(" {0} rare nonresonant particles".format(counter_x_nr))
                            print(" {0} total particles".format(counter_x))
                            print(" Projected time left for this toy: {0}".format(display_time(int((end-start)*_nonres_set_size/(counter_x_nr+1)-(end-start)))))
                            print(" Projected time left for all toys: {0}".format(display_time(int( total_estim - (end-_start) ))))
                            sys.stdout.write("\033[F")
                            sys.stdout.write("\x1b[2K")
                            sys.stdout.write("\033[F")
                            sys.stdout.write("\x1b[2K")
                            sys.stdout.write("\033[F")
                            sys.stdout.write("\x1b[2K")
                            sys.stdout.write("\033[F")
                            sys.stdout.write("\x1b[2K")
                            sys.stdout.write("\033[F")
                            sys.stdout.write("\x1b[2K")
                        if o >=100:
                            print(" Time to generate this toy set: {0}".format(display_time(int(end-start))))

                        if counter_x_nr == verbose_calls[o]:
                            o += 1

                        counter += 1

                _end = time.time()

                print(" {0} out of {1} particles chosen!".format(int(counter_x), counter))

                print(" Toy set {0} generated!".format(toy_nr))

                #Save the set

                if save:

                    part_set = {"x_part" : x_part, "counter_x_nr": counter_x_nr, "x_min": x_min, "x_max": x_max, "mode": mode}

                    pkl.dump( part_set, open("./data/true_data/true_data_toy_{0}_range({1}-{2}).pkl".format(toy_nr ,int(x_min), int(x_max)) , "wb" ) )

                    print(" Toy set {0} saved!".format(toy_nr))

                print

            #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)

            return part_set
        elif mode == "fast_binned":

            #Calculate number of bins

            nbins = int((x_max-x_min)/resolution)

            for toy_nr in range(nr_of_toys):

                print("Generating toy set {1} of {2} with {0} bins...".format(nbins, toy_nr, nr_of_toys))

                #Get the mean dq values for each bin and get value of pdf there

                dq = np.linspace(x_min, x_max ,nbins+1)

                bin_mean = []
                bin_true_height = []

                for i in range(len(dq)-1):
                    a = dq[i]
                    b = dq[i+1]
                    bin_mean.append((a+b)/2)
                    _a, _e = integrate.quad(self.total_pdf, a**2, b**2, limit = 200)
                    height = _a/(b-a)
                    bin_true_height.append(height)

                #Vary the nonres_set_size

                _nonres_set_size = int(ROOT.TRandom1().Gaus(nonres_set_size, np.sqrt(nonres_set_size)))

                #Scale true bin size according to nonres_set_size

                _area, _area_err = integrate.quad(lambda x: self.total_nonres(x ,absolut = True), self.x_min**2, self.x_max**2, limit = 3000)

                _mean_scan = _area/(self.x_max - self.x_min)

                _mean_histo_nonres = _nonres_set_size/nbins

                _ =  _mean_histo_nonres/_mean_scan

                # print(_)

                for i in range(len(bin_true_height)):
                    bin_true_height[i] = bin_true_height[i]*_

                #Start time

                start = time.time()

                bin_height = []

                #Draw random numbers

                for i in range(len(bin_mean)):
                    bin_height.append(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i]))))
                    #print(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i]))))

                #progress calls
                if verbose != 0:
                    end = time.time()
                    print(" Time to generate set: {0}s".format(end-start))

                print(" {0} bins simulated".format(nbins))

                print(" Set generated!")

                #Save the set

                if save:

                    _ = 0

                    for i in bin_height:
                        _ += i

                    part_set = {"bin_mean" : bin_mean, "bin_height": bin_height, "nbins": nbins, "bin_true_height": bin_true_height, "_nonres_set_size": nonres_set_size, "x_min": x_min, "x_max": x_max}

                    pkl.dump( part_set, open("./data/fast_binned/binned_toy_{0}.pkl".format(toy_nr) , "wb" ) )

                    print(" Set saved!")

                print

            return part_set

    def draw_plots(self, part_set, mode, x_min = 2* mmu, x_max = (mB - mK) - 0.1, resolution = 7.0): #draws and saves plots, returns nothing
        #Resolution based on experiment chosen to be ~7MeV
        #check for invalid entries

        if mode != "slim_points" and mode != "fast_binned" and mode != "true_data" and mode != "no_data":
            raise ValueError('Wrong mode entered, choose either slim_points, fast_binned, true_data or no_data' )
        if mode != self.mode:
            raise ValueError('self.mode and mode are different, set them to the same value')

        print("Generating plots")

        #Load variables

        mB = self.mB
        mK = self.mK
        mmu = self.mmu

        #Range of function in MeV

        #Choosing the number of steps greatly varies the mean of dgamma!!! 5000 -> 10000 yields a factor of 1.68
        dq = np.linspace(x_min, x_max ,self.reco_steps)

        #Translate to MeV**2

        dq2 = np.power(dq, 2)

        #calculate formfactors

        ff_0 = self.formfactor(dq2, "0")
        ff_T = self.formfactor(dq2, "T")
        ff_plus = self.formfactor(dq2, "+")

        #calculate nonresonant

        dgamma_axiv_nonres = self.axiv_nonres(dq2)
        dgamma_vec_nonres = self.vec_nonres(dq2)
        dgamma_total_nonres = self.total_nonres(dq2, absolut = True)
        dgamma_tot = self.total_pdf(dq2)


        #Plot formfactors

        plt.clf()

        plt.plot(dq2, ff_0, label = "0")
        plt.plot(dq2, ff_T, label = "T")
        plt.plot(dq2, ff_plus, label = "+")

        prepare_plot("Formfactors")

        plt.savefig("./plots/{0}/ff.png".format(mode))

        print(" ff.png created")

        #Plot nonresonant part

        plt.clf()

        plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
        plt.plot(dq, dgamma_vec_nonres, label = "vec")
        plt.plot(dq, dgamma_total_nonres, label = "total_nonres")

        prepare_plot("Nonresonant parts")

        plt.savefig("./plots/{0}/vec_axiv.png".format(mode))

        print(" vec_axiv.png created")

        #Plot all pdfs and parts

        jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]

        psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]

        total_y = np.sum([pdf(dq2, self.param_val) for pdf in self.total_pdf_list])

        tot_y = self.total_pdf(dq2)

        plt.clf()

        _ = 0
        for pdf in self.total_pdf_list:
            name = self.total_pdf_names[_]
            _pdf = pdf(dq2, param_val = self.param_val, absolut = True)
            plt.plot(dq, _pdf, label = name)
            _ += 1

        plt.plot(dq, tot_y, label = "total pdf")

        prepare_plot("All pdfs")

        # plt.ylim(0, 2*self.param_val[1])

        plt.savefig("./plots/{0}/pdf_and_parts.png".format(mode))

        print(" pdf_and_parts.png created")

            #Create histo with pdf

        if mode == "fast_binned":

            #Load particle set

            bin_mean = part_set["bin_mean"]
            bin_height = part_set["bin_height"]

            #Scale the pdf to the histo

            _sum = np.sum(bin_height)

            _mean = np.mean(bin_height)

            nbins = part_set["nbins"]
            #
            # _ = self.param_val[

            area, err = integrate.quad(self.total_pdf, x_min**2, x_max**2, limit = 300)
            # _m = area/(x_max-x_min)
            # self.param_val[0] = self.param_val[0]*_mean/_m

            plt.clf()

            # plt.ylim(0.0, 2.5*_mean)

            #Choose range in between the 2 resonances
            # plt.xlim(3150, 3650)

            #Draw the plot

            plt.hist(bin_mean, bins=nbins, range=(x_min, x_max), weights = bin_height, label = "toy data binned")

            plt.plot(dq, dgamma_tot, label = "pdf")

            prepare_plot("{0} random points generated ({1} particles)".format(len(bin_mean), int(_sum)))

            plt.savefig("./plots/fast_binned/histo.png")

            # self.param_val[0] = _

            print(" histo.png created")

            print(" All plots drawn \n")

            return

        elif mode == "slim_points" or mode == "true_data":

            #Particle set

            x_part = part_set["x_part"]

            set_size = len(x_part)

            #Histo unnormalized

            bins = int((x_max-x_min)/resolution)

            plt.clf()

            _y, _x, _ = plt.hist(x_part, bins=bins, range=(x_min, x_max), label = "toy data binned ({0} points)".format(set_size))

            _mean_histo = float(np.mean(_y))

            plt.ylim(0.0, 2.5*_mean_histo)

            prepare_plot("Binned toy data")

            plt.savefig("./plots/points/histo_raw.png")

            print(" histo_raw.png created")


            #Histo and pdf normailzed

            plt.clf()

            #Plot histo

            #Scale the pdf to the histo

            _ = []

            # __ = self.param_val[0]

            for i in range(len(_y)):
                _.append(x_min+(x_max-x_min)/bins*i)


            plt.ylim(0.0, 2.5*_mean_histo)

            plt.hist(_, bins=bins, range=(x_min, x_max), weights = _y, label = "toy data binned")

            # plt.plot(dq, _dgamma_tot, label = "pdf scaled up to match histo")

            # self.param_val[0] = __

            #Only show from 0 to twice the mean of pdf
            # plt.ylim(0, 2*self.param_val[1])

            prepare_plot("{0} random points generated according to pdf".format(set_size))

            plt.savefig("./plots/points/histo.png")

            print(" histo.png created")

            print(" All plots drawn \n")

            return

    def total_pdf(self, q2): #return absolut value of sum of all complex values (real number)

        #Calculate the pdf with the added resonances
        _sum = 0
        for pdf in self.total_pdf_list:
            _sum += pdf(q2, param_val = self.param_val)

        return np.abs(_sum)

    def resonance(self, q2, parameters, name = "", absolut = False): #returns complex value

        _mass, width, phase, scale =  parameters[name + " mass"], parameters[name + " width"], parameters[name + " phase"], parameters[name + " scale"]

        #Helper variables

        p = 0.5 * np.sqrt(q2 - 4*(mmu**2))

        p0 =  0.5 * np.sqrt(_mass**2 - 4*mmu**2)

        gamma_j = np.divide(p, q2) * _mass * width / p0

        #Calculate the resonance

        _top = np.complex(_mass * width, 0.0)


        _bottom = np.vectorize(complex)(_mass**2 - q2, -_mass*gamma_j)

        com = _top/_bottom * scale

        #Rotate by the phase

        r = np.abs(com)

        _phase = np.angle(com)

        _phase += phase

        x = np.cos(phase)*r
        y = np.sin(phase)*r

        com = np.vectorize(complex)(x, y)

        if absolut:
            return np.abs(com)
        else:
            return com

    def add_resonance(self, _mass, width, phase, scale, name, nr_of_part): #adds string to total_pdf and adds constant places for fit

        #Adds the resonace to the pdf in form of a string (to be executed later)

        self.param_val[name + " mass"] = _mass
        self.param_val[name + " width"] = width
        self.param_val[name + " phase"] = phase
        self.param_val[name + " scale"] = scale
        self.nr_of_part_list[name] = nr_of_part

        self.param_val_orig = self.param_val

        def reso_func(q2, param_val, absolut = False):
            return self.resonance(q2, parameters = param_val, name = name, absolut = absolut)
        self.total_pdf_names += [name]
        self.total_pdf_list += [reso_func]

        return

    def bifur_gauss(self, q2, parameters, name = "", absolut = False): #returns complex value

        mean, scale, sigma_L, sigma_R = parameters[name + " mass"], parameters[name + " scale"], parameters[name + " sigma_L"], parameters[name + " sigma_R"]

        #Get q out of q2

        q = np.sqrt(q2)

        #Choose the right sigma depending on side of the cusp

        q_left = np.where(q < mean, q, False)

        q_right = np.where(q >= mean, q, False)

        #Calculate the exponential part of the cusp

        _exp_left = np.exp(- np.power((q_left-mean),2) / (2 * sigma_L**2))
        _exp_right = np.exp(- np.power((q_right-mean),2) / (2 * sigma_R**2))


        #Scale so the total area under curve is 1 and the top of the cusp is continuous

        _exp = _exp_left + _exp_right

        dgamma = scale*_exp/(np.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R)

        com = np.vectorize(complex)(dgamma, 0)


        # _sum = 0
        #
        # factor = 1.0
        # for nr in self.nr_of_part_list:
        #     _sum += nr
        #
        # factor = self.nr_of_part_list[name]/_sum

        # com *= factor

        if absolut:
            return np.abs(com)
        else:
            return com

    def add_cusp(self, mean, amp, sigma_L, sigma_R, nr_of_part): #adds string to total_pdf and adds constant places for fit

        name = "cusp"

        #Add cusp to total_pdf

        def reso_func(q2, param_val, absolut = False):
            return self.bifur_gauss(q2, parameters = param_val, name = name, absolut = absolut)
        self.total_pdf_list += [reso_func]
        self.total_pdf_names += [name]

        #Save the variables of the cusp -> Used for fitting

        self.param_val[name + " mass"] = mean
        self.param_val[name + " sigma_L"] = sigma_L
        self.param_val[name + " sigma_R"] = sigma_R
        self.param_val[name + " scale"] = amp
        self.nr_of_part_list[name] = nr_of_part

        self.param_val_orig = self.param_val

    def normalize_pdf(self, value = 1.0, verbose = 1): #Normalizes pdf with a global factor in front, saves mean and global factor

        if verbose == 1:
            print("Normalizing pdf...")

        area, err = integrate.quad(self.total_pdf, self.x_min**2, self.x_max**2, limit = 3000)

        _mean = area/(self.x_max-self.x_min)

        self.param_val[0] = self.param_val[0]/area*value

        self.param_val[1] = _mean/area*value

        self.param_val_orig = self.param_val

        if verbose == 1:
            print(" pdf normalized!")
            print

    def param_list(self, orig = True): #prints a list of all parameters of the form: number, name, value

        if orig:
            print("List of parameters (true values)")
            dic = self.param_val_orig
            i = 0
        else:
            print("List of parameters (fitted values)")
            dic = self.param_val
            i = 0

        print(" Nr. Name Value")

        for key in dic.keys():
            i += 1
            print(" {0}. {1}: {2}".format(i, key, dic[key]))

        print

    def add_nonres(self):

        nr_of_part = 44000
        name = "Nonres"

        def nonres_func(q2, param_val, absolut = False):
            return self.total_nonres(q2, parameters = param_val, name = name, absolut = absolut)
        self.total_pdf_list += [nonres_func]
        print("Func added")

        # self.nr_of_part_listnr_of_part)
        self.total_pdf_names += [name]

    def log_likelihood(self, npar, apar, part_set, fitting_param_index, resolution = 7.0): #Replaced soon with TMinuit

        # self.param_val = apar

        for i in range(len(fitting_param_index)):
            self.param_val[fitting_param_index[i]] = apar[i]



        self.normalize_pdf(verbose = 0)

        # part_set = {"bin_height": bin_height, "bin_mean": bin_mean, "_x": _x, "nbins": nbins, "x_max": x_max, "x_min": x_min}
        _x = part_set["_x"]
        bin_height = part_set["bin_height"]

        ll = 0.0

        for i in range(len(bin_height)):
            area, _ = integrate.quad(self.total_pdf, _x[i]**2, _x[i+1]**2, limit = 300)
            y_true = area/(_x[i+1]-_x[i])

            ll += np.log(bin_height[i]*y_true)

        ll = -ll

        return ll

    # def fit_pdf_to_data(self, part_set, fitting_param_index = [2,3,4,5,9,13], unbinned_data = True, resolution = 0):

        # # x_part = part_set["x_part"]
        #
        # if resolution == 0:
        #     resolution = self.resolution
        #
        # if unbinned_data:
        #
        #     #Load data
        #
        #     x_part = part_set["x_part"]
        #     x_min = part_set["x_min"]
        #     x_max = part_set["x_max"]
        #
        #     #Fill data into bins
        #
        #     nbins = int((x_max-x_min)/resolution)
        #
        #     bin_height , _x, _ = plt.hist(x_part, bins=nbins, range=(x_min, x_max))
        #     plt.clf()
        #
        #     bin_mean = []
        #
        #     for i in range(len(_x)-1):
        #         bin_mean.append((_x[i]+_x[i+1])/2)
        #
        # _part_set = {"bin_height": bin_height, "bin_mean": bin_mean, "_x": _x, "nbins": nbins, "x_max": x_max, "x_min": x_min}
        #
        # x_part = arr("f", x_part)
        #
        # nPoints = len(x_part)
        #
        # name = self.param_str
        #
        # npar = len(fitting_param_index)
        #
        # vstart = arr("d", npar*[0])
        #
        # # print(len(vstart))
        #
        # amp_err = self.param_val[0]/20
        #
        # cusp_m_err = 10
        #
        # step   = arr( 'd', ( amp_err, cusp_m_err, 1.0, 0.1, 1e-7, self.param_val[9]/50, self.param_val[13] ))
        # npar = len(vstart)
        #
        # def fcn(npar, deriv, f, apar, iflag):
        #     f[0] = self.log_likelihood(npar, apar, _part_set, resolution = 7.0, fitting_param_index = fitting_param_index)
        #     # return f[0]
        #
        # gMinuit = ROOT.TMinuit(npar)
        # gMinuit.SetFCN( fcn )
        #
        # gMinuit.SetErrorDef(0.5)
        #
        # arglist = arr( 'd', npar*[0.] )
        # ierflg = ROOT.Long(1998)
        #
        # arglist[0] = 1
        # gMinuit.mnexcm( "SET ERR", arglist, 1, ierflg )
        #
        # _ = 0
        #
        # for j in fitting_param_index:
        #     # print(_, fitting_param_index[_])
        #     vstart[_] = self.param_val[fitting_param_index[_]]
        #     if _ == 1:
        #         vstart[_] -= 50
        #     gMinuit.mnparm( _, self.param_str[fitting_param_index[_]], vstart[_], step[_], 0, 0, ierflg )
        #     _ += 1
        #
        # arglist[0] = 6000
        # arglist[1] = 1.
        #
        # gMinuit.mnexcm( "MIGRAD", arglist, 2, ierflg )
        #
        # amin, edm, errdef = 0.18, 0.19, 0.20
        # nvpar, nparx, icstat = 1983, 1984, 1985