Merge pull request #133 from clbarnes/nrn2adj2

NeuronConnector (clean git history)

docs/source/api.rst

@@ -501,6 +501,11 @@ Connectivity
 ++++++++++++
 Collection of functions to work with graphs and adjacency matrices.
 
+.. autosummary::
+    :toctree: generated/
+
+    navis.NeuronConnector
+
 Graphs
 ------
 Functions to convert neurons and networkx to iGraph or networkX graphs.

docs/source/whats_new.rst

@@ -25,6 +25,8 @@ repository.
             for format specs and benchmarks)
           - new :func:`navis.read_nml` function to read single NML file (complements
             existing :func:`navis.read_nmx` files which are collections of NMLs)
+          - new :class:`navis.NeuronConnector` class for creating connectivity graphs
+         from groups neurons with consistent connector IDs.
        - Improvements:
           - made adding recordings to ``CompartmentModel`` faster
           - improved logic for splitting NBLAST across cores

navis/connectivity/__init__.py

@@ -13,8 +13,10 @@
 
 from .predict import cable_overlap
 from .matrix_utils import group_matrix
+from .adjacency import NeuronConnector
 from .cnmetrics import connectivity_sparseness
 from .similarity import connectivity_similarity, synapse_similarity
 
 __all__ = ['connectivity_sparseness', 'cable_overlap',
-           'connectivity_similarity', 'synapse_similarity']
+           'connectivity_similarity', 'synapse_similarity',
+           'NeuronConnector']

navis/connectivity/adjacency.py

@@ -0,0 +1,236 @@
+from typing import Iterable, Optional, NamedTuple
+from ..core import TreeNeuron
+import networkx as nx
+import pandas as pd
+import numpy as np
+from ..config import get_logger
+
+logger = get_logger(__name__)
+
+OTHER = "__OTHER__"
+
+
+class Edge(NamedTuple):
+    connector_id: int
+    source_name: str
+    target_name: str
+    source_node: Optional[int]
+    target_node: Optional[int]
+
+
+class NeuronConnector:
+    """Class which creates a connectivity graph from a set of neurons.
+
+    Connectivity is determined by shared IDs in the ``connectors`` table.
+
+    Add neurons with the `add_neuron` and `add_neurons` methods.
+    Alternatively, supply an iterable of neurons in the constructor.
+    Neurons must have unique names.
+
+    See the `to_(multi)digraph` method for output.
+    """
+
+    def __init__(self, nrns: Optional[Iterable[TreeNeuron]] = None) -> None:
+        self.neurons = dict()
+        self.connector_xyz = dict()
+        # connectors and the treenodes presynaptic to them
+        self.conn_inputs = dict()
+        # connectors and the treenodes postsynaptic to them
+        self.conn_outputs = dict()
+
+        if nrns is not None:
+            self.add_neurons(nrns)
+
+    def __len__(self) -> int:
+        return len(self.neurons)
+
+    def add_neurons(self, nrns: Iterable[TreeNeuron]):
+        """Add several neurons to the connector.
+
+        All neurons must have unique names.
+
+        Parameters
+        ----------
+        nrns : Iterable[TreeNeuron]
+
+        Returns
+        -------
+        Modified connector.
+        """
+        for nrn in nrns:
+            self.add_neuron(nrn)
+        return self
+
+    def add_neuron(self, nrn: TreeNeuron):
+        """Add a single neuron to the connector.
+
+        All neurons must have unique names.
+
+        Parameters
+        ----------
+        nrn : TreeNeuron
+
+        Returns
+        -------
+        Modified connector.
+        """
+        if nrn.name in self.neurons:
+            logger.warning(
+                "Neuron with name %s has already been added to NeuronConnector. "
+                "These will occupy the same node in the graph, "
+                "but have connectors from both.",
+                nrn.name
+            )
+
+        self.neurons[nrn.name] = nrn
+        if nrn.connectors is None:
+            logger.warning("Neuron with name %s has no connector information", nrn.name)
+            return self
+
+        for row in nrn.connectors.itertuples():
+            # connector_id, node_id, x, y, z, is_input
+            self.connector_xyz[row.connector_id] = (row.x, row.y, row.z)
+            if row.type == 1:
+                self.conn_outputs.setdefault(row.connector_id, []).append((nrn.name, row.node_id))
+            elif row.type == 0:
+                if row.connector_id in self.conn_inputs:
+                    logger.warning(
+                        "Connector with ID %s has multiple inputs: "
+                        "connector tables are probably inconsistent",
+                        row.connector_id
+                    )
+                self.conn_inputs[row.connector_id] = (nrn.name, row.node_id)
+
+        return self
+
+    def edges(self, include_other=True) -> Iterable[Edge]:
+        """Iterate through all synapse edges.
+
+        Parameters
+        ----------
+        include_other : bool, optional
+            Include edges for which only one partner is known, by default True.
+            If included, the name of the unknown partner will be ``"__OTHER__"``,
+            and the treenode ID will be None.
+
+        Yields
+        ------
+        tuple[int, str, str, int, int]
+            Connector ID, source name, target name, source treenode, target treenode.
+        """
+        for conn_id in set(self.conn_inputs).union(self.conn_outputs):
+            src, src_node = self.conn_inputs.get(conn_id, (OTHER, None))
+            if src_node is None and not include_other:
+                continue
+            for tgt, tgt_node in self.conn_outputs.get(conn_id, [(OTHER, None)]):
+                if tgt_node is None and not include_other:
+                    continue
+                yield Edge(conn_id, src, tgt, src_node, tgt_node)
+
+    def to_adjacency(self, include_other=True) -> pd.DataFrame:
+        """Create an adjacency matrix of neuron connectivity.
+
+        Parameters
+        ----------
+        include_other : bool, optional
+            Whether to include a node called ``"__OTHER__"``,
+            which represents all unknown partners.
+            By default True.
+            This can be helpful when calculating a neuron's input fraction,
+            but cannot be used for output fractions if synapses are polyadic.
+
+        Returns
+        -------
+        pandas.DataFrame
+            Row index is source neuron name,
+            column index is target neuron name,
+            cells are the number of synapses from source to target.
+        """
+        index = list(self.neurons)
+        if include_other:
+            index.append(OTHER)
+        data = np.zeros((len(index), len(index)), np.uint64)
+        df = pd.DataFrame(data, index, index)
+        for _, src, tgt, _, _ in self.edges(include_other):
+            df[tgt][src] += 1
+
+        return df
+
+    def to_digraph(self, include_other=True) -> nx.DiGraph:
+        """Create a graph of neuron connectivity.
+
+        Parameters
+        ----------
+        include_other : bool, optional
+            Whether to include a node called ``"__OTHER__"``,
+            which represents all unknown partners.
+            By default True.
+            This can be helpful when calculating a neuron's input fraction,
+            but cannot be used for output fractions if synapses are polyadic.
+
+        Returns
+        -------
+        nx.DiGraph
+            The graph has data ``{"connector_xyz": {connector_id: (x, y, z), ...}}``.
+            The nodes have data ``{"neuron": tree_neuron}``.
+            The edges have data ``{"connectors": data_frame, "weight": n_connectors}``,
+            where the connectors data frame has columns
+            "connector_id", "pre_node", "post_node".
+        """
+        g = nx.DiGraph()
+        g.add_nodes_from((k, {"neuron": v}) for k, v in self.neurons.items())
+        if include_other:
+            g.add_node(OTHER, neuron=None)
+
+        g.graph["connector_xyz"] = self.connector_xyz
+        headers = {
+            "connector_id": pd.UInt64Dtype(),
+            "pre_node": pd.UInt64Dtype(),
+            "post_node": pd.UInt64Dtype(),
+        }
+        edges = dict()
+        for conn_id, src, tgt, src_node, tgt_node in self.edges(include_other):
+            edges.setdefault((src, tgt), []).append([conn_id, src_node, tgt_node])
+
+        for (src, tgt), rows in edges.items():
+            df_tmp = pd.DataFrame(rows, columns=list(headers), dtype=object)
+            df = df_tmp.astype(headers, copy=False)
+            g.add_edge(src, tgt, connectors=df, weight=len(df))
+
+        return g
+
+    def to_multidigraph(self, include_other=True) -> nx.MultiDiGraph:
+        """Create a graph of neuron connectivity where each synapse is an edge.
+
+        Parameters
+        ----------
+        include_other : bool, optional
+            Whether to include a node called ``"__OTHER__"``,
+            which represents all unknown partners.
+            By default True.
+            This can be helpful when calculating a neuron's input fraction,
+            but cannot be used for output fractions if synapses are polyadic.
+
+        Returns
+        -------
+        nx.MultiDiGraph
+            The nodes have data ``{"neuron": tree_neuron}``.
+            The edges have data
+            ``{"pre_node": presyn_treenode_id, "post_node": postsyn_treenode_id, "xyz": connector_location, "connector_id": conn_id}``.
+        """
+        g = nx.MultiDiGraph()
+        g.add_nodes_from((k, {"neuron": v}) for k, v in self.neurons.items())
+        if include_other:
+            g.add_node(OTHER, neuron=None)
+
+        for conn_id, src, tgt, src_node, tgt_node in self.edges(include_other):
+            g.add_edge(
+                src,
+                tgt,
+                pre_node=src_node,
+                post_node=tgt_node,
+                xyz=self.connector_xyz[conn_id],
+                connector_id=conn_id,
+            )
+
+        return g

tests/conftest.py

@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 import os
 from typing import List
@@ -15,6 +16,11 @@ def data_dir():
     return Path(__file__).resolve().parent.parent / "navis" / "data"
 
 
+@pytest.fixture(scope="session")
+def fixture_dir():
+    return Path(__file__).resolve().parent / "fixtures"
+
+
 @pytest.fixture(
     params=["Path", "pathstr", "swcstr", "textbuffer", "rawbuffer", "DataFrame"]
 )
@@ -115,3 +121,26 @@ def treeneuron_dfs(swc_paths, synapses_paths):
         neuron.connectors = pd.read_csv(syn_path)
         out.append(neuron)
     return out
+
+
+@pytest.fixture
+def neuron_connections(fixture_dir: Path):
+    expected_jso = json.loads(
+        fixture_dir.joinpath("neuron_connector", "expected.json").read_text()
+    )
+    expected = {
+        int(pre): {
+            int(post): n for post, n in d.items()
+        } for pre, d in expected_jso.items()
+    }
+
+    nl = navis.read_json(
+        str(fixture_dir.joinpath("neuron_connector", "network.json"))
+    )
+    nrns = []
+    for nrn in nl:
+        nrn.name = f"skeleton {nrn.id}"
+        nrn.id = int(nrn.id)
+        nrns.append(nrn)
+
+    return nrns, expected

tests/fixtures/neuron_connector/expected.json

@@ -0,0 +1,50 @@
+{
+  "2582391": {
+    "2582391": 0,
+    "6557581": 10,
+    "8198416": 12,
+    "8198513": 14,
+    "8252067": 41,
+    "11051276": 55
+  },
+  "6557581": {
+    "2582391": 9,
+    "6557581": 0,
+    "8198416": 8,
+    "8198513": 9,
+    "8252067": 1,
+    "11051276": 9
+  },
+  "8198416": {
+    "2582391": 8,
+    "6557581": 8,
+    "8198416": 0,
+    "8198513": 14,
+    "8252067": 9,
+    "11051276": 10
+  },
+  "8198513": {
+    "2582391": 7,
+    "6557581": 13,
+    "8198416": 16,
+    "8198513": 0,
+    "8252067": 9,
+    "11051276": 5
+  },
+  "8252067": {
+    "2582391": 0,
+    "6557581": 3,
+    "8198416": 0,
+    "8198513": 1,
+    "8252067": 0,
+    "11051276": 1
+  },
+  "11051276": {
+    "2582391": 0,
+    "6557581": 2,
+    "8198416": 1,
+    "8198513": 1,
+    "8252067": 0,
+    "11051276": 0
+  }
+}