iteration and broadcast

Author: CarloLucibello

GNNGraphs/docs/src/guides/temporalgraph.md

@@ -85,6 +85,29 @@ GNNGraph:
   num_nodes: 10
   num_edges: 16
 ```
+## Iteration and Broadcasting
+
+Iteration and broadcasting over a temporal graph is similar to that of a vector of snapshots:
+
+```jldoctest temporal
+julia> snapshots = [rand_graph(10, 20), rand_graph(10, 14), rand_graph(10, 22)];
+
+julia> tg = TemporalSnapshotsGNNGraph(snapshots);
+
+julia> [g for g in tg] # iterate over snapshots
+3-element Vector{GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}:
+ GNNGraph(10, 20) with no data
+ GNNGraph(10, 14) with no data
+ GNNGraph(10, 22) with no data
+
+julia> f(g) = g isa GNNGraph;
+
+julia> f.(tg) # broadcast over snapshots
+3-element BitVector:
+ 1
+ 1
+ 1
+```
 
 ## Basic Queries
 

GNNGraphs/src/temporalsnapshotsgnngraph.jl

@@ -91,7 +91,13 @@ end
 
 function Base.length(tg::TemporalSnapshotsGNNGraph)
     return tg.num_snapshots
-end
+end 
+
+# Allow broadcasting over the temporal snapshots
+Base.broadcastable(tg::TemporalSnapshotsGNNGraph) = tg.snapshots
+
+Base.iterate(tg::TemporalSnapshotsGNNGraph) = Base.iterate(tg.snapshots)
+Base.iterate(tg::TemporalSnapshotsGNNGraph, i) = Base.iterate(tg.snapshots, i)
 
 function Base.setindex!(tg::TemporalSnapshotsGNNGraph, g::GNNGraph, t::Int)
     tg.snapshots[t] = g

GNNGraphs/test/temporalsnapshotsgnngraph.jl

@@ -1,3 +1,5 @@
+#TODO add graph_type = GRAPH_TYPE to all constructor calls
+
 @testset "Constructor array TemporalSnapshotsGNNGraph" begin
     snapshots = [rand_graph(10, 20) for i in 1:5]
     tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -12,6 +14,7 @@
     @test tg.num_snapshots == 5
 end
 
+
 @testset "==" begin
     snapshots = [rand_graph(10, 20) for i in 1:5]
     tsg1 = TemporalSnapshotsGNNGraph(snapshots)
@@ -41,7 +44,7 @@ end
 end
 
 @testset "getproperty" begin
-    x = rand(10)
+    x = rand(Float32, 10)
     snapshots = [rand_graph(10, 20, ndata = x) for i in 1:5]
     tsg = TemporalSnapshotsGNNGraph(snapshots)
     @test tsg.tgdata == DataStore()
@@ -111,18 +114,31 @@ end
 @testset "show" begin
     snapshots = [rand_graph(10, 20) for i in 1:5]
     tsg = TemporalSnapshotsGNNGraph(snapshots)
-    @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5) with no data"
-    @test sprint(show, MIME("text/plain"), tsg; context=:compact => true) == "TemporalSnapshotsGNNGraph(5) with no data"
+    @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5)"
+    @test sprint(show, MIME("text/plain"), tsg; context=:compact => true) == "TemporalSnapshotsGNNGraph(5)"
     @test sprint(show, MIME("text/plain"), tsg; context=:compact =>  false) == "TemporalSnapshotsGNNGraph:\n  num_nodes: [10, 10, 10, 10, 10]\n  num_edges: [20, 20, 20, 20, 20]\n  num_snapshots: 5"
-    tsg.tgdata.x=rand(4)
-    @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5) with x: 4-element data"
+    tsg.tgdata.x = rand(Float32, 4)
+    @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5)"
+end
+
+@testset "broadcastable" begin
+    snapshots = [rand_graph(10, 20) for i in 1:5]
+    tsg = TemporalSnapshotsGNNGraph(snapshots)
+    f(g) = g isa GNNGraph
+    @test f.(tsg) == trues(5)
+end
+
+@testset "iterate" begin
+    snapshots = [rand_graph(10, 20) for i in 1:5]
+    tsg = TemporalSnapshotsGNNGraph(snapshots)
+    @test [g for g in tsg] isa Vector{<:GNNGraph}
 end
 
 if TEST_GPU
     @testset "gpu" begin
-        snapshots = [rand_graph(10, 20; ndata = rand(5,10)) for i in 1:5]
+        snapshots = [rand_graph(10, 20; ndata = rand(Float32, 5,10)) for i in 1:5]
         tsg = TemporalSnapshotsGNNGraph(snapshots)
-        tsg.tgdata.x = rand(5)
+        tsg.tgdata.x = rand(Float32, 5)
         dev = CUDADevice() #TODO replace with `gpu_device()`
         tsg = tsg |> dev
         @test tsg.snapshots[1].ndata.x isa CuArray

GNNLux/docs/src_tutorials/gnn_intro.jl

@@ -220,7 +220,7 @@ visualize_embeddings(emb_init, colors = labels)
 # If you are not new to Lux, this scheme should appear familiar to you.
 
 # Note that our semi-supervised learning scenario is achieved by the following line:
-# ```
+# ```julia
 # logitcrossentropy(ŷ[:,train_mask], y[:,train_mask])
 # ```
 # While we compute node embeddings for all of our nodes, we **only make use of the training nodes for computing the loss**.