def create_random_bigraph():
dim = 4
src_num_nodes = np.random.randint(low=2, high=10)
dst_num_nodes = np.random.randint(low=2, high=10)
edges_size = np.random.randint(low=1, high=10)
edges_src = np.random.randint(low=1,
high=src_num_nodes,
size=[edges_size, 1])
edges_dst = np.random.randint(low=1,
high=dst_num_nodes,
size=[edges_size, 1])
edges = np.hstack([edges_src, edges_dst])
src_nfeat = np.random.randn(src_num_nodes, dim)
dst_nfeat = np.random.randn(dst_num_nodes, dim)
efeat = np.random.randn(len(edges), dim)
g = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'nfeat': src_nfeat},
dst_node_feat={'nfeat': dst_nfeat},
edge_feat={'efeat': efeat})
return g
def test_check_to_tensor_to_numpy(self):
src_num_nodes = 4
dst_num_nodes = 5
dim = 4
edges = [(0, 1), (1, 2), (3, 4)]
src_nfeat = np.random.randn(src_num_nodes, dim)
dst_nfeat = np.random.randn(dst_num_nodes, dim)
efeat = np.random.randn(len(edges), dim)
g1 = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'src_nfeat': src_nfeat},
dst_node_feat={'dst_nfeat': dst_nfeat},
edge_feat={'efeat': efeat})
# check inplace to tensor
self.assertFalse(g1.is_tensor())
g2 = g1.tensor(inplace=False)
g3 = g2.numpy(inplace=False)
self.assertFalse(g1.is_tensor())
self.assertTrue(g2.is_tensor())
self.assertFalse(g3.is_tensor())
def test_build_tensor_graph(self):
src_num_nodes = paddle.to_tensor(4)
dst_num_nodes = paddle.to_tensor(5)
e = np.array([(0, 1), (1, 2), (3, 4)])
edges = paddle.to_tensor(e)
g2 = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes)
def test_send_func(self):
src_num_nodes = 4
dst_num_nodes = 5
edges = [(0, 1), (1, 2), (3, 4)]
src_nfeat = np.arange(src_num_nodes).reshape(-1, 1)
dst_nfeat = np.arange(dst_num_nodes).reshape(-1, 1)
efeat = np.arange(len(edges)).reshape(-1, 1)
g1 = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'src_nfeat': src_nfeat},
dst_node_feat={'dst_nfeat': dst_nfeat},
edge_feat={'efeat': efeat})
target_src = np.array([0, 1, 3]).reshape(-1, 1)
target_dst = np.array([1, 2, 4]).reshape(-1, 1)
target_edge = np.array([0, 1, 2]).reshape(-1, 1)
g1.tensor()
src_copy = lambda sf, df, ef: {"h": sf["h"], "e": ef["e"]}
dst_copy = lambda sf, df, ef: {"h": df["h"], "e": ef["e"]}
both_copy = lambda sf, df, ef: {
"sh": sf["h"],
"dh": df["h"],
"e": ef["e"]
}
msg = g1.send(src_copy,
src_feat={"h": g1.src_node_feat["src_nfeat"]},
edge_feat={'e': g1.edge_feat['efeat']})
self.assertTrue((msg['h'].numpy() == target_src).all())
self.assertTrue((msg['e'].numpy() == target_edge).all())
msg = g1.send(dst_copy,
dst_feat={"h": g1.dst_node_feat["dst_nfeat"]},
edge_feat={'e': g1.edge_feat['efeat']})
self.assertTrue((msg['h'].numpy() == target_dst).all())
self.assertTrue((msg['e'].numpy() == target_edge).all())
msg = g1.send(both_copy,
src_feat={"h": g1.src_node_feat["src_nfeat"]},
dst_feat={"h": g1.dst_node_feat["dst_nfeat"]},
edge_feat={'e': g1.edge_feat['efeat']})
self.assertTrue((msg['sh'].numpy() == target_src).all())
self.assertTrue((msg['dh'].numpy() == target_dst).all())
self.assertTrue((msg['e'].numpy() == target_edge).all())
def test_build_graph(self):
src_num_nodes = 4
dst_num_nodes = 5
dim = 4
edges = [(0, 1), (1, 2), (3, 4)]
src_nfeat = np.random.randn(src_num_nodes, dim)
dst_nfeat = np.random.randn(dst_num_nodes, dim)
efeat = np.random.randn(len(edges), dim)
g1 = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'src_nfeat': src_nfeat},
dst_node_feat={'dst_nfeat': dst_nfeat},
edge_feat={'efeat': efeat})
def test_num_nodes_valid(self):
src_num_nodes = 1
dst_num_nodes = 5
dim = 4
edges = [(0, 1), (1, 2), (3, 4)]
src_nfeat = np.random.randn(src_num_nodes, dim)
dst_nfeat = np.random.randn(dst_num_nodes, dim)
efeat = np.random.randn(len(edges), dim)
with self.assertRaises(ValueError):
g1 = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'src_nfeat': src_nfeat},
dst_node_feat={'dst_nfeat': dst_nfeat},
edge_feat={'efeat': efeat})
def test_check_degree(self):
"""Check the de
"""
src_num_nodes = 4
dst_num_nodes = 5
edges = [(0, 1), (1, 2), (3, 4)]
g1 = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes)
indegree = np.array([0, 1, 1, 0, 1], dtype="int32")
outdegree = np.array([1, 1, 0, 1], dtype="int32")
# check degree in numpy
res = g1.indegree()
self.assertTrue(np.all(res == indegree))
# get degree from specific nodes
res = g1.indegree(nodes=[1, 2, 3, 4])
self.assertTrue(np.all(res == indegree[[1, 2, 3, 4]]))
res = g1.outdegree()
self.assertTrue(np.all(res == outdegree))
# get degree from specific nodes
res = g1.outdegree(nodes=[1, 2, 3])
self.assertTrue(np.all(res == outdegree[[1, 2, 3]]))
# check degree in Tensor
g1.tensor()
res = g1.indegree().numpy()
self.assertTrue(np.all(res == indegree))
# get degree from specific nodes
res = g1.indegree(nodes=paddle.to_tensor([1, 2, 3])).numpy()
self.assertTrue(np.all(res == indegree[[1, 2, 3]]))
res = g1.outdegree().numpy()
self.assertTrue(np.all(res == outdegree))
# get degree from specific nodes
res = g1.outdegree(nodes=paddle.to_tensor([1, 2, 3])).numpy()
self.assertTrue(np.all(res == outdegree[[1, 2, 3]]))
def test_dump_tensor_load_tensor(self):
path = './tmp'
dim = 4
src_num_nodes = np.random.randint(low=2, high=10)
dst_num_nodes = np.random.randint(low=2, high=10)
edges_size = np.random.randint(low=1, high=10)
edges_src = np.random.randint(low=1,
high=src_num_nodes,
size=[edges_size, 1])
edges_dst = np.random.randint(low=1,
high=dst_num_nodes,
size=[edges_size, 1])
edges = np.hstack([edges_src, edges_dst])
src_nfeat = np.random.randn(src_num_nodes, dim)
dst_nfeat = np.random.randn(dst_num_nodes, dim)
efeat = np.random.randn(len(edges), dim)
g = pgl.BiGraph(edges=paddle.to_tensor(edges),
src_num_nodes=paddle.to_tensor(src_num_nodes),
dst_num_nodes=paddle.to_tensor(dst_num_nodes),
src_node_feat={'nfeat': paddle.to_tensor(src_nfeat)},
dst_node_feat={'nfeat': paddle.to_tensor(dst_nfeat)},
edge_feat={'efeat': paddle.to_tensor(efeat)})
in_before = g.indegree()
g.outdegree()
g.tensor()
# Merge Graph
g.dump(path)
g2 = pgl.BiGraph.load(path)
in_after = g2.indegree()
for a, b in zip(in_before, in_after):
self.assertEqual(a, b)
del g2
del in_after
import shutil
shutil.rmtree(path)
def test_send_recv_func(self):
np.random.seed(0)
src_num_nodes = 5
dst_num_nodes = 4
edges = [(0, 1), (1, 2), (3, 3), (4, 1), (1, 0)]
src_nfeat = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6],
[4, 5, 6, 7], [5, 6, 7, 8]])
ground = np.array([[2, 3, 4, 5], [6, 8, 10, 12], [2, 3, 4, 5],
[4, 5, 6, 7]])
g = pgl.BiGraph(
edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'src_nfeat': src_nfeat},
)
g.tensor()
output = g.send_recv(g.src_node_feat['src_nfeat'], reduce_func="sum")
output = output.numpy()
self.assertTrue((ground == output).all())
def test_build_graph_without_num_nodes(self):
e = np.array([(0, 1), (1, 2), (3, 4)])
edges = paddle.to_tensor(e)
g2 = pgl.BiGraph(edges=edges)
def test_send_and_recv(self):
np.random.seed(0)
src_num_nodes = 5
dst_num_nodes = 4
edges = [(0, 1), (1, 2), (3, 3), (4, 1), (1, 0)]
src_nfeat = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6],
[4, 5, 6, 7], [5, 6, 7, 8]],
dtype="float32")
dst_nfeat = np.array(
[[2, 3, 4, 5], [3, 4, 5, 6], [4, 5, 6, 7], [5, 6, 7, 8]],
dtype="float32")
src_ground = np.array([[1, 2, 3, 4], [2, 3, 4, 5], [4, 5, 6, 7],
[5, 6, 7, 8], [2, 3, 4, 5]],
dtype="float32")
src_recv = np.array([[2, 3, 4, 5], [6, 8, 10, 12], [2, 3, 4, 5],
[4, 5, 6, 7]])
dst_ground = np.array([[3, 4, 5, 6], [4, 5, 6, 7], [5, 6, 7, 8],
[3, 4, 5, 6], [2, 3, 4, 5]],
dtype="float32")
dst_recv = np.array([[3, 4, 5, 6], [6, 8, 10, 12], [0, 0, 0, 0],
[5, 6, 7, 8], [3, 4, 5, 6]])
g = pgl.BiGraph(edges=edges,
src_num_nodes=src_num_nodes,
dst_num_nodes=dst_num_nodes,
src_node_feat={'src_nfeat': src_nfeat},
dst_node_feat={'dst_nfeat': dst_nfeat})
g.tensor()
def send_func1(src_feat, dst_feat, edge_feat):
return src_feat
def send_func2(src_feat, dst_feat, edge_feat):
return {'h': src_feat['h']}
def reduce_func(msg):
return msg.reduce_sum(msg['h'])
# test send_func1
msg1 = g.send(send_func1, src_feat={'h': g.src_node_feat['src_nfeat']})
_msg = msg1['h'].numpy()
self.assertTrue((src_ground == _msg).all())
output = g.recv(reduce_func, msg1)
output = output.numpy()
self.assertTrue((src_recv == output).all())
# test send_func2
msg2 = g.send(send_func2, src_feat={'h': g.src_node_feat['src_nfeat']})
_msg = msg2['h'].numpy()
self.assertTrue((src_ground == _msg).all())
output = g.recv(reduce_func, msg2)
output = output.numpy()
self.assertTrue((src_recv == output).all())
def send_func1_d(src_feat, dst_feat, edge_feat):
return dst_feat
def send_func2_d(src_feat, dst_feat, edge_feat):
return {'h': dst_feat['h']}
def reduce_func_d(msg):
return msg.reduce_sum(msg['h'])
# test send_func1
msg1 = g.send(send_func1_d,
dst_feat={'h': g.dst_node_feat['dst_nfeat']})
_msg = msg1['h'].numpy()
self.assertTrue((dst_ground == _msg).all())
output = g.recv(reduce_func_d, msg1, recv_mode="src")
output = output.numpy()
self.assertTrue((dst_recv == output).all())
# test send_func2
msg2 = g.send(send_func2_d,
dst_feat={'h': g.dst_node_feat['dst_nfeat']})
_msg = msg2['h'].numpy()
self.assertTrue((dst_ground == _msg).all())
output = g.recv(reduce_func_d, msg2, recv_mode="src")
output = output.numpy()
self.assertTrue((dst_recv == output).all())
def test_disjoint_tensor_graph(self):
glist = []
dim = 4
for i in range(5):
src_num_nodes = np.random.randint(low=2, high=10)
dst_num_nodes = np.random.randint(low=2, high=10)
edges_size = np.random.randint(low=1, high=10)
edges_src = np.random.randint(low=1,
high=src_num_nodes,
size=[edges_size, 1])
edges_dst = np.random.randint(low=1,
high=dst_num_nodes,
size=[edges_size, 1])
edges = np.hstack([edges_src, edges_dst])
src_nfeat = np.random.randn(src_num_nodes, dim)
dst_nfeat = np.random.randn(dst_num_nodes, dim)
efeat = np.random.randn(len(edges), dim)
g = pgl.BiGraph(
edges=paddle.to_tensor(edges),
src_num_nodes=paddle.to_tensor(src_num_nodes),
dst_num_nodes=paddle.to_tensor(dst_num_nodes),
src_node_feat={'nfeat': paddle.to_tensor(src_nfeat)},
dst_node_feat={'nfeat': paddle.to_tensor(dst_nfeat)},
edge_feat={'efeat': paddle.to_tensor(efeat)})
glist.append(g)
# Merge Graph
multi_graph = pgl.BiGraph.disjoint(glist)
# Check Graph Index
src_node_index = [
np.ones(g.src_num_nodes) * n for n, g in enumerate(glist)
]
dst_node_index = [
np.ones(g.dst_num_nodes) * n for n, g in enumerate(glist)
]
edge_index = [np.ones(g.num_edges) * n for n, g in enumerate(glist)]
src_node_index = np.concatenate(src_node_index)
dst_node_index = np.concatenate(dst_node_index)
edge_index = np.concatenate(edge_index)
self.assertTrue(
np.all(src_node_index == multi_graph.graph_src_node_id.numpy()))
self.assertTrue(
np.all(dst_node_index == multi_graph.graph_dst_node_id.numpy()))
self.assertTrue(
np.all(edge_index == multi_graph.graph_edge_id.numpy()))
multi_graph.tensor()
self.assertTrue(
np.all(src_node_index == multi_graph.graph_src_node_id.numpy()))
self.assertTrue(
np.all(dst_node_index == multi_graph.graph_dst_node_id.numpy()))
self.assertTrue(
np.all(edge_index == multi_graph.graph_edge_id.numpy()))
multi_graph = pgl.BiGraph.disjoint([glist[0]])
self.assertEqual(multi_graph.src_node_feat['nfeat'].shape,
[glist[0].src_num_nodes.numpy()[0], dim])
self.assertEqual(multi_graph.dst_node_feat['nfeat'].shape,
[glist[0].dst_num_nodes.numpy()[0], dim])
def build_graph(self, mol):
num_atoms_d, coords, features, atoms, inter_feats = mol
##################################################
# prepare distance matrix and interaction matrix #
##################################################
dist_mat = distance.cdist(coords, coords, 'euclidean')
np.fill_diagonal(dist_mat, np.inf)
inter_feats = np.array([inter_feats])
inter_feats = inter_feats / inter_feats.sum()
############################
# build atom to atom graph #
############################
num_atoms = len(coords)
dist_graph_base = dist_mat.copy()
dist_feat = dist_graph_base[dist_graph_base < self.cut_dist].reshape(
-1, 1)
dist_graph_base[dist_graph_base >= self.cut_dist] = 0.
atom_graph = coo_matrix(dist_graph_base)
a2a_edges = list(zip(atom_graph.row, atom_graph.col))
a2a_graph = pgl.Graph(a2a_edges,
num_nodes=num_atoms,
node_feat={"feat": features},
edge_feat={"dist": dist_feat})
######################
# prepare bond nodes #
######################
indices = []
bond_pair_atom_types = []
for i in range(num_atoms):
for j in range(num_atoms):
a = dist_mat[i, j]
if a < self.cut_dist:
at_i, at_j = atoms[i], atoms[j]
if i < num_atoms_d and j >= num_atoms_d and (
at_j, at_i) in pair_ids:
bond_pair_atom_types += [pair_ids.index((at_j, at_i))]
elif i >= num_atoms_d and j < num_atoms_d and (
at_i, at_j) in pair_ids:
bond_pair_atom_types += [pair_ids.index((at_i, at_j))]
else:
bond_pair_atom_types += [-1]
indices.append([i, j])
############################
# build bond to atom graph #
############################
num_bonds = len(indices)
assignment_b2a = np.zeros((num_bonds, num_atoms),
dtype=np.int64) # Maybe need too much memory
assignment_a2b = np.zeros((num_atoms, num_bonds),
dtype=np.int64) # Maybe need too much memory
for i, idx in enumerate(indices):
assignment_b2a[i, idx[1]] = 1
assignment_a2b[idx[0], i] = 1
b2a_graph = coo_matrix(assignment_b2a)
b2a_edges = list(zip(b2a_graph.row, b2a_graph.col))
b2a_graph = pgl.BiGraph(b2a_edges,
src_num_nodes=num_bonds,
dst_num_nodes=num_atoms)
############################
# build bond to bond graph #
############################
bond_graph_base = assignment_b2a @ assignment_a2b
np.fill_diagonal(bond_graph_base, 0) # eliminate self connections
bond_graph_base[range(num_bonds),
[indices.index([x[1], x[0]]) for x in indices]] = 0
x, y = np.where(bond_graph_base > 0)
num_edges = len(x)
# calculate angle
angle_feat = np.zeros_like(x, dtype=np.float32)
for i in range(num_edges):
body1 = indices[x[i]]
body2 = indices[y[i]]
a = dist_mat[body1[0], body1[1]]
b = dist_mat[body2[0], body2[1]]
c = dist_mat[body1[0], body2[1]]
if a == 0 or b == 0:
print(body1, body2)
print('One distance is zero.')
angle_feat[i] = 0.
return None, None
# exit(-1)
else:
angle_feat[i] = cos_formula(a, b, c)
# angle domain divisions
unit = 180.0 / self.num_angle
angle_index = (np.rad2deg(angle_feat) / unit).astype('int64')
angle_index = np.clip(angle_index, 0, self.num_angle - 1)
# multiple bond-to-bond graphs based on angle domains
b2b_edges_list = [[] for _ in range(self.num_angle)]
b2b_angle_list = [[] for _ in range(self.num_angle)]
for i, (ind, radian) in enumerate(zip(angle_index, angle_feat)):
b2b_edges_list[ind].append((x[i], y[i]))
b2b_angle_list[ind].append(radian)
# b2b_graph_list = [[] for _ in range(self.num_angle)]
b2b_graph_list = []
for ind in range(self.num_angle):
b2b_graph = pgl.Graph(b2b_edges_list[ind],
num_nodes=num_bonds,
edge_feat={"angle": b2b_angle_list[ind]})
b2b_graph_list.append(b2b_graph)
#########################################
# build index for inter-molecular bonds #
#########################################
bond_types = bond_pair_atom_types
type_count = [0 for _ in range(len(pair_ids))]
for type_i in bond_types:
if type_i != -1:
type_count[type_i] += 1
bond_types = np.array(bond_types)
type_count = np.array(type_count)
graphs = a2a_graph, b2a_graph, b2b_graph_list
global_feat = inter_feats, bond_types, type_count
return graphs, global_feat
def build_graph(self, mol):
num_atoms, atom_features, atom_types, atom_3dcoords, bond_features = mol
atom_features = np.array(atom_features)
dist_mat = distance.cdist(atom_3dcoords, atom_3dcoords, 'euclidean')
np.fill_diagonal(dist_mat, np.inf)
if num_atoms == 1:
return None
if len(bond_features) == 0:
print('NO BOND FEATURES,', num_atoms)
return None
# node-to-node graph
num_nodes = num_atoms
dist_graph_base = dist_mat.copy()
dist_feats = dist_graph_base[dist_graph_base < self.cut_dist].reshape(
-1, 1)
dist_graph_base[dist_graph_base >= self.cut_dist] = 0.
atom_graph = coo_matrix(dist_graph_base)
a2a_edges = list(zip(atom_graph.row, atom_graph.col))
a2a_graph = pgl.Graph(a2a_edges,
num_nodes=num_nodes,
node_feat={"feat": atom_features},
edge_feat={"dist": dist_feats})
# edge-to-node graph
indices = []
for i in range(num_atoms):
for j in range(num_atoms):
a = dist_mat[i, j]
if a < self.cut_dist:
indices.append([i, j])
num_edges = len(indices)
assignment_e2a = np.zeros((num_edges, num_nodes), dtype=np.int64)
assignment_a2e = np.zeros((num_nodes, num_edges), dtype=np.int64)
for i, idx in enumerate(indices):
assignment_e2a[i, idx[1]] = 1
assignment_a2e[idx[0], i] = 1
edge2node_graph = coo_matrix(assignment_e2a)
e2a_edges = list(zip(edge2node_graph.row, edge2node_graph.col))
e2a_graph_list = []
if self.num_dist == 1:
e2a_graph = pgl.BiGraph(e2a_edges,
src_num_nodes=num_edges,
dst_num_nodes=num_nodes,
edge_feat={"dist": dist_feats})
e2a_graph_list += [e2a_graph]
else:
dist_inds = np.clip(dist_feats, 1.0, self.cut_dist - 1e-8).astype(
np.int64) - 1
if self.num_dist == 2:
inds = np.where(dist_inds == 0)[0]
e2a_edges_sub = [e2a_edges[i] for i in inds]
dist_feat_sub = dist_feats[inds]
if len(e2a_edges_sub) == 0:
e2a_edges_sub = [(0, 0)]
dist_feat_sub = dist_feats[[0]]
e2a_graph = pgl.BiGraph(e2a_edges_sub,
src_num_nodes=num_edges,
dst_num_nodes=num_nodes,
edge_feat={"dist": dist_feat_sub})
e2a_graph_list += [e2a_graph]
inds = np.where(dist_inds >= 1)[0]
e2a_edges_sub = [e2a_edges[i] for i in inds]
dist_feat_sub = dist_feats[inds]
if len(e2a_edges_sub) == 0:
e2a_edges_sub = [(0, 0)]
dist_feat_sub = dist_feats[[0]]
e2a_graph = pgl.BiGraph(e2a_edges_sub,
src_num_nodes=num_edges,
dst_num_nodes=num_nodes,
edge_feat={"dist": dist_feat_sub})
e2a_graph_list += [e2a_graph]
else:
for k in range(self.num_dist):
inds = np.where(dist_inds == k)[0]
e2a_edges_sub = [e2a_edges[i] for i in inds]
dist_feat_sub = dist_feats[inds]
if len(e2a_edges_sub) == 0:
e2a_edges_sub = [(0, 0)]
dist_feat_sub = dist_feats[[0]]
e2a_graph = pgl.BiGraph(e2a_edges_sub,
src_num_nodes=num_edges,
dst_num_nodes=num_nodes,
edge_feat={"dist": dist_feat_sub})
e2a_graph_list += [e2a_graph]
# edge-to-edge graphs
edge_graph_base = assignment_e2a @ assignment_a2e
np.fill_diagonal(edge_graph_base, 0) # eliminate self connections
edge_graph_base[range(num_edges),
[indices.index([x[1], x[0]]) for x in indices]] = 0
x, y = np.where(edge_graph_base > 0)
# calculate angle
angle_feat = np.zeros_like(x, dtype=np.float32)
for i in range(len(x)):
body1 = indices[x[i]]
body2 = indices[y[i]]
a = dist_mat[body1[0], body1[1]]
b = dist_mat[body2[0], body2[1]]
c = dist_mat[body1[0], body2[1]]
if a == 0 or b == 0:
print(body1, body2)
print('One distance is zero.')
return None
else:
angle_feat[i] = cos_formula(a, b, c)
# angle domain divisions
unit = 180.0 / self.num_angle
angle_index = (np.rad2deg(angle_feat) / unit).astype('int64')
angle_index = np.clip(angle_index, 0, self.num_angle - 1)
e2e_edges_list = [[] for _ in range(self.num_angle)]
e2e_angle_list = [[] for _ in range(self.num_angle)]
for i, (ind, radian) in enumerate(zip(angle_index, angle_feat)):
e2e_edges_list[ind].append((x[i], y[i]))
e2e_angle_list[ind].append(radian)
e2e_graph_list = []
for ind in range(self.num_angle):
e2e_graph = pgl.Graph(e2e_edges_list[ind],
num_nodes=num_edges,
edge_feat={"angle": e2e_angle_list[ind]})
e2e_graph_list.append(e2e_graph)
return a2a_graph, e2a_graph_list, e2e_graph_list
def build_graph(self, mol):
num_atoms, atom_features, atom_2dcoords, bond_features = mol
dist_mat = distance.cdist(atom_2dcoords, atom_2dcoords, 'euclidean')
np.fill_diagonal(dist_mat, np.inf)
if num_atoms == 1:
return None
if len(bond_features) == 0:
print('NO BOND FEATURES,', num_atoms)
return None
dist_feats = []
edge_feats = []
a2a_edges = []
indices = []
# build directional graph
for i in range(num_atoms):
for j in range(num_atoms):
ii, jj = min(i, j), max(i, j)
bf = bond_features.get((ii, jj))
if bf is None:
continue
a2a_edges.append((i, j))
dist_feats.append([dist_mat[i, j]])
edge_feats.append(bf)
indices.append([i, j])
num_nodes = num_atoms
num_edges = len(indices)
# edge-to-node and node-to-edge graph
assignment_e2a = np.zeros((num_edges, num_nodes), dtype=np.int64)
assignment_a2e = np.zeros((num_nodes, num_edges), dtype=np.int64)
for i, idx in enumerate(indices):
assignment_e2a[i, idx[1]] = 1
assignment_a2e[idx[0], i] = 1
edge2node_graph = coo_matrix(assignment_e2a)
node2edge_graph = coo_matrix(assignment_a2e)
# edge-to-edge graph
edge_graph_base = assignment_e2a @ assignment_a2e
np.fill_diagonal(edge_graph_base, 0) # eliminate self connections
edge_graph_base[
range(num_edges),
[indices.index([x[1], x[0]])
for x in indices]] = 0 # eliminate connections of the same edge
x, y = np.where(edge_graph_base > 0)
angle_feats = []
for i in range(len(x)):
body1 = indices[x[i]]
body2 = indices[y[i]]
a = dist_mat[body1[0], body1[1]]
b = dist_mat[body2[0], body2[1]]
c = dist_mat[body1[0], body2[1]]
if a == 0 or b == 0:
print(body1, body2)
print('One distance is zero.')
return None
else:
angle_feats.append(cos_formula(a, b, c))
atom_features = np.array(atom_features)
bond_features = np.array(edge_feats)
# pgl graph
# a2a_edges = list(zip(range(num_nodes), range(num_nodes)))
# a2e_edges = list(zip(node2edge_graph.row, node2edge_graph.col))
# a2e_graph = pgl.BiGraph(a2e_edges, src_num_nodes=num_nodes, dst_num_nodes=num_edges)
e2a_edges = list(zip(edge2node_graph.row, edge2node_graph.col))
e2e_edges = list(zip(x, y))
e2a_graph = pgl.BiGraph(e2a_edges,
src_num_nodes=num_edges,
dst_num_nodes=num_nodes)
# print(num_nodes, num_edges, angle_feats)
# assert len(np.array(angle_feats).shape) == 2
a2a_graph = pgl.Graph(
a2a_edges,
num_nodes=num_nodes,
node_feat={"feat": atom_features},
edge_feat={"dist":
dist_feats}) # dist_feats: (num_edges_of_node, )
e2e_graph = pgl.Graph(
e2e_edges,
num_nodes=num_edges,
node_feat={"feat": bond_features},
edge_feat={"angle":
angle_feats}) # angle_feats: (num_edges_of_edge, )
return a2a_graph, e2a_graph, e2e_graph
