def explain_pagerank(model, node_idx, x, edge_index, target, include_edges=None):
data = Data(x=x, edge_index=edge_index)
g = to_networkx(data)
pagerank = nx.pagerank(g, personalization={node_idx: 1})
node_attr = np.zeros(x.shape[0])
for node, value in pagerank.items():
node_attr[node] = value
edge_mask = node_attr_to_edge(edge_index, node_attr)
return edge_mask
def to_molecule(data):
ATOM_MAP = [
'C', 'O', 'Cl', 'H', 'N', 'F', 'Br', 'S', 'P', 'I', 'Na', 'K', 'Li',
'Ca'
]
g = to_networkx(data, node_attrs=['x'])
for u, data in g.nodes(data=True):
data['name'] = ATOM_MAP[data['x'].index(1.0)]
del data['x']
return g
def test_to_networkx():
x = torch.Tensor([[1, 2], [3, 4]])
pos = torch.Tensor([[0, 0], [1, 1]])
row = torch.tensor([0, 1, 0])
col = torch.tensor([1, 0, 0])
edge_attr = torch.Tensor([1, 2, 3])
G = to_networkx(torch.stack([row, col], dim=0), x, edge_attr, pos)
assert G.nodes[0]['x'].tolist() == [1, 2]
assert G.nodes[1]['x'].tolist() == [3, 4]
assert G.nodes[0]['pos'].tolist() == [0, 0]
assert G.nodes[1]['pos'].tolist() == [1, 1]
assert networkx.to_numpy_matrix(G).tolist() == [[3, 1], [2, 0]]
edge_attr = torch.Tensor([[1, 1], [2, 2], [3, 3]])
G = to_networkx(torch.stack([row, col], dim=0), edge_attr=edge_attr)
assert G[0][0]['weight'].tolist() == [3, 3]
assert G[0][1]['weight'].tolist() == [1, 1]
assert G[1][0]['weight'].tolist() == [2, 2]
def visualization(self,
data: Data,
edge_mask: Tensor,
top_k: int,
plot_utils: PlotUtils,
words: Optional[list] = None,
node_idx: int = None,
vis_name: Optional[str] = None):
if vis_name is None:
vis_name = f"filename.png"
data = data.to('cpu')
edge_mask = edge_mask.to('cpu')
if self.explain_graph:
graph = to_networkx(data)
if words is None:
plot_utils.plot_soft_edge_mask(graph,
edge_mask,
top_k=top_k,
un_directed=True,
words=words,
figname=vis_name)
else:
plot_utils.plot_soft_edge_mask(graph,
edge_mask,
top_k=top_k,
un_directed=True,
x=x,
figname=vis_name)
else:
assert node_idx is not None, "visualization method doesn't get the target node index"
x, edge_index, y, subset, kwargs = \
self.get_subgraph(node_idx=node_idx, x=data.x, edge_index=data.edge_index, y=data.y)
new_node_idx = torch.where(subset == node_idx)[0]
new_data = Data(x=x, edge_index=edge_index)
graph = to_networkx(new_data)
plot_utils.plot_soft_edge_mask(graph,
edge_mask,
top_k=top_k,
un_directed=True,
y=y,
node_idx=new_node_idx,
figname=vis_name)
def test_networkx_vice_versa_convert():
import networkx as nx
G = nx.complete_graph(5)
assert G.is_directed() is False
data = from_networkx(G)
assert data.is_directed() is False
G = to_networkx(data)
assert G.is_directed() is True
G = nx.to_undirected(G)
assert G.is_directed() is False
def draw_pyg_graph(G, name='Lobster', path='./visualization/train_pyggraph/'):
fig = plt.figure(figsize=(12, 12))
ax = plt.subplot(111)
ax.set_title(name, fontsize=10)
nx_graph = to_networkx(G)
if not os.path.exists(path):
os.makedirs(path)
save_name = path + name + '.png'
nx.draw(nx_graph)
plt.savefig(save_name, format="PNG")
plt.close()
def explain_distance(model, node_idx, x, edge_index, target, include_edges=None):
data = Data(x=x, edge_index=edge_index)
g = to_networkx(data)
length = nx.shortest_path_length(g, target=node_idx)
def get_attr(node):
if node in length:
return 1 / (length[node] + 1)
return 0
edge_sources = edge_index[1].cpu().numpy()
return np.array([get_attr(node) for node in edge_sources])
def NC_mc_l_shapley(coalition: list,
data: Data,
local_raduis: int,
value_func: str,
node_idx: int = -1,
subgraph_building_method='zero_filling',
sample_num=1000) -> float:
""" monte carlo approximation of l_shapley where the target node is kept in both subgraph """
graph = to_networkx(data)
num_nodes = graph.number_of_nodes()
subgraph_build_func = get_graph_build_func(subgraph_building_method)
local_region = copy.copy(coalition)
for k in range(local_raduis - 1):
k_neiborhoood = []
for node in local_region:
k_neiborhoood += list(graph.neighbors(node))
local_region += k_neiborhoood
local_region = list(set(local_region))
coalition_placeholder = num_nodes
set_exclude_masks = []
set_include_masks = []
for example_idx in range(sample_num):
subset_nodes_from = [
node for node in local_region if node not in coalition
]
random_nodes_permutation = np.array(subset_nodes_from +
[coalition_placeholder])
random_nodes_permutation = np.random.permutation(
random_nodes_permutation)
split_idx = np.where(
random_nodes_permutation == coalition_placeholder)[0][0]
selected_nodes = random_nodes_permutation[:split_idx]
set_exclude_mask = np.ones(num_nodes)
set_exclude_mask[local_region] = 0.0
set_exclude_mask[selected_nodes] = 1.0
if node_idx != -1:
set_exclude_mask[node_idx] = 1.0
set_include_mask = set_exclude_mask.copy()
set_include_mask[coalition] = 1.0 # include the node_idx
set_exclude_masks.append(set_exclude_mask)
set_include_masks.append(set_include_mask)
exclude_mask = np.stack(set_exclude_masks, axis=0)
include_mask = np.stack(set_include_masks, axis=0)
marginal_contributions = \
marginal_contribution(data, exclude_mask, include_mask, value_func, subgraph_build_func)
mc_l_shapley_value = (marginal_contributions).mean().item()
return mc_l_shapley_value
def draw_graphs(glist, aids=False):
for i, g in enumerate(glist):
plt.clf()
G = to_networkx(g).to_undirected()
if aids:
label_list = aids_labels(g)
labels = {}
for j, node in enumerate(G.nodes()):
labels[node] = label_list[j]
nx.draw(G, labels=labels)
else:
nx.draw(G)
plt.savefig("graph{}.png".format(i))
def test_to_networkx():
x = torch.Tensor([[1, 2], [3, 4]])
pos = torch.Tensor([[0, 0], [1, 1]])
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
edge_attr = torch.Tensor([1, 2, 3])
data = Data(x=x, pos=pos, edge_index=edge_index, weight=edge_attr)
G = to_networkx(data, node_attrs=['x', 'pos'], edge_attrs=['weight'])
assert G.nodes[0]['x'] == [1, 2]
assert G.nodes[1]['x'] == [3, 4]
assert G.nodes[0]['pos'] == [0, 0]
assert G.nodes[1]['pos'] == [1, 1]
assert nx.to_numpy_matrix(G).tolist() == [[3, 1], [2, 0]]
def l_shapley(coalition: list,
data: Data,
local_raduis: int,
value_func: str,
subgraph_building_method='zero_filling'):
""" shapley value where players are local neighbor nodes """
graph = to_networkx(data)
num_nodes = graph.number_of_nodes()
subgraph_build_func = get_graph_build_func(subgraph_building_method)
local_region = copy.copy(coalition)
for k in range(local_raduis - 1):
k_neiborhoood = []
for node in local_region:
k_neiborhoood += list(graph.neighbors(node))
local_region += k_neiborhoood
local_region = list(set(local_region))
set_exclude_masks = []
set_include_masks = []
nodes_around = [node for node in local_region if node not in coalition]
num_nodes_around = len(nodes_around)
for subset_len in range(0, num_nodes_around + 1):
node_exclude_subsets = combinations(nodes_around, subset_len)
for node_exclude_subset in node_exclude_subsets:
set_exclude_mask = np.ones(num_nodes)
set_exclude_mask[local_region] = 0.0
if node_exclude_subset:
set_exclude_mask[list(node_exclude_subset)] = 1.0
set_include_mask = set_exclude_mask.copy()
set_include_mask[coalition] = 1.0
set_exclude_masks.append(set_exclude_mask)
set_include_masks.append(set_include_mask)
exclude_mask = np.stack(set_exclude_masks, axis=0)
include_mask = np.stack(set_include_masks, axis=0)
num_players = len(nodes_around) + 1
num_player_in_set = num_players - 1 + len(coalition) - (
1 - exclude_mask).sum(axis=1)
p = num_players
S = num_player_in_set
coeffs = torch.tensor(1.0 / comb(p, S) / (p - S + 1e-6))
marginal_contributions = \
marginal_contribution(data, exclude_mask, include_mask, value_func, subgraph_build_func)
l_shapley_value = (marginal_contributions.squeeze().cpu() *
coeffs).sum().item()
return l_shapley_value
def rdmol_to_data(mol: Mol):
assert mol.GetNumConformers() == 1
N = mol.GetNumAtoms()
pos = torch.tensor(mol.GetConformer(0).GetPositions(), dtype=torch.float)
atomic_number = []
aromatic = []
sp = []
sp2 = []
sp3 = []
num_hs = []
for atom in mol.GetAtoms():
atomic_number.append(atom.GetAtomicNum())
aromatic.append(1 if atom.GetIsAromatic() else 0)
hybridization = atom.GetHybridization()
sp.append(1 if hybridization == HybridizationType.SP else 0)
sp2.append(1 if hybridization == HybridizationType.SP2 else 0)
sp3.append(1 if hybridization == HybridizationType.SP3 else 0)
z = torch.tensor(atomic_number, dtype=torch.long)
row, col, edge_type = [], [], []
for bond in mol.GetBonds():
start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
row += [start, end]
col += [end, start]
edge_type += 2 * [BOND_TYPES[bond.GetBondType()]]
edge_index = torch.tensor([row, col], dtype=torch.long)
edge_type = torch.tensor(edge_type)
perm = (edge_index[0] * N + edge_index[1]).argsort()
edge_index = edge_index[:, perm]
edge_type = edge_type[perm]
row, col = edge_index
hs = (z == 1).to(torch.float)
num_hs = scatter(hs[row], col, dim_size=N).tolist()
smiles = Chem.MolToSmiles(mol)
data = Data(node_type=z,
pos=pos,
edge_index=edge_index,
edge_type=edge_type,
rdmol=copy.deepcopy(mol),
smiles=smiles)
data.nx = to_networkx(data, to_undirected=True)
return data
def test_from_networkx():
x = torch.randn(2, 8)
pos = torch.randn(2, 3)
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
edge_attr = torch.randn(edge_index.size(1))
perm = torch.tensor([0, 2, 1])
data = Data(x=x, pos=pos, edge_index=edge_index, edge_attr=edge_attr)
G = to_networkx(data, node_attrs=['x', 'pos'], edge_attrs=['edge_attr'])
data = from_networkx(G)
assert len(data) == 4
assert data.x.tolist() == x.tolist()
assert data.pos.tolist() == pos.tolist()
assert data.edge_index.tolist() == edge_index[:, perm].tolist()
assert data.edge_attr.tolist() == edge_attr[perm].tolist()
def gen_mst_len_dataset(size: int, min_num_node: int, max_num_node: int)-> TSPDataset:
dataset = TSPDataset(size=size, min_num_node=min_num_node, max_num_node=max_num_node, transform=Distance(cat=False))
datalist = []
for graph in dataset:
ng = to_networkx(graph, edge_attrs=["edge_attr"], to_undirected=True)
ng_mst = mst.minimum_spanning_tree(ng, "edge_attr")
mst_len = sum([edge_value["edge_attr"] for edge_value in ng_mst.edges.values()])
graph.mst_len = torch.tensor([mst_len], dtype=torch.float)
datalist.append(graph)
dataset.__data_list__ = datalist
dataset.data, dataset.slices = dataset.collate(datalist)
return dataset
def to_networkx_featured(data):
g = to_networkx(data, to_undirected=True, remove_self_loops=True)
features = {i: fea.argmax() + 1 for i, fea in enumerate(data.x)}
nx.set_node_attributes(g, values=features, name='feature')
largest_cc = max(nx.connected_components(g), key=len)
largest_cc = g.subgraph(largest_cc)
# change node index
mapping = {old: new for new, old in enumerate(largest_cc.nodes)}
largest_cc = nx.relabel_nodes(largest_cc, mapping)
# option 1: remove the checking condition, and remove isolated nodes
# option 2: for every graph, add a new node, with node feature as 0
# and connect all nodes to this single virtual node
return largest_cc # get the largest connected component
def get_subgraph(self,
node_idx: int,
x: Tensor,
edge_index: Tensor,
y: Optional[Tensor] = None,
**kwargs)\
-> Tuple[Tensor, Tensor, Tensor, List, Dict]:
r""" extract the subgraph of target node
Args:
node_idx (:obj:`int`): The node index
x (:obj:`torch.Tensor`): Node feature matrix with shape
:obj:`[num_nodes, dim_node_feature]`
edge_index (:obj:`torch.Tensor`): Graph connectivity in COO format
with shape :obj:`[2, num_edges]`
y (:obj:`torch.Tensor`, :obj:`None`): Node label matrix with shape :obj:`[num_nodes]`
(default :obj:`None`)
kwargs(:obj:`Dict`, :obj:`None`): Additional parameters
:rtype: (:class:`torch.Tensor`, :class:`torch.Tensor`, :class:`torch.Tensor`,
:obj:`List`, :class:`Dict`)
"""
num_nodes, num_edges = x.size(0), edge_index.size(1)
graph = to_networkx(data=Data(x=x, edge_index=edge_index),
to_undirected=True)
subset, edge_index, _, edge_mask = k_hop_subgraph_with_default_whole_graph(
edge_index,
node_idx,
self.num_hops,
relabel_nodes=True,
num_nodes=num_nodes,
flow=self.__flow__())
mapping = {int(v): k for k, v in enumerate(subset)}
subgraph = graph.subgraph(subset.tolist())
nx.relabel_nodes(subgraph, mapping)
x = x[subset]
for key, item in kwargs.items():
if torch.is_tensor(item) and item.size(0) == num_nodes:
item = item[subset]
elif torch.is_tensor(item) and item.size(0) == num_edges:
item = item[edge_mask]
kwargs[key] = item
if y is not None:
y = y[subset]
return x, edge_index, y, subset, kwargs
def _edge_removal(graph):
actions = set()
_, num_edges = graph.edge_index.size()
edges = graph.edge_index.t().numpy().tolist()
for i in range(num_edges//2):
tmp = to_networkx(graph, to_undirected=True)
tmp.remove_edge(*edges[i])
tmp = from_networkx(tmp)
a = graph.clone()
a.edge_index = tmp.edge_index
actions.add(a)
return actions
def test_from_networkx():
x = torch.Tensor([[1, 2], [3, 4]])
pos = torch.Tensor([[0, 0], [1, 1]])
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
edge_attr = torch.Tensor([1, 2, 3])
data = Data(x=x, pos=pos, edge_index=edge_index, edge_attr=edge_attr)
G = to_networkx(data, node_attrs=['x', 'pos'], edge_attrs=['edge_attr'])
data = from_networkx(G)
assert len(data) == 4
assert data.x.tolist() == x.tolist()
assert data.pos.tolist() == pos.tolist()
edge_index, edge_attr = coalesce(data.edge_index, data.edge_attr, 2, 2)
assert edge_index.tolist() == [[0, 0, 1], [0, 1, 0]]
assert edge_attr.tolist() == [3, 1, 2]
def __init__(self,
X: torch.Tensor,
edge_index: torch.Tensor,
num_hops: int,
n_rollout: int = 10,
min_atoms: int = 3,
c_puct: float = 10.0,
expand_atoms: int = 14,
high2low: bool = False,
node_idx: int = None,
score_func: Callable = None):
""" graph is a networkX graph """
self.X = X
self.edge_index = edge_index
self.num_hops = num_hops
self.data = Data(x=self.X, edge_index=self.edge_index)
self.graph = to_networkx(self.data, to_undirected=True)
self.data = Batch.from_data_list([self.data])
self.num_nodes = self.graph.number_of_nodes()
self.score_func = score_func
self.n_rollout = n_rollout
self.min_atoms = min_atoms
self.c_puct = c_puct
self.expand_atoms = expand_atoms
self.high2low = high2low
# extract the sub-graph and change the node indices.
if node_idx is not None:
self.ori_node_idx = node_idx
self.ori_graph = copy.copy(self.graph)
x, edge_index, subset, edge_mask, kwargs = \
self.__subgraph__(node_idx, self.X, self.edge_index, self.num_hops)
self.data = Batch.from_data_list(
[Data(x=x, edge_index=edge_index)])
self.graph = self.ori_graph.subgraph(subset.tolist())
mapping = {int(v): k for k, v in enumerate(subset)}
self.graph = nx.relabel_nodes(self.graph, mapping)
self.node_idx = torch.where(subset == self.ori_node_idx)[0]
self.num_nodes = self.graph.number_of_nodes()
self.subset = subset
self.root_coalition = sorted([node for node in range(self.num_nodes)])
self.MCTSNodeClass = partial(MCTSNode,
data=self.data,
ori_graph=self.graph,
c_puct=self.c_puct)
self.root = self.MCTSNodeClass(self.root_coalition)
self.state_map = {str(self.root.coalition): self.root}
def _data_to_nx(data, out=None, aux=None):
data = data.to('cpu')
out = out.to('cpu') if out is not None else out
aux = aux.to('cpu') if aux is not None else aux
G = to_networkx(data, node_attrs=None, edge_attrs=None)
pos = {node: p for node, p in zip(G.nodes, np.array(data.x))}
nx.set_node_attributes(G, pos, name='pos')
image = np.array(data.im.squeeze(), dtype=np.float64).transpose(2, 1, 0)
# chose colors from white to target_color
color = 1 - data.target_color
colors = th.stack([color[0]] * len(pos))
target = data.y if out is None else out
colors = 1 - target.unsqueeze(1) * colors
center = data.square_center.squeeze() if aux is None else aux.squeeze()
return image, G, center, colors
def main():
global graph
args = get_args()
random = np.random.RandomState(0)
for k in range(1, 11, 2):
print(f"=========== K-hop neighborhood sizes for k={k} ===========")
for p in np.arange(0.1, 1.1, 0.1):
dataset = get_dataset(args.dataset)
tg_dataset_prune(tg_dataset=[dataset.data],
method="random",
p=p,
random=random,
complement=False)
graph = to_networkx(data=dataset.data)
var = get_k_hop_neighborhood_size_variance_for(k)
print(f"({p}, {var})")
def load_dataset(name):
""" Load real-world datasets, available in PyTorch Geometric.
Used as a helper for DiskDataSource.
"""
task = "graph"
if name == "enzymes":
dataset = TUDataset(root="/tmp/ENZYMES", name="ENZYMES")
elif name == "proteins":
dataset = TUDataset(root="/tmp/PROTEINS", name="PROTEINS")
elif name == "cox2":
dataset = TUDataset(root="/tmp/cox2", name="COX2")
elif name == "aids":
dataset = TUDataset(root="/tmp/AIDS", name="AIDS")
elif name == "reddit-binary":
dataset = TUDataset(root="/tmp/REDDIT-BINARY", name="REDDIT-BINARY")
elif name == "imdb-binary":
dataset = TUDataset(root="/tmp/IMDB-BINARY", name="IMDB-BINARY")
elif name == "firstmm_db":
dataset = TUDataset(root="/tmp/FIRSTMM_DB", name="FIRSTMM_DB")
elif name == "dblp":
dataset = TUDataset(root="/tmp/DBLP_v1", name="DBLP_v1")
elif name == "ppi":
dataset = PPI(root="/tmp/PPI")
elif name == "qm9":
dataset = QM9(root="/tmp/QM9")
elif name == "atlas":
dataset = [g for g in nx.graph_atlas_g()[1:] if nx.is_connected(g)]
if task == "graph":
train_len = int(0.8 * len(dataset))
train, test = [], []
dataset = list(dataset)
random.shuffle(dataset)
has_name = hasattr(dataset[0], "name")
for i, graph in tqdm(enumerate(dataset)):
if not type(graph) == nx.Graph:
if has_name: del graph.name
graph = pyg_utils.to_networkx(graph).to_undirected()
if i < train_len:
train.append(graph)
else:
test.append(graph)
return train, test, task
def visualize_subgraph(self, edge_index, edge_mask, num_nodes,
threshold=None, **kwargs):
assert edge_mask.size(0) == edge_index.size(1)
if threshold is not None:
edge_mask = (edge_mask >= threshold).to(torch.float)
print(edge_mask)
data = Data(edge_index=edge_index, att=edge_mask).to('cpu')
data.num_nodes = num_nodes
G = to_networkx(data, edge_attrs=['att'])
# kwargs['with_labels'] = kwargs.get('with_labels') or True
kwargs['font_size'] = kwargs.get('font_size') or 10
node_size = kwargs.get('node_size') or 800
# kwargs['cmap'] = kwargs.get('cmap') or 'cool'
SCALE = 2
pos = nx.rescale_layout_dict(nx.kamada_kawai_layout(G), scale=SCALE)
ax = plt.gca()
ax.set_xlim((-SCALE - 0.1, SCALE + 0.1))
ax.set_ylim((-SCALE - 0.1, SCALE + 0.1))
for source, target, data in G.to_undirected().edges(data=True):
ax.annotate(
'',
xy=pos[target],
xycoords='data',
xytext=pos[source],
textcoords='data', arrowprops=dict(
arrowstyle="-",
alpha=max(data['att'], 0.1),
shrinkA=sqrt(node_size) / 2.0,
shrinkB=sqrt(node_size) / 2.0,
connectionstyle="arc3,rad=0.00",
))
nx.draw_networkx_labels(G, pos, **kwargs)
return ax, G
def save_results(base_path, queue, args, quantitative=False):
output_dir = base_path + f"/meg_output/{args['sample']}"
embedding_dir = output_dir + "/embeddings"
gexf_dir = output_dir + "/gexf_data"
if not os.path.exists(output_dir):
os.makedirs(output_dir)
os.makedirs(embedding_dir)
os.makedirs(gexf_dir)
for i, molecule in enumerate(queue):
np.save(embedding_dir + f"/{i}", molecule.pop('encoding'))
pyg = molecule.pop('pyg')
if quantitative:
g = to_networkx(pyg, to_undirected=True)
nx.write_gexf(g, f"{gexf_dir}/{i}.{molecule['prediction']['class']}.gexf")
with open(output_dir + "/data.json", "w") as outf:
json.dump(queue, outf, indent=2)
def positional_encoding(pg_graph, pos_enc_dim=20):
"""
Graph positional encoding v/ Laplacian eigenvectors (adapted from https://github.com/graphdeeplearning/benchmarking-gnns)
"""
g = dgl.DGLGraph()
g.from_networkx(nx_graph=to_networkx(pg_graph))
n = g.number_of_nodes()
# Laplacian
A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1)**-0.5,
dtype=float)
L = sp.eye(n) - N * A * N
# Eigenvectors
EigVal, EigVec = np.linalg.eig(L.toarray())
idx = EigVal.argsort() # increasing order
EigVal, EigVec = EigVal[idx], np.real(EigVec[:, idx])
pg_graph.x = torch.from_numpy(EigVec[:, 1:pos_enc_dim + 1]).float()
return pg_graph
def test_to_networkx():
x = torch.Tensor([[1, 2], [3, 4]])
pos = torch.Tensor([[0, 0], [1, 1]])
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
edge_attr = torch.Tensor([1, 2, 3])
data = Data(x=x, pos=pos, edge_index=edge_index, weight=edge_attr)
for remove_self_loops in [True, False]:
G = to_networkx(data, node_attrs=['x', 'pos'], edge_attrs=['weight'],
remove_self_loops=remove_self_loops)
assert G.nodes[0]['x'] == [1, 2]
assert G.nodes[1]['x'] == [3, 4]
assert G.nodes[0]['pos'] == [0, 0]
assert G.nodes[1]['pos'] == [1, 1]
if remove_self_loops:
assert nx.to_numpy_matrix(G).tolist() == [[0, 1], [2, 0]]
else:
assert nx.to_numpy_matrix(G).tolist() == [[3, 1], [2, 0]]
def plot_edge3d(R, ix, raw_input, features, x, y, z):
r_node = R['node'].clone()
r_edge = R['edge'].clone()
r_node[torch.isnan(r_node)] = 0
r_edge[torch.isnan(r_edge)] = 0
edge_shade = torch.norm(r_edge, dim=1)
edge_alpha = (edge_shade - edge_shade.min()) / edge_shade.max()
edge_alpha = edge_alpha.detach().cpu().numpy()
node_shade = np.linalg.norm(r_node.detach().cpu().numpy(), axis=1)
x_idx = features.index(x)
y_idx = features.index(y)
z_idx = features.index(z)
raw_input.edge_alpha = edge_alpha
raw_input.node_shade = node_shade
pos = np.array(
list(
zip(raw_input.x[:, x_idx].detach().numpy(),
raw_input.x[:, y_idx].detach().numpy(),
raw_input.x[:, z_idx].detach().numpy())))
raw_input.pos = pos
G = to_networkx(raw_input, node_attrs=["pos", "node_shade"])
if R[ix]['label'][:, 1] > 0:
title_str = "Edge and Node Relevance in 3d Space for Higgs boson Signal\n\n"
else:
title_str = "Edge and Node Relevance in 3d Space for background\n\n"
title_str += "prediction:{}".format(
R[ix]["pred"].detach().cpu().numpy().round(4)[0])
network_plot_3D(G,
45,
raw_input.y.detach(),
edge_alpha,
title_str,
zlabel=z)
def get_subgraph(self, node_idx, x, edge_index, y, **kwargs):
num_nodes, num_edges = x.size(0), edge_index.size(1)
graph = to_networkx(data=Data(x=x, edge_index=edge_index), to_undirected=True)
subset, edge_index, _, edge_mask = k_hop_subgraph_with_default_whole_graph(
node_idx, self.num_hops, edge_index, relabel_nodes=True,
num_nodes=num_nodes, flow=self.__flow__())
mapping = {int(v): k for k, v in enumerate(subset)}
subgraph = graph.subgraph(subset.tolist())
nx.relabel_nodes(subgraph, mapping)
x = x[subset]
for key, item in kwargs.items():
if torch.is_tensor(item) and item.size(0) == num_nodes:
item = item[subset]
elif torch.is_tensor(item) and item.size(0) == num_edges:
item = item[edge_mask]
kwargs[key] = item
y = y[subset]
return x, edge_index, y, subset, kwargs
def get_subgraph_list(data, n):
"""
Gets all subgroups of size n.
Parameters
----------
G : networkx graph
Returns
-------
k_paths : list of k-paths
"""
G = to_networkx(data, to_undirected=True, remove_self_loops=True)
N = G.number_of_nodes()
indices = list(range(N))
subgraphs = []
for node_set in combinations(indices, n):
G_sub = G.subgraph(node_set)
if is_connected(G_sub):
subgraphs.append(node_set)
return subgraphs
def test_from_networkx_group_attrs():
x = torch.randn(2, 2)
x1 = torch.randn(2, 4)
x2 = torch.randn(2, 8)
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
edge_attr1 = torch.randn(edge_index.size(1))
edge_attr2 = torch.randn(edge_index.size(1))
perm = torch.tensor([0, 2, 1])
data = Data(x=x, x1=x1, x2=x2, edge_index=edge_index,
edge_attr1=edge_attr1, edge_attr2=edge_attr2)
G = to_networkx(data, node_attrs=['x', 'x1', 'x2'],
edge_attrs=['edge_attr1', 'edge_attr2'])
data = from_networkx(G, group_node_attrs=['x', 'x2'], group_edge_attrs=all)
assert len(data) == 7
assert data.x.tolist() == torch.cat([x, x2], dim=-1).tolist()
assert data.x1.tolist() == x1.tolist()
assert data.x2.tolist() == x2.tolist()
assert data.edge_index.tolist() == edge_index[:, perm].tolist()
assert data.edge_attr1.tolist() == edge_attr1[perm].tolist()
assert data.edge_attr2.tolist() == edge_attr2[perm].tolist()
assert data.edge_attr.tolist() == torch.stack([edge_attr1, edge_attr2],
dim=-1)[perm].tolist()
