jaxleyverse / jaxley

Differentiable neuron simulations with biophysical detail on CPU, GPU, or TPU.
https://jaxley.readthedocs.io
Apache License 2.0
55 stars 9 forks source link

Parameter sharing with different numbers of synapses #501

Open kyralianaka opened 5 days ago

kyralianaka commented 5 days ago

The new tutorial on advanced parameter sharing suggests that something like the following would be possible:

comp = jx.Compartment()
branch = jx.Branch(comp, nseg=1)
cell = jx.Cell(branch, parents=[-1])
net = jx.Network([cell for _ in range(3)])
connect(net.cell(0), net.cell(1), IonotropicSynapse())
connect(net.cell(0) ,net.cell(2), IonotropicSynapse())
connect(net.cell(1), net.cell(2), IonotropicSynapse())

net.copy_node_property_to_edges("global_cell_index")

df = net.edges
df["controlled_by_param"] = df["pre_global_cell_index"]
net.make_trainable("IonotropicSynapse_gS")

but this will throw an error because two synapses are controlled by one pre-synaptic cell and only one synapse by another pre-synaptic cell.

This is a major limitation that prevents cells having the same (pre-)synaptic parameters but different numbers of post-synaptic partners. Maybe I am missing something as to how this parameter sharing should be done -- please let me know if so. In any case I will keep looking into it.

michaeldeistler commented 5 days ago

I think this is related to the following not working:

comp = jx.Compartment()
branch1 = jx.Branch(comp, nseg=1)
branch2 = jx.Branch(comp, nseg=2)
cell = jx.Cell([branch1, branch2], parents=[-1, 0])

cell.branch("all").make_trainable("radius")