Add PEMS-BAY dataset (#213)

* Add `PEMSBAY` in docs

* Export and include traffic

* Add common script for traffic datasets

* Modify `METRLA` using traffic

* Add `PEMSBAY` dataset

* Add `PEMSBAY ` test

Author: aurorarossi

docs/src/datasets/graphs.md

@@ -31,4 +31,5 @@ PubMed
 Reddit
 TUDataset
 METRLA
+PEMSBAY
 ```

src/MLDatasets.jl

@@ -109,6 +109,7 @@ include("graph.jl")
 # export Graph
 
 include("datasets/graphs/planetoid.jl")
+include("datasets/graphs/traffic.jl")
 # export read_planetoid_data
 include("datasets/graphs/cora.jl")
 export Cora
@@ -132,6 +133,8 @@ include("datasets/graphs/tudataset.jl")
 export TUDataset
 include("datasets/graphs/metrla.jl")
 export METRLA
+include("datasets/graphs/pemsbay.jl")
+export PEMSBAY
 
 # Meshes
 
@@ -152,6 +155,7 @@ function __init__()
     __init__reddit()
     __init__tudataset()
     __init__metrla()
+    __init__pemsbay()
 
     # misc
     __init__iris()

src/datasets/graphs/metrla.jl

@@ -26,16 +26,7 @@ struct METRLA <: AbstractDataset
 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)
+    s, t, w, x, y = processed_traffic("METR-LA", num_timesteps_in, num_timesteps_out, dir)
 
     g = Graph(; num_nodes = 207,
               edge_index = (s, t),
@@ -47,39 +38,6 @@ function METRLA(;num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir = no
     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/graphs/pemsbay.jl

@@ -0,0 +1,45 @@
+function __init__pemsbay()
+    DEPNAME = "PEMS-BAY"
+    LINK = "https://graphmining.ai/temporal_datasets/"
+    register(ManualDataDep(DEPNAME,
+                           """
+                           Dataset: $DEPNAME
+                           Website : $LINK
+                           """))
+end
+
+"""
+    PEMSBAY(; num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir=nothing)
+
+The PEMS-BAY dataset described in the [Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting](https://arxiv.org/abs/1707.01926) paper.
+It is collected by California Transportation Agencies (Cal-
+Trans) Performance Measurement System (PeMS).
+
+`PEMSBAY` is a graph with 325 nodes representing traffic sensors in the Bay Area. 
+
+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 PEMSBAY <: AbstractDataset
+    graphs::Vector{Graph}
+end
+
+function PEMSBAY(;num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir = nothing)
+    s, t, w, x, y = processed_traffic("PEMS-BAY", num_timesteps_in, num_timesteps_out, dir)
+
+    g = Graph(; num_nodes = 325,
+              edge_index = (s, t),
+              edge_data = w,
+            node_data = (features = x, targets = y)
+
+
+    )
+    return PEMSBAY([g])
+end
+
+Base.length(d::PEMSBAY) = length(d.graphs)
+Base.getindex(d::PEMSBAY, ::Colon) = d.graphs[1]
+Base.getindex(d::PEMSBAY, i) = getindex(d.graphs, i)

src/datasets/graphs/traffic.jl

@@ -0,0 +1,56 @@
+function traffic_datadir(dname ::String, dir = nothing)
+    if dname == "PEMS-BAY"
+        dir = isnothing(dir) ? datadep"PEMS-BAY" : dir
+    elseif dname == "METR-LA"
+        dir = isnothing(dir) ? datadep"METR-LA" : dir
+    end
+    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_traffic(d::String, dname::String)
+    if dname == "PEMS-BAY"
+        s="pems_"
+    elseif dname == "METR-LA"
+        s=""
+    end
+
+    adj_matrix = NPZ.npzread(joinpath(d, "$(s)adj_mat.npy"))
+    node_features = NPZ.npzread(joinpath(d, "$(s)node_values.npy"))
+
+    return adj_matrix, node_features
+end
+    
+function traffic_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
+
+function processed_traffic(dname::String, num_timesteps_in::Int, num_timesteps_out::Int, dir = nothing)
+    create_default_dir(dname)
+    d = traffic_datadir(dname, dir)
+    adj_matrix, node_values = read_traffic(d, dname)
+
+    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 = traffic_generate_task(node_values, num_timesteps_in, num_timesteps_out)
+    return s, t, w, x, y
+end

test/datasets/graphs_no_ci.jl

@@ -351,4 +351,18 @@ end
     @test g.num_edges == 1722
     @test length(g.node_data.features) == 34248
     @test length(g.node_data.targets) == 34248
+end
+
+@testset "PEMS-BAY" begin
+    data = PEMSBAY()
+    @test data isa AbstractDataset
+    @test length(data) == 1
+    g = data[1]
+    @test g === data[:]
+    @test g isa MLDatasets.Graph
+
+    @test g.num_nodes == 325
+    @test g.num_edges == 2694
+    @test length(g.node_data.features) == 52081
+    @test length(g.node_data.targets) == 52081
 end