%pylab inline

from brian import *
from brian.library.modelfitting import *
from brian.library.random_processes import *

R, tau = 1., 5 * ms
eqs = Equations('''dV/dt=(R*I-V)/tau : 1''')
eqs += OrnsteinUhlenbeck('I', 
        mu=.5, sigma=.5, tau=5 * ms)

G = NeuronGroup(N=1, model=eqs, reset=0., threshold=1., refractory=5 * ms)
Mspikes = SpikeMonitor(G)
Vstate = StateMonitor(G, 'V', record=True)
Istate = StateMonitor(G, 'I', record=True)
net = Network(G, Mspikes, Vstate, Istate)
reinit_default_clock()
net.run(1 * second)

figure(figsize=(10,4))
Vif, current, spikes = Vstate[0], Istate[0], Mspikes.spiketimes[0]
plot(Vstate.times, Vif, '-k');
vlines(spikes, 1, 3);

equations = Equations('''
    dV/dt=(R*I-V)/tau : 1
    I : 1
    R : 1
    tau : second
''')

results = modelfitting(model=equations,
    reset=0, threshold=1,
    data=spikes, input=current, dt=defaultclock.dt, 
    popsize=250, maxiter=10, cpu=1, 
    R=[0., 10.], tau=[2*ms, 50*ms],
    refractory=[0*ms, 10*ms])

print_table(results)

G = NeuronGroup(N=1, model=eqs, reset=0.,
    threshold=1.,
    refractory=results[0]['refractory'] * second)
reinit_default_clock()
G.I = TimedArray(current)
G.R, G.tau = results[0]['R'], results[0]['tau']
Mspikes = SpikeMonitor(G)
Vstate = StateMonitor(G, 'V', record=True)
Istate = StateMonitor(G, 'I', record=True)
net = Network(G, Mspikes, Vstate, Istate)
net.run(1 * second)

figure(figsize=(10,4))
V_fit, spikes_fit = Vstate[0], Mspikes.spiketimes[0]
plot(Vstate.times, Vif, 'k');
vlines(spikes, 1, 3, color='k');
plot(Vstate.times, V_fit, '-r');
vlines(spikes_fit, 1, 3, color='r');

import urllib2
path = 'https://neuralensemble.org/svn/brian/trunk/examples/modelfitting/'
# These files were saved with `numpy.savetxt`.
current = np.fromstring(
    urllib2.urlopen(path + 'current.txt').read(),
    sep='\n')
spikes = np.fromstring(
    urllib2.urlopen(path + 'spikes.txt').read(),
    sep='\n')

t = arange(current.size) * defaultclock.dt
figure(figsize=(10,4));
subplot(211);
plot(t, current, '-b');
subplot(212);
hlines([0], 0, 1);
vlines(spikes, 0, 1);
ylim(-.5, 1.5);

eqs = Equations('''
    dV/dt=(R*I-V)/tau : volt
    dVt/dt=(a*V-Vt)/taut : volt
    I : amp
    R : ohm
    a : 1
    alpha : volt
    tau : second
    taut : second
''')
threshold = '''(V>1+Vt)'''
reset = '''
V = 0
Vt += alpha
'''

results = modelfitting(model=eqs,
    reset=reset, threshold=threshold,
    data=spikes, input=current, dt=defaultclock.dt, 
    popsize=500, maxiter=5, cpu=4, 
    R=[10*Mohm, 10*Mohm, 10000.*Mohm, 10000.*Mohm], 
    a=[0., 1.], alpha=[0., 20.*mV],
    tau=[2*ms, 50*ms], taut=[20*ms, 200*ms])

print_table(results)

G = NeuronGroup(N=1, model=eqs, reset=reset,
    threshold=threshold,)
reinit_default_clock()
G.I = TimedArray(current)
for key, val in results[0].best_pos.iteritems():
    setattr(G, key, val)
Mspikes = SpikeMonitor(G)
Vstate = StateMonitor(G, 'V', record=True)
Vtstate = StateMonitor(G, 'Vt', record=True)
Istate = StateMonitor(G, 'I', record=True)
net = Network(G, Mspikes, Vstate, Vtstate, Istate)
net.run(1 * second)

figure(figsize=(10,4))
V_fit, Vt_fit, spikes_fit = Vstate[0], Vtstate[0], Mspikes.spiketimes[0]
plot(Vstate.times, V_fit, 'r');
plot(Vstate.times, 1 + Vt_fit, 'g');
vlines(spikes, 1 + Vt_fit[(spikes * 1000).astype(int)], 3, color='k');
vlines(spikes_fit, 1 + Vt_fit[(spikes_fit * 1000).astype(int)], 3, color='r');

area = 20000*umetre**2
Cm = (1*ufarad*cm**-2)*area
gl = (5e-5*siemens*cm**-2)*area 
El = -60*mV
EK = -90*mV
ENa = 50*mV
g_na = (100*msiemens*cm**-2)*area
g_kd = (30*msiemens*cm**-2)*area
VT = -63*mV
eqs = Equations('''
dv/dt = (gl*(El-v)+I-\
    g_na*(m*m*m)*h*(v-ENa)-\
    g_kd*(n*n*n*n)*(v-EK))/Cm : volt 
dm/dt = alpham*(1-m)-betam*m : 1
dn/dt = alphan*(1-n)-betan*n : 1
dh/dt = alphah*(1-h)-betah*h : 1
alpham = 0.32*(mV**-1)*(13*mV-v+VT)/ \
    (exp((13*mV-v+VT)/(4*mV))-1.)/ms : Hz
betam = 0.28*(mV**-1)*(v-VT-40*mV)/ \
    (exp((v-VT-40*mV)/(5*mV))-1)/ms : Hz
alphah = 0.128*exp((17*mV-v+VT)/(18*mV))/ms : Hz
betah = 4./(1+exp((40*mV-v+VT)/(5*mV)))/ms : Hz
alphan = 0.032*(mV**-1)*(15*mV-v+VT)/ \
    (exp((15*mV-v+VT)/(5*mV))-1.)/ms : Hz
betan = .5*exp((10*mV-v+VT)/(40*mV))/ms : Hz
I : amp
''')

G = NeuronGroup(1, model=eqs, threshold=EmpiricalThreshold(threshold=-20*mV, refractory=3*ms),
              implicit=True, freeze=True)
reinit_default_clock()
current -= mean(current)
current = 1.0e-10 + 10.0e-1 * current
G.I = TimedArray(current)
Mspikes = SpikeMonitor(G)
Vstate = StateMonitor(G, 'v', record=True)
net = Network(G, Mspikes, Vstate)
net.run(1 * second)

figure(figsize=(10,4))
Vhh, spikes = Vstate[0], Mspikes.spiketimes[0]
plot(Vstate.times, Vhh, '-k');

eqs = Equations('''
    dV/dt=(R*I-V)/tau : volt
    dVt/dt=(a*V-Vt)/taut : volt
    I : amp
    R : ohm
    a : 1
    alpha : volt
    tau : second
    taut : second
''')
threshold = '''(V>1+Vt)'''
reset = '''
V = 0
Vt += alpha
'''

results = modelfitting(model=eqs,
    reset=reset, threshold=threshold,
    data=spikes, input=current, dt=defaultclock.dt, 
    popsize=500, maxiter=5, cpu=4, 
    R=[10*Mohm, 10*Mohm, 10000.*Mohm, 10000.*Mohm], 
    a=[0., 1.], alpha=[0., 0., 20.*mV, 20.*mV],
    tau=[2*ms, 50*ms], taut=[20*ms, 200*ms])

print_table(results)

G = NeuronGroup(N=1, model=eqs, reset=reset,
    threshold=threshold,)
reinit_default_clock()
G.I = TimedArray(current)
for key, val in results[0].best_pos.iteritems():
    setattr(G, key, val)
Mspikes = SpikeMonitor(G)
Vstate = StateMonitor(G, 'V', record=True)
Vtstate = StateMonitor(G, 'Vt', record=True)
Istate = StateMonitor(G, 'I', record=True)
net = Network(G, Mspikes, Vstate, Vtstate, Istate)
net.run(1 * second)

figure(figsize=(10,4))
V_fit, Vt_fit, spikes_fit = Vstate[0], Vtstate[0], Mspikes.spiketimes[0]
plot(Vstate.times, 2+Vhh*20, 'k');
plot(Vstate.times, V_fit, 'r');
plot(Vstate.times, 1 + Vt_fit, 'g');
vlines(spikes_fit, 1 + Vt_fit[(spikes_fit * 1000).astype(int)], 3, color='r');