Fix traffic datasets (#216)

* Change datasets database and improvement

* Add new tests and fix tasks generation

* Fix

Author: aurorarossi

src/datasets/graphs/metrla.jl

@@ -1,6 +1,6 @@
 function __init__metrla()
-    DEPNAME = "METR-LA"
-    LINK = "https://graphmining.ai/temporal_datasets/"
+    DEPNAME = "METRLA"
+    LINK = "http://www-sop.inria.fr/members/Aurora.Rossi/index.html"
     register(ManualDataDep(DEPNAME,
                            """
                            Dataset: $DEPNAME
@@ -9,7 +9,7 @@ function __init__metrla()
 end
 
 """
-    METRLA(; num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir=nothing)
+    METRLA(; num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir=nothing, normalize = true)
 
 The METR-LA dataset from the [Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting](https://arxiv.org/abs/1707.01926) paper.
 
@@ -19,22 +19,22 @@ The edge weights `w` are contained as a feature array in `edge_data` and represe
 
 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.
+The target values are the traffic speed of the measurements collected by the sensors, divided into `num_timesteps_out` time steps.
+
+The `normalize` flag indicates whether the data are normalized using Z-score normalization.
 """
 struct METRLA <: AbstractDataset
     graphs::Vector{Graph}
 end
 
-function METRLA(;num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir = nothing)
-    s, t, w, x, y = processed_traffic("METR-LA", num_timesteps_in, num_timesteps_out, dir)
+function METRLA(;num_timesteps::Int = 12, dir = nothing, normalize = true)
+    s, t, w, x, y = processed_traffic("METRLA", num_timesteps, dir, normalize)
 
     g = Graph(; num_nodes = 207,
               edge_index = (s, t),
               edge_data = w,
-            node_data = (features = x, targets = y)
-
-
-    )
+            node_data = (features = x, targets = y))
+            
     return METRLA([g])
 end
 

src/datasets/graphs/pemsbay.jl

@@ -1,6 +1,6 @@
 function __init__pemsbay()
-    DEPNAME = "PEMS-BAY"
-    LINK = "https://graphmining.ai/temporal_datasets/"
+    DEPNAME = "PEMSBAY"
+    LINK = "http://www-sop.inria.fr/members/Aurora.Rossi/index.html"
     register(ManualDataDep(DEPNAME,
                            """
                            Dataset: $DEPNAME
@@ -9,7 +9,7 @@ function __init__pemsbay()
 end
 
 """
-    PEMSBAY(; num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir=nothing)
+    PEMSBAY(; num_timesteps_in::Int = 12, num_timesteps_out::Int=12, dir=nothing, normalize = true)
 
 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-
@@ -21,22 +21,22 @@ The edge weights `w` are contained as a feature array in `edge_data` and represe
 
 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.
+The target values are the traffic speed of the measurements collected by the sensors, divided into `num_timesteps_out` time steps.
+
+The `normalize` flag indicates whether the data are normalized using Z-score normalization.
 """
 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)
+function PEMSBAY(;num_timesteps::Int = 12, dir = nothing, normalize = true)
+    s, t, w, x, y = processed_traffic("PEMSBAY", num_timesteps, dir, normalize)
 
     g = Graph(; num_nodes = 325,
               edge_index = (s, t),
               edge_data = w,
-            node_data = (features = x, targets = y)
-
-
-    )
+            node_data = (features = x, targets = y))
+            
     return PEMSBAY([g])
 end
 

src/datasets/graphs/traffic.jl

@@ -1,10 +1,10 @@
 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
+    if dname == "PEMSBAY"
+        dir = isnothing(dir) ? datadep"PEMSBAY" : dir
+    elseif dname == "METRLA"
+        dir = isnothing(dir) ? datadep"METRLA" : dir
     end
-    LINK = "https://graphmining.ai/temporal_datasets/$dname.zip"
+    LINK = "http://www-sop.inria.fr/members/Aurora.Rossi/data/$dname.zip"
     if length(readdir((dir))) == 0
         DataDeps.fetch_default(LINK, dir)
         currdir = pwd()
@@ -18,39 +18,33 @@ function traffic_datadir(dname ::String, dir = nothing)
 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"))
-
+    adj_matrix = load(joinpath(d, "$(dname).jld2"), "adj_matrix")
+    node_features = load(joinpath(d, "$(dname).jld2"), "node_features")
+    node_features = permutedims(node_features,(2,3,1))
     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)]
+function traffic_generate_task(node_values::AbstractArray, num_timesteps::Int)
     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))))
+    for i in 1:size(node_values,3)-num_timesteps
+        push!(features, node_values[:,:,i:i+num_timesteps-1])
+        push!(targets, reshape(node_values[1,:,i+1:i+num_timesteps], (1, size(node_values, 2),num_timesteps)))
     end
     return features, targets
 end
 
-function processed_traffic(dname::String, num_timesteps_in::Int, num_timesteps_out::Int, dir = nothing)
+function processed_traffic(dname::String, num_timesteps::Int, dir = nothing, normalize = true)
     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
-
+    if normalize
+        node_values = (node_values .- Statistics.mean(node_values, dims=(3,1))) ./ Statistics.std(node_values, dims=(3,1)) #Z-score normalization
+    end
+    
     s, t, w = adjmatrix2edgeindex(adj_matrix; weighted = true)
-
-    x, y = traffic_generate_task(node_values, num_timesteps_in, num_timesteps_out)
+    x, y = traffic_generate_task(node_values, num_timesteps )
+    
     return s, t, w, x, y
 end

test/datasets/graphs_no_ci.jl

@@ -349,8 +349,7 @@ end
 
     @test g.num_nodes == 207
     @test g.num_edges == 1722
-    @test length(g.node_data.features) == 34248
-    @test length(g.node_data.targets) == 34248
+    @test all(g.node_data.features[1][:,:,1][2:end,1] == g.node_data.targets[1][:,:,1][1:end-1])
 end
 
 @testset "PEMS-BAY" begin
@@ -363,6 +362,5 @@ end
 
     @test g.num_nodes == 325
     @test g.num_edges == 2694
-    @test length(g.node_data.features) == 52081
-    @test length(g.node_data.targets) == 52081
+    @test all(g.node_data.features[1][:,:,1][2:end,1] == g.node_data.targets[1][:,:,1][1:end-1])
 end