#!/usr/bin/env python
import Image
import numpy
import numpy.ma
import scipy.optimize
import pylab
x = numpy.arange(0,10,.1)
y = numpy.ones_like(x)
z = y.copy()
z[40:60] = numpy.nan
y[40:60] = 2
a = numpy.ma.masked_array(y.copy(),y-1)
fitf = lambda u,x: u[0]*x+u[1]
errf = lambda u,x,y: y - fitf(u,x)
u = numpy.array([0,1])
p = scipy.optimize.leastsq(errf,u,args=(x,y))
q = scipy.optimize.leastsq(errf,u,args=(x,z))
r = scipy.optimize.leastsq(errf,u,args=(x,a))
print p,q,r
pylab.figure(figsize=(3,6))
pylab.subplot(311)
pylab.title("Fit of entire array")
pylab.plot(x,y,'ro',alpha=.5)
pylab.plot(x,fitf(p[0],x),'g--')
pylab.subplot(312)
pylab.title("Fit of array with NaNs")
pylab.plot(x,y,'ro',alpha=.5)
pylab.plot(x,fitf(q[0],x),'b--',lw=2)
pylab.subplot(313)
pylab.title("Fit of masked array")
pylab.plot(x,y,'ro',alpha=.5)
pylab.plot(x,fitf(r[0],x),'y--',lw=2)
pylab.show()