Simpler edge indexing

[jnsbck]

Currently synapse and channel parameters / states are handled differently. While channel params are referred to with global indices, synapse params are referenced on a per synapse basis. This leads to very different implementations of synapse and channel updates. This PR makes the synapse indexing global, which simplifies several aspects of the code and allows for more function reuse.

This could potentially be simplified even further, such that channels and synapses can be handled through mostly the same functions.

[jnsbck]

• rerun tests after Fix synapse clamping #492 is merged

[michaeldeistler]

Nice!

I am a bit worried about memory consumption though. If every jaxedge has every state, then there will a lot of NaNs each of which contains memory. I think that, in the worst case, memory increases by N**2, where is the number of types of synapses.

Can we have a zoom about this? E.g. next week?

[jnsbck]

True! I did not think about this. I have two thoughts though. Why not do this for nodes on a per mechanism basis as well. And is there not a more straight forward way to avoid nans while still using global indexes?
But yes, maybe a quick zoom would he great :)

[michaeldeistler]

I agree that we should also do it like this for nodes. Let's zoom next week!

[jnsbck]

fixes a bug, where dtype of bool columns was changed upon concatenation in:

comp1 = jx.Compartment()
comp2 = jx.Compartment()
comp2.insert(HH())
branch = jx.Branch([comp1, comp2)] #-> branch.nodes["HH"].dtype was `object` in this case

[jnsbck]

obsolete

[jnsbck]

group_nodes was removed in #447.

[jnsbck]

Main changes:

• rename channel_params... synapse_states to params...states (allows to reuse stuff for both synapses and channels)
• to_jax stores params/states on a per mechanism basis now -> no more NaNs. To do this some prep is needed in order to keep track of what goes where to find it again later. Hence I had to write a _prepare_for_jax method, that pre-computes indices and adds a lookup method for all states / params in the module. This should also make this easily adaptable if we add the global_param API change.
• refactor of get_all_parameters and get_all_states
• refactor of channel and synapse updates.

In general I think the code is a bit cleaner and more consolidated now (-150 line diff). Will also run regression tests. Maybe its even faster as well. I am not super happy about the _prepare_for_jax solution tbh, but it seems to work well fwiw. If you have feedback, I'd be very happy about any suggestions!

Lemme know what you think or if you have any questions.

The only thing I touched in terms of channels / synapse rewrite is the params/states rename. Does this warrant a seperate branch?

[michaeldeistler]

Thanks a lot @jnsbck!

I am already off for the holidays so will not be able to do a detailed review, but:

1. I think all your high-level explanations sound good!
2. Yes, let's make a new v1.0 branch and merge it there because it breaks all channel models.

And yes, please run regression tests and ensure that all is good! Good to go then! Thanks!

[jnsbck]

/test_regression

[github-actions]

Regression Test Results

✅ Process completed

test_runtime(num_cells=1, artificial=False, connect=False, connection_prob=0.0, voltage_solver=jaxley.stone)
🟢 build_time: (-42.62% vs 1.651s).
🟢 compile_time: (-3.17% vs 19.652s).
🟢 run_time: (-16.57% vs 2.891s).

test_runtime(num_cells=1, artificial=False, connect=False, connection_prob=0.0, voltage_solver=jax.sparse)
🟢 build_time: (-3.26% vs 0.313s).
🟢 compile_time: (-30.37% vs 3.523s).
🟢 run_time: (-16.81% vs 2.454s).

test_runtime(num_cells=10, artificial=False, connect=True, connection_prob=0.1, voltage_solver=jaxley.stone)
🟢 build_time: (-36.81% vs 3.856s).
🟢 compile_time: (-15.62% vs 29.145s).
🟢 run_time: (-19.73% vs 18.551s).

test_runtime(num_cells=10, artificial=False, connect=True, connection_prob=0.1, voltage_solver=jax.sparse)
🟢 build_time: (-31.42% vs 1.825s).
🟢 compile_time: (-22.40% vs 21.890s).
🟢 run_time: (-18.60% vs 20.018s).

test_runtime(num_cells=1000, artificial=True, connect=True, connection_prob=0.001, voltage_solver=jaxley.stone)
🟢 build_time: (-4.67% vs 113.937s).
🟢 compile_time: (-22.13% vs 45.211s).
🟢 run_time: (-21.45% vs 41.206s).

test_runtime(num_cells=1000, artificial=True, connect=True, connection_prob=0.001, voltage_solver=jax.sparse)
🟢 build_time: (-2.18% vs 108.073s).
🟢 compile_time: (-19.37% vs 49.108s).
🟢 run_time: (-17.58% vs 46.364s).