{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Channels\n",
    "\n",
    "In biological neurons, ion channels control the permeability of ions such as sodium, potassium, and calcium, which form the basis of neural signal generation and transmission. Different types of channels have their own kinetic properties and influence changes in membrane potential.\n",
    "\n",
    "Channel modeling in `braincell` has the following notable advantages:\n",
    "- Unified interface: whether it is a linear channel or a complex Markov model, all can be represented and integrated in a unified way.\n",
    "- High flexibility: supports multiple types of currents and allows user-defined modeling.\n",
    "- Compatible with `brainstate`'s automatic differentiation system: each state variable can be precisely differentiated automatically, supporting efficient training and optimization.\n",
    "- High extensibility: through inheritance and composition mechanisms, channel models can easily build subclasses, allowing parameter modification and extension.\n",
    "\n",
    "`braincell` is highly flexible. For `Channel`, `braincell` supports two ways of use: defining new ion channels, or using existing ones.\n",
    "\n",
    "## Using existing channels\n",
    "\n",
    "Let’s first explain how to use existing channels.\n",
    "\n",
    "According to [different channel functions](https://elifesciences.org/articles/22152), we classify `Channel` into the following base classes:\n",
    "- `Calcium Channels`\n",
    "- `Potassium Channels`\n",
    "- `Sodium Channels`\n",
    "- `Potassium Calcium Channels`\n",
    "- `Hyperpolarization-Activated Channels`\n",
    "- `Leakage Channels`\n",
    "\n",
    "`braincell` has already built in many specific channels under each base class, and you can use them freely.\n",
    "If you want to learn more about the built-in channels, you can check the [Ion Channel Collection](../apis/braincell.channel.rst) for details.\n",
    "\n",
    "When modeling neurons in practice, whether using existing channels or custom ones, we need to call them.\n",
    "Of course, this call is quite simple, but we should also pay attention to some usage conventions.\n",
    "\n",
    "Now let’s look at an example of modeling the `HTC` neuron:"
   ],
   "id": "7bcd9d9ab53a7700"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:56.841546Z",
     "start_time": "2025-10-13T09:12:56.838882Z"
    }
   },
   "source": [
    "import brainstate\n",
    "import braintools\n",
    "import brainunit as u\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import braincell"
   ],
   "outputs": [],
   "execution_count": 29
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:56.859633Z",
     "start_time": "2025-10-13T09:12:56.852197Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class HTC(braincell.SingleCompartment):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size,\n",
    "        gKL=0.01 * (u.mS / u.cm ** 2),\n",
    "        V_initializer=braintools.init.Constant(-65. * u.mV),\n",
    "        solver: str = 'ind_exp_euler'\n",
    "    ):\n",
    "        super().__init__(size, V_initializer=V_initializer, V_th=20. * u.mV, solver=solver)\n",
    "\n",
    "        self.area = 1e-3 / (2.9e-4 * u.cm ** 2)\n",
    "\n",
    "        self.na = braincell.ion.SodiumFixed(size, E=50. * u.mV)\n",
    "        self.na.add(INa=braincell.channel.INa_Ba2002(size, V_sh=-30 * u.mV))\n",
    "\n",
    "        self.k = braincell.ion.PotassiumFixed(size, E=-90. * u.mV)\n",
    "        self.k.add(IKL=braincell.channel.IK_Leak(size, g_max=gKL))\n",
    "        self.k.add(IDR=braincell.channel.IKDR_Ba2002(size, V_sh=-30. * u.mV, phi=0.25))\n",
    "\n",
    "        self.ca = braincell.ion.CalciumDetailed(\n",
    "            size,\n",
    "            C_rest=5e-5 * u.mM,\n",
    "            tau=10. * u.ms,\n",
    "            d=0.5 * u.um\n",
    "        )\n",
    "        self.ca.add(ICaL=braincell.channel.ICaL_IS2008(size, g_max=0.5 * (u.mS / u.cm ** 2)))\n",
    "        self.ca.add(ICaN=braincell.channel.ICaN_IS2008(size, g_max=0.5 * (u.mS / u.cm ** 2)))\n",
    "        self.ca.add(ICaT=braincell.channel.ICaT_HM1992(size, g_max=2.1 * (u.mS / u.cm ** 2)))\n",
    "        self.ca.add(ICaHT=braincell.channel.ICaHT_HM1992(size, g_max=3.0 * (u.mS / u.cm ** 2)))\n",
    "\n",
    "        self.kca = braincell.MixIons(self.k, self.ca)\n",
    "        self.kca.add(IAHP=braincell.channel.IAHP_De1994(size, g_max=0.3 * (u.mS / u.cm ** 2)))\n",
    "\n",
    "        self.Ih = braincell.channel.Ih_HM1992(size, g_max=0.01 * (u.mS / u.cm ** 2), E=-43 * u.mV)\n",
    "        self.IL = braincell.channel.IL(size, g_max=0.0075 * (u.mS / u.cm ** 2), E=-70 * u.mV)"
   ],
   "id": "5df22208c56da18a",
   "outputs": [],
   "execution_count": 30
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Through this example of the `HTC` neuron, it is clear that in practical modeling we only need to simply pass in the relevant channels.\n",
    "\n",
    "It is worth noting, however, that the channels in this example have already been explicitly defined in the `braincell` codebase.\n",
    "If you want to use other channels, you can either define your custom channels within the `braincell` codebase, or use `import` to bring in your custom channels."
   ],
   "id": "ee700b923eeb837e"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Use of `root_type`\n",
    "\n",
    "When using each channel, we need to pay attention to the use of `root_type`. Generally speaking, calcium channels depend on the `Calcium` ion, potassium channels depend on the `Potassium` ion, and sodium channels depend on the `Sodium` ion."
   ],
   "id": "54f6d86cac6bb5e1"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:56.872911Z",
     "start_time": "2025-10-13T09:12:56.870360Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.CalciumChannel.root_type",
   "id": "80d48d2760f737a0",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.ion.Calcium"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 31
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:56.898181Z",
     "start_time": "2025-10-13T09:12:56.895886Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.PotassiumChannel.root_type",
   "id": "40172773bc9debce",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.ion.Potassium"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 32
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:56.929861Z",
     "start_time": "2025-10-13T09:12:56.927338Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.SodiumChannel.root_type",
   "id": "d82f7a2b9be5dff4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.ion.Sodium"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 33
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Calcium-dependent potassium channels depend on both `Potassium` and `Calcium` ions.\n",
   "id": "5e7aba3d9696401a"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:56.968056Z",
     "start_time": "2025-10-13T09:12:56.965079Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.IAHP_De1994.root_type",
   "id": "2c6193de3d536b3b",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JointTypes[braincell.ion.Potassium, braincell.ion.Calcium]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 34
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "For some special channels, such as `IL`, they do not depend on any type of ion, so their `root_type` is the cell.\n",
   "id": "6a666d1309b03f94"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.001655Z",
     "start_time": "2025-10-13T09:12:56.998742Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.IL.root_type",
   "id": "62c1e064a340ed54",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.HHTypedNeuron"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 35
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Adding ion channels\n",
    "\n",
    "When performing single-cell modeling, we usually need to add ion channels to ions. If the use of `root_type` is not handled correctly, ion channels may not work properly. For example, sodium channels can only be added to sodium ions, and potassium channels can only be added to potassium ions. If you try to add a sodium channel to potassium ions, or add a potassium channel to calcium ions, it will result in an error.\n"
   ],
   "id": "bab63c8c58d01f6"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.041351Z",
     "start_time": "2025-10-13T09:12:57.039493Z"
    }
   },
   "cell_type": "code",
   "source": "na = braincell.ion.SodiumFixed(1)",
   "id": "462a0d861dd862c6",
   "outputs": [],
   "execution_count": 36
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.068737Z",
     "start_time": "2025-10-13T09:12:57.066632Z"
    }
   },
   "cell_type": "code",
   "source": "na.add(ina=braincell.channel.INa_HH1952(1))",
   "id": "e2944d2640832e64",
   "outputs": [],
   "execution_count": 37
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.087519Z",
     "start_time": "2025-10-13T09:12:57.084216Z"
    }
   },
   "cell_type": "code",
   "source": [
    "try:\n",
    "    na.add(ik=braincell.channel.IK_HH1952(1))\n",
    "except Exception as e:\n",
    "    print(e)"
   ],
   "id": "cdb99af405134cc",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Type does not match. IK_HH1952(\n",
      "  size=(1,),\n",
      "  name=None,\n",
      "  g_max=10. * msiemens / cmeter2,\n",
      "  phi=1.0,\n",
      "  V_sh=-45. * mvolt\n",
      ") requires a root with type of <class 'braincell.ion.Potassium'>, but the root now is <class 'braincell.ion.SodiumFixed'>.\n"
     ]
    }
   ],
   "execution_count": 38
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Simulating ion channels\n",
    "\n",
    "We can also directly simulate ion channels without adding them to ions. This allows for more flexible use of ion channels. However, when simulating ion channels, we need to provide the required information such as the neuron's membrane potential and ion concentrations. Let’s take a calcium channel as an example.\n"
   ],
   "id": "1bdced6363f4a227"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.130868Z",
     "start_time": "2025-10-13T09:12:57.127136Z"
    }
   },
   "cell_type": "code",
   "source": [
    "ca = braincell.ion.CalciumFixed(1).pack_info()\n",
    "V = -65 * u.mV\n",
    "\n",
    "# The calcium channel to be simulated\n",
    "can = braincell.channel.ICaL_IS2008(1)\n",
    "can.init_state(V, ca)"
   ],
   "id": "b261b5cfacb0b617",
   "outputs": [],
   "execution_count": 39
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:14:14.408610Z",
     "start_time": "2025-10-13T09:14:14.017548Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def run_can(i):\n",
    "    dt = 0.1 * u.ms\n",
    "    t = i * dt\n",
    "    with brainstate.environ.context(i=i, t=t, dt=dt):\n",
    "        braincell.ind_exp_euler_step(can, t, dt, V, ca)\n",
    "    return can.p.value, can.q.value\n",
    "\n",
    "\n",
    "indices = u.math.arange(1000)\n",
    "ps, qs = brainstate.transform.for_loop(run_can, indices)\n",
    "\n",
    "plt.plot(indices, ps, label='$p$')\n",
    "plt.plot(indices, qs, label='$q$')\n",
    "plt.xlabel('Time (ms)')\n",
    "plt.ylabel('State Value')\n",
    "plt.title('ICaL IS2008 State Dynamics')\n",
    "plt.legend()\n",
    "plt.show()"
   ],
   "id": "2be233ea1c7e7ab9",
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "CalciumChannel.pre_integral() takes 3 positional arguments but 5 were given",
     "output_type": "error",
     "traceback": [
      "\u001B[31m---------------------------------------------------------------------------\u001B[39m",
      "\u001B[31mTypeError\u001B[39m                                 Traceback (most recent call last)",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[41]\u001B[39m\u001B[32m, line 10\u001B[39m\n\u001B[32m      6\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m can.p.value, can.q.value\n\u001B[32m      9\u001B[39m indices = u.math.arange(\u001B[32m1000\u001B[39m)\n\u001B[32m---> \u001B[39m\u001B[32m10\u001B[39m ps, qs = \u001B[43mbrainstate\u001B[49m\u001B[43m.\u001B[49m\u001B[43mtransform\u001B[49m\u001B[43m.\u001B[49m\u001B[43mfor_loop\u001B[49m\u001B[43m(\u001B[49m\u001B[43mrun_can\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mindices\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m     12\u001B[39m plt.plot(indices, ps, label=\u001B[33m'\u001B[39m\u001B[33m$p$\u001B[39m\u001B[33m'\u001B[39m)\n\u001B[32m     13\u001B[39m plt.plot(indices, qs, label=\u001B[33m'\u001B[39m\u001B[33m$q$\u001B[39m\u001B[33m'\u001B[39m)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_loop_collect_return.py:541\u001B[39m, in \u001B[36mfor_loop\u001B[39m\u001B[34m(f, length, reverse, unroll, pbar, *xs)\u001B[39m\n\u001B[32m    465\u001B[39m \u001B[38;5;129m@set_module_as\u001B[39m(\u001B[33m'\u001B[39m\u001B[33mbrainstate.transform\u001B[39m\u001B[33m'\u001B[39m)\n\u001B[32m    466\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mfor_loop\u001B[39m(\n\u001B[32m    467\u001B[39m     f: Callable[..., Y],\n\u001B[32m   (...)\u001B[39m\u001B[32m    472\u001B[39m     pbar: Optional[ProgressBar | \u001B[38;5;28mint\u001B[39m] = \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[32m    473\u001B[39m ) -> Y:\n\u001B[32m    474\u001B[39m \u001B[38;5;250m    \u001B[39m\u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m    475\u001B[39m \u001B[33;03m    ``for-loop`` control flow with :py:class:`~.State`.\u001B[39;00m\n\u001B[32m    476\u001B[39m \n\u001B[32m   (...)\u001B[39m\u001B[32m    539\u001B[39m \u001B[33;03m        >>> results = brainstate.transform.for_loop(process_item, xs, ys, reverse=True)\u001B[39;00m\n\u001B[32m    540\u001B[39m \u001B[33;03m    \"\"\"\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m541\u001B[39m     _, ys = \u001B[43mscan\u001B[49m\u001B[43m(\u001B[49m\n\u001B[32m    542\u001B[39m \u001B[43m        \u001B[49m\u001B[43m_forloop_to_scan_fun\u001B[49m\u001B[43m(\u001B[49m\u001B[43mf\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m    543\u001B[39m \u001B[43m        \u001B[49m\u001B[43minit\u001B[49m\u001B[43m=\u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[32m    544\u001B[39m \u001B[43m        \u001B[49m\u001B[43mxs\u001B[49m\u001B[43m=\u001B[49m\u001B[43mxs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m    545\u001B[39m \u001B[43m        \u001B[49m\u001B[43mlength\u001B[49m\u001B[43m=\u001B[49m\u001B[43mlength\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m    546\u001B[39m \u001B[43m        \u001B[49m\u001B[43mreverse\u001B[49m\u001B[43m=\u001B[49m\u001B[43mreverse\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m    547\u001B[39m \u001B[43m        \u001B[49m\u001B[43munroll\u001B[49m\u001B[43m=\u001B[49m\u001B[43munroll\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m    548\u001B[39m \u001B[43m        \u001B[49m\u001B[43mpbar\u001B[49m\u001B[43m=\u001B[49m\u001B[43mpbar\u001B[49m\n\u001B[32m    549\u001B[39m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m    550\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m ys\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_loop_collect_return.py:248\u001B[39m, in \u001B[36mscan\u001B[39m\u001B[34m(f, init, xs, length, reverse, unroll, pbar)\u001B[39m\n\u001B[32m    246\u001B[39m x_avals = [jax.core.mapped_aval(length, \u001B[32m0\u001B[39m, aval) \u001B[38;5;28;01mfor\u001B[39;00m aval \u001B[38;5;129;01min\u001B[39;00m xs_avals]\n\u001B[32m    247\u001B[39m args = [init, xs_tree.unflatten(x_avals)]\n\u001B[32m--> \u001B[39m\u001B[32m248\u001B[39m stateful_fun = \u001B[43mStatefulFunction\u001B[49m\u001B[43m(\u001B[49m\u001B[43mf\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mname\u001B[49m\u001B[43m=\u001B[49m\u001B[33;43m'\u001B[39;49m\u001B[33;43mscan\u001B[39;49m\u001B[33;43m'\u001B[39;49m\u001B[43m)\u001B[49m\u001B[43m.\u001B[49m\u001B[43mmake_jaxpr\u001B[49m\u001B[43m(\u001B[49m\u001B[43m*\u001B[49m\u001B[43margs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m    249\u001B[39m state_trace = stateful_fun.get_state_trace(*args)\n\u001B[32m    250\u001B[39m all_writen_state_vals = state_trace.get_write_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_make_jaxpr.py:1082\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m   1080\u001B[39m         \u001B[38;5;28mself\u001B[39m._cached_out_shapes.pop(cache_key, \u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m   1081\u001B[39m         \u001B[38;5;28mself\u001B[39m._cached_jaxpr.pop(cache_key, \u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m-> \u001B[39m\u001B[32m1082\u001B[39m         \u001B[38;5;28;01mraise\u001B[39;00m e\n\u001B[32m   1084\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_make_jaxpr.py:1060\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m   1058\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m   1059\u001B[39m         dyn_kwargs[k] = v\n\u001B[32m-> \u001B[39m\u001B[32m1060\u001B[39m jaxpr, (out_shapes, state_shapes) = \u001B[43m_make_jaxpr\u001B[49m\u001B[43m(\u001B[49m\n\u001B[32m   1061\u001B[39m \u001B[43m    \u001B[49m\u001B[43mfunctools\u001B[49m\u001B[43m.\u001B[49m\u001B[43mpartial\u001B[49m\u001B[43m(\u001B[49m\n\u001B[32m   1062\u001B[39m \u001B[43m        \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m.\u001B[49m\u001B[43m_wrapped_fun_to_eval\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1063\u001B[39m \u001B[43m        \u001B[49m\u001B[43mcache_key\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1064\u001B[39m \u001B[43m        \u001B[49m\u001B[43mstatic_kwargs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1065\u001B[39m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1066\u001B[39m \u001B[43m    \u001B[49m\u001B[43mstatic_argnums\u001B[49m\u001B[43m=\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m.\u001B[49m\u001B[43mstatic_argnums\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1067\u001B[39m \u001B[43m    \u001B[49m\u001B[43maxis_env\u001B[49m\u001B[43m=\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m.\u001B[49m\u001B[43maxis_env\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1068\u001B[39m \u001B[43m    \u001B[49m\u001B[43mreturn_shape\u001B[49m\u001B[43m=\u001B[49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[32m   1069\u001B[39m \u001B[43m    \u001B[49m\u001B[43mabstracted_axes\u001B[49m\u001B[43m=\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m.\u001B[49m\u001B[43mabstracted_axes\u001B[49m\u001B[43m,\u001B[49m\n\u001B[32m   1070\u001B[39m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\u001B[43m(\u001B[49m\u001B[43m*\u001B[49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m*\u001B[49m\u001B[43m*\u001B[49m\u001B[43mdyn_kwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m   1072\u001B[39m \u001B[38;5;66;03m# returns\u001B[39;00m\n\u001B[32m   1073\u001B[39m \u001B[38;5;28mself\u001B[39m._cached_jaxpr_out_tree.set(cache_key, jax.tree.structure((out_shapes, state_shapes)))\n",
      "    \u001B[31m[... skipping hidden 1 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_make_jaxpr.py:1996\u001B[39m, in \u001B[36m_make_jaxpr.<locals>.make_jaxpr_f\u001B[39m\u001B[34m(*args, **kwargs)\u001B[39m\n\u001B[32m   1994\u001B[39m         jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info_)\n\u001B[32m   1995\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m1996\u001B[39m         jaxpr, out_type, consts = \u001B[43mpe\u001B[49m\u001B[43m.\u001B[49m\u001B[43mtrace_to_jaxpr_dynamic2\u001B[49m\u001B[43m(\u001B[49m\u001B[43mf\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m   1997\u001B[39m closed_jaxpr = ClosedJaxpr(jaxpr, consts)\n\u001B[32m   1998\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m return_shape:\n",
      "    \u001B[31m[... skipping hidden 5 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_make_jaxpr.py:1006\u001B[39m, in \u001B[36mStatefulFunction._wrapped_fun_to_eval\u001B[39m\u001B[34m(self, cache_key, static_kwargs, *args, **dyn_kwargs)\u001B[39m\n\u001B[32m   1004\u001B[39m \u001B[38;5;28mself\u001B[39m._cached_state_trace.set(cache_key, state_trace)\n\u001B[32m   1005\u001B[39m \u001B[38;5;28;01mwith\u001B[39;00m state_trace:\n\u001B[32m-> \u001B[39m\u001B[32m1006\u001B[39m     out = \u001B[38;5;28;43mself\u001B[39;49m\u001B[43m.\u001B[49m\u001B[43mfun\u001B[49m\u001B[43m(\u001B[49m\u001B[43m*\u001B[49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m*\u001B[49m\u001B[43m*\u001B[49m\u001B[43mdyn_kwargs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m*\u001B[49m\u001B[43m*\u001B[49m\u001B[43mstatic_kwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m   1007\u001B[39m     state_values = (\n\u001B[32m   1008\u001B[39m         state_trace.get_write_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n\u001B[32m   1009\u001B[39m         \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m.return_only_write \u001B[38;5;28;01melse\u001B[39;00m\n\u001B[32m   1010\u001B[39m         state_trace.get_state_values()\n\u001B[32m   1011\u001B[39m     )\n\u001B[32m   1012\u001B[39m state_trace.recovery_original_values()\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\brainstate\\transform\\_loop_collect_return.py:460\u001B[39m, in \u001B[36m_forloop_to_scan_fun.<locals>.scan_fun\u001B[39m\u001B[34m(carry, x)\u001B[39m\n\u001B[32m    458\u001B[39m \u001B[38;5;129m@wraps\u001B[39m(f)\n\u001B[32m    459\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mscan_fun\u001B[39m(carry, x):\n\u001B[32m--> \u001B[39m\u001B[32m460\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m carry, \u001B[43mf\u001B[49m\u001B[43m(\u001B[49m\u001B[43m*\u001B[49m\u001B[43mx\u001B[49m\u001B[43m)\u001B[49m\n",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[41]\u001B[39m\u001B[32m, line 5\u001B[39m, in \u001B[36mrun_can\u001B[39m\u001B[34m(i)\u001B[39m\n\u001B[32m      3\u001B[39m t = i * dt\n\u001B[32m      4\u001B[39m \u001B[38;5;28;01mwith\u001B[39;00m brainstate.environ.context(i=i, t=t, dt=dt):\n\u001B[32m----> \u001B[39m\u001B[32m5\u001B[39m     \u001B[43mbraincell\u001B[49m\u001B[43m.\u001B[49m\u001B[43mind_exp_euler_step\u001B[49m\u001B[43m(\u001B[49m\u001B[43mcan\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mt\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mdt\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mV\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mca\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m      6\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m can.p.value, can.q.value\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\braincell\\_integrator_exp_euler.py:229\u001B[39m, in \u001B[36mind_exp_euler_step\u001B[39m\u001B[34m(target, excluded_paths, *args)\u001B[39m\n\u001B[32m    226\u001B[39m dt = brainstate.environ.get(\u001B[33m'\u001B[39m\u001B[33mdt\u001B[39m\u001B[33m'\u001B[39m)\n\u001B[32m    228\u001B[39m \u001B[38;5;66;03m# Pre-integration hook (e.g., update gating variables)\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m229\u001B[39m \u001B[43mtarget\u001B[49m\u001B[43m.\u001B[49m\u001B[43mpre_integral\u001B[49m\u001B[43m(\u001B[49m\u001B[43m*\u001B[49m\u001B[43margs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m    231\u001B[39m \u001B[38;5;66;03m# Retrieve all states from the target module\u001B[39;00m\n\u001B[32m    232\u001B[39m all_states, diffeq_states, other_states = split_diffeq_states(target)\n",
      "\u001B[31mTypeError\u001B[39m: CalciumChannel.pre_integral() takes 3 positional arguments but 5 were given"
     ]
    }
   ],
   "execution_count": 41
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Custom Channel Modeling\n",
    "\n",
    "Now let’s explain how to use custom channels.\n",
    "\n",
    "In `braincell`, we use the `IonChannel` class as the base class for channel modeling. All specific channel models should inherit from this class.\n",
    "\n",
    "`IonChannel` provides the following interfaces to support key operations during simulation:\n",
    "- `current`: compute the ionic current generated by the channel.\n",
    "- `init_state`: initialize the channel’s state variables.\n",
    "- `reset_state`: reset the channel’s state.\n",
    "- `compute_derivative`: compute the derivatives of the channel’s internal state variables, used for the integrator.\n",
    "- `pre_integral` and `post_integral`: operations before/after integration, usually unused.\n",
    "\n",
    "`compute_derivative`, `pre_integral`, and `post_integral` are the core interfaces for integration and are typically implemented in specific channels. For a detailed introduction, see the [Differential Systems](../advanced_tutorial/differential_equation-zh.ipynb) section. These methods integrate seamlessly with `braincell`’s differential system. When customizing ion channels, you need to properly set the relevant interfaces to achieve the desired functionality.\n",
    "\n",
    "`IonChannel` contains both `Ion` and `Channel`. Therefore, we define a wrapper base class `Channel` that directly inherits from `IonChannel`. For semantic clarity, we recommend using `Channel` when developing practical channel models.\n",
    "\n",
    "Once the `Channel` base class is in place, it’s time to build specific types of channels. In practice, we want subclasses of `Channel` to implement the necessary methods that define the behavior of specific channel types.\n",
    "\n",
    "Now let’s go through how to build custom channels for different purposes.\n",
    "\n",
    "### Base Channel Modeling\n",
    "\n",
    "First, when modeling a new channel, we need to determine whether it belongs to one of the existing base channel categories.\n",
    "\n",
    "As mentioned earlier, `braincell` classifies `Channel` into the following types:\n",
    "- `Calcium Channels`\n",
    "- `Potassium Channels`\n",
    "- `Sodium Channels`\n",
    "- `Potassium Calcium Channels`\n",
    "- `Hyperpolarization-Activated Channels`\n",
    "- `Leakage Channels`\n",
    "\n",
    "All of these channels directly inherit from the `Channel` base class. If the channel you want to customize belongs to one of these classes, refer directly to the following tutorials on subclass and specific channel modeling.\n",
    "If your channel does not fit into these base classes, we will use the modeling of `Calcium Channels` as an example to show how to create a new base channel class.\n",
    "\n",
    "In the `braincell.channel` module, all ion channels ultimately inherit from the unified abstract base class `Channel`. This class defines the minimal set of interfaces that every channel must implement, providing a unified framework for general modeling.\n",
    "\n",
    "In practice, we often need to define intermediate base classes for different ion types. These type-specific base classes encapsulate channel type, input parameters, and interface structures, creating a clearer and more maintainable modeling hierarchy.\n",
    "\n",
    "Here’s an example of modeling a `Calcium Channel`:"
   ],
   "id": "f00791a9770cb002"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.287734700Z",
     "start_time": "2025-10-13T08:10:12.765814Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CalciumChannel(braincell.Channel):\n",
    "    # Specify the ion type this channel acts on\n",
    "    root_type = braincell.ion.Calcium\n",
    "\n",
    "    def pre_integral(self, V, Ca: braincell.IonInfo):\n",
    "        pass\n",
    "\n",
    "    def post_integral(self, V, Ca: braincell.IonInfo):\n",
    "        pass\n",
    "\n",
    "    def compute_derivative(self, V, Ca: braincell.IonInfo):\n",
    "        pass\n",
    "\n",
    "    def current(self, V, Ca: braincell.IonInfo):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def init_state(self, V, Ca: braincell.IonInfo, batch_size=None):\n",
    "        pass\n",
    "\n",
    "    def reset_state(self, V, Ca: braincell.IonInfo, batch_size=None):\n",
    "        pass"
   ],
   "id": "5a04727b0f525155",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The above code easily completes the modeling of `CalciumChannel`.\n",
    "\n",
    "Looking at the code, the most important part is defining `root_type = Calcium`.\n",
    "This is the core of `CalciumChannel`, which tells `braincell`:\n",
    "this channel only acts on `Calcium`, and the parameter `Ca` must be of type `IonInfo`, representing the calcium ion state.\n",
    "This ensures that when you implement interfaces such as `current` or `compute_derivative`, you can access the corresponding ion information.\n",
    "\n",
    "Once you have defined `CalciumChannel`, all subsequent specific calcium channel subclasses can inherit from it:\n",
    "at this point you no longer need to worry about the source of `IonInfo`, since `braincell` will automatically pass the `Calcium` ion to you.\n",
    "\n",
    "At the same time, when modeling a new base channel class, you need to pay attention to defining its interfaces.\n",
    "These interfaces can follow those defined in `Channel`, or they can be customized depending on actual needs.\n",
    "\n",
    "By constructing type-specific base classes like `CalciumChannel`, we can:\n",
    "- Clearly specify the ion type applicable to the model.\n",
    "- Provide a unified interface.\n",
    "- Simplify the writing of downstream models.\n",
    "- Improve code reusability and maintainability.\n",
    "\n",
    "This design also applies to sodium channels (`SodiumChannel`), potassium channels (`PotassiumChannel`), and other ion channels,\n",
    "and it is recommended as the standard approach for writing custom channel models.\n",
    "\n",
    "### Specific Channel Modeling\n",
    "\n",
    "Next, let’s use the code for `ICaT_HP1992` as an example to demonstrate how to define a custom ion channel model.\n",
    "\n",
    "The `ICaT_HP1992` ion channel is a calcium channel model, based on the calcium channel model proposed by Huguenard & Prince in 1992:\n",
    "[<A novel T-type current underlies prolonged Ca(2+)-dependent burst firing in GABAergic neurons of rat thalamic reticular nucleus>](https://pubmed.ncbi.nlm.nih.gov/1403085/).\n",
    "\n",
    "Modeling it only requires referring to the specific equations and implementing them directly in code."
   ],
   "id": "6dff80a298d20c7"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "First, we define the parameters:\n",
    "\n",
    "- $p$ and $q$ are two gating variables for activation and inactivation.\n",
    "- $\\phi_p = 5^{\\frac{T-24}{10}}$ and $\\phi_q = 3^{\\frac{T-24}{10}}$ are temperature-dependent factors ($T$ is the temperature in Celsius).\n",
    "- $E_{Ca}$ is the reversal potential of the calcium channel.\n",
    "- $p_{\\infty}$ and $q_{\\infty}$ are the steady-state activation/inactivation functions of $p$ and $q$ (voltage-dependent).\n",
    "- $\\tau_p$ and $\\tau_q$ are voltage-dependent time constants.\n",
    "- $\\phi_p$ and $\\phi_q$ are temperature correction factors.\n",
    "- $g_{max}$ is the maximum conductance.\n",
    "- $V$ is the membrane potential, and $E_{Ca}$ is the calcium reversal potential.\n",
    "- $V_{sh}$ is the membrane potential shift.\n"
   ],
   "id": "e312d06d29da49c4"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.287734700Z",
     "start_time": "2025-10-13T08:10:16.151129Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class ICaT_HP1992(CalciumChannel):\n",
    "    root_type = braincell.ion.Calcium\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        size: brainstate.typing.Size,\n",
    "        T: brainstate.typing.ArrayLike = u.celsius2kelvin(36.),\n",
    "        T_base_p: brainstate.typing.ArrayLike = 5.,\n",
    "        T_base_q: brainstate.typing.ArrayLike = 3.,\n",
    "        g_max=1.75 * (u.mS / u.cm ** 2),\n",
    "        V_sh=-3. * u.mV,\n",
    "        phi_p=None,\n",
    "        phi_q=None,\n",
    "    ):\n",
    "        super().__init__(size)\n",
    "\n",
    "        T = u.kelvin2celsius(T)\n",
    "        phi_p = T_base_p ** ((T - 24) / 10) if phi_p is None else phi_p\n",
    "        phi_q = T_base_q ** ((T - 24) / 10) if phi_q is None else phi_q\n",
    "        # parameters\n",
    "        self.phi_p = braintools.init.param(phi_p, self.varshape, allow_none=False)\n",
    "        self.phi_q = braintools.init.param(phi_q, self.varshape, allow_none=False)\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "        self.T = braintools.init.param(T, self.varshape, allow_none=False)\n",
    "        self.T_base_p = braintools.init.param(T_base_p, self.varshape, allow_none=False)\n",
    "        self.T_base_q = braintools.init.param(T_base_q, self.varshape, allow_none=False)\n",
    "        self.V_sh = braintools.init.param(V_sh, self.varshape, allow_none=False)"
   ],
   "id": "37c57b580914e4a0",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Then for $p_{\\infty}$ we have\n",
    "\n",
    "$$\n",
    "p_{\\infty} = \\frac{1}{1 + \\exp\\!\\left[-\\frac{V + 52 - V_{sh}}{7.4}\\right]}\n",
    "$$\n"
   ],
   "id": "1c4319fe0d334ad0"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.288735900Z",
     "start_time": "2025-09-27T12:52:30.754174Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def f_p_inf(self, V):\n",
    "        V = (V - self.V_sh).to_decimal(u.mV)\n",
    "        return 1. / (1. + u.math.exp(-(V + 52.) / 7.4))"
   ],
   "id": "a0409a7cb7ef15c3",
   "outputs": [],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "For $\\tau_p$ we have:\n",
    "\n",
    "$$\n",
    "\\tau_p = 3 + \\frac{1}{\\exp\\!\\left(\\frac{V + 27 - V_{sh}}{10}\\right) + \\exp\\!\\left(-\\frac{V + 102 - V_{sh}}{15}\\right)}\n",
    "$$\n"
   ],
   "id": "6149dfab2704a1d6"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.293736300Z",
     "start_time": "2025-09-27T12:52:30.809699Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def f_p_tau(self, V):\n",
    "        V = (V - self.V_sh).to_decimal(u.mV)\n",
    "        return 3. + 1. / (u.math.exp((V + 27.) / 10.) +\n",
    "                          u.math.exp(-(V + 102.) / 15.))"
   ],
   "id": "f31c2252658775",
   "outputs": [],
   "execution_count": 19
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "For $q_{\\infty}$ we have:\n",
    "\n",
    "$$\n",
    "q_{\\infty} = \\frac{1}{1 + \\exp\\!\\left(\\frac{V + 80 - V_{sh}}{5}\\right)}\n",
    "$$\n"
   ],
   "id": "64821aa73d09954f"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.295735500Z",
     "start_time": "2025-09-27T12:52:30.837381Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def f_q_inf(self, V):\n",
    "        V = (V - self.V_sh).to_decimal(u.mV)\n",
    "        return 1. / (1. + u.math.exp((V + 80.) / 5.))"
   ],
   "id": "a78ac03bccfd6f5b",
   "outputs": [],
   "execution_count": 20
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "For $\\tau_q$ we have:\n",
    "\n",
    "$$\n",
    "\\tau_q = 85 + \\frac{1}{\\exp\\!\\left(\\frac{V + 48 - V_{sh}}{4}\\right) + \\exp\\!\\left(-\\frac{V + 407 - V_{sh}}{50}\\right)}\n",
    "$$\n"
   ],
   "id": "f66f683e5820d302"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.295735500Z",
     "start_time": "2025-09-27T12:52:30.859093Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def f_q_tau(self, V):\n",
    "        V = (V - self.V_sh).to_decimal(u.mV)\n",
    "        return 85. + 1. / (u.math.exp((V + 48.) / 4.) + u.math.exp(-(V + 407.) / 50.))"
   ],
   "id": "126af3acae1b9bc7",
   "outputs": [],
   "execution_count": 21
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "At the same time, we need to set the `state`:",
   "id": "84473e09ba6c440f"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.296736900Z",
     "start_time": "2025-09-27T12:52:30.880310Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def init_state(self, V, Ca: braincell.IonInfo, batch_size: int = None):\n",
    "        self.p = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "        self.q = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))"
   ],
   "id": "622a39869c55b7bb",
   "outputs": [],
   "execution_count": 22
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.297604700Z",
     "start_time": "2025-09-27T12:52:30.903258Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def reset_state(self, V, Ca, batch_size=None):\n",
    "        self.p.value = self.f_p_inf(V)\n",
    "        self.q.value = self.f_q_inf(V)\n",
    "        if batch_size is not None:\n",
    "            assert self.p.value.shape[0] == batch_size\n",
    "            assert self.q.value.shape[0] == batch_size"
   ],
   "id": "49012f069db7fd8b",
   "outputs": [],
   "execution_count": 23
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "For $\\frac{dp}{dt}$ and $\\frac{dq}{dt}$, we have:\n",
    "\n",
    "$$\\frac{dp}{dt} = \\frac{\\phi_p \\cdot (p_{\\infty} - p)}{\\tau_p}$$\n",
    "\n",
    "$$\\frac{dq}{dt} = \\frac{\\phi_q \\cdot (q_{\\infty} - q)}{\\tau_q}$$\n"
   ],
   "id": "2d65831541c31f20"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.297604700Z",
     "start_time": "2025-09-27T12:52:30.927020Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def compute_derivative(self, V, Ca: braincell.IonInfo):\n",
    "        self.p.derivative = self.phi_p * (self.f_p_inf(V) - self.p.value) / self.f_p_tau(V) / u.ms\n",
    "        self.q.derivative = self.phi_q * (self.f_q_inf(V) - self.q.value) / self.f_q_tau(V) / u.ms"
   ],
   "id": "d73e66178714f3ab",
   "outputs": [],
   "execution_count": 24
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "For the overall $I_{CaT}$, we have:\n",
    "\n",
    "$$I_{CaT} = g_{max} \\cdot p^2 \\cdot q \\cdot (V - E_{Ca})$$\n"
   ],
   "id": "50399c876daf0424"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.298109100Z",
     "start_time": "2025-09-27T12:52:30.948303Z"
    }
   },
   "cell_type": "code",
   "source": [
    "    def current(self, V, Ca: braincell.IonInfo):\n",
    "        return self.g_max * self.p.value * self.p.value * self.q.value * (Ca.E - V)"
   ],
   "id": "a7390fbf4beeb630",
   "outputs": [],
   "execution_count": 25
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "At this point, we have completed the modeling of `ICaT_HP1992`.\n",
    "\n",
    "In `braincell`, we usually use `__init__` to define the constructor, and during class initialization, we accept relevant parameters such as `size`, `phi_p`, `phi_q`, etc.\n",
    "\n",
    "Additionally, we should use `braintools.init.param` to register `phi_p`, `phi_q`, and `g_max` as trainable parameters, which enables support for subsequent automatic differentiation, updates, and batch computations.\n",
    "\n",
    "At the same time, when defining interfaces, we need to refer to the actual situation or the specific formulas; for example, the `compute_derivative` interface follows the specific formulas.\n",
    "\n",
    "Through this modeling, the subclass channels gain the following functionalities:\n",
    "\n",
    "* Encapsulates the calcium current kinetics framework in the form of $p^2 q$.\n",
    "* Manages the initialization, updating, and current calculation of the state variables $p$ and $q$.\n",
    "* Allows concrete parameters to be defined by the subclass.\n",
    "* Supports batch computation and parameter differentiability.\n",
    "\n",
    "This modeling framework is highly versatile and flexible, suitable for various requirements encountered in practical modeling.\n"
   ],
   "id": "2449ea1b3ba51156"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Existing Channel Examples\n",
    "\n",
    "After explaining how to use existing channels and create custom channels, we introduce some specific channel examples. Whether using built-in channels or creating new custom channels, these examples provide good references and can guide you when you have questions about channel modeling. The following channel examples show specific electrophysiological properties and modeling approaches.\n",
    "\n",
    "### Calcium Channels\n",
    "\n",
    "#### Electrophysiological Properties\n",
    "\n",
    "Calcium channels are highly important membrane channels responsible for regulating the transmembrane flow of calcium ions. Calcium ions in the nervous system not only participate in membrane potential changes but also play key roles in:\n",
    "\n",
    "* Action potential generation and maintenance: Certain neuron types rely on calcium channels to trigger or sustain firing rhythms.\n",
    "* Synaptic transmission: Calcium influx triggers neurotransmitter release.\n",
    "* Intracellular signaling: Calcium can regulate protein kinase activity, gene expression, enzymatic reactions, and more.\n",
    "* Calcium-dependent channel regulation: Many potassium or mixed channels are also regulated by calcium concentration.\n",
    "\n",
    "Calcium channels typically exhibit slow activation, calcium-dependent inactivation, and high selectivity. They often work together with calcium pumps and buffering proteins to maintain calcium homeostasis.\n",
    "\n",
    "![Calcium Channel](../_static/calciumchannel.png)\n",
    "\n",
    "#### Modeling Implementation\n",
    "\n",
    "In `braincell`, calcium channels are modeled by inheriting from `CalciumChannel`, which itself inherits from `Channel` and is part of the `IonChannel` system.\n",
    "\n",
    "In practice, we need to implement relevant interfaces, such as `current(V, Ca)` and `compute_derivative(V, Ca)`.\n",
    "The `Ca` parameter is of type `IonInfo` and contains local calcium concentration information for the channel. Detailed information about `IonInfo` will be introduced in the `Ion` section.\n",
    "\n",
    "Through these interfaces, we can build calcium channels with complex dynamics, such as dual voltage and calcium regulation, calcium-dependent inactivation, and steady-state activation gating.\n",
    "\n",
    "After understanding the basic modeling, let's look at a practical example: `ICaN_IS2008`.\n",
    "\n",
    "This is a calcium-activated non-selective cation channel model proposed by Inoue & Strowbridge in 2008.\n",
    "`braincell` follows a fixed naming convention for ion channels, in the format `ChannelType_LiteratureID`. Most channels in `braincell` follow this convention, which makes them easy to reference.\n",
    "\n",
    "For the model itself, the current is described by the following formulas:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "I_{\\text{CAN}} &= g_{\\text{max}} \\cdot M([\\mathrm{Ca}^{2+}]_i) \\cdot p \\cdot (V - E) \\\\\\\\\n",
    "M([\\mathrm{Ca}^{2+}]_i) &= \\frac{[\\mathrm{Ca}^{2+}]_i}{0.2 + [\\mathrm{Ca}^{2+}]_i} \\\\\\\\\n",
    "\\frac{dp}{dt} &= \\frac{\\phi \\cdot (p_\\infty - p)}{\\tau_p} \\\\\\\\\n",
    "p_\\infty &= \\frac{1}{1 + \\exp\\left(-\\frac{V + 43}{5.2}\\right)} \\\\\\\\\n",
    "\\tau_p &= \\frac{2.7}{\\exp\\left(-\\frac{V + 55}{15}\\right) + \\exp\\left(\\frac{V + 55}{15}\\right)} + 1.6\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $M$ is the calcium-dependent activation function\n",
    "* $p$ is the voltage-gated activation variable\n",
    "* $\\phi$ is the temperature factor\n",
    "* $E$ is the reversal potential\n",
    "* $g_{\\text{max}}$ is the maximal conductance\n",
    "\n",
    "The following is the code implementation of this model:\n"
   ],
   "id": "95b3df1a05064c7f"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.298109100Z",
     "start_time": "2025-10-13T08:10:44.086622Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class ICaN_IS2008(CalciumChannel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size,\n",
    "        E=10. * u.mV,\n",
    "        g_max=1. * (u.mS / u.cm ** 2),\n",
    "        phi=1.,\n",
    "        name=None,\n",
    "    ):\n",
    "        super().__init__(size=size, name=name)\n",
    "        self.E = braintools.init.param(E, self.varshape, allow_none=False)\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "        self.phi = braintools.init.param(phi, self.varshape, allow_none=False)\n",
    "\n",
    "    def init_state(self, V, Ca: braincell.IonInfo, batch_size: int = None):\n",
    "        self.p = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "\n",
    "    def reset_state(self, V, Ca: braincell.IonInfo, batch_size=None):\n",
    "        V = V.to_decimal(u.mV)\n",
    "        self.p.value = 1.0 / (1 + u.math.exp(-(V + 43.) / 5.2))\n",
    "\n",
    "    def compute_derivative(self, V, Ca: braincell.IonInfo):\n",
    "        V = V.to_decimal(u.mV)\n",
    "        p_inf = 1.0 / (1 + u.math.exp(-(V + 43.) / 5.2))\n",
    "        tau_p = 2.7 / (u.math.exp(-(V + 55.) / 15.) + u.math.exp((V + 55.) / 15.)) + 1.6\n",
    "        self.p.derivative = self.phi * (p_inf - self.p.value) / tau_p / u.ms\n",
    "\n",
    "    def current(self, V, Ca: braincell.IonInfo):\n",
    "        M = Ca.C / (Ca.C + 0.2 * u.mM)\n",
    "        g = self.g_max * M * self.p.value\n",
    "        return g * (self.E - V)"
   ],
   "id": "5691522ec1a87111",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "In this example, we created a specific calcium channel subclass that inherits from `CalciumChannel`. Observing the code, we can see that `CalciumChannel` is very flexible when modeling a specific channel: you only need to set the relevant parameters and express the different interfaces in code according to the formulas.\n",
    "\n",
    "### Potassium Channels\n",
    "\n",
    "#### Electrophysiological Properties\n",
    "\n",
    "Potassium channels are the main hyperpolarizing channels in neurons, determining the repolarization process of action potentials and the stability of the resting membrane potential. Their functions in neurons include:\n",
    "\n",
    "* Action potential repolarization: During an action potential, sodium channels cause depolarization, and potassium channels open to allow K$^+$ efflux, returning the membrane potential to resting levels.\n",
    "* Controlling excitability: Some potassium channels have delayed activation or calcium-dependent properties, regulating firing frequency, adaptation, and sustained excitability.\n",
    "* Modulating network rhythms: Certain potassium channels play important roles in rhythmic firing and oscillations.\n",
    "\n",
    "Potassium channels have wide-ranging functions, and their kinetic modeling is crucial in computational neuroscience.\n",
    "\n",
    "![Potassium Channel](../_static/potassiumchannel.png)\n",
    "\n",
    "#### Modeling Implementation\n",
    "\n",
    "In `braincell`, potassium channels are modeled by inheriting from `PotassiumChannel`, which itself inherits from `Channel`.\n",
    "\n",
    "When implementing a specific potassium channel, similar interfaces to those used for calcium channels are typically required.\n",
    "\n",
    "Let’s look at a specific example: `IKNI_Ya1989`, which helps illustrate the approach.\n",
    "\n",
    "This model was first proposed by Yamada in 1989 and represents a classic slow non-inactivating potassium channel used to explain frequency regulation during the afterhyperpolarization (AHP) period.\n",
    "\n",
    "The model is described by the following formulas:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "I_{M} &= \\bar{g}_M \\cdot p \\cdot (V - E_K) \\\\\\\\\n",
    "\\frac{dp}{dt} &= \\frac{p_{\\infty}(V) - p}{\\tau_p(V)} \\\\\\\\\n",
    "p_{\\infty}(V) &= \\frac{1}{1 + \\exp\\left(-\\frac{V - V_{\\text{sh}} + 35}{10}\\right)} \\\\\\\\\n",
    "\\tau_p(V) &= \\frac{\\tau_{\\max}}{3.3 \\exp\\left(\\frac{V - V_{\\text{sh}} + 35}{20}\\right) + \\exp\\left(-\\frac{V - V_{\\text{sh}} + 35}{20}\\right)}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $V$ is the membrane potential\n",
    "* $E_K$ is the potassium reversal potential\n",
    "* $\\bar{g}_M$ is the maximal conductance\n",
    "* $p$ is the gating variable\n",
    "* $V_{\\text{sh}}$ is the membrane potential shift\n",
    "* $\\tau_{\\max}$ is the maximal time constant\n",
    "* $p_{\\infty}$ is the steady-state activation function\n",
    "* $\\tau_p$ is the activation time constant\n",
    "\n",
    "The following is the code implementation:\n"
   ],
   "id": "3bda188955f4e215"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.298109100Z",
     "start_time": "2025-10-13T08:10:55.931934Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from braincell.channel import PotassiumChannel\n",
    "\n",
    "\n",
    "class IKNI_Ya1989(PotassiumChannel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size,\n",
    "        g_max=0.004 * (u.mS * u.cm ** -2),\n",
    "        phi_p=1.,\n",
    "        phi_q=1.,\n",
    "        tau_max=4e3 * u.ms,\n",
    "        V_sh=0. * u.mV,\n",
    "        name=None,\n",
    "    ):\n",
    "        super().__init__(size=size, name=name)\n",
    "\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "        self.tau_max = braintools.init.param(tau_max, self.varshape, allow_none=False)\n",
    "        self.V_sh = braintools.init.param(V_sh, self.varshape, allow_none=False)\n",
    "        self.phi_p = braintools.init.param(phi_p, self.varshape, allow_none=False)\n",
    "        self.phi_q = braintools.init.param(phi_q, self.varshape, allow_none=False)\n",
    "\n",
    "    def init_state(self, V, Ca: braincell.IonInfo, batch_size: int = None):\n",
    "        self.p = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "\n",
    "    def reset_state(self, V, K: braincell.IonInfo, batch_size=None):\n",
    "        self.p.value = self.f_p_inf(V)\n",
    "        if isinstance(batch_size, int):\n",
    "            assert self.p.value.shape[0] == batch_size\n",
    "\n",
    "    def compute_derivative(self, V, K: braincell.IonInfo):\n",
    "        self.p.derivative = self.phi_p * (self.f_p_inf(V) - self.p.value) / self.f_p_tau(V)\n",
    "\n",
    "    def current(self, V, K: braincell.IonInfo):\n",
    "        return self.g_max * self.p.value * (K.E - V)\n",
    "\n",
    "    def f_p_inf(self, V):\n",
    "        V = (V - self.V_sh).to_decimal(u.mV)\n",
    "        return 1. / (1. + u.math.exp(-(V + 35.) / 10.))\n",
    "\n",
    "    def f_p_tau(self, V):\n",
    "        V = (V - self.V_sh).to_decimal(u.mV)\n",
    "        temp = V + 35.\n",
    "        return self.tau_max / (3.3 * u.math.exp(temp / 20.) + u.math.exp(-temp / 20.))"
   ],
   "id": "b602c5802ded2d16",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "From this example, we can see that `PotassiumChannel` provides a unified interface, and `IKNI_Ya1989` can easily construct a biologically meaningful ion channel model using simple mathematical expressions.\n",
    "\n",
    "In practical applications, we can further build subclasses under a subclass.\n",
    "\n",
    "For example, first construct `IKK2_pq_ss`, which inherits from `PotassiumChannel`. Later, we can construct `KK2A_HM1992` inheriting from `IKK2_pq_ss` and extend the original functionality of `IKK2_pq_ss` in `KK2A_HM1992`. This workflow facilitates large-scale modeling and improves efficiency. If you are interested in this part, you can refer to our [Ion Channel Collection](https://braincell.readthedocs.io/latest/apis/braincell.channel.html) to see how `braincell` implements it in detail.\n",
    "\n",
    "### Sodium Channels\n",
    "\n",
    "#### Electrophysiological Properties\n",
    "\n",
    "Sodium channels are indispensable excitatory channels in neurons, playing a central role in the initiation and rising phase of action potentials. When the membrane potential reaches a certain threshold, sodium channels rapidly activate, causing Na$^+$ influx and rapid depolarization of the membrane.\n",
    "\n",
    "Sodium channels have a typical dual-gate mechanism:\n",
    "\n",
    "* Activation gate: responds quickly to membrane potential changes, determining whether the channel opens.\n",
    "* Inactivation gate: responds slightly slower, closing the channel and causing inactivation.\n",
    "\n",
    "This dual-gate structure allows sodium channels to exhibit fast activation and fast inactivation dynamics, which is the fundamental mechanism for the rapid rise and fall of action potentials.\n",
    "\n",
    "![Sodium Channel](../_static/sodiumchannel.png)\n",
    "\n",
    "Sodium channels are highly sensitive to membrane potential changes, and their kinetics determine the regulation of neuronal excitability. In the classical HH model, the sodium current is typically represented as:\n",
    "\n",
    "$$\n",
    "I_{Na} = \\bar{g}_{Na} \\cdot m^3 h \\cdot (V - E_{Na})\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $m$ is the activation variable\n",
    "* $h$ is the inactivation variable\n",
    "* $g_{Na}$ is the maximal conductance\n",
    "* $E_{Na}$ is the sodium reversal potential\n",
    "\n",
    "#### Modeling Implementation\n",
    "\n",
    "In `braincell`, sodium channels are modeled by inheriting the `SodiumChannel` class, which, like the others, is a subclass of `Channel`.\n",
    "\n",
    "As with calcium and potassium channels, we only need to define the corresponding interfaces to complete the modeling.\n",
    "\n",
    "Here is a specific example: `INa_p3q_markov`.\n",
    "\n",
    "Notably, a different channel naming convention is used here. `INa_p3q_markov` is named based on the formula structure, representing a current in the form of $p^3 q$ and modeled using Markov kinetics. This is another naming convention in `braincell`: `channel type_name_current form_model`. Some ion channels adopt this convention.\n",
    "\n",
    "`INa_p3q_markov` is a sodium channel model using Markov chain kinetics, where the current is controlled by activation variable $p$ and inactivation variable $q$:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "I_{\\mathrm{Na}} &= g_{\\mathrm{max}} \\cdot p^3 \\cdot q \\cdot (V - E_{Na}) \\\\\\\\\n",
    "\\frac{dp}{dt} &= \\phi \\cdot (\\alpha_p (1 - p) - \\beta_p p) \\\\\\\\\n",
    "\\frac{dq}{dt} &= \\phi \\cdot (\\alpha_q (1 - q) - \\beta_q q)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $p$ is the activation variable\n",
    "* $q$ is the inactivation variable\n",
    "* $\\phi$ is the temperature factor\n",
    "* $g_{max}$ is the maximal conductance\n",
    "* $E_{Na}$ is the sodium reversal potential\n",
    "\n",
    "The corresponding code implementation is as follows:\n"
   ],
   "id": "aa156264b0f65567"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.299114Z",
     "start_time": "2025-09-27T12:52:31.023081Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from braincell.channel import SodiumChannel\n",
    "\n",
    "\n",
    "class INa_p3q_markov(SodiumChannel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size,\n",
    "        g_max=90. * (u.mS / u.cm ** 2),\n",
    "        phi=1.,\n",
    "        name=None,\n",
    "    ):\n",
    "        super().__init__(size=size, name=name)\n",
    "        self.phi = braintools.init.param(phi, self.varshape, allow_none=False)\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "\n",
    "    def init_state(self, V, Na: braincell.IonInfo, batch_size=None):\n",
    "        self.p = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "        self.q = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "\n",
    "    def reset_state(self, V, Na: braincell.IonInfo, batch_size=None):\n",
    "        self.p.value = self.f_p_alpha(V) / (self.f_p_alpha(V) + self.f_p_beta(V))\n",
    "        self.q.value = self.f_q_alpha(V) / (self.f_q_alpha(V) + self.f_q_beta(V))\n",
    "\n",
    "    def compute_derivative(self, V, Na: braincell.IonInfo):\n",
    "        self.p.derivative = self.phi * (\n",
    "            self.f_p_alpha(V) * (1. - self.p.value) - self.f_p_beta(V) * self.p.value) / u.ms\n",
    "        self.q.derivative = self.phi * (\n",
    "            self.f_q_alpha(V) * (1. - self.q.value) - self.f_q_beta(V) * self.q.value) / u.ms\n",
    "\n",
    "    def current(self, V, Na: braincell.IonInfo):\n",
    "        return self.g_max * self.p.value ** 3 * self.q.value * (Na.E - V)\n",
    "\n",
    "    def f_p_alpha(self, V):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def f_p_beta(self, V):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def f_q_alpha(self, V):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def f_q_beta(self, V):\n",
    "        raise NotImplementedError"
   ],
   "id": "33a49566b6afa3e0",
   "outputs": [],
   "execution_count": 28
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Like the other two types of ion channels, our sodium channel modeling retains high flexibility.\n",
    "\n",
    "### Potassium Calcium Channels\n",
    "\n",
    "#### Electrophysiological Properties\n",
    "\n",
    "Potassium calcium (KCa) channels are widely present in neurons, and their activation depends not only on membrane potential but also on intracellular calcium concentration.\n",
    "\n",
    "These channels play a crucial role in regulating neuronal excitability, afterhyperpolarization, and firing frequency. They are typically activated following calcium influx, promoting K$^+$ efflux and returning the membrane potential to rest or hyperpolarized states.\n",
    "\n",
    "#### Modeling Implementation\n",
    "\n",
    "Potassium calcium channels are commonly abbreviated as `KCa` channels in electrophysiology, and we adopt the same naming in modeling for code simplicity.\n",
    "\n",
    "In `braincell`, KCa channels are modeled by directly inheriting the `KCaChannel` class, which is also a subclass of `Channel`.\n",
    "\n",
    "As with the previously described channels, we only need to define the corresponding interfaces to complete the modeling.\n",
    "\n",
    "Here we give an example: `IAHP_De1994`.\n",
    "\n",
    "The `IAHP_De1994` model, proposed by Destexhe et al. in 1994, inherits from `KCaChannel` and is used to simulate the dynamics of slow calcium-dependent potassium channels. The expressions are as follows:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "(\\text{closed}) + n \\mathrm{Ca}_{i}^{2+} \\underset{\\beta}{\\stackrel{\\alpha}{\\rightleftharpoons}} (\\text{open})\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "I_{AHP} &= g_{\\mathrm{max}} p^2 (V - E_K) \\\\\\\\\n",
    "{dp \\over dt} &= \\phi {p_{\\infty}(V, [Ca^{2+}]_i) - p \\over \\tau_p(V, [Ca^{2+}]_i)} \\\\\\\\\n",
    "p_{\\infty} &= \\frac{\\alpha[Ca^{2+}]_i^n}{\\alpha[Ca^{2+}]_i^n + \\beta} \\\\\\\\\n",
    "\\tau_p &= \\frac{1}{\\alpha[Ca^{2+}]_i + \\beta}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $p$ is the channel state variable\n",
    "* $\\alpha, \\beta$ are the channel opening/closing rate constants\n",
    "* $n$ is the number of calcium ions bound\n",
    "* $\\phi$ is the temperature factor\n",
    "* $g_{\\text{max}}$ is the maximal conductance\n",
    "* $E_K$ is the potassium reversal potential\n",
    "\n",
    "The code implementation is as follows:\n"
   ],
   "id": "fdec02b1a538d361"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.301113600Z",
     "start_time": "2025-09-27T12:52:31.046967Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from braincell.channel import KCaChannel\n",
    "\n",
    "\n",
    "class IAHP_De1994(KCaChannel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size: brainstate.typing.Size,\n",
    "        n=2,\n",
    "        g_max=10. * (u.mS / u.cm ** 2),\n",
    "        alpha=48.,\n",
    "        beta=0.09,\n",
    "        phi=1.,\n",
    "        name=None,\n",
    "    ):\n",
    "        super().__init__(size=size, name=name)\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "        self.n = braintools.init.param(n, self.varshape, allow_none=False)\n",
    "        self.alpha = braintools.init.param(alpha, self.varshape, allow_none=False)\n",
    "        self.beta = braintools.init.param(beta, self.varshape, allow_none=False)\n",
    "        self.phi = braintools.init.param(phi, self.varshape, allow_none=False)\n",
    "\n",
    "    def init_state(self, V, K: braincell.IonInfo, Ca: braincell.IonInfo, batch_size=None):\n",
    "        self.p = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "\n",
    "    def reset_state(self, V, K: braincell.IonInfo, Ca: braincell.IonInfo, batch_size=None):\n",
    "        C2 = self.alpha * u.math.power(Ca.C / u.mM, self.n)\n",
    "        C3 = C2 + self.beta\n",
    "        if batch_size is None:\n",
    "            self.p.value = u.math.broadcast_to(C2 / C3, self.varshape)\n",
    "        else:\n",
    "            self.p.value = u.math.broadcast_to(C2 / C3, (batch_size,) + self.varshape)\n",
    "            assert self.p.value.shape[0] == batch_size\n",
    "\n",
    "    def compute_derivative(self, V, K: braincell.IonInfo, Ca: braincell.IonInfo):\n",
    "        C2 = self.alpha * u.math.power(Ca.C / u.mM, self.n)\n",
    "        C3 = C2 + self.beta\n",
    "        self.p.derivative = self.phi * (C2 / C3 - self.p.value) * C3 / u.ms\n",
    "\n",
    "    def current(self, V, K: braincell.IonInfo, Ca: braincell.IonInfo):\n",
    "        return self.g_max * self.p.value * self.p.value * (K.E - V)"
   ],
   "id": "d25373177d600967",
   "outputs": [],
   "execution_count": 29
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "In fact, the modeling approach for these examples is highly consistent, with no differences in implementation.\n",
    "\n",
    "### Hyperpolarization-Activated Channels\n",
    "\n",
    "#### Electrophysiological Properties\n",
    "\n",
    "Hyperpolarization-activated channels are cation channels activated when the membrane potential hyperpolarizes. They play key regulatory roles in pacemaker neurons of the heart and certain types of central nervous system neurons:\n",
    "\n",
    "* Regulate resting membrane potential: depolarizing currents at rest stabilize the membrane potential and prevent excessive hyperpolarization.\n",
    "* Participate in rhythmic firing: in structures like the thalamus, they help maintain intrinsic rhythmic firing.\n",
    "* Excitability regulation: they can reduce neuronal input resistance, decreasing the cell's sensitivity to inputs.\n",
    "\n",
    "#### Modeling Implementation\n",
    "\n",
    "Since `braincell` does not have many built-in hyperpolarization-activated channels, we simply let specific hyperpolarization-activated channels inherit from the `Channel` class. This does not affect practical use.\n",
    "\n",
    "Here is an example: `Ih_HM1992`.\n",
    "\n",
    "`Ih_HM1992` was proposed by Huguenard & McCormick in 1992 as a hyperpolarization-activated cation current model. Its dynamics are described by:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "I_h &= g_{\\mathrm{max}} p \\\\\\\\\n",
    "\\frac{dp}{dt} &= \\phi \\frac{p_{\\infty} - p}{\\tau_p} \\\\\\\\\n",
    "p_{\\infty} &= \\frac{1}{1+\\exp((V+75)/5.5)} \\\\\\\\\n",
    "\\tau_p &= \\frac{1}{\\exp(-0.086 V - 14.59) + \\exp(0.0701 V - 1.87)}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $p$ is the activation variable\n",
    "* $g_{\\text{max}}$ is the maximal conductance\n",
    "* $E$ is the reversal potential (about 43 mV)\n",
    "* $\\phi$ is the temperature factor\n",
    "\n",
    "The code implementation is as follows:\n"
   ],
   "id": "6b93cf68c256d745"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.303114800Z",
     "start_time": "2025-09-27T12:52:31.074072Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from braincell import Channel\n",
    "\n",
    "\n",
    "class Ih_HM1992(Channel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size,\n",
    "        g_max=10. * (u.mS / u.cm ** 2),\n",
    "        E=43. * u.mV,\n",
    "        phi=1.,\n",
    "        name=None,\n",
    "    ):\n",
    "        super().__init__(size=size, name=name)\n",
    "        self.phi = braintools.init.param(phi, self.varshape, allow_none=False)\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "        self.E = braintools.init.param(E, self.varshape, allow_none=False)\n",
    "\n",
    "    def init_state(self, V, batch_size=None):\n",
    "        self.p = braincell.DiffEqState(braintools.init.param(u.math.zeros, self.varshape, batch_size))\n",
    "\n",
    "    def reset_state(self, V, batch_size=None):\n",
    "        self.p.value = self.f_p_inf(V)\n",
    "\n",
    "    def compute_derivative(self, V):\n",
    "        self.p.derivative = self.phi * (self.f_p_inf(V) - self.p.value) / self.f_p_tau(V) / u.ms\n",
    "\n",
    "    def current(self, V):\n",
    "        return self.g_max * self.p.value * (self.E - V)\n",
    "\n",
    "    def f_p_inf(self, V):\n",
    "        V = V.to_decimal(u.mV)\n",
    "        return 1. / (1. + u.math.exp((V + 75.) / 5.5))\n",
    "\n",
    "    def f_p_tau(self, V):\n",
    "        V = V.to_decimal(u.mV)\n",
    "        return 1. / (u.math.exp(-0.086 * V - 14.59) + u.math.exp(0.0701 * V - 1.87))"
   ],
   "id": "c892325aa082afdb",
   "outputs": [],
   "execution_count": 30
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Through this example, you can see that modeling hyperpolarization-activated channels can be naturally integrated into the `braincell` framework. Using the `Channel` interface, their unique dynamics can be easily expressed.\n",
    "\n",
    "### Leakage Channels\n",
    "\n",
    "#### Electrophysiological Properties\n",
    "\n",
    "Leakage channels are a class of always-open ion channels whose permeability is not regulated by membrane potential, voltage gating, ligand binding, or other mechanisms. Electrophysiologically, leakage channels mainly maintain the resting membrane potential of neurons and are typically modeled as a linear conductance.\n",
    "\n",
    "In the classic HH model, leakage currents primarily simulate background ionic flows that are not explicitly modeled, providing a stable baseline current.\n",
    "\n",
    "#### Modeling Implementation\n",
    "\n",
    "In `braincell`, leakage channels are implemented by inheriting from `LeakageChannel`. Since leakage currents have no activation or inactivation gating variables, this type of channel does not need to implement differential processes like `compute_derivative` or `init_state`, making it relatively simple.\n",
    "\n",
    "Here is a concrete example: `IL`, the most basic linear leakage channel, with the mathematical expression:\n",
    "\n",
    "$$\n",
    "I_L = g_L (E_L - V)\n",
    "$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $g_L$ is the leakage conductance (usually constant)\n",
    "* $E_L$ is the leakage reversal potential\n",
    "* $V$ is the membrane potential\n",
    "\n",
    "The implementation code is also quite straightforward:\n"
   ],
   "id": "8621f5c9e8e66d93"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:12:57.306113900Z",
     "start_time": "2025-09-27T12:52:31.100136Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from braincell.channel import LeakageChannel\n",
    "\n",
    "\n",
    "class IL(LeakageChannel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        size,\n",
    "        g_max=0.1 * (u.mS / u.cm ** 2),\n",
    "        E=-70. * u.mV,\n",
    "        name=None,\n",
    "    ):\n",
    "        super().__init__(size=size, name=name)\n",
    "        self.E = braintools.init.param(E, self.varshape, allow_none=False)\n",
    "        self.g_max = braintools.init.param(g_max, self.varshape, allow_none=False)\n",
    "\n",
    "    def current(self, V):\n",
    "        return self.g_max * (self.E - V)"
   ],
   "id": "af8b9c1a7ff1172a",
   "outputs": [],
   "execution_count": 31
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "This `IL` class provides the simplest channel model. It does not involve any internal states or dynamics and represents a static conductance model, often used as a fundamental component in neuron models.\n",
   "id": "e7641dc75a624a03"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
