Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

allow visualizing connections #184

Merged
merged 5 commits into from
Dec 6, 2023
Merged

allow visualizing connections #184

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4b013ac
allow visualizing connections
michaeldeistler Dec 6, 2023
3242655
Refine plotting of connections
michaeldeistler Dec 6, 2023
2a39f5f
bugfixes; direction of synapse as scatter
michaeldeistler Dec 6, 2023
34f405f
updates
michaeldeistler Dec 6, 2023
7cccc8d
fix tests
michaeldeistler Dec 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion jaxley/modules/branch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ def __init__(

# Synapse indexing.
self.syn_edges = pd.DataFrame(
dict(pre_comp_index=[], post_comp_index=[], type="")
dict(global_pre_comp_index=[], global_post_comp_index=[], type="")
)
self.branch_edges = pd.DataFrame(
dict(parent_branch_index=[], child_branch_index=[])
Expand Down
2 changes: 1 addition & 1 deletion jaxley/modules/cell.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ def __init__(

# Synapse indexing.
self.syn_edges = pd.DataFrame(
dict(pre_comp_index=[], post_comp_index=[], type="")
dict(global_pre_comp_index=[], global_post_comp_index=[], type="")
)
self.branch_edges = pd.DataFrame(
dict(
Expand Down
2 changes: 1 addition & 1 deletion jaxley/modules/compartment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ def __init__(self):
)
# Synapse indexing.
self.syn_edges = pd.DataFrame(
dict(pre_comp_index=[], post_comp_index=[], type="")
dict(global_pre_comp_index=[], global_post_comp_index=[], type="")
)
self.branch_edges = pd.DataFrame(
dict(parent_branch_index=[], child_branch_index=[])
Expand Down
67 changes: 54 additions & 13 deletions jaxley/modules/network.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,10 @@ def init_conds(self, params):
return cond_params

def init_syns(self):
pre_comp_inds = []
post_comp_inds = []
global_pre_comp_inds = []
global_post_comp_inds = []
pre_locs = []
post_locs = []
pre_branch_inds = []
post_branch_inds = []
pre_cell_inds = []
Expand All @@ -210,28 +212,39 @@ def init_syns(self):
pre_cell_inds_, pre_inds, post_cell_inds_, post_inds = prepare_syn(
connectivity.conns, self.nseg
)
pre_comp_inds.append(
# Global compartment indizes.
global_pre_comp_inds.append(
self.cumsum_nbranches[pre_cell_inds_] * self.nseg + pre_inds
)
post_comp_inds.append(
global_post_comp_inds.append(
self.cumsum_nbranches[post_cell_inds_] * self.nseg + post_inds
)
pre_branch_inds.append(self.cumsum_nbranches[pre_cell_inds_])
post_branch_inds.append(self.cumsum_nbranches[post_cell_inds_])
# Local compartment inds.
pre_locs.append(np.asarray([c.pre_loc for c in connectivity.conns]))
post_locs.append(np.asarray([c.post_loc for c in connectivity.conns]))
# Local branch inds.
pre_branch_inds.append(
np.asarray([c.pre_branch_ind for c in connectivity.conns])
)
post_branch_inds.append(
np.asarray([c.post_branch_ind for c in connectivity.conns])
)
pre_cell_inds.append(pre_cell_inds_)
post_cell_inds.append(post_cell_inds_)

# Prepare synapses.
self.syn_edges = pd.DataFrame(
columns=[
"pre_comp_index",
"pre_locs",
"pre_branch_index",
"pre_cell_index",
"post_comp_index",
"post_locs",
"post_branch_index",
"post_cell_index",
"type",
"type_ind",
"global_pre_comp_index",
"global_post_comp_index",
]
)
for i, connectivity in enumerate(self.connectivities):
Expand All @@ -240,14 +253,16 @@ def init_syns(self):
self.syn_edges,
pd.DataFrame(
dict(
pre_comp_index=pre_comp_inds[i],
pre_locs=pre_locs[i],
pre_branch_index=pre_branch_inds[i],
pre_cell_index=pre_cell_inds[i],
post_comp_index=post_comp_inds[i],
post_locs=post_locs[i],
post_branch_index=post_branch_inds[i],
post_cell_index=post_cell_inds[i],
type=type(connectivity.synapse_type).__name__,
type_ind=i,
global_pre_comp_index=global_pre_comp_inds[i],
global_post_comp_index=global_post_comp_inds[i],
)
),
],
Expand Down Expand Up @@ -275,8 +290,8 @@ def _step_synapse(
voltages = u["voltages"]

grouped_syns = edges.groupby("type", sort=False, group_keys=False)
pre_syn_inds = grouped_syns["pre_comp_index"].apply(list)
post_syn_inds = grouped_syns["post_comp_index"].apply(list)
pre_syn_inds = grouped_syns["global_pre_comp_index"].apply(list)
post_syn_inds = grouped_syns["global_post_comp_index"].apply(list)
synapse_names = list(grouped_syns.indices.keys())

syn_voltage_terms = jnp.zeros_like(voltages)
Expand Down Expand Up @@ -320,7 +335,9 @@ def vis(
neuron, as read from the SWC file.
layers: Allows to plot the network in layers. Should provide the number of
neurons in each layer, e.g., [5, 10, 1] would be a network with 5 input
neurons, 10 hidden layer neurons, and 1 output neuron.
neurons, 10 hidden layer neurons, and 1 output neuron. Please note that
connections always start and end in the middle of branches, no matter
what compartment within a branch they are connected to.
options: Plotting options passed to `NetworkX.draw()`.
cols: One color in total or one color per cell.
"""
Expand All @@ -341,6 +358,30 @@ def vis(

for cell, col in zip(self.cells, cols):
fig, ax = cell.vis(detail="full", dims=dims, cols=col, fig=fig, ax=ax)

# Plot connections (i.e. synapses).
pre_locs = self.syn_edges["pre_locs"].to_numpy()
post_locs = self.syn_edges["post_locs"].to_numpy()
pre_branch = self.syn_edges["pre_branch_index"].to_numpy()
post_branch = self.syn_edges["post_branch_index"].to_numpy()
pre_cell = self.syn_edges["pre_cell_index"].to_numpy()
post_cell = self.syn_edges["post_cell_index"].to_numpy()

dims_np = np.asarray(dims)

for pre_b, post_b, pre_c, post_c in zip(
pre_branch, post_branch, pre_cell, post_cell
):
pre_coord = self.cells[pre_c].xyzr[pre_b]
middle_ind = int((len(pre_coord) - 1) * pre_locs)
pre_coord = pre_coord[middle_ind]
post_coord = self.cells[post_c].xyzr[post_b]
middle_ind = int((len(post_coord) - 1) * post_locs)
post_coord = post_coord[middle_ind]
coords = np.stack([pre_coord[dims_np], post_coord[dims_np]]).T
ax.plot(coords[0], coords[1], linewidth=3.0, c="b")
ax.scatter(post_coord[dims_np[0]], post_coord[dims_np[1]], c="b")

return fig, ax
else:
raise ValueError("detail must be in {point, full}")
Expand Down
6 changes: 3 additions & 3 deletions tests/test_swc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ def test_swc_reader_lengths():
dirname = os.path.dirname(__file__)
fname = os.path.join(dirname, "morph.swc")

_, pathlengths, _, _ = jx.utils.swc.swc_to_jaxley(fname, max_branch_len=2000.0)
_, pathlengths, _, _, _ = jx.utils.swc.swc_to_jaxley(fname, max_branch_len=2000.0)
pathlengths = np.asarray(pathlengths)[1:]

for sec in h.allsec():
Expand Down Expand Up @@ -60,7 +60,7 @@ def test_swc_radius():
dirname = os.path.dirname(__file__)
fname = os.path.join(dirname, "morph_250.swc")

_, pathlen, radius_fns, _ = jx.utils.swc.swc_to_jaxley(
_, pathlen, radius_fns, _, _ = jx.utils.swc.swc_to_jaxley(
fname, max_branch_len=2000.0, sort=False
)
jaxley_diams = []
Expand Down Expand Up @@ -128,7 +128,7 @@ def test_swc_voltages():
pathlengths_neuron = np.asarray([sec.L for sec in h.allsec()])

####################### jaxley ##################
_, pathlengths, _, _ = jx.utils.swc.swc_to_jaxley(fname, max_branch_len=2_000)
_, pathlengths, _, _, _ = jx.utils.swc.swc_to_jaxley(fname, max_branch_len=2_000)
cell = jx.read_swc(fname, nseg_per_branch, max_branch_len=2_000.0)
cell.insert(HHChannel)

Expand Down