test and fix NeuronConnector

navis/connectivity/adjacency.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional
+from typing import Iterable, Optional, NamedTuple
 from ..core import TreeNeuron
 import networkx as nx
 import pandas as pd
@@ -9,6 +9,13 @@
 
 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.
@@ -17,15 +24,18 @@ class NeuronConnector:
 
     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()
-        self.inputs = dict()
-        self.outputs = 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)
@@ -36,6 +46,8 @@ def __len__(self) -> int:
     def add_neurons(self, nrns: Iterable[TreeNeuron]):
         """Add several neurons to the connector.
 
+        All neurons must have unique names.
+
         Parameters
         ----------
         nrns : Iterable[TreeNeuron]
@@ -51,6 +63,8 @@ def add_neurons(self, nrns: Iterable[TreeNeuron]):
     def add_neuron(self, nrn: TreeNeuron):
         """Add a single neuron to the connector.
 
+        All neurons must have unique names.
+
         Parameters
         ----------
         nrn : TreeNeuron
@@ -75,20 +89,20 @@ def add_neuron(self, nrn: TreeNeuron):
         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 == 0:
-                self.outputs.setdefault(row.connector_id, []).append((nrn.name, row.node_id))
-            elif row.type == 1:
-                if row.connector_id in self.inputs:
+            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.inputs[row.connector_id] = (nrn.name, row.node_id)
+                self.conn_inputs[row.connector_id] = (nrn.name, row.node_id)
 
         return self
 
-    def edges(self, include_other=True):
+    def edges(self, include_other=True) -> Iterable[Edge]:
         """Iterate through all synapse edges.
 
         Parameters
@@ -103,14 +117,14 @@ def edges(self, include_other=True):
         tuple[int, str, str, int, int]
             Connector ID, source name, target name, source treenode, target treenode.
         """
-        for conn_id in set(self.inputs).union(self.outputs):
-            src, src_node = self.inputs.get(conn_id, (OTHER, None))
+        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.outputs.get(conn_id, [(OTHER, None)]):
+            for tgt, tgt_node in self.conn_outputs.get(conn_id, [(OTHER, None)]):
                 if tgt_node is None and not include_other:
                     continue
-                yield (conn_id, src, tgt, src_node, tgt_node)
+                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.
@@ -135,9 +149,9 @@ def to_adjacency(self, include_other=True) -> pd.DataFrame:
         if include_other:
             index.append(OTHER)
         data = np.zeros((len(index), len(index)), np.uint64)
-        df = pd.DataFrame(data, index)
+        df = pd.DataFrame(data, index, index)
         for _, src, tgt, _, _ in self.edges(include_other):
-            df[src, tgt] += 1
+            df[tgt][src] += 1
 
         return df
 
@@ -169,9 +183,9 @@ def to_digraph(self, include_other=True) -> nx.DiGraph:
 
         g.graph["connector_xyz"] = self.connector_xyz
         headers = {
-            "connector_id": pd.UInt64Dtype,
-            "pre_node": pd.UInt64Dtype,
-            "post_node": pd.UInt64Dtype,
+            "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):

tests/test_connectivity.py

@@ -1,3 +1,7 @@
+import pytest
+import pandas as pd
+import numpy as np
+
 import navis
 from navis.connectivity import NeuronConnector
 
@@ -12,11 +16,106 @@ def test_neuron_connector():
 
     adj = conn.to_adjacency()
     assert len(adj) == len(nrns) + 1
+    assert adj.to_numpy().sum() == 0
 
     dg = conn.to_digraph()
     assert dg.number_of_nodes() == len(nrns) + 1
+    assert dg.number_of_edges() == 0
 
     mdg = conn.to_multidigraph()
     assert mdg.number_of_nodes() == dg.number_of_nodes()
+    assert mdg.number_of_edges() == 0
+
+
+def path_neuron(path: list[int]):
+    nrn = navis.TreeNeuron(None)
+    nrn.name = "".join(str(n) for n in path)
+    dtypes = {
+        "node_id": np.uint64,
+        "parent_id": np.int64,
+        "x": float,
+        "y": float,
+        "z": float,
+    }
+    prev = -1
+    rows = []
+    for n in path:
+        rows.append([n, prev, 0, 0, 0])
+        prev = n
+    df = pd.DataFrame(rows, columns=list(dtypes)).astype(dtypes)
+    nrn.nodes = df
+    return nrn
+
+
+def add_connectors(
+    nrn: navis.TreeNeuron,
+    incoming: list[tuple[int, int]],
+    outgoing: list[tuple[int, int]],
+):
+    """Add connectors to a neuron.
+
+    Parameters
+    ----------
+    incoming : list[tuple[int, int]]
+        List of connector_id, node_id pairs.
+    outgoing : list[tuple[int, int]]
+        List of connector_id, node_id pairs
+    """
+    dtypes = {
+        "connector_id": np.uint64,
+        "node_id": np.uint64,
+        "x": float,
+        "y": float,
+        "z": float,
+        "type": np.uint64,
+    }
+    rows = []
+    for tc_rel, ids in [(1, incoming), (0, outgoing)]:
+        for conn, node in ids:
+            rows.append([conn, node, 0, 0, 0, tc_rel])
+    df = pd.DataFrame(rows, columns=list(dtypes)).astype(dtypes)
+    nrn.connectors = df
+
+
+@pytest.fixture
+def simple_network():
+    """
+    2 neurons, "456" and "789".
+    4 connectors, 0/1/2/3.
 
-    # todo: check edges
+        4     7
+        |     |
+    =0> 5 =1> 8 =2>
+        |     |
+        6 =3> 9
+    """
+    n456 = path_neuron([4, 5, 6])
+    add_connectors(n456, [(0, 5)], [(1, 5), (3, 6)])
+    n789 = path_neuron([7, 8, 9])
+    add_connectors(n789, [(1, 8), (3, 9)], [(2, 8)])
+    return [n456, n789]
+
+
+def test_neuron_connector_synthetic(simple_network):
+    nconn = NeuronConnector(simple_network)
+
+    adj = nconn.to_adjacency()
+    assert len(adj) == len(simple_network) + 1
+    assert adj.to_numpy().sum() == 4
+
+    expected_edges = sorted([
+        ("__OTHER__", "456"),
+        ("456", "789"),
+        ("456", "789"),
+        ("789", "__OTHER__"),
+    ])
+
+    dg = nconn.to_digraph()
+    assert dg.number_of_nodes() == len(simple_network) + 1
+    assert dg.number_of_edges() == 3
+    assert set(dg.edges()) == set(expected_edges)
+
+    mdg = nconn.to_multidigraph()
+    assert mdg.number_of_nodes() == dg.number_of_nodes()
+    assert mdg.number_of_edges() == 4
+    assert sorted(mdg.edges()) == expected_edges