jaxley 0.1.0__py3-none-any.whl

jaxley/__init__.py

@@ -0,0 +1,13 @@
+# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
+# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from jaxley.connect import (
+    connect,
+    connectivity_matrix_connect,
+    fully_connect,
+    sparse_connect,
+)
+from jaxley.integrate import integrate
+from jaxley.modules import *
+from jaxley.optimize import ParamTransform
+from jaxley.stimulus import datapoint_to_step_currents, step_current

jaxley/build_branched_tridiag.py

@@ -0,0 +1,63 @@
+# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
+# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from math import pi
+
+import jax.numpy as jnp
+from jax import lax, vmap
+
+
+def define_all_tridiags(
+    voltages: jnp.ndarray,
+    voltage_terms: jnp.asarray,
+    i_ext: jnp.ndarray,
+    num_branches: int,
+    coupling_conds_upper: float,
+    coupling_conds_lower: float,
+    summed_coupling_conds: float,
+    dt: float,
+):
+    """
+    Set up tridiagonal system for each branch.
+    """
+    voltages = jnp.reshape(voltages, (num_branches, -1))
+
+    voltage_terms = jnp.reshape(voltage_terms, (num_branches, -1))
+    i_ext = jnp.reshape(i_ext, (num_branches, -1))
+
+    lowers, diags, uppers, solves = vmap(
+        _define_tridiag_for_branch, in_axes=(0, 0, 0, None, 0, 0, 0)
+    )(
+        voltages,
+        voltage_terms,
+        i_ext,
+        dt,
+        coupling_conds_upper,
+        coupling_conds_lower,
+        summed_coupling_conds,
+    )
+
+    return (lowers, diags, uppers, solves)
+
+
+def _define_tridiag_for_branch(
+    voltages: jnp.ndarray,
+    voltage_terms: jnp.ndarray,
+    i_ext: jnp.ndarray,
+    dt: float,
+    coupling_conds_upper: float,
+    coupling_conds_lower: float,
+    summed_coupling_conds: float,
+):
+    """
+    Defines the tridiagonal system to solve for a single branch.
+    """
+
+    # Diagonal and solve.
+    a_v = 1.0 + dt * voltage_terms + dt * summed_coupling_conds
+    b_v = voltages + dt * i_ext
+
+    # Subdiagonals.
+    upper = jnp.asarray(-dt * coupling_conds_upper)
+    lower = jnp.asarray(-dt * coupling_conds_lower)
+    return lower, a_v, upper, b_v

jaxley/channels/__init__.py

@@ -0,0 +1,6 @@
+# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
+# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from jaxley.channels.channel import Channel
+from jaxley.channels.hh import HH
+from jaxley.channels.pospischil import CaL, CaT, K, Km, Leak, Na

jaxley/channels/channel.py

@@ -0,0 +1,81 @@
+# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
+# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Tuple
+
+import jax.numpy as jnp
+
+
+class Channel:
+    """Channel base class. All channels inherit from this class.
+
+    As in NEURON, a `Channel` is considered a distributed process, which means that its
+    conductances are to be specified in `S/cm2` and its currents are to be specified in
+    `uA/cm2`."""
+
+    _name = None
+    channel_params = None
+    channel_states = None
+    current_name = None
+
+    def __init__(self, name: Optional[str] = None):
+        self._name = name if name else self.__class__.__name__
+
+    @property
+    def name(self) -> Optional[str]:
+        """The name of the channel (by default, this is the class name)."""
+        return self._name
+
+    def change_name(self, new_name: str):
+        """Change the channel name.
+
+        Args:
+            new_name: The new name of the channel.
+
+        Returns:
+            Renamed channel, such that this function is chainable.
+        """
+        old_prefix = self._name + "_"
+        new_prefix = new_name + "_"
+
+        self._name = new_name
+        self.channel_params = {
+            (
+                new_prefix + key[len(old_prefix) :]
+                if key.startswith(old_prefix)
+                else key
+            ): value
+            for key, value in self.channel_params.items()
+        }
+
+        self.channel_states = {
+            (
+                new_prefix + key[len(old_prefix) :]
+                if key.startswith(old_prefix)
+                else key
+            ): value
+            for key, value in self.channel_states.items()
+        }
+        return self
+
+    def update_states(
+        self, states, dt, v, params
+    ) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]:
+        """Return the updated states."""
+        raise NotImplementedError
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Given channel states and voltage, return the current through the channel.
+
+        Args:
+            states: All states of the compartment.
+            v: Voltage of the compartment in mV.
+            params: Parameters of the channel (conductances in `S/cm2`).
+
+        Returns:
+            Current in `uA/cm2`.
+        """
+        raise NotImplementedError

jaxley/channels/hh.py

@@ -0,0 +1,98 @@
+# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
+# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Dict, Optional
+
+import jax.numpy as jnp
+
+from jaxley.channels import Channel
+from jaxley.solver_gate import save_exp, solve_gate_exponential
+
+
+class HH(Channel):
+    """Hodgkin-Huxley channel."""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gNa": 0.12,
+            f"{prefix}_gK": 0.036,
+            f"{prefix}_gLeak": 0.0003,
+            f"{prefix}_eNa": 50.0,
+            f"{prefix}_eK": -77.0,
+            f"{prefix}_eLeak": -54.3,
+        }
+        self.channel_states = {
+            f"{prefix}_m": 0.2,
+            f"{prefix}_h": 0.2,
+            f"{prefix}_n": 0.2,
+        }
+        self.current_name = f"i_HH"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """Return updated HH channel state."""
+        prefix = self._name
+        m, h, n = states[f"{prefix}_m"], states[f"{prefix}_h"], states[f"{prefix}_n"]
+        new_m = solve_gate_exponential(m, dt, *self.m_gate(v))
+        new_h = solve_gate_exponential(h, dt, *self.h_gate(v))
+        new_n = solve_gate_exponential(n, dt, *self.n_gate(v))
+        return {f"{prefix}_m": new_m, f"{prefix}_h": new_h, f"{prefix}_n": new_n}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current through HH channels."""
+        prefix = self._name
+        m, h, n = states[f"{prefix}_m"], states[f"{prefix}_h"], states[f"{prefix}_n"]
+
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gNa = params[f"{prefix}_gNa"] * (m**3) * h * 1000  # mS/cm^2
+        gK = params[f"{prefix}_gK"] * n**4 * 1000  # mS/cm^2
+        gLeak = params[f"{prefix}_gLeak"] * 1000  # mS/cm^2
+
+        return (
+            gNa * (v - params[f"{prefix}_eNa"])
+            + gK * (v - params[f"{prefix}_eK"])
+            + gLeak * (v - params[f"{prefix}_eLeak"])
+        )
+
+    def init_state(self, v, params):
+        """Initialize the state such at fixed point of gate dynamics."""
+        prefix = self._name
+        alpha_m, beta_m = self.m_gate(v)
+        alpha_h, beta_h = self.h_gate(v)
+        alpha_n, beta_n = self.n_gate(v)
+        return {
+            f"{prefix}_m": alpha_m / (alpha_m + beta_m),
+            f"{prefix}_h": alpha_h / (alpha_h + beta_h),
+            f"{prefix}_n": alpha_n / (alpha_n + beta_n),
+        }
+
+    @staticmethod
+    def m_gate(v):
+        alpha = 0.1 * _vtrap(-(v + 40), 10)
+        beta = 4.0 * save_exp(-(v + 65) / 18)
+        return alpha, beta
+
+    @staticmethod
+    def h_gate(v):
+        alpha = 0.07 * save_exp(-(v + 65) / 20)
+        beta = 1.0 / (save_exp(-(v + 35) / 10) + 1)
+        return alpha, beta
+
+    @staticmethod
+    def n_gate(v):
+        alpha = 0.01 * _vtrap(-(v + 55), 10)
+        beta = 0.125 * save_exp(-(v + 65) / 80)
+        return alpha, beta
+
+
+def _vtrap(x, y):
+    return x / (save_exp(x / y) - 1.0)

jaxley/channels/pospischil.py

@@ -0,0 +1,377 @@
+# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
+# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Dict, Optional
+
+import jax.numpy as jnp
+
+from jaxley.channels import Channel
+from jaxley.solver_gate import (
+    save_exp,
+    solve_gate_exponential,
+    solve_inf_gate_exponential,
+)
+
+# This is an implementation of Pospischil channels:
+# Leak, Na, K, Km, CaT, CaL
+# [Pospischil et al. Biological Cybernetics (2008)]
+
+__all__ = ["Leak", "Na", "K", "Km", "CaT", "CaL"]
+
+
+# Helper function
+def efun(x):
+    """x/[exp(x)-1]
+
+    Args:
+        x (float): _description_
+
+    Returns:
+        float: x/[exp(x)-1]
+    """
+    return x / (save_exp(x) - 1.0)
+
+
+class Leak(Channel):
+    """Leak current"""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gLeak": 1e-4,
+            f"{prefix}_eLeak": -70.0,
+        }
+        self.channel_states = {}
+        self.current_name = f"i_{prefix}"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """No state to update."""
+        return {}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current."""
+        prefix = self._name
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gLeak = params[f"{prefix}_gLeak"] * 1000  # mS/cm^2
+        return gLeak * (v - params[f"{prefix}_eLeak"])
+
+    def init_state(self, v, params):
+        return {}
+
+
+class Na(Channel):
+    """Sodium channel"""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gNa": 50e-3,
+            "eNa": 50.0,
+            "vt": -60.0,  # Global parameter, not prefixed with `Na`.
+        }
+        self.channel_states = {f"{prefix}_m": 0.2, f"{prefix}_h": 0.2}
+        self.current_name = f"i_Na"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """Update state."""
+        prefix = self._name
+        m, h = states[f"{prefix}_m"], states[f"{prefix}_h"]
+        new_m = solve_gate_exponential(m, dt, *self.m_gate(v, params["vt"]))
+        new_h = solve_gate_exponential(h, dt, *self.h_gate(v, params["vt"]))
+        return {f"{prefix}_m": new_m, f"{prefix}_h": new_h}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current."""
+        prefix = self._name
+        m, h = states[f"{prefix}_m"], states[f"{prefix}_h"]
+
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gNa = params[f"{prefix}_gNa"] * (m**3) * h * 1000  # mS/cm^2
+
+        current = gNa * (v - params["eNa"])
+        return current
+
+    def init_state(self, v, params):
+        """Initialize the state such at fixed point of gate dynamics."""
+        prefix = self._name
+        alpha_m, beta_m = self.m_gate(v, params["vt"])
+        alpha_h, beta_h = self.h_gate(v, params["vt"])
+        return {
+            f"{prefix}_m": alpha_m / (alpha_m + beta_m),
+            f"{prefix}_h": alpha_h / (alpha_h + beta_h),
+        }
+
+    @staticmethod
+    def m_gate(v, vt):
+        v_alpha = v - vt - 13.0
+        alpha = 0.32 * efun(-0.25 * v_alpha) / 0.25
+
+        v_beta = v - vt - 40.0
+        beta = 0.28 * efun(0.2 * v_beta) / 0.2
+        return alpha, beta
+
+    @staticmethod
+    def h_gate(v, vt):
+        v_alpha = v - vt - 17.0
+        alpha = 0.128 * save_exp(-v_alpha / 18.0)
+
+        v_beta = v - vt - 40.0
+        beta = 4.0 / (save_exp(-v_beta / 5.0) + 1.0)
+        return alpha, beta
+
+
+class K(Channel):
+    """Potassium channel"""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gK": 5e-3,
+            "eK": -90.0,
+            "vt": -60.0,  # Global parameter, not prefixed with `Na`.
+        }
+        self.channel_states = {f"{prefix}_n": 0.2}
+        self.current_name = f"i_K"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """Update state."""
+        prefix = self._name
+        n = states[f"{prefix}_n"]
+        new_n = solve_gate_exponential(n, dt, *self.n_gate(v, params["vt"]))
+        return {f"{prefix}_n": new_n}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current."""
+        prefix = self._name
+        n = states[f"{prefix}_n"]
+
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gK = params[f"{prefix}_gK"] * (n**4) * 1000  # mS/cm^2
+
+        return gK * (v - params["eK"])
+
+    def init_state(self, v, params):
+        """Initialize the state such at fixed point of gate dynamics."""
+        prefix = self._name
+        alpha_n, beta_n = self.n_gate(v, params["vt"])
+        return {f"{prefix}_n": alpha_n / (alpha_n + beta_n)}
+
+    @staticmethod
+    def n_gate(v, vt):
+        v_alpha = v - vt - 15.0
+        alpha = 0.032 * efun(-0.2 * v_alpha) / 0.2
+
+        v_beta = v - vt - 10.0
+        beta = 0.5 * save_exp(-v_beta / 40.0)
+        return alpha, beta
+
+
+class Km(Channel):
+    """Slow M Potassium channel"""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gKm": 0.004e-3,
+            f"{prefix}_taumax": 4000.0,
+            f"eK": -90.0,
+        }
+        self.channel_states = {f"{prefix}_p": 0.2}
+        self.current_name = f"i_K"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """Update state."""
+        prefix = self._name
+        p = states[f"{prefix}_p"]
+        new_p = solve_inf_gate_exponential(
+            p, dt, *self.p_gate(v, params[f"{prefix}_taumax"])
+        )
+        return {f"{prefix}_p": new_p}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current."""
+        prefix = self._name
+        p = states[f"{prefix}_p"]
+
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gKm = params[f"{prefix}_gKm"] * p * 1000  # mS/cm^2
+        return gKm * (v - params["eK"])
+
+    def init_state(self, v, params):
+        """Initialize the state such at fixed point of gate dynamics."""
+        prefix = self._name
+        alpha_p, beta_p = self.p_gate(v, params[f"{prefix}_taumax"])
+        return {f"{prefix}_p": alpha_p / (alpha_p + beta_p)}
+
+    @staticmethod
+    def p_gate(v, taumax):
+        v_p = v + 35.0
+        p_inf = 1.0 / (1.0 + save_exp(-0.1 * v_p))
+
+        tau_p = taumax / (3.3 * save_exp(0.05 * v_p) + save_exp(-0.05 * v_p))
+
+        return p_inf, tau_p
+
+
+class CaL(Channel):
+    """L-type Calcium channel"""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gCaL": 0.1e-3,
+            "eCa": 120.0,
+        }
+        self.channel_states = {f"{prefix}_q": 0.2, f"{prefix}_r": 0.2}
+        self.current_name = f"i_Ca"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """Update state."""
+        prefix = self._name
+        q, r = states[f"{prefix}_q"], states[f"{prefix}_r"]
+        new_q = solve_gate_exponential(q, dt, *self.q_gate(v))
+        new_r = solve_gate_exponential(r, dt, *self.r_gate(v))
+        return {f"{prefix}_q": new_q, f"{prefix}_r": new_r}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current."""
+        prefix = self._name
+        q, r = states[f"{prefix}_q"], states[f"{prefix}_r"]
+
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gCaL = params[f"{prefix}_gCaL"] * (q**2) * r * 1000  # mS/cm^2
+
+        return gCaL * (v - params["eCa"])
+
+    def init_state(self, v, params):
+        """Initialize the state such at fixed point of gate dynamics."""
+        prefix = self._name
+        alpha_q, beta_q = self.q_gate(v)
+        alpha_r, beta_r = self.r_gate(v)
+        return {
+            f"{prefix}_q": alpha_q / (alpha_q + beta_q),
+            f"{prefix}_r": alpha_r / (alpha_r + beta_r),
+        }
+
+    @staticmethod
+    def q_gate(v):
+        v_alpha = -v - 27.0
+        alpha = 0.055 * efun(v_alpha / 3.8) * 3.8
+
+        v_beta = -v - 75.0
+        beta = 0.94 * save_exp(v_beta / 17.0)
+        return alpha, beta
+
+    @staticmethod
+    def r_gate(v):
+        v_alpha = -v - 13.0
+        alpha = 0.000457 * save_exp(v_alpha / 50)
+
+        v_beta = -v - 15.0
+        beta = 0.0065 / (save_exp(v_beta / 28.0) + 1)
+        return alpha, beta
+
+
+class CaT(Channel):
+    """T-type Calcium channel"""
+
+    def __init__(self, name: Optional[str] = None):
+        super().__init__(name)
+        prefix = self._name
+        self.channel_params = {
+            f"{prefix}_gCaT": 0.4e-4,
+            f"{prefix}_vx": 2.0,
+            "eCa": 120.0,  # Global parameter, not prefixed with `CaT`.
+        }
+        self.channel_states = {f"{prefix}_u": 0.2}
+        self.current_name = f"i_Ca"
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt,
+        v,
+        params: Dict[str, jnp.ndarray],
+    ):
+        """Update state."""
+        prefix = self._name
+        u = states[f"{prefix}_u"]
+        new_u = solve_inf_gate_exponential(
+            u, dt, *self.u_gate(v, params[f"{prefix}_vx"])
+        )
+        return {f"{prefix}_u": new_u}
+
+    def compute_current(
+        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
+    ):
+        """Return current."""
+        prefix = self._name
+        u = states[f"{prefix}_u"]
+        s_inf = 1.0 / (1.0 + save_exp(-(v + params[f"{prefix}_vx"] + 57.0) / 6.2))
+
+        # Multiply with 1000 to convert Siemens to milli Siemens.
+        gCaT = params[f"{prefix}_gCaT"] * (s_inf**2) * u * 1000  # mS/cm^2
+
+        return gCaT * (v - params["eCa"])
+
+    def init_state(self, v, params):
+        """Initialize the state such at fixed point of gate dynamics."""
+        prefix = self._name
+        alpha_u, beta_u = self.u_gate(v, params[f"{prefix}_vx"])
+        return {f"{prefix}_u": alpha_u / (alpha_u + beta_u)}
+
+    @staticmethod
+    def u_gate(v, vx):
+        v_u1 = v + vx + 81.0
+        u_inf = 1.0 / (1.0 + save_exp(v_u1 / 4))
+
+        tau_u = (30.8 + (211.4 + save_exp((v + vx + 113.2) / 5.0))) / (
+            3.7 * (1 + save_exp((v + vx + 84.0) / 3.2))
+        )
+
+        return u_inf, tau_u