vbi 0.1.3__cp310-cp310-manylinux2014_x86_64.whl

vbi/models/cpp/wc.py

@@ -0,0 +1,216 @@
+import numpy as np
+
+try:
+    from vbi.models.cpp._src.wc_ode import WC_ode as _WC_ode
+except ImportError as e:
+    print(f"Could not import modules: {e}, probably C++ code is not compiled or properly linked.")
+
+
+################################## Wilson-Cowan ode ###########################
+###############################################################################
+
+class WC_ode(object):
+    r"""
+    **References**:
+
+    .. [WC_1972] Wilson, H.R. and Cowan, J.D. *Excitatory and inhibitory
+        interactions in localized populations of model neurons*, Biophysical
+        journal, 12: 1-24, 1972.
+    .. [WC_1973] Wilson, H.R. and Cowan, J.D  *A Mathematical Theory of the
+        Functional Dynamics of Cortical and Thalamic Nervous Tissue*
+    .. [D_2011] Daffertshofer, A. and van Wijk, B. *On the influence of
+        amplitude on the connectivity between phases*
+        Frontiers in Neuroinformatics, July, 2011
+
+    Used Eqns 11 and 12 from [WC_1972]_ in ``rhs``.  P and Q represent external
+    inputs, which when exploring the phase portrait of the local model are set
+    to constant values. However in the case of a full network, P and Q are the
+    entry point to our long range and local couplings, that is, the  activity
+    from all other nodes is the external input to the local population [WC_1973]_, [D_2011]_ .
+
+    The default parameters are taken from figure 4 of [WC_1972]_, pag. 10.
+    
+    """
+
+
+
+    def __init__(self, par={}) -> None:
+
+        self.valid_params = self.get_default_parameters().keys()
+        self.check_parameters(par)
+        self._par = self.get_default_parameters()
+        self._par.update(par)
+
+        for item in self._par.items():
+            name = item[0]
+            value = item[1]
+            setattr(self, name, value)
+
+        if self.seed is not None:
+            np.random.seed(self.seed)
+
+        self.N = self.num_nodes = np.asarray(self.weights).shape[0]
+
+    def __str__(self) -> str:
+        print("Wilson-Cowan model.")
+        print("--------------------")
+        for item in self._par.items():
+            print(f"{item[0]}, : , {item[1]}")
+        return ""
+
+    def __call__(self):
+        print("Wilson-Cowan model.")
+        return self._par
+
+    def check_parameters(self, par):
+        for key in par.keys():
+            if key not in self.valid_params:
+                raise ValueError(f"Invalid parameter: {key}")
+
+    def get_default_parameters(self):
+        par = {
+            'c_ee': 16.0,
+            'c_ei': 12.0,
+            'c_ie': 15.0,
+            'c_ii': 3.0,
+            'tau_e': 8.0,
+            'tau_i': 8.0,
+            'a_e': 1.3,
+            'a_i': 2.0,
+            'b_e': 4.0,
+            'b_i': 3.7,
+            'c_e': 1.0,
+            'c_i': 1.0,
+            'theta_e': 0.0,
+            'theta_i': 0.0,
+            'r_e': 1.0,
+            'r_i': 1.0,
+            'k_e': 0.994,
+            'k_i': 0.999,
+            'alpha_e': 1.0,
+            'alpha_i': 1.0,
+            'P': 1.25,
+            'Q': 0.0,
+            'g_e': 0.0,
+            'g_i': 0.0,
+            "method": "heun",
+            "weights": None,
+            'seed': None,
+            "t_end": 300.0,
+            "t_cut": 0.0,
+            "dt": 0.01,
+            "noise_seed": False,
+            "output": "output",
+        }
+        return par
+
+    def set_initial_state(self, seed=None):
+
+        if seed is not None:
+            np.random.seed(seed)
+        self.initial_state = np.random.rand(2*self.num_nodes)
+
+    def prepare_input(self):
+        self.noise_seed = int(self.noise_seed)
+        self.t_end = float(self.t_end)
+        self.t_cut = float(self.t_cut)
+        self.dt = float(self.dt)
+        self.P = check_sequence(self.P, self.num_nodes)
+        self.Q = check_sequence(self.Q, self.num_nodes)
+        self.c_ee = float(self.c_ee)
+        self.c_ei = float(self.c_ei)
+        self.c_ie = float(self.c_ie)
+        self.c_ii = float(self.c_ii)
+        self.tau_e = float(self.tau_e)
+        self.tau_i = float(self.tau_i)
+        self.a_e = float(self.a_e)
+        self.a_i = float(self.a_i)
+        self.b_e = float(self.b_e)
+        self.b_i = float(self.b_i)
+        self.c_e = float(self.c_e)
+        self.c_i = float(self.c_i)
+        self.theta_e = float(self.theta_e)
+        self.theta_i = float(self.theta_i)
+        self.r_e = float(self.r_e)
+        self.r_i = float(self.r_i)
+        self.k_e = float(self.k_e)
+        self.k_i = float(self.k_i)
+        self.alpha_e = float(self.alpha_e)
+        self.alpha_i = float(self.alpha_i)
+        self.g_e = float(self.g_e)
+        self.g_i = float(self.g_i)
+        self.method = str(self.method)
+        self.weights = np.asarray(self.weights)
+
+
+    def run(self, par={}, x0=None, verbose=False):
+
+        '''
+        Integrate the system of equations for the Wilson-Cowan model.
+
+        Parameters
+        ----------
+        par : dict
+            Dictionary with parameters of the model.
+        x0 : array-like
+            Initial state of the system.
+        verbose : bool
+            If True, print the integration progress.
+
+        '''
+
+        if x0 is None:
+            self.set_initial_state()
+            if verbose:
+                print("Initial state set by default.")
+        else:
+            self.initial_state = x0
+
+        for key in par.keys():
+            if key not in self.valid_params:
+                raise ValueError(f"Invalid parameter: {key}")
+            setattr(self, key, par[key]['value'])
+
+        self.prepare_input()
+
+        obj = _WC_ode(
+            self.N, self.dt, self.P, self.Q, self.initial_state, self.weights,
+            self.t_end, self.t_cut, self.c_ee, self.c_ei, self.c_ie, self.c_ii,
+            self.tau_e, self.tau_i, self.a_e, self.a_i, self.b_e, self.b_i,
+            self.c_e, self.c_i, self.theta_e, self.theta_i, self.r_e, self.r_i,
+            self.k_e, self.k_i, self.alpha_e, self.alpha_i, self.g_e, self.g_i,
+            self.noise_seed
+        )
+
+        if self.method == "euler":
+            obj.eulerIntegrate()
+        elif self.method == "heun":
+            obj.heunIntegrate()
+        elif self.method == "rk4":
+            obj.rk4Integrate()
+
+        t = np.asarray(obj.get_times())
+        x = np.asarray(obj.get_states()).T
+
+        del obj
+        return {"t": t, "x": x}
+
+
+def check_sequence(x, n):
+    '''
+    check if x is a scalar or a sequence of length n
+
+    parameters
+    ----------
+    x: scalar or sequence of length n
+    n: number of nodes
+
+    returns
+    -------
+    x: sequence of length n
+    '''
+    if isinstance(x, (np.ndarray, list, tuple)):
+        assert (len(x) == n), f" variable must be a sequence of length {n}"
+        return x
+    else:
+        return x * np.ones(n)

vbi/models/cupy/bold.py

@@ -0,0 +1,111 @@
+import numpy as np
+
+
+class BoldStephan2008:
+
+    def __init__(self, par: dict = {}) -> None:
+
+        self._par = self.get_default_parameters()
+        self.valid_parameters = list(self._par.keys())
+        self.check_parameters(par)
+        self._par.update(par)
+
+        for key, value in self._par.items():
+            setattr(self, key, value)
+
+    def _prepare(self, nn, ns, xp, n_steps, bold_decimate):
+        s = xp.zeros((2, nn, ns), dtype=self.dtype)
+        f = xp.zeros((2, nn, ns), dtype=self.dtype)
+        ftilde = xp.zeros((2, nn, ns), dtype=self.dtype)
+        vtilde = xp.zeros((2, nn, ns), dtype=self.dtype)
+        qtilde = xp.zeros((2, nn, ns), dtype=self.dtype)
+        v = xp.zeros((2, nn, ns), dtype=self.dtype)
+        q = xp.zeros((2, nn, ns), dtype=self.dtype)
+        vv = np.zeros((n_steps // bold_decimate, nn, ns), dtype="f")
+        qq = np.zeros((n_steps // bold_decimate, nn, ns), dtype="f")
+        s[0] = 1
+        f[0] = 1
+        v[0] = 1
+        q[0] = 1
+        ftilde[0] = 0
+        vtilde[0] = 0
+        qtilde[0] = 0
+
+        return {
+            "s": s,
+            "f": f,
+            "ftilde": ftilde,
+            "vtilde": vtilde,
+            "qtilde": qtilde,
+            "v": v,
+            "q": q,
+            "vv": vv,
+            "qq": qq,
+        }
+
+    def check_parameters(self, par):
+        for key in par.keys():
+            if key not in self.valid_parameters:
+                raise ValueError(f"Invalid parameter {key:s} provided.")
+
+    def get_default_parameters(self):
+
+        theta0 = 41.0
+        Eo = 0.42
+        TE = 0.05
+        epsilon = 0.36
+        r0 = 26.0
+        k1 = 4.3 * theta0 * Eo * TE
+        k2 = epsilon * r0 * Eo * TE
+        k3 = 1 - epsilon
+
+        par = {
+            "kappa": 0.7,
+            "gamma": 0.5,
+            "tau": 1.0,
+            "alpha": 0.35,
+            "epsilon": epsilon,
+            "Eo": Eo,
+            "TE": TE,
+            "vo": 0.09,
+            "r0": r0,
+            "theta0": theta0,
+            "rtol": 1e-6,
+            "atol": 1e-9,
+            "k1": k1,
+            "k2": k2,
+            "k3": k3,
+        }
+        return par
+
+    def bold_step(self, r_in, s, f, ftilde, vtilde, qtilde, v, q, dt, P):
+
+        kappa, gamma, alpha, tau, Eo = P
+        ialpha = 1 / alpha
+
+        s[1] = s[0] + dt * (r_in - kappa * s[0] - gamma * (f[0] - 1))
+        f[0] = np.clip(f[0], 1, None)
+        ftilde[1] = ftilde[0] + dt * (s[0] / f[0])
+        fv = v[0] ** ialpha  # outflow
+        vtilde[1] = vtilde[0] + dt * ((f[0] - fv) / (tau * v[0]))
+        q[0] = np.clip(q[0], 0.01, None)
+        ff = (1 - (1 - Eo) ** (1 / f[0])) / Eo  # oxygen extraction
+        qtilde[1] = qtilde[0] + dt * ((f[0] * ff - fv * q[0] / v[0]) / (tau * q[0]))
+
+        f[1] = np.exp(ftilde[1])
+        v[1] = np.exp(vtilde[1])
+        q[1] = np.exp(qtilde[1])
+
+        f[0] = f[1]
+        s[0] = s[1]
+        ftilde[0] = ftilde[1]
+        vtilde[0] = vtilde[1]
+        qtilde[0] = qtilde[1]
+        v[0] = v[1]
+        q[0] = q[1]
+
+
+class BoldTVB:
+
+    def __init__(self):
+        pass

vbi/models/cupy/ghb.py

@@ -0,0 +1,284 @@
+import tqdm
+import cupy as cp
+from copy import copy
+from vbi.models.cupy.utils import *
+
+
+class GHB_sde:
+    """
+    Generic Hopf model cupy implementation
+
+    Parameters
+    ----------
+    par: dict
+        Dictionary of parameters
+        - 'G': Global coupling
+        - 'dt': Time step
+
+    """
+
+    epsilon = 0.5
+    itaus = 1.25
+    itauf = 2.5
+    itauo = 1.02040816327
+    ialpha = 5.
+    E0 = 0.4
+    V0 = 4.
+    K1 = 2.77264
+    K2 = 0.572
+    K3 = -0.43
+
+    def __init__(self, par: dict = {}) -> None:
+
+        self.valid_params = list(self.get_default_parameters().keys())
+        self.check_parameters(par)
+        self.par_ = self.get_default_parameters()
+        self.par_.update(par)
+
+        for item in self.par_.items():
+            setattr(self, *item)
+
+        self.xp = get_module(self.engine)
+        if self.seed is not None:
+            self.xp.random.seed(self.seed)
+
+    def __call__(self):
+        print("GHB model")
+        return self.par_
+
+    def __str__(self):
+        print("GHB model")
+        print("-" * 50)
+        for item in self.par_items():
+            name = item[0]
+            value = item[1]
+            print(f"{name} : {value}")
+        return ""
+
+    def set_initial_state(self):
+        self.initial_state = set_initial_state(
+            self.nn,
+            self.num_sim,
+            self.engine,
+            self.seed,
+            self.same_initial_state,
+            self.dtype,
+        )
+
+    def check_parameters(self, par):
+        for key in par.keys():
+            assert key in self.valid_params, "Invalid parameter: " + key
+
+    def get_default_parameters(self):
+        par = {
+            "G": 25.0,
+            "t_cut": 0,
+            "dt": 0.01,
+            "eta": None,
+            "num_sim": 1,
+            "sigma": 0.1,
+            "seed": None,
+            "decimate": 1,
+            "omega": None,
+            "t_end": 10.0,
+            "engine": "cpu",
+            "weights": None,
+            "dtype": "float",
+            "method": "euler",
+            "output": "output",
+            "initial_state": None,
+            "same_initial_state": False,
+        }
+        return par
+
+    def prepare_input(self):
+        self.G = self.xp.array(self.G, dtype=self.dtype)
+        assert self.weights is not None, "weights not provided"
+        self.weights = self.xp.array(self.weights, dtype=self.dtype)
+        self.weights = self.weights.reshape(self.weights.shape+(1,))
+        self.weights = move_data(self.weights, self.engine)
+        self.nn = self.num_nodes = self.weights.shape[0]
+
+
+        if self.initial_state is None:
+            self.set_initial_state()
+        else:
+            self.initial_state = move_data(
+                self.initial_state, self.engine)
+
+
+        self.eta = prepare_vec(self.eta, self.num_sim, self.engine, self.dtype)
+        self.omega = prepare_vec(self.omega, self.num_sim, self.engine, self.dtype)
+
+    def f_sys(self, x0, t):
+
+        G = self.G
+        xp = self.xp
+        nn = self.nn
+        eta = self.eta
+        x = x0[:nn, :]
+        y = x0[nn:, :]
+        ns = self.num_sim
+        sc = self.weights
+        omega = self.omega
+
+        gx = xp.sum(sc * (x - x[:, None]), axis=1)
+        gy = xp.sum(sc * (y - y[:, None]), axis=1)
+        dxdt = xp.zeros((2 * nn, ns)).astype(self.dtype)
+
+        dxdt[:nn, :] = x * (eta - x * x - y * y) - omega * y + G * gx
+        dxdt[nn:, :] = y * (eta - x * x - y * y) + omega * x + G * gy
+
+        return dxdt
+
+    def f_fmri(self, xin, x, t):
+
+        E0 = self.E0
+        xp = self.xp
+        nn = self.num_nodes
+        ns = self.num_sim
+        itauf = self.itauf
+        itauo = self.itauo
+        itaus = self.itaus
+        ialpha = self.ialpha
+
+        dxdt = xp.zeros((4 * nn, ns)).astype(self.dtype)
+        s = x[:nn, :]
+        f = x[nn : 2 * nn, :]
+        v = x[2 * nn : 3 * nn, :]
+        q = x[3 * nn :, :]
+
+        dxdt[:nn, :] = xin[:nn, :] - itaus * s - itauf * (f - 1.0)
+        dxdt[nn : (2 * nn), :] = s
+        dxdt[(2 * nn) : (3 * nn), :] = itauo * (f - v ** (ialpha))
+        dxdt[3 * nn :, :] = (itauo) * (
+            (f * (1.0 - (1.0 - E0) ** (1.0 / f)) / E0) - (v ** (ialpha)) * (q / v)
+        )
+
+        return dxdt
+
+    def heun_sde_step(self, x0, t):
+
+        xp = self.xp
+        dt = self.dt
+        dx = self.f_sys(x0, t) * dt
+        dW = self.sigma * xp.random.normal(0, 1, size=x0.shape) * xp.sqrt(dt)
+        x1 = x0 + dx + dW
+        dx1 = self.f_sys(x1, t + dt) * dt
+        return x0 + 0.5 * (dx + dx1) + dW
+
+    def heun_ode_step(self, yin, y, t):
+
+        dt = self.dt
+        dy = self.f_fmri(yin, y, t) * dt
+        y1 = y + dy
+        dy1 = self.f_fmri(yin, y1, t + dt) * dt
+        return y + 0.5 * (dy + dy1)
+
+    def intg_fmri(self, yin, y, t):
+        """
+        Integrate one step of Balloon model
+
+        Parameters
+        ----------
+        yin: array
+            input
+        y: array [4*nn, ns]
+            state
+        t : float
+            current time
+
+        Returns
+        -------
+        bold: array [nn, ns]
+            BOLD signal
+        y: array [4*nn, ns]
+            updated state
+
+        """
+
+        V0 = self.V0
+        K1 = self.K1
+        K2 = self.K2
+        K3 = self.K3
+
+        nn = self.num_nodes
+        y = self.heun_ode_step(yin, y, t)
+        bold = V0 * (
+            K1 * (1.0 - y[(3 * nn) :, :])
+            + K2 * (1.0 - y[(3 * nn) :, :] / y[(2 * nn) : (3 * nn), :])
+            + K3 * (1.0 - y[(2 * nn) : (3 * nn), :])
+        )
+
+        return bold, y
+
+    def sync(self, engine="gpu"):
+        if engine == "gpu":
+            cp.cuda.Stream.null.synchronize()
+        else:
+            pass
+
+    def run(self, x0=None, verbose=True):
+        """
+        run ghb model
+        """
+        self.prepare_input()
+        dt = self.dt
+        xp = self.xp
+        ns = self.num_sim
+        nn = self.num_nodes
+        dec = self.decimate
+        engine = self.engine
+        t_cut = self.t_cut
+        n_steps = np.ceil(self.t_end / dt).astype(int)
+
+        y0_state = xp.zeros((4 * nn, ns)).astype(self.dtype)
+        y0_state[nn:, :] = 1.0
+        y0 = copy(self.initial_state)
+        bold = np.zeros((nn, n_steps // dec, ns)).astype(np.float32)
+
+        for it in tqdm.trange(n_steps, disable=not verbose, desc="Integrating"):
+            y0 = self.heun_sde_step(y0, it * dt)
+            bold_, y0_state = self.intg_fmri(y0, y0_state, it * dt)
+            self.sync(engine)
+            if it % dec == 0:
+                bold[:, it // dec, :] = bold_.get() if engine == "gpu" else bold_
+
+        t = np.arange(0, self.t_end, dec * dt).astype(np.float32)
+        bold = bold[:, t > t_cut, :]
+        t_bold = t[t > t_cut]
+        return {"t": t_bold, "bold": bold}
+
+
+def set_initial_state(nn, ns, engine, seed=None, same_initial_state=False, dtype=float):
+    """
+    Set initial state
+
+    Parameters
+    ----------
+    nn : int
+        number of nodes
+    ns : int
+        number of simulations
+    engine : str
+        cpu or gpu
+    same_initial_condition : bool
+        same initial condition for all simulations
+    seed : int
+        random seed
+    dtype : str
+        float: float64
+        f    : float32
+    """
+
+    if seed is not None:
+        np.random.seed(seed)
+
+    if same_initial_state:
+        y0 = np.random.rand(2 * nn)
+        y0 = repmat_vec(y0, ns, engine)
+    else:
+        y0 = np.random.rand(2 * nn, ns)
+        y0 = move_data(y0, engine)
+
+    return y0.astype(dtype)