chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,120 @@
|
||||
import itertools
|
||||
import unittest
|
||||
from collections import Counter
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import dgl.function as fn
|
||||
import networkx as nx
|
||||
import numpy as np
|
||||
import pytest
|
||||
import scipy.sparse as ssp
|
||||
from dgl import DGLError
|
||||
from utils import parametrize_idtype
|
||||
|
||||
|
||||
def create_test_heterograph(num_nodes, num_adj, idtype):
|
||||
if isinstance(num_adj, int):
|
||||
num_adj = [num_adj, num_adj + 1]
|
||||
num_adj_list = list(
|
||||
np.random.choice(np.arange(num_adj[0], num_adj[1]), num_nodes)
|
||||
)
|
||||
src = np.concatenate([[i] * num_adj_list[i] for i in range(num_nodes)])
|
||||
dst = [
|
||||
np.random.choice(num_nodes, nadj, replace=False)
|
||||
for nadj in num_adj_list
|
||||
]
|
||||
dst = np.concatenate(dst)
|
||||
return dgl.graph((src, dst), idtype=idtype)
|
||||
|
||||
|
||||
def check_sort(spm, tag_arr=None, tag_pos=None):
|
||||
if tag_arr is None:
|
||||
tag_arr = np.arange(spm.shape[0])
|
||||
else:
|
||||
tag_arr = F.asnumpy(tag_arr)
|
||||
if tag_pos is not None:
|
||||
tag_pos = F.asnumpy(tag_pos)
|
||||
for i in range(spm.shape[0]):
|
||||
row = spm.getrow(i)
|
||||
dst = row.nonzero()[1]
|
||||
if tag_pos is not None:
|
||||
tag_pos_row = tag_pos[i]
|
||||
tag_pos_ptr = tag_arr[dst[0]] if len(dst) > 0 else 0
|
||||
for j in range(len(dst) - 1):
|
||||
if tag_pos is not None and tag_arr[dst[j]] != tag_pos_ptr:
|
||||
# `tag_pos_ptr` is the expected tag value. Here we check whether the
|
||||
# tag value is equal to `tag_pos_ptr`
|
||||
return False
|
||||
if tag_arr[dst[j]] > tag_arr[dst[j + 1]]:
|
||||
# The tag should be in ascending order after sorting
|
||||
return False
|
||||
if tag_pos is not None and tag_arr[dst[j]] < tag_arr[dst[j + 1]]:
|
||||
if j + 1 != int(tag_pos_row[tag_pos_ptr + 1]):
|
||||
# The boundary of tag should be consistent with `tag_pos`
|
||||
return False
|
||||
tag_pos_ptr = tag_arr[dst[j + 1]]
|
||||
return True
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "gpu", reason="GPU sorting by tag not implemented"
|
||||
)
|
||||
@parametrize_idtype
|
||||
def test_sort_with_tag(idtype):
|
||||
num_nodes, num_adj, num_tags = 200, [20, 50], 5
|
||||
g = create_test_heterograph(num_nodes, num_adj, idtype=idtype)
|
||||
tag = F.tensor(np.random.choice(num_tags, g.num_nodes()))
|
||||
src, dst = g.edges()
|
||||
edge_tag_dst = F.gather_row(tag, F.tensor(dst))
|
||||
edge_tag_src = F.gather_row(tag, F.tensor(src))
|
||||
|
||||
for tag_type in ["node", "edge"]:
|
||||
new_g = dgl.sort_csr_by_tag(
|
||||
g, tag if tag_type == "node" else edge_tag_dst, tag_type=tag_type
|
||||
)
|
||||
old_csr = g.adj_external(scipy_fmt="csr")
|
||||
new_csr = new_g.adj_external(scipy_fmt="csr")
|
||||
assert check_sort(new_csr, tag, new_g.dstdata["_TAG_OFFSET"])
|
||||
assert not check_sort(
|
||||
old_csr, tag
|
||||
) # Check the original csr is not modified.
|
||||
|
||||
for tag_type in ["node", "edge"]:
|
||||
new_g = dgl.sort_csc_by_tag(
|
||||
g, tag if tag_type == "node" else edge_tag_src, tag_type=tag_type
|
||||
)
|
||||
old_csc = g.adj_external(transpose=True, scipy_fmt="csr")
|
||||
new_csc = new_g.adj_external(transpose=True, scipy_fmt="csr")
|
||||
assert check_sort(new_csc, tag, new_g.srcdata["_TAG_OFFSET"])
|
||||
assert not check_sort(old_csc, tag)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "gpu", reason="GPU sorting by tag not implemented"
|
||||
)
|
||||
@parametrize_idtype
|
||||
def test_sort_with_tag_bipartite(idtype):
|
||||
num_nodes, num_adj, num_tags = 200, [20, 50], 5
|
||||
g = create_test_heterograph(num_nodes, num_adj, idtype=idtype)
|
||||
g = dgl.heterograph({("_U", "_E", "_V"): g.edges()})
|
||||
utag = F.tensor(np.random.choice(num_tags, g.num_nodes("_U")))
|
||||
vtag = F.tensor(np.random.choice(num_tags, g.num_nodes("_V")))
|
||||
|
||||
new_g = dgl.sort_csr_by_tag(g, vtag)
|
||||
old_csr = g.adj_external(scipy_fmt="csr")
|
||||
new_csr = new_g.adj_external(scipy_fmt="csr")
|
||||
assert check_sort(new_csr, vtag, new_g.nodes["_U"].data["_TAG_OFFSET"])
|
||||
assert not check_sort(old_csr, vtag)
|
||||
|
||||
new_g = dgl.sort_csc_by_tag(g, utag)
|
||||
old_csc = g.adj_external(transpose=True, scipy_fmt="csr")
|
||||
new_csc = new_g.adj_external(transpose=True, scipy_fmt="csr")
|
||||
assert check_sort(new_csc, utag, new_g.nodes["_V"].data["_TAG_OFFSET"])
|
||||
assert not check_sort(old_csc, utag)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_sort_with_tag(F.int32)
|
||||
test_sort_with_tag_bipartite(F.int32)
|
||||
@@ -0,0 +1,192 @@
|
||||
##
|
||||
# Copyright 2019-2021 Contributors
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import dgl.partition
|
||||
from utils import parametrize_idtype
|
||||
|
||||
|
||||
@parametrize_idtype
|
||||
def test_to_block(idtype):
|
||||
def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
|
||||
if dst_nodes is not None:
|
||||
assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes)
|
||||
n_dst_nodes = bg.num_nodes("DST/" + ntype)
|
||||
if include_dst_in_src:
|
||||
assert F.array_equal(
|
||||
bg.srcnodes[ntype].data[dgl.NID][:n_dst_nodes],
|
||||
bg.dstnodes[ntype].data[dgl.NID],
|
||||
)
|
||||
|
||||
g = g[etype]
|
||||
bg = bg[etype]
|
||||
induced_src = bg.srcdata[dgl.NID]
|
||||
induced_dst = bg.dstdata[dgl.NID]
|
||||
induced_eid = bg.edata[dgl.EID]
|
||||
|
||||
bg_src, bg_dst = bg.all_edges(order="eid")
|
||||
src_ans, dst_ans = g.all_edges(order="eid")
|
||||
|
||||
induced_src_bg = F.gather_row(induced_src, bg_src)
|
||||
induced_dst_bg = F.gather_row(induced_dst, bg_dst)
|
||||
induced_src_ans = F.gather_row(src_ans, induced_eid)
|
||||
induced_dst_ans = F.gather_row(dst_ans, induced_eid)
|
||||
|
||||
assert F.array_equal(induced_src_bg, induced_src_ans)
|
||||
assert F.array_equal(induced_dst_bg, induced_dst_ans)
|
||||
|
||||
def checkall(g, bg, dst_nodes, include_dst_in_src=True):
|
||||
for etype in g.etypes:
|
||||
ntype = g.to_canonical_etype(etype)[2]
|
||||
if dst_nodes is not None and ntype in dst_nodes:
|
||||
check(g, bg, ntype, etype, dst_nodes[ntype], include_dst_in_src)
|
||||
else:
|
||||
check(g, bg, ntype, etype, None, include_dst_in_src)
|
||||
|
||||
# homogeneous graph
|
||||
g = dgl.graph(
|
||||
(F.tensor([1, 2], dtype=idtype), F.tensor([2, 3], dtype=idtype))
|
||||
)
|
||||
dst_nodes = F.tensor([3, 2], dtype=idtype)
|
||||
bg = dgl.to_block(g, dst_nodes=dst_nodes)
|
||||
check(g, bg, "_N", "_E", dst_nodes)
|
||||
|
||||
src_nodes = bg.srcnodes["_N"].data[dgl.NID]
|
||||
bg = dgl.to_block(g, dst_nodes=dst_nodes, src_nodes=src_nodes)
|
||||
check(g, bg, "_N", "_E", dst_nodes)
|
||||
|
||||
# heterogeneous graph
|
||||
g = dgl.heterograph(
|
||||
{
|
||||
("A", "AA", "A"): ([0, 2, 1, 3], [1, 3, 2, 4]),
|
||||
("A", "AB", "B"): ([0, 1, 3, 1], [1, 3, 5, 6]),
|
||||
("B", "BA", "A"): ([2, 3], [3, 2]),
|
||||
},
|
||||
idtype=idtype,
|
||||
device=F.ctx(),
|
||||
)
|
||||
g.nodes["A"].data["x"] = F.randn((5, 10))
|
||||
g.nodes["B"].data["x"] = F.randn((7, 5))
|
||||
g.edges["AA"].data["x"] = F.randn((4, 3))
|
||||
g.edges["AB"].data["x"] = F.randn((4, 3))
|
||||
g.edges["BA"].data["x"] = F.randn((2, 3))
|
||||
g_a = g["AA"]
|
||||
|
||||
def check_features(g, bg):
|
||||
for ntype in bg.srctypes:
|
||||
for key in g.nodes[ntype].data:
|
||||
assert F.array_equal(
|
||||
bg.srcnodes[ntype].data[key],
|
||||
F.gather_row(
|
||||
g.nodes[ntype].data[key],
|
||||
bg.srcnodes[ntype].data[dgl.NID],
|
||||
),
|
||||
)
|
||||
for ntype in bg.dsttypes:
|
||||
for key in g.nodes[ntype].data:
|
||||
assert F.array_equal(
|
||||
bg.dstnodes[ntype].data[key],
|
||||
F.gather_row(
|
||||
g.nodes[ntype].data[key],
|
||||
bg.dstnodes[ntype].data[dgl.NID],
|
||||
),
|
||||
)
|
||||
for etype in bg.canonical_etypes:
|
||||
for key in g.edges[etype].data:
|
||||
assert F.array_equal(
|
||||
bg.edges[etype].data[key],
|
||||
F.gather_row(
|
||||
g.edges[etype].data[key], bg.edges[etype].data[dgl.EID]
|
||||
),
|
||||
)
|
||||
|
||||
bg = dgl.to_block(g_a)
|
||||
check(g_a, bg, "A", "AA", None)
|
||||
check_features(g_a, bg)
|
||||
assert bg.number_of_src_nodes() == 5
|
||||
assert bg.number_of_dst_nodes() == 4
|
||||
|
||||
bg = dgl.to_block(g_a, include_dst_in_src=False)
|
||||
check(g_a, bg, "A", "AA", None, False)
|
||||
check_features(g_a, bg)
|
||||
assert bg.number_of_src_nodes() == 4
|
||||
assert bg.number_of_dst_nodes() == 4
|
||||
|
||||
dst_nodes = F.tensor([4, 3, 2, 1], dtype=idtype)
|
||||
bg = dgl.to_block(g_a, dst_nodes)
|
||||
check(g_a, bg, "A", "AA", dst_nodes)
|
||||
check_features(g_a, bg)
|
||||
|
||||
g_ab = g["AB"]
|
||||
|
||||
bg = dgl.to_block(g_ab)
|
||||
assert bg.idtype == idtype
|
||||
assert bg.num_nodes("SRC/B") == 4
|
||||
assert F.array_equal(
|
||||
bg.srcnodes["B"].data[dgl.NID], bg.dstnodes["B"].data[dgl.NID]
|
||||
)
|
||||
assert bg.num_nodes("DST/A") == 0
|
||||
checkall(g_ab, bg, None)
|
||||
check_features(g_ab, bg)
|
||||
|
||||
dst_nodes = {"B": F.tensor([5, 6, 3, 1], dtype=idtype)}
|
||||
bg = dgl.to_block(g, dst_nodes)
|
||||
assert bg.num_nodes("SRC/B") == 4
|
||||
assert F.array_equal(
|
||||
bg.srcnodes["B"].data[dgl.NID], bg.dstnodes["B"].data[dgl.NID]
|
||||
)
|
||||
assert bg.num_nodes("DST/A") == 0
|
||||
checkall(g, bg, dst_nodes)
|
||||
check_features(g, bg)
|
||||
|
||||
dst_nodes = {
|
||||
"A": F.tensor([4, 3, 2, 1], dtype=idtype),
|
||||
"B": F.tensor([3, 5, 6, 1], dtype=idtype),
|
||||
}
|
||||
bg = dgl.to_block(g, dst_nodes=dst_nodes)
|
||||
checkall(g, bg, dst_nodes)
|
||||
check_features(g, bg)
|
||||
|
||||
# test specifying lhs_nodes with include_dst_in_src
|
||||
src_nodes = {}
|
||||
for ntype in dst_nodes.keys():
|
||||
# use the previous run to get the list of source nodes
|
||||
src_nodes[ntype] = bg.srcnodes[ntype].data[dgl.NID]
|
||||
bg = dgl.to_block(g, dst_nodes=dst_nodes, src_nodes=src_nodes)
|
||||
checkall(g, bg, dst_nodes)
|
||||
check_features(g, bg)
|
||||
|
||||
# test without include_dst_in_src
|
||||
dst_nodes = {
|
||||
"A": F.tensor([4, 3, 2, 1], dtype=idtype),
|
||||
"B": F.tensor([3, 5, 6, 1], dtype=idtype),
|
||||
}
|
||||
bg = dgl.to_block(g, dst_nodes=dst_nodes, include_dst_in_src=False)
|
||||
checkall(g, bg, dst_nodes, False)
|
||||
check_features(g, bg)
|
||||
|
||||
# test specifying lhs_nodes without include_dst_in_src
|
||||
src_nodes = {}
|
||||
for ntype in dst_nodes.keys():
|
||||
# use the previous run to get the list of source nodes
|
||||
src_nodes[ntype] = bg.srcnodes[ntype].data[dgl.NID]
|
||||
bg = dgl.to_block(
|
||||
g, dst_nodes=dst_nodes, include_dst_in_src=False, src_nodes=src_nodes
|
||||
)
|
||||
checkall(g, bg, dst_nodes, False)
|
||||
check_features(g, bg)
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user