Import from .swc without altering morphology

[bantin]

We have a bit of an odd use case: we would like to import a .swc file and, for every line of the .swc file, we would like our simulation to report a voltage at that exact location in space. I think the right way to do this is to have a branch for each line of the .swc. Any chance there is a way to do this using existing functions? if not, I suppose it shouldn't be hard to write our own.

[jnsbck]

Hi, a very simple thing that might work (although I have not tested it myself), is to set max_len to some very small number in the swc_reader. This should in principle create a branch at every node.

Lemme know if that ends up working, otherwise I'll have another look :)

[michaeldeistler]

True, that's also a good idea! For future reference:

cell = jx.read_swc(fname, ncomp=1, max_branch_len=0.1)

I tried it on some small mrophologies and it only almost worked: I don't always get the exactly correct number of compartments.

[bantin]

Thanks @michaeldeistler, yeah I tried this approach and it usually gets very close -- but sometimes the number of compartments is off by one or two. I'll look into the other method you suggested.

[michaeldeistler]

Hi Ben,

this is unfortunately not a supported feature. I think the way to go is to use the graph backend for this:

1. Build the swc_graph with swc_graph = to_swc_graph("tests/swc_files/morph.swc")
2. Then, manually rename some of the graph attributes such that it is consistent with the attributes of the comp_graph
3. After the renaming, you should be able to import the modified swc_graph as a jx.Cell with cell = from_graph(comp_graph)

Let me know if this works!

All the best
Michael

[bantin]

Thanks so much @michaeldeistler and @jnsbck for your responses so far (and for being so responsive in general!).
 
Since our use case is a bit out of the norm, I figured it might help if I describe it in a bit more detail. My apologies in advance that this message is so long :) We want to build a link between multicompartment models and voltage imaging movies. Our imaging framework places basis functions along the graph of the dendritic tree, and then uses these basis functions to reconstruct the video. I am struggling to figure out the proper framework/abstraction for linking a "basis function" to a "compartment."
 
So far, we have code which takes a traced cell (we trace them by hand) and then outputs an SWC which specifies the locations of the basis functions. For technical reasons, we would like to place a basis function at every branch-point of the tree, and also at the tips of each branch (here I am using the term "branch-point" in the colloquial sense, rather than the jaxley sense). Furthermore, we would like to ensure there is a basis function at least every "N" microns along each branch. Our code outputs an SWC which satisfies these conditions. We've found that this works well for building a low-dimensional representation of our voltage videos.
 
In order to fit biophysical models to our voltage data, we need a mapping between the basis functions which I described above and a compartment in Jaxley. Broadly speaking, I can think of two ways to do it:
 

1. Build a multicompartment model in which there is a compartment located at each basis function (this is the approach I was thinking about when I asked my original question)
2. Allow Jaxley to compartmentalize the neuron, then build a linear mapping between Jaxley's compartments and the basis functions described above.
    
   My intuition is that Option 1 is the way to go, assuming that having additional compartments (at the branchpoints and the tips) won't interfere with the biophysics too much. There's probably a slight computational hit, but that's ok with us. I'm curious if your intuition agrees that Option 1 is probably better than Option 2?
    
   Assuming that we do decide to go with Option 1, I am not sure quite how to go about it. Following @michaeldeistler's advice from above, we need to add attributes to our SWC graph to match the output of build_compartment_graph. This seems straightforward enough. However, I can see that build_compartment_graph adds both compartment type nodes and branchpoint type nodes to the graph. I'm not totally sure how to replicate this behavior in our case. Maybe we want to insert zero-extent branchpoint type nodes into our SWC graph? Or maybe there is someway of having jaxley report the voltages at the branchpoint nodes and connecting this to our branchpoint basis functions.
    
   I realize I've asked quite a lengthy question, and that it might be easier to collaborate externally on this. We'd be super happy to have you guys involved in this project however you see fit. If you'd prefer to follow-up externally, just let me know and I will send you a quick email. Thanks so much!

[michaeldeistler]

Hi Ben,

thanks for the explanation, this helps a lot!

I have thought about this a bit more. I think that option 1 will be very difficult to implement, whereas option 2 will be quite straightforward.

A general note on computing the voltage at branchpoints and tips

Before getting started: In Jaxley (and in NEURON), the voltage at a tip equals the voltage at the neighboring compartment. This ensures Kirchhoff's current law (no current flows "out of the cell"). Similarly, the voltage at a branchpoint equals a weighted average of the voltage a all neighboring compartments.

This means that you will get the voltage at these points for free: your model does not have to contain compartments at the exact tips or at the exact branchpoints. For example, say you have the following branch, where x indicates an SWC traced point:

x ++++ x ++++ x ++++ x

By default, Jaxley would discretize this as follows (assuming ncomp = 4)

x ++++ x ++++ x ++++ x
aaaa--bbbb--cccc--dddd

where a, b, c, d are separate compartments. The voltage at the left tip of the branch is assumed to be exactly equal to the voltage in the aaaaa compartment. One thing to note: The second SWC traced point is not at the exact center of compartment bbbb.

Why option 1 is difficult to implement

Again, consider the following SWC file describing a single branch:

x ++++ x ++++ x ++++ x

In your case, we would want to place compartment such that the compartment center will be at the SWC points. The best we could do is:

x ++++ x ++++ x ++++ x
aaa-bbbbbb-cccccc-ddd

The means that the length of the bbbbbb compartment equals half the distance between the first to SWC points plus half the distance between the second and third SWC point. You would have to manually modify our SWC reader to achieve this, which will be quite tricky. In addition, it introduces these "half-length" compartments (aaa, ddd). That's fine, and (following the arguments above), the voltage at the left tip will be equal to the voltage in the aaa compartment.

The situation is even trickier at branchpoints though: At the branchpoint, each branch will have a "half-length" compartment (for example: aaa or ddd above) which you does not actually show up in your SWC file. The branchpoint voltage will be a weighted average of the voltage in these compartments. This will again be hard to implement, and it also introduces these "spurious" compartments which you do not even have a voltage measurement for.

How to implement Option 2

Option 2, in contrast, is straightforward to implement in Jaxley. I think the core question here is whether interpolating the voltage of two compartments is actually a good. I did some quick experiments and it looked good.

To implement it:

1. find the SWC location in cell.xyzr
2. get its location along the branch
3. linearly interpolate the voltage between the two neighboring compartments

What I would recommend

Unless you want to invest multiple days or weeks into making option 1 work, I would recommend option 2. Just choose a pretty high ncomp for every branch (e.g., by using the d_lambda rule with a pretty low value for d_lambda), and I think you will be fine. If you need help with implementing option 2, let me know!

We could also have a quick zoom call if you think it would be useful.

Michael

[bantin]

Thanks as always for the detailed response! This makes sense to me. After thinking about it more, I agree that Option 2 is probably better. Both because of the reasons you mentioned, and because it will be useful for us to be able to vary the number of basis functions which we use in the reconstruction without changing the biophysical properties of the model. Option 2 it is!

At a high level, the 3 steps you described make perfect sense. However, I don't fully understand the relationship between cell.xyzr and the morphology defined in the SWC file. As best I can tell, there isn't a 1:1 correspondence between the points in cell.xyzr and the SWC compartments. I think the operation I am not sure how to implement at the moment is: "given a point from the SWC file (or swc_graph) find the nearest compartments in the cell model.

[michaeldeistler]

As best I can tell, there isn't a 1:1 correspondence between the points in cell.xyzr and the SWC compartments.

I don't really know what you mean --- there should be a 1:1 corresponence (except that cell.xyzr sorts the coordinates into branches).

Below is a sketch for how one can achieve this:

import jaxley as jx
import numpy as np

fname = "../jaxley/tests/swc_files/morph_ca1_n120_250.swc"
cell = jx.read_swc(fname, ncomp=1)

# Reasonable default values for most models.
frequency = 100.0
d_lambda = 0.1  # Larger -> more coarse-grained.

for branch in cell.branches:
    diameter = 2 * branch.nodes["radius"].to_numpy()[0]
    c_m = branch.nodes["capacitance"].to_numpy()[0]
    r_a = branch.nodes["axial_resistivity"].to_numpy()[0]
    l = branch.nodes["length"].to_numpy()[0]

    lambda_f = 1e5 * np.sqrt(diameter / (4 * np.pi * frequency * c_m * r_a))
    ncomp = int((l / (d_lambda * lambda_f) + 0.9) / 2) * 2 + 1
    branch.set_ncomp(ncomp, initialize=False)

cell.initialize()

################################################################
################################################################
################################################################
coords = np.loadtxt(fname)

row = 55
xyzr = coords[row, [2, 3, 4, 5]]

errors = []
for b in cell.branches:
    dists = np.sum((xyzr - b.xyzr[0])**2, axis=1)
    min_error = np.min(dists)
    errors.append(min_error)
is_in_branch = np.argmin(errors)

xyz_in_branch = cell.branch(is_in_branch).xyzr[0][:, :3]
dists = np.sqrt(np.sum((xyz_in_branch[1:] - xyz_in_branch[:-1])**2, axis=1))

dists_to_swc = np.sum((xyzr[:3] - xyz_in_branch)**2, axis=1)
min_error = np.argmin(dists_to_swc)

cum_dists = np.cumsum(dists)
loc_along_branch = cum_dists[min_error] / np.max(cum_dists)

ncomps = cell.branch(is_in_branch).ncomp
locs_along_branch = np.linspace(1 / ncomps / 2, 1 - 1 / ncomps / 2, ncomps)

voltages = np.random.randn(ncomps)
v_at_swc = np.interp(loc_along_branch, locs_along_branch, voltages)

This does not yet deal with branch points. I will open a PR to record the voltage at branch points asap.

[michaeldeistler]

b.xyzr[0] does not show you the "compartment" location, but all SWC points that are included in the branch b. Therefore, min_error should be exactly 0.0 for one of the branches.

[bantin]

Ah awesome. I was seeing some cases where min_error was not exactly zero, but it looks like that was a floating point issue. I think I pretty much have everything I need here.

The next step is to build a num_swc_nodes x num_compartments matrix that can be used as part of the observation model for filtering algorithms. Will keep you posted :)

[michaeldeistler]

Great, keep us updated!

[bantin]

Here's what I have so far. I broke it up into three cases for each swc node: tip, branchpoint, and within-branch. I wasn't able to follow the logic exactly for handling the branch points, so I somewhat blindly copied what you had (though I had to change some of the indexing into the [2 x num_neighbors] array of conductances. I'd be curious if anything looks wrong with that part? Thanks again!

def get_swc_node_branch_indices(xyzr: List, swc_x: float, swc_y: float, swc_z: float) -> List[int]:
    """
    Given the xyzr array of the cell and the (x,y,z) coordinates of an swc node,
    return the indices of the two nearest compartments in the cell that correspond to the swc node.

    If the swc node is at a branchpoint, return the indices of all compartments that connect to the branchpoint.

    Args:
        xyzr: List, each entry is an (M, 4) array of (x,y,z,r) tracing along a branch
        swc_x: float
        swc_y: float
        swc_z: float
    Returns:
        branch_indices: List[int]
    """
    branch_indices = []
    for branch_idx, this_xyzr in enumerate(xyzr):
        xyz = this_xyzr[:, :3]
        is_in_branch = np.any(
            [np.allclose(xyz[i], [swc_x, swc_y, swc_z]) for i in range(xyz.shape[0])]
        )
        if is_in_branch:
            branch_indices.append(branch_idx)
    return branch_indices

def get_nearest_loc_in_branch(
        cell: jx.Cell,
        swc_xyz: np.ndarray,
        branch_idx: int,
        nearest_k: int = 1,
    ) -> np.ndarray[int]:
    """
    Given a cell, the (x,y,z) coordinates of an swc node, and
    a branch index, return the global index of the nearest compartment"""
    comp_locs = cell.branch(branch_idx).nodes[['x', 'y', 'z']].to_numpy()
    dists = np.linalg.norm(comp_locs - swc_xyz, axis=1)
    sorted_dist_inds = np.argsort(dists)
    sorted_dists = dists[sorted_dist_inds]

    nearest_comp_indices = sorted_dist_inds[:nearest_k]
    glocal_comp_idx = cell.branch(branch_idx).nodes.iloc[nearest_comp_indices]["global_comp_index"]
    global_comp_idx = np.array(glocal_comp_idx, dtype=int)

    return global_comp_idx, sorted_dists[:nearest_k]

def get_matching_branchpoint_index(
        cell: jx.Cell,
        swc_xyz: np.ndarray,
    ) -> int:
    """
    Given a cell and the (x,y,z) coordinates of an swc node,
    return the index of the matching branchpoint
    """
    branchpoint_xyzs = cell._branchpoints[['x', 'y', 'z']].to_numpy()
    matching_branch_idx = np.where(
        [np.allclose(xyz, swc_xyz) for xyz in branchpoint_xyzs]
    )[0]
    assert len(matching_branch_idx) > 0, "No matching branchpoint found for swc node at (%f, %f, %f)" % (swc_xyz[0], swc_xyz[1], swc_xyz[2])
    assert len(matching_branch_idx) == 1, "Multiple matching branchpoints found for swc node at (%f, %f, %f)" % (swc_xyz[0], swc_xyz[1], swc_xyz[2])
    return int(cell._branchpoints.index[matching_branch_idx[0]])
    
    

def build_node_compartment_matrix(cell, swc_graph):
    """
    Build a num_swc_nodes x num_compartments matrix
    that maps voltage from compartments to swc nodes.

    For SWC nodes that map to a branchpoint, we form a weighted combination
    where the weights are determined by the conductances.


    Returns:
        node_comp_matrix: np.array of shape (num_swc_nodes, num_compartments)
    """

    node_comp_matrix = np.zeros((len(swc_graph), len(cell.nodes)))
    print ("node_comp_matrix shape:", node_comp_matrix.shape)

    # step 1, for each swc node, find the corresponding branch in the cell
    for swc_node, swc_node_attrs in swc_graph.nodes(data=True):
        swc_x = swc_node_attrs["x"]
        swc_y = swc_node_attrs["y"]
        swc_z = swc_node_attrs["z"]

        branch_indices = get_swc_node_branch_indices(cell.xyzr, swc_x, swc_y, swc_z)
        assert len(branch_indices) > 0, "No branch found for swc node %d at (%f, %f, %f)" % (swc_node, swc_x, swc_y, swc_z)

        has_successors = len(list(swc_graph.successors(swc_node))) > 0
        has_multiple_branches = len(branch_indices) > 1
        swc_xyz = np.array([swc_x, swc_y, swc_z])

        # Case 1: node has no successors.
        # This means it is a tip, so we just take the nearest compartment in that branch.
        if not has_successors:
            assert not has_multiple_branches, "Node %d has no successors but maps to multiple branches: %s" % (swc_node, branch_indices)
            print("Swc node %d is a tip, mapping to branch %d" % (swc_node, branch_indices[0]))
            branch_idx = branch_indices[0]
            global_comp_idx, _ = get_nearest_loc_in_branch(cell, swc_xyz, branch_idx)
            node_comp_matrix[swc_node - 1, global_comp_idx] = 1.0

        # Case 2: node is a branch point, so we take a weighted combination of connected compartments.
        # This uses the code snippet from: https://github.com/jaxleyverse/jaxley/discussions/724
        elif has_multiple_branches:
            print('Swc node %d is a branchpoint, mapping to branches %s' % (swc_node, branch_indices))

            # get the branchpoint index associated with this swc node
            branchpoint_idx = get_matching_branchpoint_index(cell, swc_xyz)
            
            # Find all compartments that are connected to this branchpoint.
            indices = np.where(np.asarray(cell._comp_edges[["sink", "source"]]) == branchpoint_idx)[0]
            neighboring_comps = cell._comp_edges.loc[indices]
            neighboring_comps = neighboring_comps.query("type in [3, 4]")
            neighbor_comps = neighboring_comps["source"].to_numpy()


            # Compute the weights -> A branchpoint is more strongly affected by compartments that are more strongly
            # "coupled" to it, i.e., that have a higher conductance.
            # comp_edges = cell._comp_edges
            cell.set_scope("global")
            
            ordered_edge = neighboring_comps["ordered"].to_numpy()

            comp_source_g_a = 1 / cell.comp(neighbor_comps).nodes["resistive_load_out"] * cell.comp(neighbor_comps).nodes["axial_resistivity"]
            comp_sink_g_a = 1 / cell.comp(neighbor_comps).nodes["resistive_load_in"] * cell.comp(neighbor_comps).nodes["axial_resistivity"]
            comp_g_a = np.stack([comp_sink_g_a, comp_source_g_a], axis=1)

            conds_to_take = 1 - ordered_edge[:, None]
            conds_c2bp = np.take_along_axis(comp_g_a, conds_to_take, axis=1)
            weights = conds_c2bp / np.sum(conds_c2bp)

            node_comp_matrix[swc_node - 1, neighbor_comps] = np.squeeze(weights)

        # Case 3: node is an internal node with a single branch.
        # We take a weighted combination of the two nearest compartments within that branch.
        else:
            print("Swc node %d is an internal node, mapping to branch %d" % (swc_node, branch_indices[0]))
            assert not has_multiple_branches, "Node %d is an internal node but maps to multiple branches: %s" % (swc_node, branch_indices)

            branch_idx = branch_indices[0]
            # Now find the two nearest compartments in that branch
            nearest_comp_idxs, dists = get_nearest_loc_in_branch(
                cell, swc_xyz, branch_idx, nearest_k=2
            )

            weights = 1 / (dists + 1e-12)
            weights = weights / np.sum(weights)
            node_comp_matrix[swc_node - 1, nearest_comp_idxs] = weights

    return node_comp_matrix
    ```
    
    

[michaeldeistler]

On a high level, your code makes sense and looks good to me. I don't understand what you mean by of the indexing into the [2 x num_neighbors] array of conductances. though.

Michael