compartment_tree.py

import treelib3


def get_compartment_hierarchy(compartmentList):
    '''
    constructs a tree structure containing the
    compartment hierarchy excluding 2D compartments
    @param:compartmentList
    '''
    def remove_membranes(tree, nodeID):
        node = tree.get_node(nodeID)
        if node.data == 2:
            parent = tree.get_node(nodeID).bpointer
            if parent is None:
                # accounting for the case where the topmost node is a membrane
                newRoot = tree.get_node(nodeID).fpointer[0]
                tree.root = newRoot
                node = tree.get_node(newRoot)
            else:
                tree.link_past_node(nodeID)
                nodeID = parent
                node = tree.get_node(nodeID)
        for element in node.fpointer:
            removeMembranes(tree, element)

    from copy import deepcopy
    tree = treelib3.Tree()
    c2 = deepcopy(compartmentList)
    while len(c2) > 0:
        removeList = []
        for element in c2:
            if c2[element][2] in ['', None]:
                try:
                    tree.create_node(element, element, data=c2[element][0])
                except treelib3.tree.MultipleRootError:
                    # there's more than one top level element
                    tree2 = treelib3.Tree()
                    tree2.create_node('dummyRoot', 'dummyRoot', data=3)
                    tree2.paste('dummyRoot', tree)
                    tree2.create_node(
                        element, element, parent='dummyRoot',
                        data=c2[element][0])
                    tree = tree2
                removeList.append(element)
            elif tree.contains(c2[element][2]):
                tree.create_node(
                    element, element, parent=c2[element][2],
                    data=c2[element][0])
                removeList.append(element)
        for element in removeList:
            c2.pop(element)
    removeMembranes(tree, tree.root)
    return tree


def get_outside_inside_compartment(compartmentList, compartment):
    outside = compartmentList[compartment][2]
    for comp in compartmentList:
        if compartmentList[comp][2] == compartment:
            return outside, comp