export_import.py (2896B)
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(G:nx.DiGraph, posG3D, node_color, node_alpha, chosen_nodes=[], path='./', step=None): 9 10 lx = [x for x,y,z in posG3D.values()] 11 ly = [y for x,y,z in posG3D.values()] 12 lz = [z for x,y,z in posG3D.values()] 13 14 min_x = min(lx) 15 max_x = max(lx) 16 min_y = min(ly) 17 max_y = max(ly) 18 min_z = min(lz) 19 max_z = max(lz) 20 21 if step != None: 22 f = open(path + f'nodes_{step}.csv','w') 23 else: 24 f = open(path + f'nodes.csv','w') 25 26 d_node_rowID = {} 27 cc = 0 28 # for node, xyz in sorted(posG3D.items()): 29 for node in G.nodes(): 30 d_node_rowID[node] = cc 31 x = posG3D[node][0] 32 y = posG3D[node][1] 33 z = posG3D[node][2] 34 35 xn = (x-min_x)/(max_x-min_x) 36 yn = (y-min_y)/(max_y-min_y) 37 zn = (z-min_z)/(max_z-min_z) 38 39 40 if str(node) in chosen_nodes: 41 r = int(255) 42 g = int(0) 43 b = int(0) 44 alpha = 100 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 = node_alpha[node] 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() 75 76 def read_data(nodes_file, edges_file): 77 f = open(nodes_file,'r') 78 lines = f.readlines() 79 80 posG3D = {} 81 d_nodecolors = {} 82 d_hoverinfo = {} 83 d_nodesize = {} 84 85 cc = 0 86 for line in lines: 87 x = float(line.strip().split(',')[0]) 88 y = float(line.strip().split(',')[1]) 89 z = float(line.strip().split(',')[2]) 90 posG3D[str(cc)] = (x,y,z) 91 92 r = float(line.strip().split(',')[3]) 93 g = float(line.strip().split(',')[4]) 94 b = float(line.strip().split(',')[5]) 95 alpha = float(line.strip().split(',')[6]) 96 color = f'rgba({r}, {g}, {b}, {alpha})' 97 d_nodecolors[str(cc)] = color 98 99 hinfo = line.strip().split(',')[7] 100 d_hoverinfo[str(cc)] = hinfo 101 102 d_nodesize[str(cc)] = 8 103 104 cc += 1 105 106 f.close() 107 108 f = open(edges_file,'r') 109 lines = f.readlines() 110 111 G = nx.Graph() 112 113 for line in lines: 114 u = line.strip().split(',')[0] 115 v = line.strip().split(',')[1] 116 G.add_edge(u,v) 117 f.close() 118 return G, d_hoverinfo, d_nodecolors, d_nodesize, posG3D