{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dc802e0d",
   "metadata": {},
   "source": [
    "# Point-neuron networks with ``brainpy.state``\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/chaobrain/brainx/blob/main/docs/tutorials/brainpy.state.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/chaobrain/brainx/blob/main/docs/tutorials/brainpy.state.ipynb)\n",
    "\n",
    "`brainpy.state` provides point-neuron and synapse building blocks for stateful, JAX-native spiking network simulations. It is designed to work with `brainstate` transformations and `brainunit` quantities, so the same model can keep explicit units, reusable states, and efficient loop execution.\n",
    "\n",
    "In this tutorial, you will learn how to:\n",
    "\n",
    "- create a leaky integrate-and-fire population with `brainpy.state.LIF`;\n",
    "- initialize and inspect stateful modules;\n",
    "- read the four parts of a projection: `comm`, `syn`, `out`, and `post`;\n",
    "- build a compact excitatory-inhibitory network with `AlignPostProj`;\n",
    "- run the network and visualize spikes as a raster plot."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b758b5f",
   "metadata": {},
   "source": [
    "## Imports and simulation clock\n",
    "\n",
    "Most `brainpy.state` models are unit-aware. Set the global time step once before creating and running the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce5438e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import brainunit as u\n",
    "import brainstate\n",
    "import braintools\n",
    "import brainpy.state\n",
    "\n",
    "brainstate.environ.set(dt=0.1 * u.ms)\n",
    "dt = brainstate.environ.get_dt()\n",
    "print(\"dt =\", dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1ddd1f5",
   "metadata": {},
   "source": [
    "## 1. Create a LIF population\n",
    "\n",
    "`brainpy.state.LIF` is a stateful neuron population. The population keeps membrane potential and spike state internally, while the `update` call advances the dynamics by one time step.\n",
    "\n",
    "The example below creates a small population of independent LIF neurons. The parameters use physical units from `brainunit`, which makes the equations easier to read and helps catch dimension mistakes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef6297de",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_neurons = 16\n",
    "\n",
    "lif = brainpy.state.LIF(\n",
    "    num_neurons,\n",
    "    V_rest=-60.0 * u.mV,\n",
    "    V_th=-50.0 * u.mV,\n",
    "    V_reset=-65.0 * u.mV,\n",
    "    tau=10.0 * u.ms,\n",
    "    V_initializer=braintools.init.Normal(-60.0, 3.0, unit=u.mV),\n",
    "    spk_reset=\"soft\",\n",
    ")\n",
    "\n",
    "_ = brainstate.nn.init_all_states(lif)\n",
    "lif"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fb2a5c5",
   "metadata": {},
   "source": [
    "## 2. Drive the LIF population\n",
    "\n",
    "A model update is just a Python function. `brainstate.transform.for_loop` runs that function over a sequence and records the returned values. The state of `lif` is carried from one iteration to the next."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0f9eec4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_lif_population(inputs):\n",
    "    lif.reset_state()\n",
    "    indices = np.arange(len(inputs))\n",
    "\n",
    "    def step(i, current):\n",
    "        with brainstate.environ.context(i=i):\n",
    "            spike = lif(current)\n",
    "            return lif.V.value, spike\n",
    "\n",
    "    voltages, spikes = brainstate.transform.for_loop(step, indices, inputs)\n",
    "    return voltages, spikes\n",
    "\n",
    "n_steps = 600\n",
    "constant_current = 3.0 * u.mA\n",
    "inputs = np.ones(n_steps) * constant_current\n",
    "voltages, spikes = run_lif_population(inputs)\n",
    "\n",
    "times = np.arange(n_steps) * dt\n",
    "print(\"voltages:\", voltages.shape)\n",
    "print(\"spikes:\", spikes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4106f59e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 4))\n",
    "plt.plot(times, voltages[:, 0], label=\"neuron 0\")\n",
    "plt.plot(times, voltages[:, 1], label=\"neuron 1\", alpha=0.8)\n",
    "plt.axhline(-50.0 * u.mV, color=\"tab:red\", ls=\"--\", lw=0.8, label=\"threshold\")\n",
    "plt.xlabel(\"time\")\n",
    "plt.ylabel(\"membrane potential\")\n",
    "plt.title(\"LIF membrane potentials under constant current\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ba7a99e",
   "metadata": {},
   "source": [
    "## 3. Projection anatomy\n",
    "\n",
    "A projection connects a presynaptic event stream to a postsynaptic target. In `brainpy.state`, projections are built from explicit pieces:\n",
    "\n",
    "- `comm`: communication or connectivity, such as `brainstate.nn.EventFixedProb`;\n",
    "- `syn`: synaptic dynamics, such as `brainpy.state.Expon`;\n",
    "- `out`: output conversion, such as `brainpy.state.CUBA` for current-based output;\n",
    "- `post`: the postsynaptic population that receives the output.\n",
    "\n",
    "This decomposition lets you swap the connectivity, synapse, or output model without rewriting the whole network. It also makes the central projection question explicit: where should the synaptic state live?\n",
    "\n",
    "In a naive implementation, every realized connection can own its own synaptic state. For a sparse recurrent network, that means the state size scales with the number of connected pairs. This is often much larger than the number of neurons, and it becomes expensive when the same synaptic dynamics are repeated across many similar connections.\n",
    "\n",
    "`brainpy.state` avoids this by aligning synaptic state to a neuron dimension instead of a connection dimension. The projection still computes the same synaptic drive, but the temporal state is stored either on the presynaptic side or on the postsynaptic side.\n",
    "\n",
    "![AlignPre and AlignPost projection design](https://brainx.chaobrain.com/brainpy-state/_images/alignpre-alignpost.png)\n",
    "\n",
    "Figure: **AlignPre vs AlignPost.**  The synapse-dynamics block sits before the connection matrix on the left (state aligned to the presynaptic population, , works for all synapses, natural for one-to-many) and after it on the right (state aligned to the postsynaptic population, , exponential-family synapses, natural for many-to-one). (Orange = presynaptic, blue = postsynaptic, green = connection matrix.)\n",
    "\n",
    "**`AlignPre` aligns synaptic state with the presynaptic population.** Conceptually, each presynaptic neuron maintains its own trace, such as an exponentially decaying conductance or a more complex synaptic variable. After that trace is updated, `comm` distributes it through the connectivity and weights to the targets. This is useful when one source projects to many targets, or when the synapse model has dynamics that must be evaluated before the signal is communicated. Because the synapse state is attached to presynaptic neurons, memory scales with the presynaptic population size rather than the number of connections.\n",
    "\n",
    "**`AlignPost` aligns synaptic state with the postsynaptic population.** Here, presynaptic events are first communicated and accumulated for each target neuron. The synapse dynamics then evolve on the postsynaptic side, with one state per postsynaptic neuron for each projection. This is useful when many presynaptic sources converge onto the same target population. For exponential-family synapses, the summed incoming events can be merged before the synaptic state update, so the result is exact while memory scales with the postsynaptic population size.\n",
    "\n",
    "A practical way to choose between them is to ask where repeated work can be shared:\n",
    "\n",
    "- use `AlignPre` when one presynaptic source is reused across many targets, or when a general synapse model is naturally updated before communication;\n",
    "- use `AlignPost` when many inputs converge onto the same target population, especially for exponential current- or conductance-based synapses;\n",
    "- keep the four roles separate: `comm` moves spikes or traces across connectivity, `syn` stores and updates temporal state, `out` converts synaptic variables into current or conductance, and `post` receives the result.\n",
    "\n",
    "The E/I network below uses `AlignPostProj` because both excitatory and inhibitory presynaptic slices project back to the same LIF population. Each projection can therefore keep its synaptic state aligned to the target population size, while the two projections still use different weights and time constants."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a38cb88e",
   "metadata": {},
   "source": [
    "## 4. Build an excitatory-inhibitory network\n",
    "\n",
    "The network below contains one mixed LIF population. The first `n_exc` neurons are excitatory and the remaining `n_inh` neurons are inhibitory. Two post-aligned projections send spikes from each subpopulation back to the full population.\n",
    "\n",
    "`AlignPostProj` is a good default when several presynaptic sources converge onto the same postsynaptic group, because the synaptic state is aligned to the target population."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24acb564",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EINet(brainstate.nn.Module):\n",
    "    \"\"\"A compact recurrent E/I network built from brainpy.state components.\"\"\"\n",
    "\n",
    "    def __init__(self, n_exc, n_inh, prob, exc_weight, inh_weight):\n",
    "        super().__init__()\n",
    "        self.n_exc = n_exc\n",
    "        self.n_inh = n_inh\n",
    "        self.num = n_exc + n_inh\n",
    "\n",
    "        self.neurons = brainpy.state.LIF(\n",
    "            self.num,\n",
    "            V_rest=-52.0 * u.mV,\n",
    "            V_th=-50.0 * u.mV,\n",
    "            V_reset=-60.0 * u.mV,\n",
    "            tau=10.0 * u.ms,\n",
    "            V_initializer=braintools.init.Normal(-60.0, 5.0, unit=u.mV),\n",
    "            spk_reset=\"soft\",\n",
    "        )\n",
    "\n",
    "        self.exc = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_exc, self.num, prob, exc_weight),\n",
    "            syn=brainpy.state.Expon.desc(self.num, tau=2.0 * u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.neurons,\n",
    "        )\n",
    "        self.inh = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_inh, self.num, prob, inh_weight),\n",
    "            syn=brainpy.state.Expon.desc(self.num, tau=5.0 * u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.neurons,\n",
    "        )\n",
    "\n",
    "    def update(self, background_current):\n",
    "        previous_spikes = self.neurons.get_spike() != 0.0\n",
    "        self.exc(previous_spikes[:self.n_exc])\n",
    "        self.inh(previous_spikes[self.n_exc:])\n",
    "        self.neurons(background_current)\n",
    "        return self.neurons.get_spike()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1be5e84d",
   "metadata": {},
   "source": [
    "## 5. Initialize and inspect the network\n",
    "\n",
    "The connection weights are scaled by population size. This keeps the total recurrent input in a reasonable range when the connection probability or population size changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd7035d",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_exc = 320\n",
    "n_inh = 80\n",
    "prob = 0.05\n",
    "\n",
    "exc_weight = 1.0 / u.math.sqrt(prob * n_exc) * u.mS\n",
    "inh_weight = -1.2 / u.math.sqrt(prob * n_inh) * u.mS\n",
    "background_current = 3.0 * u.mA\n",
    "\n",
    "net = EINet(n_exc, n_inh, prob, exc_weight, inh_weight)\n",
    "_ = brainstate.nn.init_all_states(net)\n",
    "net"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcbb828b",
   "metadata": {},
   "source": [
    "## 6. Run the recurrent simulation\n",
    "\n",
    "Each time step uses the previous spikes to update recurrent synaptic input, then advances the neuron population. The returned array is a time-by-neuron spike matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80aee783",
   "metadata": {},
   "outputs": [],
   "source": [
    "duration = 300.0 * u.ms\n",
    "times = u.math.arange(0.0 * u.ms, duration, brainstate.environ.get_dt())\n",
    "\n",
    "spike_history = brainstate.transform.for_loop(\n",
    "    lambda t: net.update(background_current),\n",
    "    times,\n",
    "    pbar=brainstate.transform.ProgressBar(20),\n",
    ")\n",
    "\n",
    "print(\"spike_history:\", spike_history.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15f848ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_indices, neuron_indices = u.math.where(spike_history)\n",
    "\n",
    "plt.figure(figsize=(8, 4))\n",
    "plt.scatter(times[t_indices], neuron_indices, s=1, color=\"black\")\n",
    "plt.axhline(n_exc - 0.5, color=\"tab:red\", lw=0.8, label=\"E/I boundary\")\n",
    "plt.xlabel(\"time\")\n",
    "plt.ylabel(\"neuron index\")\n",
    "plt.title(\"Spike raster from a brainpy.state E/I network\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9751874f",
   "metadata": {},
   "source": [
    "## 7. What to change next\n",
    "\n",
    "Try these small modifications to understand the components:\n",
    "\n",
    "1. Change the excitatory `tau` in `brainpy.state.Expon.desc(...)` from `2.0 * u.ms` to `5.0 * u.ms`.\n",
    "2. Increase `prob` and reduce the weights to compare dense weak connectivity with sparse strong connectivity.\n",
    "3. Change `CUBA.desc()` to a conductance-based output in a model where the postsynaptic target exposes the voltage state needed by that output.\n",
    "4. Split the excitatory group into two projections with different time constants, then compare the raster.\n",
    "\n",
    "Because the network is stateful, rebuild or reset the states before each independent experiment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f940b45d",
   "metadata": {},
   "source": [
    "## Troubleshooting\n",
    "\n",
    "- **No spikes:** increase `background_current`, increase `exc_weight`, or simulate for a longer duration.\n",
    "- **Runaway activity:** reduce `exc_weight`, strengthen `inh_weight`, or shorten the excitatory synaptic time constant.\n",
    "- **Shape errors:** make sure `EventFixedProb(n_pre, n_post, ...)` matches the presynaptic slice and postsynaptic population. In this tutorial, the synaptic state uses `self.num` because `AlignPostProj` aligns it to the target population.\n",
    "- **Unit errors:** keep time constants in `u.ms`, voltages in `u.mV`, currents in `u.mA`, and conductance weights in `u.mS`.\n",
    "- **Unexpected repeated results:** call `brainstate.nn.init_all_states(model)` after constructing a model, and reset states between independent runs when needed."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}