export get_edge_weight + use weights from adjacency matrix (#86)

* When constructing a graph from an adjacency matrix, the matrix's elements are now considered edge weights. 
* export get_edge_weight

Author: CarloLucibello

Project.toml

@@ -35,9 +35,9 @@ Graphs = "1.4"
 KrylovKit = "0.5"
 LearnBase = "0.4, 0.5"
 MacroTools = "0.5"
-NearestNeighbors = "0.4"
 NNlib = "0.7"
 NNlibCUDA = "0.1"
+NearestNeighbors = "0.4"
 Reexport = "1"
 StatsBase = "0.32, 0.33"
 julia = "1.6"

docs/src/gnngraph.md

@@ -130,7 +130,15 @@ julia> g = GNNGraph(source, target, weight)
 GNNGraph:
     num_nodes = 3
     num_edges = 6
-    
+
+julia> get_edge_weight(g)
+6-element Vector{Float64}:
+ 1.0
+ 0.5
+ 2.1
+ 2.3
+ 4.0
+ 4.1
 ```
 
 ## Batches and Subgraphs

src/GNNGraphs/GNNGraphs.jl

@@ -24,9 +24,10 @@ export GNNGraph,
 
 include("query.jl")
 export adjacency_list,
-       edge_index, 
+       edge_index,
+       get_edge_weight,
        graph_indicator, 
-       has_multi_edges, 
+       has_multi_edges,
        is_directed,
        is_bidirected,
        normalized_laplacian, 
@@ -39,16 +40,16 @@ export adjacency_list,
        outneighbors 
 
 include("transform.jl")
-export add_nodes, 
-       add_edges, 
+export add_nodes,
+       add_edges,
        add_self_loops,
        getgraph,
        negative_sample,
        rand_edge_split,
-       remove_self_loops, 
+       remove_self_loops,
        remove_multi_edges,
        # from Flux
-       batch, 
+       batch,
        unbatch,
        # from SparseArrays
        blockdiag

src/GNNGraphs/convert.jl

@@ -18,21 +18,22 @@ function to_coo(A::SPARSE_T; dir=:out, num_nodes=nothing)
     if dir == :in
         s, t = t, s
     end
-    num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
+    num_nodes = isnothing(num_nodes) ? size(A, 1) : num_nodes 
     num_edges = length(s)
 
-    return (s, t, nothing), num_nodes, num_edges
+    return (s, t, v), num_nodes, num_edges
 end
 
 function to_coo(A::ADJMAT_T; dir=:out, num_nodes=nothing)
     nz = findall(!=(0), A) # vec of cartesian indexes
     s, t = ntuple(i -> map(t->t[i], nz), 2)
+    v = A[nz] 
     if dir == :in
         s, t = t, s
     end
-    num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
+    num_nodes = isnothing(num_nodes) ? size(A, 1) : num_nodes 
     num_edges = length(s)
-    return (s, t, nothing), num_nodes, num_edges
+    return (s, t, v), num_nodes, num_edges
 end
 
 function to_coo(adj_list::ADJLIST_T; dir=:out, num_nodes=nothing)
@@ -140,15 +141,17 @@ end
 
 function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing)
     s, t, eweight  = coo
-    T = T === nothing ? eltype(s) : T
-    eweight = isnothing(eweight) ? fill!(similar(s, T), 1) : eweight
+    T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T
+    eweight = eweight === nothing ? fill!(similar(s, T), 1) : eweight
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     A = sparse(s, t, eweight, num_nodes, num_nodes)
-    num_edges = length(s)
+    num_edges = nnz(A)
+    if eltype(A) != T
+        A = T.(A)
+    end
     return A, num_nodes, num_edges
 end
 
-
 @non_differentiable to_coo(x...)
 @non_differentiable to_dense(x...)
 @non_differentiable to_sparse(x...)

src/GNNGraphs/gnngraph.jl

@@ -47,7 +47,7 @@ functionality from that library.
     - An adjacency matrix
     - An adjacency list.
     - A tuple containing the source and target vectors (COO representation)
-    - A Graphs' graph.
+    - A Graphs.jl' graph.
 - `graph_type`: A keyword argument that specifies 
                 the underlying representation used by the GNNGraph. 
                 Currently supported values are 
@@ -61,9 +61,9 @@ functionality from that library.
         Possible values are `:out` and `:in`. Default `:out`.
 - `num_nodes`: The number of nodes. If not specified, inferred from `g`. Default `nothing`.
 - `graph_indicator`: For batched graphs, a vector containing the graph assigment of each node. Default `nothing`.  
-- `ndata`: Node features. A named tuple of arrays whose last dimension has size `num_nodes`.
-- `edata`: Edge features. A named tuple of arrays whose last dimension has size `num_edges`.
-- `gdata`: Graph features. A named tuple of arrays whose last dimension has size `num_graphs`. 
+- `ndata`: Node features. An array or named tuple of arrays whose last dimension has size `num_nodes`.
+- `edata`: Edge features. An array or named tuple of arrays whose last dimension has size `num_edges`.
+- `gdata`: Graph features. An array or named tuple of arrays whose last dimension has size `num_graphs`. 
 
 # Examples 
 

src/GNNGraphs/query.jl

@@ -230,7 +230,7 @@ defined as ``\hat{L} = \frac{2}{\lambda_{max}} L - I`` where ``L`` is the normal
 """
 function scaled_laplacian(g::GNNGraph, T::DataType=Float32; dir=:out)
     L = normalized_laplacian(g, T)
-    @assert issymmetric(L) "scaled_laplacian only works with symmetric matrices"
+    # @assert issymmetric(L) "scaled_laplacian only works with symmetric matrices"
     λmax = _eigmax(L)
     return  2 / λmax * L - I
 end

src/GraphNeuralNetworks.jl

@@ -22,20 +22,21 @@ using .GNNGraphs: COO_T, ADJMAT_T, SPARSE_T,
 export
     # utils
     reduce_nodes, 
-    reduce_edges, 
+    reduce_edges,
     softmax_nodes, 
     softmax_edges,
-    broadcast_nodes, 
+    broadcast_nodes,
     broadcast_edges,
     softmax_edge_neighbors,
 
     # msgpass
     apply_edges,
-    aggregate_neighbors, 
+    aggregate_neighbors,
     propagate,
-    copy_xj, 
-    copy_xi, 
+    copy_xj,
+    copy_xi,
     xi_dot_xj,
+    e_mul_xj,
 
     # layers/basic
     GNNLayer,

src/layers/conv.jl

@@ -832,16 +832,15 @@ paper. In the forward pass, takes as inputs node features `x` and edge features
 updated features `x'` and `e'` according to 
 
 ```math
-\mathbf{e}_{i\to j}'  = \phi_e([\mathbf{x}_i; \mathbf{x}_j; \mathbf{e}_{i\to j}])\\
-\mathbf{x}_{i}'  = \phi_v([\mathbf{x}_i; \square_{j\in \mathcal{N}(i)\,\mathbf{e}_{j\to i}'])
+\mathbf{e}_{i\to j}'  = \phi_e([\mathbf{x}_i; \mathbf{x}_j; \mathbf{e}_{i\to j}]),\\
+\mathbf{x}_{i}'  = \phi_v([\mathbf{x}_i; \square_{j\in \mathcal{N}(i)\,\mathbf{e}_{j\to i}']).
 ```
 
 `aggr` defines the aggregation to be performed.
 
 If the neural networks `ϕe` and  `ϕv` are not provided, they will be constructed from
 the `in` and `out` arguments instead as multi-layer perceptron with one hidden layer and `relu` 
 activations.
-````
 
 # Examples
 

test/GNNGraphs/gnngraph.jl

@@ -1,4 +1,36 @@
 @testset "GNNGraph" begin
+
+    @testset "Constructor: adjacency matrix" begin
+        A = sprand(10, 10, 0.5)
+        sA, tA, vA = findnz(A)
+        
+        g = GNNGraph(A, graph_type=GRAPH_T)
+        s, t = edge_index(g)
+        v = get_edge_weight(g)
+        @test s == sA
+        @test t == tA
+        @test v == vA
+
+        g = GNNGraph(Matrix(A), graph_type=GRAPH_T)
+        s, t = edge_index(g)
+        v = get_edge_weight(g)
+        @test s == sA
+        @test t == tA
+        @test v == vA
+
+        g = GNNGraph([0 0 0
+                      0 0 1
+                      0 1 0], graph_type=GRAPH_T)
+        @test g.num_nodes == 3
+        @test g.num_edges == 2
+        
+        g = GNNGraph([0 1 0
+                      1 0 0
+                      0 0 0], graph_type=GRAPH_T)
+        @test g.num_nodes == 3
+        @test g.num_edges == 2
+    end
+
     @testset "symmetric graph" begin
         s = [1, 1, 2, 2, 3, 3, 4, 4]
         t = [2, 4, 1, 3, 2, 4, 1, 3]
@@ -124,7 +156,7 @@
         @test adjacency_list(g, dir=:in) ==  adj_list_in
     end
 
-    @testset "Graphs constructor" begin
+    @testset "Graphs.jl constructor" begin
         lg = random_regular_graph(10, 4)
         @test !Graphs.is_directed(lg)
         g = GNNGraph(lg)

test/GNNGraphs/query.jl

@@ -92,4 +92,12 @@
         @test eltype(L) == eltype(g)
         @test L ≈ D - A
     end
+
+    @testset "adjacency_matrix" begin
+        a = sprand(5, 5, 0.5)
+        g = GNNGraph(a, graph_type=GRAPH_T)
+        A = adjacency_matrix(g, Float32)
+        @test a ≈ A
+        @test eltype(A) == Float32
+    end
 end