analyze.py

#!/usr/bin/env python
'''
    analyze.py
    Eric Ayars
    8/14/18

    An attempt to do curve fitting on 'ellipse' data and see if the 
    assignment is reasonable.

    Also demonstrates how to plot residuals. 
    Also calculates Chi-squared.

'''

from sys import argv
from pylab import *
from scipy.optimize import curve_fit   

def gaussian(x,c,w,a):
    return a*exp(-(x-c)**2/(2.0*w**2))

def lorentzian(x,c,w,a):
    return a*w**2/((x-c)**2 + w**2)

def student(x,c,w,a):
    return a*(1+(x-c)**2/w)**(-(w+1)/2.0)

def cosine2(x,c,w,a):
    return a*cos((x-c)*pi/(5*w))**2

Scale = 34.1    #size of calibration square, in mm
fn = student # pick one, put it here.

# load data
n, x, y, dx, dy = loadtxt(argv[1], unpack=True)

# convert to 'units'
x = x*0.5/Scale
y = y*0.5/Scale
dx = dx*0.5/Scale
dy = dy*0.5/Scale

# do the curve fit, starting with a guess
center = 0.65
width = 0.022
amplitude = 1.0
guess = (center, width, amplitude)

params, cov = curve_fit(fn, x,y, guess, sigma=dy)

perr = sqrt(diag(cov))
print(params)
print(perr)

# plot the data
fig = figure(1)
dataFrame = fig.add_axes((0.1,0.3,0.8,0.6))
errorbar(x,y,xerr=dx, yerr=dy, fmt='r.')
title('Data for %s' % argv[1])
ylabel('y (units)')
xticks([])

# plot curve fit
X = linspace(0,1,100)
Y = fn(X,*params)
plot(X,Y,'g-')

# residual plot
difference = y - fn(x,*params) 
residualFrame = fig.add_axes((0.1,0.1,0.8,0.18))
errorbar(x,difference, yerr=dy, fmt='r.')
yticks([])

# zero on residual plot
plot(X,[0 for i in X],'b-')
xlabel('x (units)')

# reduced chi-squared
print(sum((difference/dy)**2)/(len(difference) - 3))

savefig("pretty_graph.pdf")
show()