"""
Fit LRC data to resonance curve using Levenberg-Marquardt method.

Example demonstrates:

- Fitting of data to curve using Levenberg-Marquardt
  nonlinear least squares fit
- Calculation of fit statistics from output of fitting
  function:
  - Goodness of fit parameters
  - Confidence intervals in fit parameters
  - Correlations (covariance) between fit parameters
  Note that these confidence intervals are one-parameter estimates,
  as given in Numerical Recipes (3rd ed) Sec 15.6, eq 15.6.4
  (These uncertainties in fit parameters agree with those
  given by gnuplot.)

"""
# set numpy as numeric module (not needed in later pylabs)
# import matplotlib
# matplotlib.rcParams["numerix"]="numpy"
from pylab import *
from scipy.optimize import leastsq

# define function to be fitted
R=214.0
def resonance(p, f):
    return (R/p[0])/sqrt(1.0+p[2]**2*(f/p[1]-p[1]/f)**2)

# utilty driver to compute weighted residual
def errorfun(p, f, data, u):
    return (resonance(p,f)-data)/u

# read data
data=array([
    [ 2653, 1.08, 0.01 ],
    [ 2968, 1.28, 0.02 ],
    [ 3633, 1.92, 0.02 ],
    [ 4169, 2.86, 0.02 ],
    [ 4475, 3.70, 0.02 ],
    [ 4768, 4.90, 0.05 ],
    [ 5002, 6.05, 0.05 ],
    [ 5277, 7.45, 0.05 ],
    [ 5561, 7.80, 0.2 ],
    [ 5765, 7.00, 0.05 ],
    [ 6117, 5.55, 0.05 ],
    [ 6391, 4.55, 0.05 ],
    [ 6696, 3.76, 0.02 ],
    [ 7189, 2.95, 0.02 ],
    [ 7916, 2.24, 0.02 ],
    [ 8779, 1.75, 0.02 ],
    [ 9681, 1.45, 0.02 ],
    [ 10284, 1.27, 0.02 ]
    ])
freq=data.T[0]
vr=data.T[1]
dvr=data.T[2]


# the fit parameters
p=[214.0,5000.0,5.0]

# uncertainty calculation
v0=10.0
uvr=sqrt(dvr*dvr+vr*vr*0.0025)/v0

# do fit using Levenberg-Marquardt
p2,cov,info,mesg,success=leastsq(errorfun, p, args=(freq,vr/v0,uvr),full_output=1)

if success==1:
    print "Converged"
else:
    print "Not converged"
    print mesg


# calculate final chi square
chisq=sum(info["fvec"]*info["fvec"])


# chisq, sqrt(chisq/dof) agrees with gnuplot
print "Converged with chi squared ",chisq
print "RMS of residuals sqrt(chisq/dof) ", sqrt(chisq/(len(freq)-len(p)))
print "Reduced chisq (variance of residuals) ", chisq/(len(freq)-len(p))
print

# uncertainties are calculated as per gnuplot, "fixing" the result
# for non unit values of the reduced chisq.
eps=sqrt(chisq/(len(freq)-len(p)))
# values at min match gnuplot
print "Fitted parameters at minimum:"
for i in range(len(p2)):
    print i, p2[i], " +/- ", sqrt(cov[i,i])*sqrt(chisq/(len(freq)-len(p)))
print

print "Correlation matrix"
# correlation matrix close to gnuplot
for i in range(len(p2)):
    for j in range(i+1):
        print cov[i,j]/sqrt(cov[i,i]*cov[j,j]),
    print

# Plot

# for theory plot we need some frequencies
freqs=linspace(2000.0,11000.0,100)
# plot data
errorbar(freq, vr/v0, yerr=uvr, fmt='r+', label="Data")
# plot fit
plot(freqs, resonance(p2, freqs), 'b-', label="Fit")
xlabel("Frequency [Hz]")
ylabel("Relative amplitude")
#legend(("data","fit"))
legend()
show()