Merge pull request #77 from JuliaML/cl/tu

[TUDataset] fix in indexing and more tests

Author: CarloLucibello

Project.toml

@@ -1,6 +1,6 @@
 name = "MLDatasets"
 uuid = "eb30cadb-4394-5ae3-aed4-317e484a6458"
-version = "0.5.10"
+version = "0.5.11"
 
 [deps]
 BinDeps = "9e28174c-4ba2-5203-b857-d8d62c4213ee"

src/TUDataset/TUDataset.jl

@@ -154,14 +154,18 @@ function datadir_tudataset(name, dir = nothing)
     return d
 end
 
-function Base.getindex(data::TUDataset, i)
+Base.getindex(data::TUDataset, i::Int) = getindex(data, [i])
+
+function Base.getindex(data::TUDataset, i::AbstractVector{Int})
     node_mask = data.graph_indicator .∈ Ref(i)
-    graph_indicator = data.graph_indicator[node_mask]
 
     nodes = (1:data.num_nodes)[node_mask]
     node_labels = isnothing(data.node_labels) ? nothing : data.node_labels[node_mask]
-    nodemap = Dict(v => i for (i, v) in enumerate(nodes))
+    nodemap = Dict(v => vnew for (vnew, v) in enumerate(nodes))
 
+    graphmap = Dict(i => inew for (inew, i) in enumerate(i))
+    graph_indicator = [graphmap[i] for i in data.graph_indicator[node_mask]]
+    
     edge_mask = data.source .∈ Ref(nodes) 
     source = [nodemap[i] for i in data.source[edge_mask]]
     target = [nodemap[i] for i in data.target[edge_mask]]

test/tst_tudataset.jl

@@ -7,7 +7,7 @@ end
 
     @test data.num_nodes == 43471
     @test data.num_edges == 162088
-    @test data.num_graphs === 1113
+    @test data.num_graphs == 1113
 
     @test length(data.source) == data.num_edges 
     @test length(data.target) == data.num_edges 
@@ -22,14 +22,38 @@ end
 
     @test length(data.graph_indicator) == data.num_nodes
     @test all(sort(unique(data.graph_indicator)) .== 1:data.num_graphs)  
+
+    @testset "indexing" begin
+        d1, d2 = data[5], data[6]
+        d12 = data[5:6]
+    
+        @test d12.num_nodes == d1.num_nodes + d2.num_nodes
+        @test d12.num_edges == d1.num_edges + d2.num_edges
+        @test d12.num_graphs == d1.num_graphs + d2.num_graphs == 2
+    
+        @test length(d12.source) == d12.num_edges 
+        @test length(d12.target) == d12.num_edges 
+        @test length(d1.source) == d1.num_edges 
+        @test length(d1.target) == d1.num_edges 
+        @test length(d2.source) == d2.num_edges 
+        @test length(d2.target) == d2.num_edges 
+        
+        @test length(d12.graph_indicator) == d12.num_nodes
+        @test length(d1.graph_indicator) == d1.num_nodes
+        @test length(d2.graph_indicator) == d2.num_nodes
+    
+        @test all(d1.graph_indicator .== 1)
+        @test all(d2.graph_indicator .== 1)
+        @test sort(unique(d12.graph_indicator)) == [1,2]
+    end
 end
 
 @testset "TUDataset - QM9" begin
     data  = TUDataset("QM9")
 
     @test data.num_nodes == 2333625
     @test data.num_edges == 4823498
-    @test data.num_graphs === 129433
+    @test data.num_graphs == 129433
 
     @test length(data.source) == data.num_edges 
     @test length(data.target) == data.num_edges 
@@ -43,5 +67,40 @@ end
     @test data.graph_labels === nothing
 
     @test length(data.graph_indicator) == data.num_nodes
-    @test all(sort(unique(data.graph_indicator)) .== 1:data.num_graphs)  
+    @test all(sort(unique(data.graph_indicator)) .== 1:data.num_graphs)
+
+
+    @testset "indexing" begin
+        d1, d2 = data[5], data[6]
+        d12 = data[5:6]
+    
+        @test d12.num_nodes == d1.num_nodes + d2.num_nodes
+        @test d12.num_edges == d1.num_edges + d2.num_edges
+        @test d12.num_graphs == d1.num_graphs + d2.num_graphs == 2
+    
+        @test length(d12.source) == d12.num_edges 
+        @test length(d12.target) == d12.num_edges 
+        @test length(d1.source) == d1.num_edges 
+        @test length(d1.target) == d1.num_edges 
+        @test length(d2.source) == d2.num_edges 
+        @test length(d2.target) == d2.num_edges 
+        
+        @test length(d12.graph_indicator) == d12.num_nodes
+        @test length(d1.graph_indicator) == d1.num_nodes
+        @test length(d2.graph_indicator) == d2.num_nodes
+    
+        @test all(d1.graph_indicator .== 1)
+        @test all(d2.graph_indicator .== 1)
+        @test sort(unique(d12.graph_indicator)) == [1,2]
+
+        @test size(d12.node_attributes) == (16, d12.num_nodes)
+        @test size(d12.edge_attributes) == (4, d12.num_edges)
+        @test size(d12.graph_attributes) == (19, d12.num_graphs)
+        @test size(d1.node_attributes) == (16, d1.num_nodes)
+        @test size(d1.edge_attributes) == (4, d1.num_edges)
+        @test size(d1.graph_attributes) == (19, d1.num_graphs)
+        @test size(d2.node_attributes) == (16, d2.num_nodes)
+        @test size(d2.edge_attributes) == (4, d2.num_edges)
+        @test size(d2.graph_attributes) == (19, d2.num_graphs)
+    end
 end