{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7d0201f7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:46.750901Z",
     "iopub.status.busy": "2026-06-19T02:31:46.750583Z",
     "iopub.status.idle": "2026-06-19T02:31:52.089458Z",
     "shell.execute_reply": "2026-06-19T02:31:52.088479Z"
    },
    "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": "73185cba",
   "metadata": {},
   "source": [
    "# Creating Custom Models\n",
    "\n",
    "Guide to implementing custom neural mass models.\n",
    "\n",
    "## Model Template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3a8f622b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:52.092379Z",
     "iopub.status.busy": "2026-06-19T02:31:52.091835Z",
     "iopub.status.idle": "2026-06-19T02:31:52.097748Z",
     "shell.execute_reply": "2026-06-19T02:31:52.096889Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import brainunit as u\n",
    "import jax.numpy as jnp\n",
    "\n",
    "class MyCustomModel(brainstate.nn.Dynamics):\n",
    "    \"\"\"Custom neural mass model.\n",
    "\n",
    "    Args:\n",
    "        in_size: Number of regions/nodes\n",
    "        tau: Time constant (ms)\n",
    "        other_param: Description\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_size, tau=10.*u.ms, other_param=1.0):\n",
    "        super().__init__()\n",
    "\n",
    "        # Store parameters\n",
    "        self.in_size = in_size\n",
    "        self.tau = tau\n",
    "        self.other_param = other_param\n",
    "\n",
    "        # Initialize state variables\n",
    "        self.x = brainstate.HiddenState(\n",
    "            brainstate.init.Constant(0.),\n",
    "            sharding=brainstate.ShardingMode.REPLICATED\n",
    "        )\n",
    "\n",
    "    def init_state(self, batch_size=None):\n",
    "        \"\"\"Initialize state variables.\"\"\"\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, external_input=0.):\n",
    "        \"\"\"Update dynamics by one time step.\n",
    "\n",
    "        Args:\n",
    "            external_input: External driving input\n",
    "\n",
    "        Returns:\n",
    "            Observable output\n",
    "        \"\"\"\n",
    "        # Get current state\n",
    "        x = self.x.value\n",
    "\n",
    "        # Compute dynamics (example: dx/dt = -x/tau + input)\n",
    "        dx_dt = -x / self.tau + external_input\n",
    "\n",
    "        # Update state (Euler integration, dt assumed 1 ms)\n",
    "        dt = 1 * u.ms\n",
    "        self.x.value = x + dx_dt * dt\n",
    "\n",
    "        # Return observable\n",
    "        return self.x.value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c22b0d7",
   "metadata": {},
   "source": [
    "## Example: Custom Oscillator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e00eab8c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:52.099807Z",
     "iopub.status.busy": "2026-06-19T02:31:52.099545Z",
     "iopub.status.idle": "2026-06-19T02:31:52.105359Z",
     "shell.execute_reply": "2026-06-19T02:31:52.104745Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import brainunit as u\n",
    "import jax.numpy as jnp\n",
    "\n",
    "class DampedOscillator(brainstate.nn.Dynamics):\n",
    "    \"\"\"Damped harmonic oscillator.\n",
    "\n",
    "    Equations:\n",
    "        dx/dt = v\n",
    "        dv/dt = -omega^2 * x - 2*zeta*omega*v + F_ext\n",
    "\n",
    "    Args:\n",
    "        in_size: Number of oscillators\n",
    "        omega: Natural frequency (Hz)\n",
    "        zeta: Damping ratio (dimensionless)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_size, omega=10*u.Hz, zeta=0.1):\n",
    "        super().__init__()\n",
    "\n",
    "        self.in_size = in_size\n",
    "        self.omega = omega\n",
    "        self.zeta = zeta\n",
    "\n",
    "        # State variables: position and velocity\n",
    "        self.x = brainstate.HiddenState(brainstate.init.Constant(0.))\n",
    "        self.v = 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",
    "        self.v.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self, F_ext=0.):\n",
    "        x = self.x.value\n",
    "        v = self.v.value\n",
    "\n",
    "        # Dynamics\n",
    "        dx_dt = v\n",
    "        dv_dt = -(self.omega**2) * x - 2*self.zeta*self.omega*v + F_ext\n",
    "\n",
    "        # Euler integration\n",
    "        dt = 1 * u.ms\n",
    "        self.x.value = x + dx_dt * dt\n",
    "        self.v.value = v + dv_dt * dt\n",
    "\n",
    "        return self.x.value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9685558",
   "metadata": {},
   "source": [
    "## Adding Noise Support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "73d90cec",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:52.107380Z",
     "iopub.status.busy": "2026-06-19T02:31:52.107193Z",
     "iopub.status.idle": "2026-06-19T02:31:52.112172Z",
     "shell.execute_reply": "2026-06-19T02:31:52.111198Z"
    }
   },
   "outputs": [],
   "source": [
    "class MyModel(brainstate.nn.Dynamics):\n",
    "    def __init__(self, in_size, noise=None):\n",
    "        super().__init__()\n",
    "        self.in_size = in_size\n",
    "        self.noise = noise  # Optional noise source\n",
    "\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",
    "        # Initialize noise if present\n",
    "        if self.noise is not None:\n",
    "            self.noise.init_all_states(batch_size)\n",
    "\n",
    "    def update(self, external_input=0.):\n",
    "        x = self.x.value\n",
    "\n",
    "        # Dynamics\n",
    "        dx_dt = -x + external_input\n",
    "\n",
    "        # Add noise if present\n",
    "        if self.noise is not None:\n",
    "            noise_value = self.noise.update()\n",
    "            dx_dt += noise_value\n",
    "\n",
    "        # Integrate\n",
    "        dt = 1 * u.ms\n",
    "        self.x.value = x + dx_dt * dt\n",
    "\n",
    "        return self.x.value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf92fc3",
   "metadata": {},
   "source": [
    "## Unit Handling\n",
    "\n",
    "Ensure dimensional consistency:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "86dc2f86",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:52.114693Z",
     "iopub.status.busy": "2026-06-19T02:31:52.114449Z",
     "iopub.status.idle": "2026-06-19T02:31:52.118806Z",
     "shell.execute_reply": "2026-06-19T02:31:52.117795Z"
    }
   },
   "outputs": [],
   "source": [
    "import brainunit as u\n",
    "\n",
    "# Parameters with units\n",
    "tau = 10 * u.ms\n",
    "omega = 2 * jnp.pi * 10 * u.Hz\n",
    "\n",
    "# Computations preserve units\n",
    "frequency = 1 / tau  # Result: Hz\n",
    "period = 1 / omega  # Result: seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ded950a6",
   "metadata": {},
   "source": [
    "## Testing Your Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "13d705e2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:52.121103Z",
     "iopub.status.busy": "2026-06-19T02:31:52.120905Z",
     "iopub.status.idle": "2026-06-19T02:31:52.124500Z",
     "shell.execute_reply": "2026-06-19T02:31:52.123814Z"
    }
   },
   "outputs": [],
   "source": [
    "import pytest\n",
    "import jax.numpy as jnp\n",
    "\n",
    "def test_my_custom_model():\n",
    "    # Create model\n",
    "    model = MyCustomModel(in_size=10)\n",
    "    model.init_all_states()\n",
    "\n",
    "    # Test update\n",
    "    output = model.update(external_input=1.0)\n",
    "\n",
    "    # Check shape\n",
    "    assert output.shape == (10,)\n",
    "\n",
    "    # Test with batch\n",
    "    model.init_all_states(batch_size=32)\n",
    "    output = model.update()\n",
    "    assert output.shape == (32, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2096290c",
   "metadata": {},
   "source": [
    "## Documentation\n",
    "\n",
    "Add comprehensive docstrings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "98ed5fd9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:31:52.126636Z",
     "iopub.status.busy": "2026-06-19T02:31:52.126347Z",
     "iopub.status.idle": "2026-06-19T02:31:52.130679Z",
     "shell.execute_reply": "2026-06-19T02:31:52.129610Z"
    }
   },
   "outputs": [],
   "source": [
    "class MyModel(brainstate.nn.Dynamics):\n",
    "    \"\"\"One-line description.\n",
    "\n",
    "    Detailed explanation of the model, including:\n",
    "    - Mathematical equations\n",
    "    - Biological interpretation\n",
    "    - References to papers\n",
    "\n",
    "    Mathematical Details:\n",
    "\n",
    "    .. math::\n",
    "\n",
    "        \\\\frac{dx}{dt} = f(x, t)\n",
    "\n",
    "    Args:\n",
    "        in_size: Shape of input/number of regions\n",
    "        param1: Description with units\n",
    "        param2: Description with default behavior\n",
    "\n",
    "    Examples:\n",
    "\n",
    "        >>> model = MyModel(in_size=10)\n",
    "        >>> model.init_all_states()\n",
    "        >>> output = model.update()\n",
    "\n",
    "    References:\n",
    "        Author et al. (2020). Paper title. Journal.\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5197fb5c",
   "metadata": {},
   "source": [
    "## Contributing Your Model\n",
    "\n",
    "1. Implement model following template\n",
    "2. Add tests in `tests/`\n",
    "3. Create example notebook in `examples/`\n",
    "4. Update API documentation\n",
    "5. Submit pull request\n",
    "\n",
    "## See Also\n",
    "\n",
    "- {doc}`architecture` - Package structure\n",
    "- {doc}`extending_noise` - Custom noise processes\n",
    "- {doc}`contributing` - Contribution guidelines\n",
    "- `brainstate` documentation"
   ]
  }
 ],
 "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
}