keep vectors

Author: CarloLucibello

src/GNNGraphs/gnngraph.jl

@@ -206,7 +206,6 @@ 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

src/GNNGraphs/utils.jl

@@ -66,16 +66,17 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
         # If last array dimension is not 1, add a new dimension. 
         # 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
-
+    
+    ## 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)
+    
     sz = map(x -> x isa AbstractArray ? size(x)[end] : 0, data)
 
     if duplicate_if_needed 
@@ -88,10 +89,11 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
             end
             v
         end
-        data = NamedTuple{keys(data)}(duplicate.(values(data)))
-    else
-        @assert all(x -> x == 0 || x == n, sz) "Wrong size in last dimension for feature array."
+        data = map(duplicate, data)
     end
+    
+    @assert all(x -> x == 0 || x == n, sz) "Wrong size in last dimension for feature array."
+    
     return data
 end
 

test/GNNGraphs/gnngraph.jl

@@ -225,15 +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"
-
-
-        # Wrong need number of features
-        g = GNNGraph(erdos_renyi(10,  30), edata="ciao", graph_type=GRAPH_T)
-
+        # 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