MCPcopy
hub / github.com/dmlc/dgl / test_remove_nodes

Function test_remove_nodes

tests/python/common/transforms/test_transform.py:1727–1994  ·  view source on GitHub ↗
(idtype)

Source from the content-addressed store, hash-verified

1725
1726@parametrize_idtype
1727def test_remove_nodes(idtype):
1728 # homogeneous Graphs
1729 g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
1730 n = 0
1731 g = dgl.remove_nodes(g, n)
1732 assert g.num_nodes() == 2
1733 assert g.num_edges() == 1
1734 u, v = g.edges(form="uv", order="eid")
1735 assert F.array_equal(u, F.tensor([0], dtype=idtype))
1736 assert F.array_equal(v, F.tensor([1], dtype=idtype))
1737 g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
1738 n = [1]
1739 g = dgl.remove_nodes(g, n)
1740 assert g.num_nodes() == 2
1741 assert g.num_edges() == 0
1742 g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
1743 n = F.tensor([2], dtype=idtype)
1744 g = dgl.remove_nodes(g, n)
1745 assert g.num_nodes() == 2
1746 assert g.num_edges() == 1
1747 u, v = g.edges(form="uv", order="eid")
1748 assert F.array_equal(u, F.tensor([0], dtype=idtype))
1749 assert F.array_equal(v, F.tensor([1], dtype=idtype))
1750
1751 # invalid nid
1752 assert_fail = False
1753 try:
1754 g.remove_nodes(3)
1755 except:
1756 assert_fail = True
1757 assert assert_fail
1758
1759 # has node and edge data
1760 g = dgl.graph(([0, 0, 2], [0, 1, 2]), idtype=idtype, device=F.ctx())
1761 g.ndata["hv"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
1762 g.edata["he"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
1763 g = dgl.remove_nodes(g, F.tensor([0], dtype=idtype))
1764 assert g.num_nodes() == 2
1765 assert g.num_edges() == 1
1766 u, v = g.edges(form="uv", order="eid")
1767 assert F.array_equal(u, F.tensor([1], dtype=idtype))
1768 assert F.array_equal(v, F.tensor([1], dtype=idtype))
1769 assert F.array_equal(g.ndata["hv"], F.tensor([2, 3], dtype=idtype))
1770 assert F.array_equal(g.edata["he"], F.tensor([3], dtype=idtype))
1771
1772 # node id larger than current max node id
1773 g = dgl.heterograph(
1774 {("user", "plays", "game"): ([0, 1], [1, 2])},
1775 idtype=idtype,
1776 device=F.ctx(),
1777 )
1778 n = 0
1779 g = dgl.remove_nodes(g, n, ntype="user")
1780 assert g.num_nodes("user") == 1
1781 assert g.num_nodes("game") == 3
1782 assert g.num_edges() == 1
1783 u, v = g.edges(form="uv", order="eid")
1784 assert F.array_equal(u, F.tensor([0], dtype=idtype))

Callers

nothing calls this directly

Calls 10

remove_nodesMethod · 0.80
batch_num_nodesMethod · 0.80
batch_num_edgesMethod · 0.80
create_test_heterograph3Function · 0.70
graphMethod · 0.45
ctxMethod · 0.45
num_nodesMethod · 0.45
num_edgesMethod · 0.45
edgesMethod · 0.45
copy_toMethod · 0.45

Tested by

no test coverage detected