Skip to content

Commit 1b22675

Browse files
maedocmaedoc
authored
Merge pull request #91 from ins-amu/doc/multi-model
doc: add heterogeneous network example
2 parents 488c4c6 + cd7b60f commit 1b22675

2 files changed

Lines changed: 274 additions & 0 deletions

File tree

‎examples/heterogeneous_network.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
"""
2+
**Example**: A network with multiple neural mass models.
3+
4+
This example demonstrates how to simulate a network with different neural mass
5+
models assigned to different nodes. We'll create a network with some nodes
6+
using the Jansen-Rit (JR) model and others using the Montbrio-Pazo-Roxin (MPR)
7+
model.
8+
9+
"""
10+
import os
11+
import jax.numpy as jp
12+
import vbjax as vb
13+
import collections
14+
import matplotlib.pyplot as plt
15+
16+
# Create a directory for saving images
17+
os.makedirs('images', exist_ok=True)
18+
19+
# 1. Define the network structure
20+
n_jr_nodes = 4
21+
n_mpr_nodes = 2
22+
n_total_nodes = n_jr_nodes + n_mpr_nodes
23+
24+
# Define a connectome: JR nodes are reciprocally connected, MPR nodes too,
25+
# and there are connections from JR to MPR nodes.
26+
connectivity = jp.zeros((n_total_nodes, n_total_nodes))
27+
connectivity = connectivity.at[:n_jr_nodes, :n_jr_nodes].set(0.2)
28+
connectivity = connectivity.at[n_jr_nodes:, n_jr_nodes:].set(0.2)
29+
connectivity = connectivity.at[n_jr_nodes:, :n_jr_nodes].set(0.1)
30+
connectivity = connectivity.at[jp.diag_indices(n_total_nodes)].set(0)
31+
32+
33+
# 2. Define PyTrees for the heterogeneous network
34+
# We use namedtuples to structure the state, parameters, and noise for the
35+
# different model types. This is a flexible way to handle heterogeneous
36+
# networks and is compatible with JAX's `tree_map` and `jit`.
37+
HeteroState = collections.namedtuple('HeteroState', ['jr', 'mpr'])
38+
HeteroTheta = collections.namedtuple('HeteroTheta', ['jr', 'mpr'])
39+
40+
41+
# 3. Define the network dynamics function (dfun)
42+
# This function computes the time derivatives for each node in the network.
43+
def hetero_dfun(state, p, coupling_strength):
44+
jr_state, mpr_state = state
45+
jr_p, mpr_p = p
46+
47+
# Extract firing rates from each model type
48+
# For JR, the firing rate is proportional to the pyramidal cell activity (y0).
49+
# For MPR, the firing rate is the `r` variable.
50+
jr_r = jr_state.y0
51+
mpr_r = mpr_state.r
52+
53+
# Combine firing rates from all nodes into a single vector
54+
all_r = jp.concatenate([jr_r, mpr_r])
55+
56+
# Calculate the coupling input to each node
57+
coupling_input = coupling_strength * connectivity @ all_r
58+
59+
# Distribute the coupling input to the respective model types
60+
jr_coupling_input = coupling_input[:n_jr_nodes]
61+
# MPR's dfun expects a tuple for coupling (e.g., for different receptors)
62+
mpr_coupling_input = (coupling_input[n_jr_nodes:], 0.0)
63+
64+
# Calculate the derivatives for each model type
65+
d_jr_array = vb.jr_dfun(jr_state, jr_coupling_input, jr_p)
66+
d_mpr_array = vb.mpr_dfun(mpr_state, mpr_coupling_input, mpr_p)
67+
68+
# The dfun for each model returns a flat array. We need to repack it
69+
# into the same namedtuple structure as the state.
70+
d_jr = vb.JRState(*[d_jr_array[i] for i in range(len(vb.JRState._fields))])
71+
d_mpr = vb.MPRState(*[d_mpr_array[i] for i in range(len(vb.MPRState._fields))])
72+
73+
return HeteroState(jr=d_jr, mpr=d_mpr)
74+
75+
76+
# 4. Set up the simulation
77+
dt = 0.1
78+
duration = 1000.0
79+
t_steps = int(duration / dt)
80+
coupling_strength = 0.05
81+
noise_level = 0.01
82+
83+
# Wrap the dfun to include the coupling strength parameter
84+
dfun = lambda state, p: hetero_dfun(state, p, coupling_strength)
85+
86+
# Create the SDE loop
87+
_, loop = vb.make_sde(dt=dt, dfun=dfun, gfun=noise_level)
88+
89+
# 5. Set up initial conditions, parameters, and noise
90+
# Initial states
91+
jr_state_init = vb.JRState(*[jp.ones(n_jr_nodes) * v for v in (0.1, 0.2, 0.3, 0.4, 0.5, 0.6)])
92+
mpr_state_init = vb.MPRState(r=jp.ones(n_mpr_nodes) * 0.1, V=jp.ones(n_mpr_nodes) * -2.0)
93+
initial_state = HeteroState(jr=jr_state_init, mpr=mpr_state_init)
94+
95+
# Parameters
96+
params = HeteroTheta(jr=vb.jr_default_theta, mpr=vb.mpr_default_theta)
97+
98+
# Noise: the noise PyTree must have the same structure as the state
99+
jr_noise = vb.JRState(*[vb.randn(t_steps, n_jr_nodes) for _ in range(len(vb.JRState._fields))])
100+
mpr_noise = vb.MPRState(*[vb.randn(t_steps, n_mpr_nodes) for _ in range(len(vb.MPRState._fields))])
101+
noise = HeteroState(jr=jr_noise, mpr=mpr_noise)
102+
103+
104+
# 6. Run the simulation
105+
print("Running simulation...")
106+
result = loop(initial_state, noise, params)
107+
print("Simulation finished.")
108+
109+
110+
# 7. Plot the results
111+
print("Plotting results...")
112+
time = jp.arange(t_steps) * dt
113+
114+
plt.figure(figsize=(12, 6))
115+
116+
# Plot JR model states (y0: pyramidal cell activity)
117+
plt.subplot(2, 1, 1)
118+
plt.plot(time, result.jr.y0, alpha=0.8)
119+
plt.title('Jansen-Rit Model (Pyramidal Cell Activity)')
120+
plt.xlabel('Time (ms)')
121+
plt.ylabel('Activity')
122+
plt.grid(True)
123+
124+
# Plot MPR model states (r: firing rate)
125+
plt.subplot(2, 1, 2)
126+
plt.plot(time, result.mpr.r, alpha=0.8)
127+
plt.title('Montbrio-Pazo-Roxin Model (Firing Rate)')
128+
plt.xlabel('Time (ms)')
129+
plt.ylabel('Firing Rate')
130+
plt.grid(True)
131+
132+
plt.tight_layout()
133+
plt.savefig('images/heterogeneous_network_activity.png')
134+
plt.show()
135+
print("Plot saved to images/heterogeneous_network_activity.png")

‎examples/heterogeneous_stability.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
"""
2+
**Example**: Stability analysis of a heterogeneous network.
3+
4+
This example extends heterogeneous_network.py to demonstrate how to
5+
analyze the stability of a network's fixed point by computing the
6+
eigenvalues of the Jacobian matrix.
7+
8+
"""
9+
import os
10+
import jax
11+
import jax.numpy as jp
12+
from jax.tree_util import tree_flatten, tree_unflatten
13+
import vbjax as vb
14+
import collections
15+
import matplotlib.pyplot as plt
16+
17+
# Create a directory for saving images
18+
os.makedirs('images', exist_ok=True)
19+
20+
# 1. Define the network structure (same as heterogeneous_network.py)
21+
n_jr_nodes = 4
22+
n_mpr_nodes = 2
23+
n_total_nodes = n_jr_nodes + n_mpr_nodes
24+
25+
connectivity = jp.zeros((n_total_nodes, n_total_nodes))
26+
connectivity = connectivity.at[:n_jr_nodes, :n_jr_nodes].set(0.2)
27+
connectivity = connectivity.at[n_jr_nodes:, n_jr_nodes:].set(0.2)
28+
connectivity = connectivity.at[n_jr_nodes:, :n_jr_nodes].set(0.1)
29+
connectivity = connectivity.at[jp.diag_indices(n_total_nodes)].set(0)
30+
31+
# 2. Define PyTrees for the heterogeneous network
32+
HeteroState = collections.namedtuple('HeteroState', ['jr', 'mpr'])
33+
HeteroTheta = collections.namedtuple('HeteroTheta', ['jr', 'mpr'])
34+
35+
# 3. Define the network dynamics function (dfun)
36+
def hetero_dfun(state, p, coupling_strength):
37+
jr_state, mpr_state = state
38+
jr_p, mpr_p = p
39+
40+
jr_r = jr_state.y0
41+
mpr_r = mpr_state.r
42+
all_r = jp.concatenate([jr_r, mpr_r])
43+
44+
coupling_input = coupling_strength * connectivity @ all_r
45+
jr_coupling_input = coupling_input[:n_jr_nodes]
46+
mpr_coupling_input = (coupling_input[n_jr_nodes:], 0.0)
47+
48+
d_jr_array = vb.jr_dfun(jr_state, jr_coupling_input, jr_p)
49+
d_mpr_array = vb.mpr_dfun(mpr_state, mpr_coupling_input, mpr_p)
50+
51+
d_jr = vb.JRState(*[d_jr_array[i] for i in range(len(vb.JRState._fields))])
52+
d_mpr = vb.MPRState(*[d_mpr_array[i] for i in range(len(vb.MPRState._fields))])
53+
54+
return HeteroState(jr=d_jr, mpr=d_mpr)
55+
56+
# 4. Set up simulation parameters
57+
dt = 0.1
58+
coupling_strength = 0.05
59+
60+
# Initial states and parameters
61+
jr_state_init = vb.JRState(*[jp.ones(n_jr_nodes) * v for v in (0.1, 0.2, 0.3, 0.4, 0.5, 0.6)])
62+
mpr_state_init = vb.MPRState(r=jp.ones(n_mpr_nodes) * 0.1, V=jp.ones(n_mpr_nodes) * -2.0)
63+
initial_state = HeteroState(jr=jr_state_init, mpr=mpr_state_init)
64+
params = HeteroTheta(jr=vb.jr_default_theta, mpr=vb.mpr_default_theta)
65+
66+
# Define a deterministic dfun for finding the fixed point
67+
dfun = lambda state, p: hetero_dfun(state, p, coupling_strength)
68+
69+
# 5. Find a fixed point of the system
70+
# We'll do this by running a deterministic simulation until it settles.
71+
# This gives an approximation of a stable fixed point.
72+
print("Finding fixed point...")
73+
run_sim = vb.make_ode(dt, dfun)
74+
# Run for a bit to let transients pass away
75+
fp_state = run_sim(initial_state, params, time=1000.0)
76+
print("Fixed point found.")
77+
78+
# 6. Linearize the system and find eigenvalues
79+
# To use jax.jacobian, we need to work with flat vectors, not PyTrees.
80+
# We'll create functions to flatten the state PyTree and wrap the dfun.
81+
82+
# Get the PyTree structure definition from our initial state
83+
_, state_tree_def = tree_flatten(initial_state)
84+
85+
def flatten_state(state_pytree):
86+
"Flattens a state PyTree into a single 1D array."
87+
leaves, _ = tree_flatten(state_pytree)
88+
return jp.concatenate([leaf.ravel() for leaf in leaves])
89+
90+
def unflatten_state(flat_state_vec, tree_def):
91+
"Unflattens a 1D array back into a state PyTree."
92+
# Get the shapes of the leaves from a template PyTree
93+
template_leaves, _ = tree_flatten(tree_def.unflatten([0]*tree_def.num_leaves))
94+
leaf_shapes = [leaf.shape for leaf in template_leaves]
95+
96+
# Unflatten the vector
97+
leaf_sizes = [jp.prod(jp.array(s)) for s in leaf_shapes]
98+
indices = jp.cumsum(jp.array([0] + leaf_sizes))
99+
leaves = [
100+
flat_state_vec[indices[i]:indices[i+1]].reshape(leaf_shapes[i])
101+
for i in range(len(leaf_shapes))
102+
]
103+
return tree_unflatten(tree_def, leaves)
104+
105+
def flat_dfun(flat_state_vec, p):
106+
"A wrapper for dfun that works with flattened state vectors."
107+
state = unflatten_state(flat_state_vec, state_tree_def)
108+
d_state_pytree = dfun(state, p)
109+
return flatten_state(d_state_pytree)
110+
111+
# Now we can compute the Jacobian at the fixed point
112+
print("Computing Jacobian and eigenvalues...")
113+
fp_flat = flatten_state(fp_state)
114+
jacobian_matrix = jax.jacobian(flat_dfun)(fp_flat, params)
115+
eigenvalues = jp.linalg.eigvals(jacobian_matrix)
116+
print("Done.")
117+
118+
# 7. Plot the eigenvalues in the complex plane
119+
print("Plotting results...")
120+
plt.figure(figsize=(8, 8))
121+
plt.scatter(jp.real(eigenvalues), jp.imag(eigenvalues), c='b', marker='o')
122+
plt.axvline(0, color='r', linestyle='--', label='Stability Boundary (Re=0)')
123+
plt.title(f'Eigenvalues of the Jacobian (g={coupling_strength})')
124+
plt.xlabel('Real Part')
125+
plt.ylabel('Imaginary Part')
126+
plt.grid(True)
127+
plt.legend()
128+
plt.axis('equal')
129+
plt.savefig('images/heterogeneous_network_eigenvalues.png')
130+
plt.show()
131+
132+
print("Plot saved to images/heterogeneous_network_eigenvalues.png")
133+
134+
# Interpretation:
135+
# If all eigenvalues are in the left half-plane (real part < 0), the
136+
# fixed point is stable. If any eigenvalue has a positive real part,
137+
# the fixed point is unstable. By running this analysis for different
138+
# values of `coupling_strength`, you can find bifurcations where the
139+
# system's stability changes.

0 commit comments

Comments
 (0)