# -*- coding: utf-8 -*-
# <nbformat>3.0</nbformat>

# <headingcell level=1>

# probfit Basic Tutorial

# <markdowncell>

# [probfit](http://iminuit.github.io/probfit/) is a modeling / fitting package to be used together with [iminuit](http://iminuit.github.com/iminuit/).
# 
# This tutorial is a fast-paced introduction to the probfit features:
# 
# * built-in common models: polynomial, gaussian, ...
# * build-in common fit statistics: chi^2, binned and unbinned likelihood
# * tools to get your fits to converge and check the results: try_uml, draw, draw_residuals, ...
# * tools to help you implement your own models and fit statistics: Normalize, Extended, integrate_1d, ...
# 
# Please start this notebook with the ``ipython --pylab=inline`` option to get inline plots.

# <codecell>

# We assume you have executed this cell in all the following examples
import numpy as np
import matplotlib.pyplot as plt
import iminuit
import probfit

# <markdowncell>

# In your own code you can explicitly import what you need to save
# typing in interactive sessions, e.g.
# 
#     from iminuit import Minuit, describe
#     from probfit import gaussian, BinnedLH
# 
# We don't do this here, we only import `iminuit` and `probfit` into our
# namespace so that it is clear to you which functions and classes come
# from which package while reading the code below.

# <markdowncell>

# ## Chi^2 straight line fit
# 
# We can't really call this a fitting package without being able to fit a straight line, right?

# <codecell>

# Let's make a straight line with gaussian(mu=0, sigma=1) noise
np.random.seed(0)
x = np.linspace(0, 10, 20) 
y = 3 * x + 15 + np.random.randn(len(x))
err = np.ones(len(x))
plt.errorbar(x, y, err, fmt='.');

# <codecell>

# Let's define our line.
# First argument has to be the independent variable,
# arguments after that are shape parameters.
def line(x, m, c): # define it to be parabolic or whatever you like
    return m * x + c

# <codecell>

iminuit.describe(line)

# <codecell>

# Define a chi^2 cost function
chi2 = probfit.Chi2Regression(line, x, y, err)

# <codecell>

# Chi2Regression is just a callable object; nothing special about it
iminuit.describe(chi2)

# <codecell>

# minimize it
# yes, it gives you a heads up that you didn't give it initial value
# we can ignore it for now
minuit = iminuit.Minuit(chi2) # see iminuit tutorial on how to give initial value/range/error
minuit.migrad(); # MIGRAD is a very stable robust minimization method
# you can look at your terminal to see what it is doing;

# <codecell>

# The output above is a pretty-printed summary of the fit results from
# minuit.print_fmin()
# which was automatically called by iminuit.Minuit.migrad() after running MIGRAD.

# Let's see our results as Python dictionaries ...
print(minuit.values)
print(minuit.errors)

# <markdowncell>

# #### Parabolic error
# is calculated using the second derivative at the minimum
# This is good in most cases where the uncertainty is symmetric not much correlation
# exists. Migrad usually got this accurately but if you want ot be sure
# call `minuit.hesse()` after calling `minuit.migrad()`.
# 
# #### Minos Error
# is obtained by scanning the chi^2 or likelihood profile and find the point
# where chi^2 is increased by `minuit.errordef`. Note that in the Minuit documentation
# and output `errordef` is often called `up` ... it's the same thing.
# 
# #### What `errordef` should I use?
# 
# As explained in the Minuit documentation you should use:
# 
# * `errordef = 1` for chi^2 fits
# * `errordef = 0.5` for likelihood fits
# 
# `errordef=1` is the default, so you only have to set it to `errordef=0.5`
# if you are defining a likelihood cost function (if you don't your HESSE and MINOS errors will be incorrect).
# `probfit` helps you by defining a `default_errordef()` attribute on the
# cost function classes, which is automatically detected by the `Minuit` constructor
# and can be used to set `Minuit.errordef` correctly, so that users can't forget.
# Classes used in this tutorial:
# 
# * `probfit.Chi2Regression.get_errordef()` and `probfit.BinnedChi2.get_errordef()` return 1.
# * `probfit.BinnedLH.get_errordef()` and `probfit.UnbinnedLH.get_errordef()` return 0.5.

# <codecell>

# Let's visualize our line
chi2.draw(minuit)
# looks good;

# <codecell>

# Sometimes we want the error matrix (a.k.a. covariance matrix)
print('error matrix:')
print(minuit.matrix())
# or the correlation matrix
print('correlation matrix:')
print(minuit.matrix(correlation=True))
# or a pretty html representation
# Note that `print_matrix()` shows the correlation matrix, not the error matrix
minuit.print_matrix()

# <markdowncell>

# ## Binned Poisson likelihood fit of a Gaussian distribution
# In high energy physics, we usually want to fit a distribution to a histogram. Let's look at simple Gaussian distribution.

# <codecell>

# First let's make some example data
np.random.seed(0)
data = np.random.randn(10000) * 4 + 1
# sigma = 4 and mean = 1
plt.hist(data, bins=100, histtype='step');

# <codecell>

# Define your PDF / model
def gauss_pdf(x, mu, sigma):
    """Normalized Gaussian"""
    return 1 / np.sqrt(2 * np.pi) / sigma * np.exp(-(x - mu) ** 2 / 2. / sigma ** 2)

# <codecell>

# Build your cost function
# Here we use binned likelihood
binned_likelihood = probfit.BinnedLH(gauss_pdf, data)

# <codecell>

# Create the minuit
# and give an initial value for the sigma parameter
minuit = iminuit.Minuit(binned_likelihood, sigma=3)
# Remember: minuit.errordef is automatically set to 0.5
# as required for likelihood fits (this was explained above)
binned_likelihood.draw(minuit);

# <codecell>

minuit.migrad()
# Like in all binned fit with long zero tail. It will have to do something about the zero bin
# probfit.BinnedLH does handle them gracefully but will give you a warning;

# <codecell>

# Visually check if the fit succeeded by plotting the model over the data
binned_likelihood.draw(minuit) # uncertainty is given by symmetric Poisson;

# <codecell>

# Let's see the result
print('Value: {}'.format(minuit.values))
print('Error: {}'.format(minuit.errors))

# <codecell>

# That printout can get out of hand quickly
minuit.print_fmin()
# Also print the correlation matrix
minuit.print_matrix()

# <codecell>

# Looking at a likelihood profile is a good method
# to check that the reported errors make sense
minuit.draw_mnprofile('mu');

# <codecell>

# Plot a 2d contour error
# You can notice that it takes some time to draw
# We will this is because our PDF is defined in Python
# We will show how to speed this up later
minuit.draw_mncontour('mu', 'sigma');

# <markdowncell>

# ## Chi^2 fit of a Gaussian distribution
# 
# Let's explore another popular cost function chi^2.
# Chi^2 is bad when you have bin with 0.
# ROOT just ignore.
# ROOFIT does something I don't remember.
# But it's best to avoid using chi^2 when you have bin with 0 count.

# <codecell>

# We will use the same data as in the previous example
np.random.seed(0)
data = np.random.randn(10000) * 4 + 1
# sigma = 4 and mean = 1
plt.hist(data, bins=100, histtype='step');

# <codecell>

# We will use the same PDF as in the previous example
def gauss_pdf(x, mu, sigma):
    """Normalized Gaussian"""
    return 1 / np.sqrt(2 * np.pi) / sigma * np.exp(-(x - mu) **2 / 2. / sigma ** 2)

# <codecell>

# Binned chi^2 fit only makes sense (for now) for extended PDFs
# probfit.Extended adds a norm parameter with name 'N'
extended_gauss_pdf = probfit.Extended(gauss_pdf)

# <codecell>

# Describe the function signature
iminuit.describe(extended_gauss_pdf)

# <codecell>

# Chi^2 distribution fit is really bad for distribution with long tail
# since when bin count=0... poisson error=0 and blows up chi^2
# so give it some range
chi2 = probfit.BinnedChi2(extended_gauss_pdf, data, bound=(-7,10))
# This time we use the pedantic=False option to tell Minuit
# that we don't want warnings about parameters without initial
# value or step size.
# And print_level=0 means that no output is generated
minuit = iminuit.Minuit(chi2, sigma=1, pedantic=False, print_level=0)
minuit.migrad();

# <codecell>

# Now let's look at the results
minuit.print_fmin()
minuit.print_matrix()
chi2.draw(minuit);

# <markdowncell>

# ## Fast unbinned likelihood fit Cython
# 
# Unbinned likelihood is computationally very very expensive if you have a lot of data.
# It's now a good time that we talk about how to speed things up with [Cython](http://cython.org).

# <codecell>

# We will use the same data as in the previous example
np.random.seed(0)
data = np.random.randn(10000) * 4 + 1
# sigma = 4 and mean = 1
plt.hist(data, bins=100, histtype='step');

# <codecell>

# We want to speed things up with Cython
%load_ext cythonmagic

# <codecell>

%%cython
# Same gaussian distribution but now written in Cython
# The %%cython IPython does the following:
# * Call Cython to generate C code for a Python C extension.
# * Compile it into a Python C extension (a shared library)
# * Load it into the current namespace
# If you don't understand these things, don't worry, it basically means:
# * Get full-metal speed easily
cimport cython
from libc.math cimport exp, M_PI, sqrt
@cython.binding(True) # IMPORTANT: this tells Cython to dump the function signature
def gauss_pdf_cython(double x, double mu, double sigma):
    return 1 / sqrt(2 * M_PI) / sigma * exp(-(x - mu) ** 2 / 2. / sigma ** 2)

# <codecell>

# Define the unbinned likelihood cost function 
unbinned_likelihood = probfit.UnbinnedLH(gauss_pdf_cython, data)

# <codecell>

minuit = iminuit.Minuit(unbinned_likelihood, sigma=2, pedantic=False, print_level=0)
# Remember: minuit.errordef is automatically set to 0.5
# as required for likelihood fits (this was explained above)
minuit.migrad() # yes: amazingly fast
unbinned_likelihood.show(minuit)
minuit.print_fmin()
minuit.print_matrix() 

# <codecell>

# Remember how slow draw_mnprofile() was in the last example?
# Now it's super fast (even though the unbinned
# likelihood computation is more compute-intensive).
minuit.draw_mnprofile('mu');

# <markdowncell>

# But you really don't have to write your own gaussian, there are tons of builtin functions written in Cython for you.

# <codecell>

# Here's how you can list them
import probfit.pdf
print(dir(probfit.pdf))
print(iminuit.describe(probfit.pdf.gaussian))