Reproduce results of Jansen, Rit 1995 paper with TVB

jr95.py

import numpy as np
import tvb.simulator.lab as tvb
from tvb.basic.neotraits.api import Final, List

# setup model based on paper's parameters
model_pars = dict(
    A=3.25,
    B=22.0,
    v0=6.0,
    a=0.1, # TVB uses ms, not s
    b=0.05,
    r=0.56,
    nu_max=0.0025, # e0 in the paper
    # TVB factors C_i into J*a_i, e.g. C1 = a_1 * J
    J=135.0,
    a_1=1.0,
    a_2=0.8,
    a_3=0.25,
    a_4=0.25,
    mu=0.22, # avg of 120 and 320 pulses/s
)

# TODO fix according to paper
k1, k2 = 30.0, 30.0
connectivity = tvb.connectivity.Connectivity(
    weights=np.array([[0.0, k1], [k2, 0.0]]),
    tract_lengths=np.zeros((2, 2)),
    region_labels=np.array(['l', 'r']),
    centres=np.zeros((2, 3))
    )

# implement JR + afferent PSP
class JRPSP(tvb.models.JansenRit):
    state_variable_range = Final(
        default={"y0": np.array([-1.0, 1.0]),
                 "y1": np.array([-500.0, 500.0]),
                 "y2": np.array([-50.0, 50.0]),
                 "y3": np.array([-6.0, 6.0]),
                 "y4": np.array([-20.0, 20.0]),
                 "y5": np.array([-500.0, 500.0]),
                 "y6": np.array([-20.0, 20.0]),
                 "y7": np.array([-500.0, 500.0])})

    variables_of_interest = List(
        of=str,
        label="Variables watched by Monitors",
        choices=("y0", "y1", "y2", "y3", "y4", "y5", "y6", "y7"),
        default=("y0", "y1", "y2", "y3"))

    state_variables = tuple('y0 y1 y2 y3 y4 y5 y6 y7'.split())
    _nvar = 8
    cvar = np.array([6], dtype=np.int32)

    def dfun(self, state_variables, coupling, local_coupling=0.0):
        dy = np.zeros((8, state_variables.shape[1], 1))
        # TVB's JR is eq 6 only
        dy[:6] = super().dfun(state_variables[:6], coupling, local_coupling)
        # tack on PSP for efferent following eq 8
        # NB with this, only y12 is coupling var for TVB
        y0, y1, y2, y3, y4, y5, y6, y7 = state_variables
        a_d = self.a / 3.0
        sigm_y1_y2 = 2.0 * self.nu_max / (1.0 + np.exp(self.r * (self.v0 - (y1 - y2))))
        dy[6] = y7
        dy[7] = self.A * a_d * sigm_y1_y2 - 2.0 * a_d * y7 - self.a**2 * y6
        return dy

# factor out noise from dy4 = A a (p(t) + ...) as y4 += dt (...) + A a dW_t
# this allows us to model the noise as TVB does it, though scaling requires experiment
nsig = np.zeros((8, 2, 1))
nsig[4] = model_pars['A'] * model_pars['a'] * (.320 - .120) * 5e-3
noise = tvb.noise.Additive(nsig=nsig)

# setup simulation
sim = tvb.simulator.Simulator(
    connectivity=connectivity,
    model=JRPSP(
        variables_of_interest=("y0", "y1", "y2", "y3", "y4", "y5", "y6", "y7"),
        **{k: np.array(v) for k, v in model_pars.items()}),
    coupling=tvb.coupling.Linear(a=np.r_[1.0]),
    integrator=tvb.integrators.EulerStochastic(dt=0.1, noise=noise),
    monitors=(
        tvb.monitors.Raw(),
        tvb.monitors.ProgressLogger(period=1e2),
    )
)
sim.configure()

# run, burn in 1s transient, then keep 1s of simulation
sim.run(simulation_length=1e3) # warmup
(t, y), _ = sim.run(simulation_length=2e3)
y0, y1, y2, y3, y4, y5, y6, y7 = np.transpose(y[..., 0], (1, 0, 2))

from pylab import *
figure(); plot(t, y1 - y2); ylabel('y1(t) - y2(t)'); xlabel('t (ms)')
tight_layout(); show()