Add METR-LA dataset (#212)

* Add `adjmatrix2edgeindex`

* Add some functions for METR-LA

* Fix

* Improve with weights

* Add function

* Return weights

* Add `generate_task`

* Add pkgs and export `metrla`

* Update

* Add `metrla` tests

* Add `METRLA` in docs

* Change import

* Normal import Statistics

* Change name `generate_task`

* Fix

* Add spaces

* Change type `adj`


* Fix UD_English

Author: aurorarossi

Project.toml

@@ -24,6 +24,7 @@ Pickle = "fbb45041-c46e-462f-888f-7c521cafbc2c"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]

docs/src/datasets/graphs.md

@@ -30,4 +30,5 @@ PolBlogs
 PubMed
 Reddit
 TUDataset
+METRLA
 ```

src/MLDatasets.jl

@@ -12,6 +12,7 @@ using DelimitedFiles: readdlm
 using FileIO
 import CSV
 using LazyModules: @lazy
+using Statistics
 
 include("require.jl") # export @require
 
@@ -25,8 +26,9 @@ include("require.jl") # export @require
 @require import DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 @require import ImageShow = "4e3cecfd-b093-5904-9786-8bbb286a6a31"
 @require import Chemfiles = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
+@require import NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 
-# @lazy import NPZ # lazy imported by FileIO
+# lazy imported by FileIO
 @lazy import Pickle = "fbb45041-c46e-462f-888f-7c521cafbc2c"
 @lazy import MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 @lazy import HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
@@ -128,6 +130,8 @@ include("datasets/graphs/reddit.jl")
 export Reddit
 include("datasets/graphs/tudataset.jl")
 export TUDataset
+include("datasets/graphs/metrla.jl")
+export METRLA
 
 # Meshes
 
@@ -147,6 +151,7 @@ function __init__()
     __init__pubmed()
     __init__reddit()
     __init__tudataset()
+    __init__metrla()
 
     # misc
     __init__iris()

src/datasets/graphs/metrla.jl

@@ -0,0 +1,85 @@
+function __init__metrla()
+    DEPNAME = "METR-LA"
+    LINK = "https://graphmining.ai/temporal_datasets/"
+    register(ManualDataDep(DEPNAME,
+                           """
+                           Dataset: $DEPNAME
+                           Website : $LINK
+                           """))
+end
+
+"""
+    METRLA(; num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir=nothing)
+
+The METR-LA dataset from the [Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting](https://arxiv.org/abs/1707.01926) paper.
+
+`METRLA` is a graph with 207 nodes representing traffic sensors in Los Angeles. 
+
+The edge weights `w` are contained as a feature array in `edge_data` and represent the distance between the sensors. 
+
+The node features are the traffic speed and the time of the measurements collected by the sensors, divided into `num_timesteps_in` time steps. 
+
+The target values are the traffic speed and the time of the measurements collected by the sensors, divided into `num_timesteps_out` time steps.
+"""
+struct METRLA <: AbstractDataset
+    graphs::Vector{Graph}
+end
+
+function METRLA(;num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir = nothing)
+    create_default_dir("METR-LA")
+    d = metrla_datadir(dir)
+    adj_matrix, node_values = read_metrla(d)
+
+    node_values = permutedims(node_values,(1,3,2))
+    node_values = (node_values .- Statistics.mean(node_values, dims=(3,1))) ./ Statistics.std(node_values, dims=(3,1)) #Z-score normalization
+
+    s, t, w = adjmatrix2edgeindex(adj_matrix; weighted = true)
+
+    x, y = metrla_generate_task(node_values, num_timesteps_in, num_timesteps_out)
+
+    g = Graph(; num_nodes = 207,
+              edge_index = (s, t),
+              edge_data = w,
+            node_data = (features = x, targets = y)
+
+
+    )
+    return METRLA([g])
+end
+
+function metrla_datadir(dir = nothing)
+    dir = isnothing(dir) ? datadep"METR-LA" : dir
+    dname = "METR-LA"
+    LINK = "https://graphmining.ai/temporal_datasets/$dname.zip"
+    if length(readdir((dir))) == 0
+        DataDeps.fetch_default(LINK, dir)
+        currdir = pwd()
+        cd(dir) # Needed since `unpack` extracts in working dir
+        DataDeps.unpack(joinpath(dir, "$dname.zip"))
+        # conditions when unzipped folder is our required data dir
+        cd(currdir)
+    end
+    @assert isdir(dir)
+    return dir
+end
+
+function read_metrla(d::String)
+    adj_matrix = NPZ.npzread(joinpath(d, "adj_mat.npy"))
+    node_features = NPZ.npzread(joinpath(d, "node_values.npy"))
+    return adj_matrix, node_features
+end
+    
+function metrla_generate_task(node_values::AbstractArray, num_timesteps_in::Int, num_timesteps_out::Int)
+    indices = [(i, i + num_timesteps_in + num_timesteps_out) for i in 1:(size(node_values,1) - num_timesteps_in - num_timesteps_out)]
+    features = []
+    targets = []
+    for (i,j) in indices
+        push!(features, node_values[i:i+num_timesteps_in-1,:,:])
+        push!(targets, reshape(node_values[i+num_timesteps_in:j-1,1,:], (num_timesteps_out, 1, size(node_values, 3))))
+    end
+    return features, targets
+end
+
+Base.length(d::METRLA) = length(d.graphs)
+Base.getindex(d::METRLA, ::Colon) = d.graphs[1]
+Base.getindex(d::METRLA, i) = getindex(d.graphs, i)

src/datasets/text/udenglish.jl

@@ -22,7 +22,7 @@ function __init__udenglish()
                      """,
                      "https://raw.githubusercontent.com/UniversalDependencies/UD_English-EWT/master/" .*
                      [TRAINFILE, DEVFILE, TESTFILE],
-                     "e26845c3c78140e15d82a425388bcc58016d511616e5c2669a2e580e8ae586c0"))
+                     "2e94de3333b3b17bd06769ff8c4a81896fabec3ae51a6366f7af82625679f561"))
 end
 
 """

src/graph.jl

@@ -226,3 +226,30 @@ function edgeindex2adjlist(s, t, num_nodes; inneigs = false)
     end
     return adj
 end
+
+function adjmatrix2edgeindex(adj::AbstractMatrix{T}; weighted = true, inneigs = false) where T
+    s, t = Int[], Int[]
+    if weighted 
+        w = T[]
+    end
+    for i in 1:size(adj,1)
+        for j in 1:size(adj,2)
+            if adj[i,j] != 0
+                push!(s, i)
+                push!(t, j)
+                if weighted
+                    push!(w, adj[i,j])
+                end
+            end
+        end
+    end
+
+    if inneigs
+        s, t = t, s
+    end
+    if weighted
+        return s, t, w
+    else
+        return s, t
+    end
+end

test/datasets/graphs_no_ci.jl

@@ -338,3 +338,17 @@ end
     data = OrganicMaterialsDB(split = :test)
     @test length(data) == 2500
 end
+
+@testset "METR-LA" begin
+    data = METRLA()
+    @test data isa AbstractDataset
+    @test length(data) == 1
+    g = data[1]
+    @test g === data[:]
+    @test g isa MLDatasets.Graph
+
+    @test g.num_nodes == 207
+    @test g.num_edges == 1722
+    @test length(g.node_data.features) == 34248
+    @test length(g.node_data.targets) == 34248
+end