Add numba dfun for model KIonEx

tvb_library/tvb/simulator/models/k_ion_exchange.py

@@ -31,6 +31,7 @@
    Giovanni Rabuffo <[email protected]>, 
    Carmela Calabrese <[email protected]>,
    Jan Fousek <[email protected]>, 
+   Michiel van der Vlag <[email protected]>
 
    Under the "NextGen" Research Infrastructure Voucher SC3 associated to the HBP Flagship as a Partnering Project (PP)
    Project leader: Simona Olmi <[email protected]>
@@ -42,6 +43,8 @@
 
 import numpy
 
+from numba import guvectorize, float64, njit
+
 class KIonEx(Model):
     r"""
     KIonEx (Potassium K+ Ion exchange) mean-field model was developed in (Bandyopadhyay & Rabuffo et al. 2023). 
@@ -208,7 +211,7 @@ class KIonEx(Model):
     # Stvar is the variable where stimulus is applied.
     stvar = numpy.array([1], dtype=numpy.int32)
 
-    def dfun(self, state_variables, coupling, local_coupling=0.0):
+    def _numpy_dfun(self, state_variables, coupling, local_coupling=0.0):
         r"""
         The mean-field approximation for a population of Hodgkin-Huxley-type neurons driven by slow potassium dynamics consists of a 5D system:
 
@@ -343,3 +346,131 @@ def V_dot_form(I_Na,I_K,I_Cl,I_pump):
         derivative[4] = epsilon * (K_bath - K_o)
 
         return derivative
+
+    def dfun(self, x, c, local_coupling=0.0):
+
+        x_ = x.reshape(x.shape[:-1]).T
+        c_ = c.reshape(c.shape[:-1]).T + local_coupling * x[0]
+
+        deriv = _numba_dfun(x_, c_, self.E, self.K_bath, self.J, self.eta, self.Delta, self.c_minus, self.R_minus,
+                                 self.c_plus, self.R_plus, self.Vstar, self.Cm, self.tau_n, self.gamma, self.epsilon)
+        return deriv.T[..., numpy.newaxis]
+    # helper functions
+
+@njit
+def m_inf(V):
+    Cmna = -24.0  # mV
+    DCmna = 12.0  # mV
+    return 1.0 / (1.0 + numpy.exp((Cmna - V) / DCmna))
+
+@njit
+def n_inf(V):
+    Cnk = -19.0  # mV
+    DCnk = 18.0  # mV #Ok in the paper
+    return 1.0 / (1.0 + numpy.exp((Cnk - V) / DCnk))
+
+@njit
+def h(n):
+    return 1.1 - 1.0 / (1.0 + numpy.exp(-8.0 * (n - 0.4)))
+
+@njit
+def I_K_form(V, n, K_o, K_i):
+    g_K = 22.0  # nS  # maximal potassium conductance
+    g_Kl = 0.12  # nS  # potassium leak conductance
+    return (g_Kl + g_K * n) * (V - 26.64 * numpy.log(K_o / K_i))
+
+@njit
+def I_Na_form(V, Na_o, Na_i, n):
+    g_Na = 40.0  # nS   # maximal sodiumconductance
+    g_Nal = 0.02  # nS  # sodium leak conductance
+    return (g_Nal + g_Na * m_inf(V) * h(n)) * (V - 26.64 * numpy.log(Na_o / Na_i))
+
+@njit
+def I_Cl_form(V):
+    g_Cl = 7.5  # nS #Ok in the paper   # chloride conductance
+    Cl_i0 = 5.0  # mMol/m**3 # initial concentration of intracellular Cl
+    Cl_o0 = 112.0  # mMol/m**3 # initial concentration of extracellular Cl
+    return g_Cl * (V + 26.64 * numpy.log(Cl_o0 / Cl_i0))
+
+@njit
+def I_pump_form(Na_i, K_o):
+    rho = 250.  # 250.,#pA # maximal Na/K pump current
+    Cnap = 21.0  # mol.m**-3
+    DCnap = 2.0  # mol.m**-3
+    Ckp = 5.5  # mol.m**-3
+    DCkp = 1.0  # mol.m**-3
+    return rho * (
+            1.0 / (1.0 + numpy.exp((Cnap - Na_i) / DCnap)) * (1.0 / (1.0 + numpy.exp((Ckp - K_o) / DCkp))))
+
+# @njit
+# def V_dot_form(Cm, I_Na, I_K, I_Cl, I_pump):
+#     return (-1.0 / Cm) * (I_Na + I_K + I_Cl + I_pump)
+
+
+@guvectorize([(float64[:],) * 17], '(n),(m)' + ',()' * 14 + '->(n)', target='parallel', nopython=True)
+def _numba_dfun(state_variables, coupling, E, K_bath, J, eta, Delta, c_minus, R_minus, c_plus, R_plus, Vstar, Cm,
+                tau_n, gamma, epsilon, dx):
+    r"""
+    The mean-field approximation for a population of Hodgkin-Huxley-type neurons driven by slow potassium dynamics consists of a 5D system:
+
+    .. math::
+        \frac{dx}{dt}&=
+        \begin{cases}
+        \Delta+2R_{-}(V-c_{-})x - J r x; \  V\leq V^{\star}\\
+        \Delta+2R_{+}(V-c_{+})x - J r x; \  V> V^{\star},
+        \end{cases}\\
+        \frac{dV}{dt}&=
+        \begin{cases}
+        -\frac{1}{C_m}(I_{Cl}+I_{Na}+I_{K}+I_{pump})-R_{-}x^2+J r(E_{syn}-V)+\overline{\eta}; \  V\leq V^{\star}\\
+        -\frac{1}{C_m}(I_{Cl}+I_{Na}+I_{K}+I_{pump})-R_{+}x^2+J r(E_{syn}-V)+\overline{\eta}; \  V>V^{\star},
+        \end{cases}\\
+        \frac{dn}{dt} &= \frac{n_{\infty}(V)-n}{\tau_n}, \\
+        \frac{d \Delta [K^{+}]_{int}}{dt} &= - \frac{\gamma}{\omega_i}(I_K - 2 I_{pump}),\\
+        \frac{d[K^+]_g}{dt} &= \epsilon ([K^+]_{bath} - [K^+]_{ext}\}).\\
+
+    For details refer to (Bandyopadhyay & Rabuffo et al. 2023)
+    """
+
+    x, V, n, DKi, Kg = state_variables
+
+    Coupling_Term = coupling[0]  # This zero refers to the first element of cvar (trivial in this case)
+
+    # Constants
+    # Chn = 0.4  # dimensionless
+    # DChn = -8.0  # dimensionless
+    w_i = 2160.0  # umeter**3  # intracellular volume
+    w_o = 720.0  # umeter**3 # extracellular volume
+    Na_i0 = 16.0  # mMol/m**3 # initial concentration of intracellular Na
+    Na_o0 = 138.0  # mMol/m**3 # initial concentration of extracellular Na
+    K_i0 = 130.0  # mMol/m**3 # initial concentration of intracellular K
+    K_o0 = 4.80  # mMol/m**3 # initial concentration of extracellular K
+
+
+    beta = w_i / w_o
+    DNa_i = -DKi
+    DNa_o = -beta * DNa_i
+    DK_o = -beta * DKi
+    K_i = K_i0 + DKi
+    Na_i = Na_i0 + DNa_i
+    Na_o = Na_o0 + DNa_o
+    K_o = K_o0 + DK_o + Kg
+
+    ninf = n_inf(V)
+    I_K = I_K_form(V, n, K_o, K_i)
+    I_Na = I_Na_form(V, Na_o, Na_i, n)
+    I_Cl = I_Cl_form(V)
+    I_pump = I_pump_form(Na_i, K_o)
+
+    r = R_minus[0] * x / numpy.pi
+    Vdot = (-1.0 / Cm[0]) * (I_Na + I_K + I_Cl + I_pump)
+
+    if V <= Vstar[0]:
+        R, c = R_minus[0], c_minus[0]
+    else:
+        R, c = R_plus[0], c_plus[0]
+
+    dx[0] = Delta[0] + 2 * R * (V - c) * x - J[0] * r * x
+    dx[1] = Vdot - R * x ** 2 + eta[0] + (R_minus[0] * numpy.pi) * Coupling_Term * (E[0] - V)
+    dx[2] = (ninf - n) / tau_n[0]
+    dx[3] = -(gamma[0] / w_i) * (I_K - 2.0 * I_pump)
+    dx[4] = epsilon[0] * (K_bath[0] - K_o)