improve add_edges for graph and heterographs (#337)

* improve add_edges for graph and heterographs

* constructor test

* jldoctest

* add tests

* fix tests

Author: CarloLucibello

docs/Project.toml

@@ -17,3 +17,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 DemoCards = "0.5.0"
+Documenter = "0.27"

docs/src/heterograph.md

@@ -1,9 +1,8 @@
 # Heterogeneous Graphs
 
 Heterogeneous graphs (also called heterographs), are graphs where each node has a type,
-that we denote with symbols such as `:user` and `:movie`,
-and edges also represent different relations identified
-by a triplet of symbols, `(source_node_type, edge_type, target_node_type)`, as in `(:user, :rate, :movie)`.
+that we denote with symbols such as `:user` and `:movie`.
+Also edges have a type, such as `:rate` or `:like`, and they can connect nodes of different types. We call a triplet `(source_node_type, edge_type, target_node_type)` the type of a *relation*, e.g. `(:user, :rate, :movie)`.
 
 Different node/edge types can store different groups of features
 and this makes heterographs a very flexible modeling tools 
@@ -13,15 +12,21 @@ the type [`GNNHeteroGraph`](@ref).
 
 ## Creating a Heterograph
 
-A heterograph can be created by passing pairs of edge types and data to the constructor.
-```julia-repl
+A heterograph can be created by passing pairs of relation type and data to the constructor.
+```jldoctest
+julia> g = GNNHeteroGraph((:user, :like, :actor) => ([1,2,2,3], [1,3,2,9]),
+                          (:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
+GNNHeteroGraph:
+  num_nodes: Dict(:actor => 9, :movie => 13, :user => 3)
+  num_edges: Dict((:user, :like, :actor) => 4, (:user, :rate, :movie) => 4)
+
 julia> g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
 GNNHeteroGraph:
   num_nodes: Dict(:movie => 13, :user => 3)
   num_edges: Dict((:user, :rate, :movie) => 4)
 ```
 New relations, possibly with new node types, can be added with the function [`add_edges`](@ref).
-```julia-repl
+```jldoctest
 julia> g = add_edges(g, (:user, :like, :actor) => ([1,2,3,3,3], [3,5,1,9,4]))
 GNNHeteroGraph:
   num_nodes: Dict(:actor => 9, :movie => 13, :user => 3)
@@ -30,7 +35,7 @@ GNNHeteroGraph:
 See [`rand_heterograph`](@ref), [`rand_bipartite_heterograph`](@ref)
 for generating random heterographs. 
 
-```julia-repl
+```jldoctest
 julia> g = rand_bipartite_heterograph((10, 15), 20)
 GNNHeteroGraph:
   num_nodes: Dict(:A => 10, :B => 15)
@@ -40,7 +45,7 @@ GNNHeteroGraph:
 ## Basic Queries
 
 Basic queries are similar to those for homogeneous graphs:
-```julia-repl
+```jldoctest
 julia> g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
 GNNHeteroGraph:
   num_nodes: Dict(:movie => 13, :user => 3)
@@ -74,7 +79,7 @@ julia> g.etypes
 ## Data Features
 
 Node, edge, and graph features can be added at construction time or later using:
-```julia-repl
+```jldoctest
 # equivalent to g.ndata[:user][:x] = ...
 julia> g[:user].x = rand(Float32, 64, 3);
 
@@ -96,7 +101,7 @@ GNNHeteroGraph:
 
 ## Batching
 Similarly to graphs, also heterographs can be batched together.
-```julia-repl
+```jldoctest
 julia> gs = [rand_bipartite_heterograph((5, 10), 20) for _ in 1:32];
 
 julia> Flux.batch(gs)
@@ -108,7 +113,7 @@ GNNHeteroGraph:
 Batching is automatically performed by the [`DataLoader`](@ref) iterator
 when the `collate` option is set to `true`.
 
-```julia-repl
+```jldoctest
 using Flux: DataLoader
 
 data = [rand_bipartite_heterograph((5, 10), 20, 

docs/src/temporalgraph.md

@@ -6,7 +6,7 @@ Temporal Graphs are graphs with time varying topologies and node features. In Gr
 
 A temporal graph can be created by passing a list of snapshots to the constructor. Each snapshot is a [`GNNGraph`](@ref). 
 
-```julia-repl
+```jldoctest
 julia> snapshots = [rand_graph(10,20) for i in 1:5];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -18,14 +18,14 @@ TemporalSnapshotsGNNGraph:
 
 A new temporal graph can be created by adding or removing snapshots to an existing temporal graph. 
 
-```julia-repl
+```jldoctest
 julia> new_tg = add_snapshot(tg, 3, rand_graph(10, 16)) # add a new snapshot at time 3
 TemporalSnapshotsGNNGraph:
   num_nodes: [10, 10, 10, 10, 10, 10]
   num_edges: [20, 20, 16, 20, 20, 20]
   num_snapshots: 6
 ```
-```julia-repl
+```jldoctest
 julia> snapshots = [rand_graph(10,20), rand_graph(10,14), rand_graph(10,22)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -43,7 +43,7 @@ TemporalSnapshotsGNNGraph:
 
 See [`rand_temporal_radius_graph`](@ref) and ['rand_temporal_hyperbolic_graph'](@ref) for generating random temporal graphs. 
 
-```julia-repl
+```jldoctest
 julia> tg = rand_temporal_radius_graph(10, 3, 0.1, 0.5)
 TemporalSnapshotsGNNGraph:
   num_nodes: [10, 10, 10]
@@ -54,7 +54,7 @@ TemporalSnapshotsGNNGraph:
 ## Basic Queries
 
 Basic queries are similar to those for [`GNNGraph`](@ref)s:
-```julia-repl
+```jldoctest
 julia> snapshots = [rand_graph(10,20), rand_graph(10,14), rand_graph(10,22)];
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots)
@@ -94,7 +94,7 @@ GNNGraph:
 
 Node, edge, and graph features can be added at construction time or later using:
 
-```julia-repl
+```jldoctest
 julia> snapshots = [rand_graph(10,20; ndata = rand(3,10)), rand_graph(10,14; ndata = rand(4,10)), rand_graph(10,22; ndata = rand(5,10))]; # node features at construction time
 
 julia> tg = TemporalSnapshotsGNNGraph(snapshots);

src/GNNGraphs/datastore.jl

@@ -8,7 +8,7 @@ A container for feature arrays. The optional argument `n` enforces that
 At construction time, the `data` can be provided as any iterables of pairs
 of symbols and arrays or as keyword arguments:
 
-```julia-repl
+```jldoctest
 julia> ds = DataStore(3, x = rand(2, 3), y = rand(3))
 DataStore(3) with 2 elements:
   y = 3-element Vector{Float64}
@@ -35,7 +35,7 @@ DataStore() with 2 elements:
 The `DataStore` has an interface similar to both dictionaries and named tuples.
 Arrays can be accessed and added using either the indexing or the property syntax:
 
-```julia-repl
+```jldoctest
 julia> ds = DataStore(x = ones(2, 3), y = zeros(3))
 DataStore() with 2 elements:
   y = 3-element Vector{Float64}
@@ -57,7 +57,7 @@ The `DataStore` can be iterated over, and the keys and values can be accessed
 using `keys(ds)` and `values(ds)`. `map(f, ds)` applies the function `f`
 to each feature array:
 
-```julia-repl
+```jldoctest
 julia> ds = DataStore(a = zeros(2), b = zeros(2));
 
 julia> ds2 = map(x -> x .+ 1, ds)
@@ -142,6 +142,8 @@ function Base.show(io::IO, ds::DataStore)
         for (k, v) in getdata(ds)
             print(io, "\n  $(k) = $(summary(v))")
         end
+    else
+        print(io, " with no elements")
     end
 end
 

src/GNNGraphs/generate.jl

@@ -16,7 +16,7 @@ Additional keyword arguments will be passed to the [`GNNGraph`](@ref) constructo
 
 # Examples
 
-```juliarepl
+```jldoctest
 julia> g = rand_graph(5, 4, bidirected=false)
 GNNGraph:
     num_nodes = 5
@@ -64,7 +64,7 @@ Additional keyword arguments will be passed to the [`GNNHeteroGraph`](@ref) cons
 
 # Examples
 
-```julia-repl
+```jldoctest
 julia> g = rand_heterograph((:user => 10, :movie => 20),
                             (:user, :rate, :movie) => 30)
 GNNHeteroGraph:
@@ -170,7 +170,7 @@ to its `k` closest `points`.
 
 # Examples
 
-```juliarepl
+```jldoctest
 julia> n, k = 10, 3;
 
 julia> x = rand(3, n);
@@ -251,7 +251,7 @@ to its neighbors within a given distance `r`.
 
 # Examples
 
-```juliarepl
+```jldoctest
 julia> n, r = 10, 0.75;
 
 julia> x = rand(3, n);
@@ -331,7 +331,7 @@ If a point happens to move outside the boundary, its position is updated as if i
 
 # Example
 
-```julia-repl
+```jldoctest
 julia> n, snaps, s, r = 10, 5, 0.1, 1.5;
 
 julia> tg = rand_temporal_radius_graph(n,snaps,s,r) # complete graph at each snapshot
@@ -403,7 +403,7 @@ First, the positions of the nodes are generated with a quasi-uniform distributio
 
 # Example
 
-```julia-repl
+```jldoctest
 julia> n, snaps, α, R, speed, ζ = 10, 5, 1.0, 4.0, 0.1, 1.0;
 
 julia> thg = rand_temporal_hyperbolic_graph(n, snaps; α, R, speed, ζ)

src/GNNGraphs/gnngraph.jl

@@ -153,6 +153,8 @@ function GNNGraph(data::D;
              ndata, edata, gdata)
 end
 
+GNNGraph(; kws...) = GNNGraph(0; kws...)
+
 function (::Type{<:GNNGraph})(num_nodes::T; kws...) where {T <: Integer}
     s, t = T[], T[]
     return GNNGraph(s, t; num_nodes, kws...)

src/GNNGraphs/gnnheterograph.jl

@@ -6,24 +6,26 @@ const NDict{T} = Dict{NType, T}
 
 """
     GNNHeteroGraph(data; [ndata, edata, gdata, num_nodes])
+    GNNHeteroGraph(pairs...; [ndata, edata, gdata, num_nodes])
 
 A type representing a heterogeneous graph structure.
 It is similar to [`GNNGraph`](@ref) but nodes and edges are of different types.
 
 # Constructor Arguments
 
-- `data`: A dictionary or an iterable object that maps (source_type, edge_type, target_type)
-    triples to (source, target) index vectors.
+- `data`: A dictionary or an iterable object that maps `(source_type, edge_type, target_type)`
+          triples to `(source, target)` index vectors (or to `(source, target, weight)` if also edge weights are present).
+- `pairs`: Passing multiple relations as pairs is equivalent to passing `data=Dict(pairs...)`.
 - `ndata`: Node features. A dictionary of arrays or named tuple of arrays.
            The size of the last dimension of each array must be given by `g.num_nodes`.
-- `edata`: Edge features. A dictionary of arrays or named tuple of arrays.
-           The size of the last dimension of each array must be given by `g.num_edges`.
-- `gdata`: Graph features. An array or named tuple of arrays whose last dimension has size `num_graphs`.
+- `edata`: Edge features. A dictionary of arrays or named tuple of arrays. Default `nothing`.
+           The size of the last dimension of each array must be given by `g.num_edges`. Default `nothing`.
+- `gdata`: Graph features. An array or named tuple of arrays whose last dimension has size `num_graphs`. Default `nothing`.
 - `num_nodes`: The number of nodes for each type. If not specified, inferred from `data`. Default `nothing`.
 
 # Fields
 
-- `graph`: A dictionary that maps `(source_type, edge_type, target_type)`` triples to (source, target) index vectors.
+- `graph`: A dictionary that maps (source_type, edge_type, target_type) triples to (source, target) index vectors.
 - `num_nodes`: The number of nodes for each type.
 - `num_edges`: The number of edges for each type.
 - `ndata`: Node features.
@@ -41,15 +43,15 @@ julia> nA, nB = 10, 20;
 
 julia> num_nodes = Dict(:A => nA, :B => nB);
 
-julia> edges1 = rand(1:nA, 20), rand(1:nB, 20)
+julia> edges1 = (rand(1:nA, 20), rand(1:nB, 20))
 ([4, 8, 6, 3, 4, 7, 2, 7, 3, 2, 3, 4, 9, 4, 2, 9, 10, 1, 3, 9], [6, 4, 20, 8, 16, 7, 12, 16, 5, 4, 6, 20, 11, 19, 17, 9, 12, 2, 18, 12])
 
-julia> edges2 = rand(1:nB, 30), rand(1:nA, 30)
+julia> edges2 = (rand(1:nB, 30), rand(1:nA, 30))
 ([17, 5, 2, 4, 5, 3, 8, 7, 9, 7  …  19, 8, 20, 7, 16, 2, 9, 15, 8, 13], [1, 1, 3, 1, 1, 3, 2, 7, 4, 4  …  7, 10, 6, 3, 4, 9, 1, 5, 8, 5])
 
-julia> eindex = ((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2);
+julia> data = ((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2);
 
-julia> hg = GNNHeteroGraph(eindex; num_nodes)
+julia> hg = GNNHeteroGraph(data; num_nodes)
 GNNHeteroGraph:
   num_nodes: (:A => 10, :B => 20)
   num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)
@@ -63,7 +65,7 @@ Dict{Tuple{Symbol, Symbol, Symbol}, Int64} with 2 entries:
 julia> ndata = Dict(:A => (x = rand(2, nA), y = rand(3, num_nodes[:A])),
                     :B => rand(10, nB));
 
-julia> hg = GNNHeteroGraph(eindex; num_nodes, ndata)
+julia> hg = GNNHeteroGraph(data; num_nodes, ndata)
 GNNHeteroGraph:
     num_nodes: (:A => 10, :B => 20)
     num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)
@@ -96,9 +98,10 @@ end
 @functor GNNHeteroGraph
 
 GNNHeteroGraph(data; kws...) = GNNHeteroGraph(Dict(data); kws...)
+GNNHeteroGraph(data::Pair...; kws...) = GNNHeteroGraph(Dict(data...); kws...)
 
 function GNNHeteroGraph(data::Dict; kws...)
-    all(k -> k isa EType, keys(data)) || throw(ArgumentError("Keys of data must be tuples of the form (source_type, edge_type, target_type)"))
+    all(k -> k isa EType, keys(data)) || throw(ArgumentError("Keys of data must be tuples of the form `(source_type, edge_type, target_type)`"))
     return GNNHeteroGraph(Dict([k => v for (k, v) in pairs(data)]...); kws...)
 end
 

src/GNNGraphs/query.jl

@@ -62,7 +62,7 @@ Return `true` if there is an edge of type `edge_t` from node `i` to node `j` in
 
 # Examples
 
-```julia-repl
+```jldoctest
 julia> g = rand_bipartite_heterograph((2, 2), (4, 0), bidirected=false)
 GNNHeteroGraph:
   num_nodes: (:A => 2, :B => 2)

src/GNNGraphs/temporalsnapshotsgnngraph.jl

@@ -77,7 +77,7 @@ Return a `TemporalSnapshotsGNNGraph` created starting from `tg` by adding the sn
 
 # Examples
 
-```julia-repl
+```jldoctest
 julia> using GraphNeuralNetworks
 
 julia> snapshots = [rand_graph(10, 20) for i in 1:5];
@@ -137,7 +137,7 @@ Return a [`TemporalSnapshotsGNNGraph`](@ref) created starting from `tg` by remov
 
 # Examples
 
-```julia-repl
+```jldoctest
 julia> using GraphNeuralNetworks
 
 julia> snapshots = [rand_graph(10,20), rand_graph(10,14), rand_graph(10,22)];