make add_snapshot non mutating (#296)

CarloLucibello · web-flow · commit eefe691fbeed · 2023-06-13T07:29:27.000+02:00
* snapshot!

* doc pages

* fixes

* more tests

* fix add_snapshot

* fix docs

* fix docs
diff --git a/docs/make.jl b/docs/make.jl
@@ -1,10 +1,10 @@
 using Flux, NNlib, GraphNeuralNetworks, Graphs, SparseArrays
 using Documenter, DemoCards
 
-tutorials, tutorials_cb, tutorial_assets = makedemos("tutorials")
+# tutorials, tutorials_cb, tutorial_assets = makedemos("tutorials")
 
 assets = []
-isnothing(tutorial_assets) || push!(assets, tutorial_assets)
+# isnothing(tutorial_assets) || push!(assets, tutorial_assets)
 
 DocMeta.setdocmeta!(GraphNeuralNetworks, :DocTestSetup,
                     :(using GraphNeuralNetworks, Graphs, SparseArrays, NNlib, Flux);
@@ -20,24 +20,25 @@ makedocs(;
          format = Documenter.HTML(; mathengine, prettyurls, assets = assets),
          sitename = "GraphNeuralNetworks.jl",
          pages = ["Home" => "index.md",
-             "Graphs" => "gnngraph.md",
+             "Graphs" => ["gnngraph.md", "heterograph.md", "temporalgraph.md"],
              "Message Passing" => "messagepassing.md",
              "Model Building" => "models.md",
              "Datasets" => "datasets.md",
-             "HeteroGraphs" => "gnnheterograph.md",
-             "Tutorials" => tutorials,
+            #  "Tutorials" => tutorials,
              "API Reference" => [
                  "GNNGraph" => "api/gnngraph.md",
                  "Basic Layers" => "api/basic.md",
                  "Convolutional Layers" => "api/conv.md",
                  "Pooling Layers" => "api/pool.md",
                  "Message Passing" => "api/messagepassing.md",
+                 "Heterogeneous Graphs" => "api/heterograph.md",
+                 "Temporal Graphs" => "api/temporalgraph.md",
                  "Utils" => "api/utils.md",
              ],
              "Developer Notes" => "dev.md",
              "Summer Of Code" => "gsoc.md",
          ])
 
-tutorials_cb()
+# tutorials_cb()
 
 deploydocs(repo = "github.com/CarloLucibello/GraphNeuralNetworks.jl.git")
diff --git a/docs/src/api/basic.md b/docs/src/api/basic.md
@@ -7,7 +7,6 @@ CurrentModule = GraphNeuralNetworks
 ## Index
 
 ```@index
-Order = [:type, :function]
 Modules = [GraphNeuralNetworks]
 Pages = ["basic.md"]
 ```
diff --git a/docs/src/api/heterograph.md b/docs/src/api/heterograph.md
@@ -0,0 +1,11 @@
+# HeteroGNNGraph
+
+Documentation page for the graph type `HeteroGNNGraph` and related methods representing heterogeneous graphs,
+where  nodes and edges can have different types.
+
+
+```@autodocs
+Modules = [GraphNeuralNetworks.GNNGraphs]
+Pages   = ["gnnheterograph.jl"]
+Private = false
+```
diff --git a/docs/src/api/temporalgraph.md b/docs/src/api/temporalgraph.md
@@ -0,0 +1,10 @@
+# TemporalSnapthotsGNNGraph
+
+
+Documentation page for the graph type `TemporalSnapshotsGNNGraph` and related methods, representing time varying graphs with time varying features.
+
+```@docs
+TemporalSnapshotsGNNGraph
+add_snapshot
+remove_snapshot
+```
diff --git a/docs/src/gnngraph.md b/docs/src/gnngraph.md
@@ -1,4 +1,4 @@
-# Graphs
+# Working with GNNGraph 
 
 The fundamental graph type in GraphNeuralNetworks.jl is the [`GNNGraph`](@ref).
 A GNNGraph `g` is a directed graph with nodes labeled from 1 to `g.num_nodes`.
@@ -254,4 +254,4 @@ julia> GNNGraph(gd)
 GNNGraph:
   num_nodes: 10
   num_edges: 20
-```
+```
diff --git a/docs/src/heterograph.md b/docs/src/heterograph.md
@@ -1,9 +1,5 @@
 # Heterogeneous Graphs
 
-!!! warning
-    Heterographs support is still experimental.
-    The interface could be subject to change in the future.
-
 Heterogeneus 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
@@ -16,10 +12,3 @@ and data containers.
 In GraphNeuralNetworks.jl heterographs are implemented in 
 the type [`GNNHeteroGraph`](@ref).
 
-
-```@docs
-GNNHeteroGraph
-rand_heterograph
-```
-
-
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -61,7 +61,7 @@ opt = Adam(1f-4)
 
 Finally, we use a standard Flux training pipeline to fit our dataset.
 We use Flux's `DataLoader` to iterate over mini-batches of graphs 
-that are glued together into a single `GNNGraph` using the [`MLUtils.batch`](@ref) method. This is what happens under the hood when creating a `DataLoader` with the
+that are glued together into a single `GNNGraph` using the [`Flux.batch`](@ref) method. This is what happens under the hood when creating a `DataLoader` with the
 `collate=true` option. 
 
 ```julia
diff --git a/docs/src/temporalgraph.md b/docs/src/temporalgraph.md
@@ -0,0 +1,3 @@
+# Temporal Graphs
+
+Time varying graph topologies and node features are supported through the [`TemporalSnapshotsGNNGraph`](@ref) type.
diff --git a/src/GNNGraphs/GNNGraphs.jl b/src/GNNGraphs/GNNGraphs.jl
@@ -28,13 +28,16 @@ export GNNGraph,
        graph_features
 
 include("gnnheterograph.jl")
-export GNNHeteroGraph
-
+export GNNHeteroGraph,
+       num_edge_types,
+       num_node_types
 
 include("temporalsnapshotsgnngraph.jl")
 export TemporalSnapshotsGNNGraph,
        add_snapshot,
+       # add_snapshot!,
        remove_snapshot
+       # remove_snapshot!
 
 include("query.jl")
 export adjacency_list,
diff --git a/src/GNNGraphs/generate.jl b/src/GNNGraphs/generate.jl
@@ -60,7 +60,7 @@ Additional keyword arguments will be passed to the [`GNNHeteroGraph`](@ref) cons
 
 # Examples
 
-```juliarepl
+```julia-repl
 
 
 julia> g = rand_heterograph((:user => 10, :movie => 20),
@@ -70,7 +70,7 @@ GNNHeteroGraph:
   num_edges: ((:user, :rate, :movie) => 30,)
 ```
 """
-function rand_heteropraph end
+function rand_heterograph end
 
 # for generic iterators of pairs
 rand_heterograph(n, m; kws...) = rand_heterograph(Dict(n), Dict(m); kws...)
diff --git a/src/GNNGraphs/gnnheterograph.jl b/src/GNNGraphs/gnnheterograph.jl
@@ -3,39 +3,46 @@ const EDict{T} = Dict{Tuple{Symbol, Symbol, Symbol}, T}
 const NDict{T} = Dict{Symbol, T}
 
 """
-    GNNHeteroGraph(data; ndata, edata, gdata, num_nodes, graph_indicator, dir])
+    GNNHeteroGraph(data; [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.
 
-# Arguments
+# Constructor Arguments
 
 - `data`: A dictionary or an iterable object that maps (source_type, edge_type, target_type)
     triples to (source, target) index vectors.
-- `num_nodes`: The number of nodes for each type. If not specified, inferred from `g`. Default `nothing`.
-- `graph_indicator`: For batched graphs, a dictionary of vectors containing the graph assignment of each node. Default `nothing`.
 - `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`.
+- `num_nodes`: The number of nodes for each type. If not specified, inferred from `data`. Default `nothing`.
 
+# Fields
 
-!!! warning
-    `GNNHeteroGraph` is still experimental and not fully supported.
-    The interface could be subject to change in the future.
+- `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.
+- `edata`: Edge features.
+- `gdata`: Graph features.
+- `ntypes`: The node types.
+- `etypes`: The edge types.
 
 # Examples
 
 ```julia
-julia> using Flux, GraphNeuralNetworks
+julia> using GraphNeuralNetworks
 
-julia> num_nodes = Dict(:A => 10, :B => 20);
+julia> nA, nB = 10, 20;
 
-julia> edges1 = rand(1:num_nodes[:A], 20), rand(1:num_nodes[:B], 20)
+julia> num_nodes = Dict(:A => nA, :B => nB);
+
+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:num_nodes[:B], 30), rand(1:num_nodes[:A], 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);
@@ -51,8 +58,8 @@ Dict{Tuple{Symbol, Symbol, Symbol}, Int64} with 2 entries:
 (:B, :rel2, :A) => 30
 
 # Let's add some node features
-julia> ndata = Dict(:A => (x = rand(2, num_nodes[:A]), y = rand(3, num_nodes[:A])),
-                    :B => rand(10, num_nodes[:B]));
+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)
 GNNHeteroGraph:
@@ -86,6 +93,8 @@ end
 
 @functor GNNHeteroGraph
 
+GNNHeteroGraph(data; kws...) = GNNHeteroGraph(Dict(data); kws...)
+
 function GNNHeteroGraph(data::EDict;
                         num_nodes = nothing,
                         graph_indicator = nothing,
@@ -128,72 +137,84 @@ function GNNHeteroGraph(data::EDict;
                           ntypes, etypes)
 end
 
-function show_sorted_Dict(io::IO, d::Dict, compact::Bool)
+function show_sorted_dict(io::IO, d::Dict, compact::Bool)
     if compact
         print(io, "Dict")
     end
     print(io, "(")
     if !isempty(d)
-        if length(keys(d)) == 1
-            show(io, keys[1])
-            print(io, " => $(d[keys[1]])")
-        else
-            sorted_keys = sort!(collect(keys(d)))
-            for key in sorted_keys[1:(end - 1)]
-                show(io, key)
-                print(io, " => $(d[key]), ")
-            end
-            show(io, sorted_keys[end])
-            print(io, " => $(d[sorted_keys[end]])")
+        _keys = sort!(collect(keys(d)))
+        for key in _keys[1:end-1]
+            print(io, "$(_str(key)) => $(d[key]), ")
         end
+        print(io, "$(_str(_keys[end])) => $(d[_keys[end]])")
     end
     print(io, ")")
 end
 
 function Base.show(io::IO, g::GNNHeteroGraph)
     print(io, "GNNHeteroGraph(")
-    show_sorted_Dict(io, g.num_nodes, true)
+    show_sorted_dict(io, g.num_nodes, true)
     print(io, ", ")
-    show_sorted_Dict(io, g.num_edges, true)
+    show_sorted_dict(io, g.num_edges, true)
     print(io, ")")
 end
 
 function Base.show(io::IO, ::MIME"text/plain", g::GNNHeteroGraph)
     if get(io, :compact, false)
         print(io, "GNNHeteroGraph(")
-        show_sorted_Dict(io, g.num_nodes, true)
+        show_sorted_dict(io, g.num_nodes, true)
         print(io, ", ")
-        show_sorted_Dict(io, g.num_edges, true)
+        show_sorted_dict(io, g.num_edges, true)
         print(io, ")")
     else
-        print(io, "GNNHeteroGraph:\n num_nodes: ")
-        show_sorted_Dict(io, g.num_nodes, false)
-        print(io, "\n num_edges: ")
-        show_sorted_Dict(io, g.num_edges, false)
-        g.num_graphs > 1 && print(io, "\n num_graphs: $(g.num_graphs)")
+        print(io, "GNNHeteroGraph:\n  num_nodes: ")
+        show_sorted_dict(io, g.num_nodes, false)
+        print(io, "\n  num_edges: ")
+        show_sorted_dict(io, g.num_edges, false)
+        g.num_graphs > 1 && print(io, "\n  num_graphs: $(g.num_graphs)")
         if !isempty(g.ndata)
-            print(io, "\n ndata:")
+            print(io, "\n  ndata:")
             for k in sort(collect(keys(g.ndata)))
                 print(io, "\n\t", _str(k), "  =>  $(shortsummary(g.ndata[k]))")
             end
         end
         if !isempty(g.edata)
-            print(io, "\n edata:")
+            print(io, "\n  edata:")
             for k in sort(collect(keys(g.edata)))
                 print(io, "\n\t$k  =>  $(shortsummary(g.edata[k]))")
             end
         end
         if !isempty(g.gdata)
-            print(io, "\n gdata:\n\t")
+            print(io, "\n  gdata:\n\t")
             shortsummary(io, g.gdata)
         end
     end
 end
 
-GNNHeteroGraph(data; kws...) = GNNHeteroGraph(Dict(data); kws...)
-
 _str(s::Symbol) = ":$s"
 _str(s) = "$s"
 
 MLUtils.numobs(g::GNNHeteroGraph) = g.num_graphs
 # MLUtils.getobs(g::GNNHeteroGraph, i) = getgraph(g, i)
+
+
+"""
+    num_edge_types(g)
+
+Return the number of edge types in the graph. For [`GNNGraph`](@ref)s, this is always 1.
+For [`GNNHeteroGraph`](@ref)s, this is the number of unique edge types.
+"""
+num_edge_types(g::GNNGraph) = 1
+
+num_edge_types(g::GNNHeteroGraph) = length(g.etypes)
+
+"""
+    num_node_types(g)
+
+Return the number of node types in the graph. For [`GNNGraph`](@ref)s, this is always 1.
+For [`GNNHeteroGraph`](@ref)s, this is the number of unique node types.
+"""
+num_node_types(g::GNNGraph) = 1
+
+num_node_types(g::GNNHeteroGraph) = length(g.ntypes)
diff --git a/src/GNNGraphs/temporalsnapshotsgnngraph.jl b/src/GNNGraphs/temporalsnapshotsgnngraph.jl
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
diff --git a/src/msgpass.jl b/src/msgpass.jl
diff --git a/test/GNNGraphs/gnnheterograph.jl b/test/GNNGraphs/gnnheterograph.jl
diff --git a/test/GNNGraphs/temporalsnapshotsgnngraph.jl b/test/GNNGraphs/temporalsnapshotsgnngraph.jl

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# Temporal Graphs`
	`2`	`+`
	`3`	+Time varying graph topologies and node features are supported through the [`TemporalSnapshotsGNNGraph`](@ref) type.