{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e6a2cbed",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:07.827782Z",
     "iopub.status.busy": "2026-06-19T02:32:07.827532Z",
     "iopub.status.idle": "2026-06-19T02:32:13.200499Z",
     "shell.execute_reply": "2026-06-19T02:32:13.199437Z"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import brainmass\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "brainstate.environ.set(dt=0.1 * u.ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88a54c89",
   "metadata": {},
   "source": [
    "# Extending Noise Processes\n",
    "\n",
    "Guide to creating custom noise processes.\n",
    "\n",
    "## Noise Template\n",
    "\n",
    "### Stateless Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1fbfe674",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:13.203177Z",
     "iopub.status.busy": "2026-06-19T02:32:13.202732Z",
     "iopub.status.idle": "2026-06-19T02:32:13.208239Z",
     "shell.execute_reply": "2026-06-19T02:32:13.207090Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import brainstate\n",
    "\n",
    "class MyStatelessNoise(brainstate.nn.Dynamics):\n",
    "    \"\"\"Custom stateless noise (i.i.d. samples).\n",
    "\n",
    "    Args:\n",
    "        in_size: Shape of noise output\n",
    "        sigma: Noise amplitude\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_size, sigma=1.0):\n",
    "        super().__init__()\n",
    "        self.in_size = in_size\n",
    "        self.sigma = sigma\n",
    "\n",
    "    def init_state(self, batch_size=None):\n",
    "        # No state for stateless noise\n",
    "        pass\n",
    "\n",
    "    def update(self):\n",
    "        \"\"\"Generate noise sample.\"\"\"\n",
    "        key = jax.random.PRNGKey(0)  # Use proper key management\n",
    "        noise = jax.random.normal(key, self.in_size) * self.sigma\n",
    "        return noise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bfb9208",
   "metadata": {},
   "source": [
    "### Stateful Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1917567d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:13.210530Z",
     "iopub.status.busy": "2026-06-19T02:32:13.210319Z",
     "iopub.status.idle": "2026-06-19T02:32:13.215866Z",
     "shell.execute_reply": "2026-06-19T02:32:13.214785Z"
    }
   },
   "outputs": [],
   "source": [
    "class MyStatefulNoise(brainstate.nn.Dynamics):\n",
    "    \"\"\"Custom stateful noise process.\n",
    "\n",
    "    Args:\n",
    "        in_size: Shape of noise output\n",
    "        sigma: Noise amplitude\n",
    "        tau: Correlation time\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_size, sigma=1.0, tau=10.0):\n",
    "        super().__init__()\n",
    "        self.in_size = in_size\n",
    "        self.sigma = sigma\n",
    "        self.tau = tau\n",
    "\n",
    "        # Internal state\n",
    "        self.x = brainstate.HiddenState(brainstate.init.Constant(0.))\n",
    "\n",
    "    def init_state(self, batch_size=None):\n",
    "        shape = (batch_size, *self.in_size) if batch_size else self.in_size\n",
    "        self.x.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self):\n",
    "        x = self.x.value\n",
    "\n",
    "        # Dynamics: dx/dt = -x/tau + sigma*sqrt(2/tau)*xi\n",
    "        key = jax.random.PRNGKey(0)  # Use proper key\n",
    "        xi = jax.random.normal(key, x.shape)\n",
    "\n",
    "        dt = 1.0  # ms\n",
    "        dx = (-x / self.tau + self.sigma * jnp.sqrt(2 / self.tau) * xi) * dt\n",
    "\n",
    "        self.x.value = x + dx\n",
    "        return self.x.value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5efaa63",
   "metadata": {},
   "source": [
    "## Example: Exponential Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9e454660",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:13.217770Z",
     "iopub.status.busy": "2026-06-19T02:32:13.217560Z",
     "iopub.status.idle": "2026-06-19T02:32:13.222301Z",
     "shell.execute_reply": "2026-06-19T02:32:13.221327Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import brainstate\n",
    "\n",
    "class ExponentialNoise(brainstate.nn.Dynamics):\n",
    "    \"\"\"Exponentially distributed noise.\n",
    "\n",
    "    Args:\n",
    "        in_size: Output shape\n",
    "        rate: Rate parameter (lambda)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_size, rate=1.0):\n",
    "        super().__init__()\n",
    "        self.in_size = in_size\n",
    "        self.rate = rate\n",
    "\n",
    "    def init_state(self, batch_size=None):\n",
    "        pass\n",
    "\n",
    "    def update(self):\n",
    "        key = jax.random.PRNGKey(0)\n",
    "        return jax.random.exponential(key, shape=self.in_size) / self.rate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdc3824c",
   "metadata": {},
   "source": [
    "## Proper Key Management\n",
    "\n",
    "Use `brainstate` random key infrastructure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3aea0358",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:13.225257Z",
     "iopub.status.busy": "2026-06-19T02:32:13.225009Z",
     "iopub.status.idle": "2026-06-19T02:32:13.229765Z",
     "shell.execute_reply": "2026-06-19T02:32:13.228408Z"
    }
   },
   "outputs": [],
   "source": [
    "class ProperNoise(brainstate.nn.Dynamics):\n",
    "    def __init__(self, in_size, sigma=1.0):\n",
    "        super().__init__()\n",
    "        self.in_size = in_size\n",
    "        self.sigma = sigma\n",
    "\n",
    "    def init_state(self, batch_size=None):\n",
    "        pass\n",
    "\n",
    "    def update(self):\n",
    "        # Use brainstate's random key management\n",
    "        key = brainstate.random.split_key()\n",
    "        noise = jax.random.normal(key, self.in_size) * self.sigma\n",
    "        return noise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5898a36e",
   "metadata": {},
   "source": [
    "## Unit-Aware Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c78e9789",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:13.233114Z",
     "iopub.status.busy": "2026-06-19T02:32:13.232775Z",
     "iopub.status.idle": "2026-06-19T02:32:13.237721Z",
     "shell.execute_reply": "2026-06-19T02:32:13.236694Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainunit as u\n",
    "\n",
    "class UnitAwareNoise(brainstate.nn.Dynamics):\n",
    "    def __init__(self, in_size, sigma=1.0*u.Hz):\n",
    "        super().__init__()\n",
    "        self.in_size = in_size\n",
    "        self.sigma = sigma\n",
    "\n",
    "    def update(self):\n",
    "        key = brainstate.random.split_key()\n",
    "        noise_unitless = jax.random.normal(key, self.in_size)\n",
    "        return noise_unitless * self.sigma  # Preserves units"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b52d89",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5883a029",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:13.240607Z",
     "iopub.status.busy": "2026-06-19T02:32:13.240335Z",
     "iopub.status.idle": "2026-06-19T02:32:13.245413Z",
     "shell.execute_reply": "2026-06-19T02:32:13.244028Z"
    }
   },
   "outputs": [],
   "source": [
    "def test_my_noise():\n",
    "    noise = MyStatefulNoise(in_size=(10,), sigma=1.0, tau=10.0)\n",
    "    noise.init_all_states()\n",
    "\n",
    "    # Test output shape\n",
    "    sample = noise.update()\n",
    "    assert sample.shape == (10,)\n",
    "\n",
    "    # Test batch\n",
    "    noise.init_all_states(batch_size=32)\n",
    "    sample = noise.update()\n",
    "    assert sample.shape == (32, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57987b6",
   "metadata": {},
   "source": [
    "## See Also\n",
    "\n",
    "- {doc}`creating_models` - Custom model creation\n",
    "- {doc}`../reference/noise` - Noise API reference"
   ]
  }
 ],
 "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",
   "version": "3.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}