#!/usr/bin/env python
# coding: utf-8

# # Fit with Orthogonal distance regression
# 
# Example from http://stackoverflow.com/questions/26058792/correct-fitting-with-scipy-curve-fit-including-errors-in-x/26085136#26085136

# In[1]:


# manage data and fit
import pandas as pd
import numpy as np

# scipy ODR : orthogonal distance regression
from scipy.odr import ODR, Model, Data, RealData

# style and notebook integration of the plots
import matplotlib.pyplot as plt 
import seaborn as sns
get_ipython().run_line_magic('matplotlib', 'inline')
sns.set(style="whitegrid")


# In[2]:


def func(beta, x):
    y = beta[0]+beta[1]*x+beta[2]*x**3
    return y


# ## generate data

# In[3]:


npts = 100
np.random.seed(160810)
x = np.linspace(-3, 2, npts)
y = func([-2.3, 7.0, -4.0], x)

sigma_x = .3
sigma_y = 2

# add some noise
x += np.random.normal(scale=sigma_x, size=100)
y += np.random.normal(scale=sigma_y, size=100)


# ## Define model

# In[4]:


data = RealData(x, y, 0.3, 0.1)
model = Model(func)
odr = ODR(data, model, [1,0,0])


# ## Least squarre fit

# In[5]:


odr.set_job(fit_type=0)
lsq_output = odr.run()
print("         stop reason:", lsq_output.stopreason)
print("              params:", lsq_output.beta)
print("                info:", lsq_output.info)
print("Sum of squares error:", lsq_output.sum_square)
print("             sd_beta:", lsq_output.sd_beta)
print("      sqrt(diag(cov):", np.sqrt(np.diag(lsq_output.cov_beta)))

# check convergence and run again the algorithm if it is not reached
if lsq_output.info != 1:
    print("\nRestart ODR till convergence is reached")
    i = 1
    while lsq_output.info != 1 and i < 100:
        print("restart %3d %12.7f" % (i, lsq_output.sum_square))
        lsq_output = odr.restart()
        i += 1
    print("         stop reason:", lsq_output.stopreason)
    print("              params:", lsq_output.beta)
    print("                info:", lsq_output.info)
    print("Sum of squares error:", lsq_output.sum_square)
    print("             sd_beta:", lsq_output.sd_beta)
    print("      sqrt(diag(cov):", np.sqrt(np.diag(lsq_output.cov_beta)))


# In[6]:


lsq_output.pprint()


# ## ODR fit

# In[7]:


odr = ODR(data, model, [1,0,0])
odr.set_job(fit_type=2)
odr_output = odr.run()
print("         stop reason:", odr_output.stopreason)
print("              params:", odr_output.beta)
print("                info:", odr_output.info)
print("Sum of squares error:", odr_output.sum_square)
print("             sd_beta:", odr_output.sd_beta)
print("      sqrt(diag(cov):", np.sqrt(np.diag(odr_output.cov_beta)))

# check convergence and run again the algorithm if it is not reached
if odr_output.info != 1:
    print("\nRestart ODR till convergence is reached")
    i = 1
    while odr_output.info != 1 and i < 100:
        print("restart", i)
        odr_output = odr.restart()
        i += 1
    print("         stop reason:", odr_output.stopreason)
    print("              params:", odr_output.beta)
    print("                info:", odr_output.info)
    print("Sum of squares error:", odr_output.sum_square)
    print("             sd_beta:", odr_output.sd_beta)
    print("      sqrt(diag(cov):", np.sqrt(np.diag(odr_output.cov_beta)))


# ## Plot

# In[8]:


xn = np.linspace(-3, 2, 50)

plt.errorbar(x, y, marker="o", linestyle="", label="data", xerr=sigma_x, yerr=sigma_y, elinewidth=1, capthick=1)
plt.plot(xn, func(lsq_output.beta, xn), label='leastsq')
plt.plot(xn, func(odr_output.beta, xn), label='odr')
plt.legend()