"""
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.)

Uses: Parameter class: see scipy Cookbook/FittingData
"""
# set numpy as numeric module (not needed in later pylabs)
# import matplotlib
# matplotlib.rcParams["numerix"]="numpy"
from pylab import *
from scipy.optimize import leastsq

# From  www.scipy.org/Cookbook/FittingData :
# Modified to add uncertainties, full output
# =======================================
class Parameter:
    def __init__(self, initialvalue, name):
        self.value = initialvalue
        self.name=name
    def set(self, value):
        self.value = value
    def __call__(self):
        return self.value
 
def fit(function, parameters, x, data, u):
    def fitfun(params):
        for i,p in enumerate(parameters):
            p.set(params[i])
        return (data - function(x))/u
 
    if x is None: x = arange(data.shape[0])
    if u is None: u = ones(data.shape[0],"float")
    p = [param() for param in parameters]
    return leastsq(fitfun, p, full_output=1)
# ========================================

# define function to be fitted
def resonance(f):
    R=214.0
    return (R/R0())/sqrt(1.0+Q()**2*(f/w0()-w0()/f)**2)

# read data
freq, vr, dvr=load('lcr.dat', unpack=True)

# the fit parameters: some starting values
R0=Parameter(400.0,"R")
w0=Parameter(4000.0,"f0")
Q=Parameter(1.0,"Q")

p=[R0,w0,Q]

# for theory plot we need some frequencies
freqs=linspace(2000.0,11000.0,100)
initialplot=resonance(freqs)

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

# do fit using Levenberg-Marquardt
p2,cov,info,mesg,success=fit(resonance, p, freq, vr/v0, uvr)

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

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

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

# uncertainties are calculated as per gnuplot, "fixing" the result
# for non unit values of the reduced chisq.
# values at min match gnuplot
print "Fitted parameters at minimum, with 68% C.I.:"
for i,pmin in enumerate(p2):
    print "%2i %-10s %12f +/- %10f"%(i,p[i].name,pmin,sqrt(cov[i,i])*sqrt(chisq/dof))
print

print "Correlation matrix"
# correlation matrix close to gnuplot
print "               ",
for i in range(len(p)): print "%-10s"%(p[i].name,),
print
for i in range(len(p2)):
    print "%10s"%p[i].name,
    for j in range(i+1):
        print "%10f"%(cov[i,j]/sqrt(cov[i,i]*cov[j,j]),),
    print

# Plot

# plot data
errorbar(freq, vr/v0, yerr=uvr, fmt='r+', label="Data")
# plot fit
plot(freqs, initialplot, 'g-', label="Start")
plot(freqs, resonance(freqs), 'b-', label="Fit")
xlabel("Frequency [Hz]")
ylabel("Relative amplitude")
#legend(("data","fit"))
legend()
show()