diff --git a/README.md b/README.md index 9d9cadb4..d7c14383 100644 --- a/README.md +++ b/README.md @@ -18,26 +18,41 @@ - elegant mechanisms for parameter sharing -### Tutorials +## Getting started -Tutorials are available [on our website](https://jaxleyverse.github.io/jaxley/). We currently have tutorials on how to: -- [simulate morphologically detailed neurons](https://jaxleyverse.github.io/jaxley/tutorial/01_morph_neurons/) -- [simulate networks of such neurons](https://jaxleyverse.github.io/jaxley/tutorial/02_small_network/) -- [set parameters of cells and networks](https://jaxleyverse.github.io/jaxley/tutorial/03_setting_parameters/) -- [speed up simulations with jit and vmap](https://jaxleyverse.github.io/jaxley/tutorial/04_jit_and_vmap/) -- [define your own channels and synapses](https://jaxleyverse.github.io/jaxley/tutorial/05_channel_and_synapse_models/) -- [define groups](https://jaxleyverse.github.io/jaxley/tutorial/06_groups/) -- [read and handle SWC files](https://jaxleyverse.github.io/jaxley/tutorial/08_importing_morphologies/) -- [train biophysical models](https://jaxleyverse.github.io/jaxley/tutorial/07_gradient_descent/) +`Jaxley` allows to simulate biophysical neuron models on CPU, GPU, or TPU: +```python +import matplotlib.pyplot as plt +from jax import config + +import jaxley as jx +from jaxley.channels import HH + +config.update("jax_platform_name", "cpu") # Or "gpu" / "tpu". +cell = jx.Cell() # Define cell. +cell.insert(HH()) # Insert channels. -### Units +current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0) +cell.stimulate(current) # Stimulate with step current. +cell.record("v") # Record voltage. -`Jaxley` uses the same [units as `NEURON`](https://www.neuron.yale.edu/neuron/static/docs/units/unitchart.html). +v = jx.integrate(cell) # Run simulation. +plt.plot(v.T) # Plot voltage trace. +``` + +If you want to learn more, we recommend you to check out our tutorials on how to: +- [get started with `Jaxley`](https://jaxleyverse.github.io/jaxley/tutorial/01_morph_neurons/) +- [simulate networks of neurons](https://jaxleyverse.github.io/jaxley/tutorial/02_small_network/) +- [speed up simulations with GPUs and `jit`](https://jaxleyverse.github.io/jaxley/tutorial/04_jit_and_vmap/) +- [define your own channels and synapses](https://jaxleyverse.github.io/jaxley/tutorial/05_channel_and_synapse_models/) +- [compute the gradient and train biophysical models](https://jaxleyverse.github.io/jaxley/tutorial/07_gradient_descent/) + + +## Installation -### Installation `Jaxley` is available on [`PyPI`](https://pypi.org/project/jaxley/): ```sh pip install jaxley @@ -48,22 +63,17 @@ pip install -U "jax[cuda12]" ``` -### Feedback and Contributions +## Feedback and Contributions We welcome any feedback on how Jaxley is working for your neuron models and are happy to receive bug reports, pull requests and other feedback (see [contribute](https://github.com/jaxleyverse/jaxley/blob/main/CONTRIBUTING.md)). We wish to maintain a positive community, please read our [Code of Conduct](https://github.com/jaxleyverse/jaxley/blob/main/CODE_OF_CONDUCT.md). -### Acknowledgements - -We greatly benefited from previous toolboxes for simulating multicompartment neurons, in particular [NEURON](https://github.com/neuronsimulator/nrn). - - -### License +## License -[Apache License](https://github.com/jaxleyverse/jaxley/blob/main/LICENSE) +[Apache License Version 2.0 (Apache-2.0)](https://github.com/jaxleyverse/jaxley/blob/main/LICENSE) -### Citation +## Citation If you use `Jaxley`, consider citing the [corresponding paper](https://www.biorxiv.org/content/10.1101/2024.08.21.608979): diff --git a/docs/docs/credits.md b/docs/docs/credits.md index 8db452d5..a1c639cf 100644 --- a/docs/docs/credits.md +++ b/docs/docs/credits.md @@ -4,11 +4,16 @@ ## License -`Jaxley` is licensed under the [Apache License](https://github.com/jaxleyverse/jaxley/blob/main/LICENSE) and +`Jaxley` is licensed under the [Apache License Version 2.0 (Apache-2.0)](https://github.com/jaxleyverse/jaxley/blob/main/LICENSE) and > Copyright (C) 2024 Michael Deistler, Jakob H. Macke, Pedro J. Goncalves, Philipp Berens. ## Important dependencies and prior art -* We greatly benefited from previous toolboxes for simulating multicompartment neurons, in particular [NEURON](https://github.com/neuronsimulator/nrn) +* We greatly benefited from previous toolboxes for simulating multicompartment neurons, in particular [NEURON](https://github.com/neuronsimulator/nrn). + + +## Funding + +This work was supported by the German Research Foundation (DFG) through Germany's Excellence Strategy (EXC 2064 – Project number 390727645) and the CRC 1233 "Robust Vision", the German Federal Ministry of Education and Research (Tübingen AI Center, FKZ: 01IS18039A), the 'Certification and Foundations of Safe Machine Learning Systems in Healthcare' project funded by the Carl Zeiss Foundation, and the European Union (ERC, "DeepCoMechTome", ref. 101089288, "NextMechMod", ref. 101039115). \ No newline at end of file diff --git a/docs/docs/faq.md b/docs/docs/faq.md index e8a65c72..6506eb82 100644 --- a/docs/docs/faq.md +++ b/docs/docs/faq.md @@ -1,3 +1,8 @@ # Frequently asked questions -[How can I save and load cells and networks?](faq/question_01.md) +- [What units does `Jaxley` use?](faq/question_01.md) +- [How can I save and load cells and networks?](faq/question_02.md) + +See also the [discussion page](https://github.com/jaxleyverse/jaxley/discussions) and the [issue +tracker](https://github.com/jaxleyverse/jaxley/issues) on the `Jaxley` GitHub repository for +recent questions and problems. \ No newline at end of file diff --git a/docs/docs/faq/question_01.md b/docs/docs/faq/question_01.md index 6bd20d8f..94e98164 100644 --- a/docs/docs/faq/question_01.md +++ b/docs/docs/faq/question_01.md @@ -1,18 +1,3 @@ -# How can I save and load cells and networks? +# What units does `Jaxley` use? -All `module`s (i.e., compartments, branches, cells, and networks) in `Jaxley` can be saved and loaded with pickle: -```python -import jaxley as jx -import pickle - -# ... define network, cell, etc. -network = jx.Network([cell1, cell2]) - -# Save. -with open("path/to/file.pkl", "wb") as handle: - pickle.dump(network, handle) - -# Load. -with open("path/to/file.pkl", "rb") as handle: - network = pickle.load(handle) -``` \ No newline at end of file +`Jaxley` uses the same units as the `NEURON` simulator, which are listed [here](https://www.neuron.yale.edu/neuron/static/docs/units/unitchart.html). diff --git a/docs/docs/faq/question_02.md b/docs/docs/faq/question_02.md new file mode 100644 index 00000000..6bd20d8f --- /dev/null +++ b/docs/docs/faq/question_02.md @@ -0,0 +1,18 @@ +# How can I save and load cells and networks? + +All `module`s (i.e., compartments, branches, cells, and networks) in `Jaxley` can be saved and loaded with pickle: +```python +import jaxley as jx +import pickle + +# ... define network, cell, etc. +network = jx.Network([cell1, cell2]) + +# Save. +with open("path/to/file.pkl", "wb") as handle: + pickle.dump(network, handle) + +# Load. +with open("path/to/file.pkl", "rb") as handle: + network = pickle.load(handle) +``` \ No newline at end of file diff --git a/docs/docs/index.md b/docs/docs/index.md index 87cc9703..ae59df38 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -12,26 +12,43 @@ - elegant mechanisms for parameter sharing -## Tutorials +## Getting started -We currently have tutorials on how to: +`Jaxley` allows to simulate biophysical neuron models on CPU, GPU, or TPU: +```python +import matplotlib.pyplot as plt +from jax import config + +import jaxley as jx +from jaxley.channels import HH + +config.update("jax_platform_name", "cpu") # Or "gpu" / "tpu". + +cell = jx.Cell() # Define cell. +cell.insert(HH()) # Insert channels. + +current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0) +cell.stimulate(current) # Stimulate with step current. +cell.record("v") # Record voltage. + +v = jx.integrate(cell) # Run simulation. +plt.plot(v.T) # Plot voltage trace. +``` + +If you want to learn more, we have tutorials on how to: - [simulate morphologically detailed neurons](https://jaxleyverse.github.io/jaxley/tutorial/01_morph_neurons/) - [simulate networks of such neurons](https://jaxleyverse.github.io/jaxley/tutorial/02_small_network/) - [set parameters of cells and networks](https://jaxleyverse.github.io/jaxley/tutorial/03_setting_parameters/) -- [speed up simulations with jit and vmap](https://jaxleyverse.github.io/jaxley/tutorial/04_jit_and_vmap/) +- [speed up simulations with GPUs and jit](https://jaxleyverse.github.io/jaxley/tutorial/04_jit_and_vmap/) - [define your own channels and synapses](https://jaxleyverse.github.io/jaxley/tutorial/05_channel_and_synapse_models/) - [define groups](https://jaxleyverse.github.io/jaxley/tutorial/06_groups/) - [read and handle SWC files](https://jaxleyverse.github.io/jaxley/tutorial/08_importing_morphologies/) -- [train biophysical models](https://jaxleyverse.github.io/jaxley/tutorial/07_gradient_descent/) - - -## Units - -`Jaxley` uses the same [units as `NEURON`](https://www.neuron.yale.edu/neuron/static/docs/units/unitchart.html). +- [compute the gradient and train biophysical models](https://jaxleyverse.github.io/jaxley/tutorial/07_gradient_descent/) ## Installation + `Jaxley` is available on [`pypi`](https://pypi.org/project/jaxley/): ```sh pip install jaxley @@ -47,17 +64,12 @@ pip install -U "jax[cuda12]" We welcome any feedback on how `Jaxley` is working for your neuron models and are happy to receive bug reports, pull requests and other feedback (see [contribute](https://github.com/jaxleyverse/jaxley/blob/main/CONTRIBUTING.md)). We wish to maintain a positive community, please read our [Code of Conduct](https://github.com/jaxleyverse/jaxley/blob/main/CODE_OF_CONDUCT.md). -## Acknowledgements - -We greatly benefited from previous toolboxes for simulating multicompartment neurons, in particular [NEURON](https://github.com/neuronsimulator/nrn). - - ## License -[Apache License](https://github.com/jaxleyverse/jaxley/blob/main/LICENSE) +[Apache License Version 2.0 (Apache-2.0)](https://github.com/jaxleyverse/jaxley/blob/main/LICENSE) -### Citation +## Citation If you use `Jaxley`, consider citing the [corresponding paper](https://www.biorxiv.org/content/10.1101/2024.08.21.608979): diff --git a/tutorials/01_morph_neurons.ipynb b/tutorials/01_morph_neurons.ipynb index 363fc743..759c6812 100644 --- a/tutorials/01_morph_neurons.ipynb +++ b/tutorials/01_morph_neurons.ipynb @@ -28,7 +28,7 @@ "import matplotlib.pyplot as plt\n", "\n", "\n", - "# Build the network.\n", + "# Build the cell.\n", "comp = jx.Compartment()\n", "branch = jx.Branch(comp, nseg=4)\n", "cell = jx.Cell(branch, parents=[-1, 0, 0, 1, 1])\n", @@ -44,14 +44,14 @@ "cell.vis(ax=ax)\n", "\n", "# Stimulate.\n", - "current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, dt=0.025, t_max=10.0)\n", + "current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0)\n", "cell.branch(0).loc(0.0).stimulate(current)\n", "\n", "# Record.\n", "cell.branch(0).loc(0.0).record(\"v\")\n", "\n", "# Simulate and plot.\n", - "v = jx.integrate(net)\n", + "v = jx.integrate(cell)\n", "plt.plot(v.T)\n", "```" ] @@ -247,7 +247,7 @@ "source": [ "### Stimulate the cell\n", "\n", - "We next stimulate one of the compartments with a step current. For this, we first define the step current:" + "We next stimulate one of the compartments with a step current. For this, we first define the step current (all units are the same as for the `NEURON` simulator, which are listed [here](https://www.neuron.yale.edu/neuron/static/docs/units/unitchart.html)):" ] }, { diff --git a/tutorials/03_setting_parameters.ipynb b/tutorials/03_setting_parameters.ipynb index 434cf8fc..4e3be8c8 100644 --- a/tutorials/03_setting_parameters.ipynb +++ b/tutorials/03_setting_parameters.ipynb @@ -46,7 +46,7 @@ }, { "cell_type": "code", - "execution_count": 32, + "execution_count": 1, "id": "deb594f4", "metadata": {}, "outputs": [], @@ -72,7 +72,7 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": 2, "id": "9819bbf4", "metadata": {}, "outputs": [], @@ -105,7 +105,7 @@ "cell.branch(0).comp(0).set(\"radius\", 10.0)\n", "```\n", "\n", - "You can always inspect the current parameters by inspecting `cell.nodes`, which is a pandas Dataframe that contains all information about the cell. You can use `.set()` to set morphological parameters, channel parameters, synaptic parameters, and initial states, as outlined below:" + "You can always inspect the current parameters by inspecting `cell.nodes`, which is a pandas Dataframe that contains all information about the cell. You can use `.set()` to set morphological parameters, channel parameters, synaptic parameters, and initial states. Note that `Jaxley` uses the same units as the `NEURON` simulator, which are listed [here](https://www.neuron.yale.edu/neuron/static/docs/units/unitchart.html)." ] }, { @@ -119,12 +119,12 @@ "\n", "- `radius`: the radius of the (zylindrical) compartment (in micrometer) \n", "- `length`: the length of the zylindrical compartment (in micrometer) \n", - "- `axial_resistivity`: the resistivity of current flow between compartments (in ohm centimeter) " + "- `axial_resistivity`: the resistivity of current flow between compartments (in ohm centimeter)" ] }, { "cell_type": "code", - "execution_count": 34, + "execution_count": 3, "id": "5424211f-369d-4d82-82da-e7abc6d34c8d", "metadata": {}, "outputs": [], @@ -135,7 +135,7 @@ }, { "cell_type": "code", - "execution_count": 35, + "execution_count": 4, "id": "f9f052e9-22be-403f-9c44-73ba36d5802f", "metadata": {}, "outputs": [ @@ -233,7 +233,7 @@ "3 1.0 -70.0 " ] }, - "execution_count": 35, + "execution_count": 4, "metadata": {}, "output_type": "execute_result" } @@ -249,23 +249,23 @@ "source": [ "### Setting channel parameters\n", "\n", - "You can also modify channel parameters. Every parameter that should be modifiable has to be defined in `self.channel_params` of the channel." + "You can also modify channel parameters (again, units are listed [here](https://www.neuron.yale.edu/neuron/static/docs/units/unitchart.html)). Every parameter that should be modifiable has to be defined in `self.channel_params` of the channel." ] }, { "cell_type": "code", - "execution_count": 36, + "execution_count": 5, "id": "2614c220-fe75-47cf-96dd-6961da9ac746", "metadata": {}, "outputs": [], "source": [ "cell.insert(Na())\n", - "cell.branch(1).comp(0).set(\"Na_gNa\", 0.1)" + "cell.branch(1).comp(0).set(\"Na_gNa\", 0.1) # mS/cm^2" ] }, { "cell_type": "code", - "execution_count": 37, + "execution_count": 6, "id": "f5940f1d-202d-4a1b-b39c-10cf3c4143e5", "metadata": {}, "outputs": [ @@ -393,7 +393,7 @@ "3 1.0 -70.0 True 0.05 50.0 -60.0 0.2 0.2 " ] }, - "execution_count": 37, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } @@ -414,19 +414,10 @@ }, { "cell_type": "code", - "execution_count": 38, + "execution_count": 7, "id": "08375bd2-fde1-47b6-814c-77e1b1065990", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/michaeldeistler/Documents/phd/jaxley/jaxley/modules/base.py:1533: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n", - " self.pointer.edges = pd.concat(\n" - ] - } - ], + "outputs": [], "source": [ "from jaxley.synapses import IonotropicSynapse\n", "from jaxley.connect import fully_connect\n", @@ -436,7 +427,7 @@ "fully_connect(net.cell(0), net.cell(1), IonotropicSynapse())\n", "\n", "# Unlike for channels, you have to index into the synapse with `net.SynapseName`\n", - "net.IonotropicSynapse.set(\"IonotropicSynapse_gS\", 0.1)" + "net.IonotropicSynapse.set(\"IonotropicSynapse_gS\", 0.1) # nS" ] }, { @@ -449,7 +440,7 @@ }, { "cell_type": "code", - "execution_count": 39, + "execution_count": 8, "id": "7571ce65-ced9-43e4-a9d2-7fe0cfacb972", "metadata": {}, "outputs": [ @@ -475,10 +466,10 @@ " \n", " \n", " pre_locs\n", - " pre_branch_index\n", - " pre_cell_index\n", " post_locs\n", + " pre_branch_index\n", " post_branch_index\n", + " pre_cell_index\n", " post_cell_index\n", " type\n", " type_ind\n", @@ -496,10 +487,10 @@ " \n", " 0\n", " 0.25\n", - " 0\n", - " 0\n", " 0.25\n", + " 0\n", " 1\n", + " 0\n", " 1\n", " IonotropicSynapse\n", " 0\n", @@ -517,14 +508,14 @@ "" ], "text/plain": [ - " pre_locs pre_branch_index pre_cell_index post_locs post_branch_index \\\n", - "0 0.25 0 0 0.25 1 \n", + " pre_locs post_locs pre_branch_index post_branch_index pre_cell_index \\\n", + "0 0.25 0.25 0 1 0 \n", "\n", - " post_cell_index type type_ind global_pre_comp_index \\\n", - "0 1 IonotropicSynapse 0 0 \n", + " post_cell_index type type_ind global_pre_comp_index \\\n", + "0 1 IonotropicSynapse 0 0 \n", "\n", - " global_post_comp_index global_pre_branch_index global_post_branch_index \\\n", - "0 6 0 3 \n", + " global_post_comp_index global_pre_branch_index global_post_branch_index \\\n", + "0 6 0 3 \n", "\n", " IonotropicSynapse_gS IonotropicSynapse_e_syn IonotropicSynapse_k_minus \\\n", "0 0.1 0.0 0.025 \n", @@ -533,7 +524,7 @@ "0 0.2 " ] }, - "execution_count": 39, + "execution_count": 8, "metadata": {}, "output_type": "execute_result" } @@ -560,18 +551,18 @@ }, { "cell_type": "code", - "execution_count": 40, + "execution_count": 9, "id": "86544366-6790-4b0e-8799-251769a63b9e", "metadata": {}, "outputs": [], "source": [ - "net.set(\"v\", -72.0)\n", - "net.IonotropicSynapse.set(\"IonotropicSynapse_s\", 0.1)" + "net.set(\"v\", -72.0) # mV\n", + "net.IonotropicSynapse.set(\"IonotropicSynapse_s\", 0.1) # nS" ] }, { "cell_type": "code", - "execution_count": 41, + "execution_count": 10, "id": "f8509703-7aff-4836-979e-7e5d7f8bdaf5", "metadata": {}, "outputs": [ @@ -775,7 +766,7 @@ "7 1.0 -72.0 True 0.05 50.0 -60.0 0.2 0.2 " ] }, - "execution_count": 41, + "execution_count": 10, "metadata": {}, "output_type": "execute_result" } @@ -786,7 +777,7 @@ }, { "cell_type": "code", - "execution_count": 42, + "execution_count": 11, "id": "94564463-a5bb-443f-94e4-9bd46ebc9121", "metadata": {}, "outputs": [ @@ -812,10 +803,10 @@ " \n", " \n", " pre_locs\n", - " pre_branch_index\n", - " pre_cell_index\n", " post_locs\n", + " pre_branch_index\n", " post_branch_index\n", + " pre_cell_index\n", " post_cell_index\n", " type\n", " type_ind\n", @@ -833,10 +824,10 @@ " \n", " 0\n", " 0.25\n", - " 0\n", - " 0\n", " 0.25\n", + " 0\n", " 1\n", + " 0\n", " 1\n", " IonotropicSynapse\n", " 0\n", @@ -854,14 +845,14 @@ "" ], "text/plain": [ - " pre_locs pre_branch_index pre_cell_index post_locs post_branch_index \\\n", - "0 0.25 0 0 0.25 1 \n", + " pre_locs post_locs pre_branch_index post_branch_index pre_cell_index \\\n", + "0 0.25 0.25 0 1 0 \n", "\n", - " post_cell_index type type_ind global_pre_comp_index \\\n", - "0 1 IonotropicSynapse 0 0 \n", + " post_cell_index type type_ind global_pre_comp_index \\\n", + "0 1 IonotropicSynapse 0 0 \n", "\n", - " global_post_comp_index global_pre_branch_index global_post_branch_index \\\n", - "0 6 0 3 \n", + " global_post_comp_index global_pre_branch_index global_post_branch_index \\\n", + "0 6 0 3 \n", "\n", " IonotropicSynapse_gS IonotropicSynapse_e_syn IonotropicSynapse_k_minus \\\n", "0 0.1 0.0 0.025 \n", @@ -870,7 +861,7 @@ "0 0.1 " ] }, - "execution_count": 42, + "execution_count": 11, "metadata": {}, "output_type": "execute_result" }