Switch from error to warning for inconsistencies in TUDataset Fingerprint. (#203)

* Warning instead of error for inconsistency

* Move force consistency to TODO

* Add tests and fix graph indicator

* Apply suggestions from code review

Co-authored-by: Carlo Lucibello <[email protected]>

Co-authored-by: Carlo Lucibello <[email protected]>

Author: Dsantra92

src/datasets/graphs/tudataset.jl

@@ -63,7 +63,9 @@ function TUDataset(name; dir=nothing)
     source, target = st[:,1], st[:,2]
 
     graph_indicator = readdlm(joinpath(d, "$(name)_graph_indicator.txt"), Int) |> vec      
-    @assert all(sort(unique(graph_indicator)) .== 1:length(unique(graph_indicator)))
+    if !all(sort(unique(graph_indicator)) .== 1:length(unique(graph_indicator)))
+      @warn "Some graph indicators are not present in graph_indicator.txt. Ordering of graph and graph labels may not be consistent. Base.getindex might produce unexpected behavior for unaltered data."
+    end
 
     num_nodes = length(graph_indicator)
     num_edges = length(source)
@@ -91,6 +93,8 @@ function TUDataset(name; dir=nothing)
                         readdlm(joinpath(d, "$(name)_graph_attributes.txt"), ',', Float32)' |> collect :
                         nothing
 
+    # TODO: maybe introduce consistency in graph labels and attributes if possible
+
     # We need this two vectors sorted for efficiency in tudataset_getgraph(full_dataset, i)
     @assert issorted(graph_indicator)
     if !issorted(source)
@@ -115,7 +119,7 @@ function TUDataset(name; dir=nothing)
                         edge_attributes,
                         graph_attributes)
 
-    graphs = [tudataset_getgraph(full_dataset, i) for i in 1:num_graphs]
+    graphs = [tudataset_getgraph(full_dataset, i) for i in sort(unique(graph_indicator))]
     graph_data = (; features = graph_attributes, targets = graph_labels) |> clean_nt
     metadata = Dict{String, Any}("name" => name)
     return TUDataset(name, metadata, graphs, graph_data, num_nodes, num_edges, num_graphs)

test/datasets/graphs_no_ci.jl

@@ -313,6 +313,37 @@ end
     @test size(g.edge_data.features) == (4, g.num_edges)
 end
 
+@testset "TUDataset - Fingerprint" begin
+    @test_warn "" TUDataset("Fingerprint")
+    data = TUDataset("Fingerprint")
+
+    @test data.num_nodes == 15167
+    @test data.num_edges == 24756
+    @test data.num_graphs == 2149
+
+    @test data.num_nodes == sum(g->g.num_nodes, data.graphs)
+    @test data.num_edges == sum(g->g.num_edges, data.graphs)
+    @test data.num_edges == sum(g->length(g.edge_index[1]), data.graphs)
+    @test data.num_edges == sum(g->length(g.edge_index[2]), data.graphs)
+    @test data.num_graphs == length(data) == length(data.graphs)
+
+    i = rand(1:length(data))
+    @test_throws DimensionMismatch data[i]
+    g = data.graphs[i]
+    @test g isa Graph
+    @test all(1 .<= g.edge_index[1] .<= g.num_nodes)
+    @test all(1 .<= g.edge_index[2] .<= g.num_nodes)
+
+    # graph data
+    @test size(data.graph_data.targets) == (2800, )
+
+    # node data
+    @test size(g.node_data.features) == (2, g.num_nodes)
+
+    # edge data
+    @test size(g.edge_data.features) == (2, g.num_edges)
+end
+
 @testset "OrganicMaterialsDB" begin
     data = OrganicMaterialsDB(split=:train)
     @test length(data) == 10000