IntEnum for treenode-connector relationships

docs/source/whats_new.rst

@@ -25,6 +25,9 @@ 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)
+          - :class:`navis.NodeConnectorRelation` is an :class:`enum.IntEnum`
+            encoding relationships between (tree)nodes and connector nodes,
+            used in neurons' connector tables.
           - new :class:`navis.NeuronConnector` class for creating connectivity graphs
          from groups neurons with consistent connector IDs.
        - Improvements:

navis/core/__init__.py

@@ -12,7 +12,7 @@
 #    GNU General Public License for more details.
 
 from .volumes import Volume
-from .base import Neuron, BaseNeuron
+from .base import Neuron, BaseNeuron, NodeConnectorRelation
 from .skeleton import TreeNeuron
 from .mesh import MeshNeuron
 from .dotprop import Dotprops
@@ -25,4 +25,5 @@
 NeuronObject = Union[NeuronList, TreeNeuron, BaseNeuron, MeshNeuron]
 
 __all__ = ['Volume', 'Neuron', 'BaseNeuron', 'TreeNeuron', 'MeshNeuron',
-           'Dotprops', 'VoxelNeuron', 'NeuronList', 'make_dotprops']
+           'Dotprops', 'VoxelNeuron', 'NeuronList', 'make_dotprops',
+           'NodeConnectorRelation']

navis/core/base.py

@@ -12,6 +12,7 @@
 #    GNU General Public License for more details.
 
 import copy
+from enum import IntEnum
 import hashlib
 import numbers
 import pint
@@ -45,6 +46,43 @@
     pint.Quantity([])
 
 
+class NodeConnectorRelation(IntEnum):
+    """An integer describing a (tree)node-connector relationship.
+
+    i.e. "the (tree)node is <NodeConnectorType> the connector node"
+
+    Inspired by the `CATMAID link types`_.
+    A node PRESYNAPTIC_TO a connector is an output site.
+    A node POSTSYNAPTIC_TO a connector is an input site.
+
+    .. _`CATMAID link types`: https://github.com/catmaid/CATMAID/blob/2964e04e6e9772aff5d305e72c1b878030fe0e25/django/applications/catmaid/control/link.py#L16
+    """
+    PRESYNAPTIC_TO = 0
+    POSTSYNAPTIC_TO = 1
+    # ABUTTING = 2
+    # GAPJUNCTION_WITH = 3
+    # TIGHTJUNCTION_WITH = 4
+    # DESMOSOME_WITH = 5
+    # ATTACHMENT_TO = 6
+    # CLOSE_TO = 7
+
+    def is_output_node(self) -> bool:
+        """Whether the node is sending output via the connector"""
+        return self == NodeConnectorRelation.PRESYNAPTIC_TO
+
+    def is_input_node(self) -> bool:
+        """Whether the node is receiving input via the connector"""
+        return self == NodeConnectorRelation.POSTSYNAPTIC_TO
+
+    def is_output_connector(self) -> bool:
+        """Whether the connector is sending output to the node"""
+        return self.is_input_node()
+
+    def is_input_connector(self) -> bool:
+        """Whether the connector is receiving input from the node"""
+        return self.is_output_node()
+
+
 def Neuron(x: Union[nx.DiGraph, str, pd.DataFrame, 'TreeNeuron', 'MeshNeuron'],
            **metadata):
     """Constructor for Neuron objects. Depending on the input, either a
@@ -421,7 +459,8 @@ def presynapses(self):
             raise ValueError('No connector table found.')
         # Make an educated guess what presynapses are
         types = self.connectors['type'].unique()
-        pre = [t for t in types if 'pre' in str(t) or t in [0, "0"]]
+        pre_int = NodeConnectorRelation.PRESYNAPTIC_TO.value
+        pre = [t for t in types if 'pre' in str(t) or t in [pre_int, str(pre_int)]]
 
         if len(pre) == 0:
             logger.debug(f'Unable to find presynapses in types: {types}')
@@ -442,7 +481,8 @@ def postsynapses(self):
             raise ValueError('No connector table found.')
         # Make an educated guess what presynapses are
         types = self.connectors['type'].unique()
-        post = [t for t in types if 'post' in str(t) or t in [1, "1"]]
+        post_int = NodeConnectorRelation.POSTSYNAPTIC_TO.value
+        post = [t for t in types if 'post' in str(t) or t in [post_int, str(post_int)]]
 
         if len(post) == 0:
             logger.debug(f'Unable to find postsynapses in types: {types}')