week2/problem1/problem1.py

   1 '''
   2 AY250 HW2 Problem 1
   3 Fitting to noisy data
   4 Author:  Bryan Steinbach
   5 '''
   6
   7
   8 import numpy as np
   9 import scipy as sp
  10 from scipy import optimize,interpolate
  11 import pylab as pl
  12
  13 # Problem statement doesn't specify how many points or where they should go, beyond between 0 and 4
  14 # Pick 100 points evenly spaced
  15 n = 100
  16
  17 def model(p,x):
  18         return p[0]*np.exp(p[1]*x)+p[2]
  19
  20 sigma = 0.1
  21
  22 x = np.linspace(0,4,n)
  23
  24 def rms(x): return np.sqrt(np.dot(x,x)/x.size)
  25
  26 # The measurements
  27 ptrue = np.array([2,-0.75,0.1]) # a, alpha, k
  28 ynoise = sigma*np.random.randn(n)
  29 print "true parameters: ",ptrue
  30 print "noise standard deviation: ",np.std(ynoise), "expected ",sigma
  31 ysignal = model(ptrue,x)
  32 yd = ysignal+ynoise
  33
  34 # The analysis
  35 # Least squares fit
  36 p0 = np.array([1.,0.,min(yd)])
  37
  38 p,ier = optimize.leastsq(lambda p: model(p,x) - yd,p0)
  39
  40 yleastsq = model(p,x)
  41
  42 print "sp.optimize.leastsq ier: ",ier
  43
  44 print "leastsq parameters: ",p
  45
  46 print "leastsq rms error: ",rms(ysignal - yleastsq)
  47 print "leastsq rms resid: ",rms(yd - yleastsq)
  48
  49 # spline smooth
  50 tck = interpolate.splrep(x,yd,s=2.0)
  51 yspline = interpolate.splev(x,tck)
  52
  53 print "spline rms error: ",rms(ysignal - yspline)
  54 print "spline rms resid: ",rms(yd - yspline)
  55
  56 # poly fits
  57 npoly=3
  58 ypolys = []
  59 for i in range(npoly):
  60         poly = np.polyfit(x,yd,i)
  61
  62         ypoly = np.polyval(poly,x)
  63         ypolys.append(ypoly)
  64
  65         print "poly%d rms error: "%i,rms(ysignal - ypoly)
  66         print "poly%d rms resid: "%i,rms(yd - ypoly)
  67
  68
  69 # Plot results
  70 pl.plot(x,ysignal,label="signal")
  71 pl.plot(x,yd,label="signal+noise")
  72 pl.plot(x,yleastsq,label="least squares fit")
  73 pl.plot(x,yspline,label='spline smooth')
  74 for i in range(npoly):
  75         pl.plot(x,ypolys[i],label='poly%d'%i)
  76 pl.grid()
  77 pl.xlabel('x')
  78 pl.ylabel('y')
  79 pl.legend()
  80 pl.title('Fitting to $y = a \\exp(\\alpha x) + k + n(x)$')
  81
  82 pl.savefig('fits.png')
  83