Exponential decay fit - lmfit issues

decay_fit.py

'''
This script fits 1D noisy data (function of time) to a model function that has
2 shape parameters (lifetime or exponetial decay: `tau`, amplitude: `ampl`) 
and one `offset` parameter for the translation in time.
'''

#%%
import numpy as np
import scipy
import scipy.stats
import lmfit
from lmfit import minimize, Parameters, report_fit

import matplotlib.pyplot as plt
from matplotlib.pyplot import plot


## Time axis parameters
time_step_s = (60e-9/4096)               # time step in seconds (S.I.)
time_step_ns = time_step_s*1e9           # time step in nano-seconds
time_nbins = 2048                        # number of time bins

time_idx = np.arange(time_nbins)         # time axis in index units
time_ns = time_idx*time_step_ns          # time axis in nano-seconds


## - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
#   Generate the data
#

# IRF Definition: it is used both to generate random samples
# and for the model function employed in the fit
def exgauss(x, mu, sig, tau):
    lam = 1./tau
    return 0.5*lam * np.exp(0.5*lam * (2*mu + lam*(sig**2) - 2*x)) * \
            scipy.special.erfc((mu + lam*(sig**2) - x)/(np.sqrt(2)*sig))

irf_sig = 0.033
irf_mu = 5*irf_sig
irf_lam = 0.1
x_irf = np.arange(0, (irf_lam*15 + irf_mu)/time_step_ns)
p_irf = exgauss(x_irf, irf_mu/time_step_ns, irf_sig/time_step_ns,
                irf_lam/time_step_ns)
irf = scipy.stats.rv_discrete(name='irf', values=(x_irf, p_irf))


## Simulate the data to be fitted.
## The data points to be fitted are the histogram counts of the random values
np.random.seed(2)

num_samples = 1e6
tau = 2e-9
baseline_fraction = 0.05
offset = 2e-9

# The random data is generated from an exponetial decay convoluted by a narrow
# peak (IRF) with the addition of a constant baseline
sample_decay = np.random.exponential(scale=tau/time_step_s,
                                     size=num_samples) + offset/time_step_s
sample_decay += irf.rvs(size=num_samples)
sample_baseline = np.random.randint(low=0, high=time_nbins,
                                    size=baseline_fraction*num_samples)
sample_tot = np.hstack((sample_decay, sample_baseline))
decay_hist, bins = np.histogram(sample_tot, bins=np.arange(time_nbins + 1))


## - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
#   Fitting the data
#

def model(time, tau, ampl, baseline, offset, irf_pdf=p_irf):
    """
    Model function used to fit the data.

    Parameters:
        time (array): time axis
        tau (float): time-constant (same unit as the time axis)
        ampl (float): amplitude of the model function
        baseline (float): mean baseline value (vertical offset)
        offset (float): time axis translation (same unit as the time axis)

    Returns:
        (array) Function evaluated in each `time` point.
    """
    y = np.zeros(time.size)
    pos_range = (time >= offset)
    y[pos_range] = ampl*np.exp((-time[pos_range] + offset)/tau)
    z = np.convolve(y, irf_pdf, mode='same')
    z += baseline
    return z

def residuals(params, time, y, weights):
    """
    Returns the array of residuals for the current parameters.
    """
    tau = params['tau'].value
    baseline = params['baseline'].value
    offset = params['offset'].value
    ampl = params['ampl'].value
    return (y - model(time, tau, ampl, baseline, offset))*weights

def plot_fit(time, ydata, params, yscale='log', zoom_origin=False):
    """
    Function to plot data and model function.
    """
    plot(time, ydata, marker='.')
    plot(time, model(time, **{k: v.value for k, v in params.items()}),
         color='r', lw=2.5, alpha=0.7)
    if yscale == 'log':
        plt.yscale('log')
        plt.ylim(1)
    if zoom_origin:
        dt = time[1] - time[0]
        t0 = params['offset'] - 50*dt
        plt.xlim(t0 - 50*dt, t0 + 50*dt)


# Weights for proper scaling of the residuals
weights = 1./np.sqrt(decay_hist)

#%%
# This fails with: AttributeError: 'int' object has no attribute 'copy'
params1 = Parameters()
params1.add('tau', value=2.8)
params1.add('baseline', value=24, min=20, max=30)
params1.add('ampl', value=7000)
params1.add('offset', value=2.8, vary=False)

fit_res = minimize(residuals, params1, args=(time_ns, decay_hist, weights),
                   )#method='nelder')
report_fit(fit_res.params)
print 'Reduced Chi-square: %.3f\n' % fit_res.redchi
ci = lmfit.conf_interval(fit_res)
lmfit.printfuncs.report_ci(ci)

plot_fit(time_ns, decay_hist, params1)


#%%
# This fails with: ValueError: f(a) and f(b) must have different signs
params2 = Parameters()
params2.add('tau', value=2.8)
params2.add('baseline', value=24)
params2.add('ampl', value=7000)
params2.add('offset', value=2.8)

fit_res = minimize(residuals, params2, args=(time_ns, decay_hist, weights),
                   )#method='nelder')
report_fit(fit_res.params)
print 'Reduced Chi-square: %.3f\n' % fit_res.redchi
ci = lmfit.conf_interval(fit_res)
lmfit.printfuncs.report_ci(ci)

plot_fit(time_ns, decay_hist, params2)

#%%
# This works
params3 = Parameters()
params3.add('tau', value=2.8)
params3.add('baseline', value=24)
params3.add('ampl', value=7000)
params3.add('offset', value=2.81, vary=False)

fit_res = minimize(residuals, params3, args=(time_ns, decay_hist, weights),
                   )#method='nelder')
report_fit(fit_res.params)
print 'Reduced Chi-square: %.3f\n' % fit_res.redchi
ci = lmfit.conf_interval(fit_res)
lmfit.printfuncs.report_ci(ci)

plot_fit(time_ns, decay_hist, params3, zoom_origin=True)


#%%
# Brute force search of best offset
redchi = []
for v in np.arange(2.79, 2.83, 0.002):
    params3.add('offset', value=v, vary=False)
    fit_res = minimize(residuals, params3, args=(time_ns, decay_hist, weights))
    print ' %.3f  Reduced Chi-square: %.8f' % (v, fit_res.redchi)
    redchi.append(fit_res.redchi)