Merge pull request #150 from NeuralEnsemble/feat/add-sectionise-method

Feat/add sectionise method

neuroml/nml/helper_methods.py

@@ -1969,6 +1969,111 @@ def add_unbranched_segments(
 
         self.reorder_segment_groups()
         return self.get_segment_group(group_id)
+
+    def create_unbranched_segment_group_branches(self, root_segment_id: int, use_convention: bool=True):
+        """Organise the segments of the cell into new segment groups that each
+        form a single contiguous unbranched cell branch.
+
+        Note that the first segment (root segment) of a branch must have a proximal
+        point that connects it to the rest of the neuronal morphology. If, when
+        constructing these branches, a root segment is found that does not include
+        a proximal point, one will be added using the `get_actual_proximal` method.
+
+        No other changes will be made to any segments, or to any pre-existing
+        segment groups.
+
+        :param root_segment_id: id of segment considered the root of the tree,
+            generally the first soma segment
+        :type root_segment_id: int
+        :param use_convention: toggle using NeuroML convention for segment groups
+        :type use_convention: bool
+        :returns: modified cell with new section groups
+        :rtype: neuroml.Cell
+
+        """
+        # get morphology tree
+        morph_tree = self.get_segment_adjacency_list()
+
+        # initialise root segment and first segment group
+        seg = self.get_segment(root_segment_id)
+        group_name = f"seg_group_{len(self.morphology.segment_groups) - 1}_seg_{seg.id}"
+        new_seg_group = self.add_unbranched_segment_group(group_name)
+
+        # run recursive function
+        self.__sectionise(root_segment_id, new_seg_group, morph_tree)
+
+
+    def __sectionise(self, root_segment_id, seg_group, morph_tree):
+        """Main recursive sectionising method.
+
+        :param root_segment_id: id of root of branch
+        :type root_segment_id: int
+        :returns: TODO
+
+        """
+        # print(f"Processing element: {root_segment_id}")
+
+        try:
+            children = morph_tree[root_segment_id]
+            # keep going while there's only one child
+            # no need to use recursion here---hits Python's recursion limits in
+            # long segment groups
+            # - if there are no children, it'll go to the except block
+            # - if there are more than one children, it'll go to the next
+            # conditional
+            while len(children) == 1:
+                seg_group.add("Member", segments=root_segment_id)
+                root_segment_id = children[0]
+                children = morph_tree[root_segment_id]
+            # if there are more than one children, we've reached the end of this
+            # segment group but not of the branch. New segment groups need to start
+            # from here.
+            if len(children) > 1:
+                # this becomes the last segment of the current segment group
+                seg_group.add("Member", segments=root_segment_id)
+
+                # each child will start a new segment group
+                for child in children:
+                    seg = self.get_segment(child)
+                    # Ensure that a proximal point is set.
+                    # This is required for the first segment of unbranched segment
+                    # groups
+                    seg.proximal = self.get_actual_proximal(seg.id)
+                    group_name = f"seg_group_{len(self.morphology.segment_groups) - 1}_seg_{seg.id}"
+                    new_seg_group = self.add_unbranched_segment_group(group_name)
+
+                    self.__sectionise(child, new_seg_group, morph_tree)
+        # if there are no children, it's a leaf node, so we just add to the current
+        # seg_group and do nothing else
+        except KeyError:
+            seg_group.add("Member", segments=root_segment_id)
+
+
+    def get_segment_adjacency_list(self):
+        """Get the adjacency list of all segments in the cell morphology.
+        Returns a dict where each key is a parent segment, and the value is the
+        list of its children segments.
+
+        Segment without children (leaf segments) are not included as parents in the
+        adjacency list.
+
+        :returns: dict with parent segments as keys and their children as values
+        :rtype: dict
+
+        """
+        # create data structure holding list of children for each segment
+        child_lists = {}
+        for segment in self.morphology.segments:
+            try:
+                parent = segment.parent.segments
+                if parent not in child_lists:
+                    child_lists[parent] = []
+                child_lists[parent].append(segment.id)
+            except AttributeError:
+                print(f"Warning: Segment: {segment} has no parent")
+
+        return child_lists
+
     ''',
     class_names=("Cell"),
 )

neuroml/nml/nml.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 #
-# Generated Fri Oct  7 17:34:43 2022 by generateDS.py version 2.41.1.
+# Generated Tue Oct 11 14:23:30 2022 by generateDS.py version 2.40.13.
 # Python 3.10.7 (main, Sep  7 2022, 00:00:00) [GCC 12.2.1 20220819 (Red Hat 12.2.1-1)]
 #
 # Command line options:
@@ -47337,6 +47337,110 @@ def add_unbranched_segments(
         self.reorder_segment_groups()
         return self.get_segment_group(group_id)
 
+    def create_unbranched_segment_group_branches(
+        self, root_segment_id: int, use_convention: bool = True
+    ):
+        """Organise the segments of the cell into new segment groups that each
+        form a single contiguous unbranched cell branch.
+
+        Note that the first segment (root segment) of a branch must have a proximal
+        point that connects it to the rest of the neuronal morphology. If, when
+        constructing these branches, a root segment is found that does not include
+        a proximal point, one will be added using the `get_actual_proximal` method.
+
+        No other changes will be made to any segments, or to any pre-existing
+        segment groups.
+
+        :param root_segment_id: id of segment considered the root of the tree,
+            generally the first soma segment
+        :type root_segment_id: int
+        :param use_convention: toggle using NeuroML convention for segment groups
+        :type use_convention: bool
+        :returns: modified cell with new section groups
+        :rtype: neuroml.Cell
+
+        """
+        # get morphology tree
+        morph_tree = self.get_segment_adjacency_list()
+
+        # initialise root segment and first segment group
+        seg = self.get_segment(root_segment_id)
+        group_name = f"seg_group_{len(self.morphology.segment_groups) - 1}_seg_{seg.id}"
+        new_seg_group = self.add_unbranched_segment_group(group_name)
+
+        # run recursive function
+        self.__sectionise(root_segment_id, new_seg_group, morph_tree)
+
+    def __sectionise(self, root_segment_id, seg_group, morph_tree):
+        """Main recursive sectionising method.
+
+        :param root_segment_id: id of root of branch
+        :type root_segment_id: int
+        :returns: TODO
+
+        """
+        # print(f"Processing element: {root_segment_id}")
+
+        try:
+            children = morph_tree[root_segment_id]
+            # keep going while there's only one child
+            # no need to use recursion here---hits Python's recursion limits in
+            # long segment groups
+            # - if there are no children, it'll go to the except block
+            # - if there are more than one children, it'll go to the next
+            # conditional
+            while len(children) == 1:
+                seg_group.add("Member", segments=root_segment_id)
+                root_segment_id = children[0]
+                children = morph_tree[root_segment_id]
+            # if there are more than one children, we've reached the end of this
+            # segment group but not of the branch. New segment groups need to start
+            # from here.
+            if len(children) > 1:
+                # this becomes the last segment of the current segment group
+                seg_group.add("Member", segments=root_segment_id)
+
+                # each child will start a new segment group
+                for child in children:
+                    seg = self.get_segment(child)
+                    # Ensure that a proximal point is set.
+                    # This is required for the first segment of unbranched segment
+                    # groups
+                    seg.proximal = self.get_actual_proximal(seg.id)
+                    group_name = f"seg_group_{len(self.morphology.segment_groups) - 1}_seg_{seg.id}"
+                    new_seg_group = self.add_unbranched_segment_group(group_name)
+
+                    self.__sectionise(child, new_seg_group, morph_tree)
+        # if there are no children, it's a leaf node, so we just add to the current
+        # seg_group and do nothing else
+        except KeyError:
+            seg_group.add("Member", segments=root_segment_id)
+
+    def get_segment_adjacency_list(self):
+        """Get the adjacency list of all segments in the cell morphology.
+        Returns a dict where each key is a parent segment, and the value is the
+        list of its children segments.
+
+        Segment without children (leaf segments) are not included as parents in the
+        adjacency list.
+
+        :returns: dict with parent segments as keys and their children as values
+        :rtype: dict
+
+        """
+        # create data structure holding list of children for each segment
+        child_lists = {}
+        for segment in self.morphology.segments:
+            try:
+                parent = segment.parent.segments
+                if parent not in child_lists:
+                    child_lists[parent] = []
+                child_lists[parent].append(segment.id)
+            except AttributeError:
+                print(f"Warning: Segment: {segment} has no parent")
+
+        return child_lists
+
     # end class Cell
 
 

neuroml/test/test_nml.py

@@ -624,3 +624,69 @@ def test_optimise_segment_group(self):
         # segment group already
         cell.optimise_segment_group("all")
         self.assertEqual(1, len(cell.get_segment_group("all").includes))
+
+    def test_create_unbranched_segment_group_branches(self):
+        "Test create_unbranched_segment_group_branches"
+        cell = component_factory("Cell", id="simple_cell")  # type: neuroml.Cell
+        cell.notes = "NeuroML cell created by CellBuilder"
+
+        # Add soma segment
+        diam = 10.0
+        soma_0 = cell.add_segment(
+            prox=[0.0, 0.0, 0.0, diam],
+            dist=[0.0, 10.0, 0.0, diam],
+            name="Seg0_soma_0",
+            group_id="soma_0",
+            seg_type="soma"
+        )
+
+        # create two unbranched segment group
+        dend_group = cell.add_unbranched_segments(
+            points=[
+                [0.0, 10.0, 0.0, diam],
+                [0.0, 20.0, 0.0, diam],
+                [0.0, 30.0, 0.0, diam],
+            ],
+            parent=soma_0,
+            fraction_along=1.0,
+            group_id="dend_group_0",
+            use_convention=True,
+            seg_type="dendrite"
+        )
+        dend_group_1 = cell.add_unbranched_segments(
+            points=[
+                [0.0, 10.0, 0.0, diam],
+                [0.0, 20.0, 5.0, diam],
+                [0.0, 30.0, 5.0, diam],
+            ],
+            parent=soma_0,
+            fraction_along=1.0,
+            group_id="dend_group_1",
+            use_convention=True,
+            seg_type="dendrite"
+        )
+        cell.optimise_segment_groups()
+
+        self.assertIsNone(cell.validate(True))
+        self.assertEqual(5, len(cell.morphology.segments))
+        self.assertEqual(7, len(cell.morphology.segment_groups))
+        print("initial")
+        for g in cell.morphology.segment_groups:
+            print(g)
+            print([x.segments for x in g.members])
+
+        # remove the two unbranched groups
+        cell.morphology.segment_groups.remove(cell.get_segment_group("soma_0"))
+        cell.morphology.segment_groups.remove(dend_group)
+        cell.morphology.segment_groups.remove(dend_group_1)
+        print("after removal")
+        for g in cell.morphology.segment_groups:
+            print(g)
+            print([x.segments for x in g.members])
+
+        # create the unbranched segments
+        cell.create_unbranched_segment_group_branches(soma_0.id)
+        print("after re-creation")
+        for g in cell.morphology.segment_groups:
+            print(g)
+            print([x.segments for x in g.members])