def test_filter():
g = DGLGraph()
g.add_nodes(4)
g.add_edges([0, 1, 2, 3], [1, 2, 3, 0])
n_repr = np.zeros((4, 5))
e_repr = np.zeros((4, 5))
n_repr[[1, 3]] = 1
e_repr[[1, 3]] = 1
n_repr = F.copy_to(F.zerocopy_from_numpy(n_repr), F.ctx())
e_repr = F.copy_to(F.zerocopy_from_numpy(e_repr), F.ctx())
g.ndata['a'] = n_repr
g.edata['a'] = e_repr
def predicate(r):
return F.max(r.data['a'], 1) > 0
# full node filter
n_idx = g.filter_nodes(predicate)
assert set(F.zerocopy_to_numpy(n_idx)) == {1, 3}
# partial node filter
n_idx = g.filter_nodes(predicate, [0, 1])
assert set(F.zerocopy_to_numpy(n_idx)) == {1}
# full edge filter
e_idx = g.filter_edges(predicate)
assert set(F.zerocopy_to_numpy(e_idx)) == {1, 3}
# partial edge filter
e_idx = g.filter_edges(predicate, [0, 1])
assert set(F.zerocopy_to_numpy(e_idx)) == {1}
def test_node_frame():
g = dgl.DGLGraph()
g.add_nodes(10)
data = np.random.rand(10, 3)
new_data = data.take([0, 1, 2, 7, 8, 9], axis=0)
g.ndata['h'] = F.zerocopy_from_numpy(data)
# remove nodes
g.remove_nodes(range(3, 7))
assert F.allclose(g.ndata['h'], F.zerocopy_from_numpy(new_data))
def test_pack_traces():
traces, types = (np.array(
[[ 0, 1, -1, -1, -1, -1, -1],
[ 0, 1, 1, 3, 0, 0, 0]], dtype='int64'),
np.array([0, 0, 1, 0, 0, 1, 0], dtype='int64'))
traces = F.zerocopy_from_numpy(traces)
types = F.zerocopy_from_numpy(types)
result = dgl.sampling.pack_traces(traces, types)
assert F.array_equal(result[0], F.tensor([0, 1, 0, 1, 1, 3, 0, 0, 0], dtype=F.int64))
assert F.array_equal(result[1], F.tensor([0, 0, 0, 0, 1, 0, 0, 1, 0], dtype=F.int64))
assert F.array_equal(result[2], F.tensor([2, 7], dtype=F.int64))
assert F.array_equal(result[3], F.tensor([0, 2], dtype=F.int64))
def test_partial_edge_softmax():
g = dgl.DGLGraph()
g.add_nodes(30)
# build a complete graph
for i in range(30):
for j in range(30):
g.add_edge(i, j)
score = F.randn((300, 1))
score.requires_grad_()
grad = F.randn((300, 1))
import numpy as np
eids = np.random.choice(900, 300, replace=False).astype('int64')
eids = F.zerocopy_from_numpy(eids)
# compute partial edge softmax
y_1 = nn.edge_softmax(g, score, eids)
y_1.backward(grad)
grad_1 = score.grad
score.grad.zero_()
# compute edge softmax on edge subgraph
subg = g.edge_subgraph(eids)
y_2 = nn.edge_softmax(subg, score)
y_2.backward(grad)
grad_2 = score.grad
score.grad.zero_()
assert F.allclose(y_1, y_2)
assert F.allclose(grad_1, grad_2)
def test_partial_edge_softmax():
g = dgl.DGLGraph()
g.add_nodes(30)
# build a complete graph
for i in range(30):
for j in range(30):
g.add_edge(i, j)
score = F.randn((300, 1))
grad = F.randn((300, 1))
import numpy as np
eids = np.random.choice(900, 300, replace=False).astype('int64')
eids = F.zerocopy_from_numpy(eids)
# compute partial edge softmax
with tf.GradientTape() as tape:
tape.watch(score)
y_1 = nn.edge_softmax(g, score, eids)
grads = tape.gradient(y_1, [score])
grad_1 = grads[0]
# compute edge softmax on edge subgraph
subg = g.edge_subgraph(eids)
with tf.GradientTape() as tape:
tape.watch(score)
y_2 = nn.edge_softmax(subg, score)
grads = tape.gradient(y_2, [score])
grad_2 = grads[0]
assert F.allclose(y_1, y_2)
assert F.allclose(grad_1, grad_2)
def test_sample_neighbors_exclude_edges_homoG(dtype):
u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300,size=100, dtype=dtype)))
v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=u_nodes.shape, dtype=dtype))
g = dgl.graph((u_nodes, v_nodes))
(U, V, EID) = (0, 1, 2)
nd_b_idx = np.random.randint(low=1,high=24, dtype=dtype)
nd_e_idx = np.random.randint(low=25,high=49, dtype=dtype)
b_idx = np.random.randint(low=1,high=24, dtype=dtype)
e_idx = np.random.randint(low=25,high=49, dtype=dtype)
sampled_amount = np.random.randint(low=1,high=10, dtype=dtype)
g_edges = g.all_edges(form='all')
excluded_edges = g_edges[EID][b_idx:e_idx]
sampled_node = g_edges[V][nd_b_idx:nd_e_idx]
excluded_nodes_U = g_edges[U][b_idx:e_idx]
excluded_nodes_V = g_edges[V][b_idx:e_idx]
sg = dgl.sampling.sample_neighbors(g, sampled_node,
sampled_amount, exclude_edges=excluded_edges)
assert not np.any(F.asnumpy(sg.has_edges_between(excluded_nodes_U,excluded_nodes_V)))
def test_sample_neighbors_exclude_edges_heteroG(dtype):
d_i_d_u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300, size=100, dtype=dtype)))
d_i_d_v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=d_i_d_u_nodes.shape, dtype=dtype))
d_i_g_u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300, size=100, dtype=dtype)))
d_i_g_v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=d_i_g_u_nodes.shape, dtype=dtype))
d_t_d_u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300, size=100, dtype=dtype)))
d_t_d_v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=d_t_d_u_nodes.shape, dtype=dtype))
g = dgl.heterograph({
('drug', 'interacts', 'drug'): (d_i_d_u_nodes, d_i_d_v_nodes),
('drug', 'interacts', 'gene'): (d_i_g_u_nodes, d_i_g_v_nodes),
('drug', 'treats', 'disease'): (d_t_d_u_nodes, d_t_d_v_nodes)
})
(U, V, EID) = (0, 1, 2)
nd_b_idx = np.random.randint(low=1, high=24, dtype=dtype)
nd_e_idx = np.random.randint(low=25, high=49, dtype=dtype)
did_b_idx = np.random.randint(low=1, high=24, dtype=dtype)
did_e_idx = np.random.randint(low=25, high=49, dtype=dtype)
sampled_amount = np.random.randint(low=1, high=10, dtype=dtype)
drug_i_drug_edges = g.all_edges(form='all', etype=('drug','interacts','drug'))
excluded_d_i_d_edges = drug_i_drug_edges[EID][did_b_idx:did_e_idx]
sampled_drug_node = drug_i_drug_edges[V][nd_b_idx:nd_e_idx]
did_excluded_nodes_U = drug_i_drug_edges[U][did_b_idx:did_e_idx]
did_excluded_nodes_V = drug_i_drug_edges[V][did_b_idx:did_e_idx]
nd_b_idx = np.random.randint(low=1, high=24, dtype=dtype)
nd_e_idx = np.random.randint(low=25, high=49, dtype=dtype)
dig_b_idx = np.random.randint(low=1, high=24, dtype=dtype)
dig_e_idx = np.random.randint(low=25, high=49, dtype=dtype)
drug_i_gene_edges = g.all_edges(form='all', etype=('drug','interacts','gene'))
excluded_d_i_g_edges = drug_i_gene_edges[EID][dig_b_idx:dig_e_idx]
dig_excluded_nodes_U = drug_i_gene_edges[U][dig_b_idx:dig_e_idx]
dig_excluded_nodes_V = drug_i_gene_edges[V][dig_b_idx:dig_e_idx]
sampled_gene_node = drug_i_gene_edges[V][nd_b_idx:nd_e_idx]
nd_b_idx = np.random.randint(low=1, high=24, dtype=dtype)
nd_e_idx = np.random.randint(low=25, high=49, dtype=dtype)
dtd_b_idx = np.random.randint(low=1, high=24, dtype=dtype)
dtd_e_idx = np.random.randint(low=25, high=49, dtype=dtype)
drug_t_dis_edges = g.all_edges(form='all', etype=('drug','treats','disease'))
excluded_d_t_d_edges = drug_t_dis_edges[EID][dtd_b_idx:dtd_e_idx]
dtd_excluded_nodes_U = drug_t_dis_edges[U][dtd_b_idx:dtd_e_idx]
dtd_excluded_nodes_V = drug_t_dis_edges[V][dtd_b_idx:dtd_e_idx]
sampled_disease_node = drug_t_dis_edges[V][nd_b_idx:nd_e_idx]
excluded_edges = {('drug', 'interacts', 'drug'): excluded_d_i_d_edges,
('drug', 'interacts', 'gene'): excluded_d_i_g_edges,
('drug', 'treats', 'disease'): excluded_d_t_d_edges
}
sg = dgl.sampling.sample_neighbors(g, {'drug': sampled_drug_node,
'gene': sampled_gene_node,
'disease': sampled_disease_node},
sampled_amount, exclude_edges=excluded_edges)
assert not np.any(F.asnumpy(sg.has_edges_between(did_excluded_nodes_U,did_excluded_nodes_V,
etype=('drug','interacts','drug'))))
assert not np.any(F.asnumpy(sg.has_edges_between(dig_excluded_nodes_U,dig_excluded_nodes_V,
etype=('drug','interacts','gene'))))
assert not np.any(F.asnumpy(sg.has_edges_between(dtd_excluded_nodes_U,dtd_excluded_nodes_V,
etype=('drug','treats','disease'))))