|
|
@@ -0,0 +1,48 @@ |
|
|
|
import time, atexit, json |
|
|
|
import networkx |
|
|
|
import numpy as np |
|
|
|
import websocket |
|
|
|
from socket import error as socket_error |
|
|
|
import graphstate |
|
|
|
import util |
|
|
|
|
|
|
|
class GraphState(graphstate.GraphState, networkx.Graph): |
|
|
|
def __init__(self, *args, **kwargs): |
|
|
|
graphstate.GraphState.__init__(self, *args, **kwargs) |
|
|
|
self.connect_to_server() |
|
|
|
|
|
|
|
def connect_to_server(self, uri = "ws://localhost:5000"): |
|
|
|
""" Attempt to connect to the websocket server """ |
|
|
|
try: |
|
|
|
self.ws = websocket.create_connection(uri) |
|
|
|
atexit.register(self.shutdown) |
|
|
|
except socket_error: |
|
|
|
self.ws = None |
|
|
|
|
|
|
|
def shutdown(self): |
|
|
|
""" Close the connection to the websocket """ |
|
|
|
self.update() |
|
|
|
self.ws.close() |
|
|
|
|
|
|
|
def update(self, delay = 0.5): |
|
|
|
""" Call this function when you are ready to send data to the browser """ |
|
|
|
if not self.ws: |
|
|
|
return |
|
|
|
|
|
|
|
# Automatically perform layout if position is not provided |
|
|
|
if not all(("position" in node) for node in self.node.values()): |
|
|
|
self.layout() |
|
|
|
|
|
|
|
# Send data to browser and rate-limit |
|
|
|
self.ws.send(json.dumps(self.to_json())) |
|
|
|
time.sleep(delay) |
|
|
|
|
|
|
|
def layout(self, dim=3): |
|
|
|
""" Automatically lay out the graph """ |
|
|
|
pos = networkx.spring_layout(self, dim, 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) |
|
|
|
|
|
|
|
|