<jupyter_start><jupyter_text>Generate Graph with networkx<jupyter_code>import multiprocessing as mp
from itertools import combinations

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from scipy.spatial.distance import cdist
from torch_geometric.utils import from_networkx, to_networkx
from torch_geometric.transforms.distance import Distance
from networkx.algorithms.tree import mst
%load_ext line_profiler
pos = np.random.uniform(low=0.0, high=1.0, size=(128,10,2))
dist = cdist(pos, pos, metric='euclidean')
graph = nx.complete_graph(50)
for n in graph.nodes:
    graph.nodes[0]["pos"] = pos[n]
for u, v in graph.edges:
    graph[u][v]["weight"] = dist[u,v]
mst = nx.algorithms.tree.mst.minimum_spanning_tree(graph, algorithm="prim")
def gen_mst_graph(pos):
#     pos = np.random.uniform(low=0.0, high=1.0, size=(graph_size,2))
    graph_size, _ = pos.shape
    dist = cdist(pos, pos, metric='euclidean')
    graph = nx.complete_graph(graph_size)
    for n in graph.nodes:
        graph.nodes[n]["pos"] = pos[n]
    for u, v in graph.edges:
        graph[u][v]["weight"] = dist[u,v]
        
    mst = nx.algorithms.tree.mst.minimum_spanning_tree(graph, algorithm="prim")
    
    return mst

def gen_complete_graph(pos):
#     pos = np.random.uniform(low=0.0, high=1.0, size=(graph_size,2))
    graph_size, _ = pos.shape
    dist = cdist(pos, pos, metric='euclidean')
    graph = nx.complete_graph(graph_size)
    for n in graph.nodes:
        graph.nodes[n]["pos"] = pos[n]
    for u, v in graph.edges:
        graph[u][v]["weight"] = dist[u,v]
        
    return graph
pos = np.random.uniform(low=0.0, high=1.0, size=(128,50,2))
dist = np.stack([cdist(p, p, metric='euclidean') for p in pos])
g = gen_complete_graph(pos[0])
d=dist[0]
d
np.argsort(d)[:, 1:4]
def gen_knn(pos, k):
    dist = cdist(pos, pos, metric='euclidean')
    g_size = dist.shape[0]
    k_nn = dist.argsort()[:, 1:k+1]
    node_list = []
    edge_list = []
    for u in range(g_size):
        node_list.append([u, {"pos": pos[u]}])
        for v in k_nn[u]:
            edge_list.append([u,v, {"is_mst": 0.0, "edge_attr": dist[u,v]}])
    g = nx.Graph()
    g.add_nodes_from(node_list)
    g.add_edges_from(edge_list)
    
    return g

def mst_to_knn(graph, k):
    num_edges = graph.num_edges
    graph.is_mst = torch.ones((num_edges, 1), dtype=torch.float)
    pos = graph.pos.numpy()
    dist = cdist(pos, pos, metric='euclidean')
    knn_g = gen_knn(dist, k)
    mst_g = to_networkx(graph, edge_attrs=["is_mst", "edge_attr"], to_undirected=True)
    knn_g.update(mst_g)
    pyg = from_networkx(knn_g)
    
    pyg.pos = graph.pos
    pyg.x = graph.x
    pyg.is_mst = pyg.is_mst.unsqueeze(-1)
    distance.Distance(norm=False, cat=False)(pyg)
    
    return pyg

pos = np.random.uniform(low=0.0, high=1.0, size=(50,2))
nx.draw(to_networkx(gen_knn_graph(pos.tolist(), 3)), pos=pos)
pyg = from_networkx(gen_complete_graph(pos))
from torch_geometric.nn import knn_graph
pygknn = knn_graph(pyg.pos, k=3)
from problems.tsp.problem_tsp import gen_knn_mst_graph
val_dataset = "data/tsp/tsp50_validation_seed4321.pkl"
vd = _GNN_TSPDataset(filename=val_dataset, distribution="knn_mst")
import pickle
val_pos = None
with open(val_dataset, "rb") as f:
    val_pos = pickle.load(f)
gen_knn_mst_graph(val_pos[0])
g = gen_knn_mst_graph(10)
from nets import gnn_encoder, gnn_model
from importlib import reload
from utils import torch_load_cpu, load_problem
reload(gnn_encoder)
problem = load_problem("tsp")
model = gnn_model.GNNFFModel(embedding_dim=128, hidden_dim=128, problem=problem, n_encode_layers=3)
gnn = gnn_encoder.TransformerConv(in_channels=128, out_channels=128//4, heads=4)
data = model._init_embed(g)
x, alpha = gnn(x = data.x, edge_index=data.edge_index, prev_alpha=data.is_mst.expand_as(alpha))
alpha
g.is_mst.expand(-1,4)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14245014245014245, 'fraction_numerical': 0.02046102046102046, 'mean_word_length': 4.393854748603352, 'pattern_counts': {'":': 3, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '387189', 'license': 'MIT', 'n_tokens_mistral': 1584, 'n_tokens_neox': 1538, 'n_words': 294, 'path': 'gen_mst_dataset.ipynb', 'repo_name': 'neo-pan/attention-learn-to-route', 'size': '3861'}