extracting part of a cell, network etc.

[jnsbck]

Hi,

I have a morphology

cell = jx.read_swc(f"morphology.swc", nseg=4, max_branch_len=300.0, min_radius=5.0)
cell.insert(HH())

but I want to simulate only a specific branch. I don't care about anything else that is attached. Since cell.branch(0) returns an instance of BranchView, the following does not work.

branch = cell.branch(0)

branch.comp(0).stimulate(current)
branch.record()

voltage = jx.integrate(branch)

Is this somehow possible, since I was not able to find an easy way to do this? If not it would be awesome to have this functionality.

i.e. by adding a .copy() method to the BranchView object, that returns a new instance of Branch

branch = cell.branch(0).copy()

branch..comp(0).stimulate(current)
branch.record()

voltage = jx.integrate(branch)

Same would be applicable to Cell and Comp. Would be happy to provide a PR if this functionality is otherwise non-trivially obtained.

Thanks for your input on this. :)

[michaeldeistler]

Nope this is not possible. I agree that this would be nice but I do not think that I will find the time to do this anytime soon.

If you can make a PR this would be super awesome.

[jnsbck]

Cool, will make a PR then :)

[michaeldeistler]

Related to #354

[jnsbck]

This can be done in a limited fashion with #447 and should be improved down the line.

[jnsbck]

the following works if #487 is merged!

def infer_module_type_from_inds(idxs: pd.DataFrame) -> str:
    nuniques = idxs[["cell_index", "branch_index", "comp_index"]].nunique()
    nuniques.index = ["cell", "branch", "compartment"]
    nuniques = pd.concat([pd.Series({"network": 1}), nuniques])
    return_type = nuniques.loc[nuniques == 1].index[-1]
    return return_type


def build_module_scaffold(
    idxs: pd.DataFrame,
    return_type: Optional[str] = None,
    parent_branches: Optional[List[np.ndarray]] = None,
) -> Union[jx.Network, jx.Cell, jx.Branch, jx.Compartment]:
    """Builds a skeleton module from a DataFrame of indices.
    This is useful for instantiating a module that can be filled with data later.
    Args:
        idxs: DataFrame containing cell_index, branch_index, comp_index, i.e.
            Module.nodes or View.view.
        return_type: Type of module to return. If None, the type is inferred from the
            number of unique values in the indices. I.e. only 1 unique cell_index
                and 1 unique branch_index -> return_type = "jx.Branch".
    Returns:
        A skeleton module with the correct number of compartments, branches, cells, or
        networks."""
    return_types = ["compartment", "branch", "cell", "network"]
    build_cache = {k: [] for k in return_types}

    if return_type is None:  # infer return type from idxs
        return_type = infer_module_type_from_inds(idxs)

    comp = jx.Compartment()
    build_cache["compartment"] = [comp]

    if return_type in return_types[1:]:
        nsegs = idxs["branch_index"].value_counts().iloc[0]
        branch = jx.Branch([comp for _ in range(nsegs)])
        build_cache["branch"] = [branch]

    if return_type in return_types[2:]:
        for cell_id, cell_groups in idxs.groupby("cell_index"):
            num_branches = cell_groups["branch_index"].nunique()
            default_parents = np.arange(num_branches) - 1  # ignores morphology
            parents = (
                default_parents if parent_branches is None else parent_branches[cell_id]
            )
            cell = jx.Cell([branch] * num_branches, parents)
            build_cache["cell"].append(cell)

    if return_type in return_types[3:]:
        build_cache["network"] = [jx.Network(build_cache["cell"])]

    module = build_cache[return_type][0]
    build_cache.clear()
    return module

def to_module(view: View, reset_index: bool = True) -> "Module":
    """Extract part of a module and return a copy of its View or a new module.

    This can be used to call `jx.integrate` on part of a Module.

    Args:
        reset_index: if True, the indices of the new module are reset to start from 0.
        as_module: if True, a new module is returned instead of a View.

    Returns:
        A part of the module or a copied view of it."""
    view = deepcopy(view)
    if reset_index:
        view.nodes.reset_index(drop=True, inplace=True)
        view.edges.reset_index(drop=True, inplace=True)
        # TODO: also re-enumerate cell,branch,comp indices in nodes and edges

    testnodes = view.nodes.copy()
    testnodes.rename({"global_"+k:k for k in ["cell_index", "branch_index", "comp_index"]}, axis=1, inplace=True)
    mod_type = infer_module_type_from_inds(testnodes)
    module = build_module_scaffold(testnodes, mod_type)
    module.__dict__.update(view.__dict__)
    return module