Kuramoto dev (#239)

* kuramoto added

* variable names refactoring

* remove theta_ou_mean

* tests update

added kuramoto model to the tests

* import model in initfile

* omega init after cmat check

* renaming examples + kuramoto update

* updating examples

* removing np.mod from timeintegration

* comments adjustment

* Update example-0.5-kuramoto.ipynb

* Update timeIntegration.py

* initializing lengthMat and notebook update

* update example and loaddefaultparams: dmat

* cmat init follows other models

* mod added

---------

Co-authored-by: bramantyois <[email protected]>
Co-authored-by: akmit <[email protected]>

neurolib/models/kuramoto/__init__.py

@@ -0,0 +1 @@
+from .model import KuramotoModel

neurolib/models/kuramoto/loadDefaultParams.py

@@ -0,0 +1,62 @@
+import numpy as np
+
+from neurolib.utils.collections import dotdict
+
+
+def loadDefaultParams(Cmat=None, Dmat=None, seed=None):
+    """Load default parameters for the Kuramoto Model model
+
+    :param Cmat: Structural connectivity matrix (adjacency matrix) of coupling strengths, will be normalized to 1. If not given, then a single node simulation will be assumed, defaults to None
+    :type Cmat: numpy.ndarray, optional
+    :param Dmat: Fiber length matrix, will be used for computing the delay matrix together with the signal transmission speed parameter `signalV`, defaults to None
+    :type Dmat: numpy.ndarray, optional
+    :param seed: Seed for the random number generator, defaults to None
+    :type seed: int, optional
+
+    :return: A dictionary with the default parameters of the model
+    :rtype: dict
+    """
+    params = dotdict({})
+
+    ### runtime parameters
+
+    params.dt = 0.1 
+    params.duration = 2000  
+
+    np.random.seed(seed)  
+    params.seed = seed
+
+    # model parameters
+    params.N = 1
+    params.k = 2
+
+    # connectivity
+    if Cmat is None:
+        params.N = 1
+        params.Cmat = np.zeros((1, 1))
+        params.lengthMat = np.zeros((1, 1))
+    else:
+        params.Cmat = Cmat.copy()  # coupling matrix
+        np.fill_diagonal(params.Cmat, 0)  # no self connections
+        params.N = len(params.Cmat)  # override number of nodes
+        params.lengthMat = Dmat
+
+    params.omega = np.ones((params.N,)) * np.pi
+
+    params.signalV = 20.0
+
+    # Ornstein-Uhlenbeck process
+    params.tau_ou = 5.0  # ms Timescale of the Ornstein-Uhlenbeck noise process
+    params.sigma_ou = 0.0  # 1/ms/sqrt(ms) noise intensity
+
+    # init values
+    params.theta_init = np.random.uniform(low=0, high=2*np.pi, size=(params.N, 1))
+
+    # Ornstein-Uhlenbeck process
+    params.theta_ou = np.zeros((params.N,))
+
+    # external input
+    params.theta_ext = np.zeros((params.N,))
+
+    return params
+

neurolib/models/kuramoto/model.py

@@ -0,0 +1,35 @@
+from . import loadDefaultParams as dp
+from . import timeIntegration as ti
+
+from neurolib.models.model import Model
+
+
+class KuramotoModel(Model):
+    """
+    Kuramoto Model
+
+    Based on:
+    Kuramoto, Yoshiki (1975). H. Araki (ed.). Lecture Notes in Physics, International Symposium on Mathematical Problems in Theoretical Physics.
+    """
+
+    name = "kuramoto"
+    description = "Kuramoto Model"
+
+    init_vars = ['theta_init', 'theta_ou']
+    state_vars = ['theta', 'theta_ou'] 
+    output_vars = ['theta']
+    default_output = 'theta'
+    input_vars = ['theta_ext']
+    default_input = 'theta_ext'
+
+    def __init__(self, params=None, Cmat=None, Dmat=None, seed=None):
+        self.Cmat = Cmat
+        self.Dmat = Dmat
+        self.seed = seed
+
+        integration = ti.timeIntegration
+
+        if params is None:
+            params = dp.loadDefaultParams(Cmat=self.Cmat, Dmat=self.Dmat, seed=self.seed)
+
+        super().__init__(params=params, integration=integration)

neurolib/models/kuramoto/timeIntegration.py

@@ -0,0 +1,168 @@
+import numpy as np
+import numba
+
+from ...utils import model_utils as mu
+
+
+def timeIntegration(params):
+    """
+    setting up parameters for time integration
+    
+    :param params: Parameter dictionary of the model
+    :type params: dict
+
+    :return: Integrated activity of the model
+    :rtype: (numpy.ndarray, )
+    """ 
+    dt = params["dt"]  # Time step for the Euler intergration (ms)
+    duration = params["duration"]  # imulation duration (ms)
+    RNGseed = params["seed"]  # seed for RNG
+
+    np.random.seed(RNGseed)
+
+    # ------------------------------------------------------------------------
+    # model parameters
+    # ------------------------------------------------------------------------
+
+    N = params["N"]  # number of oscillators
+
+    omega = params["omega"]  # frequencies of oscillators
+
+    # ornstein uhlenbeck noise param
+    tau_ou = params["tau_ou"]  # noise time constant
+    sigma_ou = params["sigma_ou"]  # noise strength
+
+    # ------------------------------------------------------------------------
+    # global coupling parameters
+    # ------------------------------------------------------------------------
+
+    # Connectivity matrix and Delay
+    Cmat = params["Cmat"]
+
+    # Interareal connection delay
+    lengthMat = params["lengthMat"]
+    signalV = params["signalV"]
+    k = params["k"]  # coupling strength
+
+    if N == 1:
+        Dmat = np.zeros((N, N))
+    else:
+        # Interareal connection delays, Dmat(i,j) Connnection from jth node to ith (ms)
+        Dmat = mu.computeDelayMatrix(lengthMat, signalV)
+
+        # no self-feedback delay
+        Dmat[np.eye(len(Dmat)) == 1] = np.zeros(len(Dmat))
+    Dmat = Dmat.astype(int)
+    Dmat_ndt = np.around(Dmat / dt).astype(int)  # delay matrix in multiples of dt
+
+    # ------------------------------------------------------------------------
+    # Initialization
+    # ------------------------------------------------------------------------
+
+    t = np.arange(1, round(duration, 6) / dt + 1) * dt  # Time variable (ms)
+    sqrt_dt = np.sqrt(dt)
+
+    max_global_delay = np.max(Dmat_ndt)  # maximum global delay
+    startind = int(max_global_delay + 1)  # start simulation after delay 
+
+    # Placeholders
+    theta_ou = params['theta_ou'].copy()
+    theta = np.zeros((N, startind + len(t)))
+
+    theta_ext = mu.adjustArrayShape(params["theta_ext"], theta)
+
+    # ------------------------------------------------------------------------
+    # initial values
+    # ------------------------------------------------------------------------  
+
+    if params["theta_init"].shape[1] == 1:
+        theta_init = np.dot(params["theta_init"], np.ones((1, startind)))
+    else:
+        theta_init = params["theta_init"][:, -startind:]
+
+    # put noise to instantiated array to save memory
+    theta[:, :startind] = theta_init
+    theta[:, startind:] = np.random.standard_normal((N, len(t)))
+
+    theta_input_d = np.zeros(N)
+
+    noise_theta = 0
+
+    # ------------------------------------------------------------------------
+    # some helper variables
+    # ------------------------------------------------------------------------
+
+    k_n = k/N
+    theta_rhs = np.zeros((N,))
+
+    # ------------------------------------------------------------------------
+    # time integration
+    # ------------------------------------------------------------------------
+
+    return timeIntegration_njit_elementwise(
+        startind,
+        t, 
+        dt, 
+        sqrt_dt,
+        N,
+        omega,
+        k_n, 
+        Cmat,
+        Dmat,
+        theta,
+        theta_input_d,
+        theta_ext,
+        tau_ou,
+        sigma_ou,
+        theta_ou,
+        noise_theta,
+        theta_rhs,
+    )
+
+
+@numba.njit
+def timeIntegration_njit_elementwise(
+    startind,
+    t, 
+    dt, 
+    sqrt_dt,
+    N,
+    omega,
+    k_n, 
+    Cmat,
+    Dmat,
+    theta,
+    theta_input_d,
+    theta_ext,
+    tau_ou,
+    sigma_ou,
+    theta_ou,
+    noise_theta,
+    theta_rhs, 
+):
+    """
+    Kuramoto Model 
+    """
+    for i in range(startind, startind+len(t)):
+        # Kuramoto model
+        for n in range(N): 
+            noise_theta = theta[n, i]
+
+            theta_input_d[n] = 0.0
+
+            # adding input from other nodes
+            for m in range(N):
+                theta_input_d[n] +=  k_n * Cmat[n, m] * np.sin(theta[m, i-1-Dmat[n, m]] - theta[n, i-1])
+
+            theta_rhs[n] = omega[n] + theta_input_d[n] + theta_ou[n] + theta_ext[n, i-1]
+
+            # time integration
+            theta[n, i] = theta[n, i-1] + dt * theta_rhs[n]
+
+            # phase reset
+            theta[n, i] = np.mod(theta[n, i], 2*np.pi)
+
+            # Ornstein-Uhlenbeck
+            theta_ou[n] = theta_ou[n] - theta_ou[n] * dt / tau_ou + sigma_ou * sqrt_dt * noise_theta
+
+    return t, theta, theta_ou

tests/test_models.py

@@ -17,6 +17,7 @@
 from neurolib.utils.collections import star_dotdict
 from neurolib.utils.loadData import Dataset
 from neurolib.utils.stimulus import ZeroInput
+from neurolib.models.kuramoto import KuramotoModel
 
 class TestAln(unittest.TestCase):
     """
@@ -194,6 +195,42 @@ def test_network(self):
         logging.info("\t > Done in {:.2f} s".format(end - start))
 
 
+class TestKuramoto(unittest.TestCase):
+    """
+    Basic test for Kuramoto model.
+    """
+
+    def test_single_node(self):
+        logging.info("\t > Kuramoto: Testing single node ...")
+        start = time.time()
+        model = KuramotoModel()
+        model.params["duration"] = 2.0 * 1000
+        model.params["sigma_ou"] = 0.03
+
+        model.run()
+
+        end = time.time()
+        logging.info("\t > Done in {:.2f} s".format(end - start))
+
+    def test_network(self):
+        logging.info("\t > Kuramoto: Testing brain network (chunkwise integration and BOLD simulation) ...")
+        start = time.time()
+        ds = Dataset("gw")
+        model = KuramotoModel(Cmat=ds.Cmat, Dmat=ds.Dmat)
+        model.params["signalV"] = 4.0
+        model.params["duration"] = 10 * 1000
+        model.params["sigma_ou"] = 0.1
+        model.params["k"] = 0.6
+
+
+        # local node input parameter 
+        model.params["theta_ext"] = 0.72
+
+        model.run(chunkwise=True, append_outputs=True)
+        end = time.time()
+        logging.info("\t > Done in {:.2f} s".format(end - start))
+
+
 class TestThalamus(unittest.TestCase):
     """
     Basic test for thalamic mass model.

tests/test_stimulus.py

@@ -10,6 +10,7 @@
 from neurolib.models.fhn import FHNModel
 from neurolib.models.hopf import HopfModel
 from neurolib.models.wc import WCModel
+from neurolib.models.kuramoto import KuramotoModel
 from neurolib.utils.stimulus import (
     ConcatenatedStimulus,
     ExponentialInput,
@@ -143,6 +144,28 @@ def test_multi_node_multi_stim(self):
         self.assertTupleEqual(model_stim.shape, (5, int(model.params["duration"] / model.params["dt"])))
 
 
+class TestToKuramotoModel(unittest.TestCase):
+    def test_single_node(self):
+        model = KuramotoModel()
+        model.params["duration"] = 2 * 1000
+        stim = SinusoidalInput(amplitude=1.0, frequency=1.0)
+        model_stim = stim.to_model(model)
+        model.params["theta_ext"] = model_stim
+        model.run()
+        self.assertTrue(isinstance(model_stim, np.ndarray))
+        self.assertTupleEqual(model_stim.shape, (1, int(model.params["duration"] / model.params["dt"])))
+
+    def test_multi_node_multi_stim(self):
+        model = KuramotoModel(Cmat=np.random.rand(5, 5), Dmat=np.zeros((5, 5)))
+        model.params["duration"] = 2 * 1000
+        stim = SinusoidalInput(amplitude=1.0, frequency=1.0)
+        model_stim = stim.to_model(model)
+        model.params["theta_ext"] = model_stim
+        model.run()
+        self.assertTrue(isinstance(model_stim, np.ndarray))
+        self.assertTupleEqual(model_stim.shape, (5, int(model.params["duration"] / model.params["dt"])))
+
+
 class TestZeroInput(unittest.TestCase):
     def test_generate_input(self):
         nn = ZeroInput(n=2, seed=42).generate_input(duration=DURATION, dt=DT)