Network simulations in Jaxley

In this tutorial, you will learn how to:

• connect neurons into a network
• visualize networks
• use the .edges attribute to inspect and change synaptic parameters

Here is a code snippet which you will learn to understand in this tutorial:

import jaxley as jx
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import connect


# Define a network. `cell` is defined as in previous tutorial.
net = jx.Network([cell for _ in range(11)])

# Define synapses.
fully_connect(
    net.cell(range(10)),
    net.cell(10),
    IonotropicSynapse(),
)

# Change synaptic parameters.
net.select(edges=[0, 1]).set("IonotropicSynapse_gS", 0.1)  # nS

# Visualize the network.
net.compute_xyz()
fig, ax = plt.subplots(1, 1, figsize=(4, 4))
net.vis(ax=ax, detail="full", layers=[10, 1])  # or `detail="point"`.

In the previous tutorial, you learned how to build single cells with morphological detail, how to insert stimuli and recordings, and how to run a first simulation. In this tutorial, we will define networks of multiple cells and connect them with synapses. Let’s get started:

from jax import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import matplotlib.pyplot as plt
import numpy as np
import jax.numpy as jnp
from jax import jit

import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import fully_connect, connect

Define the network

First, we define a cell as you saw in the previous tutorial.

comp = jx.Compartment()
branch = jx.Branch(comp, ncomp=4)
cell = jx.Cell(branch, parents=[-1, 0, 0, 1, 1, 2, 2])

We can assemble multiple cells into a network by using jx.Network, which takes a list of jx.Cells. Here, we assemble 11 cells into a network:

num_cells = 11
net = jx.Network([cell for _ in range(num_cells)])

At this point, we can already visualize this network:

net.compute_xyz()
net.rotate(180)
fig, ax = plt.subplots(1, 1, figsize=(3, 6))
_ = net.vis(ax=ax, detail="full", layers=[10, 1], layer_kwargs={"within_layer_offset": 150, "between_layer_offset": 200})

Note: you can use move_to to have more control over the location of cells, e.g.: network.cell(i).move_to(x=0, y=200).

As you can see, the neurons are not connected yet. Let’s fix this by connecting neurons with synapses. We will build a network consisting of two layers: 10 neurons in the input layer and 1 neuron in the output layer.

We can use Jaxley’s fully_connect method to connect these layers:

pre = net.cell(range(10))
post = net.cell(10)
fully_connect(pre, post, IonotropicSynapse())

Let’s visualize this again:

fig, ax = plt.subplots(1, 1, figsize=(3, 6))
_ = net.vis(ax=ax, detail="full", layers=[10, 1], layer_kwargs={"within_layer_offset": 150, "between_layer_offset": 200})

As you can see, the full_connect method inserted one synapse (in blue) from every neuron in the first layer to the output neuron. The fully_connect method builds this synapse from the zero-eth compartment and zero-eth branch of the presynaptic neuron onto a random branch of the postsynaptic neuron. If you want more control over the pre- and post-synaptic branches, you can use the connect method:

pre = net.cell(0).branch(5).loc(1.0)
post = net.cell(10).branch(0).loc(0.0)
connect(pre, post, IonotropicSynapse())

fig, ax = plt.subplots(1, 1, figsize=(3, 6))
_ = net.vis(ax=ax, detail="full", layers=[10, 1], layer_kwargs={"within_layer_offset": 150, "between_layer_offset": 200})

Inspecting and changing synaptic parameters

You can inspect synaptic parameters via the .edges attribute:

net.edges

	global_edge_index	global_pre_comp_index	global_post_comp_index	type	type_ind	pre_locs	post_locs	IonotropicSynapse_gS	IonotropicSynapse_e_syn	IonotropicSynapse_k_minus	IonotropicSynapse_s	controlled_by_param
0	0	0	286	IonotropicSynapse	0	0.125	0.625	0.0001	0.0	0.025	0.2	0
1	1	28	298	IonotropicSynapse	0	0.125	0.625	0.0001	0.0	0.025	0.2	0
2	2	56	286	IonotropicSynapse	0	0.125	0.625	0.0001	0.0	0.025	0.2	0
3	3	84	295	IonotropicSynapse	0	0.125	0.875	0.0001	0.0	0.025	0.2	0
4	4	112	302	IonotropicSynapse	0	0.125	0.625	0.0001	0.0	0.025	0.2	0
5	5	140	288	IonotropicSynapse	0	0.125	0.125	0.0001	0.0	0.025	0.2	0
6	6	168	287	IonotropicSynapse	0	0.125	0.875	0.0001	0.0	0.025	0.2	0
7	7	196	305	IonotropicSynapse	0	0.125	0.375	0.0001	0.0	0.025	0.2	0
8	8	224	299	IonotropicSynapse	0	0.125	0.875	0.0001	0.0	0.025	0.2	0
9	9	252	284	IonotropicSynapse	0	0.125	0.125	0.0001	0.0	0.025	0.2	0
10	10	23	280	IonotropicSynapse	0	0.875	0.125	0.0001	0.0	0.025	0.2	0

To modify a parameter of all synapses you can again use .set():

net.set("IonotropicSynapse_gS", 0.0003)  # nS

To modify individual syanptic parameters, use the .select() method. Below, we change the values of the first two synapses:

net.select(edges=[0, 1]).set("IonotropicSynapse_gS", 0.0004)  # nS

For more details on how to flexibly set synaptic parameters (e.g., by cell type, or by pre-synaptic cell index,…), see this tutorial.

Stimulating, recording, and simulating the network

We will now set up a simulation of the network. This works exactly as it does for single neurons:

# Stimulus.
i_delay = 3.0  # ms
i_amp = 0.05  # nA
i_dur = 2.0  # ms

# Duration and step size.
dt = 0.025  # ms
t_max = 50.0  # ms

time_vec = jnp.arange(0.0, t_max + dt, dt)

As a simple example, we insert sodium, potassium, and leak into every compartment of every cell of the network.

net.insert(Na())
net.insert(K())
net.insert(Leak())

We stimulate every neuron in the input layer and record the voltage from the output neuron:

current = jx.step_current(i_delay, i_dur, i_amp, dt, t_max)
net.delete_stimuli()
for stim_ind in range(10):
    net.cell(stim_ind).branch(0).loc(0.0).stimulate(current)

net.delete_recordings()
net.cell(10).branch(0).loc(0.0).record()

Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 external_states. See `.externals` for details.
Added 1 recordings. See `.recordings` for details.

Finally, we can again run the network simulation and plot the result:

s = jx.integrate(net, delta_t=dt)

fig, ax = plt.subplots(1, 1, figsize=(4, 2))
_ = ax.plot(s.T)

That’s it! You now know how to simulate networks of morphologically detailed neurons. We recommend that you now have a look at how you can speed up your simulation. To learn more about handling synaptic parameters, we recommend to check out this tutorial.