# coding: utf-8

""" Fit a line to data. """

import numpy as np
import scipy.stats as statistics, scipy.optimize as op

__all__ = ["Line"]

class Line(object):
    """ A straight line model. """
    
    def __init__(self, x, y, yerr, outliers=True):

        assert len(x) == len(y)
        assert len(yerr) == len(y)

        self.x = x
        self.y = y
        self.yerr = yerr
        
        self.parameters = ["m", "b"]
        self.outliers = outliers
        if outliers:
            self.parameters.extend(["Po", "Yo", "Vo"])
        return None


    def optimise(self, xtol=1e-4, ftol=1e-8, maxfun=10e3, maxiter=10e3, disp=False, **kwargs):
        """
        Optimise the model parameters given the data.
        """

        slope, intercept, r_value, p_value, stderr = statistics.linregress(self.x, self.y)

        p0 = [slope, intercept]
        if self.outliers:
            p0.extend([0.5, np.median(self.y), np.std(self.y)])

        # Create the likelihood function.
        def ln_like(theta):

            m, b = theta[:2]
            line = m * self.x + b

            ivar = 1.0/self.yerr**2
            line_likelihood = -0.5 * ((line - self.y)**2 * ivar - np.log(ivar))
            if not self.outliers:
                return np.sum(line_likelihood)

            Po, Yo, Vo = theta[2:]
            # A bit of prior information here:
            if not (1 > Po > 0) or 0 > Vo:
                return -np.inf

            outlier_ivar = 1.0/(self.yerr**2 + Vo)
            outlier_likelihood = -0.5 * ((Yo - self.y)**2 * outlier_ivar - np.log(outlier_ivar))
            likelihood = np.sum(np.logaddexp(np.log(1 - Po) + line_likelihood, np.log(Po) + outlier_likelihood))
            return likelihood

        # Optimise.
        minimise = lambda theta: -ln_like(theta)
        return op.fmin(minimise, p0, xtol=xtol, ftol=ftol, maxfun=maxfun, maxiter=maxiter, disp=False, **kwargs)
        
"""
x, y = np.loadtxt("data", unpack=True)
yerr = np.array([0.1] * len(y))
line_no_outliers = Line(x, y, yerr=yerr, outliers=False)
m1, b1 = line_no_outliers.optimise()

line_with_outliers = Line(x, y, yerr=yerr, outliers=True)
m2, b2, po, yo, vo = line_with_outliers.optimise()

plt.errorbar(x, y, yerr=yerr, fmt=None, c='k')
plt.scatter(x,y, facecolor='k')
plt.plot(x, m1 * x + b1, 'r')
plt.plot(x, m2 * x + b2, 'b')
"""



@@ -1911,21 +1911,58 @@
 gs)%0A
-        %0A%22%22%22%0A
+%0Aif __name__ == %22__main__%22:%0A%0A    # DEBUG MODE%0A    
 x, y
@@ -1996,16 +1996,20 @@
 k=True)%0A
+    
 yerr = n
@@ -2019,17 +2019,18 @@
 rray(%5B0.
-1
+05
 %5D * len(
@@ -2033,16 +2033,20 @@
 len(y))%0A
+    
 line_no_
@@ -2094,16 +2094,20 @@
 =False)%0A
+    
 m1, b1 =
@@ -2136,16 +2136,20 @@
 mise()%0A%0A
+    
 line_wit
@@ -2198,16 +2198,20 @@
 s=True)%0A
+    
 m2, b2, 
@@ -2254,16 +2254,20 @@
 mise()%0A%0A
+    
 plt.erro
@@ -2301,24 +2301,28 @@
 one, c='k')%0A
+    
 plt.scatter(
@@ -2341,16 +2341,20 @@
 or='k')%0A
+    
 plt.plot
@@ -2375,16 +2375,20 @@
 1, 'r')%0A
+    
 plt.plot
@@ -2409,14 +2409,10 @@
 2, 'b')%0A
-%22%22%22%0A
 %0A%0A

Debug mode