Source code for allensdk.internal.morphology.node

# Copyright 2016 Allen Institute for Brain Science
# This file is part of Allen SDK.
#
# Allen SDK is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# Allen SDK is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Allen SDK.  If not, see <http://www.gnu.org/licenses/>.

import json
import math

def euclidean_distance(node1, node2):
    dx = node1.x - node2.x
    dy = node1.y - node2.y
    dz = node1.z - node2.z
    return math.sqrt(dx*dx + dy*dy + dz*dz)


def midpoint(node1, node2):
    px = (node1.x + node2.x) * 0.5
    py = (node1.y + node2.y) * 0.5
    pz = (node1.z + node2.z) * 0.5
    return [px, py, pz]


class Node(object): 
    """
    Represents node in SWC morphology file
    """

    def __init__(self, n, t, x, y, z, r, pn, **kwargs):
        """
        Parameters
        ----------
        n: integer
        node ID

        t: integer
        node type (SOMA, AXON, BASAL_DENDRITE or APICAL_DENDRITE)

        x: float
        x position of node

        y: float
        y position of node

        z: float
        z position of node

        r: float
        radius of node

        pn: integer
        ID of parent node (-1 if no parent)
        """
        # these values correspond to columns in an SWC file
        self.n = n
        self.t = t
        self.x = x
        self.y = y
        self.z = z
        self.radius = r
        self.parent = pn
        # 
        self.children = []  # IDs of child nodes
        self.tree_id = -1      # which unconnected graph this node belongs to
        # number of compartment that has this node as its endpoint
        # all nodes except root nodes have a compartment
        self.compartment_id = -1    

    def to_dict(self):
        """ Convert the node into a serializable dictionary """
        return {
            "id": self.n,
            "type": self.t,
            "x": self.x,
            "y": self.y,
            "z": self.z,
            "radius": self.radius,
            "parent": self.parent,
            "children": self.children,
            "tree_id": self.tree_id,
            "compartment_id": self.compartment_id
        }

    @classmethod
    def from_dict(cls, d):
        return cls(
            n = d["id"],
            t = d["type"],
            x = d["x"],
            y = d["y"],
            z = d["z"],
            r = d["radius"],
            pn = d["parent"],
        )

    def __getitem__(self, item):
        return self.to_dict()[item]

    def __str__(self):
        return self.short_string()
        return json.dumps(self.to_dict())

    def short_string(self):
        """ create string with node information in succinct, 
        single-line form """
        return "%d %d %.4f %.4f %.4f %.4f %d %s %d" % (self.n, self.t, self.x, self.y, self.z, self.radius, self.parent,
        str(self.children), self.tree_id);

# Morphology nodes have the following fields. These allow dictionary access
#   to node fields (this is for backward compatibility)
NODE_ID      = 'id'
NODE_TYPE    = 'type'
NODE_X       = 'x'
NODE_Y       = 'y'
NODE_Z       = 'z'
NODE_R       = 'radius'
NODE_PN      = 'parent'
NODE_TREE_ID = 'tree_id'     
NODE_CHILDREN = 'children'