<filename>examples/fitting_data/t1_measurements/fit_exp_leastsq.py
#! /usr/bin/env python
# fit a collection to T1 trajectories to a decaying exponential

import scipy.optimize
import numpy as np
import pickle

# read in the trajectories and times
trajs = np.load("traj.npy")
t1 = np.recfromtxt("time.dat")

# fitting function and residual calculation
def fit_func(p, x):
    A, R2 = p
    # bound A between 0.98 and 1.02 (although fits do not reflect this)
    if A > 1.02:
        A = 1.02
    if A < 0.98:
        A = 0.98

    return A * np.exp(-1.0 * np.array(x) * R2 / 1.0e6)

def residuals(p, y, x):
    err = y - fit_func(p, x)
    return err

p0 = [1.0, 0.05] # initial guess

fits = {}
# loop over the peak trajectories
for peak in trajs.dtype.names:

    print "Fitting Peak:", peak

    # get the trajectory to fit
    traj = trajs[peak]

    # fit the trajectory using leastsq (fmin, etc can also be used)
    results = scipy.optimize.leastsq(residuals, p0, args=(traj, t1))
    fits[peak] = results

# pickle the fits 
f = open("fits.pickle", 'w')
pickle.dump(fits, f)
f.close()

# output the fits nicely to file
f = open("fits.txt", 'w')
f.write("#Peak\tA\t\tR2\t\tier\n")
for k, v in fits.iteritems():
    f.write(k + "\t" + str(v[0][0]) + "\t" + str(v[0][1]) + "\t" + str(v[1]) + 
            "\n")
f.close()
