oblique manipulation tools

morph_tool/oblique.py

@@ -0,0 +1,126 @@
+"""Tools to manipulate oblique branches."""
+import numpy as np
+from morphio import SectionType, PointLevel
+from neurom.morphmath import interval_lengths
+from neurom.core import Morphology
+from neurom import iter_sections
+
+from morph_tool.sub_sectiontypes import (
+    APICAL_SUBTYPE_MAP,
+    set_apical_subtypes,
+    unset_apical_subtypes,
+)
+
+
+def get_trunk_obliques_ids(morphology):
+    """Return the ids of trunk of obliques (first oblique section off the trunk).
+
+    This function assume that set_subtypes was run to know the obliques.
+    """
+    trunk_oblique_ids = []
+    for section in iter_sections(morphology):
+        if (
+            section.type == APICAL_SUBTYPE_MAP["oblique"]
+            and section.parent.type == APICAL_SUBTYPE_MAP["trunk"]
+        ):
+            trunk_oblique_ids.append(section.id)
+    return trunk_oblique_ids
+
+
+def remove_obliques(morphology, trunk_oblique_ids=None):
+    """Remove obliqeus from morphology, from a list of trunk ids, or all is None."""
+    morphology = Morphology(morphology)
+    if trunk_oblique_ids is None:
+        trunk_oblique_ids = get_trunk_obliques_ids(morphology)
+
+    for trunk_id in trunk_oblique_ids:
+        section = morphology.section(trunk_id)
+        morphology.delete_section(section, recursive=True)
+    return morphology
+
+
+def remove_tuft(morphology):
+    """Remove obliqeus from morphology, from a list of trunk ids, or all is None."""
+    morphology = Morphology(morphology)
+    morphology = set_apical_subtypes(morphology)
+    for sec in iter_sections(morphology):
+        if sec.type == APICAL_SUBTYPE_MAP["tuft"]:
+            section = morphology.section(sec.id)
+            morphology.delete_section(section, recursive=False)
+    return unset_apical_subtypes(morphology)
+
+
+def replace_apical(morphology, diameter=4.0, length=800):
+    """Replace apical tree with stub apical consisting of 2 points."""
+    for root_section in morphology.root_sections:
+        if root_section.type == SectionType.apical_dendrite:
+            orig_point = root_section.points[0]
+            morphology.delete_section(root_section, recursive=True)
+
+    points = np.zeros((2, 3))
+    points[:, 1] = np.linspace(0, length, 2)
+    points += orig_point
+    morphology.append_root_section(PointLevel(points, 2 * [diameter]), APICAL_SUBTYPE_MAP["trunk"])
+
+
+def add_oblique(morphology, distance=100, length=200.0, diameter=2.0, direction=None):
+    """Add a stub oblique to a trunk.
+
+    Args:
+        morphology (morphio.mut.Morphology): morph to add oblique with custom5 trunk
+        distance (float): distance of oblique along trunk from trunk first point
+        length (float): lenght of oblique
+        diameter (float): diameter of oblique
+        direction (-1/1/None): left (-1), right (1) direction along x, None will be a random choice
+    """
+
+    def _copy(section, section_base):
+        """to recursively copy downstream from section_base to section"""
+        for base_child in section_base.children:
+            section.append_section(base_child)
+        for child, base_child in zip(section.children, section_base.children):
+            _copy(child, base_child)
+
+    oblique_points = np.zeros((2, 3))
+    if direction is None:
+        direction = (-1.0) ** np.random.randint(2)
+    oblique_points[:, 0] = direction * np.linspace(0, length, 2)
+    oblique_diameters = 2 * [diameter]
+    current_dist = 0
+    for section in morphology.iter():
+        if section.type == APICAL_SUBTYPE_MAP["trunk"]:
+            dists = current_dist + np.cumsum(interval_lengths(section.points, prepend_zero=True))
+            current_dist = dists[-1]
+            if dists[0] < distance <= dists[-1]:
+                if distance == dists[-1]:
+                    distance -= 0.1
+                    print("we avoid trifurcation by moving oblique back by 0.1 micron")
+
+                bif_pt = 0.5 * (
+                    section.points[dists < distance][-1] + section.points[dists > distance][0]
+                )
+                bif_diam = 0.5 * (
+                    section.diameters[dists < distance][-1] + section.diameters[dists > distance][0]
+                )
+
+                prev_points = np.append(section.points[dists < distance], [bif_pt], axis=0)
+                prev_diams = np.append(section.diameters[dists < distance], bif_diam)
+                prev_pts = PointLevel(prev_points, prev_diams)
+
+                next_points = np.insert(section.points[dists > distance], 0, bif_pt, axis=0)
+                next_diams = np.insert(section.diameters[dists > distance], 0, bif_diam)
+                next_pts = PointLevel(next_points, next_diams)
+
+                oblique_pts = PointLevel(oblique_points + next_points[0], oblique_diameters)
+
+                if section.is_root:
+                    prev_sec = morphology.append_root_section(prev_pts, APICAL_SUBTYPE_MAP["trunk"])
+                else:
+                    prev_sec = section.parent.append_section(prev_pts, APICAL_SUBTYPE_MAP["trunk"])
+                prev_sec.append_section(next_pts, APICAL_SUBTYPE_MAP["trunk"])
+                prev_sec.append_section(oblique_pts, APICAL_SUBTYPE_MAP["oblique"])
+                next_sec = prev_sec.children[0]
+
+                _copy(next_sec, section)
+                break
+    morphology.delete_section(section)

morph_tool/sub_sectiontypes.py

@@ -0,0 +1,137 @@
+"""Module to set sub section types, such as oblique for apicals, or collaterals for axons."""
+from morphio import SectionType
+from neurom import iter_sections
+from morphio.mut import Morphology
+
+from morph_tool.apical_point import apical_point_section_segment
+from morph_tool.axon_point import axon_point_section
+
+
+APICAL_SUBTYPE_MAP = {
+    "axon": SectionType.axon,
+    "basal": SectionType.basal_dendrite,
+    "apical": SectionType.apical_dendrite,
+    "trunk": SectionType.custom5,
+    "oblique": SectionType.custom6,
+    "tuft": SectionType.custom7,
+}
+REV_APICAL_SUBTYPE_MAP = {j: i for i, j in APICAL_SUBTYPE_MAP.items()}
+
+AXON_SUBTYPE_MAP = {
+    "axon": SectionType.axon,
+    "basal": SectionType.basal_dendrite,
+    "apical": SectionType.apical_dendrite,
+    "main": SectionType.custom5,
+    "collateral": SectionType.custom6,
+}
+REV_AXON_SUBTYPE_MAP = {j: i for i, j in AXON_SUBTYPE_MAP.items()}
+
+
+def apical_subtype(neuron, apical_section=None, tuft_percent=20):
+    """Return a dict of extended apical subtypes (trunk, oblique, tuft)."""
+    extended_types = dict()
+
+    apical_section = neuron.sections[
+        apical_point_section_segment(neuron, tuft_percent=tuft_percent)[0]
+    ]
+    if apical_section is not None:
+        for section in iter_sections(neuron):
+            if section.type == SectionType.apical_dendrite:
+                extended_types[section.id] = "oblique"
+
+        for section in apical_section.ipreorder():
+            extended_types[section.id] = "tuft"
+
+        for section in apical_section.iupstream():
+            extended_types[section.id] = "trunk"
+    else:
+        for section in iter_sections(neuron):
+            if section.type == SectionType.apical_dendrite:
+                extended_types[section.id] = "apical"
+
+    for section in iter_sections(neuron):
+        if section.type == SectionType.basal_dendrite:
+            extended_types[section.id] = "basal"
+        if section.type == SectionType.axon:
+            extended_types[section.id] = "axon"
+
+    return extended_types
+
+
+def axon_subtype(neuron, axonal_section=None, direction=None, bbox=None, ignore_axis=2):
+    """Return a dict of extended axonal subtypes (main, collaterals)."""
+    extended_types = dict()
+    if axonal_section is None:
+        axonal_section = neuron.sections[
+            axon_point_section(neuron, direction=direction, ignore_axis=ignore_axis)
+        ]
+        for section in iter_sections(neuron):
+            if section.type == SectionType.axon:
+                extended_types[section.id] = "collateral"
+
+        for section in axonal_section.iupstream():
+            extended_types[section.id] = "main"
+    else:
+
+        for section in iter_sections(neuron):
+            if section.type == SectionType.axon:
+                extended_types[section.id] = "axon"
+            if section.type == SectionType.basal_dendrite:
+                extended_types[section.id] = "basal"
+            if section.type == SectionType.apical_dendrite:
+                extended_types[section.id] = "apical"
+
+    return extended_types
+
+
+def set_apical_subtypes(neuron, tuft_percent=20):
+    """Set apical subtypes to a morphology.
+
+    WARNING: cannot save it to .asc file, use unset_subtype to revert
+    """
+    subtypes = apical_subtype(neuron, tuft_percent=tuft_percent)
+    sections = neuron.sections
+    for secid, subtype in subtypes.items():
+        sections[secid].morphio_section.type = APICAL_SUBTYPE_MAP[subtype]
+    return neuron
+
+
+def set_axon_subtypes(neuron, **kwargs):
+    """Set subtypes to a morphology.
+
+    WARNING: cannot save it to .asc file, use unset_subtype to revert
+    """
+    subtypes = axon_subtype(neuron, **kwargs)
+    sections = neuron.sections
+    for secid, subtype in subtypes.items():
+        sections[secid].morphio_section.type = AXON_SUBTYPE_MAP[subtype]
+    return neuron
+
+
+def _unset_subtypes(neuron, base_type):
+    """Unset subtypes to a morphology.
+
+    Args:
+        neuron (neurom/morphio): morphology to consider
+        base_type (SectionType): section type to reset all custom types.
+
+    """
+    if isinstance(neuron, Morphology):
+        for section in neuron.iter():
+            if section.type.name.startswith("custom"):
+                section.type = base_type
+    else:
+        for section in iter_sections(neuron):
+            if section.type.name.startswith("custom"):
+                section.morphio_section.type = base_type
+    return neuron
+
+
+def unset_apical_subtypes(neuron):
+    """Unset apical subtypes to a morphology."""
+    return _unset_subtypes(neuron, SectionType.apical_dendrite)
+
+
+def unset_axon_subtypes(neuron):
+    """Unset axon subtypes to a morphology."""
+    return _unset_subtypes(neuron, SectionType.axon)

tests/test_sub_sectiontypes.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+from morphio import Morphology, SectionType
+from neurom import load_morphology, iter_sections
+from morph_tool import sub_sectiontypes as tested
+
+
+DATA = Path(__file__).absolute().parent / "data"
+
+
+def test_apical_subtypes():
+    morph = load_morphology(DATA / "neuron.asc")
+    morph = tested.set_apical_subtypes(morph)
+    types = [section.type for section in iter_sections(morph)]
+    assert types[40] == SectionType.custom5
+    assert types[47] == SectionType.custom6
+    assert types[57] == SectionType.custom7
+
+    morph = tested.unset_apical_subtypes(morph)
+    types = [section.type for section in iter_sections(morph)]
+    assert types[40] == SectionType.apical_dendrite
+    assert types[47] == SectionType.apical_dendrite
+    assert types[57] == SectionType.apical_dendrite
+
+
+def test_axonal_subtypes():
+    morph = load_morphology(DATA / "neuron.asc")
+    morph = tested.set_axon_subtypes(morph)
+
+    types = [section.type for section in iter_sections(morph)]
+    assert types[79] == SectionType.custom5
+    assert types[84] == SectionType.custom6
+
+    morph = tested.unset_axon_subtypes(morph)
+    types = [section.type for section in iter_sections(morph)]
+    assert types[79] == SectionType.axon
+    assert types[84] == SectionType.axon