{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 兴奋性-抑制性神经元网络（E-I网络）的模拟与分析\n",
    "\n",
    "本实例将基于 `braincell` 框架实现一个经典的兴奋-抑制神经元网络。通过构建由 Hodgkin-Huxley模型神经元组成的 E-I 网络，你将学习如何在 `braincell` 中实现从单神经元到网络的层级建模，并分析网络的 spike 动态特性。\n",
    "\n",
    "\n",
    "## 准备工作\n",
    "首先确保已安装必要的库（`braincell`、`brainstate`、`brainunit`、`matplotlib`），并导入所需模块："
   ],
   "id": "2912659fbdc55848"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:24:08.079758Z",
     "start_time": "2025-10-13T09:24:08.077516Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import brainpy\n",
    "import brainunit as u\n",
    "import matplotlib.pyplot as plt\n",
    "import brainstate\n",
    "import braincell"
   ],
   "id": "9d0cab5ba9e88f66",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 代码详解\n",
    "\n",
    "### 参数定义\n",
    "\n",
    "首先定义关键物理参数，这些参数决定了神经元的电生理特性和网络规模："
   ],
   "id": "7c40b3a6e85f45ae"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:24:08.183332Z",
     "start_time": "2025-10-13T09:24:08.179965Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 神经元动作电位发放阈值\n",
    "V_th = -20. * u.mV\n",
    "\n",
    "# 神经元膜面积\n",
    "area = 20000 * u.um **2\n",
    "area = area.in_unit(u.cm** 2)\n",
    "\n",
    "# 膜电容\n",
    "Cm = (1 * u.uF * u.cm ** -2) * area  # 总电容 = 比电容 × 面积"
   ],
   "id": "c22585ffe4f19bdd",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 定义 HH 单神经元模型\n",
    "\n",
    "使用 `SingleCompartment` 构建基于 HH 模型的单神经元，包含钠通道 `INa` 、钾通道 `IK` 和漏电流 `IL` ，这些通道共同决定神经元的放电特性："
   ],
   "id": "12e040fbf99cff1b"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:24:08.295440Z",
     "start_time": "2025-10-13T09:24:08.292318Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class HH(braincell.SingleCompartment):\n",
    "    def __init__(self, in_size):\n",
    "        # 初始化单房室神经元\n",
    "        super().__init__(in_size, C=Cm, solver='ind_exp_euler')\n",
    "\n",
    "        # 钠离子通道（INa）\n",
    "        self.na = braincell.ion.SodiumFixed(in_size, E=50. * u.mV)\n",
    "        self.na.add_elem(\n",
    "            # 最大电导\n",
    "            INa=braincell.channel.INa_TM1991(in_size, g_max=(100. * u.mS * u.cm **-2) * area, V_sh=-63. * u.mV)\n",
    "        )\n",
    "\n",
    "        # 钾离子通道（IK）\n",
    "        self.k = braincell.ion.PotassiumFixed(in_size, E=-90 * u.mV)\n",
    "        self.k.add_elem(\n",
    "            # 最大电导\n",
    "            IK=braincell.channel.IK_TM1991(in_size, g_max=(30. * u.mS * u.cm** -2) * area, V_sh=-63. * u.mV)\n",
    "        )\n",
    "\n",
    "        # 漏电流（IL）\n",
    "        self.IL = braincell.channel.IL(\n",
    "            in_size,\n",
    "            E=-60. * u.mV,\n",
    "            g_max=(5. * u.nS * u.cm **-2) * area  # 最大电导\n",
    "        )"
   ],
   "id": "c31fb6de04535794",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 定义 E-I 网络：兴奋性与抑制性神经元的连接\n",
    "\n",
    "构建由兴奋性（E）和抑制性（I）神经元组成的网络，模拟皮层网络中常见的 E-I 平衡机制："
   ],
   "id": "a70a4069f943e1b"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:24:08.422373Z",
     "start_time": "2025-10-13T09:24:08.418238Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class EINet(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # 网络规模\n",
    "        self.n_exc = 3200\n",
    "        self.n_inh = 800\n",
    "        self.num = self.n_exc + self.n_inh  # 总神经元数：4000\n",
    "\n",
    "        # 初始化神经元群体\n",
    "        self.N = HH(self.num)\n",
    "\n",
    "        # 兴奋性突触投射\n",
    "        self.E = brainpy.state.AlignPostProj(\n",
    "            # 连接规则\n",
    "            comm=brainstate.nn.EventFixedProb(\n",
    "                self.n_exc, self.num, conn_num=0.02,  # 连接概率\n",
    "                conn_weight=6. * u.nS  # 突触权重\n",
    "            ),\n",
    "            # 突触动力学\n",
    "            syn=brainpy.state.Expon(self.num, tau=5. * u.ms),\n",
    "            # 突触后效应\n",
    "            out=brainpy.state.COBA(E=0. * u.mV),\n",
    "            post=self.N  # 投射目标为神经元群体 N\n",
    "        )\n",
    "\n",
    "        # 抑制性突触投射\n",
    "        self.I = brainpy.state.AlignPostProj(\n",
    "            # 连接规则\n",
    "            comm=brainstate.nn.EventFixedProb(\n",
    "                self.n_inh, self.num, conn_num=0.02,\n",
    "                conn_weight=67. * u.nS\n",
    "            ),\n",
    "            # 突触动力学\n",
    "            syn=brainpy.state.Expon(self.num, tau=10. * u.ms),\n",
    "            # 突触后效应\n",
    "            out=brainpy.state.COBA(E=-80. * u.mV),\n",
    "            post=self.N  # 投射目标为神经元群体 N\n",
    "        )\n",
    "\n",
    "    def update(self, t):\n",
    "        # 定义网络随时间的更新规则\n",
    "        with brainstate.environ.context(t=t):\n",
    "            # 获取当前时刻的 spike 信号\n",
    "            spk = self.N.spike.value\n",
    "\n",
    "            # 兴奋性神经元的 spike 驱动兴奋性突触投射\n",
    "            self.E(spk[:self.n_exc])\n",
    "\n",
    "            # 抑制性神经元的 spike 驱动抑制性突触投射\n",
    "            self.I(spk[self.n_exc:])\n",
    "\n",
    "            # 神经元接收突触输入后更新状态，返回新的 spike 信号\n",
    "            spk = self.N(0. * u.nA)\n"
   ],
   "id": "e72dfafb35da107c",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 运行网络模拟\n",
    "\n",
    "初始化网络并运行仿真，记录每个时刻的神经元 spike 活动："
   ],
   "id": "fa97af7731e5f5e8"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:24:09.161036Z",
     "start_time": "2025-10-13T09:24:08.524778Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 初始化 E-I 网络\n",
    "net = EINet()\n",
    "brainstate.nn.init_all_states(net)  # 初始化网络中所有神经元和突触的状态\n",
    "\n",
    "# 设置仿真参数并运行\n",
    "with brainstate.environ.context(dt=0.1 * u.ms):  # 时间步长\n",
    "    # 生成仿真时间序列\n",
    "    times = u.math.arange(0. * u.ms, 100. * u.ms, brainstate.environ.get_dt())\n",
    "\n",
    "    # 循环更新网络状态\n",
    "    spikes = brainstate.transform.for_loop(\n",
    "        net.update, times,\n",
    "        pbar=brainstate.transform.ProgressBar(10)  # 显示进度条\n",
    "    )\n"
   ],
   "id": "655509eeebf82b38",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\braintools\\surrogate.py:72: UserWarning: Explicitly requested dtype float64 requested in asarray is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.\n",
      "  z = jnp.asarray(x >= 0, dtype=x.dtype)\n",
      "Running for 1,000 iterations: 100%|██████████| 1000/1000 [00:00<00:00, 17183.01it/s]\n"
     ]
    }
   ],
   "execution_count": 12
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### 可视化网络 spike 活动\n",
    "\n",
    "将 spike 数据绘制成格栅图，直观展示神经元在时间上的发放模式："
   ],
   "id": "e595aa0208703fac"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-13T09:24:09.296059Z",
     "start_time": "2025-10-13T09:24:09.253091Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 提取 spike 发生的时间和神经元索引\n",
    "t_indices, n_indices = u.math.where(spikes)\n",
    "\n",
    "# 绘制 raster plot\n",
    "plt.scatter(times[t_indices], n_indices, s=1)\n",
    "plt.xlabel('Time (ms)')  # 横轴：时间\n",
    "plt.ylabel('Neuron index')  # 纵轴：神经元索引\n",
    "plt.show()\n",
    "\n",
    "# 此例程在jupyter中打断后无法完成可视化，暂时先将运行结果保留！"
   ],
   "id": "194cd88405441b76",
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "where requires ndarray or scalar arguments, got <class 'NoneType'> at position 0.",
     "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[13]\u001B[39m\u001B[32m, line 2\u001B[39m\n\u001B[32m      1\u001B[39m \u001B[38;5;66;03m# 提取 spike 发生的时间和神经元索引\u001B[39;00m\n\u001B[32m----> \u001B[39m\u001B[32m2\u001B[39m t_indices, n_indices = \u001B[43mu\u001B[49m\u001B[43m.\u001B[49m\u001B[43mmath\u001B[49m\u001B[43m.\u001B[49m\u001B[43mwhere\u001B[49m\u001B[43m(\u001B[49m\u001B[43mspikes\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m      4\u001B[39m \u001B[38;5;66;03m# 绘制 raster plot\u001B[39;00m\n\u001B[32m      5\u001B[39m plt.scatter(times[t_indices], n_indices, s=\u001B[32m1\u001B[39m)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\saiunit\\math\\_fun_keep_unit.py:3305\u001B[39m, in \u001B[36mwhere\u001B[39m\u001B[34m(condition, x, y, size, fill_value)\u001B[39m\n\u001B[32m   3303\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m x \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m y \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[32m   3304\u001B[39m     \u001B[38;5;28;01massert\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(fill_value, Quantity), \u001B[33m\"\u001B[39m\u001B[33mfill_value should not be a Quantity.\u001B[39m\u001B[33m\"\u001B[39m\n\u001B[32m-> \u001B[39m\u001B[32m3305\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mjnp\u001B[49m\u001B[43m.\u001B[49m\u001B[43mwhere\u001B[49m\u001B[43m(\u001B[49m\u001B[43mcondition\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43msize\u001B[49m\u001B[43m=\u001B[49m\u001B[43msize\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mfill_value\u001B[49m\u001B[43m=\u001B[49m\u001B[43mfill_value\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m   3307\u001B[39m \u001B[38;5;28;01massert\u001B[39;00m size \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m fill_value \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m, \u001B[33m\"\u001B[39m\u001B[33msize and fill_value are only supported when x and y are not None.\u001B[39m\u001B[33m\"\u001B[39m\n\u001B[32m   3308\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(x, Quantity) \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(y, Quantity):\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\jax\\_src\\numpy\\lax_numpy.py:2771\u001B[39m, in \u001B[36mwhere\u001B[39m\u001B[34m(condition, x, y, size, fill_value)\u001B[39m\n\u001B[32m   2711\u001B[39m \u001B[38;5;250m\u001B[39m\u001B[33;03m\"\"\"Select elements from two arrays based on a condition.\u001B[39;00m\n\u001B[32m   2712\u001B[39m \n\u001B[32m   2713\u001B[39m \u001B[33;03mJAX implementation of :func:`numpy.where`.\u001B[39;00m\n\u001B[32m   (...)\u001B[39m\u001B[32m   2768\u001B[39m \u001B[33;03m  Array([0, 0, 0, 0, 0, 5, 6, 7, 8, 9], dtype=int32)\u001B[39;00m\n\u001B[32m   2769\u001B[39m \u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m   2770\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m x \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m y \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m2771\u001B[39m   \u001B[43mutil\u001B[49m\u001B[43m.\u001B[49m\u001B[43mcheck_arraylike\u001B[49m\u001B[43m(\u001B[49m\u001B[33;43m\"\u001B[39;49m\u001B[33;43mwhere\u001B[39;49m\u001B[33;43m\"\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcondition\u001B[49m\u001B[43m)\u001B[49m\n\u001B[32m   2772\u001B[39m   \u001B[38;5;28;01mreturn\u001B[39;00m nonzero(condition, size=size, fill_value=fill_value)\n\u001B[32m   2773\u001B[39m \u001B[38;5;28;01melse\u001B[39;00m:\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\Document\\PyCharm\\Project\\braincell(collaborator)\\.venv\\Lib\\site-packages\\jax\\_src\\numpy\\util.py:183\u001B[39m, in \u001B[36mcheck_arraylike\u001B[39m\u001B[34m(fun_name, emit_warning, stacklevel, *args)\u001B[39m\n\u001B[32m    180\u001B[39m   warnings.warn(msg + \u001B[33m\"\u001B[39m\u001B[33m In a future JAX release this will be an error.\u001B[39m\u001B[33m\"\u001B[39m,\n\u001B[32m    181\u001B[39m                 category=\u001B[38;5;167;01mDeprecationWarning\u001B[39;00m, stacklevel=stacklevel)\n\u001B[32m    182\u001B[39m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m--> \u001B[39m\u001B[32m183\u001B[39m   \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mTypeError\u001B[39;00m(msg.format(fun_name, \u001B[38;5;28mtype\u001B[39m(arg), pos))\n",
      "\u001B[31mTypeError\u001B[39m: where requires ndarray or scalar arguments, got <class 'NoneType'> at position 0."
     ]
    }
   ],
   "execution_count": 13
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "\n",
    "## 结果解读\n",
    "\n",
    "运行代码后，你将看到一张格栅图，如下图：\n",
    "\n",
    "![raster图](../../_static/cobahh2007.png)\n",
    "\n",
    "其中每个点代表一个神经元在特定时间的 spike 发放。典型的 E-I 网络动态具有以下特征：\n",
    "- 异步发放：神经元发放时间分散，无明显同步节律。\n",
    "- 稀疏活动：大多数神经元在 100 ms 内发放次数较少。\n",
    "- 无爆发式同步：由于抑制性突触的快速反馈，避免了大规模神经元同步爆发。\n",
    "\n",
    "这些特征表明 E-I 网络通过兴奋性和抑制性的动态平衡，维持了稳定且符合生理特性的活动模式。\n",
    "\n",
    "## 扩展练习\n",
    "\n",
    "- 调整抑制性突触权重，观察网络是否出现过度同步，即爆发式活动。\n",
    "- 增加兴奋性神经元比例，分析 E-I 失衡对网络动态的影响。\n",
    "- 延长仿真时间，观察网络是否维持稳定的异步活动。\n",
    "\n",
    "通过这些扩展，你可以深入理解 E-I 平衡在维持神经网络功能中的核心作用，以及 ``braincell`` 在复杂网络建模中的灵活性。"
   ],
   "id": "55ef3973da8a763e"
  }
 ],
 "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
}