%matplotlib inline

from brian2 import *

duration = 1000*ms
t0 = duration / 2
s_ramp = 225.0 * ms # duration / 4
I_max = 40.0
def I_ramp(t, s=s_ramp):
    return I_max * exp(-(t - t0)**2/(2*s**2))

time = np.linspace(0*ms, duration, 256)
for s in linspace(50*ms, duration / 4, 6):
    plot(time / ms, I_ramp(time, s=s), label='s=%.1f ms' % (s / ms))
    
xlabel('time (ms)')
ylabel('I_ramp')
legend(loc='upper right')

n = 100
freq = 8 * Hz
I_theta = 10.0

a = 0.02
b = 0.2
# c = -50.0
d = 4.0

eqs = '''
    dv/dt = (0.04 * v**2 + 5 * v + 140 - u + I + theta) / (7*ms) : 1
    du/dt = a * (b * v - u) / (7*ms) : 1
    I = I_max * exp(-(t - t0)**2/(2*s_ramp**2)) : 1 (shared)
    theta = I_theta * sin(2 * pi * freq * t) : 1 (shared)
    c : 1
'''

neurons = NeuronGroup(n, model=eqs, threshold='v > 30', reset='v = c; u = u + d')
neurons.v = '-87.0 + 25.0 * rand()'
neurons.c = '-50.0 - 40.0 * i/n'

S = SpikeMonitor(neurons)
trace = StateMonitor(neurons, ('v', 'u', 'I', 'theta'), record=[50])

run(duration)

f, ax = subplots(3, 1, sharex=True, figsize=(13, 8))
sax, vax, rax = ax

sax.plot(S.t/ms, S.i, '.k')
sax.set(xlabel='Time (ms)', ylabel='Neuron index')

vax.plot(trace.t/ms, trace.v[0])
vax.plot(trace.t/ms, 2*(trace.theta[0]+1))
vax.set(xlabel='Time (ms)', ylabel=r'$V_m$ [Neuron 50]')

rax.plot(trace.t/ms, trace.I[0], 'r-')
rax.set(xlabel='Time (ms)', ylabel=r'$I_{RAMP}$', ylim=(0,1.1*I_max))

sax.set_xlim(0.0, duration / ms);

posterdir = '/Users/joe/projects/poster'
imagefn = 'mehta_ramp_demo.pdf'

savepath = os.path.join(posterdir, imagefn)
f.savefig(savepath)