export_vr.py (1834B)
1 #!/usr/bin/env python3.10 2 3 import networkx as nx 4 import numpy as np 5 6 import matplotlib.cm as cm 7 8 def export_data_vr(G:nx.DiGraph, posG3D, chosen_nodes, path='./', step=None): 9 node_color = {node: [1, 1, 1] for node in G.nodes} 10 11 lx = [x for x,y,z in posG3D.values()] 12 ly = [y for x,y,z in posG3D.values()] 13 lz = [z for x,y,z in posG3D.values()] 14 15 min_x = min(lx) 16 max_x = max(lx) 17 min_y = min(ly) 18 max_y = max(ly) 19 min_z = min(lz) 20 max_z = max(lz) 21 22 if step != None: 23 f = open(path + f'nodes_{step}.csv','w') 24 else: 25 f = open(path + f'nodes.csv','w') 26 27 d_node_rowID = {} 28 cc = 0 29 # for node, xyz in sorted(posG3D.items()): 30 for node in G.nodes(): 31 d_node_rowID[node] = cc 32 x = posG3D[node][0] 33 y = posG3D[node][1] 34 z = posG3D[node][2] 35 36 xn = (x-min_x)/(max_x-min_x) 37 yn = (y-min_y)/(max_y-min_y) 38 zn = (z-min_z)/(max_z-min_z) 39 40 if str(node) in chosen_nodes: 41 r = int(255) 42 g = int(0) 43 b = int(0) 44 alpha = 200 45 else: 46 r = int(255*node_color[node][0]) 47 g = int(255*node_color[node][1]) 48 b = int(255*node_color[node][2]) 49 alpha = 100 50 51 name = str(node) 52 f.write('%s,%s,%s,%s,%s,%s,%s,%s\n' %(xn,yn,zn,r,g,b,alpha,name)) 53 cc += 1 54 55 f.close() 56 57 if step != None: 58 f = open(path + f'edges_{step}.csv','w') 59 else: 60 f = open(path + f'edges.csv','w') 61 62 cc = 0 63 for u,v in G.edges(): 64 row_u = d_node_rowID[u] 65 row_v = d_node_rowID[v] 66 67 edge_color = '#999999' 68 r = int(255*node_color[node][0]) 69 g = int(255*node_color[node][1]) 70 b = int(255*node_color[node][2]) 71 72 f.write('%s,%s,%s,%s,%s,%s\n' %(row_u,row_v,r,g,b,100)) 73 74 f.close()