convenience accessors (#231)

* convenience accessors

* docstring

* tests

Author: CarloLucibello

docs/src/gnngraph.md

@@ -74,32 +74,35 @@ false
 
 One or more arrays can be associated to nodes, edges, and (sub)graphs of a `GNNGraph`.
 They will be stored in the fields `g.ndata`, `g.edata`, and `g.gdata` respectivaly.
-The data fields are `NamedTuple`s. The array they contain must have last dimension
-equal to `num_nodes` (in `ndata`), `num_edges` (in `edata`), or `num_graphs` (in `gdata`).
+The data fields are `NamedTuple`s. The arrays they contain have last dimension
+equal to `num_nodes` (in `ndata`), `num_edges` (in `edata`), or `num_graphs` (in `gdata`) respectively.
 
 ```julia
 # Create a graph with a single feature array `x` associated to nodes
-g = GNNGraph(erdos_renyi(10,  30), ndata = (; x = rand(Float32, 32, 10)))
+g = rand_graph(10,  60, ndata = (; x = rand(Float32, 32, 10)))
 
 g.ndata.x  # access the features
 
 # Equivalent definition passing directly the array
-g = GNNGraph(erdos_renyi(10,  30), ndata = rand(Float32, 32, 10))
+g = rand_graph(10,  60, ndata = rand(Float32, 32, 10))
 
 g.ndata.x  # `:x` is the default name for node features
 
+# For convinience, we can access the features through the shortcut
+g.x 
+
 # You can have multiple feature arrays
-g = GNNGraph(erdos_renyi(10,  30), ndata = (; x=rand(Float32, 32, 10), y=rand(Float32, 10)))
+g = rand_graph(10,  60, ndata = (; x=rand(Float32, 32, 10), y=rand(Float32, 10)))
 
-g.ndata.y, g.ndata.x
+g.ndata.y, g.ndata.x   # or g.x, g.y
 
 # Attach an array with edge features.
 # Since `GNNGraph`s are directed, the number of edges
 # will be double that of the original Graphs' undirected graph.
 g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 60))
 @assert g.num_edges == 60
 
-g.edata.e
+g.edata.e  # or g.e
 
 # If we pass only half of the edge features, they will be copied
 # on the reversed edges.
@@ -110,8 +113,8 @@ g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 30))
 # but replacing node data
 g′ = GNNGraph(g, ndata =(; z = ones(Float32, 16, 10)))
 
-g.ndata.z
-g.edata.e
+g′.z
+g′.e
 ```
 
 ## Edge weights

docs/src/index.md

@@ -32,7 +32,7 @@ using Flux.Data: DataLoader
 all_graphs = GNNGraph[]
 
 for _ in 1:1000
-    g = GNNGraph(random_regular_graph(10, 4),  
+    g = rand_graph(10, 40,  
             ndata=(; x = randn(Float32, 16,10)),  # input node features
             gdata=(; y = randn(Float32)))         # regression target   
     push!(all_graphs, g)
@@ -65,14 +65,14 @@ that are glued together into a single `GNNGraph` using the [`MLUtils.batch`](@re
 `collate=true` option. 
 
 ```julia
-train_graphs, test_graphs = MLUtils.split(all_graphs, at=0.8)
+train_graphs, test_graphs = MLUtils.splitobs(all_graphs, at=0.8)
 
 train_loader = DataLoader(train_graphs, 
                 batchsize=32, shuffle=true, collate=true)
 test_loader = DataLoader(test_graphs, 
                 batchsize=32, shuffle=false, collate=true)
 
-loss(g::GNNGraph) = mean((vec(model(g, g.ndata.x)) - g.gdata.y).^2)
+loss(g::GNNGraph) = mean((vec(model(g, g.x)) - g.y).^2)
 
 loss(loader) = mean(loss(g |> device) for g in loader)
 

src/GNNGraphs/gnngraph.jl

@@ -93,8 +93,8 @@ g = GNNGraph(g, ndata = (x=rand(100, g.num_nodes), y=rand(g.num_nodes)))
 # Add node features and edge features with default names `x` and `e` 
 g = GNNGraph(g, ndata = rand(100, g.num_nodes), edata = rand(16, g.num_edges))
 
-g.ndata.x
-g.ndata.e
+g.ndata.x # or just g.x
+g.ndata.e # or just g.e
 
 # Send to gpu
 g = g |> gpu
@@ -273,3 +273,21 @@ function Base.hash(g::T, h::UInt) where T<:GNNGraph
     fs = (getfield(g, k) for k in fieldnames(typeof(g)) if k !== :graph_indicator)
     return foldl((h, f) -> hash(f, h),  fs, init=hash(T, h))
 end
+
+function Base.getproperty(g::GNNGraph, s::Symbol)
+    if s in fieldnames(GNNGraph)
+        return getfield(g, s)
+    end
+    if (s in keys(g.ndata)) + (s in keys(g.edata)) + (s in keys(g.gdata)) > 1 
+        throw(ArgumentError("Ambiguous property name $s"))
+    end
+    if s in keys(g.ndata)
+        return g.ndata[s]
+    elseif s in keys(g.edata)
+        return g.edata[s]
+    elseif s in keys(g.gdata)
+        return g.gdata[s]
+    else
+        throw(ArgumentError("$(s) is not a field of GNNGraph"))
+    end
+end

test/GNNGraphs/gnngraph.jl

@@ -204,6 +204,9 @@
         @test g.ndata.x === X
         @test g.edata.e === E
         @test g.gdata.u === U
+        @test g.x === g.ndata.x
+        @test g.e === g.edata.e
+        @test g.u === g.gdata.u
 
         # Check no args
         g = GNNGraph(g)
@@ -220,6 +223,12 @@
         @test g.ndata.x2 ≈ 2X
         @test g.edata.e2 ≈ 2E
         @test g.gdata.u2 ≈ 2U
+        @test g.x === g.ndata.x
+        @test g.e === g.edata.e
+        @test g.u === g.gdata.u
+        @test g.x2 === g.ndata.x2
+        @test g.e2 === g.edata.e2
+        @test g.u2 === g.gdata.u2
 
         # Dimension checks
         @test_throws AssertionError GNNGraph(erdos_renyi(10, 30), edata=rand(29), graph_type=GRAPH_T)
@@ -249,6 +258,12 @@
         # 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)
+
+        # Error for Ambiguous getproperty
+        g = rand_graph(10, 20, ndata=rand(2,10), edata=(; x = rand(3,20)), graph_type=GRAPH_T)
+        @test size(g.ndata.x) == (2,10)
+        @test size(g.edata.x) == (3,20)
+        @test_throws ArgumentError g.x
     end
 
     @testset "MLUtils and DataLoader compat" begin