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

* nseg->ncomp for swc reader; set assign_groups=True * replace nseg with ncomp throughout the codebase
jaxleyverse · Nov 21, 2024 · 6e79424 · 6e79424
1 parent 7e71f66
commit 6e79424
Show file tree

Hide file tree

Showing 39 changed files with 1,308 additions and 2,451 deletions.
diff --git a/docs/tutorials/00_jaxley_api.ipynb b/docs/tutorials/00_jaxley_api.ipynb
diff --git a/docs/tutorials/01_morph_neurons.ipynb b/docs/tutorials/01_morph_neurons.ipynb
diff --git a/docs/tutorials/02_small_network.ipynb b/docs/tutorials/02_small_network.ipynb
diff --git a/docs/tutorials/04_jit_and_vmap.ipynb b/docs/tutorials/04_jit_and_vmap.ipynb
@@ -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",
@@ -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",
@@ -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:"
@@ -77,7 +77,7 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "b9cc56b7",
+   "id": "2f080339",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -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`."
@@ -96,7 +96,7 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "b7093b86",
+   "id": "86d4a917",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,7 +105,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3c3bd156",
+   "id": "dc16b92d",
    "metadata": {},
    "source": [
     "Next, we will import relevant libraries:"
@@ -114,7 +114,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "ad756d90",
+   "id": "bd054087",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -129,7 +129,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "54be0755",
+   "id": "9d2ae1fa",
    "metadata": {},
    "source": [
     "### Building the cell or network\n",
@@ -140,7 +140,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "bc8c6f17",
+   "id": "a869e670",
    "metadata": {},
    "outputs": [
     {
@@ -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",
@@ -174,7 +174,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "39165d99",
+   "id": "d9193627",
    "metadata": {},
    "source": [
     "### Parameter sweeps\n",
@@ -185,7 +185,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "d5088abb",
+   "id": "79a01358",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -198,7 +198,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f815d90f",
+   "id": "2f8e301a",
    "metadata": {},
    "source": [
     "The `.data_set()` method takes three arguments: \n",
@@ -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:"
@@ -219,7 +219,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "436d00a1",
+   "id": "4806598a",
    "metadata": {},
    "outputs": [
     {
@@ -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",
@@ -266,7 +266,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fdb15f95",
+   "id": "5dd3c975",
    "metadata": {},
    "source": [
     "### Speeding up for loops via `jit` compilation\n",
@@ -277,7 +277,7 @@
   {
    "cell_type": "code",
    "execution_count": 7,
-   "id": "83c96a95",
+   "id": "017e98d9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -287,7 +287,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "id": "0419f5ec",
+   "id": "d9aa805a",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -298,7 +298,7 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "id": "d197f6de",
+   "id": "27c12fe3",
    "metadata": {},
    "outputs": [
     {
@@ -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",
@@ -336,7 +336,7 @@
   {
    "cell_type": "code",
    "execution_count": 10,
-   "id": "af123748",
+   "id": "fefffaf7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -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."
@@ -355,7 +355,7 @@
   {
    "cell_type": "code",
    "execution_count": 11,
-   "id": "dfb1977e",
+   "id": "c2a22648",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -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`:"
@@ -389,7 +389,7 @@
   {
    "cell_type": "code",
    "execution_count": 12,
-   "id": "156c4333",
+   "id": "db1eced1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -399,7 +399,7 @@
   {
    "cell_type": "code",
    "execution_count": 13,
-   "id": "5ef62fd1",
+   "id": "82f34a7d",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -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",