Add support for graph operations

[jnsbck]

This PR adds support for graph methods. It implements a to_graph and from_graph method in module.

Todos:

• Add to_graph
  • Add to_graph
  • add tests
• Add from_graph
  • Add from_graph
  • add tests
    - [ ] ensure it works for views and modules, with Make View behave more like Module #351, this should be straight forward.
• implement .swc -> nx.DiGraph -> Cell functionality
  • add tests
  • ensure original tests are passing for new swc reader.

Thoughts I had while coding this up:

• We could make the graph object a more central part of jaxley and populate self.nodes and self.edges from it.
• Since all attrs from in module.edges, module.branch_edges and module.nodes, are stored as node / edge attrs in the DiGraph, this could also be used for plotting or in vis.
• Since all relevant info can be stored in DiGraph.graph,DiGraph.nodes,DiGraph.edges, this could also be used to save and share modules.
• It would be nice to have the group idendity of nodes showin in module.nodes, not sure if one group column that contains lists of groups or several boolean cols with group_name is the better solution.
• Makes only sense to be added to view once Make View behave more like Module #351 is merged, since this implements all properties that are attached to the nodes / edges for views of modules.

EDIT:

• ensure this matches with NEURON / old swc pipeline
• add tutorial
• add tests and refactor them
• update test_swc test to better match compartments with NEURON

[jnsbck]

Currently I have only implemented to_graph and it works for comp, branch, cell, net, although I have not yet added tests. The functionality is very neat and makes looking at the networks / cells really intuitive. For the following network,

comp = jx.Compartment()
branch = jx.Branch([comp for _ in range(4)])
cell = jx.Cell([branch for _ in range(5)], parents=jnp.asarray([-1, 0, 1, 2, 2]))
net = jx.Network([cell]*3)
connect(net[0,0,0], net[1,0,0], IonotropicSynapse())
connect(net[0,0,1], net[1,0,1], IonotropicSynapse())
connect(net[0,0,1], net[1,0,1], TestSynapse())
net.cell(2).add_to_group("cell2")
net.cell(2).branch(1).add_to_group("cell2brach1")


net.cell(0).insert(Na())
net.cell(0).insert(Leak())

net.cell(1).branch(1).insert(Na())
net.cell(0).insert(K())
net.compute_xyz()

net.cell(0).branch(0).loc(0.0).record()
net.cell(0).branch(0).loc(0.0).record("m")
current = jx.step_current(i_delay, i_dur, i_amp, dt, t_max)
net.cell(0).branch(2).loc(0.0).stimulate(current)
net.cell(0).branch(1).make_trainable("Na")
net.cell(1).make_trainable("K")

net.compute_xyz()
net.cell(1).move(0,30,0)
net.cell(2).move(0,-30,0)

we can just to_graph, plot it

module_graph = net.to_graph()
# plot the graph
pos = {i: (n["x"], n["y"]) for i, n in module_graph.nodes(data=True)} 
plt.figure(figsize=(8, 8))
nx.draw(module_graph, pos, with_labels=True, node_size=200, node_color="skyblue", font_size=8, font_weight="bold", font_color="black", font_family="sans-serif")
plt.show()

and look at all its properties, i.e. checking out the soma (0) and the synapse going from node 0 to node 20

print(module_graph.nodes[0])
print(module_graph.edges[(0,20)])

{'comp_index': 0, 'branch_index': 0, 'cell_index': 0, 'length': 10.0, 'radius': 1.0, 'axial_resistivity': 5000.0, 'capacitance': 1.0, 'v': -70.0, 'Na': True, 'Na_gNa': 0.05, 'Na_eNa': 35.0, 'Na_m': 0.2, 'Na_h': 0.2, 'Leak': True, 'Leak_gLeak': 5e-05, 'Leak_eLeak': -67.0, 'K': True, 'K_gK': 0.012, 'eK': -75.0, 'K_n': 0.1, 'x': 5.0, 'y': 0.0, 'z': 0.0, 'recordings': array(['v', 'm'], dtype=object)}

{'pre_locs': 0.125, 'pre_branch_index': 0, 'pre_cell_index': 0, 'post_locs': 0.125, 'post_branch_index': 0, 'post_cell_index': 1, 'type': 'IonotropicSynapse', 'type_ind': 0, 'global_pre_comp_index': 0, 'global_post_comp_index': 20, 'global_pre_branch_index': 0, 'global_post_branch_index': 5, 'IonotropicSynapse_gS': 0.5, 'IonotropicSynapse_e_syn': 0.0, 'IonotropicSynapse_k_minus': 0.025, 'IonotropicSynapse_s': 0.2, 'TestSynapse_gC': nan, 'TestSynapse_c': nan}

[jnsbck]

Both to_graph and from_graph work. If you want to give this a preliminary review, that would be awesome! Please also chime in if you have specific ideas about the format in which to best store the attributes in the graph. I think we do not have to stick with how it is stored in Module and should rather ensure that they are easy to parse and work with.

[jnsbck]

@michaeldeistler is there a scenario where there is more than one default value for an item in Module.trainable_params? I.e. the items are all dicts of form {"Na_gNa": np.array([0.1])} I think. Additionally, is there a case where there is more than one entry in the dict?"

[jnsbck]

Ok I just saw you can have multiple defaults, this code is wrong then.
But you cannot have multiple keys correct?

[michaeldeistler]

Indeed, you cannot have multiple keys.

[jnsbck]

ignore this

[michaeldeistler]

I will do a more thorough review on Monday, but for now: I really like that the methods are standalone. I guess in the long run we would try to use from_graph also for reading SWC readers?

[jnsbck]

Yes, as mentioned in the PR desc, reading swc to graph is the plan.

[michaeldeistler]

Potential use-case of to_graph:

• reordering parents to have a different "initial" parent node
• search operations such as shortest path

[michaeldeistler]

• new type -> new branch
• new type -> no radius interpolation see here
• types have to be returned and directly added to group if assign_group=True
• single point soma. Length should be 2x radius see here

[jnsbck]

It took way more time than I would have suspected or liked to spend and two mayor iterations, but its finally mostly there! Both export and import from/to graph are working now and also the swc pipeline passes the tests now. Some more tests and cleanup, but then I think you can do a review @michaeldeistler. Small benefit of this PR is also that it reduces the difference of jaxley and NEURON very slightly (by about 20%).

[jnsbck]

TLDR: I think I have a working version of the graph pipeline now, but tests are still not passing and I don't quite no why. Everything looks good to me. Help is much appreciated!

I have been trying for ages to get from_graph working. For this I have written a networkx -> jaxley converter. I think it works quite well. Running:

graph = swc_to_graph(fname)
cell = from_graph(graph, nseg=8, max_branch_len=2000.0)

reads an swc file into a networkx graph and imports the cell into jaxley. As part of from_graph the afformentioned converter is run, but it can also be run seperately via: graph = make_jaxley_compatible(graph). This produces a jx.Cell object that runs and does things. Additionally, inspecting the morphologies I get visually and comparing the compartment coordinates to NEURON / the current swc pipeline also looks good. As a bonus the compartment centers assigned via the graph are even closer to NEURON than the current ones. See the left plot.

While this all looks very promising, I have struggled to get the tests passing and I really dont know why:
For this I have replaced the hardcoded branch indices and pathlength matching in the old tests and I now just match the closest compartment to NEURON. This works quite well and I can even assign trunk_ids etc. based on the swc ids which is nice. The test should be much more robust to changes in the test_morphology or other things now! Of course it also passes for the old swc pipeline. However, when I run this test on the graph imported jx.Cell the voltages are consistently. See here:

While they look very similar, they are all slightly delayed. The MSEs also show this.

MSEs between networkX and NEURON voltages [10.14612018 10.35377231 10.84316607 11.17640196 11.06866644 11.1941421
 10.72721014 10.44691712 10.55542459  9.83932744  9.47578905 10.63712368
 13.2997599   9.5394632  10.3148209  10.06744197 10.11404793  9.99624842
 11.15663607 11.10845196 10.55436891 10.17346833 10.95011269 11.09521677
 11.24092038 10.65548104 11.36646986]
MSEs between jaxley_swc_reader and NEURON voltages [0.72412308 0.7978693  0.81918308 0.84724678 0.80464105 1.43258348
 1.47786316 0.73113843 1.14953477 1.12154859 1.18076938 1.30564702
 1.56853646 2.9105652  0.38078612 1.00738273 0.79437486 0.81636714
 0.48941009 0.47586723 0.75917783 1.04822786 1.35910234 1.56591299
 1.14104993 1.34106949 1.72583906]

Would really appreciate help. I will now be moving on to sth else for the moment though, so it is not urgent.
@michaeldeistler

Best,

Jonas

[jnsbck]

I think this is finally ready for a first round of reviews.

This has become quite the mammoth PR, but the functionality it enables is neat imo. For a rundown see the updated 08_morphologies.ipynb

All tests are passing now, which turned out to be an immense amount of work, but the imported morphologies are similar enough to NEURON both at the compartment level (x,y,z,r,l) and they also simulate correctly. I have essentially cloned the tests in tests/test_swc.py and added them to tests/test_graph.py. I have kept the tests and methods separated from everything else (although a refactor / integration might make it easier to maintain later). The import and export functionality also works correctly.

Notable changes are:

• new io submodule

  • with swc_to_graph, from_graph and to_graph functions
  • moved read_swc there from cell.py
  • added io.graph
  • fixed the _update_nodes_with_xyz since it had the same issue that caused bug in compute_xyz #410.
  • added __eq__ to Module (I needed this to test whether re-imported modules are the same as the original.)
  • test_graph.py matches voltages between NEURON and Jaxley by matching compartment centers, this makes the test much more robust. Could also be added to test_swc.py if desired (see tests/helpers.py).

• Some open questions are:

  • __eq__ is only implicitely being tested, and I dont know whether there could be a more efficient implementation
  • there is a lot of overlap between test_graph and test_swc. We could potentially move them to tests/io/ and reduce overlap.

Lemme know your thoughts. Would also be happy to go through this in person.

[manuelgloeckler]

Just noting that in the branch graph_distance_and_more, there is some preliminary API idea that works, e.g., with, e.g., "distance" as a new function. Implementation may be inefficient yet, but can be improved if we can also add the graph to the module...

It is a branch from this branch at the current stage.

[jnsbck]

Here is a little update on the current state of the PR:

1. I have rebased onto main and adapted + improved the code.
2. The following is working / not working

• WORKING: read swc into networkx graph
• WORKING: pre-processing (handling single soma, adding root branch, etc.)
• WORKING: tracing out branches
• WORKING: inserting compartments into morphology graph
• WORKING: extracting the compartment graph
• WORKING: from_graph, i.e. comp graph -> module (works for node and edge attrs, but does not attach recordings, trainables etc. I am in favour of not supporting this.)
• WORKING: to_graph, i.e. export (also only for morph)
• NOT WORKING: tests

Also added an updated tutorial.