From 2f28845c7424b64d84ee1fe81f38d6700c29ba5d Mon Sep 17 00:00:00 2001 From: Chris Barnes Date: Wed, 23 Nov 2022 14:13:21 +0000 Subject: [PATCH] SWC io: DFS-sort nodes, optionally do not reindex Previously, nodes were sorted by parent index, which would ensure that root nodes appear before non-roots but did not guarantee that all children are defined after their parent. Now, nodes are sorted in depth-first order (addressing their children in order of their index). Additionally, added the option of return_node_map=None, which does not re-index nodes before writing. This is useful when exporting a SWC files of several neurons along with where node IDs must persist their global uniqueness. --- navis/io/swc_io.py | 105 ++++++++++++++++++++++++++++++++++----------- 1 file changed, 80 insertions(+), 25 deletions(-) diff --git a/navis/io/swc_io.py b/navis/io/swc_io.py index 8ef14e38..9f8abffa 100644 --- a/navis/io/swc_io.py +++ b/navis/io/swc_io.py @@ -453,9 +453,12 @@ def write_swc(x: 'core.NeuronObject', this might drop synapses (i.e. in case of multiple pre- and/or postsynapses on a single node)! ``labels`` must be ``True`` for this to have any effect. - return_node_map : bool + return_node_map : bool, optional If True, will return a dictionary mapping the old node ID to the new reindexed node IDs in the file. + If False (default), will reindex nodes but not return the mapping. + If None, will not reindex nodes (i.e. their current IDs will be written). + Returns ------- @@ -535,7 +538,7 @@ def _write_swc(x: 'core.TreeNeuron', write_meta: Union[bool, List[str], dict] = True, labels: Union[str, dict, bool] = True, export_connectors: bool = False, - return_node_map: bool = False) -> None: + return_node_map: Optional[bool] = False) -> None: """Write single TreeNeuron to file.""" # Generate SWC table res = make_swc_table(x, @@ -589,6 +592,45 @@ def _write_swc(x: 'core.TreeNeuron', return node_map +def sort_swc(df: pd.DataFrame, roots, sort_children=True, inplace=False): + """Depth-first search tree to ensure parents are always defined before children.""" + children = defaultdict(list) + node_id_to_orig_idx = dict() + for row in df.itertuples(): + child = row.node_id + parent = row.parent_id + children[parent].append(child) + node_id_to_orig_idx[child] = row.index + + if sort_children: + to_visit = sorted(roots, reverse=True) + else: + to_visit = list(roots)[::-1] + + idx = 0 + order = np.full(len(df), np.nan) + count = 0 + while to_visit: + node_id = to_visit.pop() + order[node_id_to_orig_idx[node_id]] = count + cs = children.pop(order[-1], []) + if sort_children: + to_visit.append(sorted(sort_children, reverse=True)) + else: + to_visit.append(cs[::-1]) + count += 1 + + # undefined behaviour if any nodes are not reachable from the given roots + + if not inplace: + df = df.copy() + + df["_order"] = order + df.sort_values("_order", inplace=True) + df.drop(columns=["_order"]) + return df + + def make_swc_table(x: 'core.TreeNeuron', labels: Union[str, dict, bool] = None, export_connectors: bool = False, @@ -606,17 +648,20 @@ def make_swc_table(x: 'core.TreeNeuron', str : column name in node table dict: must be of format {node_id: 'label', ...}. - bool: if True, will generate automatic labels, if False all nodes have label "0". + bool: if True, will generate automatic labels for branches ("5") and ends ("6"), + soma where labelled ("1"), and optionally connectors (see below). + If False (or for all nodes not labelled as above) all nodes have label "0". export_connectors : bool, optional - If True, will label nodes with pre- ("7") and - postsynapse ("8"). Because only one label can be given - this might drop synapses (i.e. in case of multiple - pre- or postsynapses on a single node)! ``labels`` - must be ``True`` for this to have any effect. - return_node_map : bool + If True, will label nodes with only presynapses ("7"), + only postsynapses ("8"), or both ("9"). + This overrides branch/end/soma labels. + ``labels`` must be ``True`` for this to have any effect. + return_node_map : bool, optional If True, will return a dictionary mapping the old node ID to the new reindexed node IDs in the file. + If False, will remap IDs but not return the mapping. + If None, will not remap IDs. Returns ------- @@ -625,8 +670,8 @@ def make_swc_table(x: 'core.TreeNeuron', Only if ``return_node_map=True``. """ - # Work on a copy - swc = x.nodes.copy() + # Work on a copy sorted in depth-first order + swc = sort_swc(x.nodes, x.root, inplace=False) # Add labels swc['label'] = 0 @@ -642,31 +687,41 @@ def make_swc_table(x: 'core.TreeNeuron', if not isinstance(x.soma, type(None)): soma = utils.make_iterable(x.soma) swc.loc[swc.node_id.isin(soma), 'label'] = 1 + if export_connectors: # Add synapse label pre_ids = x.presynapses.node_id.values post_ids = x.postsynapses.node_id.values - swc.loc[swc.node_id.isin(pre_ids), 'label'] = 7 - swc.loc[swc.node_id.isin(post_ids), 'label'] = 8 - - # Sort such that the parent is always before the child - swc.sort_values('parent_id', ascending=True, inplace=True) - # Reset index - swc.reset_index(drop=True, inplace=True) + is_pre = swc["node_id"].isin(pre_ids) + swc.loc[is_pre, 'label'] = 7 - # Generate mapping - new_ids = dict(zip(swc.node_id.values, swc.index.values + 1)) + is_post = swc["node"].isin(post_ids) + swc.loc[is_post, 'label'] = 8 - swc['node_id'] = swc.node_id.map(new_ids) - # Lambda prevents potential issue with missing parents - swc['parent_id'] = swc.parent_id.map(lambda x: new_ids.get(x, -1)) + is_both = np.logical_and(is_pre, is_post) + swc.loc[is_both, 'label'] = 9 - # Get things in order + # Order columns swc = swc[['node_id', 'label', 'x', 'y', 'z', 'radius', 'parent_id']] - # Make sure radius has no `None` + # Make sure radius has no `None` or negative swc['radius'] = swc.radius.fillna(0) + swc['radius'][swc['radius'] < 0] = 0 + + if return_node_map is not None: + # remap IDs + + # Reset index + swc.reset_index(drop=True, inplace=True) + + # Generate mapping + new_ids = dict(zip(swc.node_id.values, swc.index.values + 1)) + + swc['node_id'] = swc.node_id.map(new_ids) + # Lambda prevents potential issue with missing parents + swc['parent_id'] = swc.parent_id.map(lambda x: new_ids.get(x, -1)) + # Adjust column titles swc.columns = ['PointNo', 'Label', 'X', 'Y', 'Z', 'Radius', 'Parent']