1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
#!/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()