| 1421 | |
| 1422 | @parametrize_idtype |
| 1423 | def test_add_nodes(idtype): |
| 1424 | # homogeneous Graphs |
| 1425 | g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) |
| 1426 | g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()) |
| 1427 | new_g = dgl.add_nodes(g, 1) |
| 1428 | assert g.num_nodes() == 3 |
| 1429 | assert new_g.num_nodes() == 4 |
| 1430 | assert F.array_equal(new_g.ndata["h"], F.tensor([1, 1, 1, 0], dtype=idtype)) |
| 1431 | |
| 1432 | # zero node graph |
| 1433 | g = dgl.graph(([], []), num_nodes=3, idtype=idtype, device=F.ctx()) |
| 1434 | g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()) |
| 1435 | g = dgl.add_nodes( |
| 1436 | g, 1, data={"h": F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())} |
| 1437 | ) |
| 1438 | assert g.num_nodes() == 4 |
| 1439 | assert F.array_equal(g.ndata["h"], F.tensor([1, 1, 1, 2], dtype=idtype)) |
| 1440 | |
| 1441 | # bipartite graph |
| 1442 | g = dgl.heterograph( |
| 1443 | {("user", "plays", "game"): ([0, 1], [1, 2])}, |
| 1444 | idtype=idtype, |
| 1445 | device=F.ctx(), |
| 1446 | ) |
| 1447 | g = dgl.add_nodes( |
| 1448 | g, |
| 1449 | 2, |
| 1450 | data={"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, |
| 1451 | ntype="user", |
| 1452 | ) |
| 1453 | assert g.num_nodes("user") == 4 |
| 1454 | assert g.num_nodes("game") == 3 |
| 1455 | assert F.array_equal( |
| 1456 | g.nodes["user"].data["h"], F.tensor([0, 0, 2, 2], dtype=idtype) |
| 1457 | ) |
| 1458 | g = dgl.add_nodes(g, 2, ntype="game") |
| 1459 | assert g.num_nodes("user") == 4 |
| 1460 | assert g.num_nodes("game") == 5 |
| 1461 | |
| 1462 | # heterogeneous graph |
| 1463 | g = create_test_heterograph3(idtype) |
| 1464 | g = dgl.add_nodes(g, 1, ntype="user") |
| 1465 | g = dgl.add_nodes( |
| 1466 | g, |
| 1467 | 2, |
| 1468 | data={"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}, |
| 1469 | ntype="game", |
| 1470 | ) |
| 1471 | assert g.num_nodes("user") == 4 |
| 1472 | assert g.num_nodes("game") == 4 |
| 1473 | assert g.num_nodes("developer") == 2 |
| 1474 | assert F.array_equal( |
| 1475 | g.nodes["user"].data["h"], F.tensor([1, 1, 1, 0], dtype=idtype) |
| 1476 | ) |
| 1477 | assert F.array_equal( |
| 1478 | g.nodes["game"].data["h"], F.tensor([2, 2, 2, 2], dtype=idtype) |
| 1479 | ) |
| 1480 | |