start implementing TUDataset

CarloLucibello · CarloLucibello · commit f59d1916c23c · 2021-09-03T14:47:48.000+02:00
diff --git a/docs/src/datasets/SVHN2.md b/docs/src/datasets/SVHN2.md
@@ -128,7 +128,7 @@ SVHN2.classnames
 SVHN2.convert2image
 ```
 
-## References
+### References
 
 - **Authors**: Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, Andrew Y. Ng
 
diff --git a/docs/src/datasets/UD_English.md b/docs/src/datasets/UD_English.md
@@ -11,7 +11,7 @@ devdata = UD_English.devdata()
 testdata = UD_English.devdata()
 ```
 
-## Data Size
+# Data Size
 
 |    | Train x | Train y | Test x | Test y |
 |:--:|:-------:|:-------:|:------:|:------:|
diff --git a/src/BostonHousing/BostonHousing.jl b/src/BostonHousing/BostonHousing.jl
@@ -44,7 +44,7 @@ Harrison and Rubinfeld do not mention any censoring.
 The data file stored in this repo is a copy of the This is a copy of UCI ML housing dataset. 
 https://archive.ics.uci.edu/ml/machine-learning-databases/housing/
 
-## Interface
+# Interface
 
 - [`BostonHousing.features`](@ref)
 - [`BostonHousing.targets`](@ref)
diff --git a/src/CiteSeer/CiteSeer.jl b/src/CiteSeer/CiteSeer.jl
@@ -10,11 +10,11 @@ The dataset is designed for the node classification task.
 The task is to predict the category of certain paper.
 The dataset is retrieved from Ref. [2].
 
-## Interface
+# Interface
 
 - [`CiteSeer.dataset`](@ref)
 
-## References
+# References
 
 [1]: [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://arxiv.org/abs/1707.03815)
 [2]: [Planetoid](https://github.com/kimiyoung/planetoid)
@@ -63,7 +63,7 @@ the reverse of each edge and the graph will be undirected.
 
 See also [`CiteSeer`](@ref).
 
-## Usage Examples
+# Usage Examples
 
 ```julia
 using MLDatasets: CiteSeer
diff --git a/src/Cora/Cora.jl b/src/Cora/Cora.jl
@@ -11,7 +11,7 @@ The dataset is designed for the node classification task.
 The task is to predict the category of certain paper.
 The dataset is retrieved from Ref. [2].
 
-## Statistics 
+# Statistics 
 
 - Nodes: 2708
 - Edges: 10556
@@ -24,11 +24,11 @@ The dataset is retrieved from Ref. [2].
 The split is the one used in the original paper [1] and 
 doesn't consider all nodes.
 
-## Interface
+# Interface
 
 - [`Cora.dataset`](@ref)
 
-## References
+# References
 
 [1]: [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://arxiv.org/abs/1707.03815)
 [2]: [Planetoid](https://github.com/kimiyoung/planetoid
@@ -78,7 +78,7 @@ the reverse of each edge and the graph will be undirected.
 
 See also [`Cora`](@ref).
 
-## Usage Examples
+# Usage Examples
 
 ```juliarepl
 using MLDatasets: Cora
diff --git a/src/FashionMNIST/FashionMNIST.jl b/src/FashionMNIST/FashionMNIST.jl
@@ -12,12 +12,12 @@ examples. Each example is a 28x28 grayscale image, associated
 with a label from 10 classes. It can serve as a drop-in
 replacement for MNIST.
 
-## Interface
+# Interface
 
 - [`FashionMNIST.traintensor`](@ref), [`FashionMNIST.trainlabels`](@ref), [`FashionMNIST.traindata`](@ref)
 - [`FashionMNIST.testtensor`](@ref), [`FashionMNIST.testlabels`](@ref), [`FashionMNIST.testdata`](@ref)
 
-## Utilities
+# Utilities
 
 - [`FashionMNIST.download`](@ref)
 
diff --git a/src/Iris/Iris.jl b/src/Iris/Iris.jl
@@ -13,12 +13,12 @@ The module retrieves the data from the [UCI Machine Learning Repository](https:/
 
 NOTE: no pre-defined train-test split for this dataset, `features` and `labels` return the whole dataset. 
 
-## Interface
+# Interface
 
 - [`Iris.features`](@ref)
 - [`Iris.labels`](@ref)
 
-## Utilities
+# Utilities
 
 - [`Iris.download`](@ref)
 """
diff --git a/src/MLDatasets.jl b/src/MLDatasets.jl
@@ -51,6 +51,7 @@ include("planetoid.jl")
     include("Cora/Cora.jl")
     include("PubMed/PubMed.jl")
     include("CiteSeer/CiteSeer.jl")
+include("TUDataset/TUDataset.jl")
 
 function __init__()
     # initialize optional dependencies
diff --git a/src/MNIST/MNIST.jl b/src/MNIST/MNIST.jl
@@ -12,12 +12,12 @@ used in small-scale machine learning experiments. It contains
 pixel gray-scale image that depicts a handwritten version of 1 of
 the 10 possible digits (0-9).
 
-## Interface
+# Interface
 
 - [`MNIST.traintensor`](@ref), [`MNIST.trainlabels`](@ref), [`MNIST.traindata`](@ref)
 - [`MNIST.testtensor`](@ref), [`MNIST.testlabels`](@ref), [`MNIST.testdata`](@ref)
 
-## Utilities
+# Utilities
 
 - [`MNIST.download`](@ref)
 - [`MNIST.convert2image`](@ref)
diff --git a/src/PubMed/PubMed.jl b/src/PubMed/PubMed.jl
@@ -10,11 +10,11 @@ The dataset is designed for the node classification task.
 The task is to predict the category of certain paper.
 The dataset is retrieved from Ref. [2].
 
-## Interface
+# Interface
 
 - [`PubMed.dataset`](@ref)
 
-## References
+# References
 
 [1]: [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://arxiv.org/abs/1707.03815)
 [2]: [Planetoid](https://github.com/kimiyoung/planetoid)
@@ -63,7 +63,7 @@ the reverse of each edge and the graph will be undirected.
 
 See also [`PubMed`](@ref).
 
-## Usage Examples
+# Usage Examples
 
 ```julia
 using MLDatasets: PubMed
diff --git a/src/SVHN2/SVHN2.jl b/src/SVHN2/SVHN2.jl
@@ -15,13 +15,13 @@ RGB color channel. The dataset is split up into three subsets:
 73257 digits for training, 26032 digits for testing, and 531131
 additional to use as extra training data.
 
-## Interface
+# Interface
 
 - [`SVHN2.traintensor`](@ref), [`SVHN2.trainlabels`](@ref), [`SVHN2.traindata`](@ref)
 - [`SVHN2.testtensor`](@ref), [`SVHN2.testlabels`](@ref), [`SVHN2.testdata`](@ref)
 - [`SVHN2.extratensor`](@ref), [`SVHN2.extralabels`](@ref), [`SVHN2.extradata`](@ref)
 
-## Utilities
+# Utilities
 
 - [`SVHN2.download`](@ref)
 - [`SVHN2.classnames`](@ref)
diff --git a/src/TUDataset/TUDataset.jl b/src/TUDataset/TUDataset.jl
@@ -0,0 +1,97 @@
+export TUDataset
+
+"""
+    TUDataset
+
+A variety of graph kernel benchmark datasets, *.e.g.* "IMDB-BINARY",
+"REDDIT-BINARY" or "PROTEINS", collected from the [TU Dortmund University](https://chrsmrrs.github.io/datasets).
+"""
+module TUDataset
+
+using DataDeps
+using ..MLDatasets: datafile, datadir
+using DelimitedFiles: readdlm
+
+using PyCall
+
+const DEPNAME = "TUDataset"
+# LINK = "https://github.com/shchur/gnn-benchmark/raw/master/data/npz"
+# LINK = "https://github.com/abojchevski/graph2gauss/raw/master/data/"
+const LINK = "https://www.chrsmrrs.com/graphkerneldatasets"
+const DOCS = "https://chrsmrrs.github.io/datasets"
+const DATA = "PROTEINS.zip"
+
+function __init__()
+    register(DataDep(
+        DEPNAME,
+        """
+        Dataset: The $DEPNAME dataset.
+        Website: $DOCS
+        """,
+        "$LINK/$DATA",
+        # "81de017067dc045ebdb8ffd5c0e69a209973ffdb1fe2d5b434e94d3614f3f5c7",  # if checksum omitted, will be generated by DataDeps
+        post_fetch_method = unpack
+    ))
+end
+
+struct TUData
+    source::Vector{Int}
+    target::Vector{Int}
+    graph_indicator::Vector{Int}
+    node_labels::Vector{Int}
+    edge_labels::Union{Nothing, Vector{Int}}
+    graph_labels::Vector{Int}
+    node_attributes
+    edge_attributes
+    graph_attributes
+end
+
+"""
+    dataset(name; dir=nothing)
+
+Retrieve the TUDataset dataset. The output is a named tuple with fields
+
+See [this link](https://chrsmrrs.github.io/datasets/docs/datasets/)
+for a list of the available datasets.
+"""
+function dataset(name; dir=nothing) 
+    d = datadir(DEPNAME, dir)
+    # See here for the file format https://chrsmrrs.github.io/datasets/docs/format/
+    st = readdlm(joinpath(d, name, "$(name)_A.txt"), ',', Int)
+
+    # LOAD OPTIONAL FILES IF EXIST
+    
+    if isfile(joinpath(d, name, "$(name)_edge_labels.txt"))
+        edge_labels = readdlm(joinpath(d, name, "$(name)_edge_labels.txt")) |> vec
+    else
+        edge_labels = nothing
+    end
+    if isfile(joinpath(d, name, "$(name)_node_attributes.txt"))
+        node_attributes = readdlm(joinpath(d, name, "$(name)_node_attributes.txt"), Float32)' |> collect
+    else
+        node_attributes = nothing
+    end
+    if isfile(joinpath(d, name, "$(name)_edge_attributes.txt"))
+        edge_attributes = readdlm(joinpath(d, name, "$(name)_edge_attributes.txt"), Float32)' |> collect
+    else
+        edge_attributes = nothing
+    end
+    if isfile(joinpath(d, name, "$(name)_graph_attributes.txt"))
+        graph_attributes = readdlm(joinpath(d, name, "$(name)_graph_attributes.txt"), Float32)' |> collect
+    else
+        graph_attributes = nothing
+    end
+
+    TUData(st[:,1], st[:,2], 
+            readdlm(joinpath(d, name, "$(name)_graph_indicator.txt"), Int) |> vec, 
+            readdlm(joinpath(d, name, "$(name)_node_labels.txt"), Int) |> vec,
+            edge_labels,            
+            readdlm(joinpath(d, name, "$(name)_graph_labels.txt"), Int) |> vec,
+            node_attributes, 
+            edge_attributes,
+            graph_attributes)
+end
+
+
+end #module 
+
