docstrings

Author: CarloLucibello

src/GNNGraphs/query.jl

@@ -125,6 +125,24 @@ function _get_edge_weight(g, edge_weight)
     return ew
 end
 
+"""
+    degree(g::GNNGraph, T=nothing; dir=:out, edge_weight=true)
+
+Return a vector containing the degrees of the nodes in `g`.
+
+# Arguments
+- `g`: A graph.
+- `T`: Element type of the returned vector. If `nothing`, is
+       chosen based on the graph type and will be an integer
+       if `edge_weight=false`.
+- `dir`: For `dir=:out` the degree of a node is counted based on the outgoing edges.
+         For `dir=:in`, the ingoing edges are used. If `dir=:both` we have the sum of the two.
+- `edge_weight`: If `true` and the graph contains weighted edges, the degree will 
+                be weighted. Set to `false` instead to just count the number of
+                outgoing/ingoing edges.
+                In alternative, you can also pass a vector of weights to be used
+                instead of the graph's own weights.
+"""
 function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=true)
     s, t = edge_index(g)
 

test/GNNGraphs/query.jl

@@ -47,15 +47,17 @@
         eweight = [0.1, 2.1, 1.2, 1]
         g = GNNGraph((s, t, eweight), graph_type=GRAPH_T)
         @test degree(g) ==  [2.2, 1.2, 1.0, 0.0]
+        d = degree(g, edge_weight=false)
         if GRAPH_T == :coo
-            @test degree(g, edge_weight=false) ==  [2, 1, 1, 0]
+            @test d ==  [2, 1, 1, 0]
             @test degree(g, edge_weight=nothing) ==  [2, 1, 1, 0]
         else
             # Adjacency matrix representation cannot disambiguate multiple edges
             # and edge weights
-            @test degree(g, edge_weight=false) ==  [1, 1, 1, 0]
+            @test d ==  [1, 1, 1, 0]
             @test degree(g, edge_weight=nothing) ==  [1, 1, 1, 0]
         end
+        @test eltype(d) <: Integer
         if GRAPH_T == :coo
             @test degree(g, edge_weight=2*eweight) == [4.4, 2.4, 2.0, 0.0]
         end