Commit e16667bf authored by VoVAllen's avatar VoVAllen Committed by Minjie Wang
Browse files

[Fix] Fix Transform (#876)

* add transform

* lint

* lint

* fix

* fixmx

* fix

* add test

* fix typo

* fix default num_classes

* change to non-inplace operation

* fix lint

* fix

* fix

* fix lint

* fixlint
parent 9737b270
......@@ -17,6 +17,5 @@ Transform -- Graph Transformation
laplacian_lambda_max
knn_graph
segmented_knn_graph
to_self_loop
add_self_loop
remove_self_loop
onehot_degree
......@@ -10,8 +10,8 @@ from .batched_graph import BatchedDGLGraph, unbatch
__all__ = ['line_graph', 'khop_adj', 'khop_graph', 'reverse', 'to_simple_graph', 'to_bidirected',
'laplacian_lambda_max', 'knn_graph', 'segmented_knn_graph', 'to_self_loop',
'onehot_degree', 'remove_self_loop']
'laplacian_lambda_max', 'knn_graph', 'segmented_knn_graph', 'add_self_loop',
'remove_self_loop']
def pairwise_squared_distance(x):
......@@ -405,36 +405,10 @@ def laplacian_lambda_max(g):
return rst
def onehot_degree(g, max_degree=-1, out_field='d', direction="in"):
"""Inplace add one-hot degree vector as node feature
Parameters
-----------
g: DGLGraph
max_degree: int
Maximum degree for one-hot encoding. If it's -1,
the maximum degree would be inferred from the input graph.
out_field: str
Field name for the node feature
direction: str
Either "in" or "out". Specify whether to use in degrees or out degrees.
Returns
-------
g: DGLGraph
Returns the input graph with added feature
"""
if direction == "in":
degrees = g.in_degrees()
elif direction == "out":
degrees = g.out_degrees()
else:
raise RuntimeError("Invalid Direction")
g.ndata[out_field] = F.one_hot(degrees, max_degree)
return g
def to_self_loop(g):
"""Return a new graph which contains exactly one self loop for each node.
def add_self_loop(g):
"""Return a new graph containing all the edges in the input graph plus self loops
of every nodes.
No duplicate self loop will be added for nodes already having self loops.
Self-loop edges id are not preserved. All self-loop edges would be added at the end.
Examples
......@@ -443,7 +417,7 @@ def to_self_loop(g):
>>> g = DGLGraph()
>>> g.add_nodes(5)
>>> g.add_edges([0, 1, 2], [1, 1, 2])
>>> new_g = dgl.transform.to_self_loop(g) # Nodes 0, 3, 4 don't have self-loop
>>> new_g = dgl.transform.add_self_loop(g) # Nodes 0, 3, 4 don't have self-loop
>>> new_g.edges()
(tensor([0, 0, 1, 2, 3, 4]), tensor([1, 0, 1, 2, 3, 4]))
......
......@@ -176,11 +176,12 @@ def test_laplacian_lambda_max():
assert l_max < 2 + eps
def test_to_self_loop():
def test_add_self_loop():
g = dgl.DGLGraph()
g.add_nodes(5)
g.add_edges([0, 1, 2], [1, 1, 2])
new_g = dgl.transform.to_self_loop(g) # Nodes 0, 3, 4 don't have self-loop
# Nodes 0, 3, 4 don't have self-loop
new_g = dgl.transform.add_self_loop(g)
assert F.allclose(new_g.edges()[0], F.tensor([0, 0, 1, 2, 3, 4]))
assert F.allclose(new_g.edges()[1], F.tensor([1, 0, 1, 2, 3, 4]))
......@@ -194,16 +195,6 @@ def test_remove_self_loop():
assert F.allclose(new_g.edges()[1], F.tensor([1]))
def test_onehot_degree():
g = dgl.DGLGraph()
g.add_nodes(3)
g.add_edges([0, 1, 2], [1, 1, 2])
dgl.transform.onehot_degree(g, out_field="xd")
assert F.allclose(g.ndata['xd'], F.tensor([[1, 0, 0],
[0, 0, 1],
[0, 1, 0]]))
if __name__ == '__main__':
test_line_graph()
......@@ -215,6 +206,5 @@ if __name__ == '__main__':
test_khop_adj()
test_khop_graph()
test_laplacian_lambda_max()
test_onehot_degree()
test_remove_self_loop()
test_to_self_loop()
\ No newline at end of file
test_add_self_loop()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment