finish graphs

Author: CarloLucibello

src/abstract_datasets.jl

@@ -46,7 +46,7 @@ _summary(x) = x
 _summary(x::Symbol) = ":$x"
 _summary(x::Union{Dict, AbstractArray, DataFrame}) = summary(x)
 _summary(x::Union{Tuple, NamedTuple}) = map(_summary, x)
-_summary(x::BitVector) = summary(x) * " with $(count(x)) trues"
+_summary(x::BitVector) = "$(count(x))-trues BitVector"
 
 """
     SupervisedDataset <: AbstractDataset

src/datasets/graphs/ogbdataset.jl

@@ -53,93 +53,61 @@ available for node prediction, edge prediction, or graph prediction tasks.
 julia> data = OGBDataset("ogbn-arxiv")
 dataset OGBDataset:
   name        =>    ogbn-arxiv
-  metadata    =>    Dict{String, Any} with 16 entries
+  metadata    =>    Dict{String, Any} with 17 entries
   graphs      =>    1-element Vector{MLDatasets.Graph}
-  targets     =>    nothing
-  split_idx   =>    (train = "90941-element Vector{Int64}", val = "29799-element Vector{Int64}", test = "48603-element Vector{Int64}")
+  graph_data  =>    nothing
 
 julia> data[:]
 Graph:
   num_nodes   =>    169343
   num_edges   =>    1166243
   edge_index  =>    ("1166243-element Vector{Int64}", "1166243-element Vector{Int64}")
-  node_data   =>    (year = "1×169343 Matrix{Int64}", features = "128×169343 Matrix{Float32}", label = "1×169343 Matrix{Int64}")
+  node_data   =>    (val_mask = "29799-trues BitVector", test_mask = "48603-trues BitVector", year = "169343 Vector{Int64}", features = "128×169343 Matrix{Float32}", label = "169343 Vector{Int64}", train_mask = "90941-trues BitVector")
   edge_data   =>    nothing
 
 julia> data.metadata
-Dict{String, Any} with 16 entries:
+Dict{String, Any} with 17 entries:
   "download_name"         => "arxiv"
   "num classes"           => 40
   "num tasks"             => 1
   "binary"                => false
   "url"                   => "http://snap.stanford.edu/ogb/data/nodeproppred/arxiv.zip"
   "additional node files" => "node_year"
   "is hetero"             => false
-  "path"                  => "/home/carlo/.julia/datadeps/OGBDataset/arxiv"
-  "eval metric"           => "acc"
-  "task type"             => "multiclass classification"
-  "add_inverse_edge"      => false
-  "has_node_attr"         => true
-  "additional edge files" => nothing
-  "version"               => 1
-  "has_edge_attr"         => false
-  "split"                 => "time"
+  "task level"            => "node"
+  ⋮                       => ⋮
 ```
 
 ## Edge prediction task
 
 ```julia-repl
 julia> data = OGBDataset("ogbl-collab")
-OGBDataset{Nothing}:
-  name => ogbl-collab
-  path => /home/carlo/.julia/datadeps/OGBDataset/collab
-  metadata => Dict{String, Any} with 13 entries
-  graphs => 1-element Vector{Dict}
-  labels => nothing
-  split => Dict{String, Any} with 3 entries
-
-julia> graph = data[1]  # no labels for this dataset
-Dict{String, Any} with 7 entries:
-  "edge_index"  => [150990 224882; 150990 224882; … ; 221742 135759; 207233 140615]
-  "edge_feat"   => nothing
-  "node_feat"   => Float32[-0.177486 -0.237488 … 0.004236 -0.035025; -0.10298 0.022193 … 0.031942 -0.118059; … ; 0.003879 0.062124 … 0.05208 -0.176961; -0.276317 -0.081464 … -0.201557 -0.258715]
-  "num_nodes"   => 235868
-  "edge_year"   => [2004 2002 … 2006 1984; 2004 2002 … 2006 1984]
-  "edge_weight" => [2 1 … 1 1; 2 1 … 1 1]
-  "num_edges"   => 2358104
+dataset OGBDataset:
+  name        =>    ogbl-collab
+  metadata    =>    Dict{String, Any} with 15 entries
+  graphs      =>    1-element Vector{MLDatasets.Graph}
+  graph_data  =>    nothing
+
+julia> data[:]
+Graph:
+  num_nodes   =>    235868
+  num_edges   =>    2358104
+  edge_index  =>    ("2358104-element Vector{Int64}", "2358104-element Vector{Int64}")
+  node_data   =>    (features = "128×235868 Matrix{Float32}",)
+  edge_data   =>    (year = "2×1179052 Matrix{Int64}", weight = "2×1179052 Matrix{Int64}")
 ```
 
 ## Graph prediction task
 
 ```julia-repl
 julia> data = OGBDataset("ogbg-molhiv")
-OGBDataset{Matrix{Int64}}:
-  name => ogbg-molhiv
-  path => /home/carlo/.julia/datadeps/OGBDataset/molhiv
-  metadata => Dict{String, Any} with 15 entries
-  graphs => 41127-element Vector{Dict}
-  labels => 1×41127 Matrix{Int64}
-  split => Dict{String, Any} with 3 entries
-
-julia> length(data)
-41127
-
-julia> graph, labels = data[10]
-(Dict{String, Any}("edge_index" => [-202 -201; -201 -200; … ; -198 -184; -201 -202], "node_feat" => Float32[7.0 6.0 … 7.0 7.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0], "edge_feat" => Float32[0.0 0.0 … 0.0 1.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 1.0], "num_nodes" => 20, "num_edges" => 42), [0])
-
-julia> graph, labels = data[10];
-
-julia> graph
-Dict{String, Any} with 5 entries:
-  "edge_index" => [1 2; 2 3; … ; 5 19; 2 1]
-  "edge_feat"  => Float32[0.0 0.0 … 0.0 1.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 1.0]
-  "node_feat"  => Float32[7.0 6.0 … 7.0 7.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0]
-  "num_nodes"  => 20
-  "num_edges"  => 42
-
-julia> labels
-1-element Vector{Int64}:
- 0
+dataset OGBDataset:
+  name        =>    ogbg-molhiv
+  metadata    =>    Dict{String, Any} with 17 entries
+  graphs      =>    41127-element Vector{MLDatasets.Graph}
+  graph_data  =>    (labels = "41127-element Vector{Int64}", train_mask = "32901-trues BitVector", val_mask = "4113-trues BitVector", test_mask = "4113-trues BitVector")
+
+julia> data[1]
 ```
 """
 struct OGBDataset{GD} <: AbstractDataset
@@ -309,10 +277,13 @@ function read_ogb_graph(path, metadata)
 
     splits = readdir(joinpath(path, "split"))
     @assert length(splits) == 1 # TODO check if datasets with multiple splits existin in OGB
+    
     # TODO sometimes splits are given in .pt format
+    # Use read_pytorch in src/io.jl to load them.
     split_idx = (train = read_ogb_file(joinpath(path, "split", splits[1], "train.csv"), Int; tovec=true),
                  val = read_ogb_file(joinpath(path, "split", splits[1], "valid.csv"), Int; tovec=true),
                 test = read_ogb_file(joinpath(path, "split", splits[1], "test.csv"), Int; tovec=true))
+    
     if split_idx.train !== nothing 
         split_idx.train .+= 1
     end
@@ -353,7 +324,11 @@ function read_ogb_graph(path, metadata)
         end
     end
     if metadata["task level"] == "graph"
-        graph_data = (; labels, split_idx)
+        train_mask = split_idx.train !== nothing ? indexes2mask(split_idx.train, num_graphs) : nothing
+        val_mask = split_idx.val !== nothing ? indexes2mask(split_idx.val, num_graphs) : nothing 
+        test_mask = split_idx.test !== nothing ? indexes2mask(split_idx.test, num_graphs) : nothing
+
+        graph_data = clean_nt((; labels=maybesqueeze(labels), train_mask, val_mask, test_mask))
     end
     return graphs, graph_data
 end
@@ -377,15 +352,15 @@ end
 function ogbdict2graph(d::Dict)
     edge_index = d["edge_index"][:,1], d["edge_index"][:,2] 
     num_nodes = d["num_nodes"]
-    node_data = Dict(Symbol(k[6:end]) => v for (k,v) in d if startswith(k, "node_") && v !== nothing)
-    edge_data = Dict(Symbol(k[6:end]) => v for (k,v) in d if startswith(k, "edge_") && k!="edge_index" && v !== nothing)
+    node_data = Dict(Symbol(k[6:end]) => maybesqueeze(v) for (k,v) in d if startswith(k, "node_") && v !== nothing)
+    edge_data = Dict(Symbol(k[6:end]) => maybesqueeze(v) for (k,v) in d if startswith(k, "edge_") && k!="edge_index" && v !== nothing)
     node_data = isempty(node_data) ? nothing : (; node_data...)
     edge_data = isempty(edge_data) ? nothing : (; edge_data...)
     return Graph(; num_nodes, edge_index, node_data, edge_data)
 end
 
 Base.length(data::OGBDataset) = length(data.graphs)
 Base.getindex(data::OGBDataset{Nothing}, ::Colon) = length(data.graphs) == 1 ? data.graphs[1] : data.graphs
-Base.getindex(data::OGBDataset, ::Colon) = (; data.graphs, targets=data.graph_data.labels)
+Base.getindex(data::OGBDataset, ::Colon) = (; data.graphs, data.graph_data.labels)
 Base.getindex(data::OGBDataset{Nothing}, i) = getobs(data.graphs, i)
-Base.getindex(data::OGBDataset, i) = getobs((; data.graphs, targets=data.graph_data.labels), i) 
+Base.getindex(data::OGBDataset, i) = getobs((; data.graphs, data.graph_data.labels), i) 

src/datasets/graphs/reddit.jl

@@ -72,7 +72,7 @@ function Reddit(; full=true, dir=nothing)
     nodes = graph["nodes"]
     num_edges = directed ? length(links) : length(links) * 2
     num_nodes = length(nodes)
-    num_graphs = length(graph["graph"]) # should be zero
+    @assert length(graph["graph"]) == 0 # should be zero
 
     # edges
     s = get.(links, "source", nothing) .+ 1
@@ -101,22 +101,11 @@ function Reddit(; full=true, dir=nothing)
     @assert sum(val_mask .& test_mask) == 0
     train_mask = nor.(test_mask, val_mask)
 
-    train_idx = node_idx[train_mask]
-    test_idx = node_idx[test_mask]
-    val_idx = node_idx[val_mask]
-
-    split = Dict(
-        "train" => train_idx,
-        "test" => test_idx,
-        "val"  => val_idx
-    )
-
     metadata = Dict{String, Any}("directed" => directed, "multigraph" => multigraph, 
-                "num_graphs" => num_graphs, "num_edges" => num_edges, "num_nodes" => num_nodes, 
-                "split" => split)
+                "num_edges" => num_edges, "num_nodes" => num_nodes)
     g = Graph(; num_nodes, 
         edge_index=(s, t), 
-        node_data= (; labels, features)
+        node_data= (; labels, features, train_mask, val_mask, test_mask)
     )
     return Reddit(metadata, [g])
 end

src/datasets/graphs/tudataset.jl

@@ -21,7 +21,7 @@ end
 
 A variety of graph benchmark datasets, *.e.g.* "QM9", "IMDB-BINARY",
 "REDDIT-BINARY" or "PROTEINS", collected from the [TU Dortmund University](https://chrsmrrs.github.io/datasets/).
-Retrieve from TUDataset collection the dataset `name`, where `name`
+Retrieve from the TUDataset collection the dataset `name`, where `name`
 is any of the datasets available [here](https://chrsmrrs.github.io/datasets/docs/datasets/). 
 
 A `TUDataset` object can be indexed to retrieve a specific graph or a subset of graphs.
@@ -31,16 +31,19 @@ description of the format.
 
 # Usage Example
 
-```julia
-using MLDatasets: TUDataset
-
-data = TUDataset("PROTEINS")
-
-# Access first graph
-d1 = data[1] 
-
-# Node features
-X = d1.node_attributes # (nfeatures x nnodes) matrix
+```julia-repl
+julia> data = TUDataset("PROTEINS")
+dataset TUDataset:
+  name        =>    PROTEINS
+  metadata    =>    Dict{String, Any} with 1 entry
+  graphs      =>    1113-element Vector{MLDatasets.Graph}
+  graph_data  =>    (targets = "1113-element Vector{Int64}",)
+  num_nodes   =>    43471
+  num_edges   =>    162088
+  num_graphs  =>    1113
+
+julia> data[1]
+(graphs = Graph(42, 162), targets = 1)
 ```
 """
 struct TUDataset <: AbstractDataset

src/io.jl

@@ -20,3 +20,7 @@ end
 function read_npz(path)
     return NPZ.npzread(path)
 end
+
+function read_pytorch(path)
+    return Pickle.Torch.THload(path)
+end

src/utils.jl

@@ -60,6 +60,8 @@ function mask2indexes(mask::BitVector)
     return (1:n)[mask]
 end
 
+maybesqueeze(x) = x
+maybesqueeze(x::AbstractMatrix) = size(x, 1) == 1 ? vec(x) : x
 
 """
     convert2image(d, i)

test/datasets/graphs_no_ci.jl

@@ -7,14 +7,13 @@
     @test g.num_edges == 114615892
     @test size(g.node_data.features) == (602, g.num_nodes)
     @test size(g.node_data.labels) == (g.num_nodes,)
-    @test size(data.metadata["split"]["train"]) == (153431,)
-    @test size(data.metadata["split"]["val"]) == (23831,)
-    @test size(data.metadata["split"]["test"]) == (55703,)
+    @test count(g.node_data.train_mask) == 153431
+    @test count(g.node_data.val_mask) == 23831
+    @test count(g.node_data.test_mask) == 55703
     s, t = g.edge_index
     @test length(s) == length(t) == g.num_edges
     @test minimum(s) == minimum(t) == 1
     @test maximum(s) == maximum(t) == g.num_nodes
-    @test sum(length.(values(data.metadata["split"]))) == g.num_nodes
 end
 
 @testset "Reddit_subset" begin
@@ -25,14 +24,13 @@ end
     @test g.num_edges == 23213838
     @test size(g.node_data.features) == (602, g.num_nodes)
     @test size(g.node_data.labels) == (g.num_nodes,)
-    @test size(data.metadata["split"]["train"]) == (152410,)
-    @test size(data.metadata["split"]["val"]) == (23699,)
-    @test size(data.metadata["split"]["test"]) == (55334,)
+    @test count(g.node_data.train_mask) == 152410
+    @test count(g.node_data.val_mask) == 23699
+    @test count(g.node_data.test_mask) == 55334
     s, t = g.edge_index
     @test length(s) == length(t) == g.num_edges
     @test minimum(s) == minimum(t) == 1
     @test maximum(s) == maximum(t) == g.num_nodes
-    @test sum(length.(values(data.metadata["split"]))) == g.num_nodes
 end
 
 
@@ -107,3 +105,9 @@ end
 
     @test sum(count.([g.node_data.train_mask, g.node_data.test_mask, g.node_data.val_mask])) == g.num_nodes
 end
+
+@testset "OGBDataset - ogbg-molhiv" begin
+    d = OGBDataset("ogbg-molhiv")
+    
+    @test sum(count.([d.graph_data.train_mask, d.graph_data.test_mask, d.graph_data.val_mask])) == length(d)
+end