Source code for dendrotweaks.morphology.seg_trees

# SPDX-FileCopyrightText: 2025 Poirazi Lab <dendrotweaks@dendrites.gr>
# SPDX-License-Identifier: MPL-2.0

from dendrotweaks.morphology.trees import Node, Tree
import numpy as np



class Segment(Node):
    """
    A class representing a segment.

    A segment is a part of a section used to discretize the section
    in space for the purpose of numerical simulations.

    Parameters
    ----------
    idx : int
        The index of the segment.
    parent_idx : int
        The index of the parent segment.
    sim_seg : h.Segment
        The segment object from a simulator (e.g. NEURON).
    section : Section
        The section to which the segment belongs.

    Attributes
    ----------
    _section : Section
        The section to which the segment belongs.
    _ref : h.Segment
        The segment object from a simulator (e.g. NEURON).
    """

    def __init__(self, idx, parent_idx, sim_seg, section) -> None:
        super().__init__(idx, parent_idx)
        self._section = section
        self._ref = sim_seg


    # PROPERTIES

    @property
    def domain_name(self):
        """
        The morphological or functional domain of the segment.
        """
        return self._section.domain_name

    
    @property
    def domain_color(self):
        """
        The color associated with the domain of the segment.
        """
        return self._section.domain_color


    @property
    def subtree_size(self):
        """
        The number of sections in the subtree rooted at the segment.
        """
        return self._section.subtree_size


    @property
    def Ra(self):
        """
        The axial resistance of the segment (from NEURON).
        """
        return self._section.Ra

    


    def path_distance(self, within_domain=False):
        if within_domain:
            return self._section.path_distance(
                other=self._section.domain_root.parent,
                relative_position=self.x
            )
        return self._section.path_distance(
            relative_position=self.x
        )



    @property
    def distance(self):
        return self.path_distance(within_domain=False)


    @property
    def domain_distance(self):
        return self.path_distance(within_domain=True)

    
    # PARAMETER SETTERS AND GETTERS



    def set_param_value(self, param_name, value):
        """
        Set the value of a parameter of the segment.

        Parameters
        ----------
        param_name : str
            The name of the parameter to set.
        value : float
            The value to set the parameter to.
        """
        try: # Using EAFP
            setattr(self._ref, param_name, value)
        except AttributeError:
            pass





    def get_param_value(self, param_name):
        """
        Get the value of a parameter of the segment.

        Parameters
        ----------
        param_name : str
            The name of the parameter to get.

        Returns
        -------
        float
            The value of the parameter.
        """
        if hasattr(self, param_name):
            return getattr(self, param_name)
        elif hasattr(self._ref, param_name):
            return getattr(self._ref, param_name)
        else:
            return np.nan





# -------------------------------------------------------------------
# NEURON SEGMENT
# -------------------------------------------------------------------

class NeuronSegment(Segment):
    """
    A class representing a segment for the Jaxley simulator.
    """

    def __init__(self, idx, parent_idx, sim_seg, section) -> None:
        super().__init__(idx, parent_idx, sim_seg, section)

    @property
    def x(self):
        """
        The position of the segment along the normalized section length (from NEURON).
        """
        return self._ref.x


    @property
    def area(self):
        """
        The area of the segment (from NEURON).
        """
        return self._ref.area()


    @property
    def diam(self):
        """
        The diameter of the segment (from NEURON).
        """
        return self._ref.diam



class SegmentTree(Tree):
    """
    A class representing a tree graph of segments.
    """

    def __init__(self, segments: list[Segment]) -> None:
        super().__init__(segments)

    def __repr__(self):
        return f"SegmentTree(root={self.root!r}, num_nodes={len(self._nodes)})"


    @property
    def segments(self):
        """
        The segments in the segment tree. Alias for self._nodes
        """
        return self._nodes