Fit to an expontial sum in pymc

exp_sum_pymc.ipynb

exp_sum_pymc.py

# -*- coding: utf-8 -*-
# <nbformat>3.0</nbformat>

# <codecell>

import pymc
import numpy as np

# <codecell>

# Make test data
t=np.linspace(0,20,200)
decays = np.array([2.3,9])
x=np.linspace(-1,1, 32)
amp1 = 10*np.exp(-(x+0.5)**2/0.3)
amp2 = -8*np.exp(-(x-0.5)**2/0.4)
amps=np.column_stack((amp1,amp2))
d=amps.dot(np.exp(-t/decays[:,None]))

# <codecell>

plot(x, amp1, label='amp1')
plot(x, amp2, label='amp2')
grid()
legend()

# <codecell>

#Add Noise
noise1 = np.linspace(0,0.5,32)
np.random.shuffle(noise1)
dn=d+noise1[:,None]*np.random.randn(32,t.size)
dn+=np.random.randn(1, t.size)*0.4

# <codecell>

plot(t, dn.T)
grid()
title('dn(t)')

# <codecell>

def model():
    
    sigma = pymc.Exponential('sigma', 1., value=.1*ones(32))
    eta = pymc.Normal('eta', .4, .1**-2, value=.4, observed=True)
    
    amps = pymc.Uniform('amps', -20, 20, value=zeros((32,2)))
    decays = pymc.Normal('decays', [2.3,9], .1**-2, value=[2.3,9], observed=True)
    
    base = pymc.exp(-t/decays[:,None])
    model = pymc.LinearCombination('fit', [amps], [base])
    
    global_err = pymc.Normal('errs', 0, 1/eta**2, value=zeros_like(t), observed=True)
    ch_errs=[]
    for i in range(32):
        x=pymc.Normal('err'+str(i), model[i,:]+global_err, tau=1/sigma[i]**2,
                      value=dn[i,:], observed=True)    
        ch_errs.append(x)
    
    return locals()

# <codecell>

vars = model()
pymc.MAP([vars['amps'], vars['ch_errs']]).fit(method='fmin_powell', verbose=1)

# <codecell>

m = pymc.MCMC(vars)

m.use_step_method(pymc.AdaptiveMetropolis, m.sigma)
m.use_step_method(pymc.AdaptiveMetropolis, m.amps)

%time m.sample(iter=100000, burn=50000, thin=5)

# <codecell>

plot(x, m.amps.trace().mean(0)[:,0], 'r-', label='predicted')
plot(x, m.amps.trace().mean(0)[:,1], 'r-')

plot(x, pymc.utils.hpd(m.amps.trace(), .05).reshape(32,4), 'r--')

plot(x, amp1, 'k-', label='truth')
plot(x, amp2, 'k-')
grid()
legend()

# <codecell>

plot(t, m.model.trace().mean(0).T)
grid()
title('d_predicted(t)')

# <codecell>