perturb_nodes -> remove_nodes (#454)

Author: CarloLucibello

.github/workflows/tests_GNNGraphs.yml

@@ -15,7 +15,7 @@ jobs:
       matrix:
         version:
           - '1.10' # Replace this with the minimum Julia version that your package supports.
-          - '1' #  '1' will automatically expand to the latest stable 1.x release of Julia.
+          # - '1' #  '1' will automatically expand to the latest stable 1.x release of Julia.
           # - 'pre'
         os:
           - ubuntu-latest

.github/workflows/tests_GraphNeuralNetworks.yml

@@ -15,7 +15,7 @@ jobs:
       matrix:
         version:
           - '1.10' # Replace this with the minimum Julia version that your package supports.
-          - '1' #  '1' will automatically expand to the latest stable 1.x release of Julia.
+          # - '1' #  '1' will automatically expand to the latest stable 1.x release of Julia.
           # - 'pre'
         os:
           - ubuntu-latest

GNNGraphs/Project.toml

@@ -1,7 +1,7 @@
 name = "GNNGraphs"
 uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c"
 authors = ["Carlo Lucibello and contributors"]
-version = "0.1.0"
+version = "1.0.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -46,7 +46,7 @@ SparseArrays = "1"
 Statistics = "1"
 StatsBase = "0.34"
 cuDNN = "1"
-julia = "1.9"
+julia = "1.10"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

GNNGraphs/ext/GNNGraphsCUDAExt.jl

@@ -0,0 +1,33 @@
+module GNNGraphsCUDAExt
+
+using CUDA
+using Random, Statistics, LinearAlgebra
+using GNNGraphs
+using GNNGraphs: COO_T, ADJMAT_T, SPARSE_T 
+
+const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}
+
+# Query 
+
+GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1))
+
+# Transform
+
+GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz)
+
+
+# Utils
+
+GNNGraphs.iscuarray(x::AnyCuArray) = true
+
+
+function sort_edge_index(u::AnyCuArray, v::AnyCuArray)
+    dev = get_device(u)
+    cdev = cpu_device()
+    u, v = u |> cdev, v |> cdev
+    #TODO proper cuda friendly implementation
+    sort_edge_index(u, v) |> dev
+end
+
+
+end #module

GNNGraphs/ext/GNNGraphsCUDAExt/GNNGraphsCUDAExt.jl

@@ -80,7 +80,6 @@ export add_nodes,
        perturb_edges,
        remove_nodes,
        ppr_diffusion,
-       drop_nodes,
 # from MLUtils
        batch,
        unbatch,

GNNGraphs/ext/GNNGraphsCUDAExt/query.jl

@@ -307,35 +307,27 @@ function remove_nodes(g::GNNGraph{<:COO_T}, nodes_to_remove::AbstractVector)
 end
 
 """
-    drop_nodes(g::GNNGraph{<:COO_T}, p)
+    remove_nodes(g::GNNGraph, p)
 
-Randomly drop nodes (and their associated edges) from a GNNGraph based on a given probability. 
-Dropping nodes is a technique that can be used for graph data augmentation, refering paper [DropNode](https://arxiv.org/pdf/2008.12578.pdf).
+Returns a new graph obtained by dropping nodes from `g` with independent probabilities `p`. 
 
-# Arguments
-- `g`: The input graph from which nodes (and their associated edges) will be dropped.
-- `p`: The probability of dropping each node. Default value is `0.5`.
-
-# Returns
-A modified GNNGraph with nodes (and their associated edges) dropped based on the given probability.
+# Examples
 
-# Example
 ```julia
-using GraphNeuralNetworks
-# Construct a GNNGraph
-g = GNNGraph([1, 1, 2, 2, 3], [2, 3, 1, 3, 1], num_nodes=3)
-# Drop nodes with a probability of 0.5
-g_new = drop_node(g, 0.5)
-println(g_new)
+julia> g = GNNGraph([1, 1, 2, 2, 3, 4], [1, 2, 3, 1, 3, 1])
+GNNGraph:
+  num_nodes: 4
+  num_edges: 6
+
+julia> g_new = remove_nodes(g, 0.5)
+GNNGraph:
+  num_nodes: 2
+  num_edges: 2
 ```
 """
-function drop_nodes(g::GNNGraph{<:COO_T}, p = 0.5)
-    num_nodes = g.num_nodes
-    nodes_to_remove = filter(_ -> rand() < p, 1:num_nodes)
-    
-    new_g = remove_nodes(g, nodes_to_remove)
-    
-    return new_g
+function remove_nodes(g::GNNGraph, p::AbstractFloat)
+    nodes_to_remove = filter(_ -> rand() < p, 1:g.num_nodes)
+    return remove_nodes(g, nodes_to_remove)
 end
 
 """