abp/abp/nxgraphstate.py

import networkx as nx
import graphstate
import clifford

class NXGraphState(graphstate.GraphState, nx.Graph):
    """ This is GraphState with NetworkX-like abilities """
    def __init__(self, *args, **kwargs):
        graphstate.GraphState.__init__(self, *args, **kwargs)

    def layout(self):
        """ Automatically lay out the graph """
        pos = nx.spring_layout(self, dim=3, scale=np.sqrt(self.order()))
        middle = np.average(pos.values(), axis=0)
        pos = {key: value - middle for key, value in pos.items()}
        for key, (x, y, z) in pos.items():
            self.node[key]["position"] = util.xyz(x, y, z)

    def add_vops(self):
        """ Automatically add vops if they're not present """
        for key in self.node:
            if not "vop" in self.node[key]:
                self.node[key]["vop"] = clifford.identity