{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 通道\n",
    "\n",
    "在生物神经元中，离子通道控制着钠、钾、钙等离子的通透性，是神经信号发生与传递的基础。不同类型的通道具有各自的动力学特性，并对膜电位的变化产生影响。\n",
    "\n",
    "`braincell` 的通道建模具备以下几个显著优势：\n",
    "- 接口统一：无论是线性通道还是复杂 Markov 模型，均可统一表示和集成。\n",
    "- 灵活性强：支持多种电流类型，同时支持用户自定义建模。\n",
    "- 兼容 `brainstate` 的自动微分系统：每个状态变量都能被精确地自动求导，支持高效训练与优化。\n",
    "- 可拓展性高：通过继承与组合机制，通道模型可以轻松构建子类，实现参数修改与拓展。\n",
    "\n",
    "`braincell` 具有高度的灵活性。对于 `Channel` ，`braincell` 支持两种用法，一种是自定义新的离子通道，另一种是使用现有的通道。\n",
    "\n",
    "## 现有通道调用\n",
    "\n",
    "我们先来讲解如何使用现有通道。\n",
    "\n",
    "我们根据[通道功能的不同](https://elifesciences.org/articles/22152)，将 `Channel` 分为以下几个基类：\n",
    "- `Calcium Channels`：钙离子通道\n",
    "- `Potassium Channels`：钾离子通道\n",
    "- `Sodium Channels`：钠离子通道\n",
    "- `Potassium Calcium Channels`：钙依赖性钾通道\n",
    "- `Hyperpolarization-Activated Channels`：超极化激活通道\n",
    "- `Leakage Channels`：漏电通道\n",
    "\n",
    "`braincell` 已经在每种基类中内置好了许多的具体通道，你可以任意使用。\n",
    "如果想了解更多的内置的通道，你可以通过查阅 [离子通道合集](../apis/braincell.channel.rst) 来学习使用。\n",
    "\n",
    "当在实际对神经元进行建模时，不管是现有的通道，还是自定义的通道，我们都需要进行调用。\n",
    "当然，这个调用是相当简单的。但我们同样要注意一些调用的规范。\n",
    "\n",
    "下面我们来看对 `HTC` 神经元建模的例子："
   ],
   "id": "653784fc1e9d2719"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.080971Z",
     "start_time": "2025-10-13T09:20:17.078533Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import brainstate\n",
    "import braintools\n",
    "import brainunit as u\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import braincell"
   ],
   "id": "ce159280e462071f",
   "outputs": [],
   "execution_count": 26
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.124472Z",
     "start_time": "2025-10-13T09:20:17.117206Z"
    }
   },
   "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": "1e0b3a9c1b6a2ea9",
   "outputs": [],
   "execution_count": 27
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "通过这个 `HTC` 神经元的例子可以很清晰的看出，在实际建模中，我们只需简单将相关的通道传入即可。\n",
    "\n",
    "但值得注意的是，本例中的通道都已经显式定义在 `braincell` 的代码仓库中。\n",
    "若想使用其他通道，你可以将你的自定义通道定义于`braincell` 代码仓库中，或者使用 `import` 将你自定义的通道引入。"
   ],
   "id": "bdbca8e414929ee7"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### `root_type` 的使用\n",
    "\n",
    "使用每一个通道时，我们都需要注意 `root_type` 的使用。一般来说，钙离子通道依赖于 `Calcium` 离子，钾离子通道依赖于 `Potassium` 离子，钠离子通道依赖于 `Sodium` 离子。"
   ],
   "id": "5cefbb5ddaf4e786"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.152784Z",
     "start_time": "2025-10-13T09:20:17.149717Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.CalciumChannel.root_type",
   "id": "3e93c0291342afa3",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.ion.Calcium"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 28
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.180660Z",
     "start_time": "2025-10-13T09:20:17.178086Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.PotassiumChannel.root_type",
   "id": "1e425336d82ff882",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.ion.Potassium"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.228617Z",
     "start_time": "2025-10-13T09:20:17.225659Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.SodiumChannel.root_type",
   "id": "5dd87c56c9674cd6",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.ion.Sodium"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 30
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "钙依赖性钾通道同时依赖于 `Potassium` 和 `Calcium` 离子。",
   "id": "4a2fea55c94a0282"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.265065Z",
     "start_time": "2025-10-13T09:20:17.261308Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.IAHP_De1994.root_type",
   "id": "b5289c40fe190854",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JointTypes[braincell.ion.Potassium, braincell.ion.Calcium]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 31
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "对于一些特殊的通道，如 `IL`，它们并不依赖于任何离子类型，因此其 `root_type` 为细胞。",
   "id": "76e39218a5724028"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.322346Z",
     "start_time": "2025-10-13T09:20:17.319249Z"
    }
   },
   "cell_type": "code",
   "source": "braincell.channel.IL.root_type",
   "id": "a17d1a810a96abea",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "braincell.HHTypedNeuron"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 32
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 添加离子通道\n",
    "\n",
    "在进行单细胞建模时，我们通常需要将离子通道添加到离子中。如果不注意 `root_type` 的使用，可能会导致离子通道无法正确工作。比如，钠离子通道只能添加到钠离子中，钾离子通道只能添加到钾离子中。如果你尝试将钠离子通道添加到钾离子中，或者将钾离子通道添加到钙离子中，都会导致错误。\n"
   ],
   "id": "d9d5e6d62956fe1"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.400606Z",
     "start_time": "2025-10-13T09:20:17.397987Z"
    }
   },
   "cell_type": "code",
   "source": "na = braincell.ion.SodiumFixed(1)",
   "id": "298ceeb5d51eea52",
   "outputs": [],
   "execution_count": 33
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.416369Z",
     "start_time": "2025-10-13T09:20:17.414291Z"
    }
   },
   "cell_type": "code",
   "source": "na.add(ina=braincell.channel.INa_HH1952(1))",
   "id": "81ac7c235727d928",
   "outputs": [],
   "execution_count": 34
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.468436Z",
     "start_time": "2025-10-13T09:20:17.461954Z"
    }
   },
   "cell_type": "code",
   "source": [
    "try:\n",
    "    na.add(ik=braincell.channel.IK_HH1952(1))\n",
    "except Exception as e:\n",
    "    print(e)"
   ],
   "id": "812d92f680b857a5",
   "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": 35
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 模拟离子通道\n",
    "\n",
    "我们也可以直接模拟离子通道，而不需要将其添加到离子中。这样可以更灵活地使用离子通道。然而，模拟离子通道时，我们需要添加我们需要的神经元膜电位和离子浓度等信息。我们以钙离子通道为例。"
   ],
   "id": "3d68483796a1a7cc"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.524783Z",
     "start_time": "2025-10-13T09:20:17.511793Z"
    }
   },
   "cell_type": "code",
   "source": [
    "ca = braincell.ion.CalciumFixed(1).pack_info()\n",
    "V = -65 * u.mV\n",
    "\n",
    "# 要进行模拟的钙离子通道\n",
    "can = braincell.channel.ICaL_IS2008(1)\n",
    "can.init_state(V, ca)"
   ],
   "id": "bc68edf83e1174b4",
   "outputs": [],
   "execution_count": 36
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.801903Z",
     "start_time": "2025-10-13T09:20:17.540601Z"
    }
   },
   "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": "24f2a5610f3a8300",
   "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[37]\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[37]\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": 37
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "\n",
    "## 自定义通道建模\n",
    "\n",
    "现在我们来讲解如何使用自定义通道。\n",
    "\n",
    "在 `braincell` 中，我们使用 `IonChannel` 类作为通道建模的基础类。所有具体的通道模型都应继承自该类。\n",
    "\n",
    "`IonChannel` 提供了以下接口，用于支持仿真过程中的关键操作：\n",
    "- `current`：计算通道产生的离子电流。\n",
    "- `init_state`：初始化通道的状态变量。\n",
    "- `reset_state`：重置通道状态。\n",
    "- `compute_derivative`：计算通道内部状态变量的导数，用于积分器的积分。\n",
    "- `pre_integral` 和 `post_integral`：积分前/后的操作，一般不使用。\n",
    "\n",
    "`compute_derivative`, `pre_integral` 和 `post_integral`是积分的核心接口，通常需要在具体通道中实现，详细介绍可以参考 [微分系统](../advanced_tutorial/differential_equation-zh.ipynb) 部分。这些方法都能与 `braincell` 的微分系统集成。在实际自定义离子通道的时候，我们需要设置好通道需要的相关接口，以完成特定的功能需求。\n",
    "\n",
    "`IonChannel` 包含了 `Ion` 和`Channel`。因此，我们定义了一个直接继承自 `IonChannel` 的基类 `Channel` 来进行封装。\n",
    "当然，在语义上，我们更建议你用 `Channel` 来开发实际通道模型。\n",
    "\n",
    "在完成了对 `Channel` 这个基类的建模之后，是时候对不同的具体通道进行建模了。\n",
    "从实际应用上说，我们希望 `Channel` 的子类应该实现特定的方法，以定义特定通道类型的行为。\n",
    "\n",
    "下面我们正式来讲解对于不同的目标建模通道如何完成自定义。\n",
    "\n",
    "### 基类通道建模\n",
    "\n",
    "首先，如果要对一个全新的通道进行建模，我们要先判断目标通道是否可被分类至目前已有的基类通道中。\n",
    "\n",
    "前面讲到， `braincell` 将 `Channel` 分为以下几类：\n",
    "- `Calcium Channels`：钙离子通道\n",
    "- `Potassium Channels`：钾离子通道\n",
    "- `Sodium Channels`：钠离子通道\n",
    "- `Potassium Calcium Channels`：钙依赖性钾通道\n",
    "- `Hyperpolarization-Activated Channels`：超极化激活通道\n",
    "- `Leakage Channels`：漏电通道\n",
    "\n",
    "以上这些通道都是直接继承自 `Channel` 的基类。如果你想自定义的通道属于这些类别，请直接查阅后面的子类通道建模教学与具体通道建模教学。\n",
    "如果你想自定义的通道并不属于我们为你提供的这些基类，我们将以对 `Calcium Channels` 的建模为例，讲解如何对一个离子通道基类进行建模。\n",
    "\n",
    "在 `braincell.channel` 模块中，所有离子通道都最终继承自统一的抽象基类 `Channel`。\n",
    "该类定义了所有通道必须实现的最小接口，为通用建模提供统一框架。\n",
    "\n",
    "在实际建模中，我们通常需要为不同离子类型定义各自的基类。\n",
    "这些类型专属的中间层基类可以封装通道类型、输入参数、接口结构，提供更清晰、更易维护的建模层级。\n",
    "\n",
    "我们来看下面这个 `Calcium Channels` 建模的例子："
   ],
   "id": "3941d3cec2d6a9ec"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.808868800Z",
     "start_time": "2025-10-13T09:20:12.674747Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class CalciumChannel(braincell.Channel):\n",
    "    # 指定该通道作用的离子类型\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": "d866e8e3bf93c298",
   "outputs": [],
   "execution_count": 25
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "以上代码就轻松完成了对 `CalciumChannel` 的建模。\n",
    "\n",
    "观察代码，其中最重要的就是定义 `root_type = Calcium` 。\n",
    "这是 `CalciumChannel` 的核心，它告诉 `braincell` ：\n",
    "这个通道只作用于 `Calcium` ，同时传入的参数 `Ca` 一定是 `IonInfo` 类型，用于表示钙离子的状态。\n",
    "这使得你在实现 `current`, `compute_derivative` 等接口时可以获取到该离子的信息。\n",
    "\n",
    "一旦你定义好了 `CalciumChannel`，后续所有具体的钙通道子类都可以继承它：\n",
    "此时你不再需要关心 `IonInfo` 的来源，`braincell` 会自动把 `Calcium` 类型的离子传给你。\n",
    "\n",
    "同时，在对新的基类通道建模的时候，要注意对接口的定义，这些接口可以参照 `Channel` 中定义的接口，也可以根据实际需求去进行自定义。\n",
    "\n",
    "通过构造 `CalciumChannel` 这样的类型专属通道基类，我们可以：\n",
    "- 明确模型适用的离子类型。\n",
    "- 统一接口。\n",
    "- 简化下游模型编写。\n",
    "- 提高代码复用性与可维护性。\n",
    "\n",
    "这种设计也适用于钠通道 `SodiumChannel` 、钾通道 `PotassiumChannel` 等其他离子通道，推荐作为编写自定义通道模型的标准做法。\n",
    "\n",
    "### 具体通道建模\n",
    "\n",
    "以下我们以 `ICaT_HP1992` 的代码为示例，来展示如何自定义离子通道模型。\n",
    "\n",
    " `ICaT_HP1992` 离子通道是一个钙通道模型，参考了 Huguenard & Prince 在 1992 年提出的钙通道模型：[<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",
    "对其建模只需要参考具体公式，并将公式用代码表达即可："
   ],
   "id": "bbb544a817b558cc"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "我们先定义参数：\n",
    "\n",
    "- $p$ 和 $q$ 是两个激活和失活门控变量\n",
    "- $\\phi_p = 5^{\\frac{T-24}{10}}$ 和 $\\phi_q = 3^{\\frac{T-24}{10}}$ 是温度依赖因子（$T$ 为摄氏度温度）\n",
    "- $E_{Ca}$ 是钙离子通道的反转电位\n",
    "- $p_{\\infty}$, $q_{\\infty}$ 是 $p, q$ 的稳态激活/失活函数（电压依赖）\n",
    "- $\\tau_p, \\tau_q$ 是电压依赖的时间常数\n",
    "- $\\phi_p, \\phi_q$ 是温度修正因子\n",
    "- $g_{max}$ 是最大电导\n",
    "- $V$ 是膜电位，$E_{Ca}$ 是钙离子反转电位\n",
    "- $V_{sh}$ 是膜电位偏移量\n"
   ],
   "id": "534c416b9b773bb7"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:19.350982Z",
     "start_time": "2025-10-13T09:20:19.319454Z"
    }
   },
   "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": "c8466114dd3ccb69",
   "outputs": [],
   "execution_count": 38
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "然后对$p_{\\infty}$有\n",
    "\n",
    "$$\n",
    "p_{\\infty} = \\frac{1}{1 + \\exp[-(V + 52 - V_{sh}) / 7.4]}\n",
    "$$"
   ],
   "id": "bfc656950c18084c"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:22.680670Z",
     "start_time": "2025-10-13T09:20:22.677144Z"
    }
   },
   "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": "db82fc4bc2e44adb",
   "outputs": [],
   "execution_count": 39
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "对$\\tau_p$有：\n",
    "\n",
    "$$\\tau_p = 3 + \\frac{1}{\\exp[(V + 27 - V_{sh}) / 10] + \\exp[-(V + 102 - V_{sh}) / 15]}$$"
   ],
   "id": "1a736edd88135b43"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:24.464419Z",
     "start_time": "2025-10-13T09:20:24.455949Z"
    }
   },
   "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": "ba6f03dcc637ce91",
   "outputs": [],
   "execution_count": 40
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "对$q_{\\infty}$有：\n",
    "\n",
    "$$q_{\\infty} = \\frac{1}{1 + \\exp[(V + 80 - V_{sh}) / 5]}$$"
   ],
   "id": "10a0eb821c59e44c"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:26.641526Z",
     "start_time": "2025-10-13T09:20:26.638980Z"
    }
   },
   "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": "3112590f9df916dd",
   "outputs": [],
   "execution_count": 41
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "对$\\tau_q$则有：\n",
    "\n",
    "$$\\tau_q = 85 + \\frac{1}{\\exp[(V + 48 - V_{sh}) / 4] + \\exp[-(V + 407 - V_{sh}) / 50]}$$"
   ],
   "id": "67a8adbde77955aa"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:28.345986Z",
     "start_time": "2025-10-13T09:20:28.339775Z"
    }
   },
   "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": "95d9f6dc5cb9fa3e",
   "outputs": [],
   "execution_count": 42
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "同时要对 `state` 进行设定：",
   "id": "97ccf51c3385fafa"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:29.469276Z",
     "start_time": "2025-10-13T09:20:29.466437Z"
    }
   },
   "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": "5f579e2f86f9942c",
   "outputs": [],
   "execution_count": 43
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:31.071733Z",
     "start_time": "2025-10-13T09:20:31.067498Z"
    }
   },
   "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": "7045afaafb1cd29",
   "outputs": [],
   "execution_count": 44
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "对于$\\frac{dp}{dt}$和$\\frac{dq}{dt}$则有：\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}$$"
   ],
   "id": "70f927e115cdfe07"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:32.749135Z",
     "start_time": "2025-10-13T09:20:32.729433Z"
    }
   },
   "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": "cf9537cba9000bcb",
   "outputs": [],
   "execution_count": 45
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "最后，对于整体的$I_{CaT}$有：\n",
    "\n",
    "$$I_{CaT} = g_{max} \\cdot p^2 \\cdot q \\cdot (V - E_{Ca})$$"
   ],
   "id": "f1c7feb27763f147"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:34.036259Z",
     "start_time": "2025-10-13T09:20:34.033683Z"
    }
   },
   "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": "ef8ff574f73c02e1",
   "outputs": [],
   "execution_count": 46
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "至此，我们就完成了对 `ICaT_HP1992` 的建模。\n",
    "\n",
    "在 `braincell` 中，我们通常用 `__init__` 来构造函数，并且在初始化类时接受相关参数，如 `size` 、 `phi_p`、`phi_q` 等。\n",
    "\n",
    "并且，我们要注意使用 `braintools.init.param` 把 `phi_p`、 `phi_q`、 `g_max` 注册成可变参数，这样才能支持后续的自动微分、更新和批量计算。\n",
    "\n",
    "同时，设定接口时要根据实际情况或具体参考公式，如接口 `compute_derivative` 就参考了具体公式：\n",
    "\n",
    "通过以上建模，可以使子类通道具有以下功能：\n",
    "- 封装了 $p^2 q$ 形式的钙电流动力学框架。\n",
    "- 管理状态变量 $p, q$ 的初始化、更新、计算电流。\n",
    "- 保留具体的参数由子类定义。\n",
    "- 支持批量计算和参数可微分。\n",
    "\n",
    "这种建模框架适用于在实际建模时遇到的不同需求，泛用性强，灵活性高。\n",
    "\n"
   ],
   "id": "191d5e90a59274e2"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "\n",
    "## 现有通道实例\n",
    "\n",
    "在讲解完如何调用现有通道和自定义通道之后，我们来介绍一些具体的通道实例。无论是使用现有的内置通道，还是建立新的自定义通道，这些实例都是很好的参考，能在你对通道建模有疑问的时候提供参考。以下的通道实例提供了具体的通道电生理特性与建模思路。\n",
    "\n",
    "### Calcium Channels\n",
    "\n",
    "#### 电生理特性\n",
    "\n",
    "钙离子通道是一类高度重要的膜通道，负责调节钙离子在细胞膜上的跨膜流动。钙离子在神经系统中不仅参与膜电位的变化，还在以下过程中扮演着关键角色：\n",
    "- 动作电位的产生与延续：某些神经元类型依赖钙通道触发或维持放电节律。\n",
    "- 突触传递：钙离子内流触发神经递质释放。\n",
    "- 细胞内信号转导：钙离子可调控蛋白激酶活性、基因表达、酶促反应等。\n",
    "- 钙依赖通道调控：许多钾通道或混合通道也受到钙浓度的调节。\n",
    "\n",
    "钙通道通常具有缓慢激活、钙依赖失活、以及高选择性等特点。它们常常与钙泵和缓冲蛋白共同作用，以维持钙稳态。\n",
    "\n",
    "![钙离子通道](../_static/calciumchannel.png)\n",
    "\n",
    "#### 建模实现\n",
    "\n",
    "在 `braincell` 中，钙通道的建模通过继承 `CalciumChannel` 来实现，而 `CalciumChannel` 本身又继承自 `Channel`，属于 `IonChannel` 的一部分。\n",
    "\n",
    "因此，在实际应用中我们需要实现相应的接口，如 `current(V, Ca)`、`compute_derivative(V, Ca)`等。\n",
    "其中，`Ca` 参数是一个 `IonInfo` 类型，包含通道局部的钙浓度信息。而关于 `IonInfo` 的信息，我们会在 `Ion` 部分进行详细介绍。\n",
    "\n",
    "通过这些接口，我们可以构建具有复杂动力学特性的钙通道，比如电压和钙浓度双重调控、钙依赖失活、稳态激活门控等。\n",
    "\n",
    "在了解了基本的建模后，我们来看一个实际例子： `ICaN_IS2008`。\n",
    "\n",
    "这是由 Inoue & Strowbridge 在 2008 年提出的一个钙激活的非选择性阳离子通道模型。\n",
    "`braincell` 对离子通道有一套固定的命名规则，格式为 `通道类型_文献标识`。`braincell` 中大部分通道都是采用这种形式命名的，这样可以保证各通道的便于查找。\n",
    "\n",
    "回到模型本身，该模型的电流由下列公式描述：\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",
    "其中：\n",
    "- $M$ 是钙依赖激活函数\n",
    "- $p$ 是电压门控激活变量\n",
    "- $\\phi$ 是温度因子\n",
    "- $E$ 是反转电位\n",
    "- $g_{\\text{max}}$ 是最大电导\n",
    "\n",
    "以下是该模型的代码实现："
   ],
   "id": "9aff33e8ba2671b4"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:42.381428Z",
     "start_time": "2025-10-13T09:20:42.377245Z"
    }
   },
   "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": "86381944a4e65f13",
   "outputs": [],
   "execution_count": 47
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "在这个例子中，我们建立了一个具体的钙离子通道的子类，它继承自 `CalciumChannel` 。观察代码发现，在面对具体通道的建模时，`CalciumChannel` 具有很强的灵活性，只需要设定好相关的参数，并且依照公式将不同的接口用代码表达出来即可。\n",
    "\n",
    "### Potassium Channels\n",
    "\n",
    "#### 电生理特性\n",
    "\n",
    "钾离子通道是神经元中主要的超极化通道，决定了动作电位的恢复过程与静息膜电位的稳定性。其在神经元中的功能包括：\n",
    "- 动作电位的复极：动作电位过程中，钠离子通道引起去极化，而钾通道开启导致钾外流，从而使膜电位回到静息水平。\n",
    "- 控制兴奋性：某些钾通道具有延迟激活或钙依赖性质，能调节神经元的放电频率、适应性和持续兴奋性。\n",
    "- 调控网络节律：有一些钾通道在节律放电、震荡中扮演重要角色。\n",
    "\n",
    "钾通道功能广泛，其动力学建模在计算神经科学中具有重要地位。\n",
    "\n",
    "![钾离子通道](../_static/potassiumchannel.png)\n",
    "\n",
    "#### 建模实现\n",
    "\n",
    "在 `braincell` 中，我们通过继承 `PotassiumChannel` 来实现钾通道的建模。该类继承自 `Channel`。\n",
    "\n",
    "实现具体钾通道时，与实现钙通道相同，我们通常需要设定与钙通道相似的接口。\n",
    "\n",
    "我们来看一个具体实例 `IKNI_Ya1989` ，便于更好的理解。\n",
    "\n",
    "该模型最早由 Yamada 在 1989 年提出，是一类经典的缓慢非灭活钾离子通道，用于解释 AHP 期间的频率调节机制。\n",
    "\n",
    "该模型由下列公式描述：\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",
    "其中：\n",
    "- $V$ 是膜电位\n",
    "- $E_K$ 是钾离子反转电位\n",
    "- $\\bar{g}_M$ 是最大电导\n",
    "- $p$ 是门控变量\n",
    "- $V_{\\text{sh}}$ 是膜电位偏移\n",
    "- $\\tau_{\\max}$ 是最大时间常数\n",
    "- $p_{\\infty}$ 是稳态激活函数\n",
    "- $\\tau_p$ 是激活时间常数\n",
    "\n",
    "以下是代码实现："
   ],
   "id": "6572c6ea72471bfb"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:46.012986Z",
     "start_time": "2025-10-13T09:20:46.008537Z"
    }
   },
   "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": "68ce5585d39bff2f",
   "outputs": [],
   "execution_count": 48
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "从这个例子可以看出，`PotassiumChannel` 提供了统一的接口，而 `IKNI_Ya1989` 可以通过简单的数学表达式，轻松构建一个具备具体生物意义的离子通道模型。\n",
    "\n",
    "在实际应用中，我们还可以在子类下进一步构建子类。\n",
    "\n",
    "比如，先构建 `IKK2_pq_ss` ,其继承自 `PotassiumChannel` 。而在后续工作中，我们可以构建继承自  `IKK2_pq_ss` 的 `KK2A_HM1992` ，并且在 `KK2A_HM1992` 中对原本的 `IKK2_pq_ss` 进行拓展。这样的流程方便大批量建模，提高工作效率。如果你对这一部分感兴趣，可以参考我们的 [离子通道合集](https://braincell.readthedocs.io/latest/apis/braincell.channel.html) ，观察 `braincell` 中是如何具体实现的。\n",
    "\n",
    "### Sodium Channels\n",
    "\n",
    "#### 电生理特性\n",
    "\n",
    "钠离子通道是神经元中不可或缺的兴奋性通道，它在动作电位的启动和上升相中发挥着核心作用。在膜电位达到一定阈值时，钠通道快速激活并造成钠离子内流，导致膜电位迅速去极化。\n",
    "\n",
    "钠通道具有典型的双门控机制：\n",
    "- 激活门：快速响应膜电位的变化，决定通道是否打开。\n",
    "- 失活门：稍慢地响应膜电位，关闭通道使其失活。\n",
    "\n",
    "这种双门控结构使钠通道呈现出快速激活与快速失活的动力学行为，是动作电位快速起落的根本机制。\n",
    "\n",
    "![钠离子通道](../_static/sodiumchannel.png)\n",
    "\n",
    "钠通道对膜电位变化极为敏感，其动力学特性决定了其对神经元兴奋性的调节作用。在经典的 HH 模型中，钠电流通常表示为：\n",
    "\n",
    "$$\n",
    "I_{Na} = \\bar{g}_{Na} \\cdot m^3 h \\cdot (V - E_{Na})\n",
    "$$\n",
    "\n",
    "其中：\n",
    "- $m$ 为激活变量\n",
    "- $h$ 为失活变量\n",
    "- $g_{Na}$ 是最大电导\n",
    "- $E_{Na}$ 是钠离子的反转电位\n",
    "\n",
    "#### 建模实现\n",
    "\n",
    "在 `braincell` 中，我们通过继承 `SodiumChannel` 类来实现钠离子通道的建模。该类本质上也一样是 `Channel` 的子类。\n",
    "\n",
    "与前面介绍的钙通道、钾通道一致，我们只需定义相应的接口就可以完成建模。\n",
    "\n",
    "下面我们依旧来看一个具体的例子：`INa_p3q_markov`。\n",
    "\n",
    "值得一提的是，这里我们采用了不同的通道命名方式。 `INa_p3q_markov` 是根据公式结构来命名的，表示电流为 $p^3 q$ 的形式，并采用 Markov 动力学形式。\n",
    "这里是 `braincell` 对离子通道的另一套命名规则，即 `通道类别与名称_电流形式_模型` 。有部分离子通道采用这种命名规则。\n",
    "\n",
    "`INa_p3q_markov` 是一个采用 Markov 链建模形式的钠通道模型，其中电流由激活变量 $p$ 和失活变量 $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",
    "其中：\n",
    "- `p` 是激活变量\n",
    "- `q` 是失活变量\n",
    "- `φ` 是温度因子\n",
    "- $g_{max}$ 是最大电导\n",
    "- $E_{Na}$ 是钠离子反转电位\n",
    "\n",
    "具体代码如下："
   ],
   "id": "4f16718f3e23f8cb"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:51.527071Z",
     "start_time": "2025-10-13T09:20:51.522903Z"
    }
   },
   "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": "988bf5fed27a16e0",
   "outputs": [],
   "execution_count": 49
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "与其他两种离子通道完全相同，我们的钠通道建模依旧保持了高度灵活性。\n",
    "\n",
    "### Potassium Calcium Channels\n",
    "\n",
    "#### 电生理特性\n",
    "\n",
    "钙依赖性钾通道是一类在神经元中广泛存在的钾通道，其激活不仅依赖于膜电位，还依赖于胞内钙离子浓度的变化。\n",
    "\n",
    "这类通道在神经元的兴奋-抑制调控、动作电位后超极化、放电频率调节等方面起着至关重要的作用。它们通常在钙离子内流后被激活，促进钾离子外流，从而使膜电位恢复或进入超极化状态。\n",
    "\n",
    "#### 建模实现\n",
    "\n",
    "钙依赖性钾通道，在电生理学中经常简称为 `KCa` 通道，因此在建模时我们也沿用了这个写法，便于代码保持简洁性。\n",
    "\n",
    "在 `braincell` 中，我们通过直接继承 `KCaChannel` 类来实现钠离子通道的建模。该类也是 `Channel` 的子类。\n",
    "\n",
    "与前面介绍的通道一致，我们只需定义相应的接口就可以完成建模。\n",
    "\n",
    "同样，我们还是给出一个例子 `IAHP_De1994` 。\n",
    "\n",
    "`IAHP_De1994` 模型由 Destexhe 等人在 1994 年提出，其继承自 `KCaChannel` ，用于模拟慢性钙依赖钾通道的动力学行为，表达式如下：\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}{\\left(\\alpha[Ca^{2+}]_i^n + \\beta\\right)} \\\\\n",
    "        \\tau_p &=\\frac{1}{\\left(\\alpha[Ca^{2+}]_i +\\beta\\right)}\n",
    "        \\end{aligned}\n",
    "$$\n",
    "\n",
    "其中：\n",
    "- $p$ 是通道状态变量\n",
    "- $\\alpha, \\beta$ 是通道的开启/关闭速率常数\n",
    "- $n$ 是钙离子结合数\n",
    "- $\\phi$ 是温度因子\n",
    "- $g_{\\text{max}}$ 是最大电导\n",
    "- $E_K$ 是钾离子反转电位\n",
    "\n",
    "代码实现如下："
   ],
   "id": "2a61e018fc77abfe"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:21:40.034607Z",
     "start_time": "2025-10-13T09:21:40.029378Z"
    }
   },
   "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": "38cbe62e66b639c1",
   "outputs": [],
   "execution_count": 50
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "其实这几个例子的建模方式高度统一，在实现方式上没有区别。\n",
    "\n",
    "### Hyperpolarization-Activated Channels\n",
    "\n",
    "#### 电生理特性\n",
    "\n",
    "超极化激活通道，是一类在膜电位超极化时被激活的阳离子通道。它们在心脏节律神经元和中枢神经系统中的某些类型神经元中发挥着关键的调节作用：\n",
    "- 调节静息膜电位：在静息状态下产生的去极化电流可稳定神经元的膜电位，防止过度超极化。\n",
    "- 参与节律性放电：在丘脑等结构中，帮助神经元维持固有节律放电。\n",
    "- 兴奋性调节：能够降低神经元的输入电阻，从而降低细胞对输入的敏感性。\n",
    "\n",
    "#### 建模实现\n",
    "\n",
    "由于我们在 `braincell` 中没有内置很多的超极化激活通道，所以我们直接让具体的超极化激活通道继承自 `Channel` 类，这在实际使用上并不会有任何影响。\n",
    "\n",
    "下面我们给出一个超极化激活通道的例子 `Ih_HM1992` 。\n",
    "\n",
    "`Ih_HM1992` 是由 Huguenard & McCormick 在 1992 年提出的超极化激活阳离子电流模型，其动力学由下式给出：\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",
    "其中：\n",
    "- $p$ 是激活变量\n",
    "- $g_{\\text{max}}$ 是最大电导\n",
    "- $E$ 是反转电位（约 43 mV）\n",
    "- $\\phi$ 是温度因子\n",
    "\n",
    "具体代码实现如下："
   ],
   "id": "c4e3f915b0a3d794"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:20:17.842609800Z",
     "start_time": "2025-09-27T12:52:37.121379Z"
    }
   },
   "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": "e637a7f25f9a4333",
   "outputs": [],
   "execution_count": 27
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "通过这个例子你可以看到，超极化激活通道的建模可以非常自然地融入 `braincell` 框架，借助 `Channel` 的接口，轻松表达其独特的动力学性质。\n",
    "\n",
    "### Leakage Channels\n",
    "\n",
    "#### 电生理特性\n",
    "\n",
    "漏电通道是一类始终开放的离子通道，其通透性不受膜电位、电压门控、配体结合等机制的调控。在电生理学中，漏电通道主要负责维持神经元的静息膜电位，通常被建模为一个线性电导。\n",
    "\n",
    "在经典的 HH 模型中，漏电流通道主要用于模拟未明确建模的背景离子流，提供一个稳定的背景电流。\n",
    "\n",
    "#### 建模实现\n",
    "\n",
    "在 `braincell` 中，漏电通道通过继承 `LeakageChannel` 实现。由于漏电电流无激活或失活门控变量，因此该类通道不需要实现 `compute_derivative` 或 `init_state` 等微分过程，相对比较简单。\n",
    "\n",
    "我们来看一个具体的实现：`IL`，即最基础的线性漏电通道，其数学表达式为：\n",
    "\n",
    "$$\n",
    "I_L = g_L (E_L - V)\n",
    "$$\n",
    "\n",
    "其中：\n",
    "- $g_L$ 为漏电通道电导（通常为常数）\n",
    "- $E_L$ 为漏电反转电位\n",
    "- $V$ 为膜电位\n",
    "\n",
    "实现代码也相当简单："
   ],
   "id": "dc947a1c6f73790e"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:22:07.999199Z",
     "start_time": "2025-10-13T09:22:07.995501Z"
    }
   },
   "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": "8a91f851172d2adb",
   "outputs": [],
   "execution_count": 51
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "这个 `IL` 类提供了一个最简单的通道模型，它不涉及任何内部状态或动力学，是一个静态电导模型，常作为神经元模型中的基础部分。",
   "id": "4057147d7247e452"
  }
 ],
 "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
}
