WIP: New channel models

jaxley/channels/channel.py

@@ -16,9 +16,8 @@ class Channel:
     `uA/cm2`."""
 
     _name = None
-    channel_params = None
-    channel_states = None
-    current_name = None
+    params = None
+    states = None
 
     def __init__(self, name: Optional[str] = None):
         contact = (
@@ -59,22 +58,22 @@ def change_name(self, new_name: str):
         new_prefix = new_name + "_"
 
         self._name = new_name
-        self.channel_params = {
+        self.params = {
             (
                 new_prefix + key[len(old_prefix) :]
                 if key.startswith(old_prefix)
                 else key
             ): value
-            for key, value in self.channel_params.items()
+            for key, value in self.params.items()
         }
 
-        self.channel_states = {
+        self.states = {
             (
                 new_prefix + key[len(old_prefix) :]
                 if key.startswith(old_prefix)
                 else key
             ): value
-            for key, value in self.channel_states.items()
+            for key, value in self.states.items()
         }
         return self
 

jaxley/channels/hh.py

@@ -9,72 +9,54 @@
 from jaxley.solver_gate import save_exp, solve_gate_exponential
 
 
-class HH(Channel):
-    """Hodgkin-Huxley channel."""
+class Na(Channel):
+    """Hodgkin-Huxley Sodium channel."""
 
     def __init__(self, name: Optional[str] = None):
         self.current_is_in_mA_per_cm2 = True
 
         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"
+        self.params = {"gNa": 0.12, "eNa": 50.0}
+        self.states = {"m": 0.2, "h": 0.2}
 
     def update_states(
         self,
         states: Dict[str, jnp.ndarray],
-        dt,
-        v,
+        dt: float,
+        v: jnp.ndarray,
         params: Dict[str, jnp.ndarray],
-    ):
+    ) -> 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"]
+        m, h = states["m"], states["h"]
+
         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}
+        return {"m": new_m, "h": new_h}
 
     def compute_current(
-        self, states: Dict[str, jnp.ndarray], v, params: Dict[str, jnp.ndarray]
-    ):
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> 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"]
-
-        gNa = params[f"{prefix}_gNa"] * (m**3) * h  # S/cm^2
-        gK = params[f"{prefix}_gK"] * n**4  # S/cm^2
-        gLeak = params[f"{prefix}_gLeak"]  # S/cm^2
+        m, h = states["m"], states["h"]
+        gNa, eNa = params["gNa"], params["eNa"]
 
-        return (
-            gNa * (v - params[f"{prefix}_eNa"])
-            + gK * (v - params[f"{prefix}_eK"])
-            + gLeak * (v - params[f"{prefix}_eLeak"])
-        )
+        gNa = gNa * (m**3) * h  # S/cm^2
+        return gNa * (v - eNa)
 
-    def init_state(self, states, v, params, delta_t):
+    def init_state(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+        dt: float,
+    ) -> Dict[str, jnp.ndarray]:
         """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),
-        }
+        return {"m": alpha_m / (alpha_m + beta_m), "h": alpha_h / (alpha_h + beta_h)}
 
     @staticmethod
     def m_gate(v):
@@ -88,12 +70,171 @@ def h_gate(v):
         beta = 1.0 / (save_exp(-(v + 35) / 10) + 1)
         return alpha, beta
 
+
+class K(Channel):
+    """Hodgkin-Huxley Potassium channel."""
+
+    def __init__(self, name: Optional[str] = None):
+        self.current_is_in_mA_per_cm2 = True
+
+        super().__init__(name)
+        self.params = {"gK": 0.036, "eK": -77.0}
+        self.states = {"n": 0.2}
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt: float,
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> Dict[str, jnp.ndarray]:
+        """Return updated HH channel state."""
+        n = states["n"]
+
+        new_n = solve_gate_exponential(n, dt, *self.n_gate(v))
+        return {"n": new_n}
+
+    def compute_current(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> jnp.ndarray:
+        """Return current through HH channels."""
+        n = states["n"]
+        gK, eK = params["gK"], params["eK"]
+
+        gK = gK * n**4  # S/cm^2
+        return gK * (v - eK)
+
+    def init_state(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+        dt: float,
+    ) -> Dict[str, jnp.ndarray]:
+        """Initialize the state such at fixed point of gate dynamics."""
+        alpha_n, beta_n = self.n_gate(v)
+        return {"n": alpha_n / (alpha_n + beta_n)}
+
     @staticmethod
     def n_gate(v):
         alpha = 0.01 * _vtrap(-(v + 55), 10)
         beta = 0.125 * save_exp(-(v + 65) / 80)
         return alpha, beta
 
 
+class Leak(Channel):
+    """Hodgkin-Huxley Leak channel."""
+
+    def __init__(self, name: Optional[str] = None):
+        self.current_is_in_mA_per_cm2 = True
+
+        super().__init__(name)
+        self.params = {"gLeak": 0.0003, "eLeak": -54.3}
+        self.states = {}
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt: float,
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> Dict[str, jnp.ndarray]:
+        """Return updated HH channel state."""
+        return {}
+
+    def compute_current(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> jnp.ndarray:
+        """Return current through HH channels."""
+        gLeak, eLeak = params["gLeak"], params["eLeak"]
+
+        return gLeak * (v - eLeak)
+
+    def init_state(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+        dt: float,
+    ) -> Dict[str, jnp.ndarray]:
+        """Initialize the state such at fixed point of gate dynamics."""
+        return {}
+
+
+class HH(Channel):
+    """Hodgkin-Huxley channel."""
+
+    def __init__(self, name: Optional[str] = None):
+        self.current_is_in_mA_per_cm2 = True
+
+        super().__init__(name)
+        self.Na = Na(self._name)
+        self.K = K(self._name)
+        self.Leak = Leak(self._name)
+        self.channels = [self.Na, self.K, self.Leak]
+
+        self.params = {
+            **self.Na.params,
+            **self.K.params,
+            **self.Leak.params,
+        }
+
+        self.states = {
+            **self.Na.states,
+            **self.K.states,
+            **self.Leak.states,
+        }
+
+    def change_name(self, new_name: str):
+        self._name = new_name
+        for channel in self.channels:
+            channel.change_name(new_name)
+        return self
+
+    def update_states(
+        self,
+        states: Dict[str, jnp.ndarray],
+        dt: float,
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> Dict[str, jnp.ndarray]:
+        """Return updated HH channel state."""
+        new_states = {}
+        for channel in self.channels:
+            new_states.update(channel.update_states(states, dt, v, params))
+        return new_states
+
+    def compute_current(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+    ) -> jnp.ndarray:
+        """Return current through HH channels."""
+        current = 0
+        for channel in self.channels:
+            current += channel.compute_current(states, v, params)
+        return current
+
+    def init_state(
+        self,
+        states: Dict[str, jnp.ndarray],
+        v: jnp.ndarray,
+        params: Dict[str, jnp.ndarray],
+        dt: float,
+    ) -> Dict[str, jnp.ndarray]:
+        """Initialize the state such at fixed point of gate dynamics."""
+        init_states = {}
+        for channel in self.channels:
+            init_states.update(channel.init_state(states, v, params, dt))
+        return init_states
+
+
 def _vtrap(x, y):
     return x / (save_exp(x / y) - 1.0)