Merge pull request #145 from CarloLucibello/cl/nodecheck

better printing for GNNGraphs and check that features are arrays

Author: CarloLucibello

src/GNNGraphs/gnngraph.jl

@@ -206,19 +206,19 @@ function Base.show(io::IO, g::GNNGraph)
     if !isempty(g.ndata)
         print(io, "\n    ndata:")
         for k in keys(g.ndata)
-            print(io, "\n        $k => $(size(g.ndata[k]))")
+            print(io, "\n        $k => $(summary(g.ndata[k]))")
         end
     end
     if !isempty(g.edata)
         print(io, "\n    edata:")
         for k in keys(g.edata)
-            print(io, "\n        $k => $(size(g.edata[k]))")
+            print(io, "\n        $k => $(summary(g.edata[k]))")
         end
     end
     if !isempty(g.gdata)
         print(io, "\n    gdata:")
         for k in keys(g.gdata)
-            print(io, "\n        $k => $(size(g.gdata[k]))")
+            print(io, "\n        $k => $(summary(g.gdata[k]))")
         end
     end
 end

src/GNNGraphs/utils.jl

@@ -57,36 +57,38 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
     # This had to workaround two Zygote bugs with NamedTuples
     # https://github.com/FluxML/Zygote.jl/issues/1071
     # https://github.com/FluxML/Zygote.jl/issues/1072
+
+    if n != 1
+        @assert all(x -> x isa AbstractArray, data) "Non-array features provided."
+    end
 
     if n == 1
         # If last array dimension is not 1, add a new dimension. 
-        # This is mostly usefule to reshape globale feature vectors
+        # This is mostly useful to reshape global feature vectors
         # of size D to Dx1 matrices.
-        function unsqz(v)
-            if v isa AbstractArray && size(v)[end] != 1
-                v = reshape(v, size(v)..., 1)
-            end
-            v
-        end
+        unsqz_last(v::AbstractArray) = size(v)[end] != 1 ? reshape(v, size(v)..., 1) : v
+        unsqz_last(v) = v
 
-        data = NamedTuple{keys(data)}(unsqz.(values(data)))
+        data = map(unsqz_last, data)
     end
-
-    sz = map(x -> x isa AbstractArray ? size(x)[end] : 0, data)
-    if duplicate_if_needed 
-        # Used to copy edge features on reverse edges    
-        @assert all(s -> s == 0 ||  s == n || s == n÷2, sz)  "Wrong size in last dimension for feature array."
 
+    ## Turn vectors in 1 x n matrices. 
+    # unsqz_first(v::AbstractVector) = reshape(v, 1, length(v))
+    # unsqz_first(v) = v
+    # data = map(unsqz_first, data)
+    
+    if duplicate_if_needed 
         function duplicate(v)
             if v isa AbstractArray && size(v)[end] == n÷2
                 v = cat(v, v, dims=ndims(v))
             end
             v
         end
-        data = NamedTuple{keys(data)}(duplicate.(values(data)))
-    else
-        @assert all(s -> s == 0 ||  s == n, sz) "Wrong size in last dimension for feature array."
+        data = map(duplicate, data)
     end
+    
+    @assert all(x -> x isa AbstractArray ? size(x)[end] == n : true, data)  "Wrong size in last dimension for feature array."
+    
     return data
 end
 

test/GNNGraphs/gnngraph.jl

@@ -225,10 +225,24 @@
         g = GNNGraph(erdos_renyi(10,  30), edata=e, graph_type=GRAPH_T)
         @test g.edata.e == [e; e]
 
-
-        # Attach non array data
-        g = GNNGraph(erdos_renyi(10,  30), edata="ciao", graph_type=GRAPH_T)
-        @test g.edata.e == "ciao"
+        # non-array global 
+        g = rand_graph(10,  30, gdata="ciao", graph_type=GRAPH_T)
+        @test g.gdata.u == "ciao"
+    
+        # vectors stays vectors
+        g = rand_graph(10,  30, ndata=rand(10),
+                                edata=rand(30),
+                                gdata=(u=rand(2), z=rand(1), q=1),
+                                graph_type=GRAPH_T)
+        @test size(g.ndata.x) == (10,)
+        @test size(g.edata.e) == (30,)
+        @test size(g.gdata.u) == (2, 1)
+        @test size(g.gdata.z) == (1,)
+        @test g.gdata.q === 1
+        
+        # Error for non-array ndata
+        @test_throws AssertionError rand_graph(10,  30, ndata="ciao", graph_type=GRAPH_T)
+        @test_throws AssertionError rand_graph(10,  30, ndata=1, graph_type=GRAPH_T)
     end
 
     @testset "LearnBase and DataLoader compat" begin