Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

nseg->ncomp for swc reader; set assign_groups=True #509

Merged
merged 2 commits into from
Nov 21, 2024
+1,308 −2,451
Merged

nseg->ncomp for swc reader; set assign_groups=True #509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ed46ace
nseg->ncomp for swc reader; set assign_groups=True
michaeldeistler Nov 15, 2024
a7e6dd6
replace nseg with ncomp throughout the codebase
michaeldeistler Nov 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1,231 changes: 381 additions & 850 deletions docs/tutorials/00_jaxley_api.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

142 changes: 71 additions & 71 deletions docs/tutorials/01_morph_neurons.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

156 changes: 78 additions & 78 deletions docs/tutorials/02_small_network.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

72 changes: 36 additions & 36 deletions docs/tutorials/04_jit_and_vmap.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@
"cells": [
{
"cell_type": "markdown",
"id": "99e0512a",
"id": "cfd523b5",
"metadata": {},
"source": [
"# Speeding up simulations"
]
},
{
"cell_type": "markdown",
"id": "007bf12f",
"id": "adfd37cf",
"metadata": {},
"source": [
"In this tutorial, you will learn how to:\n",
Expand Down Expand Up @@ -47,7 +47,7 @@
},
{
"cell_type": "markdown",
"id": "b596f6ba",
"id": "757dcad9",
"metadata": {},
"source": [
"In the previous tutorials, you learned how to build single cells or networks and how to change their parameters. In this tutorial, you will learn how to speed up such simulations by many orders of magnitude. This can be achieved in to ways:\n",
Expand All @@ -60,15 +60,15 @@
},
{
"cell_type": "markdown",
"id": "7b1628bf",
"id": "c813d313",
"metadata": {},
"source": [
"### Using GPU or CPU"
]
},
{
"cell_type": "markdown",
"id": "163b1c49",
"id": "f69b53c7",
"metadata": {},
"source": [
"In `Jaxley` you can set whether you want to use `gpu` or `cpu` with the following lines at the beginning of your script:"
Expand All @@ -77,7 +77,7 @@
{
"cell_type": "code",
"execution_count": 1,
"id": "b9cc56b7",
"id": "2f080339",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -87,7 +87,7 @@
},
{
"cell_type": "markdown",
"id": "611633d6",
"id": "c38c665a",
"metadata": {},
"source": [
"`JAX` (and `Jaxley`) also allow to choose between `float32` and `float64`. Especially on GPUs, `float32` will be faster, but we have experienced stability issues when simulating morphologically detailed neurons with `float32`."
Expand All @@ -96,7 +96,7 @@
{
"cell_type": "code",
"execution_count": 2,
"id": "b7093b86",
"id": "86d4a917",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -105,7 +105,7 @@
},
{
"cell_type": "markdown",
"id": "3c3bd156",
"id": "dc16b92d",
"metadata": {},
"source": [
"Next, we will import relevant libraries:"
Expand All @@ -114,7 +114,7 @@
{
"cell_type": "code",
"execution_count": 3,
"id": "ad756d90",
"id": "bd054087",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -129,7 +129,7 @@
},
{
"cell_type": "markdown",
"id": "54be0755",
"id": "9d2ae1fa",
"metadata": {},
"source": [
"### Building the cell or network\n",
Expand All @@ -140,7 +140,7 @@
{
"cell_type": "code",
"execution_count": 4,
"id": "bc8c6f17",
"id": "a869e670",
"metadata": {},
"outputs": [
{
Expand All @@ -157,7 +157,7 @@
"t_max = 10.0\n",
"\n",
"comp = jx.Compartment()\n",
"branch = jx.Branch(comp, nseg=4)\n",
"branch = jx.Branch(comp, ncomp=4)\n",
"cell = jx.Cell(branch, parents=[-1, 0, 0, 1, 1, 2, 2])\n",
"\n",
"cell.insert(Na())\n",
Expand All @@ -174,7 +174,7 @@
},
{
"cell_type": "markdown",
"id": "39165d99",
"id": "d9193627",
"metadata": {},
"source": [
"### Parameter sweeps\n",
Expand All @@ -185,7 +185,7 @@
{
"cell_type": "code",
"execution_count": 5,
"id": "d5088abb",
"id": "79a01358",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -198,7 +198,7 @@
},
{
"cell_type": "markdown",
"id": "f815d90f",
"id": "2f8e301a",
"metadata": {},
"source": [
"The `.data_set()` method takes three arguments: \n",
Expand All @@ -210,7 +210,7 @@
},
{
"cell_type": "markdown",
"id": "7e319037",
"id": "a343e454",
"metadata": {},
"source": [
"Having done this, the simplest (but least efficient) way to perform the parameter sweep is to run a for-loop over many parameter sets:"
Expand All @@ -219,7 +219,7 @@
{
"cell_type": "code",
"execution_count": 6,
"id": "436d00a1",
"id": "4806598a",
"metadata": {},
"outputs": [
{
Expand All @@ -240,15 +240,15 @@
},
{
"cell_type": "markdown",
"id": "a24cc7bd",
"id": "e0f1becb",
"metadata": {},
"source": [
"The resulting voltages have shape `(num_simulations, num_recordings, num_timesteps)`."
]
},
{
"cell_type": "markdown",
"id": "cbb72c8b",
"id": "c4345c02",
"metadata": {},
"source": [
"### Stimulus sweeps\n",
Expand All @@ -266,7 +266,7 @@
},
{
"cell_type": "markdown",
"id": "fdb15f95",
"id": "5dd3c975",
"metadata": {},
"source": [
"### Speeding up for loops via `jit` compilation\n",
Expand All @@ -277,7 +277,7 @@
{
"cell_type": "code",
"execution_count": 7,
"id": "83c96a95",
"id": "017e98d9",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -287,7 +287,7 @@
{
"cell_type": "code",
"execution_count": 8,
"id": "0419f5ec",
"id": "d9aa805a",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -298,7 +298,7 @@
{
"cell_type": "code",
"execution_count": 9,
"id": "d197f6de",
"id": "27c12fe3",
"metadata": {},
"outputs": [
{
Expand All @@ -317,15 +317,15 @@
},
{
"cell_type": "markdown",
"id": "10946557",
"id": "401d1f52",
"metadata": {},
"source": [
"`jit` compilation can be up to 10k times faster, especially for small simulations with few compartments. For very large models, the gain obtained with `jit` will be much smaller (`jit` may even provide no speed up at all)."
]
},
{
"cell_type": "markdown",
"id": "9417a2c3",
"id": "d29ff570",
"metadata": {},
"source": [
"### Speeding up with GPU parallelization via `vmap`\n",
Expand All @@ -336,7 +336,7 @@
{
"cell_type": "code",
"execution_count": 10,
"id": "af123748",
"id": "fefffaf7",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -346,7 +346,7 @@
},
{
"cell_type": "markdown",
"id": "86df1175",
"id": "fd03669d",
"metadata": {},
"source": [
"We can then run this method on __all__ parameter sets (`all_params.shape == (100, 2)`), and `Jaxley` will automatically parallelize across them. Of course, you will only get a speed-up if you have a GPU available and you specified `gpu` as device in the beginning of this tutorial."
Expand All @@ -355,7 +355,7 @@
{
"cell_type": "code",
"execution_count": 11,
"id": "dfb1977e",
"id": "c2a22648",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -364,23 +364,23 @@
},
{
"cell_type": "markdown",
"id": "96788177",
"id": "a4464e06",
"metadata": {},
"source": [
"GPU parallelization with `vmap` can give a large speed-up, which can easily be 2-3 orders of magnitude."
]
},
{
"cell_type": "markdown",
"id": "fd4a4d1c",
"id": "0df64cc1",
"metadata": {},
"source": [
"### Combining `jit` and `vmap`"
]
},
{
"cell_type": "markdown",
"id": "13ea4636",
"id": "8125f061",
"metadata": {},
"source": [
"Finally, you can also combine using `jit` and `vmap`. For example, you can run multiple batches of many parallel simulations. Each batch can be parallelized with `vmap` and simulating each batch can be compiled with `jit`:"
Expand All @@ -389,7 +389,7 @@
{
"cell_type": "code",
"execution_count": 12,
"id": "156c4333",
"id": "db1eced1",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -399,7 +399,7 @@
{
"cell_type": "code",
"execution_count": 13,
"id": "5ef62fd1",
"id": "82f34a7d",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -410,15 +410,15 @@
},
{
"cell_type": "markdown",
"id": "1c1e2a1d",
"id": "a5cca5a0",
"metadata": {},
"source": [
"That's all you have to know about `jit` and `vmap`! If you have worked through this and the previous tutorials, you should be ready to set up your first network simulations."
]
},
{
"cell_type": "markdown",
"id": "31f43b9b",
"id": "37fc2f3c",
"metadata": {},
"source": [
"### Next steps\n",
Expand Down
Loading
Loading