#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jun  2 22:30:26 2018

@author: pilou
"""
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as spo
from scipy.stats import linregress

# données expérimentales
x = np.array([1.0,2.1,3.1,3.9,5.0,6.2,6.9])
y = np.array([2.1,3.9,6.1,7.8,9.8,11.7,14.2])
# incertitudes-types sur les données expérimentales
ux = np.array([0.1,0.1,0.1,0.04,0.1,0.35,0.1])
uy = np.array([0.10,0.11,0.090,0.19,0.21,0.35,0.25])

# fonction f décrivant la courbe à ajuster aux données (toujours une droite!)
def f(x,p):
    a,b = p 
    return a*x+b
# dérivée de la fonction f par rapport à la variable de contrôle x
def Dx_f(x,p):
    a,b = p
    return a

# fonction d'écart pondérée par les erreurs
def residuals(p, y, x):
    return (y-f(x,p))/np.sqrt(uy**2 + (Dx_f(x,p)*ux)**2)

# estimation initiale des paramètres
# elle ne joue généralement aucun rôle
# néanmoins, le résultat de l'ajustement est parfois aberrant
# il faut alors choisir une meilleure estimation initiale
p0 = np.array([0,0])

# on utilise l'algorithme des moindres carrés non-linéaires 
# disponible dans la biliothèque scipy (et indirectement la
# bibliothèque Fortran MINPACK qui implémente l'algorithme
# de Levenberg-Marquardt) pour déterminer le minimum voulu
result = spo.leastsq(residuals, p0, args=(y, x), full_output=True)

# on obtient :
# les paramètres d'ajustement optimaux
print('Modélisation avec erreurs sur X et Y')
print("Paramètres optimaux (pente et ordonnée à l'origine)")
popt = result[0]
print(popt)

# la matrice de variance-covariance estimée des paramètres
pcov = result[1];
# les incertitudes-types sur ces paramètres
print('Erreurs sur les paramètres')
upopt = np.sqrt(np.abs(np.diagonal(pcov)))
print(upopt)

# calcul de la valeur du "chi2 réduit" pour les paramètres ajustés
chi2r = np.sum(np.square(residuals(popt,y,x)))/(x.size-popt.size)

#Comparaison avec la régression linéaire normale, qui donne le coefficient de correlation
sol=linregress(x,y)
pente=sol[0]
ordonnee = sol[1]
coefcor = sol[2]
coefcor2 = sol[2]**2

print("Comparaison avec une autre fonction de régression linéaire comme régressi, ne prenant pas en compte les incertitudes")
print("pente = ", pente, "ordonnee = ", ordonnee, "r^2 = ", coefcor2)


#on regarde ce que ça donne sur une courbe
plt.figure
plt.errorbar(x,y,xerr=ux,yerr=uy,fmt='b+')
plt.plot(x,f(x,popt),'r')
plt.xlabel('X (s)')
plt.ylabel('Y (m)')
plt.title('Fit avec erreurs sur X et Y')
plt.show()