IntEnum for treenode-connector relationships

navis-org · Jul 28, 2023 · 53d811a · 53d811a
1 parent 40e6f96
commit 53d811a
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 4 deletions.
diff --git a/docs/source/whats_new.rst b/docs/source/whats_new.rst
@@ -15,6 +15,11 @@ repository.
    * - Version
      - Date
      -
+     - dev
+     - ongoing
+     - - :class:`navis.NodeConnectorRelation` is an :class:`enum.IntEnum`
+         encoding relationships between (tree)nodes and connector nodes,
+         used in neurons' connector tables.
    * - 1.5.0
      - 27/07/23
      - - BREAKING: dropped support for Python 3.7

diff --git a/navis/core/__init__.py b/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']
diff --git a/navis/core/base.py b/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,27 @@
     pint.Quantity([])
 
 
+class NodeConnectorRelation(IntEnum):
+    """An integer describing a (tree)node-connector relationship.
+
+    i.e. "the (tree)node is <NodeConnectorType> the connector node"
+
+    Based on 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 Neuron(x: Union[nx.DiGraph, str, pd.DataFrame, 'TreeNeuron', 'MeshNeuron'],
            **metadata):
     """Constructor for Neuron objects. Depending on the input, either a
@@ -421,7 +443,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 +465,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}')