{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a6cf4a6f",
   "metadata": {},
   "source": [
    "# Event-driven computations with ``brainevent``\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/chaobrain/brainx/blob/main/docs/tutorials/brainevent_event-driven.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/brainevent_event-driven.ipynb)\n",
    "\n",
    "## What is ``brainevent``?\n",
    "\n",
    "``brainevent`` is the event-driven computation package in the BrainX ecosystem. It provides event representations and connectivity operators for simulations where neural activity is sparse, such as spiking neural networks and event-based recurrent models.\n",
    "\n",
    "The main reason to use ``brainevent`` is efficiency and clarity in sparse activity. At many time steps, only a small subset of neurons fires. Instead of treating every presynaptic neuron as active, a ``BinaryArray`` marks the neurons that actually fired, and event-aware multiplication visits the relevant rows or connections of the weight structure.\n",
    "\n",
    "![brainevent basic idea framework](../_static/image/brainevent_overview.png)\n",
    "\n",
    "*Figure: Basic idea of ``brainevent``. The left side shows event representation, where a binary event vector or event block records which presynaptic neurons fired. The right side summarizes common connectivity data representations, including dense matrices, CSR sparse connectivity, JIT-generated connectivity, and fixed fan-in or fan-out structures. The lower part shows the operation layer: forward event-driven projections such as ``event @ W`` and ``events @ W`` are the common path, while reverse products are mainly useful for feedback, readout, or credit-signal style computations.*\n",
    "\n",
    "This tutorial follows one small event vector through the main representations used in practice: dense matrices for clarity, CSR for explicit sparse connectivity, JIT-generated connectivity for large random graphs, and fixed connection-count structures for constrained fan-in or fan-out.\n",
    "\n",
    "You will learn how to:\n",
    "\n",
    "- represent binary spike events with ``brainevent.BinaryArray``;\n",
    "- read matrix-vector (``mv``) and matrix-matrix (``mm``) event operations;\n",
    "- build and use dense, CSR, JIT, and fixed connectivity structures;\n",
    "- place event-driven multiplication inside a ``jax.jit`` function without rebuilding connectivity inside the compiled function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fc14b0b",
   "metadata": {},
   "source": [
    "## 1. Imports\n",
    "\n",
    "The examples use small arrays so that the structure is easy to inspect. In a real simulation, the same APIs are intended for sparse, large-scale event streams."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "807e80a0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:51:59.706439Z",
     "iopub.status.busy": "2026-06-22T01:51:59.706439Z",
     "iopub.status.idle": "2026-06-22T01:52:06.305508Z",
     "shell.execute_reply": "2026-06-22T01:52:06.305508Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainevent\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ad7c2ff",
   "metadata": {},
   "source": [
    "## 2. Binary events\n",
    "\n",
    "A spike train at one time step is naturally represented as a boolean or binary vector:\n",
    "\n",
    "- `True` or `1`: the neuron fired;\n",
    "- `False` or `0`: the neuron was silent.\n",
    "\n",
    "`brainevent.BinaryArray` wraps this event vector and tells downstream operators to use event-driven kernels where possible. The object exposes array-like metadata such as shape and dtype, supports indexing, and stores the raw JAX array in `.value` for ordinary reductions or inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "acb9ae32",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:06.307518Z",
     "iopub.status.busy": "2026-06-22T01:52:06.307518Z",
     "iopub.status.idle": "2026-06-22T01:52:06.888962Z",
     "shell.execute_reply": "2026-06-22T01:52:06.888962Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Binary event representation\n",
      "  events      = BinaryArray(value=[ True False  True False  True], dtype=bool)\n",
      "  raw values  = [ True False  True False  True]\n",
      "  shape       = (5,)\n",
      "  spike count = 3\n",
      "  active ids  = [0 2 4]\n"
     ]
    }
   ],
   "source": [
    "# Five presynaptic neurons at one simulation step.\n",
    "spikes_bool = jnp.array([True, False, True, False, True])\n",
    "events = brainevent.BinaryArray(spikes_bool)\n",
    "\n",
    "print(\"Binary event representation\")\n",
    "print(f\"  events      = {events}\")\n",
    "print(f\"  raw values  = {events.value}\")\n",
    "print(f\"  shape       = {events.shape}\")\n",
    "print(f\"  spike count = {events.value.sum()}\")\n",
    "print(f\"  active ids  = {jnp.where(events.value)[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32ee7b70",
   "metadata": {},
   "source": [
    "A two-dimensional `BinaryArray` represents a batch or a short time block. In this case, matrix multiplication becomes an `mm` operation: a binary event matrix times a weight matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "103063c0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:06.890977Z",
     "iopub.status.busy": "2026-06-22T01:52:06.890977Z",
     "iopub.status.idle": "2026-06-22T01:52:06.928494Z",
     "shell.execute_reply": "2026-06-22T01:52:06.928494Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Binary event block\n",
      "  shape = (3, 5)\n",
      "  spike counts per step = [2 2 3]\n"
     ]
    }
   ],
   "source": [
    "# Three time steps, five neurons.\n",
    "event_block = brainevent.BinaryArray(jnp.array([\n",
    "    [True,  False, True,  False, False],\n",
    "    [False, True,  False, False, True ],\n",
    "    [True,  True,  False, True,  False],\n",
    "]))\n",
    "\n",
    "print(\"Binary event block\")\n",
    "print(f\"  shape = {event_block.shape}\")\n",
    "print(f\"  spike counts per step = {event_block.value.sum(axis=1)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40bfc5d2",
   "metadata": {},
   "source": [
    "## 3. Operation semantics: `mv` and `mm`\n",
    "\n",
    "`brainevent` keeps the same algebra as ordinary matrix multiplication. The forward convention in this tutorial is `x @ W`: `x` is a presynaptic event vector or event block, and `W` has shape `(n_pre, n_post)`. The implementation may call a transposed low-level kernel internally so that the event axis becomes the axis visited by the event-driven kernel. This lets the kernel skip inactive events while preserving the same user-facing `@` algebra as the dense reference.\n",
    "\n",
    "| Expression | Event input | Connectivity | Name | Result |\n",
    "| --- | --- | --- | --- | --- |\n",
    "| `x @ W` | `(n_pre,)` | `(n_pre, n_post)` | `mv` | `(n_post,)` |\n",
    "| `X @ W` | `(batch_or_time, n_pre)` | `(n_pre, n_post)` | `mm` | `(batch_or_time, n_post)` |\n",
    "\n",
    "The opposite side is also supported when the event shape matches the postsynaptic axis: `W @ y` with `y.shape == (n_post,)` returns `(n_pre,)`, and `W @ Y` with `Y.shape == (n_post, batch_or_time)` returns `(n_pre, batch_or_time)`. Physically, `W @ Y` is a target-to-source aggregation: `Y` marks active postsynaptic, target-side events or gates, and the product accumulates their weighted influence back onto each presynaptic source neuron. This can be useful for feedback, readout or credit signals, and algorithms that project postsynaptic activity back through the same connectivity. It is a different algebraic product from forward propagation, not a replacement for `x @ W`.\n",
    "\n",
    "![Matrix multiplication semantics for event-driven mv/mm operations and dense references](../_static/image/brainevent_mv_mm.png)\n",
    "\n",
    "Common pitfalls:\n",
    "\n",
    "- Store weights as `(n_pre, n_post)`: rows are presynaptic sources, columns are postsynaptic targets.\n",
    "- Treat the leading dimension of a 2-D left-hand `BinaryArray` as batch or time, not as another neuron axis.\n",
    "- Build `CSR`, fixed, and JIT connectivity objects before entering `jax.jit`. Inside a compiled function, pass the changing spike values, wrap them with `BinaryArray`, and multiply by the prebuilt connectivity object.\n",
    "\n",
    "The following sections reuse this convention for every connectivity representation.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd2977e4",
   "metadata": {},
   "source": [
    "## 4. Dense or general matrix data\n",
    "\n",
    "Dense JAX arrays are the easiest way to start because they make the expected result transparent. If `x` is a `BinaryArray` and `W` is a dense matrix, `x @ W` computes postsynaptic input by summing the rows whose presynaptic events are active.\n",
    "\n",
    "Use dense matrices for teaching, debugging, and small examples. The sparse and structured representations below keep the same multiplication syntax while changing how the connectivity is stored or generated.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5b09aa09",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:06.931694Z",
     "iopub.status.busy": "2026-06-22T01:52:06.931694Z",
     "iopub.status.idle": "2026-06-22T01:52:16.495641Z",
     "shell.execute_reply": "2026-06-22T01:52:16.495641Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dense/general matrix operations\n",
      "  dense_w shape = (5, 3)\n",
      "  mv result     = [1.4000001 0.3       1.6      ]\n",
      "  mm shape      = (3, 3)\n",
      "[[1.2        0.3        1.1       ]\n",
      " [0.3        0.7        0.5       ]\n",
      " [0.90000004 1.2        0.8       ]]\n"
     ]
    }
   ],
   "source": [
    "# Dense weights: 5 presynaptic neurons -> 3 postsynaptic neurons.\n",
    "dense_w = jnp.array([\n",
    "    [0.8, 0.0, 0.2],\n",
    "    [0.1, 0.7, 0.0],\n",
    "    [0.4, 0.3, 0.9],\n",
    "    [0.0, 0.5, 0.6],\n",
    "    [0.2, 0.0, 0.5],\n",
    "], dtype=jnp.float32)\n",
    "\n",
    "# Matrix-vector: one event vector times one weight matrix.\n",
    "dense_mv = events @ dense_w\n",
    "\n",
    "# Matrix-matrix: a block of event vectors times the same weight matrix.\n",
    "dense_mm = event_block @ dense_w\n",
    "\n",
    "print(\"Dense/general matrix operations\")\n",
    "print(f\"  dense_w shape = {dense_w.shape}\")\n",
    "print(f\"  mv result     = {dense_mv}\")\n",
    "print(f\"  mm shape      = {dense_mm.shape}\")\n",
    "print(dense_mm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ca3734",
   "metadata": {},
   "source": [
    "The `@` operator is the recommended high-level interface throughout this tutorial. It lets the same event code work with dense arrays, CSR matrices, fixed structures, and JIT connectivity. Lower-level kernels are useful when implementing custom operators or benchmarking internals, but they are not needed for ordinary tutorial code.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4957ecf1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:16.498644Z",
     "iopub.status.busy": "2026-06-22T01:52:16.497644Z",
     "iopub.status.idle": "2026-06-22T01:52:16.712299Z",
     "shell.execute_reply": "2026-06-22T01:52:16.711794Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dense reference check\n",
      "  mv matches dense reference: True\n",
      "  mm matches dense reference: True\n"
     ]
    }
   ],
   "source": [
    "# The event-driven result matches ordinary dense multiplication.\n",
    "reference_mv = events.value.astype(dense_w.dtype) @ dense_w\n",
    "reference_mm = event_block.value.astype(dense_w.dtype) @ dense_w\n",
    "\n",
    "print(\"Dense reference check\")\n",
    "print(f\"  mv matches dense reference: {jnp.allclose(dense_mv, reference_mv)}\")\n",
    "print(f\"  mm matches dense reference: {jnp.allclose(dense_mm, reference_mm)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b47beb3",
   "metadata": {},
   "source": [
    "## 5. CSR sparse connectivity\n",
    "\n",
    "Real neural connectivity is usually sparse. CSR, compressed sparse row, stores a matrix row by row:\n",
    "\n",
    "- `data`: non-zero weights;\n",
    "- `indices`: column index for each stored weight;\n",
    "- `indptr`: offsets showing where each row starts and ends in `data` and `indices`.\n",
    "\n",
    "CSR is a natural format for forward event propagation because each active presynaptic neuron selects a row, and only the non-zero outgoing connections in that row are visited.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bfc85875",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:16.714301Z",
     "iopub.status.busy": "2026-06-22T01:52:16.714301Z",
     "iopub.status.idle": "2026-06-22T01:52:18.692676Z",
     "shell.execute_reply": "2026-06-22T01:52:18.692676Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CSR representation\n",
      "  shape   = (5, 3)\n",
      "  nse     = 11\n",
      "  data    = [0.8 0.2 0.1 0.7 0.4 0.3 0.9 0.5 0.6 0.2 0.5]\n",
      "  indices = [0 2 0 1 0 1 2 1 2 0 2]\n",
      "  indptr  = [ 0  2  4  7  9 11]\n",
      "Dense view:\n",
      "[[0.8 0.  0.2]\n",
      " [0.1 0.7 0. ]\n",
      " [0.4 0.3 0.9]\n",
      " [0.  0.5 0.6]\n",
      " [0.2 0.  0.5]]\n"
     ]
    }
   ],
   "source": [
    "# Build the same 5 -> 3 matrix from coordinate triplets.\n",
    "row = jnp.array([0, 0, 1, 1, 2, 2, 2, 3, 3, 4, 4])\n",
    "col = jnp.array([0, 2, 0, 1, 0, 1, 2, 1, 2, 0, 2])\n",
    "data = jnp.array([0.8, 0.2, 0.1, 0.7, 0.4, 0.3, 0.9, 0.5, 0.6, 0.2, 0.5], dtype=jnp.float32)\n",
    "\n",
    "indptr, indices, order = brainevent.coo2csr(row, col, shape=(5, 3))\n",
    "csr_w = brainevent.CSR((data[order], indices, indptr), shape=(5, 3))\n",
    "\n",
    "print(\"CSR representation\")\n",
    "print(f\"  shape   = {csr_w.shape}\")\n",
    "print(f\"  nse     = {csr_w.nse}\")\n",
    "print(f\"  data    = {csr_w.data}\")\n",
    "print(f\"  indices = {csr_w.indices}\")\n",
    "print(f\"  indptr  = {csr_w.indptr}\")\n",
    "print(\"Dense view:\")\n",
    "print(csr_w.todense())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "456a4fdc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:18.694691Z",
     "iopub.status.busy": "2026-06-22T01:52:18.694691Z",
     "iopub.status.idle": "2026-06-22T01:52:19.141866Z",
     "shell.execute_reply": "2026-06-22T01:52:19.141361Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CSR event-driven operations\n",
      "  mv result = [1.4000001 0.3       1.6      ]\n",
      "  mm shape  = (3, 3)\n",
      "[[1.2        0.3        1.1       ]\n",
      " [0.3        0.7        0.5       ]\n",
      " [0.90000004 1.2        0.8       ]]\n",
      "\n",
      "Consistency with dense weights\n",
      "  mv matches: True\n",
      "  mm matches: True\n"
     ]
    }
   ],
   "source": [
    "csr_mv = events @ csr_w\n",
    "csr_mm = event_block @ csr_w\n",
    "\n",
    "print(\"CSR event-driven operations\")\n",
    "print(f\"  mv result = {csr_mv}\")\n",
    "print(f\"  mm shape  = {csr_mm.shape}\")\n",
    "print(csr_mm)\n",
    "print()\n",
    "print(\"Consistency with dense weights\")\n",
    "print(f\"  mv matches: {jnp.allclose(csr_mv, dense_mv)}\")\n",
    "print(f\"  mm matches: {jnp.allclose(csr_mm, dense_mm)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7a977f",
   "metadata": {},
   "source": [
    "`CSR` is usually the first sparse structure to reach for when the connection graph is known and sparse. It keeps the graph explicit, is easy to compare against a dense reference, and works with the same `events @ weights` syntax used above.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e14d90f",
   "metadata": {},
   "source": [
    "## 6. CSC and orientation\n",
    "\n",
    "CSC, compressed sparse column, stores the same idea column by column. It is useful when algorithms need efficient access by postsynaptic target or column. The forward examples below use CSR because a presynaptic event vector naturally activates rows, but it is useful to recognize both formats in API references and conversion utilities.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9d8f98f0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:19.143876Z",
     "iopub.status.busy": "2026-06-22T01:52:19.142876Z",
     "iopub.status.idle": "2026-06-22T01:52:19.269324Z",
     "shell.execute_reply": "2026-06-22T01:52:19.269324Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CSC representation\n",
      "  shape   = (5, 3)\n",
      "  nse     = 11\n",
      "Dense view:\n",
      "[[0.8 0.  0.2]\n",
      " [0.1 0.7 0. ]\n",
      " [0.4 0.3 0.9]\n",
      " [0.  0.5 0.6]\n",
      " [0.2 0.  0.5]]\n"
     ]
    }
   ],
   "source": [
    "# A CSC object can store the same dense matrix by columns.\n",
    "# Column 0: rows 0, 1, 2, 4\n",
    "# Column 1: rows 1, 2, 3\n",
    "# Column 2: rows 0, 2, 3, 4\n",
    "csc_data = jnp.array([0.8, 0.1, 0.4, 0.2, 0.7, 0.3, 0.5, 0.2, 0.9, 0.6, 0.5], dtype=jnp.float32)\n",
    "csc_indices = jnp.array([0, 1, 2, 4, 1, 2, 3, 0, 2, 3, 4])\n",
    "csc_indptr = jnp.array([0, 4, 7, 11])\n",
    "csc_w = brainevent.CSC((csc_data, csc_indices, csc_indptr), shape=(5, 3))\n",
    "\n",
    "print(\"CSC representation\")\n",
    "print(f\"  shape   = {csc_w.shape}\")\n",
    "print(f\"  nse     = {csc_w.nse}\")\n",
    "print(\"Dense view:\")\n",
    "print(csc_w.todense())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "395d4b81",
   "metadata": {},
   "source": [
    "## 7. JIT dynamic connectivity\n",
    "\n",
    "Sometimes storing the full connection matrix is still too expensive. JIT connectivity structures store the random rule, not every sampled edge. During multiplication, the operator regenerates the needed random connections from the rule and seed.\n",
    "\n",
    "Useful JIT connectivity families include:\n",
    "\n",
    "- `JITCScalarR`: Bernoulli connections with one scalar weight;\n",
    "- `JITCNormalR`: Bernoulli connections with weights sampled from a normal distribution;\n",
    "- `JITCUniformR`: Bernoulli connections with weights sampled from a uniform distribution.\n",
    "\n",
    "The `R` suffix means row-oriented, which is usually the right orientation for forward propagation from presynaptic events.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8730da11",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:19.271334Z",
     "iopub.status.busy": "2026-06-22T01:52:19.271334Z",
     "iopub.status.idle": "2026-06-22T01:52:21.536056Z",
     "shell.execute_reply": "2026-06-22T01:52:21.536056Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JIT dynamic connectivity\n",
      "  shape = (5, 3)\n",
      "  weight = 0.20000000298023224\n",
      "  probability = 0.4\n",
      "  seed = 2024\n",
      "  mv result = [0.4 0.  0. ]\n",
      "  mm result shape = (3, 3)\n",
      "[[0.2 0.  0. ]\n",
      " [0.2 0.2 0.2]\n",
      " [0.2 0.2 0.2]]\n"
     ]
    }
   ],
   "source": [
    "# Homogeneous random connections: (weight, probability, seed).\n",
    "jit_w = brainevent.JITCScalarR(\n",
    "    (0.2, 0.4, 2024),\n",
    "    shape=(5, 3),\n",
    ")\n",
    "\n",
    "jit_mv = events @ jit_w\n",
    "jit_mm = event_block @ jit_w\n",
    "\n",
    "print(\"JIT dynamic connectivity\")\n",
    "print(f\"  shape = {jit_w.shape}\")\n",
    "print(f\"  weight = {jit_w.weight}\")\n",
    "print(f\"  probability = {jit_w.prob}\")\n",
    "print(f\"  seed = {jit_w.seed}\")\n",
    "print(f\"  mv result = {jit_mv}\")\n",
    "print(f\"  mm result shape = {jit_mm.shape}\")\n",
    "print(jit_mm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "153b661f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:21.538066Z",
     "iopub.status.busy": "2026-06-22T01:52:21.538066Z",
     "iopub.status.idle": "2026-06-22T01:52:24.506631Z",
     "shell.execute_reply": "2026-06-22T01:52:24.506128Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Other JIT connectivity rules\n",
      "  normal mv = [0.25821236 0.         0.        ]\n",
      "  uniform mv = [-0.09996567  0.         -0.06874797]\n"
     ]
    }
   ],
   "source": [
    "# Distribution-based JIT connectivity uses the same multiplication pattern.\n",
    "jit_normal = brainevent.JITCNormalR(\n",
    "    (0.0, 0.1, 0.3, 7),  # (mean, std, probability, seed)\n",
    "    shape=(5, 3),\n",
    ")\n",
    "\n",
    "jit_uniform = brainevent.JITCUniformR(\n",
    "    (-0.1, 0.3, 0.3, 8),  # (low, high, probability, seed)\n",
    "    shape=(5, 3),\n",
    ")\n",
    "\n",
    "print(\"Other JIT connectivity rules\")\n",
    "print(f\"  normal mv = {events @ jit_normal}\")\n",
    "print(f\"  uniform mv = {events @ jit_uniform}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a880d2d",
   "metadata": {},
   "source": [
    "JIT connectivity uses the same multiplication pattern as dense and CSR connectivity. Keep the rule object outside performance-critical loops, then reuse it with changing event values. This keeps tutorial code close to the way a larger simulation would be written.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb29e433",
   "metadata": {},
   "source": [
    "## 8. Fixed connection-count structures\n",
    "\n",
    "Fixed connectivity structures are useful when each neuron should have a fixed number of partners.\n",
    "\n",
    "- `FixedNumPerPre`: each presynaptic neuron stores a fixed number of postsynaptic targets. This is a fixed fan-out representation.\n",
    "- `FixedNumPerPost`: each postsynaptic neuron stores a fixed number of presynaptic sources. This is a fixed fan-in representation.\n",
    "\n",
    "Both store two compact arrays: weights and integer indices. These names are the current `brainevent` API in version 0.1.0.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "07763fe0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:24.507642Z",
     "iopub.status.busy": "2026-06-22T01:52:24.507642Z",
     "iopub.status.idle": "2026-06-22T01:52:25.475255Z",
     "shell.execute_reply": "2026-06-22T01:52:25.475255Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FixedNumPerPre: fixed fan-out from each presynaptic neuron\n",
      "  shape = (5, 3)\n",
      "  data shape = (5, 2)\n",
      "  indices shape = (5, 2)\n",
      "  dense view:\n",
      "[[0.8 0.  0.2]\n",
      " [0.  0.7 0.1]\n",
      " [0.4 0.3 0. ]\n",
      " [0.  0.5 0.6]\n",
      " [0.2 0.  0.5]]\n",
      "  mv result = [1.4000001 0.3       0.7      ]\n",
      "  mm result shape = (3, 3)\n"
     ]
    }
   ],
   "source": [
    "# Fixed fan-out: each presynaptic neuron connects to two postsynaptic neurons.\n",
    "fixed_pre_indices = jnp.array([\n",
    "    [0, 2],\n",
    "    [1, 2],\n",
    "    [0, 1],\n",
    "    [1, 2],\n",
    "    [0, 2],\n",
    "], dtype=jnp.int32)\n",
    "fixed_pre_weights = jnp.array([\n",
    "    [0.8, 0.2],\n",
    "    [0.7, 0.1],\n",
    "    [0.4, 0.3],\n",
    "    [0.5, 0.6],\n",
    "    [0.2, 0.5],\n",
    "], dtype=jnp.float32)\n",
    "\n",
    "fixed_pre_w = brainevent.FixedNumPerPre(\n",
    "    (fixed_pre_weights, fixed_pre_indices),\n",
    "    shape=(5, 3),\n",
    ")\n",
    "\n",
    "fixed_pre_mv = events @ fixed_pre_w\n",
    "fixed_pre_mm = event_block @ fixed_pre_w\n",
    "\n",
    "print(\"FixedNumPerPre: fixed fan-out from each presynaptic neuron\")\n",
    "print(f\"  shape = {fixed_pre_w.shape}\")\n",
    "print(f\"  data shape = {fixed_pre_w.data.shape}\")\n",
    "print(f\"  indices shape = {fixed_pre_w.indices.shape}\")\n",
    "print(\"  dense view:\")\n",
    "print(fixed_pre_w.todense())\n",
    "print(f\"  mv result = {fixed_pre_mv}\")\n",
    "print(f\"  mm result shape = {fixed_pre_mm.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "01cb1437",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:25.477596Z",
     "iopub.status.busy": "2026-06-22T01:52:25.477596Z",
     "iopub.status.idle": "2026-06-22T01:52:26.621937Z",
     "shell.execute_reply": "2026-06-22T01:52:26.621937Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FixedNumPerPost: fixed fan-in to each postsynaptic neuron\n",
      "  shape = (5, 3)\n",
      "  data shape = (3, 2)\n",
      "  indices shape = (3, 2)\n",
      "  dense view:\n",
      "[[0.8 0.  0. ]\n",
      " [0.  0.7 0. ]\n",
      " [0.4 0.  0.9]\n",
      " [0.  0.5 0. ]\n",
      " [0.  0.  0.5]]\n",
      "  mv result = [1.2 0.  1.4]\n"
     ]
    }
   ],
   "source": [
    "# Fixed fan-in: each postsynaptic neuron receives from two presynaptic neurons.\n",
    "fixed_post_indices = jnp.array([\n",
    "    [0, 2],\n",
    "    [1, 3],\n",
    "    [2, 4],\n",
    "], dtype=jnp.int32)\n",
    "fixed_post_weights = jnp.array([\n",
    "    [0.8, 0.4],\n",
    "    [0.7, 0.5],\n",
    "    [0.9, 0.5],\n",
    "], dtype=jnp.float32)\n",
    "\n",
    "fixed_post_w = brainevent.FixedNumPerPost(\n",
    "    (fixed_post_weights, fixed_post_indices),\n",
    "    shape=(5, 3),\n",
    ")\n",
    "\n",
    "print(\"FixedNumPerPost: fixed fan-in to each postsynaptic neuron\")\n",
    "print(f\"  shape = {fixed_post_w.shape}\")\n",
    "print(f\"  data shape = {fixed_post_w.data.shape}\")\n",
    "print(f\"  indices shape = {fixed_post_w.indices.shape}\")\n",
    "print(\"  dense view:\")\n",
    "print(fixed_post_w.todense())\n",
    "print(f\"  mv result = {events @ fixed_post_w}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c102bc71",
   "metadata": {},
   "source": [
    "Fixed fan-out and fixed fan-in structures are both multiplied with the same `@` syntax. Choose the orientation that matches the model constraint you care about: fixed outgoing degree from each presynaptic neuron, or fixed incoming degree to each postsynaptic neuron.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1698579",
   "metadata": {},
   "source": [
    "## 9. JIT compilation\n",
    "\n",
    "`brainevent` is designed for JAX workflows. The usual pattern is to construct connectivity once, then compile a function that receives changing spike values.\n",
    "\n",
    "Avoid rebuilding `CSR`, fixed, or JIT connectivity objects inside the compiled function. Those objects describe the graph; the spike vector is the dynamic input that changes each step.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "25181bff",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:26.623177Z",
     "iopub.status.busy": "2026-06-22T01:52:26.623177Z",
     "iopub.status.idle": "2026-06-22T01:52:26.809407Z",
     "shell.execute_reply": "2026-06-22T01:52:26.808904Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JIT-compiled one-step forward pass\n",
      "[1.4000001 0.3       1.6      ]\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def forward_one_step(spike_values):\n",
    "    spike_events = brainevent.BinaryArray(spike_values)\n",
    "    return spike_events @ csr_w\n",
    "\n",
    "compiled_output = forward_one_step(spikes_bool)\n",
    "print(\"JIT-compiled one-step forward pass\")\n",
    "print(compiled_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e77d4f2",
   "metadata": {},
   "source": [
    "## 10. Choosing a representation\n",
    "\n",
    "| Representation | What it stores | Best for | Typical operation |\n",
    "| --- | --- | --- | --- |\n",
    "| Dense/general matrix | every weight | small examples, debugging, dense layers | `BinaryArray @ dense_w` |\n",
    "| `CSR` | non-zero weights by presynaptic row | known sparse forward connectivity | `BinaryArray @ csr_w` |\n",
    "| `CSC` | non-zero weights by postsynaptic column | column-oriented algorithms and conversions | inspect or convert by target |\n",
    "| `JITCScalarR` / `JITCNormalR` / `JITCUniformR` | random rule and seed | very large random networks | `BinaryArray @ jit_w` |\n",
    "| `FixedNumPerPre` | fixed fan-out targets per presynaptic neuron | biologically constrained outgoing degree | `BinaryArray @ fixed_pre_w` |\n",
    "| `FixedNumPerPost` | fixed fan-in sources per postsynaptic neuron | biologically constrained incoming degree | `BinaryArray @ fixed_post_w` |\n",
    "\n",
    "A useful workflow is to start with dense matrices while checking the math, switch to `CSR` when the graph is known and sparse, use JIT connectivity when random connectivity is too large to store, and use fixed structures when degree constraints are part of the model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6236a45",
   "metadata": {},
   "source": [
    "## 11. Mini example: event-driven layer\n",
    "\n",
    "The final example wraps the same ideas into a small layer. It accepts binary presynaptic activity, converts it to `BinaryArray`, and multiplies it by whichever connectivity representation you pass in.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e77b963f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-22T01:52:26.811418Z",
     "iopub.status.busy": "2026-06-22T01:52:26.810418Z",
     "iopub.status.idle": "2026-06-22T01:52:26.850941Z",
     "shell.execute_reply": "2026-06-22T01:52:26.850941Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dense\n",
      "  current = [1.4000001 0.3       1.6      ]\n",
      "  output spikes = [ True False  True]\n",
      "CSR\n",
      "  current = [1.4000001 0.3       1.6      ]\n",
      "  output spikes = [ True False  True]\n",
      "FixedNumPerPre\n",
      "  current = [1.4000001 0.3       0.7      ]\n",
      "  output spikes = [ True False  True]\n",
      "JITCScalarR\n",
      "  current = [0.4 0.  0. ]\n",
      "  output spikes = [False False False]\n"
     ]
    }
   ],
   "source": [
    "def event_layer(spike_values, weights, threshold=0.5):\n",
    "    input_events = brainevent.BinaryArray(spike_values)\n",
    "    current = input_events @ weights\n",
    "    output_events = brainevent.BinaryArray(current > threshold)\n",
    "    return current, output_events\n",
    "\n",
    "for name, weights in [\n",
    "    (\"dense\", dense_w),\n",
    "    (\"CSR\", csr_w),\n",
    "    (\"FixedNumPerPre\", fixed_pre_w),\n",
    "    (\"JITCScalarR\", jit_w),\n",
    "]:\n",
    "    current, output_events = event_layer(spikes_bool, weights)\n",
    "    print(f\"{name}\")\n",
    "    print(f\"  current = {current}\")\n",
    "    print(f\"  output spikes = {output_events.value}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}