{ "cells": [ { "cell_type": "markdown", "id": "install-title", "metadata": {}, "source": [ "# Installation\n", "\n", "`brainmass` is a JAX-native library for differentiable neural-mass modelling. This\n", "page covers installing it for CPU, GPU (CUDA 12 / 13) and TPU, the optional\n", "plotting extra, and a one-line check that the install works.\n", "\n", "## Pick your accelerator\n", "\n", "Install the package with the extra that matches your hardware. CPU is the simplest\n", "and is all you need to follow the {doc}`/getting_started/quickstart`.\n", "\n", "::::{tab-set}\n", ":::{tab-item} CPU\n", "```bash\n", "pip install -U brainmass[cpu]\n", "```\n", "Works everywhere (Linux, macOS incl. Apple Silicon, Windows). Recommended for\n", "getting started and for the notebooks in these docs.\n", ":::\n", "\n", ":::{tab-item} GPU — CUDA 12\n", "```bash\n", "pip install -U brainmass[cuda12]\n", "```\n", "For NVIDIA GPUs with a CUDA 12.x runtime. The matching `jaxlib` CUDA wheels are\n", "pulled in automatically.\n", ":::\n", "\n", ":::{tab-item} GPU — CUDA 13\n", "```bash\n", "pip install -U brainmass[cuda13]\n", "```\n", "For NVIDIA GPUs with a CUDA 13.x runtime.\n", ":::\n", "\n", ":::{tab-item} TPU\n", "```bash\n", "pip install -U brainmass[tpu]\n", "```\n", "For Google Cloud TPU VMs.\n", ":::\n", "::::\n", "\n", "## The `[viz]` plotting extra\n", "\n", "`brainmass` has **no hard dependency on matplotlib** — `import brainmass` works\n", "without it. The {mod}`brainmass.viz` helpers (used throughout the tutorials and\n", "the gallery) import matplotlib lazily, so install the **`[viz]` extra** to enable\n", "them:\n", "\n", "```bash\n", "pip install -U \"brainmass[cpu,viz]\" # CPU + plotting helpers\n", "```\n", "\n", "Combine extras freely, e.g. `brainmass[cuda12,viz]`. If you call a\n", "`brainmass.viz.*` function without matplotlib installed you get a clear error\n", "pointing you back to this extra." ] }, { "cell_type": "markdown", "id": "install-verify-md", "metadata": {}, "source": [ "## Verify the install\n", "\n", "Check the version and run one tiny simulation through the high-level\n", "{class}`~brainmass.Simulator` — if this prints a version and a trajectory shape,\n", "you are ready to go." ] }, { "cell_type": "code", "execution_count": 1, "id": "install-verify-code", "metadata": { "execution": { "iopub.execute_input": "2026-06-19T06:25:42.949950Z", "iopub.status.busy": "2026-06-19T06:25:42.949730Z", "iopub.status.idle": "2026-06-19T06:25:47.799259Z", "shell.execute_reply": "2026-06-19T06:25:47.798433Z" } }, "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" ] }, { "name": "stdout", "output_type": "stream", "text": [ "brainmass version: 0.0.6\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "trajectory shape (steps, regions): (50, 4)\n", "install OK\n" ] } ], "source": [ "import brainmass\n", "import brainunit as u\n", "\n", "print(f\"brainmass version: {brainmass.__version__}\")\n", "\n", "# One step-model, driven for 5 ms through the Simulator.\n", "node = brainmass.HopfStep(in_size=4, a=0.25)\n", "res = brainmass.Simulator(node, dt=0.1 * u.ms).run(5.0 * u.ms, monitors=[\"x\"])\n", "print(\"trajectory shape (steps, regions):\", res[\"x\"].shape)\n", "print(\"install OK\")" ] }, { "cell_type": "markdown", "id": "install-devices-md", "metadata": {}, "source": [ "To confirm which accelerator JAX is using:" ] }, { "cell_type": "code", "execution_count": 2, "id": "install-devices-code", "metadata": { "execution": { "iopub.execute_input": "2026-06-19T06:25:47.801873Z", "iopub.status.busy": "2026-06-19T06:25:47.801424Z", "iopub.status.idle": "2026-06-19T06:25:47.806571Z", "shell.execute_reply": "2026-06-19T06:25:47.805985Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[CpuDevice(id=0)]\n" ] } ], "source": [ "import jax\n", "\n", "print(jax.devices()) # CpuDevice on CPU; CudaDevice / TpuDevice with an accelerator" ] }, { "cell_type": "markdown", "id": "install-ecosystem", "metadata": {}, "source": [ "## The BrainX ecosystem\n", "\n", "`brainmass` is part of the **BrainX** ecosystem and builds directly on three of\n", "its libraries — they are installed automatically as dependencies:\n", "\n", "| Package | Role in brainmass |\n", "|---|---|\n", "| [`brainstate`](https://github.com/chaobrain/brainstate) | State management & the `transform` primitives (`jit`, `for_loop`, `grad`) the `Simulator`/`Fitter` are built on |\n", "| [`brainunit`](https://github.com/chaobrain/brainunit) | Unit-aware arrays (`u.ms`, `u.mm`, …) so quantities stay physically meaningful |\n", "| [`braintools`](https://github.com/chaobrain/braintools) | Initializers, optimizers (`braintools.optim`) and metrics (`braintools.metric`) surfaced by the `Fitter` and `viz` |\n", "\n", "You rarely import them all by hand, but knowing the split helps when reading\n", "error messages.\n", "\n", "## Install from source (development)\n", "\n", "```bash\n", "git clone https://github.com/chaobrain/brainmass.git\n", "cd brainmass\n", "pip install -e \".[cpu,viz,dev,doc]\" # editable install with test + doc tooling\n", "```\n", "\n", "## Next steps\n", "\n", "- {doc}`/getting_started/quickstart` — your first simulation, plot, network and fit in ~10 minutes.\n", "- {doc}`/getting_started/key_concepts` — the mental model behind `*Step` models, `Simulator`, `Network` and `Fitter`.\n", "- {doc}`/howto/choose_a_model` — browse the model catalogue with `brainmass.list_models()`.\n", "\n", "## See also\n", "\n", "- [JAX installation guide](https://docs.jax.dev/en/latest/installation.html)\n", "- {doc}`/faq` for troubleshooting\n", "- [GitHub issues](https://github.com/chaobrain/brainmass/issues)" ] } ], "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 }